select and asyncpg

main.py

@@ -15,8 +15,13 @@ 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 sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio import async_sessionmaker
+from sqlalchemy.ext.asyncio import async_scoped_session
+from sqlalchemy.future import select
+
 from types_aiobotocore_s3.client import S3Client
 from urllib3.util import parse_url
 from urllib3.util import Url
@@ -41,28 +46,11 @@ headers = {
     'Authorization': config['DISCORD_AUTH'],
 }
 
-def detect_text(url, vision):
-    request = {
-        'image': { 'source': { 'image_uri': url } },
-        'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
-    }
-
-    response: types.AnnotateImageResponse = vision.annotate_image(request=request)
-    vision_json = json.loads(types.AnnotateImageResponse.to_json(response))
-
-    return request
-
-def filter_message_img_urls(messages):
-    url_lists = map(
-        lambda msg: list(map(lambda attachment: (parse_url(attachment['url']), msg), msg['attachments']))
-    , messages)
-    return [item for group in url_lists for item in group]
-
 async def process_channel(
     http_session: aiohttp.ClientSession,
     storage: S3Client,
     vision: google_vision.ImageAnnotatorClient,
-    db_session_maker: sessionmaker[Session],
+    db_session: async_sessionmaker[AsyncSession],
     channel_id: int
 ):
     async def fetch_channel_messages(channel_id):
@@ -128,6 +116,28 @@ async def process_channel(
 
         return [generate_attachment(url) for url in urls]
 
+    async def upload_and_persist(attachment: Attachment, session: AsyncSession):
+        try:
+            await download_attachment(attachment=attachment)
+            # this is kinda dodge, the exception will catch if we fail to add to session
+            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())
+
+    def detect_text(url):
+        request = {
+            'image': { 'source': { 'image_uri': url } },
+            'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
+        }
+
+        response: types.AnnotateImageResponse = vision.annotate_image(request=request)
+        vision_json = json.loads(types.AnnotateImageResponse.to_json(response))
+
+        return request
+
+
     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']}")
@@ -143,32 +153,22 @@ async def process_channel(
     # super wack flatmap and filter
     attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups)))
 
-    db_session = db_session_maker()
+    print(f"found {len(attachments)} attachments for channel guild={guild['name']} name={channel['name']} id={channel['id']}");
 
-    # 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
+    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))
+        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
-    # 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()
+        # ok now we will try and upload each attachment
+        # for the ones that succeed we will add them to our session to be stored
+        downloads = [upload_and_persist(attachment=a, session=session) for a in new_attachments]
+        await asyncio.gather(*downloads)
+        await session.commit()
 
     # NOW
     # -------------------
-    # insert attachment
-    # upload to s3
     # fetch vision for relevant attachments
     # create meme from attachment and vision
     # insert meme
@@ -205,10 +205,12 @@ async def main():
 
     # db setup
     db_url = config['POSTGRES_URL']
-    db_engine = create_engine(db_url, echo=True)
-    migrate_db(db_engine, db_url)
 
-    db_session_maker = sessionmaker(bind=db_engine)
+    db_engine = create_async_engine(db_url, echo=False)
+    await migrate_db(db_engine)
+    print('db is up to date')
+
+    db_session = async_sessionmaker(db_engine, expire_on_commit=False)
 
     aioboto = aiobotocore_get_session()
     async with aioboto.create_client(
@@ -228,9 +230,11 @@ async def main():
         else:
             print("Bucket", STORAGE_BUCKET, "already exists")
 
-        channels = [process_channel(http_session, storage, vision, db_session_maker, channel_id) for channel_id in CHANNEL_IDS]
+        channels = [process_channel(http_session, storage, vision, db_session, channel_id) for channel_id in CHANNEL_IDS]
         memes = await asyncio.gather(*channels)
 
+        await db_engine.dispose()
+
 
 asyncio.run(main())