alembic: Fix compatibility with SQLAlchemy 2.0+.

SQLAlchemy 2.0 changed the way that commits/rollbacks are handled
causing the final `UPDATE` to our `alembic_version_<whatever>` tables
to be rolled back instead of committed.

We now use one connection to determine which
`alembic_version_<whatever>` table to use and another to run the
actual migrations. This prevents the erroneous rollback.

This change is compatible with both SQLAlchemy 1.4 and 2.0.
This commit is contained in:
Sean Bright
2024-03-20 12:20:40 -04:00
parent 8513881334
commit 9bdbcf0aad

View File

@@ -67,6 +67,8 @@ def run_migrations_online():
and associate a connection with the context. and associate a connection with the context.
""" """
script_location = config.get_main_option('script_location')
engine = engine_from_config( engine = engine_from_config(
config.get_section(config.config_ini_section), config.get_section(config.config_ini_section),
prefix='sqlalchemy.', prefix='sqlalchemy.',
@@ -74,14 +76,12 @@ def run_migrations_online():
logger.info('Testing for an old alembic_version table.') logger.info('Testing for an old alembic_version table.')
connection = engine.connect() with engine.connect() as connection:
context.configure( context.configure(
connection=connection, connection=connection,
target_metadata=target_metadata, target_metadata=target_metadata
version_table='alembic_version'
) )
script_location = config.get_main_option('script_location')
found = False found = False
mc = context.get_context() mc = context.get_context()
current_db_revision = mc.get_current_revision() current_db_revision = mc.get_current_revision()
@@ -124,6 +124,7 @@ def run_migrations_online():
this tree but if we still don't have an alembic_version_<tree> this tree but if we still don't have an alembic_version_<tree>
table, alembic will create it. table, alembic will create it.
""" """
with engine.connect() as connection:
context.configure( context.configure(
connection=connection, connection=connection,
target_metadata=target_metadata, target_metadata=target_metadata,
@@ -139,11 +140,8 @@ def run_migrations_online():
logger.info('Creating new alembic_version_%s table.', logger.info('Creating new alembic_version_%s table.',
script_location) script_location)
try:
with context.begin_transaction(): with context.begin_transaction():
context.run_migrations() context.run_migrations()
finally:
connection.close()
if context.is_offline_mode(): if context.is_offline_mode():