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
@@ -522,5 +523,14 @@ def create_or_replace_views(engine) -> None:
         conn.commit()
 
 
+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__":
     engine = create_database()

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"