Refresh downloaded policy DB views

Author: MaxGhenis

policyengine_us_data/db/create_database_tables.py

@@ -1,5 +1,6 @@
-import logging
 import hashlib
+import logging
+from pathlib import Path
 from typing import List, Optional
 
 from sqlalchemy import event, text, UniqueConstraint
@@ -506,14 +507,29 @@ def create_database(
     # Create validation triggers
     create_validation_triggers(engine)
 
-    # Create SQL views
+    create_or_replace_views(engine)
+
+    logger.info(f"Database and tables created successfully at {db_uri}")
+    return engine
+
+
+def create_or_replace_views(engine) -> None:
+    """Refresh SQL views so existing databases pick up schema changes."""
     with engine.connect() as conn:
+        conn.execute(text("DROP VIEW IF EXISTS stratum_domain"))
+        conn.execute(text("DROP VIEW IF EXISTS target_overview"))
         conn.execute(text(STRATUM_DOMAIN_VIEW))
         conn.execute(text(TARGET_OVERVIEW_VIEW))
         conn.commit()
 
-    logger.info(f"Database and tables created successfully at {db_uri}")
-    return engine
+
+def refresh_views_for_db_path(db_path: str | Path) -> None:
+    """Refresh SQL views for an existing SQLite database file."""
+    engine = create_engine(f"sqlite:///{Path(db_path)}")
+    try:
+        create_or_replace_views(engine)
+    finally:
+        engine.dispose()
 
 
 if __name__ == "__main__":

policyengine_us_data/storage/download_private_prerequisites.py

@@ -1,7 +1,10 @@
 import os
 
-from policyengine_us_data.utils.huggingface import download
 from pathlib import Path
+from policyengine_us_data.db.create_database_tables import (
+    refresh_views_for_db_path,
+)
+from policyengine_us_data.utils.huggingface import download
 
 FOLDER = Path(__file__).parent
 
@@ -41,3 +44,4 @@
         local_folder=FOLDER,
         version=None,
     )
+    refresh_views_for_db_path(FOLDER / "calibration" / "policy_data.db")

policyengine_us_data/tests/test_calibration/conftest.py

@@ -1,10 +1,25 @@
 """Shared fixtures for local area calibration tests."""
 
 import pytest
+from sqlalchemy import create_engine
 
+from policyengine_us_data.db.create_database_tables import (
+    create_or_replace_views,
+)
 from policyengine_us_data.storage import STORAGE_FOLDER
 
 
+@pytest.fixture(scope="session", autouse=True)
+def refresh_policy_db_views():
+    db_path = STORAGE_FOLDER / "calibration" / "policy_data.db"
+    if db_path.exists():
+        engine = create_engine(f"sqlite:///{db_path}")
+        try:
+            create_or_replace_views(engine)
+        finally:
+            engine.dispose()
+
+
 @pytest.fixture(scope="module")
 def db_uri():
     db_path = STORAGE_FOLDER / "calibration" / "policy_data.db"

policyengine_us_data/tests/test_database.py

@@ -1,4 +1,5 @@
 import hashlib
+from sqlalchemy import text
 
 import pytest
 from sqlalchemy.exc import IntegrityError
@@ -8,6 +9,7 @@
     Stratum,
     StratumConstraint,
     Target,
+    create_or_replace_views,
     create_database,
 )
 
@@ -226,6 +228,38 @@ def test_target_with_null_source(engine):
         assert retrieved.source is None
 
 
+def test_create_database_refreshes_existing_views(tmp_path):
+    db_uri = f"sqlite:///{tmp_path / 'test.db'}"
+    engine = create_database(db_uri)
+
+    with engine.connect() as conn:
+        conn.execute(text("DROP VIEW target_overview"))
+        conn.execute(
+            text(
+                """
+                CREATE VIEW target_overview AS
+                SELECT
+                    t.target_id,
+                    t.stratum_id,
+                    t.variable,
+                    t.value,
+                    t.period,
+                    t.active
+                FROM targets t
+                """
+            )
+        )
+        conn.commit()
+
+    create_or_replace_views(engine)
+
+    with engine.connect() as conn:
+        cursor = conn.execute(text("SELECT * FROM target_overview LIMIT 0"))
+        columns = [desc[0] for desc in cursor.cursor.description]
+
+    assert "reform_id" in columns
+
+
 def test_valid_geographic_hierarchy(engine):
     """CD under its correct state should succeed."""
     with Session(engine) as session: