be/alembic/migrations.py

@@ -6,10 +6,6 @@ import os
 import sys
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy import pool
-from alembic.config import Config
-from alembic.script import ScriptDirectory
-from alembic.runtime.migration import MigrationContext
-from alembic.operations import Operations
 
 # Ensure the app directory is in the Python path
 sys.path.insert(0, os.path.realpath(os.path.join(os.path.dirname(__file__), '..')))
@@ -19,9 +15,8 @@ from app.config import settings
 
 async def run_migrations():
     """Run database migrations asynchronously."""
-    # Get alembic configuration and script directory
-    alembic_cfg = Config(os.path.join(os.path.dirname(__file__), '..', 'alembic.ini'))
-    script_directory = ScriptDirectory.from_config(alembic_cfg)
+    from alembic.runtime.migration import MigrationContext
+    from alembic.operations import Operations
     
     # Create async engine
     engine = create_async_engine(
@@ -30,34 +25,33 @@ async def run_migrations():
     )
     
     async with engine.connect() as connection:
-        def get_current_rev(conn):
+        # Get current database schema version
+        def do_get_current_rev(conn):
             migration_context = MigrationContext.configure(
                 conn,
                 opts={
                     'target_metadata': DatabaseBase.metadata,
-                    'script': script_directory
+                    'compare_type': True,
+                    'compare_server_default': True
                 }
             )
             return migration_context.get_current_revision()
         
-        current_rev = await connection.run_sync(get_current_rev)
+        current_rev = await connection.run_sync(do_get_current_rev)
         
-        def upgrade_to_head(conn):
+        # Run migrations
+        def do_upgrade(conn):
             migration_context = MigrationContext.configure(
                 conn,
                 opts={
                     'target_metadata': DatabaseBase.metadata,
-                    'script': script_directory,
-                    'as_sql': False,
+                    'compare_type': True,
+                    'compare_server_default': True
                 }
             )
-            
-            with migration_context.begin_transaction():
-                migration_context._migrations_fn = script_directory._upgrade_revs(
-                    "head", current_rev
-                )
+            with Operations.context(migration_context):
                 migration_context.run_migrations()
         
-        await connection.run_sync(upgrade_to_head)
+        await connection.run_sync(do_upgrade)
     
     await engine.dispose()