attachment class

main.py

@@ -1,12 +1,13 @@
 import os
 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
-from typing import Dict, Any, List
 
 import aiohttp
 
@@ -24,6 +25,8 @@ 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
 
 # need to change the env for the google api calls
 load_dotenv()
@@ -53,7 +56,7 @@ class MemeDocument:
     id: UUID
     url: Url
     text: str
-    tags: List[str]
+    tags: list[str]
 
 def detect_text(url, vision):
     request = {
@@ -93,68 +96,85 @@ async def download_channel_attachments(session, storage, vision, channel_id):
         async with session.get(url, headers=headers) as res:
             return await res.json()
 
+
+    async def upload_object(object_name, body):
+        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,
-        url: Url,
-        message,
+        attachment: Attachment
     ):
-        # there's a big param of some kind after the image url that we remove
-        if not url.path:
-            raise RuntimeError(f"url missing path url={url}")
-
-        storage_bucket = config['UPLOAD_BUCKET']
-        if not storage_bucket:
-            raise RuntimeError(f"config missing UPLOAD_BUCKET")
-
-        (filename, ext) = os.path.splitext(os.path.basename(url.path))
-        # TODO
-        # FILTER WANTED TYPES
-
-        id = uuid4()
-
-        async with session.get(url.url) as response:
+        async with session.get(attachment.url) as response:
             if response.status == 200:
-
                 # make a webp conversion for obvious image types
-                convert_extensions = ['jpg', 'png', 'jpeg', 'webp']
                 res_bytes = await response.read()
-                object_name: str
 
-                # trim the .
-                if ext.lower()[1:] in convert_extensions:
+                (filename, ext) = attachment.file_details()
+                object_name = attachment.object_name()
+                print(f"uploading object_name={object_name} filename={filename} ext={ext}")
+
+                if attachment.convert_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)
 
-                    object_name = f"{id}.webp"
-                    print(f"uploading object_name={object_name} filename={filename} ext={ext}")
-                    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)
 
-                object_name = f"{id}{ext}"
-                print(f"uploading object_name={object_name} filename={filename} ext={ext}")
-                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']]
 
-                # add discord id to server processed msgs list
-                # if there then nothing to do
+        def generate_attachment(url: Url):
+            if not url.path:
+                return None
 
-                # return MemeDocument(
-                #     id=id,
-                #     url=self.upload_url,
-                #     text=self.text or '',
-                #     tags=[self.guild_name, self.channel_name],
-                # )
+            (filename, ext) = os.path.splitext(os.path.basename(url.path))
+            extensions = ['jpg', 'png', 'jpeg', 'webp', 'mov', 'gif', 'mp4']
+            if ext.lower()[1:] not in extensions:
+                return None
+
+            return Attachment(
+                id=uuid4(),
+                msg_id=int(message['id']),
+                message=message,
+                guild=guild,
+                url=url.url,
+                channel=channel
+            )
+
+        return [generate_attachment(url) for url in urls]
 
     channel = await fetch_channel(channel_id)
     guild = await fetch_guild_details(channel['guild_id'])
     print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}")
     messages = await fetch_channel_messages(channel_id)
-    urls_and_messages = filter_message_img_urls(messages)
 
-    tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
-    await asyncio.gather(*tasks)
+    # at this point we have the channel messages, the guild (server) and channel info
+
+    # first thing we have to do is parse out the urls in the msg attachments
+    # once we have those we can generate where we should put them in our storage
+    # should only persist them later once the processing is done
+    attachment_groups = [create_attachments(message, channel, guild) for message in messages]
+
+    # super wack flatmap and filter
+    attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups)))
+
+    # NOW
+    # -------------------
+    # insert attachment
+    # upload to s3
+    # fetch vision for relevant attachments
+    # create meme from attachment and vision
+    # insert meme
+    # add documents to meili
+
+    # tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
+    # await asyncio.gather(*tasks)
 
     # [attachment.detect_text(vision) for attachment in attachments]
     # [attachment.persist() for attachment in attachments]
@@ -184,32 +204,32 @@ async def main():
 
     # db setup
     db_url = config['POSTGRES_URL']
-    db_engine = create_engine(db_url)
+    db_engine = create_engine(db_url, echo=True)
     migrate_db(db_engine, db_url)
 
-    # 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:
-    #     storage: S3Client
+    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:
+        storage: S3Client
 
-    #     storage_bucket = config['UPLOAD_BUCKET']
-    #     assert storage_bucket
+        storage_bucket = config['STORAGE_BUCKET']
+        assert storage_bucket
 
-    #     # Make the bucket if it doesn't exist.
-    #     exists = await check_bucket_exists(storage, storage_bucket)
-    #     if not exists:
-    #         await storage.create_bucket(Bucket=storage_bucket)
-    #         print("Created bucket", storage_bucket)
-    #     else:
-    #         print("Bucket", storage_bucket, "already exists")
+        # Make the bucket if it doesn't exist.
+        exists = await check_bucket_exists(storage, storage_bucket)
+        if not exists:
+            await storage.create_bucket(Bucket=storage_bucket)
+            print("Created bucket", storage_bucket)
+        else:
+            print("Bucket", storage_bucket, "already exists")
 
-    #     async with aiohttp.ClientSession() as session:
-    #         channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS]
-    #         memes = await asyncio.gather(*channels)
+        async with aiohttp.ClientSession() as session:
+            channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS]
+            memes = await asyncio.gather(*channels)
 
 
 asyncio.run(main())