From cb5bfcf7b5cc67e927937088e51eccbc644cd70d Mon Sep 17 00:00:00 2001 From: mohamad Date: Sun, 1 Jun 2025 17:33:04 +0200 Subject: [PATCH] refactor: Separate async migration logic into dedicated module and streamline migration functions for improved clarity and maintainability --- be/alembic/env.py | 80 ++++++++++++++++------------------------ be/alembic/migrations.py | 57 ++++++++++++++++++++++++++++ be/app/main.py | 4 +- 3 files changed, 90 insertions(+), 51 deletions(-) create mode 100644 be/alembic/migrations.py diff --git a/be/alembic/env.py b/be/alembic/env.py index eead375..d75e983 100644 --- a/be/alembic/env.py +++ b/be/alembic/env.py @@ -1,7 +1,6 @@ from logging.config import fileConfig import os import sys -import asyncio from sqlalchemy import engine_from_config from sqlalchemy import pool from sqlalchemy.ext.asyncio import create_async_engine @@ -16,8 +15,7 @@ import app.models # Ensure all models are loaded and registered to app.database from app.database import Base as DatabaseBase # Explicitly get Base from database.py from app.config import settings # Import settings to get DATABASE_URL -# this is the Alembic Config object, which provides -# access to the values within the .ini file in use. +# Get alembic config config = context.config # Set the sqlalchemy.url from your application settings @@ -26,7 +24,6 @@ if not settings.DATABASE_URL: raise ValueError("DATABASE_URL not found in settings for Alembic.") config.set_main_option('sqlalchemy.url', settings.DATABASE_URL) - # Interpret the config file for Python logging. # This line sets up loggers basically. if config.config_file_name is not None: @@ -43,50 +40,11 @@ target_metadata = DatabaseBase.metadata # Use metadata from app.database.Base # my_important_option = config.get_main_option("my_important_option") # ... etc. -def do_run_migrations(connection): +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode.""" + url = config.get_main_option("sqlalchemy.url") context.configure( - connection=connection, - target_metadata=target_metadata, - compare_type=True, - compare_server_default=True - ) - with context.begin_transaction(): - context.run_migrations() - -async def run_migrations_online_async(alembic_cfg=None): - """Run migrations in 'online' mode asynchronously.""" - # If running from Alembic directly - if alembic_cfg: - if alembic_cfg.config_file_name is not None: - fileConfig(alembic_cfg.config_file_name) - - connectable = create_async_engine( - settings.DATABASE_URL, - poolclass=pool.NullPool, - ) - - async with connectable.connect() as connection: - await connection.run_sync(do_run_migrations) - - await connectable.dispose() - -def run_migrations_offline(): - """Run migrations in 'offline' mode. - - 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. - - """ - if not settings.DATABASE_URL: - raise ValueError("DATABASE_URL not found in settings for Alembic.") - - context.configure( - url=settings.DATABASE_URL, + url=url, target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}, @@ -95,8 +53,32 @@ def run_migrations_offline(): with context.begin_transaction(): context.run_migrations() +async def run_migrations_online() -> None: + """Run migrations in 'online' mode.""" + connectable = create_async_engine( + settings.DATABASE_URL, + poolclass=pool.NullPool, + ) + + async with connectable.connect() as connection: + await connection.run_sync(_run_migrations) + + await connectable.dispose() + +def _run_migrations(connection): + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + compare_server_default=True + ) + + with context.begin_transaction(): + context.run_migrations() + # This section only runs when executing alembic commands directly (not when imported) if context.is_offline_mode(): run_migrations_offline() -elif 'ALEMBIC_CONFIG' in os.environ: # Only run if called from alembic command - asyncio.run(run_migrations_online_async(context.config)) +else: + import asyncio + asyncio.run(run_migrations_online()) diff --git a/be/alembic/migrations.py b/be/alembic/migrations.py new file mode 100644 index 0000000..f943597 --- /dev/null +++ b/be/alembic/migrations.py @@ -0,0 +1,57 @@ +""" +Async migrations handler for FastAPI application. +This file is separate from env.py to avoid Alembic context issues. +""" +import os +import sys +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy import pool + +# Ensure the app directory is in the Python path +sys.path.insert(0, os.path.realpath(os.path.join(os.path.dirname(__file__), '..'))) + +from app.database import Base as DatabaseBase +from app.config import settings + +async def run_migrations(): + """Run database migrations asynchronously.""" + from alembic.runtime.migration import MigrationContext + from alembic.operations import Operations + + # Create async engine + engine = create_async_engine( + settings.DATABASE_URL, + poolclass=pool.NullPool, + ) + + async with engine.connect() as connection: + # Get current database schema version + def do_get_current_rev(): + migration_context = MigrationContext.configure( + connection, + opts={ + 'target_metadata': DatabaseBase.metadata, + 'compare_type': True, + 'compare_server_default': True + } + ) + return migration_context.get_current_revision() + + current_rev = await connection.run_sync(do_get_current_rev) + + # Run migrations + def do_upgrade(): + migration_context = MigrationContext.configure( + connection, + opts={ + 'target_metadata': DatabaseBase.metadata, + 'compare_type': True, + 'compare_server_default': True + } + ) + with Operations.context(migration_context): + migration_context.run_migrations() + + await connection.run_sync(do_upgrade) + + await engine.dispose() \ No newline at end of file diff --git a/be/app/main.py b/be/app/main.py index d4d1230..e56f0a3 100644 --- a/be/app/main.py +++ b/be/app/main.py @@ -228,8 +228,8 @@ async def run_migrations(): sys.path.insert(0, alembic_path) # Import and run migrations - from env import run_migrations_online_async - await run_migrations_online_async() + from migrations import run_migrations as run_db_migrations + await run_db_migrations() logger.info("Database migrations completed successfully.") except Exception as e: -- 2.45.2