change adoption to separate step

scripts/location_db_adopt.py

@@ -6,6 +6,8 @@ from pathlib import Path
 
 from alembic import command
 from alembic.config import Config
+from alembic.script import ScriptDirectory
+from alembic.util.exc import CommandError
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 if str(PROJECT_ROOT) not in sys.path:
@@ -43,6 +45,24 @@ def _make_alembic_config(database_url: str) -> Config:
     return config
 
 
+def _expected_head_revision(alembic_config: Config) -> str:
+    script = ScriptDirectory.from_config(alembic_config)
+    heads = script.get_heads()
+    if len(heads) != 1:
+        raise LocationDatabaseAdoptionError(
+            f"Expected exactly one Alembic head for location DB, got {len(heads)}"
+        )
+    return heads[0]
+
+
+def _is_known_revision(alembic_config: Config, revision: str) -> bool:
+    script = ScriptDirectory.from_config(alembic_config)
+    try:
+        return script.get_revision(revision) is not None
+    except CommandError:
+        return False
+
+
 def _location_table_exists(database_path: Path) -> bool:
     conn = sqlite3.connect(database_path)
     try:
@@ -117,6 +137,8 @@ def validate_legacy_location_db(database_url: str) -> None:
 
 def validate_location_runtime_db(database_url: str) -> None:
     database_path = _database_path_from_url(database_url)
+    alembic_config = _make_alembic_config(database_url)
+    expected_revision = _expected_head_revision(alembic_config)
     if not database_path.exists():
         raise LocationDatabaseAdoptionError(
             "Location DB file was not found. Run 'python scripts/location_db_adopt.py' "
@@ -131,30 +153,36 @@ def validate_location_runtime_db(database_url: str) -> None:
         )
 
     current_revision = _fetch_alembic_revision(database_path)
-    if current_revision != LOCATION_BASELINE_REVISION:
+    if current_revision != expected_revision:
         raise LocationDatabaseAdoptionError(
             "Location DB revision mismatch. Refusing to start the app: "
-            f"expected {LOCATION_BASELINE_REVISION}, got {current_revision}"
+            f"expected {expected_revision}, got {current_revision}"
         )
 
 
 def adopt_or_initialize_location_db(database_url: str) -> str:
     database_path = _database_path_from_url(database_url)
     alembic_config = _make_alembic_config(database_url)
+    expected_revision = _expected_head_revision(alembic_config)
 
     if database_path.exists():
         if _alembic_version_table_exists(database_path):
             current_revision = _fetch_alembic_revision(database_path)
-            if current_revision != LOCATION_BASELINE_REVISION:
+            if current_revision == expected_revision:
+                return "already_managed"
+            if not _is_known_revision(alembic_config, current_revision):
                 raise LocationDatabaseAdoptionError(
                     "Location DB is already Alembic-managed but revision does not match "
-                    f"the expected baseline: expected {LOCATION_BASELINE_REVISION}, "
-                    f"got {current_revision}"
+                    f"a known migration revision: got {current_revision}"
                 )
-            return "already_managed"
+            command.upgrade(alembic_config, "head")
+            return "upgraded"
 
         validate_legacy_location_db(database_url)
         command.stamp(alembic_config, LOCATION_BASELINE_REVISION)
+        if LOCATION_BASELINE_REVISION != expected_revision:
+            command.upgrade(alembic_config, "head")
+            return "upgraded"
         return "adopted"
 
     database_path.parent.mkdir(parents=True, exist_ok=True)