diff --git a/rag/alembic/env.py b/rag/alembic/env.py new file mode 100644 index 0000000..0df2b4e --- /dev/null +++ b/rag/alembic/env.py @@ -0,0 +1,46 @@ +from logging.config import fileConfig +from sqlalchemy import pool +from sqlalchemy.ext.asyncio import create_async_engine +from alembic import context +import asyncio + +config = context.config + +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +target_metadata = None + +def run_migrations_offline(): + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + with context.begin_transaction(): + context.run_migrations() + +def run_migrations_online(): + connectable = create_async_engine( + config.get_main_option("sqlalchemy.url"), + poolclass=pool.NullPool, + ) + + async def run_async_migrations(): + async with connectable.connect() as connection: + await connection.run_sync(do_run_migrations) + await connectable.dispose() + + def do_run_migrations(connection): + context.configure(connection=connection, target_metadata=target_metadata) + with context.begin_transaction(): + context.run_migrations() + + asyncio.run(run_async_migrations()) + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() \ No newline at end of file