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

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,14 +76,12 @@ def run_migrations_online():
 
     logger.info('Testing for an old alembic_version table.')
 
-    connection = engine.connect()
+    with engine.connect() as connection:
         context.configure(
             connection=connection,
-                target_metadata=target_metadata,
-                version_table='alembic_version'
+            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()
@@ -124,6 +124,7 @@ def run_migrations_online():
     this tree but if we still don't have an alembic_version_<tree>
     table, alembic will create it.
     """
+    with engine.connect() as connection:
         context.configure(
             connection=connection,
             target_metadata=target_metadata,
@@ -139,11 +140,8 @@ def run_migrations_online():
             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():