about to filter

db/attachment.py

@@ -1,9 +1,15 @@
+from urllib.parse import urlparse
 from uuid import UUID
 import os
 
 from typing import Any, Optional
 
+from urllib3.util import parse_url
+from urllib3.util import Url
+
 from sqlalchemy import ForeignKey
+from sqlalchemy import UniqueConstraint
+from sqlalchemy import BigInteger
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
 from sqlalchemy.orm import DeclarativeBase
@@ -18,7 +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)
-    msg_id: Mapped[int] = mapped_column(nullable=False, index=True)
+    msg_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
 
     url: Mapped[str]
 
@@ -27,15 +33,21 @@ 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]:
-        return os.path.splitext(os.path.basename(self.url))
+        u = urlparse(self.url)
+        if not u.path:
+            raise RuntimeError(f"attachment url missing path {self.url}")
+        return os.path.splitext(os.path.basename(u.path))
 
     def object_name(self) -> str:
+        print(self.id)
         (filename, ext) = self.file_details()
         if self.convert_to_webp():
             return f"{self.id}.webp"
 
-        return f"{id}{ext}"
+        return f"{self.id}{ext}"
 
     def convert_to_webp(self) -> bool:
         (filename, ext) = self.file_details()

main.py

@@ -1,38 +1,31 @@
-import os
+from logging import warn
-import asyncio
-import json
-import itertools
-
-from urllib3.util import parse_url
-from pathlib import Path
-from io import BytesIO
-from dataclasses import dataclass
-from uuid import uuid4, UUID
-
 import aiohttp
+import asyncio
+import itertools
+import json
+import meilisearch
+import os
 
 from aiobotocore.session import get_session as aiobotocore_get_session
-from types_aiobotocore_s3.client import S3Client
-
-from PIL import Image
-from dotenv import dotenv_values, load_dotenv
-
-from google.cloud import vision as google_vision
-from google.protobuf.json_format import MessageToDict, MessageToJson
-from google.cloud.vision_v1 import types
-import meilisearch
-from urllib3.util import Url, url
-
-from sqlalchemy import create_engine
-from db.migrate import migrate_db
 from db.attachment import Attachment
 from db.meme import Meme
+from db.migrate import migrate_db
+from dotenv import dotenv_values, load_dotenv
+from google.cloud import vision as google_vision
+from google.cloud.vision_v1 import types
+from io import BytesIO
+from PIL import Image
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, Session
+from types_aiobotocore_s3.client import S3Client
+from urllib3.util import parse_url
+from urllib3.util import Url
+from uuid import uuid4
 
 # need to change the env for the google api calls
 load_dotenv()
 config = dotenv_values(".env")
 
-DATA_DIR = './data/memestream'
 MEILI_INDEX = 'memes'
 MEILI_URL = 'http://localhost:7700'
 
@@ -41,23 +34,13 @@ CHANNEL_IDS = [
     1095060717117374687, # aella memes
 ]
 
+STORAGE_BUCKET = config['STORAGE_BUCKET']
+assert STORAGE_BUCKET
+
 headers = {
     'Authorization': config['DISCORD_AUTH'],
 }
 
-# TODO
-# [ ] Backfill last N images from a channel
-# [ ] Set oldest_img_id for channel
-# [ ] save img as json + imagefile
-# [ ] check extensions (ignore mp4 maybe)
-
-@dataclass
-class MemeDocument:
-    id: UUID
-    url: Url
-    text: str
-    tags: list[str]
-
 def detect_text(url, vision):
     request = {
         'image': { 'source': { 'image_uri': url } },
@@ -75,15 +58,21 @@ def filter_message_img_urls(messages):
     , messages)
     return [item for group in url_lists for item in group]
 
-async def download_channel_attachments(session, storage, vision, channel_id):
+async def process_channel(
+    http_session: aiohttp.ClientSession,
+    storage: S3Client,
+    vision: google_vision.ImageAnnotatorClient,
+    db_session_maker: sessionmaker[Session],
+    channel_id: int
+):
     async def fetch_channel_messages(channel_id):
         url = f"https://discord.com/api/v9/channels/{channel_id}/messages"
-        async with session.get(url, headers=headers) as res:
+        async with http_session.get(url, headers=headers) as res:
             return await res.json()
 
     async def fetch_guild_details(guild_id):
         url = f"https://discord.com/api/v9/guilds/{guild_id}"
-        async with session.get(url, headers=headers) as res:
+        async with http_session.get(url, headers=headers) as res:
             data = await res.json()
 
             keys_to_remove = ['emojis', 'stickers', 'roles']
@@ -93,38 +82,28 @@ async def download_channel_attachments(session, storage, vision, channel_id):
 
     async def fetch_channel(channel_id):
         url = f"https://discord.com/api/v9/channels/{channel_id}"
-        async with session.get(url, headers=headers) as res:
+        async with http_session.get(url, headers=headers) as res:
             return await res.json()
 
-
+    async def download_attachment(attachment: Attachment):
-    async def upload_object(object_name, body):
+        async with http_session.get(attachment.url) as response:
-        storage_bucket = config['STORAGE_BUCKET']
-        if not storage_bucket:
-            raise RuntimeError(f"config missing STORAGE_BUCKET")
-        return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=body)
-
-    async def download_attachment(
-        session: aiohttp.ClientSession,
-        storage: S3Client,
-        attachment: Attachment
-    ):
-        async with session.get(attachment.url) as response:
             if response.status == 200:
-                # make a webp conversion for obvious image types
-                res_bytes = await response.read()
-
                 (filename, ext) = attachment.file_details()
                 object_name = attachment.object_name()
-                print(f"uploading object_name={object_name} filename={filename} ext={ext}")
 
+                print(f"uploading object_name={object_name} filename={filename} ext={ext}")
+                res_bytes = await response.read()
+
+                # make a webp conversion for obvious image types
                 if attachment.convert_to_webp():
+                    # have to read the whole img before converting to webp
                     image = Image.open(BytesIO(res_bytes))
                     webp_output = BytesIO()
                     image.save(webp_output, format='WebP')
                     webp_output.seek(0)
-                    return await upload_object(object_name=object_name, body=webp_output)
+                    return await storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=webp_output)
 
-                return await upload_object(object_name=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):
         urls = [parse_url(attachment['url']) for attachment in message['attachments']]
@@ -164,6 +143,28 @@ async def download_channel_attachments(session, storage, vision, channel_id):
     # super wack flatmap and filter
     attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups)))
 
+    db_session = db_session_maker()
+
+    # filter attachments that already exist in the db
+    # dirty way is just check if msg_id exists in db, don't care about specific attachments
+
+    # ok now we will try and upload each attachment
+    # for the ones that succeed we will add them to our session to be stored
+    async def upload_and_persist(attachment: Attachment):
+        try:
+            await download_attachment(attachment=attachment)
+            # this is kinda dodge, the exception will catch if we fail to add to session
+            db_session.add(attachment)
+        except Exception as e:
+            warn(e)
+            warn(f'deleting attachment upload {attachment.object_name()}')
+            await storage.delete_object(Bucket=STORAGE_BUCKET, Key=attachment.object_name())
+
+    attachments = [upload_and_persist(attachment=a) for a in attachments]
+    await asyncio.gather(*attachments)
+
+    db_session.commit()
+
     # NOW
     # -------------------
     # insert attachment
@@ -207,29 +208,33 @@ async def main():
     db_engine = create_engine(db_url, echo=True)
     migrate_db(db_engine, db_url)
 
+    db_session_maker = sessionmaker(bind=db_engine)
+
     aioboto = aiobotocore_get_session()
     async with aioboto.create_client(
         's3',
         endpoint_url=config['STORAGE_URL'],
         aws_access_key_id=config['STORAGE_ACCESS_KEY'],
         aws_secret_access_key=config['STORAGE_SECRET_KEY'],
-    ) as storage:
+    ) as storage, aiohttp.ClientSession() as http_session:
+
         storage: S3Client
 
-        storage_bucket = config['STORAGE_BUCKET']
-        assert storage_bucket
-
         # Make the bucket if it doesn't exist.
-        exists = await check_bucket_exists(storage, storage_bucket)
+        exists = await check_bucket_exists(storage, STORAGE_BUCKET)
         if not exists:
-            await storage.create_bucket(Bucket=storage_bucket)
+            await storage.create_bucket(Bucket=STORAGE_BUCKET)
-            print("Created bucket", storage_bucket)
+            print("Created bucket", STORAGE_BUCKET)
         else:
-            print("Bucket", storage_bucket, "already exists")
+            print("Bucket", STORAGE_BUCKET, "already exists")
 
-        async with aiohttp.ClientSession() as session:
+        channels = [process_channel(http_session, storage, vision, db_session_maker, channel_id) for channel_id in CHANNEL_IDS]
-            channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS]
+        memes = await asyncio.gather(*channels)
-            memes = await asyncio.gather(*channels)
 
 
 asyncio.run(main())
+
+# TODO
+# [ ] Backfill last N images from a channel
+# [ ] Set oldest_img_id for channel
+# logger

migrations/versions/50b5b3980665_create_memes_and_attachments_table.py

@@ -1,8 +1,8 @@
 """Create Memes and Attachments table
 
-Revision ID: e04bee35feab
+Revision ID: 50b5b3980665
 Revises: 
-Create Date: 2024-04-08 10:45:47.138438
+Create Date: 2024-04-08 18:30:57.222252
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = 'e04bee35feab'
+revision: str = '50b5b3980665'
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -22,17 +22,18 @@ 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.Integer(), nullable=False),
+    sa.Column('msg_id', sa.BigInteger(), nullable=False),
-    sa.Column('attributes', sa.JSON(), nullable=False),
+    sa.Column('url', sa.String(), nullable=False),
-    sa.Column('vision', sa.JSON(), nullable=False),
     sa.Column('message', sa.JSON(), nullable=False),
     sa.Column('guild', sa.JSON(), nullable=False),
     sa.Column('channel', sa.JSON(), nullable=False),
-    sa.PrimaryKeyConstraint('id')
+    sa.Column('vision', sa.JSON(), nullable=True),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('msg_id', 'url', name='ix_attachments_msg_id_url_uc')
     )
     op.create_index(op.f('ix_attachments_msg_id'), 'attachments', ['msg_id'], unique=False)
     op.create_table('memes',
-    sa.Column('id', sa.Uuid(), nullable=False),
+    sa.Column('id', sa.Integer(), nullable=False),
     sa.Column('url', sa.String(), nullable=False),
     sa.Column('text', sa.String(), nullable=False),
     sa.Column('tags', sa.JSON(), nullable=False),

migrations/versions/9159293e4f30_allow_vision_to_be_nullable.py

@@ -1,36 +0,0 @@
-"""allow vision to be nullable
-
-Revision ID: 9159293e4f30
-Revises: e04bee35feab
-Create Date: 2024-04-08 10:50:21.447942
-
-"""
-from typing import Sequence, Union
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision: str = '9159293e4f30'
-down_revision: Union[str, None] = 'e04bee35feab'
-branch_labels: Union[str, Sequence[str], None] = None
-depends_on: Union[str, Sequence[str], None] = None
-
-
-def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.alter_column('attachments', 'vision',
-               existing_type=postgresql.JSON(astext_type=sa.Text()),
-               nullable=True)
-    op.drop_column('attachments', 'attributes')
-    # ### end Alembic commands ###
-
-
-def downgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column('attachments', sa.Column('attributes', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=False))
-    op.alter_column('attachments', 'vision',
-               existing_type=postgresql.JSON(astext_type=sa.Text()),
-               nullable=False)
-    # ### end Alembic commands ###