From 5dcabd51f7a814f98619024bb94a5297fba4300a Mon Sep 17 00:00:00 2001 From: mohamad Date: Sun, 1 Jun 2025 17:45:32 +0200 Subject: [PATCH] refactor: Introduce migration function to streamline upgrade steps in Alembic migrations --- be/alembic/migrations.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/be/alembic/migrations.py b/be/alembic/migrations.py index f0d31ef..b535a37 100644 --- a/be/alembic/migrations.py +++ b/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)