diff --git a/db/migrate.py b/db/migrate.py index ca5fd09..dae6fe4 100644 --- a/db/migrate.py +++ b/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")) diff --git a/main.py b/main.py index 053a0be..c5b64e7 100644 --- a/main.py +++ b/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()) diff --git a/migrations/README b/migrations/README index 98e4f9c..e0d0858 100644 --- a/migrations/README +++ b/migrations/README @@ -1 +1 @@ -Generic single-database configuration. \ No newline at end of file +Generic single-database configuration with an async dbapi. \ No newline at end of file diff --git a/migrations/env.py b/migrations/env.py index 8606572..14140c3 100644 --- a/migrations/env.py +++ b/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() \ No newline at end of file diff --git a/migrations/versions/50b5b3980665_create_memes_and_attachments_table.py b/migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py similarity index 94% rename from migrations/versions/50b5b3980665_create_memes_and_attachments_table.py rename to migrations/versions/7ff6690dc082_create_memes_and_attachments_table.py index 416283a..9594093 100644 --- a/migrations/versions/50b5b3980665_create_memes_and_attachments_table.py +++ b/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 diff --git a/requirements.txt b/requirements.txt index ba6e93a..4a5e600 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,6 +6,6 @@ pillow aiobotocore types-aiobotocore-s3 -sqlalchemy -alembic -psycopg2-binary \ No newline at end of file +sqlalchemy[asyncio] +asyncpg +alembic \ No newline at end of file