From 6efa793ec2eaa02fa60894c994dbb761984f0702 Mon Sep 17 00:00:00 2001
From: Takashi Kajinami <kajinamit@oss.nttdata.com>
Date: Fri, 19 Apr 2024 16:03:48 +0900
Subject: [PATCH] Fix db upgrade with SQLAlchemy 2.0

DB schema upgrade is stuck when SQLAlchemy 2.0 is used because all
operations are not executed within a transaction.

Change-Id: Ief7ea81e78344ea32fb2a295f11a5d5a873251a8
---
 aodh/storage/impl_sqlalchemy.py        | 22 ++++++++++++----------
 aodh/storage/sqlalchemy/alembic/env.py | 17 ++++++++++++++---
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/aodh/storage/impl_sqlalchemy.py b/aodh/storage/impl_sqlalchemy.py
index d70db0778..ba929e7ce 100644
--- a/aodh/storage/impl_sqlalchemy.py
+++ b/aodh/storage/impl_sqlalchemy.py
@@ -156,17 +156,19 @@ class Connection(base.Connection):
     def upgrade(self, nocreate=False):
         cfg = self._get_alembic_config()
         cfg.conf = self.conf
-        if nocreate:
-            command.upgrade(cfg, "head")
-        else:
-            engine = enginefacade.writer.get_engine()
-            ctxt = migration.MigrationContext.configure(engine.connect())
-            current_version = ctxt.get_current_revision()
-            if current_version is None:
-                models.Base.metadata.create_all(engine, checkfirst=False)
-                command.stamp(cfg, "head")
-            else:
+        engine = enginefacade.writer.get_engine()
+        with engine.connect() as conn, conn.begin():
+            cfg.attributes['connection'] = conn
+            if nocreate:
                 command.upgrade(cfg, "head")
+            else:
+                ctxt = migration.MigrationContext.configure(conn)
+                current_version = ctxt.get_current_revision()
+                if current_version is None:
+                    models.Base.metadata.create_all(conn, checkfirst=False)
+                    command.stamp(cfg, "head")
+                else:
+                    command.upgrade(cfg, "head")
 
     def clear(self):
         engine = enginefacade.writer.get_engine()
diff --git a/aodh/storage/sqlalchemy/alembic/env.py b/aodh/storage/sqlalchemy/alembic/env.py
index 11b207831..69804bad8 100644
--- a/aodh/storage/sqlalchemy/alembic/env.py
+++ b/aodh/storage/sqlalchemy/alembic/env.py
@@ -68,10 +68,21 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
-    engine = enginefacade.writer.get_engine()
-    with engine.connect() as connection:
+    connectable = config.attributes.get('connection', None)
+
+    if connectable is None:
+        engine = enginefacade.writer.get_engine()
+        with engine.connect() as connection:
+            context.configure(
+                connection=connection,
+                target_metadata=target_metadata
+            )
+
+            with context.begin_transaction():
+                context.run_migrations()
+    else:
         context.configure(
-            connection=connection,
+            connection=connectable,
             target_metadata=target_metadata
         )