Add discord.py thread models & migration

3 files changed, 86 insertions, 6 deletions

diff --git a/alembic/versions/563a15b2a76e_add_support_for_threads.py b/alembic/versions/563a15b2a76e_add_support_for_threads.py
new file mode 100644
index 0000000..6f21e92
--- /dev/null
+++ b/alembic/versions/563a15b2a76e_add_support_for_threads.py
@@ -0,0 +1,47 @@
+"""Add support for threads
+
+Revision ID: 563a15b2a76e
+Revises: d6c9452c3940
+Create Date: 2021-10-31 14:46:49.926646
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '563a15b2a76e'
+down_revision = 'd6c9452c3940'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        'threads',
+        sa.Column('id', sa.String(), nullable=False),
+        sa.Column('parent_channel_id', sa.String(), nullable=False),
+        sa.Column('created_at', sa.DateTime(), nullable=True),
+        sa.Column('name', sa.String(), nullable=False),
+        sa.Column('archived', sa.Boolean(), nullable=False),
+        sa.Column('auto_archive_duration', sa.Integer(), nullable=False),
+        sa.Column('locked', sa.Boolean(), nullable=False),
+        sa.Column('type', sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(['parent_channel_id'], ['channels.id'], ondelete='CASCADE'),
+        sa.PrimaryKeyConstraint('id')
+    )
+    op.add_column('messages', sa.Column('thread_id', sa.String(), nullable=True))
+    op.create_index(op.f('ix_messages_thread_id'), 'messages', ['thread_id'], unique=False)
+    op.create_index(op.f('ix_threads_type'), 'threads', ['type'], unique=False)
+    op.create_foreign_key(None, 'messages', 'threads', ['thread_id'], ['id'], ondelete='CASCADE')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'messages', type_='foreignkey')
+    op.drop_index(op.f('ix_messages_thread_id'), table_name='messages')
+    op.drop_column('messages', 'thread_id')
+    op.drop_table('threads')
+    # ### end Alembic commands ###
diff --git a/metricity/bot.py b/metricity/bot.py
index a33c40b..3d4f3cd 100644
--- a/metricity/bot.py
+++ b/metricity/bot.py
@@ -314,12 +314,21 @@ async def on_message(message: DiscordMessage) -> None:
     if cat_id in BotConfig.ignore_categories:
         return
 
-    await Message.create(
-        id=str(message.id),
-        channel_id=str(message.channel.id),
-        author_id=str(message.author.id),
-        created_at=message.created_at
-    )
+    args = {
+        "id": str(message.id),
+        "channel_id": str(message.channel.id),
+        "author_id": str(message.author.id),
+        "created_at": message.created_at
+    }
+
+    if isinstance(message.channel, ThreadChannel):
+        thread = message.channel
+        args["channel_id"] = str(thread.parent_id)
+        args["thread_id"] = str(thread.id)
+        if not await Thread.get(str(thread.id)):
+            await insert_thread(thread)
+
+    await Message.create(**args)
 
 
 @bot.event
diff --git a/metricity/models.py b/metricity/models.py
index 1f7e6ef..87be19b 100644
--- a/metricity/models.py
+++ b/metricity/models.py
@@ -33,6 +33,25 @@ class Channel(db.Model):
     is_staff = db.Column(db.Boolean, nullable=False)
 
 
+class Thread(db.Model):
+    """Database model representing a Thread channel."""
+
+    __tablename__ = "threads"
+
+    id = db.Column(db.String, primary_key=True)
+    parent_channel_id = db.Column(
+        db.String,
+        db.ForeignKey("channels.id", ondelete="CASCADE"),
+        nullable=False
+    )
+    created_at = db.Column(TZDateTime(), default=datetime.now(timezone.utc))
+    name = db.Column(db.String, nullable=False)
+    archived = db.Column(db.Boolean, default=False, nullable=False)
+    auto_archive_duration = db.Column(db.Integer, nullable=False)
+    locked = db.Column(db.Boolean, default=False, nullable=False)
+    type = db.Column(db.String, nullable=False, index=True)
+
+
 class User(db.Model):
     """Database model representing a Discord user."""
 
@@ -84,6 +103,11 @@ class Message(db.Model):
         db.ForeignKey("channels.id", ondelete="CASCADE"),
         index=True
     )
+    thread_id = db.Column(
+        db.String,
+        db.ForeignKey("threads.id", ondelete="CASCADE"),
+        index=True
+    )
     author_id = db.Column(
         db.String,
         db.ForeignKey("users.id", ondelete="CASCADE"),