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 1944c9d72e8a87f502aa6029c063ca1b7805d432)

diff --git a/contrib/ast-db-manage/env.py b/contrib/ast-db-manage/env.py
index 5036f4ccfe5c6e87b465add9cbe03d042a90d5cf..797d9ce3e3b53b91d5e6a880edf230f24ed5d9cc 100644 (file)
--- a/contrib/ast-db-manage/env.py
+++ b/contrib/ast-db-manage/env.py
@@ -67,6 +67,8 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
+    script_location = config.get_main_option('script_location')
+
     engine = engine_from_config(
                 config.get_section(config.config_ini_section),
                 prefix='sqlalchemy.',
@@ -74,76 +76,72 @@ def run_migrations_online():
 
     logger.info('Testing for an old alembic_version table.')
 
-    connection = engine.connect()
-    context.configure(
-                connection=connection,
-                target_metadata=target_metadata,
-                version_table='alembic_version'
-                )
-
-    script_location = config.get_main_option('script_location')
-    found = False
-    mc = context.get_context()
-    current_db_revision = mc.get_current_revision()
-    script = ScriptDirectory.from_config(config)
-    """ If there was an existing alembic_version table, we need to
-    check that it's current revision is in the history for the tree
-    we're working with.
-    """
-    for x in script.iterate_revisions('head', 'base'):
-        if x.revision == current_db_revision:
-            """ An alembic_versions table was found and it belongs to
-            this alembic tree
-            """
-            logger.info(
-                ('An old alembic_version table at revision %s was '
-                 'found for %s.  Renaming to alembic_version_%s.'),
-                        current_db_revision, script_location,
-                        script_location)
-            op = Operations(mc)
-            try:
-                with context.begin_transaction():
-                    op.rename_table(
-                        'alembic_version', 'alembic_version_%s'
-                        % script_location)
-                found = True
-            except:
-                logger.error(('Unable to rename alembic_version to '
-                             'alembic_version_%s.'),
-                             script_location)
-                connection.close()
-                return
-
-            break
-
-    if not found:
-        logger.info('Didn\'t find an old alembic_version table.')
-    logger.info('Trying alembic_version_%s.' % script_location)
+    with engine.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata
+        )
+
+        found = False
+        mc = context.get_context()
+        current_db_revision = mc.get_current_revision()
+        script = ScriptDirectory.from_config(config)
+        """ If there was an existing alembic_version table, we need to
+        check that it's current revision is in the history for the tree
+        we're working with.
+        """
+        for x in script.iterate_revisions('head', 'base'):
+            if x.revision == current_db_revision:
+                """ An alembic_versions table was found and it belongs to
+                this alembic tree
+                """
+                logger.info(
+                    ('An old alembic_version table at revision %s was '
+                     'found for %s.  Renaming to alembic_version_%s.'),
+                            current_db_revision, script_location,
+                            script_location)
+                op = Operations(mc)
+                try:
+                    with context.begin_transaction():
+                        op.rename_table(
+                            'alembic_version', 'alembic_version_%s'
+                            % script_location)
+                    found = True
+                except:
+                    logger.error(('Unable to rename alembic_version to '
+                                'alembic_version_%s.'),
+                                script_location)
+                    connection.close()
+                    return
+
+                break
+
+        if not found:
+            logger.info('Didn\'t find an old alembic_version table.')
+        logger.info('Trying alembic_version_%s.' % script_location)
 
     """ 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>
     table, alembic will create it.
     """
-    context.configure(
-                connection=connection,
-                target_metadata=target_metadata,
-                version_table='alembic_version_' + script_location
-                )
-    mc = context.get_context()
-    current_db_revision = mc.get_current_revision()
-    if current_db_revision:
-        logger.info(
-            'Using the alembic_version_%s table at revision %s.',
-            script_location, current_db_revision)
-    else:
-        logger.info('Creating new alembic_version_%s table.',
-                    script_location)
-
-    try:
+    with engine.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table='alembic_version_' + script_location
+        )
+        mc = context.get_context()
+        current_db_revision = mc.get_current_revision()
+        if current_db_revision:
+            logger.info(
+                'Using the alembic_version_%s table at revision %s.',
+                script_location, current_db_revision)
+        else:
+            logger.info('Creating new alembic_version_%s table.',
+                        script_location)
+
         with context.begin_transaction():
             context.run_migrations()
-    finally:
-        connection.close()
 
 
 if context.is_offline_mode():