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 9bdbcf0aad)

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'
-                )
+    with engine.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata
+        )
 
-    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
+        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
+                break
 
-    if not found:
-        logger.info('Didn\'t find an old alembic_version table.')
-    logger.info('Trying alembic_version_%s.' % script_location)
+        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)
+    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)
 
-    try:
         with context.begin_transaction():
             context.run_migrations()
-    finally:
-        connection.close()
 
 
 if context.is_offline_mode():