select and asyncpg

db/migrate.py

@@ -1,16 +1,13 @@
-from sqlalchemy import create_engine, Column, Integer, String, Engine
 from alembic.config import Config
 from alembic import command
-from dotenv import dotenv_values
+import asyncio
+from sqlalchemy.ext.asyncio import AsyncEngine
+from alembic import command, config
 
-# Function to perform the database migration
-def migrate_db(engine: Engine, db_url):
-    # Configure Alembic
-    alembic_cfg = Config()
-    alembic_cfg.set_main_option("script_location", "migrations")
-    alembic_cfg.set_main_option("sqlalchemy.url", db_url)
+def run_upgrade(connection, cfg):
+    cfg.attributes["connection"] = connection
+    command.upgrade(cfg, "head")
 
-    # Run migrations using Alembic
-    command.upgrade(alembic_cfg, "head")
-
-    print('db up to date')
+async def migrate_db(async_engine: AsyncEngine):
+    async with async_engine.begin() as conn:
+        await conn.run_sync(run_upgrade, config.Config("alembic.ini"))

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())
 

migrations/README

@@ -1 +1 @@
-Generic single-database configuration.
+Generic single-database configuration with an async dbapi.

migrations/env.py

@@ -1,86 +1,57 @@
+import asyncio
 from logging.config import fileConfig
 
+from alembic import context
 from sqlalchemy import engine_from_config
 from sqlalchemy import pool
-
-from alembic import context
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 from dotenv import dotenv_values
-
 from db import Base
 
-# this is the Alembic Config object, which provides
-# access to the values within the .ini file in use.
 config = context.config
 
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
 env = dotenv_values(".env")
 assert env['POSTGRES_URL']
 config.set_main_option('sqlalchemy.url', env['POSTGRES_URL'])
 
-# Interpret the config file for Python logging.
-# This line sets up loggers basically.
-if config.config_file_name is not None:
-    fileConfig(config.config_file_name)
-
-# add your model's MetaData object here
-# for 'autogenerate' support
-# from myapp import mymodel
-# target_metadata = mymodel.Base.metadata
 target_metadata = Base.metadata
 
-# other values from the config, defined by the needs of env.py,
-# can be acquired:
-# my_important_option = config.get_main_option("my_important_option")
-# ... etc.
+def run_migrations_online():
+    connectable = context.config.attributes.get("connection", None)
+    if connectable is None:
+        connectable = AsyncEngine(
+            engine_from_config(
+                context.config.get_section(context.config.config_ini_section),
+                prefix="sqlalchemy.",
+                poolclass=pool.NullPool,
+                future=True
+            )
+        )
+
+    if isinstance(connectable, AsyncEngine):
+        asyncio.run(run_async_migrations(connectable))
+    else:
+        do_run_migrations(connectable)
 
 
-def run_migrations_offline() -> None:
-    """Run migrations in 'offline' mode.
+async def run_async_migrations(connectable):
+    async with connectable.connect() as connection:
+        await connection.run_sync(do_run_migrations)
+    await connectable.dispose()
 
-    This configures the context with just a URL
-    and not an Engine, though an Engine is acceptable
-    here as well.  By skipping the Engine creation
-    we don't even need a DBAPI to be available.
 
-    Calls to context.execute() here emit the given string to the
-    script output.
-
-    """
-    url = config.get_main_option("sqlalchemy.url")
+def do_run_migrations(connection):
     context.configure(
-        url=url,
+        connection=connection,
         target_metadata=target_metadata,
-        literal_binds=True,
-        dialect_opts={"paramstyle": "named"},
+        compare_type=True,
     )
-
     with context.begin_transaction():
         context.run_migrations()
 
 
-def run_migrations_online() -> None:
-    """Run migrations in 'online' mode.
-
-    In this scenario we need to create an Engine
-    and associate a connection with the context.
-
-    """
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section, {}),
-        prefix="sqlalchemy.",
-        poolclass=pool.NullPool,
-    )
-
-    with connectable.connect() as connection:
-        context.configure(
-            connection=connection, target_metadata=target_metadata
-        )
-
-        with context.begin_transaction():
-            context.run_migrations()
-
-
-if context.is_offline_mode():
-    run_migrations_offline()
-else:
-    run_migrations_online()
+run_migrations_online()

migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py

@@ -1,8 +1,8 @@
 """Create Memes and Attachments table
 
-Revision ID: 50b5b3980665
+Revision ID: 7ff6690dc082
 Revises: 
-Create Date: 2024-04-08 18:30:57.222252
+Create Date: 2024-04-08 21:29:55.858921
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = '50b5b3980665'
+revision: str = '7ff6690dc082'
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None

requirements.txt

@@ -6,6 +6,6 @@ pillow
 aiobotocore
 types-aiobotocore-s3
 
-sqlalchemy
+sqlalchemy[asyncio]
+asyncpg
 alembic
-psycopg2-binary