From bb033050403eefc128fb9d55d0350a64286974a7 Mon Sep 17 00:00:00 2001
From: Alessio Siniscalchi <a.siniscalchi@bopen.eu>
Date: Thu, 25 Jul 2024 18:00:04 +0200
Subject: [PATCH] fix creation of alembic table

---
 cads_broker/database.py   | 4 ++--
 tests/test_02_database.py | 5 ++++-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/cads_broker/database.py b/cads_broker/database.py
index f2be84c..4ce2462 100644
--- a/cads_broker/database.py
+++ b/cads_broker/database.py
@@ -879,9 +879,9 @@ def init_database(connection_string: str, force: bool = False) -> sa.engine.Engi
         sqlalchemy_utils.create_database(engine.url)
         # cleanup and create the schema
         BaseModel.metadata.drop_all(engine)
+        alembic.command.stamp(alembic_cfg, "head")
         cacholote.init_database(connection_string, force)
         BaseModel.metadata.create_all(engine)
-        alembic.command.stamp(alembic_cfg, "head")
     else:
         # check the structure is empty or incomplete
         query = sa.text(
@@ -894,9 +894,9 @@ def init_database(connection_string: str, force: bool = False) -> sa.engine.Engi
     if force:
         # cleanup and create the schema
         BaseModel.metadata.drop_all(engine)
+        alembic.command.stamp(alembic_cfg, "head")
         cacholote.init_database(connection_string, force)
         BaseModel.metadata.create_all(engine)
-        alembic.command.stamp(alembic_cfg, "head")
     else:
         # update db structure
         cacholote.init_database(connection_string, force)
diff --git a/tests/test_02_database.py b/tests/test_02_database.py
index f288220..d295e1e 100644
--- a/tests/test_02_database.py
+++ b/tests/test_02_database.py
@@ -977,11 +977,13 @@ def test_init_database(postgresql: Connection[str]) -> None:
     # start with an empty db structure
     expected_tables_at_beginning: set[str] = set()
     assert set(conn.execute(query).scalars()) == expected_tables_at_beginning  # type: ignore
-
+    conn.close()
     # verify create structure
     db.init_database(connection_string, force=True)
+    conn = engine.connect()
     expected_tables_complete = (
         set(db.BaseModel.metadata.tables)
+        .union({"alembic_version"})
         .union({"alembic_version_cacholote"})
         .union(set(cacholote.database.Base.metadata.tables))
     )
@@ -1032,6 +1034,7 @@ def test_init_database_with_password(postgresql2: Connection[str]) -> None:
     db.init_database(connection_string, force=True)
     expected_tables_complete = (
         set(db.BaseModel.metadata.tables)
+        .union({"alembic_version"})
         .union({"alembic_version_cacholote"})
         .union(set(cacholote.database.Base.metadata.tables))
     )