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.

(cherry picked from commit 8715a700e2)
This commit is contained in:
Sean Bright
2024-03-20 12:20:40 -04:00
committed by Asterisk Development Team
parent 2680d5be35
commit 396cc55965

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,76 +76,72 @@ 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() script = ScriptDirectory.from_config(config)
script = ScriptDirectory.from_config(config) """ If there was an existing alembic_version table, we need to
""" If there was an existing alembic_version table, we need to check that it's current revision is in the history for the tree
check that it's current revision is in the history for the tree we're working with.
we're working with. """
""" for x in script.iterate_revisions('head', 'base'):
for x in script.iterate_revisions('head', 'base'): if x.revision == current_db_revision:
if x.revision == current_db_revision: """ An alembic_versions table was found and it belongs to
""" An alembic_versions table was found and it belongs to this alembic tree
this alembic tree """
""" logger.info(
logger.info( ('An old alembic_version table at revision %s was '
('An old alembic_version table at revision %s was ' 'found for %s. Renaming to alembic_version_%s.'),
'found for %s. Renaming to alembic_version_%s.'), current_db_revision, script_location,
current_db_revision, script_location, script_location)
script_location) op = Operations(mc)
op = Operations(mc) try:
try: with context.begin_transaction():
with context.begin_transaction(): op.rename_table(
op.rename_table( 'alembic_version', 'alembic_version_%s'
'alembic_version', 'alembic_version_%s' % script_location)
% script_location) found = True
found = True except:
except: logger.error(('Unable to rename alembic_version to '
logger.error(('Unable to rename alembic_version to ' 'alembic_version_%s.'),
'alembic_version_%s.'), script_location)
script_location) connection.close()
connection.close() return
return
break break
if not found: if not found:
logger.info('Didn\'t find an old alembic_version table.') logger.info('Didn\'t find an old alembic_version table.')
logger.info('Trying alembic_version_%s.' % script_location) logger.info('Trying alembic_version_%s.' % script_location)
""" We MAY have an alembic_version table that doesn't belong to """ We MAY have an alembic_version table that doesn't belong to
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.
""" """
context.configure( with engine.connect() as connection:
connection=connection, context.configure(
target_metadata=target_metadata, connection=connection,
version_table='alembic_version_' + script_location target_metadata=target_metadata,
) version_table='alembic_version_' + script_location
mc = context.get_context() )
current_db_revision = mc.get_current_revision() mc = context.get_context()
if current_db_revision: current_db_revision = mc.get_current_revision()
logger.info( if current_db_revision:
'Using the alembic_version_%s table at revision %s.', logger.info(
script_location, current_db_revision) 'Using the alembic_version_%s table at revision %s.',
else: script_location, current_db_revision)
logger.info('Creating new alembic_version_%s table.', else:
script_location) logger.info('Creating new alembic_version_%s table.',
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():