Merge pull request 'refactor: Improve Alembic migration functions by integrating configuration and script directory handling for enhanced migration context management' (#31) from ph4 into prod
Reviewed-on: #31
This commit is contained in:
commit
26315cd407
@ -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()
|
Loading…
Reference in New Issue
Block a user