fix creation of alembic table

ecmwf-projects · Jul 25, 2024 · bb03305 · bb03305
1 parent aa53cac
commit bb03305
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git 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
@@ -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))
     )