refactor: Update Alembic migration functions to support asynchronous execution and streamline migration handling in application startup

be/alembic/env.py

@@ -1,15 +1,12 @@
 from logging.config import fileConfig
 import os
 import sys
-import asyncio # Add this import
-
+import asyncio
 from sqlalchemy import engine_from_config
 from sqlalchemy import pool
-from sqlalchemy.ext.asyncio import create_async_engine # Add this specific import
-
+from sqlalchemy.ext.asyncio import create_async_engine
 from alembic import context
 
-
 # Ensure the 'app' directory is in the Python path
 # Adjust the path if your project structure is different
 sys.path.insert(0, os.path.realpath(os.path.join(os.path.dirname(__file__), '..')))
@@ -46,8 +43,29 @@ 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):
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        compare_type=True,
+        compare_server_default=True
+    )
+    with context.begin_transaction():
+        context.run_migrations()
 
-def run_migrations_offline() -> None:
+async def run_migrations_online_async():
+    """Run migrations in 'online' mode asynchronously."""
+    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
@@ -70,38 +88,8 @@ def run_migrations_offline() -> None:
     with context.begin_transaction():
         context.run_migrations()
 
-
-async def run_migrations_online_async() -> None:  # Renamed and make async
-    """Run migrations in 'online' mode.
-
-    In this scenario we need to create an Engine
-    and associate a connection with the context.
-
-    """
-    # connectable here will be an AsyncEngine if the URL is asyncpg
-    db_url = config.get_main_option("sqlalchemy.url") # Get the async URL
-    if not db_url:
-        raise ValueError("Database URL is not configured in Alembic.")
-
-    connectable = create_async_engine(db_url, poolclass=pool.NullPool)
-
-    async with connectable.connect() as connection: # Use async with
-        # Pass target_metadata to the run_sync callback
-        await connection.run_sync(do_run_migrations, target_metadata)
-
-    await connectable.dispose() # Dispose of the async engine
-
-def do_run_migrations(connection, metadata):
-    """Helper function to configure and run migrations within a sync callback."""
-    context.configure(
-        connection=connection,
-        target_metadata=metadata
-        # Include other options like compare_type=True, compare_server_default=True if needed
-    )
-    with context.begin_transaction():
-        context.run_migrations()
-
 if context.is_offline_mode():
     run_migrations_offline()
 else:
-    asyncio.run(run_migrations_online_async()) # Call the new async function
+    # Don't run migrations here - they will be run from FastAPI
+    pass

be/app/main.py

@@ -13,6 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine
 from alembic.config import Config
 from alembic import command
 import os
+import sys
 
 from app.api.api_router import api_router
 from app.config import settings
@@ -232,8 +233,12 @@ async def run_migrations():
             raise ValueError("DATABASE_URL is not configured in settings.")
         alembic_cfg.set_main_option('sqlalchemy.url', settings.DATABASE_URL)
         
-        # Run the migration
-        command.upgrade(alembic_cfg, "head")
+        # Import the async migration function from env.py
+        sys.path.insert(0, script_location)
+        from env import run_migrations_online_async
+        
+        # Run the migration asynchronously
+        await run_migrations_online_async()
         
         logger.info("Database migrations completed successfully.")
     except Exception as e: