From 796e9d2bd866f9ddad4849a83014b63012c9e862 Mon Sep 17 00:00:00 2001 From: Nathan Rashleigh Date: Mon, 8 Apr 2024 23:06:42 +1000 Subject: [PATCH] uhhh --- db/attachment.py | 3 +-- main.py | 23 ++++++++++++++----- ...9c1_create_memes_and_attachments_table.py} | 12 ++++------ 3 files changed, 23 insertions(+), 15 deletions(-) rename migrations/versions/{7ff6690dc082_create_memes_and_attachments_table.py => 1dc4564189c1_create_memes_and_attachments_table.py} (76%) diff --git a/db/attachment.py b/db/attachment.py index 46579f9..9de190e 100644 --- a/db/attachment.py +++ b/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: diff --git a/main.py b/main.py index c5b64e7..d5a3a17 100644 --- a/main.py +++ b/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 diff --git a/migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py b/migrations/versions/1dc4564189c1_create_memes_and_attachments_table.py similarity index 76% rename from migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py rename to migrations/versions/1dc4564189c1_create_memes_and_attachments_table.py index 9594093..618b016 100644 --- a/migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py +++ b/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 ###