This commit is contained in:
Nathan Rashleigh committed 2024-04-08 23:06:42 +10:00
1 parent e16a193cb7
commit 796e9d2bd8
3 files changed
+23 -15

No files matched your search

+1 -2
View File
@@ -24,6 +24,7 @@ class Attachment(Base):
# can be multiple attachments on a msg # can be multiple attachments on a msg
# so we make our own ids for attachments # so we make our own ids for attachments
id: Mapped[UUID] = mapped_column(primary_key=True) 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) msg_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
url: Mapped[str] url: Mapped[str]
@@ -33,8 +34,6 @@ class Attachment(Base):
channel: Mapped[dict[str, Any]] channel: Mapped[dict[str, Any]]
vision: Mapped[Optional[dict[str, Any]]] = mapped_column(default=None) 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]: def file_details(self) -> tuple[str, str]:
u = urlparse(self.url) u = urlparse(self.url)
if not u.path: if not u.path:
+17 -6
View File
@@ -7,6 +7,7 @@ import meilisearch
import os import os
from aiobotocore.session import get_session as aiobotocore_get_session from aiobotocore.session import get_session as aiobotocore_get_session
from db import attachment
from db.attachment import Attachment from db.attachment import Attachment
from db.meme import Meme from db.meme import Meme
from db.migrate import migrate_db 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) return await storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=res_bytes)
def create_attachments(message, channel, guild): 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: if not url.path:
return None return None
@@ -107,6 +108,7 @@ async def process_channel(
return Attachment( return Attachment(
id=uuid4(), id=uuid4(),
attachment_id=int(attachment_id),
msg_id=int(message['id']), msg_id=int(message['id']),
message=message, message=message,
guild=guild, guild=guild,
@@ -114,7 +116,7 @@ async def process_channel(
channel=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): async def upload_and_persist(attachment: Attachment, session: AsyncSession):
try: 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']}"); 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: async with db_session() as session:
results = await session.execute(select(Attachment.msg_id).distinct()) # this is a bit weird
msg_ids = [row[0] for row in results.fetchall()] # we find the union of attachment_ids in the db and our list
new_attachments = list(filter(lambda a: a.msg_id not in msg_ids, attachments)) 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']}"); 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 # ok now we will try and upload each attachment
@@ -1,8 +1,8 @@
"""Create Memes and Attachments table """Create Memes and Attachments table
Revision ID: 7ff6690dc082 Revision ID: 1dc4564189c1
Revises: Revises:
Create Date: 2024-04-08 21:29:55.858921 Create Date: 2024-04-08 22:42:46.378215
""" """
from typing import Sequence, Union from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '7ff6690dc082' revision: str = '1dc4564189c1'
down_revision: Union[str, None] = None down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: 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! ### # ### commands auto generated by Alembic - please adjust! ###
op.create_table('attachments', op.create_table('attachments',
sa.Column('id', sa.Uuid(), nullable=False), 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('url', sa.String(), nullable=False),
sa.Column('message', sa.JSON(), nullable=False), sa.Column('message', sa.JSON(), nullable=False),
sa.Column('guild', sa.JSON(), nullable=False), sa.Column('guild', sa.JSON(), nullable=False),
sa.Column('channel', sa.JSON(), nullable=False), sa.Column('channel', sa.JSON(), nullable=False),
sa.Column('vision', sa.JSON(), nullable=True), sa.Column('vision', sa.JSON(), nullable=True),
sa.PrimaryKeyConstraint('id'), 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', op.create_table('memes',
sa.Column('id', sa.Integer(), nullable=False), sa.Column('id', sa.Integer(), nullable=False),
sa.Column('url', sa.String(), nullable=False), sa.Column('url', sa.String(), nullable=False),
@@ -45,6 +44,5 @@ def upgrade() -> None:
def downgrade() -> None: def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
op.drop_table('memes') op.drop_table('memes')
op.drop_index(op.f('ix_attachments_msg_id'), table_name='attachments')
op.drop_table('attachments') op.drop_table('attachments')
# ### end Alembic commands ### # ### end Alembic commands ###