uhhh

db/attachment.py

@@ -24,6 +24,7 @@ class Attachment(Base):
     # can be multiple attachments on a msg
     # so we make our own ids for attachments
     id: Mapped[UUID] = mapped_column(primary_key=True)
+    attachment_id: Mapped[int] = mapped_column(BigInteger, nullable=False, unique=True)
     msg_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
 
     url: Mapped[str]
@@ -33,8 +34,6 @@ class Attachment(Base):
     channel: Mapped[dict[str, Any]]
     vision: Mapped[Optional[dict[str, Any]]] = mapped_column(default=None)
 
-    __table_args__ = (UniqueConstraint('msg_id', 'url', name='ix_attachments_msg_id_url_uc'),)
-
     def file_details(self) -> tuple[str, str]:
         u = urlparse(self.url)
         if not u.path:

main.py

@@ -7,6 +7,7 @@ import meilisearch
 import os
 
 from aiobotocore.session import get_session as aiobotocore_get_session
+from db import attachment
 from db.attachment import Attachment
 from db.meme import Meme
 from db.migrate import migrate_db
@@ -94,9 +95,9 @@ async def process_channel(
                 return await storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=res_bytes)
 
     def create_attachments(message, channel, guild):
-        urls = [parse_url(attachment['url']) for attachment in message['attachments']]
+        urls = [(parse_url(attachment['url']), attachment['id']) for attachment in message['attachments']]
 
-        def generate_attachment(url: Url):
+        def generate_attachment(url: Url, attachment_id: str):
             if not url.path:
                 return None
 
@@ -107,6 +108,7 @@ async def process_channel(
 
             return Attachment(
                 id=uuid4(),
+                attachment_id=int(attachment_id),
                 msg_id=int(message['id']),
                 message=message,
                 guild=guild,
@@ -114,7 +116,7 @@ async def process_channel(
                 channel=channel
             )
 
-        return [generate_attachment(url) for url in urls]
+        return [generate_attachment(*url) for url in urls]
 
     async def upload_and_persist(attachment: Attachment, session: AsyncSession):
         try:
@@ -155,10 +157,19 @@ async def process_channel(
 
     print(f"found {len(attachments)} attachments for channel guild={guild['name']} name={channel['name']} id={channel['id']}");
 
+    attachment_ids = list(map(lambda a: a.attachment_id, attachments))
+
     async with db_session() as session:
-        results = await session.execute(select(Attachment.msg_id).distinct())
-        msg_ids = [row[0] for row in results.fetchall()]
-        new_attachments = list(filter(lambda a: a.msg_id not in msg_ids, attachments))
+        # this is a bit weird
+        # we find the union of attachment_ids in the db and our list
+        results = await session.execute(
+            select(Attachment.attachment_id)
+                .filter(Attachment.attachment_id.in_(attachment_ids)))
+
+        existing_ids = [row[0] for row in results.fetchall()]
+
+        # and then filter out msgs that are already processed
+        new_attachments = list(filter(lambda a: a.attachment_id not in existing_ids, attachments))
         print(f"{len(new_attachments)} attachments to download for channel guild={guild['name']} name={channel['name']} id={channel['id']}");
 
         # ok now we will try and upload each attachment

migrations/versions/1dc4564189c1_create_memes_and_attachments_table.py

@@ -1,8 +1,8 @@
 """Create Memes and Attachments table
 
-Revision ID: 7ff6690dc082
+Revision ID: 1dc4564189c1
 Revises: 
-Create Date: 2024-04-08 21:29:55.858921
+Create Date: 2024-04-08 22:42:46.378215
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = '7ff6690dc082'
+revision: str = '1dc4564189c1'
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -22,16 +22,15 @@ def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
     op.create_table('attachments',
     sa.Column('id', sa.Uuid(), nullable=False),
-    sa.Column('msg_id', sa.BigInteger(), nullable=False),
+    sa.Column('attachment_id', sa.BigInteger(), nullable=False),
     sa.Column('url', sa.String(), nullable=False),
     sa.Column('message', sa.JSON(), nullable=False),
     sa.Column('guild', sa.JSON(), nullable=False),
     sa.Column('channel', sa.JSON(), nullable=False),
     sa.Column('vision', sa.JSON(), nullable=True),
     sa.PrimaryKeyConstraint('id'),
-    sa.UniqueConstraint('msg_id', 'url', name='ix_attachments_msg_id_url_uc')
+    sa.UniqueConstraint('attachment_id')
     )
-    op.create_index(op.f('ix_attachments_msg_id'), 'attachments', ['msg_id'], unique=False)
     op.create_table('memes',
     sa.Column('id', sa.Integer(), nullable=False),
     sa.Column('url', sa.String(), nullable=False),
@@ -45,6 +44,5 @@ def upgrade() -> None:
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
     op.drop_table('memes')
-    op.drop_index(op.f('ix_attachments_msg_id'), table_name='attachments')
     op.drop_table('attachments')
     # ### end Alembic commands ###