Fix post-reform calibration test compatibility

Author: MaxGhenis

policyengine_us_data/calibration/unified_matrix_builder.py

@@ -19,6 +19,7 @@
 from scipy import sparse
 from sqlalchemy import create_engine, text
 
+from policyengine_us_data.db.create_database_tables import refresh_sql_views
 from policyengine_us_data.storage import STORAGE_FOLDER
 from policyengine_us_data.utils.census import STATE_NAME_TO_FIPS
 from policyengine_us_data.calibration.calibration_utils import (
@@ -928,6 +929,8 @@ def __init__(
     ):
         self.db_uri = db_uri
         self.engine = create_engine(db_uri)
+        # Existing SQLite checkpoints may carry an older target_overview view.
+        refresh_sql_views(self.engine)
         self.time_period = time_period
         self.dataset_path = dataset_path
         self._entity_rel_cache = None

policyengine_us_data/calibration/validate_staging.py

@@ -41,6 +41,7 @@
 from policyengine_us_data.calibration.sanity_checks import (
     run_sanity_checks,
 )
+from policyengine_us_data.db.create_database_tables import refresh_sql_views
 
 logger = logging.getLogger(__name__)
 
@@ -537,6 +538,7 @@ def _validate_single_area(
     from sqlalchemy import create_engine as _create_engine
 
     engine = _create_engine(f"sqlite:///{db_path}")
+    refresh_sql_views(engine)
 
     logger.info("Loading sim from %s", h5_path)
     try:
@@ -1015,6 +1017,7 @@ def main(argv=None):
     from policyengine_us import Microsimulation
 
     engine = create_engine(f"sqlite:///{args.db_path}")
+    refresh_sql_views(engine)
 
     all_targets = _query_all_active_targets(engine, args.period)
     logger.info("Loaded %d active targets from DB", len(all_targets))

policyengine_us_data/db/create_database_tables.py

@@ -353,6 +353,16 @@ def validate_parent_child_constraints(mapper, connection, target: Stratum):
 """
 
 
+def refresh_sql_views(engine) -> None:
+    """Recreate derived SQL views so existing DBs pick up schema changes."""
+    with engine.connect() as conn:
+        conn.execute(text("DROP VIEW IF EXISTS target_overview"))
+        conn.execute(text("DROP VIEW IF EXISTS stratum_domain"))
+        conn.execute(text(STRATUM_DOMAIN_VIEW))
+        conn.execute(text(TARGET_OVERVIEW_VIEW))
+        conn.commit()
+
+
 def create_validation_triggers(engine) -> None:
     """Create SQL triggers that validate fields against field_valid_values.
 
@@ -506,11 +516,8 @@ def create_database(
     # Create validation triggers
     create_validation_triggers(engine)
 
-    # Create SQL views
-    with engine.connect() as conn:
-        conn.execute(text(STRATUM_DOMAIN_VIEW))
-        conn.execute(text(TARGET_OVERVIEW_VIEW))
-        conn.commit()
+    # Recreate SQL views so existing DB files do not keep stale definitions.
+    refresh_sql_views(engine)
 
     logger.info(f"Database and tables created successfully at {db_uri}")
     return engine

policyengine_us_data/tests/test_schema_views_and_lookups.py

@@ -11,13 +11,15 @@
 import tempfile
 import unittest
 
+from sqlalchemy import text
 from sqlmodel import Session
 
 from policyengine_us_data.db.create_database_tables import (
     Stratum,
     StratumConstraint,
     Target,
     create_database,
+    refresh_sql_views,
 )
 from policyengine_us_data.utils.db import get_geographic_strata
 from policyengine_us_data.calibration.calibration_utils import (
@@ -399,6 +401,24 @@ def test_reform_id_passthrough(self):
         self.assertEqual(len(matches), 1)
         self.assertEqual(matches[0][reform_idx], 1)
 
+    def test_refresh_sql_views_updates_existing_target_overview(self):
+        """Refreshing views updates stale target_overview definitions."""
+        with self.engine.connect() as conn:
+            conn.execute(text("DROP VIEW IF EXISTS target_overview"))
+            conn.execute(
+                text(
+                    "CREATE VIEW target_overview AS "
+                    "SELECT target_id, stratum_id, variable, value, period, active "
+                    "FROM targets"
+                )
+            )
+            conn.commit()
+
+        refresh_sql_views(self.engine)
+
+        cols = self._overview_columns()
+        self.assertIn("reform_id", cols)
+
     # ----------------------------------------------------------------
     # get_geographic_strata()
     # ----------------------------------------------------------------