From 8f1da5d440b5747f2a00bdb716b41c1c570120b7 Mon Sep 17 00:00:00 2001
From: mohamad <mohamad>
Date: Sun, 1 Jun 2025 17:42:17 +0200
Subject: [PATCH] refactor: Improve Alembic migration functions by integrating
 configuration and script directory handling for enhanced migration context
 management

---
 be/alembic/migrations.py | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/be/alembic/migrations.py b/be/alembic/migrations.py
index 9bfc5d0..f0d31ef 100644
--- a/be/alembic/migrations.py
+++ b/be/alembic/migrations.py
@@ -6,6 +6,10 @@ 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__), '..')))
@@ -15,8 +19,9 @@ from app.config import settings
 
 async def run_migrations():
     """Run database migrations asynchronously."""
-    from alembic.runtime.migration import MigrationContext
-    from alembic.operations import Operations
+    # 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)
     
     # Create async engine
     engine = create_async_engine(
@@ -25,33 +30,34 @@ async def run_migrations():
     )
     
     async with engine.connect() as connection:
-        # Get current database schema version
-        def do_get_current_rev(conn):
+        def get_current_rev(conn):
             migration_context = MigrationContext.configure(
                 conn,
                 opts={
                     'target_metadata': DatabaseBase.metadata,
-                    'compare_type': True,
-                    'compare_server_default': True
+                    'script': script_directory
                 }
             )
             return migration_context.get_current_revision()
         
-        current_rev = await connection.run_sync(do_get_current_rev)
+        current_rev = await connection.run_sync(get_current_rev)
         
-        # Run migrations
-        def do_upgrade(conn):
+        def upgrade_to_head(conn):
             migration_context = MigrationContext.configure(
                 conn,
                 opts={
                     'target_metadata': DatabaseBase.metadata,
-                    'compare_type': True,
-                    'compare_server_default': True
+                    'script': script_directory,
+                    'as_sql': False,
                 }
             )
-            with Operations.context(migration_context):
+            
+            with migration_context.begin_transaction():
+                migration_context._migrations_fn = script_directory._upgrade_revs(
+                    "head", current_rev
+                )
                 migration_context.run_migrations()
         
-        await connection.run_sync(do_upgrade)
+        await connection.run_sync(upgrade_to_head)
     
     await engine.dispose() 
\ No newline at end of file