Merge pull request 'refactor: Introduce migration function to streamline upgrade steps in Alembic migrations' (#32) from ph4 into prod

Reviewed-on: #32

be/alembic/migrations.py

@@ -17,6 +17,16 @@ 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
 
+def _get_migration_fn(script_directory, current_rev):
+    """Create a migration function that knows how to upgrade from current revision."""
+    def migration_fn(rev, context):
+        # Get all upgrade steps from current revision to head
+        revisions = script_directory._upgrade_revs("head", current_rev)
+        for revision in revisions:
+            script = script_directory.get_revision(revision)
+            script.module.upgrade(context)
+    return migration_fn
+
 async def run_migrations():
     """Run database migrations asynchronously."""
     # Get alembic configuration and script directory
@@ -52,10 +62,10 @@ async def run_migrations():
                 }
             )
             
+            # Set the migration function
+            migration_context._migrations_fn = _get_migration_fn(script_directory, current_rev)
+            
             with migration_context.begin_transaction():
-                migration_context._migrations_fn = script_directory._upgrade_revs(
-                    "head", current_rev
-                )
                 migration_context.run_migrations()
         
         await connection.run_sync(upgrade_to_head)