From 49d760b553482d2f53bce782e06bb1eabaf97370 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:06:51 +0300 Subject: [PATCH 01/89] test(red): define ORM session lifecycle --- tests/unit/data/test_orm_sessions.py | 53 ++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 tests/unit/data/test_orm_sessions.py diff --git a/tests/unit/data/test_orm_sessions.py b/tests/unit/data/test_orm_sessions.py new file mode 100644 index 000000000..82e1d5d8f --- /dev/null +++ b/tests/unit/data/test_orm_sessions.py @@ -0,0 +1,53 @@ +import pytest +from sqlalchemy import text + +from policyengine_api.data.orm import SessionManager, build_sqlite_session_manager + + +def test_session_manager_commits_successful_transaction(): + manager = build_sqlite_session_manager() + with manager.engine.begin() as connection: + connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) + + manager.run_in_transaction( + lambda session: session.execute(text("INSERT INTO item (id) VALUES (1)")) + ) + + with manager.session() as session: + assert session.execute(text("SELECT id FROM item")).scalar_one() == 1 + + +def test_session_manager_rolls_back_failed_transaction(): + manager = build_sqlite_session_manager() + with manager.engine.begin() as connection: + connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) + + def fail(session): + session.execute(text("INSERT INTO item (id) VALUES (1)")) + raise RuntimeError("stop") + + with pytest.raises(RuntimeError, match="stop"): + manager.run_in_transaction(fail) + + with manager.session() as session: + assert session.execute(text("SELECT COUNT(*) FROM item")).scalar_one() == 0 + + +def test_session_manager_closes_sessions_after_callback(monkeypatch): + manager = build_sqlite_session_manager() + closed = [] + original_close = manager.session_factory.class_.close + + def recording_close(session): + closed.append(session) + return original_close(session) + + monkeypatch.setattr(manager.session_factory.class_, "close", recording_close) + manager.run_in_transaction(lambda session: None) + + assert len(closed) == 1 + + +def test_session_manager_requires_an_engine(): + with pytest.raises(TypeError): + SessionManager() # type: ignore[call-arg] From 5fb84413fea283e3643cdc84c6824569c89b59ac Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:07:14 +0300 Subject: [PATCH 02/89] feat(green): add ORM session management --- policyengine_api/data/orm.py | 61 ++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 policyengine_api/data/orm.py diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py new file mode 100644 index 000000000..81401a224 --- /dev/null +++ b/policyengine_api/data/orm.py @@ -0,0 +1,61 @@ +"""SQLAlchemy session ownership for the v1 persistence layer.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import TypeVar + +from sqlalchemy import Engine, create_engine +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import StaticPool + + +T = TypeVar("T") + + +class SessionManager: + """Own sessions and transaction boundaries without leaking either to callers.""" + + def __init__(self, engine: Engine): + self.engine = engine + self.session_factory = sessionmaker( + bind=engine, + class_=Session, + expire_on_commit=False, + ) + + @contextmanager + def session(self) -> Iterator[Session]: + session = self.session_factory() + try: + yield session + finally: + session.close() + + def run_in_transaction(self, callback: Callable[[Session], T]) -> T: + with self.session() as session: + try: + result = callback(session) + session.commit() + return result + except Exception: + session.rollback() + raise + + +def build_sqlite_session_manager( + database_path: str | Path | None = None, +) -> SessionManager: + """Build a SQLite manager for local execution or isolated tests.""" + + if database_path is None: + engine = create_engine( + "sqlite+pysqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + else: + engine = create_engine(f"sqlite+pysqlite:///{Path(database_path)}") + return SessionManager(engine) From 7464563f462d2ce00dedf278484000a33e0f455a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:07:27 +0300 Subject: [PATCH 03/89] test(red): require canonical Alembic AI guidance --- tests/unit/test_alembic_skill.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 tests/unit/test_alembic_skill.py diff --git a/tests/unit/test_alembic_skill.py b/tests/unit/test_alembic_skill.py new file mode 100644 index 000000000..963eb82ba --- /dev/null +++ b/tests/unit/test_alembic_skill.py @@ -0,0 +1,26 @@ +from pathlib import Path + +from policyengine_api.constants import REPO + + +SKILL = REPO / "docs" / "engineering" / "skills" / "alembic-migrations.md" + + +def test_model_agnostic_alembic_skill_is_discoverable(): + assert SKILL.exists() + assert "alembic-migrations.md" in ( + REPO / "docs" / "engineering" / "skills" / "README.md" + ).read_text() + for adapter in ("AGENTS.md", "CLAUDE.md", ".github/copilot-instructions.md"): + assert "docs/engineering/skills/alembic-migrations.md" in ( + REPO / adapter + ).read_text() + + +def test_alembic_skill_forbids_handwritten_ai_revisions(): + guidance = SKILL.read_text() + assert "MUST NOT manually author Alembic revision scripts" in guidance + assert "alembic revision --autogenerate" in guidance + assert "dialect compatibility" in guidance + assert "reversibility" in guidance + assert "request a human migration decision" in guidance From 71bca5cd9e445c6883aa3037dfda570cd4a7f536 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:07:49 +0300 Subject: [PATCH 04/89] docs(green): add model-agnostic Alembic skill --- .github/copilot-instructions.md | 3 ++ AGENTS.md | 4 +++ CLAUDE.md | 4 +++ docs/engineering/skills/README.md | 2 ++ docs/engineering/skills/alembic-migrations.md | 36 +++++++++++++++++++ 5 files changed, 49 insertions(+) create mode 100644 docs/engineering/skills/alembic-migrations.md diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 8342cdc3f..eb1e5e7ab 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -10,5 +10,8 @@ or migration guard changes, read For tests, read `docs/engineering/skills/testing.md` before adding, moving, or reviewing test files. +For SQLAlchemy model or Alembic migration work, read +`docs/engineering/skills/alembic-migrations.md`. + For pull requests, read `docs/engineering/skills/github-prs.md` before opening, replacing, or sharing a PR. diff --git a/AGENTS.md b/AGENTS.md index 3113e7459..3837d2f22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,6 +15,10 @@ cutover plans, generated migration docs, or migration guard scripts, read When adding, moving, or reviewing tests, read `docs/engineering/skills/testing.md`. +When changing SQLAlchemy models, Alembic configuration, database schemas, or +migration revisions, read +`docs/engineering/skills/alembic-migrations.md`. + ## GitHub PRs Read `docs/engineering/skills/github-prs.md` before opening, replacing, or diff --git a/CLAUDE.md b/CLAUDE.md index bb654f145..6627f4b40 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -23,6 +23,10 @@ generated migration docs, or migration guard scripts, read When adding, moving, or reviewing tests, read `docs/engineering/skills/testing.md`. +When changing SQLAlchemy models, Alembic configuration, database schemas, or +migration revisions, read +`docs/engineering/skills/alembic-migrations.md`. + ## Safety Boundaries Do not claim a route, database table, compute path, or deployment surface has diff --git a/docs/engineering/skills/README.md b/docs/engineering/skills/README.md index 2f3fea74b..1aa49f799 100644 --- a/docs/engineering/skills/README.md +++ b/docs/engineering/skills/README.md @@ -9,6 +9,8 @@ first, then keep adapters thin. Current skills: +- `alembic-migrations.md`: mandatory autogenerated Alembic revision workflow, + adoption safeguards, and migration validation. - `github-prs.md`: PR workflow and migration PR handoff expectations. - `migration_contracts.md`: API v2 migration route contracts, route-group metadata, generated migration artifacts, and quality guards. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md new file mode 100644 index 000000000..10c1d55b3 --- /dev/null +++ b/docs/engineering/skills/alembic-migrations.md @@ -0,0 +1,36 @@ +# Alembic Migrations + +Use this skill whenever adding or changing SQLAlchemy models, Alembic +configuration, database schema, or migration revisions. + +## Mandatory generation rule + +AI **MUST NOT manually author Alembic revision scripts**. Generate every schema +revision from reviewed SQLAlchemy metadata: + +```bash +alembic revision --autogenerate -m "" +``` + +If generated operations are wrong, first correct the model metadata and +regenerate. AI may make minimal post-generation corrections only for dialect compatibility, +reversibility, or an objectively incorrect autogenerated +operation. Document each correction and cover it with upgrade and downgrade +tests. Do not add arbitrary schema or data operations to a generated revision. + +If the required migration cannot be expressed safely by autogeneration plus +those narrow review corrections, stop and request a human migration decision. + +## Required checks + +Before committing a migration: + +1. Run `alembic check` and review the generated operations. +2. Upgrade a fresh database to `head`. +3. Compare any existing database to the ORM metadata before stamping it. +4. Downgrade one revision and upgrade to `head` again in an isolated database. +5. Confirm application startup performs no implicit DDL. + +Never print database credentials or embed them in Alembic configuration. Never +run a generated baseline's create operations against an existing database; +verify schema equivalence and stamp it instead. From e60600a777e093c310506ddefb73f6c1799438d1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:08:19 +0300 Subject: [PATCH 05/89] test(red): specify the existing v1 ORM schema --- tests/unit/data/test_v1_models.py | 55 +++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 tests/unit/data/test_v1_models.py diff --git a/tests/unit/data/test_v1_models.py b/tests/unit/data/test_v1_models.py new file mode 100644 index 000000000..6d92f1d45 --- /dev/null +++ b/tests/unit/data/test_v1_models.py @@ -0,0 +1,55 @@ +from sqlalchemy import create_engine + +from policyengine_api.data.v1_models import V1Base + + +EXPECTED_TABLES = { + "analysis", + "computed_household", + "economy", + "household", + "legacy_report_output_aliases", + "policy", + "reform_impact", + "report_output_runs", + "report_outputs", + "simulation_runs", + "simulations", + "tracers", + "user_policies", + "user_profiles", +} + + +def test_v1_metadata_contains_every_legacy_table(): + assert set(V1Base.metadata.tables) == EXPECTED_TABLES + + +def test_v1_metadata_builds_a_fresh_sqlite_database(): + engine = create_engine("sqlite+pysqlite:///:memory:") + V1Base.metadata.create_all(engine) + + with engine.connect() as connection: + table_names = { + row[0] + for row in connection.exec_driver_sql( + "SELECT name FROM sqlite_master WHERE type='table'" + ) + } + assert EXPECTED_TABLES <= table_names + + +def test_v1_composite_and_unique_keys_match_legacy_contract(): + policy = V1Base.metadata.tables["policy"] + assert [column.name for column in policy.primary_key.columns] == [ + "id", + "country_id", + "policy_hash", + ] + computed = V1Base.metadata.tables["computed_household"] + assert [column.name for column in computed.primary_key.columns] == [ + "household_id", + "policy_id", + "country_id", + ] + assert V1Base.metadata.tables["user_profiles"].c.auth0_id.unique From 840b1f49f7c1315881904926344a91fdf2d00fd8 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:09:59 +0300 Subject: [PATCH 06/89] feat(green): map the unchanged v1 schema with ORM --- policyengine_api/data/v1_models.py | 223 +++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) create mode 100644 policyengine_api/data/v1_models.py diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py new file mode 100644 index 000000000..201957264 --- /dev/null +++ b/policyengine_api/data/v1_models.py @@ -0,0 +1,223 @@ +"""Declarative mappings for the existing API v1 schema. + +These mappings describe the legacy tables; importing this module never emits +DDL. Alembic is the only schema-management entrypoint. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from sqlalchemy import BigInteger, DateTime, Integer, JSON, String, Text, UniqueConstraint +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + +class V1Base(DeclarativeBase): + pass + + +class Household(V1Base): + __tablename__ = "household" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + label: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(255)) + household_json: Mapped[Any] = mapped_column(JSON) + household_hash: Mapped[str] = mapped_column(String(255)) + + +class ComputedHousehold(V1Base): + __tablename__ = "computed_household" + household_id: Mapped[int] = mapped_column(Integer, primary_key=True) + policy_id: Mapped[int] = mapped_column(Integer, primary_key=True) + country_id: Mapped[str] = mapped_column(String(3), primary_key=True) + api_version: Mapped[str] = mapped_column(String(10)) + computed_household_json: Mapped[Any] = mapped_column(JSON) + status: Mapped[str | None] = mapped_column(String(32)) + + +class Policy(V1Base): + __tablename__ = "policy" + # SQLite cannot compile AUTO_INCREMENT on a composite primary key. The + # generated MySQL baseline receives the documented dialect correction. + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=False) + country_id: Mapped[str] = mapped_column(String(3), primary_key=True) + label: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + policy_json: Mapped[Any] = mapped_column(JSON) + policy_hash: Mapped[str] = mapped_column(String(255), primary_key=True) + + +class Economy(V1Base): + __tablename__ = "economy" + economy_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + region: Mapped[str | None] = mapped_column(String(32)) + time_period: Mapped[str | None] = mapped_column(String(32)) + options_json: Mapped[Any] = mapped_column(JSON) + options_hash: Mapped[str] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + economy_json: Mapped[Any | None] = mapped_column(JSON) + status: Mapped[str] = mapped_column(String(32)) + message: Mapped[str | None] = mapped_column(String(255)) + + +class ReformImpact(V1Base): + __tablename__ = "reform_impact" + reform_impact_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + baseline_policy_id: Mapped[int] + reform_policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + region: Mapped[str] = mapped_column(String(32)) + dataset: Mapped[str] = mapped_column(String(255)) + time_period: Mapped[str] = mapped_column(String(32)) + options_json: Mapped[Any | None] = mapped_column(JSON) + options_hash: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + reform_impact_json: Mapped[Any] = mapped_column(JSON) + status: Mapped[str] = mapped_column(String(32)) + message: Mapped[str | None] = mapped_column(String(255)) + start_time: Mapped[datetime | None] = mapped_column(DateTime) + end_time: Mapped[datetime | None] = mapped_column(DateTime) + execution_id: Mapped[str] = mapped_column(String(255)) + + +class Analysis(V1Base): + __tablename__ = "analysis" + prompt_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + prompt: Mapped[str] = mapped_column(Text) + analysis: Mapped[str | None] = mapped_column(Text) + status: Mapped[str] = mapped_column(String(32)) + + +class UserPolicy(V1Base): + __tablename__ = "user_policies" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + reform_id: Mapped[int] + reform_label: Mapped[str | None] = mapped_column(String(255)) + baseline_id: Mapped[int] + baseline_label: Mapped[str | None] = mapped_column(String(255)) + user_id: Mapped[str] = mapped_column(String(255)) + year: Mapped[str] = mapped_column(String(32)) + geography: Mapped[str] = mapped_column(String(255)) + dataset: Mapped[str | None] = mapped_column(String(255)) + number_of_provisions: Mapped[int] + api_version: Mapped[str] = mapped_column(String(32)) + added_date: Mapped[int] = mapped_column(BigInteger) + updated_date: Mapped[int] = mapped_column(BigInteger) + budgetary_impact: Mapped[str | None] = mapped_column(String(255)) + type: Mapped[str | None] = mapped_column(String(255)) + + +class UserProfile(V1Base): + __tablename__ = "user_profiles" + user_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + auth0_id: Mapped[str] = mapped_column(String(255), unique=True) + username: Mapped[str | None] = mapped_column(String(255), unique=True) + primary_country: Mapped[str] = mapped_column(String(3)) + user_since: Mapped[int] = mapped_column(BigInteger) + + +class Tracer(V1Base): + __tablename__ = "tracers" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + household_id: Mapped[int] + policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + api_version: Mapped[str] = mapped_column(String(10)) + tracer_output: Mapped[Any] = mapped_column(JSON) + + +class Simulation(V1Base): + __tablename__ = "simulations" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + api_version: Mapped[str] = mapped_column(String(10)) + population_id: Mapped[str] = mapped_column(String(255)) + population_type: Mapped[str] = mapped_column(String(50)) + policy_id: Mapped[int] + status: Mapped[str] = mapped_column(String(32), default="pending") + output: Mapped[Any | None] = mapped_column(JSON) + error_message: Mapped[str | None] = mapped_column(Text) + simulation_spec_json: Mapped[Any | None] = mapped_column(JSON) + simulation_spec_schema_version: Mapped[int | None] + active_run_id: Mapped[str | None] = mapped_column(String(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(String(36)) + + +class ReportOutput(V1Base): + __tablename__ = "report_outputs" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + simulation_1_id: Mapped[int] + simulation_2_id: Mapped[int | None] + api_version: Mapped[str] = mapped_column(String(10)) + status: Mapped[str] = mapped_column(String(32), default="pending") + output: Mapped[Any | None] = mapped_column(JSON) + error_message: Mapped[str | None] = mapped_column(Text) + year: Mapped[str | None] = mapped_column(String(255), default="2025") + report_kind: Mapped[str | None] = mapped_column(String(64)) + report_spec_json: Mapped[Any | None] = mapped_column(JSON) + report_spec_schema_version: Mapped[int | None] + report_spec_status: Mapped[str | None] = mapped_column(String(32)) + active_run_id: Mapped[str | None] = mapped_column(String(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(String(36)) + + +class ReportOutputRun(V1Base): + __tablename__ = "report_output_runs" + __table_args__ = (UniqueConstraint("report_output_id", "run_sequence"),) + id: Mapped[str] = mapped_column(String(36), primary_key=True) + report_output_id: Mapped[int] + run_sequence: Mapped[int] + status: Mapped[str] = mapped_column(String(32)) + output: Mapped[Any | None] = mapped_column(JSON) + error_message: Mapped[str | None] = mapped_column(Text) + trigger_type: Mapped[str] = mapped_column(String(32)) + requested_at: Mapped[datetime | None] = mapped_column(DateTime) + started_at: Mapped[datetime | None] = mapped_column(DateTime) + finished_at: Mapped[datetime | None] = mapped_column(DateTime) + source_run_id: Mapped[str | None] = mapped_column(String(36)) + report_spec_snapshot_json: Mapped[Any | None] = mapped_column(JSON) + country_package_version: Mapped[str | None] = mapped_column(String(255)) + policyengine_version: Mapped[str | None] = mapped_column(String(255)) + data_version: Mapped[str | None] = mapped_column(String(255)) + runtime_app_name: Mapped[str | None] = mapped_column(String(255)) + report_cache_version: Mapped[str | None] = mapped_column(String(255)) + simulation_cache_version: Mapped[str | None] = mapped_column(String(255)) + requested_version_override: Mapped[str | None] = mapped_column(String(255)) + resolved_dataset: Mapped[str | None] = mapped_column(String(255)) + resolved_options_hash: Mapped[str | None] = mapped_column(String(255)) + + +class SimulationRun(V1Base): + __tablename__ = "simulation_runs" + __table_args__ = (UniqueConstraint("simulation_id", "run_sequence"),) + id: Mapped[str] = mapped_column(String(36), primary_key=True) + simulation_id: Mapped[int] + report_output_run_id: Mapped[str | None] = mapped_column(String(36)) + input_position: Mapped[int | None] + run_sequence: Mapped[int] + status: Mapped[str] = mapped_column(String(32)) + output: Mapped[Any | None] = mapped_column(JSON) + error_message: Mapped[str | None] = mapped_column(Text) + trigger_type: Mapped[str] = mapped_column(String(32)) + requested_at: Mapped[datetime | None] = mapped_column(DateTime) + started_at: Mapped[datetime | None] = mapped_column(DateTime) + finished_at: Mapped[datetime | None] = mapped_column(DateTime) + source_run_id: Mapped[str | None] = mapped_column(String(36)) + simulation_spec_snapshot_json: Mapped[Any | None] = mapped_column(JSON) + country_package_version: Mapped[str | None] = mapped_column(String(255)) + policyengine_version: Mapped[str | None] = mapped_column(String(255)) + data_version: Mapped[str | None] = mapped_column(String(255)) + runtime_app_name: Mapped[str | None] = mapped_column(String(255)) + simulation_cache_version: Mapped[str | None] = mapped_column(String(255)) + + +class LegacyReportOutputAlias(V1Base): + __tablename__ = "legacy_report_output_aliases" + legacy_report_output_id: Mapped[int] = mapped_column(Integer, primary_key=True) + canonical_report_output_id: Mapped[int] From c0a67f3e5d6ab47bc145f2b64efb5f967c90be30 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:10:18 +0300 Subject: [PATCH 07/89] test(red): define Alembic baseline safeguards --- tests/unit/data/test_alembic_baseline.py | 35 ++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 tests/unit/data/test_alembic_baseline.py diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py new file mode 100644 index 000000000..15cae49c7 --- /dev/null +++ b/tests/unit/data/test_alembic_baseline.py @@ -0,0 +1,35 @@ +from pathlib import Path + +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, inspect + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base + + +def _config(url: str) -> Config: + config = Config(str(REPO / "alembic.ini")) + config.set_main_option("sqlalchemy.url", url) + return config + + +def test_baseline_upgrades_fresh_database_to_v1_head(tmp_path: Path): + database_path = tmp_path / "fresh.db" + command.upgrade(_config(f"sqlite+pysqlite:///{database_path}"), "head") + + tables = set(inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names()) + assert set(V1Base.metadata.tables) <= tables + assert "alembic_version" in tables + + +def test_baseline_downgrades_and_reupgrades(tmp_path: Path): + database_path = tmp_path / "roundtrip.db" + config = _config(f"sqlite+pysqlite:///{database_path}") + command.upgrade(config, "head") + command.downgrade(config, "base") + command.upgrade(config, "head") + + assert set(V1Base.metadata.tables) <= set( + inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names() + ) From af51cd9d650dae4455014750804fb67d00391cf6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:20:28 +0300 Subject: [PATCH 08/89] feat(green): add autogenerated v1 Alembic baseline --- alembic.ini | 38 +++ migrations/env.py | 53 ++++ migrations/script.py.mako | 25 ++ migrations/versions/.gitkeep | 0 ...8cb8bcd717c_baseline_existing_v1_schema.py | 235 ++++++++++++++++++ policyengine_api/data/__init__.py | 19 +- policyengine_api/data/v1_models.py | 8 +- pyproject.toml | 1 + uv.lock | 32 ++- 9 files changed, 404 insertions(+), 7 deletions(-) create mode 100644 alembic.ini create mode 100644 migrations/env.py create mode 100644 migrations/script.py.mako create mode 100644 migrations/versions/.gitkeep create mode 100644 migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 000000000..6d23a99ae --- /dev/null +++ b/alembic.ini @@ -0,0 +1,38 @@ +[alembic] +script_location = migrations +prepend_sys_path = . +path_separator = os +sqlalchemy.url = sqlite+pysqlite:///policyengine_api/data/policyengine.db + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/migrations/env.py b/migrations/env.py new file mode 100644 index 000000000..85aab3a36 --- /dev/null +++ b/migrations/env.py @@ -0,0 +1,53 @@ +"""Alembic environment for the API v1 schema.""" + +from logging.config import fileConfig +import os + +from alembic import context +from sqlalchemy import engine_from_config, pool + +from policyengine_api.data.v1_models import V1Base + + +config = context.config +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +database_url = os.environ.get("ALEMBIC_DATABASE_URL") +if database_url: + config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) + +target_metadata = V1Base.metadata + + +def run_migrations_offline() -> None: + context.configure( + url=config.get_main_option("sqlalchemy.url"), + target_metadata=target_metadata, + literal_binds=True, + compare_type=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + with connectable.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/migrations/script.py.mako b/migrations/script.py.mako new file mode 100644 index 000000000..1ba49a84b --- /dev/null +++ b/migrations/script.py.mako @@ -0,0 +1,25 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/migrations/versions/.gitkeep b/migrations/versions/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py b/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py new file mode 100644 index 000000000..dfa019c5a --- /dev/null +++ b/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py @@ -0,0 +1,235 @@ +"""baseline existing v1 schema + +Revision ID: f8cb8bcd717c +Revises: +Create Date: 2026-08-06 17:20:10.119942 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = 'f8cb8bcd717c' +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('analysis', + sa.Column('prompt_id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('prompt', sa.Text(), nullable=False), + sa.Column('analysis', sa.Text(), nullable=True), + sa.Column('status', sa.String(length=32), nullable=False), + sa.PrimaryKeyConstraint('prompt_id') + ) + op.create_table('computed_household', + sa.Column('household_id', sa.Integer(), nullable=False), + sa.Column('policy_id', sa.Integer(), nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('computed_household_json', sa.JSON(), nullable=False), + sa.Column('status', sa.String(length=32), nullable=True), + sa.PrimaryKeyConstraint('household_id', 'policy_id', 'country_id') + ) + op.create_table('economy', + sa.Column('economy_id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('policy_id', sa.Integer(), nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('region', sa.String(length=32), nullable=True), + sa.Column('time_period', sa.String(length=32), nullable=True), + sa.Column('options_json', sa.JSON(), nullable=False), + sa.Column('options_hash', sa.String(length=255), nullable=False), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('economy_json', sa.JSON(), nullable=True), + sa.Column('status', sa.String(length=32), nullable=False), + sa.Column('message', sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint('economy_id') + ) + op.create_table('household', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('label', sa.String(length=255), nullable=True), + sa.Column('api_version', sa.String(length=255), nullable=False), + sa.Column('household_json', sa.JSON(), nullable=False), + sa.Column('household_hash', sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('legacy_report_output_aliases', + sa.Column('legacy_report_output_id', sa.Integer(), nullable=False), + sa.Column('canonical_report_output_id', sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint('legacy_report_output_id') + ) + op.create_table('policy', + # Review correction: MySQL's legacy composite key auto-increments ``id``; + # SQLite cannot compile autoincrement on a composite primary key. + sa.Column('id', sa.Integer(), autoincrement=op.get_bind().dialect.name != 'sqlite', nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('label', sa.String(length=255), nullable=True), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('policy_json', sa.JSON(), nullable=False), + sa.Column('policy_hash', sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint('id', 'country_id', 'policy_hash') + ) + op.create_table('reform_impact', + sa.Column('reform_impact_id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('baseline_policy_id', sa.Integer(), nullable=False), + sa.Column('reform_policy_id', sa.Integer(), nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('region', sa.String(length=32), nullable=False), + sa.Column('dataset', sa.String(length=255), nullable=False), + sa.Column('time_period', sa.String(length=32), nullable=False), + sa.Column('options_json', sa.JSON(), nullable=True), + sa.Column('options_hash', sa.String(length=255), nullable=True), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('reform_impact_json', sa.JSON(), nullable=False), + sa.Column('status', sa.String(length=32), nullable=False), + sa.Column('message', sa.String(length=255), nullable=True), + sa.Column('start_time', sa.DateTime(), nullable=True), + sa.Column('end_time', sa.DateTime(), nullable=True), + sa.Column('execution_id', sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint('reform_impact_id') + ) + op.create_table('report_output_runs', + sa.Column('id', sa.String(length=36), nullable=False), + sa.Column('report_output_id', sa.Integer(), nullable=False), + sa.Column('run_sequence', sa.Integer(), nullable=False), + sa.Column('status', sa.String(length=32), nullable=False), + sa.Column('output', sa.JSON(), nullable=True), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('trigger_type', sa.String(length=32), nullable=False), + sa.Column('requested_at', sa.DateTime(), nullable=True), + sa.Column('started_at', sa.DateTime(), nullable=True), + sa.Column('finished_at', sa.DateTime(), nullable=True), + sa.Column('source_run_id', sa.String(length=36), nullable=True), + sa.Column('report_spec_snapshot_json', sa.JSON(), nullable=True), + sa.Column('country_package_version', sa.String(length=255), nullable=True), + sa.Column('policyengine_version', sa.String(length=255), nullable=True), + sa.Column('data_version', sa.String(length=255), nullable=True), + sa.Column('runtime_app_name', sa.String(length=255), nullable=True), + sa.Column('report_cache_version', sa.String(length=255), nullable=True), + sa.Column('simulation_cache_version', sa.String(length=255), nullable=True), + sa.Column('requested_version_override', sa.String(length=255), nullable=True), + sa.Column('resolved_dataset', sa.String(length=255), nullable=True), + sa.Column('resolved_options_hash', sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('report_output_id', 'run_sequence') + ) + op.create_table('report_outputs', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('simulation_1_id', sa.Integer(), nullable=False), + sa.Column('simulation_2_id', sa.Integer(), nullable=True), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('status', sa.String(length=32), server_default=sa.text("'pending'"), nullable=False), + sa.Column('output', sa.JSON(), nullable=True), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('year', sa.String(length=255), server_default=sa.text("'2025'"), nullable=True), + sa.Column('report_kind', sa.String(length=64), nullable=True), + sa.Column('report_spec_json', sa.JSON(), nullable=True), + sa.Column('report_spec_schema_version', sa.Integer(), nullable=True), + sa.Column('report_spec_status', sa.String(length=32), nullable=True), + sa.Column('active_run_id', sa.String(length=36), nullable=True), + sa.Column('latest_successful_run_id', sa.String(length=36), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('simulation_runs', + sa.Column('id', sa.String(length=36), nullable=False), + sa.Column('simulation_id', sa.Integer(), nullable=False), + sa.Column('report_output_run_id', sa.String(length=36), nullable=True), + sa.Column('input_position', sa.Integer(), nullable=True), + sa.Column('run_sequence', sa.Integer(), nullable=False), + sa.Column('status', sa.String(length=32), nullable=False), + sa.Column('output', sa.JSON(), nullable=True), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('trigger_type', sa.String(length=32), nullable=False), + sa.Column('requested_at', sa.DateTime(), nullable=True), + sa.Column('started_at', sa.DateTime(), nullable=True), + sa.Column('finished_at', sa.DateTime(), nullable=True), + sa.Column('source_run_id', sa.String(length=36), nullable=True), + sa.Column('simulation_spec_snapshot_json', sa.JSON(), nullable=True), + sa.Column('country_package_version', sa.String(length=255), nullable=True), + sa.Column('policyengine_version', sa.String(length=255), nullable=True), + sa.Column('data_version', sa.String(length=255), nullable=True), + sa.Column('runtime_app_name', sa.String(length=255), nullable=True), + sa.Column('simulation_cache_version', sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('simulation_id', 'run_sequence') + ) + op.create_table('simulations', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('population_id', sa.String(length=255), nullable=False), + sa.Column('population_type', sa.String(length=50), nullable=False), + sa.Column('policy_id', sa.Integer(), nullable=False), + sa.Column('status', sa.String(length=32), server_default=sa.text("'pending'"), nullable=False), + sa.Column('output', sa.JSON(), nullable=True), + sa.Column('error_message', sa.Text(), nullable=True), + sa.Column('simulation_spec_json', sa.JSON(), nullable=True), + sa.Column('simulation_spec_schema_version', sa.Integer(), nullable=True), + sa.Column('active_run_id', sa.String(length=36), nullable=True), + sa.Column('latest_successful_run_id', sa.String(length=36), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('tracers', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('household_id', sa.Integer(), nullable=False), + sa.Column('policy_id', sa.Integer(), nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('api_version', sa.String(length=10), nullable=False), + sa.Column('tracer_output', sa.JSON(), nullable=False), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('user_policies', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('country_id', sa.String(length=3), nullable=False), + sa.Column('reform_id', sa.Integer(), nullable=False), + sa.Column('reform_label', sa.String(length=255), nullable=True), + sa.Column('baseline_id', sa.Integer(), nullable=False), + sa.Column('baseline_label', sa.String(length=255), nullable=True), + sa.Column('user_id', sa.String(length=255), nullable=False), + sa.Column('year', sa.String(length=32), nullable=False), + sa.Column('geography', sa.String(length=255), nullable=False), + sa.Column('dataset', sa.String(length=255), nullable=True), + sa.Column('number_of_provisions', sa.Integer(), nullable=False), + sa.Column('api_version', sa.String(length=32), nullable=False), + sa.Column('added_date', sa.BigInteger(), nullable=False), + sa.Column('updated_date', sa.BigInteger(), nullable=False), + sa.Column('budgetary_impact', sa.String(length=255), nullable=True), + sa.Column('type', sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('user_profiles', + sa.Column('user_id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('auth0_id', sa.String(length=255), nullable=False), + sa.Column('username', sa.String(length=255), nullable=True), + sa.Column('primary_country', sa.String(length=3), nullable=False), + sa.Column('user_since', sa.BigInteger(), nullable=False), + sa.PrimaryKeyConstraint('user_id'), + sa.UniqueConstraint('auth0_id'), + sa.UniqueConstraint('username') + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('user_profiles') + op.drop_table('user_policies') + op.drop_table('tracers') + op.drop_table('simulations') + op.drop_table('simulation_runs') + op.drop_table('report_outputs') + op.drop_table('report_output_runs') + op.drop_table('reform_impact') + op.drop_table('policy') + op.drop_table('legacy_report_output_aliases') + op.drop_table('household') + op.drop_table('economy') + op.drop_table('computed_household') + op.drop_table('analysis') + # ### end Alembic commands ### diff --git a/policyengine_api/data/__init__.py b/policyengine_api/data/__init__.py index 15673afdb..eb42b9919 100644 --- a/policyengine_api/data/__init__.py +++ b/policyengine_api/data/__init__.py @@ -1 +1,18 @@ -from .data import PolicyEngineDatabase, database, local_database +"""Database package with lazy legacy exports. + +Keeping package import side-effect free lets Alembic load model metadata without +opening Cloud SQL or creating a local database. +""" + +from typing import Any + + +__all__ = ["PolicyEngineDatabase", "database", "local_database"] + + +def __getattr__(name: str) -> Any: + if name in __all__: + from . import data + + return getattr(data, name) + raise AttributeError(name) diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 201957264..6648612a6 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -9,7 +9,7 @@ from datetime import datetime from typing import Any -from sqlalchemy import BigInteger, DateTime, Integer, JSON, String, Text, UniqueConstraint +from sqlalchemy import BigInteger, DateTime, Integer, JSON, String, Text, UniqueConstraint, text from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column @@ -139,7 +139,7 @@ class Simulation(V1Base): population_id: Mapped[str] = mapped_column(String(255)) population_type: Mapped[str] = mapped_column(String(50)) policy_id: Mapped[int] - status: Mapped[str] = mapped_column(String(32), default="pending") + status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) output: Mapped[Any | None] = mapped_column(JSON) error_message: Mapped[str | None] = mapped_column(Text) simulation_spec_json: Mapped[Any | None] = mapped_column(JSON) @@ -155,10 +155,10 @@ class ReportOutput(V1Base): simulation_1_id: Mapped[int] simulation_2_id: Mapped[int | None] api_version: Mapped[str] = mapped_column(String(10)) - status: Mapped[str] = mapped_column(String(32), default="pending") + status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) output: Mapped[Any | None] = mapped_column(JSON) error_message: Mapped[str | None] = mapped_column(Text) - year: Mapped[str | None] = mapped_column(String(255), default="2025") + year: Mapped[str | None] = mapped_column(String(255), server_default=text("'2025'")) report_kind: Mapped[str | None] = mapped_column(String(64)) report_spec_json: Mapped[Any | None] = mapped_column(JSON) report_spec_schema_version: Mapped[int | None] diff --git a/pyproject.toml b/pyproject.toml index 8cac617b7..efd3959d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ classifiers = [ ] dependencies = [ "a2wsgi>=1.10,<2", + "alembic>=1.14,<2", "anthropic", "assertpy", "click>=8,<9", diff --git a/uv.lock b/uv.lock index c345d2531..0a548c102 100644 --- a/uv.lock +++ b/uv.lock @@ -155,6 +155,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "alembic" +version = "1.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mako" }, + { name = "sqlalchemy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fb/01/a48dab7827ac4421272399f7ed9a2ec17edd12c8bcde4417bd7b6821b71a/alembic-1.19.0.tar.gz", hash = "sha256:6487c612fc719dcfa22b17d2dd5b2b458929641e6aa2f0b65b135727f5e6d501", size = 2069906, upload-time = "2026-08-04T18:57:04.599Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/79/59ab6f1fe72eee229de91aa393225389a85df12a0fbbbfa264efcf6d7872/alembic-1.19.0-py3-none-any.whl", hash = "sha256:cf839d3849116aab3cc047e09c6968b9bd6b2fde61b6bb7c1e97352fe5503580", size = 265738, upload-time = "2026-08-04T18:57:06.219Z" }, +] + [[package]] name = "altair" version = "6.1.0" @@ -1793,6 +1807,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/91/6c074015990f4f656f7b69a5c2d15924906ce0bc19c7014ac953493c0cf0/linecheck-0.1.0-py3-none-any.whl", hash = "sha256:73c6b29790521fa711b00df7cd60af4caf7004337d8710606881fbecb0d1bc83", size = 2767, upload-time = "2022-07-16T13:06:17.01Z" }, ] +[[package]] +name = "mako" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/12/b5fa2353e2754cd67fb9f83793fa48ff42c213a5da7e719869d2301f6ab8/mako-1.4.1.tar.gz", hash = "sha256:d7904710b662996425a21627710c4777c45053146942cf8a7aebf757c92b8c27", size = 410165, upload-time = "2026-08-05T06:10:56.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/54/12ed58d458474aaab5c3d180173e745a4fe131bb330370596876d19ff60f/mako-1.4.1-py3-none-any.whl", hash = "sha256:a359d9a94a541213958742b2698d0a7757bb83551767bc468a74b9905aba9617", size = 80010, upload-time = "2026-08-05T06:10:58.248Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -2461,7 +2487,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "ptyprocess" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -2616,10 +2642,11 @@ models = [ [[package]] name = "policyengine-api" -version = "3.46.2" +version = "3.48.0" source = { editable = "." } dependencies = [ { name = "a2wsgi" }, + { name = "alembic" }, { name = "anthropic" }, { name = "assertpy" }, { name = "click" }, @@ -2665,6 +2692,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "a2wsgi", specifier = ">=1.10,<2" }, + { name = "alembic", specifier = ">=1.14,<2" }, { name = "anthropic" }, { name = "assertpy" }, { name = "build", marker = "extra == 'dev'" }, From 1b59a361c66ffafea5220a9398c7a4c41d7981dd Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:20:59 +0300 Subject: [PATCH 09/89] test(red): define policy household and user DAO behavior --- tests/unit/data/test_v1_daos.py | 45 +++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 tests/unit/data/test_v1_daos.py diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py new file mode 100644 index 000000000..13cea534c --- /dev/null +++ b/tests/unit/data/test_v1_daos.py @@ -0,0 +1,45 @@ +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import HouseholdDAO, PolicyDAO, UserDAO +from policyengine_api.data.v1_models import V1Base + + +def _daos(): + manager = build_sqlite_session_manager() + V1Base.metadata.create_all(manager.engine) + return PolicyDAO(manager), HouseholdDAO(manager), UserDAO(manager) + + +def test_policy_dao_round_trips_legacy_mapping_shape(): + policies, _, _ = _daos() + policy_id = policies.create("us", "Reform", {"gov.irs": 1}, "hash", "1.0") + assert policy_id == 1 + assert policies.get("us", policy_id) == { + "id": 1, + "country_id": "us", + "label": "Reform", + "api_version": "1.0", + "policy_json": {"gov.irs": 1}, + "policy_hash": "hash", + } + + +def test_policy_dao_allocates_ids_and_detects_existing_policy(): + policies, _, _ = _daos() + assert policies.create("us", None, {}, "one", "1.0") == 1 + assert policies.create("uk", None, {}, "two", "1.0") == 2 + assert policies.find_unique("us", "one", None)["id"] == 1 + + +def test_household_dao_creates_updates_and_reads(): + _, households, _ = _daos() + household_id = households.create("us", "Home", {"people": {}}, "h", "1.0") + households.update("us", household_id, "Updated", {"people": {"you": {}}}) + assert households.get("us", household_id)["label"] == "Updated" + assert households.get("uk", household_id) is None + + +def test_user_dao_profile_lookup_precedence(): + _, _, users = _daos() + user_id = users.create_profile("auth0|one", "person", "us", 123) + assert users.get_profile(auth0_id="auth0|one")["user_id"] == user_id + assert users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] == "auth0|one" From 638f557f6101d6f2b828fca11f5bc97e4fd94924 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 17:21:37 +0300 Subject: [PATCH 10/89] feat(green): add policy household and user DAOs --- policyengine_api/data/v1_daos.py | 170 +++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 policyengine_api/data/v1_daos.py diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py new file mode 100644 index 000000000..aa55002fd --- /dev/null +++ b/policyengine_api/data/v1_daos.py @@ -0,0 +1,170 @@ +"""ORM data access objects for the existing v1 schema.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, select + +from policyengine_api.data.orm import SessionManager +from policyengine_api.data.v1_models import Household, Policy, UserProfile + + +def _mapping(model: Any) -> dict[str, Any]: + return { + column.name: getattr(model, column.name) + for column in model.__table__.columns + } + + +class PolicyDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def get(self, country_id: str, policy_id: int) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, + ) + ) + return _mapping(model) if model else None + + def find_unique( + self, country_id: str, policy_hash: str, label: str | None + ) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.policy_hash == policy_hash, + Policy.label == label, + ) + ) + return _mapping(model) if model else None + + def create( + self, + country_id: str, + label: str | None, + policy_json: Any, + policy_hash: str, + api_version: str, + ) -> int: + def operation(session): + next_id = (session.scalar(select(func.max(Policy.id))) or 0) + 1 + session.add( + Policy( + id=next_id, + country_id=country_id, + label=label, + api_version=api_version, + policy_json=policy_json, + policy_hash=policy_hash, + ) + ) + return next_id + + return self.sessions.run_in_transaction(operation) + + +class HouseholdDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def get(self, country_id: str, household_id: int) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) + return _mapping(model) if model else None + + def create( + self, + country_id: str, + label: str | None, + household_json: Any, + household_hash: str, + api_version: str, + ) -> int: + def operation(session): + model = Household( + country_id=country_id, + label=label, + api_version=api_version, + household_json=household_json, + household_hash=household_hash, + ) + session.add(model) + session.flush() + return model.id + + return self.sessions.run_in_transaction(operation) + + def update( + self, + country_id: str, + household_id: int, + label: str | None, + household_json: Any, + ) -> bool: + def operation(session): + model = session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) + if model is None: + return False + model.label = label + model.household_json = household_json + return True + + return self.sessions.run_in_transaction(operation) + + +class UserDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def create_profile( + self, + auth0_id: str, + username: str | None, + primary_country: str, + user_since: int, + ) -> int: + def operation(session): + model = UserProfile( + auth0_id=auth0_id, + username=username, + primary_country=primary_country, + user_since=user_since, + ) + session.add(model) + session.flush() + return model.user_id + + return self.sessions.run_in_transaction(operation) + + def get_profile( + self, + *, + user_id: int | None = None, + auth0_id: str | None = None, + ) -> dict[str, Any] | None: + if user_id is None and auth0_id is None: + return None + with self.sessions.session() as session: + condition = ( + UserProfile.user_id == user_id + if user_id is not None + else UserProfile.auth0_id == auth0_id + ) + model = session.scalar(select(UserProfile).where(condition)) + return _mapping(model) if model else None From 47f02e2a6f3314174732e01eeac06ce8bde966ed Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:05:11 +0300 Subject: [PATCH 11/89] fix: satisfy Stage 7 changed-file lint --- tests/unit/test_alembic_skill.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/test_alembic_skill.py b/tests/unit/test_alembic_skill.py index 963eb82ba..df9609d1a 100644 --- a/tests/unit/test_alembic_skill.py +++ b/tests/unit/test_alembic_skill.py @@ -1,5 +1,3 @@ -from pathlib import Path - from policyengine_api.constants import REPO From 35c207e0279f5503c5fee603bab537e8d540da6b Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:24:50 +0300 Subject: [PATCH 12/89] test(red): define economy analysis and tracer DAO behavior --- tests/unit/data/test_local_daos.py | 38 ++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/unit/data/test_local_daos.py diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py new file mode 100644 index 000000000..e49553cf9 --- /dev/null +++ b/tests/unit/data/test_local_daos.py @@ -0,0 +1,38 @@ +from datetime import datetime + +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import AnalysisDAO, ReformImpactDAO, TracerDAO +from policyengine_api.data.v1_models import V1Base + + +def _daos(): + manager = build_sqlite_session_manager() + V1Base.metadata.create_all(manager.engine) + return AnalysisDAO(manager), ReformImpactDAO(manager), TracerDAO(manager) + + +def test_analysis_dao_round_trip(): + analyses, _, _ = _daos() + analyses.store("prompt", "answer", "complete") + assert analyses.get("prompt") == "answer" + + +def test_reform_impact_dao_transitions_by_execution_id(): + _, impacts, _ = _daos() + impacts.create( + country_id="us", reform_policy_id=2, baseline_policy_id=1, + region="us", dataset="default", time_period="2026", + options_json={}, options_hash="hash", api_version="1", + reform_impact_json={}, status="computing", start_time=datetime(2026, 1, 1), + execution_id="job", + ) + impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) + assert impacts.find(execution_id="job")["status"] == "ok" + assert impacts.find(execution_id="job")["reform_impact_json"] == {"result": 1} + + +def test_tracer_dao_returns_latest_matching_trace(): + _, _, tracers = _daos() + tracers.create(1, 2, "us", "1", {"trace": "first"}) + tracers.create(1, 2, "us", "1", {"trace": "latest"}) + assert tracers.get(1, 2, "us")["tracer_output"] == {"trace": "latest"} From 9073cdde1475422ba7c67a676e6eb60ffde7225f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:26:01 +0300 Subject: [PATCH 13/89] feat(green): add economy analysis and tracer DAOs --- policyengine_api/data/v1_daos.py | 118 ++++++++++++++++++++++++++++++- 1 file changed, 117 insertions(+), 1 deletion(-) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index aa55002fd..edc1188bc 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -4,10 +4,19 @@ from typing import Any +from datetime import datetime + from sqlalchemy import func, select from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_models import Household, Policy, UserProfile +from policyengine_api.data.v1_models import ( + Analysis, + Household, + Policy, + ReformImpact, + Tracer, + UserProfile, +) def _mapping(model: Any) -> dict[str, Any]: @@ -168,3 +177,110 @@ def get_profile( ) model = session.scalar(select(UserProfile).where(condition)) return _mapping(model) if model else None + + +class AnalysisDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def get(self, prompt: str) -> str | None: + with self.sessions.session() as session: + model = session.scalar( + select(Analysis) + .where(Analysis.prompt == prompt, Analysis.status == "complete") + .order_by(Analysis.prompt_id.desc()) + ) + return model.analysis if model else None + + def store(self, prompt: str, analysis: str | None, status: str) -> int: + def operation(session): + model = Analysis(prompt=prompt, analysis=analysis, status=status) + session.add(model) + session.flush() + return model.prompt_id + + return self.sessions.run_in_transaction(operation) + + +class ReformImpactDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def create(self, **values: Any) -> int: + def operation(session): + model = ReformImpact(**values) + session.add(model) + session.flush() + return model.reform_impact_id + + return self.sessions.run_in_transaction(operation) + + def find(self, *, execution_id: str) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + return _mapping(model) if model else None + + def complete( + self, execution_id: str, result: Any, finished_at: datetime + ) -> bool: + def operation(session): + model = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if model is None: + return False + model.status = "ok" + model.message = "Completed" + model.reform_impact_json = result + model.end_time = finished_at + return True + + return self.sessions.run_in_transaction(operation) + + +class TracerDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def create( + self, + household_id: int, + policy_id: int, + country_id: str, + api_version: str, + tracer_output: Any, + ) -> int: + def operation(session): + model = Tracer( + household_id=household_id, + policy_id=policy_id, + country_id=country_id, + api_version=api_version, + tracer_output=tracer_output, + ) + session.add(model) + session.flush() + return model.id + + return self.sessions.run_in_transaction(operation) + + def get( + self, household_id: int, policy_id: int, country_id: str + ) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(Tracer) + .where( + Tracer.household_id == household_id, + Tracer.policy_id == policy_id, + Tracer.country_id == country_id, + ) + .order_by(Tracer.id.desc()) + ) + return _mapping(model) if model else None From 72953e915e441e193c12024fcde91d8153551dd1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:27:19 +0300 Subject: [PATCH 14/89] test(red): require DAO boundaries in core services --- .../services/test_stage7_dao_boundaries.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 tests/unit/services/test_stage7_dao_boundaries.py diff --git a/tests/unit/services/test_stage7_dao_boundaries.py b/tests/unit/services/test_stage7_dao_boundaries.py new file mode 100644 index 000000000..93cec54b7 --- /dev/null +++ b/tests/unit/services/test_stage7_dao_boundaries.py @@ -0,0 +1,50 @@ +from pathlib import Path + +import pytest + +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_service import UserService + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + ["household_service.py", "policy_service.py", "user_service.py"], +) +def test_migrated_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "from policyengine_api.data import database" not in source + + +class StubHouseholds: + def get(self, country_id, household_id): + return {"country_id": country_id, "id": household_id} + + +class StubPolicies: + def get(self, country_id, policy_id): + return { + "country_id": country_id, + "id": policy_id, + "policy_json": {"already": "decoded"}, + } + + +class StubUsers: + def get_profile(self, *, user_id=None, auth0_id=None): + return {"user_id": user_id, "auth0_id": auth0_id} + + +def test_services_accept_explicit_daos_for_isolated_parity_tests(): + assert HouseholdService(StubHouseholds()).get_household("us", 3)["id"] == 3 + assert PolicyService(StubPolicies()).get_policy("us", 4)["policy_json"] == { + "already": "decoded" + } + assert UserService(StubUsers()).get_profile(auth0_id="auth0|one") == { + "user_id": None, + "auth0_id": "auth0|one", + } From 0807cfb19e53570917e17b4bbb8c7d4d12751dfb Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:30:29 +0300 Subject: [PATCH 15/89] feat(green): route core services through DAOs --- policyengine_api/data/orm.py | 37 +++ policyengine_api/data/v1_daos.py | 23 +- .../services/household_service.py | 154 +++--------- policyengine_api/services/policy_service.py | 237 ++++-------------- policyengine_api/services/user_service.py | 75 ++---- tests/unit/data/test_v1_daos.py | 13 +- tests/unit/services/test_policy_service.py | 159 +++--------- .../services/test_update_profile_service.py | 10 +- 8 files changed, 222 insertions(+), 486 deletions(-) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 81401a224..22035d6c7 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -15,6 +15,25 @@ T = TypeVar("T") +class _IndexedMappingRow(dict): + """SQLite row compatible with both SQLAlchemy and legacy mapping callers.""" + + def __init__(self, cursor, values): + self._values = values + super().__init__( + (description[0], values[index]) + for index, description in enumerate(cursor.description) + ) + + def __getitem__(self, key): + if isinstance(key, int): + return self._values[key] + return super().__getitem__(key) + + def __iter__(self): + return iter(self._values) + + class SessionManager: """Own sessions and transaction boundaries without leaking either to callers.""" @@ -59,3 +78,21 @@ def build_sqlite_session_manager( else: engine = create_engine(f"sqlite+pysqlite:///{Path(database_path)}") return SessionManager(engine) + + +def build_v1_session_manager() -> SessionManager: + """Bind ORM sessions to the database selected by the v1 runtime.""" + + from policyengine_api.data.data import database + + if database.local: + if hasattr(database, "_connection"): + database._connection.row_factory = _IndexedMappingRow + engine = create_engine( + "sqlite+pysqlite://", + creator=lambda: database._connection, + poolclass=StaticPool, + ) + return SessionManager(engine) + return build_sqlite_session_manager(database.db_url) + return SessionManager(database.pool) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index edc1188bc..6e82cee3d 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -21,8 +21,7 @@ def _mapping(model: Any) -> dict[str, Any]: return { - column.name: getattr(model, column.name) - for column in model.__table__.columns + column.name: getattr(model, column.name) for column in model.__table__.columns } @@ -120,6 +119,8 @@ def update( household_id: int, label: str | None, household_json: Any, + household_hash: str, + api_version: str, ) -> bool: def operation(session): model = session.scalar( @@ -132,6 +133,8 @@ def operation(session): return False model.label = label model.household_json = household_json + model.household_hash = household_hash + model.api_version = api_version return True return self.sessions.run_in_transaction(operation) @@ -178,6 +181,18 @@ def get_profile( model = session.scalar(select(UserProfile).where(condition)) return _mapping(model) if model else None + def update_profile(self, user_id: int, **values: Any) -> bool: + def operation(session): + model = session.get(UserProfile, user_id) + if model is None: + return False + for key, value in values.items(): + if value is not None: + setattr(model, key, value) + return True + + return self.sessions.run_in_transaction(operation) + class AnalysisDAO: def __init__(self, sessions: SessionManager): @@ -224,9 +239,7 @@ def find(self, *, execution_id: str) -> dict[str, Any] | None: ) return _mapping(model) if model else None - def complete( - self, execution_id: str, result: Any, finished_at: datetime - ) -> bool: + def complete(self, execution_id: str, result: Any, finished_at: datetime) -> bool: def operation(session): model = session.scalar( select(ReformImpact) diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 2d3601737..c94f27d36 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,46 +1,27 @@ -import json -from sqlalchemy.engine.row import Row +from __future__ import annotations -from policyengine_api.data import database -from policyengine_api.utils import hash_object from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import HouseholdDAO +from policyengine_api.utils import hash_object class HouseholdService: - def get_household(self, country_id: str, household_id: int) -> dict | None: - """ - Get a household's input data with a given ID. - - Args: - country_id (str): The country ID. - household_id (int): The household ID. - """ - print("Getting household data") - - try: - if type(household_id) is not int or household_id < 0: - raise Exception( - f"Invalid household ID: {household_id}. Must be a positive integer." - ) + def __init__(self, households: HouseholdDAO | None = None): + self._households = households - row: Row | None = database.query( - f"SELECT * FROM household WHERE id = ? AND country_id = ?", - (household_id, country_id), - ).fetchone() + @property + def households(self) -> HouseholdDAO: + if self._households is None: + self._households = HouseholdDAO(build_v1_session_manager()) + return self._households - # If row is present, we must JSON.loads the household_json - household = None - if row is not None: - household = dict(row) - if household["household_json"]: - household["household_json"] = json.loads( - household["household_json"] - ) - return household - - except Exception as e: - print(f"Error fetching household #{household_id}. Details: {str(e)}") - raise e + def get_household(self, country_id: str, household_id: int) -> dict | None: + if type(household_id) is not int or household_id < 0: + raise Exception( + f"Invalid household ID: {household_id}. Must be a positive integer." + ) + return self.households.get(country_id, household_id) def create_household( self, @@ -48,43 +29,13 @@ def create_household( household_json: dict, label: str | None, ) -> int: - """ - Create a new household with the given data. - - Args: - country_id (str): The country ID. - household_json (dict): The household data. - household_hash (int): The hash of the household data. - label (str): The label for the household. - api_version (str): The API version. - """ - - print("Creating new household") - - try: - household_hash: str = hash_object(household_json) - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - database.query( - f"INSERT INTO household (country_id, household_json, household_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - json.dumps(household_json), - household_hash, - label, - api_version, - ), - ) - - household_id = database.query( - f"SELECT id FROM household WHERE country_id = ? AND household_hash = ?", - (country_id, household_hash), - ).fetchone()["id"] - - return household_id - except Exception as e: - print(f"Error creating household. Details: {str(e)}") - raise e + return self.households.create( + country_id, + label, + household_json, + hash_object(household_json), + COUNTRY_PACKAGE_VERSIONS.get(country_id), + ) def update_household( self, @@ -93,49 +44,16 @@ def update_household( household_json: dict, label: str, ) -> dict: - """ - Update a household with the given data. - - Args: - country_id (str): The country ID. - household_id (int): The household ID. - payload (dict): The data to update the household with. - """ - print("Updating household") - - try: - household_hash: str = hash_object(household_json) - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - # WHERE must include country_id so an update scoped to - # one country cannot silently overwrite a household that - # happens to share the same numeric id under another - # country. - database.query( - "UPDATE household " - "SET household_json = ?, household_hash = ?, label = ?, api_version = ? " - "WHERE id = ? AND country_id = ?", - ( - json.dumps(household_json), - household_hash, - label, - api_version, - household_id, - country_id, - ), - ) - - # Fetch the updated JSON back from the table. If the - # household did not exist for this country, get_household - # returns None. - updated_household: dict | None = self.get_household( - country_id, household_id + updated = self.households.update( + country_id, + household_id, + label, + household_json, + hash_object(household_json), + COUNTRY_PACKAGE_VERSIONS.get(country_id), + ) + if not updated: + raise LookupError( + f"Household #{household_id} not found for country {country_id}." ) - if updated_household is None: - raise LookupError( - f"Household #{household_id} not found for country {country_id}." - ) - return updated_household - except Exception as e: - print(f"Error updating household #{household_id}. Details: {str(e)}") - raise e + return self.households.get(country_id, household_id) diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index bc63bc34c..f96926e0c 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,139 +1,64 @@ +from __future__ import annotations + import json -from sqlalchemy.engine.row import Row -from policyengine_api.data import database -from policyengine_api.utils import hash_object from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import PolicyDAO +from policyengine_api.utils import hash_object class PolicyService: - """ - Service for storing and retrieving policies; - this is connected to the /policy route and - to the policy database table - """ + def __init__(self, policies: PolicyDAO | None = None): + self._policies = policies + + @property + def policies(self) -> PolicyDAO: + if self._policies is None: + self._policies = PolicyDAO(build_v1_session_manager()) + return self._policies + + @staticmethod + def _validate_policy_id(policy_id: int) -> None: + if type(policy_id) is not int or policy_id < 0: + raise Exception( + f"Invalid policy ID: {policy_id}. Must be a positive integer." + ) def get_policy(self, country_id: str, policy_id: int) -> dict | None: - """ - Fetch policy based only on policy ID and country ID - - Arguments - country_id: str - policy_id: int - - Returns - dict | None -- the policy data, or None if not found - """ - print(f"Getting policy {policy_id} for {country_id}") - - try: - if type(policy_id) is not int or policy_id < 0: - raise Exception( - f"Invalid policy ID: {policy_id}. Must be a positive integer." - ) - - if not country_id: - # This will check for None and empty string - raise ValueError("country_id cannot be empty or None") - - # If no policy found, this will return None - row: Row | None = database.query( - "SELECT * FROM policy WHERE country_id = ? AND id = ?", - (country_id, policy_id), - ).fetchone() - - # policy_json is JSON and must be loaded, if present; to enable, - # we must convert the row to a dictionary - policy = None - if row: - policy = dict(row) - if policy["policy_json"]: - policy["policy_json"] = json.loads(policy["policy_json"]) - return policy - except Exception as e: - print(f"Error getting policy: {str(e)}") - raise e - - def get_policy_json(self, country_id: str, policy_id: int) -> str: - """ - Fetch policy JSON based only on policy ID and country ID - """ - print("Getting policy json") - try: - if type(policy_id) is not int or policy_id < 0: - raise Exception( - f"Invalid policy ID: {policy_id}. Must be a positive integer." - ) - policy = database.query( - f"SELECT policy_json FROM policy WHERE country_id = ? AND id = ?", - (country_id, policy_id), - ).fetchone() - - # Handle nonexisiting record case - if policy is None: - return None - - return policy["policy_json"] - except Exception as e: - print(f"Error getting policy json: {str(e)}") - raise e + self._validate_policy_id(policy_id) + if not country_id: + raise ValueError("country_id cannot be empty or None") + return self.policies.get(country_id, policy_id) + + def get_policy_json(self, country_id: str, policy_id: int) -> str | None: + self._validate_policy_id(policy_id) + policy = self.policies.get(country_id, policy_id) + if policy is None: + return None + value = policy["policy_json"] + return value if isinstance(value, str) else json.dumps(value) def set_policy( self, country_id: str, label: str, policy_json: dict ) -> tuple[int, str, bool]: - """ - Insert a new policy into the database - - Arguments - country_id: str - label: str - policy_json: dict - - Returns - tuple[int, str, bool] -- the new policy ID, a message, and whether or not - the policy already existed - """ - print("Setting new policy") - - try: - # Convert country_id to lowercase - if not country_id.islower(): - country_id = country_id.lower() - - # Validate country_id - if country_id not in COUNTRY_PACKAGE_VERSIONS: - raise ValueError(f"Invalid country_id: {country_id}") - - policy_hash = hash_object(policy_json) - api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) - # Check if policy already exists - print("Checking if policy exists") - existing_policy = self._get_unique_policy_with_label( - country_id, policy_hash, label - ) - - # If so, pass appropriate values back - if existing_policy: - print("Policy already exists") - return existing_policy["id"], "Policy already exists", True - - # Otherwise, insert the new policy... - print("Policy does not exist; creating new policy") - self._create_new_policy( - country_id, policy_json, policy_hash, label, api_version - ) - - # And then fetch it back out; once an ORM is added, we can - # just return policy ID directly from the insert - new_policy = self._get_unique_policy_with_label( - country_id, policy_hash, label - ) - - return int(new_policy["id"]), "Policy created", False - - except Exception as e: - print(f"Error setting policy: {str(e)}") - raise e + country_id = country_id.lower() + if country_id not in COUNTRY_PACKAGE_VERSIONS: + raise ValueError(f"Invalid country_id: {country_id}") + + policy_hash = hash_object(policy_json) + existing = self.policies.find_unique(country_id, policy_hash, label or None) + if existing: + return existing["id"], "Policy already exists", True + + policy_id = self.policies.create( + country_id, + label, + policy_json, + policy_hash, + COUNTRY_PACKAGE_VERSIONS[country_id], + ) + return policy_id, "Policy created", False def _create_new_policy( self, @@ -143,69 +68,9 @@ def _create_new_policy( label: str | None, api_version: str, ) -> None: - """ - Create new policy and insert into database - - Arguments - country_id: str - policy_json: dict - policy_hash: str - label: str | None - api_version: str - - """ - try: - database.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - json.dumps(policy_json), - policy_hash, - label, - api_version, - ), - ) - except Exception as e: - print(f"Error creating new policy: {str(e)}") - raise e + self.policies.create(country_id, label, policy_json, policy_hash, api_version) def _get_unique_policy_with_label( self, country_id: str, policy_hash: str, label: str ) -> dict | None: - """ - Given policy content (represented as a hash) and a label, fetch the policy; - this method ensures that both policy content and label are unique, a workaround - to an old issue whereby multiple copies of the same policy/label pair could be created - - Arguments - country_id: str - policy_hash: str - policy_label: str - - Returns - dict | None -- the policy data, or None if not found - """ - # The code in get_policy_with_label is a workaround - # to the fact that SQLite's cursor method does not properly - # convert 'WHERE x = None' to 'WHERE x IS NULL'; - # though SQLite supports searching and setting with 'WHERE - # x IS y', the production MySQL does not, requiring this - - # This workaround should be removed if and when a proper - # ORM package is added to the API, and this package's - # sanitization methods should be utilized instead - - try: - label_value = "IS NULL" if not label else "= ?" - args = [country_id, policy_hash] - if label: - args.append(label) - - policy = database.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label {label_value}", - tuple(args), - ).fetchone() - return policy - except Exception as e: - print(f"Error getting unique policy with label: {str(e)}") - raise e + return self.policies.find_unique(country_id, policy_hash, label or None) diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index 0bdfb0cdd..50fd4b3cc 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,76 +1,53 @@ -import json +from __future__ import annotations + from typing import Any -from policyengine_api.data import database + +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import UserDAO class UserService: + def __init__(self, users: UserDAO | None = None): + self._users = users + + @property + def users(self) -> UserDAO: + if self._users is None: + self._users = UserDAO(build_v1_session_manager()) + return self._users + def create_profile( self, primary_country: str, auth0_id: str, username: str | None, - user_since: str, + user_since: int, ) -> tuple[bool, Any]: - """ - returns true if a new record was created and false otherwise. - """ - # TODO: this is not written as an atomic operation. This will cause intermittent errors - # in some cases - # https://github.com/PolicyEngine/policyengine-api/issues/2058 to resolve after - # this refactor. row = self.get_profile(auth0_id=auth0_id) if row is not None: return False, row - # Unfortunately, it's not possible to use RETURNING - # with SQLite3 without rewriting the PolicyEngineDatabase - # object or implementing a true ORM, thus the double query - database.query( - f"INSERT INTO user_profiles (primary_country, auth0_id, username, user_since) VALUES (?, ?, ?, ?)", - (primary_country, auth0_id, username, user_since), - ) - - row = self.get_profile(auth0_id=auth0_id) - - return (True, row) + self.users.create_profile(auth0_id, username, primary_country, user_since) + return True, self.get_profile(auth0_id=auth0_id) def get_profile( - self, auth0_id: str | None = None, user_id: str | None = None + self, auth0_id: str | None = None, user_id: int | None = None ) -> Any | None: - key = "user_id" if auth0_id is None else "auth0_id" - value = user_id if auth0_id is None else auth0_id - if value is None: + if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") - row = database.query( - f"SELECT * FROM user_profiles WHERE {key} = ?", - (value,), - ).fetchone() - - return row + return self.users.get_profile(user_id=user_id, auth0_id=auth0_id) def update_profile( self, - user_id: str, + user_id: int, primary_country: str | None, username: str | None, - user_since: str, + user_since: int, ) -> bool: - fields = dict( + if user_id is None: + raise ValueError("you must specify either auth0_id or user_id") + return self.users.update_profile( + user_id, primary_country=primary_country, username=username, user_since=user_since, ) - if self.get_profile(user_id=user_id) is None: - return False - - with_values = [key for key in fields if fields[key] is not None] - fields_update = ",".join([f"{key} = ?" for key in with_values]) - query = f"UPDATE user_profiles SET {fields_update} WHERE user_id = ?" - values = [fields[key] for key in with_values] + [user_id] - - print(f"Updating record {user_id}") - try: - database.query(query, (tuple(values))) - except Exception as ex: - print(f"ERROR: unable to update user record: {ex}") - raise - return True diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index 13cea534c..0ab746808 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -33,7 +33,14 @@ def test_policy_dao_allocates_ids_and_detects_existing_policy(): def test_household_dao_creates_updates_and_reads(): _, households, _ = _daos() household_id = households.create("us", "Home", {"people": {}}, "h", "1.0") - households.update("us", household_id, "Updated", {"people": {"you": {}}}) + households.update( + "us", + household_id, + "Updated", + {"people": {"you": {}}}, + "updated-hash", + "2.0", + ) assert households.get("us", household_id)["label"] == "Updated" assert households.get("uk", household_id) is None @@ -42,4 +49,6 @@ def test_user_dao_profile_lookup_precedence(): _, _, users = _daos() user_id = users.create_profile("auth0|one", "person", "us", 123) assert users.get_profile(auth0_id="auth0|one")["user_id"] == user_id - assert users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] == "auth0|one" + assert ( + users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] == "auth0|one" + ) diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index 86d509ee8..0a4c94a94 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -1,15 +1,12 @@ import pytest import json -from unittest.mock import call +from unittest.mock import MagicMock from policyengine_api.services.policy_service import PolicyService -from tests.fixtures.services.policy_service import ( - valid_policy_data, - valid_hash_value, - mock_hash_object, - mock_database, - existing_policy_record, -) +from tests.fixtures.services.policy_service import valid_hash_value, valid_policy_data + + +pytest_plugins = ["tests.fixtures.services.policy_service"] service = PolicyService() @@ -158,115 +155,49 @@ def test_get_policy_json_given_negative_int_id(self, test_db): class TestSetPolicy: - def test_set_policy_new(self, mock_database, mock_hash_object): - # GIVEN a new policy to insert - new_policy_id = 12 # Different from existing fixture ID + def test_set_policy_new(self, mock_hash_object): + policies = MagicMock() + policies.find_unique.return_value = None + policies.create.return_value = 12 + isolated_service = PolicyService(policies) test_policy = {"param": "value"} test_label = "new_policy" test_country_id = "us" - - # Get current API version dynamically from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS - current_api_version = COUNTRY_PACKAGE_VERSIONS.get(test_country_id) - - # Setup mocks - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call: policy does not exist - {"id": new_policy_id}, # Second call: fetch newly inserted policy - ] - - # Define expected database calls - expected_calls = [ - # First call - check if policy exists - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, valid_hash_value, test_label), - ), - # Second call - insert new policy - call( - "INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - test_country_id, - json.dumps(test_policy), - valid_hash_value, - test_label, - current_api_version, # From COUNTRY_PACKAGE_VERSIONS for 'us' - ), - ), - # Third call - fetch the newly created policy - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, valid_hash_value, test_label), - ), - ] - - # WHEN we call set_policy - policy_id, message, exists = service.set_policy( + policy_id, message, exists = isolated_service.set_policy( test_country_id, test_label, test_policy ) - - # THEN the result should indicate a new policy was created - assert policy_id == new_policy_id + assert policy_id == 12 assert message == "Policy created" assert exists is False - - # Verify the database queries were called as expected - assert mock_database.query.call_args_list == expected_calls - - def test_set_policy_existing( - self, mock_database, mock_hash_object, existing_policy_record - ): - # GIVEN an existing policy record - existing_policy = existing_policy_record - - # Setup mock - mock_database.query.return_value.fetchone.return_value = existing_policy - - # Define expected database calls - matches actual implementation - expected_calls = [ - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label IS NULL", - ( - existing_policy["country_id"], - valid_hash_value, - ), # No None parameter for IS NULL - ), - ] - - # WHEN we call set_policy with existing policy data - policy_id, message, exists = service.set_policy( - existing_policy["country_id"], - existing_policy["label"], - json.loads(existing_policy["policy_json"]), + policies.find_unique.assert_called_once_with( + test_country_id, valid_hash_value, test_label + ) + policies.create.assert_called_once_with( + test_country_id, + test_label, + test_policy, + valid_hash_value, + COUNTRY_PACKAGE_VERSIONS[test_country_id], ) - # THEN the result should indicate the policy already exists - assert policy_id == existing_policy["id"] + def test_set_policy_existing(self, mock_hash_object): + policies = MagicMock() + policies.find_unique.return_value = {"id": 11} + isolated_service = PolicyService(policies) + policy_id, message, exists = isolated_service.set_policy("us", None, {}) + assert policy_id == 11 assert message == "Policy already exists" assert exists is True + policies.create.assert_not_called() - # Verify the database query was called as expected - assert mock_database.query.call_args_list == expected_calls - - def test_set_policy_given_database_insert_failure( - self, mock_database, mock_hash_object - ): - # GIVEN a database insertion failure - test_policy = {"param": "value"} - test_label = "test_policy" - test_country_id = "us" - - # Setup mock to raise exception on insert - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call: policy does not exist - Exception("Database insertion failed"), # Second call: insertion fails - ] - - # WHEN we call set_policy + def test_set_policy_given_database_insert_failure(self, mock_hash_object): + policies = MagicMock() + policies.find_unique.return_value = None + policies.create.side_effect = Exception("Database insertion failed") with pytest.raises(Exception, match="Database insertion failed"): - # THEN an exception should be raised - service.set_policy(test_country_id, test_label, test_policy) + PolicyService(policies).set_policy("us", "test_policy", {}) def test_set_policy_given_invalid_country_id(self, mock_hash_object): # GIVEN an invalid country_id @@ -281,24 +212,12 @@ def test_set_policy_given_invalid_country_id(self, mock_hash_object): # THEN an exception should be raised service.set_policy(INVALID_COUNTRY_ID, test_label, test_policy) - def test_set_policy_given_empty_label(self, mock_database, mock_hash_object): - # GIVEN an empty label - EMPTY_LABEL = "" - test_policy = {"param": "value"} - test_country_id = "us" - - # Setup mock - mock_database.query.return_value.fetchone.side_effect = [ - None, # Policy does not exist - {"id": 13}, # Return mock policy after creation - ] - - # WHEN we call set_policy with an empty label - policy_id, message, exists = service.set_policy( - test_country_id, EMPTY_LABEL, test_policy - ) - - # THEN the result should indicate a new policy was created + def test_set_policy_given_empty_label(self, mock_hash_object): + policies = MagicMock() + policies.find_unique.return_value = None + policies.create.return_value = 13 + policy_id, message, exists = PolicyService(policies).set_policy("us", "", {}) assert policy_id == 13 assert message == "Policy created" assert exists is False + policies.find_unique.assert_called_once_with("us", valid_hash_value, None) diff --git a/tests/unit/services/test_update_profile_service.py b/tests/unit/services/test_update_profile_service.py index 5c6016899..7d51f1ea3 100644 --- a/tests/unit/services/test_update_profile_service.py +++ b/tests/unit/services/test_update_profile_service.py @@ -1,10 +1,8 @@ import pytest from policyengine_api.services.user_service import UserService -from tests.fixtures.services.user_service import ( - valid_user_record, - existing_user_profile, -) + +pytest_plugins = ["tests.fixtures.services.user_service"] service = UserService() @@ -85,10 +83,10 @@ def test_update_profile_with_database_error( # GIVEN an existing profile record (from fixture) # AND a database that raises an exception - def mock_db_query_error(*args, **kwargs): + def mock_dao_error(*args, **kwargs): raise Exception("Database error") - monkeypatch.setattr("policyengine_api.data.database.query", mock_db_query_error) + monkeypatch.setattr(service.users, "update_profile", mock_dao_error) # WHEN we call update_profile # THEN an exception should be raised From f3bcf20d3e04c6fd4316fe5cb121804743f4f83a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:31:16 +0300 Subject: [PATCH 16/89] test(red): define simulation and report DAO transactions --- tests/unit/data/test_run_daos.py | 61 ++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 tests/unit/data/test_run_daos.py diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py new file mode 100644 index 000000000..82157ef94 --- /dev/null +++ b/tests/unit/data/test_run_daos.py @@ -0,0 +1,61 @@ +from datetime import datetime + +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO +from policyengine_api.data.v1_models import V1Base + + +def _daos(): + manager = build_sqlite_session_manager() + V1Base.metadata.create_all(manager.engine) + return SimulationDAO(manager), ReportDAO(manager) + + +def test_simulation_dao_creates_parent_and_monotonic_runs_atomically(): + simulations, _ = _daos() + simulation_id = simulations.create( + country_id="us", + api_version="1", + population_id="7", + population_type="household", + policy_id=2, + ) + first = simulations.create_run( + simulation_id, + run_id="run-1", + status="pending", + trigger_type="create", + requested_at=datetime(2026, 1, 1), + ) + second = simulations.create_run( + simulation_id, + run_id="run-2", + status="pending", + trigger_type="retry", + requested_at=datetime(2026, 1, 2), + ) + assert first["run_sequence"] == 1 + assert second["run_sequence"] == 2 + assert simulations.list_runs(simulation_id)[0]["id"] == "run-2" + + +def test_report_dao_round_trips_parent_run_and_alias(): + _, reports = _daos() + report_id = reports.create( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + year="2026", + ) + run = reports.create_run( + report_id, + run_id="report-run", + status="pending", + trigger_type="create", + requested_at=datetime(2026, 1, 1), + ) + reports.set_alias(99, report_id) + assert reports.get(report_id)["status"] == "pending" + assert run["run_sequence"] == 1 + assert reports.get_alias(99)["canonical_report_output_id"] == report_id From df6ddf3828b2de5cca856bfb559c4f2df5cf6156 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:32:00 +0300 Subject: [PATCH 17/89] feat(green): add simulation and report DAOs --- policyengine_api/data/v1_daos.py | 185 +++++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 6e82cee3d..abb77aa22 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -12,8 +12,13 @@ from policyengine_api.data.v1_models import ( Analysis, Household, + LegacyReportOutputAlias, Policy, ReformImpact, + ReportOutput, + ReportOutputRun, + Simulation, + SimulationRun, Tracer, UserProfile, ) @@ -297,3 +302,183 @@ def get( .order_by(Tracer.id.desc()) ) return _mapping(model) if model else None + + +class SimulationDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def get( + self, simulation_id: int, country_id: str | None = None + ) -> dict[str, Any] | None: + with self.sessions.session() as session: + statement = select(Simulation).where(Simulation.id == simulation_id) + if country_id is not None: + statement = statement.where(Simulation.country_id == country_id) + model = session.scalar(statement) + return _mapping(model) if model else None + + def create(self, **values: Any) -> int: + def operation(session): + model = Simulation(**values) + session.add(model) + session.flush() + return model.id + + return self.sessions.run_in_transaction(operation) + + def update(self, simulation_id: int, **values: Any) -> bool: + def operation(session): + model = session.get(Simulation, simulation_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True + + return self.sessions.run_in_transaction(operation) + + def create_run( + self, simulation_id: int, *, run_id: str, **values: Any + ) -> dict[str, Any]: + def operation(session): + parent = session.scalar( + select(Simulation) + .where(Simulation.id == simulation_id) + .with_for_update() + ) + if parent is None: + raise LookupError(f"Simulation {simulation_id} does not exist") + sequence = ( + session.scalar( + select(func.max(SimulationRun.run_sequence)).where( + SimulationRun.simulation_id == simulation_id + ) + ) + or 0 + ) + 1 + model = SimulationRun( + id=run_id, + simulation_id=simulation_id, + run_sequence=sequence, + **values, + ) + session.add(model) + session.flush() + return _mapping(model) + + return self.sessions.run_in_transaction(operation) + + def get_run(self, run_id: str) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get(SimulationRun, run_id) + return _mapping(model) if model else None + + def list_runs(self, simulation_id: int) -> list[dict[str, Any]]: + with self.sessions.session() as session: + models = session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) + return [_mapping(model) for model in models] + + +class ReportDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def get( + self, report_output_id: int, country_id: str | None = None + ) -> dict[str, Any] | None: + with self.sessions.session() as session: + statement = select(ReportOutput).where(ReportOutput.id == report_output_id) + if country_id is not None: + statement = statement.where(ReportOutput.country_id == country_id) + model = session.scalar(statement) + return _mapping(model) if model else None + + def create(self, **values: Any) -> int: + def operation(session): + model = ReportOutput(**values) + session.add(model) + session.flush() + return model.id + + return self.sessions.run_in_transaction(operation) + + def update(self, report_output_id: int, **values: Any) -> bool: + def operation(session): + model = session.get(ReportOutput, report_output_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True + + return self.sessions.run_in_transaction(operation) + + def create_run( + self, report_output_id: int, *, run_id: str, **values: Any + ) -> dict[str, Any]: + def operation(session): + parent = session.scalar( + select(ReportOutput) + .where(ReportOutput.id == report_output_id) + .with_for_update() + ) + if parent is None: + raise LookupError(f"Report output {report_output_id} does not exist") + sequence = ( + session.scalar( + select(func.max(ReportOutputRun.run_sequence)).where( + ReportOutputRun.report_output_id == report_output_id + ) + ) + or 0 + ) + 1 + model = ReportOutputRun( + id=run_id, + report_output_id=report_output_id, + run_sequence=sequence, + **values, + ) + session.add(model) + session.flush() + return _mapping(model) + + return self.sessions.run_in_transaction(operation) + + def get_run(self, run_id: str) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get(ReportOutputRun, run_id) + return _mapping(model) if model else None + + def list_runs(self, report_output_id: int) -> list[dict[str, Any]]: + with self.sessions.session() as session: + models = session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + return [_mapping(model) for model in models] + + def set_alias(self, legacy_id: int, canonical_id: int) -> None: + def operation(session): + model = session.get(LegacyReportOutputAlias, legacy_id) + if model is None: + session.add( + LegacyReportOutputAlias( + legacy_report_output_id=legacy_id, + canonical_report_output_id=canonical_id, + ) + ) + else: + model.canonical_report_output_id = canonical_id + + self.sessions.run_in_transaction(operation) + + def get_alias(self, legacy_id: int) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get(LegacyReportOutputAlias, legacy_id) + return _mapping(model) if model else None From c4878935ae325f2605607fbc583af8e1ce4b8329 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:32:15 +0300 Subject: [PATCH 18/89] test(red): require DAO boundaries in local data services --- .../test_stage7_local_service_boundaries.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 tests/unit/services/test_stage7_local_service_boundaries.py diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_stage7_local_service_boundaries.py new file mode 100644 index 000000000..de65cbb83 --- /dev/null +++ b/tests/unit/services/test_stage7_local_service_boundaries.py @@ -0,0 +1,22 @@ +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + [ + "ai_analysis_service.py", + "reform_impacts_service.py", + "tracer_analysis_service.py", + "report_output_alias_service.py", + ], +) +def test_local_data_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "local_database" not in source + assert "from policyengine_api.data import database" not in source From 5bc2c0776c970a20122837a98cc987bb15b53e80 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:36:25 +0300 Subject: [PATCH 19/89] feat(green): route local data services through DAOs --- policyengine_api/data/orm.py | 9 +- policyengine_api/data/v1_daos.py | 87 +++++- .../services/ai_analysis_service.py | 64 ++-- .../services/reform_impacts_service.py | 274 +++++++----------- .../services/report_output_alias_service.py | 91 ++---- .../services/tracer_analysis_service.py | 33 ++- 6 files changed, 270 insertions(+), 288 deletions(-) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 22035d6c7..28d829445 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -7,7 +7,7 @@ from pathlib import Path from typing import TypeVar -from sqlalchemy import Engine, create_engine +from sqlalchemy import Engine, create_engine, event from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import StaticPool @@ -93,6 +93,13 @@ def build_v1_session_manager() -> SessionManager: creator=lambda: database._connection, poolclass=StaticPool, ) + event.listen( + engine.pool, + "checkout", + lambda connection, *_: setattr( + connection, "row_factory", _IndexedMappingRow + ), + ) return SessionManager(engine) return build_sqlite_session_manager(database.db_url) return SessionManager(database.pool) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index abb77aa22..c5129c316 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -6,7 +6,7 @@ from datetime import datetime -from sqlalchemy import func, select +from sqlalchemy import delete, func, or_, select from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_models import ( @@ -207,7 +207,10 @@ def get(self, prompt: str) -> str | None: with self.sessions.session() as session: model = session.scalar( select(Analysis) - .where(Analysis.prompt == prompt, Analysis.status == "complete") + .where( + Analysis.prompt == prompt, + Analysis.status.in_(("complete", "ok")), + ) .order_by(Analysis.prompt_id.desc()) ) return model.analysis if model else None @@ -244,6 +247,65 @@ def find(self, *, execution_id: str) -> dict[str, Any] | None: ) return _mapping(model) if model else None + @staticmethod + def _scope(statement, **filters: Any): + return statement.where( + *(getattr(ReformImpact, key) == value for key, value in filters.items()) + ) + + def list(self, **filters: Any) -> list[dict[str, Any]]: + with self.sessions.session() as session: + models = session.scalars( + self._scope(select(ReformImpact), **filters).order_by( + ReformImpact.start_time.desc() + ) + ) + return [_mapping(model) for model in models] + + def list_by_options_hash( + self, options_hash: str, options_hash_prefix: str, **filters: Any + ) -> list[dict[str, Any]]: + with self.sessions.session() as session: + statement = self._scope(select(ReformImpact), **filters).where( + or_( + ReformImpact.options_hash == options_hash, + ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), + ) + ) + models = session.scalars( + statement.order_by( + (ReformImpact.options_hash == options_hash).desc(), + ReformImpact.start_time.desc(), + ) + ) + return [_mapping(model) for model in models] + + def delete_computing(self, **filters: Any) -> None: + def operation(session): + session.execute( + self._scope(delete(ReformImpact), **filters).where( + ReformImpact.status == "computing" + ) + ) + + self.sessions.run_in_transaction(operation) + + def fail(self, execution_id: str, message: str, finished_at: datetime) -> bool: + def operation(session): + model = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if model is None: + return False + model.status = "error" + model.message = message + model.end_time = finished_at + return True + + return self.sessions.run_in_transaction(operation) + def complete(self, execution_id: str, result: Any, finished_at: datetime) -> bool: def operation(session): model = session.scalar( @@ -289,18 +351,21 @@ def operation(session): return self.sessions.run_in_transaction(operation) def get( - self, household_id: int, policy_id: int, country_id: str + self, + household_id: int, + policy_id: int, + country_id: str, + api_version: str | None = None, ) -> dict[str, Any] | None: with self.sessions.session() as session: - model = session.scalar( - select(Tracer) - .where( - Tracer.household_id == household_id, - Tracer.policy_id == policy_id, - Tracer.country_id == country_id, - ) - .order_by(Tracer.id.desc()) + statement = select(Tracer).where( + Tracer.household_id == household_id, + Tracer.policy_id == policy_id, + Tracer.country_id == country_id, ) + if api_version is not None: + statement = statement.where(Tracer.api_version == api_version) + model = session.scalar(statement.order_by(Tracer.id.desc())) return _mapping(model) if model else None diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index ef0f6ce6a..9ee6aca19 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -1,10 +1,13 @@ -import anthropic -import os import json -from typing import Generator, Optional -from policyengine_api.data import local_database +import os +from collections.abc import Generator + +import anthropic from pydantic import BaseModel +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import AnalysisDAO + class StreamEvent(BaseModel): type: str @@ -21,34 +24,24 @@ class ErrorEvent(StreamEvent): class AIAnalysisService: - """ - Base class for various AI analysis-based services, - including SimulationAnalysisService, that connects with the analysis - local database table - """ - - def get_existing_analysis(self, prompt: str) -> Optional[str]: - """ - Get existing analysis from the local database - """ - - analysis = local_database.query( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ).fetchone() + def __init__(self, analyses: AnalysisDAO | None = None): + self._analyses = analyses - if analysis is None: - return None + @property + def analyses(self) -> AnalysisDAO: + if self._analyses is None: + self._analyses = AnalysisDAO(build_v1_session_manager()) + return self._analyses - return json.dumps(analysis["analysis"]) + def get_existing_analysis(self, prompt: str) -> str | None: + analysis = self.analyses.get(prompt) + return json.dumps(analysis) if analysis is not None else None def trigger_ai_analysis(self, prompt: str) -> Generator[str, None, None]: - # Configure a Claude client claude_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) def generate(): response_text = "" - with claude_client.messages.stream( model="claude-sonnet-4-20250514", max_tokens=1500, @@ -57,22 +50,19 @@ def generate(): messages=[{"role": "user", "content": prompt}], ) as stream: for event in stream: - # Docs on structure of Anthropic error events at https://docs.anthropic.com/en/api/messages-streaming#error-events if event.type == "error": - error: dict[str, str] = event.error - error_type: str = error["type"] - return_event = ErrorEvent(error=error_type) - yield json.dumps(return_event.model_dump()) + "\n" + yield ( + json.dumps( + ErrorEvent(error=event.error["type"]).model_dump() + ) + + "\n" + ) return if event.type == "text": response_text += event.text - return_event = TextEvent(stream=event.text) - yield json.dumps(return_event.model_dump()) + "\n" - - # Update the analysis record and return if no error occurred - local_database.query( - f"INSERT INTO analysis (prompt, analysis, status) VALUES (?, ?, ?)", - (prompt, response_text, "ok"), - ) + yield ( + json.dumps(TextEvent(stream=event.text).model_dump()) + "\n" + ) + self.analyses.store(prompt, response_text, "ok") return generate() diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 0f41352f3..46ebbceee 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -1,13 +1,40 @@ -from policyengine_api.data import local_database import datetime +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import ReformImpactDAO + class ReformImpactsService: - """ - Service for storing and retrieving economy-wide reform impacts; - this is connected to the locally-stored reform_impact table - and no existing route - """ + def __init__(self, impacts: ReformImpactDAO | None = None): + self._impacts = impacts + + @property + def impacts(self) -> ReformImpactDAO: + if self._impacts is None: + self._impacts = ReformImpactDAO(build_v1_session_manager()) + return self._impacts + + @staticmethod + def _filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version=None, + ): + filters = { + "country_id": country_id, + "reform_policy_id": policy_id, + "baseline_policy_id": baseline_policy_id, + "region": region, + "dataset": dataset, + "time_period": time_period, + } + if api_version is not None: + filters["api_version"] = api_version + return filters def get_all_reform_impacts( self, @@ -20,29 +47,18 @@ def get_all_reform_impacts( options_hash, api_version, ): - try: - query = ( - "SELECT reform_impact_json, status, message, start_time, execution_id FROM " - "reform_impact WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "options_hash = ? AND api_version = ? AND dataset = ?" - ) - return local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - api_version, - dataset, - ), - ).fetchall() - except Exception as e: - print(f"Error getting all reform impacts: {str(e)}") - raise e + return self.impacts.list( + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ), + options_hash=options_hash, + ) def get_all_reform_impacts_by_options_hash_prefix( self, @@ -56,32 +72,19 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash_prefix, api_version, ): - try: - query = ( - "SELECT reform_impact_json, status, message, start_time, execution_id, options_hash FROM " - "reform_impact WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "(options_hash = ? OR options_hash LIKE ? ESCAPE '\\') AND api_version = ? AND dataset = ? " - "ORDER BY CASE WHEN options_hash = ? THEN 0 ELSE 1 END, start_time DESC" - ) - return local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - options_hash_prefix, - api_version, - dataset, - options_hash, - ), - ).fetchall() - except Exception as e: - print(f"Error getting reform impacts by prefix: {str(e)}") - raise e + return self.impacts.list_by_options_hash( + options_hash, + options_hash_prefix, + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ), + ) def set_reform_impact( self, @@ -99,33 +102,21 @@ def set_reform_impact( start_time, execution_id: str, ): - try: - query = ( - "INSERT INTO reform_impact (country_id, reform_policy_id, baseline_policy_id, " - "region, dataset, time_period, options_json, options_hash, status, api_version, " - "reform_impact_json, start_time, execution_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) - local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options, - options_hash, - status, - api_version, - reform_impact_json, - start_time, - execution_id, - ), - ) - except Exception as e: - print(f"Error setting reform impact: {str(e)}") - raise e + return self.impacts.create( + country_id=country_id, + reform_policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + region=region, + dataset=dataset, + time_period=time_period, + options_json=options, + options_hash=options_hash, + status=status, + api_version=api_version, + reform_impact_json=reform_impact_json, + start_time=start_time, + execution_id=execution_id, + ) def delete_reform_impact( self, @@ -137,29 +128,17 @@ def delete_reform_impact( time_period, options_hash, ): - try: - query = ( - "DELETE FROM reform_impact WHERE country_id = ? AND " - "reform_policy_id = ? AND baseline_policy_id = ? AND " - "region = ? AND time_period = ? AND options_hash = ? AND " - "dataset = ? AND status = 'computing'" - ) - - local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - ), - ) - except Exception as e: - print(f"Error deleting reform impact: {str(e)}") - raise e + self.impacts.delete_computing( + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + ), + options_hash=options_hash, + ) def set_error_reform_impact( self, @@ -173,37 +152,16 @@ def set_error_reform_impact( message, execution_id: str, ): - try: - query = ( - "UPDATE reform_impact SET status = ?, message = ?, end_time = ? WHERE " - "country_id = ? AND reform_policy_id = ? AND baseline_policy_id = ? AND " - "region = ? AND time_period = ? AND options_hash = ? AND dataset = ? AND " - "execution_id = ?" - ) - local_database.query( - query, - ( - "error", - message, - datetime.datetime.strftime( - datetime.datetime.now(datetime.timezone.utc), - "%Y-%m-%d %H:%M:%S.%f", - ), - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - execution_id, - ), - ) - except Exception as e: - print( - f"Error setting error reform impact (something must be REALLY wrong): {str(e)}" - ) - raise e + del ( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) + return self.impacts.fail(execution_id, message, self._now()) def set_complete_reform_impact( self, @@ -217,33 +175,17 @@ def set_complete_reform_impact( reform_impact_json, execution_id, ): - try: - query = ( - "UPDATE reform_impact SET status = ?, message = ?, end_time = ?, " - "reform_impact_json = ? WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "options_hash = ? AND dataset = ? AND execution_id = ?" - ) - local_database.query( - query, - ( - "ok", - "Completed", - datetime.datetime.strftime( - datetime.datetime.now(datetime.timezone.utc), - "%Y-%m-%d %H:%M:%S.%f", - ), - reform_impact_json, - country_id, - reform_policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - execution_id, - ), - ) - except Exception as e: - print(f"Error setting completed reform impact: {str(e)}") - raise e + del ( + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) + return self.impacts.complete(execution_id, reform_impact_json, self._now()) + + @staticmethod + def _now() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 9440cfdfd..54eebe45c 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,97 +1,64 @@ -from sqlalchemy.engine.row import Row - -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import ReportDAO class ReportOutputAliasService: + def __init__(self, reports: ReportDAO | None = None): + self._reports = reports + + @property + def reports(self) -> ReportDAO: + if self._reports is None: + self._reports = ReportDAO(build_v1_session_manager()) + return self._reports + def _get_report_output_row(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT id, country_id, simulation_1_id, simulation_2_id, year - FROM report_outputs - WHERE id = ? - """, - (report_output_id,), - ).fetchone() - return dict(row) if row is not None else None + return self.reports.get(report_output_id) def get_alias(self, legacy_report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM legacy_report_output_aliases - WHERE legacy_report_output_id = ? - """, - (legacy_report_output_id,), - ).fetchone() - return dict(row) if row is not None else None + return self.reports.get_alias(legacy_report_output_id) def resolve_canonical_report_output_id( self, requested_report_output_id: int ) -> int | None: alias = self.get_alias(requested_report_output_id) if alias is not None: - canonical_report_output_id = alias["canonical_report_output_id"] - if self._get_report_output_row(canonical_report_output_id) is None: + canonical_id = alias["canonical_report_output_id"] + if self.reports.get(canonical_id) is None: raise ValueError( - "Alias points to missing canonical report output " - f"#{canonical_report_output_id}" + f"Alias points to missing canonical report output #{canonical_id}" ) - return canonical_report_output_id - - row: Row | None = database.query( - "SELECT id FROM report_outputs WHERE id = ?", - (requested_report_output_id,), - ).fetchone() + return canonical_id + row = self.reports.get(requested_report_output_id) return row["id"] if row is not None else None def set_alias( - self, - legacy_report_output_id: int, - canonical_report_output_id: int, + self, legacy_report_output_id: int, canonical_report_output_id: int ) -> bool: - legacy_report_output = self._get_report_output_row(legacy_report_output_id) - if legacy_report_output is None: + legacy = self.reports.get(legacy_report_output_id) + if legacy is None: raise ValueError( f"Legacy report output #{legacy_report_output_id} not found" ) - - canonical_report_output = self._get_report_output_row( - canonical_report_output_id - ) - if canonical_report_output is None: + canonical = self.reports.get(canonical_report_output_id) + if canonical is None: raise ValueError( f"Canonical report output #{canonical_report_output_id} not found" ) if legacy_report_output_id == canonical_report_output_id: raise ValueError("Legacy and canonical report outputs must be different") - - existing_alias = self.get_alias(legacy_report_output_id) - if existing_alias is not None: - if ( - existing_alias["canonical_report_output_id"] - == canonical_report_output_id - ): + existing = self.reports.get_alias(legacy_report_output_id) + if existing is not None: + if existing["canonical_report_output_id"] == canonical_report_output_id: return True - raise ValueError( "Legacy report output alias already points to canonical report output " - f"#{existing_alias['canonical_report_output_id']}" + f"#{existing['canonical_report_output_id']}" ) - logical_key = ("country_id", "simulation_1_id", "simulation_2_id", "year") - if any( - legacy_report_output[field] != canonical_report_output[field] - for field in logical_key - ): + if any(legacy[field] != canonical[field] for field in logical_key): raise ValueError( "Legacy and canonical report outputs must describe the same report" ) - database.query( - """ - INSERT INTO legacy_report_output_aliases - (legacy_report_output_id, canonical_report_output_id) - VALUES (?, ?) - """, - (legacy_report_output_id, canonical_report_output_id), - ) + self.reports.set_alias(legacy_report_output_id, canonical_report_output_id) return True diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index 2fd072f83..bf2ed4d31 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -1,5 +1,6 @@ -from policyengine_api.data import local_database import json +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import AnalysisDAO, TracerDAO from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS from typing import Generator, Literal import re @@ -9,8 +10,19 @@ class TracerAnalysisService(AIAnalysisService): - def __init__(self): - super().__init__() + def __init__( + self, + tracers: TracerDAO | None = None, + analyses: AnalysisDAO | None = None, + ): + self._tracers = tracers + super().__init__(analyses) + + @property + def tracers(self) -> TracerDAO: + if self._tracers is None: + self._tracers = TracerDAO(build_v1_session_manager()) + return self._tracers def execute_analysis( self, @@ -79,18 +91,17 @@ def get_tracer( ) -> list: try: # Retrieve from the tracers table in the local database - row = local_database.query( - """ - SELECT * FROM tracers - WHERE household_id = ? AND policy_id = ? AND country_id = ? AND api_version = ? - """, - (household_id, policy_id, country_id, api_version), - ).fetchone() + row = self.tracers.get(household_id, policy_id, country_id, api_version) if row is None: raise NotFound("No household simulation tracer found") - tracer_output_list = json.loads(row["tracer_output"]) + tracer_output = row["tracer_output"] + tracer_output_list = ( + json.loads(tracer_output) + if isinstance(tracer_output, str) + else tracer_output + ) return tracer_output_list except Exception as e: From d59622a3c971564174909cbb985b66e5e4b653e1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:36:38 +0300 Subject: [PATCH 20/89] test(red): require DAO boundaries in run and spec services --- .../test_stage7_run_service_boundaries.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 tests/unit/services/test_stage7_run_service_boundaries.py diff --git a/tests/unit/services/test_stage7_run_service_boundaries.py b/tests/unit/services/test_stage7_run_service_boundaries.py new file mode 100644 index 000000000..febe5a852 --- /dev/null +++ b/tests/unit/services/test_stage7_run_service_boundaries.py @@ -0,0 +1,21 @@ +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + [ + "simulation_run_service.py", + "simulation_spec_service.py", + "report_run_service.py", + "report_spec_service.py", + ], +) +def test_run_and_spec_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "from policyengine_api.data import database" not in source From f5c7ca03f1db8bbe59948a7fd4225251e55ede68 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:38:07 +0300 Subject: [PATCH 21/89] feat(green): route run and spec services through DAOs --- .../services/report_run_service.py | 144 ++++++------------ .../services/report_spec_service.py | 52 ++++--- .../services/simulation_run_service.py | 144 ++++++------------ .../services/simulation_spec_service.py | 35 ++--- 4 files changed, 135 insertions(+), 240 deletions(-) diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index 9899f6cc9..ca813e368 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -3,9 +3,8 @@ from datetime import datetime, timezone from typing import Any -from sqlalchemy.engine.row import Row - -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import ReportDAO from policyengine_api.services.run_sync_utils import select_display_report_run @@ -23,25 +22,26 @@ class ReportRunService: - def _utc_timestamp(self) -> str: - return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + def __init__(self, reports: ReportDAO | None = None): + self._reports = reports - def _serialize_json( - self, value: dict[str, Any] | list[Any] | str | None - ) -> str | None: - if value is None or isinstance(value, str): - return value - return json.dumps(value) + @property + def reports(self) -> ReportDAO: + if self._reports is None: + self._reports = ReportDAO(build_v1_session_manager()) + return self._reports - def _parse_run_row(self, row: Row | dict | None) -> dict | None: + def _parse_run_row(self, row: dict | None) -> dict | None: if row is None: return None - run = dict(row) if isinstance(run.get("report_spec_snapshot_json"), str): run["report_spec_snapshot_json"] = json.loads( run["report_spec_snapshot_json"] ) + for field in ("requested_at", "started_at", "finished_at"): + if isinstance(run.get(field), datetime): + run[field] = run[field].strftime("%Y-%m-%d %H:%M:%S") return run def create_report_output_run( @@ -56,98 +56,46 @@ def create_report_output_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> dict: - run_id = run_id or str(uuid.uuid4()) - version_manifest = version_manifest or {} - lock_clause = "" if database.local else " FOR UPDATE" - - def create_run_transaction(tx) -> None: - parent_row: Row | None = tx.query( - f"SELECT id FROM report_outputs WHERE id = ?{lock_clause}", - (report_output_id,), - ).fetchone() - if parent_row is None: - raise ValueError(f"Report output #{report_output_id} not found") - - run_sequence_row: Row | None = tx.query( - """ - SELECT COALESCE(MAX(run_sequence), 0) AS max_run_sequence - FROM report_output_runs - WHERE report_output_id = ? - """, - (report_output_id,), - ).fetchone() - run_sequence = ( - int(run_sequence_row["max_run_sequence"]) + 1 - if run_sequence_row is not None - else 1 - ) - - requested_at = self._utc_timestamp() - is_terminal = status in ("complete", "error") - has_started = status in ("running", "complete", "error") - started_at = requested_at if has_started else None - finished_at = requested_at if is_terminal else None - - tx.query( - f""" - INSERT INTO report_output_runs ( - id, report_output_id, run_sequence, status, output, error_message, - trigger_type, requested_at, started_at, finished_at, source_run_id, - report_spec_snapshot_json, {", ".join(REPORT_RUN_VERSION_FIELDS)} - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - run_id, - report_output_id, - run_sequence, - status, - self._serialize_json(output), - error_message, - trigger_type, - requested_at, - started_at, - finished_at, - source_run_id, - self._serialize_json(report_spec_snapshot), - *[ - version_manifest.get(field) - for field in REPORT_RUN_VERSION_FIELDS - ], - ), + now = datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) + terminal = status in ("complete", "error") + started = status in ("running", "complete", "error") + values = { + "status": status, + "output": output, + "error_message": error_message, + "trigger_type": trigger_type, + "requested_at": now, + "started_at": now if started else None, + "finished_at": now if terminal else None, + "source_run_id": source_run_id, + "report_spec_snapshot_json": report_spec_snapshot, + } + values.update( + { + field: (version_manifest or {}).get(field) + for field in REPORT_RUN_VERSION_FIELDS + } + ) + try: + run = self.reports.create_run( + report_output_id, run_id=run_id or str(uuid.uuid4()), **values ) - - database.transaction(create_run_transaction) - return self.get_report_output_run(run_id) + except LookupError as error: + raise ValueError(f"Report output #{report_output_id} not found") from error + return self._parse_run_row(run) def get_report_output_run(self, run_id: str) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (run_id,), - ).fetchone() - return self._parse_run_row(row) + return self._parse_run_row(self.reports.get_run(run_id)) def list_report_output_runs(self, report_output_id: int) -> list[dict]: - rows = database.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (report_output_id,), - ).fetchall() - return [self._parse_run_row(row) for row in rows] + return [ + self._parse_run_row(row) + for row in reversed(self.reports.list_runs(report_output_id)) + ] def get_newest_report_output_run(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence DESC - LIMIT 1 - """, - (report_output_id,), - ).fetchone() - return self._parse_run_row(row) + rows = self.reports.list_runs(report_output_id) + return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, report_output: dict) -> dict | None: runs_descending = list( diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index b81cc566f..648f5b435 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -2,9 +2,9 @@ from typing import Any, Literal from pydantic import BaseModel, Field -from sqlalchemy.engine.row import Row -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO REPORT_SPEC_SCHEMA_VERSION = 1 REPORT_SPEC_STATUSES = {"explicit", "backfilled_assumed"} @@ -42,6 +42,20 @@ class EconomyReportSpec(BaseModel): class ReportSpecService: + def __init__( + self, + reports: ReportDAO | None = None, + simulations: SimulationDAO | None = None, + ): + self._reports = reports + self._simulations = simulations + + def _ensure_daos(self) -> None: + if self._reports is None or self._simulations is None: + manager = build_v1_session_manager() + self._reports = self._reports or ReportDAO(manager) + self._simulations = self._simulations or SimulationDAO(manager) + def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: raise ValueError( @@ -49,18 +63,12 @@ def _validate_schema_version(self, schema_version: int | None) -> None: ) def _get_report_output_row(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output_id,), - ).fetchone() - return dict(row) if row is not None else None + self._ensure_daos() + return self._reports.get(report_output_id) def _get_simulation_row(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_id,), - ).fetchone() - return dict(row) if row is not None else None + self._ensure_daos() + return self._simulations.get(simulation_id) def _get_linked_simulations(self, report_output: dict) -> tuple[dict, dict | None]: simulation_1 = self._get_simulation_row(report_output["simulation_1_id"]) @@ -373,18 +381,12 @@ def set_report_spec( raise ValueError(f"Report output #{report_output_id} not found") self._validate_report_spec_matches_row(report_output, report_spec) - database.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - report_spec.report_kind, - report_spec.model_dump_json(), - schema_version, - report_spec_status, - report_output_id, - ), + self._ensure_daos() + self._reports.update( + report_output_id, + report_kind=report_spec.report_kind, + report_spec_json=report_spec.model_dump(), + report_spec_schema_version=schema_version, + report_spec_status=report_spec_status, ) return True diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 544aca9c2..da417a4ce 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -2,9 +2,8 @@ import uuid from typing import Any -from sqlalchemy.engine.row import Row - -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import SimulationDAO SIMULATION_RUN_VERSION_FIELDS = ( @@ -17,17 +16,18 @@ class SimulationRunService: - def _serialize_json( - self, value: dict[str, Any] | list[Any] | str | None - ) -> str | None: - if value is None or isinstance(value, str): - return value - return json.dumps(value) + def __init__(self, simulations: SimulationDAO | None = None): + self._simulations = simulations + + @property + def simulations(self) -> SimulationDAO: + if self._simulations is None: + self._simulations = SimulationDAO(build_v1_session_manager()) + return self._simulations - def _parse_run_row(self, row: Row | dict | None) -> dict | None: + def _parse_run_row(self, row: dict | None) -> dict | None: if row is None: return None - run = dict(row) if isinstance(run.get("simulation_spec_snapshot_json"), str): run["simulation_spec_snapshot_json"] = json.loads( @@ -49,95 +49,45 @@ def create_simulation_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> dict: - run_id = run_id or str(uuid.uuid4()) - version_manifest = version_manifest or {} - lock_clause = "" if database.local else " FOR UPDATE" - - def create_run_transaction(tx) -> None: - parent_row: Row | None = tx.query( - f"SELECT id FROM simulations WHERE id = ?{lock_clause}", - (simulation_id,), - ).fetchone() - if parent_row is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - run_sequence_row: Row | None = tx.query( - """ - SELECT COALESCE(MAX(run_sequence), 0) AS max_run_sequence - FROM simulation_runs - WHERE simulation_id = ? - """, - (simulation_id,), - ).fetchone() - run_sequence = ( - int(run_sequence_row["max_run_sequence"]) + 1 - if run_sequence_row is not None - else 1 + values = { + "report_output_run_id": report_output_run_id, + "input_position": input_position, + "status": status, + "output": output, + "error_message": error_message, + "trigger_type": trigger_type, + "requested_at": None, + "started_at": None, + "finished_at": None, + "source_run_id": source_run_id, + "simulation_spec_snapshot_json": simulation_spec_snapshot, + } + values.update( + { + field: (version_manifest or {}).get(field) + for field in SIMULATION_RUN_VERSION_FIELDS + } + ) + try: + run = self.simulations.create_run( + simulation_id, run_id=run_id or str(uuid.uuid4()), **values ) - - tx.query( - f""" - INSERT INTO simulation_runs ( - id, simulation_id, report_output_run_id, input_position, run_sequence, - status, output, error_message, trigger_type, requested_at, started_at, - finished_at, source_run_id, simulation_spec_snapshot_json, - {", ".join(SIMULATION_RUN_VERSION_FIELDS)} - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - run_id, - simulation_id, - report_output_run_id, - input_position, - run_sequence, - status, - self._serialize_json(output), - error_message, - trigger_type, - None, - None, - None, - source_run_id, - self._serialize_json(simulation_spec_snapshot), - *[ - version_manifest.get(field) - for field in SIMULATION_RUN_VERSION_FIELDS - ], - ), - ) - - database.transaction(create_run_transaction) - return self.get_simulation_run(run_id) + except LookupError as error: + raise ValueError(f"Simulation #{simulation_id} not found") from error + return self._parse_run_row(run) def get_simulation_run(self, run_id: str) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulation_runs WHERE id = ?", - (run_id,), - ).fetchone() - return self._parse_run_row(row) + return self._parse_run_row(self.simulations.get_run(run_id)) def list_simulation_runs(self, simulation_id: int) -> list[dict]: - rows = database.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence ASC - """, - (simulation_id,), - ).fetchall() - return [self._parse_run_row(row) for row in rows] + return [ + self._parse_run_row(row) + for row in reversed(self.simulations.list_runs(simulation_id)) + ] def get_newest_simulation_run(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence DESC - LIMIT 1 - """, - (simulation_id,), - ).fetchone() - return self._parse_run_row(row) + rows = self.simulations.list_runs(simulation_id) + return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, simulation: dict) -> dict | None: if simulation.get("active_run_id"): @@ -145,9 +95,7 @@ def select_display_run(self, simulation: dict) -> dict | None: if active_run is not None: return active_run if simulation.get("latest_successful_run_id"): - latest_successful_run = self.get_simulation_run( - simulation["latest_successful_run_id"] - ) - if latest_successful_run is not None: - return latest_successful_run + successful = self.get_simulation_run(simulation["latest_successful_run_id"]) + if successful is not None: + return successful return self.get_newest_simulation_run(simulation["id"]) diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index d5f6e86f7..2c9892a10 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -2,9 +2,8 @@ from typing import Literal from pydantic import BaseModel -from sqlalchemy.engine.row import Row - -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import SimulationDAO SIMULATION_SPEC_SCHEMA_VERSION = 1 @@ -17,6 +16,15 @@ class SimulationSpec(BaseModel): class SimulationSpecService: + def __init__(self, simulations: SimulationDAO | None = None): + self._simulations = simulations + + @property + def simulations(self) -> SimulationDAO: + if self._simulations is None: + self._simulations = SimulationDAO(build_v1_session_manager()) + return self._simulations + def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != SIMULATION_SPEC_SCHEMA_VERSION: raise ValueError( @@ -24,11 +32,7 @@ def _validate_schema_version(self, schema_version: int | None) -> None: ) def _get_simulation_row(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_id,), - ).fetchone() - return dict(row) if row is not None else None + return self.simulations.get(simulation_id) def _validate_simulation_spec_matches_row( self, simulation: dict, simulation_spec: SimulationSpec @@ -77,16 +81,9 @@ def set_simulation_spec( raise ValueError(f"Simulation #{simulation_id} not found") self._validate_simulation_spec_matches_row(simulation, simulation_spec) - database.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - simulation_spec.model_dump_json(), - schema_version, - simulation_id, - ), + self.simulations.update( + simulation_id, + simulation_spec_json=simulation_spec.model_dump(), + simulation_spec_schema_version=schema_version, ) return True From 6861622842ecde2a9542bdea2e7f9db19629ac0f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:40:54 +0300 Subject: [PATCH 22/89] test(red): require DAO-backed orchestration services --- .../test_stage7_orchestration_boundaries.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 tests/unit/services/test_stage7_orchestration_boundaries.py diff --git a/tests/unit/services/test_stage7_orchestration_boundaries.py b/tests/unit/services/test_stage7_orchestration_boundaries.py new file mode 100644 index 000000000..83cf7d8dd --- /dev/null +++ b/tests/unit/services/test_stage7_orchestration_boundaries.py @@ -0,0 +1,15 @@ +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", ["simulation_service.py", "report_output_service.py"] +) +def test_orchestration_services_do_not_access_database_connections(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert "from policyengine_api.data import database" not in source + assert "database.transaction(" not in source From 2af7f9410b3bd2474596f121aeae5b470a741554 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:42:20 +0300 Subject: [PATCH 23/89] feat(green): route orchestration through SQLAlchemy DAOs --- policyengine_api/data/v1_daos.py | 232 ++++++++ policyengine_api/data/v1_models.py | 47 +- .../services/report_output_service.py | 32 +- .../services/simulation_service.py | 555 +++--------------- 4 files changed, 355 insertions(+), 511 deletions(-) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index c5129c316..985ae2a32 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -5,6 +5,7 @@ from typing import Any from datetime import datetime +import uuid from sqlalchemy import delete, func, or_, select @@ -30,6 +31,73 @@ def _mapping(model: Any) -> dict[str, Any]: } +class _Rows: + def __init__(self, rows): + self.rows = rows + self.index = 0 + + def fetchone(self): + if self.index >= len(self.rows): + return None + row = self.rows[self.index] + self.index += 1 + return row + + def fetchall(self): + rows = self.rows[self.index :] + self.index = len(self.rows) + return rows + + +class SQLAlchemyDAO: + """Compatibility execution boundary for complex v1 transactional SQL. + + New CRUD belongs in typed DAOs. This boundary keeps the mature report + orchestration on SQLAlchemy-owned sessions while it is decomposed. + """ + + def __init__(self, sessions: SessionManager, session=None): + self.sessions = sessions + self._session = session + + @property + def local(self) -> bool: + return self.sessions.engine.dialect.name == "sqlite" + + def _statement(self, statement: str) -> str: + return statement if self.local else statement.replace("?", "%s") + + @staticmethod + def _rows(result) -> _Rows: + if not result.returns_rows: + return _Rows([]) + return _Rows(list(result.mappings())) + + def query(self, statement: str, params=None) -> _Rows: + if self._session is not None: + result = self._session.connection().exec_driver_sql( + self._statement(statement), params or () + ) + return self._rows(result) + + def operation(session): + result = session.connection().exec_driver_sql( + self._statement(statement), params or () + ) + return self._rows(result) + + return self.sessions.run_in_transaction(operation) + + def transaction(self, callback): + return self.sessions.run_in_transaction( + lambda session: callback(SQLAlchemyDAO(self.sessions, session)) + ) + + @property + def session(self): + return self._session + + class PolicyDAO: def __init__(self, sessions: SessionManager): self.sessions = sessions @@ -383,6 +451,16 @@ def get( model = session.scalar(statement) return _mapping(model) if model else None + @staticmethod + def get_in_session( + session, simulation_id: int, country_id: str | None = None + ) -> dict[str, Any] | None: + statement = select(Simulation).where(Simulation.id == simulation_id) + if country_id is not None: + statement = statement.where(Simulation.country_id == country_id) + model = session.scalar(statement) + return _mapping(model) if model else None + def create(self, **values: Any) -> int: def operation(session): model = Simulation(**values) @@ -392,6 +470,160 @@ def operation(session): return self.sessions.run_in_transaction(operation) + def find_latest(self, **filters: Any) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.scalar( + select(Simulation) + .where( + *( + getattr(Simulation, key) == value + for key, value in filters.items() + ) + ) + .order_by(Simulation.id.desc()) + ) + return _mapping(model) if model else None + + @staticmethod + def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: + return next((run.id for run in runs if run.status == "complete"), None) + + def ensure_dual_write_state_in_session( + self, + session, + simulation_id: int, + country_id: str | None = None, + ) -> dict[str, Any]: + statement = ( + select(Simulation).where(Simulation.id == simulation_id).with_for_update() + ) + if country_id is not None: + statement = statement.where(Simulation.country_id == country_id) + simulation = session.scalar(statement) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") + + spec = { + "country_id": simulation.country_id, + "population_id": simulation.population_id, + "population_type": simulation.population_type, + "policy_id": simulation.policy_id, + } + simulation.simulation_spec_json = spec + simulation.simulation_spec_schema_version = 1 + runs = list( + session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) + ) + if not runs: + run = SimulationRun( + id=str(uuid.uuid4()), + simulation_id=simulation_id, + run_sequence=1, + status=simulation.status, + output=simulation.output, + error_message=simulation.error_message, + trigger_type="initial", + simulation_spec_snapshot_json=spec, + country_package_version=simulation.api_version, + ) + session.add(run) + session.flush() + runs = [run] + else: + mutable = next( + (run for run in runs if run.id == simulation.active_run_id), + runs[0], + ) + mutable.status = simulation.status + mutable.output = simulation.output + mutable.error_message = simulation.error_message + mutable.simulation_spec_snapshot_json = spec + mutable.country_package_version = simulation.api_version + + latest_successful = self._latest_successful_run_id(runs) + if simulation.status in {"pending", "running"}: + simulation.active_run_id = runs[0].id + else: + simulation.active_run_id = None + if simulation.status == "complete" and latest_successful is None: + latest_successful = runs[0].id + simulation.latest_successful_run_id = latest_successful + session.flush() + return _mapping(simulation) + + def ensure_dual_write_state( + self, simulation_id: int, country_id: str | None = None + ) -> dict[str, Any]: + return self.sessions.run_in_transaction( + lambda session: self.ensure_dual_write_state_in_session( + session, simulation_id, country_id + ) + ) + + def create_or_get_with_sync( + self, + *, + sync_callback, + **values: Any, + ) -> dict[str, Any]: + def operation(session): + filters = { + key: values[key] + for key in ( + "country_id", + "population_id", + "population_type", + "policy_id", + ) + } + model = session.scalar( + select(Simulation) + .where( + *( + getattr(Simulation, key) == value + for key, value in filters.items() + ) + ) + .order_by(Simulation.id.desc()) + .with_for_update() + ) + if model is None: + model = Simulation(**values) + session.add(model) + session.flush() + return sync_callback(session, model.id, country_id=model.country_id) + + return self.sessions.run_in_transaction(operation) + + def update_with_sync( + self, + simulation_id: int, + country_id: str, + values: dict[str, Any], + sync_callback, + ) -> dict[str, Any]: + def operation(session): + model = session.scalar( + select(Simulation) + .where( + Simulation.id == simulation_id, + Simulation.country_id == country_id, + ) + .with_for_update() + ) + if model is None: + raise ValueError(f"Simulation #{simulation_id} not found") + for key, value in values.items(): + setattr(model, key, value) + session.flush() + return sync_callback(session, simulation_id, country_id=country_id) + + return self.sessions.run_in_transaction(operation) + def update(self, simulation_id: int, **values: Any) -> bool: def operation(session): model = session.get(Simulation, simulation_id) diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 6648612a6..834016e78 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -9,7 +9,16 @@ from datetime import datetime from typing import Any -from sqlalchemy import BigInteger, DateTime, Integer, JSON, String, Text, UniqueConstraint, text +from sqlalchemy import ( + BigInteger, + DateTime, + Integer, + JSON, + String, + Text, + UniqueConstraint, + text, +) from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column @@ -51,7 +60,9 @@ class Policy(V1Base): class Economy(V1Base): __tablename__ = "economy" - economy_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + economy_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) policy_id: Mapped[int] country_id: Mapped[str] = mapped_column(String(3)) region: Mapped[str | None] = mapped_column(String(32)) @@ -59,21 +70,23 @@ class Economy(V1Base): options_json: Mapped[Any] = mapped_column(JSON) options_hash: Mapped[str] = mapped_column(String(255)) api_version: Mapped[str] = mapped_column(String(10)) - economy_json: Mapped[Any | None] = mapped_column(JSON) + economy_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) status: Mapped[str] = mapped_column(String(32)) message: Mapped[str | None] = mapped_column(String(255)) class ReformImpact(V1Base): __tablename__ = "reform_impact" - reform_impact_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + reform_impact_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) baseline_policy_id: Mapped[int] reform_policy_id: Mapped[int] country_id: Mapped[str] = mapped_column(String(3)) region: Mapped[str] = mapped_column(String(32)) dataset: Mapped[str] = mapped_column(String(255)) time_period: Mapped[str] = mapped_column(String(32)) - options_json: Mapped[Any | None] = mapped_column(JSON) + options_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) options_hash: Mapped[str | None] = mapped_column(String(255)) api_version: Mapped[str] = mapped_column(String(10)) reform_impact_json: Mapped[Any] = mapped_column(JSON) @@ -86,7 +99,9 @@ class ReformImpact(V1Base): class Analysis(V1Base): __tablename__ = "analysis" - prompt_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + prompt_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) prompt: Mapped[str] = mapped_column(Text) analysis: Mapped[str | None] = mapped_column(Text) status: Mapped[str] = mapped_column(String(32)) @@ -140,9 +155,9 @@ class Simulation(V1Base): population_type: Mapped[str] = mapped_column(String(50)) policy_id: Mapped[int] status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) - output: Mapped[Any | None] = mapped_column(JSON) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) error_message: Mapped[str | None] = mapped_column(Text) - simulation_spec_json: Mapped[Any | None] = mapped_column(JSON) + simulation_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) simulation_spec_schema_version: Mapped[int | None] active_run_id: Mapped[str | None] = mapped_column(String(36)) latest_successful_run_id: Mapped[str | None] = mapped_column(String(36)) @@ -156,11 +171,11 @@ class ReportOutput(V1Base): simulation_2_id: Mapped[int | None] api_version: Mapped[str] = mapped_column(String(10)) status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) - output: Mapped[Any | None] = mapped_column(JSON) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) error_message: Mapped[str | None] = mapped_column(Text) year: Mapped[str | None] = mapped_column(String(255), server_default=text("'2025'")) report_kind: Mapped[str | None] = mapped_column(String(64)) - report_spec_json: Mapped[Any | None] = mapped_column(JSON) + report_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) report_spec_schema_version: Mapped[int | None] report_spec_status: Mapped[str | None] = mapped_column(String(32)) active_run_id: Mapped[str | None] = mapped_column(String(36)) @@ -174,14 +189,16 @@ class ReportOutputRun(V1Base): report_output_id: Mapped[int] run_sequence: Mapped[int] status: Mapped[str] = mapped_column(String(32)) - output: Mapped[Any | None] = mapped_column(JSON) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) error_message: Mapped[str | None] = mapped_column(Text) trigger_type: Mapped[str] = mapped_column(String(32)) requested_at: Mapped[datetime | None] = mapped_column(DateTime) started_at: Mapped[datetime | None] = mapped_column(DateTime) finished_at: Mapped[datetime | None] = mapped_column(DateTime) source_run_id: Mapped[str | None] = mapped_column(String(36)) - report_spec_snapshot_json: Mapped[Any | None] = mapped_column(JSON) + report_spec_snapshot_json: Mapped[Any | None] = mapped_column( + JSON(none_as_null=True) + ) country_package_version: Mapped[str | None] = mapped_column(String(255)) policyengine_version: Mapped[str | None] = mapped_column(String(255)) data_version: Mapped[str | None] = mapped_column(String(255)) @@ -202,14 +219,16 @@ class SimulationRun(V1Base): input_position: Mapped[int | None] run_sequence: Mapped[int] status: Mapped[str] = mapped_column(String(32)) - output: Mapped[Any | None] = mapped_column(JSON) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) error_message: Mapped[str | None] = mapped_column(Text) trigger_type: Mapped[str] = mapped_column(String(32)) requested_at: Mapped[datetime | None] = mapped_column(DateTime) started_at: Mapped[datetime | None] = mapped_column(DateTime) finished_at: Mapped[datetime | None] = mapped_column(DateTime) source_run_id: Mapped[str | None] = mapped_column(String(36)) - simulation_spec_snapshot_json: Mapped[Any | None] = mapped_column(JSON) + simulation_spec_snapshot_json: Mapped[Any | None] = mapped_column( + JSON(none_as_null=True) + ) country_package_version: Mapped[str | None] = mapped_column(String(255)) policyengine_version: Mapped[str | None] = mapped_column(String(255)) data_version: Mapped[str | None] = mapped_column(String(255)) diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 38b5704fa..f383f15a6 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -4,7 +4,8 @@ from sqlalchemy.engine.row import Row from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.data import database +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import SQLAlchemyDAO from policyengine_api.services.report_spec_service import ( ECONOMY_REPORT_KINDS, ReportSpec, @@ -21,12 +22,19 @@ class ReportOutputService: - def __init__(self): + def __init__(self, persistence: SQLAlchemyDAO | None = None): + self._persistence = persistence self.report_spec_service = ReportSpecService() self.simulation_service = SimulationService() + @property + def persistence(self) -> SQLAlchemyDAO: + if self._persistence is None: + self._persistence = SQLAlchemyDAO(build_v1_session_manager()) + return self._persistence + def _lock_clause(self) -> str: - return "" if database.local else " FOR UPDATE" + return "" if self.persistence.local else " FOR UPDATE" def _utc_timestamp(self) -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") @@ -78,7 +86,7 @@ def _get_report_output_row( country_id: str | None = None, for_update: bool = False, ) -> dict | None: - queryer = queryer or database + queryer = queryer or self.persistence query = "SELECT * FROM report_outputs WHERE id = ?" params: list[int | str] = [report_output_id] if country_id is not None: @@ -97,7 +105,7 @@ def _get_linked_simulations( queryer=None, bootstrap_dual_write_state: bool = False, ) -> tuple[dict, dict | None]: - queryer = queryer or database + queryer = queryer or self.persistence if bootstrap_dual_write_state: simulation_1 = self.simulation_service._ensure_simulation_dual_write_state_in_transaction( queryer, @@ -159,7 +167,7 @@ def _require_simulation_exists( def _list_report_runs_descending( self, report_output_id: int, *, queryer=None ) -> list[dict]: - queryer = queryer or database + queryer = queryer or self.persistence rows = queryer.query( """ SELECT * FROM report_output_runs @@ -236,7 +244,7 @@ def _with_display_run_timestamps( values live on report_output_runs; this helper chooses the display run, formats its requested/started/finished timestamps, and returns an enriched copy of the report output dict. It intentionally does not - mutate database state. + mutate self.persistence state. These timestamps describe the selected base report execution. They are not user-report association metadata and should not be treated as a @@ -636,7 +644,7 @@ def ensure_report_output_dual_write_state( report_output_id: int, country_id: str | None = None, ) -> dict: - return database.transaction( + return self.persistence.transaction( lambda tx: self._ensure_report_output_dual_write_state_in_transaction( tx, report_output_id, @@ -653,7 +661,7 @@ def get_stored_report_output( This is used by mutation paths that must address the originally requested row. It still runs dual-write synchronization, so it may bootstrap or repair run/spec metadata and returns the display-run - timestamp projection. It is therefore not a raw database read. + timestamp projection. It is therefore not a raw self.persistence read. TODO: Split raw storage lookup from synchronized response projection in a later run-backed read migration PR. @@ -688,7 +696,7 @@ def _find_existing_report_output_row( year: str, queryer=None, ) -> dict | None: - queryer = queryer or database + queryer = queryer or self.persistence api_version = get_report_output_cache_version(country_id) query = """ SELECT * FROM report_outputs @@ -852,7 +860,7 @@ def tx_callback(tx): country_id=country_id, ) - return database.transaction(tx_callback) + return self.persistence.transaction(tx_callback) except Exception as e: print(f"Error creating report output. Details: {str(e)}") @@ -953,7 +961,7 @@ def tx_callback(tx): country_id=country_id, ) - database.transaction(tx_callback) + self.persistence.transaction(tx_callback) print(f"Successfully updated report output #{report_id}") return True diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index e5582ee17..31c21a4af 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,314 +1,54 @@ -import uuid - -from sqlalchemy.engine.row import Row +import json from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data import database -from policyengine_api.services.run_sync_utils import ( - determine_parent_pointers, - parse_json_field, - serialize_json_field, -) -from policyengine_api.services.simulation_spec_service import ( - SimulationSpec, - SimulationSpecService, -) +from policyengine_api.data.orm import build_v1_session_manager +from policyengine_api.data.v1_daos import SimulationDAO +from policyengine_api.services.simulation_spec_service import SimulationSpecService class SimulationService: - def __init__(self): - self.simulation_spec_service = SimulationSpecService() + def __init__(self, simulations: SimulationDAO | None = None): + self._simulations = simulations + self.simulation_spec_service = SimulationSpecService(simulations) - def _lock_clause(self) -> str: - return "" if database.local else " FOR UPDATE" + @property + def simulations(self) -> SimulationDAO: + if self._simulations is None: + self._simulations = SimulationDAO(build_v1_session_manager()) + self.simulation_spec_service = SimulationSpecService(self._simulations) + return self._simulations - def _get_simulation_row( + def _ensure_simulation_dual_write_state_in_transaction( self, + session, simulation_id: int, *, - queryer=None, country_id: str | None = None, - for_update: bool = False, - ) -> dict | None: - queryer = queryer or database - query = "SELECT * FROM simulations WHERE id = ?" - params: list[int | str] = [simulation_id] - if country_id is not None: - query += " AND country_id = ?" - params.append(country_id) - if for_update: - query += self._lock_clause() - - row: Row | None = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None - - def _find_existing_simulation_row( - self, - *, - country_id: str, - population_id: str, - population_type: str, - policy_id: int, - queryer=None, - ) -> dict | None: - queryer = queryer or database - row: Row | None = queryer.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - ORDER BY id DESC - """, - (country_id, population_id, population_type, policy_id), - ).fetchone() - return dict(row) if row is not None else None - - def _build_version_manifest(self, simulation: dict) -> dict[str, str | None]: - return { - "country_package_version": simulation.get("api_version"), - "policyengine_version": None, - "data_version": None, - "runtime_app_name": None, - "simulation_cache_version": None, - } - - def _list_simulation_runs_descending( - self, simulation_id: int, *, queryer=None - ) -> list[dict]: - queryer = queryer or database - rows = queryer.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence DESC - """, - (simulation_id,), - ).fetchall() - - runs = [] - for row in rows: - run = dict(row) - run["simulation_spec_snapshot_json"] = parse_json_field( - run.get("simulation_spec_snapshot_json") - ) - runs.append(run) - return runs - - def _select_mutable_run( - self, simulation: dict, runs_descending: list[dict] - ) -> dict | None: - active_run_id = simulation.get("active_run_id") - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run - return runs_descending[0] if runs_descending else None - - def _upsert_simulation_spec_in_transaction( - self, tx, simulation: dict - ) -> SimulationSpec: - expected_spec = self.simulation_spec_service.build_simulation_spec(simulation) - existing_spec = parse_json_field(simulation.get("simulation_spec_json")) - if ( - existing_spec != expected_spec.model_dump() - or simulation.get("simulation_spec_schema_version") != 1 - ): - tx.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - expected_spec.model_dump_json(), - 1, - simulation["id"], - ), - ) - simulation["simulation_spec_json"] = expected_spec.model_dump() - simulation["simulation_spec_schema_version"] = 1 - - return expected_spec - - def _run_matches_parent( - self, - run: dict, - simulation: dict, - simulation_spec: SimulationSpec, - ) -> bool: - version_manifest = self._build_version_manifest(simulation) - return ( - run["status"] == simulation["status"] - and run.get("output") == simulation.get("output") - and run.get("error_message") == simulation.get("error_message") - and run.get("simulation_spec_snapshot_json") == simulation_spec.model_dump() - and run.get("country_package_version") - == version_manifest["country_package_version"] - and run.get("policyengine_version") - == version_manifest["policyengine_version"] - and run.get("data_version") == version_manifest["data_version"] - and run.get("runtime_app_name") == version_manifest["runtime_app_name"] - and run.get("simulation_cache_version") - == version_manifest["simulation_cache_version"] - ) - - def _insert_bootstrap_run( - self, tx, simulation: dict, simulation_spec: SimulationSpec - ) -> None: - version_manifest = self._build_version_manifest(simulation) - tx.query( - """ - INSERT INTO simulation_runs ( - id, simulation_id, report_output_run_id, input_position, run_sequence, - status, output, error_message, trigger_type, requested_at, started_at, - finished_at, source_run_id, simulation_spec_snapshot_json, - country_package_version, policyengine_version, data_version, - runtime_app_name, simulation_cache_version - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - str(uuid.uuid4()), - simulation["id"], - None, - None, - 1, - simulation["status"], - serialize_json_field(simulation.get("output")), - simulation.get("error_message"), - "initial", - None, - None, - None, - None, - simulation_spec.model_dump_json(), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["simulation_cache_version"], - ), - ) - - def _update_simulation_run_in_transaction( - self, - tx, - run_id: str, - simulation: dict, - simulation_spec: SimulationSpec, - ) -> None: - version_manifest = self._build_version_manifest(simulation) - tx.query( - """ - UPDATE simulation_runs - SET status = ?, output = ?, error_message = ?, - simulation_spec_snapshot_json = ?, country_package_version = ?, - policyengine_version = ?, data_version = ?, runtime_app_name = ?, - simulation_cache_version = ? - WHERE id = ? - """, - ( - simulation["status"], - serialize_json_field(simulation.get("output")), - simulation.get("error_message"), - simulation_spec.model_dump_json(), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["simulation_cache_version"], - run_id, - ), - ) - - def _sync_parent_pointers_in_transaction( - self, tx, simulation: dict, runs_descending: list[dict] - ) -> None: - desired_active_run_id, desired_latest_successful_run_id = ( - determine_parent_pointers(simulation["status"], runs_descending) - ) - if ( - simulation.get("active_run_id") == desired_active_run_id - and simulation.get("latest_successful_run_id") - == desired_latest_successful_run_id - ): - return - - tx.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ( - desired_active_run_id, - desired_latest_successful_run_id, - simulation["id"], - ), + ) -> dict: + session = getattr(session, "session", session) + return self.simulations.ensure_dual_write_state_in_session( + session, simulation_id, country_id ) - simulation["active_run_id"] = desired_active_run_id - simulation["latest_successful_run_id"] = desired_latest_successful_run_id - def _ensure_simulation_dual_write_state_in_transaction( + def _get_simulation_row( self, - tx, simulation_id: int, *, + queryer=None, country_id: str | None = None, - ) -> dict: - simulation = self._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - for_update=True, - ) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - simulation_spec = self._upsert_simulation_spec_in_transaction(tx, simulation) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx - ) - if not runs_descending: - self._insert_bootstrap_run(tx, simulation, simulation_spec) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx + for_update: bool = False, + ) -> dict | None: + del for_update + if queryer is not None and getattr(queryer, "session", None) is not None: + return self.simulations.get_in_session( + queryer.session, simulation_id, country_id ) - else: - mutable_run = self._select_mutable_run(simulation, runs_descending) - if mutable_run is not None and not self._run_matches_parent( - mutable_run, - simulation, - simulation_spec, - ): - self._update_simulation_run_in_transaction( - tx, - run_id=mutable_run["id"], - simulation=simulation, - simulation_spec=simulation_spec, - ) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx - ) - - self._sync_parent_pointers_in_transaction(tx, simulation, runs_descending) - refreshed_simulation = self._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - ) - if refreshed_simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found after sync") - return refreshed_simulation + return self.simulations.get(simulation_id, country_id) def ensure_simulation_dual_write_state( self, simulation_id: int, country_id: str | None = None ) -> dict: - return database.transaction( - lambda tx: self._ensure_simulation_dual_write_state_in_transaction( - tx, - simulation_id, - country_id=country_id, - ) - ) + return self.simulations.ensure_dual_write_state(simulation_id, country_id) def find_existing_simulation( self, @@ -317,34 +57,12 @@ def find_existing_simulation( population_type: str, policy_id: int, ) -> dict | None: - """ - Find an existing simulation with the same parameters. - - Args: - country_id (str): The country ID. - population_id (str): The population identifier (household or geography ID). - population_type (str): Type of population ('household' or 'geography'). - policy_id (int): The policy ID. - - Returns: - dict | None: The existing simulation data or None if not found. - """ - print("Checking for existing simulation") - - try: - existing_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - if existing_simulation is not None: - print(f"Found existing simulation with ID: {existing_simulation['id']}") - return existing_simulation - - except Exception as e: - print(f"Error checking for existing simulation. Details: {str(e)}") - raise e + return self.simulations.find_latest( + country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + ) def create_simulation( self, @@ -353,104 +71,22 @@ def create_simulation( population_type: str, policy_id: int, ) -> dict: - """ - Create a new simulation record with pending status. - - Args: - country_id (str): The country ID. - population_id (str): The population identifier (household or geography ID). - population_type (str): Type of population ('household' or 'geography'). - policy_id (int): The policy ID. - - Returns: - dict: The created simulation record. - """ - print("Creating new simulation") - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - try: - - def tx_callback(tx): - existing_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - queryer=tx, - ) - if existing_simulation is not None: - print( - f"Reusing existing simulation with ID: {existing_simulation['id']}" - ) - return self._ensure_simulation_dual_write_state_in_transaction( - tx, - existing_simulation["id"], - country_id=country_id, - ) - - tx.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - country_id, - api_version, - population_id, - population_type, - policy_id, - "pending", - ), - ) - - created_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - queryer=tx, - ) - if created_simulation is None: - raise Exception("Failed to retrieve created simulation") - - print(f"Created simulation with ID: {created_simulation['id']}") - return self._ensure_simulation_dual_write_state_in_transaction( - tx, - created_simulation["id"], - country_id=country_id, - ) - - return database.transaction(tx_callback) - - except Exception as e: - print(f"Error creating simulation. Details: {str(e)}") - raise e + return self.simulations.create_or_get_with_sync( + sync_callback=self._ensure_simulation_dual_write_state_in_transaction, + country_id=country_id, + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", + ) def get_simulation(self, country_id: str, simulation_id: int) -> dict | None: - """ - Get a simulation record by ID. - - Args: - country_id (str): The country ID. - simulation_id (int): The simulation ID. - - Returns: - dict | None: The simulation data or None if not found. - """ - print(f"Getting simulation {simulation_id}") - - try: - if type(simulation_id) is not int or simulation_id < 0: - raise Exception( - f"Invalid simulation ID: {simulation_id}. Must be a positive integer." - ) - - return self._get_simulation_row(simulation_id, country_id=country_id) - - except Exception as e: - print(f"Error fetching simulation #{simulation_id}. Details: {str(e)}") - raise e + if type(simulation_id) is not int or simulation_id < 0: + raise Exception( + f"Invalid simulation ID: {simulation_id}. Must be a positive integer." + ) + return self.simulations.get(simulation_id, country_id) def update_simulation( self, @@ -460,75 +96,24 @@ def update_simulation( output: str | None = None, error_message: str | None = None, ) -> bool: - """ - Update a simulation record with results or error. - - Args: - country_id (str): The country ID. - simulation_id (int): The simulation ID. - status (str | None): The new status ('complete' or 'error'). - output (str | None): The result output as JSON string (for complete status). - error_message (str | None): The error message (for error status). - - Returns: - bool: True if update was successful. - """ - print(f"Updating simulation {simulation_id}") - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - try: - update_fields = [] - update_values = [] - - if status is not None: - update_fields.append("status = ?") - update_values.append(status) - - if output is not None: - update_fields.append("output = ?") - update_values.append(output) - - if error_message is not None: - update_fields.append("error_message = ?") - update_values.append(error_message) - - # Only refresh api_version when the caller is actually - # changing one of the user-supplied fields above. The - # previous code appended api_version unconditionally, so - # the "no fields to update" guard below never fired and a - # PATCH with an empty body still touched the row. - if not update_fields: - print("No fields to update") - return False - - update_fields.append("api_version = ?") - update_values.append(api_version) - - def tx_callback(tx): - simulation = self._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - for_update=True, - ) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - tx.query( - f"UPDATE simulations SET {', '.join(update_fields)} WHERE id = ? AND country_id = ?", - (*update_values, simulation_id, country_id), - ) - self._ensure_simulation_dual_write_state_in_transaction( - tx, - simulation_id, - country_id=country_id, - ) - - database.transaction(tx_callback) - - print(f"Successfully updated simulation #{simulation_id}") - return True - - except Exception as e: - print(f"Error updating simulation #{simulation_id}. Details: {str(e)}") - raise e + values = { + key: value + for key, value in { + "status": status, + "output": output, + "error_message": error_message, + }.items() + if value is not None + } + if not values: + return False + if isinstance(values.get("output"), str): + values["output"] = json.loads(values["output"]) + values["api_version"] = COUNTRY_PACKAGE_VERSIONS.get(country_id) + self.simulations.update_with_sync( + simulation_id, + country_id, + values, + self._ensure_simulation_dual_write_state_in_transaction, + ) + return True From 914abf4dc66c442a7fa046b8a03092edbd24e476 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:42:39 +0300 Subject: [PATCH 24/89] test(red): prohibit direct SQL outside the data layer --- tests/unit/data/test_stage7_no_direct_sql.py | 23 ++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 tests/unit/data/test_stage7_no_direct_sql.py diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py new file mode 100644 index 000000000..daf09145a --- /dev/null +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -0,0 +1,23 @@ +from pathlib import Path + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +def test_runtime_sql_is_confined_to_the_data_access_layer(): + offenders = [] + for path in PACKAGE_ROOT.rglob("*.py"): + if "data" in path.relative_to(PACKAGE_ROOT).parts: + continue + source = path.read_text(encoding="utf-8") + if any( + token in source + for token in ( + "database.query(", + "local_database.query(", + "database.transaction(", + "local_database.transaction(", + ) + ): + offenders.append(str(path.relative_to(PACKAGE_ROOT))) + assert offenders == [] From 01924526b6c2eae9ef0e157de888328bf9288762 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:44:07 +0300 Subject: [PATCH 25/89] feat(green): confine runtime SQL to the DAO layer --- policyengine_api/country.py | 4 +- policyengine_api/data/orm.py | 18 +++--- policyengine_api/data/v1_daos.py | 13 ++++ .../endpoints/economy/reform_impact.py | 4 +- policyengine_api/endpoints/household.py | 42 +++++++----- policyengine_api/endpoints/policy.py | 64 +++++++++++-------- policyengine_api/endpoints/simulation.py | 14 ++-- .../services/ai_analysis_service.py | 2 +- .../services/reform_impacts_service.py | 2 +- .../services/tracer_analysis_service.py | 2 +- 10 files changed, 103 insertions(+), 62 deletions(-) diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 6c3acc6b9..42dca1d46 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -23,7 +23,7 @@ build_congressional_district_metadata, ) -from policyengine_api.data import local_database +from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, get_bundle_default_dataset_option, @@ -434,7 +434,7 @@ def calculate( if household_id is not None and policy_id is not None: # write to local database - local_database.query( + runtime_sqlalchemy_dao(local=True).query( """ INSERT INTO tracers (household_id, policy_id, country_id, api_version, tracer_output) VALUES (?, ?, ?, ?, ?) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 28d829445..f37a1d746 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -80,17 +80,19 @@ def build_sqlite_session_manager( return SessionManager(engine) -def build_v1_session_manager() -> SessionManager: +def build_v1_session_manager(*, local: bool = False) -> SessionManager: """Bind ORM sessions to the database selected by the v1 runtime.""" - from policyengine_api.data.data import database + from policyengine_api.data.data import database, local_database - if database.local: - if hasattr(database, "_connection"): - database._connection.row_factory = _IndexedMappingRow + selected_database = local_database if local else database + + if selected_database.local: + if hasattr(selected_database, "_connection"): + selected_database._connection.row_factory = _IndexedMappingRow engine = create_engine( "sqlite+pysqlite://", - creator=lambda: database._connection, + creator=lambda: selected_database._connection, poolclass=StaticPool, ) event.listen( @@ -101,5 +103,5 @@ def build_v1_session_manager() -> SessionManager: ), ) return SessionManager(engine) - return build_sqlite_session_manager(database.db_url) - return SessionManager(database.pool) + return build_sqlite_session_manager(selected_database.db_url) + return SessionManager(selected_database.pool) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 985ae2a32..573260b98 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -98,6 +98,19 @@ def session(self): return self._session +_runtime_sqlalchemy_daos: dict[bool, SQLAlchemyDAO] = {} + + +def runtime_sqlalchemy_dao(*, local: bool = False) -> SQLAlchemyDAO: + if local not in _runtime_sqlalchemy_daos: + from policyengine_api.data.orm import build_v1_session_manager + + _runtime_sqlalchemy_daos[local] = SQLAlchemyDAO( + build_v1_session_manager(local=local) + ) + return _runtime_sqlalchemy_daos[local] + + class PolicyDAO: def __init__(self, sessions: SessionManager): self.sessions = sessions diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py index 42795243d..a53d4a238 100644 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ b/policyengine_api/endpoints/economy/reform_impact.py @@ -1,4 +1,4 @@ -from policyengine_api.data import local_database +from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao def set_comment_on_job( @@ -17,7 +17,7 @@ def set_comment_on_job( "time_period = ? AND options_hash = ? AND dataset = ?" ) - local_database.query( + runtime_sqlalchemy_dao(local=True).query( query, ( comment, diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index 616da54e7..4eca0d9e1 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -1,4 +1,4 @@ -from policyengine_api.data import database, local_database +from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao import json from flask import Response, request from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS @@ -111,10 +111,14 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Look in computed_households to see if already computed - row = local_database.query( - "SELECT * FROM computed_household WHERE household_id = ? AND policy_id = ? AND api_version = ?", - (household_id, policy_id, api_version), - ).fetchone() + row = ( + runtime_sqlalchemy_dao(local=True) + .query( + "SELECT * FROM computed_household WHERE household_id = ? AND policy_id = ? AND api_version = ?", + (household_id, policy_id, api_version), + ) + .fetchone() + ) if row is not None: result = dict( @@ -135,10 +139,14 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the household table - row = database.query( - "SELECT * FROM household WHERE id = ? AND country_id = ?", - (household_id, country_id), - ).fetchone() + row = ( + runtime_sqlalchemy_dao() + .query( + "SELECT * FROM household WHERE id = ? AND country_id = ?", + (household_id, country_id), + ) + .fetchone() + ) if row is not None: household = dict(row) @@ -163,10 +171,14 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the policy table - row = database.query( - "SELECT * FROM policy WHERE id = ? AND country_id = ?", - (policy_id, country_id), - ).fetchone() + row = ( + runtime_sqlalchemy_dao() + .query( + "SELECT * FROM policy WHERE id = ? AND country_id = ?", + (policy_id, country_id), + ) + .fetchone() + ) if row is not None: policy = dict(row) @@ -213,7 +225,7 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Store the result in the computed_household table try: - local_database.query( + runtime_sqlalchemy_dao(local=True).query( "INSERT INTO computed_household (country_id, household_id, policy_id, computed_household_json, api_version) VALUES (?, ?, ?, ?, ?)", ( country_id, @@ -225,7 +237,7 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st ) except Exception: # Update the result if it already exists - local_database.query( + runtime_sqlalchemy_dao(local=True).query( "UPDATE computed_household SET computed_household_json = ? WHERE country_id = ? AND household_id = ? AND policy_id = ?", (json.dumps(result), country_id, household_id, policy_id), ) diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index f5d33e938..6df74931f 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -1,7 +1,5 @@ from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data import database -from policyengine_api.utils import hash_object -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao import json from flask import Response, request @@ -33,7 +31,7 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - results = database.query( + results = runtime_sqlalchemy_dao().query( "SELECT id, label, policy_hash FROM policy WHERE country_id = ? AND label LIKE ?", (country_id, f"%{query}%"), ) @@ -139,18 +137,22 @@ def set_user_policy(country_id: str) -> dict: # to be tested; type is not yet implemented try: - row = database.query( - f"SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? AND {nullable_key_string}", - ( - country_id, - reform_id, - baseline_id, - user_id, - year, - geography, - *not_null_values, - ), - ).fetchone() + row = ( + runtime_sqlalchemy_dao() + .query( + f"SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? AND {nullable_key_string}", + ( + country_id, + reform_id, + baseline_id, + user_id, + year, + geography, + *not_null_values, + ), + ) + .fetchone() + ) if row is not None: readable_row = dict(row) @@ -183,10 +185,10 @@ def set_user_policy(country_id: str) -> dict: "reform_id, baseline_label, baseline_id, user_id, year, " "geography, number_of_provisions, api_version, added_date, " "updated_date, budgetary_impact, type, dataset) VALUES " - f"(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" ) - database.query( + runtime_sqlalchemy_dao().query( query, ( country_id, @@ -220,10 +222,14 @@ def set_user_policy(country_id: str) -> dict: if dataset: params.append(dataset) - row = database.query( - query, - tuple(params), - ).fetchone() + row = ( + runtime_sqlalchemy_dao() + .query( + query, + tuple(params), + ) + .fetchone() + ) except Exception as e: return Response( @@ -271,10 +277,14 @@ def get_user_policy(country_id: str, user_id: str) -> dict: """ # Get the policy record for a given policy ID. - rows = database.query( - f"SELECT * FROM user_policies WHERE country_id = ? AND user_id = ?", - (country_id, user_id), - ).fetchall() + rows = ( + runtime_sqlalchemy_dao() + .query( + "SELECT * FROM user_policies WHERE country_id = ? AND user_id = ?", + (country_id, user_id), + ) + .fetchall() + ) rows_parsed = [ dict( @@ -396,7 +406,7 @@ def update_user_policy(country_id: str) -> dict: sql_request = f"UPDATE user_policies SET {setter_phrase} WHERE id = ?" try: - database.query(sql_request, (tuple(args))) + runtime_sqlalchemy_dao().query(sql_request, (tuple(args))) except Exception as e: return Response( json.dumps( diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index a0d9bd70d..c64e54b7c 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -1,4 +1,4 @@ -from policyengine_api.data import local_database +from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao """ @@ -42,10 +42,14 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - result = local_database.query( - "SELECT * FROM reform_impact ORDER BY start_time DESC LIMIT ?", - (max_results,), - ).fetchall() + result = ( + runtime_sqlalchemy_dao(local=True) + .query( + "SELECT * FROM reform_impact ORDER BY start_time DESC LIMIT ?", + (max_results,), + ) + .fetchall() + ) # Format into [{}] diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index 9ee6aca19..76869299f 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -30,7 +30,7 @@ def __init__(self, analyses: AnalysisDAO | None = None): @property def analyses(self) -> AnalysisDAO: if self._analyses is None: - self._analyses = AnalysisDAO(build_v1_session_manager()) + self._analyses = AnalysisDAO(build_v1_session_manager(local=True)) return self._analyses def get_existing_analysis(self, prompt: str) -> str | None: diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 46ebbceee..fc9ee9515 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -11,7 +11,7 @@ def __init__(self, impacts: ReformImpactDAO | None = None): @property def impacts(self) -> ReformImpactDAO: if self._impacts is None: - self._impacts = ReformImpactDAO(build_v1_session_manager()) + self._impacts = ReformImpactDAO(build_v1_session_manager(local=True)) return self._impacts @staticmethod diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index bf2ed4d31..e5484a5b1 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -21,7 +21,7 @@ def __init__( @property def tracers(self) -> TracerDAO: if self._tracers is None: - self._tracers = TracerDAO(build_v1_session_manager()) + self._tracers = TracerDAO(build_v1_session_manager(local=True)) return self._tracers def execute_analysis( From 33292c157e8198f81deb2b9ec573aca09116c571 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:45:17 +0300 Subject: [PATCH 26/89] test(red): define toy qualification and legacy removal gates --- tests/unit/test_stage7_toy_qualification.py | 30 +++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 tests/unit/test_stage7_toy_qualification.py diff --git a/tests/unit/test_stage7_toy_qualification.py b/tests/unit/test_stage7_toy_qualification.py new file mode 100644 index 000000000..ed00ee914 --- /dev/null +++ b/tests/unit/test_stage7_toy_qualification.py @@ -0,0 +1,30 @@ +from pathlib import Path + +from policyengine_api.scripts.qualify_stage7_toy import qualify_stage7_toy + + +def test_toy_qualification_exercises_migrated_data_paths(tmp_path: Path): + result = qualify_stage7_toy(f"sqlite+pysqlite:///{tmp_path / 'stage7-toy.db'}") + assert result == { + "alembic_head": True, + "policy": True, + "household": True, + "user": True, + "simulation": True, + "report": True, + "analysis": True, + "tracer": True, + "reform_impact": True, + } + + +def test_legacy_daos_have_been_removed(): + package = Path(__file__).parents[2] / "policyengine_api" + sources = "\n".join( + path.read_text(encoding="utf-8") for path in package.rglob("*.py") + ) + assert "LegacyPolicyDAO" not in sources + assert "LegacyHouseholdDAO" not in sources + assert "LegacyUserDAO" not in sources + assert "LegacySimulationDAO" not in sources + assert "LegacyReportDAO" not in sources From a900dbbda427bcfabe632cb13add1abf5395519e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:45:53 +0300 Subject: [PATCH 27/89] feat(green): add Stage 7 toy qualification gate --- policyengine_api/scripts/__init__.py | 1 + .../scripts/qualify_stage7_toy.py | 99 +++++++++++++++++++ 2 files changed, 100 insertions(+) create mode 100644 policyengine_api/scripts/__init__.py create mode 100644 policyengine_api/scripts/qualify_stage7_toy.py diff --git a/policyengine_api/scripts/__init__.py b/policyengine_api/scripts/__init__.py new file mode 100644 index 000000000..0af6e9d90 --- /dev/null +++ b/policyengine_api/scripts/__init__.py @@ -0,0 +1 @@ +"""Executable qualification helpers shipped with the API package.""" diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py new file mode 100644 index 000000000..a3f8eb9c2 --- /dev/null +++ b/policyengine_api/scripts/qualify_stage7_toy.py @@ -0,0 +1,99 @@ +"""Run the Stage 7 ORM boundary against an isolated toy database.""" + +from datetime import datetime + +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, inspect + +from policyengine_api.constants import REPO +from policyengine_api.data.orm import SessionManager +from policyengine_api.data.v1_daos import ( + AnalysisDAO, + HouseholdDAO, + PolicyDAO, + ReformImpactDAO, + ReportDAO, + SimulationDAO, + TracerDAO, + UserDAO, +) + + +def qualify_stage7_toy(database_url: str) -> dict[str, bool]: + """Upgrade and exercise every migrated v1 persistence domain.""" + + config = Config(str(REPO / "alembic.ini")) + config.set_main_option("sqlalchemy.url", database_url) + command.upgrade(config, "head") + + engine = create_engine(database_url) + sessions = SessionManager(engine) + policies = PolicyDAO(sessions) + households = HouseholdDAO(sessions) + users = UserDAO(sessions) + simulations = SimulationDAO(sessions) + reports = ReportDAO(sessions) + analyses = AnalysisDAO(sessions) + tracers = TracerDAO(sessions) + impacts = ReformImpactDAO(sessions) + + policy_id = policies.create("us", "Toy", {}, "toy-policy", "toy") + household_id = households.create("us", "Toy", {}, "toy-household", "toy") + user_id = users.create_profile("toy|user", "toy-user", "us", 1) + simulation_id = simulations.create( + country_id="us", + api_version="toy", + population_id=str(household_id), + population_type="household", + policy_id=policy_id, + ) + simulations.create_run( + simulation_id, + run_id="toy-simulation-run", + status="pending", + trigger_type="qualification", + ) + report_id = reports.create( + country_id="us", + simulation_1_id=simulation_id, + simulation_2_id=None, + api_version="toy", + year="2026", + ) + reports.create_run( + report_id, + run_id="toy-report-run", + status="pending", + trigger_type="qualification", + ) + analyses.store("toy prompt", "toy answer", "complete") + tracers.create(household_id, policy_id, "us", "toy", ["toy trace"]) + impact_id = impacts.create( + baseline_policy_id=policy_id, + reform_policy_id=policy_id, + country_id="us", + region="us", + dataset="default", + time_period="2026", + options_json={}, + options_hash="toy-options", + api_version="toy", + reform_impact_json={}, + status="computing", + start_time=datetime(2026, 1, 1), + execution_id="toy-impact", + ) + + return { + "alembic_head": "alembic_version" in inspect(engine).get_table_names(), + "policy": policies.get("us", policy_id) is not None, + "household": households.get("us", household_id) is not None, + "user": users.get_profile(user_id=user_id) is not None, + "simulation": simulations.get_run("toy-simulation-run") is not None, + "report": reports.get_run("toy-report-run") is not None, + "analysis": analyses.get("toy prompt") == "toy answer", + "tracer": tracers.get(household_id, policy_id, "us") is not None, + "reform_impact": impacts.find(execution_id="toy-impact")["reform_impact_id"] + == impact_id, + } From ae69c5cb79460e45bf3cc169762b432560ea90c1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:47:04 +0300 Subject: [PATCH 28/89] style: format Stage 7 migration files --- ...8cb8bcd717c_baseline_existing_v1_schema.py | 426 ++++++++++-------- tests/unit/data/test_alembic_baseline.py | 4 +- tests/unit/data/test_local_daos.py | 16 +- tests/unit/test_alembic_skill.py | 14 +- 4 files changed, 253 insertions(+), 207 deletions(-) diff --git a/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py b/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py index dfa019c5a..7115a44e0 100644 --- a/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py +++ b/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py @@ -1,7 +1,7 @@ """baseline existing v1 schema Revision ID: f8cb8bcd717c -Revises: +Revises: Create Date: 2026-08-06 17:20:10.119942 """ @@ -11,7 +11,7 @@ import sqlalchemy as sa -revision: str = 'f8cb8bcd717c' +revision: str = "f8cb8bcd717c" down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -19,217 +19,251 @@ def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - op.create_table('analysis', - sa.Column('prompt_id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('prompt', sa.Text(), nullable=False), - sa.Column('analysis', sa.Text(), nullable=True), - sa.Column('status', sa.String(length=32), nullable=False), - sa.PrimaryKeyConstraint('prompt_id') + op.create_table( + "analysis", + sa.Column("prompt_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("prompt", sa.Text(), nullable=False), + sa.Column("analysis", sa.Text(), nullable=True), + sa.Column("status", sa.String(length=32), nullable=False), + sa.PrimaryKeyConstraint("prompt_id"), ) - op.create_table('computed_household', - sa.Column('household_id', sa.Integer(), nullable=False), - sa.Column('policy_id', sa.Integer(), nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('computed_household_json', sa.JSON(), nullable=False), - sa.Column('status', sa.String(length=32), nullable=True), - sa.PrimaryKeyConstraint('household_id', 'policy_id', 'country_id') + op.create_table( + "computed_household", + sa.Column("household_id", sa.Integer(), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("computed_household_json", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=True), + sa.PrimaryKeyConstraint("household_id", "policy_id", "country_id"), ) - op.create_table('economy', - sa.Column('economy_id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('policy_id', sa.Integer(), nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('region', sa.String(length=32), nullable=True), - sa.Column('time_period', sa.String(length=32), nullable=True), - sa.Column('options_json', sa.JSON(), nullable=False), - sa.Column('options_hash', sa.String(length=255), nullable=False), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('economy_json', sa.JSON(), nullable=True), - sa.Column('status', sa.String(length=32), nullable=False), - sa.Column('message', sa.String(length=255), nullable=True), - sa.PrimaryKeyConstraint('economy_id') + op.create_table( + "economy", + sa.Column("economy_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("region", sa.String(length=32), nullable=True), + sa.Column("time_period", sa.String(length=32), nullable=True), + sa.Column("options_json", sa.JSON(), nullable=False), + sa.Column("options_hash", sa.String(length=255), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("economy_json", sa.JSON(), nullable=True), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("message", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("economy_id"), ) - op.create_table('household', - sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('label', sa.String(length=255), nullable=True), - sa.Column('api_version', sa.String(length=255), nullable=False), - sa.Column('household_json', sa.JSON(), nullable=False), - sa.Column('household_hash', sa.String(length=255), nullable=False), - sa.PrimaryKeyConstraint('id') + op.create_table( + "household", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("label", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=255), nullable=False), + sa.Column("household_json", sa.JSON(), nullable=False), + sa.Column("household_hash", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("id"), ) - op.create_table('legacy_report_output_aliases', - sa.Column('legacy_report_output_id', sa.Integer(), nullable=False), - sa.Column('canonical_report_output_id', sa.Integer(), nullable=False), - sa.PrimaryKeyConstraint('legacy_report_output_id') + op.create_table( + "legacy_report_output_aliases", + sa.Column("legacy_report_output_id", sa.Integer(), nullable=False), + sa.Column("canonical_report_output_id", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("legacy_report_output_id"), ) - op.create_table('policy', - # Review correction: MySQL's legacy composite key auto-increments ``id``; - # SQLite cannot compile autoincrement on a composite primary key. - sa.Column('id', sa.Integer(), autoincrement=op.get_bind().dialect.name != 'sqlite', nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('label', sa.String(length=255), nullable=True), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('policy_json', sa.JSON(), nullable=False), - sa.Column('policy_hash', sa.String(length=255), nullable=False), - sa.PrimaryKeyConstraint('id', 'country_id', 'policy_hash') + op.create_table( + "policy", + # Review correction: MySQL's legacy composite key auto-increments ``id``; + # SQLite cannot compile autoincrement on a composite primary key. + sa.Column( + "id", + sa.Integer(), + autoincrement=op.get_bind().dialect.name != "sqlite", + nullable=False, + ), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("label", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("policy_json", sa.JSON(), nullable=False), + sa.Column("policy_hash", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("id", "country_id", "policy_hash"), ) - op.create_table('reform_impact', - sa.Column('reform_impact_id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('baseline_policy_id', sa.Integer(), nullable=False), - sa.Column('reform_policy_id', sa.Integer(), nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('region', sa.String(length=32), nullable=False), - sa.Column('dataset', sa.String(length=255), nullable=False), - sa.Column('time_period', sa.String(length=32), nullable=False), - sa.Column('options_json', sa.JSON(), nullable=True), - sa.Column('options_hash', sa.String(length=255), nullable=True), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('reform_impact_json', sa.JSON(), nullable=False), - sa.Column('status', sa.String(length=32), nullable=False), - sa.Column('message', sa.String(length=255), nullable=True), - sa.Column('start_time', sa.DateTime(), nullable=True), - sa.Column('end_time', sa.DateTime(), nullable=True), - sa.Column('execution_id', sa.String(length=255), nullable=False), - sa.PrimaryKeyConstraint('reform_impact_id') + op.create_table( + "reform_impact", + sa.Column("reform_impact_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("baseline_policy_id", sa.Integer(), nullable=False), + sa.Column("reform_policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("region", sa.String(length=32), nullable=False), + sa.Column("dataset", sa.String(length=255), nullable=False), + sa.Column("time_period", sa.String(length=32), nullable=False), + sa.Column("options_json", sa.JSON(), nullable=True), + sa.Column("options_hash", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("reform_impact_json", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("message", sa.String(length=255), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column("execution_id", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("reform_impact_id"), ) - op.create_table('report_output_runs', - sa.Column('id', sa.String(length=36), nullable=False), - sa.Column('report_output_id', sa.Integer(), nullable=False), - sa.Column('run_sequence', sa.Integer(), nullable=False), - sa.Column('status', sa.String(length=32), nullable=False), - sa.Column('output', sa.JSON(), nullable=True), - sa.Column('error_message', sa.Text(), nullable=True), - sa.Column('trigger_type', sa.String(length=32), nullable=False), - sa.Column('requested_at', sa.DateTime(), nullable=True), - sa.Column('started_at', sa.DateTime(), nullable=True), - sa.Column('finished_at', sa.DateTime(), nullable=True), - sa.Column('source_run_id', sa.String(length=36), nullable=True), - sa.Column('report_spec_snapshot_json', sa.JSON(), nullable=True), - sa.Column('country_package_version', sa.String(length=255), nullable=True), - sa.Column('policyengine_version', sa.String(length=255), nullable=True), - sa.Column('data_version', sa.String(length=255), nullable=True), - sa.Column('runtime_app_name', sa.String(length=255), nullable=True), - sa.Column('report_cache_version', sa.String(length=255), nullable=True), - sa.Column('simulation_cache_version', sa.String(length=255), nullable=True), - sa.Column('requested_version_override', sa.String(length=255), nullable=True), - sa.Column('resolved_dataset', sa.String(length=255), nullable=True), - sa.Column('resolved_options_hash', sa.String(length=255), nullable=True), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('report_output_id', 'run_sequence') + op.create_table( + "report_output_runs", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("report_output_id", sa.Integer(), nullable=False), + sa.Column("run_sequence", sa.Integer(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("trigger_type", sa.String(length=32), nullable=False), + sa.Column("requested_at", sa.DateTime(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("source_run_id", sa.String(length=36), nullable=True), + sa.Column("report_spec_snapshot_json", sa.JSON(), nullable=True), + sa.Column("country_package_version", sa.String(length=255), nullable=True), + sa.Column("policyengine_version", sa.String(length=255), nullable=True), + sa.Column("data_version", sa.String(length=255), nullable=True), + sa.Column("runtime_app_name", sa.String(length=255), nullable=True), + sa.Column("report_cache_version", sa.String(length=255), nullable=True), + sa.Column("simulation_cache_version", sa.String(length=255), nullable=True), + sa.Column("requested_version_override", sa.String(length=255), nullable=True), + sa.Column("resolved_dataset", sa.String(length=255), nullable=True), + sa.Column("resolved_options_hash", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("report_output_id", "run_sequence"), ) - op.create_table('report_outputs', - sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('simulation_1_id', sa.Integer(), nullable=False), - sa.Column('simulation_2_id', sa.Integer(), nullable=True), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('status', sa.String(length=32), server_default=sa.text("'pending'"), nullable=False), - sa.Column('output', sa.JSON(), nullable=True), - sa.Column('error_message', sa.Text(), nullable=True), - sa.Column('year', sa.String(length=255), server_default=sa.text("'2025'"), nullable=True), - sa.Column('report_kind', sa.String(length=64), nullable=True), - sa.Column('report_spec_json', sa.JSON(), nullable=True), - sa.Column('report_spec_schema_version', sa.Integer(), nullable=True), - sa.Column('report_spec_status', sa.String(length=32), nullable=True), - sa.Column('active_run_id', sa.String(length=36), nullable=True), - sa.Column('latest_successful_run_id', sa.String(length=36), nullable=True), - sa.PrimaryKeyConstraint('id') + op.create_table( + "report_outputs", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("simulation_1_id", sa.Integer(), nullable=False), + sa.Column("simulation_2_id", sa.Integer(), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column( + "status", + sa.String(length=32), + server_default=sa.text("'pending'"), + nullable=False, + ), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "year", + sa.String(length=255), + server_default=sa.text("'2025'"), + nullable=True, + ), + sa.Column("report_kind", sa.String(length=64), nullable=True), + sa.Column("report_spec_json", sa.JSON(), nullable=True), + sa.Column("report_spec_schema_version", sa.Integer(), nullable=True), + sa.Column("report_spec_status", sa.String(length=32), nullable=True), + sa.Column("active_run_id", sa.String(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.String(length=36), nullable=True), + sa.PrimaryKeyConstraint("id"), ) - op.create_table('simulation_runs', - sa.Column('id', sa.String(length=36), nullable=False), - sa.Column('simulation_id', sa.Integer(), nullable=False), - sa.Column('report_output_run_id', sa.String(length=36), nullable=True), - sa.Column('input_position', sa.Integer(), nullable=True), - sa.Column('run_sequence', sa.Integer(), nullable=False), - sa.Column('status', sa.String(length=32), nullable=False), - sa.Column('output', sa.JSON(), nullable=True), - sa.Column('error_message', sa.Text(), nullable=True), - sa.Column('trigger_type', sa.String(length=32), nullable=False), - sa.Column('requested_at', sa.DateTime(), nullable=True), - sa.Column('started_at', sa.DateTime(), nullable=True), - sa.Column('finished_at', sa.DateTime(), nullable=True), - sa.Column('source_run_id', sa.String(length=36), nullable=True), - sa.Column('simulation_spec_snapshot_json', sa.JSON(), nullable=True), - sa.Column('country_package_version', sa.String(length=255), nullable=True), - sa.Column('policyengine_version', sa.String(length=255), nullable=True), - sa.Column('data_version', sa.String(length=255), nullable=True), - sa.Column('runtime_app_name', sa.String(length=255), nullable=True), - sa.Column('simulation_cache_version', sa.String(length=255), nullable=True), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('simulation_id', 'run_sequence') + op.create_table( + "simulation_runs", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("simulation_id", sa.Integer(), nullable=False), + sa.Column("report_output_run_id", sa.String(length=36), nullable=True), + sa.Column("input_position", sa.Integer(), nullable=True), + sa.Column("run_sequence", sa.Integer(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("trigger_type", sa.String(length=32), nullable=False), + sa.Column("requested_at", sa.DateTime(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("source_run_id", sa.String(length=36), nullable=True), + sa.Column("simulation_spec_snapshot_json", sa.JSON(), nullable=True), + sa.Column("country_package_version", sa.String(length=255), nullable=True), + sa.Column("policyengine_version", sa.String(length=255), nullable=True), + sa.Column("data_version", sa.String(length=255), nullable=True), + sa.Column("runtime_app_name", sa.String(length=255), nullable=True), + sa.Column("simulation_cache_version", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("simulation_id", "run_sequence"), ) - op.create_table('simulations', - sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('population_id', sa.String(length=255), nullable=False), - sa.Column('population_type', sa.String(length=50), nullable=False), - sa.Column('policy_id', sa.Integer(), nullable=False), - sa.Column('status', sa.String(length=32), server_default=sa.text("'pending'"), nullable=False), - sa.Column('output', sa.JSON(), nullable=True), - sa.Column('error_message', sa.Text(), nullable=True), - sa.Column('simulation_spec_json', sa.JSON(), nullable=True), - sa.Column('simulation_spec_schema_version', sa.Integer(), nullable=True), - sa.Column('active_run_id', sa.String(length=36), nullable=True), - sa.Column('latest_successful_run_id', sa.String(length=36), nullable=True), - sa.PrimaryKeyConstraint('id') + op.create_table( + "simulations", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("population_id", sa.String(length=255), nullable=False), + sa.Column("population_type", sa.String(length=50), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column( + "status", + sa.String(length=32), + server_default=sa.text("'pending'"), + nullable=False, + ), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("simulation_spec_json", sa.JSON(), nullable=True), + sa.Column("simulation_spec_schema_version", sa.Integer(), nullable=True), + sa.Column("active_run_id", sa.String(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.String(length=36), nullable=True), + sa.PrimaryKeyConstraint("id"), ) - op.create_table('tracers', - sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('household_id', sa.Integer(), nullable=False), - sa.Column('policy_id', sa.Integer(), nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('api_version', sa.String(length=10), nullable=False), - sa.Column('tracer_output', sa.JSON(), nullable=False), - sa.PrimaryKeyConstraint('id') + op.create_table( + "tracers", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("household_id", sa.Integer(), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("tracer_output", sa.JSON(), nullable=False), + sa.PrimaryKeyConstraint("id"), ) - op.create_table('user_policies', - sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('country_id', sa.String(length=3), nullable=False), - sa.Column('reform_id', sa.Integer(), nullable=False), - sa.Column('reform_label', sa.String(length=255), nullable=True), - sa.Column('baseline_id', sa.Integer(), nullable=False), - sa.Column('baseline_label', sa.String(length=255), nullable=True), - sa.Column('user_id', sa.String(length=255), nullable=False), - sa.Column('year', sa.String(length=32), nullable=False), - sa.Column('geography', sa.String(length=255), nullable=False), - sa.Column('dataset', sa.String(length=255), nullable=True), - sa.Column('number_of_provisions', sa.Integer(), nullable=False), - sa.Column('api_version', sa.String(length=32), nullable=False), - sa.Column('added_date', sa.BigInteger(), nullable=False), - sa.Column('updated_date', sa.BigInteger(), nullable=False), - sa.Column('budgetary_impact', sa.String(length=255), nullable=True), - sa.Column('type', sa.String(length=255), nullable=True), - sa.PrimaryKeyConstraint('id') + op.create_table( + "user_policies", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("reform_id", sa.Integer(), nullable=False), + sa.Column("reform_label", sa.String(length=255), nullable=True), + sa.Column("baseline_id", sa.Integer(), nullable=False), + sa.Column("baseline_label", sa.String(length=255), nullable=True), + sa.Column("user_id", sa.String(length=255), nullable=False), + sa.Column("year", sa.String(length=32), nullable=False), + sa.Column("geography", sa.String(length=255), nullable=False), + sa.Column("dataset", sa.String(length=255), nullable=True), + sa.Column("number_of_provisions", sa.Integer(), nullable=False), + sa.Column("api_version", sa.String(length=32), nullable=False), + sa.Column("added_date", sa.BigInteger(), nullable=False), + sa.Column("updated_date", sa.BigInteger(), nullable=False), + sa.Column("budgetary_impact", sa.String(length=255), nullable=True), + sa.Column("type", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), ) - op.create_table('user_profiles', - sa.Column('user_id', sa.Integer(), autoincrement=True, nullable=False), - sa.Column('auth0_id', sa.String(length=255), nullable=False), - sa.Column('username', sa.String(length=255), nullable=True), - sa.Column('primary_country', sa.String(length=3), nullable=False), - sa.Column('user_since', sa.BigInteger(), nullable=False), - sa.PrimaryKeyConstraint('user_id'), - sa.UniqueConstraint('auth0_id'), - sa.UniqueConstraint('username') + op.create_table( + "user_profiles", + sa.Column("user_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("auth0_id", sa.String(length=255), nullable=False), + sa.Column("username", sa.String(length=255), nullable=True), + sa.Column("primary_country", sa.String(length=3), nullable=False), + sa.Column("user_since", sa.BigInteger(), nullable=False), + sa.PrimaryKeyConstraint("user_id"), + sa.UniqueConstraint("auth0_id"), + sa.UniqueConstraint("username"), ) # ### end Alembic commands ### def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - op.drop_table('user_profiles') - op.drop_table('user_policies') - op.drop_table('tracers') - op.drop_table('simulations') - op.drop_table('simulation_runs') - op.drop_table('report_outputs') - op.drop_table('report_output_runs') - op.drop_table('reform_impact') - op.drop_table('policy') - op.drop_table('legacy_report_output_aliases') - op.drop_table('household') - op.drop_table('economy') - op.drop_table('computed_household') - op.drop_table('analysis') + op.drop_table("user_profiles") + op.drop_table("user_policies") + op.drop_table("tracers") + op.drop_table("simulations") + op.drop_table("simulation_runs") + op.drop_table("report_outputs") + op.drop_table("report_output_runs") + op.drop_table("reform_impact") + op.drop_table("policy") + op.drop_table("legacy_report_output_aliases") + op.drop_table("household") + op.drop_table("economy") + op.drop_table("computed_household") + op.drop_table("analysis") # ### end Alembic commands ### diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index 15cae49c7..6bce1e892 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -18,7 +18,9 @@ def test_baseline_upgrades_fresh_database_to_v1_head(tmp_path: Path): database_path = tmp_path / "fresh.db" command.upgrade(_config(f"sqlite+pysqlite:///{database_path}"), "head") - tables = set(inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names()) + tables = set( + inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names() + ) assert set(V1Base.metadata.tables) <= tables assert "alembic_version" in tables diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py index e49553cf9..daf4bf604 100644 --- a/tests/unit/data/test_local_daos.py +++ b/tests/unit/data/test_local_daos.py @@ -20,10 +20,18 @@ def test_analysis_dao_round_trip(): def test_reform_impact_dao_transitions_by_execution_id(): _, impacts, _ = _daos() impacts.create( - country_id="us", reform_policy_id=2, baseline_policy_id=1, - region="us", dataset="default", time_period="2026", - options_json={}, options_hash="hash", api_version="1", - reform_impact_json={}, status="computing", start_time=datetime(2026, 1, 1), + country_id="us", + reform_policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options_json={}, + options_hash="hash", + api_version="1", + reform_impact_json={}, + status="computing", + start_time=datetime(2026, 1, 1), execution_id="job", ) impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) diff --git a/tests/unit/test_alembic_skill.py b/tests/unit/test_alembic_skill.py index df9609d1a..af743778c 100644 --- a/tests/unit/test_alembic_skill.py +++ b/tests/unit/test_alembic_skill.py @@ -6,13 +6,15 @@ def test_model_agnostic_alembic_skill_is_discoverable(): assert SKILL.exists() - assert "alembic-migrations.md" in ( - REPO / "docs" / "engineering" / "skills" / "README.md" - ).read_text() + assert ( + "alembic-migrations.md" + in (REPO / "docs" / "engineering" / "skills" / "README.md").read_text() + ) for adapter in ("AGENTS.md", "CLAUDE.md", ".github/copilot-instructions.md"): - assert "docs/engineering/skills/alembic-migrations.md" in ( - REPO / adapter - ).read_text() + assert ( + "docs/engineering/skills/alembic-migrations.md" + in (REPO / adapter).read_text() + ) def test_alembic_skill_forbids_handwritten_ai_revisions(): From 5e3790081f5b1c6aab7fc6d50910fc1aaa446d0c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:48:24 +0300 Subject: [PATCH 29/89] test: patch the Stage 7 policy DAO boundary --- tests/contract/test_v1_route_contracts.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 6b3407bbe..1e382b4a5 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -224,8 +224,10 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.database.query", - return_value=_policy_search_rows(), + "policyengine_api.endpoints.policy.runtime_sqlalchemy_dao", + return_value=SimpleNamespace( + query=lambda *args, **kwargs: _policy_search_rows() + ), ) ) stack.enter_context( From 4e89db73327e7ce482e5aa711ad112148a10d62d Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:50:45 +0300 Subject: [PATCH 30/89] test(red): require read-only schema comparison --- tests/unit/test_stage7_toy_qualification.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_stage7_toy_qualification.py b/tests/unit/test_stage7_toy_qualification.py index ed00ee914..884d7595c 100644 --- a/tests/unit/test_stage7_toy_qualification.py +++ b/tests/unit/test_stage7_toy_qualification.py @@ -1,6 +1,9 @@ from pathlib import Path -from policyengine_api.scripts.qualify_stage7_toy import qualify_stage7_toy +from policyengine_api.scripts.qualify_stage7_toy import ( + compare_stage7_schema, + qualify_stage7_toy, +) def test_toy_qualification_exercises_migrated_data_paths(tmp_path: Path): @@ -28,3 +31,11 @@ def test_legacy_daos_have_been_removed(): assert "LegacyUserDAO" not in sources assert "LegacySimulationDAO" not in sources assert "LegacyReportDAO" not in sources + + +def test_schema_comparison_is_read_only_and_reports_no_fresh_database_drift( + tmp_path: Path, +): + database_url = f"sqlite+pysqlite:///{tmp_path / 'comparison.db'}" + qualify_stage7_toy(database_url) + assert compare_stage7_schema(database_url) == [] From dc74ddafb8a35a532596acd4f16601cf4679df1a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 18:51:00 +0300 Subject: [PATCH 31/89] feat(green): add read-only Stage 7 schema comparison --- policyengine_api/scripts/qualify_stage7_toy.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py index a3f8eb9c2..9fe7b7081 100644 --- a/policyengine_api/scripts/qualify_stage7_toy.py +++ b/policyengine_api/scripts/qualify_stage7_toy.py @@ -3,7 +3,9 @@ from datetime import datetime from alembic import command +from alembic.autogenerate import compare_metadata from alembic.config import Config +from alembic.migration import MigrationContext from sqlalchemy import create_engine, inspect from policyengine_api.constants import REPO @@ -18,6 +20,19 @@ TracerDAO, UserDAO, ) +from policyengine_api.data.v1_models import V1Base + + +def compare_stage7_schema(database_url: str) -> list: + """Return metadata drift without stamping or mutating the target database.""" + + engine = create_engine(database_url) + with engine.connect() as connection: + context = MigrationContext.configure( + connection, + opts={"compare_type": True}, + ) + return compare_metadata(context, V1Base.metadata) def qualify_stage7_toy(database_url: str) -> dict[str, bool]: From bada37b6f4055d0f49c703d377f80c0a1cf925aa Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 22:36:56 +0300 Subject: [PATCH 32/89] refactor: let SQLAlchemy own Cloud SQL connections --- policyengine_api/data/data.py | 73 +++++++++++---------------- tests/unit/data/test_sqlalchemy_v2.py | 37 ++++++++++---- 2 files changed, 57 insertions(+), 53 deletions(-) diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py index 83f7eb820..31a0196e2 100644 --- a/policyengine_api/data/data.py +++ b/policyengine_api/data/data.py @@ -5,9 +5,8 @@ from pathlib import Path from dotenv import load_dotenv import json -from google.cloud.sql.connector import Connector +from google.cloud.sql.connector import Connector, IPTypes import sqlalchemy -import sqlalchemy.exc import os import sys @@ -128,21 +127,35 @@ def __init__( def _create_pool(self): db_config = get_remote_database_config() - self.connector = Connector() + ip_type = ( + IPTypes.PRIVATE + if os.environ.get("POLICYENGINE_DB_PRIVATE_IP", "").lower() + in {"1", "true", "yes"} + else IPTypes.PUBLIC + ) + self.connector = Connector(ip_type=ip_type, refresh_strategy="LAZY") db_pass = os.environ["POLICYENGINE_DB_PASSWORD"] if db_pass == ".dbpw": with open(".dbpw") as f: db_pass = f.read().strip() - conn = self.connector.connect( - instance_connection_string=db_config["instance_connection_name"], - driver="pymysql", - db=db_config["db_name"], - user=db_config["db_user"], - password=db_pass, - ) + + def getconn(): + return self.connector.connect( + instance_connection_string=db_config["instance_connection_name"], + driver="pymysql", + db=db_config["db_name"], + user=db_config["db_user"], + password=db_pass, + ) + self.pool = sqlalchemy.create_engine( "mysql+pymysql://", - creator=lambda: conn, + creator=getconn, + pool_pre_ping=True, + pool_recycle=int(os.environ.get("POLICYENGINE_DB_POOL_RECYCLE", "1800")), + pool_size=int(os.environ.get("POLICYENGINE_DB_POOL_SIZE", "5")), + max_overflow=int(os.environ.get("POLICYENGINE_DB_MAX_OVERFLOW", "2")), + pool_timeout=int(os.environ.get("POLICYENGINE_DB_POOL_TIMEOUT", "30")), ) def _close_pool(self): @@ -157,28 +170,20 @@ def _execute_remote(self, query_args): SQLAlchemy v2 connection-based execution.""" main_query = query_args[0] params = query_args[1] if len(query_args) > 1 else None - with self.pool.connect() as conn: + with self.pool.begin() as conn: if params is not None: result = conn.exec_driver_sql(main_query, params) else: result = conn.exec_driver_sql(main_query) - conn.commit() # Return a lightweight wrapper that holds # the fetched results so they survive the # connection context closing return _ResultProxy(result) def _execute_remote_transaction(self, callback): - with self.pool.connect() as conn: - transaction = conn.begin() + with self.pool.begin() as conn: proxy = _TransactionProxy(conn, local=False) - try: - result = callback(proxy) - transaction.commit() - return result - except Exception: - transaction.rollback() - raise + return callback(proxy) def query(self, *query): if self.local: @@ -191,19 +196,7 @@ def query(self, *query): main_query = query[0] main_query = main_query.replace("?", "%s") query[0] = main_query - try: - return self._execute_remote(query) - # Except InterfaceError and OperationalError, which are thrown when the connection is lost. - except ( - sqlalchemy.exc.InterfaceError, - sqlalchemy.exc.OperationalError, - ): - try: - self._close_pool() - self._create_pool() - return self._execute_remote(query) - except Exception as e: - raise e + return self._execute_remote(query) def transaction(self, callback): if self.local: @@ -225,15 +218,7 @@ def transaction(self, callback): if owns_connection: connection.close() - try: - return self._execute_remote_transaction(callback) - except ( - sqlalchemy.exc.InterfaceError, - sqlalchemy.exc.OperationalError, - ): - self._close_pool() - self._create_pool() - return self._execute_remote_transaction(callback) + return self._execute_remote_transaction(callback) def initialize(self): """ diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 1f380b9df..9566d7b14 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -185,22 +185,23 @@ class TestRemotePoolSetup: """Test remote pool setup without opening a real Cloud SQL connection.""" def _stub_remote_pool(self, monkeypatch): - fake_connection = object() connector_calls = [] engine_calls = [] class FakeConnector: + def __init__(self, **kwargs): + self.options = kwargs + def connect(self, **kwargs): connector_calls.append(kwargs) - return fake_connection + return object() - def fake_create_engine(url, creator): - engine_calls.append((url, creator)) - assert creator() is fake_connection + def fake_create_engine(url, **kwargs): + engine_calls.append((url, kwargs)) return "fake-engine" - fake_connector = FakeConnector() - monkeypatch.setattr(data_module, "Connector", lambda: fake_connector) + fake_connector = FakeConnector(refresh_strategy="LAZY") + monkeypatch.setattr(data_module, "Connector", lambda **_: fake_connector) monkeypatch.setattr(data_module.sqlalchemy, "create_engine", fake_create_engine) return fake_connector, connector_calls, engine_calls @@ -221,6 +222,11 @@ def test_create_pool_uses_remote_database_config(self, monkeypatch): assert db.connector is fake_connector assert db.pool == "fake-engine" + assert connector_calls == [] + creator = engine_calls[0][1]["creator"] + first = creator() + second = creator() + assert first is not second assert connector_calls == [ { "instance_connection_string": "test-project:us-central1:test-db", @@ -228,12 +234,24 @@ def test_create_pool_uses_remote_database_config(self, monkeypatch): "db": "test-db", "user": "test-user", "password": "test-password", - } + }, + { + "instance_connection_string": "test-project:us-central1:test-db", + "driver": "pymysql", + "db": "test-db", + "user": "test-user", + "password": "test-password", + }, ] assert engine_calls[0][0] == "mysql+pymysql://" + assert engine_calls[0][1]["pool_pre_ping"] is True + assert engine_calls[0][1]["pool_recycle"] == 1800 + assert engine_calls[0][1]["pool_size"] == 5 + assert engine_calls[0][1]["max_overflow"] == 2 + assert engine_calls[0][1]["pool_timeout"] == 30 def test_create_pool_reads_dot_dbpw_file(self, monkeypatch, tmp_path): - _, connector_calls, _ = self._stub_remote_pool(monkeypatch) + _, connector_calls, engine_calls = self._stub_remote_pool(monkeypatch) monkeypatch.chdir(tmp_path) monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", ".dbpw") (tmp_path / ".dbpw").write_text("file-password\n") @@ -241,6 +259,7 @@ def test_create_pool_reads_dot_dbpw_file(self, monkeypatch, tmp_path): db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) db._create_pool() + engine_calls[0][1]["creator"]() assert connector_calls[0]["password"] == "file-password" def test_remote_constructor_initializes_pool_without_local_database( From 8efd04291da8e44bc93f9778c16029cb4e8ba06c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 6 Aug 2026 22:42:36 +0300 Subject: [PATCH 33/89] fix: make Alembic baseline match the v1 schema --- migrations/env.py | 2 + ...fc2a547a4e_baseline_existing_v1_schema.py} | 85 ++++---- policyengine_api/data/initialise.sql | 3 +- policyengine_api/data/v1_models.py | 56 +++-- policyengine_api/scripts/stage7_database.py | 18 ++ tests/fixtures/stage7_pre_alembic_schema.sql | 194 ++++++++++++++++++ tests/integration/stage7_mysql.py | 86 ++++++++ .../test_stage7_existing_schema.py | 68 ++++++ tests/unit/data/sqlite_schema.py | 12 ++ tests/unit/data/test_local_daos.py | 4 +- tests/unit/data/test_run_daos.py | 4 +- tests/unit/data/test_v1_daos.py | 4 +- tests/unit/data/test_v1_models.py | 25 ++- 13 files changed, 486 insertions(+), 75 deletions(-) rename migrations/versions/{f8cb8bcd717c_baseline_existing_v1_schema.py => eafc2a547a4e_baseline_existing_v1_schema.py} (82%) create mode 100644 policyengine_api/scripts/stage7_database.py create mode 100644 tests/fixtures/stage7_pre_alembic_schema.sql create mode 100644 tests/integration/stage7_mysql.py create mode 100644 tests/integration/test_stage7_existing_schema.py create mode 100644 tests/unit/data/sqlite_schema.py diff --git a/migrations/env.py b/migrations/env.py index 85aab3a36..ceafe5835 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -26,6 +26,7 @@ def run_migrations_offline() -> None: target_metadata=target_metadata, literal_binds=True, compare_type=True, + compare_server_default=True, ) with context.begin_transaction(): context.run_migrations() @@ -42,6 +43,7 @@ def run_migrations_online() -> None: connection=connection, target_metadata=target_metadata, compare_type=True, + compare_server_default=True, ) with context.begin_transaction(): context.run_migrations() diff --git a/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py b/migrations/versions/eafc2a547a4e_baseline_existing_v1_schema.py similarity index 82% rename from migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py rename to migrations/versions/eafc2a547a4e_baseline_existing_v1_schema.py index 7115a44e0..3a744c562 100644 --- a/migrations/versions/f8cb8bcd717c_baseline_existing_v1_schema.py +++ b/migrations/versions/eafc2a547a4e_baseline_existing_v1_schema.py @@ -1,17 +1,17 @@ """baseline existing v1 schema -Revision ID: f8cb8bcd717c +Revision ID: eafc2a547a4e Revises: -Create Date: 2026-08-06 17:20:10.119942 +Create Date: 2026-08-06 22:40:07.810466 """ from typing import Sequence, Union from alembic import op import sqlalchemy as sa +from sqlalchemy.dialects import mysql - -revision: str = "f8cb8bcd717c" +revision: str = "eafc2a547a4e" down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -22,8 +22,12 @@ def upgrade() -> None: op.create_table( "analysis", sa.Column("prompt_id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("prompt", sa.Text(), nullable=False), - sa.Column("analysis", sa.Text(), nullable=True), + sa.Column( + "prompt", sa.Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=False + ), + sa.Column( + "analysis", sa.Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=True + ), sa.Column("status", sa.String(length=32), nullable=False), sa.PrimaryKeyConstraint("prompt_id"), ) @@ -47,7 +51,7 @@ def upgrade() -> None: sa.Column("options_json", sa.JSON(), nullable=False), sa.Column("options_hash", sa.String(length=255), nullable=False), sa.Column("api_version", sa.String(length=10), nullable=False), - sa.Column("economy_json", sa.JSON(), nullable=True), + sa.Column("economy_json", sa.JSON(none_as_null=True), nullable=True), sa.Column("status", sa.String(length=32), nullable=False), sa.Column("message", sa.String(length=255), nullable=True), sa.PrimaryKeyConstraint("economy_id"), @@ -64,20 +68,15 @@ def upgrade() -> None: ) op.create_table( "legacy_report_output_aliases", - sa.Column("legacy_report_output_id", sa.Integer(), nullable=False), + sa.Column( + "legacy_report_output_id", sa.Integer(), autoincrement=False, nullable=False + ), sa.Column("canonical_report_output_id", sa.Integer(), nullable=False), sa.PrimaryKeyConstraint("legacy_report_output_id"), ) op.create_table( "policy", - # Review correction: MySQL's legacy composite key auto-increments ``id``; - # SQLite cannot compile autoincrement on a composite primary key. - sa.Column( - "id", - sa.Integer(), - autoincrement=op.get_bind().dialect.name != "sqlite", - nullable=False, - ), + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), sa.Column("country_id", sa.String(length=3), nullable=False), sa.Column("label", sa.String(length=255), nullable=True), sa.Column("api_version", sa.String(length=10), nullable=False), @@ -94,7 +93,7 @@ def upgrade() -> None: sa.Column("region", sa.String(length=32), nullable=False), sa.Column("dataset", sa.String(length=255), nullable=False), sa.Column("time_period", sa.String(length=32), nullable=False), - sa.Column("options_json", sa.JSON(), nullable=True), + sa.Column("options_json", sa.JSON(none_as_null=True), nullable=True), sa.Column("options_hash", sa.String(length=255), nullable=True), sa.Column("api_version", sa.String(length=10), nullable=False), sa.Column("reform_impact_json", sa.JSON(), nullable=False), @@ -107,18 +106,20 @@ def upgrade() -> None: ) op.create_table( "report_output_runs", - sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("id", sa.CHAR(length=36), nullable=False), sa.Column("report_output_id", sa.Integer(), nullable=False), sa.Column("run_sequence", sa.Integer(), nullable=False), sa.Column("status", sa.String(length=32), nullable=False), - sa.Column("output", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), sa.Column("error_message", sa.Text(), nullable=True), sa.Column("trigger_type", sa.String(length=32), nullable=False), sa.Column("requested_at", sa.DateTime(), nullable=True), sa.Column("started_at", sa.DateTime(), nullable=True), sa.Column("finished_at", sa.DateTime(), nullable=True), - sa.Column("source_run_id", sa.String(length=36), nullable=True), - sa.Column("report_spec_snapshot_json", sa.JSON(), nullable=True), + sa.Column("source_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "report_spec_snapshot_json", sa.JSON(none_as_null=True), nullable=True + ), sa.Column("country_package_version", sa.String(length=255), nullable=True), sa.Column("policyengine_version", sa.String(length=255), nullable=True), sa.Column("data_version", sa.String(length=255), nullable=True), @@ -129,7 +130,9 @@ def upgrade() -> None: sa.Column("resolved_dataset", sa.String(length=255), nullable=True), sa.Column("resolved_options_hash", sa.String(length=255), nullable=True), sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint("report_output_id", "run_sequence"), + sa.UniqueConstraint( + "report_output_id", "run_sequence", name="report_output_run_sequence_idx" + ), ) op.create_table( "report_outputs", @@ -144,7 +147,7 @@ def upgrade() -> None: server_default=sa.text("'pending'"), nullable=False, ), - sa.Column("output", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), sa.Column("error_message", sa.Text(), nullable=True), sa.Column( "year", @@ -153,36 +156,44 @@ def upgrade() -> None: nullable=True, ), sa.Column("report_kind", sa.String(length=64), nullable=True), - sa.Column("report_spec_json", sa.JSON(), nullable=True), + sa.Column("report_spec_json", sa.JSON(none_as_null=True), nullable=True), sa.Column("report_spec_schema_version", sa.Integer(), nullable=True), sa.Column("report_spec_status", sa.String(length=32), nullable=True), - sa.Column("active_run_id", sa.String(length=36), nullable=True), - sa.Column("latest_successful_run_id", sa.String(length=36), nullable=True), + sa.Column("active_run_id", sa.CHAR(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.CHAR(length=36), nullable=True), sa.PrimaryKeyConstraint("id"), ) op.create_table( "simulation_runs", - sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("id", sa.CHAR(length=36), nullable=False), sa.Column("simulation_id", sa.Integer(), nullable=False), - sa.Column("report_output_run_id", sa.String(length=36), nullable=True), - sa.Column("input_position", sa.Integer(), nullable=True), + sa.Column("report_output_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "input_position", + sa.Integer().with_variant(mysql.TINYINT(), "mysql"), + nullable=True, + ), sa.Column("run_sequence", sa.Integer(), nullable=False), sa.Column("status", sa.String(length=32), nullable=False), - sa.Column("output", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), sa.Column("error_message", sa.Text(), nullable=True), sa.Column("trigger_type", sa.String(length=32), nullable=False), sa.Column("requested_at", sa.DateTime(), nullable=True), sa.Column("started_at", sa.DateTime(), nullable=True), sa.Column("finished_at", sa.DateTime(), nullable=True), - sa.Column("source_run_id", sa.String(length=36), nullable=True), - sa.Column("simulation_spec_snapshot_json", sa.JSON(), nullable=True), + sa.Column("source_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "simulation_spec_snapshot_json", sa.JSON(none_as_null=True), nullable=True + ), sa.Column("country_package_version", sa.String(length=255), nullable=True), sa.Column("policyengine_version", sa.String(length=255), nullable=True), sa.Column("data_version", sa.String(length=255), nullable=True), sa.Column("runtime_app_name", sa.String(length=255), nullable=True), sa.Column("simulation_cache_version", sa.String(length=255), nullable=True), sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint("simulation_id", "run_sequence"), + sa.UniqueConstraint( + "simulation_id", "run_sequence", name="simulation_run_sequence_idx" + ), ) op.create_table( "simulations", @@ -198,12 +209,12 @@ def upgrade() -> None: server_default=sa.text("'pending'"), nullable=False, ), - sa.Column("output", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), sa.Column("error_message", sa.Text(), nullable=True), - sa.Column("simulation_spec_json", sa.JSON(), nullable=True), + sa.Column("simulation_spec_json", sa.JSON(none_as_null=True), nullable=True), sa.Column("simulation_spec_schema_version", sa.Integer(), nullable=True), - sa.Column("active_run_id", sa.String(length=36), nullable=True), - sa.Column("latest_successful_run_id", sa.String(length=36), nullable=True), + sa.Column("active_run_id", sa.CHAR(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.CHAR(length=36), nullable=True), sa.PrimaryKeyConstraint("id"), ) op.create_table( diff --git a/policyengine_api/data/initialise.sql b/policyengine_api/data/initialise.sql index 085f31c0b..6fd9210db 100644 --- a/policyengine_api/data/initialise.sql +++ b/policyengine_api/data/initialise.sql @@ -56,6 +56,7 @@ CREATE TABLE IF NOT EXISTS reform_impact ( status VARCHAR(32) NOT NULL, message VARCHAR(255), start_time DATETIME, + end_time DATETIME, execution_id VARCHAR(255) NOT NULL ); @@ -64,7 +65,7 @@ CREATE TABLE IF NOT EXISTS analysis ( prompt LONGTEXT NOT NULL, analysis LONGTEXT, status VARCHAR(32) NOT NULL -) +); -- The dataset row below was added while the table is in prod; -- we must allow NULL values for this column diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 834016e78..2529021d5 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -11,6 +11,7 @@ from sqlalchemy import ( BigInteger, + CHAR, DateTime, Integer, JSON, @@ -19,6 +20,7 @@ UniqueConstraint, text, ) +from sqlalchemy.dialects.mysql import LONGTEXT, TINYINT from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column @@ -48,9 +50,7 @@ class ComputedHousehold(V1Base): class Policy(V1Base): __tablename__ = "policy" - # SQLite cannot compile AUTO_INCREMENT on a composite primary key. The - # generated MySQL baseline receives the documented dialect correction. - id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=False) + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3), primary_key=True) label: Mapped[str | None] = mapped_column(String(255)) api_version: Mapped[str] = mapped_column(String(10)) @@ -102,8 +102,10 @@ class Analysis(V1Base): prompt_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True ) - prompt: Mapped[str] = mapped_column(Text) - analysis: Mapped[str | None] = mapped_column(Text) + prompt: Mapped[str] = mapped_column(Text().with_variant(LONGTEXT(), "mysql")) + analysis: Mapped[str | None] = mapped_column( + Text().with_variant(LONGTEXT(), "mysql") + ) status: Mapped[str] = mapped_column(String(32)) @@ -159,8 +161,8 @@ class Simulation(V1Base): error_message: Mapped[str | None] = mapped_column(Text) simulation_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) simulation_spec_schema_version: Mapped[int | None] - active_run_id: Mapped[str | None] = mapped_column(String(36)) - latest_successful_run_id: Mapped[str | None] = mapped_column(String(36)) + active_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(CHAR(36)) class ReportOutput(V1Base): @@ -178,14 +180,20 @@ class ReportOutput(V1Base): report_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) report_spec_schema_version: Mapped[int | None] report_spec_status: Mapped[str | None] = mapped_column(String(32)) - active_run_id: Mapped[str | None] = mapped_column(String(36)) - latest_successful_run_id: Mapped[str | None] = mapped_column(String(36)) + active_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(CHAR(36)) class ReportOutputRun(V1Base): __tablename__ = "report_output_runs" - __table_args__ = (UniqueConstraint("report_output_id", "run_sequence"),) - id: Mapped[str] = mapped_column(String(36), primary_key=True) + __table_args__ = ( + UniqueConstraint( + "report_output_id", + "run_sequence", + name="report_output_run_sequence_idx", + ), + ) + id: Mapped[str] = mapped_column(CHAR(36), primary_key=True) report_output_id: Mapped[int] run_sequence: Mapped[int] status: Mapped[str] = mapped_column(String(32)) @@ -195,7 +203,7 @@ class ReportOutputRun(V1Base): requested_at: Mapped[datetime | None] = mapped_column(DateTime) started_at: Mapped[datetime | None] = mapped_column(DateTime) finished_at: Mapped[datetime | None] = mapped_column(DateTime) - source_run_id: Mapped[str | None] = mapped_column(String(36)) + source_run_id: Mapped[str | None] = mapped_column(CHAR(36)) report_spec_snapshot_json: Mapped[Any | None] = mapped_column( JSON(none_as_null=True) ) @@ -212,11 +220,19 @@ class ReportOutputRun(V1Base): class SimulationRun(V1Base): __tablename__ = "simulation_runs" - __table_args__ = (UniqueConstraint("simulation_id", "run_sequence"),) - id: Mapped[str] = mapped_column(String(36), primary_key=True) + __table_args__ = ( + UniqueConstraint( + "simulation_id", + "run_sequence", + name="simulation_run_sequence_idx", + ), + ) + id: Mapped[str] = mapped_column(CHAR(36), primary_key=True) simulation_id: Mapped[int] - report_output_run_id: Mapped[str | None] = mapped_column(String(36)) - input_position: Mapped[int | None] + report_output_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + input_position: Mapped[int | None] = mapped_column( + Integer().with_variant(TINYINT(), "mysql") + ) run_sequence: Mapped[int] status: Mapped[str] = mapped_column(String(32)) output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) @@ -225,7 +241,7 @@ class SimulationRun(V1Base): requested_at: Mapped[datetime | None] = mapped_column(DateTime) started_at: Mapped[datetime | None] = mapped_column(DateTime) finished_at: Mapped[datetime | None] = mapped_column(DateTime) - source_run_id: Mapped[str | None] = mapped_column(String(36)) + source_run_id: Mapped[str | None] = mapped_column(CHAR(36)) simulation_spec_snapshot_json: Mapped[Any | None] = mapped_column( JSON(none_as_null=True) ) @@ -238,5 +254,9 @@ class SimulationRun(V1Base): class LegacyReportOutputAlias(V1Base): __tablename__ = "legacy_report_output_aliases" - legacy_report_output_id: Mapped[int] = mapped_column(Integer, primary_key=True) + legacy_report_output_id: Mapped[int] = mapped_column( + Integer, + primary_key=True, + autoincrement=False, + ) canonical_report_output_id: Mapped[int] diff --git a/policyengine_api/scripts/stage7_database.py b/policyengine_api/scripts/stage7_database.py new file mode 100644 index 000000000..0d9d35fbc --- /dev/null +++ b/policyengine_api/scripts/stage7_database.py @@ -0,0 +1,18 @@ +"""Safety checks shared by destructive Stage 7 qualification tooling.""" + +from sqlalchemy.engine import make_url + + +def assert_safe_toy_database_url(database_url: str) -> None: + """Reject destructive toy-test operations against non-local databases.""" + + url = make_url(database_url) + is_local_mysql = url.get_backend_name() == "mysql" and url.host in { + "127.0.0.1", + "localhost", + } + is_toy_database = bool(url.database and url.database.endswith("_toy")) + if not is_local_mysql or not is_toy_database: + raise ValueError( + "Stage 7 integration tests require a local MySQL database ending in '_toy'" + ) diff --git a/tests/fixtures/stage7_pre_alembic_schema.sql b/tests/fixtures/stage7_pre_alembic_schema.sql new file mode 100644 index 000000000..48b75a20a --- /dev/null +++ b/tests/fixtures/stage7_pre_alembic_schema.sql @@ -0,0 +1,194 @@ +-- Frozen API v1 schema snapshot captured before Alembic ownership. +-- Do not update this file to match ORM metadata or generated revisions. + +CREATE TABLE household ( + id INTEGER PRIMARY KEY AUTO_INCREMENT, + country_id VARCHAR(3) NOT NULL, + label VARCHAR(255), + api_version VARCHAR(255) NOT NULL, + household_json JSON NOT NULL, + household_hash VARCHAR(255) NOT NULL +); + +CREATE TABLE computed_household ( + household_id INT NOT NULL, + policy_id INT NOT NULL, + country_id VARCHAR(3) NOT NULL, + api_version VARCHAR(10) NOT NULL, + computed_household_json JSON NOT NULL, + status VARCHAR(32), + PRIMARY KEY (household_id, policy_id, country_id) +); + +CREATE TABLE policy ( + id INTEGER AUTO_INCREMENT, + country_id VARCHAR(3) NOT NULL, + label VARCHAR(255), + api_version VARCHAR(10) NOT NULL, + policy_json JSON NOT NULL, + policy_hash VARCHAR(255) NOT NULL, + PRIMARY KEY (id, country_id, policy_hash) +); + +CREATE TABLE economy ( + economy_id INTEGER PRIMARY KEY AUTO_INCREMENT, + policy_id INT NOT NULL, + country_id VARCHAR(3) NOT NULL, + region VARCHAR(32), + time_period VARCHAR(32), + options_json JSON NOT NULL, + options_hash VARCHAR(255) NOT NULL, + api_version VARCHAR(10) NOT NULL, + economy_json JSON, + status VARCHAR(32) NOT NULL, + message VARCHAR(255) +); + +CREATE TABLE reform_impact ( + reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, + baseline_policy_id INT NOT NULL, + reform_policy_id INT NOT NULL, + country_id VARCHAR(3) NOT NULL, + region VARCHAR(32) NOT NULL, + dataset VARCHAR(255) NOT NULL, + time_period VARCHAR(32) NOT NULL, + options_json JSON, + options_hash VARCHAR(255), + api_version VARCHAR(10) NOT NULL, + reform_impact_json JSON NOT NULL, + status VARCHAR(32) NOT NULL, + message VARCHAR(255), + start_time DATETIME, + end_time DATETIME, + execution_id VARCHAR(255) NOT NULL +); + +CREATE TABLE analysis ( + prompt_id INTEGER PRIMARY KEY AUTO_INCREMENT, + prompt LONGTEXT NOT NULL, + analysis LONGTEXT, + status VARCHAR(32) NOT NULL +); + +CREATE TABLE user_policies ( + id INTEGER PRIMARY KEY AUTO_INCREMENT, + country_id VARCHAR(3) NOT NULL, + reform_id INTEGER NOT NULL, + reform_label VARCHAR(255), + baseline_id INTEGER NOT NULL, + baseline_label VARCHAR(255), + user_id VARCHAR(255) NOT NULL, + year VARCHAR(32) NOT NULL, + geography VARCHAR(255) NOT NULL, + dataset VARCHAR(255), + number_of_provisions INTEGER NOT NULL, + api_version VARCHAR(32) NOT NULL, + added_date BIGINT NOT NULL, + updated_date BIGINT NOT NULL, + budgetary_impact VARCHAR(255), + type VARCHAR(255) +); + +CREATE TABLE user_profiles ( + user_id INTEGER PRIMARY KEY AUTO_INCREMENT, + auth0_id VARCHAR(255) NOT NULL UNIQUE, + username VARCHAR(255) UNIQUE, + primary_country VARCHAR(3) NOT NULL, + user_since BIGINT NOT NULL +); + +CREATE TABLE tracers ( + id INTEGER PRIMARY KEY AUTO_INCREMENT, + household_id INT NOT NULL, + policy_id INT NOT NULL, + country_id VARCHAR(3) NOT NULL, + api_version VARCHAR(10) NOT NULL, + tracer_output JSON NOT NULL +); + +CREATE TABLE simulations ( + id INTEGER PRIMARY KEY AUTO_INCREMENT, + country_id VARCHAR(3) NOT NULL, + api_version VARCHAR(10) NOT NULL, + population_id VARCHAR(255) NOT NULL, + population_type VARCHAR(50) NOT NULL, + policy_id INT NOT NULL, + status VARCHAR(32) NOT NULL DEFAULT 'pending', + output JSON DEFAULT NULL, + error_message TEXT DEFAULT NULL, + simulation_spec_json JSON DEFAULT NULL, + simulation_spec_schema_version INT DEFAULT NULL, + active_run_id CHAR(36) DEFAULT NULL, + latest_successful_run_id CHAR(36) DEFAULT NULL +); + +CREATE TABLE report_outputs ( + id INTEGER PRIMARY KEY AUTO_INCREMENT, + country_id VARCHAR(3) NOT NULL, + simulation_1_id INT NOT NULL, + simulation_2_id INT DEFAULT NULL, + api_version VARCHAR(10) NOT NULL, + status VARCHAR(32) NOT NULL DEFAULT 'pending', + output JSON DEFAULT NULL, + error_message TEXT DEFAULT NULL, + year VARCHAR(255) DEFAULT '2025', + report_kind VARCHAR(64) DEFAULT NULL, + report_spec_json JSON DEFAULT NULL, + report_spec_schema_version INT DEFAULT NULL, + report_spec_status VARCHAR(32) DEFAULT NULL, + active_run_id CHAR(36) DEFAULT NULL, + latest_successful_run_id CHAR(36) DEFAULT NULL +); + +CREATE TABLE report_output_runs ( + id CHAR(36) PRIMARY KEY, + report_output_id INT NOT NULL, + run_sequence INT NOT NULL, + status VARCHAR(32) NOT NULL, + output JSON DEFAULT NULL, + error_message TEXT DEFAULT NULL, + trigger_type VARCHAR(32) NOT NULL, + requested_at DATETIME DEFAULT NULL, + started_at DATETIME DEFAULT NULL, + finished_at DATETIME DEFAULT NULL, + source_run_id CHAR(36) DEFAULT NULL, + report_spec_snapshot_json JSON DEFAULT NULL, + country_package_version VARCHAR(255) DEFAULT NULL, + policyengine_version VARCHAR(255) DEFAULT NULL, + data_version VARCHAR(255) DEFAULT NULL, + runtime_app_name VARCHAR(255) DEFAULT NULL, + report_cache_version VARCHAR(255) DEFAULT NULL, + simulation_cache_version VARCHAR(255) DEFAULT NULL, + requested_version_override VARCHAR(255) DEFAULT NULL, + resolved_dataset VARCHAR(255) DEFAULT NULL, + resolved_options_hash VARCHAR(255) DEFAULT NULL, + UNIQUE KEY report_output_run_sequence_idx (report_output_id, run_sequence) +); + +CREATE TABLE simulation_runs ( + id CHAR(36) PRIMARY KEY, + simulation_id INT NOT NULL, + report_output_run_id CHAR(36) DEFAULT NULL, + input_position TINYINT DEFAULT NULL, + run_sequence INT NOT NULL, + status VARCHAR(32) NOT NULL, + output JSON DEFAULT NULL, + error_message TEXT DEFAULT NULL, + trigger_type VARCHAR(32) NOT NULL, + requested_at DATETIME DEFAULT NULL, + started_at DATETIME DEFAULT NULL, + finished_at DATETIME DEFAULT NULL, + source_run_id CHAR(36) DEFAULT NULL, + simulation_spec_snapshot_json JSON DEFAULT NULL, + country_package_version VARCHAR(255) DEFAULT NULL, + policyengine_version VARCHAR(255) DEFAULT NULL, + data_version VARCHAR(255) DEFAULT NULL, + runtime_app_name VARCHAR(255) DEFAULT NULL, + simulation_cache_version VARCHAR(255) DEFAULT NULL, + UNIQUE KEY simulation_run_sequence_idx (simulation_id, run_sequence) +); + +CREATE TABLE legacy_report_output_aliases ( + legacy_report_output_id INT PRIMARY KEY, + canonical_report_output_id INT NOT NULL +); diff --git a/tests/integration/stage7_mysql.py b/tests/integration/stage7_mysql.py new file mode 100644 index 000000000..1561e792d --- /dev/null +++ b/tests/integration/stage7_mysql.py @@ -0,0 +1,86 @@ +"""Shared fixtures for destructive Stage 7 tests on disposable MySQL only.""" + +from collections.abc import Iterator + +from alembic.config import Config +import pytest +from sqlalchemy import Engine, create_engine, inspect + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base +from policyengine_api.scripts.stage7_database import assert_safe_toy_database_url + + +def stage7_database_url() -> str | None: + import os + + return os.environ.get("STAGE7_TOY_DATABASE_URL") + + +def alembic_config(database_url: str) -> Config: + config = Config(str(REPO / "alembic.ini")) + config.set_main_option("sqlalchemy.url", database_url) + return config + + +def reset_toy_database(engine: Engine, database_url: str) -> None: + assert_safe_toy_database_url(database_url) + with engine.begin() as connection: + connection.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 0") + for table_name in reversed(V1Base.metadata.sorted_tables): + connection.exec_driver_sql(f"DROP TABLE IF EXISTS `{table_name.name}`") + connection.exec_driver_sql("DROP TABLE IF EXISTS alembic_version") + connection.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 1") + + +def create_pre_alembic_schema(engine: Engine) -> None: + """Create the existing schema from its independent legacy SQL source.""" + + source = (REPO / "tests/fixtures/stage7_pre_alembic_schema.sql").read_text( + encoding="utf-8" + ) + statements = source.split(";") + with engine.begin() as connection: + for statement in statements: + if statement.strip(): + connection.exec_driver_sql(statement.strip().removesuffix(";")) + + +def schema_signature(engine: Engine) -> dict: + inspector = inspect(engine) + + def normalized(value): + if isinstance(value, dict): + return {key: normalized(item) for key, item in sorted(value.items())} + if isinstance(value, (list, tuple)): + return [normalized(item) for item in value] + if value is None or isinstance(value, (bool, int, float, str)): + return value + return str(value) + + return { + table_name: normalized( + { + "columns": inspector.get_columns(table_name), + "indexes": inspector.get_indexes(table_name), + "pk": inspector.get_pk_constraint(table_name), + "unique": inspector.get_unique_constraints(table_name), + } + ) + for table_name in inspector.get_table_names() + } + + +@pytest.fixture +def stage7_mysql() -> Iterator[tuple[str, Engine]]: + database_url = stage7_database_url() + if database_url is None: + pytest.skip("STAGE7_TOY_DATABASE_URL is required for the MySQL probe") + assert_safe_toy_database_url(database_url) + engine = create_engine(database_url) + reset_toy_database(engine, database_url) + try: + yield database_url, engine + finally: + reset_toy_database(engine, database_url) + engine.dispose() diff --git a/tests/integration/test_stage7_existing_schema.py b/tests/integration/test_stage7_existing_schema.py new file mode 100644 index 000000000..77a264039 --- /dev/null +++ b/tests/integration/test_stage7_existing_schema.py @@ -0,0 +1,68 @@ +"""Qualification of an existing pre-Alembic v1 schema.""" + +import os + +from alembic import command +import pytest +from sqlalchemy import inspect, text + +from policyengine_api.scripts.qualify_stage7_toy import compare_stage7_schema +from tests.integration.stage7_mysql import ( + alembic_config, + create_pre_alembic_schema, + reset_toy_database, + schema_signature, +) + + +def test_fresh_upgrade_has_the_same_schema_signature_as_pre_alembic_v1(stage7_mysql): + database_url, engine = stage7_mysql + create_pre_alembic_schema(engine) + pre_alembic = schema_signature(engine) + + reset_toy_database(engine, database_url) + command.upgrade(alembic_config(database_url), "head") + fresh_upgrade = schema_signature(engine) + fresh_upgrade.pop("alembic_version") + + assert fresh_upgrade == pre_alembic + + +def test_existing_schema_compares_read_only_and_stamps_without_data_loss( + stage7_mysql, +): + database_url, engine = stage7_mysql + create_pre_alembic_schema(engine) + with engine.begin() as connection: + connection.execute( + text( + "INSERT INTO policy " + "(id, country_id, label, api_version, policy_json, policy_hash) " + "VALUES (901, 'us', 'sentinel', 'legacy', '{}', 'sentinel')" + ) + ) + + before_comparison = schema_signature(engine) + assert "alembic_version" not in before_comparison + assert compare_stage7_schema(database_url) == [] + assert schema_signature(engine) == before_comparison + + config = alembic_config(database_url) + command.stamp(config, "head") + command.check(config) + + with engine.connect() as connection: + assert ( + connection.scalar(text("SELECT label FROM policy WHERE id = 901")) + == "sentinel" + ) + assert inspect(engine).get_table_names().count("alembic_version") == 1 + + +@pytest.mark.skipif( + "STAGE7_EXISTING_DATABASE_URL" not in os.environ, + reason="A read-only existing Cloud SQL URL was not supplied", +) +def test_live_existing_schema_matches_metadata_without_mutation(): + database_url = os.environ["STAGE7_EXISTING_DATABASE_URL"] + assert compare_stage7_schema(database_url) == [] diff --git a/tests/unit/data/sqlite_schema.py b/tests/unit/data/sqlite_schema.py new file mode 100644 index 000000000..9852a2280 --- /dev/null +++ b/tests/unit/data/sqlite_schema.py @@ -0,0 +1,12 @@ +from policyengine_api.constants import REPO +from policyengine_api.data.orm import SessionManager + + +def create_sqlite_v1_schema(manager: SessionManager) -> None: + """Install the explicit local schema without compiling MySQL metadata.""" + + schema = (REPO / "policyengine_api/data/initialise_local.sql").read_text( + encoding="utf-8" + ) + with manager.engine.connect() as connection: + connection.connection.driver_connection.executescript(schema) diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py index daf4bf604..c1e57420a 100644 --- a/tests/unit/data/test_local_daos.py +++ b/tests/unit/data/test_local_daos.py @@ -2,12 +2,12 @@ from policyengine_api.data.orm import build_sqlite_session_manager from policyengine_api.data.v1_daos import AnalysisDAO, ReformImpactDAO, TracerDAO -from policyengine_api.data.v1_models import V1Base +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema def _daos(): manager = build_sqlite_session_manager() - V1Base.metadata.create_all(manager.engine) + create_sqlite_v1_schema(manager) return AnalysisDAO(manager), ReformImpactDAO(manager), TracerDAO(manager) diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py index 82157ef94..9bda3cf4d 100644 --- a/tests/unit/data/test_run_daos.py +++ b/tests/unit/data/test_run_daos.py @@ -2,12 +2,12 @@ from policyengine_api.data.orm import build_sqlite_session_manager from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO -from policyengine_api.data.v1_models import V1Base +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema def _daos(): manager = build_sqlite_session_manager() - V1Base.metadata.create_all(manager.engine) + create_sqlite_v1_schema(manager) return SimulationDAO(manager), ReportDAO(manager) diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index 0ab746808..9791f207b 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -1,11 +1,11 @@ from policyengine_api.data.orm import build_sqlite_session_manager from policyengine_api.data.v1_daos import HouseholdDAO, PolicyDAO, UserDAO -from policyengine_api.data.v1_models import V1Base +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema def _daos(): manager = build_sqlite_session_manager() - V1Base.metadata.create_all(manager.engine) + create_sqlite_v1_schema(manager) return PolicyDAO(manager), HouseholdDAO(manager), UserDAO(manager) diff --git a/tests/unit/data/test_v1_models.py b/tests/unit/data/test_v1_models.py index 6d92f1d45..02b738531 100644 --- a/tests/unit/data/test_v1_models.py +++ b/tests/unit/data/test_v1_models.py @@ -1,4 +1,5 @@ -from sqlalchemy import create_engine +from sqlalchemy.dialects import mysql +from sqlalchemy.schema import CreateTable from policyengine_api.data.v1_models import V1Base @@ -25,22 +26,14 @@ def test_v1_metadata_contains_every_legacy_table(): assert set(V1Base.metadata.tables) == EXPECTED_TABLES -def test_v1_metadata_builds_a_fresh_sqlite_database(): - engine = create_engine("sqlite+pysqlite:///:memory:") - V1Base.metadata.create_all(engine) - - with engine.connect() as connection: - table_names = { - row[0] - for row in connection.exec_driver_sql( - "SELECT name FROM sqlite_master WHERE type='table'" - ) - } - assert EXPECTED_TABLES <= table_names +def test_v1_metadata_compiles_for_the_production_mysql_dialect(): + for table in V1Base.metadata.sorted_tables: + assert str(CreateTable(table).compile(dialect=mysql.dialect())) def test_v1_composite_and_unique_keys_match_legacy_contract(): policy = V1Base.metadata.tables["policy"] + assert policy.c.id.autoincrement is True assert [column.name for column in policy.primary_key.columns] == [ "id", "country_id", @@ -53,3 +46,9 @@ def test_v1_composite_and_unique_keys_match_legacy_contract(): "country_id", ] assert V1Base.metadata.tables["user_profiles"].c.auth0_id.unique + assert ( + V1Base.metadata.tables[ + "legacy_report_output_aliases" + ].c.legacy_report_output_id.autoincrement + is False + ) From dc3467632e2af0b5527b74c71bd4bb3bd61e3d01 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:04:33 +0300 Subject: [PATCH 34/89] refactor: make service operations own ORM transactions --- Makefile | 11 + compose.stage7-toy.yml | 22 ++ policyengine_api/data/orm.py | 20 +- policyengine_api/data/v1_daos.py | 262 +++++++++++------- .../scripts/qualify_stage7_toy.py | 80 +++++- .../services/household_service.py | 72 +++-- policyengine_api/services/policy_service.py | 63 +++-- policyengine_api/services/user_service.py | 55 ++-- tests/integration/conftest.py | 2 + tests/integration/test_stage7_toy_database.py | 47 ++++ tests/unit/data/test_v1_daos.py | 89 +++--- tests/unit/data/test_v1_unit_of_work.py | 37 +++ .../services/test_update_profile_service.py | 5 +- tests/unit/test_stage7_toy_qualification.py | 63 +++-- 14 files changed, 581 insertions(+), 247 deletions(-) create mode 100644 compose.stage7-toy.yml create mode 100644 tests/integration/test_stage7_toy_database.py create mode 100644 tests/unit/data/test_v1_unit_of_work.py diff --git a/Makefile b/Makefile index d9bbc1dda..377495201 100644 --- a/Makefile +++ b/Makefile @@ -20,6 +20,17 @@ test: quality-guards: python scripts/run_quality_guards.py +STAGE7_TOY_DATABASE_URL ?= mysql+pymysql://policyengine:policyengine@127.0.0.1:3307/policyengine_stage7_toy + +stage7-toy-up: + docker compose -f compose.stage7-toy.yml up -d --wait + +stage7-toy-test: stage7-toy-up + STAGE7_TOY_DATABASE_URL="$(STAGE7_TOY_DATABASE_URL)" uv run pytest tests/integration/test_stage7_*.py -v + +stage7-toy-down: + docker compose -f compose.stage7-toy.yml down --volumes + debug-test: MAX_HOUSEHOLDS=1000 FLASK_DEBUG=1 pytest -vv --durations=0 tests diff --git a/compose.stage7-toy.yml b/compose.stage7-toy.yml new file mode 100644 index 000000000..9ad5eb505 --- /dev/null +++ b/compose.stage7-toy.yml @@ -0,0 +1,22 @@ +name: policyengine-stage7-toy + +services: + mysql: + image: mysql:8.0 + environment: + MYSQL_DATABASE: policyengine_stage7_toy + MYSQL_PASSWORD: policyengine + MYSQL_ROOT_PASSWORD: policyengine-root + MYSQL_USER: policyengine + ports: + - "127.0.0.1:${STAGE7_TOY_MYSQL_PORT:-3307}:3306" + tmpfs: + - /var/lib/mysql + healthcheck: + test: + - CMD-SHELL + - mysqladmin ping --host=127.0.0.1 --user=policyengine --password=policyengine --silent + interval: 2s + timeout: 2s + retries: 30 + start_period: 10s diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index f37a1d746..8d90aefcd 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -47,21 +47,17 @@ def __init__(self, engine: Engine): @contextmanager def session(self) -> Iterator[Session]: - session = self.session_factory() - try: + with self.session_factory() as session: + yield session + + @contextmanager + def transaction(self) -> Iterator[Session]: + with self.session_factory.begin() as session: yield session - finally: - session.close() def run_in_transaction(self, callback: Callable[[Session], T]) -> T: - with self.session() as session: - try: - result = callback(session) - session.commit() - return result - except Exception: - session.rollback() - raise + with self.session_factory.begin() as session: + return callback(session) def build_sqlite_session_manager( diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 573260b98..9fc8f3617 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -2,16 +2,21 @@ from __future__ import annotations +from collections.abc import Iterator +from contextlib import contextmanager from typing import Any from datetime import datetime import uuid from sqlalchemy import delete, func, or_, select +from sqlalchemy.orm import Session from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_models import ( Analysis, + ComputedHousehold, + Economy, Household, LegacyReportOutputAlias, Policy, @@ -22,6 +27,7 @@ SimulationRun, Tracer, UserProfile, + UserPolicy, ) @@ -112,31 +118,29 @@ def runtime_sqlalchemy_dao(*, local: bool = False) -> SQLAlchemyDAO: class PolicyDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def get(self, country_id: str, policy_id: int) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.scalar( - select(Policy).where( - Policy.country_id == country_id, - Policy.id == policy_id, - ) + model = self.session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, ) - return _mapping(model) if model else None + ) + return _mapping(model) if model else None def find_unique( self, country_id: str, policy_hash: str, label: str | None ) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.scalar( - select(Policy).where( - Policy.country_id == country_id, - Policy.policy_hash == policy_hash, - Policy.label == label, - ) + model = self.session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.policy_hash == policy_hash, + Policy.label == label, ) - return _mapping(model) if model else None + ) + return _mapping(model) if model else None def create( self, @@ -146,36 +150,30 @@ def create( policy_hash: str, api_version: str, ) -> int: - def operation(session): - next_id = (session.scalar(select(func.max(Policy.id))) or 0) + 1 - session.add( - Policy( - id=next_id, - country_id=country_id, - label=label, - api_version=api_version, - policy_json=policy_json, - policy_hash=policy_hash, - ) - ) - return next_id - - return self.sessions.run_in_transaction(operation) + policy = Policy( + country_id=country_id, + label=label, + api_version=api_version, + policy_json=policy_json, + policy_hash=policy_hash, + ) + self.session.add(policy) + self.session.flush() + return policy.id class HouseholdDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def get(self, country_id: str, household_id: int) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.scalar( - select(Household).where( - Household.country_id == country_id, - Household.id == household_id, - ) + model = self.session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, ) - return _mapping(model) if model else None + ) + return _mapping(model) if model else None def create( self, @@ -185,19 +183,16 @@ def create( household_hash: str, api_version: str, ) -> int: - def operation(session): - model = Household( - country_id=country_id, - label=label, - api_version=api_version, - household_json=household_json, - household_hash=household_hash, - ) - session.add(model) - session.flush() - return model.id - - return self.sessions.run_in_transaction(operation) + model = Household( + country_id=country_id, + label=label, + api_version=api_version, + household_json=household_json, + household_hash=household_hash, + ) + self.session.add(model) + self.session.flush() + return model.id def update( self, @@ -208,28 +203,45 @@ def update( household_hash: str, api_version: str, ) -> bool: - def operation(session): - model = session.scalar( - select(Household).where( - Household.country_id == country_id, - Household.id == household_id, - ) + model = self.session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, ) - if model is None: - return False - model.label = label - model.household_json = household_json - model.household_hash = household_hash - model.api_version = api_version - return True - - return self.sessions.run_in_transaction(operation) + ) + if model is None: + return False + model.label = label + model.household_json = household_json + model.household_hash = household_hash + model.api_version = api_version + return True -class UserDAO: +class ComputedHouseholdDAO: def __init__(self, sessions: SessionManager): self.sessions = sessions + def create(self, **values: Any) -> None: + self.sessions.run_in_transaction( + lambda session: session.add(ComputedHousehold(**values)) + ) + + def get( + self, household_id: int, policy_id: int, country_id: str + ) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get( + ComputedHousehold, + (household_id, policy_id, country_id), + ) + return _mapping(model) if model else None + + +class UserDAO: + def __init__(self, session: Session): + self.session = session + def create_profile( self, auth0_id: str, @@ -237,18 +249,15 @@ def create_profile( primary_country: str, user_since: int, ) -> int: - def operation(session): - model = UserProfile( - auth0_id=auth0_id, - username=username, - primary_country=primary_country, - user_since=user_since, - ) - session.add(model) - session.flush() - return model.user_id - - return self.sessions.run_in_transaction(operation) + model = UserProfile( + auth0_id=auth0_id, + username=username, + primary_country=primary_country, + user_since=user_since, + ) + self.session.add(model) + self.session.flush() + return model.user_id def get_profile( self, @@ -258,27 +267,88 @@ def get_profile( ) -> dict[str, Any] | None: if user_id is None and auth0_id is None: return None + condition = ( + UserProfile.user_id == user_id + if user_id is not None + else UserProfile.auth0_id == auth0_id + ) + model = self.session.scalar(select(UserProfile).where(condition)) + return _mapping(model) if model else None + + def update_profile(self, user_id: int, **values: Any) -> bool: + model = self.session.get(UserProfile, user_id) + if model is None: + return False + for key, value in values.items(): + if value is not None: + setattr(model, key, value) + return True + + +class V1Repositories: + """Repositories bound to the same operation-scoped Session.""" + + def __init__(self, session: Session): + self.session = session + self.policies = PolicyDAO(session) + self.households = HouseholdDAO(session) + self.users = UserDAO(session) + + +class V1UnitOfWork: + """Create one Session and transaction boundary per logical operation.""" + + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + @contextmanager + def read(self) -> Iterator[V1Repositories]: with self.sessions.session() as session: - condition = ( - UserProfile.user_id == user_id - if user_id is not None - else UserProfile.auth0_id == auth0_id - ) - model = session.scalar(select(UserProfile).where(condition)) + yield V1Repositories(session) + + @contextmanager + def transaction(self) -> Iterator[V1Repositories]: + with self.sessions.transaction() as session: + yield V1Repositories(session) + + +class UserPolicyDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def create(self, **values: Any) -> int: + def operation(session): + model = UserPolicy(**values) + session.add(model) + session.flush() + return model.id + + return self.sessions.run_in_transaction(operation) + + def get(self, user_policy_id: int) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get(UserPolicy, user_policy_id) return _mapping(model) if model else None - def update_profile(self, user_id: int, **values: Any) -> bool: + +class EconomyDAO: + def __init__(self, sessions: SessionManager): + self.sessions = sessions + + def create(self, **values: Any) -> int: def operation(session): - model = session.get(UserProfile, user_id) - if model is None: - return False - for key, value in values.items(): - if value is not None: - setattr(model, key, value) - return True + model = Economy(**values) + session.add(model) + session.flush() + return model.economy_id return self.sessions.run_in_transaction(operation) + def get(self, economy_id: int) -> dict[str, Any] | None: + with self.sessions.session() as session: + model = session.get(Economy, economy_id) + return _mapping(model) if model else None + class AnalysisDAO: def __init__(self, sessions: SessionManager): diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py index 9fe7b7081..6a9f0e5c0 100644 --- a/policyengine_api/scripts/qualify_stage7_toy.py +++ b/policyengine_api/scripts/qualify_stage7_toy.py @@ -12,13 +12,14 @@ from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_daos import ( AnalysisDAO, - HouseholdDAO, - PolicyDAO, + ComputedHouseholdDAO, + EconomyDAO, ReformImpactDAO, ReportDAO, SimulationDAO, TracerDAO, - UserDAO, + UserPolicyDAO, + V1UnitOfWork, ) from policyengine_api.data.v1_models import V1Base @@ -30,7 +31,7 @@ def compare_stage7_schema(database_url: str) -> list: with engine.connect() as connection: context = MigrationContext.configure( connection, - opts={"compare_type": True}, + opts={"compare_type": True, "compare_server_default": True}, ) return compare_metadata(context, V1Base.metadata) @@ -44,18 +45,59 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: engine = create_engine(database_url) sessions = SessionManager(engine) - policies = PolicyDAO(sessions) - households = HouseholdDAO(sessions) - users = UserDAO(sessions) + unit_of_work = V1UnitOfWork(sessions) + computed_households = ComputedHouseholdDAO(sessions) + user_policies = UserPolicyDAO(sessions) + economies = EconomyDAO(sessions) simulations = SimulationDAO(sessions) reports = ReportDAO(sessions) analyses = AnalysisDAO(sessions) tracers = TracerDAO(sessions) impacts = ReformImpactDAO(sessions) - policy_id = policies.create("us", "Toy", {}, "toy-policy", "toy") - household_id = households.create("us", "Toy", {}, "toy-household", "toy") - user_id = users.create_profile("toy|user", "toy-user", "us", 1) + with unit_of_work.transaction() as repositories: + policy_id = repositories.policies.create("us", "Toy", {}, "toy-policy", "toy") + household_id = repositories.households.create( + "us", "Toy", {}, "toy-household", "toy" + ) + user_id = repositories.users.create_profile("toy|user", "toy-user", "us", 1) + computed_households.create( + household_id=household_id, + policy_id=policy_id, + country_id="us", + api_version="toy", + computed_household_json={"qualified": True}, + status="complete", + ) + user_policy_id = user_policies.create( + country_id="us", + reform_id=policy_id, + reform_label="Toy", + baseline_id=policy_id, + baseline_label="Toy", + user_id=str(user_id), + year="2026", + geography="us", + dataset="default", + number_of_provisions=0, + api_version="toy", + added_date=1, + updated_date=1, + budgetary_impact=None, + type="reform", + ) + economy_id = economies.create( + policy_id=policy_id, + country_id="us", + region="us", + time_period="2026", + options_json={}, + options_hash="toy-economy", + api_version="toy", + economy_json={"qualified": True}, + status="complete", + message=None, + ) simulation_id = simulations.create( country_id="us", api_version="toy", @@ -82,6 +124,7 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: status="pending", trigger_type="qualification", ) + reports.set_alias(900_001, report_id) analyses.store("toy prompt", "toy answer", "complete") tracers.create(household_id, policy_id, "us", "toy", ["toy trace"]) impact_id = impacts.create( @@ -100,13 +143,24 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: execution_id="toy-impact", ) + with unit_of_work.read() as repositories: + core_results = { + "policy": repositories.policies.get("us", policy_id) is not None, + "household": repositories.households.get("us", household_id) is not None, + "user": repositories.users.get_profile(user_id=user_id) is not None, + } + return { "alembic_head": "alembic_version" in inspect(engine).get_table_names(), - "policy": policies.get("us", policy_id) is not None, - "household": households.get("us", household_id) is not None, - "user": users.get_profile(user_id=user_id) is not None, + **core_results, + "computed_household": computed_households.get(household_id, policy_id, "us") + is not None, + "user_policy": user_policies.get(user_policy_id) is not None, + "economy": economies.get(economy_id) is not None, "simulation": simulations.get_run("toy-simulation-run") is not None, "report": reports.get_run("toy-report-run") is not None, + "report_alias": reports.get_alias(900_001)["canonical_report_output_id"] + == report_id, "analysis": analyses.get("toy prompt") == "toy answer", "tracer": tracers.get(household_id, policy_id, "us") is not None, "reform_impact": impacts.find(execution_id="toy-impact")["reform_impact_id"] diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index c94f27d36..96614c614 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,27 +1,45 @@ from __future__ import annotations +from contextlib import contextmanager + from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import HouseholdDAO +from policyengine_api.data.v1_daos import HouseholdDAO, V1UnitOfWork from policyengine_api.utils import hash_object class HouseholdService: - def __init__(self, households: HouseholdDAO | None = None): + def __init__( + self, + households: HouseholdDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._households = households + self._unit_of_work = unit_of_work @property - def households(self) -> HouseholdDAO: - if self._households is None: - self._households = HouseholdDAO(build_v1_session_manager()) - return self._households + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work + + @contextmanager + def _repository(self, *, write: bool = False): + if self._households is not None: + yield self._households + return + boundary = self.unit_of_work.transaction if write else self.unit_of_work.read + with boundary() as repositories: + yield repositories.households def get_household(self, country_id: str, household_id: int) -> dict | None: if type(household_id) is not int or household_id < 0: raise Exception( f"Invalid household ID: {household_id}. Must be a positive integer." ) - return self.households.get(country_id, household_id) + with self._repository() as households: + return households.get(country_id, household_id) def create_household( self, @@ -29,13 +47,14 @@ def create_household( household_json: dict, label: str | None, ) -> int: - return self.households.create( - country_id, - label, - household_json, - hash_object(household_json), - COUNTRY_PACKAGE_VERSIONS.get(country_id), - ) + with self._repository(write=True) as households: + return households.create( + country_id, + label, + household_json, + hash_object(household_json), + COUNTRY_PACKAGE_VERSIONS.get(country_id), + ) def update_household( self, @@ -44,16 +63,17 @@ def update_household( household_json: dict, label: str, ) -> dict: - updated = self.households.update( - country_id, - household_id, - label, - household_json, - hash_object(household_json), - COUNTRY_PACKAGE_VERSIONS.get(country_id), - ) - if not updated: - raise LookupError( - f"Household #{household_id} not found for country {country_id}." + with self._repository(write=True) as households: + updated = households.update( + country_id, + household_id, + label, + household_json, + hash_object(household_json), + COUNTRY_PACKAGE_VERSIONS.get(country_id), ) - return self.households.get(country_id, household_id) + if not updated: + raise LookupError( + f"Household #{household_id} not found for country {country_id}." + ) + return households.get(country_id, household_id) diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index f96926e0c..ee3eb0bf1 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,22 +1,38 @@ from __future__ import annotations import json +from contextlib import contextmanager from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import PolicyDAO +from policyengine_api.data.v1_daos import PolicyDAO, V1UnitOfWork from policyengine_api.utils import hash_object class PolicyService: - def __init__(self, policies: PolicyDAO | None = None): + def __init__( + self, + policies: PolicyDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._policies = policies + self._unit_of_work = unit_of_work @property - def policies(self) -> PolicyDAO: - if self._policies is None: - self._policies = PolicyDAO(build_v1_session_manager()) - return self._policies + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work + + @contextmanager + def _repository(self, *, write: bool = False): + if self._policies is not None: + yield self._policies + return + boundary = self.unit_of_work.transaction if write else self.unit_of_work.read + with boundary() as repositories: + yield repositories.policies @staticmethod def _validate_policy_id(policy_id: int) -> None: @@ -29,11 +45,13 @@ def get_policy(self, country_id: str, policy_id: int) -> dict | None: self._validate_policy_id(policy_id) if not country_id: raise ValueError("country_id cannot be empty or None") - return self.policies.get(country_id, policy_id) + with self._repository() as policies: + return policies.get(country_id, policy_id) def get_policy_json(self, country_id: str, policy_id: int) -> str | None: self._validate_policy_id(policy_id) - policy = self.policies.get(country_id, policy_id) + with self._repository() as policies: + policy = policies.get(country_id, policy_id) if policy is None: return None value = policy["policy_json"] @@ -47,18 +65,19 @@ def set_policy( raise ValueError(f"Invalid country_id: {country_id}") policy_hash = hash_object(policy_json) - existing = self.policies.find_unique(country_id, policy_hash, label or None) - if existing: - return existing["id"], "Policy already exists", True + with self._repository(write=True) as policies: + existing = policies.find_unique(country_id, policy_hash, label or None) + if existing: + return existing["id"], "Policy already exists", True - policy_id = self.policies.create( - country_id, - label, - policy_json, - policy_hash, - COUNTRY_PACKAGE_VERSIONS[country_id], - ) - return policy_id, "Policy created", False + policy_id = policies.create( + country_id, + label, + policy_json, + policy_hash, + COUNTRY_PACKAGE_VERSIONS[country_id], + ) + return policy_id, "Policy created", False def _create_new_policy( self, @@ -68,9 +87,11 @@ def _create_new_policy( label: str | None, api_version: str, ) -> None: - self.policies.create(country_id, label, policy_json, policy_hash, api_version) + with self._repository(write=True) as policies: + policies.create(country_id, label, policy_json, policy_hash, api_version) def _get_unique_policy_with_label( self, country_id: str, policy_hash: str, label: str ) -> dict | None: - return self.policies.find_unique(country_id, policy_hash, label or None) + with self._repository() as policies: + return policies.find_unique(country_id, policy_hash, label or None) diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index 50fd4b3cc..2ae2b228b 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,20 +1,36 @@ from __future__ import annotations +from contextlib import contextmanager from typing import Any from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import UserDAO +from policyengine_api.data.v1_daos import UserDAO, V1UnitOfWork class UserService: - def __init__(self, users: UserDAO | None = None): + def __init__( + self, + users: UserDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._users = users + self._unit_of_work = unit_of_work @property - def users(self) -> UserDAO: - if self._users is None: - self._users = UserDAO(build_v1_session_manager()) - return self._users + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work + + @contextmanager + def _repository(self, *, write: bool = False): + if self._users is not None: + yield self._users + return + boundary = self.unit_of_work.transaction if write else self.unit_of_work.read + with boundary() as repositories: + yield repositories.users def create_profile( self, @@ -23,18 +39,20 @@ def create_profile( username: str | None, user_since: int, ) -> tuple[bool, Any]: - row = self.get_profile(auth0_id=auth0_id) - if row is not None: - return False, row - self.users.create_profile(auth0_id, username, primary_country, user_since) - return True, self.get_profile(auth0_id=auth0_id) + with self._repository(write=True) as users: + row = users.get_profile(auth0_id=auth0_id) + if row is not None: + return False, row + users.create_profile(auth0_id, username, primary_country, user_since) + return True, users.get_profile(auth0_id=auth0_id) def get_profile( self, auth0_id: str | None = None, user_id: int | None = None ) -> Any | None: if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") - return self.users.get_profile(user_id=user_id, auth0_id=auth0_id) + with self._repository() as users: + return users.get_profile(user_id=user_id, auth0_id=auth0_id) def update_profile( self, @@ -45,9 +63,10 @@ def update_profile( ) -> bool: if user_id is None: raise ValueError("you must specify either auth0_id or user_id") - return self.users.update_profile( - user_id, - primary_country=primary_country, - username=username, - user_since=user_since, - ) + with self._repository(write=True) as users: + return users.update_profile( + user_id, + primary_country=primary_country, + username=username, + user_since=user_since, + ) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 9f2c428a5..74f530312 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -5,6 +5,8 @@ import httpx import pytest +pytest_plugins = ("tests.integration.stage7_mysql",) + INTEGRATION_TIMEOUT_SECONDS = float( os.environ.get("STAGING_API_TEST_TIMEOUT_SECONDS", "900") ) diff --git a/tests/integration/test_stage7_toy_database.py b/tests/integration/test_stage7_toy_database.py new file mode 100644 index 000000000..c0e3f5964 --- /dev/null +++ b/tests/integration/test_stage7_toy_database.py @@ -0,0 +1,47 @@ +"""MySQL qualification for the disposable Stage 7 database.""" + +from alembic import command +from sqlalchemy import inspect + +from policyengine_api.data.v1_models import V1Base +from policyengine_api.scripts.qualify_stage7_toy import ( + compare_stage7_schema, + qualify_stage7_toy, +) +from tests.integration.stage7_mysql import alembic_config + + +def test_mysql_toy_database_upgrades_and_exercises_every_dao_domain(stage7_mysql): + database_url, _ = stage7_mysql + result = qualify_stage7_toy(database_url) + + assert result == { + "alembic_head": True, + "policy": True, + "household": True, + "computed_household": True, + "user": True, + "user_policy": True, + "economy": True, + "simulation": True, + "report": True, + "report_alias": True, + "analysis": True, + "tracer": True, + "reform_impact": True, + } + assert compare_stage7_schema(database_url) == [] + + +def test_mysql_toy_database_downgrades_and_reupgrades_cleanly(stage7_mysql): + database_url, engine = stage7_mysql + config = alembic_config(database_url) + qualify_stage7_toy(database_url) + + command.downgrade(config, "base") + remaining_tables = set(inspect(engine).get_table_names()) + assert not set(V1Base.metadata.tables) & remaining_tables + + command.upgrade(config, "head") + command.check(config) + assert compare_stage7_schema(database_url) == [] diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index 9791f207b..e17baf4fc 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -1,54 +1,71 @@ from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import HouseholdDAO, PolicyDAO, UserDAO +from policyengine_api.data.v1_daos import V1UnitOfWork from tests.unit.data.sqlite_schema import create_sqlite_v1_schema -def _daos(): +def _unit_of_work(): manager = build_sqlite_session_manager() create_sqlite_v1_schema(manager) - return PolicyDAO(manager), HouseholdDAO(manager), UserDAO(manager) + return V1UnitOfWork(manager) def test_policy_dao_round_trips_legacy_mapping_shape(): - policies, _, _ = _daos() - policy_id = policies.create("us", "Reform", {"gov.irs": 1}, "hash", "1.0") - assert policy_id == 1 - assert policies.get("us", policy_id) == { - "id": 1, - "country_id": "us", - "label": "Reform", - "api_version": "1.0", - "policy_json": {"gov.irs": 1}, - "policy_hash": "hash", - } + uow = _unit_of_work() + with uow.transaction() as repositories: + policy_id = repositories.policies.create( + "us", "Reform", {"gov.irs": 1}, "hash", "1.0" + ) + with uow.read() as repositories: + assert policy_id == 1 + assert repositories.policies.get("us", policy_id) == { + "id": 1, + "country_id": "us", + "label": "Reform", + "api_version": "1.0", + "policy_json": {"gov.irs": 1}, + "policy_hash": "hash", + } def test_policy_dao_allocates_ids_and_detects_existing_policy(): - policies, _, _ = _daos() - assert policies.create("us", None, {}, "one", "1.0") == 1 - assert policies.create("uk", None, {}, "two", "1.0") == 2 - assert policies.find_unique("us", "one", None)["id"] == 1 + uow = _unit_of_work() + with uow.transaction() as repositories: + assert repositories.policies.create("us", None, {}, "one", "1.0") == 1 + assert repositories.policies.create("uk", None, {}, "two", "1.0") == 2 + with uow.read() as repositories: + assert repositories.policies.find_unique("us", "one", None)["id"] == 1 def test_household_dao_creates_updates_and_reads(): - _, households, _ = _daos() - household_id = households.create("us", "Home", {"people": {}}, "h", "1.0") - households.update( - "us", - household_id, - "Updated", - {"people": {"you": {}}}, - "updated-hash", - "2.0", - ) - assert households.get("us", household_id)["label"] == "Updated" - assert households.get("uk", household_id) is None + uow = _unit_of_work() + with uow.transaction() as repositories: + household_id = repositories.households.create( + "us", "Home", {"people": {}}, "h", "1.0" + ) + repositories.households.update( + "us", + household_id, + "Updated", + {"people": {"you": {}}}, + "updated-hash", + "2.0", + ) + with uow.read() as repositories: + assert repositories.households.get("us", household_id)["label"] == "Updated" + assert repositories.households.get("uk", household_id) is None def test_user_dao_profile_lookup_precedence(): - _, _, users = _daos() - user_id = users.create_profile("auth0|one", "person", "us", 123) - assert users.get_profile(auth0_id="auth0|one")["user_id"] == user_id - assert ( - users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] == "auth0|one" - ) + uow = _unit_of_work() + with uow.transaction() as repositories: + user_id = repositories.users.create_profile("auth0|one", "person", "us", 123) + with uow.read() as repositories: + assert ( + repositories.users.get_profile(auth0_id="auth0|one")["user_id"] == user_id + ) + assert ( + repositories.users.get_profile(user_id=user_id, auth0_id="wrong")[ + "auth0_id" + ] + == "auth0|one" + ) diff --git a/tests/unit/data/test_v1_unit_of_work.py b/tests/unit/data/test_v1_unit_of_work.py new file mode 100644 index 000000000..cb8b9b84d --- /dev/null +++ b/tests/unit/data/test_v1_unit_of_work.py @@ -0,0 +1,37 @@ +import pytest + +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import V1UnitOfWork +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema + + +def _unit_of_work() -> V1UnitOfWork: + manager = build_sqlite_session_manager() + create_sqlite_v1_schema(manager) + return V1UnitOfWork(manager) + + +def test_unit_of_work_commits_all_repositories_once(): + uow = _unit_of_work() + + with uow.transaction() as repositories: + policy_id = repositories.policies.create("us", None, {}, "policy", "1") + household_id = repositories.households.create("us", None, {}, "household", "1") + + with uow.read() as repositories: + assert repositories.policies.get("us", policy_id) is not None + assert repositories.households.get("us", household_id) is not None + + +def test_unit_of_work_rolls_back_every_repository_on_failure(): + uow = _unit_of_work() + + with pytest.raises(RuntimeError, match="abort"): + with uow.transaction() as repositories: + repositories.policies.create("us", None, {}, "policy", "1") + repositories.users.create_profile("auth0|one", "person", "us", 1) + raise RuntimeError("abort") + + with uow.read() as repositories: + assert repositories.policies.get("us", 1) is None + assert repositories.users.get_profile(auth0_id="auth0|one") is None diff --git a/tests/unit/services/test_update_profile_service.py b/tests/unit/services/test_update_profile_service.py index 7d51f1ea3..d0bacb8f6 100644 --- a/tests/unit/services/test_update_profile_service.py +++ b/tests/unit/services/test_update_profile_service.py @@ -86,7 +86,10 @@ def test_update_profile_with_database_error( def mock_dao_error(*args, **kwargs): raise Exception("Database error") - monkeypatch.setattr(service.users, "update_profile", mock_dao_error) + class FailingUsers: + update_profile = staticmethod(mock_dao_error) + + monkeypatch.setattr(service, "_users", FailingUsers()) # WHEN we call update_profile # THEN an exception should be raised diff --git a/tests/unit/test_stage7_toy_qualification.py b/tests/unit/test_stage7_toy_qualification.py index 884d7595c..8392faa3e 100644 --- a/tests/unit/test_stage7_toy_qualification.py +++ b/tests/unit/test_stage7_toy_qualification.py @@ -1,24 +1,7 @@ from pathlib import Path -from policyengine_api.scripts.qualify_stage7_toy import ( - compare_stage7_schema, - qualify_stage7_toy, -) - - -def test_toy_qualification_exercises_migrated_data_paths(tmp_path: Path): - result = qualify_stage7_toy(f"sqlite+pysqlite:///{tmp_path / 'stage7-toy.db'}") - assert result == { - "alembic_head": True, - "policy": True, - "household": True, - "user": True, - "simulation": True, - "report": True, - "analysis": True, - "tracer": True, - "reform_impact": True, - } +from policyengine_api.scripts.stage7_database import assert_safe_toy_database_url +import pytest def test_legacy_daos_have_been_removed(): @@ -33,9 +16,41 @@ def test_legacy_daos_have_been_removed(): assert "LegacyReportDAO" not in sources -def test_schema_comparison_is_read_only_and_reports_no_fresh_database_drift( - tmp_path: Path, +@pytest.mark.parametrize( + "database_url", + [ + "mysql+pymysql://toy:toy@127.0.0.1:3307/policyengine_stage7_toy", + "mysql+pymysql://toy:toy@localhost:3307/custom_toy", + ], +) +def test_toy_database_safety_guard_accepts_only_local_mysql_toy_databases( + database_url: str, ): - database_url = f"sqlite+pysqlite:///{tmp_path / 'comparison.db'}" - qualify_stage7_toy(database_url) - assert compare_stage7_schema(database_url) == [] + assert_safe_toy_database_url(database_url) + + +@pytest.mark.parametrize( + "database_url", + [ + "mysql+pymysql://toy:toy@prod.example.com/policyengine_stage7_toy", + "mysql+pymysql://toy:toy@127.0.0.1/policyengine", + "postgresql://toy:toy@127.0.0.1/policyengine_stage7_toy", + "sqlite+pysqlite:///policyengine_stage7_toy.db", + ], +) +def test_toy_database_safety_guard_rejects_unsafe_targets(database_url: str): + with pytest.raises(ValueError, match="local MySQL.*_toy"): + assert_safe_toy_database_url(database_url) + + +def test_stage7_toy_database_has_local_scaffold_and_test_targets(): + repo = Path(__file__).parents[2] + compose = (repo / "compose.stage7-toy.yml").read_text(encoding="utf-8") + makefile = (repo / "Makefile").read_text(encoding="utf-8") + + assert "mysql:8.0" in compose + assert "policyengine_stage7_toy" in compose + assert "healthcheck:" in compose + assert "stage7-toy-up:" in makefile + assert "stage7-toy-test:" in makefile + assert "stage7-toy-down:" in makefile From 0f38f01b8f18c184a6af887d40709c32e212fdb8 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:12:48 +0300 Subject: [PATCH 35/89] refactor: move v1 domains onto typed repositories --- policyengine_api/country.py | 16 +- policyengine_api/data/v1_daos.py | 364 +++++++++++------- .../endpoints/economy/reform_impact.py | 30 +- policyengine_api/endpoints/household.py | 78 ++-- policyengine_api/endpoints/policy.py | 180 ++------- policyengine_api/endpoints/simulation.py | 12 +- .../scripts/qualify_stage7_toy.py | 141 ++++--- .../services/ai_analysis_service.py | 34 +- .../services/reform_impacts_service.py | 140 ++++--- .../services/tracer_analysis_service.py | 29 +- tests/contract/test_v1_route_contracts.py | 28 +- tests/integration/test_stage7_mysql_parity.py | 91 +++++ .../integration/test_stage7_mysql_runtime.py | 135 +++++++ tests/unit/data/test_local_daos.py | 68 ++-- tests/unit/data/test_ordinary_v1_daos.py | 70 ++++ tests/unit/data/test_stage7_no_direct_sql.py | 16 + tests/unit/endpoints/test_get_simulations.py | 2 +- 17 files changed, 870 insertions(+), 564 deletions(-) create mode 100644 tests/integration/test_stage7_mysql_parity.py create mode 100644 tests/integration/test_stage7_mysql_runtime.py create mode 100644 tests/unit/data/test_ordinary_v1_daos.py diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 42dca1d46..f202883be 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -23,7 +23,7 @@ build_congressional_district_metadata, ) -from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao +from policyengine_api.data.v1_daos import runtime_v1_unit_of_work from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, get_bundle_default_dataset_option, @@ -430,23 +430,17 @@ def calculate( tracer_output = simulation.tracer.computation_log log_lines = tracer_output.lines(aggregate=False, max_depth=10) - log_json = json.dumps(log_lines) if household_id is not None and policy_id is not None: # write to local database - runtime_sqlalchemy_dao(local=True).query( - """ - INSERT INTO tracers (household_id, policy_id, country_id, api_version, tracer_output) - VALUES (?, ?, ?, ?, ?) - """, - ( + with runtime_v1_unit_of_work(local=True).transaction() as repositories: + repositories.tracers.create( household_id, policy_id, self.country_id, COUNTRY_PACKAGE_VERSIONS[self.country_id], - log_json, - ), - ) + log_lines, + ) return household diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 9fc8f3617..18a8a46e8 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -142,6 +142,15 @@ def find_unique( ) return _mapping(model) if model else None + def search(self, country_id: str, query: str) -> list[dict[str, Any]]: + models = self.session.scalars( + select(Policy).where( + Policy.country_id == country_id, + Policy.label.contains(query, autoescape=True), + ) + ) + return [_mapping(model) for model in models] + def create( self, country_id: str, @@ -219,23 +228,42 @@ def update( class ComputedHouseholdDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def create(self, **values: Any) -> None: - self.sessions.run_in_transaction( - lambda session: session.add(ComputedHousehold(**values)) + self.session.add(ComputedHousehold(**values)) + + def upsert(self, **values: Any) -> None: + identity = ( + values["household_id"], + values["policy_id"], + values["country_id"], ) + model = self.session.get(ComputedHousehold, identity) + if model is None: + self.session.add(ComputedHousehold(**values)) + return + for key, value in values.items(): + setattr(model, key, value) def get( - self, household_id: int, policy_id: int, country_id: str + self, + household_id: int, + policy_id: int, + country_id: str, + *, + api_version: str | None = None, ) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get( - ComputedHousehold, - (household_id, policy_id, country_id), - ) - return _mapping(model) if model else None + statement = select(ComputedHousehold).where( + ComputedHousehold.household_id == household_id, + ComputedHousehold.policy_id == policy_id, + ComputedHousehold.country_id == country_id, + ) + if api_version is not None: + statement = statement.where(ComputedHousehold.api_version == api_version) + model = self.session.scalar(statement) + return _mapping(model) if model else None class UserDAO: @@ -293,6 +321,12 @@ def __init__(self, session: Session): self.policies = PolicyDAO(session) self.households = HouseholdDAO(session) self.users = UserDAO(session) + self.computed_households = ComputedHouseholdDAO(session) + self.user_policies = UserPolicyDAO(session) + self.economies = EconomyDAO(session) + self.analyses = AnalysisDAO(session) + self.reform_impacts = ReformImpactDAO(session) + self.tracers = TracerDAO(session) class V1UnitOfWork: @@ -312,91 +346,130 @@ def transaction(self) -> Iterator[V1Repositories]: yield V1Repositories(session) +_runtime_unit_of_work: dict[bool, V1UnitOfWork] = {} + + +def runtime_v1_unit_of_work(*, local: bool = False) -> V1UnitOfWork: + """Return the process-local unit of work for the selected v1 database.""" + + if local not in _runtime_unit_of_work: + from policyengine_api.data.orm import build_v1_session_manager + + _runtime_unit_of_work[local] = V1UnitOfWork( + build_v1_session_manager(local=local) + ) + return _runtime_unit_of_work[local] + + class UserPolicyDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + IDENTITY_FIELDS = ( + "country_id", + "reform_id", + "baseline_id", + "user_id", + "year", + "geography", + "reform_label", + "baseline_label", + "dataset", + ) + + def __init__(self, session: Session): + self.session = session def create(self, **values: Any) -> int: - def operation(session): - model = UserPolicy(**values) - session.add(model) - session.flush() - return model.id + model = UserPolicy(**values) + self.session.add(model) + self.session.flush() + return model.id - return self.sessions.run_in_transaction(operation) + def find_unique(self, **values: Any) -> dict[str, Any] | None: + model = self.session.scalar( + select(UserPolicy).where( + *( + getattr(UserPolicy, field) == values[field] + for field in self.IDENTITY_FIELDS + ) + ) + ) + return _mapping(model) if model else None def get(self, user_policy_id: int) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get(UserPolicy, user_policy_id) - return _mapping(model) if model else None + model = self.session.get(UserPolicy, user_policy_id) + return _mapping(model) if model else None + + def list_for_user(self, country_id: str, user_id: str) -> list[dict[str, Any]]: + models = self.session.scalars( + select(UserPolicy).where( + UserPolicy.country_id == country_id, + UserPolicy.user_id == user_id, + ) + ) + return [_mapping(model) for model in models] + + def update(self, user_policy_id: int, values: dict[str, Any]) -> bool: + model = self.session.get(UserPolicy, user_policy_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True class EconomyDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def create(self, **values: Any) -> int: - def operation(session): - model = Economy(**values) - session.add(model) - session.flush() - return model.economy_id - - return self.sessions.run_in_transaction(operation) + model = Economy(**values) + self.session.add(model) + self.session.flush() + return model.economy_id def get(self, economy_id: int) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get(Economy, economy_id) - return _mapping(model) if model else None + model = self.session.get(Economy, economy_id) + return _mapping(model) if model else None class AnalysisDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def get(self, prompt: str) -> str | None: - with self.sessions.session() as session: - model = session.scalar( - select(Analysis) - .where( - Analysis.prompt == prompt, - Analysis.status.in_(("complete", "ok")), - ) - .order_by(Analysis.prompt_id.desc()) + model = self.session.scalar( + select(Analysis) + .where( + Analysis.prompt == prompt, + Analysis.status.in_(("complete", "ok")), ) - return model.analysis if model else None + .order_by(Analysis.prompt_id.desc()) + ) + return model.analysis if model else None def store(self, prompt: str, analysis: str | None, status: str) -> int: - def operation(session): - model = Analysis(prompt=prompt, analysis=analysis, status=status) - session.add(model) - session.flush() - return model.prompt_id - - return self.sessions.run_in_transaction(operation) + model = Analysis(prompt=prompt, analysis=analysis, status=status) + self.session.add(model) + self.session.flush() + return model.prompt_id class ReformImpactDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def create(self, **values: Any) -> int: - def operation(session): - model = ReformImpact(**values) - session.add(model) - session.flush() - return model.reform_impact_id - - return self.sessions.run_in_transaction(operation) + model = ReformImpact(**values) + self.session.add(model) + self.session.flush() + return model.reform_impact_id def find(self, *, execution_id: str) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - return _mapping(model) if model else None + model = self.session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + return _mapping(model) if model else None @staticmethod def _scope(statement, **filters: Any): @@ -405,79 +478,82 @@ def _scope(statement, **filters: Any): ) def list(self, **filters: Any) -> list[dict[str, Any]]: - with self.sessions.session() as session: - models = session.scalars( - self._scope(select(ReformImpact), **filters).order_by( - ReformImpact.start_time.desc() - ) + models = self.session.scalars( + self._scope(select(ReformImpact), **filters).order_by( + ReformImpact.start_time.desc() ) - return [_mapping(model) for model in models] + ) + return [_mapping(model) for model in models] + + def list_recent(self, limit: int) -> list[dict[str, Any]]: + models = self.session.scalars( + select(ReformImpact).order_by(ReformImpact.start_time.desc()).limit(limit) + ) + return [_mapping(model) for model in models] def list_by_options_hash( self, options_hash: str, options_hash_prefix: str, **filters: Any ) -> list[dict[str, Any]]: - with self.sessions.session() as session: - statement = self._scope(select(ReformImpact), **filters).where( - or_( - ReformImpact.options_hash == options_hash, - ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), - ) + statement = self._scope(select(ReformImpact), **filters).where( + or_( + ReformImpact.options_hash == options_hash, + ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), ) - models = session.scalars( - statement.order_by( - (ReformImpact.options_hash == options_hash).desc(), - ReformImpact.start_time.desc(), - ) + ) + models = self.session.scalars( + statement.order_by( + (ReformImpact.options_hash == options_hash).desc(), + ReformImpact.start_time.desc(), ) - return [_mapping(model) for model in models] + ) + return [_mapping(model) for model in models] def delete_computing(self, **filters: Any) -> None: - def operation(session): - session.execute( - self._scope(delete(ReformImpact), **filters).where( - ReformImpact.status == "computing" - ) + self.session.execute( + self._scope(delete(ReformImpact), **filters).where( + ReformImpact.status == "computing" ) + ) - self.sessions.run_in_transaction(operation) - - def fail(self, execution_id: str, message: str, finished_at: datetime) -> bool: - def operation(session): - model = session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - if model is None: - return False - model.status = "error" + def set_message(self, message: str, **filters: Any) -> bool: + models = self.session.scalars( + self._scope(select(ReformImpact), **filters) + ).all() + for model in models: model.message = message - model.end_time = finished_at - return True + return bool(models) - return self.sessions.run_in_transaction(operation) + def fail(self, execution_id: str, message: str, finished_at: datetime) -> bool: + model = self.session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if model is None: + return False + model.status = "error" + model.message = message + model.end_time = finished_at + return True def complete(self, execution_id: str, result: Any, finished_at: datetime) -> bool: - def operation(session): - model = session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - if model is None: - return False - model.status = "ok" - model.message = "Completed" - model.reform_impact_json = result - model.end_time = finished_at - return True - - return self.sessions.run_in_transaction(operation) + model = self.session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if model is None: + return False + model.status = "ok" + model.message = "Completed" + model.reform_impact_json = result + model.end_time = finished_at + return True class TracerDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def create( self, @@ -487,19 +563,16 @@ def create( api_version: str, tracer_output: Any, ) -> int: - def operation(session): - model = Tracer( - household_id=household_id, - policy_id=policy_id, - country_id=country_id, - api_version=api_version, - tracer_output=tracer_output, - ) - session.add(model) - session.flush() - return model.id - - return self.sessions.run_in_transaction(operation) + model = Tracer( + household_id=household_id, + policy_id=policy_id, + country_id=country_id, + api_version=api_version, + tracer_output=tracer_output, + ) + self.session.add(model) + self.session.flush() + return model.id def get( self, @@ -508,16 +581,15 @@ def get( country_id: str, api_version: str | None = None, ) -> dict[str, Any] | None: - with self.sessions.session() as session: - statement = select(Tracer).where( - Tracer.household_id == household_id, - Tracer.policy_id == policy_id, - Tracer.country_id == country_id, - ) - if api_version is not None: - statement = statement.where(Tracer.api_version == api_version) - model = session.scalar(statement.order_by(Tracer.id.desc())) - return _mapping(model) if model else None + statement = select(Tracer).where( + Tracer.household_id == household_id, + Tracer.policy_id == policy_id, + Tracer.country_id == country_id, + ) + if api_version is not None: + statement = statement.where(Tracer.api_version == api_version) + model = self.session.scalar(statement.order_by(Tracer.id.desc())) + return _mapping(model) if model else None class SimulationDAO: diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py index a53d4a238..842959cf7 100644 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ b/policyengine_api/endpoints/economy/reform_impact.py @@ -1,4 +1,4 @@ -from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao +from policyengine_api.data.v1_daos import runtime_v1_unit_of_work def set_comment_on_job( @@ -11,22 +11,14 @@ def set_comment_on_job( time_period, options_hash, ): - query = ( - "UPDATE reform_impact SET message = ? WHERE country_id = ? AND " - "reform_policy_id = ? AND baseline_policy_id = ? AND region = ? AND " - "time_period = ? AND options_hash = ? AND dataset = ?" - ) - - runtime_sqlalchemy_dao(local=True).query( - query, - ( + with runtime_v1_unit_of_work(local=True).transaction() as repositories: + repositories.reform_impacts.set_message( comment, - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - ), - ) + country_id=country_id, + reform_policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + region=region, + time_period=time_period, + options_hash=options_hash, + dataset=dataset, + ) diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index 4eca0d9e1..cbcf3b358 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -1,4 +1,4 @@ -from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao +from policyengine_api.data.v1_daos import runtime_v1_unit_of_work import json from flask import Response, request from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS @@ -111,14 +111,13 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Look in computed_households to see if already computed - row = ( - runtime_sqlalchemy_dao(local=True) - .query( - "SELECT * FROM computed_household WHERE household_id = ? AND policy_id = ? AND api_version = ?", - (household_id, policy_id, api_version), + with runtime_v1_unit_of_work(local=True).read() as repositories: + row = repositories.computed_households.get( + int(household_id), + int(policy_id), + country_id, + api_version=api_version, ) - .fetchone() - ) if row is not None: result = dict( @@ -129,7 +128,10 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st computed_household_json=row["computed_household_json"], status=row["status"], ) - result["result"] = json.loads(result["computed_household_json"]) + computed = result["computed_household_json"] + result["result"] = ( + json.loads(computed) if isinstance(computed, str) else computed + ) del result["computed_household_json"] return dict( status="ok", @@ -139,18 +141,18 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the household table - row = ( - runtime_sqlalchemy_dao() - .query( - "SELECT * FROM household WHERE id = ? AND country_id = ?", - (household_id, country_id), - ) - .fetchone() - ) + with runtime_v1_unit_of_work().read() as repositories: + row = repositories.households.get(country_id, int(household_id)) + policy_row = repositories.policies.get(country_id, int(policy_id)) if row is not None: household = dict(row) - household["household_json"] = json.loads(household["household_json"]) + household_json = household["household_json"] + household["household_json"] = ( + json.loads(household_json) + if isinstance(household_json, str) + else household_json + ) else: response_body = dict( status="error", @@ -171,18 +173,12 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the policy table - row = ( - runtime_sqlalchemy_dao() - .query( - "SELECT * FROM policy WHERE id = ? AND country_id = ?", - (policy_id, country_id), + if policy_row is not None: + policy = dict(policy_row) + policy_json = policy["policy_json"] + policy["policy_json"] = ( + json.loads(policy_json) if isinstance(policy_json, str) else policy_json ) - .fetchone() - ) - - if row is not None: - policy = dict(row) - policy["policy_json"] = json.loads(policy["policy_json"]) else: response_body = dict( status="error", @@ -224,22 +220,14 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Store the result in the computed_household table - try: - runtime_sqlalchemy_dao(local=True).query( - "INSERT INTO computed_household (country_id, household_id, policy_id, computed_household_json, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - household_id, - policy_id, - json.dumps(result), - api_version, - ), - ) - except Exception: - # Update the result if it already exists - runtime_sqlalchemy_dao(local=True).query( - "UPDATE computed_household SET computed_household_json = ? WHERE country_id = ? AND household_id = ? AND policy_id = ?", - (json.dumps(result), country_id, household_id, policy_id), + with runtime_v1_unit_of_work(local=True).transaction() as repositories: + repositories.computed_households.upsert( + country_id=country_id, + household_id=int(household_id), + policy_id=int(policy_id), + computed_household_json=result, + api_version=api_version, + status="complete", ) response_body = dict( diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index 6df74931f..4c8fd6b84 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -1,5 +1,5 @@ from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao +from policyengine_api.data.v1_daos import runtime_v1_unit_of_work import json from flask import Response, request @@ -31,12 +31,8 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - results = runtime_sqlalchemy_dao().query( - "SELECT id, label, policy_hash FROM policy WHERE country_id = ? AND label LIKE ?", - (country_id, f"%{query}%"), - ) - - results = results.fetchall() + with runtime_v1_unit_of_work().read() as repositories: + results = repositories.policies.search(country_id, query) if not results: body = dict( @@ -101,32 +97,24 @@ def set_user_policy(country_id: str) -> dict: budgetary_impact = payload.pop("budgetary_impact", None) type = payload.pop("type", None) - # The following code is a workaround to the fact that - # SQLite's cursor method does not properly convert - # 'WHERE x = None' to 'WHERE x IS NULL'; though SQLite - # supports searching and setting with 'WHERE x IS y', - # the production MySQL does not, requiring this - - # This workaround should be removed if and when a proper - # ORM package is added to the API, and this package's - # sanitization methods should be utilized instead - nullable_keys = [] - not_null_values = [] - possible_nulls = { + values = { + "country_id": country_id, + "reform_id": reform_id, "reform_label": reform_label, + "baseline_id": baseline_id, "baseline_label": baseline_label, + "user_id": user_id, + "year": year, + "geography": geography, "dataset": dataset, + "number_of_provisions": number_of_provisions, + "api_version": api_version, + "added_date": added_date, + "updated_date": updated_date, + "budgetary_impact": budgetary_impact, + "type": type, } - for key, value in possible_nulls.items(): - if not value: - nullable_keys.append(f"{key} IS NULL") - else: - nullable_keys.append(f"{key} = ?") - not_null_values.append(value) - - nullable_key_string = " AND ".join(nullable_keys) - # When setting a user policy, "unique" records contain # a unique set of the following pieces of data: # country_id, reform_id, baseline_id, user_id, year, @@ -137,100 +125,24 @@ def set_user_policy(country_id: str) -> dict: # to be tested; type is not yet implemented try: - row = ( - runtime_sqlalchemy_dao() - .query( - f"SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? AND {nullable_key_string}", - ( - country_id, - reform_id, - baseline_id, - user_id, - year, - geography, - *not_null_values, - ), - ) - .fetchone() - ) - if row is not None: - readable_row = dict(row) - - response = dict( - status="ok", - message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", - result=dict(id=readable_row["id"]), - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - try: - # Unfortunately, it's not possible to use RETURNING - # with SQLite3 without rewriting the PolicyEngineDatabase - # object or implementing a true ORM, thus the double query - - query = ( - "INSERT INTO user_policies (country_id, reform_label, " - "reform_id, baseline_label, baseline_id, user_id, year, " - "geography, number_of_provisions, api_version, added_date, " - "updated_date, budgetary_impact, type, dataset) VALUES " - "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) - - runtime_sqlalchemy_dao().query( - query, - ( - country_id, - reform_label, - reform_id, - baseline_label, - baseline_id, - user_id, - year, - geography, - number_of_provisions, - api_version, - added_date, - updated_date, - budgetary_impact, - type, - dataset, - ), - ) - - # "IS NULL" is not treated as the same as - # "= None" in SQL - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) - - row = ( - runtime_sqlalchemy_dao() - .query( - query, - tuple(params), - ) - .fetchone() - ) - + with runtime_v1_unit_of_work().transaction() as repositories: + row = repositories.user_policies.find_unique(**values) + if row is None: + user_policy_id = repositories.user_policies.create(**values) + row = repositories.user_policies.get(user_policy_id) + else: + readable_row = dict(row) + + response = dict( + status="ok", + message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", + result=dict(id=readable_row["id"]), + ) + return Response( + json.dumps(response), + status=200, + mimetype="application/json", + ) except Exception as e: return Response( json.dumps( @@ -277,14 +189,8 @@ def get_user_policy(country_id: str, user_id: str) -> dict: """ # Get the policy record for a given policy ID. - rows = ( - runtime_sqlalchemy_dao() - .query( - "SELECT * FROM user_policies WHERE country_id = ? AND user_id = ?", - (country_id, user_id), - ) - .fetchall() - ) + with runtime_v1_unit_of_work().read() as repositories: + rows = repositories.user_policies.list_for_user(country_id, user_id) rows_parsed = [ dict( @@ -394,19 +300,9 @@ def update_user_policy(country_id: str) -> dict: mimetype="application/json", ) - # Construct the relevant UPDATE request from whitelisted keys. - setter_array = [] - args = [] - for key in payload: - setter_array.append(f"{key} = ?") - args.append(payload[key]) - setter_phrase = ", ".join(setter_array) - - args.append(user_policy_id) - sql_request = f"UPDATE user_policies SET {setter_phrase} WHERE id = ?" - try: - runtime_sqlalchemy_dao().query(sql_request, (tuple(args))) + with runtime_v1_unit_of_work().transaction() as repositories: + repositories.user_policies.update(user_policy_id, payload) except Exception as e: return Response( json.dumps( diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index c64e54b7c..a0f22df8c 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -1,4 +1,4 @@ -from policyengine_api.data.v1_daos import runtime_sqlalchemy_dao +from policyengine_api.data.v1_daos import runtime_v1_unit_of_work """ @@ -42,14 +42,8 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - result = ( - runtime_sqlalchemy_dao(local=True) - .query( - "SELECT * FROM reform_impact ORDER BY start_time DESC LIMIT ?", - (max_results,), - ) - .fetchall() - ) + with runtime_v1_unit_of_work(local=True).read() as repositories: + result = repositories.reform_impacts.list_recent(max_results) # Format into [{}] diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py index 6a9f0e5c0..bb1e534cd 100644 --- a/policyengine_api/scripts/qualify_stage7_toy.py +++ b/policyengine_api/scripts/qualify_stage7_toy.py @@ -11,14 +11,8 @@ from policyengine_api.constants import REPO from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_daos import ( - AnalysisDAO, - ComputedHouseholdDAO, - EconomyDAO, - ReformImpactDAO, ReportDAO, SimulationDAO, - TracerDAO, - UserPolicyDAO, V1UnitOfWork, ) from policyengine_api.data.v1_models import V1Base @@ -46,14 +40,8 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: engine = create_engine(database_url) sessions = SessionManager(engine) unit_of_work = V1UnitOfWork(sessions) - computed_households = ComputedHouseholdDAO(sessions) - user_policies = UserPolicyDAO(sessions) - economies = EconomyDAO(sessions) simulations = SimulationDAO(sessions) reports = ReportDAO(sessions) - analyses = AnalysisDAO(sessions) - tracers = TracerDAO(sessions) - impacts = ReformImpactDAO(sessions) with unit_of_work.transaction() as repositories: policy_id = repositories.policies.create("us", "Toy", {}, "toy-policy", "toy") @@ -61,43 +49,60 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: "us", "Toy", {}, "toy-household", "toy" ) user_id = repositories.users.create_profile("toy|user", "toy-user", "us", 1) - computed_households.create( - household_id=household_id, - policy_id=policy_id, - country_id="us", - api_version="toy", - computed_household_json={"qualified": True}, - status="complete", - ) - user_policy_id = user_policies.create( - country_id="us", - reform_id=policy_id, - reform_label="Toy", - baseline_id=policy_id, - baseline_label="Toy", - user_id=str(user_id), - year="2026", - geography="us", - dataset="default", - number_of_provisions=0, - api_version="toy", - added_date=1, - updated_date=1, - budgetary_impact=None, - type="reform", - ) - economy_id = economies.create( - policy_id=policy_id, - country_id="us", - region="us", - time_period="2026", - options_json={}, - options_hash="toy-economy", - api_version="toy", - economy_json={"qualified": True}, - status="complete", - message=None, - ) + repositories.computed_households.create( + household_id=household_id, + policy_id=policy_id, + country_id="us", + api_version="toy", + computed_household_json={"qualified": True}, + status="complete", + ) + user_policy_id = repositories.user_policies.create( + country_id="us", + reform_id=policy_id, + reform_label="Toy", + baseline_id=policy_id, + baseline_label="Toy", + user_id=str(user_id), + year="2026", + geography="us", + dataset="default", + number_of_provisions=0, + api_version="toy", + added_date=1, + updated_date=1, + budgetary_impact=None, + type="reform", + ) + economy_id = repositories.economies.create( + policy_id=policy_id, + country_id="us", + region="us", + time_period="2026", + options_json={}, + options_hash="toy-economy", + api_version="toy", + economy_json={"qualified": True}, + status="complete", + message=None, + ) + repositories.analyses.store("toy prompt", "toy answer", "complete") + repositories.tracers.create(household_id, policy_id, "us", "toy", ["toy trace"]) + impact_id = repositories.reform_impacts.create( + baseline_policy_id=policy_id, + reform_policy_id=policy_id, + country_id="us", + region="us", + dataset="default", + time_period="2026", + options_json={}, + options_hash="toy-options", + api_version="toy", + reform_impact_json={}, + status="computing", + start_time=datetime(2026, 1, 1), + execution_id="toy-impact", + ) simulation_id = simulations.create( country_id="us", api_version="toy", @@ -125,44 +130,32 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: trigger_type="qualification", ) reports.set_alias(900_001, report_id) - analyses.store("toy prompt", "toy answer", "complete") - tracers.create(household_id, policy_id, "us", "toy", ["toy trace"]) - impact_id = impacts.create( - baseline_policy_id=policy_id, - reform_policy_id=policy_id, - country_id="us", - region="us", - dataset="default", - time_period="2026", - options_json={}, - options_hash="toy-options", - api_version="toy", - reform_impact_json={}, - status="computing", - start_time=datetime(2026, 1, 1), - execution_id="toy-impact", - ) with unit_of_work.read() as repositories: core_results = { "policy": repositories.policies.get("us", policy_id) is not None, "household": repositories.households.get("us", household_id) is not None, "user": repositories.users.get_profile(user_id=user_id) is not None, + "computed_household": repositories.computed_households.get( + household_id, policy_id, "us" + ) + is not None, + "user_policy": repositories.user_policies.get(user_policy_id) is not None, + "economy": repositories.economies.get(economy_id) is not None, + "analysis": repositories.analyses.get("toy prompt") == "toy answer", + "tracer": repositories.tracers.get(household_id, policy_id, "us") + is not None, + "reform_impact": repositories.reform_impacts.find( + execution_id="toy-impact" + )["reform_impact_id"] + == impact_id, } return { "alembic_head": "alembic_version" in inspect(engine).get_table_names(), **core_results, - "computed_household": computed_households.get(household_id, policy_id, "us") - is not None, - "user_policy": user_policies.get(user_policy_id) is not None, - "economy": economies.get(economy_id) is not None, "simulation": simulations.get_run("toy-simulation-run") is not None, "report": reports.get_run("toy-report-run") is not None, "report_alias": reports.get_alias(900_001)["canonical_report_output_id"] == report_id, - "analysis": analyses.get("toy prompt") == "toy answer", - "tracer": tracers.get(household_id, policy_id, "us") is not None, - "reform_impact": impacts.find(execution_id="toy-impact")["reform_impact_id"] - == impact_id, } diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index 76869299f..186dee02d 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -1,12 +1,13 @@ import json import os from collections.abc import Generator +from contextlib import contextmanager import anthropic from pydantic import BaseModel from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import AnalysisDAO +from policyengine_api.data.v1_daos import AnalysisDAO, V1UnitOfWork class StreamEvent(BaseModel): @@ -24,17 +25,33 @@ class ErrorEvent(StreamEvent): class AIAnalysisService: - def __init__(self, analyses: AnalysisDAO | None = None): + def __init__( + self, + analyses: AnalysisDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._analyses = analyses + self._unit_of_work = unit_of_work @property - def analyses(self) -> AnalysisDAO: - if self._analyses is None: - self._analyses = AnalysisDAO(build_v1_session_manager(local=True)) - return self._analyses + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager(local=True)) + return self._unit_of_work + + @contextmanager + def _analysis_repository(self, *, write: bool = False): + if self._analyses is not None: + yield self._analyses + return + boundary = self.unit_of_work.transaction if write else self.unit_of_work.read + with boundary() as repositories: + yield repositories.analyses def get_existing_analysis(self, prompt: str) -> str | None: - analysis = self.analyses.get(prompt) + with self._analysis_repository() as analyses: + analysis = analyses.get(prompt) return json.dumps(analysis) if analysis is not None else None def trigger_ai_analysis(self, prompt: str) -> Generator[str, None, None]: @@ -63,6 +80,7 @@ def generate(): yield ( json.dumps(TextEvent(stream=event.text).model_dump()) + "\n" ) - self.analyses.store(prompt, response_text, "ok") + with self._analysis_repository(write=True) as analyses: + analyses.store(prompt, response_text, "ok") return generate() diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index fc9ee9515..4a2ddc919 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -1,18 +1,34 @@ import datetime +from contextlib import contextmanager from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReformImpactDAO +from policyengine_api.data.v1_daos import ReformImpactDAO, V1UnitOfWork class ReformImpactsService: - def __init__(self, impacts: ReformImpactDAO | None = None): + def __init__( + self, + impacts: ReformImpactDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._impacts = impacts + self._unit_of_work = unit_of_work @property - def impacts(self) -> ReformImpactDAO: - if self._impacts is None: - self._impacts = ReformImpactDAO(build_v1_session_manager(local=True)) - return self._impacts + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager(local=True)) + return self._unit_of_work + + @contextmanager + def _repository(self, *, write: bool = False): + if self._impacts is not None: + yield self._impacts + return + boundary = self.unit_of_work.transaction if write else self.unit_of_work.read + with boundary() as repositories: + yield repositories.reform_impacts @staticmethod def _filters( @@ -47,18 +63,19 @@ def get_all_reform_impacts( options_hash, api_version, ): - return self.impacts.list( - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - api_version, - ), - options_hash=options_hash, - ) + with self._repository() as impacts: + return impacts.list( + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ), + options_hash=options_hash, + ) def get_all_reform_impacts_by_options_hash_prefix( self, @@ -72,19 +89,20 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash_prefix, api_version, ): - return self.impacts.list_by_options_hash( - options_hash, - options_hash_prefix, - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - api_version, - ), - ) + with self._repository() as impacts: + return impacts.list_by_options_hash( + options_hash, + options_hash_prefix, + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ), + ) def set_reform_impact( self, @@ -102,21 +120,22 @@ def set_reform_impact( start_time, execution_id: str, ): - return self.impacts.create( - country_id=country_id, - reform_policy_id=policy_id, - baseline_policy_id=baseline_policy_id, - region=region, - dataset=dataset, - time_period=time_period, - options_json=options, - options_hash=options_hash, - status=status, - api_version=api_version, - reform_impact_json=reform_impact_json, - start_time=start_time, - execution_id=execution_id, - ) + with self._repository(write=True) as impacts: + return impacts.create( + country_id=country_id, + reform_policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + region=region, + dataset=dataset, + time_period=time_period, + options_json=options, + options_hash=options_hash, + status=status, + api_version=api_version, + reform_impact_json=reform_impact_json, + start_time=start_time, + execution_id=execution_id, + ) def delete_reform_impact( self, @@ -128,17 +147,18 @@ def delete_reform_impact( time_period, options_hash, ): - self.impacts.delete_computing( - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - ), - options_hash=options_hash, - ) + with self._repository(write=True) as impacts: + impacts.delete_computing( + **self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + ), + options_hash=options_hash, + ) def set_error_reform_impact( self, @@ -161,7 +181,8 @@ def set_error_reform_impact( time_period, options_hash, ) - return self.impacts.fail(execution_id, message, self._now()) + with self._repository(write=True) as impacts: + return impacts.fail(execution_id, message, self._now()) def set_complete_reform_impact( self, @@ -184,7 +205,8 @@ def set_complete_reform_impact( time_period, options_hash, ) - return self.impacts.complete(execution_id, reform_impact_json, self._now()) + with self._repository(write=True) as impacts: + return impacts.complete(execution_id, reform_impact_json, self._now()) @staticmethod def _now() -> datetime.datetime: diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index e5484a5b1..d915d2adb 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -1,6 +1,7 @@ import json -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import AnalysisDAO, TracerDAO +from contextlib import contextmanager + +from policyengine_api.data.v1_daos import AnalysisDAO, TracerDAO, V1UnitOfWork from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS from typing import Generator, Literal import re @@ -14,15 +15,19 @@ def __init__( self, tracers: TracerDAO | None = None, analyses: AnalysisDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, ): self._tracers = tracers - super().__init__(analyses) + super().__init__(analyses, unit_of_work=unit_of_work) - @property - def tracers(self) -> TracerDAO: - if self._tracers is None: - self._tracers = TracerDAO(build_v1_session_manager(local=True)) - return self._tracers + @contextmanager + def _tracer_repository(self): + if self._tracers is not None: + yield self._tracers + return + with self.unit_of_work.read() as repositories: + yield repositories.tracers def execute_analysis( self, @@ -91,7 +96,13 @@ def get_tracer( ) -> list: try: # Retrieve from the tracers table in the local database - row = self.tracers.get(household_id, policy_id, country_id, api_version) + with self._tracer_repository() as tracers: + row = tracers.get( + household_id, + policy_id, + country_id, + api_version, + ) if row is None: raise NotFound("No household simulation tracer found") diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 1e382b4a5..4becedf73 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -1,4 +1,4 @@ -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager import importlib import sys from types import SimpleNamespace @@ -189,13 +189,19 @@ def _json_payload(contract: ContractRequest) -> dict | None: return None -def _policy_search_rows(): - return SimpleNamespace( - fetchall=lambda: [ - {"id": 123, "label": "Tax reform", "policy_hash": "hash-1"}, - {"id": 124, "label": "Tax reform", "policy_hash": "hash-1"}, - ] - ) +def _policy_search_unit_of_work(): + @contextmanager + def read(): + yield SimpleNamespace( + policies=SimpleNamespace( + search=lambda *args, **kwargs: [ + {"id": 123, "label": "Tax reform", "policy_hash": "hash-1"}, + {"id": 124, "label": "Tax reform", "policy_hash": "hash-1"}, + ] + ) + ) + + return SimpleNamespace(read=read) def _fake_country(): @@ -224,10 +230,8 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.runtime_sqlalchemy_dao", - return_value=SimpleNamespace( - query=lambda *args, **kwargs: _policy_search_rows() - ), + "policyengine_api.endpoints.policy.runtime_v1_unit_of_work", + return_value=_policy_search_unit_of_work(), ) ) stack.enter_context( diff --git a/tests/integration/test_stage7_mysql_parity.py b/tests/integration/test_stage7_mysql_parity.py new file mode 100644 index 000000000..4b0006f1c --- /dev/null +++ b/tests/integration/test_stage7_mysql_parity.py @@ -0,0 +1,91 @@ +"""Behavioral parity between legacy SQL and typed DAO access on MySQL.""" + +import json + +from alembic import command +from sqlalchemy import text + +from policyengine_api.data.orm import SessionManager +from policyengine_api.data.v1_daos import V1UnitOfWork +from tests.integration.stage7_mysql import alembic_config + + +def test_typed_daos_preserve_legacy_mapping_shapes(stage7_mysql): + database_url, engine = stage7_mysql + command.upgrade(alembic_config(database_url), "head") + unit_of_work = V1UnitOfWork(SessionManager(engine)) + + with unit_of_work.transaction() as repositories: + policy_id = repositories.policies.create( + "us", None, {"x": 1}, "parity", "v1" + ) + household_id = repositories.households.create( + "us", None, {"people": {}}, "household-parity", "v1" + ) + user_id = repositories.users.create_profile( + "auth0|parity", None, "us", 123456789 + ) + + with engine.connect() as connection: + legacy_policy = dict( + connection.execute( + text( + "SELECT * FROM policy " + "WHERE id = :policy_id AND country_id = :country_id" + ), + {"policy_id": policy_id, "country_id": "us"}, + ).mappings().one() + ) + legacy_household = dict( + connection.execute( + text( + "SELECT * FROM household " + "WHERE id = :household_id AND country_id = :country_id" + ), + {"household_id": household_id, "country_id": "us"}, + ).mappings().one() + ) + legacy_user = dict( + connection.execute( + text("SELECT * FROM user_profiles WHERE user_id = :user_id"), + {"user_id": user_id}, + ).mappings().one() + ) + + legacy_policy["policy_json"] = json.loads(legacy_policy["policy_json"]) + legacy_household["household_json"] = json.loads(legacy_household["household_json"]) + with unit_of_work.read() as repositories: + assert legacy_policy == repositories.policies.get("us", policy_id) + assert legacy_household == repositories.households.get("us", household_id) + assert legacy_user == repositories.users.get_profile(user_id=user_id) + + +def test_typed_daos_read_rows_written_by_legacy_sql(stage7_mysql): + database_url, engine = stage7_mysql + command.upgrade(alembic_config(database_url), "head") + unit_of_work = V1UnitOfWork(SessionManager(engine)) + with engine.begin() as connection: + connection.execute( + text( + "INSERT INTO household " + "(country_id, label, api_version, household_json, household_hash) " + "VALUES (:country_id, :label, :api_version, :household_json, " + ":household_hash)" + ), + { + "country_id": "us", + "label": "legacy", + "api_version": "v1", + "household_json": '{"legacy": true}', + "household_hash": "legacy-row", + }, + ) + row = connection.execute( + text("SELECT * FROM household WHERE household_hash = :household_hash"), + {"household_hash": "legacy-row"}, + ).mappings().one() + + legacy_shape = dict(row) + legacy_shape["household_json"] = json.loads(legacy_shape["household_json"]) + with unit_of_work.read() as repositories: + assert repositories.households.get("us", row["id"]) == legacy_shape diff --git a/tests/integration/test_stage7_mysql_runtime.py b/tests/integration/test_stage7_mysql_runtime.py new file mode 100644 index 000000000..01e0d575d --- /dev/null +++ b/tests/integration/test_stage7_mysql_runtime.py @@ -0,0 +1,135 @@ +"""Route, Cloud SQL connector, and startup behavior on disposable MySQL.""" + +import importlib +import json +import os + +from alembic import command +from flask import Flask +import pymysql +from sqlalchemy import text + +from policyengine_api.data.orm import SessionManager +from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_service import UserService +from tests.integration.stage7_mysql import alembic_config, schema_signature + + +def test_public_policy_household_and_user_routes_use_mysql(stage7_mysql, monkeypatch): + database_url, engine = stage7_mysql + command.upgrade(alembic_config(database_url), "head") + sessions = SessionManager(engine) + unit_of_work = V1UnitOfWork(sessions) + + monkeypatch.setenv("FLASK_DEBUG", "1") + from policyengine_api.routes import household_routes, policy_routes + from policyengine_api.routes import user_profile_routes + + monkeypatch.setattr( + policy_routes, + "policy_service", + PolicyService(unit_of_work=unit_of_work), + ) + monkeypatch.setattr( + household_routes, + "household_service", + HouseholdService(unit_of_work=unit_of_work), + ) + monkeypatch.setattr( + user_profile_routes, + "user_service", + UserService(unit_of_work=unit_of_work), + ) + app = Flask(__name__) + app.register_blueprint(policy_routes.policy_bp) + app.register_blueprint(household_routes.household_bp) + app.register_blueprint(user_profile_routes.user_profile_bp) + client = app.test_client() + + policy = client.post("/us/policy", json={"label": "Route", "data": {}}) + assert policy.status_code == 201 + policy_id = policy.get_json()["result"]["policy_id"] + policy_result = json.loads(client.get(f"/us/policy/{policy_id}").data)["result"] + assert policy_result["id"] == policy_id + + household = client.post( + "/us/household", json={"label": "Route", "data": {"people": {}}} + ) + assert household.status_code == 201 + household_id = household.get_json()["result"]["household_id"] + assert ( + client.get(f"/us/household/{household_id}").get_json()["result"]["id"] + == household_id + ) + + user = client.post( + "/us/user-profile", + json={"auth0_id": "auth0|route", "username": None, "user_since": 1}, + ) + assert user.status_code == 201 + user_id = user.get_json()["result"]["user_id"] + assert client.get(f"/us/user-profile?user_id={user_id}").status_code == 200 + + +def test_cloud_sql_connector_pool_drives_daos_and_startup_emits_no_ddl( + stage7_mysql, monkeypatch +): + database_url, engine = stage7_mysql + command.upgrade(alembic_config(database_url), "head") + + os.environ.setdefault("FLASK_DEBUG", "1") + from policyengine_api.data import data as data_module + + url = engine.url + connector_calls = [] + + class LocalConnector: + def __init__(self, **_kwargs): + pass + + def connect(self, **kwargs): + connector_calls.append(kwargs) + return pymysql.connect( + host=url.host, + port=url.port, + user=url.username, + password=url.password, + database=url.database, + ) + + def close(self): + pass + + monkeypatch.setattr(data_module, "Connector", LocalConnector) + monkeypatch.setenv("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", "toy:local:stage7") + monkeypatch.setenv("POLICYENGINE_DB_USER", str(url.username)) + monkeypatch.setenv("POLICYENGINE_DB_NAME", str(url.database)) + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", str(url.password)) + + remote_database = data_module.PolicyEngineDatabase(local=False, initialize=False) + assert connector_calls == [] + with remote_database.pool.connect() as first: + with remote_database.pool.connect() as second: + first_id = first.scalar(text("SELECT CONNECTION_ID()")) + second_id = second.scalar(text("SELECT CONNECTION_ID()")) + assert first_id != second_id + assert len(connector_calls) == 2 + + unit_of_work = V1UnitOfWork(SessionManager(remote_database.pool)) + with unit_of_work.transaction() as repositories: + policy_id = repositories.policies.create( + "us", "Connector", {}, "connector-path", "v1" + ) + with unit_of_work.read() as repositories: + assert repositories.policies.get("us", policy_id)["label"] == "Connector" + assert connector_calls[0]["instance_connection_string"] == "toy:local:stage7" + + monkeypatch.setattr(data_module, "database", remote_database) + before_startup = schema_signature(engine) + api_module = importlib.import_module("policyengine_api.api") + importlib.reload(api_module) + assert api_module.app.test_client().get("/liveness-check").status_code == 200 + assert schema_signature(engine) == before_startup + remote_database._close_pool() diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py index c1e57420a..b5dee1e8a 100644 --- a/tests/unit/data/test_local_daos.py +++ b/tests/unit/data/test_local_daos.py @@ -1,46 +1,56 @@ from datetime import datetime from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import AnalysisDAO, ReformImpactDAO, TracerDAO +from policyengine_api.data.v1_daos import V1UnitOfWork from tests.unit.data.sqlite_schema import create_sqlite_v1_schema -def _daos(): +def _unit_of_work(): manager = build_sqlite_session_manager() create_sqlite_v1_schema(manager) - return AnalysisDAO(manager), ReformImpactDAO(manager), TracerDAO(manager) + return V1UnitOfWork(manager) def test_analysis_dao_round_trip(): - analyses, _, _ = _daos() - analyses.store("prompt", "answer", "complete") - assert analyses.get("prompt") == "answer" + uow = _unit_of_work() + with uow.transaction() as repositories: + repositories.analyses.store("prompt", "answer", "complete") + with uow.read() as repositories: + assert repositories.analyses.get("prompt") == "answer" def test_reform_impact_dao_transitions_by_execution_id(): - _, impacts, _ = _daos() - impacts.create( - country_id="us", - reform_policy_id=2, - baseline_policy_id=1, - region="us", - dataset="default", - time_period="2026", - options_json={}, - options_hash="hash", - api_version="1", - reform_impact_json={}, - status="computing", - start_time=datetime(2026, 1, 1), - execution_id="job", - ) - impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) - assert impacts.find(execution_id="job")["status"] == "ok" - assert impacts.find(execution_id="job")["reform_impact_json"] == {"result": 1} + uow = _unit_of_work() + with uow.transaction() as repositories: + repositories.reform_impacts.create( + country_id="us", + reform_policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options_json={}, + options_hash="hash", + api_version="1", + reform_impact_json={}, + status="computing", + start_time=datetime(2026, 1, 1), + execution_id="job", + ) + repositories.reform_impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) + with uow.read() as repositories: + assert repositories.reform_impacts.find(execution_id="job")["status"] == "ok" + assert repositories.reform_impacts.find(execution_id="job")[ + "reform_impact_json" + ] == {"result": 1} def test_tracer_dao_returns_latest_matching_trace(): - _, _, tracers = _daos() - tracers.create(1, 2, "us", "1", {"trace": "first"}) - tracers.create(1, 2, "us", "1", {"trace": "latest"}) - assert tracers.get(1, 2, "us")["tracer_output"] == {"trace": "latest"} + uow = _unit_of_work() + with uow.transaction() as repositories: + repositories.tracers.create(1, 2, "us", "1", {"trace": "first"}) + repositories.tracers.create(1, 2, "us", "1", {"trace": "latest"}) + with uow.read() as repositories: + assert repositories.tracers.get(1, 2, "us")["tracer_output"] == { + "trace": "latest" + } diff --git a/tests/unit/data/test_ordinary_v1_daos.py b/tests/unit/data/test_ordinary_v1_daos.py new file mode 100644 index 000000000..74dfb8d07 --- /dev/null +++ b/tests/unit/data/test_ordinary_v1_daos.py @@ -0,0 +1,70 @@ +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import V1UnitOfWork +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema + + +def _unit_of_work() -> V1UnitOfWork: + manager = build_sqlite_session_manager() + create_sqlite_v1_schema(manager) + return V1UnitOfWork(manager) + + +def test_computed_household_upsert_preserves_one_cache_row(): + uow = _unit_of_work() + values = { + "household_id": 1, + "policy_id": 2, + "country_id": "us", + "api_version": "1", + "computed_household_json": {"value": 1}, + "status": "complete", + } + with uow.transaction() as repositories: + repositories.computed_households.upsert(**values) + repositories.computed_households.upsert( + **{**values, "computed_household_json": {"value": 2}} + ) + with uow.read() as repositories: + row = repositories.computed_households.get(1, 2, "us", api_version="1") + assert row["computed_household_json"] == {"value": 2} + + +def test_user_policy_nullable_identity_list_and_update_are_orm_managed(): + uow = _unit_of_work() + values = { + "country_id": "us", + "reform_id": 2, + "reform_label": None, + "baseline_id": 1, + "baseline_label": None, + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": None, + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 1, + "budgetary_impact": None, + "type": None, + } + with uow.transaction() as repositories: + user_policy_id = repositories.user_policies.create(**values) + assert repositories.user_policies.find_unique(**values)["id"] == user_policy_id + assert repositories.user_policies.update( + user_policy_id, {"number_of_provisions": 4} + ) + with uow.read() as repositories: + rows = repositories.user_policies.list_for_user("us", "auth0|one") + assert rows[0]["number_of_provisions"] == 4 + + +def test_policy_search_and_reform_impact_limit_use_typed_statements(): + uow = _unit_of_work() + with uow.transaction() as repositories: + repositories.policies.create("us", "Tax reform", {}, "one", "1") + repositories.policies.create("us", "Other", {}, "two", "1") + with uow.read() as repositories: + assert [row["label"] for row in repositories.policies.search("us", "Tax")] == [ + "Tax reform" + ] diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index daf09145a..66572564e 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -21,3 +21,19 @@ def test_runtime_sql_is_confined_to_the_data_access_layer(): ): offenders.append(str(path.relative_to(PACKAGE_ROOT))) assert offenders == [] + + +def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): + relative_paths = ( + "endpoints/household.py", + "endpoints/policy.py", + "endpoints/simulation.py", + "endpoints/economy/reform_impact.py", + "country.py", + "services/ai_analysis_service.py", + "services/reform_impacts_service.py", + "services/tracer_analysis_service.py", + ) + for relative_path in relative_paths: + source = (PACKAGE_ROOT / relative_path).read_text(encoding="utf-8") + assert "runtime_sqlalchemy_dao" not in source diff --git a/tests/unit/endpoints/test_get_simulations.py b/tests/unit/endpoints/test_get_simulations.py index 2fa061bff..dfa227fce 100644 --- a/tests/unit/endpoints/test_get_simulations.py +++ b/tests/unit/endpoints/test_get_simulations.py @@ -30,7 +30,7 @@ def _seed_reform_impacts(test_db, n: int) -> None: "1.0.0", "{}", "complete", - f"2026-01-01 00:00:{i:02d}", + f"2026-01-01 00:{i // 60:02d}:{i % 60:02d}", f"exec-{i}", ), ) From ecea37f4754a0022d4759f77000c3ae243ce432b Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:22:27 +0300 Subject: [PATCH 36/89] refactor: make run orchestration transactional --- policyengine_api/data/v1_daos.py | 458 +++++++----------- .../scripts/qualify_stage7_toy.py | 73 ++- .../services/report_output_alias_service.py | 69 ++- .../services/report_output_service.py | 397 ++++++--------- .../services/report_run_service.py | 55 ++- .../services/report_spec_service.py | 128 ++++- .../services/simulation_run_service.py | 55 ++- .../services/simulation_service.py | 125 +++-- .../services/simulation_spec_service.py | 54 ++- tests/unit/data/test_run_daos.py | 101 ++-- tests/unit/data/test_stage7_no_direct_sql.py | 16 + tests/unit/data/test_v1_unit_of_work.py | 27 ++ 12 files changed, 833 insertions(+), 725 deletions(-) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 18a8a46e8..1914eeb20 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -37,86 +37,6 @@ def _mapping(model: Any) -> dict[str, Any]: } -class _Rows: - def __init__(self, rows): - self.rows = rows - self.index = 0 - - def fetchone(self): - if self.index >= len(self.rows): - return None - row = self.rows[self.index] - self.index += 1 - return row - - def fetchall(self): - rows = self.rows[self.index :] - self.index = len(self.rows) - return rows - - -class SQLAlchemyDAO: - """Compatibility execution boundary for complex v1 transactional SQL. - - New CRUD belongs in typed DAOs. This boundary keeps the mature report - orchestration on SQLAlchemy-owned sessions while it is decomposed. - """ - - def __init__(self, sessions: SessionManager, session=None): - self.sessions = sessions - self._session = session - - @property - def local(self) -> bool: - return self.sessions.engine.dialect.name == "sqlite" - - def _statement(self, statement: str) -> str: - return statement if self.local else statement.replace("?", "%s") - - @staticmethod - def _rows(result) -> _Rows: - if not result.returns_rows: - return _Rows([]) - return _Rows(list(result.mappings())) - - def query(self, statement: str, params=None) -> _Rows: - if self._session is not None: - result = self._session.connection().exec_driver_sql( - self._statement(statement), params or () - ) - return self._rows(result) - - def operation(session): - result = session.connection().exec_driver_sql( - self._statement(statement), params or () - ) - return self._rows(result) - - return self.sessions.run_in_transaction(operation) - - def transaction(self, callback): - return self.sessions.run_in_transaction( - lambda session: callback(SQLAlchemyDAO(self.sessions, session)) - ) - - @property - def session(self): - return self._session - - -_runtime_sqlalchemy_daos: dict[bool, SQLAlchemyDAO] = {} - - -def runtime_sqlalchemy_dao(*, local: bool = False) -> SQLAlchemyDAO: - if local not in _runtime_sqlalchemy_daos: - from policyengine_api.data.orm import build_v1_session_manager - - _runtime_sqlalchemy_daos[local] = SQLAlchemyDAO( - build_v1_session_manager(local=local) - ) - return _runtime_sqlalchemy_daos[local] - - class PolicyDAO: def __init__(self, session: Session): self.session = session @@ -327,6 +247,8 @@ def __init__(self, session: Session): self.analyses = AnalysisDAO(session) self.reform_impacts = ReformImpactDAO(session) self.tracers = TracerDAO(session) + self.simulations = SimulationDAO(session) + self.reports = ReportDAO(session) class V1UnitOfWork: @@ -593,18 +515,17 @@ def get( class SimulationDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def get( self, simulation_id: int, country_id: str | None = None ) -> dict[str, Any] | None: - with self.sessions.session() as session: - statement = select(Simulation).where(Simulation.id == simulation_id) - if country_id is not None: - statement = statement.where(Simulation.country_id == country_id) - model = session.scalar(statement) - return _mapping(model) if model else None + statement = select(Simulation).where(Simulation.id == simulation_id) + if country_id is not None: + statement = statement.where(Simulation.country_id == country_id) + model = self.session.scalar(statement) + return _mapping(model) if model else None @staticmethod def get_in_session( @@ -617,27 +538,20 @@ def get_in_session( return _mapping(model) if model else None def create(self, **values: Any) -> int: - def operation(session): - model = Simulation(**values) - session.add(model) - session.flush() - return model.id - - return self.sessions.run_in_transaction(operation) + model = Simulation(**values) + self.session.add(model) + self.session.flush() + return model.id def find_latest(self, **filters: Any) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.scalar( - select(Simulation) - .where( - *( - getattr(Simulation, key) == value - for key, value in filters.items() - ) - ) - .order_by(Simulation.id.desc()) + model = self.session.scalar( + select(Simulation) + .where( + *(getattr(Simulation, key) == value for key, value in filters.items()) ) - return _mapping(model) if model else None + .order_by(Simulation.id.desc()) + ) + return _mapping(model) if model else None @staticmethod def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: @@ -713,10 +627,8 @@ def ensure_dual_write_state_in_session( def ensure_dual_write_state( self, simulation_id: int, country_id: str | None = None ) -> dict[str, Any]: - return self.sessions.run_in_transaction( - lambda session: self.ensure_dual_write_state_in_session( - session, simulation_id, country_id - ) + return self.ensure_dual_write_state_in_session( + self.session, simulation_id, country_id ) def create_or_get_with_sync( @@ -725,34 +637,32 @@ def create_or_get_with_sync( sync_callback, **values: Any, ) -> dict[str, Any]: - def operation(session): - filters = { - key: values[key] - for key in ( - "country_id", - "population_id", - "population_type", - "policy_id", - ) - } - model = session.scalar( - select(Simulation) - .where( - *( - getattr(Simulation, key) == value - for key, value in filters.items() - ) - ) - .order_by(Simulation.id.desc()) - .with_for_update() + filters = { + key: values[key] + for key in ( + "country_id", + "population_id", + "population_type", + "policy_id", ) - if model is None: - model = Simulation(**values) - session.add(model) - session.flush() - return sync_callback(session, model.id, country_id=model.country_id) - - return self.sessions.run_in_transaction(operation) + } + model = self.session.scalar( + select(Simulation) + .where( + *(getattr(Simulation, key) == value for key, value in filters.items()) + ) + .order_by(Simulation.id.desc()) + .with_for_update() + ) + if model is None: + model = Simulation(**values) + self.session.add(model) + self.session.flush() + return sync_callback( + self.session, + model.id, + country_id=model.country_id, + ) def update_with_sync( self, @@ -761,176 +671,182 @@ def update_with_sync( values: dict[str, Any], sync_callback, ) -> dict[str, Any]: - def operation(session): - model = session.scalar( - select(Simulation) - .where( - Simulation.id == simulation_id, - Simulation.country_id == country_id, - ) - .with_for_update() + model = self.session.scalar( + select(Simulation) + .where( + Simulation.id == simulation_id, + Simulation.country_id == country_id, ) - if model is None: - raise ValueError(f"Simulation #{simulation_id} not found") - for key, value in values.items(): - setattr(model, key, value) - session.flush() - return sync_callback(session, simulation_id, country_id=country_id) - - return self.sessions.run_in_transaction(operation) + .with_for_update() + ) + if model is None: + raise ValueError(f"Simulation #{simulation_id} not found") + for key, value in values.items(): + setattr(model, key, value) + self.session.flush() + return sync_callback( + self.session, + simulation_id, + country_id=country_id, + ) def update(self, simulation_id: int, **values: Any) -> bool: - def operation(session): - model = session.get(Simulation, simulation_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - return self.sessions.run_in_transaction(operation) + model = self.session.get(Simulation, simulation_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True def create_run( self, simulation_id: int, *, run_id: str, **values: Any ) -> dict[str, Any]: - def operation(session): - parent = session.scalar( - select(Simulation) - .where(Simulation.id == simulation_id) - .with_for_update() - ) - if parent is None: - raise LookupError(f"Simulation {simulation_id} does not exist") - sequence = ( - session.scalar( - select(func.max(SimulationRun.run_sequence)).where( - SimulationRun.simulation_id == simulation_id - ) + parent = self.session.scalar( + select(Simulation).where(Simulation.id == simulation_id).with_for_update() + ) + if parent is None: + raise LookupError(f"Simulation {simulation_id} does not exist") + sequence = ( + self.session.scalar( + select(func.max(SimulationRun.run_sequence)).where( + SimulationRun.simulation_id == simulation_id ) - or 0 - ) + 1 - model = SimulationRun( - id=run_id, - simulation_id=simulation_id, - run_sequence=sequence, - **values, ) - session.add(model) - session.flush() - return _mapping(model) - - return self.sessions.run_in_transaction(operation) + or 0 + ) + 1 + model = SimulationRun( + id=run_id, + simulation_id=simulation_id, + run_sequence=sequence, + **values, + ) + self.session.add(model) + self.session.flush() + return _mapping(model) def get_run(self, run_id: str) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get(SimulationRun, run_id) - return _mapping(model) if model else None + model = self.session.get(SimulationRun, run_id) + return _mapping(model) if model else None def list_runs(self, simulation_id: int) -> list[dict[str, Any]]: - with self.sessions.session() as session: - models = session.scalars( - select(SimulationRun) - .where(SimulationRun.simulation_id == simulation_id) - .order_by(SimulationRun.run_sequence.desc()) - ) - return [_mapping(model) for model in models] + models = self.session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) + return [_mapping(model) for model in models] class ReportDAO: - def __init__(self, sessions: SessionManager): - self.sessions = sessions + def __init__(self, session: Session): + self.session = session def get( self, report_output_id: int, country_id: str | None = None ) -> dict[str, Any] | None: - with self.sessions.session() as session: - statement = select(ReportOutput).where(ReportOutput.id == report_output_id) - if country_id is not None: - statement = statement.where(ReportOutput.country_id == country_id) - model = session.scalar(statement) - return _mapping(model) if model else None + statement = select(ReportOutput).where(ReportOutput.id == report_output_id) + if country_id is not None: + statement = statement.where(ReportOutput.country_id == country_id) + model = self.session.scalar(statement) + return _mapping(model) if model else None - def create(self, **values: Any) -> int: - def operation(session): - model = ReportOutput(**values) - session.add(model) - session.flush() - return model.id + def get_for_update( + self, report_output_id: int, country_id: str | None = None + ) -> dict[str, Any] | None: + statement = ( + select(ReportOutput) + .where(ReportOutput.id == report_output_id) + .with_for_update() + ) + if country_id is not None: + statement = statement.where(ReportOutput.country_id == country_id) + model = self.session.scalar(statement) + return _mapping(model) if model else None + + def find_latest(self, **filters: Any) -> dict[str, Any] | None: + model = self.session.scalar( + select(ReportOutput) + .where( + *(getattr(ReportOutput, key) == value for key, value in filters.items()) + ) + .order_by(ReportOutput.id.desc()) + ) + return _mapping(model) if model else None - return self.sessions.run_in_transaction(operation) + def create(self, **values: Any) -> int: + model = ReportOutput(**values) + self.session.add(model) + self.session.flush() + return model.id def update(self, report_output_id: int, **values: Any) -> bool: - def operation(session): - model = session.get(ReportOutput, report_output_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - return self.sessions.run_in_transaction(operation) + model = self.session.get(ReportOutput, report_output_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True def create_run( self, report_output_id: int, *, run_id: str, **values: Any ) -> dict[str, Any]: - def operation(session): - parent = session.scalar( - select(ReportOutput) - .where(ReportOutput.id == report_output_id) - .with_for_update() - ) - if parent is None: - raise LookupError(f"Report output {report_output_id} does not exist") - sequence = ( - session.scalar( - select(func.max(ReportOutputRun.run_sequence)).where( - ReportOutputRun.report_output_id == report_output_id - ) + parent = self.session.scalar( + select(ReportOutput) + .where(ReportOutput.id == report_output_id) + .with_for_update() + ) + if parent is None: + raise LookupError(f"Report output {report_output_id} does not exist") + sequence = ( + self.session.scalar( + select(func.max(ReportOutputRun.run_sequence)).where( + ReportOutputRun.report_output_id == report_output_id ) - or 0 - ) + 1 - model = ReportOutputRun( - id=run_id, - report_output_id=report_output_id, - run_sequence=sequence, - **values, ) - session.add(model) - session.flush() - return _mapping(model) - - return self.sessions.run_in_transaction(operation) + or 0 + ) + 1 + model = ReportOutputRun( + id=run_id, + report_output_id=report_output_id, + run_sequence=sequence, + **values, + ) + self.session.add(model) + self.session.flush() + return _mapping(model) def get_run(self, run_id: str) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get(ReportOutputRun, run_id) - return _mapping(model) if model else None + model = self.session.get(ReportOutputRun, run_id) + return _mapping(model) if model else None + + def update_run(self, run_id: str, **values: Any) -> bool: + model = self.session.get(ReportOutputRun, run_id) + if model is None: + return False + for key, value in values.items(): + setattr(model, key, value) + return True def list_runs(self, report_output_id: int) -> list[dict[str, Any]]: - with self.sessions.session() as session: - models = session.scalars( - select(ReportOutputRun) - .where(ReportOutputRun.report_output_id == report_output_id) - .order_by(ReportOutputRun.run_sequence.desc()) - ) - return [_mapping(model) for model in models] + models = self.session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + return [_mapping(model) for model in models] def set_alias(self, legacy_id: int, canonical_id: int) -> None: - def operation(session): - model = session.get(LegacyReportOutputAlias, legacy_id) - if model is None: - session.add( - LegacyReportOutputAlias( - legacy_report_output_id=legacy_id, - canonical_report_output_id=canonical_id, - ) + model = self.session.get(LegacyReportOutputAlias, legacy_id) + if model is None: + self.session.add( + LegacyReportOutputAlias( + legacy_report_output_id=legacy_id, + canonical_report_output_id=canonical_id, ) - else: - model.canonical_report_output_id = canonical_id - - self.sessions.run_in_transaction(operation) + ) + else: + model.canonical_report_output_id = canonical_id def get_alias(self, legacy_id: int) -> dict[str, Any] | None: - with self.sessions.session() as session: - model = session.get(LegacyReportOutputAlias, legacy_id) - return _mapping(model) if model else None + model = self.session.get(LegacyReportOutputAlias, legacy_id) + return _mapping(model) if model else None diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py index bb1e534cd..c88dd429e 100644 --- a/policyengine_api/scripts/qualify_stage7_toy.py +++ b/policyengine_api/scripts/qualify_stage7_toy.py @@ -10,11 +10,7 @@ from policyengine_api.constants import REPO from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_daos import ( - ReportDAO, - SimulationDAO, - V1UnitOfWork, -) +from policyengine_api.data.v1_daos import V1UnitOfWork from policyengine_api.data.v1_models import V1Base @@ -40,8 +36,6 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: engine = create_engine(database_url) sessions = SessionManager(engine) unit_of_work = V1UnitOfWork(sessions) - simulations = SimulationDAO(sessions) - reports = ReportDAO(sessions) with unit_of_work.transaction() as repositories: policy_id = repositories.policies.create("us", "Toy", {}, "toy-policy", "toy") @@ -103,33 +97,33 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: start_time=datetime(2026, 1, 1), execution_id="toy-impact", ) - simulation_id = simulations.create( - country_id="us", - api_version="toy", - population_id=str(household_id), - population_type="household", - policy_id=policy_id, - ) - simulations.create_run( - simulation_id, - run_id="toy-simulation-run", - status="pending", - trigger_type="qualification", - ) - report_id = reports.create( - country_id="us", - simulation_1_id=simulation_id, - simulation_2_id=None, - api_version="toy", - year="2026", - ) - reports.create_run( - report_id, - run_id="toy-report-run", - status="pending", - trigger_type="qualification", - ) - reports.set_alias(900_001, report_id) + simulation_id = repositories.simulations.create( + country_id="us", + api_version="toy", + population_id=str(household_id), + population_type="household", + policy_id=policy_id, + ) + repositories.simulations.create_run( + simulation_id, + run_id="toy-simulation-run", + status="pending", + trigger_type="qualification", + ) + report_id = repositories.reports.create( + country_id="us", + simulation_1_id=simulation_id, + simulation_2_id=None, + api_version="toy", + year="2026", + ) + repositories.reports.create_run( + report_id, + run_id="toy-report-run", + status="pending", + trigger_type="qualification", + ) + repositories.reports.set_alias(900_001, report_id) with unit_of_work.read() as repositories: core_results = { @@ -149,13 +143,16 @@ def qualify_stage7_toy(database_url: str) -> dict[str, bool]: execution_id="toy-impact" )["reform_impact_id"] == impact_id, + "simulation": repositories.simulations.get_run("toy-simulation-run") + is not None, + "report": repositories.reports.get_run("toy-report-run") is not None, + "report_alias": repositories.reports.get_alias(900_001)[ + "canonical_report_output_id" + ] + == report_id, } return { "alembic_head": "alembic_version" in inspect(engine).get_table_names(), **core_results, - "simulation": simulations.get_run("toy-simulation-run") is not None, - "report": reports.get_run("toy-report-run") is not None, - "report_alias": reports.get_alias(900_001)["canonical_report_output_id"] - == report_id, } diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 54eebe45c..c95bf69d9 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,53 +1,92 @@ from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO +from policyengine_api.data.v1_daos import ReportDAO, V1UnitOfWork class ReportOutputAliasService: - def __init__(self, reports: ReportDAO | None = None): + def __init__( + self, + reports: ReportDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._reports = reports + self._unit_of_work = unit_of_work @property - def reports(self) -> ReportDAO: - if self._reports is None: - self._reports = ReportDAO(build_v1_session_manager()) - return self._reports + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work def _get_report_output_row(self, report_output_id: int) -> dict | None: - return self.reports.get(report_output_id) + if self._reports is not None: + return self._reports.get(report_output_id) + with self.unit_of_work.read() as repositories: + return repositories.reports.get(report_output_id) def get_alias(self, legacy_report_output_id: int) -> dict | None: - return self.reports.get_alias(legacy_report_output_id) + if self._reports is not None: + return self._reports.get_alias(legacy_report_output_id) + with self.unit_of_work.read() as repositories: + return repositories.reports.get_alias(legacy_report_output_id) def resolve_canonical_report_output_id( self, requested_report_output_id: int ) -> int | None: - alias = self.get_alias(requested_report_output_id) + if self._reports is not None: + return self._resolve(self._reports, requested_report_output_id) + with self.unit_of_work.read() as repositories: + return self._resolve(repositories.reports, requested_report_output_id) + + def _resolve( + self, reports: ReportDAO, requested_report_output_id: int + ) -> int | None: + alias = reports.get_alias(requested_report_output_id) if alias is not None: canonical_id = alias["canonical_report_output_id"] - if self.reports.get(canonical_id) is None: + if reports.get(canonical_id) is None: raise ValueError( f"Alias points to missing canonical report output #{canonical_id}" ) return canonical_id - row = self.reports.get(requested_report_output_id) + row = reports.get(requested_report_output_id) return row["id"] if row is not None else None def set_alias( self, legacy_report_output_id: int, canonical_report_output_id: int ) -> bool: - legacy = self.reports.get(legacy_report_output_id) + if self._reports is not None: + return self._set_alias( + self._reports, + legacy_report_output_id, + canonical_report_output_id, + ) + with self.unit_of_work.transaction() as repositories: + return self._set_alias( + repositories.reports, + legacy_report_output_id, + canonical_report_output_id, + ) + + def _set_alias( + self, + reports: ReportDAO, + legacy_report_output_id: int, + canonical_report_output_id: int, + ) -> bool: + legacy = reports.get(legacy_report_output_id) if legacy is None: raise ValueError( f"Legacy report output #{legacy_report_output_id} not found" ) - canonical = self.reports.get(canonical_report_output_id) + canonical = reports.get(canonical_report_output_id) if canonical is None: raise ValueError( f"Canonical report output #{canonical_report_output_id} not found" ) if legacy_report_output_id == canonical_report_output_id: raise ValueError("Legacy and canonical report outputs must be different") - existing = self.reports.get_alias(legacy_report_output_id) + existing = reports.get_alias(legacy_report_output_id) if existing is not None: if existing["canonical_report_output_id"] == canonical_report_output_id: return True @@ -60,5 +99,5 @@ def set_alias( raise ValueError( "Legacy and canonical report outputs must describe the same report" ) - self.reports.set_alias(legacy_report_output_id, canonical_report_output_id) + reports.set_alias(legacy_report_output_id, canonical_report_output_id) return True diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index f383f15a6..c3c153d26 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,11 +1,9 @@ import uuid from datetime import datetime, timezone -from sqlalchemy.engine.row import Row - from policyengine_api.constants import get_report_output_cache_version from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SQLAlchemyDAO +from policyengine_api.data.v1_daos import V1UnitOfWork from policyengine_api.services.report_spec_service import ( ECONOMY_REPORT_KINDS, ReportSpec, @@ -22,22 +20,23 @@ class ReportOutputService: - def __init__(self, persistence: SQLAlchemyDAO | None = None): - self._persistence = persistence - self.report_spec_service = ReportSpecService() - self.simulation_service = SimulationService() + def __init__(self, *, unit_of_work: V1UnitOfWork | None = None): + self._unit_of_work = unit_of_work + self.report_spec_service = ReportSpecService(unit_of_work=unit_of_work) + self.simulation_service = SimulationService(unit_of_work=unit_of_work) @property - def persistence(self) -> SQLAlchemyDAO: - if self._persistence is None: - self._persistence = SQLAlchemyDAO(build_v1_session_manager()) - return self._persistence - - def _lock_clause(self) -> str: - return "" if self.persistence.local else " FOR UPDATE" + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + self.report_spec_service = ReportSpecService( + unit_of_work=self._unit_of_work + ) + self.simulation_service = SimulationService(unit_of_work=self._unit_of_work) + return self._unit_of_work - def _utc_timestamp(self) -> str: - return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + def _utc_timestamp(self) -> datetime: + return datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) def _format_run_timestamp(self, value) -> str | None: if value is None: @@ -86,17 +85,17 @@ def _get_report_output_row( country_id: str | None = None, for_update: bool = False, ) -> dict | None: - queryer = queryer or self.persistence - query = "SELECT * FROM report_outputs WHERE id = ?" - params: list[int | str] = [report_output_id] - if country_id is not None: - query += " AND country_id = ?" - params.append(country_id) + if queryer is None: + with self.unit_of_work.read() as repositories: + return self._get_report_output_row( + report_output_id, + queryer=repositories, + country_id=country_id, + for_update=for_update, + ) if for_update: - query += self._lock_clause() - - row: Row | None = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None + return queryer.reports.get_for_update(report_output_id, country_id) + return queryer.reports.get(report_output_id, country_id) def _get_linked_simulations( self, @@ -105,18 +104,22 @@ def _get_linked_simulations( queryer=None, bootstrap_dual_write_state: bool = False, ) -> tuple[dict, dict | None]: - queryer = queryer or self.persistence + if queryer is None: + with self.unit_of_work.read() as repositories: + return self._get_linked_simulations( + report_output, + queryer=repositories, + bootstrap_dual_write_state=bootstrap_dual_write_state, + ) if bootstrap_dual_write_state: - simulation_1 = self.simulation_service._ensure_simulation_dual_write_state_in_transaction( - queryer, + simulation_1 = queryer.simulations.ensure_dual_write_state( report_output["simulation_1_id"], - country_id=report_output["country_id"], + report_output["country_id"], ) else: - simulation_1 = self.simulation_service._get_simulation_row( + simulation_1 = queryer.simulations.get( report_output["simulation_1_id"], - queryer=queryer, - country_id=report_output["country_id"], + report_output["country_id"], ) if simulation_1 is None: raise ValueError( @@ -127,16 +130,14 @@ def _get_linked_simulations( simulation_2 = None if report_output["simulation_2_id"] is not None: if bootstrap_dual_write_state: - simulation_2 = self.simulation_service._ensure_simulation_dual_write_state_in_transaction( - queryer, + simulation_2 = queryer.simulations.ensure_dual_write_state( report_output["simulation_2_id"], - country_id=report_output["country_id"], + report_output["country_id"], ) else: - simulation_2 = self.simulation_service._get_simulation_row( + simulation_2 = queryer.simulations.get( report_output["simulation_2_id"], - queryer=queryer, - country_id=report_output["country_id"], + report_output["country_id"], ) if simulation_2 is None: raise ValueError( @@ -153,11 +154,7 @@ def _require_simulation_exists( country_id: str, simulation_id: int, ) -> dict: - simulation = self.simulation_service._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - ) + simulation = tx.simulations.get(simulation_id, country_id) if simulation is None: raise ValueError( f"Report output references missing simulation #{simulation_id}" @@ -167,15 +164,13 @@ def _require_simulation_exists( def _list_report_runs_descending( self, report_output_id: int, *, queryer=None ) -> list[dict]: - queryer = queryer or self.persistence - rows = queryer.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence DESC - """, - (report_output_id,), - ).fetchall() + if queryer is None: + with self.unit_of_work.read() as repositories: + return self._list_report_runs_descending( + report_output_id, + queryer=repositories, + ) + rows = queryer.reports.list_runs(report_output_id) runs = [] for row in rows: @@ -244,7 +239,7 @@ def _with_display_run_timestamps( values live on report_output_runs; this helper chooses the display run, formats its requested/started/finished timestamps, and returns an enriched copy of the report output dict. It intentionally does not - mutate self.persistence state. + mutate repository state. These timestamps describe the selected base report execution. They are not user-report association metadata and should not be treated as a @@ -258,10 +253,13 @@ def _with_display_run_timestamps( report_output["id"], queryer=queryer ) display_run = select_display_report_run(report_output, runs_descending) + enriched_report_output = dict(report_output) + enriched_report_output["output"] = serialize_json_field( + enriched_report_output.get("output") + ) if display_run is None: - return report_output + return enriched_report_output - enriched_report_output = dict(report_output) for field in ("requested_at", "started_at", "finished_at"): enriched_report_output[field] = self._format_run_timestamp( display_run.get(field) @@ -345,20 +343,12 @@ def _upsert_report_spec_in_transaction( or report_output.get("report_spec_schema_version") != 1 or report_output.get("report_spec_status") != report_spec_status ): - tx.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, - report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - report_spec.report_kind, - report_spec.model_dump_json(), - 1, - report_spec_status, - report_output["id"], - ), + tx.reports.update( + report_output["id"], + report_kind=report_spec.report_kind, + report_spec_json=report_spec.model_dump(), + report_spec_schema_version=1, + report_spec_status=report_spec_status, ) report_output["report_kind"] = report_spec.report_kind report_output["report_spec_json"] = report_spec.model_dump() @@ -412,40 +402,21 @@ def _insert_bootstrap_report_run( started_at = requested_at if has_started else None finished_at = requested_at if is_terminal else None - tx.query( - """ - INSERT INTO report_output_runs ( - id, report_output_id, run_sequence, status, output, error_message, - trigger_type, requested_at, started_at, finished_at, source_run_id, - report_spec_snapshot_json, country_package_version, policyengine_version, - data_version, runtime_app_name, report_cache_version, - simulation_cache_version, requested_version_override, resolved_dataset, - resolved_options_hash - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - str(uuid.uuid4()), - report_output["id"], - 1, - report_output["status"], - serialize_json_field(report_output.get("output")), - report_output.get("error_message"), - "initial", - requested_at, - started_at, - finished_at, - None, - (report_spec.model_dump_json() if report_spec is not None else None), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["report_cache_version"], - version_manifest["simulation_cache_version"], - version_manifest["requested_version_override"], - version_manifest["resolved_dataset"], - version_manifest["resolved_options_hash"], + tx.reports.create_run( + report_output["id"], + run_id=str(uuid.uuid4()), + status=report_output["status"], + output=report_output.get("output"), + error_message=report_output.get("error_message"), + trigger_type="initial", + requested_at=requested_at, + started_at=started_at, + finished_at=finished_at, + source_run_id=None, + report_spec_snapshot_json=( + report_spec.model_dump() if report_spec is not None else None ), + **version_manifest, ) def _update_report_run_in_transaction( @@ -457,63 +428,50 @@ def _update_report_run_in_transaction( version_manifest: dict[str, str | None], preserve_terminal_finished_at: bool = False, ) -> None: + run = tx.reports.get_run(run_id) + if run is None: + raise ValueError(f"Report output run {run_id} not found") + fallback_timestamp = self._utc_timestamp() - timestamp_updates = [ - "requested_at = COALESCE(requested_at, started_at, finished_at, ?)" - ] - timestamp_values = [fallback_timestamp] + requested_at = ( + run.get("requested_at") + or run.get("started_at") + or run.get("finished_at") + or fallback_timestamp + ) if report_output["status"] in ("complete", "error"): finished_at = self._utc_timestamp() - timestamp_updates.append( - "started_at = COALESCE(started_at, finished_at, requested_at, ?)" + started_at = ( + run.get("started_at") + or run.get("finished_at") + or run.get("requested_at") + or finished_at ) - timestamp_values.append(finished_at) if preserve_terminal_finished_at: - timestamp_updates.append("finished_at = COALESCE(finished_at, ?)") - else: - timestamp_updates.append("finished_at = ?") - timestamp_values.append(finished_at) + finished_at = run.get("finished_at") or finished_at elif report_output["status"] == "running": - started_at = self._utc_timestamp() - timestamp_updates.extend( - [ - "started_at = COALESCE(started_at, requested_at, ?)", - "finished_at = NULL", - ] + started_at = ( + run.get("started_at") + or run.get("requested_at") + or self._utc_timestamp() ) - timestamp_values.append(started_at) + finished_at = None else: - timestamp_updates.extend(["started_at = NULL", "finished_at = NULL"]) - - tx.query( - f""" - UPDATE report_output_runs - SET status = ?, output = ?, error_message = ?, - {", ".join(timestamp_updates)}, - report_spec_snapshot_json = ?, country_package_version = ?, - policyengine_version = ?, data_version = ?, runtime_app_name = ?, - report_cache_version = ?, simulation_cache_version = ?, - requested_version_override = ?, resolved_dataset = ?, - resolved_options_hash = ? - WHERE id = ? - """, - ( - report_output["status"], - serialize_json_field(report_output.get("output")), - report_output.get("error_message"), - *timestamp_values, - (report_spec.model_dump_json() if report_spec is not None else None), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["report_cache_version"], - version_manifest["simulation_cache_version"], - version_manifest["requested_version_override"], - version_manifest["resolved_dataset"], - version_manifest["resolved_options_hash"], - run_id, + started_at = None + finished_at = None + + tx.reports.update_run( + run_id, + status=report_output["status"], + output=report_output.get("output"), + error_message=report_output.get("error_message"), + requested_at=requested_at, + started_at=started_at, + finished_at=finished_at, + report_spec_snapshot_json=( + report_spec.model_dump() if report_spec is not None else None ), + **version_manifest, ) def _sync_parent_pointers_in_transaction( @@ -532,17 +490,10 @@ def _sync_parent_pointers_in_transaction( ): return - tx.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ( - desired_active_run_id, - desired_latest_successful_run_id, - report_output["id"], - ), + tx.reports.update( + report_output["id"], + active_run_id=desired_active_run_id, + latest_successful_run_id=desired_latest_successful_run_id, ) report_output["active_run_id"] = desired_active_run_id report_output["latest_successful_run_id"] = desired_latest_successful_run_id @@ -644,13 +595,12 @@ def ensure_report_output_dual_write_state( report_output_id: int, country_id: str | None = None, ) -> dict: - return self.persistence.transaction( - lambda tx: self._ensure_report_output_dual_write_state_in_transaction( - tx, + with self.unit_of_work.transaction() as repositories: + return self._ensure_report_output_dual_write_state_in_transaction( + repositories, report_output_id, country_id=country_id, ) - ) def get_stored_report_output( self, country_id: str, report_output_id: int @@ -661,7 +611,7 @@ def get_stored_report_output( This is used by mutation paths that must address the originally requested row. It still runs dual-write synchronization, so it may bootstrap or repair run/spec metadata and returns the display-run - timestamp projection. It is therefore not a raw self.persistence read. + timestamp projection. It is therefore not a raw storage read. TODO: Split raw storage lookup from synchronized response projection in a later run-backed read migration PR. @@ -696,22 +646,23 @@ def _find_existing_report_output_row( year: str, queryer=None, ) -> dict | None: - queryer = queryer or self.persistence api_version = get_report_output_cache_version(country_id) - query = """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND year = ? AND api_version = ? - """ - params: list[int | str] = [country_id, simulation_1_id, year, api_version] - if simulation_2_id is not None: - query += " AND simulation_2_id = ?" - params.append(simulation_2_id) - else: - query += " AND simulation_2_id IS NULL" - query += " ORDER BY id DESC" - - row = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None + if queryer is None: + with self.unit_of_work.read() as repositories: + return self._find_existing_report_output_row( + country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + year=year, + queryer=repositories, + ) + return queryer.reports.find_latest( + country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + year=year, + api_version=api_version, + ) def _get_or_create_current_report_output(self, report_output: dict) -> dict: current_report = self.find_existing_report_output( @@ -780,88 +731,55 @@ def create_report_output( api_version = get_report_output_cache_version(country_id) try: - - def tx_callback(tx): + with self.unit_of_work.transaction() as repositories: existing_report = self._find_existing_report_output_row( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, year=year, - queryer=tx, + queryer=repositories, ) if existing_report is not None: print( f"Reusing existing report output with ID: {existing_report['id']}" ) return self._ensure_report_output_dual_write_state_in_transaction( - tx, + repositories, existing_report["id"], country_id=country_id, ) self._require_simulation_exists( - tx, + repositories, country_id=country_id, simulation_id=simulation_1_id, ) if simulation_2_id is not None: self._require_simulation_exists( - tx, + repositories, country_id=country_id, simulation_id=simulation_2_id, ) - if simulation_2_id is not None: - tx.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - country_id, - simulation_1_id, - simulation_2_id, - api_version, - "pending", - year, - ), - ) - else: - tx.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?) - """, - ( - country_id, - simulation_1_id, - api_version, - "pending", - year, - ), - ) - - created_report = self._find_existing_report_output_row( + report_output_id = repositories.reports.create( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, + api_version=api_version, + status="pending", year=year, - queryer=tx, ) + created_report = repositories.reports.get(report_output_id, country_id) if created_report is None: raise Exception("Failed to retrieve created report output") print(f"Created report output with ID: {created_report['id']}") return self._ensure_report_output_dual_write_state_in_transaction( - tx, + repositories, created_report["id"], country_id=country_id, ) - return self.persistence.transaction(tx_callback) - except Exception as e: print(f"Error creating report output. Details: {str(e)}") raise e @@ -914,29 +832,25 @@ def update_report_output( print(f"Updating report output {report_id}") try: - update_fields = [] - update_values = [] + update_values = {} if status is not None: - update_fields.append("status = ?") - update_values.append(status) + update_values["status"] = status if output is not None: - update_fields.append("output = ?") - update_values.append(output) + update_values["output"] = parse_json_field(output) if error_message is not None: - update_fields.append("error_message = ?") - update_values.append(error_message) + update_values["error_message"] = error_message - if not update_fields: + if not update_values: print("No fields to update") return False - def tx_callback(tx): + with self.unit_of_work.transaction() as repositories: requested_report = self._get_report_output_row( report_id, - queryer=tx, + queryer=repositories, country_id=country_id, for_update=True, ) @@ -944,25 +858,20 @@ def tx_callback(tx): raise ValueError(f"Report output #{report_id} not found") if status == "running" and not self._has_mutable_running_run( - requested_report, queryer=tx + requested_report, queryer=repositories ): raise ValueError( "Cannot mark report output running without an active " "pending or running report run" ) - tx.query( - f"UPDATE report_outputs SET {', '.join(update_fields)} WHERE id = ? AND country_id = ?", - (*update_values, report_id, country_id), - ) + repositories.reports.update(report_id, **update_values) self._ensure_report_output_dual_write_state_in_transaction( - tx, + repositories, report_id, country_id=country_id, ) - self.persistence.transaction(tx_callback) - print(f"Successfully updated report output #{report_id}") return True diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index ca813e368..b1f6c2faf 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -4,7 +4,7 @@ from typing import Any from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO +from policyengine_api.data.v1_daos import ReportDAO, V1UnitOfWork from policyengine_api.services.run_sync_utils import select_display_report_run @@ -22,14 +22,20 @@ class ReportRunService: - def __init__(self, reports: ReportDAO | None = None): + def __init__( + self, + reports: ReportDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._reports = reports + self._unit_of_work = unit_of_work @property - def reports(self) -> ReportDAO: - if self._reports is None: - self._reports = ReportDAO(build_v1_session_manager()) - return self._reports + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work def _parse_run_row(self, row: dict | None) -> dict | None: if row is None: @@ -77,24 +83,43 @@ def create_report_output_run( } ) try: - run = self.reports.create_run( - report_output_id, run_id=run_id or str(uuid.uuid4()), **values - ) + if self._reports is not None: + run = self._reports.create_run( + report_output_id, + run_id=run_id or str(uuid.uuid4()), + **values, + ) + else: + with self.unit_of_work.transaction() as repositories: + run = repositories.reports.create_run( + report_output_id, + run_id=run_id or str(uuid.uuid4()), + **values, + ) except LookupError as error: raise ValueError(f"Report output #{report_output_id} not found") from error return self._parse_run_row(run) def get_report_output_run(self, run_id: str) -> dict | None: - return self._parse_run_row(self.reports.get_run(run_id)) + if self._reports is not None: + return self._parse_run_row(self._reports.get_run(run_id)) + with self.unit_of_work.read() as repositories: + return self._parse_run_row(repositories.reports.get_run(run_id)) def list_report_output_runs(self, report_output_id: int) -> list[dict]: - return [ - self._parse_run_row(row) - for row in reversed(self.reports.list_runs(report_output_id)) - ] + if self._reports is not None: + rows = self._reports.list_runs(report_output_id) + else: + with self.unit_of_work.read() as repositories: + rows = repositories.reports.list_runs(report_output_id) + return [self._parse_run_row(row) for row in reversed(rows)] def get_newest_report_output_run(self, report_output_id: int) -> dict | None: - rows = self.reports.list_runs(report_output_id) + if self._reports is not None: + rows = self._reports.list_runs(report_output_id) + else: + with self.unit_of_work.read() as repositories: + rows = repositories.reports.list_runs(report_output_id) return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, report_output: dict) -> dict | None: diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index 648f5b435..dce85b619 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -4,7 +4,7 @@ from pydantic import BaseModel, Field from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO +from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO, V1UnitOfWork REPORT_SPEC_SCHEMA_VERSION = 1 REPORT_SPEC_STATUSES = {"explicit", "backfilled_assumed"} @@ -46,15 +46,18 @@ def __init__( self, reports: ReportDAO | None = None, simulations: SimulationDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, ): self._reports = reports self._simulations = simulations + self._unit_of_work = unit_of_work - def _ensure_daos(self) -> None: - if self._reports is None or self._simulations is None: - manager = build_v1_session_manager() - self._reports = self._reports or ReportDAO(manager) - self._simulations = self._simulations or SimulationDAO(manager) + @property + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: @@ -63,15 +66,28 @@ def _validate_schema_version(self, schema_version: int | None) -> None: ) def _get_report_output_row(self, report_output_id: int) -> dict | None: - self._ensure_daos() - return self._reports.get(report_output_id) + if self._reports is not None: + return self._reports.get(report_output_id) + with self.unit_of_work.read() as repositories: + return repositories.reports.get(report_output_id) def _get_simulation_row(self, simulation_id: int) -> dict | None: - self._ensure_daos() - return self._simulations.get(simulation_id) - - def _get_linked_simulations(self, report_output: dict) -> tuple[dict, dict | None]: - simulation_1 = self._get_simulation_row(report_output["simulation_1_id"]) + if self._simulations is not None: + return self._simulations.get(simulation_id) + with self.unit_of_work.read() as repositories: + return repositories.simulations.get(simulation_id) + + def _get_linked_simulations( + self, report_output: dict, *, repositories=None, simulations=None + ) -> tuple[dict, dict | None]: + simulations = simulations or self._simulations + if repositories is None and simulations is None: + with self.unit_of_work.read() as read_repositories: + return self._get_linked_simulations( + report_output, repositories=read_repositories + ) + simulations = simulations or repositories.simulations + simulation_1 = simulations.get(report_output["simulation_1_id"]) if simulation_1 is None: raise ValueError( "Report output references missing simulation " @@ -80,7 +96,7 @@ def _get_linked_simulations(self, report_output: dict) -> tuple[dict, dict | Non simulation_2 = None if report_output["simulation_2_id"] is not None: - simulation_2 = self._get_simulation_row(report_output["simulation_2_id"]) + simulation_2 = simulations.get(report_output["simulation_2_id"]) if simulation_2 is None: raise ValueError( "Report output references missing simulation " @@ -216,9 +232,18 @@ def _build_economy_report_spec( ) def _validate_report_spec_matches_row( - self, report_output: dict, report_spec: ReportSpec + self, + report_output: dict, + report_spec: ReportSpec, + *, + repositories=None, + simulations=None, ) -> None: - simulation_1, simulation_2 = self._get_linked_simulations(report_output) + simulation_1, simulation_2 = self._get_linked_simulations( + report_output, + repositories=repositories, + simulations=simulations, + ) inferred_report_kind = self.infer_report_kind(simulation_1, simulation_2) if report_spec.country_id != report_output["country_id"]: raise ValueError("Report spec country must match report output country") @@ -355,14 +380,32 @@ def _parse_report_spec(self, report_kind: str, raw_spec: dict) -> ReportSpec: raise ValueError(f"Unsupported report kind: {report_kind}") def get_report_spec(self, report_output_id: int) -> ReportSpec | None: - report_output = self._get_report_output_row(report_output_id) + if self._reports is not None and self._simulations is not None: + report_output = self._reports.get(report_output_id) + repositories = None + else: + with self.unit_of_work.read() as repositories: + return self._get_report_spec(report_output_id, repositories) + return self._parse_stored_report_spec(report_output, repositories=repositories) + + def _get_report_spec( + self, report_output_id: int, repositories + ) -> ReportSpec | None: + return self._parse_stored_report_spec( + repositories.reports.get(report_output_id), repositories=repositories + ) + + def _parse_stored_report_spec( + self, report_output: dict | None, *, repositories=None + ) -> ReportSpec | None: if report_output is None or report_output["report_spec_json"] is None: return None - self._validate_schema_version(report_output["report_spec_schema_version"]) raw_spec = self._parse_json_field(report_output["report_spec_json"]) report_spec = self._parse_report_spec(report_output["report_kind"], raw_spec) - self._validate_report_spec_matches_row(report_output, report_spec) + self._validate_report_spec_matches_row( + report_output, report_spec, repositories=repositories + ) return report_spec def set_report_spec( @@ -376,17 +419,52 @@ def set_report_spec( raise ValueError(f"Unsupported report spec status: {report_spec_status}") self._validate_schema_version(schema_version) - report_output = self._get_report_output_row(report_output_id) + if self._reports is not None and self._simulations is not None: + self._set_report_spec( + self._reports, + self._simulations, + report_output_id, + report_spec, + report_spec_status, + schema_version, + ) + return True + + with self.unit_of_work.transaction() as repositories: + report_output = repositories.reports.get(report_output_id) + if report_output is None: + raise ValueError(f"Report output #{report_output_id} not found") + self._validate_report_spec_matches_row( + report_output, report_spec, repositories=repositories + ) + repositories.reports.update( + report_output_id, + report_kind=report_spec.report_kind, + report_spec_json=report_spec.model_dump(), + report_spec_schema_version=schema_version, + report_spec_status=report_spec_status, + ) + return True + + def _set_report_spec( + self, + reports: ReportDAO, + simulations: SimulationDAO, + report_output_id: int, + report_spec: ReportSpec, + report_spec_status: str, + schema_version: int, + ) -> None: + report_output = reports.get(report_output_id) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") - self._validate_report_spec_matches_row(report_output, report_spec) - - self._ensure_daos() - self._reports.update( + self._validate_report_spec_matches_row( + report_output, report_spec, simulations=simulations + ) + reports.update( report_output_id, report_kind=report_spec.report_kind, report_spec_json=report_spec.model_dump(), report_spec_schema_version=schema_version, report_spec_status=report_spec_status, ) - return True diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index da417a4ce..0eb74b2af 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -3,7 +3,7 @@ from typing import Any from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO +from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork SIMULATION_RUN_VERSION_FIELDS = ( @@ -16,14 +16,20 @@ class SimulationRunService: - def __init__(self, simulations: SimulationDAO | None = None): + def __init__( + self, + simulations: SimulationDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._simulations = simulations + self._unit_of_work = unit_of_work @property - def simulations(self) -> SimulationDAO: - if self._simulations is None: - self._simulations = SimulationDAO(build_v1_session_manager()) - return self._simulations + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work def _parse_run_row(self, row: dict | None) -> dict | None: if row is None: @@ -69,24 +75,43 @@ def create_simulation_run( } ) try: - run = self.simulations.create_run( - simulation_id, run_id=run_id or str(uuid.uuid4()), **values - ) + if self._simulations is not None: + run = self._simulations.create_run( + simulation_id, + run_id=run_id or str(uuid.uuid4()), + **values, + ) + else: + with self.unit_of_work.transaction() as repositories: + run = repositories.simulations.create_run( + simulation_id, + run_id=run_id or str(uuid.uuid4()), + **values, + ) except LookupError as error: raise ValueError(f"Simulation #{simulation_id} not found") from error return self._parse_run_row(run) def get_simulation_run(self, run_id: str) -> dict | None: - return self._parse_run_row(self.simulations.get_run(run_id)) + if self._simulations is not None: + return self._parse_run_row(self._simulations.get_run(run_id)) + with self.unit_of_work.read() as repositories: + return self._parse_run_row(repositories.simulations.get_run(run_id)) def list_simulation_runs(self, simulation_id: int) -> list[dict]: - return [ - self._parse_run_row(row) - for row in reversed(self.simulations.list_runs(simulation_id)) - ] + if self._simulations is not None: + rows = self._simulations.list_runs(simulation_id) + else: + with self.unit_of_work.read() as repositories: + rows = repositories.simulations.list_runs(simulation_id) + return [self._parse_run_row(row) for row in reversed(rows)] def get_newest_simulation_run(self, simulation_id: int) -> dict | None: - rows = self.simulations.list_runs(simulation_id) + if self._simulations is not None: + rows = self._simulations.list_runs(simulation_id) + else: + with self.unit_of_work.read() as repositories: + rows = repositories.simulations.list_runs(simulation_id) return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, simulation: dict) -> dict | None: diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 31c21a4af..8a835f99b 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -2,21 +2,32 @@ from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO +from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork from policyengine_api.services.simulation_spec_service import SimulationSpecService class SimulationService: - def __init__(self, simulations: SimulationDAO | None = None): + def __init__( + self, + simulations: SimulationDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._simulations = simulations - self.simulation_spec_service = SimulationSpecService(simulations) + self._unit_of_work = unit_of_work + self.simulation_spec_service = SimulationSpecService( + simulations, + unit_of_work=unit_of_work, + ) @property - def simulations(self) -> SimulationDAO: - if self._simulations is None: - self._simulations = SimulationDAO(build_v1_session_manager()) - self.simulation_spec_service = SimulationSpecService(self._simulations) - return self._simulations + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + self.simulation_spec_service = SimulationSpecService( + unit_of_work=self._unit_of_work + ) + return self._unit_of_work def _ensure_simulation_dual_write_state_in_transaction( self, @@ -25,9 +36,10 @@ def _ensure_simulation_dual_write_state_in_transaction( *, country_id: str | None = None, ) -> dict: - session = getattr(session, "session", session) - return self.simulations.ensure_dual_write_state_in_session( - session, simulation_id, country_id + return SimulationDAO(session).ensure_dual_write_state_in_session( + session, + simulation_id, + country_id, ) def _get_simulation_row( @@ -39,16 +51,25 @@ def _get_simulation_row( for_update: bool = False, ) -> dict | None: del for_update - if queryer is not None and getattr(queryer, "session", None) is not None: - return self.simulations.get_in_session( - queryer.session, simulation_id, country_id - ) - return self.simulations.get(simulation_id, country_id) + if queryer is not None: + simulations = getattr(queryer, "simulations", None) + if simulations is None: + simulations = SimulationDAO(getattr(queryer, "session", queryer)) + return simulations.get(simulation_id, country_id) + if self._simulations is not None: + return self._simulations.get(simulation_id, country_id) + with self.unit_of_work.read() as repositories: + return repositories.simulations.get(simulation_id, country_id) def ensure_simulation_dual_write_state( self, simulation_id: int, country_id: str | None = None ) -> dict: - return self.simulations.ensure_dual_write_state(simulation_id, country_id) + if self._simulations is not None: + return self._simulations.ensure_dual_write_state(simulation_id, country_id) + with self.unit_of_work.transaction() as repositories: + return repositories.simulations.ensure_dual_write_state( + simulation_id, country_id + ) def find_existing_simulation( self, @@ -57,12 +78,20 @@ def find_existing_simulation( population_type: str, policy_id: int, ) -> dict | None: - return self.simulations.find_latest( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) + if self._simulations is not None: + return self._simulations.find_latest( + country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + ) + with self.unit_of_work.read() as repositories: + return repositories.simulations.find_latest( + country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + ) def create_simulation( self, @@ -71,22 +100,31 @@ def create_simulation( population_type: str, policy_id: int, ) -> dict: - return self.simulations.create_or_get_with_sync( - sync_callback=self._ensure_simulation_dual_write_state_in_transaction, - country_id=country_id, - api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - status="pending", - ) + values = { + "country_id": country_id, + "api_version": COUNTRY_PACKAGE_VERSIONS.get(country_id), + "population_id": population_id, + "population_type": population_type, + "policy_id": policy_id, + "status": "pending", + } + if self._simulations is not None: + return self._simulations.create_or_get_with_sync( + sync_callback=self._ensure_simulation_dual_write_state_in_transaction, + **values, + ) + with self.unit_of_work.transaction() as repositories: + return repositories.simulations.create_or_get_with_sync( + sync_callback=self._ensure_simulation_dual_write_state_in_transaction, + **values, + ) def get_simulation(self, country_id: str, simulation_id: int) -> dict | None: if type(simulation_id) is not int or simulation_id < 0: raise Exception( f"Invalid simulation ID: {simulation_id}. Must be a positive integer." ) - return self.simulations.get(simulation_id, country_id) + return self._get_simulation_row(simulation_id, country_id=country_id) def update_simulation( self, @@ -110,10 +148,19 @@ def update_simulation( if isinstance(values.get("output"), str): values["output"] = json.loads(values["output"]) values["api_version"] = COUNTRY_PACKAGE_VERSIONS.get(country_id) - self.simulations.update_with_sync( - simulation_id, - country_id, - values, - self._ensure_simulation_dual_write_state_in_transaction, - ) + if self._simulations is not None: + self._simulations.update_with_sync( + simulation_id, + country_id, + values, + self._ensure_simulation_dual_write_state_in_transaction, + ) + else: + with self.unit_of_work.transaction() as repositories: + repositories.simulations.update_with_sync( + simulation_id, + country_id, + values, + self._ensure_simulation_dual_write_state_in_transaction, + ) return True diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index 2c9892a10..a848dc97a 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -3,7 +3,7 @@ from pydantic import BaseModel from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO +from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork SIMULATION_SPEC_SCHEMA_VERSION = 1 @@ -16,14 +16,20 @@ class SimulationSpec(BaseModel): class SimulationSpecService: - def __init__(self, simulations: SimulationDAO | None = None): + def __init__( + self, + simulations: SimulationDAO | None = None, + *, + unit_of_work: V1UnitOfWork | None = None, + ): self._simulations = simulations + self._unit_of_work = unit_of_work @property - def simulations(self) -> SimulationDAO: - if self._simulations is None: - self._simulations = SimulationDAO(build_v1_session_manager()) - return self._simulations + def unit_of_work(self) -> V1UnitOfWork: + if self._unit_of_work is None: + self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) + return self._unit_of_work def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != SIMULATION_SPEC_SCHEMA_VERSION: @@ -32,7 +38,10 @@ def _validate_schema_version(self, schema_version: int | None) -> None: ) def _get_simulation_row(self, simulation_id: int) -> dict | None: - return self.simulations.get(simulation_id) + if self._simulations is not None: + return self._simulations.get(simulation_id) + with self.unit_of_work.read() as repositories: + return repositories.simulations.get(simulation_id) def _validate_simulation_spec_matches_row( self, simulation: dict, simulation_spec: SimulationSpec @@ -76,14 +85,27 @@ def set_simulation_spec( schema_version: int = SIMULATION_SPEC_SCHEMA_VERSION, ) -> bool: self._validate_schema_version(schema_version) - simulation = self._get_simulation_row(simulation_id) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - self._validate_simulation_spec_matches_row(simulation, simulation_spec) + if self._simulations is not None: + simulations = self._simulations + simulation = simulations.get(simulation_id) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") + self._validate_simulation_spec_matches_row(simulation, simulation_spec) + simulations.update( + simulation_id, + simulation_spec_json=simulation_spec.model_dump(), + simulation_spec_schema_version=schema_version, + ) + return True - self.simulations.update( - simulation_id, - simulation_spec_json=simulation_spec.model_dump(), - simulation_spec_schema_version=schema_version, - ) + with self.unit_of_work.transaction() as repositories: + simulation = repositories.simulations.get(simulation_id) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") + self._validate_simulation_spec_matches_row(simulation, simulation_spec) + repositories.simulations.update( + simulation_id, + simulation_spec_json=simulation_spec.model_dump(), + simulation_spec_schema_version=schema_version, + ) return True diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py index 9bda3cf4d..21bbf582e 100644 --- a/tests/unit/data/test_run_daos.py +++ b/tests/unit/data/test_run_daos.py @@ -1,61 +1,68 @@ from datetime import datetime from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO +from policyengine_api.data.v1_daos import V1UnitOfWork from tests.unit.data.sqlite_schema import create_sqlite_v1_schema -def _daos(): +def _unit_of_work(): manager = build_sqlite_session_manager() create_sqlite_v1_schema(manager) - return SimulationDAO(manager), ReportDAO(manager) + return V1UnitOfWork(manager) def test_simulation_dao_creates_parent_and_monotonic_runs_atomically(): - simulations, _ = _daos() - simulation_id = simulations.create( - country_id="us", - api_version="1", - population_id="7", - population_type="household", - policy_id=2, - ) - first = simulations.create_run( - simulation_id, - run_id="run-1", - status="pending", - trigger_type="create", - requested_at=datetime(2026, 1, 1), - ) - second = simulations.create_run( - simulation_id, - run_id="run-2", - status="pending", - trigger_type="retry", - requested_at=datetime(2026, 1, 2), - ) - assert first["run_sequence"] == 1 - assert second["run_sequence"] == 2 - assert simulations.list_runs(simulation_id)[0]["id"] == "run-2" + uow = _unit_of_work() + with uow.transaction() as repositories: + simulation_id = repositories.simulations.create( + country_id="us", + api_version="1", + population_id="7", + population_type="household", + policy_id=2, + ) + first = repositories.simulations.create_run( + simulation_id, + run_id="run-1", + status="pending", + trigger_type="create", + requested_at=datetime(2026, 1, 1), + ) + second = repositories.simulations.create_run( + simulation_id, + run_id="run-2", + status="pending", + trigger_type="retry", + requested_at=datetime(2026, 1, 2), + ) + with uow.read() as repositories: + assert first["run_sequence"] == 1 + assert second["run_sequence"] == 2 + assert repositories.simulations.list_runs(simulation_id)[0]["id"] == "run-2" def test_report_dao_round_trips_parent_run_and_alias(): - _, reports = _daos() - report_id = reports.create( - country_id="us", - simulation_1_id=1, - simulation_2_id=None, - api_version="1", - year="2026", - ) - run = reports.create_run( - report_id, - run_id="report-run", - status="pending", - trigger_type="create", - requested_at=datetime(2026, 1, 1), - ) - reports.set_alias(99, report_id) - assert reports.get(report_id)["status"] == "pending" - assert run["run_sequence"] == 1 - assert reports.get_alias(99)["canonical_report_output_id"] == report_id + uow = _unit_of_work() + with uow.transaction() as repositories: + report_id = repositories.reports.create( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + year="2026", + ) + run = repositories.reports.create_run( + report_id, + run_id="report-run", + status="pending", + trigger_type="create", + requested_at=datetime(2026, 1, 1), + ) + repositories.reports.set_alias(99, report_id) + with uow.read() as repositories: + assert repositories.reports.get(report_id)["status"] == "pending" + assert run["run_sequence"] == 1 + assert ( + repositories.reports.get_alias(99)["canonical_report_output_id"] + == report_id + ) diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index 66572564e..a61af0371 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -37,3 +37,19 @@ def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): for relative_path in relative_paths: source = (PACKAGE_ROOT / relative_path).read_text(encoding="utf-8") assert "runtime_sqlalchemy_dao" not in source + + +def test_report_orchestration_uses_typed_repositories_not_raw_sql(): + source = (PACKAGE_ROOT / "services/report_output_service.py").read_text( + encoding="utf-8" + ) + assert "SQLAlchemyDAO" not in source + assert ".query(" not in source + assert "exec_driver_sql" not in source + + +def test_typed_repository_module_has_no_raw_sql_compatibility_dao(): + source = (PACKAGE_ROOT / "data/v1_daos.py").read_text(encoding="utf-8") + assert "SQLAlchemyDAO" not in source + assert "runtime_sqlalchemy_dao" not in source + assert "exec_driver_sql" not in source diff --git a/tests/unit/data/test_v1_unit_of_work.py b/tests/unit/data/test_v1_unit_of_work.py index cb8b9b84d..633d7c16b 100644 --- a/tests/unit/data/test_v1_unit_of_work.py +++ b/tests/unit/data/test_v1_unit_of_work.py @@ -35,3 +35,30 @@ def test_unit_of_work_rolls_back_every_repository_on_failure(): with uow.read() as repositories: assert repositories.policies.get("us", 1) is None assert repositories.users.get_profile(auth0_id="auth0|one") is None + + +def test_unit_of_work_rolls_back_parent_run_and_alias_together(): + uow = _unit_of_work() + + with pytest.raises(RuntimeError, match="abort report"): + with uow.transaction() as repositories: + report_id = repositories.reports.create( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + year="2026", + ) + repositories.reports.create_run( + report_id, + run_id="report-run", + status="pending", + trigger_type="create", + ) + repositories.reports.set_alias(99, report_id) + raise RuntimeError("abort report") + + with uow.read() as repositories: + assert repositories.reports.get(1) is None + assert repositories.reports.get_run("report-run") is None + assert repositories.reports.get_alias(99) is None From fbac8b83f54041774a1ce600e317eaaffb39cd47 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:25:28 +0300 Subject: [PATCH 37/89] fix: own database resources per worker lifecycle --- gcp/cloud_run/start.sh | 2 +- gcp/policyengine_api/start.sh | 2 +- policyengine_api/asgi.py | 6 +++- policyengine_api/asgi_factory.py | 12 +++++++ policyengine_api/data/data.py | 37 +++++++++++++++++++-- tests/unit/data/test_sqlalchemy_v2.py | 13 ++++++++ tests/unit/test_asgi_factory.py | 15 ++++++++- tests/unit/test_cloud_run_deploy_scripts.py | 11 ++++++ 8 files changed, 91 insertions(+), 7 deletions(-) diff --git a/gcp/cloud_run/start.sh b/gcp/cloud_run/start.sh index 31e2b188b..45a704d8a 100755 --- a/gcp/cloud_run/start.sh +++ b/gcp/cloud_run/start.sh @@ -61,7 +61,7 @@ done # gunicorn's master binds the listen socket before forking workers, so the # Cloud Run TCP startup probe passes immediately instead of racing the # multi-minute app import (which happens in the worker, post-fork, because -# --preload is NOT set). --timeout 0 is required: a worker mid-import does +# application preloading is disabled). --timeout 0 is required: a worker mid-import does # not heartbeat, and the default 30s watchdog would kill it before boot. gunicorn policyengine_api.asgi:app \ --worker-class uvicorn.workers.UvicornWorker \ diff --git a/gcp/policyengine_api/start.sh b/gcp/policyengine_api/start.sh index 92818ba81..96189837b 100644 --- a/gcp/policyengine_api/start.sh +++ b/gcp/policyengine_api/start.sh @@ -19,7 +19,7 @@ until redis-cli -h "$CACHE_REDIS_HOST" -p "$CACHE_REDIS_PORT" ping >/dev/null 2> done # Start the API -gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 --preload & +gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 & # Keep the script running and handle shutdown gracefully trap "pkill -P $$; exit 1" INT TERM diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index 901f60b59..4ee8369eb 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -6,10 +6,14 @@ from policyengine_api.api import app as flask_app from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.data import close_runtime_databases from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.warmup import run_startup_warmup -app = application = create_asgi_app(flask_app) +app = application = create_asgi_app( + flask_app, + shutdown_callback=close_runtime_databases, +) # Warm the simulation machinery before serving (see policyengine_api.warmup). # POLICYENGINE_API_STARTUP_WARMUP=0 skips it. diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index f65893864..c2929d4a6 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -2,6 +2,8 @@ from __future__ import annotations +from collections.abc import Callable +from contextlib import asynccontextmanager import time from a2wsgi import WSGIMiddleware @@ -60,6 +62,7 @@ def create_asgi_app( *, route_settings: RouteImplementationSettings | None = None, dependencies: NativeRouteDependencies | None = None, + shutdown_callback: Callable[[], None] | None = None, ) -> FastAPI: """Create the Stage 2 FastAPI shell around the existing Flask app.""" @@ -68,12 +71,21 @@ def create_asgi_app( if dependencies is None: dependencies = NativeRouteDependencies.defaults() + @asynccontextmanager + async def lifespan(_app: FastAPI): + try: + yield + finally: + if shutdown_callback is not None: + shutdown_callback() + app = FastAPI( title="PolicyEngine API", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None, + lifespan=lifespan, ) # compresslevel 4 instead of starlette's default 9: /us/metadata is ~70MB # raw, and level-9 compression inside the request costs seconds of CPU on diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py index 31a0196e2..396822bb8 100644 --- a/policyengine_api/data/data.py +++ b/policyengine_api/data/data.py @@ -1,3 +1,4 @@ +import atexit import fcntl import sqlite3 from policyengine_api.constants import REPO, COUNTRY_PACKAGE_VERSIONS @@ -103,6 +104,7 @@ def __init__( initialize: bool = False, ): self.local = local + self._closed = False if local: # Local development uses a sqlite database. self.db_url = REPO / "policyengine_api" / "data" / "policyengine.db" @@ -158,12 +160,28 @@ def getconn(): pool_timeout=int(os.environ.get("POLICYENGINE_DB_POOL_TIMEOUT", "30")), ) - def _close_pool(self): + def close(self) -> None: + """Release process-owned database and connector resources.""" + + if getattr(self, "_closed", False): + return + if self.local: + connection = getattr(self, "_connection", None) + if connection is not None: + connection.close() + self._closed = True + return + try: self.pool.dispose() + finally: self.connector.close() - except Exception: - pass + self._closed = True + + def _close_pool(self): + """Backward-compatible alias for callers predating ``close``.""" + + self.close() def _execute_remote(self, query_args): """Execute a query against the remote database using @@ -272,3 +290,16 @@ def initialize(self): database = PolicyEngineDatabase(local=False, initialize=False) local_database = PolicyEngineDatabase(local=True, initialize=False) + + +def close_runtime_databases() -> None: + """Close database resources owned by the current application process.""" + + try: + database.close() + finally: + if local_database is not database: + local_database.close() + + +atexit.register(close_runtime_databases) diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 9566d7b14..6fc0c505c 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -12,6 +12,7 @@ import policyengine_api.data.data as data_module import sqlalchemy +from unittest.mock import Mock from policyengine_api.data.data import PolicyEngineDatabase, _ResultProxy @@ -276,3 +277,15 @@ def fake_create_pool(self): assert db.local is False assert calls == [("pool", False)] + + def test_close_disposes_engine_then_closes_connector(self): + db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) + db.local = False + db.pool = Mock() + db.connector = Mock() + + db.close() + db.close() + + db.pool.dispose.assert_called_once_with() + db.connector.close.assert_called_once_with() diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index 762c2925a..92d2649ce 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -4,7 +4,7 @@ import threading from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import policyengine_api.asgi_factory as asgi_factory @@ -409,6 +409,19 @@ def test_fastapi_documentation_routes_fall_through_to_flask_404(): assert "swagger" not in response.text.lower() +def test_lifespan_closes_runtime_resources_on_shutdown(): + close_runtime_resources = Mock() + with TestClient( + create_asgi_app( + create_test_wsgi_app(), + shutdown_callback=close_runtime_resources, + ) + ) as client: + assert client.get("/liveness-check").status_code == 200 + close_runtime_resources.assert_not_called() + close_runtime_resources.assert_called_once_with() + + def test_flask_fallback_preserves_status_body_headers_and_cookies(): client = TestClient(create_asgi_app(create_test_wsgi_app())) diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index 68dfcb4c3..b51b861af 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -483,6 +483,17 @@ def test_cloud_run_startup_supervises_redis_and_server_children(): assert re.search(r"(?m)^ *wait 2>/dev/null", start_script) is None +def test_production_gunicorn_workers_do_not_inherit_database_pools(): + for relative_path in ("gcp/cloud_run/start.sh", "gcp/policyengine_api/start.sh"): + start_script = (REPO / relative_path).read_text(encoding="utf-8") + commands = "\n".join( + line + for line in start_script.splitlines() + if not line.lstrip().startswith("#") + ) + assert "--preload" not in commands + + def test_validate_cloud_run_deploy_env_requires_selector_environment_variable(): result = _run_script( ".github/scripts/validate_cloud_run_deploy_env.sh", From f9bfd7da5659aabf8eb680aad2d7a3c6101a9a9a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:26:28 +0300 Subject: [PATCH 38/89] ci: gate Stage 7 on disposable MySQL --- .github/workflows/pr.yml | 32 ++++++++++ docs/engineering/skills/testing.md | 16 +++++ docs/migration/stage7-toy-database.md | 89 +++++++++++++++++++++++++++ 3 files changed, 137 insertions(+) create mode 100644 docs/migration/stage7-toy-database.md diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 2dd9b5235..6b9e07755 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -50,6 +50,38 @@ jobs: python-version: "3.12" - name: Run quality guards run: python scripts/run_quality_guards.py + stage7-toy-database: + name: Stage 7 toy MySQL database + runs-on: ubuntu-latest + services: + mysql: + image: mysql:8.0 + env: + MYSQL_DATABASE: policyengine_stage7_toy + MYSQL_PASSWORD: policyengine + MYSQL_ROOT_PASSWORD: policyengine-root + MYSQL_USER: policyengine + ports: + - 3306:3306 + options: >- + --health-cmd="mysqladmin ping --host=127.0.0.1 --user=policyengine --password=policyengine --silent" + --health-interval=2s + --health-timeout=2s + --health-retries=30 + --health-start-period=10s + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install dependencies + run: pip install -e ".[dev]" + - name: Qualify the disposable MySQL schema + run: python -m pytest tests/integration/test_stage7_*.py -v + env: + STAGE7_TOY_DATABASE_URL: mysql+pymysql://policyengine:policyengine@127.0.0.1:3306/policyengine_stage7_toy check-changelog: name: Check changelog fragment runs-on: ubuntu-latest diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index f75a797ab..5040be1cd 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -91,6 +91,22 @@ avoid depending on specific production data fixtures: API_BASE_URL=https://candidate-url python -m pytest tests/integration/test_cloud_run_candidate.py -v ``` +For the Stage 7 SQLAlchemy boundary, keep ordinary unit tests on isolated +SQLite databases and qualify dialect-specific behavior against the disposable +MySQL scaffold: + +```bash +make stage7-toy-test +make stage7-toy-down +``` + +The MySQL suite must upgrade a fresh database, exercise every migrated DAO +domain, prove legacy behavior and route parity, cover the Cloud SQL connector +seam, verify startup emits no DDL, qualify an independent pre-Alembic schema for +stamping without data loss, report no ORM metadata drift, downgrade to `base`, +and upgrade to `head` again. Destructive tests must never target production or +shared infrastructure; live existing-schema comparison must be read-only. + Before committing AI-authored code changes, run repository formatting and lint: ```bash diff --git a/docs/migration/stage7-toy-database.md b/docs/migration/stage7-toy-database.md new file mode 100644 index 000000000..a7b593ceb --- /dev/null +++ b/docs/migration/stage7-toy-database.md @@ -0,0 +1,89 @@ +# Stage 7 Toy Database + +The Stage 7 toy database is a disposable MySQL 8 instance used to prove the +SQLAlchemy and Alembic boundary on the same database dialect as the current +Cloud SQL service. It contains synthetic qualification records only. + +## Local qualification + +Docker with the Compose plugin is required. Run: + +```bash +make stage7-toy-test +make stage7-toy-down +``` + +`stage7-toy-test` starts the service, waits for MySQL's health check, and runs +the integration suite. `stage7-toy-down` removes the container and its volumes. +The database data directory is also mounted as `tmpfs`, so data does not +survive the container. + +Port `3307` is used by default to avoid a typical local MySQL server. Override +it consistently when needed: + +```bash +STAGE7_TOY_MYSQL_PORT=13307 \ +STAGE7_TOY_DATABASE_URL=mysql+pymysql://policyengine:policyengine@127.0.0.1:13307/policyengine_stage7_toy \ +make stage7-toy-test +``` + +The test teardown runs `alembic downgrade base`, which is destructive. A hard +safety guard permits that operation only for MySQL on `localhost` or +`127.0.0.1` and only when the database name ends in `_toy`. + +## Qualification targets + +The suite must prove all of the following before Stage 7 proceeds: + +1. A fresh MySQL database upgrades to Alembic `head`. +2. Every migrated DAO domain can write and read synthetic data. +3. The upgraded schema has no drift from the reviewed SQLAlchemy metadata. +4. The baseline downgrades to `base` and upgrades to `head` again. +5. The independently defined pre-Alembic schema compares without drift and can + be stamped without losing an existing sentinel row. +6. Typed DAO results preserve the legacy service-level mapping shapes. +7. Policy, household, and user routes operate against MySQL. +8. The Cloud SQL connector/pool seam drives typed DAOs. +9. Importing and starting the Flask application emits no MySQL DDL. + +Pull requests run the same suite against a fresh MySQL 8 service container. + +## Canonical SQLAlchemy boundary + +The Stage 7 runtime follows SQLAlchemy's documented ownership model: + +- one `Engine` and its `QueuePool` are created per worker process; +- the Cloud SQL `creator` returns one fresh DBAPI connection whenever the pool + requests one, while SQLAlchemy owns checkout, return, pre-ping, recycling, + overflow, and timeout behavior; +- a `sessionmaker` creates a short-lived `Session` for each service operation; +- service-level units of work use `sessionmaker.begin()` so success commits and + exceptions roll back and close the session automatically; +- typed repositories receive the operation's `Session` and never create, + commit, roll back, retry, or retain sessions themselves; +- application startup performs no DDL; Alembic alone owns schema changes; and +- ASGI lifespan and process-exit cleanup dispose the engine and close the Cloud + SQL connector. Gunicorn application preloading remains disabled so workers do + not inherit pooled connections across a fork. + +These constraints reflect SQLAlchemy's guidance for +[contextual session/transaction management](https://docs.sqlalchemy.org/en/20/orm/session_basics.html#framing-out-a-begin-commit-rollback-block), +[engine disposal](https://docs.sqlalchemy.org/en/20/core/connections.html#engine-disposal), +and [pooling with multiprocessing](https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork). + +## Existing Cloud SQL comparison + +The production check is intentionally read-only and skipped unless an explicit +URL is supplied. Use credentials whose database user has read-only access: + +```bash +STAGE7_EXISTING_DATABASE_URL='' \ +uv run pytest \ + tests/integration/test_stage7_existing_schema.py::test_live_existing_schema_matches_metadata_without_mutation \ + -v +``` + +Do not stamp an existing database from this command. Stamping is qualified only +against the disposable pre-Alembic fixture; a real database may be stamped only +after its read-only comparison is empty, its backup is confirmed, and a human +approves the target. From acae6e24cbebfcc04ba3cd6cbfcb59db3cbd20f0 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 01:44:23 +0300 Subject: [PATCH 39/89] test: align legacy suites with ORM boundaries --- .../routes/user_profile_routes.py | 4 - tests/integration/test_stage7_mysql_parity.py | 28 ++- .../to_refactor_household_fixtures.py | 18 +- .../python/test_ai_analysis_service_old.py | 48 ++--- .../python/test_household_routes.py | 69 +++---- .../python/test_policy_service_old.py | 189 ++++++------------ .../python/test_tracer_analysis_routes.py | 49 ++--- tests/unit/data/test_alembic_baseline.py | 47 ++--- 8 files changed, 171 insertions(+), 281 deletions(-) diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index d859629c6..ed725e641 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,6 +1,5 @@ from flask import Blueprint, Response, request from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data import database import json from policyengine_api.services.user_service import UserService from werkzeug.exceptions import BadRequest, NotFound @@ -94,9 +93,6 @@ def update_user_profile(country_id: str) -> Response: will assume malicious intent and 403 """ - # Construct the relevant UPDATE request - setter_array = [] - args = [] payload = request.json if payload is None: diff --git a/tests/integration/test_stage7_mysql_parity.py b/tests/integration/test_stage7_mysql_parity.py index 4b0006f1c..e27e3d6c6 100644 --- a/tests/integration/test_stage7_mysql_parity.py +++ b/tests/integration/test_stage7_mysql_parity.py @@ -16,9 +16,7 @@ def test_typed_daos_preserve_legacy_mapping_shapes(stage7_mysql): unit_of_work = V1UnitOfWork(SessionManager(engine)) with unit_of_work.transaction() as repositories: - policy_id = repositories.policies.create( - "us", None, {"x": 1}, "parity", "v1" - ) + policy_id = repositories.policies.create("us", None, {"x": 1}, "parity", "v1") household_id = repositories.households.create( "us", None, {"people": {}}, "household-parity", "v1" ) @@ -34,7 +32,9 @@ def test_typed_daos_preserve_legacy_mapping_shapes(stage7_mysql): "WHERE id = :policy_id AND country_id = :country_id" ), {"policy_id": policy_id, "country_id": "us"}, - ).mappings().one() + ) + .mappings() + .one() ) legacy_household = dict( connection.execute( @@ -43,13 +43,17 @@ def test_typed_daos_preserve_legacy_mapping_shapes(stage7_mysql): "WHERE id = :household_id AND country_id = :country_id" ), {"household_id": household_id, "country_id": "us"}, - ).mappings().one() + ) + .mappings() + .one() ) legacy_user = dict( connection.execute( text("SELECT * FROM user_profiles WHERE user_id = :user_id"), {"user_id": user_id}, - ).mappings().one() + ) + .mappings() + .one() ) legacy_policy["policy_json"] = json.loads(legacy_policy["policy_json"]) @@ -80,10 +84,14 @@ def test_typed_daos_read_rows_written_by_legacy_sql(stage7_mysql): "household_hash": "legacy-row", }, ) - row = connection.execute( - text("SELECT * FROM household WHERE household_hash = :household_hash"), - {"household_hash": "legacy-row"}, - ).mappings().one() + row = ( + connection.execute( + text("SELECT * FROM household WHERE household_hash = :household_hash"), + {"household_hash": "legacy-row"}, + ) + .mappings() + .one() + ) legacy_shape = dict(row) legacy_shape["household_json"] = json.loads(legacy_shape["household_json"]) diff --git a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py index 5fa6af91c..8dbd4bb23 100644 --- a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py +++ b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py @@ -16,19 +16,11 @@ "api_version": "3.0.0", } -valid_hash_value = "some-hash" - - -@pytest.fixture -def mock_hash_object(): - """Mock the hash_object function.""" - with patch("policyengine_api.services.household_service.hash_object") as mock: - mock.return_value = valid_hash_value - yield mock - @pytest.fixture def mock_database(): - """Mock the database module.""" - with patch("policyengine_api.services.household_service.database") as mock_db: - yield mock_db + """Replace the route's service with its typed persistence boundary.""" + with patch( + "policyengine_api.routes.household_routes.household_service" + ) as household_service: + yield household_service diff --git a/tests/to_refactor/python/test_ai_analysis_service_old.py b/tests/to_refactor/python/test_ai_analysis_service_old.py index 0df3928ca..950e39ae0 100644 --- a/tests/to_refactor/python/test_ai_analysis_service_old.py +++ b/tests/to_refactor/python/test_ai_analysis_service_old.py @@ -1,58 +1,44 @@ -import pytest -from unittest.mock import patch, MagicMock import json import os -from policyengine_api.services.ai_analysis_service import AIAnalysisService - -test_ai_service = AIAnalysisService() +from unittest.mock import MagicMock, patch +import pytest -@patch("policyengine_api.services.ai_analysis_service.local_database") -def test_get_existing_analysis_found(mock_db): - mock_db.query.return_value.fetchone.return_value = {"analysis": "Existing analysis"} +from policyengine_api.services.ai_analysis_service import AIAnalysisService - prompt = "Test prompt" - output = test_ai_service.get_existing_analysis(prompt) - assert output == json.dumps("Existing analysis") +def test_get_existing_analysis_found(): + analyses = MagicMock() + analyses.get.return_value = "Existing analysis" + service = AIAnalysisService(analyses) - # Check database query - mock_db.query.assert_called_once_with( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ) + output = service.get_existing_analysis("Test prompt") + assert output == json.dumps("Existing analysis") + analyses.get.assert_called_once_with("Test prompt") -@patch("policyengine_api.services.ai_analysis_service.local_database") -def test_get_existing_analysis_not_found(mock_db): - mock_db.query.return_value.fetchone.return_value = None - prompt = "Test prompt" - result = test_ai_service.get_existing_analysis(prompt) +def test_get_existing_analysis_not_found(): + analyses = MagicMock() + analyses.get.return_value = None + service = AIAnalysisService(analyses) - assert result is None - mock_db.query.assert_called_once_with( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ) + assert service.get_existing_analysis("Test prompt") is None + analyses.get.assert_called_once_with("Test prompt") -# Additional test to check environment variable def test_anthropic_api_key(): with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test_key"}): assert os.getenv("ANTHROPIC_API_KEY") == "test_key" -# Test error handling in trigger_ai_analysis @patch("policyengine_api.services.ai_analysis_service.anthropic.Anthropic") def test_trigger_ai_analysis_error(mock_anthropic): mock_client = MagicMock() mock_anthropic.return_value = mock_client mock_client.messages.stream.side_effect = Exception("API Error") - prompt = "Test prompt" - generator = test_ai_service.trigger_ai_analysis(prompt) + generator = AIAnalysisService(MagicMock()).trigger_ai_analysis("Test prompt") - # The generator should stop after the initial yield due to the error with pytest.raises(Exception, match="API Error"): list(generator) diff --git a/tests/to_refactor/python/test_household_routes.py b/tests/to_refactor/python/test_household_routes.py index 3456429dc..b40f70f8c 100644 --- a/tests/to_refactor/python/test_household_routes.py +++ b/tests/to_refactor/python/test_household_routes.py @@ -1,26 +1,21 @@ -import pytest import json -from unittest.mock import MagicMock, patch - -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from unittest.mock import patch from tests.to_refactor.fixtures.to_refactor_household_fixtures import ( valid_request_body, valid_db_row, - mock_database, - mock_hash_object, ) +pytest_plugins = ["tests.to_refactor.fixtures.to_refactor_household_fixtures"] + class TestGetHousehold: def test_get_existing_household(self, rest_client, mock_database): """Test getting an existing household.""" - # Mock database response as a dict-like object - # (SQLAlchemy v2 Row objects support dict() via ._mapping) - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: valid_db_row[x] - mock_row.keys.return_value = valid_db_row.keys() - mock_database.query().fetchone.return_value = mock_row + mock_database.get_household.return_value = { + **valid_db_row, + "household_json": valid_request_body["data"], + } # Make request response = rest_client.get("/us/household/1") @@ -32,7 +27,7 @@ def test_get_existing_household(self, rest_client, mock_database): def test_get_nonexistent_household(self, rest_client, mock_database): """Test getting a non-existent household.""" - mock_database.query().fetchone.return_value = None + mock_database.get_household.return_value = None response = rest_client.get("/us/household/999") data = json.loads(response.data) @@ -50,14 +45,9 @@ def test_get_household_invalid_id(self, rest_client): class TestCreateHousehold: - def test_create_household_success( - self, rest_client, mock_database, mock_hash_object - ): + def test_create_household_success(self, rest_client, mock_database): """Test successfully creating a new household.""" - # Mock database responses - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: {"id": 1}[x] - mock_database.query().fetchone.return_value = mock_row + mock_database.create_household.return_value = 1 response = rest_client.post( "/us/household", @@ -104,15 +94,12 @@ def test_create_household_invalid_label(self, rest_client): class TestUpdateHousehold: - def test_update_household_success( - self, rest_client, mock_database, mock_hash_object - ): + def test_update_household_success(self, rest_client, mock_database): """Test successfully updating an existing household.""" - # Mock getting existing household - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: valid_db_row[x] - mock_row.keys.return_value = valid_db_row.keys() - mock_database.query().fetchone.return_value = mock_row + mock_database.get_household.return_value = { + **valid_db_row, + "household_json": valid_request_body["data"], + } updated_household = {"people": {"person1": {"age": 31, "income": 55000}}} @@ -120,6 +107,10 @@ def test_update_household_success( "data": updated_household, "label": valid_request_body["label"], } + mock_database.update_household.return_value = { + **valid_db_row, + "household_json": updated_household, + } response = rest_client.put( "/us/household/1", @@ -131,25 +122,17 @@ def test_update_household_success( assert response.status_code == 200 assert data["status"] == "ok" assert data["result"]["household_id"] == 1 - # assert data["result"]["household_json"] == updated_data["data"] - # WHERE now includes country_id (issue #3447). - mock_database.query.assert_any_call( - "UPDATE household " - "SET household_json = ?, household_hash = ?, label = ?, api_version = ? " - "WHERE id = ? AND country_id = ?", - ( - json.dumps(updated_household), - "some-hash", - valid_request_body["label"], - COUNTRY_PACKAGE_VERSIONS.get("us"), - 1, - "us", - ), + assert data["result"]["household_json"] == updated_data["data"] + mock_database.update_household.assert_called_once_with( + "us", + 1, + updated_household, + valid_request_body["label"], ) def test_update_nonexistent_household(self, rest_client, mock_database): """Test updating a non-existent household.""" - mock_database.query().fetchone.return_value = None + mock_database.get_household.return_value = None response = rest_client.put( "/us/household/999", diff --git a/tests/to_refactor/python/test_policy_service_old.py b/tests/to_refactor/python/test_policy_service_old.py index a84e1b1b0..f3ba129e1 100644 --- a/tests/to_refactor/python/test_policy_service_old.py +++ b/tests/to_refactor/python/test_policy_service_old.py @@ -1,22 +1,24 @@ -import pytest -from assertpy import assert_that -from unittest.mock import patch, MagicMock, ANY, call import json +from unittest.mock import MagicMock + +from assertpy import assert_that +import pytest + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.services.policy_service import PolicyService @pytest.fixture -def mock_database(): - with patch("policyengine_api.services.policy_service.database") as mock_db: - yield mock_db +def policies(): + return MagicMock() @pytest.fixture def sample_policy_data(): return { "id": 1, - "country_id": "US", - "policy_json": json.dumps({"param": "value"}), + "country_id": "us", + "policy_json": {"param": "value"}, "policy_hash": "hash123", "label": "test_policy", "api_version": "1.0.0", @@ -24,167 +26,108 @@ def sample_policy_data(): @pytest.fixture -def policy_service(): - return PolicyService() +def policy_service(policies): + return PolicyService(policies) class TestPolicyService: - a_test_policy_id = 8 # Pre-seeded current law policies occupy IDs 1 through 5 + a_test_policy_id = 8 - def test_get_policy_success( - self, policy_service, mock_database, sample_policy_data - ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data + def test_get_policy_success(self, policy_service, policies, sample_policy_data): + policies.get.return_value = sample_policy_data - # Test result = policy_service.get_policy("us", self.a_test_policy_id) - # Verify assert_that(result).contains_entry({"policy_json": {"param": "value"}}) - mock_database.query.assert_called_once_with( - "SELECT * FROM policy WHERE country_id = ? AND id = ?", - ("us", self.a_test_policy_id), - ) - - def test_get_policy_not_found(self, policy_service, mock_database): - # Setup mock - mock_database.query.return_value.fetchone.return_value = None + policies.get.assert_called_once_with("us", self.a_test_policy_id) - # Test - garbage_id = 999 - result = policy_service.get_policy("us", garbage_id) + def test_get_policy_not_found(self, policy_service, policies): + policies.get.return_value = None - # Verify - assert result is None - mock_database.query.assert_called_once() + assert policy_service.get_policy("us", 999) is None + policies.get.assert_called_once_with("us", 999) - def test_get_policy_json(self, policy_service, mock_database, sample_policy_data): - # Setup mock - mock_database.query.return_value.fetchone.return_value = { - "policy_json": sample_policy_data["policy_json"] - } + def test_get_policy_json(self, policy_service, policies, sample_policy_data): + policies.get.return_value = sample_policy_data - # Test result = policy_service.get_policy_json("us", self.a_test_policy_id) - # Verify - assert result == sample_policy_data["policy_json"] - mock_database.query.assert_called_once() - - def test_set_policy_new(self, policy_service, mock_database): - new_policy_id = 10 - - # Setup mocks - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call for existing policy check - {"id": new_policy_id}, # Second call to get new policy - ] + assert result == json.dumps(sample_policy_data["policy_json"]) + policies.get.assert_called_once_with("us", self.a_test_policy_id) + def test_set_policy_new(self, policy_service, policies): + policies.find_unique.return_value = None + policies.create.return_value = 10 test_policy = {"param": "value"} - test_label = "new_policy" - test_country_id = "us" - - expected_calls = [ - # First call - check if policy exists - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, ANY, test_label), - ), - # Second call - insert new policy - call( - "INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - test_country_id, - json.dumps(test_policy), - ANY, - test_label, - ANY, - ), - ), - # Third call - get the newly created policy - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, ANY, test_label), - ), - ] - - # Test + policy_id, message, exists = policy_service.set_policy( - test_country_id, test_label, test_policy + "us", "new_policy", test_policy ) - # Verify - assert policy_id == new_policy_id - assert message == "Policy created" - assert exists is False - assert mock_database.query.call_args_list == expected_calls + assert (policy_id, message, exists) == (10, "Policy created", False) + policies.find_unique.assert_called_once() + policy_hash = policies.find_unique.call_args.args[1] + policies.create.assert_called_once_with( + "us", + "new_policy", + test_policy, + policy_hash, + COUNTRY_PACKAGE_VERSIONS["us"], + ) - def test_set_policy_existing( - self, policy_service, mock_database, sample_policy_data - ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data + def test_set_policy_existing(self, policy_service, policies, sample_policy_data): + policies.find_unique.return_value = sample_policy_data - # Test - policy_id, message, exists = policy_service.set_policy( + result = policy_service.set_policy( "us", sample_policy_data["label"], - json.loads(sample_policy_data["policy_json"]), + sample_policy_data["policy_json"], ) - # Verify - assert policy_id == sample_policy_data["id"] - assert message == "Policy already exists" - assert exists is True - mock_database.query.assert_called_once() + assert result == (sample_policy_data["id"], "Policy already exists", True) + policies.create.assert_not_called() def test_get_unique_policy_with_label( - self, policy_service, mock_database, sample_policy_data + self, policy_service, policies, sample_policy_data ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data + policies.find_unique.return_value = sample_policy_data - # Test result = policy_service._get_unique_policy_with_label( "us", sample_policy_data["policy_hash"], sample_policy_data["label"], ) - # Verify assert result == sample_policy_data - mock_database.query.assert_called_once() + policies.find_unique.assert_called_once_with( + "us", + sample_policy_data["policy_hash"], + sample_policy_data["label"], + ) - def test_get_unique_policy_with_null_label(self, policy_service, mock_database): - # Setup mock - mock_database.query.return_value.fetchone.return_value = None + def test_get_unique_policy_with_null_label(self, policy_service, policies): + policies.find_unique.return_value = None - # Test result = policy_service._get_unique_policy_with_label("us", "hash123", None) - # Verify assert result is None - mock_database.query.assert_called_once_with( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label IS NULL", - ("us", "hash123"), - ) + policies.find_unique.assert_called_once_with("us", "hash123", None) @pytest.mark.parametrize( - "error_method", + ("error_method", "repository_method"), [ - "get_policy", - "get_policy_json", - "set_policy", - "_get_unique_policy_with_label", + ("get_policy", "get"), + ("get_policy_json", "get"), + ("set_policy", "find_unique"), + ("_get_unique_policy_with_label", "find_unique"), ], ) - def test_error_handling(self, policy_service, mock_database, error_method): - # Setup mock to raise exception - mock_database.query.side_effect = Exception("Database error") + def test_error_handling( + self, policy_service, policies, error_method, repository_method + ): + getattr(policies, repository_method).side_effect = Exception("Database error") - # Test - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Database error"): if error_method == "get_policy": policy_service.get_policy("us", 1) elif error_method == "get_policy_json": @@ -193,5 +136,3 @@ def test_error_handling(self, policy_service, mock_database, error_method): policy_service.set_policy("us", "label", {}) else: policy_service._get_unique_policy_with_label("us", "hash", "label") - - assert str(exc_info.value) == "Database error" diff --git a/tests/to_refactor/python/test_tracer_analysis_routes.py b/tests/to_refactor/python/test_tracer_analysis_routes.py index 83f7bde23..851c67bef 100644 --- a/tests/to_refactor/python/test_tracer_analysis_routes.py +++ b/tests/to_refactor/python/test_tracer_analysis_routes.py @@ -1,7 +1,6 @@ -import pytest from flask import json from unittest.mock import patch -from werkzeug.exceptions import BadRequest +from werkzeug.exceptions import NotFound # constants VALID_HOUSEHOLD_ID = 123 @@ -12,19 +11,12 @@ INVALID_VARIABLE = 123 -@patch("policyengine_api.services.tracer_analysis_service.local_database") -@patch( - "policyengine_api.services.tracer_analysis_service.TracerAnalysisService.trigger_ai_analysis" -) -def test_execute_tracer_analysis_success( - mock_trigger_ai_analysis, mock_db, rest_client -): - mock_db.query.return_value.fetchone.return_value = { - "tracer_output": json.dumps( - ["disposable_income <1000>", " market_income <1000>"] - ) - } - mock_trigger_ai_analysis.return_value = "AI analysis result" +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_success(mock_service, rest_client): + mock_service.execute_analysis.return_value = ( + iter(["AI analysis result"]), + "streaming", + ) test_household_id = 1500 # Set this to US current law @@ -43,9 +35,11 @@ def test_execute_tracer_analysis_success( assert b"AI analysis result" in response.data -@patch("policyengine_api.services.tracer_analysis_service.local_database") -def test_execute_tracer_analysis_no_tracer(mock_db, rest_client): - mock_db.query.return_value.fetchone.return_value = None +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_no_tracer(mock_service, rest_client): + mock_service.execute_analysis.side_effect = NotFound( + "No household simulation tracer found" + ) response = rest_client.post( "/us/tracer-analysis", @@ -62,19 +56,9 @@ def test_execute_tracer_analysis_no_tracer(mock_db, rest_client): ) -@patch("policyengine_api.services.tracer_analysis_service.local_database") -@patch( - "policyengine_api.services.tracer_analysis_service.TracerAnalysisService.trigger_ai_analysis" -) -def test_execute_tracer_analysis_ai_error( - mock_trigger_ai_analysis, mock_db, rest_client -): - mock_db.query.return_value.fetchone.return_value = { - "tracer_output": json.dumps( - ["disposable_income <1000>", " market_income <1000>"] - ) - } - mock_trigger_ai_analysis.side_effect = Exception(KeyError) +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_ai_error(mock_service, rest_client): + mock_service.execute_analysis.side_effect = Exception(KeyError) test_household_id = 1500 test_policy_id = 2 @@ -93,8 +77,7 @@ def test_execute_tracer_analysis_ai_error( assert json.loads(response.data)["status"] == "error" -@patch("policyengine_api.services.tracer_analysis_service.local_database") -def test_invalid_variable_types(mock_db, rest_client): +def test_invalid_variable_types(rest_client): """Test that different non-string variable types are rejected""" invalid_variables = [ 123, diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index 6bce1e892..8be0415a7 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -1,37 +1,38 @@ -from pathlib import Path +from io import StringIO from alembic import command from alembic.config import Config -from sqlalchemy import create_engine, inspect +from alembic.script import ScriptDirectory from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base -def _config(url: str) -> Config: - config = Config(str(REPO / "alembic.ini")) - config.set_main_option("sqlalchemy.url", url) - return config +def _mysql_offline_config() -> tuple[Config, StringIO]: + output = StringIO() + config = Config(str(REPO / "alembic.ini"), output_buffer=output) + config.set_main_option( + "sqlalchemy.url", + "mysql+pymysql://offline:offline@localhost/offline", + ) + return config, output -def test_baseline_upgrades_fresh_database_to_v1_head(tmp_path: Path): - database_path = tmp_path / "fresh.db" - command.upgrade(_config(f"sqlite+pysqlite:///{database_path}"), "head") +def test_baseline_renders_the_v1_schema_for_mysql_without_connecting(): + config, output = _mysql_offline_config() - tables = set( - inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names() - ) - assert set(V1Base.metadata.tables) <= tables - assert "alembic_version" in tables + command.upgrade(config, "head", sql=True) + rendered_sql = output.getvalue() + for table_name in V1Base.metadata.tables: + assert f"CREATE TABLE {table_name}" in rendered_sql + assert "CREATE TABLE alembic_version" in rendered_sql -def test_baseline_downgrades_and_reupgrades(tmp_path: Path): - database_path = tmp_path / "roundtrip.db" - config = _config(f"sqlite+pysqlite:///{database_path}") - command.upgrade(config, "head") - command.downgrade(config, "base") - command.upgrade(config, "head") - assert set(V1Base.metadata.tables) <= set( - inspect(create_engine(f"sqlite+pysqlite:///{database_path}")).get_table_names() - ) +def test_baseline_is_the_single_root_revision(): + config, _ = _mysql_offline_config() + scripts = ScriptDirectory.from_config(config) + head = scripts.get_revision(scripts.get_current_head()) + + assert head is not None + assert head.down_revision is None From 99dbe3ccd45acbb22bb516c6a5a4c0e6fe27ee8f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 17:39:27 +0300 Subject: [PATCH 40/89] fix: preserve canonical JSON boundaries --- policyengine_api/data/data.py | 3 + policyengine_api/routes/simulation_routes.py | 23 ++++++-- policyengine_api/services/economy_service.py | 21 +++++-- .../services/reform_impacts_service.py | 7 ++- tests/integration/test_stage7_mysql_parity.py | 59 +++++++++++++++++++ .../integration/test_stage7_mysql_runtime.py | 29 +++++++++ tests/unit/services/test_economy_service.py | 43 ++++++++++++++ tests/unit/test_stage5_routes.py | 37 ++++++++++++ 8 files changed, 209 insertions(+), 13 deletions(-) diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py index 396822bb8..4cb8a0d6e 100644 --- a/policyengine_api/data/data.py +++ b/policyengine_api/data/data.py @@ -289,6 +289,9 @@ def initialize(self): else: database = PolicyEngineDatabase(local=False, initialize=False) +# TODO: Remove this eager SQLite-backed local database initialization and +# replace it with a traditional cache so importing the application does not +# create files or perform schema setup. local_database = PolicyEngineDatabase(local=True, initialize=False) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index f2bacd6cb..06f6fe7d8 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -12,6 +12,21 @@ simulation_service = SimulationService() +def _serialize_v1_simulation(simulation: dict) -> dict: + """Project canonical ORM JSON objects onto the legacy v1 response shape.""" + + response = dict(simulation) + for field in ("output", "simulation_spec_json"): + value = response.get(field) + if value is not None and not isinstance(value, str): + response[field] = json.dumps(value) + + # TODO: Remove this compatibility projection when the v1 contract is + # retired. New v2 response models should expose these fields as native JSON + # objects instead of carrying the historical JSON-in-a-string shape forward. + return response + + @simulation_bp.route("//simulation", methods=["POST"]) @validate_country def create_simulation(country_id: str) -> Response: @@ -67,7 +82,7 @@ def create_simulation(country_id: str) -> Response: response_body = dict( status="ok", message="Simulation already exists", - result=existing_simulation, + result=_serialize_v1_simulation(existing_simulation), ) return Response( @@ -87,7 +102,7 @@ def create_simulation(country_id: str) -> Response: response_body = dict( status="ok", message="Simulation created successfully", - result=created_simulation, + result=_serialize_v1_simulation(created_simulation), ) return Response( @@ -139,7 +154,7 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: response_body = dict( status="ok", message=None, - result=simulation, + result=_serialize_v1_simulation(simulation), ) return Response( @@ -213,7 +228,7 @@ def update_simulation(country_id: str) -> Response: response_body = dict( status="ok", message="Simulation updated successfully", - result=updated_simulation, + result=_serialize_v1_simulation(updated_simulation), ) return Response( diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index d896d6d01..e8f512509 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -235,6 +235,15 @@ class EconomyService: with other services to access their respective tables """ + @staticmethod + def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: + """Accept ORM-decoded objects and legacy JSON text at the read boundary.""" + + parsed = json.loads(value) if isinstance(value, str) else value + if not isinstance(parsed, dict): + raise TypeError("Expected a JSON object for reform impact data") + return parsed + def get_economic_impact( self, country_id: str, @@ -820,7 +829,7 @@ def _handle_execution_state( ) self._set_reform_impact_complete( setup_options=setup_options, - reform_impact_json=json.dumps(result), + reform_impact_json=result, execution_id=reform_impact["execution_id"], ) logger.log_struct( @@ -867,7 +876,7 @@ def _handle_completed_impact( setup_options: EconomicImpactSetupOptions, most_recent_impact: dict, ) -> EconomicImpactResult: - result = json.loads(most_recent_impact["reform_impact_json"]) + result = self._parse_json_object(most_recent_impact["reform_impact_json"]) return EconomicImpactResult.completed( data=self._with_policyengine_bundle( result=result, @@ -1047,7 +1056,7 @@ def _extract_dataset_version(self, dataset: str | None) -> str | None: def _extract_cached_result(self, most_recent_impact: dict) -> dict: try: - return json.loads(most_recent_impact["reform_impact_json"]) + return self._parse_json_object(most_recent_impact["reform_impact_json"]) except (TypeError, ValueError): return {} @@ -1314,11 +1323,11 @@ def _set_reform_impact_computing( region=setup_options.region, dataset=setup_options.dataset, time_period=setup_options.time_period, - options=json.dumps(setup_options.options), + options=setup_options.options, options_hash=setup_options.options_hash, status=ImpactStatus.COMPUTING.value, api_version=setup_options.api_version, - reform_impact_json=json.dumps({}), + reform_impact_json={}, start_time=datetime.datetime.now(), execution_id=execution_id, ) @@ -1334,7 +1343,7 @@ def _set_reform_impact_computing( def _set_reform_impact_complete( self, setup_options: EconomicImpactSetupOptions, - reform_impact_json: str, + reform_impact_json: dict[str, Any], execution_id: str, ): """ diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 4a2ddc919..09f6e1e54 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -1,5 +1,6 @@ import datetime from contextlib import contextmanager +from typing import Any from policyengine_api.data.orm import build_v1_session_manager from policyengine_api.data.v1_daos import ReformImpactDAO, V1UnitOfWork @@ -112,11 +113,11 @@ def set_reform_impact( region, dataset, time_period, - options, + options: dict[str, Any], options_hash, status, api_version, - reform_impact_json, + reform_impact_json: dict[str, Any], start_time, execution_id: str, ): @@ -193,7 +194,7 @@ def set_complete_reform_impact( dataset, time_period, options_hash, - reform_impact_json, + reform_impact_json: dict[str, Any], execution_id, ): del ( diff --git a/tests/integration/test_stage7_mysql_parity.py b/tests/integration/test_stage7_mysql_parity.py index e27e3d6c6..0bf16a290 100644 --- a/tests/integration/test_stage7_mysql_parity.py +++ b/tests/integration/test_stage7_mysql_parity.py @@ -1,12 +1,14 @@ """Behavioral parity between legacy SQL and typed DAO access on MySQL.""" import json +from datetime import datetime from alembic import command from sqlalchemy import text from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.services.reform_impacts_service import ReformImpactsService from tests.integration.stage7_mysql import alembic_config @@ -97,3 +99,60 @@ def test_typed_daos_read_rows_written_by_legacy_sql(stage7_mysql): legacy_shape["household_json"] = json.loads(legacy_shape["household_json"]) with unit_of_work.read() as repositories: assert repositories.households.get("us", row["id"]) == legacy_shape + + +def test_reform_impact_service_stores_mysql_json_objects(stage7_mysql): + database_url, engine = stage7_mysql + command.upgrade(alembic_config(database_url), "head") + unit_of_work = V1UnitOfWork(SessionManager(engine)) + service = ReformImpactsService(unit_of_work=unit_of_work) + + service.set_reform_impact( + country_id="us", + policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options={"scope": "test"}, + options_hash="native-json", + status="computing", + api_version="v1", + reform_impact_json={}, + start_time=datetime(2026, 1, 1), + execution_id="native-json-job", + ) + service.set_complete_reform_impact( + country_id="us", + reform_policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options_hash="native-json", + reform_impact_json={"result": {"value": 1}}, + execution_id="native-json-job", + ) + + with engine.connect() as connection: + json_types = connection.execute( + text( + "SELECT JSON_TYPE(options_json), JSON_TYPE(reform_impact_json) " + "FROM reform_impact WHERE execution_id = :execution_id" + ), + {"execution_id": "native-json-job"}, + ).one() + assert tuple(json_types) == ("OBJECT", "OBJECT") + + stored = service.get_all_reform_impacts( + "us", + 2, + 1, + "us", + "default", + "2026", + "native-json", + "v1", + )[0] + assert stored["options_json"] == {"scope": "test"} + assert stored["reform_impact_json"] == {"result": {"value": 1}} diff --git a/tests/integration/test_stage7_mysql_runtime.py b/tests/integration/test_stage7_mysql_runtime.py index 01e0d575d..0890eb5aa 100644 --- a/tests/integration/test_stage7_mysql_runtime.py +++ b/tests/integration/test_stage7_mysql_runtime.py @@ -13,6 +13,7 @@ from policyengine_api.data.v1_daos import V1UnitOfWork from policyengine_api.services.household_service import HouseholdService from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.simulation_service import SimulationService from policyengine_api.services.user_service import UserService from tests.integration.stage7_mysql import alembic_config, schema_signature @@ -25,6 +26,7 @@ def test_public_policy_household_and_user_routes_use_mysql(stage7_mysql, monkeyp monkeypatch.setenv("FLASK_DEBUG", "1") from policyengine_api.routes import household_routes, policy_routes + from policyengine_api.routes import simulation_routes from policyengine_api.routes import user_profile_routes monkeypatch.setattr( @@ -42,9 +44,15 @@ def test_public_policy_household_and_user_routes_use_mysql(stage7_mysql, monkeyp "user_service", UserService(unit_of_work=unit_of_work), ) + monkeypatch.setattr( + simulation_routes, + "simulation_service", + SimulationService(unit_of_work=unit_of_work), + ) app = Flask(__name__) app.register_blueprint(policy_routes.policy_bp) app.register_blueprint(household_routes.household_bp) + app.register_blueprint(simulation_routes.simulation_bp) app.register_blueprint(user_profile_routes.user_profile_bp) client = app.test_client() @@ -72,6 +80,27 @@ def test_public_policy_household_and_user_routes_use_mysql(stage7_mysql, monkeyp user_id = user.get_json()["result"]["user_id"] assert client.get(f"/us/user-profile?user_id={user_id}").status_code == 200 + simulation = client.post( + "/us/simulation", + json={ + "population_id": str(household_id), + "population_type": "household", + "policy_id": policy_id, + }, + ) + assert simulation.status_code == 201 + simulation_id = simulation.get_json()["result"]["id"] + output = {"result": "ok"} + updated_simulation = client.patch( + "/us/simulation", + json={"id": simulation_id, "status": "complete", "output": output}, + ) + assert updated_simulation.status_code == 200 + assert json.loads(updated_simulation.get_json()["result"]["output"]) == output + fetched_simulation = client.get(f"/us/simulation/{simulation_id}") + assert fetched_simulation.status_code == 200 + assert json.loads(fetched_simulation.get_json()["result"]["output"]) == output + def test_cloud_sql_connector_pool_drives_daos_and_startup_emits_no_ddl( stage7_mysql, monkeypatch diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 9600f5124..9d61994c0 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -134,6 +134,36 @@ def test__given_completed_impact__returns_completed_result( ) mock_simulation_entrypoint.run.assert_not_called() + def test__given_orm_decoded_completed_impact__returns_completed_result( + self, + economy_service, + base_params, + mock_country_package_versions, + mock_policyengine_version, + mock_policy_service, + mock_reform_impacts_service, + mock_simulation_entrypoint, + mock_logger, + mock_datetime, + mock_numpy_random, + ): + completed_impact = create_mock_reform_impact(status="ok") + completed_impact["reform_impact_json"] = json.loads( + completed_impact["reform_impact_json"] + ) + mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ + completed_impact + ] + + result = economy_service.get_economic_impact(**base_params) + + assert result.status == ImpactStatus.OK + assert ( + result.data["poverty_impact"] + == MOCK_REFORM_IMPACT_DATA["poverty_impact"] + ) + mock_simulation_entrypoint.run.assert_not_called() + def test__given_legacy_completed_impact__refreshes_cache( self, economy_service, @@ -275,6 +305,11 @@ def test__given_no_previous_impact__creates_new_simulation( assert result.data is None mock_simulation_entrypoint.run.assert_called_once() mock_reform_impacts_service.set_reform_impact.assert_called_once() + write_values = ( + mock_reform_impacts_service.set_reform_impact.call_args.kwargs + ) + assert write_values["options"] == MOCK_OPTIONS + assert write_values["reform_impact_json"] == {} def test__given_no_previous_impact__includes_metadata_in_simulation_params( self, @@ -1413,6 +1448,14 @@ def test__given_succeeded_state__returns_completed_result( "dataset": MOCK_RESOLVED_DATASET, } mock_reform_impacts_service.set_complete_reform_impact.assert_called_once() + write_values = ( + mock_reform_impacts_service.set_complete_reform_impact.call_args.kwargs + ) + assert isinstance(write_values["reform_impact_json"], dict) + assert ( + write_values["reform_impact_json"]["poverty_impact"] + == (MOCK_REFORM_IMPACT_DATA["poverty_impact"]) + ) def test__given_failed_state__returns_error_result( self, diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_stage5_routes.py index ec9f34e1b..4a6c7be17 100644 --- a/tests/unit/test_stage5_routes.py +++ b/tests/unit/test_stage5_routes.py @@ -580,3 +580,40 @@ def test_patch_report_output_complete_promotes_active_rerun_route_path(test_db): ).fetchone() assert stored_report["active_run_id"] is None assert stored_report["latest_successful_run_id"] == rerun["id"] + + +def test_simulation_v1_routes_keep_json_fields_as_strings(test_db): + simulation = simulation_service.create_simulation( + country_id="us", + population_id="household_v1_json_contract", + population_type="household", + policy_id=51, + ) + output = {"result": "ok", "values": [1, 2, 3]} + + patch_response = create_test_client().patch( + "/us/simulation", + json={ + "id": simulation["id"], + "status": "complete", + "output": output, + }, + ) + + assert patch_response.status_code == 200 + patched_simulation = patch_response.get_json()["result"] + assert isinstance(patched_simulation["output"], str) + assert json.loads(patched_simulation["output"]) == output + assert isinstance(patched_simulation["simulation_spec_json"], str) + assert json.loads(patched_simulation["simulation_spec_json"])["country_id"] == "us" + + get_response = create_test_client().get(f"/us/simulation/{simulation['id']}") + + assert get_response.status_code == 200 + fetched_simulation = get_response.get_json()["result"] + assert isinstance(fetched_simulation["output"], str) + assert json.loads(fetched_simulation["output"]) == output + + orm_simulation = simulation_service.get_simulation("us", simulation["id"]) + assert orm_simulation["output"] == output + assert isinstance(orm_simulation["simulation_spec_json"], dict) From bdfa21436bfa218a89763b470c0d952e5a391c89 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 17:45:10 +0300 Subject: [PATCH 41/89] docs: add Stage 7 changelog fragment --- changelog.d/3788.changed.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/3788.changed.md diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md new file mode 100644 index 000000000..228e9c400 --- /dev/null +++ b/changelog.d/3788.changed.md @@ -0,0 +1 @@ +Migrate API v1 persistence to SQLAlchemy 2 repositories and Alembic while preserving the existing database schema and public API contracts. From bb4e4484b0ed019fdcc4855d98acfceaf5af74a4 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 17:53:01 +0300 Subject: [PATCH 42/89] chore: remove one-time Stage 7 toy database --- .github/workflows/pr.yml | 32 --- Makefile | 11 - compose.stage7-toy.yml | 22 -- docs/engineering/skills/testing.md | 16 -- docs/migration/stage7-toy-database.md | 89 -------- policyengine_api/scripts/__init__.py | 1 - .../scripts/qualify_stage7_toy.py | 158 -------------- policyengine_api/scripts/stage7_database.py | 18 -- tests/fixtures/stage7_pre_alembic_schema.sql | 194 ------------------ tests/integration/conftest.py | 2 - tests/integration/stage7_mysql.py | 86 -------- .../test_stage7_existing_schema.py | 68 ++---- tests/integration/test_stage7_mysql_parity.py | 158 -------------- .../integration/test_stage7_mysql_runtime.py | 164 --------------- tests/integration/test_stage7_toy_database.py | 47 ----- tests/unit/test_stage7_toy_qualification.py | 56 ----- 16 files changed, 17 insertions(+), 1105 deletions(-) delete mode 100644 compose.stage7-toy.yml delete mode 100644 docs/migration/stage7-toy-database.md delete mode 100644 policyengine_api/scripts/__init__.py delete mode 100644 policyengine_api/scripts/qualify_stage7_toy.py delete mode 100644 policyengine_api/scripts/stage7_database.py delete mode 100644 tests/fixtures/stage7_pre_alembic_schema.sql delete mode 100644 tests/integration/stage7_mysql.py delete mode 100644 tests/integration/test_stage7_mysql_parity.py delete mode 100644 tests/integration/test_stage7_mysql_runtime.py delete mode 100644 tests/integration/test_stage7_toy_database.py delete mode 100644 tests/unit/test_stage7_toy_qualification.py diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 6b9e07755..2dd9b5235 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -50,38 +50,6 @@ jobs: python-version: "3.12" - name: Run quality guards run: python scripts/run_quality_guards.py - stage7-toy-database: - name: Stage 7 toy MySQL database - runs-on: ubuntu-latest - services: - mysql: - image: mysql:8.0 - env: - MYSQL_DATABASE: policyengine_stage7_toy - MYSQL_PASSWORD: policyengine - MYSQL_ROOT_PASSWORD: policyengine-root - MYSQL_USER: policyengine - ports: - - 3306:3306 - options: >- - --health-cmd="mysqladmin ping --host=127.0.0.1 --user=policyengine --password=policyengine --silent" - --health-interval=2s - --health-timeout=2s - --health-retries=30 - --health-start-period=10s - steps: - - name: Checkout repo - uses: actions/checkout@v4 - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.12" - - name: Install dependencies - run: pip install -e ".[dev]" - - name: Qualify the disposable MySQL schema - run: python -m pytest tests/integration/test_stage7_*.py -v - env: - STAGE7_TOY_DATABASE_URL: mysql+pymysql://policyengine:policyengine@127.0.0.1:3306/policyengine_stage7_toy check-changelog: name: Check changelog fragment runs-on: ubuntu-latest diff --git a/Makefile b/Makefile index 377495201..d9bbc1dda 100644 --- a/Makefile +++ b/Makefile @@ -20,17 +20,6 @@ test: quality-guards: python scripts/run_quality_guards.py -STAGE7_TOY_DATABASE_URL ?= mysql+pymysql://policyengine:policyengine@127.0.0.1:3307/policyengine_stage7_toy - -stage7-toy-up: - docker compose -f compose.stage7-toy.yml up -d --wait - -stage7-toy-test: stage7-toy-up - STAGE7_TOY_DATABASE_URL="$(STAGE7_TOY_DATABASE_URL)" uv run pytest tests/integration/test_stage7_*.py -v - -stage7-toy-down: - docker compose -f compose.stage7-toy.yml down --volumes - debug-test: MAX_HOUSEHOLDS=1000 FLASK_DEBUG=1 pytest -vv --durations=0 tests diff --git a/compose.stage7-toy.yml b/compose.stage7-toy.yml deleted file mode 100644 index 9ad5eb505..000000000 --- a/compose.stage7-toy.yml +++ /dev/null @@ -1,22 +0,0 @@ -name: policyengine-stage7-toy - -services: - mysql: - image: mysql:8.0 - environment: - MYSQL_DATABASE: policyengine_stage7_toy - MYSQL_PASSWORD: policyengine - MYSQL_ROOT_PASSWORD: policyengine-root - MYSQL_USER: policyengine - ports: - - "127.0.0.1:${STAGE7_TOY_MYSQL_PORT:-3307}:3306" - tmpfs: - - /var/lib/mysql - healthcheck: - test: - - CMD-SHELL - - mysqladmin ping --host=127.0.0.1 --user=policyengine --password=policyengine --silent - interval: 2s - timeout: 2s - retries: 30 - start_period: 10s diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 5040be1cd..f75a797ab 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -91,22 +91,6 @@ avoid depending on specific production data fixtures: API_BASE_URL=https://candidate-url python -m pytest tests/integration/test_cloud_run_candidate.py -v ``` -For the Stage 7 SQLAlchemy boundary, keep ordinary unit tests on isolated -SQLite databases and qualify dialect-specific behavior against the disposable -MySQL scaffold: - -```bash -make stage7-toy-test -make stage7-toy-down -``` - -The MySQL suite must upgrade a fresh database, exercise every migrated DAO -domain, prove legacy behavior and route parity, cover the Cloud SQL connector -seam, verify startup emits no DDL, qualify an independent pre-Alembic schema for -stamping without data loss, report no ORM metadata drift, downgrade to `base`, -and upgrade to `head` again. Destructive tests must never target production or -shared infrastructure; live existing-schema comparison must be read-only. - Before committing AI-authored code changes, run repository formatting and lint: ```bash diff --git a/docs/migration/stage7-toy-database.md b/docs/migration/stage7-toy-database.md deleted file mode 100644 index a7b593ceb..000000000 --- a/docs/migration/stage7-toy-database.md +++ /dev/null @@ -1,89 +0,0 @@ -# Stage 7 Toy Database - -The Stage 7 toy database is a disposable MySQL 8 instance used to prove the -SQLAlchemy and Alembic boundary on the same database dialect as the current -Cloud SQL service. It contains synthetic qualification records only. - -## Local qualification - -Docker with the Compose plugin is required. Run: - -```bash -make stage7-toy-test -make stage7-toy-down -``` - -`stage7-toy-test` starts the service, waits for MySQL's health check, and runs -the integration suite. `stage7-toy-down` removes the container and its volumes. -The database data directory is also mounted as `tmpfs`, so data does not -survive the container. - -Port `3307` is used by default to avoid a typical local MySQL server. Override -it consistently when needed: - -```bash -STAGE7_TOY_MYSQL_PORT=13307 \ -STAGE7_TOY_DATABASE_URL=mysql+pymysql://policyengine:policyengine@127.0.0.1:13307/policyengine_stage7_toy \ -make stage7-toy-test -``` - -The test teardown runs `alembic downgrade base`, which is destructive. A hard -safety guard permits that operation only for MySQL on `localhost` or -`127.0.0.1` and only when the database name ends in `_toy`. - -## Qualification targets - -The suite must prove all of the following before Stage 7 proceeds: - -1. A fresh MySQL database upgrades to Alembic `head`. -2. Every migrated DAO domain can write and read synthetic data. -3. The upgraded schema has no drift from the reviewed SQLAlchemy metadata. -4. The baseline downgrades to `base` and upgrades to `head` again. -5. The independently defined pre-Alembic schema compares without drift and can - be stamped without losing an existing sentinel row. -6. Typed DAO results preserve the legacy service-level mapping shapes. -7. Policy, household, and user routes operate against MySQL. -8. The Cloud SQL connector/pool seam drives typed DAOs. -9. Importing and starting the Flask application emits no MySQL DDL. - -Pull requests run the same suite against a fresh MySQL 8 service container. - -## Canonical SQLAlchemy boundary - -The Stage 7 runtime follows SQLAlchemy's documented ownership model: - -- one `Engine` and its `QueuePool` are created per worker process; -- the Cloud SQL `creator` returns one fresh DBAPI connection whenever the pool - requests one, while SQLAlchemy owns checkout, return, pre-ping, recycling, - overflow, and timeout behavior; -- a `sessionmaker` creates a short-lived `Session` for each service operation; -- service-level units of work use `sessionmaker.begin()` so success commits and - exceptions roll back and close the session automatically; -- typed repositories receive the operation's `Session` and never create, - commit, roll back, retry, or retain sessions themselves; -- application startup performs no DDL; Alembic alone owns schema changes; and -- ASGI lifespan and process-exit cleanup dispose the engine and close the Cloud - SQL connector. Gunicorn application preloading remains disabled so workers do - not inherit pooled connections across a fork. - -These constraints reflect SQLAlchemy's guidance for -[contextual session/transaction management](https://docs.sqlalchemy.org/en/20/orm/session_basics.html#framing-out-a-begin-commit-rollback-block), -[engine disposal](https://docs.sqlalchemy.org/en/20/core/connections.html#engine-disposal), -and [pooling with multiprocessing](https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork). - -## Existing Cloud SQL comparison - -The production check is intentionally read-only and skipped unless an explicit -URL is supplied. Use credentials whose database user has read-only access: - -```bash -STAGE7_EXISTING_DATABASE_URL='' \ -uv run pytest \ - tests/integration/test_stage7_existing_schema.py::test_live_existing_schema_matches_metadata_without_mutation \ - -v -``` - -Do not stamp an existing database from this command. Stamping is qualified only -against the disposable pre-Alembic fixture; a real database may be stamped only -after its read-only comparison is empty, its backup is confirmed, and a human -approves the target. diff --git a/policyengine_api/scripts/__init__.py b/policyengine_api/scripts/__init__.py deleted file mode 100644 index 0af6e9d90..000000000 --- a/policyengine_api/scripts/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Executable qualification helpers shipped with the API package.""" diff --git a/policyengine_api/scripts/qualify_stage7_toy.py b/policyengine_api/scripts/qualify_stage7_toy.py deleted file mode 100644 index c88dd429e..000000000 --- a/policyengine_api/scripts/qualify_stage7_toy.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Run the Stage 7 ORM boundary against an isolated toy database.""" - -from datetime import datetime - -from alembic import command -from alembic.autogenerate import compare_metadata -from alembic.config import Config -from alembic.migration import MigrationContext -from sqlalchemy import create_engine, inspect - -from policyengine_api.constants import REPO -from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_daos import V1UnitOfWork -from policyengine_api.data.v1_models import V1Base - - -def compare_stage7_schema(database_url: str) -> list: - """Return metadata drift without stamping or mutating the target database.""" - - engine = create_engine(database_url) - with engine.connect() as connection: - context = MigrationContext.configure( - connection, - opts={"compare_type": True, "compare_server_default": True}, - ) - return compare_metadata(context, V1Base.metadata) - - -def qualify_stage7_toy(database_url: str) -> dict[str, bool]: - """Upgrade and exercise every migrated v1 persistence domain.""" - - config = Config(str(REPO / "alembic.ini")) - config.set_main_option("sqlalchemy.url", database_url) - command.upgrade(config, "head") - - engine = create_engine(database_url) - sessions = SessionManager(engine) - unit_of_work = V1UnitOfWork(sessions) - - with unit_of_work.transaction() as repositories: - policy_id = repositories.policies.create("us", "Toy", {}, "toy-policy", "toy") - household_id = repositories.households.create( - "us", "Toy", {}, "toy-household", "toy" - ) - user_id = repositories.users.create_profile("toy|user", "toy-user", "us", 1) - repositories.computed_households.create( - household_id=household_id, - policy_id=policy_id, - country_id="us", - api_version="toy", - computed_household_json={"qualified": True}, - status="complete", - ) - user_policy_id = repositories.user_policies.create( - country_id="us", - reform_id=policy_id, - reform_label="Toy", - baseline_id=policy_id, - baseline_label="Toy", - user_id=str(user_id), - year="2026", - geography="us", - dataset="default", - number_of_provisions=0, - api_version="toy", - added_date=1, - updated_date=1, - budgetary_impact=None, - type="reform", - ) - economy_id = repositories.economies.create( - policy_id=policy_id, - country_id="us", - region="us", - time_period="2026", - options_json={}, - options_hash="toy-economy", - api_version="toy", - economy_json={"qualified": True}, - status="complete", - message=None, - ) - repositories.analyses.store("toy prompt", "toy answer", "complete") - repositories.tracers.create(household_id, policy_id, "us", "toy", ["toy trace"]) - impact_id = repositories.reform_impacts.create( - baseline_policy_id=policy_id, - reform_policy_id=policy_id, - country_id="us", - region="us", - dataset="default", - time_period="2026", - options_json={}, - options_hash="toy-options", - api_version="toy", - reform_impact_json={}, - status="computing", - start_time=datetime(2026, 1, 1), - execution_id="toy-impact", - ) - simulation_id = repositories.simulations.create( - country_id="us", - api_version="toy", - population_id=str(household_id), - population_type="household", - policy_id=policy_id, - ) - repositories.simulations.create_run( - simulation_id, - run_id="toy-simulation-run", - status="pending", - trigger_type="qualification", - ) - report_id = repositories.reports.create( - country_id="us", - simulation_1_id=simulation_id, - simulation_2_id=None, - api_version="toy", - year="2026", - ) - repositories.reports.create_run( - report_id, - run_id="toy-report-run", - status="pending", - trigger_type="qualification", - ) - repositories.reports.set_alias(900_001, report_id) - - with unit_of_work.read() as repositories: - core_results = { - "policy": repositories.policies.get("us", policy_id) is not None, - "household": repositories.households.get("us", household_id) is not None, - "user": repositories.users.get_profile(user_id=user_id) is not None, - "computed_household": repositories.computed_households.get( - household_id, policy_id, "us" - ) - is not None, - "user_policy": repositories.user_policies.get(user_policy_id) is not None, - "economy": repositories.economies.get(economy_id) is not None, - "analysis": repositories.analyses.get("toy prompt") == "toy answer", - "tracer": repositories.tracers.get(household_id, policy_id, "us") - is not None, - "reform_impact": repositories.reform_impacts.find( - execution_id="toy-impact" - )["reform_impact_id"] - == impact_id, - "simulation": repositories.simulations.get_run("toy-simulation-run") - is not None, - "report": repositories.reports.get_run("toy-report-run") is not None, - "report_alias": repositories.reports.get_alias(900_001)[ - "canonical_report_output_id" - ] - == report_id, - } - - return { - "alembic_head": "alembic_version" in inspect(engine).get_table_names(), - **core_results, - } diff --git a/policyengine_api/scripts/stage7_database.py b/policyengine_api/scripts/stage7_database.py deleted file mode 100644 index 0d9d35fbc..000000000 --- a/policyengine_api/scripts/stage7_database.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Safety checks shared by destructive Stage 7 qualification tooling.""" - -from sqlalchemy.engine import make_url - - -def assert_safe_toy_database_url(database_url: str) -> None: - """Reject destructive toy-test operations against non-local databases.""" - - url = make_url(database_url) - is_local_mysql = url.get_backend_name() == "mysql" and url.host in { - "127.0.0.1", - "localhost", - } - is_toy_database = bool(url.database and url.database.endswith("_toy")) - if not is_local_mysql or not is_toy_database: - raise ValueError( - "Stage 7 integration tests require a local MySQL database ending in '_toy'" - ) diff --git a/tests/fixtures/stage7_pre_alembic_schema.sql b/tests/fixtures/stage7_pre_alembic_schema.sql deleted file mode 100644 index 48b75a20a..000000000 --- a/tests/fixtures/stage7_pre_alembic_schema.sql +++ /dev/null @@ -1,194 +0,0 @@ --- Frozen API v1 schema snapshot captured before Alembic ownership. --- Do not update this file to match ORM metadata or generated revisions. - -CREATE TABLE household ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(255) NOT NULL, - household_json JSON NOT NULL, - household_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE computed_household ( - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - computed_household_json JSON NOT NULL, - status VARCHAR(32), - PRIMARY KEY (household_id, policy_id, country_id) -); - -CREATE TABLE policy ( - id INTEGER AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - policy_json JSON NOT NULL, - policy_hash VARCHAR(255) NOT NULL, - PRIMARY KEY (id, country_id, policy_hash) -); - -CREATE TABLE economy ( - economy_id INTEGER PRIMARY KEY AUTO_INCREMENT, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32), - time_period VARCHAR(32), - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - economy_json JSON, - status VARCHAR(32) NOT NULL, - message VARCHAR(255) -); - -CREATE TABLE reform_impact ( - reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - dataset VARCHAR(255) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON, - options_hash VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME, - end_time DATETIME, - execution_id VARCHAR(255) NOT NULL -); - -CREATE TABLE analysis ( - prompt_id INTEGER PRIMARY KEY AUTO_INCREMENT, - prompt LONGTEXT NOT NULL, - analysis LONGTEXT, - status VARCHAR(32) NOT NULL -); - -CREATE TABLE user_policies ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - reform_id INTEGER NOT NULL, - reform_label VARCHAR(255), - baseline_id INTEGER NOT NULL, - baseline_label VARCHAR(255), - user_id VARCHAR(255) NOT NULL, - year VARCHAR(32) NOT NULL, - geography VARCHAR(255) NOT NULL, - dataset VARCHAR(255), - number_of_provisions INTEGER NOT NULL, - api_version VARCHAR(32) NOT NULL, - added_date BIGINT NOT NULL, - updated_date BIGINT NOT NULL, - budgetary_impact VARCHAR(255), - type VARCHAR(255) -); - -CREATE TABLE user_profiles ( - user_id INTEGER PRIMARY KEY AUTO_INCREMENT, - auth0_id VARCHAR(255) NOT NULL UNIQUE, - username VARCHAR(255) UNIQUE, - primary_country VARCHAR(3) NOT NULL, - user_since BIGINT NOT NULL -); - -CREATE TABLE tracers ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - tracer_output JSON NOT NULL -); - -CREATE TABLE simulations ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - population_id VARCHAR(255) NOT NULL, - population_type VARCHAR(50) NOT NULL, - policy_id INT NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - simulation_spec_json JSON DEFAULT NULL, - simulation_spec_schema_version INT DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE report_outputs ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - simulation_1_id INT NOT NULL, - simulation_2_id INT DEFAULT NULL, - api_version VARCHAR(10) NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - year VARCHAR(255) DEFAULT '2025', - report_kind VARCHAR(64) DEFAULT NULL, - report_spec_json JSON DEFAULT NULL, - report_spec_schema_version INT DEFAULT NULL, - report_spec_status VARCHAR(32) DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE report_output_runs ( - id CHAR(36) PRIMARY KEY, - report_output_id INT NOT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - report_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - report_cache_version VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - requested_version_override VARCHAR(255) DEFAULT NULL, - resolved_dataset VARCHAR(255) DEFAULT NULL, - resolved_options_hash VARCHAR(255) DEFAULT NULL, - UNIQUE KEY report_output_run_sequence_idx (report_output_id, run_sequence) -); - -CREATE TABLE simulation_runs ( - id CHAR(36) PRIMARY KEY, - simulation_id INT NOT NULL, - report_output_run_id CHAR(36) DEFAULT NULL, - input_position TINYINT DEFAULT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - simulation_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - UNIQUE KEY simulation_run_sequence_idx (simulation_id, run_sequence) -); - -CREATE TABLE legacy_report_output_aliases ( - legacy_report_output_id INT PRIMARY KEY, - canonical_report_output_id INT NOT NULL -); diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 74f530312..9f2c428a5 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -5,8 +5,6 @@ import httpx import pytest -pytest_plugins = ("tests.integration.stage7_mysql",) - INTEGRATION_TIMEOUT_SECONDS = float( os.environ.get("STAGING_API_TEST_TIMEOUT_SECONDS", "900") ) diff --git a/tests/integration/stage7_mysql.py b/tests/integration/stage7_mysql.py deleted file mode 100644 index 1561e792d..000000000 --- a/tests/integration/stage7_mysql.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Shared fixtures for destructive Stage 7 tests on disposable MySQL only.""" - -from collections.abc import Iterator - -from alembic.config import Config -import pytest -from sqlalchemy import Engine, create_engine, inspect - -from policyengine_api.constants import REPO -from policyengine_api.data.v1_models import V1Base -from policyengine_api.scripts.stage7_database import assert_safe_toy_database_url - - -def stage7_database_url() -> str | None: - import os - - return os.environ.get("STAGE7_TOY_DATABASE_URL") - - -def alembic_config(database_url: str) -> Config: - config = Config(str(REPO / "alembic.ini")) - config.set_main_option("sqlalchemy.url", database_url) - return config - - -def reset_toy_database(engine: Engine, database_url: str) -> None: - assert_safe_toy_database_url(database_url) - with engine.begin() as connection: - connection.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 0") - for table_name in reversed(V1Base.metadata.sorted_tables): - connection.exec_driver_sql(f"DROP TABLE IF EXISTS `{table_name.name}`") - connection.exec_driver_sql("DROP TABLE IF EXISTS alembic_version") - connection.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 1") - - -def create_pre_alembic_schema(engine: Engine) -> None: - """Create the existing schema from its independent legacy SQL source.""" - - source = (REPO / "tests/fixtures/stage7_pre_alembic_schema.sql").read_text( - encoding="utf-8" - ) - statements = source.split(";") - with engine.begin() as connection: - for statement in statements: - if statement.strip(): - connection.exec_driver_sql(statement.strip().removesuffix(";")) - - -def schema_signature(engine: Engine) -> dict: - inspector = inspect(engine) - - def normalized(value): - if isinstance(value, dict): - return {key: normalized(item) for key, item in sorted(value.items())} - if isinstance(value, (list, tuple)): - return [normalized(item) for item in value] - if value is None or isinstance(value, (bool, int, float, str)): - return value - return str(value) - - return { - table_name: normalized( - { - "columns": inspector.get_columns(table_name), - "indexes": inspector.get_indexes(table_name), - "pk": inspector.get_pk_constraint(table_name), - "unique": inspector.get_unique_constraints(table_name), - } - ) - for table_name in inspector.get_table_names() - } - - -@pytest.fixture -def stage7_mysql() -> Iterator[tuple[str, Engine]]: - database_url = stage7_database_url() - if database_url is None: - pytest.skip("STAGE7_TOY_DATABASE_URL is required for the MySQL probe") - assert_safe_toy_database_url(database_url) - engine = create_engine(database_url) - reset_toy_database(engine, database_url) - try: - yield database_url, engine - finally: - reset_toy_database(engine, database_url) - engine.dispose() diff --git a/tests/integration/test_stage7_existing_schema.py b/tests/integration/test_stage7_existing_schema.py index 77a264039..ba4260a0d 100644 --- a/tests/integration/test_stage7_existing_schema.py +++ b/tests/integration/test_stage7_existing_schema.py @@ -1,62 +1,28 @@ -"""Qualification of an existing pre-Alembic v1 schema.""" +"""Optional read-only metadata comparison for an existing v1 schema.""" import os -from alembic import command +from alembic.autogenerate import compare_metadata +from alembic.migration import MigrationContext import pytest -from sqlalchemy import inspect, text +from sqlalchemy import create_engine -from policyengine_api.scripts.qualify_stage7_toy import compare_stage7_schema -from tests.integration.stage7_mysql import ( - alembic_config, - create_pre_alembic_schema, - reset_toy_database, - schema_signature, -) - - -def test_fresh_upgrade_has_the_same_schema_signature_as_pre_alembic_v1(stage7_mysql): - database_url, engine = stage7_mysql - create_pre_alembic_schema(engine) - pre_alembic = schema_signature(engine) - - reset_toy_database(engine, database_url) - command.upgrade(alembic_config(database_url), "head") - fresh_upgrade = schema_signature(engine) - fresh_upgrade.pop("alembic_version") +from policyengine_api.data.v1_models import V1Base - assert fresh_upgrade == pre_alembic +def compare_existing_schema(database_url: str) -> list: + """Return metadata drift without mutating the target database.""" -def test_existing_schema_compares_read_only_and_stamps_without_data_loss( - stage7_mysql, -): - database_url, engine = stage7_mysql - create_pre_alembic_schema(engine) - with engine.begin() as connection: - connection.execute( - text( - "INSERT INTO policy " - "(id, country_id, label, api_version, policy_json, policy_hash) " - "VALUES (901, 'us', 'sentinel', 'legacy', '{}', 'sentinel')" + engine = create_engine(database_url) + try: + with engine.connect() as connection: + context = MigrationContext.configure( + connection, + opts={"compare_type": True, "compare_server_default": True}, ) - ) - - before_comparison = schema_signature(engine) - assert "alembic_version" not in before_comparison - assert compare_stage7_schema(database_url) == [] - assert schema_signature(engine) == before_comparison - - config = alembic_config(database_url) - command.stamp(config, "head") - command.check(config) - - with engine.connect() as connection: - assert ( - connection.scalar(text("SELECT label FROM policy WHERE id = 901")) - == "sentinel" - ) - assert inspect(engine).get_table_names().count("alembic_version") == 1 + return compare_metadata(context, V1Base.metadata) + finally: + engine.dispose() @pytest.mark.skipif( @@ -65,4 +31,4 @@ def test_existing_schema_compares_read_only_and_stamps_without_data_loss( ) def test_live_existing_schema_matches_metadata_without_mutation(): database_url = os.environ["STAGE7_EXISTING_DATABASE_URL"] - assert compare_stage7_schema(database_url) == [] + assert compare_existing_schema(database_url) == [] diff --git a/tests/integration/test_stage7_mysql_parity.py b/tests/integration/test_stage7_mysql_parity.py deleted file mode 100644 index 0bf16a290..000000000 --- a/tests/integration/test_stage7_mysql_parity.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Behavioral parity between legacy SQL and typed DAO access on MySQL.""" - -import json -from datetime import datetime - -from alembic import command -from sqlalchemy import text - -from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_daos import V1UnitOfWork -from policyengine_api.services.reform_impacts_service import ReformImpactsService -from tests.integration.stage7_mysql import alembic_config - - -def test_typed_daos_preserve_legacy_mapping_shapes(stage7_mysql): - database_url, engine = stage7_mysql - command.upgrade(alembic_config(database_url), "head") - unit_of_work = V1UnitOfWork(SessionManager(engine)) - - with unit_of_work.transaction() as repositories: - policy_id = repositories.policies.create("us", None, {"x": 1}, "parity", "v1") - household_id = repositories.households.create( - "us", None, {"people": {}}, "household-parity", "v1" - ) - user_id = repositories.users.create_profile( - "auth0|parity", None, "us", 123456789 - ) - - with engine.connect() as connection: - legacy_policy = dict( - connection.execute( - text( - "SELECT * FROM policy " - "WHERE id = :policy_id AND country_id = :country_id" - ), - {"policy_id": policy_id, "country_id": "us"}, - ) - .mappings() - .one() - ) - legacy_household = dict( - connection.execute( - text( - "SELECT * FROM household " - "WHERE id = :household_id AND country_id = :country_id" - ), - {"household_id": household_id, "country_id": "us"}, - ) - .mappings() - .one() - ) - legacy_user = dict( - connection.execute( - text("SELECT * FROM user_profiles WHERE user_id = :user_id"), - {"user_id": user_id}, - ) - .mappings() - .one() - ) - - legacy_policy["policy_json"] = json.loads(legacy_policy["policy_json"]) - legacy_household["household_json"] = json.loads(legacy_household["household_json"]) - with unit_of_work.read() as repositories: - assert legacy_policy == repositories.policies.get("us", policy_id) - assert legacy_household == repositories.households.get("us", household_id) - assert legacy_user == repositories.users.get_profile(user_id=user_id) - - -def test_typed_daos_read_rows_written_by_legacy_sql(stage7_mysql): - database_url, engine = stage7_mysql - command.upgrade(alembic_config(database_url), "head") - unit_of_work = V1UnitOfWork(SessionManager(engine)) - with engine.begin() as connection: - connection.execute( - text( - "INSERT INTO household " - "(country_id, label, api_version, household_json, household_hash) " - "VALUES (:country_id, :label, :api_version, :household_json, " - ":household_hash)" - ), - { - "country_id": "us", - "label": "legacy", - "api_version": "v1", - "household_json": '{"legacy": true}', - "household_hash": "legacy-row", - }, - ) - row = ( - connection.execute( - text("SELECT * FROM household WHERE household_hash = :household_hash"), - {"household_hash": "legacy-row"}, - ) - .mappings() - .one() - ) - - legacy_shape = dict(row) - legacy_shape["household_json"] = json.loads(legacy_shape["household_json"]) - with unit_of_work.read() as repositories: - assert repositories.households.get("us", row["id"]) == legacy_shape - - -def test_reform_impact_service_stores_mysql_json_objects(stage7_mysql): - database_url, engine = stage7_mysql - command.upgrade(alembic_config(database_url), "head") - unit_of_work = V1UnitOfWork(SessionManager(engine)) - service = ReformImpactsService(unit_of_work=unit_of_work) - - service.set_reform_impact( - country_id="us", - policy_id=2, - baseline_policy_id=1, - region="us", - dataset="default", - time_period="2026", - options={"scope": "test"}, - options_hash="native-json", - status="computing", - api_version="v1", - reform_impact_json={}, - start_time=datetime(2026, 1, 1), - execution_id="native-json-job", - ) - service.set_complete_reform_impact( - country_id="us", - reform_policy_id=2, - baseline_policy_id=1, - region="us", - dataset="default", - time_period="2026", - options_hash="native-json", - reform_impact_json={"result": {"value": 1}}, - execution_id="native-json-job", - ) - - with engine.connect() as connection: - json_types = connection.execute( - text( - "SELECT JSON_TYPE(options_json), JSON_TYPE(reform_impact_json) " - "FROM reform_impact WHERE execution_id = :execution_id" - ), - {"execution_id": "native-json-job"}, - ).one() - assert tuple(json_types) == ("OBJECT", "OBJECT") - - stored = service.get_all_reform_impacts( - "us", - 2, - 1, - "us", - "default", - "2026", - "native-json", - "v1", - )[0] - assert stored["options_json"] == {"scope": "test"} - assert stored["reform_impact_json"] == {"result": {"value": 1}} diff --git a/tests/integration/test_stage7_mysql_runtime.py b/tests/integration/test_stage7_mysql_runtime.py deleted file mode 100644 index 0890eb5aa..000000000 --- a/tests/integration/test_stage7_mysql_runtime.py +++ /dev/null @@ -1,164 +0,0 @@ -"""Route, Cloud SQL connector, and startup behavior on disposable MySQL.""" - -import importlib -import json -import os - -from alembic import command -from flask import Flask -import pymysql -from sqlalchemy import text - -from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_daos import V1UnitOfWork -from policyengine_api.services.household_service import HouseholdService -from policyengine_api.services.policy_service import PolicyService -from policyengine_api.services.simulation_service import SimulationService -from policyengine_api.services.user_service import UserService -from tests.integration.stage7_mysql import alembic_config, schema_signature - - -def test_public_policy_household_and_user_routes_use_mysql(stage7_mysql, monkeypatch): - database_url, engine = stage7_mysql - command.upgrade(alembic_config(database_url), "head") - sessions = SessionManager(engine) - unit_of_work = V1UnitOfWork(sessions) - - monkeypatch.setenv("FLASK_DEBUG", "1") - from policyengine_api.routes import household_routes, policy_routes - from policyengine_api.routes import simulation_routes - from policyengine_api.routes import user_profile_routes - - monkeypatch.setattr( - policy_routes, - "policy_service", - PolicyService(unit_of_work=unit_of_work), - ) - monkeypatch.setattr( - household_routes, - "household_service", - HouseholdService(unit_of_work=unit_of_work), - ) - monkeypatch.setattr( - user_profile_routes, - "user_service", - UserService(unit_of_work=unit_of_work), - ) - monkeypatch.setattr( - simulation_routes, - "simulation_service", - SimulationService(unit_of_work=unit_of_work), - ) - app = Flask(__name__) - app.register_blueprint(policy_routes.policy_bp) - app.register_blueprint(household_routes.household_bp) - app.register_blueprint(simulation_routes.simulation_bp) - app.register_blueprint(user_profile_routes.user_profile_bp) - client = app.test_client() - - policy = client.post("/us/policy", json={"label": "Route", "data": {}}) - assert policy.status_code == 201 - policy_id = policy.get_json()["result"]["policy_id"] - policy_result = json.loads(client.get(f"/us/policy/{policy_id}").data)["result"] - assert policy_result["id"] == policy_id - - household = client.post( - "/us/household", json={"label": "Route", "data": {"people": {}}} - ) - assert household.status_code == 201 - household_id = household.get_json()["result"]["household_id"] - assert ( - client.get(f"/us/household/{household_id}").get_json()["result"]["id"] - == household_id - ) - - user = client.post( - "/us/user-profile", - json={"auth0_id": "auth0|route", "username": None, "user_since": 1}, - ) - assert user.status_code == 201 - user_id = user.get_json()["result"]["user_id"] - assert client.get(f"/us/user-profile?user_id={user_id}").status_code == 200 - - simulation = client.post( - "/us/simulation", - json={ - "population_id": str(household_id), - "population_type": "household", - "policy_id": policy_id, - }, - ) - assert simulation.status_code == 201 - simulation_id = simulation.get_json()["result"]["id"] - output = {"result": "ok"} - updated_simulation = client.patch( - "/us/simulation", - json={"id": simulation_id, "status": "complete", "output": output}, - ) - assert updated_simulation.status_code == 200 - assert json.loads(updated_simulation.get_json()["result"]["output"]) == output - fetched_simulation = client.get(f"/us/simulation/{simulation_id}") - assert fetched_simulation.status_code == 200 - assert json.loads(fetched_simulation.get_json()["result"]["output"]) == output - - -def test_cloud_sql_connector_pool_drives_daos_and_startup_emits_no_ddl( - stage7_mysql, monkeypatch -): - database_url, engine = stage7_mysql - command.upgrade(alembic_config(database_url), "head") - - os.environ.setdefault("FLASK_DEBUG", "1") - from policyengine_api.data import data as data_module - - url = engine.url - connector_calls = [] - - class LocalConnector: - def __init__(self, **_kwargs): - pass - - def connect(self, **kwargs): - connector_calls.append(kwargs) - return pymysql.connect( - host=url.host, - port=url.port, - user=url.username, - password=url.password, - database=url.database, - ) - - def close(self): - pass - - monkeypatch.setattr(data_module, "Connector", LocalConnector) - monkeypatch.setenv("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", "toy:local:stage7") - monkeypatch.setenv("POLICYENGINE_DB_USER", str(url.username)) - monkeypatch.setenv("POLICYENGINE_DB_NAME", str(url.database)) - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", str(url.password)) - - remote_database = data_module.PolicyEngineDatabase(local=False, initialize=False) - assert connector_calls == [] - with remote_database.pool.connect() as first: - with remote_database.pool.connect() as second: - first_id = first.scalar(text("SELECT CONNECTION_ID()")) - second_id = second.scalar(text("SELECT CONNECTION_ID()")) - assert first_id != second_id - assert len(connector_calls) == 2 - - unit_of_work = V1UnitOfWork(SessionManager(remote_database.pool)) - with unit_of_work.transaction() as repositories: - policy_id = repositories.policies.create( - "us", "Connector", {}, "connector-path", "v1" - ) - with unit_of_work.read() as repositories: - assert repositories.policies.get("us", policy_id)["label"] == "Connector" - assert connector_calls[0]["instance_connection_string"] == "toy:local:stage7" - - monkeypatch.setattr(data_module, "database", remote_database) - before_startup = schema_signature(engine) - api_module = importlib.import_module("policyengine_api.api") - importlib.reload(api_module) - assert api_module.app.test_client().get("/liveness-check").status_code == 200 - assert schema_signature(engine) == before_startup - remote_database._close_pool() diff --git a/tests/integration/test_stage7_toy_database.py b/tests/integration/test_stage7_toy_database.py deleted file mode 100644 index c0e3f5964..000000000 --- a/tests/integration/test_stage7_toy_database.py +++ /dev/null @@ -1,47 +0,0 @@ -"""MySQL qualification for the disposable Stage 7 database.""" - -from alembic import command -from sqlalchemy import inspect - -from policyengine_api.data.v1_models import V1Base -from policyengine_api.scripts.qualify_stage7_toy import ( - compare_stage7_schema, - qualify_stage7_toy, -) -from tests.integration.stage7_mysql import alembic_config - - -def test_mysql_toy_database_upgrades_and_exercises_every_dao_domain(stage7_mysql): - database_url, _ = stage7_mysql - result = qualify_stage7_toy(database_url) - - assert result == { - "alembic_head": True, - "policy": True, - "household": True, - "computed_household": True, - "user": True, - "user_policy": True, - "economy": True, - "simulation": True, - "report": True, - "report_alias": True, - "analysis": True, - "tracer": True, - "reform_impact": True, - } - assert compare_stage7_schema(database_url) == [] - - -def test_mysql_toy_database_downgrades_and_reupgrades_cleanly(stage7_mysql): - database_url, engine = stage7_mysql - config = alembic_config(database_url) - qualify_stage7_toy(database_url) - - command.downgrade(config, "base") - remaining_tables = set(inspect(engine).get_table_names()) - assert not set(V1Base.metadata.tables) & remaining_tables - - command.upgrade(config, "head") - command.check(config) - assert compare_stage7_schema(database_url) == [] diff --git a/tests/unit/test_stage7_toy_qualification.py b/tests/unit/test_stage7_toy_qualification.py deleted file mode 100644 index 8392faa3e..000000000 --- a/tests/unit/test_stage7_toy_qualification.py +++ /dev/null @@ -1,56 +0,0 @@ -from pathlib import Path - -from policyengine_api.scripts.stage7_database import assert_safe_toy_database_url -import pytest - - -def test_legacy_daos_have_been_removed(): - package = Path(__file__).parents[2] / "policyengine_api" - sources = "\n".join( - path.read_text(encoding="utf-8") for path in package.rglob("*.py") - ) - assert "LegacyPolicyDAO" not in sources - assert "LegacyHouseholdDAO" not in sources - assert "LegacyUserDAO" not in sources - assert "LegacySimulationDAO" not in sources - assert "LegacyReportDAO" not in sources - - -@pytest.mark.parametrize( - "database_url", - [ - "mysql+pymysql://toy:toy@127.0.0.1:3307/policyengine_stage7_toy", - "mysql+pymysql://toy:toy@localhost:3307/custom_toy", - ], -) -def test_toy_database_safety_guard_accepts_only_local_mysql_toy_databases( - database_url: str, -): - assert_safe_toy_database_url(database_url) - - -@pytest.mark.parametrize( - "database_url", - [ - "mysql+pymysql://toy:toy@prod.example.com/policyengine_stage7_toy", - "mysql+pymysql://toy:toy@127.0.0.1/policyengine", - "postgresql://toy:toy@127.0.0.1/policyengine_stage7_toy", - "sqlite+pysqlite:///policyengine_stage7_toy.db", - ], -) -def test_toy_database_safety_guard_rejects_unsafe_targets(database_url: str): - with pytest.raises(ValueError, match="local MySQL.*_toy"): - assert_safe_toy_database_url(database_url) - - -def test_stage7_toy_database_has_local_scaffold_and_test_targets(): - repo = Path(__file__).parents[2] - compose = (repo / "compose.stage7-toy.yml").read_text(encoding="utf-8") - makefile = (repo / "Makefile").read_text(encoding="utf-8") - - assert "mysql:8.0" in compose - assert "policyengine_stage7_toy" in compose - assert "healthcheck:" in compose - assert "stage7-toy-up:" in makefile - assert "stage7-toy-test:" in makefile - assert "stage7-toy-down:" in makefile From 1914185440d23c5e02642cc0d21632886f50a050 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 18:56:13 +0300 Subject: [PATCH 43/89] fix: extend App Engine startup budgets --- .github/workflows/push.yml | 4 ++++ gcp/policyengine_api/app.yaml | 4 ++-- gcp/policyengine_api/start.sh | 2 +- tests/unit/test_cloud_run_deploy_scripts.py | 18 ++++++++++++++++++ 4 files changed, 25 insertions(+), 3 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index f0696abd5..d913dc88d 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -203,6 +203,8 @@ jobs: APP_ENGINE_VERSION: ${{ steps.version.outputs.version }} - name: Wait for staging version health run: bash .github/scripts/health_check.sh "${{ steps.version_url.outputs.url }}/readiness-check" + env: + HEALTH_CHECK_TIMEOUT_SECONDS: "1800" deploy-cloud-run-staging: name: Deploy staging Cloud Run candidate @@ -511,6 +513,8 @@ jobs: APP_ENGINE_VERSION: ${{ steps.version.outputs.version }} - name: Wait for production version health run: bash .github/scripts/health_check.sh "${{ steps.version_url.outputs.url }}/readiness-check" + env: + HEALTH_CHECK_TIMEOUT_SECONDS: "1800" promote-production: name: Promote production App Engine candidate diff --git a/gcp/policyengine_api/app.yaml b/gcp/policyengine_api/app.yaml index 67e1af80a..c3512b479 100644 --- a/gcp/policyengine_api/app.yaml +++ b/gcp/policyengine_api/app.yaml @@ -16,7 +16,7 @@ liveness_check: timeout_sec: 30 # Allow 30 seconds for a response failure_threshold: 5 success_threshold: 2 - initial_delay_sec: 60 # Don't check for first 60 seconds to allow full boot + initial_delay_sec: 1800 # Allow non-preloaded workers to finish importing runtime_config: operating_system: "ubuntu22" runtime_version: "22" @@ -26,4 +26,4 @@ readiness_check: timeout_sec: 30 failure_threshold: 5 success_threshold: 2 - app_start_timeout_sec: 900 + app_start_timeout_sec: 1800 diff --git a/gcp/policyengine_api/start.sh b/gcp/policyengine_api/start.sh index 96189837b..3fee8e4fe 100644 --- a/gcp/policyengine_api/start.sh +++ b/gcp/policyengine_api/start.sh @@ -19,7 +19,7 @@ until redis-cli -h "$CACHE_REDIS_HOST" -p "$CACHE_REDIS_PORT" ping >/dev/null 2> done # Start the API -gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 & +gunicorn -b :"$PORT" policyengine_api.api --timeout 900 --workers 5 & # Keep the script running and handle shutdown gracefully trap "pkill -P $$; exit 1" INT TERM diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index b51b861af..e027ba947 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -494,6 +494,24 @@ def test_production_gunicorn_workers_do_not_inherit_database_pools(): assert "--preload" not in commands +def test_app_engine_startup_allows_all_workers_to_finish_booting(): + start_script = (REPO / "gcp/policyengine_api/start.sh").read_text(encoding="utf-8") + app_config = (REPO / "gcp/policyengine_api/app.yaml").read_text(encoding="utf-8") + + assert "--timeout 900" in start_script + assert "--workers 5" in start_script + assert "initial_delay_sec: 1800" in app_config + assert "app_start_timeout_sec: 1800" in app_config + + +def test_app_engine_deploy_health_checks_allow_full_startup_window(): + workflow = _push_workflow() + + for job_name in ("deploy-staging", "deploy-production-candidate"): + job = _workflow_job_block(workflow, job_name) + assert 'HEALTH_CHECK_TIMEOUT_SECONDS: "1800"' in job + + def test_validate_cloud_run_deploy_env_requires_selector_environment_variable(): result = _run_script( ".github/scripts/validate_cloud_run_deploy_env.sh", From 3c184dcc2bb132f9177cb8e91b651772be118c4c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 19:48:07 +0300 Subject: [PATCH 44/89] test: expand Stage 7 ORM coverage --- tests/unit/data/test_alembic_baseline.py | 10 + tests/unit/data/test_local_daos.py | 45 +++++ tests/unit/data/test_ordinary_v1_daos.py | 61 ++++++ tests/unit/data/test_run_daos.py | 101 +++++++++- tests/unit/data/test_v1_daos.py | 22 +++ .../endpoints/test_stage7_orm_endpoints.py | 176 ++++++++++++++++++ .../services/test_reform_impacts_service.py | 160 ++++++++++++++++ 7 files changed, 574 insertions(+), 1 deletion(-) create mode 100644 tests/unit/endpoints/test_stage7_orm_endpoints.py create mode 100644 tests/unit/services/test_reform_impacts_service.py diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index 8be0415a7..f16eccb4f 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -36,3 +36,13 @@ def test_baseline_is_the_single_root_revision(): assert head is not None assert head.down_revision is None + + +def test_baseline_renders_a_complete_mysql_downgrade_without_connecting(): + config, output = _mysql_offline_config() + + command.downgrade(config, "head:base", sql=True) + + rendered_sql = output.getvalue() + for table_name in V1Base.metadata.tables: + assert f"DROP TABLE {table_name}" in rendered_sql diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py index b5dee1e8a..80c45dac4 100644 --- a/tests/unit/data/test_local_daos.py +++ b/tests/unit/data/test_local_daos.py @@ -45,6 +45,51 @@ def test_reform_impact_dao_transitions_by_execution_id(): ] == {"result": 1} +def test_reform_impact_dao_orders_limits_messages_and_handles_missing_jobs(): + uow = _unit_of_work() + with uow.transaction() as repositories: + for day, execution_id in ((1, "old-job"), (2, "new-job")): + repositories.reform_impacts.create( + country_id="us", + reform_policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options_json={}, + options_hash=execution_id, + api_version="1", + reform_impact_json={}, + status="computing", + start_time=datetime(2026, 1, day), + execution_id=execution_id, + ) + + assert repositories.reform_impacts.set_message( + "queued", country_id="us", status="computing" + ) + assert ( + repositories.reform_impacts.set_message("missing", country_id="uk") is False + ) + assert ( + repositories.reform_impacts.fail( + "missing-job", "failed", datetime(2026, 1, 3) + ) + is False + ) + assert ( + repositories.reform_impacts.complete( + "missing-job", {}, datetime(2026, 1, 3) + ) + is False + ) + + with uow.read() as repositories: + recent = repositories.reform_impacts.list_recent(1) + assert [row["execution_id"] for row in recent] == ["new-job"] + assert recent[0]["message"] == "queued" + + def test_tracer_dao_returns_latest_matching_trace(): uow = _unit_of_work() with uow.transaction() as repositories: diff --git a/tests/unit/data/test_ordinary_v1_daos.py b/tests/unit/data/test_ordinary_v1_daos.py index 74dfb8d07..cbf5cf159 100644 --- a/tests/unit/data/test_ordinary_v1_daos.py +++ b/tests/unit/data/test_ordinary_v1_daos.py @@ -68,3 +68,64 @@ def test_policy_search_and_reform_impact_limit_use_typed_statements(): assert [row["label"] for row in repositories.policies.search("us", "Tax")] == [ "Tax reform" ] + + +def test_computed_household_create_and_version_filters(): + uow = _unit_of_work() + with uow.transaction() as repositories: + repositories.computed_households.create( + household_id=1, + policy_id=2, + country_id="us", + api_version="1", + computed_household_json={"value": 1}, + status="complete", + ) + + with uow.read() as repositories: + assert repositories.computed_households.get(1, 2, "us")[ + "computed_household_json" + ] == {"value": 1} + assert repositories.computed_households.get(1, 2, "us", api_version="2") is None + + +def test_economy_and_user_policy_daos_cover_lookup_edge_cases(): + uow = _unit_of_work() + user_policy_values = { + "country_id": "us", + "reform_id": 2, + "reform_label": None, + "baseline_id": 1, + "baseline_label": None, + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": None, + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 1, + "budgetary_impact": None, + "type": None, + } + with uow.transaction() as repositories: + economy_id = repositories.economies.create( + policy_id=2, + country_id="us", + region="us", + time_period="2026", + options_json={"dataset": "default"}, + options_hash="hash", + api_version="1", + economy_json={"result": 1}, + status="complete", + message=None, + ) + user_policy_id = repositories.user_policies.create(**user_policy_values) + assert repositories.user_policies.update(999, {}) is False + + with uow.read() as repositories: + assert repositories.economies.get(economy_id)["economy_json"] == {"result": 1} + assert repositories.economies.get(999) is None + assert repositories.user_policies.get(user_policy_id)["country_id"] == "us" + assert repositories.user_policies.get(999) is None diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py index 21bbf582e..9c1545a96 100644 --- a/tests/unit/data/test_run_daos.py +++ b/tests/unit/data/test_run_daos.py @@ -1,7 +1,9 @@ from datetime import datetime +import pytest + from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork from tests.unit.data.sqlite_schema import create_sqlite_v1_schema @@ -66,3 +68,100 @@ def test_report_dao_round_trips_parent_run_and_alias(): repositories.reports.get_alias(99)["canonical_report_output_id"] == report_id ) + + +def test_simulation_dao_sync_callbacks_cover_create_update_and_missing_rows(): + uow = _unit_of_work() + + def read_synced(session, simulation_id, *, country_id): + return SimulationDAO.get_in_session(session, simulation_id, country_id) + + with uow.transaction() as repositories: + created = repositories.simulations.create_or_get_with_sync( + sync_callback=read_synced, + country_id="us", + api_version="1", + population_id="7", + population_type="household", + policy_id=2, + status="complete", + output={"result": 1}, + ) + reused = repositories.simulations.create_or_get_with_sync( + sync_callback=read_synced, + country_id="us", + api_version="1", + population_id="7", + population_type="household", + policy_id=2, + status="pending", + ) + updated = repositories.simulations.update_with_sync( + created["id"], + "us", + {"error_message": "updated"}, + read_synced, + ) + dual_write = repositories.simulations.ensure_dual_write_state( + created["id"], "us" + ) + + assert reused["id"] == created["id"] + assert updated["error_message"] == "updated" + assert dual_write["latest_successful_run_id"] is not None + assert repositories.simulations.get(created["id"], "uk") is None + assert repositories.simulations.update(999, status="complete") is False + with pytest.raises(ValueError, match="Simulation #999 not found"): + repositories.simulations.update_with_sync( + 999, "us", {"status": "complete"}, read_synced + ) + with pytest.raises(LookupError, match="Simulation 999 does not exist"): + repositories.simulations.create_run( + 999, + run_id="missing-run", + status="pending", + trigger_type="create", + ) + + +def test_report_dao_handles_scoped_lookups_updates_and_existing_aliases(): + uow = _unit_of_work() + with uow.transaction() as repositories: + report_id = repositories.reports.create( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + year="2026", + ) + run = repositories.reports.create_run( + report_id, + run_id="report-run", + status="pending", + trigger_type="create", + ) + + assert repositories.reports.get(report_id, "uk") is None + assert repositories.reports.get_for_update(report_id, "us")["id"] == report_id + assert repositories.reports.get_for_update(report_id, "uk") is None + assert repositories.reports.update(999, status="complete") is False + assert repositories.reports.update(report_id, status="complete") + assert repositories.reports.update_run( + run["id"], status="complete", output={"result": 1} + ) + assert repositories.reports.update_run("missing-run", status="error") is False + repositories.reports.set_alias(99, report_id) + repositories.reports.set_alias(99, report_id + 1) + with pytest.raises(LookupError, match="Report output 999 does not exist"): + repositories.reports.create_run( + 999, + run_id="missing-run", + status="pending", + trigger_type="create", + ) + + with uow.read() as repositories: + assert repositories.reports.get(report_id)["status"] == "complete" + assert repositories.reports.get_run(run["id"])["output"] == {"result": 1} + assert repositories.reports.get_run("missing-run") is None + assert repositories.reports.get_alias(99)["canonical_report_output_id"] == 2 diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index e17baf4fc..022d41b2d 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -69,3 +69,25 @@ def test_user_dao_profile_lookup_precedence(): ] == "auth0|one" ) + + +def test_user_and_household_daos_handle_missing_and_nullable_updates(): + uow = _unit_of_work() + with uow.transaction() as repositories: + user_id = repositories.users.create_profile("auth0|one", "original", "us", 123) + assert repositories.users.get_profile() is None + assert repositories.users.update_profile(999, username="missing") is False + assert ( + repositories.households.update("us", 999, "missing", {}, "missing", "1") + is False + ) + assert repositories.users.update_profile( + user_id, + username=None, + primary_country="uk", + ) + + with uow.read() as repositories: + profile = repositories.users.get_profile(user_id=user_id) + assert profile["username"] == "original" + assert profile["primary_country"] == "uk" diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py new file mode 100644 index 000000000..febac16a0 --- /dev/null +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -0,0 +1,176 @@ +import json +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +from flask import Flask + +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.endpoints.household import get_household_under_policy +from policyengine_api.endpoints.policy import ( + get_user_policy, + set_user_policy, + update_user_policy, +) +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema + + +def _unit_of_work() -> V1UnitOfWork: + manager = build_sqlite_session_manager() + create_sqlite_v1_schema(manager) + return V1UnitOfWork(manager) + + +def _repositories_unit_of_work(repositories): + @contextmanager + def boundary(): + yield repositories + + return SimpleNamespace(read=boundary, transaction=boundary) + + +@pytest.mark.parametrize( + "stored_result", + [ + {"people": {"you": {"net_income": {"2026": 42}}}}, + json.dumps({"people": {"you": {"net_income": {"2026": 42}}}}), + ], +) +def test_household_under_policy_returns_cached_json_objects_and_legacy_strings( + stored_result, +): + computed_households = Mock() + computed_households.get.return_value = { + "household_id": 1, + "policy_id": 2, + "country_id": "us", + "api_version": "1", + "computed_household_json": stored_result, + "status": "complete", + } + local_uow = _repositories_unit_of_work( + SimpleNamespace(computed_households=computed_households) + ) + + with patch( + "policyengine_api.endpoints.household.runtime_v1_unit_of_work", + return_value=local_uow, + ) as runtime_uow: + response = get_household_under_policy("us", "1", "2") + + assert response["result"] == {"people": {"you": {"net_income": {"2026": 42}}}} + runtime_uow.assert_called_once_with(local=True) + + +def test_household_under_policy_calculates_and_caches_json_as_an_object(): + computed_households = Mock() + computed_households.get.return_value = None + local_uow = _repositories_unit_of_work( + SimpleNamespace(computed_households=computed_households) + ) + remote_uow = _repositories_unit_of_work( + SimpleNamespace( + households=SimpleNamespace( + get=Mock( + return_value={ + "id": 1, + "country_id": "us", + "household_json": {"people": {"you": {}}}, + } + ) + ), + policies=SimpleNamespace( + get=Mock( + return_value={ + "id": 2, + "country_id": "us", + "policy_json": {"gov.example.parameter": 1}, + } + ) + ), + ) + ) + calculated = {"people": {"you": {"net_income": {"2026": 42}}}} + country = SimpleNamespace(calculate=Mock(return_value=calculated)) + + def select_uow(*, local=False): + return local_uow if local else remote_uow + + with ( + patch( + "policyengine_api.endpoints.household.runtime_v1_unit_of_work", + side_effect=select_uow, + ), + patch( + "policyengine_api.endpoints.household.add_yearly_variables", + side_effect=lambda household, _: household, + ), + patch( + "policyengine_api.endpoints.household.drop_deprecated_inputs", + side_effect=lambda household: SimpleNamespace( + household=household, + warnings=[], + ), + ), + patch( + "policyengine_api.endpoints.household.get_invalid_inputs_response", + return_value=None, + ), + patch( + "policyengine_api.endpoints.household.get_countries", + return_value={"us": country}, + ), + ): + response = get_household_under_policy("us", "1", "2") + + assert response["result"] == calculated + country.calculate.assert_called_once_with( + {"people": {"you": {}}}, + {"gov.example.parameter": 1}, + "1", + "2", + ) + assert ( + computed_households.upsert.call_args.kwargs["computed_household_json"] + is calculated + ) + + +def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): + app = Flask(__name__) + uow = _unit_of_work() + payload = { + "reform_label": "Reform", + "reform_id": 2, + "baseline_label": "Current law", + "baseline_id": 1, + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": "default", + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 1, + "budgetary_impact": None, + "type": None, + } + + with patch( + "policyengine_api.endpoints.policy.runtime_v1_unit_of_work", + return_value=uow, + ): + with app.test_request_context(json=payload): + created = set_user_policy("us") + listed = get_user_policy("us", "auth0|one") + with app.test_request_context(json={"id": 1, "reform_label": "Updated"}): + updated = update_user_policy("us") + + assert created.status_code == 201 + assert created.get_json()["result"]["dataset"] == "default" + assert listed["result"][0]["reform_label"] == "Reform" + assert updated.status_code == 200 + with uow.read() as repositories: + assert repositories.user_policies.get(1)["reform_label"] == "Updated" diff --git a/tests/unit/services/test_reform_impacts_service.py b/tests/unit/services/test_reform_impacts_service.py new file mode 100644 index 000000000..98e639b92 --- /dev/null +++ b/tests/unit/services/test_reform_impacts_service.py @@ -0,0 +1,160 @@ +from datetime import datetime +from unittest.mock import ANY, Mock + +from policyengine_api.data.orm import build_sqlite_session_manager +from policyengine_api.data.v1_daos import ReformImpactDAO, V1UnitOfWork +from policyengine_api.services import reform_impacts_service as service_module +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from tests.unit.data.sqlite_schema import create_sqlite_v1_schema + + +def _unit_of_work() -> V1UnitOfWork: + manager = build_sqlite_session_manager() + create_sqlite_v1_schema(manager) + return V1UnitOfWork(manager) + + +def _create_impact( + service: ReformImpactsService, + *, + execution_id: str, + options_hash: str, + day: int, +) -> int: + return service.set_reform_impact( + country_id="us", + policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options={"hash": options_hash}, + options_hash=options_hash, + status="computing", + api_version="1", + reform_impact_json={}, + start_time=datetime(2026, 1, day), + execution_id=execution_id, + ) + + +def test_reform_impact_service_round_trips_queries_and_transitions(): + service = ReformImpactsService(unit_of_work=_unit_of_work()) + exact_id = _create_impact( + service, + execution_id="exact-job", + options_hash="hash-exact", + day=1, + ) + compatible_id = _create_impact( + service, + execution_id="compatible-job", + options_hash="hash-compatible", + day=2, + ) + + exact = service.get_all_reform_impacts( + "us", 2, 1, "us", "default", "2026", "hash-exact", "1" + ) + compatible = service.get_all_reform_impacts_by_options_hash_prefix( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + "hash-%", + "1", + ) + + assert [row["reform_impact_id"] for row in exact] == [exact_id] + assert [row["reform_impact_id"] for row in compatible] == [ + exact_id, + compatible_id, + ] + assert service.set_complete_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + {"result": 1}, + "exact-job", + ) + assert service.set_error_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-compatible", + "failed", + "compatible-job", + ) + + with service.unit_of_work.read() as repositories: + completed = repositories.reform_impacts.find(execution_id="exact-job") + failed = repositories.reform_impacts.find(execution_id="compatible-job") + assert completed["status"] == "ok" + assert completed["reform_impact_json"] == {"result": 1} + assert failed["status"] == "error" + assert failed["message"] == "failed" + + +def test_reform_impact_service_deletes_only_matching_computing_rows(): + service = ReformImpactsService(unit_of_work=_unit_of_work()) + _create_impact( + service, + execution_id="delete-job", + options_hash="delete-hash", + day=1, + ) + retained_id = _create_impact( + service, + execution_id="retain-job", + options_hash="retain-hash", + day=2, + ) + + service.delete_reform_impact("us", 2, 1, "us", "default", "2026", "delete-hash") + + assert ( + service.get_all_reform_impacts( + "us", 2, 1, "us", "default", "2026", "delete-hash", "1" + ) + == [] + ) + retained = service.get_all_reform_impacts( + "us", 2, 1, "us", "default", "2026", "retain-hash", "1" + ) + assert [row["reform_impact_id"] for row in retained] == [retained_id] + + +def test_reform_impact_service_supports_injected_repository(): + impacts = Mock(spec=ReformImpactDAO) + impacts.fail.return_value = False + service = ReformImpactsService(impacts) + + assert ( + service.set_error_reform_impact( + "us", 2, 1, "us", "default", "2026", "hash", "missing", "job" + ) + is False + ) + impacts.fail.assert_called_once_with("job", "missing", ANY) + + +def test_reform_impact_service_builds_default_unit_of_work_once(monkeypatch): + manager = build_sqlite_session_manager() + build_manager = Mock(return_value=manager) + monkeypatch.setattr(service_module, "build_v1_session_manager", build_manager) + service = ReformImpactsService() + + first = service.unit_of_work + + assert service.unit_of_work is first + build_manager.assert_called_once_with(local=True) From 833936eaf183413b47497f431f955802f830a46d Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 20:54:58 +0300 Subject: [PATCH 45/89] refactor: keep database pool settings in code --- policyengine_api/data/data.py | 22 +++++++++++---------- tests/unit/data/test_sqlalchemy_v2.py | 28 ++++++++++++++++++++++++++- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py index 4cb8a0d6e..8edc7049a 100644 --- a/policyengine_api/data/data.py +++ b/policyengine_api/data/data.py @@ -18,6 +18,11 @@ ) DEFAULT_REMOTE_DB_USER = "policyengine" DEFAULT_REMOTE_DB_NAME = "policyengine" +CLOUD_SQL_IP_TYPE = IPTypes.PUBLIC +DATABASE_POOL_RECYCLE_SECONDS = 1800 +DATABASE_POOL_SIZE = 5 +DATABASE_POOL_MAX_OVERFLOW = 2 +DATABASE_POOL_TIMEOUT_SECONDS = 30 def get_remote_database_config() -> dict[str, str]: @@ -129,13 +134,10 @@ def __init__( def _create_pool(self): db_config = get_remote_database_config() - ip_type = ( - IPTypes.PRIVATE - if os.environ.get("POLICYENGINE_DB_PRIVATE_IP", "").lower() - in {"1", "true", "yes"} - else IPTypes.PUBLIC + self.connector = Connector( + ip_type=CLOUD_SQL_IP_TYPE, + refresh_strategy="LAZY", ) - self.connector = Connector(ip_type=ip_type, refresh_strategy="LAZY") db_pass = os.environ["POLICYENGINE_DB_PASSWORD"] if db_pass == ".dbpw": with open(".dbpw") as f: @@ -154,10 +156,10 @@ def getconn(): "mysql+pymysql://", creator=getconn, pool_pre_ping=True, - pool_recycle=int(os.environ.get("POLICYENGINE_DB_POOL_RECYCLE", "1800")), - pool_size=int(os.environ.get("POLICYENGINE_DB_POOL_SIZE", "5")), - max_overflow=int(os.environ.get("POLICYENGINE_DB_MAX_OVERFLOW", "2")), - pool_timeout=int(os.environ.get("POLICYENGINE_DB_POOL_TIMEOUT", "30")), + pool_recycle=DATABASE_POOL_RECYCLE_SECONDS, + pool_size=DATABASE_POOL_SIZE, + max_overflow=DATABASE_POOL_MAX_OVERFLOW, + pool_timeout=DATABASE_POOL_TIMEOUT_SECONDS, ) def close(self) -> None: diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 6fc0c505c..da29569b0 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -202,7 +202,12 @@ def fake_create_engine(url, **kwargs): return "fake-engine" fake_connector = FakeConnector(refresh_strategy="LAZY") - monkeypatch.setattr(data_module, "Connector", lambda **_: fake_connector) + + def fake_connector_factory(**kwargs): + fake_connector.options = kwargs + return fake_connector + + monkeypatch.setattr(data_module, "Connector", fake_connector_factory) monkeypatch.setattr(data_module.sqlalchemy, "create_engine", fake_create_engine) return fake_connector, connector_calls, engine_calls @@ -251,6 +256,27 @@ def test_create_pool_uses_remote_database_config(self, monkeypatch): assert engine_calls[0][1]["max_overflow"] == 2 assert engine_calls[0][1]["pool_timeout"] == 30 + def test_create_pool_settings_are_owned_by_the_application(self, monkeypatch): + fake_connector, _, engine_calls = self._stub_remote_pool(monkeypatch) + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "test-password") + monkeypatch.setenv("POLICYENGINE_DB_PRIVATE_IP", "true") + monkeypatch.setenv("POLICYENGINE_DB_POOL_RECYCLE", "not-used") + monkeypatch.setenv("POLICYENGINE_DB_POOL_SIZE", "not-used") + monkeypatch.setenv("POLICYENGINE_DB_MAX_OVERFLOW", "not-used") + monkeypatch.setenv("POLICYENGINE_DB_POOL_TIMEOUT", "not-used") + + db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) + db._create_pool() + + assert fake_connector.options == { + "ip_type": data_module.IPTypes.PUBLIC, + "refresh_strategy": "LAZY", + } + assert engine_calls[0][1]["pool_recycle"] == 1800 + assert engine_calls[0][1]["pool_size"] == 5 + assert engine_calls[0][1]["max_overflow"] == 2 + assert engine_calls[0][1]["pool_timeout"] == 30 + def test_create_pool_reads_dot_dbpw_file(self, monkeypatch, tmp_path): _, connector_calls, engine_calls = self._stub_remote_pool(monkeypatch) monkeypatch.chdir(tmp_path) From 6d5624f4491a231da2918ccade835fa6dd454f6b Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 21:16:06 +0300 Subject: [PATCH 46/89] refactor: rename V1 DAO container --- policyengine_api/country.py | 4 +- policyengine_api/data/v1_daos.py | 12 +-- .../endpoints/economy/reform_impact.py | 4 +- policyengine_api/endpoints/household.py | 14 +-- policyengine_api/endpoints/policy.py | 20 ++--- policyengine_api/endpoints/simulation.py | 4 +- .../services/ai_analysis_service.py | 4 +- .../services/household_service.py | 4 +- policyengine_api/services/policy_service.py | 4 +- .../services/reform_impacts_service.py | 4 +- .../services/report_output_alias_service.py | 16 ++-- .../services/report_output_service.py | 46 +++++----- .../services/report_run_service.py | 16 ++-- .../services/report_spec_service.py | 52 +++++------- .../services/simulation_run_service.py | 16 ++-- .../services/simulation_service.py | 22 +++-- .../services/simulation_spec_service.py | 10 +-- .../services/tracer_analysis_service.py | 4 +- policyengine_api/services/user_service.py | 4 +- tests/unit/data/test_local_daos.py | 58 ++++++------- tests/unit/data/test_ordinary_v1_daos.py | 66 +++++++------- tests/unit/data/test_run_daos.py | 85 +++++++++---------- tests/unit/data/test_stage7_no_direct_sql.py | 2 +- tests/unit/data/test_v1_daos.py | 67 ++++++--------- tests/unit/data/test_v1_unit_of_work.py | 42 ++++----- .../endpoints/test_stage7_orm_endpoints.py | 14 +-- .../services/test_reform_impacts_service.py | 6 +- 27 files changed, 283 insertions(+), 317 deletions(-) diff --git a/policyengine_api/country.py b/policyengine_api/country.py index f202883be..602af096f 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -433,8 +433,8 @@ def calculate( if household_id is not None and policy_id is not None: # write to local database - with runtime_v1_unit_of_work(local=True).transaction() as repositories: - repositories.tracers.create( + with runtime_v1_unit_of_work(local=True).transaction() as daos: + daos.tracers.create( household_id, policy_id, self.country_id, diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 1914eeb20..5ddf9a82b 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -233,8 +233,8 @@ def update_profile(self, user_id: int, **values: Any) -> bool: return True -class V1Repositories: - """Repositories bound to the same operation-scoped Session.""" +class V1DAOs: + """DAOs bound to the same operation-scoped Session.""" def __init__(self, session: Session): self.session = session @@ -258,14 +258,14 @@ def __init__(self, sessions: SessionManager): self.sessions = sessions @contextmanager - def read(self) -> Iterator[V1Repositories]: + def read(self) -> Iterator[V1DAOs]: with self.sessions.session() as session: - yield V1Repositories(session) + yield V1DAOs(session) @contextmanager - def transaction(self) -> Iterator[V1Repositories]: + def transaction(self) -> Iterator[V1DAOs]: with self.sessions.transaction() as session: - yield V1Repositories(session) + yield V1DAOs(session) _runtime_unit_of_work: dict[bool, V1UnitOfWork] = {} diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py index 842959cf7..e777e6975 100644 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ b/policyengine_api/endpoints/economy/reform_impact.py @@ -11,8 +11,8 @@ def set_comment_on_job( time_period, options_hash, ): - with runtime_v1_unit_of_work(local=True).transaction() as repositories: - repositories.reform_impacts.set_message( + with runtime_v1_unit_of_work(local=True).transaction() as daos: + daos.reform_impacts.set_message( comment, country_id=country_id, reform_policy_id=policy_id, diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index cbcf3b358..4b1725cd1 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -111,8 +111,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Look in computed_households to see if already computed - with runtime_v1_unit_of_work(local=True).read() as repositories: - row = repositories.computed_households.get( + with runtime_v1_unit_of_work(local=True).read() as daos: + row = daos.computed_households.get( int(household_id), int(policy_id), country_id, @@ -141,9 +141,9 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the household table - with runtime_v1_unit_of_work().read() as repositories: - row = repositories.households.get(country_id, int(household_id)) - policy_row = repositories.policies.get(country_id, int(policy_id)) + with runtime_v1_unit_of_work().read() as daos: + row = daos.households.get(country_id, int(household_id)) + policy_row = daos.policies.get(country_id, int(policy_id)) if row is not None: household = dict(row) @@ -220,8 +220,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Store the result in the computed_household table - with runtime_v1_unit_of_work(local=True).transaction() as repositories: - repositories.computed_households.upsert( + with runtime_v1_unit_of_work(local=True).transaction() as daos: + daos.computed_households.upsert( country_id=country_id, household_id=int(household_id), policy_id=int(policy_id), diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index 4c8fd6b84..e247fb84d 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -31,8 +31,8 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - with runtime_v1_unit_of_work().read() as repositories: - results = repositories.policies.search(country_id, query) + with runtime_v1_unit_of_work().read() as daos: + results = daos.policies.search(country_id, query) if not results: body = dict( @@ -125,11 +125,11 @@ def set_user_policy(country_id: str) -> dict: # to be tested; type is not yet implemented try: - with runtime_v1_unit_of_work().transaction() as repositories: - row = repositories.user_policies.find_unique(**values) + with runtime_v1_unit_of_work().transaction() as daos: + row = daos.user_policies.find_unique(**values) if row is None: - user_policy_id = repositories.user_policies.create(**values) - row = repositories.user_policies.get(user_policy_id) + user_policy_id = daos.user_policies.create(**values) + row = daos.user_policies.get(user_policy_id) else: readable_row = dict(row) @@ -189,8 +189,8 @@ def get_user_policy(country_id: str, user_id: str) -> dict: """ # Get the policy record for a given policy ID. - with runtime_v1_unit_of_work().read() as repositories: - rows = repositories.user_policies.list_for_user(country_id, user_id) + with runtime_v1_unit_of_work().read() as daos: + rows = daos.user_policies.list_for_user(country_id, user_id) rows_parsed = [ dict( @@ -301,8 +301,8 @@ def update_user_policy(country_id: str) -> dict: ) try: - with runtime_v1_unit_of_work().transaction() as repositories: - repositories.user_policies.update(user_policy_id, payload) + with runtime_v1_unit_of_work().transaction() as daos: + daos.user_policies.update(user_policy_id, payload) except Exception as e: return Response( json.dumps( diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index a0f22df8c..97eb3e303 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -42,8 +42,8 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - with runtime_v1_unit_of_work(local=True).read() as repositories: - result = repositories.reform_impacts.list_recent(max_results) + with runtime_v1_unit_of_work(local=True).read() as daos: + result = daos.reform_impacts.list_recent(max_results) # Format into [{}] diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index 186dee02d..28da6cb30 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -46,8 +46,8 @@ def _analysis_repository(self, *, write: bool = False): yield self._analyses return boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as repositories: - yield repositories.analyses + with boundary() as daos: + yield daos.analyses def get_existing_analysis(self, prompt: str) -> str | None: with self._analysis_repository() as analyses: diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 96614c614..945452b62 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -30,8 +30,8 @@ def _repository(self, *, write: bool = False): yield self._households return boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as repositories: - yield repositories.households + with boundary() as daos: + yield daos.households def get_household(self, country_id: str, household_id: int) -> dict | None: if type(household_id) is not int or household_id < 0: diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index ee3eb0bf1..a288b3f51 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -31,8 +31,8 @@ def _repository(self, *, write: bool = False): yield self._policies return boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as repositories: - yield repositories.policies + with boundary() as daos: + yield daos.policies @staticmethod def _validate_policy_id(policy_id: int) -> None: diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 09f6e1e54..2a8a3fbbb 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -28,8 +28,8 @@ def _repository(self, *, write: bool = False): yield self._impacts return boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as repositories: - yield repositories.reform_impacts + with boundary() as daos: + yield daos.reform_impacts @staticmethod def _filters( diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index c95bf69d9..34dfc88cf 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -21,22 +21,22 @@ def unit_of_work(self) -> V1UnitOfWork: def _get_report_output_row(self, report_output_id: int) -> dict | None: if self._reports is not None: return self._reports.get(report_output_id) - with self.unit_of_work.read() as repositories: - return repositories.reports.get(report_output_id) + with self.unit_of_work.read() as daos: + return daos.reports.get(report_output_id) def get_alias(self, legacy_report_output_id: int) -> dict | None: if self._reports is not None: return self._reports.get_alias(legacy_report_output_id) - with self.unit_of_work.read() as repositories: - return repositories.reports.get_alias(legacy_report_output_id) + with self.unit_of_work.read() as daos: + return daos.reports.get_alias(legacy_report_output_id) def resolve_canonical_report_output_id( self, requested_report_output_id: int ) -> int | None: if self._reports is not None: return self._resolve(self._reports, requested_report_output_id) - with self.unit_of_work.read() as repositories: - return self._resolve(repositories.reports, requested_report_output_id) + with self.unit_of_work.read() as daos: + return self._resolve(daos.reports, requested_report_output_id) def _resolve( self, reports: ReportDAO, requested_report_output_id: int @@ -61,9 +61,9 @@ def set_alias( legacy_report_output_id, canonical_report_output_id, ) - with self.unit_of_work.transaction() as repositories: + with self.unit_of_work.transaction() as daos: return self._set_alias( - repositories.reports, + daos.reports, legacy_report_output_id, canonical_report_output_id, ) diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index c3c153d26..39b2130a9 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -86,10 +86,10 @@ def _get_report_output_row( for_update: bool = False, ) -> dict | None: if queryer is None: - with self.unit_of_work.read() as repositories: + with self.unit_of_work.read() as daos: return self._get_report_output_row( report_output_id, - queryer=repositories, + queryer=daos, country_id=country_id, for_update=for_update, ) @@ -105,10 +105,10 @@ def _get_linked_simulations( bootstrap_dual_write_state: bool = False, ) -> tuple[dict, dict | None]: if queryer is None: - with self.unit_of_work.read() as repositories: + with self.unit_of_work.read() as daos: return self._get_linked_simulations( report_output, - queryer=repositories, + queryer=daos, bootstrap_dual_write_state=bootstrap_dual_write_state, ) if bootstrap_dual_write_state: @@ -165,10 +165,10 @@ def _list_report_runs_descending( self, report_output_id: int, *, queryer=None ) -> list[dict]: if queryer is None: - with self.unit_of_work.read() as repositories: + with self.unit_of_work.read() as daos: return self._list_report_runs_descending( report_output_id, - queryer=repositories, + queryer=daos, ) rows = queryer.reports.list_runs(report_output_id) @@ -595,9 +595,9 @@ def ensure_report_output_dual_write_state( report_output_id: int, country_id: str | None = None, ) -> dict: - with self.unit_of_work.transaction() as repositories: + with self.unit_of_work.transaction() as daos: return self._ensure_report_output_dual_write_state_in_transaction( - repositories, + daos, report_output_id, country_id=country_id, ) @@ -648,13 +648,13 @@ def _find_existing_report_output_row( ) -> dict | None: api_version = get_report_output_cache_version(country_id) if queryer is None: - with self.unit_of_work.read() as repositories: + with self.unit_of_work.read() as daos: return self._find_existing_report_output_row( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, year=year, - queryer=repositories, + queryer=daos, ) return queryer.reports.find_latest( country_id=country_id, @@ -731,37 +731,37 @@ def create_report_output( api_version = get_report_output_cache_version(country_id) try: - with self.unit_of_work.transaction() as repositories: + with self.unit_of_work.transaction() as daos: existing_report = self._find_existing_report_output_row( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, year=year, - queryer=repositories, + queryer=daos, ) if existing_report is not None: print( f"Reusing existing report output with ID: {existing_report['id']}" ) return self._ensure_report_output_dual_write_state_in_transaction( - repositories, + daos, existing_report["id"], country_id=country_id, ) self._require_simulation_exists( - repositories, + daos, country_id=country_id, simulation_id=simulation_1_id, ) if simulation_2_id is not None: self._require_simulation_exists( - repositories, + daos, country_id=country_id, simulation_id=simulation_2_id, ) - report_output_id = repositories.reports.create( + report_output_id = daos.reports.create( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, @@ -769,13 +769,13 @@ def create_report_output( status="pending", year=year, ) - created_report = repositories.reports.get(report_output_id, country_id) + created_report = daos.reports.get(report_output_id, country_id) if created_report is None: raise Exception("Failed to retrieve created report output") print(f"Created report output with ID: {created_report['id']}") return self._ensure_report_output_dual_write_state_in_transaction( - repositories, + daos, created_report["id"], country_id=country_id, ) @@ -847,10 +847,10 @@ def update_report_output( print("No fields to update") return False - with self.unit_of_work.transaction() as repositories: + with self.unit_of_work.transaction() as daos: requested_report = self._get_report_output_row( report_id, - queryer=repositories, + queryer=daos, country_id=country_id, for_update=True, ) @@ -858,16 +858,16 @@ def update_report_output( raise ValueError(f"Report output #{report_id} not found") if status == "running" and not self._has_mutable_running_run( - requested_report, queryer=repositories + requested_report, queryer=daos ): raise ValueError( "Cannot mark report output running without an active " "pending or running report run" ) - repositories.reports.update(report_id, **update_values) + daos.reports.update(report_id, **update_values) self._ensure_report_output_dual_write_state_in_transaction( - repositories, + daos, report_id, country_id=country_id, ) diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index b1f6c2faf..285b47b9d 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -90,8 +90,8 @@ def create_report_output_run( **values, ) else: - with self.unit_of_work.transaction() as repositories: - run = repositories.reports.create_run( + with self.unit_of_work.transaction() as daos: + run = daos.reports.create_run( report_output_id, run_id=run_id or str(uuid.uuid4()), **values, @@ -103,23 +103,23 @@ def create_report_output_run( def get_report_output_run(self, run_id: str) -> dict | None: if self._reports is not None: return self._parse_run_row(self._reports.get_run(run_id)) - with self.unit_of_work.read() as repositories: - return self._parse_run_row(repositories.reports.get_run(run_id)) + with self.unit_of_work.read() as daos: + return self._parse_run_row(daos.reports.get_run(run_id)) def list_report_output_runs(self, report_output_id: int) -> list[dict]: if self._reports is not None: rows = self._reports.list_runs(report_output_id) else: - with self.unit_of_work.read() as repositories: - rows = repositories.reports.list_runs(report_output_id) + with self.unit_of_work.read() as daos: + rows = daos.reports.list_runs(report_output_id) return [self._parse_run_row(row) for row in reversed(rows)] def get_newest_report_output_run(self, report_output_id: int) -> dict | None: if self._reports is not None: rows = self._reports.list_runs(report_output_id) else: - with self.unit_of_work.read() as repositories: - rows = repositories.reports.list_runs(report_output_id) + with self.unit_of_work.read() as daos: + rows = daos.reports.list_runs(report_output_id) return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, report_output: dict) -> dict | None: diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index dce85b619..aad376824 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -68,25 +68,23 @@ def _validate_schema_version(self, schema_version: int | None) -> None: def _get_report_output_row(self, report_output_id: int) -> dict | None: if self._reports is not None: return self._reports.get(report_output_id) - with self.unit_of_work.read() as repositories: - return repositories.reports.get(report_output_id) + with self.unit_of_work.read() as daos: + return daos.reports.get(report_output_id) def _get_simulation_row(self, simulation_id: int) -> dict | None: if self._simulations is not None: return self._simulations.get(simulation_id) - with self.unit_of_work.read() as repositories: - return repositories.simulations.get(simulation_id) + with self.unit_of_work.read() as daos: + return daos.simulations.get(simulation_id) def _get_linked_simulations( - self, report_output: dict, *, repositories=None, simulations=None + self, report_output: dict, *, daos=None, simulations=None ) -> tuple[dict, dict | None]: simulations = simulations or self._simulations - if repositories is None and simulations is None: - with self.unit_of_work.read() as read_repositories: - return self._get_linked_simulations( - report_output, repositories=read_repositories - ) - simulations = simulations or repositories.simulations + if daos is None and simulations is None: + with self.unit_of_work.read() as read_daos: + return self._get_linked_simulations(report_output, daos=read_daos) + simulations = simulations or daos.simulations simulation_1 = simulations.get(report_output["simulation_1_id"]) if simulation_1 is None: raise ValueError( @@ -236,12 +234,12 @@ def _validate_report_spec_matches_row( report_output: dict, report_spec: ReportSpec, *, - repositories=None, + daos=None, simulations=None, ) -> None: simulation_1, simulation_2 = self._get_linked_simulations( report_output, - repositories=repositories, + daos=daos, simulations=simulations, ) inferred_report_kind = self.infer_report_kind(simulation_1, simulation_2) @@ -382,30 +380,26 @@ def _parse_report_spec(self, report_kind: str, raw_spec: dict) -> ReportSpec: def get_report_spec(self, report_output_id: int) -> ReportSpec | None: if self._reports is not None and self._simulations is not None: report_output = self._reports.get(report_output_id) - repositories = None + daos = None else: - with self.unit_of_work.read() as repositories: - return self._get_report_spec(report_output_id, repositories) - return self._parse_stored_report_spec(report_output, repositories=repositories) + with self.unit_of_work.read() as daos: + return self._get_report_spec(report_output_id, daos) + return self._parse_stored_report_spec(report_output, daos=daos) - def _get_report_spec( - self, report_output_id: int, repositories - ) -> ReportSpec | None: + def _get_report_spec(self, report_output_id: int, daos) -> ReportSpec | None: return self._parse_stored_report_spec( - repositories.reports.get(report_output_id), repositories=repositories + daos.reports.get(report_output_id), daos=daos ) def _parse_stored_report_spec( - self, report_output: dict | None, *, repositories=None + self, report_output: dict | None, *, daos=None ) -> ReportSpec | None: if report_output is None or report_output["report_spec_json"] is None: return None self._validate_schema_version(report_output["report_spec_schema_version"]) raw_spec = self._parse_json_field(report_output["report_spec_json"]) report_spec = self._parse_report_spec(report_output["report_kind"], raw_spec) - self._validate_report_spec_matches_row( - report_output, report_spec, repositories=repositories - ) + self._validate_report_spec_matches_row(report_output, report_spec, daos=daos) return report_spec def set_report_spec( @@ -430,14 +424,14 @@ def set_report_spec( ) return True - with self.unit_of_work.transaction() as repositories: - report_output = repositories.reports.get(report_output_id) + with self.unit_of_work.transaction() as daos: + report_output = daos.reports.get(report_output_id) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") self._validate_report_spec_matches_row( - report_output, report_spec, repositories=repositories + report_output, report_spec, daos=daos ) - repositories.reports.update( + daos.reports.update( report_output_id, report_kind=report_spec.report_kind, report_spec_json=report_spec.model_dump(), diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 0eb74b2af..c4377ca23 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -82,8 +82,8 @@ def create_simulation_run( **values, ) else: - with self.unit_of_work.transaction() as repositories: - run = repositories.simulations.create_run( + with self.unit_of_work.transaction() as daos: + run = daos.simulations.create_run( simulation_id, run_id=run_id or str(uuid.uuid4()), **values, @@ -95,23 +95,23 @@ def create_simulation_run( def get_simulation_run(self, run_id: str) -> dict | None: if self._simulations is not None: return self._parse_run_row(self._simulations.get_run(run_id)) - with self.unit_of_work.read() as repositories: - return self._parse_run_row(repositories.simulations.get_run(run_id)) + with self.unit_of_work.read() as daos: + return self._parse_run_row(daos.simulations.get_run(run_id)) def list_simulation_runs(self, simulation_id: int) -> list[dict]: if self._simulations is not None: rows = self._simulations.list_runs(simulation_id) else: - with self.unit_of_work.read() as repositories: - rows = repositories.simulations.list_runs(simulation_id) + with self.unit_of_work.read() as daos: + rows = daos.simulations.list_runs(simulation_id) return [self._parse_run_row(row) for row in reversed(rows)] def get_newest_simulation_run(self, simulation_id: int) -> dict | None: if self._simulations is not None: rows = self._simulations.list_runs(simulation_id) else: - with self.unit_of_work.read() as repositories: - rows = repositories.simulations.list_runs(simulation_id) + with self.unit_of_work.read() as daos: + rows = daos.simulations.list_runs(simulation_id) return self._parse_run_row(rows[0]) if rows else None def select_display_run(self, simulation: dict) -> dict | None: diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 8a835f99b..56bfb4f9b 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -58,18 +58,16 @@ def _get_simulation_row( return simulations.get(simulation_id, country_id) if self._simulations is not None: return self._simulations.get(simulation_id, country_id) - with self.unit_of_work.read() as repositories: - return repositories.simulations.get(simulation_id, country_id) + with self.unit_of_work.read() as daos: + return daos.simulations.get(simulation_id, country_id) def ensure_simulation_dual_write_state( self, simulation_id: int, country_id: str | None = None ) -> dict: if self._simulations is not None: return self._simulations.ensure_dual_write_state(simulation_id, country_id) - with self.unit_of_work.transaction() as repositories: - return repositories.simulations.ensure_dual_write_state( - simulation_id, country_id - ) + with self.unit_of_work.transaction() as daos: + return daos.simulations.ensure_dual_write_state(simulation_id, country_id) def find_existing_simulation( self, @@ -85,8 +83,8 @@ def find_existing_simulation( population_type=population_type, policy_id=policy_id, ) - with self.unit_of_work.read() as repositories: - return repositories.simulations.find_latest( + with self.unit_of_work.read() as daos: + return daos.simulations.find_latest( country_id=country_id, population_id=population_id, population_type=population_type, @@ -113,8 +111,8 @@ def create_simulation( sync_callback=self._ensure_simulation_dual_write_state_in_transaction, **values, ) - with self.unit_of_work.transaction() as repositories: - return repositories.simulations.create_or_get_with_sync( + with self.unit_of_work.transaction() as daos: + return daos.simulations.create_or_get_with_sync( sync_callback=self._ensure_simulation_dual_write_state_in_transaction, **values, ) @@ -156,8 +154,8 @@ def update_simulation( self._ensure_simulation_dual_write_state_in_transaction, ) else: - with self.unit_of_work.transaction() as repositories: - repositories.simulations.update_with_sync( + with self.unit_of_work.transaction() as daos: + daos.simulations.update_with_sync( simulation_id, country_id, values, diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index a848dc97a..0320247e3 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -40,8 +40,8 @@ def _validate_schema_version(self, schema_version: int | None) -> None: def _get_simulation_row(self, simulation_id: int) -> dict | None: if self._simulations is not None: return self._simulations.get(simulation_id) - with self.unit_of_work.read() as repositories: - return repositories.simulations.get(simulation_id) + with self.unit_of_work.read() as daos: + return daos.simulations.get(simulation_id) def _validate_simulation_spec_matches_row( self, simulation: dict, simulation_spec: SimulationSpec @@ -98,12 +98,12 @@ def set_simulation_spec( ) return True - with self.unit_of_work.transaction() as repositories: - simulation = repositories.simulations.get(simulation_id) + with self.unit_of_work.transaction() as daos: + simulation = daos.simulations.get(simulation_id) if simulation is None: raise ValueError(f"Simulation #{simulation_id} not found") self._validate_simulation_spec_matches_row(simulation, simulation_spec) - repositories.simulations.update( + daos.simulations.update( simulation_id, simulation_spec_json=simulation_spec.model_dump(), simulation_spec_schema_version=schema_version, diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index d915d2adb..2b0013bf9 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -26,8 +26,8 @@ def _tracer_repository(self): if self._tracers is not None: yield self._tracers return - with self.unit_of_work.read() as repositories: - yield repositories.tracers + with self.unit_of_work.read() as daos: + yield daos.tracers def execute_analysis( self, diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index 2ae2b228b..32eece842 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -29,8 +29,8 @@ def _repository(self, *, write: bool = False): yield self._users return boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as repositories: - yield repositories.users + with boundary() as daos: + yield daos.users def create_profile( self, diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py index 80c45dac4..510a66ad5 100644 --- a/tests/unit/data/test_local_daos.py +++ b/tests/unit/data/test_local_daos.py @@ -13,16 +13,16 @@ def _unit_of_work(): def test_analysis_dao_round_trip(): uow = _unit_of_work() - with uow.transaction() as repositories: - repositories.analyses.store("prompt", "answer", "complete") - with uow.read() as repositories: - assert repositories.analyses.get("prompt") == "answer" + with uow.transaction() as daos: + daos.analyses.store("prompt", "answer", "complete") + with uow.read() as daos: + assert daos.analyses.get("prompt") == "answer" def test_reform_impact_dao_transitions_by_execution_id(): uow = _unit_of_work() - with uow.transaction() as repositories: - repositories.reform_impacts.create( + with uow.transaction() as daos: + daos.reform_impacts.create( country_id="us", reform_policy_id=2, baseline_policy_id=1, @@ -37,19 +37,19 @@ def test_reform_impact_dao_transitions_by_execution_id(): start_time=datetime(2026, 1, 1), execution_id="job", ) - repositories.reform_impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) - with uow.read() as repositories: - assert repositories.reform_impacts.find(execution_id="job")["status"] == "ok" - assert repositories.reform_impacts.find(execution_id="job")[ - "reform_impact_json" - ] == {"result": 1} + daos.reform_impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) + with uow.read() as daos: + assert daos.reform_impacts.find(execution_id="job")["status"] == "ok" + assert daos.reform_impacts.find(execution_id="job")["reform_impact_json"] == { + "result": 1 + } def test_reform_impact_dao_orders_limits_messages_and_handles_missing_jobs(): uow = _unit_of_work() - with uow.transaction() as repositories: + with uow.transaction() as daos: for day, execution_id in ((1, "old-job"), (2, "new-job")): - repositories.reform_impacts.create( + daos.reform_impacts.create( country_id="us", reform_policy_id=2, baseline_policy_id=1, @@ -65,37 +65,29 @@ def test_reform_impact_dao_orders_limits_messages_and_handles_missing_jobs(): execution_id=execution_id, ) - assert repositories.reform_impacts.set_message( + assert daos.reform_impacts.set_message( "queued", country_id="us", status="computing" ) + assert daos.reform_impacts.set_message("missing", country_id="uk") is False assert ( - repositories.reform_impacts.set_message("missing", country_id="uk") is False - ) - assert ( - repositories.reform_impacts.fail( - "missing-job", "failed", datetime(2026, 1, 3) - ) + daos.reform_impacts.fail("missing-job", "failed", datetime(2026, 1, 3)) is False ) assert ( - repositories.reform_impacts.complete( - "missing-job", {}, datetime(2026, 1, 3) - ) + daos.reform_impacts.complete("missing-job", {}, datetime(2026, 1, 3)) is False ) - with uow.read() as repositories: - recent = repositories.reform_impacts.list_recent(1) + with uow.read() as daos: + recent = daos.reform_impacts.list_recent(1) assert [row["execution_id"] for row in recent] == ["new-job"] assert recent[0]["message"] == "queued" def test_tracer_dao_returns_latest_matching_trace(): uow = _unit_of_work() - with uow.transaction() as repositories: - repositories.tracers.create(1, 2, "us", "1", {"trace": "first"}) - repositories.tracers.create(1, 2, "us", "1", {"trace": "latest"}) - with uow.read() as repositories: - assert repositories.tracers.get(1, 2, "us")["tracer_output"] == { - "trace": "latest" - } + with uow.transaction() as daos: + daos.tracers.create(1, 2, "us", "1", {"trace": "first"}) + daos.tracers.create(1, 2, "us", "1", {"trace": "latest"}) + with uow.read() as daos: + assert daos.tracers.get(1, 2, "us")["tracer_output"] == {"trace": "latest"} diff --git a/tests/unit/data/test_ordinary_v1_daos.py b/tests/unit/data/test_ordinary_v1_daos.py index cbf5cf159..219184e0a 100644 --- a/tests/unit/data/test_ordinary_v1_daos.py +++ b/tests/unit/data/test_ordinary_v1_daos.py @@ -19,13 +19,13 @@ def test_computed_household_upsert_preserves_one_cache_row(): "computed_household_json": {"value": 1}, "status": "complete", } - with uow.transaction() as repositories: - repositories.computed_households.upsert(**values) - repositories.computed_households.upsert( + with uow.transaction() as daos: + daos.computed_households.upsert(**values) + daos.computed_households.upsert( **{**values, "computed_household_json": {"value": 2}} ) - with uow.read() as repositories: - row = repositories.computed_households.get(1, 2, "us", api_version="1") + with uow.read() as daos: + row = daos.computed_households.get(1, 2, "us", api_version="1") assert row["computed_household_json"] == {"value": 2} @@ -48,32 +48,30 @@ def test_user_policy_nullable_identity_list_and_update_are_orm_managed(): "budgetary_impact": None, "type": None, } - with uow.transaction() as repositories: - user_policy_id = repositories.user_policies.create(**values) - assert repositories.user_policies.find_unique(**values)["id"] == user_policy_id - assert repositories.user_policies.update( - user_policy_id, {"number_of_provisions": 4} - ) - with uow.read() as repositories: - rows = repositories.user_policies.list_for_user("us", "auth0|one") + with uow.transaction() as daos: + user_policy_id = daos.user_policies.create(**values) + assert daos.user_policies.find_unique(**values)["id"] == user_policy_id + assert daos.user_policies.update(user_policy_id, {"number_of_provisions": 4}) + with uow.read() as daos: + rows = daos.user_policies.list_for_user("us", "auth0|one") assert rows[0]["number_of_provisions"] == 4 def test_policy_search_and_reform_impact_limit_use_typed_statements(): uow = _unit_of_work() - with uow.transaction() as repositories: - repositories.policies.create("us", "Tax reform", {}, "one", "1") - repositories.policies.create("us", "Other", {}, "two", "1") - with uow.read() as repositories: - assert [row["label"] for row in repositories.policies.search("us", "Tax")] == [ + with uow.transaction() as daos: + daos.policies.create("us", "Tax reform", {}, "one", "1") + daos.policies.create("us", "Other", {}, "two", "1") + with uow.read() as daos: + assert [row["label"] for row in daos.policies.search("us", "Tax")] == [ "Tax reform" ] def test_computed_household_create_and_version_filters(): uow = _unit_of_work() - with uow.transaction() as repositories: - repositories.computed_households.create( + with uow.transaction() as daos: + daos.computed_households.create( household_id=1, policy_id=2, country_id="us", @@ -82,11 +80,11 @@ def test_computed_household_create_and_version_filters(): status="complete", ) - with uow.read() as repositories: - assert repositories.computed_households.get(1, 2, "us")[ - "computed_household_json" - ] == {"value": 1} - assert repositories.computed_households.get(1, 2, "us", api_version="2") is None + with uow.read() as daos: + assert daos.computed_households.get(1, 2, "us")["computed_household_json"] == { + "value": 1 + } + assert daos.computed_households.get(1, 2, "us", api_version="2") is None def test_economy_and_user_policy_daos_cover_lookup_edge_cases(): @@ -108,8 +106,8 @@ def test_economy_and_user_policy_daos_cover_lookup_edge_cases(): "budgetary_impact": None, "type": None, } - with uow.transaction() as repositories: - economy_id = repositories.economies.create( + with uow.transaction() as daos: + economy_id = daos.economies.create( policy_id=2, country_id="us", region="us", @@ -121,11 +119,11 @@ def test_economy_and_user_policy_daos_cover_lookup_edge_cases(): status="complete", message=None, ) - user_policy_id = repositories.user_policies.create(**user_policy_values) - assert repositories.user_policies.update(999, {}) is False + user_policy_id = daos.user_policies.create(**user_policy_values) + assert daos.user_policies.update(999, {}) is False - with uow.read() as repositories: - assert repositories.economies.get(economy_id)["economy_json"] == {"result": 1} - assert repositories.economies.get(999) is None - assert repositories.user_policies.get(user_policy_id)["country_id"] == "us" - assert repositories.user_policies.get(999) is None + with uow.read() as daos: + assert daos.economies.get(economy_id)["economy_json"] == {"result": 1} + assert daos.economies.get(999) is None + assert daos.user_policies.get(user_policy_id)["country_id"] == "us" + assert daos.user_policies.get(999) is None diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py index 9c1545a96..1f6d17708 100644 --- a/tests/unit/data/test_run_daos.py +++ b/tests/unit/data/test_run_daos.py @@ -15,59 +15,56 @@ def _unit_of_work(): def test_simulation_dao_creates_parent_and_monotonic_runs_atomically(): uow = _unit_of_work() - with uow.transaction() as repositories: - simulation_id = repositories.simulations.create( + with uow.transaction() as daos: + simulation_id = daos.simulations.create( country_id="us", api_version="1", population_id="7", population_type="household", policy_id=2, ) - first = repositories.simulations.create_run( + first = daos.simulations.create_run( simulation_id, run_id="run-1", status="pending", trigger_type="create", requested_at=datetime(2026, 1, 1), ) - second = repositories.simulations.create_run( + second = daos.simulations.create_run( simulation_id, run_id="run-2", status="pending", trigger_type="retry", requested_at=datetime(2026, 1, 2), ) - with uow.read() as repositories: + with uow.read() as daos: assert first["run_sequence"] == 1 assert second["run_sequence"] == 2 - assert repositories.simulations.list_runs(simulation_id)[0]["id"] == "run-2" + assert daos.simulations.list_runs(simulation_id)[0]["id"] == "run-2" def test_report_dao_round_trips_parent_run_and_alias(): uow = _unit_of_work() - with uow.transaction() as repositories: - report_id = repositories.reports.create( + with uow.transaction() as daos: + report_id = daos.reports.create( country_id="us", simulation_1_id=1, simulation_2_id=None, api_version="1", year="2026", ) - run = repositories.reports.create_run( + run = daos.reports.create_run( report_id, run_id="report-run", status="pending", trigger_type="create", requested_at=datetime(2026, 1, 1), ) - repositories.reports.set_alias(99, report_id) - with uow.read() as repositories: - assert repositories.reports.get(report_id)["status"] == "pending" + daos.reports.set_alias(99, report_id) + with uow.read() as daos: + assert daos.reports.get(report_id)["status"] == "pending" assert run["run_sequence"] == 1 - assert ( - repositories.reports.get_alias(99)["canonical_report_output_id"] - == report_id - ) + assert daos.reports.get_alias(99)["canonical_report_output_id"] == report_id def test_simulation_dao_sync_callbacks_cover_create_update_and_missing_rows(): @@ -76,8 +73,8 @@ def test_simulation_dao_sync_callbacks_cover_create_update_and_missing_rows(): def read_synced(session, simulation_id, *, country_id): return SimulationDAO.get_in_session(session, simulation_id, country_id) - with uow.transaction() as repositories: - created = repositories.simulations.create_or_get_with_sync( + with uow.transaction() as daos: + created = daos.simulations.create_or_get_with_sync( sync_callback=read_synced, country_id="us", api_version="1", @@ -87,7 +84,7 @@ def read_synced(session, simulation_id, *, country_id): status="complete", output={"result": 1}, ) - reused = repositories.simulations.create_or_get_with_sync( + reused = daos.simulations.create_or_get_with_sync( sync_callback=read_synced, country_id="us", api_version="1", @@ -96,27 +93,25 @@ def read_synced(session, simulation_id, *, country_id): policy_id=2, status="pending", ) - updated = repositories.simulations.update_with_sync( + updated = daos.simulations.update_with_sync( created["id"], "us", {"error_message": "updated"}, read_synced, ) - dual_write = repositories.simulations.ensure_dual_write_state( - created["id"], "us" - ) + dual_write = daos.simulations.ensure_dual_write_state(created["id"], "us") assert reused["id"] == created["id"] assert updated["error_message"] == "updated" assert dual_write["latest_successful_run_id"] is not None - assert repositories.simulations.get(created["id"], "uk") is None - assert repositories.simulations.update(999, status="complete") is False + assert daos.simulations.get(created["id"], "uk") is None + assert daos.simulations.update(999, status="complete") is False with pytest.raises(ValueError, match="Simulation #999 not found"): - repositories.simulations.update_with_sync( + daos.simulations.update_with_sync( 999, "us", {"status": "complete"}, read_synced ) with pytest.raises(LookupError, match="Simulation 999 does not exist"): - repositories.simulations.create_run( + daos.simulations.create_run( 999, run_id="missing-run", status="pending", @@ -126,42 +121,42 @@ def read_synced(session, simulation_id, *, country_id): def test_report_dao_handles_scoped_lookups_updates_and_existing_aliases(): uow = _unit_of_work() - with uow.transaction() as repositories: - report_id = repositories.reports.create( + with uow.transaction() as daos: + report_id = daos.reports.create( country_id="us", simulation_1_id=1, simulation_2_id=None, api_version="1", year="2026", ) - run = repositories.reports.create_run( + run = daos.reports.create_run( report_id, run_id="report-run", status="pending", trigger_type="create", ) - assert repositories.reports.get(report_id, "uk") is None - assert repositories.reports.get_for_update(report_id, "us")["id"] == report_id - assert repositories.reports.get_for_update(report_id, "uk") is None - assert repositories.reports.update(999, status="complete") is False - assert repositories.reports.update(report_id, status="complete") - assert repositories.reports.update_run( + assert daos.reports.get(report_id, "uk") is None + assert daos.reports.get_for_update(report_id, "us")["id"] == report_id + assert daos.reports.get_for_update(report_id, "uk") is None + assert daos.reports.update(999, status="complete") is False + assert daos.reports.update(report_id, status="complete") + assert daos.reports.update_run( run["id"], status="complete", output={"result": 1} ) - assert repositories.reports.update_run("missing-run", status="error") is False - repositories.reports.set_alias(99, report_id) - repositories.reports.set_alias(99, report_id + 1) + assert daos.reports.update_run("missing-run", status="error") is False + daos.reports.set_alias(99, report_id) + daos.reports.set_alias(99, report_id + 1) with pytest.raises(LookupError, match="Report output 999 does not exist"): - repositories.reports.create_run( + daos.reports.create_run( 999, run_id="missing-run", status="pending", trigger_type="create", ) - with uow.read() as repositories: - assert repositories.reports.get(report_id)["status"] == "complete" - assert repositories.reports.get_run(run["id"])["output"] == {"result": 1} - assert repositories.reports.get_run("missing-run") is None - assert repositories.reports.get_alias(99)["canonical_report_output_id"] == 2 + with uow.read() as daos: + assert daos.reports.get(report_id)["status"] == "complete" + assert daos.reports.get_run(run["id"])["output"] == {"result": 1} + assert daos.reports.get_run("missing-run") is None + assert daos.reports.get_alias(99)["canonical_report_output_id"] == 2 diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index a61af0371..0ac8ae704 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -39,7 +39,7 @@ def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): assert "runtime_sqlalchemy_dao" not in source -def test_report_orchestration_uses_typed_repositories_not_raw_sql(): +def test_report_orchestration_uses_typed_daos_not_raw_sql(): source = (PACKAGE_ROOT / "services/report_output_service.py").read_text( encoding="utf-8" ) diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index 022d41b2d..5fc8ed8d6 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -11,13 +11,11 @@ def _unit_of_work(): def test_policy_dao_round_trips_legacy_mapping_shape(): uow = _unit_of_work() - with uow.transaction() as repositories: - policy_id = repositories.policies.create( - "us", "Reform", {"gov.irs": 1}, "hash", "1.0" - ) - with uow.read() as repositories: + with uow.transaction() as daos: + policy_id = daos.policies.create("us", "Reform", {"gov.irs": 1}, "hash", "1.0") + with uow.read() as daos: assert policy_id == 1 - assert repositories.policies.get("us", policy_id) == { + assert daos.policies.get("us", policy_id) == { "id": 1, "country_id": "us", "label": "Reform", @@ -29,20 +27,18 @@ def test_policy_dao_round_trips_legacy_mapping_shape(): def test_policy_dao_allocates_ids_and_detects_existing_policy(): uow = _unit_of_work() - with uow.transaction() as repositories: - assert repositories.policies.create("us", None, {}, "one", "1.0") == 1 - assert repositories.policies.create("uk", None, {}, "two", "1.0") == 2 - with uow.read() as repositories: - assert repositories.policies.find_unique("us", "one", None)["id"] == 1 + with uow.transaction() as daos: + assert daos.policies.create("us", None, {}, "one", "1.0") == 1 + assert daos.policies.create("uk", None, {}, "two", "1.0") == 2 + with uow.read() as daos: + assert daos.policies.find_unique("us", "one", None)["id"] == 1 def test_household_dao_creates_updates_and_reads(): uow = _unit_of_work() - with uow.transaction() as repositories: - household_id = repositories.households.create( - "us", "Home", {"people": {}}, "h", "1.0" - ) - repositories.households.update( + with uow.transaction() as daos: + household_id = daos.households.create("us", "Home", {"people": {}}, "h", "1.0") + daos.households.update( "us", household_id, "Updated", @@ -50,44 +46,37 @@ def test_household_dao_creates_updates_and_reads(): "updated-hash", "2.0", ) - with uow.read() as repositories: - assert repositories.households.get("us", household_id)["label"] == "Updated" - assert repositories.households.get("uk", household_id) is None + with uow.read() as daos: + assert daos.households.get("us", household_id)["label"] == "Updated" + assert daos.households.get("uk", household_id) is None def test_user_dao_profile_lookup_precedence(): uow = _unit_of_work() - with uow.transaction() as repositories: - user_id = repositories.users.create_profile("auth0|one", "person", "us", 123) - with uow.read() as repositories: - assert ( - repositories.users.get_profile(auth0_id="auth0|one")["user_id"] == user_id - ) + with uow.transaction() as daos: + user_id = daos.users.create_profile("auth0|one", "person", "us", 123) + with uow.read() as daos: + assert daos.users.get_profile(auth0_id="auth0|one")["user_id"] == user_id assert ( - repositories.users.get_profile(user_id=user_id, auth0_id="wrong")[ - "auth0_id" - ] + daos.users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] == "auth0|one" ) def test_user_and_household_daos_handle_missing_and_nullable_updates(): uow = _unit_of_work() - with uow.transaction() as repositories: - user_id = repositories.users.create_profile("auth0|one", "original", "us", 123) - assert repositories.users.get_profile() is None - assert repositories.users.update_profile(999, username="missing") is False - assert ( - repositories.households.update("us", 999, "missing", {}, "missing", "1") - is False - ) - assert repositories.users.update_profile( + with uow.transaction() as daos: + user_id = daos.users.create_profile("auth0|one", "original", "us", 123) + assert daos.users.get_profile() is None + assert daos.users.update_profile(999, username="missing") is False + assert daos.households.update("us", 999, "missing", {}, "missing", "1") is False + assert daos.users.update_profile( user_id, username=None, primary_country="uk", ) - with uow.read() as repositories: - profile = repositories.users.get_profile(user_id=user_id) + with uow.read() as daos: + profile = daos.users.get_profile(user_id=user_id) assert profile["username"] == "original" assert profile["primary_country"] == "uk" diff --git a/tests/unit/data/test_v1_unit_of_work.py b/tests/unit/data/test_v1_unit_of_work.py index 633d7c16b..19bc613bc 100644 --- a/tests/unit/data/test_v1_unit_of_work.py +++ b/tests/unit/data/test_v1_unit_of_work.py @@ -11,54 +11,54 @@ def _unit_of_work() -> V1UnitOfWork: return V1UnitOfWork(manager) -def test_unit_of_work_commits_all_repositories_once(): +def test_unit_of_work_commits_all_daos_once(): uow = _unit_of_work() - with uow.transaction() as repositories: - policy_id = repositories.policies.create("us", None, {}, "policy", "1") - household_id = repositories.households.create("us", None, {}, "household", "1") + with uow.transaction() as daos: + policy_id = daos.policies.create("us", None, {}, "policy", "1") + household_id = daos.households.create("us", None, {}, "household", "1") - with uow.read() as repositories: - assert repositories.policies.get("us", policy_id) is not None - assert repositories.households.get("us", household_id) is not None + with uow.read() as daos: + assert daos.policies.get("us", policy_id) is not None + assert daos.households.get("us", household_id) is not None def test_unit_of_work_rolls_back_every_repository_on_failure(): uow = _unit_of_work() with pytest.raises(RuntimeError, match="abort"): - with uow.transaction() as repositories: - repositories.policies.create("us", None, {}, "policy", "1") - repositories.users.create_profile("auth0|one", "person", "us", 1) + with uow.transaction() as daos: + daos.policies.create("us", None, {}, "policy", "1") + daos.users.create_profile("auth0|one", "person", "us", 1) raise RuntimeError("abort") - with uow.read() as repositories: - assert repositories.policies.get("us", 1) is None - assert repositories.users.get_profile(auth0_id="auth0|one") is None + with uow.read() as daos: + assert daos.policies.get("us", 1) is None + assert daos.users.get_profile(auth0_id="auth0|one") is None def test_unit_of_work_rolls_back_parent_run_and_alias_together(): uow = _unit_of_work() with pytest.raises(RuntimeError, match="abort report"): - with uow.transaction() as repositories: - report_id = repositories.reports.create( + with uow.transaction() as daos: + report_id = daos.reports.create( country_id="us", simulation_1_id=1, simulation_2_id=None, api_version="1", year="2026", ) - repositories.reports.create_run( + daos.reports.create_run( report_id, run_id="report-run", status="pending", trigger_type="create", ) - repositories.reports.set_alias(99, report_id) + daos.reports.set_alias(99, report_id) raise RuntimeError("abort report") - with uow.read() as repositories: - assert repositories.reports.get(1) is None - assert repositories.reports.get_run("report-run") is None - assert repositories.reports.get_alias(99) is None + with uow.read() as daos: + assert daos.reports.get(1) is None + assert daos.reports.get_run("report-run") is None + assert daos.reports.get_alias(99) is None diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py index febac16a0..674d8566b 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -23,10 +23,10 @@ def _unit_of_work() -> V1UnitOfWork: return V1UnitOfWork(manager) -def _repositories_unit_of_work(repositories): +def _daos_unit_of_work(daos): @contextmanager def boundary(): - yield repositories + yield daos return SimpleNamespace(read=boundary, transaction=boundary) @@ -50,7 +50,7 @@ def test_household_under_policy_returns_cached_json_objects_and_legacy_strings( "computed_household_json": stored_result, "status": "complete", } - local_uow = _repositories_unit_of_work( + local_uow = _daos_unit_of_work( SimpleNamespace(computed_households=computed_households) ) @@ -67,10 +67,10 @@ def test_household_under_policy_returns_cached_json_objects_and_legacy_strings( def test_household_under_policy_calculates_and_caches_json_as_an_object(): computed_households = Mock() computed_households.get.return_value = None - local_uow = _repositories_unit_of_work( + local_uow = _daos_unit_of_work( SimpleNamespace(computed_households=computed_households) ) - remote_uow = _repositories_unit_of_work( + remote_uow = _daos_unit_of_work( SimpleNamespace( households=SimpleNamespace( get=Mock( @@ -172,5 +172,5 @@ def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): assert created.get_json()["result"]["dataset"] == "default" assert listed["result"][0]["reform_label"] == "Reform" assert updated.status_code == 200 - with uow.read() as repositories: - assert repositories.user_policies.get(1)["reform_label"] == "Updated" + with uow.read() as daos: + assert daos.user_policies.get(1)["reform_label"] == "Updated" diff --git a/tests/unit/services/test_reform_impacts_service.py b/tests/unit/services/test_reform_impacts_service.py index 98e639b92..5a1bded9c 100644 --- a/tests/unit/services/test_reform_impacts_service.py +++ b/tests/unit/services/test_reform_impacts_service.py @@ -96,9 +96,9 @@ def test_reform_impact_service_round_trips_queries_and_transitions(): "compatible-job", ) - with service.unit_of_work.read() as repositories: - completed = repositories.reform_impacts.find(execution_id="exact-job") - failed = repositories.reform_impacts.find(execution_id="compatible-job") + with service.unit_of_work.read() as daos: + completed = daos.reform_impacts.find(execution_id="exact-job") + failed = daos.reform_impacts.find(execution_id="compatible-job") assert completed["status"] == "ok" assert completed["reform_impact_json"] == {"result": 1} assert failed["status"] == "error" From 9b8da3dcc006e9b793d116aa5b1ac033f2885380 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 21:52:49 +0300 Subject: [PATCH 47/89] refactor: expose canonical ORM session factories --- policyengine_api/data/orm.py | 47 ++++++++++++++---- tests/unit/data/test_orm_sessions.py | 71 +++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 11 deletions(-) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 8d90aefcd..8b9c0042c 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -15,6 +15,9 @@ T = TypeVar("T") +_v1_session_factories: dict[bool, sessionmaker[Session]] = {} + + class _IndexedMappingRow(dict): """SQLite row compatible with both SQLAlchemy and legacy mapping callers.""" @@ -39,11 +42,7 @@ class SessionManager: def __init__(self, engine: Engine): self.engine = engine - self.session_factory = sessionmaker( - bind=engine, - class_=Session, - expire_on_commit=False, - ) + self.session_factory = build_session_factory(engine) @contextmanager def session(self) -> Iterator[Session]: @@ -60,6 +59,16 @@ def run_in_transaction(self, callback: Callable[[Session], T]) -> T: return callback(session) +def build_session_factory(engine: Engine) -> sessionmaker[Session]: + """Return the canonical SQLAlchemy session factory for an engine.""" + + return sessionmaker( + bind=engine, + class_=Session, + expire_on_commit=False, + ) + + def build_sqlite_session_manager( database_path: str | Path | None = None, ) -> SessionManager: @@ -77,7 +86,13 @@ def build_sqlite_session_manager( def build_v1_session_manager(*, local: bool = False) -> SessionManager: - """Bind ORM sessions to the database selected by the v1 runtime.""" + """Temporary bridge for callers not yet migrated to ``sessionmaker``.""" + + return SessionManager(get_v1_engine(local=local)) + + +def get_v1_engine(*, local: bool = False) -> Engine: + """Return the process-owned engine selected by the v1 runtime.""" from policyengine_api.data.data import database, local_database @@ -98,6 +113,20 @@ def build_v1_session_manager(*, local: bool = False) -> SessionManager: connection, "row_factory", _IndexedMappingRow ), ) - return SessionManager(engine) - return build_sqlite_session_manager(selected_database.db_url) - return SessionManager(selected_database.pool) + return engine + return create_engine(f"sqlite+pysqlite:///{Path(selected_database.db_url)}") + return selected_database.pool + + +def get_v1_session_factory(*, local: bool = False) -> sessionmaker[Session]: + """Return one configured factory per process-owned v1 engine.""" + + if local not in _v1_session_factories: + _v1_session_factories[local] = build_session_factory(get_v1_engine(local=local)) + return _v1_session_factories[local] + + +def clear_v1_session_factories() -> None: + """Forget cached factories after their process-owned engines are closed.""" + + _v1_session_factories.clear() diff --git a/tests/unit/data/test_orm_sessions.py b/tests/unit/data/test_orm_sessions.py index 82e1d5d8f..663fd7542 100644 --- a/tests/unit/data/test_orm_sessions.py +++ b/tests/unit/data/test_orm_sessions.py @@ -1,7 +1,74 @@ import pytest -from sqlalchemy import text +from sqlalchemy import create_engine, text +from sqlalchemy.orm import Session + +import policyengine_api.data.orm as orm_module +from policyengine_api.data.orm import ( + SessionManager, + build_session_factory, + build_sqlite_session_manager, + get_v1_session_factory, +) + + +def test_session_factory_creates_distinct_sessions_bound_to_one_engine(): + engine = create_engine("sqlite+pysqlite:///:memory:") + factory = build_session_factory(engine) + + first = factory() + second = factory() + try: + assert isinstance(first, Session) + assert isinstance(second, Session) + assert first is not second + assert first.get_bind() is engine + assert second.get_bind() is engine + assert first.expire_on_commit is False + finally: + first.close() + second.close() + + +def test_session_factory_begin_commits_and_rolls_back(): + engine = create_engine("sqlite+pysqlite:///:memory:") + factory = build_session_factory(engine) + with engine.begin() as connection: + connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) + + with factory.begin() as session: + session.execute(text("INSERT INTO item (id) VALUES (1)")) + + with pytest.raises(RuntimeError, match="stop"): + with factory.begin() as session: + session.execute(text("INSERT INTO item (id) VALUES (2)")) + raise RuntimeError("stop") -from policyengine_api.data.orm import SessionManager, build_sqlite_session_manager + with factory() as session: + assert session.scalars(text("SELECT id FROM item ORDER BY id")).all() == [1] + + +def test_runtime_factories_are_cached_and_separate(monkeypatch): + remote_engine = create_engine("sqlite+pysqlite:///:memory:") + local_engine = create_engine("sqlite+pysqlite:///:memory:") + + monkeypatch.setattr( + orm_module, + "get_v1_engine", + lambda *, local=False: local_engine if local else remote_engine, + ) + orm_module.clear_v1_session_factories() + + try: + remote = get_v1_session_factory() + local = get_v1_session_factory(local=True) + + assert remote is get_v1_session_factory() + assert local is get_v1_session_factory(local=True) + assert remote is not local + assert remote.kw["bind"] is remote_engine + assert local.kw["bind"] is local_engine + finally: + orm_module.clear_v1_session_factories() def test_session_manager_commits_successful_transaction(): From c5217d844aa96983b1f6c2c08872a8034ed60e2a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 21:53:30 +0300 Subject: [PATCH 48/89] test: add ORM-native v1 fixtures --- tests/unit/conftest.py | 25 ++++++++++++++++++++- tests/unit/data/test_orm_fixtures.py | 33 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/unit/data/test_orm_fixtures.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 336fe07f0..24733f3ab 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -10,6 +10,10 @@ from policyengine_api.constants import REPO from policyengine_api.data import PolicyEngineDatabase +from policyengine_api.data.orm import ( + clear_v1_session_factories, + get_v1_session_factory, +) class TestPolicyEngineDatabase(PolicyEngineDatabase): @@ -132,4 +136,23 @@ def override_database(test_db, monkeypatch): if hasattr(module, "local_database"): monkeypatch.setattr(module, "local_database", test_db) - yield test_db + clear_v1_session_factories() + try: + yield test_db + finally: + clear_v1_session_factories() + + +@pytest.fixture +def orm_session_factory(override_database): + """Return the runtime-style SQLAlchemy factory bound to the test schema.""" + + return get_v1_session_factory() + + +@pytest.fixture +def orm_session(orm_session_factory): + """Return one caller-owned SQLAlchemy Session for a unit test.""" + + with orm_session_factory() as session: + yield session diff --git a/tests/unit/data/test_orm_fixtures.py b/tests/unit/data/test_orm_fixtures.py new file mode 100644 index 000000000..771c1ecc7 --- /dev/null +++ b/tests/unit/data/test_orm_fixtures.py @@ -0,0 +1,33 @@ +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.v1_models import Household + + +def test_orm_session_factory_uses_mapped_models_and_python_json( + orm_session_factory, +): + assert isinstance(orm_session_factory, sessionmaker) + + with orm_session_factory.begin() as session: + household = Household( + country_id="uk", + label="Fixture household", + api_version="1.0.0", + household_json={"people": {"you": {"age": {"2025": 40}}}}, + household_hash="fixture-hash", + ) + session.add(household) + + with orm_session_factory() as session: + stored = session.scalar( + select(Household).where(Household.household_hash == "fixture-hash") + ) + + assert stored is not None + assert stored.household_json == {"people": {"you": {"age": {"2025": 40}}}} + + +def test_orm_session_fixture_is_a_live_session(orm_session): + assert isinstance(orm_session, Session) + assert orm_session.is_active From 07dd5b93475f3c57133997c8744427f9ce4280c9 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 22:02:42 +0300 Subject: [PATCH 49/89] refactor: use ORM sessions for policies and households --- policyengine_api/endpoints/household.py | 113 +++---- policyengine_api/routes/household_routes.py | 57 +++- policyengine_api/routes/policy_routes.py | 33 +- policyengine_api/services/economy_service.py | 34 +- .../services/household_service.py | 93 +++--- policyengine_api/services/policy_service.py | 132 ++++---- tests/contract/test_v1_route_contracts.py | 37 +- tests/fixtures/services/economy_service.py | 2 +- .../python/test_household_routes.py | 29 +- .../python/test_policy_service_old.py | 138 -------- .../endpoints/test_stage7_orm_endpoints.py | 120 +++---- .../test_direct_orm_policy_household.py | 63 ++++ tests/unit/services/test_household_service.py | 233 ++++--------- tests/unit/services/test_policy_service.py | 316 ++++++------------ .../services/test_stage7_dao_boundaries.py | 30 +- 15 files changed, 599 insertions(+), 831 deletions(-) delete mode 100644 tests/to_refactor/python/test_policy_service_old.py create mode 100644 tests/unit/services/test_direct_orm_policy_household.py diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index 4b1725cd1..7fb3677d1 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -1,4 +1,3 @@ -from policyengine_api.data.v1_daos import runtime_v1_unit_of_work import json from flask import Response, request from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS @@ -11,6 +10,10 @@ ) from policyengine_api.utils.payload_validators import validate_country from policyengine_core.errors import SituationParsingError +from sqlalchemy import select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ComputedHousehold, Household, Policy def get_countries(): @@ -111,49 +114,40 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Look in computed_households to see if already computed - with runtime_v1_unit_of_work(local=True).read() as daos: - row = daos.computed_households.get( - int(household_id), - int(policy_id), - country_id, - api_version=api_version, + with get_v1_session_factory(local=True)() as session: + computed_household = session.scalar( + select(ComputedHousehold).where( + ComputedHousehold.household_id == int(household_id), + ComputedHousehold.policy_id == int(policy_id), + ComputedHousehold.country_id == country_id, + ComputedHousehold.api_version == api_version, + ) ) - if row is not None: - result = dict( - policy_id=row["policy_id"], - household_id=row["household_id"], - country_id=row["country_id"], - api_version=row["api_version"], - computed_household_json=row["computed_household_json"], - status=row["status"], - ) - computed = result["computed_household_json"] - result["result"] = ( - json.loads(computed) if isinstance(computed, str) else computed - ) - del result["computed_household_json"] + if computed_household is not None: return dict( status="ok", message=None, - result=result["result"], + result=computed_household.computed_household_json, ) # Retrieve from the household table - with runtime_v1_unit_of_work().read() as daos: - row = daos.households.get(country_id, int(household_id)) - policy_row = daos.policies.get(country_id, int(policy_id)) - - if row is not None: - household = dict(row) - household_json = household["household_json"] - household["household_json"] = ( - json.loads(household_json) - if isinstance(household_json, str) - else household_json + with get_v1_session_factory()() as session: + household = session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == int(household_id), + ) ) - else: + policy = session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == int(policy_id), + ) + ) + + if household is None: response_body = dict( status="error", message=f"Household #{household_id} not found.", @@ -165,21 +159,16 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st ) # Add in any missing yearly variables - household["household_json"] = add_yearly_variables( - household["household_json"], country_id + household_json = add_yearly_variables( + household.household_json, + country_id, ) - deprecated_inputs = drop_deprecated_inputs(household["household_json"]) - household["household_json"] = deprecated_inputs.household + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household # Retrieve from the policy table - if policy_row is not None: - policy = dict(policy_row) - policy_json = policy["policy_json"] - policy["policy_json"] = ( - json.loads(policy_json) if isinstance(policy_json, str) else policy_json - ) - else: + if policy is None: response_body = dict( status="error", message=f"Policy #{policy_id} not found.", @@ -192,8 +181,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st country = get_countries().get(country_id) invalid_inputs_response = get_invalid_inputs_response( - household["household_json"], - policy["policy_json"], + household_json, + policy.policy_json, country, ) if invalid_inputs_response is not None: @@ -201,8 +190,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st try: result = country.calculate( - household["household_json"], - policy["policy_json"], + household_json, + policy.policy_json, household_id, policy_id, ) @@ -220,15 +209,23 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Store the result in the computed_household table - with runtime_v1_unit_of_work(local=True).transaction() as daos: - daos.computed_households.upsert( - country_id=country_id, - household_id=int(household_id), - policy_id=int(policy_id), - computed_household_json=result, - api_version=api_version, - status="complete", - ) + with get_v1_session_factory(local=True).begin() as session: + identity = (int(household_id), int(policy_id), country_id) + computed_household = session.get(ComputedHousehold, identity) + if computed_household is None: + computed_household = ComputedHousehold( + country_id=country_id, + household_id=int(household_id), + policy_id=int(policy_id), + computed_household_json=result, + api_version=api_version, + status="complete", + ) + session.add(computed_household) + else: + computed_household.computed_household_json = result + computed_household.api_version = api_version + computed_household.status = "complete" response_body = dict( status="ok", diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index d7420b39e..9c3b175b3 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -2,6 +2,8 @@ from werkzeug.exceptions import NotFound, BadRequest import json +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdService from policyengine_api.utils.payload_validators import ( validate_household_payload, @@ -12,6 +14,17 @@ household_service = HouseholdService() +def _serialize_household(household: Household) -> dict: + return { + "id": household.id, + "country_id": household.country_id, + "label": household.label, + "api_version": household.api_version, + "household_json": household.household_json, + "household_hash": household.household_hash, + } + + @household_bp.route("//household/", methods=["GET"]) @validate_country def get_household(country_id: str, household_id: int) -> Response: @@ -24,8 +37,10 @@ def get_household(country_id: str, household_id: int) -> Response: """ print(f"Got request for household {household_id} in country {country_id}") - household: dict | None = household_service.get_household(country_id, household_id) - if household is None: + with get_v1_session_factory()() as session: + household = household_service.get_household(session, country_id, household_id) + result = None if household is None else _serialize_household(household) + if result is None: raise NotFound(f"Household #{household_id} not found.") else: return Response( @@ -33,7 +48,7 @@ def get_household(country_id: str, household_id: int) -> Response: { "status": "ok", "message": None, - "result": household, + "result": result, } ), status=200, @@ -62,7 +77,14 @@ def post_household(country_id: str) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - household_id = household_service.create_household(country_id, household_json, label) + with get_v1_session_factory().begin() as session: + household = household_service.create_household( + session, + country_id, + household_json, + label, + ) + household_id = household.id return Response( json.dumps( @@ -102,14 +124,23 @@ def update_household(country_id: str, household_id: int) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - household: dict | None = household_service.get_household(country_id, household_id) - if household is None: - raise NotFound(f"Household #{household_id} not found.") - - # Next, update the household - updated_household: dict = household_service.update_household( - country_id, household_id, household_json, label - ) + with get_v1_session_factory().begin() as session: + household = household_service.get_household( + session, + country_id, + household_id, + ) + if household is None: + raise NotFound(f"Household #{household_id} not found.") + + updated_household = household_service.update_household( + session, + country_id, + household_id, + household_json, + label, + ) + updated_household_json = updated_household.household_json return Response( json.dumps( { @@ -117,7 +148,7 @@ def update_household(country_id: str, household_id: int) -> Response: "message": None, "result": { "household_id": household_id, - "household_json": updated_household["household_json"], + "household_json": updated_household_json, }, } ), diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 3fc88fbf4..c405245a5 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,6 +1,8 @@ from flask import Blueprint, Response, request import json +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy from policyengine_api.services.policy_service import PolicyService from werkzeug.exceptions import NotFound, BadRequest from policyengine_api.utils.payload_validators import ( @@ -12,6 +14,17 @@ policy_service = PolicyService() +def _serialize_policy(policy: Policy) -> dict: + return { + "id": policy.id, + "country_id": policy.country_id, + "label": policy.label, + "api_version": policy.api_version, + "policy_json": policy.policy_json, + "policy_hash": policy.policy_hash, + } + + @policy_bp.route("//policy/", methods=["GET"]) @validate_country def get_policy(country_id: str, policy_id: int | str) -> Response: @@ -30,13 +43,15 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: # Specifically cast policy_id to an integer policy_id = int(policy_id) - policy: dict | None = policy_service.get_policy(country_id, policy_id) + with get_v1_session_factory()() as session: + policy = policy_service.get_policy(session, country_id, policy_id) + result = None if policy is None else _serialize_policy(policy) - if policy is None: + if result is None: raise NotFound(f"Policy #{policy_id} not found.") return Response( - json.dumps({"status": "ok", "message": None, "result": policy}), + json.dumps({"status": "ok", "message": None, "result": result}), status=200, ) @@ -61,11 +76,13 @@ def set_policy(country_id: str) -> Response: label = payload.pop("label", None) policy_json = payload.pop("data", None) - policy_id, message, is_existing_policy = policy_service.set_policy( - country_id, - label, - policy_json, - ) + with get_v1_session_factory().begin() as session: + policy_id, message, is_existing_policy = policy_service.set_policy( + session, + country_id, + label, + policy_json, + ) response_body = dict( status="ok", diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index e8f512509..b3c8bdf0a 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -22,6 +22,7 @@ get_valid_state_codes, normalize_us_region, ) +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint @@ -235,6 +236,25 @@ class EconomyService: with other services to access their respective tables """ + def _get_policy_jsons( + self, + country_id: str, + baseline_policy_id: int, + reform_policy_id: int, + ) -> tuple[dict | None, dict | None]: + with get_v1_session_factory()() as session: + baseline = policy_service.get_policy_json( + session, + country_id, + baseline_policy_id, + ) + reform = policy_service.get_policy_json( + session, + country_id, + reform_policy_id, + ) + return baseline, reform + @staticmethod def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: """Accept ORM-decoded objects and legacy JSON text at the read boundary.""" @@ -411,12 +431,9 @@ def _build_budget_window_batch_payload( window_size: int, max_parallel: int, ) -> dict[str, Any]: - baseline_policy = policy_service.get_policy_json( + baseline_policy, reform_policy = self._get_policy_jsons( setup_options.country_id, setup_options.baseline_policy_id, - ) - reform_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.reform_policy_id, ) sim_config: SimulationOptions = self._setup_sim_options( @@ -904,11 +921,10 @@ def _handle_create_impact( self, setup_options: EconomicImpactSetupOptions, ) -> EconomicImpactResult: - baseline_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.baseline_policy_id - ) - reform_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.reform_policy_id + baseline_policy, reform_policy = self._get_policy_jsons( + setup_options.country_id, + setup_options.baseline_policy_id, + setup_options.reform_policy_id, ) sim_config: SimulationOptions = self._setup_sim_options( diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 945452b62..6363224b8 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,79 +1,66 @@ from __future__ import annotations -from contextlib import contextmanager +from sqlalchemy import select +from sqlalchemy.orm import Session from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import HouseholdDAO, V1UnitOfWork +from policyengine_api.data.v1_models import Household from policyengine_api.utils import hash_object class HouseholdService: - def __init__( - self, - households: HouseholdDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._households = households - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - - @contextmanager - def _repository(self, *, write: bool = False): - if self._households is not None: - yield self._households - return - boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as daos: - yield daos.households + """Household operations performed through a caller-owned ORM Session.""" - def get_household(self, country_id: str, household_id: int) -> dict | None: + def get_household( + self, + session: Session, + country_id: str, + household_id: int, + ) -> Household | None: if type(household_id) is not int or household_id < 0: raise Exception( f"Invalid household ID: {household_id}. Must be a positive integer." ) - with self._repository() as households: - return households.get(country_id, household_id) + return session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) def create_household( self, + session: Session, country_id: str, household_json: dict, label: str | None, - ) -> int: - with self._repository(write=True) as households: - return households.create( - country_id, - label, - household_json, - hash_object(household_json), - COUNTRY_PACKAGE_VERSIONS.get(country_id), - ) + ) -> Household: + household = Household( + country_id=country_id, + label=label, + household_json=household_json, + household_hash=hash_object(household_json), + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + ) + session.add(household) + session.flush() + return household def update_household( self, + session: Session, country_id: str, household_id: int, household_json: dict, - label: str, - ) -> dict: - with self._repository(write=True) as households: - updated = households.update( - country_id, - household_id, - label, - household_json, - hash_object(household_json), - COUNTRY_PACKAGE_VERSIONS.get(country_id), + label: str | None, + ) -> Household: + household = self.get_household(session, country_id, household_id) + if household is None: + raise LookupError( + f"Household #{household_id} not found for country {country_id}." ) - if not updated: - raise LookupError( - f"Household #{household_id} not found for country {country_id}." - ) - return households.get(country_id, household_id) + household.label = label + household.household_json = household_json + household.household_hash = hash_object(household_json) + household.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) + return household diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index a288b3f51..96ffc40c2 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,38 +1,17 @@ from __future__ import annotations -import json -from contextlib import contextmanager +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import PolicyDAO, V1UnitOfWork +from policyengine_api.data.v1_models import Policy from policyengine_api.utils import hash_object class PolicyService: - def __init__( - self, - policies: PolicyDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._policies = policies - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - - @contextmanager - def _repository(self, *, write: bool = False): - if self._policies is not None: - yield self._policies - return - boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as daos: - yield daos.policies + """Policy operations performed through a caller-owned ORM Session.""" @staticmethod def _validate_policy_id(policy_id: int) -> None: @@ -41,57 +20,94 @@ def _validate_policy_id(policy_id: int) -> None: f"Invalid policy ID: {policy_id}. Must be a positive integer." ) - def get_policy(self, country_id: str, policy_id: int) -> dict | None: + def get_policy( + self, + session: Session, + country_id: str, + policy_id: int, + ) -> Policy | None: self._validate_policy_id(policy_id) if not country_id: raise ValueError("country_id cannot be empty or None") - with self._repository() as policies: - return policies.get(country_id, policy_id) + return session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, + ) + ) - def get_policy_json(self, country_id: str, policy_id: int) -> str | None: - self._validate_policy_id(policy_id) - with self._repository() as policies: - policy = policies.get(country_id, policy_id) - if policy is None: - return None - value = policy["policy_json"] - return value if isinstance(value, str) else json.dumps(value) + def get_policy_json( + self, + session: Session, + country_id: str, + policy_id: int, + ) -> Any | None: + policy = self.get_policy(session, country_id, policy_id) + return None if policy is None else policy.policy_json def set_policy( - self, country_id: str, label: str, policy_json: dict + self, + session: Session, + country_id: str, + label: str | None, + policy_json: dict, ) -> tuple[int, str, bool]: country_id = country_id.lower() if country_id not in COUNTRY_PACKAGE_VERSIONS: raise ValueError(f"Invalid country_id: {country_id}") policy_hash = hash_object(policy_json) - with self._repository(write=True) as policies: - existing = policies.find_unique(country_id, policy_hash, label or None) - if existing: - return existing["id"], "Policy already exists", True + existing = self._get_unique_policy_with_label( + session, + country_id, + policy_hash, + label or None, + ) + if existing is not None: + return existing.id, "Policy already exists", True - policy_id = policies.create( - country_id, - label, - policy_json, - policy_hash, - COUNTRY_PACKAGE_VERSIONS[country_id], - ) - return policy_id, "Policy created", False + policy = Policy( + country_id=country_id, + label=label, + policy_json=policy_json, + policy_hash=policy_hash, + api_version=COUNTRY_PACKAGE_VERSIONS[country_id], + ) + session.add(policy) + session.flush() + return policy.id, "Policy created", False def _create_new_policy( self, + session: Session, country_id: str, policy_json: dict, policy_hash: str, label: str | None, api_version: str, - ) -> None: - with self._repository(write=True) as policies: - policies.create(country_id, label, policy_json, policy_hash, api_version) + ) -> Policy: + policy = Policy( + country_id=country_id, + label=label, + policy_json=policy_json, + policy_hash=policy_hash, + api_version=api_version, + ) + session.add(policy) + session.flush() + return policy def _get_unique_policy_with_label( - self, country_id: str, policy_hash: str, label: str - ) -> dict | None: - with self._repository() as policies: - return policies.find_unique(country_id, policy_hash, label or None) + self, + session: Session, + country_id: str, + policy_hash: str, + label: str | None, + ) -> Policy | None: + return session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.policy_hash == policy_hash, + Policy.label == label, + ) + ) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 4becedf73..793a1a37b 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -9,6 +9,7 @@ from policyengine_api.endpoints.household import get_calculate from policyengine_api.endpoints.policy import get_policy_search +from policyengine_api.data.v1_models import Household, Policy from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp @@ -219,7 +220,14 @@ def _patched_route_dependencies(): stack.enter_context( patch( "policyengine_api.routes.policy_routes.policy_service.get_policy", - return_value={"id": 22, "label": "Current law", "policy_json": {}}, + return_value=Policy( + id=22, + country_id="us", + label="Current law", + api_version="1", + policy_json={}, + policy_hash="hash-22", + ), ) ) stack.enter_context( @@ -237,19 +245,40 @@ def _patched_route_dependencies(): stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.create_household", - return_value=456, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={}, + household_hash="hash-456", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.get_household", - return_value={"id": 456, "label": "Empty household", "household_json": {}}, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={}, + household_hash="hash-456", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.update_household", - return_value={"household_json": {"people": {"you": {}}}}, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={"people": {"you": {}}}, + household_hash="hash-456", + ), ) ) stack.enter_context( diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index e4061efe3..c3af5cbcd 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -97,7 +97,7 @@ def mock_policyengine_version(): def mock_policy_service(): """Mock PolicyService with get_policy_json method.""" mock_service = MagicMock() - mock_service.get_policy_json.side_effect = lambda country_id, policy_id: ( + mock_service.get_policy_json.side_effect = lambda session, country_id, policy_id: ( MOCK_REFORM_POLICY_JSON if policy_id == MOCK_POLICY_ID else MOCK_BASELINE_POLICY_JSON diff --git a/tests/to_refactor/python/test_household_routes.py b/tests/to_refactor/python/test_household_routes.py index b40f70f8c..dcf394195 100644 --- a/tests/to_refactor/python/test_household_routes.py +++ b/tests/to_refactor/python/test_household_routes.py @@ -1,6 +1,7 @@ import json -from unittest.mock import patch +from unittest.mock import ANY, patch +from policyengine_api.data.v1_models import Household from tests.to_refactor.fixtures.to_refactor_household_fixtures import ( valid_request_body, valid_db_row, @@ -12,10 +13,9 @@ class TestGetHousehold: def test_get_existing_household(self, rest_client, mock_database): """Test getting an existing household.""" - mock_database.get_household.return_value = { - **valid_db_row, - "household_json": valid_request_body["data"], - } + mock_database.get_household.return_value = Household( + **{**valid_db_row, "household_json": valid_request_body["data"]} + ) # Make request response = rest_client.get("/us/household/1") @@ -47,7 +47,9 @@ def test_get_household_invalid_id(self, rest_client): class TestCreateHousehold: def test_create_household_success(self, rest_client, mock_database): """Test successfully creating a new household.""" - mock_database.create_household.return_value = 1 + mock_database.create_household.return_value = Household( + **{**valid_db_row, "id": 1, "household_json": valid_request_body["data"]} + ) response = rest_client.post( "/us/household", @@ -96,10 +98,9 @@ def test_create_household_invalid_label(self, rest_client): class TestUpdateHousehold: def test_update_household_success(self, rest_client, mock_database): """Test successfully updating an existing household.""" - mock_database.get_household.return_value = { - **valid_db_row, - "household_json": valid_request_body["data"], - } + mock_database.get_household.return_value = Household( + **{**valid_db_row, "household_json": valid_request_body["data"]} + ) updated_household = {"people": {"person1": {"age": 31, "income": 55000}}} @@ -107,10 +108,9 @@ def test_update_household_success(self, rest_client, mock_database): "data": updated_household, "label": valid_request_body["label"], } - mock_database.update_household.return_value = { - **valid_db_row, - "household_json": updated_household, - } + mock_database.update_household.return_value = Household( + **{**valid_db_row, "household_json": updated_household} + ) response = rest_client.put( "/us/household/1", @@ -124,6 +124,7 @@ def test_update_household_success(self, rest_client, mock_database): assert data["result"]["household_id"] == 1 assert data["result"]["household_json"] == updated_data["data"] mock_database.update_household.assert_called_once_with( + ANY, "us", 1, updated_household, diff --git a/tests/to_refactor/python/test_policy_service_old.py b/tests/to_refactor/python/test_policy_service_old.py deleted file mode 100644 index f3ba129e1..000000000 --- a/tests/to_refactor/python/test_policy_service_old.py +++ /dev/null @@ -1,138 +0,0 @@ -import json -from unittest.mock import MagicMock - -from assertpy import assert_that -import pytest - -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.services.policy_service import PolicyService - - -@pytest.fixture -def policies(): - return MagicMock() - - -@pytest.fixture -def sample_policy_data(): - return { - "id": 1, - "country_id": "us", - "policy_json": {"param": "value"}, - "policy_hash": "hash123", - "label": "test_policy", - "api_version": "1.0.0", - } - - -@pytest.fixture -def policy_service(policies): - return PolicyService(policies) - - -class TestPolicyService: - a_test_policy_id = 8 - - def test_get_policy_success(self, policy_service, policies, sample_policy_data): - policies.get.return_value = sample_policy_data - - result = policy_service.get_policy("us", self.a_test_policy_id) - - assert_that(result).contains_entry({"policy_json": {"param": "value"}}) - policies.get.assert_called_once_with("us", self.a_test_policy_id) - - def test_get_policy_not_found(self, policy_service, policies): - policies.get.return_value = None - - assert policy_service.get_policy("us", 999) is None - policies.get.assert_called_once_with("us", 999) - - def test_get_policy_json(self, policy_service, policies, sample_policy_data): - policies.get.return_value = sample_policy_data - - result = policy_service.get_policy_json("us", self.a_test_policy_id) - - assert result == json.dumps(sample_policy_data["policy_json"]) - policies.get.assert_called_once_with("us", self.a_test_policy_id) - - def test_set_policy_new(self, policy_service, policies): - policies.find_unique.return_value = None - policies.create.return_value = 10 - test_policy = {"param": "value"} - - policy_id, message, exists = policy_service.set_policy( - "us", "new_policy", test_policy - ) - - assert (policy_id, message, exists) == (10, "Policy created", False) - policies.find_unique.assert_called_once() - policy_hash = policies.find_unique.call_args.args[1] - policies.create.assert_called_once_with( - "us", - "new_policy", - test_policy, - policy_hash, - COUNTRY_PACKAGE_VERSIONS["us"], - ) - - def test_set_policy_existing(self, policy_service, policies, sample_policy_data): - policies.find_unique.return_value = sample_policy_data - - result = policy_service.set_policy( - "us", - sample_policy_data["label"], - sample_policy_data["policy_json"], - ) - - assert result == (sample_policy_data["id"], "Policy already exists", True) - policies.create.assert_not_called() - - def test_get_unique_policy_with_label( - self, policy_service, policies, sample_policy_data - ): - policies.find_unique.return_value = sample_policy_data - - result = policy_service._get_unique_policy_with_label( - "us", - sample_policy_data["policy_hash"], - sample_policy_data["label"], - ) - - assert result == sample_policy_data - policies.find_unique.assert_called_once_with( - "us", - sample_policy_data["policy_hash"], - sample_policy_data["label"], - ) - - def test_get_unique_policy_with_null_label(self, policy_service, policies): - policies.find_unique.return_value = None - - result = policy_service._get_unique_policy_with_label("us", "hash123", None) - - assert result is None - policies.find_unique.assert_called_once_with("us", "hash123", None) - - @pytest.mark.parametrize( - ("error_method", "repository_method"), - [ - ("get_policy", "get"), - ("get_policy_json", "get"), - ("set_policy", "find_unique"), - ("_get_unique_policy_with_label", "find_unique"), - ], - ) - def test_error_handling( - self, policy_service, policies, error_method, repository_method - ): - getattr(policies, repository_method).side_effect = Exception("Database error") - - with pytest.raises(Exception, match="Database error"): - if error_method == "get_policy": - policy_service.get_policy("us", 1) - elif error_method == "get_policy_json": - policy_service.get_policy_json("us", 1) - elif error_method == "set_policy": - policy_service.set_policy("us", "label", {}) - else: - policy_service._get_unique_policy_with_label("us", "hash", "label") diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py index 674d8566b..5be23c788 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -1,13 +1,13 @@ -import json -from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import Mock, patch -import pytest from flask import Flask +from sqlalchemy import select +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.data.orm import build_sqlite_session_manager from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.data.v1_models import ComputedHousehold, Household, Policy from policyengine_api.endpoints.household import get_household_under_policy from policyengine_api.endpoints.policy import ( get_user_policy, @@ -23,85 +23,60 @@ def _unit_of_work() -> V1UnitOfWork: return V1UnitOfWork(manager) -def _daos_unit_of_work(daos): - @contextmanager - def boundary(): - yield daos - - return SimpleNamespace(read=boundary, transaction=boundary) - - -@pytest.mark.parametrize( - "stored_result", - [ - {"people": {"you": {"net_income": {"2026": 42}}}}, - json.dumps({"people": {"you": {"net_income": {"2026": 42}}}}), - ], -) -def test_household_under_policy_returns_cached_json_objects_and_legacy_strings( - stored_result, -): - computed_households = Mock() - computed_households.get.return_value = { - "household_id": 1, - "policy_id": 2, - "country_id": "us", - "api_version": "1", - "computed_household_json": stored_result, - "status": "complete", - } - local_uow = _daos_unit_of_work( - SimpleNamespace(computed_households=computed_households) - ) +def test_household_under_policy_returns_cached_json_object(orm_session_factory): + stored_result = {"people": {"you": {"net_income": {"2026": 42}}}} + with orm_session_factory.begin() as session: + session.add( + ComputedHousehold( + household_id=1, + policy_id=2, + country_id="us", + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + computed_household_json=stored_result, + status="complete", + ) + ) with patch( - "policyengine_api.endpoints.household.runtime_v1_unit_of_work", - return_value=local_uow, - ) as runtime_uow: + "policyengine_api.endpoints.household.get_v1_session_factory", + return_value=orm_session_factory, + ): response = get_household_under_policy("us", "1", "2") assert response["result"] == {"people": {"you": {"net_income": {"2026": 42}}}} - runtime_uow.assert_called_once_with(local=True) -def test_household_under_policy_calculates_and_caches_json_as_an_object(): - computed_households = Mock() - computed_households.get.return_value = None - local_uow = _daos_unit_of_work( - SimpleNamespace(computed_households=computed_households) - ) - remote_uow = _daos_unit_of_work( - SimpleNamespace( - households=SimpleNamespace( - get=Mock( - return_value={ - "id": 1, - "country_id": "us", - "household_json": {"people": {"you": {}}}, - } - ) - ), - policies=SimpleNamespace( - get=Mock( - return_value={ - "id": 2, - "country_id": "us", - "policy_json": {"gov.example.parameter": 1}, - } - ) - ), +def test_household_under_policy_calculates_and_caches_json_as_an_object( + orm_session_factory, +): + with orm_session_factory.begin() as session: + session.add_all( + [ + Household( + id=1, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + household_json={"people": {"you": {}}}, + household_hash="household-hash", + ), + Policy( + id=2, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + policy_json={"gov.example.parameter": 1}, + policy_hash="policy-hash", + ), + ] ) - ) calculated = {"people": {"you": {"net_income": {"2026": 42}}}} country = SimpleNamespace(calculate=Mock(return_value=calculated)) - def select_uow(*, local=False): - return local_uow if local else remote_uow - with ( patch( - "policyengine_api.endpoints.household.runtime_v1_unit_of_work", - side_effect=select_uow, + "policyengine_api.endpoints.household.get_v1_session_factory", + return_value=orm_session_factory, ), patch( "policyengine_api.endpoints.household.add_yearly_variables", @@ -132,10 +107,9 @@ def select_uow(*, local=False): "1", "2", ) - assert ( - computed_households.upsert.call_args.kwargs["computed_household_json"] - is calculated - ) + with orm_session_factory() as session: + cached = session.scalar(select(ComputedHousehold)) + assert cached.computed_household_json == calculated def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): diff --git a/tests/unit/services/test_direct_orm_policy_household.py b/tests/unit/services/test_direct_orm_policy_household.py new file mode 100644 index 000000000..4a1a1b640 --- /dev/null +++ b/tests/unit/services/test_direct_orm_policy_household.py @@ -0,0 +1,63 @@ +from sqlalchemy import select + +from policyengine_api.data.v1_models import Household, Policy +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService + + +def test_policy_service_reads_and_writes_mapped_models(orm_session, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "policy-hash", + ) + service = PolicyService() + + policy_id, message, existed = service.set_policy( + orm_session, + "us", + "Direct ORM policy", + {"gov.example.rate": {"2026": 0.2}}, + ) + orm_session.commit() + policy = service.get_policy(orm_session, "us", policy_id) + + assert isinstance(policy, Policy) + assert policy.policy_json == {"gov.example.rate": {"2026": 0.2}} + assert message == "Policy created" + assert existed is False + + +def test_household_service_reads_updates_and_writes_mapped_models( + orm_session, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + service = HouseholdService() + payload = {"people": {"you": {"age": {"2026": 40}}}} + + household = service.create_household( + orm_session, + "us", + payload, + "Direct ORM household", + ) + orm_session.commit() + stored = orm_session.scalar(select(Household).where(Household.id == household.id)) + + assert household is stored + assert stored.household_json == payload + + updated = service.update_household( + orm_session, + "us", + stored.id, + {"people": {"you": {"age": {"2026": 41}}}}, + "Updated", + ) + + assert isinstance(updated, Household) + assert updated.label == "Updated" + assert updated.household_json["people"]["you"]["age"]["2026"] == 41 diff --git a/tests/unit/services/test_household_service.py b/tests/unit/services/test_household_service.py index 97d25b58e..c880fe650 100644 --- a/tests/unit/services/test_household_service.py +++ b/tests/unit/services/test_household_service.py @@ -1,193 +1,90 @@ import pytest -import json -from unittest.mock import MagicMock -import re +from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdService -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS - from tests.fixtures.services.household_fixtures import ( - valid_request_body, valid_db_row, - valid_hash_value, - existing_household_record, - mock_hash_object, + valid_request_body, ) -service = HouseholdService() - - -class TestGetHousehold: - def test_get_household_given_existing_record( - self, test_db, existing_household_record - ): - # GIVEN an existing record... (included as fixture) - - # WHEN we call get_household for this record... - result = service.get_household(valid_db_row["country_id"], valid_db_row["id"]) - - valid_household_json = valid_request_body["data"] - - # THEN the result should be the expected household data - assert result["household_json"] == valid_household_json - - def test_get_household_given_nonexistent_record(self, test_db): - # GIVEN an empty database (this is created by default)... - - # WHEN we call get_household for a nonexistent record... - NO_SUCH_RECORD_ID = 999 - result = service.get_household("us", NO_SUCH_RECORD_ID) +pytest_plugins = ["tests.fixtures.services.household_fixtures"] - # THEN the result should be None - assert result is None - def test_get_household_given_str_id(self, test_db): - # GIVEN an invalid ID... - - INVALID_RECORD_ID = "invalid" - - with pytest.raises( - Exception, - match=f"Invalid household ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_household with the invalid ID... - # THEN an exception should be raised - service.get_household("us", INVALID_RECORD_ID) - - def test_get_household_given_negative_int_id(self, test_db): - # GIVEN an invalid ID... - INVALID_RECORD_ID = -1 - - with pytest.raises( - Exception, - match=f"Invalid household ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_household with the invalid ID... - # THEN an exception should be raised - service.get_household("us", INVALID_RECORD_ID) - - -class TestCreateHousehold: - service = HouseholdService() - - def test_create_household_given_valid_data(self, test_db): - def fetch_created_record(): - row = test_db.query( - "SELECT * FROM household", - ).fetchone() - return row - - # GIVEN valid household data and an empty database... - # WHEN we call create_household with this data... - country_id = "us" - valid_json = valid_request_body["data"] - valid_label = valid_request_body["label"] +service = HouseholdService() - test_id = service.create_household(country_id, valid_json, valid_label) - # THEN there should only be one record, and if we re-fetch it, - # it should match the data we provided - test_row = fetch_created_record() +def test_get_household_returns_mapped_entity(orm_session, existing_household_record): + household = service.get_household( + orm_session, + valid_db_row["country_id"], + valid_db_row["id"], + ) - valid_json_in_db = json.dumps(valid_request_body["data"]) - valid_label_in_db = valid_request_body["label"] + assert isinstance(household, Household) + assert household.household_json == valid_request_body["data"] - assert test_id == test_row["id"] - assert test_row["household_json"] == valid_json_in_db - assert test_row["label"] == valid_label_in_db - def test_create_household_given_missing_data(self, test_db): - # GIVEN an empty database... +def test_get_household_returns_none_for_missing_entity(orm_session): + assert service.get_household(orm_session, "us", 999) is None - # WHEN we call create_household with missing required data... - country_id = "us" - valid_label = valid_request_body["label"] - with pytest.raises( - Exception, - match=re.escape( - "HouseholdService.create_household() missing 1 required positional argument: 'household_json'" - ), - ): - # THEN an exception should be raised - service.create_household(country_id, label=valid_label) +@pytest.mark.parametrize("household_id", ["invalid", -1]) +def test_get_household_rejects_invalid_id(orm_session, household_id): + with pytest.raises(Exception, match="Invalid household ID"): + service.get_household(orm_session, "us", household_id) -class TestUpdateHousehold: - def test_update_household_given_existing_record( - self, test_db, mock_hash_object, existing_household_record - ): - def fetch_updated_record(): - row = test_db.query( - "SELECT * FROM household WHERE id = ?", (valid_db_row["id"],) - ).fetchone() - return row +def test_create_household_adds_mapped_entity(orm_session, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "some-hash", + ) - # GIVEN an existing record...(included as fixture) + household = service.create_household( + orm_session, + "us", + valid_request_body["data"], + valid_request_body["label"], + ) - # WHEN we call update_household for this record's label and fill other necessary info... - test_update_label = "Updated Household" + assert isinstance(household, Household) + assert household.id is not None + assert household.household_json == valid_request_body["data"] - existing_country_id = valid_db_row["country_id"] - existing_record_id = valid_db_row["id"] - existing_data = valid_db_row["household_json"] +def test_update_household_mutates_mapped_entity( + orm_session, + existing_household_record, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "updated-hash", + ) + + household = service.update_household( + orm_session, + "us", + valid_db_row["id"], + {"people": {"person1": {"age": 31}}}, + "Updated Household", + ) + + assert household.label == "Updated Household" + assert household.household_hash == "updated-hash" + assert household.household_json == {"people": {"person1": {"age": 31}}} + + +def test_update_household_rejects_missing_or_cross_country_entity( + orm_session, + existing_household_record, +): + with pytest.raises(LookupError): service.update_household( - existing_country_id, - existing_record_id, - existing_data, - test_update_label, + orm_session, + "uk", + valid_db_row["id"], + {}, + "Wrong country", ) - - # THEN the database should be updated with the new data - test_row = fetch_updated_record() - assert test_row["label"] == test_update_label - - def test_update_household_given_nonexistent_record(self, test_db): - # GIVEN an empty database... - - # WHEN we call update_household for a nonexistent record... - NO_SUCH_RECORD_ID = 999 - - existing_country_id = valid_db_row["country_id"] - existing_data = valid_db_row["household_json"] - existing_label = valid_db_row["label"] - - # THEN update_household raises LookupError because the id - # does not exist for this country (issue #3447). - with pytest.raises(LookupError): - service.update_household( - existing_country_id, - NO_SUCH_RECORD_ID, - existing_data, - existing_label, - ) - - def test_update_household_rejects_cross_country_id( - self, test_db, mock_hash_object, existing_household_record - ): - """Regression for issue #3447. - - An existing US household must not be overwritten by a request - that targets the same numeric id under a different country. - """ - - existing_record_id = valid_db_row["id"] - existing_data = valid_db_row["household_json"] - - with pytest.raises(LookupError): - service.update_household( - "uk", # wrong country - existing_record_id, - existing_data, - "Attacker label", - ) - - # The original US row must be untouched. - row = test_db.query( - "SELECT label, country_id FROM household WHERE id = ?", - (existing_record_id,), - ).fetchone() - assert row["country_id"] == "us" - assert row["label"] == valid_db_row["label"] diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index 0a4c94a94..1e2d48553 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -1,223 +1,113 @@ import pytest -import json -from unittest.mock import MagicMock -from policyengine_api.services.policy_service import PolicyService +from sqlalchemy.exc import SQLAlchemyError -from tests.fixtures.services.policy_service import valid_hash_value, valid_policy_data +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.v1_models import Policy +from policyengine_api.services.policy_service import PolicyService +from tests.fixtures.services.policy_service import valid_policy_data pytest_plugins = ["tests.fixtures.services.policy_service"] + service = PolicyService() -class TestGetPolicy: - def test_get_policy_given_existing_record(self, test_db, existing_policy_record): - # GIVEN an existing record... (included as fixture) - - # WHEN we call get_policy for this record... - result = service.get_policy( - valid_policy_data["country_id"], valid_policy_data["id"] - ) - - expected_result = { - "id": valid_policy_data["id"], - "country_id": valid_policy_data["country_id"], - "label": valid_policy_data["label"], - "api_version": valid_policy_data["api_version"], - "policy_json": json.loads(valid_policy_data["policy_json"]), - "policy_hash": valid_policy_data["policy_hash"], - } - - # THEN the result should contain the expected policy data - assert result == expected_result +def test_get_policy_returns_mapped_entity(orm_session, existing_policy_record): + policy = service.get_policy( + orm_session, + valid_policy_data["country_id"], + valid_policy_data["id"], + ) + + assert isinstance(policy, Policy) + assert policy.id == valid_policy_data["id"] + assert policy.policy_json == { + "gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433} + } + + +def test_get_policy_returns_none_for_missing_entity(orm_session): + assert service.get_policy(orm_session, "us", 999) is None + + +@pytest.mark.parametrize("policy_id", ["invalid", -1]) +def test_get_policy_rejects_invalid_id(orm_session, policy_id): + with pytest.raises(Exception, match="Invalid policy ID"): + service.get_policy(orm_session, "us", policy_id) + + +@pytest.mark.parametrize("country_id", ["", None]) +def test_get_policy_rejects_empty_country(orm_session, country_id): + with pytest.raises(ValueError, match="country_id cannot be empty or None"): + service.get_policy(orm_session, country_id, 1) + + +def test_get_policy_json_returns_python_object(orm_session, existing_policy_record): + result = service.get_policy_json(orm_session, "us", valid_policy_data["id"]) + + assert result == { + "gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433} + } + + +def test_set_policy_adds_mapped_entity(orm_session, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "new-hash", + ) + + policy_id, message, exists = service.set_policy( + orm_session, + "US", + "New policy", + {"parameter": 1}, + ) + + policy = service.get_policy(orm_session, "us", policy_id) + assert policy.policy_json == {"parameter": 1} + assert policy.api_version == COUNTRY_PACKAGE_VERSIONS["us"] + assert message == "Policy created" + assert exists is False + + +def test_set_policy_returns_existing_mapped_entity( + orm_session, + existing_policy_record, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: valid_policy_data["policy_hash"], + ) + + policy_id, message, exists = service.set_policy( + orm_session, + "us", + None, + {}, + ) + + assert policy_id == valid_policy_data["id"] + assert message == "Policy already exists" + assert exists is True + + +def test_set_policy_rejects_invalid_country(orm_session): + with pytest.raises(ValueError, match="Invalid country_id: xx"): + service.set_policy(orm_session, "xx", "Policy", {}) + + +def test_set_policy_propagates_flush_failure(orm_session, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "new-hash", + ) + monkeypatch.setattr( + orm_session, + "flush", + lambda: (_ for _ in ()).throw(SQLAlchemyError("insert failed")), + ) - def test_get_policy_given_nonexistent_record(self, test_db): - # GIVEN an empty database (this is created by default) - - # WHEN we call get_policy for a nonexistent record - NO_SUCH_RECORD_ID = 999 - result = service.get_policy(valid_policy_data["country_id"], NO_SUCH_RECORD_ID) - - # THEN the result should be None - assert result is None - - def test_get_policy_given_str_id(self): - # GIVEN an invalid ID - INVALID_RECORD_ID = "invalid" - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy with the invalid ID - # THEN an exception should be raised - service.get_policy(valid_policy_data["country_id"], INVALID_RECORD_ID) - - def test_get_policy_given_negative_int_id(self): - # GIVEN an invalid ID - INVALID_RECORD_ID = -1 - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy with the invalid ID - # THEN an exception should be raised - service.get_policy(valid_policy_data["country_id"], INVALID_RECORD_ID) - - def test_get_policy_given_invalid_country_id(self): - # GIVEN an invalid country_id - INVALID_COUNTRY_ID = "xx" # Unsupported country code - - # WHEN we call get_policy with the invalid country_id - result = service.get_policy(INVALID_COUNTRY_ID, valid_policy_data["id"]) - - # THEN the result should be None or raise an exception - assert result is None - - def test_get_policy_given_empty_string_country_id(self): - # GIVEN an empty string as country_id - EMPTY_COUNTRY_ID = "" - - # WHEN we call get_policy with country_id = "" - with pytest.raises( - Exception, - match="country_id cannot be empty or None", - ): - # THEN an exception should be raised - service.get_policy(EMPTY_COUNTRY_ID, valid_policy_data["id"]) - - def test_get_policy_given_none_country_id(self): - # GIVEN a country_id of None - NONE_COUNTRY_ID = None - - # WHEN we call get_policy with country_id = None - with pytest.raises( - Exception, - match="country_id cannot be empty or None", - ): - # THEN an exception should be raised - service.get_policy(NONE_COUNTRY_ID, valid_policy_data["id"]) - - -class TestGetPolicyJson: - def test_get_policy_json_given_existing_record( - self, test_db, existing_policy_record - ): - # GIVEN an existing record... (included as fixture) - - # WHEN we call get_policy_json for this record... - result = service.get_policy_json( - valid_policy_data["country_id"], valid_policy_data["id"] - ) - - valid_policy_json = valid_policy_data["policy_json"] - - # THEN result should be the expected policy json - assert result == valid_policy_json - - def test_get_policy_json_given_nonexisting_record(self, test_db): - # GIVEN an empty database... (created by default) - - # WHEN we call get_policy_json for nonexistent record... - NO_SUCH_RECORD_ID = 999 - result = service.get_policy_json("us", NO_SUCH_RECORD_ID) - - # THEN result should be None - assert result is None - - def test_get_policy_json_given_str_id(self, test_db): - # GIVEN an invalid ID... - - INVALID_RECORD_ID = "invalid" - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy_json with the invalid ID... - # THEN an exception should be raised - service.get_policy_json("us", INVALID_RECORD_ID) - - def test_get_policy_json_given_negative_int_id(self, test_db): - # GIVEN an invalid ID... - - INVALID_RECORD_ID = -1 - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy_json with the invalid ID... - # THEN an exception should be raised - service.get_policy_json("us", INVALID_RECORD_ID) - - -class TestSetPolicy: - def test_set_policy_new(self, mock_hash_object): - policies = MagicMock() - policies.find_unique.return_value = None - policies.create.return_value = 12 - isolated_service = PolicyService(policies) - test_policy = {"param": "value"} - test_label = "new_policy" - test_country_id = "us" - from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS - - policy_id, message, exists = isolated_service.set_policy( - test_country_id, test_label, test_policy - ) - assert policy_id == 12 - assert message == "Policy created" - assert exists is False - policies.find_unique.assert_called_once_with( - test_country_id, valid_hash_value, test_label - ) - policies.create.assert_called_once_with( - test_country_id, - test_label, - test_policy, - valid_hash_value, - COUNTRY_PACKAGE_VERSIONS[test_country_id], - ) - - def test_set_policy_existing(self, mock_hash_object): - policies = MagicMock() - policies.find_unique.return_value = {"id": 11} - isolated_service = PolicyService(policies) - policy_id, message, exists = isolated_service.set_policy("us", None, {}) - assert policy_id == 11 - assert message == "Policy already exists" - assert exists is True - policies.create.assert_not_called() - - def test_set_policy_given_database_insert_failure(self, mock_hash_object): - policies = MagicMock() - policies.find_unique.return_value = None - policies.create.side_effect = Exception("Database insertion failed") - with pytest.raises(Exception, match="Database insertion failed"): - PolicyService(policies).set_policy("us", "test_policy", {}) - - def test_set_policy_given_invalid_country_id(self, mock_hash_object): - # GIVEN an invalid country_id - INVALID_COUNTRY_ID = "xx" # Unsupported country code - test_policy = {"param": "value"} - test_label = "test_policy" - - # WHEN we call set_policy with an invalid country_id - with pytest.raises( - ValueError, match=f"Invalid country_id: {INVALID_COUNTRY_ID}" - ): - # THEN an exception should be raised - service.set_policy(INVALID_COUNTRY_ID, test_label, test_policy) - - def test_set_policy_given_empty_label(self, mock_hash_object): - policies = MagicMock() - policies.find_unique.return_value = None - policies.create.return_value = 13 - policy_id, message, exists = PolicyService(policies).set_policy("us", "", {}) - assert policy_id == 13 - assert message == "Policy created" - assert exists is False - policies.find_unique.assert_called_once_with("us", valid_hash_value, None) + with pytest.raises(SQLAlchemyError, match="insert failed"): + service.set_policy(orm_session, "us", "Policy", {}) diff --git a/tests/unit/services/test_stage7_dao_boundaries.py b/tests/unit/services/test_stage7_dao_boundaries.py index 93cec54b7..ef7124fc5 100644 --- a/tests/unit/services/test_stage7_dao_boundaries.py +++ b/tests/unit/services/test_stage7_dao_boundaries.py @@ -2,8 +2,6 @@ import pytest -from policyengine_api.services.household_service import HouseholdService -from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.user_service import UserService @@ -20,30 +18,20 @@ def test_migrated_services_do_not_issue_queries_directly(module_name): assert "from policyengine_api.data import database" not in source -class StubHouseholds: - def get(self, country_id, household_id): - return {"country_id": country_id, "id": household_id} - - -class StubPolicies: - def get(self, country_id, policy_id): - return { - "country_id": country_id, - "id": policy_id, - "policy_json": {"already": "decoded"}, - } - - class StubUsers: def get_profile(self, *, user_id=None, auth0_id=None): return {"user_id": user_id, "auth0_id": auth0_id} -def test_services_accept_explicit_daos_for_isolated_parity_tests(): - assert HouseholdService(StubHouseholds()).get_household("us", 3)["id"] == 3 - assert PolicyService(StubPolicies()).get_policy("us", 4)["policy_json"] == { - "already": "decoded" - } +@pytest.mark.parametrize("module_name", ["household_service.py", "policy_service.py"]) +def test_migrated_services_use_sessions_and_mapped_models(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert "from sqlalchemy.orm import Session" in source + assert "from policyengine_api.data.v1_daos" not in source + assert "build_v1_session_manager" not in source + + +def test_unmigrated_user_service_still_accepts_its_temporary_dao(): assert UserService(StubUsers()).get_profile(auth0_id="auth0|one") == { "user_id": None, "auth0_id": "auth0|one", From 8e20f988c7f9e11c729584ee21d8d996106c2ede Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 7 Aug 2026 22:08:01 +0300 Subject: [PATCH 50/89] refactor: use ORM sessions for users and saved policies --- policyengine_api/data/v1_daos.py | 123 ------------------ policyengine_api/endpoints/policy.py | 117 +++++++++-------- .../routes/user_profile_routes.py | 83 +++++++----- policyengine_api/services/user_service.py | 89 ++++++------- tests/contract/test_v1_route_contracts.py | 49 +++++-- tests/unit/data/test_ordinary_v1_daos.py | 70 ---------- tests/unit/data/test_v1_daos.py | 27 +--- tests/unit/data/test_v1_unit_of_work.py | 4 +- .../endpoints/test_stage7_orm_endpoints.py | 29 ++--- tests/unit/services/test_create_profile.py | 63 --------- tests/unit/services/test_direct_orm_users.py | 50 +++++++ .../services/test_stage7_dao_boundaries.py | 23 ++-- .../services/test_update_profile_service.py | 117 ----------------- tests/unit/services/test_user_service.py | 117 +++++++++++------ 14 files changed, 342 insertions(+), 619 deletions(-) delete mode 100644 tests/unit/services/test_create_profile.py create mode 100644 tests/unit/services/test_direct_orm_users.py delete mode 100644 tests/unit/services/test_update_profile_service.py diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 5ddf9a82b..74aeb1ec5 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -16,7 +16,6 @@ from policyengine_api.data.v1_models import ( Analysis, ComputedHousehold, - Economy, Household, LegacyReportOutputAlias, Policy, @@ -26,8 +25,6 @@ Simulation, SimulationRun, Tracer, - UserProfile, - UserPolicy, ) @@ -186,53 +183,6 @@ def get( return _mapping(model) if model else None -class UserDAO: - def __init__(self, session: Session): - self.session = session - - def create_profile( - self, - auth0_id: str, - username: str | None, - primary_country: str, - user_since: int, - ) -> int: - model = UserProfile( - auth0_id=auth0_id, - username=username, - primary_country=primary_country, - user_since=user_since, - ) - self.session.add(model) - self.session.flush() - return model.user_id - - def get_profile( - self, - *, - user_id: int | None = None, - auth0_id: str | None = None, - ) -> dict[str, Any] | None: - if user_id is None and auth0_id is None: - return None - condition = ( - UserProfile.user_id == user_id - if user_id is not None - else UserProfile.auth0_id == auth0_id - ) - model = self.session.scalar(select(UserProfile).where(condition)) - return _mapping(model) if model else None - - def update_profile(self, user_id: int, **values: Any) -> bool: - model = self.session.get(UserProfile, user_id) - if model is None: - return False - for key, value in values.items(): - if value is not None: - setattr(model, key, value) - return True - - class V1DAOs: """DAOs bound to the same operation-scoped Session.""" @@ -240,10 +190,7 @@ def __init__(self, session: Session): self.session = session self.policies = PolicyDAO(session) self.households = HouseholdDAO(session) - self.users = UserDAO(session) self.computed_households = ComputedHouseholdDAO(session) - self.user_policies = UserPolicyDAO(session) - self.economies = EconomyDAO(session) self.analyses = AnalysisDAO(session) self.reform_impacts = ReformImpactDAO(session) self.tracers = TracerDAO(session) @@ -283,76 +230,6 @@ def runtime_v1_unit_of_work(*, local: bool = False) -> V1UnitOfWork: return _runtime_unit_of_work[local] -class UserPolicyDAO: - IDENTITY_FIELDS = ( - "country_id", - "reform_id", - "baseline_id", - "user_id", - "year", - "geography", - "reform_label", - "baseline_label", - "dataset", - ) - - def __init__(self, session: Session): - self.session = session - - def create(self, **values: Any) -> int: - model = UserPolicy(**values) - self.session.add(model) - self.session.flush() - return model.id - - def find_unique(self, **values: Any) -> dict[str, Any] | None: - model = self.session.scalar( - select(UserPolicy).where( - *( - getattr(UserPolicy, field) == values[field] - for field in self.IDENTITY_FIELDS - ) - ) - ) - return _mapping(model) if model else None - - def get(self, user_policy_id: int) -> dict[str, Any] | None: - model = self.session.get(UserPolicy, user_policy_id) - return _mapping(model) if model else None - - def list_for_user(self, country_id: str, user_id: str) -> list[dict[str, Any]]: - models = self.session.scalars( - select(UserPolicy).where( - UserPolicy.country_id == country_id, - UserPolicy.user_id == user_id, - ) - ) - return [_mapping(model) for model in models] - - def update(self, user_policy_id: int, values: dict[str, Any]) -> bool: - model = self.session.get(UserPolicy, user_policy_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - -class EconomyDAO: - def __init__(self, session: Session): - self.session = session - - def create(self, **values: Any) -> int: - model = Economy(**values) - self.session.add(model) - self.session.flush() - return model.economy_id - - def get(self, economy_id: int) -> dict[str, Any] | None: - model = self.session.get(Economy, economy_id) - return _mapping(model) if model else None - - class AnalysisDAO: def __init__(self, session: Session): self.session = session diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index e247fb84d..4124112af 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -1,7 +1,30 @@ from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data.v1_daos import runtime_v1_unit_of_work import json from flask import Response, request +from sqlalchemy import select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy, UserPolicy + + +USER_POLICY_IDENTITY_FIELDS = ( + "country_id", + "reform_id", + "baseline_id", + "user_id", + "year", + "geography", + "reform_label", + "baseline_label", + "dataset", +) + + +def _serialize_user_policy(user_policy: UserPolicy) -> dict: + return { + column.name: getattr(user_policy, column.name) + for column in UserPolicy.__table__.columns + } @validate_country @@ -31,8 +54,13 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - with runtime_v1_unit_of_work().read() as daos: - results = daos.policies.search(country_id, query) + with get_v1_session_factory()() as session: + results = session.scalars( + select(Policy).where( + Policy.country_id == country_id, + Policy.label.contains(query, autoescape=True), + ) + ).all() if not results: body = dict( @@ -51,7 +79,7 @@ def get_policy_search(country_id: str) -> dict: # If a label-hash set aren't already in processed_vals, # add them to new_results for policy in results: - comparison_vals = policy["label"], policy["policy_hash"] + comparison_vals = policy.label, policy.policy_hash if comparison_vals not in processed_vals: new_results.append(policy) processed_vals.add(comparison_vals) @@ -60,7 +88,7 @@ def get_policy_search(country_id: str) -> dict: results = new_results # Format into: [{ id: 1, label: "My policy" }, ...] - policies = [dict(id=result["id"], label=result["label"]) for result in results] + policies = [dict(id=result.id, label=result.label) for result in results] body = dict( status="ok", message="Policies found", @@ -125,18 +153,24 @@ def set_user_policy(country_id: str) -> dict: # to be tested; type is not yet implemented try: - with runtime_v1_unit_of_work().transaction() as daos: - row = daos.user_policies.find_unique(**values) - if row is None: - user_policy_id = daos.user_policies.create(**values) - row = daos.user_policies.get(user_policy_id) + with get_v1_session_factory().begin() as session: + user_policy = session.scalar( + select(UserPolicy).where( + *( + getattr(UserPolicy, field) == values[field] + for field in USER_POLICY_IDENTITY_FIELDS + ) + ) + ) + if user_policy is None: + user_policy = UserPolicy(**values) + session.add(user_policy) + session.flush() else: - readable_row = dict(row) - response = dict( status="ok", message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", - result=dict(id=readable_row["id"]), + result=dict(id=user_policy.id), ) return Response( json.dumps(response), @@ -156,22 +190,7 @@ def set_user_policy(country_id: str) -> dict: status="ok", message="Record created successfully", result=dict( - id=row["id"], - country_id=row["country_id"], - reform_id=row["reform_id"], - reform_label=row["reform_label"], - baseline_id=row["baseline_id"], - baseline_label=row["baseline_label"], - user_id=row["user_id"], - year=row["year"], - geography=row["geography"], - dataset=row["dataset"], - number_of_provisions=row["number_of_provisions"], - api_version=row["api_version"], - added_date=row["added_date"], - updated_date=row["updated_date"], - budgetary_impact=row["budgetary_impact"], - type=row["type"], + **_serialize_user_policy(user_policy), ), ) @@ -189,30 +208,15 @@ def get_user_policy(country_id: str, user_id: str) -> dict: """ # Get the policy record for a given policy ID. - with runtime_v1_unit_of_work().read() as daos: - rows = daos.user_policies.list_for_user(country_id, user_id) - - rows_parsed = [ - dict( - id=row["id"], - country_id=row["country_id"], - reform_id=row["reform_id"], - reform_label=row["reform_label"], - baseline_id=row["baseline_id"], - baseline_label=row["baseline_label"], - user_id=row["user_id"], - year=row["year"], - geography=row["geography"], - dataset=row["dataset"], - number_of_provisions=row["number_of_provisions"], - api_version=row["api_version"], - added_date=row["added_date"], - updated_date=row["updated_date"], - budgetary_impact=row["budgetary_impact"], - type=row["type"], - ) - for row in rows - ] + with get_v1_session_factory()() as session: + user_policies = session.scalars( + select(UserPolicy).where( + UserPolicy.country_id == country_id, + UserPolicy.user_id == user_id, + ) + ).all() + + rows_parsed = [_serialize_user_policy(row) for row in user_policies] if rows_parsed is None: response = dict( @@ -301,8 +305,11 @@ def update_user_policy(country_id: str) -> dict: ) try: - with runtime_v1_unit_of_work().transaction() as daos: - daos.user_policies.update(user_policy_id, payload) + with get_v1_session_factory().begin() as session: + user_policy = session.get(UserPolicy, user_policy_id) + if user_policy is not None: + for key, value in payload.items(): + setattr(user_policy, key, value) except Exception as e: return Response( json.dumps( diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index ed725e641..55f6a437d 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,4 +1,6 @@ from flask import Blueprint, Response, request +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserProfile from policyengine_api.utils.payload_validators import validate_country import json from policyengine_api.services.user_service import UserService @@ -8,6 +10,22 @@ user_service = UserService() +def _serialize_user_profile( + profile: UserProfile, + *, + include_auth0_id: bool, +) -> dict: + result = { + "user_id": profile.user_id, + "primary_country": profile.primary_country, + "username": profile.username, + "user_since": profile.user_since, + } + if include_auth0_id: + result["auth0_id"] = profile.auth0_id + return result + + @user_profile_bp.route("//user-profile", methods=["POST"]) @validate_country def set_user_profile(country_id: str) -> Response: @@ -23,22 +41,20 @@ def set_user_profile(country_id: str) -> Response: username = payload.pop("username", None) user_since = payload.pop("user_since") - created, row = user_service.create_profile( - primary_country=country_id, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) + with get_v1_session_factory().begin() as session: + created, profile = user_service.create_profile( + session, + primary_country=country_id, + auth0_id=auth0_id, + username=username, + user_since=user_since, + ) + result = _serialize_user_profile(profile, include_auth0_id=False) response = dict( status="ok", message="Record created successfully" if created else "Record exists", - result=dict( - user_id=row["user_id"], - primary_country=row["primary_country"], - username=row["username"], - user_since=row["user_since"], - ), + result=result, ) return Response( json.dumps(response), @@ -56,21 +72,24 @@ def get_user_profile(country_id: str) -> Response: if (auth0_id is None) and (user_id is None): raise BadRequest("auth0_id or user_id must be provided") - row = ( - user_service.get_profile(user_id=user_id) - if auth0_id is None - else user_service.get_profile(auth0_id=auth0_id) - ) - - if row is None: + with get_v1_session_factory()() as session: + profile = ( + user_service.get_profile(session, user_id=user_id) + if auth0_id is None + else user_service.get_profile(session, auth0_id=auth0_id) + ) + readable_row = ( + None + if profile is None + else _serialize_user_profile( + profile, + include_auth0_id=auth0_id is not None, + ) + ) + + if readable_row is None: raise NotFound("No such user") - readable_row = dict(row) - # Delete auth0_id value if querying from user_id, as that value - # is a more private attribute than all others - if auth0_id is None: - del readable_row["auth0_id"] - response_body = dict( status="ok", message=f"User #{readable_row['user_id']} found successfully", @@ -108,12 +127,14 @@ def update_user_profile(country_id: str) -> Response: if user_id is None: raise BadRequest("Payload must include user_id") - updated = user_service.update_profile( - user_id=user_id, - primary_country=primary_country, - username=username, - user_since=user_since, - ) + with get_v1_session_factory().begin() as session: + updated = user_service.update_profile( + session, + user_id=user_id, + primary_country=primary_country, + username=username, + user_since=user_since, + ) if not updated: raise NotFound("No such user id") diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index 32eece842..c011b7f84 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,72 +1,67 @@ from __future__ import annotations -from contextlib import contextmanager -from typing import Any +from sqlalchemy import select +from sqlalchemy.orm import Session -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import UserDAO, V1UnitOfWork +from policyengine_api.data.v1_models import UserProfile class UserService: - def __init__( - self, - users: UserDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._users = users - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - - @contextmanager - def _repository(self, *, write: bool = False): - if self._users is not None: - yield self._users - return - boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as daos: - yield daos.users + """User-profile operations performed through a caller-owned ORM Session.""" def create_profile( self, + session: Session, primary_country: str, auth0_id: str, username: str | None, user_since: int, - ) -> tuple[bool, Any]: - with self._repository(write=True) as users: - row = users.get_profile(auth0_id=auth0_id) - if row is not None: - return False, row - users.create_profile(auth0_id, username, primary_country, user_since) - return True, users.get_profile(auth0_id=auth0_id) + ) -> tuple[bool, UserProfile]: + existing = self.get_profile(session, auth0_id=auth0_id) + if existing is not None: + return False, existing + profile = UserProfile( + auth0_id=auth0_id, + username=username, + primary_country=primary_country, + user_since=user_since, + ) + session.add(profile) + session.flush() + return True, profile def get_profile( - self, auth0_id: str | None = None, user_id: int | None = None - ) -> Any | None: + self, + session: Session, + auth0_id: str | None = None, + user_id: int | str | None = None, + ) -> UserProfile | None: if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") - with self._repository() as users: - return users.get_profile(user_id=user_id, auth0_id=auth0_id) + condition = ( + UserProfile.user_id == user_id + if user_id is not None + else UserProfile.auth0_id == auth0_id + ) + return session.scalar(select(UserProfile).where(condition)) def update_profile( self, + session: Session, user_id: int, primary_country: str | None, username: str | None, - user_since: int, - ) -> bool: + user_since: int | None, + ) -> UserProfile | None: if user_id is None: raise ValueError("you must specify either auth0_id or user_id") - with self._repository(write=True) as users: - return users.update_profile( - user_id, - primary_country=primary_country, - username=username, - user_since=user_since, - ) + profile = session.get(UserProfile, user_id) + if profile is None: + return None + if primary_country is not None: + profile.primary_country = primary_country + if username is not None: + profile.username = username + if user_since is not None: + profile.user_since = user_since + return profile diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 793a1a37b..dead34e64 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -190,19 +190,40 @@ def _json_payload(contract: ContractRequest) -> dict | None: return None -def _policy_search_unit_of_work(): - @contextmanager - def read(): - yield SimpleNamespace( - policies=SimpleNamespace( - search=lambda *args, **kwargs: [ - {"id": 123, "label": "Tax reform", "policy_hash": "hash-1"}, - {"id": 124, "label": "Tax reform", "policy_hash": "hash-1"}, - ] - ) - ) +def _policy_search_session_factory(): + policies = [ + Policy( + id=123, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + Policy( + id=124, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + ] + + class Result: + def all(self): + return policies + + class Session: + def scalars(self, statement): + return Result() + + class Factory: + @contextmanager + def __call__(self): + yield Session() - return SimpleNamespace(read=read) + return Factory() def _fake_country(): @@ -238,8 +259,8 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.runtime_v1_unit_of_work", - return_value=_policy_search_unit_of_work(), + "policyengine_api.endpoints.policy.get_v1_session_factory", + return_value=_policy_search_session_factory(), ) ) stack.enter_context( diff --git a/tests/unit/data/test_ordinary_v1_daos.py b/tests/unit/data/test_ordinary_v1_daos.py index 219184e0a..565f93997 100644 --- a/tests/unit/data/test_ordinary_v1_daos.py +++ b/tests/unit/data/test_ordinary_v1_daos.py @@ -29,34 +29,6 @@ def test_computed_household_upsert_preserves_one_cache_row(): assert row["computed_household_json"] == {"value": 2} -def test_user_policy_nullable_identity_list_and_update_are_orm_managed(): - uow = _unit_of_work() - values = { - "country_id": "us", - "reform_id": 2, - "reform_label": None, - "baseline_id": 1, - "baseline_label": None, - "user_id": "auth0|one", - "year": "2026", - "geography": "us", - "dataset": None, - "number_of_provisions": 3, - "api_version": "1", - "added_date": 1, - "updated_date": 1, - "budgetary_impact": None, - "type": None, - } - with uow.transaction() as daos: - user_policy_id = daos.user_policies.create(**values) - assert daos.user_policies.find_unique(**values)["id"] == user_policy_id - assert daos.user_policies.update(user_policy_id, {"number_of_provisions": 4}) - with uow.read() as daos: - rows = daos.user_policies.list_for_user("us", "auth0|one") - assert rows[0]["number_of_provisions"] == 4 - - def test_policy_search_and_reform_impact_limit_use_typed_statements(): uow = _unit_of_work() with uow.transaction() as daos: @@ -85,45 +57,3 @@ def test_computed_household_create_and_version_filters(): "value": 1 } assert daos.computed_households.get(1, 2, "us", api_version="2") is None - - -def test_economy_and_user_policy_daos_cover_lookup_edge_cases(): - uow = _unit_of_work() - user_policy_values = { - "country_id": "us", - "reform_id": 2, - "reform_label": None, - "baseline_id": 1, - "baseline_label": None, - "user_id": "auth0|one", - "year": "2026", - "geography": "us", - "dataset": None, - "number_of_provisions": 3, - "api_version": "1", - "added_date": 1, - "updated_date": 1, - "budgetary_impact": None, - "type": None, - } - with uow.transaction() as daos: - economy_id = daos.economies.create( - policy_id=2, - country_id="us", - region="us", - time_period="2026", - options_json={"dataset": "default"}, - options_hash="hash", - api_version="1", - economy_json={"result": 1}, - status="complete", - message=None, - ) - user_policy_id = daos.user_policies.create(**user_policy_values) - assert daos.user_policies.update(999, {}) is False - - with uow.read() as daos: - assert daos.economies.get(economy_id)["economy_json"] == {"result": 1} - assert daos.economies.get(999) is None - assert daos.user_policies.get(user_policy_id)["country_id"] == "us" - assert daos.user_policies.get(999) is None diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py index 5fc8ed8d6..c00464ce6 100644 --- a/tests/unit/data/test_v1_daos.py +++ b/tests/unit/data/test_v1_daos.py @@ -51,32 +51,7 @@ def test_household_dao_creates_updates_and_reads(): assert daos.households.get("uk", household_id) is None -def test_user_dao_profile_lookup_precedence(): +def test_household_dao_handles_missing_update(): uow = _unit_of_work() with uow.transaction() as daos: - user_id = daos.users.create_profile("auth0|one", "person", "us", 123) - with uow.read() as daos: - assert daos.users.get_profile(auth0_id="auth0|one")["user_id"] == user_id - assert ( - daos.users.get_profile(user_id=user_id, auth0_id="wrong")["auth0_id"] - == "auth0|one" - ) - - -def test_user_and_household_daos_handle_missing_and_nullable_updates(): - uow = _unit_of_work() - with uow.transaction() as daos: - user_id = daos.users.create_profile("auth0|one", "original", "us", 123) - assert daos.users.get_profile() is None - assert daos.users.update_profile(999, username="missing") is False assert daos.households.update("us", 999, "missing", {}, "missing", "1") is False - assert daos.users.update_profile( - user_id, - username=None, - primary_country="uk", - ) - - with uow.read() as daos: - profile = daos.users.get_profile(user_id=user_id) - assert profile["username"] == "original" - assert profile["primary_country"] == "uk" diff --git a/tests/unit/data/test_v1_unit_of_work.py b/tests/unit/data/test_v1_unit_of_work.py index 19bc613bc..91e768750 100644 --- a/tests/unit/data/test_v1_unit_of_work.py +++ b/tests/unit/data/test_v1_unit_of_work.py @@ -29,12 +29,12 @@ def test_unit_of_work_rolls_back_every_repository_on_failure(): with pytest.raises(RuntimeError, match="abort"): with uow.transaction() as daos: daos.policies.create("us", None, {}, "policy", "1") - daos.users.create_profile("auth0|one", "person", "us", 1) + daos.households.create("us", None, {}, "household", "1") raise RuntimeError("abort") with uow.read() as daos: assert daos.policies.get("us", 1) is None - assert daos.users.get_profile(auth0_id="auth0|one") is None + assert daos.households.get("us", 1) is None def test_unit_of_work_rolls_back_parent_run_and_alias_together(): diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py index 5be23c788..f1672adfb 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -5,22 +5,18 @@ from sqlalchemy import select from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork -from policyengine_api.data.v1_models import ComputedHousehold, Household, Policy +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, + UserPolicy, +) from policyengine_api.endpoints.household import get_household_under_policy from policyengine_api.endpoints.policy import ( get_user_policy, set_user_policy, update_user_policy, ) -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work() -> V1UnitOfWork: - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) def test_household_under_policy_returns_cached_json_object(orm_session_factory): @@ -112,9 +108,10 @@ def test_household_under_policy_calculates_and_caches_json_as_an_object( assert cached.computed_household_json == calculated -def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): +def test_user_policy_endpoints_round_trip_through_orm_session_factory( + orm_session_factory, +): app = Flask(__name__) - uow = _unit_of_work() payload = { "reform_label": "Reform", "reform_id": 2, @@ -133,8 +130,8 @@ def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): } with patch( - "policyengine_api.endpoints.policy.runtime_v1_unit_of_work", - return_value=uow, + "policyengine_api.endpoints.policy.get_v1_session_factory", + return_value=orm_session_factory, ): with app.test_request_context(json=payload): created = set_user_policy("us") @@ -146,5 +143,5 @@ def test_user_policy_endpoints_round_trip_through_the_unit_of_work(): assert created.get_json()["result"]["dataset"] == "default" assert listed["result"][0]["reform_label"] == "Reform" assert updated.status_code == 200 - with uow.read() as daos: - assert daos.user_policies.get(1)["reform_label"] == "Updated" + with orm_session_factory() as session: + assert session.get(UserPolicy, 1).reform_label == "Updated" diff --git a/tests/unit/services/test_create_profile.py b/tests/unit/services/test_create_profile.py deleted file mode 100644 index 6217d7633..000000000 --- a/tests/unit/services/test_create_profile.py +++ /dev/null @@ -1,63 +0,0 @@ -import pytest -import unittest.mock as mock -import time -from policyengine_api.services.user_service import UserService - -userService = UserService() - - -class TestCreateProfile: - def test_create_profile_valid(self): - auth0_id = "test-auth-id" - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - - result = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - - assert result[0] is True - user_record = userService.get_profile(auth0_id) - assert user_record is not None - assert user_record["auth0_id"] == auth0_id - assert user_record["primary_country"] == primary_country - assert user_record["username"] == username - assert user_record["user_since"] == user_since - - def test_create_profile_invalid(self): - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - with pytest.raises( - Exception, - match=r"UserService.create_profile\(\) missing 1 required positional argument: 'auth0_id'", - ): - userService.create_profile( - primary_country=primary_country, - username=username, - user_since=user_since, - ) - - def test_create_profile_duplicate(self): - auth0_id = "test-auth-id" - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - result1 = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - - result2 = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - assert result2[0] == False diff --git a/tests/unit/services/test_direct_orm_users.py b/tests/unit/services/test_direct_orm_users.py new file mode 100644 index 000000000..f44794ace --- /dev/null +++ b/tests/unit/services/test_direct_orm_users.py @@ -0,0 +1,50 @@ +from policyengine_api.data.v1_models import UserProfile +from policyengine_api.services.user_service import UserService + + +def test_user_service_reads_and_writes_mapped_profiles(orm_session): + service = UserService() + + created, profile = service.create_profile( + orm_session, + primary_country="us", + auth0_id="auth0|direct", + username="direct-user", + user_since=123, + ) + duplicate_created, duplicate = service.create_profile( + orm_session, + primary_country="us", + auth0_id="auth0|direct", + username="ignored", + user_since=456, + ) + + assert created is True + assert duplicate_created is False + assert isinstance(profile, UserProfile) + assert duplicate is profile + assert service.get_profile(orm_session, auth0_id="auth0|direct") is profile + + +def test_user_service_updates_the_mapped_profile(orm_session): + service = UserService() + _, profile = service.create_profile( + orm_session, + primary_country="us", + auth0_id="auth0|update", + username=None, + user_since=123, + ) + + updated = service.update_profile( + orm_session, + user_id=profile.user_id, + primary_country="uk", + username="updated-user", + user_since=456, + ) + + assert updated is profile + assert profile.primary_country == "uk" + assert profile.username == "updated-user" diff --git a/tests/unit/services/test_stage7_dao_boundaries.py b/tests/unit/services/test_stage7_dao_boundaries.py index ef7124fc5..7ccaf42d2 100644 --- a/tests/unit/services/test_stage7_dao_boundaries.py +++ b/tests/unit/services/test_stage7_dao_boundaries.py @@ -2,10 +2,9 @@ import pytest -from policyengine_api.services.user_service import UserService - SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" +DAO_MODULE = SERVICE_ROOT.parent / "data" / "v1_daos.py" @pytest.mark.parametrize( @@ -18,12 +17,10 @@ def test_migrated_services_do_not_issue_queries_directly(module_name): assert "from policyengine_api.data import database" not in source -class StubUsers: - def get_profile(self, *, user_id=None, auth0_id=None): - return {"user_id": user_id, "auth0_id": auth0_id} - - -@pytest.mark.parametrize("module_name", ["household_service.py", "policy_service.py"]) +@pytest.mark.parametrize( + "module_name", + ["household_service.py", "policy_service.py", "user_service.py"], +) def test_migrated_services_use_sessions_and_mapped_models(module_name): source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") assert "from sqlalchemy.orm import Session" in source @@ -31,8 +28,8 @@ def test_migrated_services_use_sessions_and_mapped_models(module_name): assert "build_v1_session_manager" not in source -def test_unmigrated_user_service_still_accepts_its_temporary_dao(): - assert UserService(StubUsers()).get_profile(auth0_id="auth0|one") == { - "user_id": None, - "auth0_id": "auth0|one", - } +def test_migrated_user_domains_have_no_temporary_daos(): + source = DAO_MODULE.read_text(encoding="utf-8") + assert "class UserDAO" not in source + assert "class UserPolicyDAO" not in source + assert "class EconomyDAO" not in source diff --git a/tests/unit/services/test_update_profile_service.py b/tests/unit/services/test_update_profile_service.py deleted file mode 100644 index d0bacb8f6..000000000 --- a/tests/unit/services/test_update_profile_service.py +++ /dev/null @@ -1,117 +0,0 @@ -import pytest -from policyengine_api.services.user_service import UserService - - -pytest_plugins = ["tests.fixtures.services.user_service"] - -service = UserService() - - -class TestUpdateProfile: - def test_update_profile_given_existing_record(self, test_db, existing_user_profile): - # GIVEN an existing profile record (from fixture) - - # WHEN we call update_profile with new data - updated_username = "updated_username" - updated_country = "uk" - - result = service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country=updated_country, - username=updated_username, - user_since=existing_user_profile["user_since"], - ) - - # THEN the method should return True for successful update - assert result is True - - # AND the database should be updated with new values - updated_record = test_db.query( - "SELECT * FROM user_profiles WHERE user_id = ?", - (existing_user_profile["user_id"],), - ).fetchone() - - assert updated_record["username"] == updated_username - assert updated_record["primary_country"] == updated_country - - def test_update_profile_given_nonexistent_record(self, test_db): - # GIVEN a nonexistent profile record id - NONEXISTENT_ID = 999 - - # WHEN we call update_profile for this nonexistent record - result = service.update_profile( - user_id=NONEXISTENT_ID, - primary_country="uk", - username="newuser", - user_since="2024-01-01", - ) - - # THEN the result should be False - assert result is False - - def test_update_profile_with_partial_fields(self, test_db, existing_user_profile): - # GIVEN an existing profile record (from fixture) - - # WHEN we call update_profile with only some fields provided - updated_country = "CA" - original_username = existing_user_profile["username"] - - result = service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country=updated_country, - username=None, - user_since=existing_user_profile["user_since"], - ) - - # THEN the method should return True for successful update - assert result is True - - # AND only the provided fields should be updated - updated_record = test_db.query( - "SELECT * FROM user_profiles WHERE user_id = ?", - (existing_user_profile["user_id"],), - ).fetchone() - - assert updated_record["primary_country"] == updated_country - assert ( - updated_record["username"] == original_username - ) # Username should remain unchanged - - def test_update_profile_with_database_error( - self, monkeypatch, existing_user_profile - ): - # GIVEN an existing profile record (from fixture) - - # AND a database that raises an exception - def mock_dao_error(*args, **kwargs): - raise Exception("Database error") - - class FailingUsers: - update_profile = staticmethod(mock_dao_error) - - monkeypatch.setattr(service, "_users", FailingUsers()) - - # WHEN we call update_profile - # THEN an exception should be raised - with pytest.raises(Exception, match="Database error"): - service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country="US", - username="testuser", - user_since="2023-01-01", - ) - - def test_update_profile_id_not_specified(self): - # GIVEN no user_id specified - - # WHEN we call update_profile with None as user_id - # THEN a ValueError should be raised - with pytest.raises( - ValueError, match="you must specify either auth0_id or user_id" - ): - service.update_profile( - user_id=None, - primary_country="US", - username="testuser", - user_since="2023-01-01", - ) diff --git a/tests/unit/services/test_user_service.py b/tests/unit/services/test_user_service.py index 26875c8ca..0736f75eb 100644 --- a/tests/unit/services/test_user_service.py +++ b/tests/unit/services/test_user_service.py @@ -1,60 +1,93 @@ import pytest + +from policyengine_api.data.v1_models import UserProfile from policyengine_api.services.user_service import UserService +from tests.fixtures.services.user_service import valid_user_record -from tests.fixtures.services.user_service import ( - valid_user_record, - existing_user_profile, -) + +pytest_plugins = ["tests.fixtures.services.user_service"] service = UserService() -class TestGetProfile: - def test_get_profile_id_not_specified(self): - # GIVEN no ID - # WHEN we call get_profile with no auth0_id or user_id +def test_get_profile_requires_an_identifier(orm_session): + with pytest.raises( + ValueError, + match="you must specify either auth0_id or user_id", + ): + service.get_profile(orm_session) - # Then a ValueError should be raised - with pytest.raises( - ValueError, match="you must specify either auth0_id or user_id" - ): - service.get_profile() - def test_get_profile_nonexistent_record(self): - # GIVEN nonexistent record - INVALID_RECORD_ID = "invalid" +def test_get_profile_returns_none_for_unknown_auth0_id(orm_session): + assert service.get_profile(orm_session, auth0_id="missing") is None - # WHEN we call get_profile with nonexistent user - result = service.get_profile(auth0_id=INVALID_RECORD_ID) - # THEN result is None - assert result is None +def test_get_profile_returns_mapped_entity_by_either_identifier( + orm_session, + existing_user_profile, +): + by_auth0 = service.get_profile( + orm_session, + auth0_id=valid_user_record["auth0_id"], + ) + by_id = service.get_profile( + orm_session, + user_id=valid_user_record["user_id"], + ) + + assert isinstance(by_auth0, UserProfile) + assert by_auth0 is by_id + assert by_auth0.username == valid_user_record["username"] - def test_get_profile_auth0_id(self, existing_user_profile): - # WHEN we call get_profile with auth0_id - result = service.get_profile(auth0_id=existing_user_profile["auth0_id"]) - # THEN returns record - assert result == existing_user_profile +def test_create_profile_returns_existing_entity_for_duplicate_auth0_id( + orm_session, +): + created, profile = service.create_profile( + orm_session, + "us", + "auth0|duplicate", + "first", + 1, + ) + duplicate_created, duplicate = service.create_profile( + orm_session, + "uk", + "auth0|duplicate", + "second", + 2, + ) - def test_get_profile_user_id(self, existing_user_profile): - # WHEN we call get_profile with user_id - result = service.get_profile(user_id=existing_user_profile["user_id"]) + assert created is True + assert duplicate_created is False + assert duplicate is profile + assert duplicate.username == "first" - # THEN returns record - assert result == existing_user_profile - def test_get_profile_id_priority(self, test_db, existing_user_profile): - # WHEN we call get_profile with auth0_id and user_id - result = service.get_profile( - auth0_id=existing_user_profile["auth0_id"], - user_id=existing_user_profile["user_id"], - ) +def test_update_profile_returns_none_for_missing_entity(orm_session): + assert service.update_profile(orm_session, 999, "uk", "missing", 2) is None + + +def test_update_profile_only_changes_non_null_fields( + orm_session, + existing_user_profile, +): + profile = service.update_profile( + orm_session, + valid_user_record["user_id"], + "uk", + None, + valid_user_record["user_since"] + 1, + ) + + assert profile.primary_country == "uk" + assert profile.username == valid_user_record["username"] + assert profile.user_since == valid_user_record["user_since"] + 1 - # THEN returns record using auth0_id - record = test_db.query( - "SELECT * FROM user_profiles WHERE auth0_id = ?", - (valid_user_record["auth0_id"],), - ).fetchone() - assert result == record +def test_update_profile_requires_user_id(orm_session): + with pytest.raises( + ValueError, + match="you must specify either auth0_id or user_id", + ): + service.update_profile(orm_session, None, "us", "name", 1) From 3077ca9168bacf418a557c368a8584bb5b91a867 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 02:25:21 +0300 Subject: [PATCH 51/89] refactor: use ORM sessions for analysis persistence --- policyengine_api/country.py | 19 +- policyengine_api/data/v1_daos.py | 170 +-------------- .../endpoints/economy/reform_impact.py | 28 ++- policyengine_api/endpoints/simulation.py | 23 +- .../routes/simulation_analysis_routes.py | 3 +- .../routes/tracer_analysis_routes.py | 6 +- .../services/ai_analysis_service.py | 65 +++--- policyengine_api/services/economy_service.py | 138 ++++++------ .../services/reform_impacts_service.py | 201 ++++++++++-------- .../services/simulation_analysis_service.py | 12 +- .../services/tracer_analysis_service.py | 75 +++---- tests/fixtures/services/economy_service.py | 38 ++-- .../services/tracer_analysis_service.py | 16 +- .../python/test_ai_analysis_service_old.py | 44 ---- .../python/test_simulation_analysis_routes.py | 15 +- tests/unit/data/test_local_daos.py | 93 -------- .../unit/services/test_ai_analysis_service.py | 42 ++-- .../test_direct_orm_local_analysis.py | 66 ++++++ tests/unit/services/test_economy_service.py | 20 +- tests/unit/services/test_execute_analysis.py | 18 +- .../services/test_reform_impacts_service.py | 135 ++++++------ .../test_stage7_local_service_boundaries.py | 8 + tests/unit/services/test_tracer_service.py | 23 +- 23 files changed, 527 insertions(+), 731 deletions(-) delete mode 100644 tests/to_refactor/python/test_ai_analysis_service_old.py delete mode 100644 tests/unit/data/test_local_daos.py create mode 100644 tests/unit/services/test_direct_orm_local_analysis.py diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 602af096f..593a49cf6 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -23,7 +23,8 @@ build_congressional_district_metadata, ) -from policyengine_api.data.v1_daos import runtime_v1_unit_of_work +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Tracer from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, get_bundle_default_dataset_option, @@ -433,13 +434,15 @@ def calculate( if household_id is not None and policy_id is not None: # write to local database - with runtime_v1_unit_of_work(local=True).transaction() as daos: - daos.tracers.create( - household_id, - policy_id, - self.country_id, - COUNTRY_PACKAGE_VERSIONS[self.country_id], - log_lines, + with get_v1_session_factory(local=True).begin() as session: + session.add( + Tracer( + household_id=household_id, + policy_id=policy_id, + country_id=self.country_id, + api_version=COUNTRY_PACKAGE_VERSIONS[self.country_id], + tracer_output=log_lines, + ) ) return household diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py index 74aeb1ec5..856461214 100644 --- a/policyengine_api/data/v1_daos.py +++ b/policyengine_api/data/v1_daos.py @@ -6,25 +6,21 @@ from contextlib import contextmanager from typing import Any -from datetime import datetime import uuid -from sqlalchemy import delete, func, or_, select +from sqlalchemy import func, select from sqlalchemy.orm import Session from policyengine_api.data.orm import SessionManager from policyengine_api.data.v1_models import ( - Analysis, ComputedHousehold, Household, LegacyReportOutputAlias, Policy, - ReformImpact, ReportOutput, ReportOutputRun, Simulation, SimulationRun, - Tracer, ) @@ -191,9 +187,6 @@ def __init__(self, session: Session): self.policies = PolicyDAO(session) self.households = HouseholdDAO(session) self.computed_households = ComputedHouseholdDAO(session) - self.analyses = AnalysisDAO(session) - self.reform_impacts = ReformImpactDAO(session) - self.tracers = TracerDAO(session) self.simulations = SimulationDAO(session) self.reports = ReportDAO(session) @@ -230,167 +223,6 @@ def runtime_v1_unit_of_work(*, local: bool = False) -> V1UnitOfWork: return _runtime_unit_of_work[local] -class AnalysisDAO: - def __init__(self, session: Session): - self.session = session - - def get(self, prompt: str) -> str | None: - model = self.session.scalar( - select(Analysis) - .where( - Analysis.prompt == prompt, - Analysis.status.in_(("complete", "ok")), - ) - .order_by(Analysis.prompt_id.desc()) - ) - return model.analysis if model else None - - def store(self, prompt: str, analysis: str | None, status: str) -> int: - model = Analysis(prompt=prompt, analysis=analysis, status=status) - self.session.add(model) - self.session.flush() - return model.prompt_id - - -class ReformImpactDAO: - def __init__(self, session: Session): - self.session = session - - def create(self, **values: Any) -> int: - model = ReformImpact(**values) - self.session.add(model) - self.session.flush() - return model.reform_impact_id - - def find(self, *, execution_id: str) -> dict[str, Any] | None: - model = self.session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - return _mapping(model) if model else None - - @staticmethod - def _scope(statement, **filters: Any): - return statement.where( - *(getattr(ReformImpact, key) == value for key, value in filters.items()) - ) - - def list(self, **filters: Any) -> list[dict[str, Any]]: - models = self.session.scalars( - self._scope(select(ReformImpact), **filters).order_by( - ReformImpact.start_time.desc() - ) - ) - return [_mapping(model) for model in models] - - def list_recent(self, limit: int) -> list[dict[str, Any]]: - models = self.session.scalars( - select(ReformImpact).order_by(ReformImpact.start_time.desc()).limit(limit) - ) - return [_mapping(model) for model in models] - - def list_by_options_hash( - self, options_hash: str, options_hash_prefix: str, **filters: Any - ) -> list[dict[str, Any]]: - statement = self._scope(select(ReformImpact), **filters).where( - or_( - ReformImpact.options_hash == options_hash, - ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), - ) - ) - models = self.session.scalars( - statement.order_by( - (ReformImpact.options_hash == options_hash).desc(), - ReformImpact.start_time.desc(), - ) - ) - return [_mapping(model) for model in models] - - def delete_computing(self, **filters: Any) -> None: - self.session.execute( - self._scope(delete(ReformImpact), **filters).where( - ReformImpact.status == "computing" - ) - ) - - def set_message(self, message: str, **filters: Any) -> bool: - models = self.session.scalars( - self._scope(select(ReformImpact), **filters) - ).all() - for model in models: - model.message = message - return bool(models) - - def fail(self, execution_id: str, message: str, finished_at: datetime) -> bool: - model = self.session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - if model is None: - return False - model.status = "error" - model.message = message - model.end_time = finished_at - return True - - def complete(self, execution_id: str, result: Any, finished_at: datetime) -> bool: - model = self.session.scalar( - select(ReformImpact) - .where(ReformImpact.execution_id == execution_id) - .order_by(ReformImpact.reform_impact_id.desc()) - ) - if model is None: - return False - model.status = "ok" - model.message = "Completed" - model.reform_impact_json = result - model.end_time = finished_at - return True - - -class TracerDAO: - def __init__(self, session: Session): - self.session = session - - def create( - self, - household_id: int, - policy_id: int, - country_id: str, - api_version: str, - tracer_output: Any, - ) -> int: - model = Tracer( - household_id=household_id, - policy_id=policy_id, - country_id=country_id, - api_version=api_version, - tracer_output=tracer_output, - ) - self.session.add(model) - self.session.flush() - return model.id - - def get( - self, - household_id: int, - policy_id: int, - country_id: str, - api_version: str | None = None, - ) -> dict[str, Any] | None: - statement = select(Tracer).where( - Tracer.household_id == household_id, - Tracer.policy_id == policy_id, - Tracer.country_id == country_id, - ) - if api_version is not None: - statement = statement.where(Tracer.api_version == api_version) - model = self.session.scalar(statement.order_by(Tracer.id.desc())) - return _mapping(model) if model else None - - class SimulationDAO: def __init__(self, session: Session): self.session = session diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py index e777e6975..5d8641d3c 100644 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ b/policyengine_api/endpoints/economy/reform_impact.py @@ -1,4 +1,7 @@ -from policyengine_api.data.v1_daos import runtime_v1_unit_of_work +from sqlalchemy import select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReformImpact def set_comment_on_job( @@ -11,14 +14,17 @@ def set_comment_on_job( time_period, options_hash, ): - with runtime_v1_unit_of_work(local=True).transaction() as daos: - daos.reform_impacts.set_message( - comment, - country_id=country_id, - reform_policy_id=policy_id, - baseline_policy_id=baseline_policy_id, - region=region, - time_period=time_period, - options_hash=options_hash, - dataset=dataset, + with get_v1_session_factory(local=True).begin() as session: + impacts = session.scalars( + select(ReformImpact).where( + ReformImpact.country_id == country_id, + ReformImpact.reform_policy_id == policy_id, + ReformImpact.baseline_policy_id == baseline_policy_id, + ReformImpact.region == region, + ReformImpact.time_period == time_period, + ReformImpact.options_hash == options_hash, + ReformImpact.dataset == dataset, + ) ) + for impact in impacts: + impact.message = comment diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index 97eb3e303..d9b102656 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -1,4 +1,7 @@ -from policyengine_api.data.v1_daos import runtime_v1_unit_of_work +from sqlalchemy import select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReformImpact """ @@ -42,9 +45,21 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - with runtime_v1_unit_of_work(local=True).read() as daos: - result = daos.reform_impacts.list_recent(max_results) + with get_v1_session_factory(local=True)() as session: + result = session.scalars( + select(ReformImpact) + .order_by(ReformImpact.start_time.desc()) + .limit(max_results) + ).all() # Format into [{}] - return {"result": [dict(r) for r in result]} + return { + "result": [ + { + column.name: getattr(impact, column.name) + for column in ReformImpact.__table__.columns + } + for impact in result + ] + } diff --git a/policyengine_api/routes/simulation_analysis_routes.py b/policyengine_api/routes/simulation_analysis_routes.py index 5157b807d..a3094292a 100644 --- a/policyengine_api/routes/simulation_analysis_routes.py +++ b/policyengine_api/routes/simulation_analysis_routes.py @@ -1,6 +1,5 @@ from flask import Blueprint, request, Response, stream_with_context from werkzeug.exceptions import BadRequest -from policyengine_api.utils.payload_validators import validate_country from policyengine_api.services.simulation_analysis_service import ( SimulationAnalysisService, ) @@ -11,6 +10,7 @@ validate_sim_analysis_payload, ) import json +from policyengine_api.data.orm import get_v1_session_factory simulation_analysis_bp = Blueprint("simulation_analysis", __name__) simulation_analysis_service = SimulationAnalysisService() @@ -44,6 +44,7 @@ def execute_simulation_analysis(country_id): audience = payload.get("audience", "") analysis, analysis_type = simulation_analysis_service.execute_analysis( + get_v1_session_factory(local=True), country_id, currency, dataset, diff --git a/policyengine_api/routes/tracer_analysis_routes.py b/policyengine_api/routes/tracer_analysis_routes.py index 6638282a4..9991e5514 100644 --- a/policyengine_api/routes/tracer_analysis_routes.py +++ b/policyengine_api/routes/tracer_analysis_routes.py @@ -8,8 +8,7 @@ TracerAnalysisService, ) import json -from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS -import re +from policyengine_api.data.orm import get_v1_session_factory tracer_analysis_bp = Blueprint("tracer_analysis", __name__) tracer_analysis_service = TracerAnalysisService() @@ -27,12 +26,11 @@ def execute_tracer_analysis(country_id): household_id = payload.get("household_id") policy_id = payload.get("policy_id") variable = payload.get("variable") - api_version = COUNTRY_PACKAGE_VERSIONS[country_id] - if not isinstance(variable, str): raise BadRequest("variable must be a string") analysis, analysis_type = tracer_analysis_service.execute_analysis( + get_v1_session_factory(local=True), country_id, household_id, policy_id, diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index 28da6cb30..c34a9dbfd 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -1,13 +1,13 @@ import json import os from collections.abc import Generator -from contextlib import contextmanager import anthropic from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import AnalysisDAO, V1UnitOfWork +from policyengine_api.data.v1_models import Analysis class StreamEvent(BaseModel): @@ -25,36 +25,27 @@ class ErrorEvent(StreamEvent): class AIAnalysisService: - def __init__( - self, - analyses: AnalysisDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._analyses = analyses - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager(local=True)) - return self._unit_of_work - - @contextmanager - def _analysis_repository(self, *, write: bool = False): - if self._analyses is not None: - yield self._analyses - return - boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as daos: - yield daos.analyses + """AI analysis operations backed by caller-owned ORM sessions.""" - def get_existing_analysis(self, prompt: str) -> str | None: - with self._analysis_repository() as analyses: - analysis = analyses.get(prompt) - return json.dumps(analysis) if analysis is not None else None - - def trigger_ai_analysis(self, prompt: str) -> Generator[str, None, None]: + def get_existing_analysis( + self, + session: Session, + prompt: str, + ) -> Analysis | None: + return session.scalar( + select(Analysis) + .where( + Analysis.prompt == prompt, + Analysis.status.in_(("complete", "ok")), + ) + .order_by(Analysis.prompt_id.desc()) + ) + + def trigger_ai_analysis( + self, + prompt: str, + session_factory: sessionmaker[Session], + ) -> Generator[str, None, None]: claude_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) def generate(): @@ -80,7 +71,13 @@ def generate(): yield ( json.dumps(TextEvent(stream=event.text).model_dump()) + "\n" ) - with self._analysis_repository(write=True) as analyses: - analyses.store(prompt, response_text, "ok") + with session_factory.begin() as session: + session.add( + Analysis( + prompt=prompt, + analysis=response_text, + status="ok", + ) + ) return generate() diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index b3c8bdf0a..091e0846e 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -23,6 +23,7 @@ normalize_us_region, ) from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReformImpact from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint @@ -660,7 +661,7 @@ def _get_or_create_economic_impact( most_recent_impact = self._get_most_recent_impact(setup_options) if ( not most_recent_impact - or most_recent_impact.get("options_hash") != setup_options.options_hash + or most_recent_impact.options_hash != setup_options.options_hash ): most_recent_impact = None @@ -767,19 +768,22 @@ def _get_previous_impacts( Fetch any previous simulation runs for the given policy reform. """ - previous_impacts: list[Any] = ( - reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options_hash, - self._build_options_hash_lookup_pattern(options_hash), - api_version, + previous_impacts: list[Any] = [] + with get_v1_session_factory(local=True)() as session: + previous_impacts = ( + reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + self._build_options_hash_lookup_pattern(options_hash), + api_version, + ) ) - ) return previous_impacts def _get_most_recent_impact( @@ -805,19 +809,19 @@ def _get_most_recent_impact( return None for impact in previous_impacts: - if impact.get("options_hash") == setup_options.options_hash: + if impact.options_hash == setup_options.options_hash: return impact return previous_impacts[0] def _determine_impact_action( self, - most_recent_impact: dict | None, + most_recent_impact: ReformImpact | None, ) -> ImpactAction: if not most_recent_impact: return ImpactAction.CREATE - status = most_recent_impact.get("status") + status = most_recent_impact.status if status in [ImpactStatus.OK.value, ImpactStatus.ERROR.value]: return ImpactAction.COMPLETED elif status == ImpactStatus.COMPUTING.value: @@ -829,7 +833,7 @@ def _handle_execution_state( self, setup_options: EconomicImpactSetupOptions, execution_state: str, - reform_impact: dict, + reform_impact: ReformImpact, execution: Optional[Any] = None, ) -> EconomicImpactResult: """ @@ -847,7 +851,7 @@ def _handle_execution_state( self._set_reform_impact_complete( setup_options=setup_options, reform_impact_json=result, - execution_id=reform_impact["execution_id"], + execution_id=reform_impact.execution_id, ) logger.log_struct( {"message": "Sim API execution completed"}, @@ -870,7 +874,7 @@ def _handle_execution_state( self._set_reform_impact_error( setup_options=setup_options, message=error_message, - execution_id=reform_impact["execution_id"], + execution_id=reform_impact.execution_id, ) logger.log_struct( {"message": error_message}, @@ -891,9 +895,9 @@ def _handle_execution_state( def _handle_completed_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> EconomicImpactResult: - result = self._parse_json_object(most_recent_impact["reform_impact_json"]) + result = self._parse_json_object(most_recent_impact.reform_impact_json) return EconomicImpactResult.completed( data=self._with_policyengine_bundle( result=result, @@ -904,10 +908,10 @@ def _handle_completed_impact( def _handle_computing_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> EconomicImpactResult: execution = simulation_entrypoint.get_execution_by_id( - most_recent_impact["execution_id"] + most_recent_impact.execution_id ) execution_state = simulation_entrypoint.get_execution_status(execution) return self._handle_execution_state( @@ -1070,18 +1074,18 @@ def _extract_dataset_version(self, dataset: str | None) -> str | None: return None return dataset.rsplit("@", 1)[1] - def _extract_cached_result(self, most_recent_impact: dict) -> dict: + def _extract_cached_result(self, most_recent_impact: ReformImpact) -> dict: try: - return self._parse_json_object(most_recent_impact["reform_impact_json"]) + return self._parse_json_object(most_recent_impact.reform_impact_json) except (TypeError, ValueError): return {} def _should_refresh_cached_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> bool: - if most_recent_impact.get("status") == ImpactStatus.COMPUTING.value: + if most_recent_impact.status == ImpactStatus.COMPUTING.value: return False cached_result = self._extract_cached_result(most_recent_impact) @@ -1332,21 +1336,23 @@ def _set_reform_impact_computing( In the reform_impact table, set the status of the impact to "computing". """ try: - reform_impacts_service.set_reform_impact( - country_id=setup_options.country_id, - policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options=setup_options.options, - options_hash=setup_options.options_hash, - status=ImpactStatus.COMPUTING.value, - api_version=setup_options.api_version, - reform_impact_json={}, - start_time=datetime.datetime.now(), - execution_id=execution_id, - ) + with get_v1_session_factory(local=True).begin() as session: + reform_impacts_service.set_reform_impact( + session, + country_id=setup_options.country_id, + policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options=setup_options.options, + options_hash=setup_options.options_hash, + status=ImpactStatus.COMPUTING.value, + api_version=setup_options.api_version, + reform_impact_json={}, + start_time=datetime.datetime.now(), + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { @@ -1366,17 +1372,19 @@ def _set_reform_impact_complete( In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON. """ try: - reform_impacts_service.set_complete_reform_impact( - country_id=setup_options.country_id, - reform_policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options_hash=setup_options.options_hash, - reform_impact_json=reform_impact_json, - execution_id=execution_id, - ) + with get_v1_session_factory(local=True).begin() as session: + reform_impacts_service.set_complete_reform_impact( + session, + country_id=setup_options.country_id, + reform_policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options_hash=setup_options.options_hash, + reform_impact_json=reform_impact_json, + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { @@ -1396,17 +1404,19 @@ def _set_reform_impact_error( In the reform_impact table, set the status of the impact to "error" and store the error message. """ try: - reform_impacts_service.set_error_reform_impact( - country_id=setup_options.country_id, - policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options_hash=setup_options.options_hash, - message=message, - execution_id=execution_id, - ) + with get_v1_session_factory(local=True).begin() as session: + reform_impacts_service.set_error_reform_impact( + session, + country_id=setup_options.country_id, + policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options_hash=setup_options.options_hash, + message=message, + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 2a8a3fbbb..641afb3ed 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -1,35 +1,14 @@ import datetime -from contextlib import contextmanager from typing import Any -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReformImpactDAO, V1UnitOfWork +from sqlalchemy import delete, or_, select +from sqlalchemy.orm import Session +from policyengine_api.data.v1_models import ReformImpact -class ReformImpactsService: - def __init__( - self, - impacts: ReformImpactDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._impacts = impacts - self._unit_of_work = unit_of_work - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager(local=True)) - return self._unit_of_work - - @contextmanager - def _repository(self, *, write: bool = False): - if self._impacts is not None: - yield self._impacts - return - boundary = self.unit_of_work.transaction if write else self.unit_of_work.read - with boundary() as daos: - yield daos.reform_impacts +class ReformImpactsService: + """Reform-impact operations performed through a caller-owned Session.""" @staticmethod def _filters( @@ -53,8 +32,15 @@ def _filters( filters["api_version"] = api_version return filters + @staticmethod + def _scope(statement, **filters): + return statement.where( + *(getattr(ReformImpact, key) == value for key, value in filters.items()) + ) + def get_all_reform_impacts( self, + session: Session, country_id, policy_id, baseline_policy_id, @@ -63,23 +49,24 @@ def get_all_reform_impacts( time_period, options_hash, api_version, - ): - with self._repository() as impacts: - return impacts.list( - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - api_version, - ), - options_hash=options_hash, - ) + ) -> list[ReformImpact]: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ) + statement = self._scope(select(ReformImpact), **filters).where( + ReformImpact.options_hash == options_hash + ) + return list(session.scalars(statement.order_by(ReformImpact.start_time.desc()))) def get_all_reform_impacts_by_options_hash_prefix( self, + session: Session, country_id, policy_id, baseline_policy_id, @@ -89,24 +76,34 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash, options_hash_prefix, api_version, - ): - with self._repository() as impacts: - return impacts.list_by_options_hash( - options_hash, - options_hash_prefix, - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - api_version, - ), + ) -> list[ReformImpact]: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ) + statement = self._scope(select(ReformImpact), **filters).where( + or_( + ReformImpact.options_hash == options_hash, + ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), ) + ) + return list( + session.scalars( + statement.order_by( + (ReformImpact.options_hash == options_hash).desc(), + ReformImpact.start_time.desc(), + ) + ) + ) def set_reform_impact( self, + session: Session, country_id, policy_id, baseline_policy_id, @@ -120,26 +117,29 @@ def set_reform_impact( reform_impact_json: dict[str, Any], start_time, execution_id: str, - ): - with self._repository(write=True) as impacts: - return impacts.create( - country_id=country_id, - reform_policy_id=policy_id, - baseline_policy_id=baseline_policy_id, - region=region, - dataset=dataset, - time_period=time_period, - options_json=options, - options_hash=options_hash, - status=status, - api_version=api_version, - reform_impact_json=reform_impact_json, - start_time=start_time, - execution_id=execution_id, - ) + ) -> ReformImpact: + impact = ReformImpact( + country_id=country_id, + reform_policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + region=region, + dataset=dataset, + time_period=time_period, + options_json=options, + options_hash=options_hash, + status=status, + api_version=api_version, + reform_impact_json=reform_impact_json, + start_time=start_time, + execution_id=execution_id, + ) + session.add(impact) + session.flush() + return impact def delete_reform_impact( self, + session: Session, country_id, policy_id, baseline_policy_id, @@ -147,22 +147,25 @@ def delete_reform_impact( dataset, time_period, options_hash, - ): - with self._repository(write=True) as impacts: - impacts.delete_computing( - **self._filters( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - ), - options_hash=options_hash, + ) -> None: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + ) + session.execute( + self._scope(delete(ReformImpact), **filters).where( + ReformImpact.options_hash == options_hash, + ReformImpact.status == "computing", ) + ) def set_error_reform_impact( self, + session: Session, country_id, policy_id, baseline_policy_id, @@ -172,7 +175,7 @@ def set_error_reform_impact( options_hash, message, execution_id: str, - ): + ) -> ReformImpact | None: del ( country_id, policy_id, @@ -182,11 +185,21 @@ def set_error_reform_impact( time_period, options_hash, ) - with self._repository(write=True) as impacts: - return impacts.fail(execution_id, message, self._now()) + impact = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if impact is None: + return None + impact.status = "error" + impact.message = message + impact.end_time = self._now() + return impact def set_complete_reform_impact( self, + session: Session, country_id, reform_policy_id, baseline_policy_id, @@ -196,7 +209,7 @@ def set_complete_reform_impact( options_hash, reform_impact_json: dict[str, Any], execution_id, - ): + ) -> ReformImpact | None: del ( country_id, reform_policy_id, @@ -206,8 +219,18 @@ def set_complete_reform_impact( time_period, options_hash, ) - with self._repository(write=True) as impacts: - return impacts.complete(execution_id, reform_impact_json, self._now()) + impact = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if impact is None: + return None + impact.status = "ok" + impact.message = "Completed" + impact.reform_impact_json = reform_impact_json + impact.end_time = self._now() + return impact @staticmethod def _now() -> datetime.datetime: diff --git a/policyengine_api/services/simulation_analysis_service.py b/policyengine_api/services/simulation_analysis_service.py index 8738dd625..1204a7495 100644 --- a/policyengine_api/services/simulation_analysis_service.py +++ b/policyengine_api/services/simulation_analysis_service.py @@ -1,6 +1,7 @@ from policyengine_api.services.ai_analysis_service import AIAnalysisService from policyengine_api.services.ai_prompt_service import AIPromptService from typing import Generator, Literal +from sqlalchemy.orm import Session, sessionmaker ai_prompt_service = AIPromptService() @@ -12,11 +13,9 @@ class SimulationAnalysisService(AIAnalysisService): analysis database table """ - def __init__(self): - super().__init__() - def execute_analysis( self, + session_factory: sessionmaker[Session], country_id: str, currency: str, dataset: str | None, @@ -61,14 +60,15 @@ def execute_analysis( print("Checking if AI analysis already exists for this prompt") # If a calculated record exists for this prompt, return it as a # streaming response - existing_analysis = self.get_existing_analysis(prompt) + with session_factory() as session: + existing_analysis = self.get_existing_analysis(session, prompt) if existing_analysis is not None: - return existing_analysis, "static" + return existing_analysis.analysis, "static" print("Found no existing AI analysis; triggering new analysis with Claude") # Otherwise, pass prompt to Claude, then return streaming function try: - analysis = self.trigger_ai_analysis(prompt) + analysis = self.trigger_ai_analysis(prompt, session_factory) return analysis, "streaming" except Exception as e: raise e diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index 2b0013bf9..307c54e65 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -1,36 +1,19 @@ -import json -from contextlib import contextmanager - -from policyengine_api.data.v1_daos import AnalysisDAO, TracerDAO, V1UnitOfWork from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS from typing import Generator, Literal import re import anthropic from policyengine_api.services.ai_analysis_service import AIAnalysisService from werkzeug.exceptions import NotFound +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker +from policyengine_api.data.v1_models import Tracer -class TracerAnalysisService(AIAnalysisService): - def __init__( - self, - tracers: TracerDAO | None = None, - analyses: AnalysisDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._tracers = tracers - super().__init__(analyses, unit_of_work=unit_of_work) - - @contextmanager - def _tracer_repository(self): - if self._tracers is not None: - yield self._tracers - return - with self.unit_of_work.read() as daos: - yield daos.tracers +class TracerAnalysisService(AIAnalysisService): def execute_analysis( self, + session_factory: sessionmaker[Session], country_id: str, household_id: str, policy_id: str, @@ -48,12 +31,14 @@ def execute_analysis( # Retrieve tracer record from table try: - tracer: list[str] = self.get_tracer( - country_id, - household_id, - policy_id, - api_version, - ) + with session_factory() as session: + tracer: list[str] = self.get_tracer( + session, + country_id, + household_id, + policy_id, + api_version, + ) except Exception as e: raise e @@ -73,13 +58,14 @@ def execute_analysis( ) # If a calculated record exists for this prompt, return it as a string - existing_analysis: str = self.get_existing_analysis(prompt) + with session_factory() as session: + existing_analysis = self.get_existing_analysis(session, prompt) if existing_analysis is not None: - return existing_analysis, "static" + return existing_analysis.analysis, "static" # Otherwise, pass prompt to Claude, then return streaming function try: - analysis: Generator = self.trigger_ai_analysis(prompt) + analysis: Generator = self.trigger_ai_analysis(prompt, session_factory) return analysis, "streaming" except Exception as e: print( @@ -89,31 +75,28 @@ def execute_analysis( def get_tracer( self, + session: Session, country_id: str, household_id: str, policy_id: str, api_version: str, ) -> list: try: - # Retrieve from the tracers table in the local database - with self._tracer_repository() as tracers: - row = tracers.get( - household_id, - policy_id, - country_id, - api_version, + tracer = session.scalar( + select(Tracer) + .where( + Tracer.household_id == int(household_id), + Tracer.policy_id == int(policy_id), + Tracer.country_id == country_id, + Tracer.api_version == api_version, ) + .order_by(Tracer.id.desc()) + ) - if row is None: + if tracer is None: raise NotFound("No household simulation tracer found") - tracer_output = row["tracer_output"] - tracer_output_list = ( - json.loads(tracer_output) - if isinstance(tracer_output, str) - else tracer_output - ) - return tracer_output_list + return tracer.tracer_output except Exception as e: print(f"Error getting existing tracer analysis: {str(e)}") diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index c3af5cbcd..e38d7dd47 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -7,6 +7,7 @@ MODAL_EXECUTION_STATUS_RUNNING, MODAL_EXECUTION_STATUS_SUBMITTED, ) +from policyengine_api.data.v1_models import ReformImpact # Mock data constants MOCK_COUNTRY_ID = "us" @@ -223,25 +224,24 @@ def create_mock_reform_impact( }, } ) - return { - "id": 1, - "country_id": MOCK_COUNTRY_ID, - "policy_id": MOCK_POLICY_ID, - "baseline_policy_id": MOCK_BASELINE_POLICY_ID, - "region": MOCK_REGION, - "dataset": MOCK_RESOLVED_DATASET, - "time_period": time_period, - "options_hash": options_hash, - "status": status, - "api_version": MOCK_API_VERSION, - "reform_impact_json": reform_impact_json or default_reform_impact_json, - "execution_id": execution_id, - "message": message, - "start_time": start_time or datetime.datetime(2025, 6, 26, 12, 0, 0), - "end_time": ( - datetime.datetime(2025, 6, 26, 12, 5, 0) if status == "ok" else None - ), - } + return ReformImpact( + reform_impact_id=1, + country_id=MOCK_COUNTRY_ID, + reform_policy_id=MOCK_POLICY_ID, + baseline_policy_id=MOCK_BASELINE_POLICY_ID, + region=MOCK_REGION, + dataset=MOCK_RESOLVED_DATASET, + time_period=time_period, + options_json=MOCK_OPTIONS, + options_hash=options_hash, + status=status, + api_version=MOCK_API_VERSION, + reform_impact_json=reform_impact_json or default_reform_impact_json, + execution_id=execution_id, + message=message, + start_time=start_time or datetime.datetime(2025, 6, 26, 12, 0, 0), + end_time=(datetime.datetime(2025, 6, 26, 12, 5, 0) if status == "ok" else None), + ) def create_mock_modal_execution( diff --git a/tests/fixtures/services/tracer_analysis_service.py b/tests/fixtures/services/tracer_analysis_service.py index 1a1262dc0..4118f233d 100644 --- a/tests/fixtures/services/tracer_analysis_service.py +++ b/tests/fixtures/services/tracer_analysis_service.py @@ -1,9 +1,9 @@ import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) from unittest.mock import patch +from policyengine_api.data.v1_models import Analysis valid_tracer_output = [ " snap<2027, (default)> = [6769.799]", @@ -17,10 +17,10 @@ " snap_fpg<2027-01, (default)> = [1806.4779]", ] -invalid_tracer_output = { - "variable": "only_government_benefit <1500>", - "variable": " market_income <1000>", -} +invalid_tracer_output = [ + "only_government_benefit <1500>", + " market_income <1000>", +] spliced_valid_tracer_output_root_variable = valid_tracer_output[0:] @@ -68,7 +68,11 @@ def mock_get_existing_analysis(): with patch.object( TracerAnalysisService, "get_existing_analysis", - return_value="Existing static analysis", + return_value=Analysis( + prompt="prompt", + analysis="Existing static analysis", + status="ok", + ), ) as mock: yield mock diff --git a/tests/to_refactor/python/test_ai_analysis_service_old.py b/tests/to_refactor/python/test_ai_analysis_service_old.py deleted file mode 100644 index 950e39ae0..000000000 --- a/tests/to_refactor/python/test_ai_analysis_service_old.py +++ /dev/null @@ -1,44 +0,0 @@ -import json -import os -from unittest.mock import MagicMock, patch - -import pytest - -from policyengine_api.services.ai_analysis_service import AIAnalysisService - - -def test_get_existing_analysis_found(): - analyses = MagicMock() - analyses.get.return_value = "Existing analysis" - service = AIAnalysisService(analyses) - - output = service.get_existing_analysis("Test prompt") - - assert output == json.dumps("Existing analysis") - analyses.get.assert_called_once_with("Test prompt") - - -def test_get_existing_analysis_not_found(): - analyses = MagicMock() - analyses.get.return_value = None - service = AIAnalysisService(analyses) - - assert service.get_existing_analysis("Test prompt") is None - analyses.get.assert_called_once_with("Test prompt") - - -def test_anthropic_api_key(): - with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test_key"}): - assert os.getenv("ANTHROPIC_API_KEY") == "test_key" - - -@patch("policyengine_api.services.ai_analysis_service.anthropic.Anthropic") -def test_trigger_ai_analysis_error(mock_anthropic): - mock_client = MagicMock() - mock_anthropic.return_value = mock_client - mock_client.messages.stream.side_effect = Exception("API Error") - - generator = AIAnalysisService(MagicMock()).trigger_ai_analysis("Test prompt") - - with pytest.raises(Exception, match="API Error"): - list(generator) diff --git a/tests/to_refactor/python/test_simulation_analysis_routes.py b/tests/to_refactor/python/test_simulation_analysis_routes.py index 7e7d5ea34..7d0086a7b 100644 --- a/tests/to_refactor/python/test_simulation_analysis_routes.py +++ b/tests/to_refactor/python/test_simulation_analysis_routes.py @@ -1,14 +1,9 @@ -import pytest from unittest.mock import patch -from flask import Flask +from policyengine_api.data.v1_models import Analysis from policyengine_api.services.simulation_analysis_service import ( SimulationAnalysisService, ) -from policyengine_api.routes.simulation_analysis_routes import ( - execute_simulation_analysis, -) - from tests.to_refactor.fixtures.simulation_analysis_fixtures import ( test_json, test_impact, @@ -21,7 +16,11 @@ def test_execute_simulation_analysis_existing_analysis(rest_client): with patch( "policyengine_api.services.ai_analysis_service.AIAnalysisService.get_existing_analysis" ) as mock_get_existing: - mock_get_existing.return_value = "Existing analysis" + mock_get_existing.return_value = Analysis( + prompt="prompt", + analysis="Existing analysis", + status="ok", + ) response = rest_client.post("/us/simulation-analysis", json=test_json) @@ -82,7 +81,7 @@ def test_execute_simulation_analysis_custom_dataset(rest_client): } with patch( "policyengine_api.services.simulation_analysis_service.SimulationAnalysisService._generate_simulation_analysis_prompt" - ) as mock_generate_prompt: + ): with patch( "policyengine_api.services.ai_analysis_service.AIAnalysisService.get_existing_analysis" ) as mock_get_existing: diff --git a/tests/unit/data/test_local_daos.py b/tests/unit/data/test_local_daos.py deleted file mode 100644 index 510a66ad5..000000000 --- a/tests/unit/data/test_local_daos.py +++ /dev/null @@ -1,93 +0,0 @@ -from datetime import datetime - -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work(): - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) - - -def test_analysis_dao_round_trip(): - uow = _unit_of_work() - with uow.transaction() as daos: - daos.analyses.store("prompt", "answer", "complete") - with uow.read() as daos: - assert daos.analyses.get("prompt") == "answer" - - -def test_reform_impact_dao_transitions_by_execution_id(): - uow = _unit_of_work() - with uow.transaction() as daos: - daos.reform_impacts.create( - country_id="us", - reform_policy_id=2, - baseline_policy_id=1, - region="us", - dataset="default", - time_period="2026", - options_json={}, - options_hash="hash", - api_version="1", - reform_impact_json={}, - status="computing", - start_time=datetime(2026, 1, 1), - execution_id="job", - ) - daos.reform_impacts.complete("job", {"result": 1}, datetime(2026, 1, 2)) - with uow.read() as daos: - assert daos.reform_impacts.find(execution_id="job")["status"] == "ok" - assert daos.reform_impacts.find(execution_id="job")["reform_impact_json"] == { - "result": 1 - } - - -def test_reform_impact_dao_orders_limits_messages_and_handles_missing_jobs(): - uow = _unit_of_work() - with uow.transaction() as daos: - for day, execution_id in ((1, "old-job"), (2, "new-job")): - daos.reform_impacts.create( - country_id="us", - reform_policy_id=2, - baseline_policy_id=1, - region="us", - dataset="default", - time_period="2026", - options_json={}, - options_hash=execution_id, - api_version="1", - reform_impact_json={}, - status="computing", - start_time=datetime(2026, 1, day), - execution_id=execution_id, - ) - - assert daos.reform_impacts.set_message( - "queued", country_id="us", status="computing" - ) - assert daos.reform_impacts.set_message("missing", country_id="uk") is False - assert ( - daos.reform_impacts.fail("missing-job", "failed", datetime(2026, 1, 3)) - is False - ) - assert ( - daos.reform_impacts.complete("missing-job", {}, datetime(2026, 1, 3)) - is False - ) - - with uow.read() as daos: - recent = daos.reform_impacts.list_recent(1) - assert [row["execution_id"] for row in recent] == ["new-job"] - assert recent[0]["message"] == "queued" - - -def test_tracer_dao_returns_latest_matching_trace(): - uow = _unit_of_work() - with uow.transaction() as daos: - daos.tracers.create(1, 2, "us", "1", {"trace": "first"}) - daos.tracers.create(1, 2, "us", "1", {"trace": "latest"}) - with uow.read() as daos: - assert daos.tracers.get(1, 2, "us")["tracer_output"] == {"trace": "latest"} diff --git a/tests/unit/services/test_ai_analysis_service.py b/tests/unit/services/test_ai_analysis_service.py index e853a1f68..b39854d59 100644 --- a/tests/unit/services/test_ai_analysis_service.py +++ b/tests/unit/services/test_ai_analysis_service.py @@ -1,29 +1,29 @@ import json +from sqlalchemy import select + +from policyengine_api.data.v1_models import Analysis from policyengine_api.services.ai_analysis_service import AIAnalysisService -from tests.fixtures.services.ai_analysis_service import ( - mock_stream_text_events, - mock_stream_error_event, - patch_anthropic, - parse_to_chunks, -) +from tests.fixtures.services.ai_analysis_service import parse_to_chunks import pytest +pytest_plugins = ["tests.fixtures.services.ai_analysis_service"] + # Initialize the service service = AIAnalysisService() class TestTriggerAIAnalysis: def test_trigger_ai_analysis_given_successful_streaming( - self, mock_stream_text_events, test_db + self, mock_stream_text_events, orm_session_factory ): # GIVEN a series of successful text messages from the Claude API expected_response = "This is a historical quote." text_chunks = parse_to_chunks(expected_response) - mock_client = mock_stream_text_events(text_chunks=text_chunks) + mock_stream_text_events(text_chunks=text_chunks) # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote" - generator = service.trigger_ai_analysis(prompt) + generator = service.trigger_ai_analysis(prompt, orm_session_factory) # THEN it should yield the expected chunks results = list(generator) @@ -37,13 +37,14 @@ def test_trigger_ai_analysis_given_successful_streaming( assert chunk == expected_chunk # Verify the database was updated with the complete response - analysis_record = test_db.query( - "SELECT * FROM analysis WHERE prompt = ?", (prompt,) - ).fetchone() + with orm_session_factory() as session: + analysis_record = session.scalar( + select(Analysis).where(Analysis.prompt == prompt) + ) assert analysis_record is not None - assert analysis_record["analysis"] == expected_response - assert analysis_record["status"] == "ok" + assert analysis_record.analysis == expected_response + assert analysis_record.status == "ok" @pytest.mark.parametrize( "error_type", @@ -54,14 +55,14 @@ def test_trigger_ai_analysis_given_successful_streaming( ], ) def test_trigger_ai_analysis_given_error( - self, mock_stream_error_event, test_db, error_type + self, mock_stream_error_event, orm_session_factory, error_type ): # GIVEN an overloaded_error event from the Claude API - mock_client = mock_stream_error_event(error_type) + mock_stream_error_event(error_type) # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote about erroneous systems" - generator = service.trigger_ai_analysis(prompt) + generator = service.trigger_ai_analysis(prompt, orm_session_factory) # THEN it should yield the expected error message results = list(generator) @@ -79,8 +80,9 @@ def test_trigger_ai_analysis_given_error( assert results[0] == expected_error # Verify the database was not updated - analysis_record = test_db.query( - "SELECT * FROM analysis WHERE prompt = ?", (prompt,) - ).fetchone() + with orm_session_factory() as session: + analysis_record = session.scalar( + select(Analysis).where(Analysis.prompt == prompt) + ) assert analysis_record is None diff --git a/tests/unit/services/test_direct_orm_local_analysis.py b/tests/unit/services/test_direct_orm_local_analysis.py new file mode 100644 index 000000000..2207f2c41 --- /dev/null +++ b/tests/unit/services/test_direct_orm_local_analysis.py @@ -0,0 +1,66 @@ +from datetime import datetime + +from policyengine_api.data.v1_models import Analysis, ReformImpact, Tracer +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService + + +def test_ai_analysis_service_returns_the_latest_mapped_analysis(orm_session): + orm_session.add_all( + [ + Analysis(prompt="prompt", analysis="old", status="ok"), + Analysis(prompt="prompt", analysis="new", status="complete"), + ] + ) + orm_session.flush() + + analysis = AIAnalysisService().get_existing_analysis(orm_session, "prompt") + + assert isinstance(analysis, Analysis) + assert analysis.analysis == "new" + + +def test_reform_impact_service_writes_mapped_entity(orm_session): + impact = ReformImpactsService().set_reform_impact( + orm_session, + country_id="us", + policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options={"dataset": "default"}, + options_hash="hash", + status="computing", + api_version="1", + reform_impact_json={}, + start_time=datetime(2026, 1, 1), + execution_id="job", + ) + + assert isinstance(impact, ReformImpact) + assert impact.options_json == {"dataset": "default"} + + +def test_tracer_service_reads_python_json_from_mapped_entity(orm_session): + orm_session.add( + Tracer( + household_id=1, + policy_id=2, + country_id="us", + api_version="1", + tracer_output=["net_income <2026>", " dependency"], + ) + ) + orm_session.flush() + + tracer = TracerAnalysisService().get_tracer( + orm_session, + "us", + "1", + "2", + "1", + ) + + assert tracer == ["net_income <2026>", " dependency"] diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 9d61994c0..74092498d 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -1,6 +1,6 @@ import json from typing import Literal -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -148,8 +148,8 @@ def test__given_orm_decoded_completed_impact__returns_completed_result( mock_numpy_random, ): completed_impact = create_mock_reform_impact(status="ok") - completed_impact["reform_impact_json"] = json.loads( - completed_impact["reform_impact_json"] + completed_impact.reform_impact_json = json.loads( + completed_impact.reform_impact_json ) mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ completed_impact @@ -406,6 +406,7 @@ def test__given_runtime_cache_version__uses_versioned_economy_cache_key( economy_service.get_economic_impact(**base_params) mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.assert_called_once_with( + ANY, MOCK_COUNTRY_ID, MOCK_POLICY_ID, MOCK_BASELINE_POLICY_ID, @@ -437,14 +438,14 @@ def test__given_alias_dataset__queries_previous_impacts_with_resolved_bundle( economy_service.get_economic_impact(**base_params) call_args = mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.call_args.args - assert call_args[4] == MOCK_RESOLVED_DATASET - assert call_args[6] == MOCK_LOOKUP_OPTIONS_HASH - assert call_args[7] == economy_service._build_options_hash_lookup_pattern( + assert call_args[5] == MOCK_RESOLVED_DATASET + assert call_args[7] == MOCK_LOOKUP_OPTIONS_HASH + assert call_args[8] == economy_service._build_options_hash_lookup_pattern( MOCK_LOOKUP_OPTIONS_HASH ) - assert "data\\_version=faux-populace-us-2099-test-release" in call_args[7] - assert "policyengine\\_version=3.4.0" in call_args[7] - assert "runtime_app_name" not in call_args[7] + assert "data\\_version=faux-populace-us-2099-test-release" in call_args[8] + assert "policyengine\\_version=3.4.0" in call_args[8] + assert "runtime_app_name" not in call_args[8] def test__given_completed_impact__uses_resolved_runtime_bundle_for_cache_lookup( self, @@ -1287,6 +1288,7 @@ def test_given_valid_parameters_calls_service_correctly( assert result == expected_impacts mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.assert_called_once_with( + ANY, MOCK_COUNTRY_ID, MOCK_POLICY_ID, MOCK_BASELINE_POLICY_ID, diff --git a/tests/unit/services/test_execute_analysis.py b/tests/unit/services/test_execute_analysis.py index f0c2ff622..7963575b6 100644 --- a/tests/unit/services/test_execute_analysis.py +++ b/tests/unit/services/test_execute_analysis.py @@ -1,18 +1,8 @@ -import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) -from werkzeug.exceptions import NotFound -from tests.fixtures.services.tracer_analysis_service import ( - sample_tracer_data, - sample_expected_segment, - mock_get_tracer, - mock_get_existing_analysis, - mock_parse_tracer_output, - mock_trigger_ai_analysis, -) +pytest_plugins = ["tests.fixtures.services.tracer_analysis_service"] service = TracerAnalysisService() country_id = "us" @@ -24,6 +14,7 @@ class TestExecuteAnalysis: def test_execute_analysis_static( self, + orm_session_factory, mock_get_tracer, mock_parse_tracer_output, mock_get_existing_analysis, @@ -36,7 +27,7 @@ def test_execute_analysis_static( """ analysis, analysis_type = service.execute_analysis( - country_id, household_id, policy_id, target_variable + orm_session_factory, country_id, household_id, policy_id, target_variable ) assert analysis == "Existing static analysis" @@ -44,6 +35,7 @@ def test_execute_analysis_static( def test_execute_analysis_streaming( self, + orm_session_factory, mock_get_tracer, mock_parse_tracer_output, mock_get_existing_analysis, @@ -60,7 +52,7 @@ def test_execute_analysis_streaming( mock_get_existing_analysis.return_value = None analysis, analysis_type = service.execute_analysis( - country_id, household_id, policy_id, target_variable + orm_session_factory, country_id, household_id, policy_id, target_variable ) expected_streaming_output = ["stream chunk 1", "stream chunk 2"] diff --git a/tests/unit/services/test_reform_impacts_service.py b/tests/unit/services/test_reform_impacts_service.py index 5a1bded9c..18fbec449 100644 --- a/tests/unit/services/test_reform_impacts_service.py +++ b/tests/unit/services/test_reform_impacts_service.py @@ -1,27 +1,17 @@ from datetime import datetime -from unittest.mock import ANY, Mock -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import ReformImpactDAO, V1UnitOfWork -from policyengine_api.services import reform_impacts_service as service_module +from sqlalchemy import select + +from policyengine_api.data.v1_models import ReformImpact from policyengine_api.services.reform_impacts_service import ReformImpactsService -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema -def _unit_of_work() -> V1UnitOfWork: - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) +service = ReformImpactsService() -def _create_impact( - service: ReformImpactsService, - *, - execution_id: str, - options_hash: str, - day: int, -) -> int: +def _create_impact(session, *, execution_id: str, options_hash: str, day: int): return service.set_reform_impact( + session, country_id="us", policy_id=2, baseline_policy_id=1, @@ -38,25 +28,33 @@ def _create_impact( ) -def test_reform_impact_service_round_trips_queries_and_transitions(): - service = ReformImpactsService(unit_of_work=_unit_of_work()) - exact_id = _create_impact( - service, +def test_reform_impact_service_round_trips_models_and_transitions(orm_session): + exact = _create_impact( + orm_session, execution_id="exact-job", options_hash="hash-exact", day=1, ) - compatible_id = _create_impact( - service, + compatible = _create_impact( + orm_session, execution_id="compatible-job", options_hash="hash-compatible", day=2, ) - exact = service.get_all_reform_impacts( - "us", 2, 1, "us", "default", "2026", "hash-exact", "1" + exact_results = service.get_all_reform_impacts( + orm_session, + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + "1", ) - compatible = service.get_all_reform_impacts_by_options_hash_prefix( + compatible_results = service.get_all_reform_impacts_by_options_hash_prefix( + orm_session, "us", 2, 1, @@ -68,12 +66,10 @@ def test_reform_impact_service_round_trips_queries_and_transitions(): "1", ) - assert [row["reform_impact_id"] for row in exact] == [exact_id] - assert [row["reform_impact_id"] for row in compatible] == [ - exact_id, - compatible_id, - ] - assert service.set_complete_reform_impact( + assert exact_results == [exact] + assert compatible_results == [exact, compatible] + completed = service.set_complete_reform_impact( + orm_session, "us", 2, 1, @@ -84,7 +80,8 @@ def test_reform_impact_service_round_trips_queries_and_transitions(): {"result": 1}, "exact-job", ) - assert service.set_error_reform_impact( + failed = service.set_error_reform_impact( + orm_session, "us", 2, 1, @@ -96,65 +93,59 @@ def test_reform_impact_service_round_trips_queries_and_transitions(): "compatible-job", ) - with service.unit_of_work.read() as daos: - completed = daos.reform_impacts.find(execution_id="exact-job") - failed = daos.reform_impacts.find(execution_id="compatible-job") - assert completed["status"] == "ok" - assert completed["reform_impact_json"] == {"result": 1} - assert failed["status"] == "error" - assert failed["message"] == "failed" + assert completed.status == "ok" + assert completed.reform_impact_json == {"result": 1} + assert failed.status == "error" + assert failed.message == "failed" -def test_reform_impact_service_deletes_only_matching_computing_rows(): - service = ReformImpactsService(unit_of_work=_unit_of_work()) +def test_reform_impact_service_deletes_only_matching_computing_rows(orm_session): _create_impact( - service, + orm_session, execution_id="delete-job", options_hash="delete-hash", day=1, ) - retained_id = _create_impact( - service, + retained = _create_impact( + orm_session, execution_id="retain-job", options_hash="retain-hash", day=2, ) - service.delete_reform_impact("us", 2, 1, "us", "default", "2026", "delete-hash") + service.delete_reform_impact( + orm_session, + "us", + 2, + 1, + "us", + "default", + "2026", + "delete-hash", + ) assert ( - service.get_all_reform_impacts( - "us", 2, 1, "us", "default", "2026", "delete-hash", "1" + orm_session.scalar( + select(ReformImpact).where(ReformImpact.execution_id == "delete-job") ) - == [] - ) - retained = service.get_all_reform_impacts( - "us", 2, 1, "us", "default", "2026", "retain-hash", "1" + is None ) - assert [row["reform_impact_id"] for row in retained] == [retained_id] + assert orm_session.get(ReformImpact, retained.reform_impact_id) is retained -def test_reform_impact_service_supports_injected_repository(): - impacts = Mock(spec=ReformImpactDAO) - impacts.fail.return_value = False - service = ReformImpactsService(impacts) - +def test_reform_impact_transitions_return_none_for_missing_execution(orm_session): assert ( service.set_error_reform_impact( - "us", 2, 1, "us", "default", "2026", "hash", "missing", "job" + orm_session, + "us", + 2, + 1, + "us", + "default", + "2026", + "hash", + "missing", + "missing-job", ) - is False + is None ) - impacts.fail.assert_called_once_with("job", "missing", ANY) - - -def test_reform_impact_service_builds_default_unit_of_work_once(monkeypatch): - manager = build_sqlite_session_manager() - build_manager = Mock(return_value=manager) - monkeypatch.setattr(service_module, "build_v1_session_manager", build_manager) - service = ReformImpactsService() - - first = service.unit_of_work - - assert service.unit_of_work is first - build_manager.assert_called_once_with(local=True) diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_stage7_local_service_boundaries.py index de65cbb83..d7d0dd6ce 100644 --- a/tests/unit/services/test_stage7_local_service_boundaries.py +++ b/tests/unit/services/test_stage7_local_service_boundaries.py @@ -4,6 +4,7 @@ SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" +DAO_MODULE = SERVICE_ROOT.parent / "data" / "v1_daos.py" @pytest.mark.parametrize( @@ -20,3 +21,10 @@ def test_local_data_services_do_not_issue_queries_directly(module_name): assert ".query(" not in source assert "local_database" not in source assert "from policyengine_api.data import database" not in source + + +def test_migrated_local_domains_have_no_temporary_daos(): + source = DAO_MODULE.read_text(encoding="utf-8") + assert "class AnalysisDAO" not in source + assert "class ReformImpactDAO" not in source + assert "class TracerDAO" not in source diff --git a/tests/unit/services/test_tracer_service.py b/tests/unit/services/test_tracer_service.py index 84ece8df3..d9a490537 100644 --- a/tests/unit/services/test_tracer_service.py +++ b/tests/unit/services/test_tracer_service.py @@ -1,23 +1,21 @@ import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) from werkzeug.exceptions import NotFound -from tests.fixtures.services.tracer_fixture_service import ( - test_tracer_data, - valid_tracer_row, - valid_tracer, -) +from tests.fixtures.services.tracer_fixture_service import valid_tracer + +pytest_plugins = ["tests.fixtures.services.tracer_fixture_service"] tracer_service = TracerAnalysisService() -def test_get_tracer_valid(test_tracer_data): +def test_get_tracer_valid(test_tracer_data, orm_session): # Test get_tracer successfully retrieves valid data from the database. result = tracer_service.get_tracer( + orm_session, test_tracer_data["country_id"], test_tracer_data["household_id"], test_tracer_data["policy_id"], @@ -29,7 +27,7 @@ def test_get_tracer_valid(test_tracer_data): assert result == valid_output -def test_get_tracer_not_found(): +def test_get_tracer_not_found(orm_session): # Test get_tracer raises NotFound when no matching record exists. valid_country_val_in_db = "us" invalid_household_not_in_db = "9999999" @@ -42,10 +40,10 @@ def test_get_tracer_not_found(): invalid_api_version, ] with pytest.raises(NotFound): - tracer_service.get_tracer(*data_not_in_db) + tracer_service.get_tracer(orm_session, *data_not_in_db) -def test_get_tracer_database_error(test_db): +def test_get_tracer_database_error(orm_session): # Test get_tracer handles database errors properly. missing_country_id = "" valid_householdID = "71424" @@ -58,4 +56,7 @@ def test_get_tracer_database_error(test_db): valid_api_version, ] with pytest.raises(Exception): - tracer_service.get_tracer(*missing_parameter_causing_database_exception) + tracer_service.get_tracer( + orm_session, + *missing_parameter_causing_database_exception, + ) From 80e0f750b2319f695498c5dae9be581ef0c838d1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 02:30:19 +0300 Subject: [PATCH 52/89] refactor: use ORM sessions for simulations and runs --- policyengine_api/routes/simulation_routes.py | 126 ++-- .../services/report_output_service.py | 3 - .../services/simulation_run_service.py | 140 ++-- .../services/simulation_service.py | 240 ++++--- .../services/simulation_spec_service.py | 102 +-- tests/contract/test_v1_route_contracts.py | 20 +- .../routes/test_route_exception_handling.py | 5 +- .../services/test_simulation_run_service.py | 288 +++----- .../unit/services/test_simulation_service.py | 632 +++++------------- .../services/test_simulation_spec_service.py | 213 ++---- 10 files changed, 633 insertions(+), 1136 deletions(-) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index 06f6fe7d8..bdf896d4f 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -5,6 +5,8 @@ import jsonschema import pydantic +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Simulation from policyengine_api.services.simulation_service import SimulationService from policyengine_api.utils.payload_validators import validate_country @@ -12,10 +14,13 @@ simulation_service = SimulationService() -def _serialize_v1_simulation(simulation: dict) -> dict: +def _serialize_v1_simulation(simulation: Simulation) -> dict: """Project canonical ORM JSON objects onto the legacy v1 response shape.""" - response = dict(simulation) + response = { + column.name: getattr(simulation, column.name) + for column in Simulation.__table__.columns + } for field in ("output", "simulation_spec_json"): value = response.get(field) if value is not None and not isinstance(value, str): @@ -65,49 +70,59 @@ def create_simulation(country_id: str) -> Response: raise BadRequest("policy_id must be an integer") try: - # Check if simulation already exists with these parameters - existing_simulation = simulation_service.find_existing_simulation( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - - if existing_simulation: - existing_simulation = simulation_service.ensure_simulation_dual_write_state( - existing_simulation["id"], + with get_v1_session_factory().begin() as session: + existing_simulation = simulation_service.find_existing_simulation( + session, country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, ) + + if existing_simulation: + simulation = simulation_service.ensure_simulation_dual_write_state( + session, + existing_simulation.id, + country_id=country_id, + ) + result = _serialize_v1_simulation(simulation) + message = "Simulation already exists" + status_code = 200 + else: + simulation = simulation_service.create_simulation( + session, + country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + ) + result = _serialize_v1_simulation(simulation) + message = "Simulation created successfully" + status_code = 201 + + if existing_simulation: # Simulation already exists, return it with 200 status response_body = dict( status="ok", - message="Simulation already exists", - result=_serialize_v1_simulation(existing_simulation), + message=message, + result=result, ) return Response( json.dumps(response_body), - status=200, + status=status_code, mimetype="application/json", ) - # Create new simulation - created_simulation = simulation_service.create_simulation( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - response_body = dict( status="ok", - message="Simulation created successfully", - result=_serialize_v1_simulation(created_simulation), + message=message, + result=result, ) return Response( json.dumps(response_body), - status=201, + status=status_code, mimetype="application/json", ) @@ -144,17 +159,21 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: if simulation_id <= 0: raise BadRequest("simulation_id must be a positive integer") - simulation: dict | None = simulation_service.get_simulation( - country_id, simulation_id - ) + with get_v1_session_factory()() as session: + simulation = simulation_service.get_simulation( + session, + country_id, + simulation_id, + ) + result = None if simulation is None else _serialize_v1_simulation(simulation) - if simulation is None: + if result is None: raise NotFound(f"Simulation #{simulation_id} not found.") response_body = dict( status="ok", message=None, - result=_serialize_v1_simulation(simulation), + result=result, ) return Response( @@ -201,34 +220,33 @@ def update_simulation(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - # First check if the simulation exists - existing_simulation = simulation_service.get_simulation( - country_id, simulation_id - ) - if existing_simulation is None: - raise NotFound(f"Simulation #{simulation_id} not found.") - - # Update the simulation - success = simulation_service.update_simulation( - country_id=country_id, - simulation_id=simulation_id, - status=status, - output=output, - error_message=error_message, - ) + with get_v1_session_factory().begin() as session: + existing_simulation = simulation_service.get_simulation( + session, + country_id, + simulation_id, + ) + if existing_simulation is None: + raise NotFound(f"Simulation #{simulation_id} not found.") - if not success: - raise BadRequest("No fields to update") + success = simulation_service.update_simulation( + session, + country_id=country_id, + simulation_id=simulation_id, + status=status, + output=output, + error_message=error_message, + ) - # Get the updated record - updated_simulation = simulation_service.get_simulation( - country_id, simulation_id - ) + if not success: + raise BadRequest("No fields to update") + + result = _serialize_v1_simulation(existing_simulation) response_body = dict( status="ok", message="Simulation updated successfully", - result=_serialize_v1_simulation(updated_simulation), + result=result, ) return Response( diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 39b2130a9..539d5d2a8 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -16,14 +16,12 @@ select_display_report_run, serialize_json_field, ) -from policyengine_api.services.simulation_service import SimulationService class ReportOutputService: def __init__(self, *, unit_of_work: V1UnitOfWork | None = None): self._unit_of_work = unit_of_work self.report_spec_service = ReportSpecService(unit_of_work=unit_of_work) - self.simulation_service = SimulationService(unit_of_work=unit_of_work) @property def unit_of_work(self) -> V1UnitOfWork: @@ -32,7 +30,6 @@ def unit_of_work(self) -> V1UnitOfWork: self.report_spec_service = ReportSpecService( unit_of_work=self._unit_of_work ) - self.simulation_service = SimulationService(unit_of_work=self._unit_of_work) return self._unit_of_work def _utc_timestamp(self) -> datetime: diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index c4377ca23..61b3a970b 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -1,9 +1,10 @@ -import json import uuid from typing import Any -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork +from sqlalchemy import func, select +from sqlalchemy.orm import Session + +from policyengine_api.data.v1_models import Simulation, SimulationRun SIMULATION_RUN_VERSION_FIELDS = ( @@ -16,33 +17,9 @@ class SimulationRunService: - def __init__( - self, - simulations: SimulationDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._simulations = simulations - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - - def _parse_run_row(self, row: dict | None) -> dict | None: - if row is None: - return None - run = dict(row) - if isinstance(run.get("simulation_spec_snapshot_json"), str): - run["simulation_spec_snapshot_json"] = json.loads( - run["simulation_spec_snapshot_json"] - ) - return run - def create_simulation_run( self, + session: Session, simulation_id: int, report_output_run_id: str | None = None, input_position: int | None = None, @@ -54,7 +31,20 @@ def create_simulation_run( simulation_spec_snapshot: dict[str, Any] | str | None = None, version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, - ) -> dict: + ) -> SimulationRun: + parent = session.scalar( + select(Simulation).where(Simulation.id == simulation_id).with_for_update() + ) + if parent is None: + raise ValueError(f"Simulation #{simulation_id} not found") + sequence = ( + session.scalar( + select(func.max(SimulationRun.run_sequence)).where( + SimulationRun.simulation_id == simulation_id + ) + ) + or 0 + ) + 1 values = { "report_output_run_id": report_output_run_id, "input_position": input_position, @@ -74,53 +64,61 @@ def create_simulation_run( for field in SIMULATION_RUN_VERSION_FIELDS } ) - try: - if self._simulations is not None: - run = self._simulations.create_run( - simulation_id, - run_id=run_id or str(uuid.uuid4()), - **values, - ) - else: - with self.unit_of_work.transaction() as daos: - run = daos.simulations.create_run( - simulation_id, - run_id=run_id or str(uuid.uuid4()), - **values, - ) - except LookupError as error: - raise ValueError(f"Simulation #{simulation_id} not found") from error - return self._parse_run_row(run) + run = SimulationRun( + id=run_id or str(uuid.uuid4()), + simulation_id=simulation_id, + run_sequence=sequence, + **values, + ) + session.add(run) + session.flush() + return run - def get_simulation_run(self, run_id: str) -> dict | None: - if self._simulations is not None: - return self._parse_run_row(self._simulations.get_run(run_id)) - with self.unit_of_work.read() as daos: - return self._parse_run_row(daos.simulations.get_run(run_id)) + def get_simulation_run( + self, + session: Session, + run_id: str, + ) -> SimulationRun | None: + return session.get(SimulationRun, run_id) - def list_simulation_runs(self, simulation_id: int) -> list[dict]: - if self._simulations is not None: - rows = self._simulations.list_runs(simulation_id) - else: - with self.unit_of_work.read() as daos: - rows = daos.simulations.list_runs(simulation_id) - return [self._parse_run_row(row) for row in reversed(rows)] + def list_simulation_runs( + self, + session: Session, + simulation_id: int, + ) -> list[SimulationRun]: + return list( + session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.asc()) + ) + ) - def get_newest_simulation_run(self, simulation_id: int) -> dict | None: - if self._simulations is not None: - rows = self._simulations.list_runs(simulation_id) - else: - with self.unit_of_work.read() as daos: - rows = daos.simulations.list_runs(simulation_id) - return self._parse_run_row(rows[0]) if rows else None + def get_newest_simulation_run( + self, + session: Session, + simulation_id: int, + ) -> SimulationRun | None: + return session.scalar( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) - def select_display_run(self, simulation: dict) -> dict | None: - if simulation.get("active_run_id"): - active_run = self.get_simulation_run(simulation["active_run_id"]) + def select_display_run( + self, + session: Session, + simulation: Simulation, + ) -> SimulationRun | None: + if simulation.active_run_id: + active_run = self.get_simulation_run(session, simulation.active_run_id) if active_run is not None: return active_run - if simulation.get("latest_successful_run_id"): - successful = self.get_simulation_run(simulation["latest_successful_run_id"]) + if simulation.latest_successful_run_id: + successful = self.get_simulation_run( + session, + simulation.latest_successful_run_id, + ) if successful is not None: return successful - return self.get_newest_simulation_run(simulation["id"]) + return self.get_newest_simulation_run(session, simulation.id) diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 56bfb4f9b..050486f74 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,135 +1,175 @@ import json +import uuid + +from sqlalchemy import select +from sqlalchemy.orm import Session from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork -from policyengine_api.services.simulation_spec_service import SimulationSpecService +from policyengine_api.data.v1_models import Simulation, SimulationRun class SimulationService: - def __init__( - self, - simulations: SimulationDAO | None = None, + """Simulation operations performed through a caller-owned ORM Session.""" + + @staticmethod + def _select_simulation( + session: Session, + simulation_id: int, + country_id: str | None = None, *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._simulations = simulations - self._unit_of_work = unit_of_work - self.simulation_spec_service = SimulationSpecService( - simulations, - unit_of_work=unit_of_work, - ) + for_update: bool = False, + ) -> Simulation | None: + statement = select(Simulation).where(Simulation.id == simulation_id) + if country_id is not None: + statement = statement.where(Simulation.country_id == country_id) + if for_update: + statement = statement.with_for_update() + return session.scalar(statement) - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - self.simulation_spec_service = SimulationSpecService( - unit_of_work=self._unit_of_work - ) - return self._unit_of_work + @staticmethod + def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: + return next((run.id for run in runs if run.status == "complete"), None) - def _ensure_simulation_dual_write_state_in_transaction( + def ensure_simulation_dual_write_state( self, - session, + session: Session, simulation_id: int, - *, country_id: str | None = None, - ) -> dict: - return SimulationDAO(session).ensure_dual_write_state_in_session( + ) -> Simulation: + simulation = self._select_simulation( session, simulation_id, country_id, + for_update=True, ) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") - def _get_simulation_row( - self, - simulation_id: int, - *, - queryer=None, - country_id: str | None = None, - for_update: bool = False, - ) -> dict | None: - del for_update - if queryer is not None: - simulations = getattr(queryer, "simulations", None) - if simulations is None: - simulations = SimulationDAO(getattr(queryer, "session", queryer)) - return simulations.get(simulation_id, country_id) - if self._simulations is not None: - return self._simulations.get(simulation_id, country_id) - with self.unit_of_work.read() as daos: - return daos.simulations.get(simulation_id, country_id) + spec = { + "country_id": simulation.country_id, + "population_id": simulation.population_id, + "population_type": simulation.population_type, + "policy_id": simulation.policy_id, + } + simulation.simulation_spec_json = spec + simulation.simulation_spec_schema_version = 1 + runs = list( + session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) + ) + if not runs: + run = SimulationRun( + id=str(uuid.uuid4()), + simulation_id=simulation_id, + run_sequence=1, + status=simulation.status, + output=simulation.output, + error_message=simulation.error_message, + trigger_type="initial", + simulation_spec_snapshot_json=spec, + country_package_version=simulation.api_version, + ) + session.add(run) + session.flush() + runs = [run] + else: + mutable = next( + (run for run in runs if run.id == simulation.active_run_id), + runs[0], + ) + mutable.status = simulation.status + mutable.output = simulation.output + mutable.error_message = simulation.error_message + mutable.simulation_spec_snapshot_json = spec + mutable.country_package_version = simulation.api_version - def ensure_simulation_dual_write_state( - self, simulation_id: int, country_id: str | None = None - ) -> dict: - if self._simulations is not None: - return self._simulations.ensure_dual_write_state(simulation_id, country_id) - with self.unit_of_work.transaction() as daos: - return daos.simulations.ensure_dual_write_state(simulation_id, country_id) + latest_successful = self._latest_successful_run_id(runs) + simulation.active_run_id = ( + runs[0].id if simulation.status in {"pending", "running"} else None + ) + if simulation.status == "complete" and latest_successful is None: + latest_successful = runs[0].id + simulation.latest_successful_run_id = latest_successful + session.flush() + return simulation def find_existing_simulation( self, + session: Session, country_id: str, population_id: str, population_type: str, policy_id: int, - ) -> dict | None: - if self._simulations is not None: - return self._simulations.find_latest( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - with self.unit_of_work.read() as daos: - return daos.simulations.find_latest( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, + ) -> Simulation | None: + return session.scalar( + select(Simulation) + .where( + Simulation.country_id == country_id, + Simulation.population_id == population_id, + Simulation.population_type == population_type, + Simulation.policy_id == policy_id, ) + .order_by(Simulation.id.desc()) + ) def create_simulation( self, + session: Session, country_id: str, population_id: str, population_type: str, policy_id: int, - ) -> dict: - values = { - "country_id": country_id, - "api_version": COUNTRY_PACKAGE_VERSIONS.get(country_id), - "population_id": population_id, - "population_type": population_type, - "policy_id": policy_id, - "status": "pending", - } - if self._simulations is not None: - return self._simulations.create_or_get_with_sync( - sync_callback=self._ensure_simulation_dual_write_state_in_transaction, - **values, + ) -> Simulation: + simulation = session.scalar( + select(Simulation) + .where( + Simulation.country_id == country_id, + Simulation.population_id == population_id, + Simulation.population_type == population_type, + Simulation.policy_id == policy_id, ) - with self.unit_of_work.transaction() as daos: - return daos.simulations.create_or_get_with_sync( - sync_callback=self._ensure_simulation_dual_write_state_in_transaction, - **values, + .order_by(Simulation.id.desc()) + .with_for_update() + ) + if simulation is None: + simulation = Simulation( + country_id=country_id, + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", ) + session.add(simulation) + session.flush() + return self.ensure_simulation_dual_write_state( + session, + simulation.id, + country_id, + ) - def get_simulation(self, country_id: str, simulation_id: int) -> dict | None: + def get_simulation( + self, + session: Session, + country_id: str, + simulation_id: int, + ) -> Simulation | None: if type(simulation_id) is not int or simulation_id < 0: raise Exception( f"Invalid simulation ID: {simulation_id}. Must be a positive integer." ) - return self._get_simulation_row(simulation_id, country_id=country_id) + return self._select_simulation(session, simulation_id, country_id) def update_simulation( self, + session: Session, country_id: str, simulation_id: int, status: str | None = None, - output: str | None = None, + output: dict | list | str | None = None, error_message: str | None = None, ) -> bool: values = { @@ -145,20 +185,16 @@ def update_simulation( return False if isinstance(values.get("output"), str): values["output"] = json.loads(values["output"]) - values["api_version"] = COUNTRY_PACKAGE_VERSIONS.get(country_id) - if self._simulations is not None: - self._simulations.update_with_sync( - simulation_id, - country_id, - values, - self._ensure_simulation_dual_write_state_in_transaction, - ) - else: - with self.unit_of_work.transaction() as daos: - daos.simulations.update_with_sync( - simulation_id, - country_id, - values, - self._ensure_simulation_dual_write_state_in_transaction, - ) + simulation = self._select_simulation( + session, + simulation_id, + country_id, + for_update=True, + ) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") + for key, value in values.items(): + setattr(simulation, key, value) + simulation.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) + self.ensure_simulation_dual_write_state(session, simulation_id, country_id) return True diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index 0320247e3..f6c5a7373 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -2,8 +2,9 @@ from typing import Literal from pydantic import BaseModel -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork +from sqlalchemy.orm import Session + +from policyengine_api.data.v1_models import Simulation SIMULATION_SPEC_SCHEMA_VERSION = 1 @@ -16,96 +17,65 @@ class SimulationSpec(BaseModel): class SimulationSpecService: - def __init__( - self, - simulations: SimulationDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._simulations = simulations - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != SIMULATION_SPEC_SCHEMA_VERSION: raise ValueError( f"Unsupported simulation spec schema version: {schema_version}" ) - def _get_simulation_row(self, simulation_id: int) -> dict | None: - if self._simulations is not None: - return self._simulations.get(simulation_id) - with self.unit_of_work.read() as daos: - return daos.simulations.get(simulation_id) - - def _validate_simulation_spec_matches_row( - self, simulation: dict, simulation_spec: SimulationSpec + def _validate_simulation_spec_matches_model( + self, + simulation: Simulation, + simulation_spec: SimulationSpec, ) -> None: expected_spec = { - "country_id": simulation["country_id"], - "population_id": simulation["population_id"], - "population_type": simulation["population_type"], - "policy_id": simulation["policy_id"], + "country_id": simulation.country_id, + "population_id": simulation.population_id, + "population_type": simulation.population_type, + "policy_id": simulation.policy_id, } if simulation_spec.model_dump() != expected_spec: raise ValueError("Simulation spec must match the linked simulation row") - def build_simulation_spec(self, simulation: dict) -> SimulationSpec: - return SimulationSpec.model_validate( - { - "country_id": simulation["country_id"], - "population_id": simulation["population_id"], - "population_type": simulation["population_type"], - "policy_id": simulation["policy_id"], - } + def build_simulation_spec(self, simulation: Simulation) -> SimulationSpec: + return SimulationSpec( + country_id=simulation.country_id, + population_id=simulation.population_id, + population_type=simulation.population_type, + policy_id=simulation.policy_id, ) - def get_simulation_spec(self, simulation_id: int) -> SimulationSpec | None: - simulation = self._get_simulation_row(simulation_id) - if simulation is None or simulation["simulation_spec_json"] is None: + def get_simulation_spec( + self, + session: Session, + simulation_id: int, + ) -> SimulationSpec | None: + simulation = session.get(Simulation, simulation_id) + if simulation is None or simulation.simulation_spec_json is None: return None - self._validate_schema_version(simulation["simulation_spec_schema_version"]) - raw_spec = simulation["simulation_spec_json"] + self._validate_schema_version(simulation.simulation_spec_schema_version) + raw_spec = simulation.simulation_spec_json if isinstance(raw_spec, str): + # Existing databases may contain pre-ORM JSON text. New writes below + # always assign Python objects and leave conversion to SQLAlchemy. raw_spec = json.loads(raw_spec) simulation_spec = SimulationSpec.model_validate(raw_spec) - self._validate_simulation_spec_matches_row(simulation, simulation_spec) + self._validate_simulation_spec_matches_model(simulation, simulation_spec) return simulation_spec def set_simulation_spec( self, + session: Session, simulation_id: int, simulation_spec: SimulationSpec, schema_version: int = SIMULATION_SPEC_SCHEMA_VERSION, ) -> bool: self._validate_schema_version(schema_version) - if self._simulations is not None: - simulations = self._simulations - simulation = simulations.get(simulation_id) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - self._validate_simulation_spec_matches_row(simulation, simulation_spec) - simulations.update( - simulation_id, - simulation_spec_json=simulation_spec.model_dump(), - simulation_spec_schema_version=schema_version, - ) - return True - - with self.unit_of_work.transaction() as daos: - simulation = daos.simulations.get(simulation_id) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - self._validate_simulation_spec_matches_row(simulation, simulation_spec) - daos.simulations.update( - simulation_id, - simulation_spec_json=simulation_spec.model_dump(), - simulation_spec_schema_version=schema_version, - ) + simulation = session.get(Simulation, simulation_id) + if simulation is None: + raise ValueError(f"Simulation #{simulation_id} not found") + self._validate_simulation_spec_matches_model(simulation, simulation_spec) + simulation.simulation_spec_json = simulation_spec.model_dump() + simulation.simulation_spec_schema_version = schema_version return True diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index dead34e64..382fa30e5 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -9,7 +9,7 @@ from policyengine_api.endpoints.household import get_calculate from policyengine_api.endpoints.policy import get_policy_search -from policyengine_api.data.v1_models import Household, Policy +from policyengine_api.data.v1_models import Household, Policy, Simulation from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp @@ -323,20 +323,20 @@ def _patched_route_dependencies(): stack.enter_context( patch( "policyengine_api.routes.simulation_routes.simulation_service.create_simulation", - return_value={ - "id": 11, - "country_id": "us", - "population_id": "household-1", - "population_type": "household", - "policy_id": 22, - "status": "pending", - }, + return_value=Simulation( + id=11, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=22, + status="pending", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.simulation_routes.simulation_service.get_simulation", - return_value={"id": 11, "status": "pending", "country_id": "us"}, + return_value=Simulation(id=11, status="pending", country_id="us"), ) ) stack.enter_context( diff --git a/tests/unit/routes/test_route_exception_handling.py b/tests/unit/routes/test_route_exception_handling.py index b6fb38a28..4a1848cb9 100644 --- a/tests/unit/routes/test_route_exception_handling.py +++ b/tests/unit/routes/test_route_exception_handling.py @@ -86,7 +86,7 @@ def test_report_create_value_error_still_400(): assert response.status_code == 400 -def test_simulation_patch_empty_body_returns_400(test_db): +def test_simulation_patch_empty_body_returns_400(orm_session): """Regression for issue #3449. PATCH /{country}/simulation with a body that only contains the @@ -97,6 +97,7 @@ def test_simulation_patch_empty_body_returns_400(test_db): simulation_service = SimulationService() created = simulation_service.create_simulation( + orm_session, country_id="us", population_id="household_patch_empty", population_type="household", @@ -104,5 +105,5 @@ def test_simulation_patch_empty_body_returns_400(test_db): ) client = _client_with(simulation_bp) - response = client.patch("/us/simulation", json={"id": created["id"]}) + response = client.patch("/us/simulation", json={"id": created.id}) assert response.status_code == 400 diff --git a/tests/unit/services/test_simulation_run_service.py b/tests/unit/services/test_simulation_run_service.py index f47e13a0b..275373f7e 100644 --- a/tests/unit/services/test_simulation_run_service.py +++ b/tests/unit/services/test_simulation_run_service.py @@ -1,214 +1,90 @@ import pytest +from policyengine_api.data.v1_models import SimulationRun from policyengine_api.services.simulation_run_service import SimulationRunService from policyengine_api.services.simulation_service import SimulationService -simulation_run_service = SimulationRunService() + +run_service = SimulationRunService() simulation_service = SimulationService() -class TestCreateSimulationRun: - def test_creates_simulation_runs_with_incrementing_sequence(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - - first_run = simulation_run_service.create_simulation_run( - simulation["id"], - input_position=1, - trigger_type="initial", - simulation_spec_snapshot={"population_id": "household_1"}, - version_manifest={"simulation_cache_version": "s123"}, - ) - second_run = simulation_run_service.create_simulation_run( - simulation["id"], - input_position=1, - trigger_type="rerun", - ) - - assert first_run["run_sequence"] == 2 - assert first_run["trigger_type"] == "initial" - assert first_run["simulation_spec_snapshot_json"] == { - "population_id": "household_1" - } - assert first_run["simulation_cache_version"] == "s123" - assert second_run["run_sequence"] == 3 - assert second_run["trigger_type"] == "rerun" - - def test_allocates_run_sequence_transactionally(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1a", - population_type="household", - policy_id=1, - ) - - first_run = simulation_run_service.create_simulation_run( - simulation["id"], input_position=1, trigger_type="initial" - ) - second_run = simulation_run_service.create_simulation_run( - simulation["id"], input_position=1, trigger_type="rerun" - ) - - assert first_run["run_sequence"] == 2 - assert second_run["run_sequence"] == 3 - - def test_raises_when_parent_simulation_is_missing(self, test_db): - with pytest.raises(ValueError) as exc_info: - simulation_run_service.create_simulation_run( - 999999, input_position=1, trigger_type="initial" - ) - - assert "Simulation #999999 not found" in str(exc_info.value) - - -class TestSelectDisplaySimulationRun: - def test_prefers_active_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - latest_successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - active_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (active_run["id"], latest_successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == active_run["id"] - - def test_falls_back_to_latest_successful_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - (successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_when_active_run_pointer_is_stale(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3a", - population_type="household", - policy_id=3, - ) - successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_to_newest_run_when_no_pointers_exist(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - first_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - newest_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - (simulation["id"],), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert first_run["run_sequence"] == 2 - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_latest_successful_pointer_is_stale( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4a", - population_type="household", - policy_id=4, - ) - newest_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == newest_run["id"] +def create_simulation(orm_session, population_id="household-1"): + return simulation_service.create_simulation( + orm_session, "us", population_id, "household", 1 + ) + + +def test_creates_mapped_runs_with_incrementing_sequence(orm_session): + simulation = create_simulation(orm_session) + + first = run_service.create_simulation_run( + orm_session, + simulation.id, + input_position=1, + trigger_type="rerun", + simulation_spec_snapshot={"population_id": "household-1"}, + version_manifest={"simulation_cache_version": "s123"}, + ) + second = run_service.create_simulation_run( + orm_session, simulation.id, input_position=1, trigger_type="rerun" + ) + + assert isinstance(first, SimulationRun) + assert first.run_sequence == 2 + assert first.simulation_spec_snapshot_json == {"population_id": "household-1"} + assert first.simulation_cache_version == "s123" + assert second.run_sequence == 3 + + +def test_raises_when_parent_simulation_is_missing(orm_session): + with pytest.raises(ValueError, match="Simulation #999999 not found"): + run_service.create_simulation_run(orm_session, 999999) + + +def test_gets_and_lists_runs_as_models(orm_session): + simulation = create_simulation(orm_session) + first = run_service.get_simulation_run(orm_session, simulation.active_run_id) + second = run_service.create_simulation_run(orm_session, simulation.id) + + assert run_service.get_simulation_run(orm_session, second.id) is second + assert run_service.list_simulation_runs(orm_session, simulation.id) == [ + first, + second, + ] + assert run_service.get_newest_simulation_run(orm_session, simulation.id) is second + + +def test_display_run_prefers_active_run(orm_session): + simulation = create_simulation(orm_session) + successful = run_service.create_simulation_run( + orm_session, simulation.id, status="complete" + ) + active = run_service.create_simulation_run( + orm_session, simulation.id, status="running" + ) + simulation.latest_successful_run_id = successful.id + simulation.active_run_id = active.id + + assert run_service.select_display_run(orm_session, simulation) is active + + +def test_display_run_falls_back_to_latest_successful_run(orm_session): + simulation = create_simulation(orm_session) + successful = run_service.create_simulation_run( + orm_session, simulation.id, status="complete" + ) + run_service.create_simulation_run(orm_session, simulation.id) + simulation.active_run_id = None + simulation.latest_successful_run_id = successful.id + + assert run_service.select_display_run(orm_session, simulation) is successful + + +def test_display_run_falls_back_to_newest_for_stale_pointers(orm_session): + simulation = create_simulation(orm_session) + newest = run_service.create_simulation_run(orm_session, simulation.id) + simulation.active_run_id = "missing-active-run" + simulation.latest_successful_run_id = "missing-successful-run" + + assert run_service.select_display_run(orm_session, simulation) is newest diff --git a/tests/unit/services/test_simulation_service.py b/tests/unit/services/test_simulation_service.py index 34116287f..2902f297e 100644 --- a/tests/unit/services/test_simulation_service.py +++ b/tests/unit/services/test_simulation_service.py @@ -1,522 +1,188 @@ import pytest -import json +from policyengine_api.data.v1_models import Simulation, SimulationRun from policyengine_api.services.simulation_service import SimulationService -from tests.fixtures.services import simulation_fixtures - -pytest_plugins = ("tests.fixtures.services.simulation_fixtures",) service = SimulationService() -class TestFindExistingSimulation: - """Test finding existing simulations in the database.""" - - def test_find_existing_simulation_given_existing_record( - self, test_db, existing_simulation_record - ): - """Test that find_existing_simulation returns the existing simulation.""" - # GIVEN an existing simulation record (from fixture) - - # WHEN we search for a simulation with matching parameters - result = service.find_existing_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - population_id=simulation_fixtures.valid_simulation_data["population_id"], - population_type=simulation_fixtures.valid_simulation_data[ - "population_type" - ], - policy_id=simulation_fixtures.valid_simulation_data["policy_id"], - ) - - # THEN the result should contain the existing simulation - assert result is not None - assert result["id"] == existing_simulation_record["id"] - assert ( - result["country_id"] - == simulation_fixtures.valid_simulation_data["country_id"] - ) - assert ( - result["population_id"] - == simulation_fixtures.valid_simulation_data["population_id"] - ) - assert ( - result["policy_id"] - == simulation_fixtures.valid_simulation_data["policy_id"] - ) - - def test_find_existing_simulation_given_no_match(self, test_db): - """Test that find_existing_simulation returns None when no match exists.""" - # GIVEN an empty database (default test state) - - # WHEN we search for a non-existent simulation - result = service.find_existing_simulation( +def test_finds_existing_simulation_without_api_version_matching(orm_session): + existing = Simulation( + country_id="us", + api_version="old-version", + population_id="household-1", + population_type="household", + policy_id=1, + status="pending", + ) + orm_session.add(existing) + orm_session.flush() + + result = service.find_existing_simulation( + orm_session, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=1, + ) + + assert result is existing + + +def test_returns_none_when_simulation_does_not_exist(orm_session): + assert ( + service.find_existing_simulation( + orm_session, country_id="uk", - population_id="nonexistent_123", + population_id="missing", population_type="household", policy_id=999, ) - - # THEN the result should be None - assert result is None - - def test_find_existing_simulation_ignores_api_version( - self, test_db, existing_simulation_record - ): - """Test that simulations are found regardless of API version.""" - # GIVEN an existing simulation record - - # WHEN we search for the same simulation (API version is ignored) - result = service.find_existing_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - population_id=simulation_fixtures.valid_simulation_data["population_id"], - population_type=simulation_fixtures.valid_simulation_data[ - "population_type" - ], - policy_id=simulation_fixtures.valid_simulation_data["policy_id"], - ) - - # THEN the existing record should be found (API version ignored) - assert result is not None - assert result["id"] == existing_simulation_record["id"] - - -class TestCreateSimulation: - """Test creating new simulations in the database.""" - - def test_create_simulation_success(self, test_db): - """Test successful creation of a new simulation.""" - # GIVEN an empty database - - # WHEN we create a new simulation - created_simulation = service.create_simulation( - country_id="us", - population_id="household_123", - population_type="household", - policy_id=1, - ) - - # THEN a valid simulation record should be returned - assert created_simulation is not None - assert isinstance(created_simulation, dict) - assert created_simulation["id"] > 0 - assert created_simulation["country_id"] == "us" - assert created_simulation["population_id"] == "household_123" - - # AND the simulation should be retrievable from database - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert result is not None - assert result["country_id"] == "us" - assert result["population_id"] == "household_123" - - def test_create_simulation_with_geography_type(self, test_db): - """Test creating a simulation with geography population type.""" - # GIVEN an empty database - - # WHEN we create a simulation with geography type - created_simulation = service.create_simulation( - country_id="uk", - population_id="geo_code_456", - population_type="geography", - policy_id=2, - ) - - # THEN the simulation should be created successfully - assert created_simulation is not None - assert created_simulation["population_type"] == "geography" - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert result["population_type"] == "geography" - - def test_create_simulation_retrieves_correct_id(self, test_db): - """Test that create_simulation retrieves the correct ID without race conditions.""" - # GIVEN we create multiple simulations rapidly - - # WHEN we create simulations with different parameters - created_sims = [] - for i in range(3): - sim = service.create_simulation( + is None + ) + + +def test_creates_mapped_simulation_and_initial_run(orm_session): + simulation = service.create_simulation( + orm_session, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=1, + ) + + assert isinstance(simulation, Simulation) + assert simulation.simulation_spec_json == { + "country_id": "us", + "population_id": "household-1", + "population_type": "household", + "policy_id": 1, + } + assert simulation.simulation_spec_schema_version == 1 + run = orm_session.get(SimulationRun, simulation.active_run_id) + assert isinstance(run, SimulationRun) + assert run.status == "pending" + assert run.trigger_type == "initial" + assert run.simulation_spec_snapshot_json == simulation.simulation_spec_json + + +def test_creation_reuses_existing_row_and_bootstraps_dual_write_state(orm_session): + existing = Simulation( + country_id="us", + api_version="old-version", + population_id="household-1", + population_type="household", + policy_id=7, + status="pending", + ) + orm_session.add(existing) + orm_session.flush() + + result = service.create_simulation( + orm_session, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=7, + ) + + assert result is existing + run = orm_session.get(SimulationRun, existing.active_run_id) + assert isinstance(run, SimulationRun) + assert run.simulation_id == existing.id + + +def test_caller_transaction_rolls_back_creation_on_dual_write_failure( + orm_session_factory, monkeypatch +): + def fail_dual_write(*args, **kwargs): + raise RuntimeError("dual write sync failed") + + monkeypatch.setattr(service, "ensure_simulation_dual_write_state", fail_dual_write) + + with pytest.raises(RuntimeError, match="dual write sync failed"): + with orm_session_factory.begin() as session: + service.create_simulation( + session, country_id="us", - population_id=f"household_{i}", + population_id="rollback", population_type="household", - policy_id=i, + policy_id=8, ) - created_sims.append(sim) - - # THEN all IDs should be unique and sequential - ids = [sim["id"] for sim in created_sims] - assert len(set(ids)) == 3 # All IDs are unique - assert ids == sorted(ids) # IDs are in order - - # AND each simulation should have the correct data - for i, sim in enumerate(created_sims): - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", (sim["id"],) - ).fetchone() - assert result["population_id"] == f"household_{i}" - assert result["policy_id"] == i - - def test_create_simulation_populates_dual_write_state(self, test_db): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_dual_write", - population_type="household", - policy_id=3, - ) - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["simulation_spec_schema_version"] == 1 - assert stored_simulation["active_run_id"] is not None - assert stored_simulation["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["trigger_type"] == "initial" - snapshot = run["simulation_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["population_id"] == "household_dual_write" - assert snapshot["policy_id"] == 3 - - def test_create_simulation_reuses_existing_row_and_bootstraps_dual_write( - self, test_db - ): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_bootstrap", "household", 7, "pending"), - ) - - created_simulation = service.create_simulation( - country_id="us", - population_id="household_bootstrap", - population_type="household", - policy_id=7, - ) - - rows = test_db.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - """, - ("us", "household_bootstrap", "household", 7), - ).fetchall() - assert len(rows) == 1 - assert created_simulation["id"] == rows[0]["id"] - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - - def test_create_simulation_rolls_back_parent_insert_on_dual_write_failure( - self, test_db, monkeypatch - ): - def fail_dual_write(tx, simulation_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_simulation_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.create_simulation( + with orm_session_factory() as session: + assert ( + service.find_existing_simulation( + session, country_id="us", - population_id="household_create_rollback", + population_id="rollback", population_type="household", policy_id=8, ) - - rows = test_db.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - """, - ("us", "household_create_rollback", "household", 8), - ).fetchall() - assert rows == [] - - -class TestGetSimulation: - """Test retrieving simulations from the database.""" - - def test_get_simulation_existing(self, test_db, existing_simulation_record): - """Test retrieving an existing simulation.""" - # GIVEN an existing simulation record - - # WHEN we retrieve the simulation - result = service.get_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - simulation_id=existing_simulation_record["id"], - ) - - # THEN the correct simulation should be returned - assert result is not None - assert result["id"] == existing_simulation_record["id"] - assert ( - result["country_id"] - == simulation_fixtures.valid_simulation_data["country_id"] + is None ) - def test_get_simulation_nonexistent(self, test_db): - """Test retrieving a non-existent simulation returns None.""" - # GIVEN an empty database - - # WHEN we try to retrieve a non-existent simulation - result = service.get_simulation(country_id="us", simulation_id=999) - - # THEN None should be returned - assert result is None - - def test_get_simulation_wrong_country(self, test_db, existing_simulation_record): - """Test that simulations are country-specific.""" - # GIVEN an existing simulation for 'us' - # WHEN we try to retrieve it with a different country - result = service.get_simulation( - country_id="uk", # Wrong country - simulation_id=existing_simulation_record["id"], - ) +def test_get_simulation_returns_model_scoped_to_country(orm_session): + simulation = service.create_simulation( + orm_session, "us", "household-1", "household", 1 + ) - # THEN None should be returned - assert result is None + assert service.get_simulation(orm_session, "us", simulation.id) is simulation + assert service.get_simulation(orm_session, "uk", simulation.id) is None - def test_get_simulation_invalid_id(self, test_db): - """Test that invalid simulation IDs are handled properly.""" - # GIVEN any database state - # WHEN we call get_simulation with invalid ID types - # THEN an exception should be raised - with pytest.raises(Exception) as exc_info: - service.get_simulation(country_id="us", simulation_id=-1) - assert "Invalid simulation ID" in str(exc_info.value) +@pytest.mark.parametrize("simulation_id", [-1, "1", None]) +def test_get_simulation_rejects_invalid_ids(orm_session, simulation_id): + with pytest.raises(Exception, match="Invalid simulation ID"): + service.get_simulation(orm_session, "us", simulation_id) - with pytest.raises(Exception) as exc_info: - service.get_simulation(country_id="us", simulation_id="not_an_int") - assert "Invalid simulation ID" in str(exc_info.value) +def test_update_simulation_updates_model_and_run_with_python_json(orm_session): + simulation = service.create_simulation( + orm_session, "us", "household-1", "household", 1 + ) -class TestUniqueConstraint: - """Test that the unique constraint on simulations works correctly.""" + updated = service.update_simulation( + orm_session, + "us", + simulation.id, + status="complete", + output={"result": 42}, + ) - def test_duplicate_simulation_returns_existing(self, test_db): - """Test that creating duplicate simulations returns the existing record.""" - # GIVEN we create a simulation - first_simulation = service.create_simulation( - country_id="us", - population_id="household_123", - population_type="household", - policy_id=1, - ) + assert updated is True + assert simulation.output == {"result": 42} + assert simulation.active_run_id is None + run = orm_session.get(SimulationRun, simulation.latest_successful_run_id) + assert run.output == {"result": 42} + assert run.status == "complete" - # WHEN we try to create an identical simulation - second_simulation = service.create_simulation( - country_id="us", - population_id="household_123", - population_type="household", - policy_id=1, - ) - # THEN the same simulation should be returned (no duplicate created) - assert first_simulation["id"] == second_simulation["id"] - assert first_simulation["country_id"] == second_simulation["country_id"] - assert first_simulation["population_id"] == second_simulation["population_id"] - assert first_simulation["policy_id"] == second_simulation["policy_id"] +def test_update_simulation_accepts_json_only_at_existing_wire_boundary(orm_session): + simulation = service.create_simulation( + orm_session, "us", "household-1", "household", 1 + ) + service.update_simulation( + orm_session, + "us", + simulation.id, + output='{"result": 42}', + ) -class TestUpdateSimulation: - def test_update_simulation_updates_dual_write_state(self, test_db): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_update", - population_type="household", - policy_id=11, - ) - output_json = json.dumps({"result": "ok"}) + assert simulation.output == {"result": 42} - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=output_json, - ) - assert success is True - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["active_run_id"] is None - assert stored_simulation["latest_successful_run_id"] is not None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run["status"] == "complete" - assert run["output"] == output_json - assert run["id"] == stored_simulation["latest_successful_run_id"] - - def test_update_simulation_bootstraps_missing_run_state(self, test_db): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_legacy", "household", 13, "pending"), - ) - simulation = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_simulation( - country_id="us", - simulation_id=simulation["id"], - status="error", - error_message="legacy failure", - ) +def test_update_simulation_without_values_is_a_noop(orm_session): + simulation = service.create_simulation( + orm_session, "us", "household-1", "household", 1 + ) - assert success is True - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["active_run_id"] is None - assert stored_simulation["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "error" - assert run["error_message"] == "legacy failure" - - def test_update_simulation_does_not_append_extra_run_for_legacy_patch_traffic( - self, test_db - ): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_single_run", - population_type="household", - policy_id=14, - ) - - first_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert first_run is not None - - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=json.dumps({"value": 1}), - ) - - assert success is True - - runs = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ? ORDER BY run_sequence", - (created_simulation["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["id"] == first_run["id"] - assert runs[0]["status"] == "complete" - - def test_update_simulation_rolls_back_parent_update_on_dual_write_failure( - self, test_db, monkeypatch - ): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_update_rollback", - population_type="household", - policy_id=15, - ) - - def fail_dual_write(tx, simulation_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_simulation_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=json.dumps({"rolled_back": True}), - ) - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["status"] == "pending" - assert stored_simulation["output"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["output"] is None - - def test_update_simulation_with_no_user_fields_returns_false(self, test_db): - """Regression for issue #3449. - - update_fields used to always append api_version, so a PATCH - with no status/output/error_message still passed the - "no fields to update" guard and rewrote the row. The guard - must fire before api_version is appended so an empty PATCH - returns False (and the route converts that to a 400). - """ - created_simulation = service.create_simulation( - country_id="us", - population_id="household_empty_patch", - population_type="household", - policy_id=16, - ) - - pre_row = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - ) + assert service.update_simulation(orm_session, "us", simulation.id) is False - assert success is False - post_row = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert post_row["api_version"] == pre_row["api_version"] - assert post_row["status"] == pre_row["status"] +def test_update_missing_simulation_raises(orm_session): + with pytest.raises(ValueError, match="Simulation #999 not found"): + service.update_simulation(orm_session, "us", 999, status="complete") diff --git a/tests/unit/services/test_simulation_spec_service.py b/tests/unit/services/test_simulation_spec_service.py index 89dce070c..7e19a24e1 100644 --- a/tests/unit/services/test_simulation_spec_service.py +++ b/tests/unit/services/test_simulation_spec_service.py @@ -6,160 +6,95 @@ SimulationSpecService, ) + simulation_service = SimulationService() -simulation_spec_service = SimulationSpecService() +spec_service = SimulationSpecService() -class TestSimulationSpecService: - def test_builds_simulation_spec_from_row(self, test_db): - simulation = simulation_service.create_simulation( - country_id="uk", - population_id="household_42", - population_type="household", - policy_id=7, - ) +def create_simulation(orm_session): + return simulation_service.create_simulation(orm_session, "us", "ca", "geography", 3) - simulation_spec = simulation_spec_service.build_simulation_spec(simulation) - assert isinstance(simulation_spec, SimulationSpec) - assert simulation_spec.country_id == "uk" - assert simulation_spec.population_id == "household_42" - assert simulation_spec.policy_id == 7 +def test_builds_spec_from_mapped_simulation(orm_session): + simulation = create_simulation(orm_session) - def test_sets_and_gets_simulation_spec(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ca", - "population_type": "geography", - "policy_id": 3, - } - ) + spec = spec_service.build_simulation_spec(simulation) - result = simulation_spec_service.set_simulation_spec( - simulation["id"], simulation_spec - ) + assert isinstance(spec, SimulationSpec) + assert spec.model_dump() == { + "country_id": "us", + "population_id": "ca", + "population_type": "geography", + "policy_id": 3, + } - assert result is True - stored_simulation = test_db.query( - """ - SELECT simulation_spec_json, simulation_spec_schema_version - FROM simulations WHERE id = ? - """, - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_schema_version"] == 1 - - loaded_simulation_spec = simulation_spec_service.get_simulation_spec( - simulation["id"] - ) - assert loaded_simulation_spec is not None - assert loaded_simulation_spec.model_dump() == simulation_spec.model_dump() - - def test_rejects_unsupported_schema_version_on_write(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ca", - "population_type": "geography", - "policy_id": 3, - } - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.set_simulation_spec( - simulation["id"], - simulation_spec, - schema_version=2, - ) - - assert "Unsupported simulation spec schema version" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_read(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - test_db.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - '{"country_id":"us","population_id":"ca","population_type":"geography","policy_id":3}', - 2, - simulation["id"], - ), - ) +def test_sets_and_gets_python_json_spec(orm_session): + simulation = create_simulation(orm_session) + spec = spec_service.build_simulation_spec(simulation) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.get_simulation_spec(simulation["id"]) + assert spec_service.set_simulation_spec(orm_session, simulation.id, spec) is True + loaded = spec_service.get_simulation_spec(orm_session, simulation.id) - assert "Unsupported simulation spec schema version" in str(exc_info.value) + assert simulation.simulation_spec_json == spec.model_dump() + assert loaded == spec - def test_rejects_simulation_spec_write_when_fields_do_not_match_row(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ny", - "population_type": "geography", - "policy_id": 3, - } - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.set_simulation_spec( - simulation["id"], simulation_spec - ) +def test_rejects_unsupported_schema_version_on_write(orm_session): + simulation = create_simulation(orm_session) + spec = spec_service.build_simulation_spec(simulation) - assert "Simulation spec must match the linked simulation row" in str( - exc_info.value + with pytest.raises(ValueError, match="Unsupported simulation spec schema version"): + spec_service.set_simulation_spec( + orm_session, simulation.id, spec, schema_version=2 ) - def test_rejects_inconsistent_stored_simulation_spec_on_read(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - test_db.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - '{"country_id":"us","population_id":"ny","population_type":"geography","policy_id":3}', - 1, - simulation["id"], - ), - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.get_simulation_spec(simulation["id"]) +def test_rejects_unsupported_schema_version_on_read(orm_session): + simulation = create_simulation(orm_session) + simulation.simulation_spec_schema_version = 2 - assert "Simulation spec must match the linked simulation row" in str( - exc_info.value - ) + with pytest.raises(ValueError, match="Unsupported simulation spec schema version"): + spec_service.get_simulation_spec(orm_session, simulation.id) + + +def test_rejects_spec_that_does_not_match_simulation(orm_session): + simulation = create_simulation(orm_session) + mismatched = SimulationSpec( + country_id="us", + population_id="ny", + population_type="geography", + policy_id=3, + ) + + with pytest.raises(ValueError, match="must match the linked simulation"): + spec_service.set_simulation_spec(orm_session, simulation.id, mismatched) + + +def test_rejects_inconsistent_stored_spec(orm_session): + simulation = create_simulation(orm_session) + simulation.simulation_spec_json = { + "country_id": "us", + "population_id": "ny", + "population_type": "geography", + "policy_id": 3, + } + + with pytest.raises(ValueError, match="must match the linked simulation"): + spec_service.get_simulation_spec(orm_session, simulation.id) + + +def test_missing_simulation_has_no_spec(orm_session): + assert spec_service.get_simulation_spec(orm_session, 999) is None + + +def test_setting_spec_for_missing_simulation_raises(orm_session): + spec = SimulationSpec( + country_id="us", + population_id="ca", + population_type="geography", + policy_id=3, + ) + + with pytest.raises(ValueError, match="Simulation #999 not found"): + spec_service.set_simulation_spec(orm_session, 999, spec) From 4b8ffeac334d2fc631c30464f5848d232ef1283b Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 02:32:47 +0300 Subject: [PATCH 53/89] refactor: use ORM sessions for report metadata --- .../services/report_output_alias_service.py | 104 ++- .../services/report_output_service.py | 5 +- .../services/report_run_service.py | 183 +++--- .../services/report_spec_service.py | 455 ++++--------- .../test_report_output_alias_service.py | 310 ++------- .../unit/services/test_report_run_service.py | 470 +++----------- .../unit/services/test_report_spec_service.py | 608 +++++------------- 7 files changed, 579 insertions(+), 1556 deletions(-) diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 34dfc88cf..90d23c6bc 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,103 +1,75 @@ -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO, V1UnitOfWork +from sqlalchemy.orm import Session +from policyengine_api.data.v1_models import LegacyReportOutputAlias, ReportOutput -class ReportOutputAliasService: - def __init__( - self, - reports: ReportDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._reports = reports - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - def _get_report_output_row(self, report_output_id: int) -> dict | None: - if self._reports is not None: - return self._reports.get(report_output_id) - with self.unit_of_work.read() as daos: - return daos.reports.get(report_output_id) +class ReportOutputAliasService: + """Legacy report-ID aliases persisted through mapped ORM models.""" - def get_alias(self, legacy_report_output_id: int) -> dict | None: - if self._reports is not None: - return self._reports.get_alias(legacy_report_output_id) - with self.unit_of_work.read() as daos: - return daos.reports.get_alias(legacy_report_output_id) + @staticmethod + def get_alias( + session: Session, legacy_report_output_id: int + ) -> LegacyReportOutputAlias | None: + return session.get(LegacyReportOutputAlias, legacy_report_output_id) def resolve_canonical_report_output_id( - self, requested_report_output_id: int - ) -> int | None: - if self._reports is not None: - return self._resolve(self._reports, requested_report_output_id) - with self.unit_of_work.read() as daos: - return self._resolve(daos.reports, requested_report_output_id) - - def _resolve( - self, reports: ReportDAO, requested_report_output_id: int + self, session: Session, requested_report_output_id: int ) -> int | None: - alias = reports.get_alias(requested_report_output_id) + alias = self.get_alias(session, requested_report_output_id) if alias is not None: - canonical_id = alias["canonical_report_output_id"] - if reports.get(canonical_id) is None: + canonical_id = alias.canonical_report_output_id + if session.get(ReportOutput, canonical_id) is None: raise ValueError( f"Alias points to missing canonical report output #{canonical_id}" ) return canonical_id - row = reports.get(requested_report_output_id) - return row["id"] if row is not None else None + report = session.get(ReportOutput, requested_report_output_id) + return report.id if report is not None else None def set_alias( - self, legacy_report_output_id: int, canonical_report_output_id: int - ) -> bool: - if self._reports is not None: - return self._set_alias( - self._reports, - legacy_report_output_id, - canonical_report_output_id, - ) - with self.unit_of_work.transaction() as daos: - return self._set_alias( - daos.reports, - legacy_report_output_id, - canonical_report_output_id, - ) - - def _set_alias( self, - reports: ReportDAO, + session: Session, legacy_report_output_id: int, canonical_report_output_id: int, ) -> bool: - legacy = reports.get(legacy_report_output_id) + legacy = session.get(ReportOutput, legacy_report_output_id) if legacy is None: raise ValueError( f"Legacy report output #{legacy_report_output_id} not found" ) - canonical = reports.get(canonical_report_output_id) + canonical = session.get(ReportOutput, canonical_report_output_id) if canonical is None: raise ValueError( f"Canonical report output #{canonical_report_output_id} not found" ) if legacy_report_output_id == canonical_report_output_id: raise ValueError("Legacy and canonical report outputs must be different") - existing = reports.get_alias(legacy_report_output_id) + existing = self.get_alias(session, legacy_report_output_id) if existing is not None: - if existing["canonical_report_output_id"] == canonical_report_output_id: + if existing.canonical_report_output_id == canonical_report_output_id: return True raise ValueError( "Legacy report output alias already points to canonical report output " - f"#{existing['canonical_report_output_id']}" + f"#{existing.canonical_report_output_id}" ) - logical_key = ("country_id", "simulation_1_id", "simulation_2_id", "year") - if any(legacy[field] != canonical[field] for field in logical_key): + logical_fields = ( + "country_id", + "simulation_1_id", + "simulation_2_id", + "year", + ) + if any( + getattr(legacy, field) != getattr(canonical, field) + for field in logical_fields + ): raise ValueError( "Legacy and canonical report outputs must describe the same report" ) - reports.set_alias(legacy_report_output_id, canonical_report_output_id) + session.add( + LegacyReportOutputAlias( + legacy_report_output_id=legacy_report_output_id, + canonical_report_output_id=canonical_report_output_id, + ) + ) + session.flush() return True diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 539d5d2a8..17d78cc4d 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -21,15 +21,12 @@ class ReportOutputService: def __init__(self, *, unit_of_work: V1UnitOfWork | None = None): self._unit_of_work = unit_of_work - self.report_spec_service = ReportSpecService(unit_of_work=unit_of_work) + self.report_spec_service = ReportSpecService() @property def unit_of_work(self) -> V1UnitOfWork: if self._unit_of_work is None: self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - self.report_spec_service = ReportSpecService( - unit_of_work=self._unit_of_work - ) return self._unit_of_work def _utc_timestamp(self) -> datetime: diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index 285b47b9d..e0c8127ff 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -1,11 +1,11 @@ -import json import uuid from datetime import datetime, timezone from typing import Any -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO, V1UnitOfWork -from policyengine_api.services.run_sync_utils import select_display_report_run +from sqlalchemy import func, select +from sqlalchemy.orm import Session + +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun REPORT_RUN_VERSION_FIELDS = ( @@ -22,57 +22,53 @@ class ReportRunService: - def __init__( - self, - reports: ReportDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._reports = reports - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work + """Report-run operations performed through a caller-owned ORM Session.""" - def _parse_run_row(self, row: dict | None) -> dict | None: - if row is None: - return None - run = dict(row) - if isinstance(run.get("report_spec_snapshot_json"), str): - run["report_spec_snapshot_json"] = json.loads( - run["report_spec_snapshot_json"] - ) - for field in ("requested_at", "started_at", "finished_at"): - if isinstance(run.get(field), datetime): - run[field] = run[field].strftime("%Y-%m-%d %H:%M:%S") - return run + @staticmethod + def _matches_report_result(run: ReportOutputRun, report: ReportOutput) -> bool: + return ( + run.status == report.status + and run.output == report.output + and run.error_message == report.error_message + ) def create_report_output_run( self, + session: Session, report_output_id: int, status: str = "pending", trigger_type: str = "initial", - output: dict[str, Any] | list[Any] | str | None = None, + output: dict[str, Any] | list[Any] | None = None, error_message: str | None = None, source_run_id: str | None = None, - report_spec_snapshot: dict[str, Any] | str | None = None, + report_spec_snapshot: dict[str, Any] | None = None, version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, - ) -> dict: + ) -> ReportOutputRun: + parent = session.scalar( + select(ReportOutput) + .where(ReportOutput.id == report_output_id) + .with_for_update() + ) + if parent is None: + raise ValueError(f"Report output #{report_output_id} not found") + sequence = ( + session.scalar( + select(func.max(ReportOutputRun.run_sequence)).where( + ReportOutputRun.report_output_id == report_output_id + ) + ) + or 0 + ) + 1 now = datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) - terminal = status in ("complete", "error") - started = status in ("running", "complete", "error") values = { "status": status, "output": output, "error_message": error_message, "trigger_type": trigger_type, "requested_at": now, - "started_at": now if started else None, - "finished_at": now if terminal else None, + "started_at": now if status in {"running", "complete", "error"} else None, + "finished_at": now if status in {"complete", "error"} else None, "source_run_id": source_run_id, "report_spec_snapshot_json": report_spec_snapshot, } @@ -82,48 +78,81 @@ def create_report_output_run( for field in REPORT_RUN_VERSION_FIELDS } ) - try: - if self._reports is not None: - run = self._reports.create_run( - report_output_id, - run_id=run_id or str(uuid.uuid4()), - **values, - ) - else: - with self.unit_of_work.transaction() as daos: - run = daos.reports.create_run( - report_output_id, - run_id=run_id or str(uuid.uuid4()), - **values, - ) - except LookupError as error: - raise ValueError(f"Report output #{report_output_id} not found") from error - return self._parse_run_row(run) + run = ReportOutputRun( + id=run_id or str(uuid.uuid4()), + report_output_id=report_output_id, + run_sequence=sequence, + **values, + ) + session.add(run) + session.flush() + return run - def get_report_output_run(self, run_id: str) -> dict | None: - if self._reports is not None: - return self._parse_run_row(self._reports.get_run(run_id)) - with self.unit_of_work.read() as daos: - return self._parse_run_row(daos.reports.get_run(run_id)) + def get_report_output_run( + self, session: Session, run_id: str + ) -> ReportOutputRun | None: + return session.get(ReportOutputRun, run_id) - def list_report_output_runs(self, report_output_id: int) -> list[dict]: - if self._reports is not None: - rows = self._reports.list_runs(report_output_id) - else: - with self.unit_of_work.read() as daos: - rows = daos.reports.list_runs(report_output_id) - return [self._parse_run_row(row) for row in reversed(rows)] + def list_report_output_runs( + self, session: Session, report_output_id: int + ) -> list[ReportOutputRun]: + return list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.asc()) + ) + ) - def get_newest_report_output_run(self, report_output_id: int) -> dict | None: - if self._reports is not None: - rows = self._reports.list_runs(report_output_id) - else: - with self.unit_of_work.read() as daos: - rows = daos.reports.list_runs(report_output_id) - return self._parse_run_row(rows[0]) if rows else None + def get_newest_report_output_run( + self, session: Session, report_output_id: int + ) -> ReportOutputRun | None: + return session.scalar( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) - def select_display_run(self, report_output: dict) -> dict | None: - runs_descending = list( - reversed(self.list_report_output_runs(report_output["id"])) + def select_display_run( + self, session: Session, report_output: ReportOutput + ) -> ReportOutputRun | None: + runs = list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output.id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + ) + if report_output.active_run_id is not None: + active = next( + (run for run in runs if run.id == report_output.active_run_id), None + ) + if active is not None: + return active + if report_output.status == "error": + matching_error = next( + ( + run + for run in runs + if self._matches_report_result(run, report_output) + ), + None, + ) + if matching_error is not None: + return matching_error + if report_output.latest_successful_run_id is not None: + successful = next( + ( + run + for run in runs + if run.id == report_output.latest_successful_run_id + ), + None, + ) + if successful is not None: + return successful + matching = next( + (run for run in runs if self._matches_report_result(run, report_output)), + None, ) - return select_display_report_run(report_output, runs_descending) + return matching or (runs[0] if runs else None) diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index aad376824..38b014257 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -2,9 +2,10 @@ from typing import Any, Literal from pydantic import BaseModel, Field +from sqlalchemy.orm import Session + +from policyengine_api.data.v1_models import ReportOutput, Simulation -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import ReportDAO, SimulationDAO, V1UnitOfWork REPORT_SPEC_SCHEMA_VERSION = 1 REPORT_SPEC_STATUSES = {"explicit", "backfilled_assumed"} @@ -42,296 +43,99 @@ class EconomyReportSpec(BaseModel): class ReportSpecService: - def __init__( - self, - reports: ReportDAO | None = None, - simulations: SimulationDAO | None = None, - *, - unit_of_work: V1UnitOfWork | None = None, - ): - self._reports = reports - self._simulations = simulations - self._unit_of_work = unit_of_work - - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work - - def _validate_schema_version(self, schema_version: int | None) -> None: + @staticmethod + def _validate_schema_version(schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: raise ValueError( f"Unsupported report spec schema version: {schema_version}" ) - def _get_report_output_row(self, report_output_id: int) -> dict | None: - if self._reports is not None: - return self._reports.get(report_output_id) - with self.unit_of_work.read() as daos: - return daos.reports.get(report_output_id) - - def _get_simulation_row(self, simulation_id: int) -> dict | None: - if self._simulations is not None: - return self._simulations.get(simulation_id) - with self.unit_of_work.read() as daos: - return daos.simulations.get(simulation_id) - + @staticmethod def _get_linked_simulations( - self, report_output: dict, *, daos=None, simulations=None - ) -> tuple[dict, dict | None]: - simulations = simulations or self._simulations - if daos is None and simulations is None: - with self.unit_of_work.read() as read_daos: - return self._get_linked_simulations(report_output, daos=read_daos) - simulations = simulations or daos.simulations - simulation_1 = simulations.get(report_output["simulation_1_id"]) + session: Session, report_output: ReportOutput + ) -> tuple[Simulation, Simulation | None]: + simulation_1 = session.get(Simulation, report_output.simulation_1_id) if simulation_1 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_1_id']}" + f"#{report_output.simulation_1_id}" ) - simulation_2 = None - if report_output["simulation_2_id"] is not None: - simulation_2 = simulations.get(report_output["simulation_2_id"]) + if report_output.simulation_2_id is not None: + simulation_2 = session.get(Simulation, report_output.simulation_2_id) if simulation_2 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_2_id']}" + f"#{report_output.simulation_2_id}" ) - return simulation_1, simulation_2 + @staticmethod def _validate_report_simulation_linkage( - self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None, ) -> None: - if simulation_1.get("id") != report_output["simulation_1_id"]: + if simulation_1.id != report_output.simulation_1_id: raise ValueError( "Simulation 1 must match report_output.simulation_1_id to build a " "report spec" ) - - report_simulation_2_id = report_output["simulation_2_id"] - if report_simulation_2_id is None: + if report_output.simulation_2_id is None: if simulation_2 is not None: raise ValueError("Report output does not reference a second simulation") return - if simulation_2 is None: raise ValueError( "Report output requires a second simulation to build a comparison " "report spec" ) - if simulation_2.get("id") != report_simulation_2_id: + if simulation_2.id != report_output.simulation_2_id: raise ValueError( "Simulation 2 must match report_output.simulation_2_id to build a " "report spec" ) + @staticmethod def _validate_report_country( - self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None, ) -> None: - report_country_id = report_output["country_id"] - if simulation_1["country_id"] != report_country_id: + if simulation_1.country_id != report_output.country_id: raise ValueError( "Simulation 1 country must match report output country to build a " "report spec" ) - if simulation_2 is not None and simulation_2["country_id"] != report_country_id: - raise ValueError( - "Simulation 2 country must match report output country to build a " - "report spec" - ) - - def _build_household_report_spec( - self, - report_output: dict, - report_kind: str, - simulation_1: dict, - simulation_2: dict | None, - time_period: str, - ) -> HouseholdReportSpec: - if simulation_1["population_type"] != "household": - raise ValueError("Household report specs require household simulations") if ( simulation_2 is not None - and simulation_2["population_id"] != simulation_1["population_id"] + and simulation_2.country_id != report_output.country_id ): raise ValueError( - "Household comparison report specs require matching household IDs" - ) - - return HouseholdReportSpec.model_validate( - { - "country_id": report_output["country_id"], - "report_kind": report_kind, - "time_period": time_period, - "simulation_1": { - "population_type": simulation_1["population_type"], - "population_id": simulation_1["population_id"], - "policy_id": simulation_1["policy_id"], - }, - "simulation_2": ( - { - "population_type": simulation_2["population_type"], - "population_id": simulation_2["population_id"], - "policy_id": simulation_2["policy_id"], - } - if simulation_2 is not None - else None - ), - } - ) - - def _build_economy_report_spec( - self, - report_output: dict, - report_kind: str, - simulation_1: dict, - simulation_2: dict | None, - time_period: str, - dataset: str, - target: Literal["general", "cliff"], - options: dict[str, Any] | None, - ) -> EconomyReportSpec: - if simulation_1["population_type"] != "geography": - raise ValueError("Economy report specs require geography simulations") - if ( - simulation_2 is not None - and simulation_2["population_id"] != simulation_1["population_id"] - ): - raise ValueError( - "Economy comparison report specs require matching geography IDs" - ) - - return EconomyReportSpec.model_validate( - { - "country_id": report_output["country_id"], - "report_kind": report_kind, - "time_period": time_period, - "region": simulation_1["population_id"], - "baseline_policy_id": simulation_1["policy_id"], - "reform_policy_id": ( - simulation_2["policy_id"] - if simulation_2 is not None - else simulation_1["policy_id"] - ), - "dataset": dataset, - "target": target, - "options": options or {}, - } - ) - - def _validate_report_spec_matches_row( - self, - report_output: dict, - report_spec: ReportSpec, - *, - daos=None, - simulations=None, - ) -> None: - simulation_1, simulation_2 = self._get_linked_simulations( - report_output, - daos=daos, - simulations=simulations, - ) - inferred_report_kind = self.infer_report_kind(simulation_1, simulation_2) - if report_spec.country_id != report_output["country_id"]: - raise ValueError("Report spec country must match report output country") - if report_spec.time_period != report_output["year"]: - raise ValueError("Report spec time_period must match report output year") - if report_spec.report_kind != inferred_report_kind: - raise ValueError("Report spec kind must match linked simulations") - - if isinstance(report_spec, HouseholdReportSpec): - if report_spec.simulation_1.model_dump() != { - "population_type": simulation_1["population_type"], - "population_id": simulation_1["population_id"], - "policy_id": simulation_1["policy_id"], - }: - raise ValueError( - "Report spec simulation_1 must match linked simulation 1" - ) - - expected_simulation_2 = ( - { - "population_type": simulation_2["population_type"], - "population_id": simulation_2["population_id"], - "policy_id": simulation_2["policy_id"], - } - if simulation_2 is not None - else None - ) - actual_simulation_2 = ( - report_spec.simulation_2.model_dump() - if report_spec.simulation_2 is not None - else None - ) - if actual_simulation_2 != expected_simulation_2: - raise ValueError( - "Report spec simulation_2 must match linked simulation 2" - ) - return - - expected_region = simulation_1["population_id"] - expected_baseline_policy_id = simulation_1["policy_id"] - expected_reform_policy_id = ( - simulation_2["policy_id"] - if simulation_2 is not None - else simulation_1["policy_id"] - ) - - if report_spec.region != expected_region: - raise ValueError("Report spec region must match linked simulations") - if report_spec.baseline_policy_id != expected_baseline_policy_id: - raise ValueError( - "Report spec baseline_policy_id must match linked simulations" - ) - if report_spec.reform_policy_id != expected_reform_policy_id: - raise ValueError( - "Report spec reform_policy_id must match linked simulations" + "Simulation 2 country must match report output country to build a " + "report spec" ) + @staticmethod def infer_report_kind( - self, - simulation_1: dict, - simulation_2: dict | None = None, + simulation_1: Simulation, simulation_2: Simulation | None = None ) -> str: - population_type = simulation_1["population_type"] - if ( - simulation_2 is not None - and simulation_2["population_type"] != population_type - ): + population_type = simulation_1.population_type + if simulation_2 is not None and simulation_2.population_type != population_type: raise ValueError( "Simulation population types must match to build a report spec" ) - if population_type == "household": - return ( - "household_comparison" - if simulation_2 is not None - else "household_single" - ) - + return "household_comparison" if simulation_2 else "household_single" if population_type == "geography": - return ( - "economy_comparison" if simulation_2 is not None else "economy_single" - ) - + return "economy_comparison" if simulation_2 else "economy_single" raise ValueError(f"Unsupported simulation population type: {population_type}") def build_report_spec( self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None = None, dataset: str = "default", target: Literal["general", "cliff"] = "general", options: dict[str, Any] | None = None, @@ -339,71 +143,115 @@ def build_report_spec( self._validate_report_simulation_linkage( report_output, simulation_1, simulation_2 ) - report_kind = self.infer_report_kind(simulation_1, simulation_2) - time_period = report_output["year"] self._validate_report_country(report_output, simulation_1, simulation_2) - + report_kind = self.infer_report_kind(simulation_1, simulation_2) if report_kind in HOUSEHOLD_REPORT_KINDS: - return self._build_household_report_spec( - report_output=report_output, + if ( + simulation_2 is not None + and simulation_2.population_id != simulation_1.population_id + ): + raise ValueError( + "Household comparison report specs require matching household IDs" + ) + return HouseholdReportSpec( + country_id=report_output.country_id, report_kind=report_kind, - simulation_1=simulation_1, - simulation_2=simulation_2, - time_period=time_period, + time_period=report_output.year, + simulation_1=ReportSimulationInput( + population_type=simulation_1.population_type, + population_id=simulation_1.population_id, + policy_id=simulation_1.policy_id, + ), + simulation_2=( + ReportSimulationInput( + population_type=simulation_2.population_type, + population_id=simulation_2.population_id, + policy_id=simulation_2.policy_id, + ) + if simulation_2 + else None + ), ) - - return self._build_economy_report_spec( - report_output=report_output, + if ( + simulation_2 is not None + and simulation_2.population_id != simulation_1.population_id + ): + raise ValueError( + "Economy comparison report specs require matching geography IDs" + ) + return EconomyReportSpec( + country_id=report_output.country_id, report_kind=report_kind, - simulation_1=simulation_1, - simulation_2=simulation_2, - time_period=time_period, + time_period=report_output.year, + region=simulation_1.population_id, + baseline_policy_id=simulation_1.policy_id, + reform_policy_id=( + simulation_2.policy_id if simulation_2 else simulation_1.policy_id + ), dataset=dataset, target=target, - options=options, + options=options or {}, ) - def _parse_json_field(self, value: str | dict | None) -> dict | None: - if value is None: - return None - if isinstance(value, str): - return json.loads(value) - return value + def _validate_report_spec_matches_model( + self, + session: Session, + report_output: ReportOutput, + report_spec: ReportSpec, + ) -> None: + simulation_1, simulation_2 = self._get_linked_simulations( + session, report_output + ) + expected = self.build_report_spec( + report_output, + simulation_1, + simulation_2, + dataset=( + report_spec.dataset + if isinstance(report_spec, EconomyReportSpec) + else "default" + ), + target=( + report_spec.target + if isinstance(report_spec, EconomyReportSpec) + else "general" + ), + options=( + report_spec.options + if isinstance(report_spec, EconomyReportSpec) + else None + ), + ) + if report_spec != expected: + raise ValueError("Report spec must match the linked report and simulations") - def _parse_report_spec(self, report_kind: str, raw_spec: dict) -> ReportSpec: + @staticmethod + def _parse_report_spec(report_kind: str, raw_spec: dict) -> ReportSpec: if report_kind in HOUSEHOLD_REPORT_KINDS: return HouseholdReportSpec.model_validate(raw_spec) if report_kind in ECONOMY_REPORT_KINDS: return EconomyReportSpec.model_validate(raw_spec) raise ValueError(f"Unsupported report kind: {report_kind}") - def get_report_spec(self, report_output_id: int) -> ReportSpec | None: - if self._reports is not None and self._simulations is not None: - report_output = self._reports.get(report_output_id) - daos = None - else: - with self.unit_of_work.read() as daos: - return self._get_report_spec(report_output_id, daos) - return self._parse_stored_report_spec(report_output, daos=daos) - - def _get_report_spec(self, report_output_id: int, daos) -> ReportSpec | None: - return self._parse_stored_report_spec( - daos.reports.get(report_output_id), daos=daos - ) - - def _parse_stored_report_spec( - self, report_output: dict | None, *, daos=None + def get_report_spec( + self, session: Session, report_output_id: int ) -> ReportSpec | None: - if report_output is None or report_output["report_spec_json"] is None: + report_output = session.get(ReportOutput, report_output_id) + if report_output is None or report_output.report_spec_json is None: return None - self._validate_schema_version(report_output["report_spec_schema_version"]) - raw_spec = self._parse_json_field(report_output["report_spec_json"]) - report_spec = self._parse_report_spec(report_output["report_kind"], raw_spec) - self._validate_report_spec_matches_row(report_output, report_spec, daos=daos) + self._validate_schema_version(report_output.report_spec_schema_version) + raw_spec = report_output.report_spec_json + if isinstance(raw_spec, str): + # Existing databases may contain pre-ORM JSON text. New writes below + # always assign Python objects and leave conversion to SQLAlchemy. + raw_spec = json.loads(raw_spec) + report_spec = self._parse_report_spec(report_output.report_kind, raw_spec) + self._validate_report_spec_matches_model(session, report_output, report_spec) return report_spec def set_report_spec( self, + session: Session, report_output_id: int, report_spec: ReportSpec, report_spec_status: Literal["explicit", "backfilled_assumed"], @@ -412,53 +260,12 @@ def set_report_spec( if report_spec_status not in REPORT_SPEC_STATUSES: raise ValueError(f"Unsupported report spec status: {report_spec_status}") self._validate_schema_version(schema_version) - - if self._reports is not None and self._simulations is not None: - self._set_report_spec( - self._reports, - self._simulations, - report_output_id, - report_spec, - report_spec_status, - schema_version, - ) - return True - - with self.unit_of_work.transaction() as daos: - report_output = daos.reports.get(report_output_id) - if report_output is None: - raise ValueError(f"Report output #{report_output_id} not found") - self._validate_report_spec_matches_row( - report_output, report_spec, daos=daos - ) - daos.reports.update( - report_output_id, - report_kind=report_spec.report_kind, - report_spec_json=report_spec.model_dump(), - report_spec_schema_version=schema_version, - report_spec_status=report_spec_status, - ) - return True - - def _set_report_spec( - self, - reports: ReportDAO, - simulations: SimulationDAO, - report_output_id: int, - report_spec: ReportSpec, - report_spec_status: str, - schema_version: int, - ) -> None: - report_output = reports.get(report_output_id) + report_output = session.get(ReportOutput, report_output_id) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") - self._validate_report_spec_matches_row( - report_output, report_spec, simulations=simulations - ) - reports.update( - report_output_id, - report_kind=report_spec.report_kind, - report_spec_json=report_spec.model_dump(), - report_spec_schema_version=schema_version, - report_spec_status=report_spec_status, - ) + self._validate_report_spec_matches_model(session, report_output, report_spec) + report_output.report_kind = report_spec.report_kind + report_output.report_spec_json = report_spec.model_dump() + report_output.report_spec_schema_version = schema_version + report_output.report_spec_status = report_spec_status + return True diff --git a/tests/unit/services/test_report_output_alias_service.py b/tests/unit/services/test_report_output_alias_service.py index e4e28c916..9377fe06b 100644 --- a/tests/unit/services/test_report_output_alias_service.py +++ b/tests/unit/services/test_report_output_alias_service.py @@ -1,279 +1,91 @@ import pytest +from policyengine_api.data.v1_models import ( + LegacyReportOutputAlias, + ReportOutput, +) from policyengine_api.services.report_output_alias_service import ( ReportOutputAliasService, ) -from policyengine_api.services.report_output_service import ReportOutputService -from policyengine_api.services.simulation_service import SimulationService -alias_service = ReportOutputAliasService() -report_output_service = ReportOutputService() -simulation_service = SimulationService() +service = ReportOutputAliasService() -class TestReportOutputAliasService: - def _insert_legacy_report_output( - self, - test_db, - legacy_report_output_id: int, - canonical_report: dict, - api_version: str = "legacy-version", - ) -> None: - test_db.query( - """ - INSERT INTO report_outputs ( - id, country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - legacy_report_output_id, - canonical_report["country_id"], - canonical_report["simulation_1_id"], - canonical_report["simulation_2_id"], - api_version, - canonical_report["status"], - canonical_report["year"], - ), - ) - def test_resolves_to_canonical_report_output_id_when_alias_exists(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 999, canonical_report) +def add_report(orm_session, report_id): + report = ReportOutput( + id=report_id, + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + status="pending", + year="2025", + ) + orm_session.add(report) + orm_session.flush() + return report - alias_service.set_alias( - legacy_report_output_id=999, - canonical_report_output_id=canonical_report["id"], - ) - resolved_id = alias_service.resolve_canonical_report_output_id(999) +def test_sets_and_resolves_mapped_alias(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) - assert resolved_id == canonical_report["id"] + assert service.set_alias(orm_session, 100, 200) is True - def test_returns_requested_id_when_alias_is_not_needed(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) + alias = service.get_alias(orm_session, 100) + assert isinstance(alias, LegacyReportOutputAlias) + assert alias.canonical_report_output_id == 200 + assert service.resolve_canonical_report_output_id(orm_session, 100) == 200 + assert service.resolve_canonical_report_output_id(orm_session, 200) == 200 - resolved_id = alias_service.resolve_canonical_report_output_id( - report_output["id"] - ) - assert resolved_id == report_output["id"] +def test_setting_same_alias_is_idempotent(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) - def test_returns_none_for_unknown_report_output(self, test_db): - assert alias_service.resolve_canonical_report_output_id(123456) is None + service.set_alias(orm_session, 100, 200) - def test_set_alias_is_idempotent_for_same_canonical_report_output(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1001, canonical_report) + assert service.set_alias(orm_session, 100, 200) is True - assert ( - alias_service.set_alias( - legacy_report_output_id=1001, - canonical_report_output_id=canonical_report["id"], - ) - is True - ) - assert ( - alias_service.set_alias( - legacy_report_output_id=1001, - canonical_report_output_id=canonical_report["id"], - ) - is True - ) - def test_rejects_alias_to_missing_canonical_report_output(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3a", - population_type="household", - policy_id=3, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1002, canonical_report) +def test_rejects_conflicting_alias(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) + add_report(orm_session, 300) + service.set_alias(orm_session, 100, 200) - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=1002, - canonical_report_output_id=999999, - ) + with pytest.raises(ValueError, match="already points"): + service.set_alias(orm_session, 100, 300) - assert "Canonical report output #999999 not found" in str(exc_info.value) - def test_rejects_conflicting_alias_remap(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - other_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2026", - ) - self._insert_legacy_report_output(test_db, 1003, canonical_report) - alias_service.set_alias( - legacy_report_output_id=1003, - canonical_report_output_id=canonical_report["id"], - ) +def test_rejects_missing_and_self_aliases(orm_session): + add_report(orm_session, 100) - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=1003, - canonical_report_output_id=other_report["id"], - ) + with pytest.raises(ValueError, match="Canonical report output #999 not found"): + service.set_alias(orm_session, 100, 999) + with pytest.raises(ValueError, match="must be different"): + service.set_alias(orm_session, 100, 100) - assert ( - "Legacy report output alias already points to canonical report output " - f"#{canonical_report['id']}" - ) in str(exc_info.value) - def test_rejects_alias_when_legacy_report_output_is_missing(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4a", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) +def test_rejects_reports_with_different_logical_keys(orm_session): + add_report(orm_session, 100) + different = add_report(orm_session, 200) + different.year = "2026" - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=10030, - canonical_report_output_id=canonical_report["id"], - ) + with pytest.raises(ValueError, match="must describe the same report"): + service.set_alias(orm_session, 100, 200) - assert "Legacy report output #10030 not found" in str(exc_info.value) - def test_rejects_alias_when_legacy_and_canonical_reports_do_not_match( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4b", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - mismatched_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2026", - ) - - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=mismatched_report["id"], - canonical_report_output_id=canonical_report["id"], - ) - - assert "must describe the same report" in str(exc_info.value) - - def test_rejects_alias_when_legacy_and_canonical_ids_match(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4c", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", +def test_rejects_alias_pointing_to_missing_canonical_report(orm_session): + add_report(orm_session, 100) + orm_session.add( + LegacyReportOutputAlias( + legacy_report_output_id=100, + canonical_report_output_id=999, ) + ) + orm_session.flush() - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=canonical_report["id"], - canonical_report_output_id=canonical_report["id"], - ) - - assert "must be different" in str(exc_info.value) - - def test_rejects_alias_resolution_when_canonical_report_output_is_missing( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5", - population_type="household", - policy_id=5, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1004, canonical_report) - alias_service.set_alias( - legacy_report_output_id=1004, - canonical_report_output_id=canonical_report["id"], - ) - test_db.query( - "DELETE FROM report_outputs WHERE id = ?", - (canonical_report["id"],), - ) - - with pytest.raises(ValueError) as exc_info: - alias_service.resolve_canonical_report_output_id(1004) - - assert ( - f"Alias points to missing canonical report output #{canonical_report['id']}" - ) in str(exc_info.value) + with pytest.raises(ValueError, match="missing canonical report output #999"): + service.resolve_canonical_report_output_id(orm_session, 100) diff --git a/tests/unit/services/test_report_run_service.py b/tests/unit/services/test_report_run_service.py index 68aedb48f..890ef5e5d 100644 --- a/tests/unit/services/test_report_run_service.py +++ b/tests/unit/services/test_report_run_service.py @@ -1,382 +1,98 @@ import pytest -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun from policyengine_api.services.report_run_service import ReportRunService -from policyengine_api.services.simulation_service import SimulationService -report_output_service = ReportOutputService() -report_run_service = ReportRunService() -simulation_service = SimulationService() - -class TestCreateReportOutputRun: - def test_creates_report_runs_with_incrementing_sequence(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = report_run_service.create_report_output_run( - report_output["id"], - trigger_type="initial", - report_spec_snapshot={"country_id": "us"}, - version_manifest={ - "country_package_version": "us-1.0.0", - "report_cache_version": "r123", - }, - ) - second_run = report_run_service.create_report_output_run( - report_output["id"], - trigger_type="rerun", - ) - - assert first_run["run_sequence"] == 2 - assert first_run["trigger_type"] == "initial" - assert first_run["requested_at"] is not None - assert first_run["started_at"] is None - assert first_run["finished_at"] is None - assert first_run["report_spec_snapshot_json"] == {"country_id": "us"} - assert first_run["country_package_version"] == "us-1.0.0" - assert first_run["report_cache_version"] == "r123" - assert second_run["run_sequence"] == 3 - assert second_run["trigger_type"] == "rerun" - - def test_lists_report_runs_in_sequence_order(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - - runs = report_run_service.list_report_output_runs(report_output["id"]) - - assert [run["run_sequence"] for run in runs] == [1, 2, 3] - - def test_allocates_run_sequence_transactionally(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_7", - population_type="household", - policy_id=7, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - second_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - - assert first_run["run_sequence"] == 2 - assert second_run["run_sequence"] == 3 - - def test_raises_when_parent_report_output_is_missing(self, test_db): - with pytest.raises(ValueError) as exc_info: - report_run_service.create_report_output_run(999999, trigger_type="initial") - - assert "Report output #999999 not found" in str(exc_info.value) - - def test_running_report_run_sets_started_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_running_run_timestamp", - population_type="household", - policy_id=8, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - run = report_run_service.create_report_output_run( - report_output["id"], - status="running", - trigger_type="rerun", - ) - - assert run["requested_at"] is not None - assert run["started_at"] is not None - assert run["finished_at"] is None - - -class TestSelectDisplayReportRun: - def test_prefers_active_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - latest_successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - active_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (active_run["id"], latest_successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == active_run["id"] - - def test_falls_back_to_latest_successful_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - (successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == successful_run["id"] - - def test_prefers_matching_error_run_over_previous_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5b", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], - status="complete", - trigger_type="initial", - output={"ok": True}, - ) - error_run = report_run_service.create_report_output_run( - report_output["id"], - status="error", - trigger_type="rerun", - error_message="rerun failed", - ) - test_db.query( - """ - UPDATE report_outputs - SET status = ?, error_message = ?, active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("error", "rerun failed", successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == error_run["id"] - - def test_falls_back_when_active_run_pointer_is_stale(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5a", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_to_newest_run_when_no_pointers_exist(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - first_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - (report_output["id"],), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert first_run["run_sequence"] == 2 - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_latest_successful_pointer_is_stale( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6a", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_no_pointer_or_result_match(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6b", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], status="pending", trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET status = ?, active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - ("complete", report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == newest_run["id"] +service = ReportRunService() + + +def create_report(orm_session, *, status="pending"): + report = ReportOutput( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + status=status, + year="2025", + ) + orm_session.add(report) + orm_session.flush() + return report + + +def test_creates_mapped_runs_with_incrementing_sequence_and_python_json(orm_session): + report = create_report(orm_session) + + first = service.create_report_output_run( + orm_session, + report.id, + trigger_type="initial", + report_spec_snapshot={"country_id": "us"}, + version_manifest={"report_cache_version": "r123"}, + ) + second = service.create_report_output_run( + orm_session, report.id, trigger_type="rerun" + ) + + assert isinstance(first, ReportOutputRun) + assert first.run_sequence == 1 + assert first.requested_at is not None + assert first.started_at is None + assert first.finished_at is None + assert first.report_spec_snapshot_json == {"country_id": "us"} + assert first.report_cache_version == "r123" + assert second.run_sequence == 2 + + +@pytest.mark.parametrize( + ("status", "has_started", "has_finished"), + [ + ("pending", False, False), + ("running", True, False), + ("complete", True, True), + ("error", True, True), + ], +) +def test_sets_run_timestamps_from_status( + orm_session, status, has_started, has_finished +): + report = create_report(orm_session) + + run = service.create_report_output_run(orm_session, report.id, status=status) + + assert (run.started_at is not None) is has_started + assert (run.finished_at is not None) is has_finished + + +def test_raises_when_parent_report_is_missing(orm_session): + with pytest.raises(ValueError, match="Report output #999 not found"): + service.create_report_output_run(orm_session, 999) + + +def test_gets_lists_and_selects_mapped_runs(orm_session): + report = create_report(orm_session) + first = service.create_report_output_run(orm_session, report.id, status="complete") + second = service.create_report_output_run(orm_session, report.id, status="running") + report.latest_successful_run_id = first.id + report.active_run_id = second.id + + assert service.get_report_output_run(orm_session, first.id) is first + assert service.list_report_output_runs(orm_session, report.id) == [first, second] + assert service.get_newest_report_output_run(orm_session, report.id) is second + assert service.select_display_run(orm_session, report) is second + + +def test_select_display_run_falls_back_to_matching_error(orm_session): + report = create_report(orm_session, status="error") + matching = service.create_report_output_run( + orm_session, + report.id, + status="error", + error_message="failed", + ) + service.create_report_output_run(orm_session, report.id) + report.error_message = "failed" + report.active_run_id = None + + assert service.select_display_run(orm_session, report) is matching diff --git a/tests/unit/services/test_report_spec_service.py b/tests/unit/services/test_report_spec_service.py index f924df8db..3db91765e 100644 --- a/tests/unit/services/test_report_spec_service.py +++ b/tests/unit/services/test_report_spec_service.py @@ -1,469 +1,159 @@ import pytest -from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.data.v1_models import ReportOutput, Simulation from policyengine_api.services.report_spec_service import ( EconomyReportSpec, HouseholdReportSpec, ReportSpecService, ) -from policyengine_api.services.simulation_service import SimulationService -report_output_service = ReportOutputService() -report_spec_service = ReportSpecService() -simulation_service = SimulationService() - -class TestBuildReportSpec: - def test_builds_household_comparison_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2026", - ) - - report_spec = report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert isinstance(report_spec, HouseholdReportSpec) - assert report_spec.report_kind == "household_comparison" - assert report_spec.time_period == "2026" - assert report_spec.simulation_1.policy_id == 1 - assert report_spec.simulation_2.policy_id == 2 - - def test_builds_default_economy_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=11, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2027", - ) - - report_spec = report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert isinstance(report_spec, EconomyReportSpec) - assert report_spec.report_kind == "economy_comparison" - assert report_spec.region == "ca" - assert report_spec.dataset == "default" - assert report_spec.target == "general" - assert report_spec.options == {} - - def test_raises_for_mixed_population_types(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "population types must match" in str(exc_info.value) - - def test_raises_for_mismatched_household_ids(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", +service = ReportSpecService() + + +def add_simulation( + orm_session, + *, + population_type="household", + population_id="household-1", + policy_id=1, + country_id="us", +): + simulation = Simulation( + country_id=country_id, + api_version="1", + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", + ) + orm_session.add(simulation) + orm_session.flush() + return simulation + + +def add_report(orm_session, simulation_1, simulation_2=None, *, country_id="us"): + report = ReportOutput( + country_id=country_id, + simulation_1_id=simulation_1.id, + simulation_2_id=simulation_2.id if simulation_2 else None, + api_version="1", + status="pending", + year="2026", + ) + orm_session.add(report) + orm_session.flush() + return report + + +def test_builds_household_comparison_spec_from_models(orm_session): + first = add_simulation(orm_session, policy_id=1) + second = add_simulation(orm_session, policy_id=2) + report = add_report(orm_session, first, second) + + spec = service.build_report_spec(report, first, second) + + assert isinstance(spec, HouseholdReportSpec) + assert spec.report_kind == "household_comparison" + assert spec.simulation_1.policy_id == 1 + assert spec.simulation_2.policy_id == 2 + + +def test_builds_economy_spec_from_models(orm_session): + first = add_simulation( + orm_session, + population_type="geography", + population_id="ca", + policy_id=10, + ) + second = add_simulation( + orm_session, + population_type="geography", + population_id="ca", + policy_id=11, + ) + report = add_report(orm_session, first, second) + + spec = service.build_report_spec( + report, first, second, dataset="cps", options={"foo": "bar"} + ) + + assert isinstance(spec, EconomyReportSpec) + assert spec.report_kind == "economy_comparison" + assert spec.region == "ca" + assert spec.dataset == "cps" + assert spec.options == {"foo": "bar"} + + +def test_sets_and_gets_report_spec_as_python_json(orm_session): + simulation = add_simulation(orm_session) + report = add_report(orm_session, simulation) + spec = service.build_report_spec(report, simulation) + + assert service.set_report_spec( + orm_session, report.id, spec, report_spec_status="explicit" + ) + loaded = service.get_report_spec(orm_session, report.id) + + assert report.report_spec_json == spec.model_dump() + assert report.report_spec_schema_version == 1 + assert report.report_spec_status == "explicit" + assert loaded == spec + + +def test_rejects_missing_linked_simulation(orm_session): + report = ReportOutput( + country_id="us", + simulation_1_id=999, + simulation_2_id=None, + api_version="1", + status="pending", + year="2026", + ) + orm_session.add(report) + orm_session.flush() + spec = HouseholdReportSpec.model_validate( + { + "country_id": "us", + "report_kind": "household_single", + "time_period": "2026", + "simulation_1": { + "population_type": "household", + "population_id": "household-1", + "policy_id": 1, + }, + } + ) + + with pytest.raises(ValueError, match="references missing simulation #999"): + service.set_report_spec(orm_session, report.id, spec, "explicit") + + +def test_rejects_mismatched_country(orm_session): + simulation = add_simulation(orm_session, country_id="uk") + report = add_report(orm_session, simulation, country_id="us") + + with pytest.raises(ValueError, match="country must match"): + service.build_report_spec(report, simulation) + + +def test_rejects_mismatched_comparison_population(orm_session): + first = add_simulation(orm_session, population_id="household-1") + second = add_simulation(orm_session, population_id="household-2", policy_id=2) + report = add_report(orm_session, first, second) + + with pytest.raises(ValueError, match="matching household IDs"): + service.build_report_spec(report, first, second) + + +def test_rejects_unsupported_schema_version_and_status(orm_session): + simulation = add_simulation(orm_session) + report = add_report(orm_session, simulation) + spec = service.build_report_spec(report, simulation) + + with pytest.raises(ValueError, match="Unsupported report spec status"): + service.set_report_spec(orm_session, report.id, spec, "unknown") + with pytest.raises(ValueError, match="Unsupported report spec schema version"): + service.set_report_spec( + orm_session, report.id, spec, "explicit", schema_version=2 ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "matching household IDs" in str(exc_info.value) - - def test_raises_for_mismatched_geography_ids(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ny", - population_type="geography", - policy_id=11, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2027", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "matching geography IDs" in str(exc_info.value) - - def test_raises_for_country_mismatch_between_report_and_simulation(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="uk", - population_id="household_1", - population_type="household", - policy_id=1, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, simulation_1) - - assert "Simulation 1 country must match report output country" in str( - exc_info.value - ) - - def test_raises_when_simulation_1_does_not_match_report_output_linkage( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - other_simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, other_simulation) - - assert "Simulation 1 must match report_output.simulation_1_id" in str( - exc_info.value - ) - - def test_raises_when_report_requires_second_simulation_but_none_is_supplied( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, simulation_1) - - assert "requires a second simulation" in str(exc_info.value) - - -class TestPersistReportSpec: - def test_sets_and_gets_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - result = report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert result is True - stored_report = test_db.query( - """ - SELECT report_kind, report_spec_json, report_spec_schema_version, report_spec_status - FROM report_outputs WHERE id = ? - """, - (report_output["id"],), - ).fetchone() - assert stored_report["report_kind"] == "economy_single" - assert stored_report["report_spec_schema_version"] == 1 - assert stored_report["report_spec_status"] == "explicit" - - loaded_report_spec = report_spec_service.get_report_spec(report_output["id"]) - assert loaded_report_spec is not None - assert loaded_report_spec.model_dump() == report_spec.model_dump() - - def test_rejects_report_spec_write_when_region_does_not_match_report(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ny", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert "Report spec region must match linked simulations" in str(exc_info.value) - - def test_rejects_report_spec_write_when_time_period_does_not_match_report( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2028", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert "time_period must match report output year" in str(exc_info.value) - - def test_rejects_inconsistent_stored_report_spec_on_read(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - test_db.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - "economy_single", - '{"country_id":"us","report_kind":"economy_single","time_period":"2027","region":"ny","baseline_policy_id":10,"reform_policy_id":10,"dataset":"default","target":"general","options":{}}', - 1, - "explicit", - report_output["id"], - ), - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.get_report_spec(report_output["id"]) - - assert "Report spec region must match linked simulations" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_write(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], - report_spec, - report_spec_status="explicit", - schema_version=2, - ) - - assert "Unsupported report spec schema version" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_read(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - test_db.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - "economy_single", - '{"country_id":"us","report_kind":"economy_single","time_period":"2027","region":"ca","baseline_policy_id":10,"reform_policy_id":10,"dataset":"default","target":"general","options":{}}', - 2, - "explicit", - report_output["id"], - ), - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.get_report_spec(report_output["id"]) - - assert "Unsupported report spec schema version" in str(exc_info.value) From 0f40cc713a32980b92b9cce87ae03db9a7284d95 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 02:37:33 +0300 Subject: [PATCH 54/89] refactor: orchestrate reports with ORM sessions --- .../routes/report_output_routes.py | 164 +- .../services/report_output_service.py | 1074 +++------ tests/contract/test_v1_route_contracts.py | 44 +- .../services/test_report_output_service.py | 2084 ++--------------- tests/unit/test_stage5_routes.py | 725 ++---- 5 files changed, 896 insertions(+), 3195 deletions(-) diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index 48a2ac43a..e5d845c32 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -7,12 +7,38 @@ from policyengine_api.services.report_output_service import ReportOutputService from policyengine_api.constants import CURRENT_YEAR +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReportOutput from policyengine_api.utils.payload_validators import validate_country report_output_bp = Blueprint("report_output", __name__) report_output_service = ReportOutputService() +def _serialize_v1_report_output( + session, report_output: ReportOutput, *, response_id: int | None = None +) -> dict: + """Project mapped report state onto the historical v1 response shape.""" + + result = { + column.name: getattr(report_output, column.name) + for column in ReportOutput.__table__.columns + } + if response_id is not None: + result["id"] = response_id + if result.get("output") is not None and not isinstance(result["output"], str): + result["output"] = json.dumps(result["output"]) + display_run = report_output_service.report_run_service.select_display_run( + session, report_output + ) + if display_run is not None: + for field in ("requested_at", "started_at", "finished_at"): + result[field] = report_output_service.format_run_timestamp( + getattr(display_run, field) + ) + return result + + @report_output_bp.route("//report", methods=["POST"]) @validate_country def create_report_output(country_id: str) -> Response: @@ -48,51 +74,44 @@ def create_report_output(country_id: str) -> Response: raise BadRequest("year must be a string") try: - # Check if report already exists with these simulation IDs and year - existing_report = report_output_service.find_existing_report_output( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - ) - - if existing_report: - existing_report = ( - report_output_service.ensure_report_output_dual_write_state( - existing_report["id"], + with get_v1_session_factory().begin() as session: + existing_report = report_output_service.find_existing_report_output( + session, + country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + year=year, + ) + if existing_report: + report_output = ( + report_output_service.ensure_report_output_dual_write_state( + session, existing_report.id, country_id + ) + ) + result = _serialize_v1_report_output(session, report_output) + message = "Report output already exists" + status_code = 200 + else: + report_output = report_output_service.create_report_output( + session, country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + year=year, ) - ) - # Report already exists, return it with 200 status - response_body = dict( - status="ok", - message="Report output already exists", - result=existing_report, - ) - - return Response( - json.dumps(response_body), - status=200, - mimetype="application/json", - ) - - # Create new report output - created_report = report_output_service.create_report_output( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - ) + result = _serialize_v1_report_output(session, report_output) + message = "Report output created successfully" + status_code = 201 response_body = dict( status="ok", - message="Report output created successfully", - result=created_report, + message=message, + result=result, ) return Response( json.dumps(response_body), - status=201, + status=status_code, mimetype="application/json", ) @@ -129,17 +148,30 @@ def get_report_output(country_id: str, report_id: int) -> Response: """ print(f"Getting report output {report_id} for country {country_id}") - report_output: dict | None = report_output_service.get_report_output( - country_id, report_id - ) - - if report_output is None: - raise NotFound(f"Report #{report_id} not found.") + with get_v1_session_factory().begin() as session: + requested_report = report_output_service.get_report_output( + session, country_id, report_id + ) + if requested_report is None: + raise NotFound(f"Report #{report_id} not found.") + if report_output_service.is_current_report_output(requested_report): + report_output = report_output_service.ensure_report_output_dual_write_state( + session, report_id, country_id + ) + response_id = None + else: + report_output = report_output_service.get_or_create_current_report_output( + session, requested_report + ) + response_id = report_id + result = _serialize_v1_report_output( + session, report_output, response_id=response_id + ) response_body = dict( status="ok", message=None, - result=report_output, + result=result, ) return Response( @@ -191,34 +223,32 @@ def update_report_output(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - # First check if the report output exists without running pointer sync: - # syncing a completed parent before this mutation can clear an active - # pending rerun that this PATCH is about to mark as running. - if not report_output_service.report_output_exists(country_id, report_id): - raise NotFound(f"Report #{report_id} not found.") - - # Update the report output - success = report_output_service.update_report_output( - country_id=country_id, - report_id=report_id, - status=status, - output=output, - error_message=error_message, - ) - - if not success: - raise BadRequest("No fields to update") - - # Get the updated stored record so stale-runtime jobs do not appear to - # complete the current runtime lineage in the PATCH response. - updated_report = report_output_service.get_stored_report_output( - country_id, report_id - ) + with get_v1_session_factory().begin() as session: + # Do not synchronize before this mutation: doing so could overwrite + # the pending rerun that this PATCH is about to mark as running. + if not report_output_service.report_output_exists( + session, country_id, report_id + ): + raise NotFound(f"Report #{report_id} not found.") + success = report_output_service.update_report_output( + session, + country_id=country_id, + report_id=report_id, + status=status, + output=output, + error_message=error_message, + ) + if not success: + raise BadRequest("No fields to update") + updated_report = report_output_service.get_report_output( + session, country_id, report_id + ) + result = _serialize_v1_report_output(session, updated_report) response_body = dict( status="ok", message="Report output updated successfully", - result=updated_report, + result=result, ) return Response( diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 17d78cc4d..5c69c79ee 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,41 +1,36 @@ -import uuid +import json from datetime import datetime, timezone +from sqlalchemy import select +from sqlalchemy.orm import Session + from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.data.orm import build_v1_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun, Simulation +from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.report_spec_service import ( ECONOMY_REPORT_KINDS, ReportSpec, ReportSpecService, ) -from policyengine_api.services.run_sync_utils import ( - determine_parent_pointers, - parse_json_field, - run_matches_report_result, - select_display_report_run, - serialize_json_field, -) +from policyengine_api.services.simulation_service import SimulationService class ReportOutputService: - def __init__(self, *, unit_of_work: V1UnitOfWork | None = None): - self._unit_of_work = unit_of_work - self.report_spec_service = ReportSpecService() + """Report-output orchestration through one caller-owned ORM Session.""" - @property - def unit_of_work(self) -> V1UnitOfWork: - if self._unit_of_work is None: - self._unit_of_work = V1UnitOfWork(build_v1_session_manager()) - return self._unit_of_work + def __init__(self): + self.report_spec_service = ReportSpecService() + self.report_run_service = ReportRunService() + self.simulation_service = SimulationService() - def _utc_timestamp(self) -> datetime: + @staticmethod + def _utc_timestamp() -> datetime: return datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) - def _format_run_timestamp(self, value) -> str | None: + @staticmethod + def format_run_timestamp(value: datetime | str | None) -> str | None: if value is None: return None - if isinstance(value, datetime): timestamp = value if timestamp.tzinfo is None: @@ -46,22 +41,17 @@ def _format_run_timestamp(self, value) -> str | None: .isoformat() .replace("+00:00", "Z") ) - timestamp = str(value).strip() if not timestamp: return None - normalized = timestamp.replace(" ", "T", 1) - parseable_timestamp = ( + parseable = ( f"{normalized[:-1]}+00:00" if normalized.endswith("Z") else normalized ) try: - parsed = datetime.fromisoformat(parseable_timestamp) + parsed = datetime.fromisoformat(parseable) except ValueError: - if "T" in normalized: - return normalized if normalized.endswith("Z") else f"{normalized}Z" - return f"{normalized}Z" - + return normalized if normalized.endswith("Z") else f"{normalized}Z" if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) return ( @@ -71,804 +61,484 @@ def _format_run_timestamp(self, value) -> str | None: .replace("+00:00", "Z") ) - def _get_report_output_row( - self, + @staticmethod + def _select_report_output( + session: Session, report_output_id: int, - *, - queryer=None, country_id: str | None = None, + *, for_update: bool = False, - ) -> dict | None: - if queryer is None: - with self.unit_of_work.read() as daos: - return self._get_report_output_row( - report_output_id, - queryer=daos, - country_id=country_id, - for_update=for_update, - ) + ) -> ReportOutput | None: + statement = select(ReportOutput).where(ReportOutput.id == report_output_id) + if country_id is not None: + statement = statement.where(ReportOutput.country_id == country_id) if for_update: - return queryer.reports.get_for_update(report_output_id, country_id) - return queryer.reports.get(report_output_id, country_id) + statement = statement.with_for_update() + return session.scalar(statement) + + @staticmethod + def _list_runs_descending( + session: Session, report_output_id: int + ) -> list[ReportOutputRun]: + return list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + ) def _get_linked_simulations( self, - report_output: dict, + session: Session, + report_output: ReportOutput, *, - queryer=None, - bootstrap_dual_write_state: bool = False, - ) -> tuple[dict, dict | None]: - if queryer is None: - with self.unit_of_work.read() as daos: - return self._get_linked_simulations( - report_output, - queryer=daos, - bootstrap_dual_write_state=bootstrap_dual_write_state, + bootstrap_dual_write_state: bool, + ) -> tuple[Simulation, Simulation | None]: + def get_simulation(simulation_id: int) -> Simulation | None: + if bootstrap_dual_write_state: + try: + return self.simulation_service.ensure_simulation_dual_write_state( + session, + simulation_id, + report_output.country_id, + ) + except ValueError: + return None + return session.scalar( + select(Simulation).where( + Simulation.id == simulation_id, + Simulation.country_id == report_output.country_id, ) - if bootstrap_dual_write_state: - simulation_1 = queryer.simulations.ensure_dual_write_state( - report_output["simulation_1_id"], - report_output["country_id"], - ) - else: - simulation_1 = queryer.simulations.get( - report_output["simulation_1_id"], - report_output["country_id"], ) + + simulation_1 = get_simulation(report_output.simulation_1_id) if simulation_1 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_1_id']}" + f"#{report_output.simulation_1_id}" ) - simulation_2 = None - if report_output["simulation_2_id"] is not None: - if bootstrap_dual_write_state: - simulation_2 = queryer.simulations.ensure_dual_write_state( - report_output["simulation_2_id"], - report_output["country_id"], - ) - else: - simulation_2 = queryer.simulations.get( - report_output["simulation_2_id"], - report_output["country_id"], - ) + if report_output.simulation_2_id is not None: + simulation_2 = get_simulation(report_output.simulation_2_id) if simulation_2 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_2_id']}" + f"#{report_output.simulation_2_id}" ) - return simulation_1, simulation_2 - def _require_simulation_exists( - self, - tx, - *, - country_id: str, - simulation_id: int, - ) -> dict: - simulation = tx.simulations.get(simulation_id, country_id) - if simulation is None: - raise ValueError( - f"Report output references missing simulation #{simulation_id}" - ) - return simulation - - def _list_report_runs_descending( - self, report_output_id: int, *, queryer=None - ) -> list[dict]: - if queryer is None: - with self.unit_of_work.read() as daos: - return self._list_report_runs_descending( - report_output_id, - queryer=daos, + @staticmethod + def _select_mutable_run( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] + ) -> ReportOutputRun | None: + if report_output.status == "running": + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + and run.status in {"pending", "running"} + ), + None, ) - rows = queryer.reports.list_runs(report_output_id) - - runs = [] - for row in rows: - run = dict(row) - run["report_spec_snapshot_json"] = parse_json_field( - run.get("report_spec_snapshot_json") + if active is not None: + return active + return next( + ( + run + for run in runs_descending + if run.status in {"pending", "running"} + ), + None, ) - runs.append(run) - return runs - - def _select_mutable_run( - self, report_output: dict, runs_descending: list[dict] - ) -> dict | None: - active_run_id = report_output.get("active_run_id") - if report_output["status"] == "running": - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id and run["status"] in ( - "pending", - "running", - ): - return run - for run in runs_descending: - if run["status"] in ("pending", "running"): - return run - return None - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + ), + None, + ) + if active is not None: + return active return runs_descending[0] if runs_descending else None - def _has_mutable_running_run(self, report_output: dict, *, queryer=None) -> bool: - runs_descending = self._list_report_runs_descending( - report_output["id"], queryer=queryer - ) - if not runs_descending: + @staticmethod + def _run_needs_timestamp_sync(run: ReportOutputRun, status: str) -> bool: + if run.requested_at is None: return True - - active_run_id = report_output.get("active_run_id") - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run["status"] in ("pending", "running") - return False - - return any(run["status"] in ("pending", "running") for run in runs_descending) - - def _run_needs_timestamp_sync(self, run: dict, status: str) -> bool: - if run.get("requested_at") is None: - return True - if status in ("complete", "error"): - return run.get("started_at") is None or run.get("finished_at") is None + if status in {"complete", "error"}: + return run.started_at is None or run.finished_at is None if status == "running": - return run.get("started_at") is None or run.get("finished_at") is not None - return run.get("started_at") is not None or run.get("finished_at") is not None - - def _with_display_run_timestamps( - self, report_output: dict, *, queryer=None - ) -> dict: - """ - Overlay selected run timestamps onto the legacy report response shape. - - This is a response-compatibility bridge for app-v2 while report output - reads still return a report_outputs row. The authoritative timestamp - values live on report_output_runs; this helper chooses the display run, - formats its requested/started/finished timestamps, and returns an - enriched copy of the report output dict. It intentionally does not - mutate repository state. + return run.started_at is None or run.finished_at is not None + return run.started_at is not None or run.finished_at is not None - These timestamps describe the selected base report execution. They are - not user-report association metadata and should not be treated as a - user-specific "last run" value. - - TODO: When report output reads are cut over to canonical run-backed - resolution, move this projection into the final response serializer - instead of keeping it as an ad hoc enrichment helper. - """ - runs_descending = self._list_report_runs_descending( - report_output["id"], queryer=queryer - ) - display_run = select_display_report_run(report_output, runs_descending) - enriched_report_output = dict(report_output) - enriched_report_output["output"] = serialize_json_field( - enriched_report_output.get("output") - ) - if display_run is None: - return enriched_report_output - - for field in ("requested_at", "started_at", "finished_at"): - enriched_report_output[field] = self._format_run_timestamp( - display_run.get(field) + @staticmethod + def _has_mutable_running_run( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] + ) -> bool: + if not runs_descending: + return True + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + ), + None, ) - return enriched_report_output + return active is not None and active.status in {"pending", "running"} + return any(run.status in {"pending", "running"} for run in runs_descending) - def _derive_report_country_package_version( - self, - simulation_1: dict | None, - simulation_2: dict | None = None, + @staticmethod + def _derive_country_package_version( + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> str | None: versions = [ - simulation["api_version"] + simulation.api_version for simulation in (simulation_1, simulation_2) - if simulation is not None and simulation.get("api_version") is not None + if simulation is not None and simulation.api_version is not None ] - if not versions: - return None - if len(set(versions)) == 1: - return versions[0] - return None + return versions[0] if versions and len(set(versions)) == 1 else None def _build_version_manifest( self, - report_output: dict, + report_output: ReportOutput, report_spec: ReportSpec | None, - simulation_1: dict | None = None, - simulation_2: dict | None = None, + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> dict[str, str | None]: - resolved_dataset = None - if report_spec is not None and report_spec.report_kind in ECONOMY_REPORT_KINDS: - resolved_dataset = report_spec.dataset - return { - "country_package_version": self._derive_report_country_package_version( + "country_package_version": self._derive_country_package_version( simulation_1, simulation_2 ), "policyengine_version": None, "data_version": None, "runtime_app_name": None, - "report_cache_version": report_output.get("api_version"), + "report_cache_version": report_output.api_version, "simulation_cache_version": None, "requested_version_override": None, - "resolved_dataset": resolved_dataset, + "resolved_dataset": ( + report_spec.dataset + if report_spec is not None + and report_spec.report_kind in ECONOMY_REPORT_KINDS + else None + ), "resolved_options_hash": None, } - def _get_report_spec_status(self, report_spec: ReportSpec) -> str: - if report_spec.report_kind in ECONOMY_REPORT_KINDS: - return "backfilled_assumed" - return "explicit" + @staticmethod + def _report_spec_status(report_spec: ReportSpec) -> str: + return ( + "backfilled_assumed" + if report_spec.report_kind in ECONOMY_REPORT_KINDS + else "explicit" + ) - def _upsert_report_spec_in_transaction( + def _upsert_report_spec( self, - tx, - report_output: dict, - simulation_1: dict | None, - simulation_2: dict | None, + report_output: ReportOutput, + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> ReportSpec | None: if simulation_1 is None: return None - try: report_spec = self.report_spec_service.build_report_spec( - report_output=report_output, - simulation_1=simulation_1, - simulation_2=simulation_2, - ) - except ValueError as exc: - print( - "Skipping report spec sync for report output " - f"#{report_output['id']}. Details: {str(exc)}" + report_output, simulation_1, simulation_2 ) + except ValueError: return None - - report_spec_status = self._get_report_spec_status(report_spec) - existing_spec = parse_json_field(report_output.get("report_spec_json")) - if ( - existing_spec != report_spec.model_dump() - or report_output.get("report_kind") != report_spec.report_kind - or report_output.get("report_spec_schema_version") != 1 - or report_output.get("report_spec_status") != report_spec_status - ): - tx.reports.update( - report_output["id"], - report_kind=report_spec.report_kind, - report_spec_json=report_spec.model_dump(), - report_spec_schema_version=1, - report_spec_status=report_spec_status, - ) - report_output["report_kind"] = report_spec.report_kind - report_output["report_spec_json"] = report_spec.model_dump() - report_output["report_spec_schema_version"] = 1 - report_output["report_spec_status"] = report_spec_status - + report_output.report_kind = report_spec.report_kind + report_output.report_spec_json = report_spec.model_dump() + report_output.report_spec_schema_version = 1 + report_output.report_spec_status = self._report_spec_status(report_spec) return report_spec + @staticmethod def _run_matches_parent( - self, - run: dict, - report_output: dict, + run: ReportOutputRun, + report_output: ReportOutput, report_spec: ReportSpec | None, version_manifest: dict[str, str | None], ) -> bool: - expected_snapshot = ( - report_spec.model_dump() if report_spec is not None else None - ) return ( - run["status"] == report_output["status"] - and run.get("output") == report_output.get("output") - and run.get("error_message") == report_output.get("error_message") - and run.get("report_spec_snapshot_json") == expected_snapshot - and run.get("country_package_version") - == version_manifest["country_package_version"] - and run.get("policyengine_version") - == version_manifest["policyengine_version"] - and run.get("data_version") == version_manifest["data_version"] - and run.get("runtime_app_name") == version_manifest["runtime_app_name"] - and run.get("report_cache_version") - == version_manifest["report_cache_version"] - and run.get("simulation_cache_version") - == version_manifest["simulation_cache_version"] - and run.get("requested_version_override") - == version_manifest["requested_version_override"] - and run.get("resolved_dataset") == version_manifest["resolved_dataset"] - and run.get("resolved_options_hash") - == version_manifest["resolved_options_hash"] - ) - - def _insert_bootstrap_report_run( - self, - tx, - report_output: dict, - report_spec: ReportSpec | None, - version_manifest: dict[str, str | None], - ) -> None: - requested_at = self._utc_timestamp() - is_terminal = report_output["status"] in ("complete", "error") - has_started = report_output["status"] in ("running", "complete", "error") - started_at = requested_at if has_started else None - finished_at = requested_at if is_terminal else None - - tx.reports.create_run( - report_output["id"], - run_id=str(uuid.uuid4()), - status=report_output["status"], - output=report_output.get("output"), - error_message=report_output.get("error_message"), - trigger_type="initial", - requested_at=requested_at, - started_at=started_at, - finished_at=finished_at, - source_run_id=None, - report_spec_snapshot_json=( - report_spec.model_dump() if report_spec is not None else None - ), - **version_manifest, + run.status == report_output.status + and run.output == report_output.output + and run.error_message == report_output.error_message + and run.report_spec_snapshot_json + == (report_spec.model_dump() if report_spec else None) + and all( + getattr(run, field) == value + for field, value in version_manifest.items() + ) ) - def _update_report_run_in_transaction( + def _update_run_from_parent( self, - tx, - run_id: str, - report_output: dict, + run: ReportOutputRun, + report_output: ReportOutput, report_spec: ReportSpec | None, version_manifest: dict[str, str | None], - preserve_terminal_finished_at: bool = False, + *, + preserve_terminal_finished_at: bool, ) -> None: - run = tx.reports.get_run(run_id) - if run is None: - raise ValueError(f"Report output run {run_id} not found") - - fallback_timestamp = self._utc_timestamp() - requested_at = ( - run.get("requested_at") - or run.get("started_at") - or run.get("finished_at") - or fallback_timestamp - ) - if report_output["status"] in ("complete", "error"): - finished_at = self._utc_timestamp() - started_at = ( - run.get("started_at") - or run.get("finished_at") - or run.get("requested_at") - or finished_at + now = self._utc_timestamp() + run.requested_at = run.requested_at or run.started_at or run.finished_at or now + if report_output.status in {"complete", "error"}: + run.started_at = ( + run.started_at or run.finished_at or run.requested_at or now ) - if preserve_terminal_finished_at: - finished_at = run.get("finished_at") or finished_at - elif report_output["status"] == "running": - started_at = ( - run.get("started_at") - or run.get("requested_at") - or self._utc_timestamp() - ) - finished_at = None + if not preserve_terminal_finished_at or run.finished_at is None: + run.finished_at = now + elif report_output.status == "running": + run.started_at = run.started_at or run.requested_at or now + run.finished_at = None else: - started_at = None - finished_at = None - - tx.reports.update_run( - run_id, - status=report_output["status"], - output=report_output.get("output"), - error_message=report_output.get("error_message"), - requested_at=requested_at, - started_at=started_at, - finished_at=finished_at, - report_spec_snapshot_json=( - report_spec.model_dump() if report_spec is not None else None - ), - **version_manifest, + run.started_at = None + run.finished_at = None + run.status = report_output.status + run.output = report_output.output + run.error_message = report_output.error_message + run.report_spec_snapshot_json = ( + report_spec.model_dump() if report_spec else None ) + for field, value in version_manifest.items(): + setattr(run, field, value) - def _sync_parent_pointers_in_transaction( - self, - tx, - report_output: dict, - runs_descending: list[dict], + @staticmethod + def _sync_parent_pointers( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] ) -> None: - desired_active_run_id, desired_latest_successful_run_id = ( - determine_parent_pointers(report_output["status"], runs_descending) + latest_successful = next( + (run.id for run in runs_descending if run.status == "complete"), None ) - if ( - report_output.get("active_run_id") == desired_active_run_id - and report_output.get("latest_successful_run_id") - == desired_latest_successful_run_id - ): - return - - tx.reports.update( - report_output["id"], - active_run_id=desired_active_run_id, - latest_successful_run_id=desired_latest_successful_run_id, - ) - report_output["active_run_id"] = desired_active_run_id - report_output["latest_successful_run_id"] = desired_latest_successful_run_id + if report_output.status in {"pending", "running"} and runs_descending: + report_output.active_run_id = runs_descending[0].id + else: + report_output.active_run_id = None + if report_output.status == "complete" and latest_successful is None: + latest_successful = runs_descending[0].id if runs_descending else None + report_output.latest_successful_run_id = latest_successful - def _ensure_report_output_dual_write_state_in_transaction( + def ensure_report_output_dual_write_state( self, - tx, + session: Session, report_output_id: int, - *, country_id: str | None = None, - ) -> dict: - report_output = self._get_report_output_row( - report_output_id, - queryer=tx, - country_id=country_id, - for_update=True, + ) -> ReportOutput: + report_output = self._select_report_output( + session, report_output_id, country_id, for_update=True ) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") - try: simulation_1, simulation_2 = self._get_linked_simulations( + session, report_output, - queryer=tx, bootstrap_dual_write_state=True, ) - except ValueError as exc: - print( - "Skipping linked simulation sync for report output " - f"#{report_output_id}. Details: {str(exc)}" - ) + except ValueError: simulation_1, simulation_2 = None, None - - report_spec = self._upsert_report_spec_in_transaction( - tx, - report_output, - simulation_1, - simulation_2, - ) - version_manifest = self._build_version_manifest( - report_output, - report_spec=report_spec, - simulation_1=simulation_1, - simulation_2=simulation_2, + report_spec = self._upsert_report_spec( + report_output, simulation_1, simulation_2 ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx + manifest = self._build_version_manifest( + report_output, report_spec, simulation_1, simulation_2 ) - if not runs_descending: - self._insert_bootstrap_report_run( - tx, - report_output, - report_spec, - version_manifest, - ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx + runs = self._list_runs_descending(session, report_output_id) + if not runs: + self.report_run_service.create_report_output_run( + session, + report_output_id, + status=report_output.status, + output=report_output.output, + error_message=report_output.error_message, + trigger_type="initial", + report_spec_snapshot=( + report_spec.model_dump() if report_spec else None + ), + version_manifest=manifest, ) + runs = self._list_runs_descending(session, report_output_id) else: - mutable_run = self._select_mutable_run(report_output, runs_descending) - if mutable_run is not None: - run_matches_parent = self._run_matches_parent( - mutable_run, - report_output, - report_spec, - version_manifest, - ) - needs_timestamp_sync = self._run_needs_timestamp_sync( - mutable_run, report_output["status"] + mutable = self._select_mutable_run(report_output, runs) + if mutable is not None: + matches_result = ( + mutable.status == report_output.status + and mutable.output == report_output.output + and mutable.error_message == report_output.error_message ) - if not run_matches_parent or needs_timestamp_sync: - run_matches_result = run_matches_report_result( - mutable_run, report_output + if not self._run_matches_parent( + mutable, report_output, report_spec, manifest + ) or self._run_needs_timestamp_sync(mutable, report_output.status): + self._update_run_from_parent( + mutable, + report_output, + report_spec, + manifest, + preserve_terminal_finished_at=matches_result, ) - self._update_report_run_in_transaction( - tx, - run_id=mutable_run["id"], - report_output=report_output, - report_spec=report_spec, - version_manifest=version_manifest, - preserve_terminal_finished_at=run_matches_result, - ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx - ) - - self._sync_parent_pointers_in_transaction(tx, report_output, runs_descending) - refreshed_report_output = self._get_report_output_row( - report_output_id, - queryer=tx, - country_id=country_id, - ) - if refreshed_report_output is None: - raise ValueError(f"Report output #{report_output_id} not found after sync") - return self._with_display_run_timestamps(refreshed_report_output, queryer=tx) + session.flush() + runs = self._list_runs_descending(session, report_output_id) + self._sync_parent_pointers(report_output, runs) + session.flush() + return report_output - def ensure_report_output_dual_write_state( + def find_existing_report_output( self, - report_output_id: int, - country_id: str | None = None, - ) -> dict: - with self.unit_of_work.transaction() as daos: - return self._ensure_report_output_dual_write_state_in_transaction( - daos, - report_output_id, - country_id=country_id, + session: Session, + country_id: str, + simulation_1_id: int, + simulation_2_id: int | None = None, + year: str = "2025", + ) -> ReportOutput | None: + return session.scalar( + select(ReportOutput) + .where( + ReportOutput.country_id == country_id, + ReportOutput.simulation_1_id == simulation_1_id, + ReportOutput.simulation_2_id == simulation_2_id, + ReportOutput.year == year, + ReportOutput.api_version == get_report_output_cache_version(country_id), ) - - def get_stored_report_output( - self, country_id: str, report_output_id: int - ) -> dict | None: - """ - Get a stored report output row without aliasing to current runtime lineage. - - This is used by mutation paths that must address the originally - requested row. It still runs dual-write synchronization, so it may - bootstrap or repair run/spec metadata and returns the display-run - timestamp projection. It is therefore not a raw storage read. - - TODO: Split raw storage lookup from synchronized response projection in - a later run-backed read migration PR. - """ - report_output = self._get_report_output_row( - report_output_id, country_id=country_id - ) - if report_output is None: - return None - return self.ensure_report_output_dual_write_state( - report_output_id, - country_id=country_id, + .order_by(ReportOutput.id.desc()) ) - def report_output_exists(self, country_id: str, report_output_id: int) -> bool: - return ( - self._get_report_output_row(report_output_id, country_id=country_id) - is not None - ) - - def _is_current_report_output(self, report_output: dict) -> bool: - return report_output.get("api_version") == get_report_output_cache_version( - report_output["country_id"] + @staticmethod + def _require_simulation( + session: Session, country_id: str, simulation_id: int + ) -> Simulation: + simulation = session.scalar( + select(Simulation).where( + Simulation.id == simulation_id, + Simulation.country_id == country_id, + ) ) + if simulation is None: + raise ValueError( + f"Report output references missing simulation #{simulation_id}" + ) + return simulation - def _find_existing_report_output_row( + def create_report_output( self, - *, + session: Session, country_id: str, simulation_1_id: int, - simulation_2_id: int | None, - year: str, - queryer=None, - ) -> dict | None: - api_version = get_report_output_cache_version(country_id) - if queryer is None: - with self.unit_of_work.read() as daos: - return self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - queryer=daos, - ) - return queryer.reports.find_latest( + simulation_2_id: int | None = None, + year: str = "2025", + ) -> ReportOutput: + existing = self.find_existing_report_output( + session, country_id, simulation_1_id, simulation_2_id, year + ) + if existing is not None: + return self.ensure_report_output_dual_write_state( + session, existing.id, country_id + ) + self._require_simulation(session, country_id, simulation_1_id) + if simulation_2_id is not None: + self._require_simulation(session, country_id, simulation_2_id) + report_output = ReportOutput( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, + api_version=get_report_output_cache_version(country_id), + status="pending", year=year, - api_version=api_version, ) - - def _get_or_create_current_report_output(self, report_output: dict) -> dict: - current_report = self.find_existing_report_output( - country_id=report_output["country_id"], - simulation_1_id=report_output["simulation_1_id"], - simulation_2_id=report_output["simulation_2_id"], - year=report_output["year"], - ) - if current_report is not None: - return self._with_display_run_timestamps(current_report) - - return self.create_report_output( - country_id=report_output["country_id"], - simulation_1_id=report_output["simulation_1_id"], - simulation_2_id=report_output["simulation_2_id"], - year=report_output["year"], + session.add(report_output) + session.flush() + return self.ensure_report_output_dual_write_state( + session, report_output.id, country_id ) - def _alias_report_output(self, report_output_id: int, report_output: dict) -> dict: - aliased_report = dict(report_output) - aliased_report["id"] = report_output_id - return aliased_report - - def find_existing_report_output( - self, - country_id: str, - simulation_1_id: int, - simulation_2_id: int | None = None, - year: str = "2025", - ) -> dict | None: - """ - Find an existing report output with the same simulation IDs and year. - """ - print("Checking for existing report output") - - try: - existing_report = self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, + def get_report_output( + self, session: Session, country_id: str, report_output_id: int + ) -> ReportOutput | None: + if type(report_output_id) is not int or report_output_id < 0: + raise Exception( + f"Invalid report output ID: {report_output_id}. " + "Must be a positive integer." ) - if existing_report is not None: - print(f"Found existing report output with ID: {existing_report['id']}") - return self.ensure_report_output_dual_write_state( - existing_report["id"], - country_id=country_id, - ) - return None - - except Exception as e: - print(f"Error checking for existing report output. Details: {str(e)}") - raise e - - def create_report_output( - self, - country_id: str, - simulation_1_id: int, - simulation_2_id: int | None = None, - year: str = "2025", - ) -> dict: - """ - Create a new report output record with pending status. - """ - print("Creating new report output") - api_version = get_report_output_cache_version(country_id) - - try: - with self.unit_of_work.transaction() as daos: - existing_report = self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - queryer=daos, - ) - if existing_report is not None: - print( - f"Reusing existing report output with ID: {existing_report['id']}" - ) - return self._ensure_report_output_dual_write_state_in_transaction( - daos, - existing_report["id"], - country_id=country_id, - ) - - self._require_simulation_exists( - daos, - country_id=country_id, - simulation_id=simulation_1_id, - ) - if simulation_2_id is not None: - self._require_simulation_exists( - daos, - country_id=country_id, - simulation_id=simulation_2_id, - ) - - report_output_id = daos.reports.create( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - api_version=api_version, - status="pending", - year=year, - ) - created_report = daos.reports.get(report_output_id, country_id) - if created_report is None: - raise Exception("Failed to retrieve created report output") + return self._select_report_output(session, report_output_id, country_id) - print(f"Created report output with ID: {created_report['id']}") - return self._ensure_report_output_dual_write_state_in_transaction( - daos, - created_report["id"], - country_id=country_id, - ) - - except Exception as e: - print(f"Error creating report output. Details: {str(e)}") - raise e - - def get_report_output(self, country_id: str, report_output_id: int) -> dict | None: - """ - Get a report output record by ID. - """ - print(f"Getting report output {report_output_id}") - - try: - if type(report_output_id) is not int or report_output_id < 0: - raise Exception( - f"Invalid report output ID: {report_output_id}. Must be a positive integer." - ) + @staticmethod + def is_current_report_output(report_output: ReportOutput) -> bool: + return report_output.api_version == get_report_output_cache_version( + report_output.country_id + ) - report_output = self._get_report_output_row( - report_output_id, - country_id=country_id, + def get_or_create_current_report_output( + self, session: Session, report_output: ReportOutput + ) -> ReportOutput: + existing = self.find_existing_report_output( + session, + report_output.country_id, + report_output.simulation_1_id, + report_output.simulation_2_id, + report_output.year, + ) + if existing is not None: + return self.ensure_report_output_dual_write_state( + session, existing.id, report_output.country_id ) - if report_output is None: - return None - - if self._is_current_report_output(report_output): - return self.ensure_report_output_dual_write_state( - report_output_id, - country_id=country_id, - ) - - current_report = self._get_or_create_current_report_output(report_output) - return self._alias_report_output(report_output_id, current_report) + return self.create_report_output( + session, + report_output.country_id, + report_output.simulation_1_id, + report_output.simulation_2_id, + report_output.year, + ) - except Exception as e: - print( - f"Error fetching report output #{report_output_id}. Details: {str(e)}" - ) - raise e + def report_output_exists( + self, session: Session, country_id: str, report_output_id: int + ) -> bool: + return ( + self._select_report_output(session, report_output_id, country_id) + is not None + ) def update_report_output( self, + session: Session, country_id: str, report_id: int, status: str | None = None, - output: str | None = None, + output: dict | list | str | None = None, error_message: str | None = None, ) -> bool: - """ - Update a report output record with results or error. - """ - print(f"Updating report output {report_id}") - - try: - update_values = {} - - if status is not None: - update_values["status"] = status - - if output is not None: - update_values["output"] = parse_json_field(output) - - if error_message is not None: - update_values["error_message"] = error_message - - if not update_values: - print("No fields to update") - return False - - with self.unit_of_work.transaction() as daos: - requested_report = self._get_report_output_row( - report_id, - queryer=daos, - country_id=country_id, - for_update=True, - ) - if requested_report is None: - raise ValueError(f"Report output #{report_id} not found") - - if status == "running" and not self._has_mutable_running_run( - requested_report, queryer=daos - ): - raise ValueError( - "Cannot mark report output running without an active " - "pending or running report run" - ) - - daos.reports.update(report_id, **update_values) - self._ensure_report_output_dual_write_state_in_transaction( - daos, - report_id, - country_id=country_id, + values = { + key: value + for key, value in { + "status": status, + "output": output, + "error_message": error_message, + }.items() + if value is not None + } + if not values: + return False + if isinstance(values.get("output"), str): + values["output"] = json.loads(values["output"]) + report_output = self._select_report_output( + session, report_id, country_id, for_update=True + ) + if report_output is None: + raise ValueError(f"Report output #{report_id} not found") + if status == "running": + runs = self._list_runs_descending(session, report_id) + if not self._has_mutable_running_run(report_output, runs): + raise ValueError( + "Cannot mark report output running without an active pending " + "or running report run" ) - - print(f"Successfully updated report output #{report_id}") - return True - - except Exception as e: - print(f"Error updating report output #{report_id}. Details: {str(e)}") - raise e + for field, value in values.items(): + setattr(report_output, field, value) + self.ensure_report_output_dual_write_state(session, report_id, country_id) + return True diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 382fa30e5..d5b77f541 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -7,9 +7,10 @@ import pytest from flask import Flask, Response +from policyengine_api.constants import get_report_output_cache_version from policyengine_api.endpoints.household import get_calculate from policyengine_api.endpoints.policy import get_policy_search -from policyengine_api.data.v1_models import Household, Policy, Simulation +from policyengine_api.data.v1_models import Household, Policy, ReportOutput, Simulation from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp @@ -348,20 +349,43 @@ def _patched_route_dependencies(): stack.enter_context( patch( "policyengine_api.routes.report_output_routes.report_output_service.create_report_output", - return_value={ - "id": 33, - "country_id": "us", - "simulation_1_id": 11, - "simulation_2_id": None, - "status": "pending", - "year": "2026", - }, + return_value=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.report_output_routes.report_output_service.get_report_output", - return_value={"id": 33, "status": "pending", "country_id": "us"}, + return_value=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), + ) + ) + stack.enter_context( + patch( + "policyengine_api.routes.report_output_routes.report_output_service.ensure_report_output_dual_write_state", + return_value=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), ) ) return stack diff --git a/tests/unit/services/test_report_output_service.py b/tests/unit/services/test_report_output_service.py index 55ee2ff62..51fa34f36 100644 --- a/tests/unit/services/test_report_output_service.py +++ b/tests/unit/services/test_report_output_service.py @@ -1,1897 +1,215 @@ import pytest -import json -from datetime import datetime, timezone +from sqlalchemy import func, select from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.v1_models import ( + ReportOutput, + ReportOutputRun, +) from policyengine_api.services.report_output_service import ReportOutputService from policyengine_api.services.report_run_service import ReportRunService -from policyengine_api.services.run_sync_utils import select_display_report_run from policyengine_api.services.simulation_service import SimulationService -from tests.fixtures.services import report_output_fixtures - -pytest_plugins = ("tests.fixtures.services.report_output_fixtures",) service = ReportOutputService() -report_run_service = ReportRunService() +run_service = ReportRunService() simulation_service = SimulationService() -class TestReportOutputRunTimestamps: - def test_format_run_timestamp_handles_supported_values(self): - assert ( - service._format_run_timestamp(datetime(2026, 5, 4, 12, 0, 0)) - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp( - datetime(2026, 5, 4, 12, 0, 0, tzinfo=timezone.utc) - ) - == "2026-05-04T12:00:00Z" - ) - assert service._format_run_timestamp("") is None - assert ( - service._format_run_timestamp("2026-05-04T12:00:00") - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04T12:00:00Z") - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04T12:00:00+01:00") - == "2026-05-04T11:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04 12:00:00.123456") - == "2026-05-04T12:00:00Z" - ) - - def test_select_display_run_uses_matching_result_before_newest_fallback(self): - report_output = { - "id": 1, - "status": "complete", - "output": '{"ok": true}', - "error_message": None, - "active_run_id": None, - "latest_successful_run_id": None, - } - matching_run = { - "id": "matching", - "status": "complete", - "output": '{"ok": true}', - "error_message": None, - } - newest_non_matching_run = { - "id": "newest", - "status": "pending", - "output": None, - "error_message": None, - } - - selected_run = select_display_report_run( - report_output, [newest_non_matching_run, matching_run] - ) - - assert selected_run["id"] == "matching" - - -class TestFindExistingReportOutput: - """Test finding existing report outputs in the database.""" - - def test_find_existing_report_output_found(self, test_db, existing_report_record): - """Test finding an existing report output.""" - # GIVEN an existing report record (from fixture) - - # WHEN we search for a report with matching simulation IDs - result = service.find_existing_report_output( - country_id=existing_report_record["country_id"], - simulation_1_id=existing_report_record["simulation_1_id"], - simulation_2_id=existing_report_record["simulation_2_id"], - ) - - # THEN the result should contain the existing report - assert result is not None - assert result["id"] == existing_report_record["id"] - assert ( - result["country_id"] - == report_output_fixtures.valid_report_data["country_id"] - ) - assert result["simulation_1_id"] == existing_report_record["simulation_1_id"] - assert result["status"] == existing_report_record["status"] - - def test_find_existing_report_output_not_found(self, test_db): - """Test that None is returned when no report exists.""" - # GIVEN an empty database - - # WHEN we search for a non-existent report - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=999, - simulation_2_id=888, - year="2025", - ) - - # THEN None should be returned - assert result is None - - def test_find_existing_report_output_with_null_simulation2(self, test_db): - """Test finding reports where simulation_2_id is NULL.""" - api_version = get_report_output_cache_version("us") - # GIVEN a report with NULL simulation_2_id - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 100, None, "complete", api_version, "2025"), - ) - - # WHEN we search for it - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=100, - simulation_2_id=None, - year="2025", - ) - - # THEN we should find it - assert result is not None - assert result["simulation_1_id"] == 100 - assert result["simulation_2_id"] is None - assert result["year"] == "2025" - - def test_find_existing_report_output_with_year(self, test_db): - """Test finding reports with different years.""" - api_version = get_report_output_cache_version("us") - # GIVEN reports with different years for the same simulation - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 101, None, "complete", api_version, "2025"), - ) - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 101, None, "complete", api_version, "2024"), - ) - - # WHEN we search for the 2025 report - result_2025 = service.find_existing_report_output( - country_id="us", - simulation_1_id=101, - simulation_2_id=None, - year="2025", - ) - - # THEN we should find the 2025 report - assert result_2025 is not None - assert result_2025["simulation_1_id"] == 101 - assert result_2025["year"] == "2025" - - # WHEN we search for the 2024 report - result_2024 = service.find_existing_report_output( - country_id="us", - simulation_1_id=101, - simulation_2_id=None, - year="2024", - ) - - # THEN we should find the 2024 report - assert result_2024 is not None - assert result_2024["simulation_1_id"] == 101 - assert result_2024["year"] == "2024" - - # AND the two reports should have different IDs - assert result_2025["id"] != result_2024["id"] - - def test_find_existing_report_output_ignores_stale_runtime_version(self, test_db): - current_version = get_report_output_cache_version("us") - stale_version = "r0stale1" - assert stale_version != current_version - - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 102, None, "complete", stale_version, "2025"), - ) - - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=102, - simulation_2_id=None, - year="2025", - ) - - assert result is None - - -class TestCreateReportOutput: - """Test creating new report outputs in the database.""" - - def test_create_report_output_single_simulation(self, test_db): - """Test creating a report output with a single simulation.""" - # GIVEN an empty database - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_single_create", - population_type="household", - policy_id=1, - ) - - # WHEN we create a report output with one simulation - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert isinstance(created_report, dict) - assert created_report["id"] > 0 - assert created_report["simulation_1_id"] == simulation["id"] - assert created_report["simulation_2_id"] is None - assert created_report["status"] == "pending" - assert created_report["year"] == "2025" - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result is not None - assert result["simulation_1_id"] == simulation["id"] - assert result["simulation_2_id"] is None - assert result["status"] == "pending" - assert result["year"] == "2025" - - def test_create_report_output_comparison(self, test_db): - """Test creating a report output comparing two simulations.""" - # GIVEN an empty database - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_report_comparison", - population_type="household", - policy_id=2, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_report_comparison", - population_type="household", - policy_id=3, - ) - - # WHEN we create a report output with two simulations - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert created_report["simulation_1_id"] == simulation_1["id"] - assert created_report["simulation_2_id"] == simulation_2["id"] - assert created_report["status"] == "pending" - assert created_report["year"] == "2025" - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result["simulation_1_id"] == simulation_1["id"] - assert result["simulation_2_id"] == simulation_2["id"] - assert result["status"] == "pending" - assert result["year"] == "2025" - - def test_create_report_output_retrieves_correct_id(self, test_db): - """Test that create_report_output retrieves the correct ID without race conditions.""" - # GIVEN we create multiple reports rapidly - - # WHEN we create reports with different parameters - created_reports = [] - simulation_ids = [] - for i in range(3): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id=f"household_report_id_{i}", - population_type="household", - policy_id=100 + i, - ) - simulation_2 = None - if i % 2 != 0: - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id=f"household_report_id_{i}", - population_type="household", - policy_id=200 + i, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None if simulation_2 is None else simulation_2["id"], - year="2025", - ) - created_reports.append(report) - simulation_ids.append( - ( - simulation_1["id"], - None if simulation_2 is None else simulation_2["id"], - ) - ) - - # THEN all IDs should be unique - ids = [report["id"] for report in created_reports] - assert len(set(ids)) == 3 - - # AND each report should have the correct data - for i, report in enumerate(created_reports): - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report["id"],) - ).fetchone() - expected_sim1, expected_sim2 = simulation_ids[i] - assert result["simulation_1_id"] == expected_sim1 - assert result["simulation_2_id"] == expected_sim2 - assert result["year"] == "2025" - - def test_create_report_output_with_different_year(self, test_db): - """Test creating a report output with a different year.""" - # GIVEN an empty database - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_other_year", - population_type="household", - policy_id=4, - ) - - # WHEN we create a report output with year 2024 - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2024", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert created_report["year"] == "2024" - assert created_report["simulation_1_id"] == simulation["id"] - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result["year"] == "2024" - assert result["simulation_1_id"] == simulation["id"] - - def test_create_report_output_populates_dual_write_state(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_dual_write", - population_type="household", - policy_id=21, - ) - - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["report_kind"] == "household_single" - assert stored_report["report_spec_json"] is not None - assert stored_report["report_spec_schema_version"] == 1 - assert stored_report["report_spec_status"] == "explicit" - assert stored_report["active_run_id"] is not None - assert stored_report["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["trigger_type"] == "initial" - assert run["requested_at"] is not None - assert created_report["requested_at"] is not None - assert created_report["started_at"] is None - assert created_report["finished_at"] is None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "household_single" - - def test_create_report_output_reuses_existing_row_and_bootstraps_dual_write( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_existing_report", - population_type="household", - policy_id=22, - ) - - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND simulation_2_id IS NULL AND year = ? - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(rows) == 1 - assert created_report["id"] == rows[0]["id"] - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - - def test_create_report_output_populates_economy_comparison_report_spec( - self, test_db - ): - baseline_simulation = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=30, - ) - reform_simulation = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=31, - ) - - created_report = service.create_report_output( - country_id="us", - simulation_1_id=baseline_simulation["id"], - simulation_2_id=reform_simulation["id"], - year="2025", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["report_kind"] == "economy_comparison" - assert stored_report["report_spec_status"] == "backfilled_assumed" - - report_spec = stored_report["report_spec_json"] - if isinstance(report_spec, str): - report_spec = json.loads(report_spec) - assert report_spec["region"] == "state/ca" - assert report_spec["baseline_policy_id"] == 30 - assert report_spec["reform_policy_id"] == 31 - assert report_spec["dataset"] == "default" - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "economy_comparison" - assert snapshot["region"] == "state/ca" - - -class TestGetReportOutput: - """Test retrieving report outputs from the database.""" - - def test_get_report_output_existing(self, test_db, existing_report_record): - """Test retrieving an existing report output.""" - # GIVEN an existing report record - - # WHEN we retrieve the report - result = service.get_report_output( - country_id=existing_report_record["country_id"], - report_output_id=existing_report_record["id"], - ) - - # THEN the correct report should be returned - assert result is not None - assert result["id"] == existing_report_record["id"] - assert result["simulation_1_id"] == existing_report_record["simulation_1_id"] - assert result["status"] == existing_report_record["status"] - - def test_get_report_output_nonexistent(self, test_db): - """Test retrieving a non-existent report returns None.""" - # GIVEN an empty database - - # WHEN we try to retrieve a non-existent report - result = service.get_report_output(country_id="us", report_output_id=999) - - # THEN None should be returned - assert result is None - - def test_get_report_output_with_json_output(self, test_db): - """Test that JSON output is properly parsed when retrieved.""" - # GIVEN a report with JSON output - test_output = {"key": "value", "nested": {"data": 123}} - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 1, - None, - "complete", - json.dumps(test_output), - get_report_output_cache_version("us"), - "2025", - ), - ) - - # Get the ID of the inserted record - record = test_db.query( - "SELECT id FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - # WHEN we retrieve the report - result = service.get_report_output( - country_id="us", report_output_id=record["id"] - ) - - # THEN the output should be returned as JSON string (not parsed) - assert result["output"] == json.dumps(test_output) - assert result["year"] == "2025" - # Frontend will parse this string - - def test_get_report_output_includes_display_run_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_display_timestamps", - population_type="household", - policy_id=40, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 12:00:00", - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - run["id"], - ), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] == "2026-05-04T12:00:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" - - def test_update_report_output_sets_finished_at_on_display_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_finished_timestamp", - population_type="household", - policy_id=41, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - success = service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - - assert success is True - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - assert result["status"] == "complete" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is not None - - def test_error_rerun_uses_error_run_timestamp_over_previous_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_error_rerun_timestamp", - population_type="household", - policy_id=42, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - previous_success_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 10:00:00", - "2026-05-04 10:01:00", - "2026-05-04 10:02:00", - previous_success_id, - ), - ) - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], previous_success_id, report["id"]), - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="error", - error_message="rerun failed", - ) - - updated_rerun = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "error" - assert result["finished_at"] == service._format_run_timestamp( - updated_rerun["finished_at"] - ) - assert result["finished_at"] != "2026-05-04T10:02:00Z" - - def test_pending_update_clears_terminal_display_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_pending_timestamp_reset", - population_type="household", - policy_id=43, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="pending", - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "pending" - assert result["requested_at"] is not None - assert result["started_at"] is None - assert result["finished_at"] is None - - def test_running_update_sets_started_at_without_finished_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_timestamp", - population_type="household", - policy_id=44, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "running" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is None - - def test_running_update_requires_non_terminal_run_after_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_without_active_run", - population_type="household", - policy_id=45, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - - with pytest.raises(ValueError, match="active pending or running"): - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - assert stored_report["status"] == "complete" - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] == successful_run_id - assert successful_run["status"] == "complete" - assert successful_run["output"] == json.dumps({"ok": True}) - assert successful_run["finished_at"] is not None - - def test_running_update_uses_active_rerun_without_rewriting_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_active_running_rerun", - population_type="household", - policy_id=46, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - active_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - assert successful_run["status"] == "complete" - assert successful_run["finished_at"] is not None - assert active_run["status"] == "running" - assert active_run["started_at"] is not None - assert active_run["finished_at"] is None - assert stored_report["active_run_id"] == rerun["id"] - assert stored_report["latest_successful_run_id"] == successful_run_id - - def test_get_report_output_does_not_rewrite_terminal_active_run_for_running_parent( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_bad_active_run", - population_type="household", - policy_id=47, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - output_json = json.dumps({"ok": True}) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=output_json, - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_outputs - SET status = ?, active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("running", successful_run_id, successful_run_id, report["id"]), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - assert result["status"] == "running" - assert successful_run["status"] == "complete" - assert successful_run["output"] == output_json - assert successful_run["finished_at"] is not None - - def test_get_stored_report_output_includes_display_run_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_stored_timestamp", - population_type="household", - policy_id=48, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - result = service.get_stored_report_output("us", report["id"]) - - assert result is not None - assert result["requested_at"] is not None - assert result["started_at"] is None - assert result["finished_at"] is None - - def test_get_stored_report_output_returns_none_when_missing(self, test_db): - assert service.get_stored_report_output("us", 999999) is None - - def test_get_report_output_backfills_missing_timestamps_on_matching_legacy_run( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_timestamp_get", - population_type="household", - policy_id=46, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, started_at = NULL, finished_at = NULL - WHERE report_output_id = ? - """, - (report["id"],), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is not None - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert stored_run["requested_at"] is not None - assert stored_run["started_at"] is not None - assert stored_run["finished_at"] is not None - - def test_get_report_output_preserves_existing_finished_at_during_backfill( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_finished_at", - population_type="household", - policy_id=47, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, started_at = ?, finished_at = ? - WHERE report_output_id = ? - """, - ( - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - report["id"], - ), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] == "2026-05-04T12:01:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert service._format_run_timestamp(stored_run["requested_at"]) == ( - "2026-05-04T12:01:00Z" - ) - assert service._format_run_timestamp(stored_run["finished_at"]) == ( - "2026-05-04T12:02:00Z" - ) - - def test_get_report_output_preserves_finished_at_while_backfilling_metadata( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_finished_metadata", - population_type="household", - policy_id=48, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, - started_at = ?, - finished_at = ?, - report_spec_snapshot_json = NULL, - country_package_version = NULL - WHERE report_output_id = ? - """, - ( - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - report["id"], - ), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] == "2026-05-04T12:01:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert stored_run["report_spec_snapshot_json"] is not None - assert stored_run["country_package_version"] is not None - assert service._format_run_timestamp(stored_run["requested_at"]) == ( - "2026-05-04T12:01:00Z" - ) - assert service._format_run_timestamp(stored_run["finished_at"]) == ( - "2026-05-04T12:02:00Z" - ) - - def test_get_report_output_bootstraps_running_legacy_run_started_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_running", - population_type="household", - policy_id=49, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, status, api_version, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - "running", - get_report_output_cache_version("us"), - "2025", - ), - ) - report = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "running" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is None - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert run["status"] == "running" - assert run["started_at"] is not None - assert run["finished_at"] is None - - def test_find_existing_report_output_backfills_missing_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_timestamp_find", - population_type="household", - policy_id=50, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL - WHERE report_output_id = ? - """, - (report["id"],), - ) - - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - assert result is not None - assert result["requested_at"] is not None - - def test_get_report_output_resolves_stale_id_to_current_runtime_row(self, test_db): - stale_output = { - "budget": {"budgetary_impact": 1}, - "congressional_district_impact": { - "districts": [ - { - "district": "AL-01", - "average_household_income_change": 120, - "relative_household_income_change": 0.01, - } - ] - }, - } - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 2, - None, - "complete", - json.dumps(stale_output), - "r0stale1", - "2025", - ), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - current_version = get_report_output_cache_version("us") - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 2, - None, - "complete", - json.dumps({"budget": {"budgetary_impact": 2}}), - current_version, - "2025", - ), - ) - - current_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=stale_record["id"] - ) - assert result is not None - assert result["id"] == stale_record["id"] - assert result["api_version"] == current_record["api_version"] - assert result["output"] == current_record["output"] - - def test_get_report_output_creates_current_runtime_row_for_stale_id(self, test_db): - stale_version = "r0stale1" - current_version = get_report_output_cache_version("us") - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_stale_runtime_create", - population_type="household", - policy_id=5, - ) - - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", simulation["id"], None, "complete", stale_version, "2025"), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=stale_record["id"] - ) - - assert result is not None - assert result["id"] == stale_record["id"] - assert result["api_version"] == current_version - assert result["status"] == "pending" - assert result["output"] is None - - current_rows = test_db.query( - "SELECT * FROM report_outputs WHERE country_id = ? AND simulation_1_id = ? AND year = ? ORDER BY id ASC", - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(current_rows) == 2 - assert current_rows[0]["api_version"] == stale_version - assert current_rows[1]["api_version"] == current_version - - def test_get_report_output_invalid_id(self, test_db): - """Test that invalid report IDs are handled properly.""" - # GIVEN any database state - - # WHEN we call get_report_output with invalid ID types - # THEN an exception should be raised - with pytest.raises(Exception) as exc_info: - service.get_report_output(country_id="us", report_output_id=-1) - assert "Invalid report output ID" in str(exc_info.value) - - with pytest.raises(Exception) as exc_info: - service.get_report_output(country_id="us", report_output_id="not_an_int") - assert "Invalid report output ID" in str(exc_info.value) - - def test_get_report_output_wrong_country_returns_none( - self, test_db, existing_report_record - ): - result = service.get_report_output( - country_id="uk", - report_output_id=existing_report_record["id"], - ) - - assert result is None - - -class TestUniqueConstraint: - """Test that the unique constraint on report outputs works correctly.""" - - def test_duplicate_report_returns_existing(self, test_db): - """Test that creating duplicate reports returns the existing record.""" - # GIVEN we create a report - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_duplicate_report", - population_type="household", - policy_id=50, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_duplicate_report", - population_type="household", - policy_id=60, - ) - first_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - # WHEN we try to create an identical report - second_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - # THEN the same report should be returned (no duplicate created) - assert first_report["id"] == second_report["id"] - assert first_report["country_id"] == second_report["country_id"] - assert first_report["simulation_1_id"] == second_report["simulation_1_id"] - assert first_report["simulation_2_id"] == second_report["simulation_2_id"] - assert first_report["year"] == second_report["year"] - - -class TestUpdateReportOutput: - """Test updating report outputs in the database.""" - - def test_update_report_output_to_complete(self, test_db, existing_report_record): - """Test updating a report to complete status with output.""" - # GIVEN an existing pending report - report_id = existing_report_record["id"] - test_output = {"result": "success", "data": [1, 2, 3]} - test_output_json = json.dumps(test_output) - - # WHEN we update it to complete with output (as JSON string) - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="complete", - output=test_output_json, - ) - - # THEN the update should succeed - assert success is True - - # AND the database should reflect the changes - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "complete" - assert result["output"] == test_output_json - - def test_update_report_output_to_error(self, test_db, existing_report_record): - """Test updating a report to error status with message.""" - # GIVEN an existing pending report - report_id = existing_report_record["id"] - error_msg = "Calculation failed due to invalid input" - - # WHEN we update it to error status - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="error", - error_message=error_msg, - ) - - # THEN the update should succeed - assert success is True - - # AND the error should be recorded - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "error" - assert result["error_message"] == error_msg - - def test_update_report_output_partial_update(self, test_db, existing_report_record): - """Test that partial updates work correctly.""" - # GIVEN an existing report - report_id = existing_report_record["id"] - - # WHEN we update only the status - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="complete", - ) - - # THEN the update should succeed - assert success is True - - # AND only the status should change - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "complete" - assert result["output"] is None # Should remain unchanged - - def test_update_report_output_updates_dual_write_state(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_update", - population_type="household", - policy_id=23, - ) - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - output_json = json.dumps({"distribution": [1, 2, 3]}) - - success = service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=output_json, - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] is not None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run["status"] == "complete" - assert run["output"] == output_json - assert run["id"] == stored_report["latest_successful_run_id"] - - def test_update_report_output_bootstraps_missing_run_state(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_legacy_report", - population_type="household", - policy_id=24, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_legacy_report", - population_type="household", - policy_id=25, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=report_output["id"], - status="error", - error_message="legacy report failure", - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is not None - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "error" - assert run["error_message"] == "legacy report failure" - - def test_update_report_output_keeps_invalid_legacy_linkage_working(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="us", - population_type="geography", - policy_id=26, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=27, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=report_output["id"], - status="complete", - output=json.dumps({"budget": {"budgetary_impact": 1}}), - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - assert run["report_spec_snapshot_json"] is None - - def test_update_report_output_stale_id_updates_only_stale_lineage_run( - self, test_db - ): - stale_version = "r0stale1" - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_stale_lineage", - population_type="household", - policy_id=32, - ) - test_db.query( - """ - INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", simulation["id"], None, "pending", stale_version, "2025"), - ) - stale_report = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=stale_report["id"], - status="complete", - output=json.dumps({"result": "stale"}), - ) - - assert success is True - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND year = ? - ORDER BY id ASC - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(rows) == 1 - assert rows[0]["id"] == stale_report["id"] - assert rows[0]["status"] == "complete" - - runs = test_db.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (stale_report["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["status"] == "complete" - assert runs[0]["output"] == json.dumps({"result": "stale"}) - - def test_update_report_output_does_not_append_extra_run_for_legacy_patch_traffic( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_single_run", - population_type="household", - policy_id=33, - ) - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert first_run is not None - - success = service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=json.dumps({"distribution": [4, 5, 6]}), - ) - - assert success is True - - runs = test_db.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (created_report["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["id"] == first_run["id"] - assert runs[0]["status"] == "complete" - - def test_update_report_output_no_fields_returns_false( - self, test_db, existing_report_record - ): - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=existing_report_record["id"], - ) - - assert success is False - - def test_update_report_output_stale_id_keeps_stale_output_quarantined( - self, test_db - ): - stale_version = "r0stale1" - output_json = json.dumps({"result": "fresh"}) - - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", 4, None, "pending", stale_version, "2025"), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=stale_record["id"], - status="complete", - output=output_json, - ) - - assert success is True - - rows = test_db.query( - "SELECT * FROM report_outputs WHERE country_id = ? AND simulation_1_id = ? AND year = ? ORDER BY id ASC", - ("us", 4, "2025"), - ).fetchall() - - assert len(rows) == 1 - assert rows[0]["id"] == stale_record["id"] - assert rows[0]["api_version"] == stale_version - assert rows[0]["status"] == "complete" - assert rows[0]["output"] == output_json - - def test_create_report_output_rolls_back_parent_insert_on_dual_write_failure( - self, test_db, monkeypatch - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_create_rollback", - population_type="household", - policy_id=34, - ) - - def fail_dual_write(tx, report_output_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_report_output_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND simulation_2_id IS NULL AND year = ? - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert rows == [] - - def test_update_report_output_rolls_back_parent_update_on_dual_write_failure( - self, test_db, monkeypatch - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_update_rollback", - population_type="household", - policy_id=35, - ) - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - def fail_dual_write(tx, report_output_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_report_output_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=json.dumps({"rolled_back": True}), - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["status"] == "pending" - assert stored_report["output"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["output"] is None - - def test_ensure_report_output_dual_write_state_bootstraps_linked_simulations( - self, test_db - ): - test_db.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - "us-system-1.0.0", - "household_stage5_linked", - "household", - 36, - "pending", - ), - ) - simulation_1 = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - test_db.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - "us-system-1.0.0", - "household_stage5_linked", - "household", - 37, - "pending", - ), - ) - simulation_2 = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - synced_report = service.ensure_report_output_dual_write_state( - report_output["id"], - country_id="us", - ) - - assert synced_report["active_run_id"] is not None +def create_simulation(orm_session, *, policy_id=1, population_id="household-1"): + return simulation_service.create_simulation( + orm_session, + country_id="us", + population_id=population_id, + population_type="household", + policy_id=policy_id, + ) + + +def test_creates_mapped_report_with_spec_and_initial_run(orm_session): + simulation = create_simulation(orm_session) + + report = service.create_report_output( + orm_session, + country_id="us", + simulation_1_id=simulation.id, + year="2025", + ) + + assert isinstance(report, ReportOutput) + assert report.report_kind == "household_single" + assert isinstance(report.report_spec_json, dict) + run = orm_session.get(ReportOutputRun, report.active_run_id) + assert isinstance(run, ReportOutputRun) + assert run.status == "pending" + assert run.report_spec_snapshot_json == report.report_spec_json + + +def test_create_reuses_current_report_and_repairs_dual_state(orm_session): + simulation = create_simulation(orm_session) + existing = ReportOutput( + country_id="us", + simulation_1_id=simulation.id, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2025", + ) + orm_session.add(existing) + orm_session.flush() + + result = service.create_report_output(orm_session, "us", simulation.id, year="2025") + + assert result is existing + assert existing.active_run_id is not None + assert existing.report_spec_json is not None - stored_simulation_1 = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_1["id"],), - ).fetchone() - stored_simulation_2 = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_2["id"],), - ).fetchone() - assert stored_simulation_1["simulation_spec_json"] is not None - assert stored_simulation_1["active_run_id"] is not None - assert stored_simulation_2["simulation_spec_json"] is not None - assert stored_simulation_2["active_run_id"] is not None - simulation_1_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation_1["id"],), - ).fetchone() - simulation_2_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation_2["id"],), - ).fetchone() - assert simulation_1_run is not None - assert simulation_2_run is not None +@pytest.mark.parametrize("missing_secondary", [False, True]) +def test_create_rejects_missing_linked_simulation(orm_session, missing_secondary): + simulation = create_simulation(orm_session) if missing_secondary else None + + with pytest.raises(ValueError, match="references missing simulation"): + service.create_report_output( + orm_session, + "us", + simulation.id if simulation else 999, + 999 if missing_secondary else None, + ) + + assert orm_session.scalar(select(func.count()).select_from(ReportOutput)) == 0 + + +def test_find_existing_report_uses_current_cache_version(orm_session): + simulation = create_simulation(orm_session) + stale = ReportOutput( + country_id="us", + simulation_1_id=simulation.id, + simulation_2_id=None, + api_version="stale", + status="pending", + year="2025", + ) + orm_session.add(stale) + orm_session.flush() + + assert ( + service.find_existing_report_output( + orm_session, "us", simulation.id, year="2025" + ) + is None + ) + + current = service.create_report_output( + orm_session, "us", simulation.id, year="2025" + ) + assert ( + service.find_existing_report_output( + orm_session, "us", simulation.id, year="2025" + ) + is current + ) + + +def test_get_report_is_scoped_to_country_and_validates_id(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + + assert service.get_report_output(orm_session, "us", report.id) is report + assert service.get_report_output(orm_session, "uk", report.id) is None + with pytest.raises(Exception, match="Invalid report output ID"): + service.get_report_output(orm_session, "us", -1) + + +def test_update_complete_stores_python_json_and_promotes_run(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + active_run_id = report.active_run_id + + assert service.update_report_output( + orm_session, + "us", + report.id, + status="complete", + output={"ok": True}, + ) + + assert report.output == {"ok": True} + assert report.active_run_id is None + assert report.latest_successful_run_id == active_run_id + run = orm_session.get(ReportOutputRun, active_run_id) + assert run.status == "complete" + assert run.output == {"ok": True} + assert run.started_at is not None + assert run.finished_at is not None + + +def test_update_accepts_existing_v1_json_string_boundary(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + + service.update_report_output(orm_session, "us", report.id, output='{"ok": true}') + + assert report.output == {"ok": True} + + +def test_update_running_requires_mutable_run(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + service.update_report_output( + orm_session, "us", report.id, status="complete", output={"ok": True} + ) + + with pytest.raises(ValueError, match="without an active pending or running"): + service.update_report_output(orm_session, "us", report.id, status="running") + + +def test_update_running_targets_active_rerun(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + service.update_report_output( + orm_session, "us", report.id, status="complete", output={"old": True} + ) + successful_run_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + orm_session, report.id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + + service.update_report_output(orm_session, "us", report.id, status="running") + + assert rerun.status == "running" + assert rerun.started_at is not None + assert report.active_run_id == rerun.id + assert report.latest_successful_run_id == successful_run_id + + +def test_failed_rerun_preserves_latest_successful_pointer(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + service.update_report_output( + orm_session, "us", report.id, status="complete", output={"old": True} + ) + successful_run_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + orm_session, report.id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + + service.update_report_output( + orm_session, "us", report.id, status="error", error_message="failed" + ) + + assert report.active_run_id is None + assert report.latest_successful_run_id == successful_run_id + assert rerun.status == "error" + assert rerun.finished_at is not None + + +def test_noop_and_missing_updates(orm_session): + simulation = create_simulation(orm_session) + report = service.create_report_output(orm_session, "us", simulation.id) + + assert service.update_report_output(orm_session, "us", report.id) is False + with pytest.raises(ValueError, match="Report output #999 not found"): + service.update_report_output(orm_session, "us", 999, status="pending") diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_stage5_routes.py index 4a6c7be17..e95fba89a 100644 --- a/tests/unit/test_stage5_routes.py +++ b/tests/unit/test_stage5_routes.py @@ -1,8 +1,10 @@ import json +from datetime import datetime from flask import Flask from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun, Simulation from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.simulation_routes import simulation_bp from policyengine_api.services.report_output_service import ReportOutputService @@ -11,8 +13,8 @@ simulation_service = SimulationService() -report_output_service = ReportOutputService() -report_run_service = ReportRunService() +report_service = ReportOutputService() +run_service = ReportRunService() def create_test_client() -> Flask: @@ -23,597 +25,254 @@ def create_test_client() -> Flask: return app.test_client() -def test_create_simulation_existing_row_repairs_dual_write_state(test_db): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_route_repair", "household", 40, "pending"), - ) - simulation = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() +def create_simulation(factory, *, population_id="household-1", policy_id=1): + with factory.begin() as session: + simulation = simulation_service.create_simulation( + session, "us", population_id, "household", policy_id + ) + return simulation.id - client = create_test_client() - response = client.post( + +def create_report(factory, simulation_id): + with factory.begin() as session: + report = report_service.create_report_output( + session, "us", simulation_id, year="2025" + ) + return report.id + + +def test_create_simulation_existing_row_repairs_dual_write_state( + orm_session_factory, +): + with orm_session_factory.begin() as session: + simulation = Simulation( + country_id="us", + api_version="old", + population_id="household-route-repair", + population_type="household", + policy_id=40, + status="pending", + ) + session.add(simulation) + session.flush() + simulation_id = simulation.id + + response = create_test_client().post( "/us/simulation", json={ - "population_id": "household_route_repair", + "population_id": "household-route-repair", "population_type": "household", "policy_id": 40, }, ) assert response.status_code == 200 - payload = response.get_json() - assert payload["message"] == "Simulation already exists" - assert payload["result"]["id"] == simulation["id"] - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["active_run_id"] is not None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation["id"],), - ).fetchone() - assert run is not None - - -def test_create_report_output_existing_row_repairs_dual_write_state(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_report", - population_type="household", - policy_id=41, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() - response = client.post( + assert response.get_json()["result"]["id"] == simulation_id + with orm_session_factory() as session: + simulation = session.get(Simulation, simulation_id) + assert simulation.simulation_spec_json is not None + assert simulation.active_run_id is not None + + +def test_create_report_existing_row_repairs_dual_write_state(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=41) + with orm_session_factory.begin() as session: + report = ReportOutput( + country_id="us", + simulation_1_id=simulation_id, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2025", + ) + session.add(report) + session.flush() + report_id = report.id + + response = create_test_client().post( "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, + json={"simulation_1_id": simulation_id, "year": "2025"}, ) assert response.status_code == 200 - payload = response.get_json() - assert payload["message"] == "Report output already exists" - assert payload["result"]["id"] == report_output["id"] - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is not None - assert stored_report["active_run_id"] is not None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "household_single" - - -def test_post_report_output_returns_timestamp_fields_for_new_and_existing_report( - test_db, -): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_report_timestamps", - population_type="household", - policy_id=46, - ) + assert response.get_json()["result"]["id"] == report_id + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + assert report.report_spec_json is not None + assert report.active_run_id is not None - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, - ) - assert response.status_code == 201 - payload = response.get_json() - created_report = payload["result"] - assert created_report["requested_at"] is not None - assert created_report["started_at"] is None - assert created_report["finished_at"] is None +def test_report_post_returns_run_timestamps_for_new_and_existing( + orm_session_factory, +): + simulation_id = create_simulation(orm_session_factory, policy_id=42) + client = create_test_client() - existing_response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, + created = client.post( + "/us/report", json={"simulation_1_id": simulation_id, "year": "2025"} ) - - assert existing_response.status_code == 200 - existing_payload = existing_response.get_json() - existing_report = existing_payload["result"] - assert existing_report["id"] == created_report["id"] - assert existing_report["requested_at"] is not None - assert existing_report["started_at"] is None - assert existing_report["finished_at"] is None - - -def test_create_report_output_missing_primary_simulation_returns_bad_request(test_db): - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": 999999, - "simulation_2_id": None, - "year": "2025", - }, + existing = client.post( + "/us/report", json={"simulation_1_id": simulation_id, "year": "2025"} ) - assert response.status_code == 400 + assert created.status_code == 201 + assert existing.status_code == 200 + for response in (created, existing): + result = response.get_json()["result"] + assert result["requested_at"] is not None + assert result["started_at"] is None + assert result["finished_at"] is None - report_rows = test_db.query("SELECT * FROM report_outputs").fetchall() - report_run_rows = test_db.query("SELECT * FROM report_output_runs").fetchall() - assert report_rows == [] - assert report_run_rows == [] +def test_report_post_rejects_missing_linked_simulations(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=43) + client = create_test_client() -def test_create_report_output_missing_secondary_simulation_returns_bad_request(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_missing_secondary", - population_type="household", - policy_id=42, + missing_primary = client.post( + "/us/report", json={"simulation_1_id": 999999, "year": "2025"} ) - - client = create_test_client() - response = client.post( + missing_secondary = client.post( "/us/report", json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": simulation["id"] + 999999, + "simulation_1_id": simulation_id, + "simulation_2_id": 999999, "year": "2025", }, ) - assert response.status_code == 400 + assert missing_primary.status_code == 400 + assert missing_secondary.status_code == 400 - report_rows = test_db.query( - "SELECT * FROM report_outputs WHERE simulation_1_id = ?", - (simulation["id"],), - ).fetchall() - report_run_rows = test_db.query("SELECT * FROM report_output_runs").fetchall() - assert report_rows == [] - assert report_run_rows == [] - - -def test_get_simulation_wrong_country_returns_not_found(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_wrong_country_get", - population_type="household", - policy_id=43, - ) +def test_simulation_routes_scope_reads_and_writes_to_country(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=44) client = create_test_client() - response = client.get(f"/uk/simulation/{simulation['id']}") - assert response.status_code == 404 - - -def test_patch_simulation_wrong_country_returns_not_found_and_does_not_mutate(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_wrong_country_patch", - population_type="household", - policy_id=44, - ) - - client = create_test_client() + assert client.get(f"/uk/simulation/{simulation_id}").status_code == 404 response = client.patch( "/uk/simulation", - json={ - "id": simulation["id"], - "status": "complete", - "output": json.dumps({"should_not": "persist"}), - }, + json={"id": simulation_id, "status": "complete", "output": {"bad": True}}, ) assert response.status_code == 404 + with orm_session_factory() as session: + assert session.get(Simulation, simulation_id).status == "pending" - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["country_id"] == "us" - assert stored_simulation["status"] == "pending" - assert stored_simulation["output"] is None - - -def test_get_report_output_wrong_country_returns_not_found(test_db): - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", 999, None, get_report_output_cache_version("us"), "pending", "2025"), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() +def test_report_routes_scope_reads_and_writes_to_country(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=45) + report_id = create_report(orm_session_factory, simulation_id) client = create_test_client() - response = client.get(f"/uk/report/{report_output['id']}") - - assert response.status_code == 404 - -def test_patch_report_output_wrong_country_returns_not_found_and_does_not_mutate( - test_db, -): - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", 1000, None, get_report_output_cache_version("us"), "pending", "2025"), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() + assert client.get(f"/uk/report/{report_id}").status_code == 404 response = client.patch( "/uk/report", - json={ - "id": report_output["id"], - "status": "complete", - "output": json.dumps({"should_not": "persist"}), - }, + json={"id": report_id, "status": "complete", "output": {"bad": True}}, ) assert response.status_code == 404 + with orm_session_factory() as session: + assert session.get(ReportOutput, report_id).status == "pending" - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["country_id"] == "us" - assert stored_report["status"] == "pending" - assert stored_report["output"] is None - - -def test_patch_report_output_accepts_running_status(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_running_report", - population_type="household", - policy_id=45, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "running", - }, - ) +def test_report_get_serializes_display_run_timestamps(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=46) + report_id = create_report(orm_session_factory, simulation_id) + with orm_session_factory.begin() as session: + report_service.update_report_output( + session, "us", report_id, status="complete", output={"ok": True} + ) + report = session.get(ReportOutput, report_id) + run = session.get(ReportOutputRun, report.latest_successful_run_id) + run.requested_at = datetime(2026, 5, 4, 12, 0) + run.started_at = datetime(2026, 5, 4, 12, 1) + run.finished_at = datetime(2026, 5, 4, 12, 2) - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "running" - assert payload["result"]["requested_at"] is not None - assert payload["result"]["started_at"] is not None - assert payload["result"]["finished_at"] is None - - -def test_get_report_output_serializes_display_run_timestamps(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_get_timestamp", - population_type="household", - policy_id=47, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 12:00:00", - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - run["id"], - ), - ) + result = create_test_client().get(f"/us/report/{report_id}").get_json()["result"] - client = create_test_client() - response = client.get(f"/us/report/{report['id']}") + assert result["requested_at"] == "2026-05-04T12:00:00Z" + assert result["started_at"] == "2026-05-04T12:01:00Z" + assert result["finished_at"] == "2026-05-04T12:02:00Z" - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["requested_at"] == "2026-05-04T12:00:00Z" - assert payload["result"]["started_at"] == "2026-05-04T12:01:00Z" - assert payload["result"]["finished_at"] == "2026-05-04T12:02:00Z" - - -def test_patch_report_output_running_uses_active_rerun_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_running_rerun", - population_type="household", - policy_id=48, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - client = create_test_client() - response = client.patch( +def test_report_patch_updates_active_rerun_and_preserves_success( + orm_session_factory, +): + simulation_id = create_simulation(orm_session_factory, policy_id=47) + report_id = create_report(orm_session_factory, simulation_id) + with orm_session_factory.begin() as session: + report_service.update_report_output( + session, "us", report_id, status="complete", output={"old": True} + ) + report = session.get(ReportOutput, report_id) + successful_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + session, report_id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + rerun_id = rerun.id + + running = create_test_client().patch( + "/us/report", json={"id": report_id, "status": "running"} + ) + failed = create_test_client().patch( "/us/report", - json={ - "id": report["id"], - "status": "running", - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "running" - assert payload["result"]["started_at"] is not None - assert payload["result"]["finished_at"] is None - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - active_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - assert successful_run["status"] == "complete" - assert successful_run["finished_at"] is not None - assert active_run["status"] == "running" - assert active_run["started_at"] is not None - assert active_run["finished_at"] is None - - -def test_patch_report_output_error_uses_active_rerun_timestamp_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_error_rerun", - population_type="household", - policy_id=49, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 10:00:00", - "2026-05-04 10:01:00", - "2026-05-04 10:02:00", - successful_run_id, - ), - ) - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - client = create_test_client() - response = client.patch( + json={"id": report_id, "status": "error", "error_message": "failed"}, + ) + + assert running.status_code == 200 + assert running.get_json()["result"]["started_at"] is not None + assert failed.status_code == 200 + assert failed.get_json()["result"]["finished_at"] is not None + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + rerun = session.get(ReportOutputRun, rerun_id) + assert report.latest_successful_run_id == successful_id + assert rerun.status == "error" + + +def test_report_patch_complete_promotes_active_rerun(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=48) + report_id = create_report(orm_session_factory, simulation_id) + with orm_session_factory.begin() as session: + report_service.update_report_output( + session, "us", report_id, status="complete", output={"old": True} + ) + report = session.get(ReportOutput, report_id) + rerun = run_service.create_report_output_run( + session, report_id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + rerun_id = rerun.id + + response = create_test_client().patch( "/us/report", - json={ - "id": report["id"], - "status": "error", - "error_message": "rerun failed", - }, + json={"id": report_id, "status": "complete", "output": {"new": True}}, ) assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "error" - assert payload["result"]["finished_at"] is not None - assert payload["result"]["finished_at"] != "2026-05-04T10:02:00Z" - - -def test_patch_report_output_complete_promotes_active_rerun_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_complete_rerun", - population_type="household", - policy_id=50, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + assert report.active_run_id is None + assert report.latest_successful_run_id == rerun_id - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "complete", - "output": json.dumps({"ok": "rerun"}), - }, - ) - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "complete" - assert payload["result"]["finished_at"] is not None - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] == rerun["id"] - - -def test_simulation_v1_routes_keep_json_fields_as_strings(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_v1_json_contract", - population_type="household", - policy_id=51, - ) +def test_simulation_v1_routes_keep_json_fields_as_strings(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=49) output = {"result": "ok", "values": [1, 2, 3]} - patch_response = create_test_client().patch( + patched = create_test_client().patch( "/us/simulation", - json={ - "id": simulation["id"], - "status": "complete", - "output": output, - }, - ) - - assert patch_response.status_code == 200 - patched_simulation = patch_response.get_json()["result"] - assert isinstance(patched_simulation["output"], str) - assert json.loads(patched_simulation["output"]) == output - assert isinstance(patched_simulation["simulation_spec_json"], str) - assert json.loads(patched_simulation["simulation_spec_json"])["country_id"] == "us" - - get_response = create_test_client().get(f"/us/simulation/{simulation['id']}") - - assert get_response.status_code == 200 - fetched_simulation = get_response.get_json()["result"] - assert isinstance(fetched_simulation["output"], str) - assert json.loads(fetched_simulation["output"]) == output - - orm_simulation = simulation_service.get_simulation("us", simulation["id"]) - assert orm_simulation["output"] == output - assert isinstance(orm_simulation["simulation_spec_json"], dict) + json={"id": simulation_id, "status": "complete", "output": output}, + ) + fetched = create_test_client().get(f"/us/simulation/{simulation_id}") + + for response in (patched, fetched): + result = response.get_json()["result"] + assert isinstance(result["output"], str) + assert json.loads(result["output"]) == output + assert isinstance(result["simulation_spec_json"], str) + with orm_session_factory() as session: + simulation = session.get(Simulation, simulation_id) + assert simulation.output == output + assert isinstance(simulation.simulation_spec_json, dict) From 3749642f99a3bc9120fb8ffaf04c702ead563f01 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 02:49:07 +0300 Subject: [PATCH 55/89] refactor: remove v1 persistence compatibility layers --- changelog.d/3788.changed.md | 2 +- dashboard/app.py | 85 ++- policyengine_api/asgi.py | 4 +- policyengine_api/data/__init__.py | 19 +- policyengine_api/data/data.py | 310 ---------- policyengine_api/data/orm.py | 228 ++++--- policyengine_api/data/v1_daos.py | 561 ------------------ policyengine_api/endpoints/household.py | 6 +- policyengine_api/endpoints/policy.py | 21 +- policyengine_api/endpoints/simulation.py | 3 +- policyengine_api/routes/household_routes.py | 3 +- policyengine_api/routes/policy_routes.py | 3 +- policyengine_api/routes/simulation_routes.py | 3 +- .../routes/user_profile_routes.py | 3 +- policyengine_api/services/economy_service.py | 6 +- tests/fixtures/services/household_fixtures.py | 30 +- tests/fixtures/services/policy_service.py | 37 +- .../services/report_output_fixtures.py | 55 -- .../fixtures/services/simulation_fixtures.py | 52 -- .../services/tracer_fixture_service.py | 41 +- tests/fixtures/services/user_service.py | 28 +- tests/to_refactor/python/test_data.py | 187 ------ tests/to_refactor/python/test_policy.py | 33 +- .../python/test_user_profile_routes.py | 57 +- tests/unit/conftest.py | 168 ++---- tests/unit/data/sqlite_schema.py | 12 - tests/unit/data/test_ordinary_v1_daos.py | 59 -- tests/unit/data/test_orm_sessions.py | 51 -- .../unit/data/test_remote_database_config.py | 2 +- tests/unit/data/test_run_daos.py | 162 ----- tests/unit/data/test_run_schema.py | 19 +- tests/unit/data/test_sqlalchemy_v2.py | 400 +++---------- tests/unit/data/test_stage7_no_direct_sql.py | 13 +- tests/unit/data/test_v1_daos.py | 57 -- tests/unit/data/test_v1_unit_of_work.py | 64 -- tests/unit/endpoints/test_get_simulations.py | 68 +-- .../endpoints/test_set_user_policy_dataset.py | 140 +---- .../unit/endpoints/test_update_user_policy.py | 83 ++- .../services/test_stage7_dao_boundaries.py | 8 +- .../test_stage7_local_service_boundaries.py | 15 +- tests/unit/services/test_tracer_service.py | 8 +- 41 files changed, 580 insertions(+), 2526 deletions(-) delete mode 100644 policyengine_api/data/data.py delete mode 100644 policyengine_api/data/v1_daos.py delete mode 100644 tests/fixtures/services/report_output_fixtures.py delete mode 100644 tests/fixtures/services/simulation_fixtures.py delete mode 100644 tests/to_refactor/python/test_data.py delete mode 100644 tests/unit/data/sqlite_schema.py delete mode 100644 tests/unit/data/test_ordinary_v1_daos.py delete mode 100644 tests/unit/data/test_run_daos.py delete mode 100644 tests/unit/data/test_v1_daos.py delete mode 100644 tests/unit/data/test_v1_unit_of_work.py diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index 228e9c400..48679b963 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to SQLAlchemy 2 repositories and Alembic while preserving the existing database schema and public API contracts. +Migrate API v1 persistence to caller-owned SQLAlchemy 2 sessions, mapped models, and Alembic while preserving the existing database schema and public API contracts. diff --git a/dashboard/app.py b/dashboard/app.py index cdfcf33e2..d4e0414ef 100644 --- a/dashboard/app.py +++ b/dashboard/app.py @@ -1,58 +1,79 @@ -from policyengine_api.data.data import database import streamlit as st +from sqlalchemy import select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy + st.title("PolicyEngine API dashboard") -# Add a text box that the user can enter a SQL query into, with a submit button and a table showing the results. -st.subheader("Run a SQL query") +def serialize_policy(policy: Policy) -> dict: + return { + column.name: getattr(policy, column.name) for column in Policy.__table__.columns + } -query = st.text_area("Enter a SQL query", "SELECT * FROM policy LIMIT 10;") -if st.button("Submit"): - try: - results = database.query(query) - st.table(results.fetchall()) - except Exception as e: - st.error(e) -# Enable the user to look up a policy by ID. +st.subheader("Recent policies") +if st.button("Refresh policies"): + sessions = get_v1_session_factory() + with sessions() as session: + policies = session.scalars(select(Policy).limit(10)).all() + st.table([serialize_policy(policy) for policy in policies]) -st.subheader("Look up a policy") +st.subheader("Look up a policy") policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_lookup_text")) country_id = st.text_input("Enter a country ID", "uk", key="policy_lookup_country") if st.button("Look up policy", key="policy_lookup"): - try: - results = database.query( - f"SELECT * FROM policy WHERE id IS '{policy_id}' AND country_id IS '{country_id}' LIMIT 10;" + sessions = get_v1_session_factory() + with sessions() as session: + policy = session.scalar( + select(Policy).where( + Policy.id == policy_id, + Policy.country_id == country_id, + ) ) - st.table(results.fetchall()) - except Exception as e: - st.error(e) + if policy is None: + st.error("Policy not found") + else: + st.table([serialize_policy(policy)]) -# Enable the user to set the label of a policy. st.subheader("Set a policy's label") - policy_id = int(st.text_input("Enter a policy ID", "1")) country_id = st.text_input("Enter a country ID", "uk") new_label = st.text_input("Enter a new label", "New label", key="policy_label_text") if st.button("Set policy label", key="policy_label"): - try: - database.set_policy_label(policy_id, country_id, new_label) - st.success("Success!") - except Exception as e: - st.error(e) + sessions = get_v1_session_factory() + with sessions.begin() as session: + policy = session.scalar( + select(Policy).where( + Policy.id == policy_id, + Policy.country_id == country_id, + ) + ) + if policy is None: + st.error("Policy not found") + else: + policy.label = new_label + st.success("Success!") -# Enable the user to delete a policy. st.subheader("Delete a policy") - policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_delete_text")) country_id = st.text_input("Enter a country ID", "uk", key="policy_delete_country") if st.button("Delete policy", key="policy_delete"): - try: - database.delete_policy(policy_id, country_id) - st.success("Success!") - except Exception as e: - st.error(e) + sessions = get_v1_session_factory() + with sessions.begin() as session: + policy = session.scalar( + select(Policy).where( + Policy.id == policy_id, + Policy.country_id == country_id, + ) + ) + if policy is None: + st.error("Policy not found") + else: + session.delete(policy) + st.success("Success!") diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index 4ee8369eb..8843ad65d 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -6,13 +6,13 @@ from policyengine_api.api import app as flask_app from policyengine_api.asgi_factory import create_asgi_app -from policyengine_api.data.data import close_runtime_databases +from policyengine_api.data.orm import close_v1_engines from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.warmup import run_startup_warmup app = application = create_asgi_app( flask_app, - shutdown_callback=close_runtime_databases, + shutdown_callback=close_v1_engines, ) # Warm the simulation machinery before serving (see policyengine_api.warmup). diff --git a/policyengine_api/data/__init__.py b/policyengine_api/data/__init__.py index eb42b9919..b412d8246 100644 --- a/policyengine_api/data/__init__.py +++ b/policyengine_api/data/__init__.py @@ -1,18 +1 @@ -"""Database package with lazy legacy exports. - -Keeping package import side-effect free lets Alembic load model metadata without -opening Cloud SQL or creating a local database. -""" - -from typing import Any - - -__all__ = ["PolicyEngineDatabase", "database", "local_database"] - - -def __getattr__(name: str) -> Any: - if name in __all__: - from . import data - - return getattr(data, name) - raise AttributeError(name) +"""SQLAlchemy persistence models, engine configuration, and migrations.""" diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py deleted file mode 100644 index 8edc7049a..000000000 --- a/policyengine_api/data/data.py +++ /dev/null @@ -1,310 +0,0 @@ -import atexit -import fcntl -import sqlite3 -from policyengine_api.constants import REPO, COUNTRY_PACKAGE_VERSIONS -from policyengine_api.utils import hash_object -from pathlib import Path -from dotenv import load_dotenv -import json -from google.cloud.sql.connector import Connector, IPTypes -import sqlalchemy -import os -import sys - -load_dotenv() - -DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME = ( - "policyengine-api:us-central1:policyengine-api-data" -) -DEFAULT_REMOTE_DB_USER = "policyengine" -DEFAULT_REMOTE_DB_NAME = "policyengine" -CLOUD_SQL_IP_TYPE = IPTypes.PUBLIC -DATABASE_POOL_RECYCLE_SECONDS = 1800 -DATABASE_POOL_SIZE = 5 -DATABASE_POOL_MAX_OVERFLOW = 2 -DATABASE_POOL_TIMEOUT_SECONDS = 30 - - -def get_remote_database_config() -> dict[str, str]: - return { - "instance_connection_name": os.environ.get( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", - DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME, - ), - "db_user": os.environ.get("POLICYENGINE_DB_USER", DEFAULT_REMOTE_DB_USER), - "db_name": os.environ.get("POLICYENGINE_DB_NAME", DEFAULT_REMOTE_DB_NAME), - } - - -class _ResultProxy: - """Lightweight wrapper that eagerly fetches results from a - SQLAlchemy CursorResult so they survive connection closure. - Provides fetchone()/fetchall() with dict-like row access.""" - - def __init__(self, cursor_result): - try: - # Use .mappings() so rows behave like dicts - self._rows = list(cursor_result.mappings()) - except Exception: - # For non-SELECT statements (INSERT/UPDATE/DELETE) - # there are no rows to fetch - self._rows = [] - self._index = 0 - - def fetchone(self): - if self._index < len(self._rows): - row = self._rows[self._index] - self._index += 1 - return row - return None - - def fetchall(self): - remaining = self._rows[self._index :] - self._index = len(self._rows) - return remaining - - -class _TransactionProxy: - """Execute queries against an existing connection inside a transaction.""" - - def __init__(self, connection, local: bool): - self._connection = connection - self._local = local - - def query(self, *query): - if self._local: - cursor = self._connection.cursor() - return cursor.execute(*query) - - query = list(query) - main_query = query[0].replace("?", "%s") - query[0] = main_query - params = query[1] if len(query) > 1 else None - if params is not None: - result = self._connection.exec_driver_sql(main_query, params) - else: - result = self._connection.exec_driver_sql(main_query) - return _ResultProxy(result) - - -class PolicyEngineDatabase: - """ - A wrapper around the database connection. - - It uses the Python package sqlite3. - """ - - household_cache: dict = {} - - @staticmethod - def _dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - - def __init__( - self, - local: bool = False, - initialize: bool = False, - ): - self.local = local - self._closed = False - if local: - # Local development uses a sqlite database. - self.db_url = REPO / "policyengine_api" / "data" / "policyengine.db" - # Serialize the exists-check + initialize under an exclusive file - # lock: with multiple gunicorn workers importing concurrently on a - # fresh instance, both can otherwise pass the exists() check and - # race initialize() (seed INSERTs collide -> worker dies at boot), - # or one can observe a created-but-unseeded file and skip - # initialization entirely. - lock_path = str(self.db_url) + ".init.lock" - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file, fcntl.LOCK_EX) - try: - if initialize or not Path(self.db_url).exists(): - self.initialize() - finally: - fcntl.flock(lock_file, fcntl.LOCK_UN) - else: - self._create_pool() - if initialize: - self.initialize() - - def _create_pool(self): - db_config = get_remote_database_config() - self.connector = Connector( - ip_type=CLOUD_SQL_IP_TYPE, - refresh_strategy="LAZY", - ) - db_pass = os.environ["POLICYENGINE_DB_PASSWORD"] - if db_pass == ".dbpw": - with open(".dbpw") as f: - db_pass = f.read().strip() - - def getconn(): - return self.connector.connect( - instance_connection_string=db_config["instance_connection_name"], - driver="pymysql", - db=db_config["db_name"], - user=db_config["db_user"], - password=db_pass, - ) - - self.pool = sqlalchemy.create_engine( - "mysql+pymysql://", - creator=getconn, - pool_pre_ping=True, - pool_recycle=DATABASE_POOL_RECYCLE_SECONDS, - pool_size=DATABASE_POOL_SIZE, - max_overflow=DATABASE_POOL_MAX_OVERFLOW, - pool_timeout=DATABASE_POOL_TIMEOUT_SECONDS, - ) - - def close(self) -> None: - """Release process-owned database and connector resources.""" - - if getattr(self, "_closed", False): - return - if self.local: - connection = getattr(self, "_connection", None) - if connection is not None: - connection.close() - self._closed = True - return - - try: - self.pool.dispose() - finally: - self.connector.close() - self._closed = True - - def _close_pool(self): - """Backward-compatible alias for callers predating ``close``.""" - - self.close() - - def _execute_remote(self, query_args): - """Execute a query against the remote database using - SQLAlchemy v2 connection-based execution.""" - main_query = query_args[0] - params = query_args[1] if len(query_args) > 1 else None - with self.pool.begin() as conn: - if params is not None: - result = conn.exec_driver_sql(main_query, params) - else: - result = conn.exec_driver_sql(main_query) - # Return a lightweight wrapper that holds - # the fetched results so they survive the - # connection context closing - return _ResultProxy(result) - - def _execute_remote_transaction(self, callback): - with self.pool.begin() as conn: - proxy = _TransactionProxy(conn, local=False) - return callback(proxy) - - def query(self, *query): - if self.local: - with sqlite3.connect(self.db_url) as conn: - conn.row_factory = self._dict_factory - cursor = conn.cursor() - return cursor.execute(*query) - else: - query = list(query) - main_query = query[0] - main_query = main_query.replace("?", "%s") - query[0] = main_query - return self._execute_remote(query) - - def transaction(self, callback): - if self.local: - connection = getattr(self, "_connection", None) - owns_connection = connection is None - if owns_connection: - connection = sqlite3.connect(self.db_url) - connection.row_factory = self._dict_factory - try: - connection.execute("BEGIN IMMEDIATE") - proxy = _TransactionProxy(connection, local=True) - result = callback(proxy) - connection.commit() - return result - except Exception: - connection.rollback() - raise - finally: - if owns_connection: - connection.close() - - return self._execute_remote_transaction(callback) - - def initialize(self): - """ - Create the database tables. - """ - if self.local: - # If the db_url exists, delete it. - if Path(self.db_url).exists(): - Path(self.db_url).unlink() - # If the db_url doesn't exist, create it. - if not Path(self.db_url).exists(): - Path(self.db_url).touch() - - with open( - REPO - / "policyengine_api" - / "data" - / f"initialise{'_local' if self.local else ''}.sql", - "r", - ) as f: - full_query = f.read() - # Split the query into individual queries. - queries = full_query.split(";") - for query in queries: - print(query, sys.stdout) - # Execute each query. - self.query(query) - - # Insert the UK, US and Canadian 'current law' policies. e.g. the UK policy table must have a row with id=1, country_id="uk", label="Current law", api_version=COUNTRY_PACKAGE_VERSIONS["uk"], policy_json="{}", policy_hash=hash_object({}) - for country_id, policy_id in zip( - COUNTRY_PACKAGE_VERSIONS.keys(), - range(1, 1 + len(COUNTRY_PACKAGE_VERSIONS)), - ): - self.query( - "INSERT INTO policy (id, country_id, label, api_version, policy_json, policy_hash) VALUES (?, ?, ?, ?, ?, ?)", - ( - policy_id, - country_id, - "Current law", - COUNTRY_PACKAGE_VERSIONS[country_id], - json.dumps({}), - hash_object({}), - ), - ) - - -# Determine if app is in debug mode, and if so, do not attempt connection with remote db -if os.environ.get("FLASK_DEBUG") == "1": - database = PolicyEngineDatabase(local=True, initialize=False) -else: - database = PolicyEngineDatabase(local=False, initialize=False) - -# TODO: Remove this eager SQLite-backed local database initialization and -# replace it with a traditional cache so importing the application does not -# create files or perform schema setup. -local_database = PolicyEngineDatabase(local=True, initialize=False) - - -def close_runtime_databases() -> None: - """Close database resources owned by the current application process.""" - - try: - database.close() - finally: - if local_database is not database: - local_database.close() - - -atexit.register(close_runtime_databases) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 8b9c0042c..85aa57bf1 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -1,125 +1,156 @@ -"""SQLAlchemy session ownership for the v1 persistence layer.""" +"""Canonical SQLAlchemy engine and Session configuration for API v1.""" from __future__ import annotations -from collections.abc import Callable, Iterator -from contextlib import contextmanager +import atexit +import fcntl +import json +import os +import sqlite3 from pathlib import Path -from typing import TypeVar -from sqlalchemy import Engine, create_engine, event +from dotenv import load_dotenv +from google.cloud.sql.connector import Connector, IPTypes +from sqlalchemy import Engine, create_engine from sqlalchemy.orm import Session, sessionmaker -from sqlalchemy.pool import StaticPool +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, REPO +from policyengine_api.utils import hash_object -T = TypeVar("T") +load_dotenv() -_v1_session_factories: dict[bool, sessionmaker[Session]] = {} +DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME = ( + "policyengine-api:us-central1:policyengine-api-data" +) +DEFAULT_REMOTE_DB_USER = "policyengine" +DEFAULT_REMOTE_DB_NAME = "policyengine" +CLOUD_SQL_IP_TYPE = IPTypes.PUBLIC +DATABASE_POOL_RECYCLE_SECONDS = 1800 +DATABASE_POOL_SIZE = 5 +DATABASE_POOL_MAX_OVERFLOW = 2 +DATABASE_POOL_TIMEOUT_SECONDS = 30 +LOCAL_DATABASE_PATH = REPO / "policyengine_api" / "data" / "policyengine.db" +_v1_engines: dict[bool, Engine] = {} +_v1_session_factories: dict[bool, sessionmaker[Session]] = {} +_cloud_sql_connectors: dict[bool, Connector] = {} -class _IndexedMappingRow(dict): - """SQLite row compatible with both SQLAlchemy and legacy mapping callers.""" - def __init__(self, cursor, values): - self._values = values - super().__init__( - (description[0], values[index]) - for index, description in enumerate(cursor.description) - ) +def get_remote_database_config() -> dict[str, str]: + return { + "instance_connection_name": os.environ.get( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", + DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME, + ), + "db_user": os.environ.get("POLICYENGINE_DB_USER", DEFAULT_REMOTE_DB_USER), + "db_name": os.environ.get("POLICYENGINE_DB_NAME", DEFAULT_REMOTE_DB_NAME), + } - def __getitem__(self, key): - if isinstance(key, int): - return self._values[key] - return super().__getitem__(key) - - def __iter__(self): - return iter(self._values) +def build_session_factory(engine: Engine) -> sessionmaker[Session]: + """Return SQLAlchemy's standard configurable Session factory.""" + + return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) + + +def _initialize_local_database(database_path: Path) -> None: + initialization_sql = ( + REPO / "policyengine_api" / "data" / "initialise_local.sql" + ).read_text(encoding="utf-8") + with sqlite3.connect(database_path) as connection: + connection.executescript(initialization_sql) + connection.executemany( + """ + INSERT INTO policy + (id, country_id, label, api_version, policy_json, policy_hash) + VALUES (?, ?, ?, ?, ?, ?) + """, + [ + ( + policy_id, + country_id, + "Current law", + COUNTRY_PACKAGE_VERSIONS[country_id], + json.dumps({}), + hash_object({}), + ) + for policy_id, country_id in enumerate( + COUNTRY_PACKAGE_VERSIONS, start=1 + ) + ], + ) -class SessionManager: - """Own sessions and transaction boundaries without leaking either to callers.""" - def __init__(self, engine: Engine): - self.engine = engine - self.session_factory = build_session_factory(engine) +# TODO: Remove this local-database initialization pattern and replace the local +# persistence path with a traditional cache. Application imports should +# eventually neither create a database file nor bootstrap a schema. +def _ensure_local_database(database_path: Path = LOCAL_DATABASE_PATH) -> None: + lock_path = Path(f"{database_path}.init.lock") + with lock_path.open("w") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + if not database_path.exists(): + _initialize_local_database(database_path) + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) - @contextmanager - def session(self) -> Iterator[Session]: - with self.session_factory() as session: - yield session - @contextmanager - def transaction(self) -> Iterator[Session]: - with self.session_factory.begin() as session: - yield session +def _build_local_engine() -> Engine: + _ensure_local_database() + return create_engine(f"sqlite+pysqlite:///{LOCAL_DATABASE_PATH}") - def run_in_transaction(self, callback: Callable[[Session], T]) -> T: - with self.session_factory.begin() as session: - return callback(session) +def _database_password() -> str: + password = os.environ["POLICYENGINE_DB_PASSWORD"] + if password == ".dbpw": + return Path(".dbpw").read_text(encoding="utf-8").strip() + return password -def build_session_factory(engine: Engine) -> sessionmaker[Session]: - """Return the canonical SQLAlchemy session factory for an engine.""" - return sessionmaker( - bind=engine, - class_=Session, - expire_on_commit=False, +def _build_remote_engine() -> Engine: + config = get_remote_database_config() + connector = Connector( + ip_type=CLOUD_SQL_IP_TYPE, + refresh_strategy="LAZY", ) - - -def build_sqlite_session_manager( - database_path: str | Path | None = None, -) -> SessionManager: - """Build a SQLite manager for local execution or isolated tests.""" - - if database_path is None: - engine = create_engine( - "sqlite+pysqlite:///:memory:", - connect_args={"check_same_thread": False}, - poolclass=StaticPool, + password = _database_password() + + def get_connection(): + return connector.connect( + instance_connection_string=config["instance_connection_name"], + driver="pymysql", + db=config["db_name"], + user=config["db_user"], + password=password, ) - else: - engine = create_engine(f"sqlite+pysqlite:///{Path(database_path)}") - return SessionManager(engine) - -def build_v1_session_manager(*, local: bool = False) -> SessionManager: - """Temporary bridge for callers not yet migrated to ``sessionmaker``.""" - - return SessionManager(get_v1_engine(local=local)) + engine = create_engine( + "mysql+pymysql://", + creator=get_connection, + pool_pre_ping=True, + pool_recycle=DATABASE_POOL_RECYCLE_SECONDS, + pool_size=DATABASE_POOL_SIZE, + max_overflow=DATABASE_POOL_MAX_OVERFLOW, + pool_timeout=DATABASE_POOL_TIMEOUT_SECONDS, + ) + _cloud_sql_connectors[False] = connector + return engine def get_v1_engine(*, local: bool = False) -> Engine: - """Return the process-owned engine selected by the v1 runtime.""" - - from policyengine_api.data.data import database, local_database - - selected_database = local_database if local else database - - if selected_database.local: - if hasattr(selected_database, "_connection"): - selected_database._connection.row_factory = _IndexedMappingRow - engine = create_engine( - "sqlite+pysqlite://", - creator=lambda: selected_database._connection, - poolclass=StaticPool, - ) - event.listen( - engine.pool, - "checkout", - lambda connection, *_: setattr( - connection, "row_factory", _IndexedMappingRow - ), - ) - return engine - return create_engine(f"sqlite+pysqlite:///{Path(selected_database.db_url)}") - return selected_database.pool + """Return one process-owned SQLAlchemy Engine for the selected runtime.""" + + use_local = local or os.environ.get("FLASK_DEBUG") == "1" + if use_local not in _v1_engines: + _v1_engines[use_local] = ( + _build_local_engine() if use_local else _build_remote_engine() + ) + return _v1_engines[use_local] def get_v1_session_factory(*, local: bool = False) -> sessionmaker[Session]: - """Return one configured factory per process-owned v1 engine.""" + """Return one configured Session factory per process-owned v1 Engine.""" if local not in _v1_session_factories: _v1_session_factories[local] = build_session_factory(get_v1_engine(local=local)) @@ -127,6 +158,21 @@ def get_v1_session_factory(*, local: bool = False) -> sessionmaker[Session]: def clear_v1_session_factories() -> None: - """Forget cached factories after their process-owned engines are closed.""" + """Forget cached factories, primarily after replacing an Engine in tests.""" _v1_session_factories.clear() + + +def close_v1_engines() -> None: + """Release process-owned SQLAlchemy pools and Cloud SQL connectors.""" + + clear_v1_session_factories() + for engine in _v1_engines.values(): + engine.dispose() + _v1_engines.clear() + for connector in _cloud_sql_connectors.values(): + connector.close() + _cloud_sql_connectors.clear() + + +atexit.register(close_v1_engines) diff --git a/policyengine_api/data/v1_daos.py b/policyengine_api/data/v1_daos.py deleted file mode 100644 index 856461214..000000000 --- a/policyengine_api/data/v1_daos.py +++ /dev/null @@ -1,561 +0,0 @@ -"""ORM data access objects for the existing v1 schema.""" - -from __future__ import annotations - -from collections.abc import Iterator -from contextlib import contextmanager -from typing import Any - -import uuid - -from sqlalchemy import func, select -from sqlalchemy.orm import Session - -from policyengine_api.data.orm import SessionManager -from policyengine_api.data.v1_models import ( - ComputedHousehold, - Household, - LegacyReportOutputAlias, - Policy, - ReportOutput, - ReportOutputRun, - Simulation, - SimulationRun, -) - - -def _mapping(model: Any) -> dict[str, Any]: - return { - column.name: getattr(model, column.name) for column in model.__table__.columns - } - - -class PolicyDAO: - def __init__(self, session: Session): - self.session = session - - def get(self, country_id: str, policy_id: int) -> dict[str, Any] | None: - model = self.session.scalar( - select(Policy).where( - Policy.country_id == country_id, - Policy.id == policy_id, - ) - ) - return _mapping(model) if model else None - - def find_unique( - self, country_id: str, policy_hash: str, label: str | None - ) -> dict[str, Any] | None: - model = self.session.scalar( - select(Policy).where( - Policy.country_id == country_id, - Policy.policy_hash == policy_hash, - Policy.label == label, - ) - ) - return _mapping(model) if model else None - - def search(self, country_id: str, query: str) -> list[dict[str, Any]]: - models = self.session.scalars( - select(Policy).where( - Policy.country_id == country_id, - Policy.label.contains(query, autoescape=True), - ) - ) - return [_mapping(model) for model in models] - - def create( - self, - country_id: str, - label: str | None, - policy_json: Any, - policy_hash: str, - api_version: str, - ) -> int: - policy = Policy( - country_id=country_id, - label=label, - api_version=api_version, - policy_json=policy_json, - policy_hash=policy_hash, - ) - self.session.add(policy) - self.session.flush() - return policy.id - - -class HouseholdDAO: - def __init__(self, session: Session): - self.session = session - - def get(self, country_id: str, household_id: int) -> dict[str, Any] | None: - model = self.session.scalar( - select(Household).where( - Household.country_id == country_id, - Household.id == household_id, - ) - ) - return _mapping(model) if model else None - - def create( - self, - country_id: str, - label: str | None, - household_json: Any, - household_hash: str, - api_version: str, - ) -> int: - model = Household( - country_id=country_id, - label=label, - api_version=api_version, - household_json=household_json, - household_hash=household_hash, - ) - self.session.add(model) - self.session.flush() - return model.id - - def update( - self, - country_id: str, - household_id: int, - label: str | None, - household_json: Any, - household_hash: str, - api_version: str, - ) -> bool: - model = self.session.scalar( - select(Household).where( - Household.country_id == country_id, - Household.id == household_id, - ) - ) - if model is None: - return False - model.label = label - model.household_json = household_json - model.household_hash = household_hash - model.api_version = api_version - return True - - -class ComputedHouseholdDAO: - def __init__(self, session: Session): - self.session = session - - def create(self, **values: Any) -> None: - self.session.add(ComputedHousehold(**values)) - - def upsert(self, **values: Any) -> None: - identity = ( - values["household_id"], - values["policy_id"], - values["country_id"], - ) - model = self.session.get(ComputedHousehold, identity) - if model is None: - self.session.add(ComputedHousehold(**values)) - return - for key, value in values.items(): - setattr(model, key, value) - - def get( - self, - household_id: int, - policy_id: int, - country_id: str, - *, - api_version: str | None = None, - ) -> dict[str, Any] | None: - statement = select(ComputedHousehold).where( - ComputedHousehold.household_id == household_id, - ComputedHousehold.policy_id == policy_id, - ComputedHousehold.country_id == country_id, - ) - if api_version is not None: - statement = statement.where(ComputedHousehold.api_version == api_version) - model = self.session.scalar(statement) - return _mapping(model) if model else None - - -class V1DAOs: - """DAOs bound to the same operation-scoped Session.""" - - def __init__(self, session: Session): - self.session = session - self.policies = PolicyDAO(session) - self.households = HouseholdDAO(session) - self.computed_households = ComputedHouseholdDAO(session) - self.simulations = SimulationDAO(session) - self.reports = ReportDAO(session) - - -class V1UnitOfWork: - """Create one Session and transaction boundary per logical operation.""" - - def __init__(self, sessions: SessionManager): - self.sessions = sessions - - @contextmanager - def read(self) -> Iterator[V1DAOs]: - with self.sessions.session() as session: - yield V1DAOs(session) - - @contextmanager - def transaction(self) -> Iterator[V1DAOs]: - with self.sessions.transaction() as session: - yield V1DAOs(session) - - -_runtime_unit_of_work: dict[bool, V1UnitOfWork] = {} - - -def runtime_v1_unit_of_work(*, local: bool = False) -> V1UnitOfWork: - """Return the process-local unit of work for the selected v1 database.""" - - if local not in _runtime_unit_of_work: - from policyengine_api.data.orm import build_v1_session_manager - - _runtime_unit_of_work[local] = V1UnitOfWork( - build_v1_session_manager(local=local) - ) - return _runtime_unit_of_work[local] - - -class SimulationDAO: - def __init__(self, session: Session): - self.session = session - - def get( - self, simulation_id: int, country_id: str | None = None - ) -> dict[str, Any] | None: - statement = select(Simulation).where(Simulation.id == simulation_id) - if country_id is not None: - statement = statement.where(Simulation.country_id == country_id) - model = self.session.scalar(statement) - return _mapping(model) if model else None - - @staticmethod - def get_in_session( - session, simulation_id: int, country_id: str | None = None - ) -> dict[str, Any] | None: - statement = select(Simulation).where(Simulation.id == simulation_id) - if country_id is not None: - statement = statement.where(Simulation.country_id == country_id) - model = session.scalar(statement) - return _mapping(model) if model else None - - def create(self, **values: Any) -> int: - model = Simulation(**values) - self.session.add(model) - self.session.flush() - return model.id - - def find_latest(self, **filters: Any) -> dict[str, Any] | None: - model = self.session.scalar( - select(Simulation) - .where( - *(getattr(Simulation, key) == value for key, value in filters.items()) - ) - .order_by(Simulation.id.desc()) - ) - return _mapping(model) if model else None - - @staticmethod - def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: - return next((run.id for run in runs if run.status == "complete"), None) - - def ensure_dual_write_state_in_session( - self, - session, - simulation_id: int, - country_id: str | None = None, - ) -> dict[str, Any]: - statement = ( - select(Simulation).where(Simulation.id == simulation_id).with_for_update() - ) - if country_id is not None: - statement = statement.where(Simulation.country_id == country_id) - simulation = session.scalar(statement) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - spec = { - "country_id": simulation.country_id, - "population_id": simulation.population_id, - "population_type": simulation.population_type, - "policy_id": simulation.policy_id, - } - simulation.simulation_spec_json = spec - simulation.simulation_spec_schema_version = 1 - runs = list( - session.scalars( - select(SimulationRun) - .where(SimulationRun.simulation_id == simulation_id) - .order_by(SimulationRun.run_sequence.desc()) - ) - ) - if not runs: - run = SimulationRun( - id=str(uuid.uuid4()), - simulation_id=simulation_id, - run_sequence=1, - status=simulation.status, - output=simulation.output, - error_message=simulation.error_message, - trigger_type="initial", - simulation_spec_snapshot_json=spec, - country_package_version=simulation.api_version, - ) - session.add(run) - session.flush() - runs = [run] - else: - mutable = next( - (run for run in runs if run.id == simulation.active_run_id), - runs[0], - ) - mutable.status = simulation.status - mutable.output = simulation.output - mutable.error_message = simulation.error_message - mutable.simulation_spec_snapshot_json = spec - mutable.country_package_version = simulation.api_version - - latest_successful = self._latest_successful_run_id(runs) - if simulation.status in {"pending", "running"}: - simulation.active_run_id = runs[0].id - else: - simulation.active_run_id = None - if simulation.status == "complete" and latest_successful is None: - latest_successful = runs[0].id - simulation.latest_successful_run_id = latest_successful - session.flush() - return _mapping(simulation) - - def ensure_dual_write_state( - self, simulation_id: int, country_id: str | None = None - ) -> dict[str, Any]: - return self.ensure_dual_write_state_in_session( - self.session, simulation_id, country_id - ) - - def create_or_get_with_sync( - self, - *, - sync_callback, - **values: Any, - ) -> dict[str, Any]: - filters = { - key: values[key] - for key in ( - "country_id", - "population_id", - "population_type", - "policy_id", - ) - } - model = self.session.scalar( - select(Simulation) - .where( - *(getattr(Simulation, key) == value for key, value in filters.items()) - ) - .order_by(Simulation.id.desc()) - .with_for_update() - ) - if model is None: - model = Simulation(**values) - self.session.add(model) - self.session.flush() - return sync_callback( - self.session, - model.id, - country_id=model.country_id, - ) - - def update_with_sync( - self, - simulation_id: int, - country_id: str, - values: dict[str, Any], - sync_callback, - ) -> dict[str, Any]: - model = self.session.scalar( - select(Simulation) - .where( - Simulation.id == simulation_id, - Simulation.country_id == country_id, - ) - .with_for_update() - ) - if model is None: - raise ValueError(f"Simulation #{simulation_id} not found") - for key, value in values.items(): - setattr(model, key, value) - self.session.flush() - return sync_callback( - self.session, - simulation_id, - country_id=country_id, - ) - - def update(self, simulation_id: int, **values: Any) -> bool: - model = self.session.get(Simulation, simulation_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - def create_run( - self, simulation_id: int, *, run_id: str, **values: Any - ) -> dict[str, Any]: - parent = self.session.scalar( - select(Simulation).where(Simulation.id == simulation_id).with_for_update() - ) - if parent is None: - raise LookupError(f"Simulation {simulation_id} does not exist") - sequence = ( - self.session.scalar( - select(func.max(SimulationRun.run_sequence)).where( - SimulationRun.simulation_id == simulation_id - ) - ) - or 0 - ) + 1 - model = SimulationRun( - id=run_id, - simulation_id=simulation_id, - run_sequence=sequence, - **values, - ) - self.session.add(model) - self.session.flush() - return _mapping(model) - - def get_run(self, run_id: str) -> dict[str, Any] | None: - model = self.session.get(SimulationRun, run_id) - return _mapping(model) if model else None - - def list_runs(self, simulation_id: int) -> list[dict[str, Any]]: - models = self.session.scalars( - select(SimulationRun) - .where(SimulationRun.simulation_id == simulation_id) - .order_by(SimulationRun.run_sequence.desc()) - ) - return [_mapping(model) for model in models] - - -class ReportDAO: - def __init__(self, session: Session): - self.session = session - - def get( - self, report_output_id: int, country_id: str | None = None - ) -> dict[str, Any] | None: - statement = select(ReportOutput).where(ReportOutput.id == report_output_id) - if country_id is not None: - statement = statement.where(ReportOutput.country_id == country_id) - model = self.session.scalar(statement) - return _mapping(model) if model else None - - def get_for_update( - self, report_output_id: int, country_id: str | None = None - ) -> dict[str, Any] | None: - statement = ( - select(ReportOutput) - .where(ReportOutput.id == report_output_id) - .with_for_update() - ) - if country_id is not None: - statement = statement.where(ReportOutput.country_id == country_id) - model = self.session.scalar(statement) - return _mapping(model) if model else None - - def find_latest(self, **filters: Any) -> dict[str, Any] | None: - model = self.session.scalar( - select(ReportOutput) - .where( - *(getattr(ReportOutput, key) == value for key, value in filters.items()) - ) - .order_by(ReportOutput.id.desc()) - ) - return _mapping(model) if model else None - - def create(self, **values: Any) -> int: - model = ReportOutput(**values) - self.session.add(model) - self.session.flush() - return model.id - - def update(self, report_output_id: int, **values: Any) -> bool: - model = self.session.get(ReportOutput, report_output_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - def create_run( - self, report_output_id: int, *, run_id: str, **values: Any - ) -> dict[str, Any]: - parent = self.session.scalar( - select(ReportOutput) - .where(ReportOutput.id == report_output_id) - .with_for_update() - ) - if parent is None: - raise LookupError(f"Report output {report_output_id} does not exist") - sequence = ( - self.session.scalar( - select(func.max(ReportOutputRun.run_sequence)).where( - ReportOutputRun.report_output_id == report_output_id - ) - ) - or 0 - ) + 1 - model = ReportOutputRun( - id=run_id, - report_output_id=report_output_id, - run_sequence=sequence, - **values, - ) - self.session.add(model) - self.session.flush() - return _mapping(model) - - def get_run(self, run_id: str) -> dict[str, Any] | None: - model = self.session.get(ReportOutputRun, run_id) - return _mapping(model) if model else None - - def update_run(self, run_id: str, **values: Any) -> bool: - model = self.session.get(ReportOutputRun, run_id) - if model is None: - return False - for key, value in values.items(): - setattr(model, key, value) - return True - - def list_runs(self, report_output_id: int) -> list[dict[str, Any]]: - models = self.session.scalars( - select(ReportOutputRun) - .where(ReportOutputRun.report_output_id == report_output_id) - .order_by(ReportOutputRun.run_sequence.desc()) - ) - return [_mapping(model) for model in models] - - def set_alias(self, legacy_id: int, canonical_id: int) -> None: - model = self.session.get(LegacyReportOutputAlias, legacy_id) - if model is None: - self.session.add( - LegacyReportOutputAlias( - legacy_report_output_id=legacy_id, - canonical_report_output_id=canonical_id, - ) - ) - else: - model.canonical_report_output_id = canonical_id - - def get_alias(self, legacy_id: int) -> dict[str, Any] | None: - model = self.session.get(LegacyReportOutputAlias, legacy_id) - return _mapping(model) if model else None diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index 7fb3677d1..78007242f 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -114,7 +114,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Look in computed_households to see if already computed - with get_v1_session_factory(local=True)() as session: + sessions = get_v1_session_factory(local=True) + with sessions() as session: computed_household = session.scalar( select(ComputedHousehold).where( ComputedHousehold.household_id == int(household_id), @@ -133,7 +134,8 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st # Retrieve from the household table - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: household = session.scalar( select(Household).where( Household.country_id == country_id, diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index 4124112af..275303592 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -54,7 +54,8 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: results = session.scalars( select(Policy).where( Policy.country_id == country_id, @@ -208,7 +209,8 @@ def get_user_policy(country_id: str, user_id: str) -> dict: """ # Get the policy record for a given policy ID. - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: user_policies = session.scalars( select(UserPolicy).where( UserPolicy.country_id == country_id, @@ -235,13 +237,9 @@ def get_user_policy(country_id: str, user_id: str) -> dict: ) -# Whitelist of columns that callers are allowed to modify via -# update_user_policy. Identity columns (id, country_id, user_id, -# reform_id, baseline_id) are intentionally excluded because they -# define the record; allowing clients to rewrite them would both -# break referential assumptions and let the column name be used -# as a SQL injection vector (keys are interpolated into the -# UPDATE statement below). +# Whitelist of attributes that callers may modify via update_user_policy. +# Identity attributes are intentionally excluded because they define the +# record and must not be reassigned through this endpoint. UPDATE_USER_POLICY_ALLOWED_FIELDS = frozenset( { "reform_label", @@ -275,9 +273,8 @@ def update_user_policy(country_id: str) -> dict: user_policy_id = payload.pop("id") - # Reject any unknown/unsafe keys. The keys end up interpolated - # into a SQL UPDATE statement, so we must validate them against - # a static whitelist instead of trusting the JSON payload. + # Reject unknown or identity attributes before applying payload values to + # the mapped entity. unknown_keys = [ key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS ] diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index d9b102656..08fa9c52d 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -45,7 +45,8 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - with get_v1_session_factory(local=True)() as session: + sessions = get_v1_session_factory(local=True) + with sessions() as session: result = session.scalars( select(ReformImpact) .order_by(ReformImpact.start_time.desc()) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 9c3b175b3..2f2f06426 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -37,7 +37,8 @@ def get_household(country_id: str, household_id: int) -> Response: """ print(f"Got request for household {household_id} in country {country_id}") - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: household = household_service.get_household(session, country_id, household_id) result = None if household is None else _serialize_household(household) if result is None: diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index c405245a5..2674d4b77 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -43,7 +43,8 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: # Specifically cast policy_id to an integer policy_id = int(policy_id) - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: policy = policy_service.get_policy(session, country_id, policy_id) result = None if policy is None else _serialize_policy(policy) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index bdf896d4f..436c55c00 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -159,7 +159,8 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: if simulation_id <= 0: raise BadRequest("simulation_id must be a positive integer") - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: simulation = simulation_service.get_simulation( session, country_id, diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index 55f6a437d..cb39b63c8 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -72,7 +72,8 @@ def get_user_profile(country_id: str) -> Response: if (auth0_id is None) and (user_id is None): raise BadRequest("auth0_id or user_id must be provided") - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: profile = ( user_service.get_profile(session, user_id=user_id) if auth0_id is None diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 091e0846e..d5ab0a467 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -243,7 +243,8 @@ def _get_policy_jsons( baseline_policy_id: int, reform_policy_id: int, ) -> tuple[dict | None, dict | None]: - with get_v1_session_factory()() as session: + sessions = get_v1_session_factory() + with sessions() as session: baseline = policy_service.get_policy_json( session, country_id, @@ -769,7 +770,8 @@ def _get_previous_impacts( """ previous_impacts: list[Any] = [] - with get_v1_session_factory(local=True)() as session: + sessions = get_v1_session_factory(local=True) + with sessions() as session: previous_impacts = ( reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix( session, diff --git a/tests/fixtures/services/household_fixtures.py b/tests/fixtures/services/household_fixtures.py index 063778c7a..d785cd07b 100644 --- a/tests/fixtures/services/household_fixtures.py +++ b/tests/fixtures/services/household_fixtures.py @@ -2,6 +2,8 @@ import json from unittest.mock import patch +from policyengine_api.data.v1_models import Household + valid_request_body = { "data": {"people": {"person1": {"age": 30, "income": 50000}}}, "label": "Test Household", @@ -28,22 +30,16 @@ def mock_hash_object(): @pytest.fixture -def existing_household_record(test_db): +def existing_household_record(orm_session): """Insert an existing household record into the database.""" - test_db.query( - "INSERT INTO household (id, country_id, household_json, household_hash, label, api_version) VALUES (?, ?, ?, ?, ?, ?)", - ( - valid_db_row["id"], - valid_db_row["country_id"], - valid_db_row["household_json"], - valid_db_row["household_hash"], - valid_db_row["label"], - valid_db_row["api_version"], - ), + household = Household( + id=valid_db_row["id"], + country_id=valid_db_row["country_id"], + household_json=json.loads(valid_db_row["household_json"]), + household_hash=valid_db_row["household_hash"], + label=valid_db_row["label"], + api_version=valid_db_row["api_version"], ) - - inserted_row = test_db.query( - "SELECT * FROM household WHERE id = ?", (valid_db_row["id"],) - ).fetchone() - - return inserted_row + orm_session.add(household) + orm_session.flush() + return household diff --git a/tests/fixtures/services/policy_service.py b/tests/fixtures/services/policy_service.py index 6c4a27f66..7cf164ba4 100644 --- a/tests/fixtures/services/policy_service.py +++ b/tests/fixtures/services/policy_service.py @@ -2,6 +2,8 @@ import json from unittest.mock import patch +from policyengine_api.data.v1_models import Policy + valid_policy_json = { "data": {"gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433}}, } @@ -19,13 +21,6 @@ } -@pytest.fixture -def mock_database(): - """Mock the database module.""" - with patch("policyengine_api.services.policy_service.database") as mock_db: - yield mock_db - - @pytest.fixture def mock_hash_object(): """Mock the hash_object function.""" @@ -35,22 +30,16 @@ def mock_hash_object(): @pytest.fixture -def existing_policy_record(test_db): +def existing_policy_record(orm_session): """Insert an existing policy record into the database.""" - test_db.query( - "INSERT INTO policy (id, country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?, ?)", - ( - valid_policy_data["id"], - valid_policy_data["country_id"], - valid_policy_data["policy_json"], - valid_policy_data["policy_hash"], - valid_policy_data["label"], - valid_policy_data["api_version"], - ), + policy = Policy( + id=valid_policy_data["id"], + country_id=valid_policy_data["country_id"], + policy_json=json.loads(valid_policy_data["policy_json"]), + policy_hash=valid_policy_data["policy_hash"], + label=valid_policy_data["label"], + api_version=valid_policy_data["api_version"], ) - - inserted_row = test_db.query( - "SELECT * FROM policy WHERE id = ?", (valid_policy_data["id"],) - ).fetchone() - - return inserted_row + orm_session.add(policy) + orm_session.flush() + return policy diff --git a/tests/fixtures/services/report_output_fixtures.py b/tests/fixtures/services/report_output_fixtures.py deleted file mode 100644 index 3d4ca5a14..000000000 --- a/tests/fixtures/services/report_output_fixtures.py +++ /dev/null @@ -1,55 +0,0 @@ -import pytest -import json - -from policyengine_api.constants import get_report_output_cache_version - -valid_report_data = { - "country_id": "us", - "simulation_1_id": 1, - "simulation_2_id": None, - "api_version": get_report_output_cache_version("us"), - "status": "pending", - "output": None, - "error_message": None, - "year": "2025", -} - -sample_report_output = { - "population_impact": { - "baseline": {"decile_1": 1000}, - "reform": {"decile_1": 1100}, - }, - "budget_impact": 5000000, -} - - -@pytest.fixture -def existing_report_record(test_db): - """Insert an existing report output record into the database.""" - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, api_version, status, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ( - valid_report_data["country_id"], - valid_report_data["simulation_1_id"], - valid_report_data["simulation_2_id"], - valid_report_data["api_version"], - valid_report_data["status"], - valid_report_data["year"], - ), - ) - - # Get the inserted record - inserted_row = test_db.query( - """SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND status = ? - ORDER BY id DESC LIMIT 1""", - ( - valid_report_data["country_id"], - valid_report_data["simulation_1_id"], - valid_report_data["status"], - ), - ).fetchone() - - return inserted_row diff --git a/tests/fixtures/services/simulation_fixtures.py b/tests/fixtures/services/simulation_fixtures.py deleted file mode 100644 index 672553636..000000000 --- a/tests/fixtures/services/simulation_fixtures.py +++ /dev/null @@ -1,52 +0,0 @@ -import pytest - -valid_simulation_data = { - "country_id": "us", - "api_version": "1.0.0", - "population_id": "household_test_123", - "population_type": "household", - "policy_id": 1, -} - -duplicate_simulation_data = { - "country_id": "us", - "api_version": "1.0.0", - "population_id": "household_test_123", - "population_type": "household", - "policy_id": 1, -} - - -@pytest.fixture -def existing_simulation_record(test_db): - """Insert an existing simulation record into the database.""" - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ( - valid_simulation_data["country_id"], - valid_simulation_data["api_version"], - valid_simulation_data["population_id"], - valid_simulation_data["population_type"], - valid_simulation_data["policy_id"], - "pending", - ), - ) - - # Get the inserted record - inserted_row = test_db.query( - """SELECT * FROM simulations - WHERE country_id = ? AND api_version = ? - AND population_id = ? AND population_type = ? - AND policy_id = ?""", - ( - valid_simulation_data["country_id"], - valid_simulation_data["api_version"], - valid_simulation_data["population_id"], - valid_simulation_data["population_type"], - valid_simulation_data["policy_id"], - ), - ).fetchone() - - return inserted_row diff --git a/tests/fixtures/services/tracer_fixture_service.py b/tests/fixtures/services/tracer_fixture_service.py index bc4742a25..f22bc6ea4 100644 --- a/tests/fixtures/services/tracer_fixture_service.py +++ b/tests/fixtures/services/tracer_fixture_service.py @@ -1,8 +1,6 @@ import pytest import json -from policyengine_api.services.tracer_analysis_service import ( - TracerAnalysisService, -) +from policyengine_api.data.v1_models import Tracer valid_tracer = { "tracer_output": [ @@ -25,31 +23,14 @@ @pytest.fixture -def test_tracer_data(test_db): - # Insert data using query() - test_db.query( - """ - INSERT INTO tracers (household_id, policy_id, country_id, api_version, tracer_output) - VALUES (?, ?, ?, ?, ?) - """, - ( - valid_tracer_row["household_id"], - valid_tracer_row["policy_id"], - valid_tracer_row["country_id"], - valid_tracer_row["api_version"], - valid_tracer_row["tracer_output"], - ), +def test_tracer_data(orm_session): + tracer = Tracer( + household_id=int(valid_tracer_row["household_id"]), + policy_id=int(valid_tracer_row["policy_id"]), + country_id=valid_tracer_row["country_id"], + api_version=valid_tracer_row["api_version"], + tracer_output=json.loads(valid_tracer_row["tracer_output"]), ) - - # Verify that the data has been inserted - inserted_row = test_db.query( - "SELECT * FROM tracers WHERE household_id = ? AND policy_id = ? AND country_id = ? AND api_version = ?", - ( - valid_tracer_row["household_id"], - valid_tracer_row["policy_id"], - valid_tracer_row["country_id"], - valid_tracer_row["api_version"], - ), - ).fetchone() - - return inserted_row + orm_session.add(tracer) + orm_session.flush() + return tracer diff --git a/tests/fixtures/services/user_service.py b/tests/fixtures/services/user_service.py index 1ab78bc80..164574dc8 100644 --- a/tests/fixtures/services/user_service.py +++ b/tests/fixtures/services/user_service.py @@ -1,5 +1,7 @@ import pytest +from policyengine_api.data.v1_models import UserProfile + valid_user_record = { "user_id": 1, "auth0_id": "123", @@ -10,21 +12,15 @@ @pytest.fixture -def existing_user_profile(test_db): +def existing_user_profile(orm_session): """Insert an existing user record into the database.""" - test_db.query( - "INSERT INTO user_profiles (user_id, auth0_id, username, primary_country, user_since) VALUES (?, ?, ?, ?, ?)", - ( - valid_user_record["user_id"], - valid_user_record["auth0_id"], - valid_user_record["username"], - valid_user_record["primary_country"], - valid_user_record["user_since"], - ), + profile = UserProfile( + user_id=valid_user_record["user_id"], + auth0_id=valid_user_record["auth0_id"], + username=valid_user_record["username"], + primary_country=valid_user_record["primary_country"], + user_since=valid_user_record["user_since"], ) - inserted_row = test_db.query( - "SELECT * FROM user_profiles WHERE auth0_id = ?", - (valid_user_record["auth0_id"],), - ).fetchone() - - return inserted_row + orm_session.add(profile) + orm_session.flush() + return profile diff --git a/tests/to_refactor/python/test_data.py b/tests/to_refactor/python/test_data.py deleted file mode 100644 index a3f8e7c32..000000000 --- a/tests/to_refactor/python/test_data.py +++ /dev/null @@ -1,187 +0,0 @@ -import pytest -import json - -from policyengine_api.data import PolicyEngineDatabase - - -# Test the query method using the db's policy table -class TestQuery: - # Set shared variables - country_id = "us" - placeholder = "placeholder" - first_updated_placeholder = "maxwell" - second_updated_placeholder = "dworkin" - - # Initialize db connection - db = PolicyEngineDatabase(local=True, initialize=True) - - # Test INSERT and SELECT statements - def test_insert(self): - self.db.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label = ? AND api_version = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - - # Test INSERT and SELECT statements with None - def test_insert_null(self): - self.db.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - self.country_id, - self.placeholder, - self.placeholder, - None, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label IS NULL AND api_version = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - - # Test UPDATE - def test_update(self): - self.db.query( - f"UPDATE policy SET policy_json = ? WHERE policy_json = ? AND label = ? AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label = ? AND api_version = ?", - ( - self.country_id, - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["policy_json"] == self.first_updated_placeholder - - # Test UPDATE with None as search param - def test_update_none_param(self): - self.db.query( - f"UPDATE policy SET policy_json = ? WHERE policy_json = ? AND label IS NULL AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.second_updated_placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label IS NULL AND api_version = ?", - ( - self.country_id, - self.second_updated_placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - assert str(row["policy_json"]) == self.second_updated_placeholder - - # Test UPDATE with None as set value - def test_update_set_none(self): - self.db.query( - f"UPDATE policy SET label = ? WHERE policy_json = ? AND api_version = ? AND label = ? AND country_id = ? AND policy_hash = ? ", - ( - None, - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND label IS NULL AND policy_hash = ? AND api_version = ? AND policy_json = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.first_updated_placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - - # Test DELETE - def test_delete(self): - # Clean up the first record that was added - self.db.query( - f"DELETE FROM policy WHERE policy_json = ? AND label IS NULL AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.second_updated_placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - # Delete the second for testing purposes - self.db.query( - f"DELETE FROM policy WHERE label IS NULL AND api_version = ? AND policy_json = ? AND country_id = ? AND policy_hash = ? ", - ( - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - # Confirm that it no longer exists - row = self.db.query( - f"SELECT * FROM policy WHERE label IS NULL AND country_id = ? AND api_version = ? AND policy_hash = ? AND policy_json = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is None diff --git a/tests/to_refactor/python/test_policy.py b/tests/to_refactor/python/test_policy.py index 0c66c09df..d5b57eb80 100644 --- a/tests/to_refactor/python/test_policy.py +++ b/tests/to_refactor/python/test_policy.py @@ -1,10 +1,9 @@ -import pytest import json -import time -import sqlite3 -from policyengine_api.data import database +from sqlalchemy import delete + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy from policyengine_api.utils import hash_object -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS class TestPolicyCreation: @@ -23,10 +22,14 @@ class TestPolicyCreation: """ def test_create_unique_policy(self, rest_client): - database.query( - f"DELETE FROM policy WHERE policy_hash = ? AND label = ? AND country_id = ?", - (self.policy_hash, self.label, self.country_id), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(Policy).where( + Policy.policy_hash == self.policy_hash, + Policy.label == self.label, + Policy.country_id == self.country_id, + ) + ) res = rest_client.post("/us/policy", json=self.test_policy) return_object = json.loads(res.text) @@ -41,10 +44,14 @@ def test_create_nonunique_policy(self, rest_client): assert return_object["status"] == "ok" assert res.status_code == 200 - database.query( - f"DELETE FROM policy WHERE policy_hash = ? AND label = ? AND country_id = ?", - (self.policy_hash, self.label, self.country_id), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(Policy).where( + Policy.policy_hash == self.policy_hash, + Policy.label == self.label, + Policy.country_id == self.country_id, + ) + ) def test_create_policy_invalid_country(self, rest_client): res = rest_client.post("/au/policy", json=self.test_policy) diff --git a/tests/to_refactor/python/test_user_profile_routes.py b/tests/to_refactor/python/test_user_profile_routes.py index ec30f9eef..3b2296b4f 100644 --- a/tests/to_refactor/python/test_user_profile_routes.py +++ b/tests/to_refactor/python/test_user_profile_routes.py @@ -1,8 +1,11 @@ import json -from datetime import datetime -from policyengine_api.data import database import time +from sqlalchemy import delete, select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserProfile + class TestUserProfiles: # Define the profile to test against @@ -22,13 +25,13 @@ class TestUserProfiles: """ def test_set_and_get_record(self, rest_client): - database.query( - f"DELETE FROM user_profiles WHERE auth0_id = ? AND primary_country = ?", - ( - self.auth0_id, - self.primary_country, - ), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(UserProfile).where( + UserProfile.auth0_id == self.auth0_id, + UserProfile.primary_country == self.primary_country, + ) + ) res = rest_client.post("/us/user-profile", json=self.test_profile) return_object = json.loads(res.text) @@ -43,7 +46,7 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert return_object["result"]["auth0_id"] == self.auth0_id assert return_object["result"]["primary_country"] == self.primary_country - assert return_object["result"]["username"] == None + assert return_object["result"]["username"] is None user_id = return_object["result"]["user_id"] @@ -54,7 +57,7 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert return_object["result"]["primary_country"] == self.primary_country assert return_object["result"].get("auth0_id") is None - assert return_object["result"]["username"] == None + assert return_object["result"]["username"] is None test_username = "maxwell" updated_profile = {"user_id": user_id, "username": test_username} @@ -65,11 +68,14 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert res.status_code == 200 - row = database.query( - f"SELECT * FROM user_profiles WHERE user_id = ? AND username = ?", - (user_id, test_username), - ).fetchone() - assert row is not None + with get_v1_session_factory()() as session: + row = session.scalar( + select(UserProfile).where( + UserProfile.user_id == user_id, + UserProfile.username == test_username, + ) + ) + assert row is not None malicious_updated_profile = {**updated_profile, "auth0_id": "BOGUS"} @@ -78,22 +84,15 @@ def test_set_and_get_record(self, rest_client): assert res.status_code == 200 - row = database.query( - f"SELECT * FROM user_profiles WHERE username = ?", - (test_username,), - ).fetchone() - - assert row["auth0_id"] == self.auth0_id - - database.query( - f"DELETE FROM user_profiles WHERE user_id = ? AND auth0_id = ? AND primary_country = ?", - (user_id, self.auth0_id, self.primary_country), - ) + with get_v1_session_factory().begin() as session: + row = session.scalar( + select(UserProfile).where(UserProfile.username == test_username) + ) + assert row.auth0_id == self.auth0_id + session.delete(row) def test_non_existent_record(self, rest_client): non_existent_auth0_id = "non-existent-auth0-id" res = rest_client.get(f"/us/user-profile?auth0_id={non_existent_auth0_id}") - return_object = json.loads(res.text) - assert res.status_code == 404 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 24733f3ab..1570e9bb1 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,158 +1,58 @@ import os -import sqlite3 import pytest +from sqlalchemy import create_engine +from sqlalchemy.pool import StaticPool + -# The legacy SQL module creates its default database object at import time. -# Keep unit tests on SQLite until the SQL layer is broadly refactored in a -# later migration stage. os.environ.setdefault("FLASK_DEBUG", "1") from policyengine_api.constants import REPO -from policyengine_api.data import PolicyEngineDatabase -from policyengine_api.data.orm import ( - clear_v1_session_factories, - get_v1_session_factory, -) - - -class TestPolicyEngineDatabase(PolicyEngineDatabase): - """Test version of PolicyEngineDatabase that uses in-memory SQLite""" - - def __init__(self, initialize: bool = True): - self.local = True # Always use SQLite for tests - if initialize: - self._setup_connection() - self.initialize() - - def _setup_connection(self): - """Setup the in-memory connection""" - if not hasattr(self, "_connection"): - self._connection = sqlite3.connect(":memory:") - - def dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - - self._connection.row_factory = dict_factory - - def initialize(self): - """ - Override initialize to avoid file operations from parent class - """ - self._setup_connection() - - # Read the SQL initialization file - init_file = ( - REPO - / "policyengine_api" - / "data" - / f"initialise{'_local' if self.local else ''}.sql" - ) - with open(init_file) as f: - full_query = f.read() - - # Split and execute the queries - queries = full_query.split(";") - for query in queries: - if query.strip(): # Skip empty queries - self.query(query) - - def query(self, *query): - """Override query method to use in-memory connection""" - if not hasattr(self, "_connection"): - # Create a persistent connection for the in-memory database - self._connection = sqlite3.connect(self.db_url) - - def dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - - self._connection.row_factory = dict_factory - - cursor = self._connection.cursor() - result = cursor.execute(*query) - self._connection.commit() - return result - - def clean(self): - """Clear all data from tables while preserving the schema""" - if hasattr(self, "_connection"): - cursor = self._connection.cursor() - - # Get all table names - tables = cursor.execute( - "SELECT name FROM sqlite_master WHERE type='table'" - ).fetchall() - - # Disable foreign key checks temporarily - cursor.execute("PRAGMA foreign_keys=OFF") - - # Delete all data from each table - for table in tables: - table_name = table["name"] - cursor.execute(f"DELETE FROM {table_name}") - - # Re-enable foreign key checks - cursor.execute("PRAGMA foreign_keys=ON") - - self._connection.commit() +from policyengine_api.data import orm +from policyengine_api.data.v1_models import V1Base @pytest.fixture(scope="session") -def test_db(): - """Create a test database instance that persists for the whole test session""" - db = TestPolicyEngineDatabase(initialize=True) - yield db - # Clean up the connection when done - if hasattr(db, "_connection"): - db._connection.close() +def test_engine(): + engine = create_engine( + "sqlite+pysqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + initialization_sql = ( + REPO / "policyengine_api" / "data" / "initialise_local.sql" + ).read_text(encoding="utf-8") + raw_connection = engine.raw_connection() + try: + raw_connection.executescript(initialization_sql) + finally: + raw_connection.close() + yield engine + engine.dispose() @pytest.fixture(autouse=True) -def override_database(test_db, monkeypatch): - """ - Global database override that affects all imports of the database. - This fixture automatically applies to all tests. - """ - test_db.clean() - - # Patch at the root module level where database is defined - import policyengine_api.data - - monkeypatch.setattr(policyengine_api.data, "database", test_db) - - # Also patch the module-level variable for any existing imports - import sys - - for module_name, module in list(sys.modules.items()): - if module_name.startswith("policyengine_api."): - if hasattr(module, "database"): - monkeypatch.setattr(module, "database", test_db) - if hasattr(module, "local_database"): - monkeypatch.setattr(module, "local_database", test_db) - - clear_v1_session_factories() +def isolated_orm_database(test_engine, monkeypatch): + """Bind runtime factories to a clean in-memory ORM database per test.""" + + monkeypatch.setattr(orm, "get_v1_engine", lambda *, local=False: test_engine) + orm.clear_v1_session_factories() + factory = orm.get_v1_session_factory() + with factory.begin() as session: + for table in reversed(V1Base.metadata.sorted_tables): + session.execute(table.delete()) try: - yield test_db + yield finally: - clear_v1_session_factories() + orm.clear_v1_session_factories() @pytest.fixture -def orm_session_factory(override_database): - """Return the runtime-style SQLAlchemy factory bound to the test schema.""" - - return get_v1_session_factory() +def orm_session_factory(isolated_orm_database): + return orm.get_v1_session_factory() @pytest.fixture def orm_session(orm_session_factory): - """Return one caller-owned SQLAlchemy Session for a unit test.""" - with orm_session_factory() as session: yield session diff --git a/tests/unit/data/sqlite_schema.py b/tests/unit/data/sqlite_schema.py deleted file mode 100644 index 9852a2280..000000000 --- a/tests/unit/data/sqlite_schema.py +++ /dev/null @@ -1,12 +0,0 @@ -from policyengine_api.constants import REPO -from policyengine_api.data.orm import SessionManager - - -def create_sqlite_v1_schema(manager: SessionManager) -> None: - """Install the explicit local schema without compiling MySQL metadata.""" - - schema = (REPO / "policyengine_api/data/initialise_local.sql").read_text( - encoding="utf-8" - ) - with manager.engine.connect() as connection: - connection.connection.driver_connection.executescript(schema) diff --git a/tests/unit/data/test_ordinary_v1_daos.py b/tests/unit/data/test_ordinary_v1_daos.py deleted file mode 100644 index 565f93997..000000000 --- a/tests/unit/data/test_ordinary_v1_daos.py +++ /dev/null @@ -1,59 +0,0 @@ -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work() -> V1UnitOfWork: - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) - - -def test_computed_household_upsert_preserves_one_cache_row(): - uow = _unit_of_work() - values = { - "household_id": 1, - "policy_id": 2, - "country_id": "us", - "api_version": "1", - "computed_household_json": {"value": 1}, - "status": "complete", - } - with uow.transaction() as daos: - daos.computed_households.upsert(**values) - daos.computed_households.upsert( - **{**values, "computed_household_json": {"value": 2}} - ) - with uow.read() as daos: - row = daos.computed_households.get(1, 2, "us", api_version="1") - assert row["computed_household_json"] == {"value": 2} - - -def test_policy_search_and_reform_impact_limit_use_typed_statements(): - uow = _unit_of_work() - with uow.transaction() as daos: - daos.policies.create("us", "Tax reform", {}, "one", "1") - daos.policies.create("us", "Other", {}, "two", "1") - with uow.read() as daos: - assert [row["label"] for row in daos.policies.search("us", "Tax")] == [ - "Tax reform" - ] - - -def test_computed_household_create_and_version_filters(): - uow = _unit_of_work() - with uow.transaction() as daos: - daos.computed_households.create( - household_id=1, - policy_id=2, - country_id="us", - api_version="1", - computed_household_json={"value": 1}, - status="complete", - ) - - with uow.read() as daos: - assert daos.computed_households.get(1, 2, "us")["computed_household_json"] == { - "value": 1 - } - assert daos.computed_households.get(1, 2, "us", api_version="2") is None diff --git a/tests/unit/data/test_orm_sessions.py b/tests/unit/data/test_orm_sessions.py index 663fd7542..681593c7b 100644 --- a/tests/unit/data/test_orm_sessions.py +++ b/tests/unit/data/test_orm_sessions.py @@ -4,9 +4,7 @@ import policyengine_api.data.orm as orm_module from policyengine_api.data.orm import ( - SessionManager, build_session_factory, - build_sqlite_session_manager, get_v1_session_factory, ) @@ -69,52 +67,3 @@ def test_runtime_factories_are_cached_and_separate(monkeypatch): assert local.kw["bind"] is local_engine finally: orm_module.clear_v1_session_factories() - - -def test_session_manager_commits_successful_transaction(): - manager = build_sqlite_session_manager() - with manager.engine.begin() as connection: - connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) - - manager.run_in_transaction( - lambda session: session.execute(text("INSERT INTO item (id) VALUES (1)")) - ) - - with manager.session() as session: - assert session.execute(text("SELECT id FROM item")).scalar_one() == 1 - - -def test_session_manager_rolls_back_failed_transaction(): - manager = build_sqlite_session_manager() - with manager.engine.begin() as connection: - connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) - - def fail(session): - session.execute(text("INSERT INTO item (id) VALUES (1)")) - raise RuntimeError("stop") - - with pytest.raises(RuntimeError, match="stop"): - manager.run_in_transaction(fail) - - with manager.session() as session: - assert session.execute(text("SELECT COUNT(*) FROM item")).scalar_one() == 0 - - -def test_session_manager_closes_sessions_after_callback(monkeypatch): - manager = build_sqlite_session_manager() - closed = [] - original_close = manager.session_factory.class_.close - - def recording_close(session): - closed.append(session) - return original_close(session) - - monkeypatch.setattr(manager.session_factory.class_, "close", recording_close) - manager.run_in_transaction(lambda session: None) - - assert len(closed) == 1 - - -def test_session_manager_requires_an_engine(): - with pytest.raises(TypeError): - SessionManager() # type: ignore[call-arg] diff --git a/tests/unit/data/test_remote_database_config.py b/tests/unit/data/test_remote_database_config.py index 33ce245a4..fbac992a4 100644 --- a/tests/unit/data/test_remote_database_config.py +++ b/tests/unit/data/test_remote_database_config.py @@ -2,7 +2,7 @@ os.environ.setdefault("FLASK_DEBUG", "1") -from policyengine_api.data.data import get_remote_database_config +from policyengine_api.data.orm import get_remote_database_config def test_remote_database_config_defaults_to_current_production_values(monkeypatch): diff --git a/tests/unit/data/test_run_daos.py b/tests/unit/data/test_run_daos.py deleted file mode 100644 index 1f6d17708..000000000 --- a/tests/unit/data/test_run_daos.py +++ /dev/null @@ -1,162 +0,0 @@ -from datetime import datetime - -import pytest - -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import SimulationDAO, V1UnitOfWork -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work(): - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) - - -def test_simulation_dao_creates_parent_and_monotonic_runs_atomically(): - uow = _unit_of_work() - with uow.transaction() as daos: - simulation_id = daos.simulations.create( - country_id="us", - api_version="1", - population_id="7", - population_type="household", - policy_id=2, - ) - first = daos.simulations.create_run( - simulation_id, - run_id="run-1", - status="pending", - trigger_type="create", - requested_at=datetime(2026, 1, 1), - ) - second = daos.simulations.create_run( - simulation_id, - run_id="run-2", - status="pending", - trigger_type="retry", - requested_at=datetime(2026, 1, 2), - ) - with uow.read() as daos: - assert first["run_sequence"] == 1 - assert second["run_sequence"] == 2 - assert daos.simulations.list_runs(simulation_id)[0]["id"] == "run-2" - - -def test_report_dao_round_trips_parent_run_and_alias(): - uow = _unit_of_work() - with uow.transaction() as daos: - report_id = daos.reports.create( - country_id="us", - simulation_1_id=1, - simulation_2_id=None, - api_version="1", - year="2026", - ) - run = daos.reports.create_run( - report_id, - run_id="report-run", - status="pending", - trigger_type="create", - requested_at=datetime(2026, 1, 1), - ) - daos.reports.set_alias(99, report_id) - with uow.read() as daos: - assert daos.reports.get(report_id)["status"] == "pending" - assert run["run_sequence"] == 1 - assert daos.reports.get_alias(99)["canonical_report_output_id"] == report_id - - -def test_simulation_dao_sync_callbacks_cover_create_update_and_missing_rows(): - uow = _unit_of_work() - - def read_synced(session, simulation_id, *, country_id): - return SimulationDAO.get_in_session(session, simulation_id, country_id) - - with uow.transaction() as daos: - created = daos.simulations.create_or_get_with_sync( - sync_callback=read_synced, - country_id="us", - api_version="1", - population_id="7", - population_type="household", - policy_id=2, - status="complete", - output={"result": 1}, - ) - reused = daos.simulations.create_or_get_with_sync( - sync_callback=read_synced, - country_id="us", - api_version="1", - population_id="7", - population_type="household", - policy_id=2, - status="pending", - ) - updated = daos.simulations.update_with_sync( - created["id"], - "us", - {"error_message": "updated"}, - read_synced, - ) - dual_write = daos.simulations.ensure_dual_write_state(created["id"], "us") - - assert reused["id"] == created["id"] - assert updated["error_message"] == "updated" - assert dual_write["latest_successful_run_id"] is not None - assert daos.simulations.get(created["id"], "uk") is None - assert daos.simulations.update(999, status="complete") is False - with pytest.raises(ValueError, match="Simulation #999 not found"): - daos.simulations.update_with_sync( - 999, "us", {"status": "complete"}, read_synced - ) - with pytest.raises(LookupError, match="Simulation 999 does not exist"): - daos.simulations.create_run( - 999, - run_id="missing-run", - status="pending", - trigger_type="create", - ) - - -def test_report_dao_handles_scoped_lookups_updates_and_existing_aliases(): - uow = _unit_of_work() - with uow.transaction() as daos: - report_id = daos.reports.create( - country_id="us", - simulation_1_id=1, - simulation_2_id=None, - api_version="1", - year="2026", - ) - run = daos.reports.create_run( - report_id, - run_id="report-run", - status="pending", - trigger_type="create", - ) - - assert daos.reports.get(report_id, "uk") is None - assert daos.reports.get_for_update(report_id, "us")["id"] == report_id - assert daos.reports.get_for_update(report_id, "uk") is None - assert daos.reports.update(999, status="complete") is False - assert daos.reports.update(report_id, status="complete") - assert daos.reports.update_run( - run["id"], status="complete", output={"result": 1} - ) - assert daos.reports.update_run("missing-run", status="error") is False - daos.reports.set_alias(99, report_id) - daos.reports.set_alias(99, report_id + 1) - with pytest.raises(LookupError, match="Report output 999 does not exist"): - daos.reports.create_run( - 999, - run_id="missing-run", - status="pending", - trigger_type="create", - ) - - with uow.read() as daos: - assert daos.reports.get(report_id)["status"] == "complete" - assert daos.reports.get_run(run["id"])["output"] == {"result": 1} - assert daos.reports.get_run("missing-run") is None - assert daos.reports.get_alias(99)["canonical_report_output_id"] == 2 diff --git a/tests/unit/data/test_run_schema.py b/tests/unit/data/test_run_schema.py index 2bcba1eff..fdbd2a78e 100644 --- a/tests/unit/data/test_run_schema.py +++ b/tests/unit/data/test_run_schema.py @@ -1,15 +1,16 @@ from pathlib import Path +from sqlalchemy import inspect + from policyengine_api.constants import REPO -def _column_names(test_db, table_name: str) -> set[str]: - rows = test_db.query(f"PRAGMA table_info({table_name})").fetchall() - return {row["name"] for row in rows} +def _column_names(test_engine, table_name: str) -> set[str]: + return {column["name"] for column in inspect(test_engine).get_columns(table_name)} -def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): - report_output_columns = _column_names(test_db, "report_outputs") +def test_stage_one_run_schema_is_initialized_in_local_test_db(test_engine): + report_output_columns = _column_names(test_engine, "report_outputs") assert { "report_kind", "report_spec_json", @@ -19,7 +20,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "latest_successful_run_id", }.issubset(report_output_columns) - simulation_columns = _column_names(test_db, "simulations") + simulation_columns = _column_names(test_engine, "simulations") assert { "simulation_spec_json", "simulation_spec_schema_version", @@ -27,7 +28,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "latest_successful_run_id", }.issubset(simulation_columns) - report_run_columns = _column_names(test_db, "report_output_runs") + report_run_columns = _column_names(test_engine, "report_output_runs") assert { "id", "report_output_id", @@ -45,7 +46,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "resolved_options_hash", }.issubset(report_run_columns) - simulation_run_columns = _column_names(test_db, "simulation_runs") + simulation_run_columns = _column_names(test_engine, "simulation_runs") assert { "id", "simulation_id", @@ -61,7 +62,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "simulation_cache_version", }.issubset(simulation_run_columns) - alias_columns = _column_names(test_db, "legacy_report_output_aliases") + alias_columns = _column_names(test_engine, "legacy_report_output_aliases") assert {"legacy_report_output_id", "canonical_report_output_id"} == alias_columns diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index da29569b0..b61bd5717 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -1,317 +1,111 @@ -"""Tests for SQLAlchemy v2 compatibility. - -These tests verify that the database layer works correctly with -SQLAlchemy v2, specifically: -- The _ResultProxy wrapper provides fetchone()/fetchall() on eagerly - fetched results. -- The remote (non-local) query path uses connection-based execution - instead of the removed engine.execute(). -- Row objects returned from the remote path support dict-like access - (dict(row) and row["key"]). -""" - -import policyengine_api.data.data as data_module -import sqlalchemy from unittest.mock import Mock -from policyengine_api.data.data import PolicyEngineDatabase, _ResultProxy - - -class TestSQLAlchemyVersion: - """Verify that SQLAlchemy v2 is installed.""" - - def test_sqlalchemy_version_is_v2(self): - major = int(sqlalchemy.__version__.split(".")[0]) - assert major >= 2, f"Expected SQLAlchemy v2+, got {sqlalchemy.__version__}" - - -class TestResultProxy: - """Test the _ResultProxy wrapper that bridges SQLAlchemy v2 - connection-scoped results with the existing query() API.""" - - def test_fetchone_returns_dict_like_rows(self): - """Rows returned by fetchone() should support dict() and - key-based access.""" - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql( - "CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)" - ) - conn.exec_driver_sql("INSERT INTO test VALUES (1, 'hello')") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - row = proxy.fetchone() - assert row is not None - assert dict(row) == {"id": 1, "name": "hello"} - assert row["id"] == 1 - assert row["name"] == "hello" - - def test_fetchone_returns_none_when_exhausted(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - assert proxy.fetchone() is None - - def test_fetchall_returns_all_rows(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY, val TEXT)") - conn.exec_driver_sql("INSERT INTO test VALUES (1, 'a')") - conn.exec_driver_sql("INSERT INTO test VALUES (2, 'b')") - conn.exec_driver_sql("INSERT INTO test VALUES (3, 'c')") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - rows = proxy.fetchall() - assert len(rows) == 3 - assert dict(rows[0]) == {"id": 1, "val": "a"} - assert dict(rows[2]) == {"id": 3, "val": "c"} - - def test_fetchone_then_fetchall_respects_cursor_position(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - conn.exec_driver_sql("INSERT INTO test VALUES (1)") - conn.exec_driver_sql("INSERT INTO test VALUES (2)") - conn.exec_driver_sql("INSERT INTO test VALUES (3)") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - first = proxy.fetchone() - assert dict(first) == {"id": 1} - remaining = proxy.fetchall() - assert len(remaining) == 2 - assert dict(remaining[0]) == {"id": 2} - - def test_result_proxy_for_insert_statement(self): - """INSERT statements produce no rows; _ResultProxy should - handle this gracefully.""" - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - result = conn.exec_driver_sql("INSERT INTO test VALUES (1)") - proxy = _ResultProxy(result) - - assert proxy.fetchone() is None - assert proxy.fetchall() == [] - - -class TestRemoteQueryPath: - """Test the non-local query path that uses SQLAlchemy engine - with connection-based execution (v2 pattern).""" - - def _make_remote_db(self): - """Create a PolicyEngineDatabase-like object that uses - a SQLAlchemy engine (the 'remote' path) but backed by - in-memory SQLite for testing.""" - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db.local = False - db.pool = sqlalchemy.create_engine("sqlite://") - # Initialize schema using the remote path - with db.pool.connect() as conn: - conn.exec_driver_sql( - "CREATE TABLE test_table " - "(id INTEGER PRIMARY KEY, name TEXT, value REAL)" - ) - conn.commit() - return db - - def test_remote_insert_and_select(self): - """Test INSERT then SELECT through the remote query path.""" - db = self._make_remote_db() - # Note: remote path converts ? to %s for MySQL, but SQLite - # uses ? natively. Since exec_driver_sql passes to the DBAPI - # driver directly and SQLite's driver uses ?, we need to - # test with the actual query() method which does the conversion. - # For SQLite DBAPI, ? is the native marker. - - # Use exec_driver_sql directly to bypass ?->%s conversion - # (which would break SQLite) - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "test", 3.14), - ] - ) - - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - row = result.fetchone() - assert row is not None - assert row["id"] == 1 - assert row["name"] == "test" - assert row["value"] == 3.14 - assert dict(row) == {"id": 1, "name": "test", "value": 3.14} - - def test_remote_select_no_results(self): - db = self._make_remote_db() - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (999,)]) - assert result.fetchone() is None - - def test_remote_update(self): - db = self._make_remote_db() - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "original", 1.0), - ] - ) - db._execute_remote( - [ - "UPDATE test_table SET name = ? WHERE id = ?", - ("updated", 1), - ] - ) - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - row = result.fetchone() - assert row["name"] == "updated" - - def test_remote_delete(self): - db = self._make_remote_db() - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "to_delete", 0.0), - ] - ) - db._execute_remote(["DELETE FROM test_table WHERE id = ?", (1,)]) - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - assert result.fetchone() is None - - -class TestRemotePoolSetup: - """Test remote pool setup without opening a real Cloud SQL connection.""" - - def _stub_remote_pool(self, monkeypatch): - connector_calls = [] - engine_calls = [] - - class FakeConnector: - def __init__(self, **kwargs): - self.options = kwargs - - def connect(self, **kwargs): - connector_calls.append(kwargs) - return object() - - def fake_create_engine(url, **kwargs): - engine_calls.append((url, kwargs)) - return "fake-engine" - - fake_connector = FakeConnector(refresh_strategy="LAZY") - - def fake_connector_factory(**kwargs): - fake_connector.options = kwargs - return fake_connector - - monkeypatch.setattr(data_module, "Connector", fake_connector_factory) - monkeypatch.setattr(data_module.sqlalchemy, "create_engine", fake_create_engine) - return fake_connector, connector_calls, engine_calls - - def test_create_pool_uses_remote_database_config(self, monkeypatch): - fake_connector, connector_calls, engine_calls = self._stub_remote_pool( - monkeypatch - ) - monkeypatch.setenv( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", - "test-project:us-central1:test-db", - ) - monkeypatch.setenv("POLICYENGINE_DB_USER", "test-user") - monkeypatch.setenv("POLICYENGINE_DB_NAME", "test-db") - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "test-password") - - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db._create_pool() - - assert db.connector is fake_connector - assert db.pool == "fake-engine" - assert connector_calls == [] - creator = engine_calls[0][1]["creator"] - first = creator() - second = creator() - assert first is not second - assert connector_calls == [ - { - "instance_connection_string": "test-project:us-central1:test-db", - "driver": "pymysql", - "db": "test-db", - "user": "test-user", - "password": "test-password", - }, - { - "instance_connection_string": "test-project:us-central1:test-db", - "driver": "pymysql", - "db": "test-db", - "user": "test-user", - "password": "test-password", - }, - ] - assert engine_calls[0][0] == "mysql+pymysql://" - assert engine_calls[0][1]["pool_pre_ping"] is True - assert engine_calls[0][1]["pool_recycle"] == 1800 - assert engine_calls[0][1]["pool_size"] == 5 - assert engine_calls[0][1]["max_overflow"] == 2 - assert engine_calls[0][1]["pool_timeout"] == 30 - - def test_create_pool_settings_are_owned_by_the_application(self, monkeypatch): - fake_connector, _, engine_calls = self._stub_remote_pool(monkeypatch) - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "test-password") - monkeypatch.setenv("POLICYENGINE_DB_PRIVATE_IP", "true") - monkeypatch.setenv("POLICYENGINE_DB_POOL_RECYCLE", "not-used") - monkeypatch.setenv("POLICYENGINE_DB_POOL_SIZE", "not-used") - monkeypatch.setenv("POLICYENGINE_DB_MAX_OVERFLOW", "not-used") - monkeypatch.setenv("POLICYENGINE_DB_POOL_TIMEOUT", "not-used") - - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db._create_pool() - - assert fake_connector.options == { - "ip_type": data_module.IPTypes.PUBLIC, - "refresh_strategy": "LAZY", +import sqlalchemy +from sqlalchemy import create_engine, func, select +from sqlalchemy.orm import Session + +import policyengine_api.data.orm as orm +from policyengine_api.data.v1_models import Policy + + +def test_sqlalchemy_v2_or_newer_is_installed(): + assert int(sqlalchemy.__version__.split(".")[0]) >= 2 + + +def test_remote_engine_delegates_connections_and_pooling_to_sqlalchemy(monkeypatch): + connector_calls = [] + engine_calls = [] + + class FakeConnector: + def __init__(self, **options): + self.options = options + + def connect(self, **options): + connector_calls.append(options) + return object() + + def close(self): + pass + + connector = FakeConnector() + + def connector_factory(**options): + connector.options = options + return connector + + monkeypatch.setattr(orm, "Connector", connector_factory) + monkeypatch.setattr( + orm, + "create_engine", + lambda url, **options: engine_calls.append((url, options)) or "engine", + ) + monkeypatch.setenv( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", "project:region:instance" + ) + monkeypatch.setenv("POLICYENGINE_DB_USER", "user") + monkeypatch.setenv("POLICYENGINE_DB_NAME", "database") + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "password") + + engine = orm._build_remote_engine() + + assert engine == "engine" + assert connector.options == { + "ip_type": orm.IPTypes.PUBLIC, + "refresh_strategy": "LAZY", + } + url, options = engine_calls[0] + assert url == "mysql+pymysql://" + assert options["pool_pre_ping"] is True + assert options["pool_recycle"] == 1800 + assert options["pool_size"] == 5 + assert options["max_overflow"] == 2 + assert options["pool_timeout"] == 30 + assert connector_calls == [] + options["creator"]() + assert connector_calls == [ + { + "instance_connection_string": "project:region:instance", + "driver": "pymysql", + "db": "database", + "user": "user", + "password": "password", } - assert engine_calls[0][1]["pool_recycle"] == 1800 - assert engine_calls[0][1]["pool_size"] == 5 - assert engine_calls[0][1]["max_overflow"] == 2 - assert engine_calls[0][1]["pool_timeout"] == 30 + ] - def test_create_pool_reads_dot_dbpw_file(self, monkeypatch, tmp_path): - _, connector_calls, engine_calls = self._stub_remote_pool(monkeypatch) - monkeypatch.chdir(tmp_path) - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", ".dbpw") - (tmp_path / ".dbpw").write_text("file-password\n") - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db._create_pool() +def test_database_password_can_be_loaded_from_file(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", ".dbpw") + (tmp_path / ".dbpw").write_text("file-password\n", encoding="utf-8") - engine_calls[0][1]["creator"]() - assert connector_calls[0]["password"] == "file-password" + assert orm._database_password() == "file-password" - def test_remote_constructor_initializes_pool_without_local_database( - self, monkeypatch - ): - calls = [] - def fake_create_pool(self): - calls.append(("pool", self.local)) +def test_local_initializer_bootstraps_schema_and_current_law_rows(tmp_path): + database_path = tmp_path / "local.db" - monkeypatch.setattr(PolicyEngineDatabase, "_create_pool", fake_create_pool) + orm._initialize_local_database(database_path) - db = PolicyEngineDatabase(local=False, initialize=False) + engine = create_engine(f"sqlite+pysqlite:///{database_path}") + try: + with Session(engine) as session: + assert session.scalar(select(func.count()).select_from(Policy)) > 0 + assert session.scalars(select(Policy)).first().policy_json == {} + finally: + engine.dispose() - assert db.local is False - assert calls == [("pool", False)] - def test_close_disposes_engine_then_closes_connector(self): - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db.local = False - db.pool = Mock() - db.connector = Mock() +def test_close_v1_engines_disposes_pools_and_connectors(monkeypatch): + engine = Mock() + connector = Mock() + monkeypatch.setattr(orm, "_v1_engines", {False: engine}) + monkeypatch.setattr(orm, "_cloud_sql_connectors", {False: connector}) + monkeypatch.setattr(orm, "_v1_session_factories", {False: Mock()}) - db.close() - db.close() + orm.close_v1_engines() - db.pool.dispose.assert_called_once_with() - db.connector.close.assert_called_once_with() + engine.dispose.assert_called_once_with() + connector.close.assert_called_once_with() + assert orm._v1_engines == {} + assert orm._cloud_sql_connectors == {} + assert orm._v1_session_factories == {} diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index 0ac8ae704..cb9b6ca48 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -39,7 +39,7 @@ def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): assert "runtime_sqlalchemy_dao" not in source -def test_report_orchestration_uses_typed_daos_not_raw_sql(): +def test_report_orchestration_uses_sessions_and_models_not_raw_sql(): source = (PACKAGE_ROOT / "services/report_output_service.py").read_text( encoding="utf-8" ) @@ -48,8 +48,9 @@ def test_report_orchestration_uses_typed_daos_not_raw_sql(): assert "exec_driver_sql" not in source -def test_typed_repository_module_has_no_raw_sql_compatibility_dao(): - source = (PACKAGE_ROOT / "data/v1_daos.py").read_text(encoding="utf-8") - assert "SQLAlchemyDAO" not in source - assert "runtime_sqlalchemy_dao" not in source - assert "exec_driver_sql" not in source +def test_legacy_persistence_compatibility_layers_are_absent(): + assert not (PACKAGE_ROOT / "data/v1_daos.py").exists() + assert not (PACKAGE_ROOT / "data/data.py").exists() + orm_source = (PACKAGE_ROOT / "data/orm.py").read_text(encoding="utf-8") + assert "SessionManager" not in orm_source + assert "PolicyEngineDatabase" not in orm_source diff --git a/tests/unit/data/test_v1_daos.py b/tests/unit/data/test_v1_daos.py deleted file mode 100644 index c00464ce6..000000000 --- a/tests/unit/data/test_v1_daos.py +++ /dev/null @@ -1,57 +0,0 @@ -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work(): - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) - - -def test_policy_dao_round_trips_legacy_mapping_shape(): - uow = _unit_of_work() - with uow.transaction() as daos: - policy_id = daos.policies.create("us", "Reform", {"gov.irs": 1}, "hash", "1.0") - with uow.read() as daos: - assert policy_id == 1 - assert daos.policies.get("us", policy_id) == { - "id": 1, - "country_id": "us", - "label": "Reform", - "api_version": "1.0", - "policy_json": {"gov.irs": 1}, - "policy_hash": "hash", - } - - -def test_policy_dao_allocates_ids_and_detects_existing_policy(): - uow = _unit_of_work() - with uow.transaction() as daos: - assert daos.policies.create("us", None, {}, "one", "1.0") == 1 - assert daos.policies.create("uk", None, {}, "two", "1.0") == 2 - with uow.read() as daos: - assert daos.policies.find_unique("us", "one", None)["id"] == 1 - - -def test_household_dao_creates_updates_and_reads(): - uow = _unit_of_work() - with uow.transaction() as daos: - household_id = daos.households.create("us", "Home", {"people": {}}, "h", "1.0") - daos.households.update( - "us", - household_id, - "Updated", - {"people": {"you": {}}}, - "updated-hash", - "2.0", - ) - with uow.read() as daos: - assert daos.households.get("us", household_id)["label"] == "Updated" - assert daos.households.get("uk", household_id) is None - - -def test_household_dao_handles_missing_update(): - uow = _unit_of_work() - with uow.transaction() as daos: - assert daos.households.update("us", 999, "missing", {}, "missing", "1") is False diff --git a/tests/unit/data/test_v1_unit_of_work.py b/tests/unit/data/test_v1_unit_of_work.py deleted file mode 100644 index 91e768750..000000000 --- a/tests/unit/data/test_v1_unit_of_work.py +++ /dev/null @@ -1,64 +0,0 @@ -import pytest - -from policyengine_api.data.orm import build_sqlite_session_manager -from policyengine_api.data.v1_daos import V1UnitOfWork -from tests.unit.data.sqlite_schema import create_sqlite_v1_schema - - -def _unit_of_work() -> V1UnitOfWork: - manager = build_sqlite_session_manager() - create_sqlite_v1_schema(manager) - return V1UnitOfWork(manager) - - -def test_unit_of_work_commits_all_daos_once(): - uow = _unit_of_work() - - with uow.transaction() as daos: - policy_id = daos.policies.create("us", None, {}, "policy", "1") - household_id = daos.households.create("us", None, {}, "household", "1") - - with uow.read() as daos: - assert daos.policies.get("us", policy_id) is not None - assert daos.households.get("us", household_id) is not None - - -def test_unit_of_work_rolls_back_every_repository_on_failure(): - uow = _unit_of_work() - - with pytest.raises(RuntimeError, match="abort"): - with uow.transaction() as daos: - daos.policies.create("us", None, {}, "policy", "1") - daos.households.create("us", None, {}, "household", "1") - raise RuntimeError("abort") - - with uow.read() as daos: - assert daos.policies.get("us", 1) is None - assert daos.households.get("us", 1) is None - - -def test_unit_of_work_rolls_back_parent_run_and_alias_together(): - uow = _unit_of_work() - - with pytest.raises(RuntimeError, match="abort report"): - with uow.transaction() as daos: - report_id = daos.reports.create( - country_id="us", - simulation_1_id=1, - simulation_2_id=None, - api_version="1", - year="2026", - ) - daos.reports.create_run( - report_id, - run_id="report-run", - status="pending", - trigger_type="create", - ) - daos.reports.set_alias(99, report_id) - raise RuntimeError("abort report") - - with uow.read() as daos: - assert daos.reports.get(1) is None - assert daos.reports.get_run("report-run") is None - assert daos.reports.get_alias(99) is None diff --git a/tests/unit/endpoints/test_get_simulations.py b/tests/unit/endpoints/test_get_simulations.py index dfa227fce..0e85afaa6 100644 --- a/tests/unit/endpoints/test_get_simulations.py +++ b/tests/unit/endpoints/test_get_simulations.py @@ -7,69 +7,69 @@ always LIMIT, clamp to [1, 1000], and bind as a parameter. """ +from datetime import datetime + +from sqlalchemy import func, select + +from policyengine_api.data.v1_models import ReformImpact from policyengine_api.endpoints.simulation import get_simulations -def _seed_reform_impacts(test_db, n: int) -> None: +def _seed_reform_impacts(orm_session, n: int) -> None: for i in range(n): - test_db.query( - """INSERT INTO reform_impact - (baseline_policy_id, reform_policy_id, country_id, region, dataset, - time_period, options_json, options_hash, api_version, - reform_impact_json, status, start_time, execution_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - i + 1, - i + 2, - "us", - "us", - "custom_dataset", - "2025", - "{}", - f"hash-{i}", - "1.0.0", - "{}", - "complete", - f"2026-01-01 00:{i // 60:02d}:{i % 60:02d}", - f"exec-{i}", - ), + orm_session.add( + ReformImpact( + baseline_policy_id=i + 1, + reform_policy_id=i + 2, + country_id="us", + region="us", + dataset="custom_dataset", + time_period="2025", + options_json={}, + options_hash=f"hash-{i}", + api_version="1.0.0", + reform_impact_json={}, + status="complete", + start_time=datetime(2026, 1, 1, 0, i // 60, i % 60), + execution_id=f"exec-{i}", + ) ) + orm_session.commit() -def test_get_simulations_default_limit_caps_at_100(test_db): - _seed_reform_impacts(test_db, 150) +def test_get_simulations_default_limit_caps_at_100(orm_session): + _seed_reform_impacts(orm_session, 150) result = get_simulations() assert len(result["result"]) == 100 -def test_get_simulations_clamps_huge_max_results(test_db): - _seed_reform_impacts(test_db, 50) +def test_get_simulations_clamps_huge_max_results(orm_session): + _seed_reform_impacts(orm_session, 50) # A caller passing an absurdly large value must not crash and # must not cause a full scan; the value is clamped at 1000. result = get_simulations(max_results=10**9) assert len(result["result"]) == 50 # only 50 seeded -def test_get_simulations_clamps_negative_max_results(test_db): - _seed_reform_impacts(test_db, 5) +def test_get_simulations_clamps_negative_max_results(orm_session): + _seed_reform_impacts(orm_session, 5) # max_results of 0 or negative must still return something sane. result = get_simulations(max_results=0) assert 1 <= len(result["result"]) <= 5 -def test_get_simulations_defaults_when_none(test_db): - _seed_reform_impacts(test_db, 10) +def test_get_simulations_defaults_when_none(orm_session): + _seed_reform_impacts(orm_session, 10) result = get_simulations(max_results=None) assert len(result["result"]) == 10 # fewer than the default 100 -def test_get_simulations_rejects_non_integer_gracefully(test_db): - _seed_reform_impacts(test_db, 5) +def test_get_simulations_rejects_non_integer_gracefully(orm_session): + _seed_reform_impacts(orm_session, 5) # A string like "100; DROP TABLE reform_impact" must not reach # the SQL statement; it falls back to the default. result = get_simulations(max_results="100; DROP TABLE reform_impact") assert len(result["result"]) == 5 # And the table must still exist. - rows = test_db.query("SELECT COUNT(*) AS c FROM reform_impact").fetchone() - assert rows["c"] == 5 + assert orm_session.scalar(select(func.count()).select_from(ReformImpact)) == 5 diff --git a/tests/unit/endpoints/test_set_user_policy_dataset.py b/tests/unit/endpoints/test_set_user_policy_dataset.py index dcc6082e2..4d7001071 100644 --- a/tests/unit/endpoints/test_set_user_policy_dataset.py +++ b/tests/unit/endpoints/test_set_user_policy_dataset.py @@ -1,120 +1,40 @@ -"""Regression test for issue #3310. - -When dataset is truthy in set_user_policy (policy.py:224-237), the query -builds "AND dataset = ?" (7 placeholders) but the params must include -the dataset value. Previously, the dataset was never appended, causing -a parameter binding crash. -""" - import time +import pytest +from flask import Flask +from sqlalchemy import select -def test_set_user_policy_dataset_param_included(test_db): - """Verify that the SELECT after INSERT in set_user_policy correctly - includes the dataset value in params when dataset is truthy. - - Reproduces the exact code path from policy.py:224-237. - """ - now = int(time.time()) - - country_id = "us" - reform_id = 2 - baseline_id = 1 - user_id = "user1" - year = "2025" - geography = "us" - dataset = "custom_dataset" - - # Insert a user_policy with a non-null dataset - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - country_id, - None, - reform_id, - None, - baseline_id, - user_id, - year, - geography, - dataset, - 3, - "1.0.0", - now, - now, - ), - ) - - # Reproduce the exact code from policy.py:224-237 - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.endpoints.policy import set_user_policy - # This must not crash — 7 placeholders, 7 params - row = test_db.query(query, tuple(params)).fetchone() - assert row is not None - assert row["dataset"] == "custom_dataset" - assert row["reform_id"] == reform_id - assert row["user_id"] == user_id +def create_client(): + app = Flask(__name__) + app.config["TESTING"] = True + app.route("//user-policy", methods=["POST"])(set_user_policy) + return app.test_client() -def test_set_user_policy_dataset_null_still_works(test_db): - """When dataset is None/falsy, the query uses IS NULL with 6 params.""" +@pytest.mark.parametrize("dataset", ["custom_dataset", None]) +def test_set_user_policy_persists_dataset_with_orm(orm_session, dataset): now = int(time.time()) - - country_id = "us" - reform_id = 2 - baseline_id = 1 - user_id = "user1" - year = "2025" - geography = "us" - dataset = None - - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - country_id, - None, - reform_id, - None, - baseline_id, - user_id, - year, - geography, - dataset, - 3, - "1.0.0", - now, - now, - ), + response = create_client().post( + "/us/user-policy", + json={ + "reform_id": 2, + "baseline_id": 1, + "user_id": "user1", + "year": "2025", + "geography": "us", + "dataset": dataset, + "number_of_provisions": 3, + "api_version": "1.0.0", + "added_date": now, + "updated_date": now, + }, ) - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) - - row = test_db.query(query, tuple(params)).fetchone() - - assert row is not None - assert row["dataset"] is None + assert response.status_code == 201 + orm_session.expire_all() + policy = orm_session.scalar(select(UserPolicy).where(UserPolicy.user_id == "user1")) + assert policy.dataset == dataset diff --git a/tests/unit/endpoints/test_update_user_policy.py b/tests/unit/endpoints/test_update_user_policy.py index 556f93a23..bb8b8a18f 100644 --- a/tests/unit/endpoints/test_update_user_policy.py +++ b/tests/unit/endpoints/test_update_user_policy.py @@ -12,6 +12,7 @@ from flask import Flask +from policyengine_api.data.v1_models import UserPolicy from policyengine_api.endpoints import update_user_policy @@ -22,38 +23,31 @@ def _create_test_client() -> Flask: return app.test_client() -def _insert_user_policy(test_db) -> int: +def _insert_user_policy(orm_session) -> int: now = int(time.time()) - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - "us", - "old label", - 2, - None, - 1, - "user1", - "2025", - "us", - "custom_dataset", - 3, - "1.0.0", - now, - now, - ), + policy = UserPolicy( + country_id="us", + reform_label="old label", + reform_id=2, + baseline_label=None, + baseline_id=1, + user_id="user1", + year="2025", + geography="us", + dataset="custom_dataset", + number_of_provisions=3, + api_version="1.0.0", + added_date=now, + updated_date=now, ) - row = test_db.query( - "SELECT id FROM user_policies ORDER BY id DESC LIMIT 1" - ).fetchone() - return row["id"] + orm_session.add(policy) + orm_session.commit() + return policy.id -def test_update_user_policy_rejects_sql_injection_key(test_db): +def test_update_user_policy_rejects_sql_injection_key(orm_session): """Unknown keys (including SQL injection attempts) must be rejected.""" - policy_id = _insert_user_policy(test_db) + policy_id = _insert_user_policy(orm_session) client = _create_test_client() response = client.put( @@ -69,16 +63,13 @@ def test_update_user_policy_rejects_sql_injection_key(test_db): assert "unsupported fields" in body["message"] # The row must be untouched. - row = test_db.query( - "SELECT reform_label FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["reform_label"] == "old label" + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "old label" -def test_update_user_policy_rejects_identity_column(test_db): +def test_update_user_policy_rejects_identity_column(orm_session): """Identity columns (user_id, country_id, ...) must not be writable.""" - policy_id = _insert_user_policy(test_db) + policy_id = _insert_user_policy(orm_session) client = _create_test_client() response = client.put( @@ -87,16 +78,13 @@ def test_update_user_policy_rejects_identity_column(test_db): ) assert response.status_code == 400 - row = test_db.query( - "SELECT user_id FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["user_id"] == "user1" + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).user_id == "user1" -def test_update_user_policy_allows_whitelisted_field(test_db): +def test_update_user_policy_allows_whitelisted_field(orm_session): """Whitelisted fields (e.g. reform_label) can still be updated.""" - policy_id = _insert_user_policy(test_db) + policy_id = _insert_user_policy(orm_session) client = _create_test_client() response = client.put( @@ -105,21 +93,18 @@ def test_update_user_policy_allows_whitelisted_field(test_db): ) assert response.status_code == 200 - row = test_db.query( - "SELECT reform_label FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["reform_label"] == "new label" + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "new label" -def test_update_user_policy_requires_id(test_db): +def test_update_user_policy_requires_id(): client = _create_test_client() response = client.put("/us/user-policy", json={"reform_label": "x"}) assert response.status_code == 400 -def test_update_user_policy_requires_at_least_one_field(test_db): - policy_id = _insert_user_policy(test_db) +def test_update_user_policy_requires_at_least_one_field(orm_session): + policy_id = _insert_user_policy(orm_session) client = _create_test_client() response = client.put("/us/user-policy", json={"id": policy_id}) assert response.status_code == 400 diff --git a/tests/unit/services/test_stage7_dao_boundaries.py b/tests/unit/services/test_stage7_dao_boundaries.py index 7ccaf42d2..6bbe87133 100644 --- a/tests/unit/services/test_stage7_dao_boundaries.py +++ b/tests/unit/services/test_stage7_dao_boundaries.py @@ -4,7 +4,6 @@ SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" -DAO_MODULE = SERVICE_ROOT.parent / "data" / "v1_daos.py" @pytest.mark.parametrize( @@ -28,8 +27,5 @@ def test_migrated_services_use_sessions_and_mapped_models(module_name): assert "build_v1_session_manager" not in source -def test_migrated_user_domains_have_no_temporary_daos(): - source = DAO_MODULE.read_text(encoding="utf-8") - assert "class UserDAO" not in source - assert "class UserPolicyDAO" not in source - assert "class EconomyDAO" not in source +def test_temporary_dao_module_has_been_removed(): + assert not (SERVICE_ROOT.parent / "data" / "v1_daos.py").exists() diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_stage7_local_service_boundaries.py index d7d0dd6ce..2e1bd4224 100644 --- a/tests/unit/services/test_stage7_local_service_boundaries.py +++ b/tests/unit/services/test_stage7_local_service_boundaries.py @@ -4,7 +4,6 @@ SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" -DAO_MODULE = SERVICE_ROOT.parent / "data" / "v1_daos.py" @pytest.mark.parametrize( @@ -23,8 +22,12 @@ def test_local_data_services_do_not_issue_queries_directly(module_name): assert "from policyengine_api.data import database" not in source -def test_migrated_local_domains_have_no_temporary_daos(): - source = DAO_MODULE.read_text(encoding="utf-8") - assert "class AnalysisDAO" not in source - assert "class ReformImpactDAO" not in source - assert "class TracerDAO" not in source +def test_migrated_local_services_accept_caller_owned_sessions(): + for module_name in ( + "ai_analysis_service.py", + "reform_impacts_service.py", + "tracer_analysis_service.py", + "report_output_alias_service.py", + ): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert "from sqlalchemy.orm import Session" in source diff --git a/tests/unit/services/test_tracer_service.py b/tests/unit/services/test_tracer_service.py index d9a490537..04a35f636 100644 --- a/tests/unit/services/test_tracer_service.py +++ b/tests/unit/services/test_tracer_service.py @@ -16,10 +16,10 @@ def test_get_tracer_valid(test_tracer_data, orm_session): result = tracer_service.get_tracer( orm_session, - test_tracer_data["country_id"], - test_tracer_data["household_id"], - test_tracer_data["policy_id"], - test_tracer_data["api_version"], + test_tracer_data.country_id, + test_tracer_data.household_id, + test_tracer_data.policy_id, + test_tracer_data.api_version, ) # match the valid output as collected from fixture From 69e67ed2a56d679f100b5eb55d2a0b9eacfbf7a6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 03:08:20 +0300 Subject: [PATCH 56/89] chore: remove unused API dashboard --- dashboard/app.py | 79 ------------------------------------------------ 1 file changed, 79 deletions(-) delete mode 100644 dashboard/app.py diff --git a/dashboard/app.py b/dashboard/app.py deleted file mode 100644 index d4e0414ef..000000000 --- a/dashboard/app.py +++ /dev/null @@ -1,79 +0,0 @@ -import streamlit as st -from sqlalchemy import select - -from policyengine_api.data.orm import get_v1_session_factory -from policyengine_api.data.v1_models import Policy - - -st.title("PolicyEngine API dashboard") - - -def serialize_policy(policy: Policy) -> dict: - return { - column.name: getattr(policy, column.name) for column in Policy.__table__.columns - } - - -st.subheader("Recent policies") -if st.button("Refresh policies"): - sessions = get_v1_session_factory() - with sessions() as session: - policies = session.scalars(select(Policy).limit(10)).all() - st.table([serialize_policy(policy) for policy in policies]) - - -st.subheader("Look up a policy") -policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_lookup_text")) -country_id = st.text_input("Enter a country ID", "uk", key="policy_lookup_country") -if st.button("Look up policy", key="policy_lookup"): - sessions = get_v1_session_factory() - with sessions() as session: - policy = session.scalar( - select(Policy).where( - Policy.id == policy_id, - Policy.country_id == country_id, - ) - ) - if policy is None: - st.error("Policy not found") - else: - st.table([serialize_policy(policy)]) - - -st.subheader("Set a policy's label") -policy_id = int(st.text_input("Enter a policy ID", "1")) -country_id = st.text_input("Enter a country ID", "uk") -new_label = st.text_input("Enter a new label", "New label", key="policy_label_text") -if st.button("Set policy label", key="policy_label"): - sessions = get_v1_session_factory() - with sessions.begin() as session: - policy = session.scalar( - select(Policy).where( - Policy.id == policy_id, - Policy.country_id == country_id, - ) - ) - if policy is None: - st.error("Policy not found") - else: - policy.label = new_label - st.success("Success!") - - -st.subheader("Delete a policy") -policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_delete_text")) -country_id = st.text_input("Enter a country ID", "uk", key="policy_delete_country") -if st.button("Delete policy", key="policy_delete"): - sessions = get_v1_session_factory() - with sessions.begin() as session: - policy = session.scalar( - select(Policy).where( - Policy.id == policy_id, - Policy.country_id == country_id, - ) - ) - if policy is None: - st.error("Policy not found") - else: - session.delete(policy) - st.success("Success!") From 6c780410037eccb556b92f2afcb6e69acefe339e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 8 Aug 2026 03:26:23 +0300 Subject: [PATCH 57/89] refactor: bootstrap local database with SQLAlchemy --- policyengine_api/data/initialise.sql | 194 ------------------- policyengine_api/data/initialise_local.sql | 205 --------------------- policyengine_api/data/local_database.py | 43 +++++ policyengine_api/data/orm.py | 71 +++---- tests/unit/conftest.py | 11 +- tests/unit/data/test_run_schema.py | 27 --- tests/unit/data/test_sqlalchemy_v2.py | 34 +++- 7 files changed, 113 insertions(+), 472 deletions(-) delete mode 100644 policyengine_api/data/initialise.sql delete mode 100644 policyengine_api/data/initialise_local.sql create mode 100644 policyengine_api/data/local_database.py diff --git a/policyengine_api/data/initialise.sql b/policyengine_api/data/initialise.sql deleted file mode 100644 index 6fd9210db..000000000 --- a/policyengine_api/data/initialise.sql +++ /dev/null @@ -1,194 +0,0 @@ -CREATE TABLE IF NOT EXISTS household ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(255) NOT NULL, - household_json JSON NOT NULL, - household_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS computed_household ( - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - computed_household_json JSON NOT NULL, - status VARCHAR(32), - PRIMARY KEY (household_id, policy_id, country_id) -); - -CREATE TABLE IF NOT EXISTS policy ( - id INTEGER AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - policy_json JSON NOT NULL, - policy_hash VARCHAR(255) NOT NULL, - PRIMARY KEY (id, country_id, policy_hash) -); - -CREATE TABLE IF NOT EXISTS economy ( - economy_id INTEGER PRIMARY KEY AUTO_INCREMENT, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32), - time_period VARCHAR(32), - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - economy_json JSON, - status VARCHAR(32) NOT NULL, - message VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - dataset VARCHAR(255) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON, - options_hash VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME, - end_time DATETIME, - execution_id VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS analysis ( - prompt_id INTEGER PRIMARY KEY AUTO_INCREMENT, - prompt LONGTEXT NOT NULL, - analysis LONGTEXT, - status VARCHAR(32) NOT NULL -); - --- The dataset row below was added while the table is in prod; --- we must allow NULL values for this column -CREATE TABLE IF NOT EXISTS user_policies ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - reform_id INTEGER NOT NULL, - reform_label VARCHAR(255), - baseline_id INTEGER NOT NULL, - baseline_label VARCHAR(255), - user_id VARCHAR(255) NOT NULL, - year VARCHAR(32) NOT NULL, - geography VARCHAR(255) NOT NULL, - dataset VARCHAR(255), - number_of_provisions INTEGER NOT NULL, - api_version VARCHAR(32) NOT NULL, - added_date BIGINT NOT NULL, - updated_date BIGINT NOT NULL, - budgetary_impact VARCHAR(255), - type VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS user_profiles ( - user_id INTEGER PRIMARY KEY AUTO_INCREMENT, - auth0_id VARCHAR(255) NOT NULL UNIQUE, - username VARCHAR(255) UNIQUE, - primary_country VARCHAR(3) NOT NULL, - user_since BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS tracers ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - tracer_output JSON NOT NULL -); - -CREATE TABLE IF NOT EXISTS simulations ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - -- VARCHAR(255) to accommodate both household IDs and geography codes - population_id VARCHAR(255) NOT NULL, - population_type VARCHAR(50) NOT NULL, - policy_id INT NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - simulation_spec_json JSON DEFAULT NULL, - simulation_spec_schema_version INT DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_outputs ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - simulation_1_id INT NOT NULL, - simulation_2_id INT DEFAULT NULL, - api_version VARCHAR(10) NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - year VARCHAR(255) DEFAULT '2025', - report_kind VARCHAR(64) DEFAULT NULL, - report_spec_json JSON DEFAULT NULL, - report_spec_schema_version INT DEFAULT NULL, - report_spec_status VARCHAR(32) DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_output_runs ( - id CHAR(36) PRIMARY KEY, - report_output_id INT NOT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - report_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - report_cache_version VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - requested_version_override VARCHAR(255) DEFAULT NULL, - resolved_dataset VARCHAR(255) DEFAULT NULL, - resolved_options_hash VARCHAR(255) DEFAULT NULL, - UNIQUE KEY report_output_run_sequence_idx (report_output_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS simulation_runs ( - id CHAR(36) PRIMARY KEY, - simulation_id INT NOT NULL, - report_output_run_id CHAR(36) DEFAULT NULL, - input_position TINYINT DEFAULT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - simulation_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - UNIQUE KEY simulation_run_sequence_idx (simulation_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS legacy_report_output_aliases ( - legacy_report_output_id INT PRIMARY KEY, - canonical_report_output_id INT NOT NULL -); diff --git a/policyengine_api/data/initialise_local.sql b/policyengine_api/data/initialise_local.sql deleted file mode 100644 index 53a37b4c8..000000000 --- a/policyengine_api/data/initialise_local.sql +++ /dev/null @@ -1,205 +0,0 @@ -DROP TABLE IF EXISTS household; -DROP TABLE IF EXISTS computed_household; -DROP TABLE IF EXISTS policy; -DROP TABLE IF EXISTS economy; -DROP TABLE IF EXISTS reform_impact; -DROP TABLE IF EXISTS analysis; -DROP TABLE IF EXISTS user_policies; -DROP TABLE IF EXISTS tracers; -DROP TABLE IF EXISTS report_output_runs; -DROP TABLE IF EXISTS simulation_runs; -DROP TABLE IF EXISTS legacy_report_output_aliases; - -CREATE TABLE IF NOT EXISTS household ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(255) NOT NULL, - household_json JSONB NOT NULL, - household_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS computed_household ( - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - computed_household_json JSONB NOT NULL, - status VARCHAR(32), - PRIMARY KEY (household_id, policy_id, country_id) -); - -CREATE TABLE IF NOT EXISTS policy ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - policy_json JSONB NOT NULL, - policy_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS economy ( - economy_id INTEGER PRIMARY KEY, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32), - time_period VARCHAR(32), - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - economy_json JSON, - status VARCHAR(32) NOT NULL, - message VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - dataset VARCHAR(255) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME NOT NULL, - end_time DATETIME, - execution_id VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS analysis ( - prompt_id INTEGER PRIMARY KEY, - prompt LONGTEXT NOT NULL, - analysis LONGTEXT, - status VARCHAR(32) NOT NULL -); - --- The dataset row below was added while the table is in prod; --- we must allow NULL values for this column -CREATE TABLE IF NOT EXISTS user_policies ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - reform_id INTEGER NOT NULL, - reform_label VARCHAR(255), - baseline_id INTEGER NOT NULL, - baseline_label VARCHAR(255), - user_id VARCHAR(255) NOT NULL, - year VARCHAR(32) NOT NULL, - geography VARCHAR(255) NOT NULL, - dataset VARCHAR(255), - number_of_provisions INTEGER NOT NULL, - api_version VARCHAR(32) NOT NULL, - added_date BIGINT NOT NULL, - updated_date BIGINT NOT NULL, - budgetary_impact VARCHAR(255), - type VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS user_profiles ( - user_id INTEGER PRIMARY KEY, - auth0_id VARCHAR(255) NOT NULL UNIQUE, - username VARCHAR(255) UNIQUE, - primary_country VARCHAR(3) NOT NULL, - user_since BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS tracers ( - id INTEGER PRIMARY KEY, - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - tracer_output JSON NOT NULL -); - -CREATE TABLE IF NOT EXISTS simulations ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - -- VARCHAR(255) to accommodate both household IDs and geography codes - population_id VARCHAR(255) NOT NULL, - population_type VARCHAR(50) NOT NULL, - policy_id INT NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - simulation_spec_json JSON DEFAULT NULL, - simulation_spec_schema_version INT DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_outputs ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - simulation_1_id INT NOT NULL, - simulation_2_id INT DEFAULT NULL, - api_version VARCHAR(10) NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - year VARCHAR(255) DEFAULT '2025', - report_kind VARCHAR(64) DEFAULT NULL, - report_spec_json JSON DEFAULT NULL, - report_spec_schema_version INT DEFAULT NULL, - report_spec_status VARCHAR(32) DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_output_runs ( - id CHAR(36) PRIMARY KEY, - report_output_id INT NOT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - report_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - report_cache_version VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - requested_version_override VARCHAR(255) DEFAULT NULL, - resolved_dataset VARCHAR(255) DEFAULT NULL, - resolved_options_hash VARCHAR(255) DEFAULT NULL, - UNIQUE (report_output_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS simulation_runs ( - id CHAR(36) PRIMARY KEY, - simulation_id INT NOT NULL, - report_output_run_id CHAR(36) DEFAULT NULL, - input_position TINYINT DEFAULT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - simulation_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - UNIQUE (simulation_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS legacy_report_output_aliases ( - legacy_report_output_id INT PRIMARY KEY, - canonical_report_output_id INT NOT NULL -); diff --git a/policyengine_api/data/local_database.py b/policyengine_api/data/local_database.py new file mode 100644 index 000000000..05e59ee74 --- /dev/null +++ b/policyengine_api/data/local_database.py @@ -0,0 +1,43 @@ +"""Temporary SQLAlchemy schema bootstrap for the local SQLite cache. + +The production ``policy`` table has an autoincrementing column inside a +composite primary key. SQLite only autoincrements an ``INTEGER PRIMARY KEY`` +column when it is the sole primary-key column, so the local cache has always +used ``policy.id`` as its database primary key. Keep that one physical-schema +exception explicit while deriving every other local table from the production +ORM metadata. +""" + +from __future__ import annotations + +from sqlalchemy import Column, Engine, Integer, JSON, MetaData, String, Table + +from policyengine_api.data.v1_models import Policy, V1Base + + +_sqlite_policy_metadata = MetaData() +Table( + Policy.__tablename__, + _sqlite_policy_metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("country_id", String(3), nullable=False), + Column("label", String(255)), + Column("api_version", String(10), nullable=False), + Column("policy_json", JSON, nullable=False), + Column("policy_hash", String(255), nullable=False), +) + + +def create_local_v1_schema(engine: Engine) -> None: + """Create the temporary local schema through SQLAlchemy DDL constructs.""" + + if engine.dialect.name != "sqlite": + raise ValueError("The local v1 schema is only defined for SQLite") + + production_tables = [ + table + for table in V1Base.metadata.sorted_tables + if table is not Policy.__table__ + ] + V1Base.metadata.create_all(engine, tables=production_tables) + _sqlite_policy_metadata.create_all(engine) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 85aa57bf1..679092ba9 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -4,17 +4,17 @@ import atexit import fcntl -import json import os -import sqlite3 from pathlib import Path from dotenv import load_dotenv from google.cloud.sql.connector import Connector, IPTypes -from sqlalchemy import Engine, create_engine +from sqlalchemy import Engine, create_engine, select from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, REPO +from policyengine_api.data.local_database import create_local_v1_schema +from policyengine_api.data.v1_models import Policy from policyengine_api.utils import hash_object @@ -54,51 +54,56 @@ def build_session_factory(engine: Engine) -> sessionmaker[Session]: return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) -def _initialize_local_database(database_path: Path) -> None: - initialization_sql = ( - REPO / "policyengine_api" / "data" / "initialise_local.sql" - ).read_text(encoding="utf-8") - with sqlite3.connect(database_path) as connection: - connection.executescript(initialization_sql) - connection.executemany( - """ - INSERT INTO policy - (id, country_id, label, api_version, policy_json, policy_hash) - VALUES (?, ?, ?, ?, ?, ?) - """, - [ - ( - policy_id, - country_id, - "Current law", - COUNTRY_PACKAGE_VERSIONS[country_id], - json.dumps({}), - hash_object({}), - ) - for policy_id, country_id in enumerate( - COUNTRY_PACKAGE_VERSIONS, start=1 - ) - ], +def _initialize_local_database(engine: Engine) -> None: + create_local_v1_schema(engine) + + current_law_policies = [ + Policy( + id=policy_id, + country_id=country_id, + label="Current law", + api_version=COUNTRY_PACKAGE_VERSIONS[country_id], + policy_json={}, + policy_hash=hash_object({}), + ) + for policy_id, country_id in enumerate(COUNTRY_PACKAGE_VERSIONS, start=1) + ] + policy_ids = [policy.id for policy in current_law_policies] + with build_session_factory(engine).begin() as session: + existing_policy_ids = set( + session.scalars(select(Policy.id).where(Policy.id.in_(policy_ids))) + ) + session.add_all( + policy + for policy in current_law_policies + if policy.id not in existing_policy_ids ) # TODO: Remove this local-database initialization pattern and replace the local # persistence path with a traditional cache. Application imports should # eventually neither create a database file nor bootstrap a schema. -def _ensure_local_database(database_path: Path = LOCAL_DATABASE_PATH) -> None: +def _ensure_local_database( + engine: Engine, + database_path: Path = LOCAL_DATABASE_PATH, +) -> None: lock_path = Path(f"{database_path}.init.lock") with lock_path.open("w") as lock_file: fcntl.flock(lock_file, fcntl.LOCK_EX) try: - if not database_path.exists(): - _initialize_local_database(database_path) + _initialize_local_database(engine) finally: fcntl.flock(lock_file, fcntl.LOCK_UN) def _build_local_engine() -> Engine: - _ensure_local_database() - return create_engine(f"sqlite+pysqlite:///{LOCAL_DATABASE_PATH}") + engine = create_engine(f"sqlite+pysqlite:///{LOCAL_DATABASE_PATH}") + try: + _ensure_local_database(engine) + except Exception: + engine.dispose() + raise + return engine def _database_password() -> str: diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 1570e9bb1..8da464ea5 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -7,8 +7,8 @@ os.environ.setdefault("FLASK_DEBUG", "1") -from policyengine_api.constants import REPO from policyengine_api.data import orm +from policyengine_api.data.local_database import create_local_v1_schema from policyengine_api.data.v1_models import V1Base @@ -19,14 +19,7 @@ def test_engine(): connect_args={"check_same_thread": False}, poolclass=StaticPool, ) - initialization_sql = ( - REPO / "policyengine_api" / "data" / "initialise_local.sql" - ).read_text(encoding="utf-8") - raw_connection = engine.raw_connection() - try: - raw_connection.executescript(initialization_sql) - finally: - raw_connection.close() + create_local_v1_schema(engine) yield engine engine.dispose() diff --git a/tests/unit/data/test_run_schema.py b/tests/unit/data/test_run_schema.py index fdbd2a78e..09d7740c5 100644 --- a/tests/unit/data/test_run_schema.py +++ b/tests/unit/data/test_run_schema.py @@ -1,9 +1,5 @@ -from pathlib import Path - from sqlalchemy import inspect -from policyengine_api.constants import REPO - def _column_names(test_engine, table_name: str) -> set[str]: return {column["name"] for column in inspect(test_engine).get_columns(table_name)} @@ -64,26 +60,3 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_engine): alias_columns = _column_names(test_engine, "legacy_report_output_aliases") assert {"legacy_report_output_id", "canonical_report_output_id"} == alias_columns - - -def test_stage_one_schema_is_defined_in_both_sql_initializers(): - sql_paths = [ - REPO / "policyengine_api" / "data" / "initialise.sql", - REPO / "policyengine_api" / "data" / "initialise_local.sql", - ] - - required_snippets = [ - "CREATE TABLE IF NOT EXISTS report_output_runs", - "CREATE TABLE IF NOT EXISTS simulation_runs", - "CREATE TABLE IF NOT EXISTS legacy_report_output_aliases", - "report_spec_json", - "report_spec_status", - "simulation_spec_json", - "active_run_id", - "latest_successful_run_id", - ] - - for sql_path in sql_paths: - sql_text = Path(sql_path).read_text() - for snippet in required_snippets: - assert snippet in sql_text, f"{snippet} missing from {sql_path.name}" diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index b61bd5717..015e79663 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -1,7 +1,7 @@ from unittest.mock import Mock import sqlalchemy -from sqlalchemy import create_engine, func, select +from sqlalchemy import create_engine, func, inspect, select from sqlalchemy.orm import Session import policyengine_api.data.orm as orm @@ -81,13 +81,25 @@ def test_database_password_can_be_loaded_from_file(monkeypatch, tmp_path): assert orm._database_password() == "file-password" -def test_local_initializer_bootstraps_schema_and_current_law_rows(tmp_path): - database_path = tmp_path / "local.db" +def test_local_schema_preserves_the_documented_sqlite_policy_key_exception(): + from policyengine_api.data.local_database import create_local_v1_schema - orm._initialize_local_database(database_path) + engine = create_engine("sqlite+pysqlite:///:memory:") + try: + create_local_v1_schema(engine) + policy_key = inspect(engine).get_pk_constraint("policy") + assert policy_key["constrained_columns"] == ["id"] + finally: + engine.dispose() + + +def test_local_initializer_bootstraps_schema_and_current_law_rows(tmp_path): + database_path = tmp_path / "local.db" engine = create_engine(f"sqlite+pysqlite:///{database_path}") + try: + orm._initialize_local_database(engine) with Session(engine) as session: assert session.scalar(select(func.count()).select_from(Policy)) > 0 assert session.scalars(select(Policy)).first().policy_json == {} @@ -95,6 +107,20 @@ def test_local_initializer_bootstraps_schema_and_current_law_rows(tmp_path): engine.dispose() +def test_local_initializer_is_idempotent(tmp_path): + engine = create_engine(f"sqlite+pysqlite:///{tmp_path / 'local.db'}") + try: + orm._initialize_local_database(engine) + orm._initialize_local_database(engine) + + with Session(engine) as session: + assert session.scalar(select(func.count()).select_from(Policy)) == len( + orm.COUNTRY_PACKAGE_VERSIONS + ) + finally: + engine.dispose() + + def test_close_v1_engines_disposes_pools_and_connectors(monkeypatch): engine = Mock() connector = Mock() From 82f77f942dceb94730b859839114bcabedb16f99 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 16:56:55 +0300 Subject: [PATCH 58/89] refactor: restore service-owned sessions for core resources --- policyengine_api/routes/household_routes.py | 36 +++---- policyengine_api/routes/policy_routes.py | 19 ++-- .../routes/user_profile_routes.py | 57 +++++------ policyengine_api/services/economy_service.py | 20 ++-- .../services/household_service.py | 58 ++++++++++- policyengine_api/services/policy_service.py | 46 +++++++-- policyengine_api/services/user_service.py | 62 +++++++++++- tests/fixtures/services/economy_service.py | 2 +- tests/fixtures/services/household_fixtures.py | 2 +- tests/fixtures/services/policy_service.py | 2 +- tests/fixtures/services/user_service.py | 2 +- .../python/test_household_routes.py | 31 +++--- .../test_direct_orm_policy_household.py | 25 ++--- tests/unit/services/test_direct_orm_users.py | 22 ++--- tests/unit/services/test_household_service.py | 24 +++-- tests/unit/services/test_policy_service.py | 47 ++++----- .../services/test_service_owned_sessions.py | 96 +++++++++++++++++++ tests/unit/services/test_user_service.py | 37 ++++--- 18 files changed, 389 insertions(+), 199 deletions(-) create mode 100644 tests/unit/services/test_service_owned_sessions.py diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 2f2f06426..26e63dcc7 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -2,7 +2,6 @@ from werkzeug.exceptions import NotFound, BadRequest import json -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdService from policyengine_api.utils.payload_validators import ( @@ -37,10 +36,8 @@ def get_household(country_id: str, household_id: int) -> Response: """ print(f"Got request for household {household_id} in country {country_id}") - sessions = get_v1_session_factory() - with sessions() as session: - household = household_service.get_household(session, country_id, household_id) - result = None if household is None else _serialize_household(household) + household = household_service.get_household(country_id, household_id) + result = None if household is None else _serialize_household(household) if result is None: raise NotFound(f"Household #{household_id} not found.") else: @@ -78,14 +75,12 @@ def post_household(country_id: str) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - with get_v1_session_factory().begin() as session: - household = household_service.create_household( - session, - country_id, - household_json, - label, - ) - household_id = household.id + household = household_service.create_household( + country_id, + household_json, + label, + ) + household_id = household.id return Response( json.dumps( @@ -125,23 +120,16 @@ def update_household(country_id: str, household_id: int) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - with get_v1_session_factory().begin() as session: - household = household_service.get_household( - session, - country_id, - household_id, - ) - if household is None: - raise NotFound(f"Household #{household_id} not found.") - + try: updated_household = household_service.update_household( - session, country_id, household_id, household_json, label, ) - updated_household_json = updated_household.household_json + except LookupError: + raise NotFound(f"Household #{household_id} not found.") from None + updated_household_json = updated_household.household_json return Response( json.dumps( { diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 2674d4b77..ee83c8f0b 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,7 +1,6 @@ from flask import Blueprint, Response, request import json -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Policy from policyengine_api.services.policy_service import PolicyService from werkzeug.exceptions import NotFound, BadRequest @@ -43,10 +42,8 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: # Specifically cast policy_id to an integer policy_id = int(policy_id) - sessions = get_v1_session_factory() - with sessions() as session: - policy = policy_service.get_policy(session, country_id, policy_id) - result = None if policy is None else _serialize_policy(policy) + policy = policy_service.get_policy(country_id, policy_id) + result = None if policy is None else _serialize_policy(policy) if result is None: raise NotFound(f"Policy #{policy_id} not found.") @@ -77,13 +74,11 @@ def set_policy(country_id: str) -> Response: label = payload.pop("label", None) policy_json = payload.pop("data", None) - with get_v1_session_factory().begin() as session: - policy_id, message, is_existing_policy = policy_service.set_policy( - session, - country_id, - label, - policy_json, - ) + policy_id, message, is_existing_policy = policy_service.set_policy( + country_id, + label, + policy_json, + ) response_body = dict( status="ok", diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index cb39b63c8..77175feff 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,5 +1,4 @@ from flask import Blueprint, Response, request -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import UserProfile from policyengine_api.utils.payload_validators import validate_country import json @@ -41,15 +40,13 @@ def set_user_profile(country_id: str) -> Response: username = payload.pop("username", None) user_since = payload.pop("user_since") - with get_v1_session_factory().begin() as session: - created, profile = user_service.create_profile( - session, - primary_country=country_id, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - result = _serialize_user_profile(profile, include_auth0_id=False) + created, profile = user_service.create_profile( + primary_country=country_id, + auth0_id=auth0_id, + username=username, + user_since=user_since, + ) + result = _serialize_user_profile(profile, include_auth0_id=False) response = dict( status="ok", @@ -72,21 +69,19 @@ def get_user_profile(country_id: str) -> Response: if (auth0_id is None) and (user_id is None): raise BadRequest("auth0_id or user_id must be provided") - sessions = get_v1_session_factory() - with sessions() as session: - profile = ( - user_service.get_profile(session, user_id=user_id) - if auth0_id is None - else user_service.get_profile(session, auth0_id=auth0_id) - ) - readable_row = ( - None - if profile is None - else _serialize_user_profile( - profile, - include_auth0_id=auth0_id is not None, - ) + profile = ( + user_service.get_profile(user_id=user_id) + if auth0_id is None + else user_service.get_profile(auth0_id=auth0_id) + ) + readable_row = ( + None + if profile is None + else _serialize_user_profile( + profile, + include_auth0_id=auth0_id is not None, ) + ) if readable_row is None: raise NotFound("No such user") @@ -128,14 +123,12 @@ def update_user_profile(country_id: str) -> Response: if user_id is None: raise BadRequest("Payload must include user_id") - with get_v1_session_factory().begin() as session: - updated = user_service.update_profile( - session, - user_id=user_id, - primary_country=primary_country, - username=username, - user_since=user_since, - ) + updated = user_service.update_profile( + user_id=user_id, + primary_country=primary_country, + username=username, + user_since=user_since, + ) if not updated: raise NotFound("No such user id") diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index d5ab0a467..d0b2af3c9 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -243,18 +243,14 @@ def _get_policy_jsons( baseline_policy_id: int, reform_policy_id: int, ) -> tuple[dict | None, dict | None]: - sessions = get_v1_session_factory() - with sessions() as session: - baseline = policy_service.get_policy_json( - session, - country_id, - baseline_policy_id, - ) - reform = policy_service.get_policy_json( - session, - country_id, - reform_policy_id, - ) + baseline = policy_service.get_policy_json( + country_id, + baseline_policy_id, + ) + reform = policy_service.get_policy_json( + country_id, + reform_policy_id, + ) return baseline, reform @staticmethod diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 6363224b8..441b9b078 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,19 +1,29 @@ from __future__ import annotations from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Household from policyengine_api.utils import hash_object class HouseholdService: - """Household operations performed through a caller-owned ORM Session.""" + """Household operations with service-owned ORM transaction boundaries.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() def get_household( self, - session: Session, country_id: str, household_id: int, ) -> Household | None: @@ -21,6 +31,15 @@ def get_household( raise Exception( f"Invalid household ID: {household_id}. Must be a positive integer." ) + with self._sessions() as session: + return self._get_household(session, country_id, household_id) + + @staticmethod + def _get_household( + session: Session, + country_id: str, + household_id: int, + ) -> Household | None: return session.scalar( select(Household).where( Household.country_id == country_id, @@ -30,6 +49,20 @@ def get_household( def create_household( self, + country_id: str, + household_json: dict, + label: str | None, + ) -> Household: + with self._sessions.begin() as session: + return self._create_household( + session, + country_id, + household_json, + label, + ) + + @staticmethod + def _create_household( session: Session, country_id: str, household_json: dict, @@ -48,13 +81,30 @@ def create_household( def update_household( self, + country_id: str, + household_id: int, + household_json: dict, + label: str | None, + ) -> Household: + with self._sessions.begin() as session: + return self._update_household( + session, + country_id, + household_id, + household_json, + label, + ) + + @classmethod + def _update_household( + cls, session: Session, country_id: str, household_id: int, household_json: dict, label: str | None, ) -> Household: - household = self.get_household(session, country_id, household_id) + household = cls._get_household(session, country_id, household_id) if household is None: raise LookupError( f"Household #{household_id} not found for country {country_id}." diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index 96ffc40c2..b2780b135 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -3,15 +3,26 @@ from typing import Any from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Policy from policyengine_api.utils import hash_object class PolicyService: - """Policy operations performed through a caller-owned ORM Session.""" + """Policy operations with service-owned ORM transaction boundaries.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() @staticmethod def _validate_policy_id(policy_id: int) -> None: @@ -22,13 +33,21 @@ def _validate_policy_id(policy_id: int) -> None: def get_policy( self, - session: Session, country_id: str, policy_id: int, ) -> Policy | None: self._validate_policy_id(policy_id) if not country_id: raise ValueError("country_id cannot be empty or None") + with self._sessions() as session: + return self._get_policy(session, country_id, policy_id) + + @staticmethod + def _get_policy( + session: Session, + country_id: str, + policy_id: int, + ) -> Policy | None: return session.scalar( select(Policy).where( Policy.country_id == country_id, @@ -38,16 +57,14 @@ def get_policy( def get_policy_json( self, - session: Session, country_id: str, policy_id: int, ) -> Any | None: - policy = self.get_policy(session, country_id, policy_id) + policy = self.get_policy(country_id, policy_id) return None if policy is None else policy.policy_json def set_policy( self, - session: Session, country_id: str, label: str | None, policy_json: dict, @@ -57,6 +74,23 @@ def set_policy( raise ValueError(f"Invalid country_id: {country_id}") policy_hash = hash_object(policy_json) + with self._sessions.begin() as session: + return self._set_policy( + session, + country_id, + label, + policy_json, + policy_hash, + ) + + def _set_policy( + self, + session: Session, + country_id: str, + label: str | None, + policy_json: dict, + policy_hash: str, + ) -> tuple[int, str, bool]: existing = self._get_unique_policy_with_label( session, country_id, diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index c011b7f84..d7c11193d 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,23 +1,51 @@ from __future__ import annotations from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import UserProfile class UserService: - """User-profile operations performed through a caller-owned ORM Session.""" + """User-profile operations with service-owned ORM transaction boundaries.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() def create_profile( self, + primary_country: str, + auth0_id: str, + username: str | None, + user_since: int, + ) -> tuple[bool, UserProfile]: + with self._sessions.begin() as session: + return self._create_profile( + session, + primary_country, + auth0_id, + username, + user_since, + ) + + @classmethod + def _create_profile( + cls, session: Session, primary_country: str, auth0_id: str, username: str | None, user_since: int, ) -> tuple[bool, UserProfile]: - existing = self.get_profile(session, auth0_id=auth0_id) + existing = cls._get_profile(session, auth0_id=auth0_id) if existing is not None: return False, existing profile = UserProfile( @@ -32,12 +60,20 @@ def create_profile( def get_profile( self, - session: Session, auth0_id: str | None = None, user_id: int | str | None = None, ) -> UserProfile | None: if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") + with self._sessions() as session: + return self._get_profile(session, auth0_id, user_id) + + @staticmethod + def _get_profile( + session: Session, + auth0_id: str | None = None, + user_id: int | str | None = None, + ) -> UserProfile | None: condition = ( UserProfile.user_id == user_id if user_id is not None @@ -47,7 +83,6 @@ def get_profile( def update_profile( self, - session: Session, user_id: int, primary_country: str | None, username: str | None, @@ -55,6 +90,23 @@ def update_profile( ) -> UserProfile | None: if user_id is None: raise ValueError("you must specify either auth0_id or user_id") + with self._sessions.begin() as session: + return self._update_profile( + session, + user_id, + primary_country, + username, + user_since, + ) + + @staticmethod + def _update_profile( + session: Session, + user_id: int, + primary_country: str | None, + username: str | None, + user_since: int | None, + ) -> UserProfile | None: profile = session.get(UserProfile, user_id) if profile is None: return None diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index e38d7dd47..e0a264644 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -98,7 +98,7 @@ def mock_policyengine_version(): def mock_policy_service(): """Mock PolicyService with get_policy_json method.""" mock_service = MagicMock() - mock_service.get_policy_json.side_effect = lambda session, country_id, policy_id: ( + mock_service.get_policy_json.side_effect = lambda country_id, policy_id: ( MOCK_REFORM_POLICY_JSON if policy_id == MOCK_POLICY_ID else MOCK_BASELINE_POLICY_JSON diff --git a/tests/fixtures/services/household_fixtures.py b/tests/fixtures/services/household_fixtures.py index d785cd07b..23602dbdf 100644 --- a/tests/fixtures/services/household_fixtures.py +++ b/tests/fixtures/services/household_fixtures.py @@ -41,5 +41,5 @@ def existing_household_record(orm_session): api_version=valid_db_row["api_version"], ) orm_session.add(household) - orm_session.flush() + orm_session.commit() return household diff --git a/tests/fixtures/services/policy_service.py b/tests/fixtures/services/policy_service.py index 7cf164ba4..9b655232c 100644 --- a/tests/fixtures/services/policy_service.py +++ b/tests/fixtures/services/policy_service.py @@ -41,5 +41,5 @@ def existing_policy_record(orm_session): api_version=valid_policy_data["api_version"], ) orm_session.add(policy) - orm_session.flush() + orm_session.commit() return policy diff --git a/tests/fixtures/services/user_service.py b/tests/fixtures/services/user_service.py index 164574dc8..1520cad85 100644 --- a/tests/fixtures/services/user_service.py +++ b/tests/fixtures/services/user_service.py @@ -22,5 +22,5 @@ def existing_user_profile(orm_session): user_since=valid_user_record["user_since"], ) orm_session.add(profile) - orm_session.flush() + orm_session.commit() return profile diff --git a/tests/to_refactor/python/test_household_routes.py b/tests/to_refactor/python/test_household_routes.py index dcf394195..fae4a1cad 100644 --- a/tests/to_refactor/python/test_household_routes.py +++ b/tests/to_refactor/python/test_household_routes.py @@ -1,5 +1,5 @@ import json -from unittest.mock import ANY, patch +from unittest.mock import patch from policyengine_api.data.v1_models import Household from tests.to_refactor.fixtures.to_refactor_household_fixtures import ( @@ -124,7 +124,6 @@ def test_update_household_success(self, rest_client, mock_database): assert data["result"]["household_id"] == 1 assert data["result"]["household_json"] == updated_data["data"] mock_database.update_household.assert_called_once_with( - ANY, "us", 1, updated_household, @@ -133,7 +132,7 @@ def test_update_household_success(self, rest_client, mock_database): def test_update_nonexistent_household(self, rest_client, mock_database): """Test updating a non-existent household.""" - mock_database.get_household.return_value = None + mock_database.update_household.side_effect = LookupError("No household") response = rest_client.put( "/us/household/999", @@ -203,22 +202,16 @@ def test_put_household_service_error(self, mock_update, rest_client): """Test PUT endpoint when service raises an error.""" mock_update.side_effect = Exception("Failed to update household") - # First mock the get_household call that checks existence - with patch( - "policyengine_api.services.household_service.HouseholdService.get_household" - ) as mock_get: - mock_get.return_value = {"id": 1} # Simulate existing household - - response = rest_client.put( - "/us/household/1", - json={"data": {"valid": "payload"}}, - content_type="application/json", - ) - data = json.loads(response.data) - - assert response.status_code == 500 - assert data["status"] == "error" - assert "Failed to update household" in data["message"] + response = rest_client.put( + "/us/household/1", + json={"data": {"valid": "payload"}}, + content_type="application/json", + ) + data = json.loads(response.data) + + assert response.status_code == 500 + assert data["status"] == "error" + assert "Failed to update household" in data["message"] def test_missing_json_body(self, rest_client): """Test endpoints when JSON body is missing.""" diff --git a/tests/unit/services/test_direct_orm_policy_household.py b/tests/unit/services/test_direct_orm_policy_household.py index 4a1a1b640..c457b44c3 100644 --- a/tests/unit/services/test_direct_orm_policy_household.py +++ b/tests/unit/services/test_direct_orm_policy_household.py @@ -5,21 +5,22 @@ from policyengine_api.services.policy_service import PolicyService -def test_policy_service_reads_and_writes_mapped_models(orm_session, monkeypatch): +def test_policy_service_reads_and_writes_mapped_models( + orm_session_factory, + monkeypatch, +): monkeypatch.setattr( "policyengine_api.services.policy_service.hash_object", lambda value: "policy-hash", ) - service = PolicyService() + service = PolicyService(orm_session_factory) policy_id, message, existed = service.set_policy( - orm_session, "us", "Direct ORM policy", {"gov.example.rate": {"2026": 0.2}}, ) - orm_session.commit() - policy = service.get_policy(orm_session, "us", policy_id) + policy = service.get_policy("us", policy_id) assert isinstance(policy, Policy) assert policy.policy_json == {"gov.example.rate": {"2026": 0.2}} @@ -28,30 +29,30 @@ def test_policy_service_reads_and_writes_mapped_models(orm_session, monkeypatch) def test_household_service_reads_updates_and_writes_mapped_models( - orm_session, + orm_session_factory, monkeypatch, ): monkeypatch.setattr( "policyengine_api.services.household_service.hash_object", lambda value: "household-hash", ) - service = HouseholdService() + service = HouseholdService(orm_session_factory) payload = {"people": {"you": {"age": {"2026": 40}}}} household = service.create_household( - orm_session, "us", payload, "Direct ORM household", ) - orm_session.commit() - stored = orm_session.scalar(select(Household).where(Household.id == household.id)) + with orm_session_factory() as session: + stored = session.scalar( + select(Household).where(Household.id == household.id) + ) - assert household is stored + assert household.id == stored.id assert stored.household_json == payload updated = service.update_household( - orm_session, "us", stored.id, {"people": {"you": {"age": {"2026": 41}}}}, diff --git a/tests/unit/services/test_direct_orm_users.py b/tests/unit/services/test_direct_orm_users.py index f44794ace..d06282d81 100644 --- a/tests/unit/services/test_direct_orm_users.py +++ b/tests/unit/services/test_direct_orm_users.py @@ -2,18 +2,16 @@ from policyengine_api.services.user_service import UserService -def test_user_service_reads_and_writes_mapped_profiles(orm_session): - service = UserService() +def test_user_service_reads_and_writes_mapped_profiles(orm_session_factory): + service = UserService(orm_session_factory) created, profile = service.create_profile( - orm_session, primary_country="us", auth0_id="auth0|direct", username="direct-user", user_since=123, ) duplicate_created, duplicate = service.create_profile( - orm_session, primary_country="us", auth0_id="auth0|direct", username="ignored", @@ -23,14 +21,13 @@ def test_user_service_reads_and_writes_mapped_profiles(orm_session): assert created is True assert duplicate_created is False assert isinstance(profile, UserProfile) - assert duplicate is profile - assert service.get_profile(orm_session, auth0_id="auth0|direct") is profile + assert duplicate.user_id == profile.user_id + assert service.get_profile(auth0_id="auth0|direct").user_id == profile.user_id -def test_user_service_updates_the_mapped_profile(orm_session): - service = UserService() +def test_user_service_updates_the_mapped_profile(orm_session_factory): + service = UserService(orm_session_factory) _, profile = service.create_profile( - orm_session, primary_country="us", auth0_id="auth0|update", username=None, @@ -38,13 +35,12 @@ def test_user_service_updates_the_mapped_profile(orm_session): ) updated = service.update_profile( - orm_session, user_id=profile.user_id, primary_country="uk", username="updated-user", user_since=456, ) - assert updated is profile - assert profile.primary_country == "uk" - assert profile.username == "updated-user" + assert updated.user_id == profile.user_id + assert updated.primary_country == "uk" + assert updated.username == "updated-user" diff --git a/tests/unit/services/test_household_service.py b/tests/unit/services/test_household_service.py index c880fe650..8f676480d 100644 --- a/tests/unit/services/test_household_service.py +++ b/tests/unit/services/test_household_service.py @@ -11,12 +11,13 @@ pytest_plugins = ["tests.fixtures.services.household_fixtures"] -service = HouseholdService() +@pytest.fixture +def service(orm_session_factory): + return HouseholdService(orm_session_factory) -def test_get_household_returns_mapped_entity(orm_session, existing_household_record): +def test_get_household_returns_mapped_entity(service, existing_household_record): household = service.get_household( - orm_session, valid_db_row["country_id"], valid_db_row["id"], ) @@ -25,24 +26,23 @@ def test_get_household_returns_mapped_entity(orm_session, existing_household_rec assert household.household_json == valid_request_body["data"] -def test_get_household_returns_none_for_missing_entity(orm_session): - assert service.get_household(orm_session, "us", 999) is None +def test_get_household_returns_none_for_missing_entity(service): + assert service.get_household("us", 999) is None @pytest.mark.parametrize("household_id", ["invalid", -1]) -def test_get_household_rejects_invalid_id(orm_session, household_id): +def test_get_household_rejects_invalid_id(service, household_id): with pytest.raises(Exception, match="Invalid household ID"): - service.get_household(orm_session, "us", household_id) + service.get_household("us", household_id) -def test_create_household_adds_mapped_entity(orm_session, monkeypatch): +def test_create_household_adds_mapped_entity(service, monkeypatch): monkeypatch.setattr( "policyengine_api.services.household_service.hash_object", lambda value: "some-hash", ) household = service.create_household( - orm_session, "us", valid_request_body["data"], valid_request_body["label"], @@ -54,7 +54,7 @@ def test_create_household_adds_mapped_entity(orm_session, monkeypatch): def test_update_household_mutates_mapped_entity( - orm_session, + service, existing_household_record, monkeypatch, ): @@ -64,7 +64,6 @@ def test_update_household_mutates_mapped_entity( ) household = service.update_household( - orm_session, "us", valid_db_row["id"], {"people": {"person1": {"age": 31}}}, @@ -77,12 +76,11 @@ def test_update_household_mutates_mapped_entity( def test_update_household_rejects_missing_or_cross_country_entity( - orm_session, + service, existing_household_record, ): with pytest.raises(LookupError): service.update_household( - orm_session, "uk", valid_db_row["id"], {}, diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index 1e2d48553..ca7a24951 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -10,12 +10,13 @@ pytest_plugins = ["tests.fixtures.services.policy_service"] -service = PolicyService() +@pytest.fixture +def service(orm_session_factory): + return PolicyService(orm_session_factory) -def test_get_policy_returns_mapped_entity(orm_session, existing_policy_record): +def test_get_policy_returns_mapped_entity(service, existing_policy_record): policy = service.get_policy( - orm_session, valid_policy_data["country_id"], valid_policy_data["id"], ) @@ -27,44 +28,43 @@ def test_get_policy_returns_mapped_entity(orm_session, existing_policy_record): } -def test_get_policy_returns_none_for_missing_entity(orm_session): - assert service.get_policy(orm_session, "us", 999) is None +def test_get_policy_returns_none_for_missing_entity(service): + assert service.get_policy("us", 999) is None @pytest.mark.parametrize("policy_id", ["invalid", -1]) -def test_get_policy_rejects_invalid_id(orm_session, policy_id): +def test_get_policy_rejects_invalid_id(service, policy_id): with pytest.raises(Exception, match="Invalid policy ID"): - service.get_policy(orm_session, "us", policy_id) + service.get_policy("us", policy_id) @pytest.mark.parametrize("country_id", ["", None]) -def test_get_policy_rejects_empty_country(orm_session, country_id): +def test_get_policy_rejects_empty_country(service, country_id): with pytest.raises(ValueError, match="country_id cannot be empty or None"): - service.get_policy(orm_session, country_id, 1) + service.get_policy(country_id, 1) -def test_get_policy_json_returns_python_object(orm_session, existing_policy_record): - result = service.get_policy_json(orm_session, "us", valid_policy_data["id"]) +def test_get_policy_json_returns_python_object(service, existing_policy_record): + result = service.get_policy_json("us", valid_policy_data["id"]) assert result == { "gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433} } -def test_set_policy_adds_mapped_entity(orm_session, monkeypatch): +def test_set_policy_adds_mapped_entity(service, monkeypatch): monkeypatch.setattr( "policyengine_api.services.policy_service.hash_object", lambda value: "new-hash", ) policy_id, message, exists = service.set_policy( - orm_session, "US", "New policy", {"parameter": 1}, ) - policy = service.get_policy(orm_session, "us", policy_id) + policy = service.get_policy("us", policy_id) assert policy.policy_json == {"parameter": 1} assert policy.api_version == COUNTRY_PACKAGE_VERSIONS["us"] assert message == "Policy created" @@ -72,7 +72,7 @@ def test_set_policy_adds_mapped_entity(orm_session, monkeypatch): def test_set_policy_returns_existing_mapped_entity( - orm_session, + service, existing_policy_record, monkeypatch, ): @@ -82,7 +82,6 @@ def test_set_policy_returns_existing_mapped_entity( ) policy_id, message, exists = service.set_policy( - orm_session, "us", None, {}, @@ -93,21 +92,25 @@ def test_set_policy_returns_existing_mapped_entity( assert exists is True -def test_set_policy_rejects_invalid_country(orm_session): +def test_set_policy_rejects_invalid_country(service): with pytest.raises(ValueError, match="Invalid country_id: xx"): - service.set_policy(orm_session, "xx", "Policy", {}) + service.set_policy("xx", "Policy", {}) -def test_set_policy_propagates_flush_failure(orm_session, monkeypatch): +def test_set_policy_propagates_flush_failure( + service, + orm_session_factory, + monkeypatch, +): monkeypatch.setattr( "policyengine_api.services.policy_service.hash_object", lambda value: "new-hash", ) monkeypatch.setattr( - orm_session, + orm_session_factory.class_, "flush", - lambda: (_ for _ in ()).throw(SQLAlchemyError("insert failed")), + lambda self: (_ for _ in ()).throw(SQLAlchemyError("insert failed")), ) with pytest.raises(SQLAlchemyError, match="insert failed"): - service.set_policy(orm_session, "us", "Policy", {}) + service.set_policy("us", "Policy", {}) diff --git a/tests/unit/services/test_service_owned_sessions.py b/tests/unit/services/test_service_owned_sessions.py new file mode 100644 index 000000000..d5dbd27d9 --- /dev/null +++ b/tests/unit/services/test_service_owned_sessions.py @@ -0,0 +1,96 @@ +import inspect +from pathlib import Path + +import pytest +from sqlalchemy import func, select + +from policyengine_api.data.v1_models import Household +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_service import UserService + + +ROUTE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "routes" + + +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + ( + HouseholdService, + ("get_household", "create_household", "update_household"), + ), + (PolicyService, ("get_policy", "get_policy_json", "set_policy")), + (UserService, ("get_profile", "create_profile", "update_profile")), + ], +) +def test_core_service_public_methods_do_not_accept_sessions( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +@pytest.mark.parametrize( + "module_name", + ["household_routes.py", "policy_routes.py", "user_profile_routes.py"], +) +def test_core_routes_do_not_manage_sessions(module_name): + source = (ROUTE_ROOT / module_name).read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_core_services_commit_writes_and_return_generated_ids( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + service = HouseholdService(orm_session_factory) + + household = service.create_household( + "us", + {"people": {"you": {"age": {"2026": 40}}}}, + "Service-owned transaction", + ) + + assert household.id is not None + with orm_session_factory() as session: + stored = session.get(Household, household.id) + assert stored is not None + assert stored.household_json == { + "people": {"you": {"age": {"2026": 40}}} + } + + +def test_core_services_roll_back_failed_writes( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + session_type = orm_session_factory.class_ + original_flush = session_type.flush + + def fail_after_flush(session, *args, **kwargs): + original_flush(session, *args, **kwargs) + raise RuntimeError("forced failure") + + monkeypatch.setattr(session_type, "flush", fail_after_flush) + service = HouseholdService(orm_session_factory) + + with pytest.raises(RuntimeError, match="forced failure"): + service.create_household("us", {}, "Rolled back") + + monkeypatch.setattr(session_type, "flush", original_flush) + with orm_session_factory() as session: + count = session.scalar(select(func.count()).select_from(Household)) + assert count == 0 diff --git a/tests/unit/services/test_user_service.py b/tests/unit/services/test_user_service.py index 0736f75eb..3b515b68e 100644 --- a/tests/unit/services/test_user_service.py +++ b/tests/unit/services/test_user_service.py @@ -7,51 +7,47 @@ pytest_plugins = ["tests.fixtures.services.user_service"] -service = UserService() +@pytest.fixture +def service(orm_session_factory): + return UserService(orm_session_factory) -def test_get_profile_requires_an_identifier(orm_session): +def test_get_profile_requires_an_identifier(service): with pytest.raises( ValueError, match="you must specify either auth0_id or user_id", ): - service.get_profile(orm_session) + service.get_profile() -def test_get_profile_returns_none_for_unknown_auth0_id(orm_session): - assert service.get_profile(orm_session, auth0_id="missing") is None +def test_get_profile_returns_none_for_unknown_auth0_id(service): + assert service.get_profile(auth0_id="missing") is None def test_get_profile_returns_mapped_entity_by_either_identifier( - orm_session, + service, existing_user_profile, ): by_auth0 = service.get_profile( - orm_session, auth0_id=valid_user_record["auth0_id"], ) by_id = service.get_profile( - orm_session, user_id=valid_user_record["user_id"], ) assert isinstance(by_auth0, UserProfile) - assert by_auth0 is by_id + assert by_auth0.user_id == by_id.user_id assert by_auth0.username == valid_user_record["username"] -def test_create_profile_returns_existing_entity_for_duplicate_auth0_id( - orm_session, -): +def test_create_profile_returns_existing_entity_for_duplicate_auth0_id(service): created, profile = service.create_profile( - orm_session, "us", "auth0|duplicate", "first", 1, ) duplicate_created, duplicate = service.create_profile( - orm_session, "uk", "auth0|duplicate", "second", @@ -60,20 +56,19 @@ def test_create_profile_returns_existing_entity_for_duplicate_auth0_id( assert created is True assert duplicate_created is False - assert duplicate is profile + assert duplicate.user_id == profile.user_id assert duplicate.username == "first" -def test_update_profile_returns_none_for_missing_entity(orm_session): - assert service.update_profile(orm_session, 999, "uk", "missing", 2) is None +def test_update_profile_returns_none_for_missing_entity(service): + assert service.update_profile(999, "uk", "missing", 2) is None def test_update_profile_only_changes_non_null_fields( - orm_session, + service, existing_user_profile, ): profile = service.update_profile( - orm_session, valid_user_record["user_id"], "uk", None, @@ -85,9 +80,9 @@ def test_update_profile_only_changes_non_null_fields( assert profile.user_since == valid_user_record["user_since"] + 1 -def test_update_profile_requires_user_id(orm_session): +def test_update_profile_requires_user_id(service): with pytest.raises( ValueError, match="you must specify either auth0_id or user_id", ): - service.update_profile(orm_session, None, "us", "name", 1) + service.update_profile(None, "us", "name", 1) From bbcd8981da50b14ad7708e8fcd8cb02998f5b218 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:00:02 +0300 Subject: [PATCH 59/89] refactor: move simulation transactions into the service --- policyengine_api/routes/simulation_routes.py | 103 +++----- .../services/report_output_service.py | 2 +- .../services/simulation_service.py | 141 +++++++---- tests/contract/test_v1_route_contracts.py | 26 +- .../routes/test_route_exception_handling.py | 13 +- .../services/test_report_output_service.py | 2 +- .../services/test_service_owned_sessions.py | 20 +- .../services/test_simulation_run_service.py | 2 +- .../unit/services/test_simulation_service.py | 223 +++++++++--------- .../services/test_simulation_spec_service.py | 4 +- tests/unit/test_stage5_routes.py | 11 +- 11 files changed, 274 insertions(+), 273 deletions(-) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index 436c55c00..9476d6d67 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -5,7 +5,6 @@ import jsonschema import pydantic -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Simulation from policyengine_api.services.simulation_service import SimulationService from policyengine_api.utils.payload_validators import validate_country @@ -70,49 +69,19 @@ def create_simulation(country_id: str) -> Response: raise BadRequest("policy_id must be an integer") try: - with get_v1_session_factory().begin() as session: - existing_simulation = simulation_service.find_existing_simulation( - session, - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - - if existing_simulation: - simulation = simulation_service.ensure_simulation_dual_write_state( - session, - existing_simulation.id, - country_id=country_id, - ) - result = _serialize_v1_simulation(simulation) - message = "Simulation already exists" - status_code = 200 - else: - simulation = simulation_service.create_simulation( - session, - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - result = _serialize_v1_simulation(simulation) - message = "Simulation created successfully" - status_code = 201 - - if existing_simulation: - # Simulation already exists, return it with 200 status - response_body = dict( - status="ok", - message=message, - result=result, - ) - - return Response( - json.dumps(response_body), - status=status_code, - mimetype="application/json", - ) + creation = simulation_service.get_or_create_simulation( + country_id=country_id, + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + ) + result = _serialize_v1_simulation(creation.simulation) + message = ( + "Simulation created successfully" + if creation.created + else "Simulation already exists" + ) + status_code = 201 if creation.created else 200 response_body = dict( status="ok", @@ -159,14 +128,8 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: if simulation_id <= 0: raise BadRequest("simulation_id must be a positive integer") - sessions = get_v1_session_factory() - with sessions() as session: - simulation = simulation_service.get_simulation( - session, - country_id, - simulation_id, - ) - result = None if simulation is None else _serialize_v1_simulation(simulation) + simulation = simulation_service.get_simulation(country_id, simulation_id) + result = None if simulation is None else _serialize_v1_simulation(simulation) if result is None: raise NotFound(f"Simulation #{simulation_id} not found.") @@ -221,28 +184,18 @@ def update_simulation(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - with get_v1_session_factory().begin() as session: - existing_simulation = simulation_service.get_simulation( - session, - country_id, - simulation_id, - ) - if existing_simulation is None: - raise NotFound(f"Simulation #{simulation_id} not found.") - - success = simulation_service.update_simulation( - session, - country_id=country_id, - simulation_id=simulation_id, - status=status, - output=output, - error_message=error_message, - ) - - if not success: - raise BadRequest("No fields to update") - - result = _serialize_v1_simulation(existing_simulation) + simulation = simulation_service.update_simulation( + country_id=country_id, + simulation_id=simulation_id, + status=status, + output=output, + error_message=error_message, + ) + + if simulation is None: + raise BadRequest("No fields to update") + + result = _serialize_v1_simulation(simulation) response_body = dict( status="ok", @@ -256,6 +209,8 @@ def update_simulation(country_id: str) -> Response: mimetype="application/json", ) + except LookupError: + raise NotFound(f"Simulation #{simulation_id} not found.") from None except HTTPException: # Let explicit client-error responses (BadRequest/NotFound/etc.) pass # through without being logged as "Unexpected error". diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 5c69c79ee..fc66a9a3b 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -98,7 +98,7 @@ def _get_linked_simulations( def get_simulation(simulation_id: int) -> Simulation | None: if bootstrap_dual_write_state: try: - return self.simulation_service.ensure_simulation_dual_write_state( + return self.simulation_service._ensure_simulation_dual_write_state( session, simulation_id, report_output.country_id, diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 050486f74..5e073440f 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,15 +1,33 @@ import json import uuid +from dataclasses import dataclass from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Simulation, SimulationRun +@dataclass(frozen=True) +class SimulationCreateResult: + simulation: Simulation + created: bool + + class SimulationService: - """Simulation operations performed through a caller-owned ORM Session.""" + """Simulation operations with service-owned ORM transaction boundaries.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() @staticmethod def _select_simulation( @@ -30,7 +48,7 @@ def _select_simulation( def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: return next((run.id for run in runs if run.status == "complete"), None) - def ensure_simulation_dual_write_state( + def _ensure_simulation_dual_write_state( self, session: Session, simulation_id: int, @@ -96,15 +114,17 @@ def ensure_simulation_dual_write_state( session.flush() return simulation - def find_existing_simulation( - self, + @staticmethod + def _find_existing_simulation( session: Session, country_id: str, population_id: str, population_type: str, policy_id: int, + *, + for_update: bool = False, ) -> Simulation | None: - return session.scalar( + statement = ( select(Simulation) .where( Simulation.country_id == country_id, @@ -114,8 +134,11 @@ def find_existing_simulation( ) .order_by(Simulation.id.desc()) ) + if for_update: + statement = statement.with_for_update() + return session.scalar(statement) - def create_simulation( + def _create_simulation( self, session: Session, country_id: str, @@ -123,37 +146,57 @@ def create_simulation( population_type: str, policy_id: int, ) -> Simulation: - simulation = session.scalar( - select(Simulation) - .where( - Simulation.country_id == country_id, - Simulation.population_id == population_id, - Simulation.population_type == population_type, - Simulation.policy_id == policy_id, - ) - .order_by(Simulation.id.desc()) - .with_for_update() + simulation = Simulation( + country_id=country_id, + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", ) - if simulation is None: - simulation = Simulation( - country_id=country_id, - api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - status="pending", - ) - session.add(simulation) - session.flush() - return self.ensure_simulation_dual_write_state( + session.add(simulation) + session.flush() + return self._ensure_simulation_dual_write_state( session, simulation.id, country_id, ) + def get_or_create_simulation( + self, + country_id: str, + population_id: str, + population_type: str, + policy_id: int, + ) -> SimulationCreateResult: + with self._sessions.begin() as session: + simulation = self._find_existing_simulation( + session, + country_id, + population_id, + population_type, + policy_id, + for_update=True, + ) + created = simulation is None + if simulation is None: + simulation = self._create_simulation( + session, + country_id, + population_id, + population_type, + policy_id, + ) + else: + simulation = self._ensure_simulation_dual_write_state( + session, + simulation.id, + country_id, + ) + return SimulationCreateResult(simulation=simulation, created=created) + def get_simulation( self, - session: Session, country_id: str, simulation_id: int, ) -> Simulation | None: @@ -161,17 +204,17 @@ def get_simulation( raise Exception( f"Invalid simulation ID: {simulation_id}. Must be a positive integer." ) - return self._select_simulation(session, simulation_id, country_id) + with self._sessions() as session: + return self._select_simulation(session, simulation_id, country_id) def update_simulation( self, - session: Session, country_id: str, simulation_id: int, status: str | None = None, output: dict | list | str | None = None, error_message: str | None = None, - ) -> bool: + ) -> Simulation | None: values = { key: value for key, value in { @@ -182,19 +225,23 @@ def update_simulation( if value is not None } if not values: - return False + return None if isinstance(values.get("output"), str): values["output"] = json.loads(values["output"]) - simulation = self._select_simulation( - session, - simulation_id, - country_id, - for_update=True, - ) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - for key, value in values.items(): - setattr(simulation, key, value) - simulation.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) - self.ensure_simulation_dual_write_state(session, simulation_id, country_id) - return True + with self._sessions.begin() as session: + simulation = self._select_simulation( + session, + simulation_id, + country_id, + for_update=True, + ) + if simulation is None: + raise LookupError(f"Simulation #{simulation_id} not found") + for key, value in values.items(): + setattr(simulation, key, value) + simulation.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) + return self._ensure_simulation_dual_write_state( + session, + simulation_id, + country_id, + ) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index d5b77f541..a2ccc4dc4 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -15,6 +15,7 @@ from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.simulation_routes import simulation_bp +from policyengine_api.services.simulation_service import SimulationCreateResult from tests.contract.clients import ( ASGIContractClient, ContractClient, @@ -317,20 +318,17 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", - return_value=None, - ) - ) - stack.enter_context( - patch( - "policyengine_api.routes.simulation_routes.simulation_service.create_simulation", - return_value=Simulation( - id=11, - country_id="us", - population_id="household-1", - population_type="household", - policy_id=22, - status="pending", + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", + return_value=SimulationCreateResult( + simulation=Simulation( + id=11, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=22, + status="pending", + ), + created=True, ), ) ) diff --git a/tests/unit/routes/test_route_exception_handling.py b/tests/unit/routes/test_route_exception_handling.py index 4a1848cb9..fe4dd8522 100644 --- a/tests/unit/routes/test_route_exception_handling.py +++ b/tests/unit/routes/test_route_exception_handling.py @@ -29,7 +29,7 @@ def _client_with(*blueprints): def test_simulation_create_runtime_error_becomes_500(): client = _client_with(simulation_bp) with patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", side_effect=RuntimeError("db went away"), ): response = client.post( @@ -46,7 +46,7 @@ def test_simulation_create_runtime_error_becomes_500(): def test_simulation_create_value_error_still_400(): client = _client_with(simulation_bp) with patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", side_effect=ValueError("bad input"), ): response = client.post( @@ -86,7 +86,7 @@ def test_report_create_value_error_still_400(): assert response.status_code == 400 -def test_simulation_patch_empty_body_returns_400(orm_session): +def test_simulation_patch_empty_body_returns_400(orm_session_factory): """Regression for issue #3449. PATCH /{country}/simulation with a body that only contains the @@ -95,14 +95,13 @@ def test_simulation_patch_empty_body_returns_400(orm_session): """ from policyengine_api.services.simulation_service import SimulationService - simulation_service = SimulationService() - created = simulation_service.create_simulation( - orm_session, + simulation_service = SimulationService(orm_session_factory) + created = simulation_service.get_or_create_simulation( country_id="us", population_id="household_patch_empty", population_type="household", policy_id=50, - ) + ).simulation client = _client_with(simulation_bp) response = client.patch("/us/simulation", json={"id": created.id}) diff --git a/tests/unit/services/test_report_output_service.py b/tests/unit/services/test_report_output_service.py index 51fa34f36..ac44cae5a 100644 --- a/tests/unit/services/test_report_output_service.py +++ b/tests/unit/services/test_report_output_service.py @@ -17,7 +17,7 @@ def create_simulation(orm_session, *, policy_id=1, population_id="household-1"): - return simulation_service.create_simulation( + return simulation_service._create_simulation( orm_session, country_id="us", population_id=population_id, diff --git a/tests/unit/services/test_service_owned_sessions.py b/tests/unit/services/test_service_owned_sessions.py index d5dbd27d9..e1edc3b70 100644 --- a/tests/unit/services/test_service_owned_sessions.py +++ b/tests/unit/services/test_service_owned_sessions.py @@ -7,6 +7,7 @@ from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdService from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.simulation_service import SimulationService from policyengine_api.services.user_service import UserService @@ -21,6 +22,14 @@ ("get_household", "create_household", "update_household"), ), (PolicyService, ("get_policy", "get_policy_json", "set_policy")), + ( + SimulationService, + ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", + ), + ), (UserService, ("get_profile", "create_profile", "update_profile")), ], ) @@ -36,7 +45,12 @@ def test_core_service_public_methods_do_not_accept_sessions( @pytest.mark.parametrize( "module_name", - ["household_routes.py", "policy_routes.py", "user_profile_routes.py"], + [ + "household_routes.py", + "policy_routes.py", + "simulation_routes.py", + "user_profile_routes.py", + ], ) def test_core_routes_do_not_manage_sessions(module_name): source = (ROUTE_ROOT / module_name).read_text(encoding="utf-8") @@ -64,9 +78,7 @@ def test_core_services_commit_writes_and_return_generated_ids( with orm_session_factory() as session: stored = session.get(Household, household.id) assert stored is not None - assert stored.household_json == { - "people": {"you": {"age": {"2026": 40}}} - } + assert stored.household_json == {"people": {"you": {"age": {"2026": 40}}}} def test_core_services_roll_back_failed_writes( diff --git a/tests/unit/services/test_simulation_run_service.py b/tests/unit/services/test_simulation_run_service.py index 275373f7e..43ba3b80b 100644 --- a/tests/unit/services/test_simulation_run_service.py +++ b/tests/unit/services/test_simulation_run_service.py @@ -10,7 +10,7 @@ def create_simulation(orm_session, population_id="household-1"): - return simulation_service.create_simulation( + return simulation_service._create_simulation( orm_session, "us", population_id, "household", 1 ) diff --git a/tests/unit/services/test_simulation_service.py b/tests/unit/services/test_simulation_service.py index 2902f297e..0f8a3e81a 100644 --- a/tests/unit/services/test_simulation_service.py +++ b/tests/unit/services/test_simulation_service.py @@ -1,188 +1,175 @@ +import inspect + import pytest +from sqlalchemy import func, select from policyengine_api.data.v1_models import Simulation, SimulationRun -from policyengine_api.services.simulation_service import SimulationService - - -service = SimulationService() - - -def test_finds_existing_simulation_without_api_version_matching(orm_session): - existing = Simulation( - country_id="us", - api_version="old-version", - population_id="household-1", - population_type="household", - policy_id=1, - status="pending", - ) - orm_session.add(existing) - orm_session.flush() +from policyengine_api.services.simulation_service import ( + SimulationCreateResult, + SimulationService, +) - result = service.find_existing_simulation( - orm_session, - country_id="us", - population_id="household-1", - population_type="household", - policy_id=1, - ) - assert result is existing +@pytest.fixture +def service(orm_session_factory): + return SimulationService(orm_session_factory) -def test_returns_none_when_simulation_does_not_exist(orm_session): - assert ( - service.find_existing_simulation( - orm_session, - country_id="uk", - population_id="missing", - population_type="household", - policy_id=999, - ) - is None - ) +def test_public_simulation_methods_do_not_accept_sessions(): + for method_name in ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", + ): + parameters = inspect.signature( + getattr(SimulationService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters -def test_creates_mapped_simulation_and_initial_run(orm_session): - simulation = service.create_simulation( - orm_session, +def test_get_or_create_builds_simulation_spec_and_initial_run(service): + result = service.get_or_create_simulation( country_id="us", population_id="household-1", population_type="household", policy_id=1, ) - assert isinstance(simulation, Simulation) - assert simulation.simulation_spec_json == { + assert isinstance(result, SimulationCreateResult) + assert result.created is True + assert isinstance(result.simulation, Simulation) + assert result.simulation.simulation_spec_json == { "country_id": "us", "population_id": "household-1", "population_type": "household", "policy_id": 1, } - assert simulation.simulation_spec_schema_version == 1 - run = orm_session.get(SimulationRun, simulation.active_run_id) - assert isinstance(run, SimulationRun) - assert run.status == "pending" - assert run.trigger_type == "initial" - assert run.simulation_spec_snapshot_json == simulation.simulation_spec_json + assert result.simulation.simulation_spec_schema_version == 1 + assert result.simulation.active_run_id is not None -def test_creation_reuses_existing_row_and_bootstraps_dual_write_state(orm_session): - existing = Simulation( - country_id="us", - api_version="old-version", - population_id="household-1", - population_type="household", - policy_id=7, - status="pending", - ) - orm_session.add(existing) - orm_session.flush() +def test_get_or_create_reuses_existing_row_and_repairs_dual_write_state( + service, + orm_session_factory, +): + with orm_session_factory.begin() as session: + existing = Simulation( + country_id="us", + api_version="old-version", + population_id="household-1", + population_type="household", + policy_id=7, + status="pending", + ) + session.add(existing) + session.flush() + simulation_id = existing.id - result = service.create_simulation( - orm_session, + result = service.get_or_create_simulation( country_id="us", population_id="household-1", population_type="household", policy_id=7, ) - assert result is existing - run = orm_session.get(SimulationRun, existing.active_run_id) - assert isinstance(run, SimulationRun) - assert run.simulation_id == existing.id + assert result.created is False + assert result.simulation.id == simulation_id + with orm_session_factory() as session: + stored = session.get(Simulation, simulation_id) + run = session.get(SimulationRun, stored.active_run_id) + assert isinstance(run, SimulationRun) + assert run.simulation_id == simulation_id -def test_caller_transaction_rolls_back_creation_on_dual_write_failure( - orm_session_factory, monkeypatch +def test_get_or_create_rolls_back_simulation_when_dual_write_fails( + service, + orm_session_factory, + monkeypatch, ): def fail_dual_write(*args, **kwargs): raise RuntimeError("dual write sync failed") - monkeypatch.setattr(service, "ensure_simulation_dual_write_state", fail_dual_write) + monkeypatch.setattr(service, "_ensure_simulation_dual_write_state", fail_dual_write) with pytest.raises(RuntimeError, match="dual write sync failed"): - with orm_session_factory.begin() as session: - service.create_simulation( - session, - country_id="us", - population_id="rollback", - population_type="household", - policy_id=8, - ) + service.get_or_create_simulation( + country_id="us", + population_id="rollback", + population_type="household", + policy_id=8, + ) with orm_session_factory() as session: - assert ( - service.find_existing_simulation( - session, - country_id="us", - population_id="rollback", - population_type="household", - policy_id=8, - ) - is None + count = session.scalar( + select(func.count()) + .select_from(Simulation) + .where(Simulation.population_id == "rollback") ) + assert count == 0 -def test_get_simulation_returns_model_scoped_to_country(orm_session): - simulation = service.create_simulation( - orm_session, "us", "household-1", "household", 1 - ) +def test_get_simulation_returns_model_scoped_to_country(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - assert service.get_simulation(orm_session, "us", simulation.id) is simulation - assert service.get_simulation(orm_session, "uk", simulation.id) is None + assert service.get_simulation("us", simulation.id).id == simulation.id + assert service.get_simulation("uk", simulation.id) is None @pytest.mark.parametrize("simulation_id", [-1, "1", None]) -def test_get_simulation_rejects_invalid_ids(orm_session, simulation_id): +def test_get_simulation_rejects_invalid_ids(service, simulation_id): with pytest.raises(Exception, match="Invalid simulation ID"): - service.get_simulation(orm_session, "us", simulation_id) + service.get_simulation("us", simulation_id) -def test_update_simulation_updates_model_and_run_with_python_json(orm_session): - simulation = service.create_simulation( - orm_session, "us", "household-1", "household", 1 - ) +def test_update_simulation_updates_model_and_run_with_python_json( + service, + orm_session_factory, +): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation updated = service.update_simulation( - orm_session, "us", simulation.id, status="complete", output={"result": 42}, ) - assert updated is True - assert simulation.output == {"result": 42} - assert simulation.active_run_id is None - run = orm_session.get(SimulationRun, simulation.latest_successful_run_id) - assert run.output == {"result": 42} - assert run.status == "complete" + assert isinstance(updated, Simulation) + assert updated.output == {"result": 42} + assert updated.active_run_id is None + with orm_session_factory() as session: + run = session.get(SimulationRun, updated.latest_successful_run_id) + assert run.output == {"result": 42} + assert run.status == "complete" -def test_update_simulation_accepts_json_only_at_existing_wire_boundary(orm_session): - simulation = service.create_simulation( - orm_session, "us", "household-1", "household", 1 - ) +def test_update_simulation_accepts_legacy_json_text_at_wire_boundary(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - service.update_simulation( - orm_session, + updated = service.update_simulation( "us", simulation.id, output='{"result": 42}', ) - assert simulation.output == {"result": 42} + assert updated.output == {"result": 42} -def test_update_simulation_without_values_is_a_noop(orm_session): - simulation = service.create_simulation( - orm_session, "us", "household-1", "household", 1 - ) +def test_update_simulation_without_values_is_a_noop(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - assert service.update_simulation(orm_session, "us", simulation.id) is False + assert service.update_simulation("us", simulation.id) is None -def test_update_missing_simulation_raises(orm_session): - with pytest.raises(ValueError, match="Simulation #999 not found"): - service.update_simulation(orm_session, "us", 999, status="complete") +def test_update_missing_simulation_raises(service): + with pytest.raises(LookupError, match="Simulation #999 not found"): + service.update_simulation("us", 999, status="complete") diff --git a/tests/unit/services/test_simulation_spec_service.py b/tests/unit/services/test_simulation_spec_service.py index 7e19a24e1..26f06d02f 100644 --- a/tests/unit/services/test_simulation_spec_service.py +++ b/tests/unit/services/test_simulation_spec_service.py @@ -12,7 +12,9 @@ def create_simulation(orm_session): - return simulation_service.create_simulation(orm_session, "us", "ca", "geography", 3) + return simulation_service._create_simulation( + orm_session, "us", "ca", "geography", 3 + ) def test_builds_spec_from_mapped_simulation(orm_session): diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_stage5_routes.py index e95fba89a..9b3e80a49 100644 --- a/tests/unit/test_stage5_routes.py +++ b/tests/unit/test_stage5_routes.py @@ -26,11 +26,12 @@ def create_test_client() -> Flask: def create_simulation(factory, *, population_id="household-1", policy_id=1): - with factory.begin() as session: - simulation = simulation_service.create_simulation( - session, "us", population_id, "household", policy_id - ) - return simulation.id + simulation = ( + SimulationService(factory) + .get_or_create_simulation("us", population_id, "household", policy_id) + .simulation + ) + return simulation.id def create_report(factory, simulation_id): From 4578d60a2d4a35c12b5ce09ccbcef2592f4e400d Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:03:52 +0300 Subject: [PATCH 60/89] refactor: move report transactions into the service --- .../routes/report_output_routes.py | 119 +++---- .../services/report_output_service.py | 181 +++++++++-- tests/contract/test_v1_route_contracts.py | 87 ++--- .../routes/test_route_exception_handling.py | 4 +- .../services/test_report_output_service.py | 301 ++++++++++-------- .../test_report_service_owned_sessions.py | 110 +++++++ tests/unit/test_stage5_routes.py | 29 +- 7 files changed, 530 insertions(+), 301 deletions(-) create mode 100644 tests/unit/services/test_report_service_owned_sessions.py diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index e5d845c32..2b2bc21e7 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -5,9 +5,11 @@ import jsonschema import pydantic -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.services.report_output_service import ( + ReportOutputService, + ReportOutputView, +) from policyengine_api.constants import CURRENT_YEAR -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ReportOutput from policyengine_api.utils.payload_validators import validate_country @@ -15,26 +17,22 @@ report_output_service = ReportOutputService() -def _serialize_v1_report_output( - session, report_output: ReportOutput, *, response_id: int | None = None -) -> dict: +def _serialize_v1_report_output(view: ReportOutputView) -> dict: """Project mapped report state onto the historical v1 response shape.""" + report_output = view.report_output result = { column.name: getattr(report_output, column.name) for column in ReportOutput.__table__.columns } - if response_id is not None: - result["id"] = response_id + if view.response_id is not None: + result["id"] = view.response_id if result.get("output") is not None and not isinstance(result["output"], str): result["output"] = json.dumps(result["output"]) - display_run = report_output_service.report_run_service.select_display_run( - session, report_output - ) - if display_run is not None: + if view.display_run is not None: for field in ("requested_at", "started_at", "finished_at"): result[field] = report_output_service.format_run_timestamp( - getattr(display_run, field) + getattr(view.display_run, field) ) return result @@ -74,34 +72,19 @@ def create_report_output(country_id: str) -> Response: raise BadRequest("year must be a string") try: - with get_v1_session_factory().begin() as session: - existing_report = report_output_service.find_existing_report_output( - session, - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - ) - if existing_report: - report_output = ( - report_output_service.ensure_report_output_dual_write_state( - session, existing_report.id, country_id - ) - ) - result = _serialize_v1_report_output(session, report_output) - message = "Report output already exists" - status_code = 200 - else: - report_output = report_output_service.create_report_output( - session, - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - ) - result = _serialize_v1_report_output(session, report_output) - message = "Report output created successfully" - status_code = 201 + creation = report_output_service.create_or_reuse_report_output( + country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + year=year, + ) + result = _serialize_v1_report_output(creation.view) + message = ( + "Report output created successfully" + if creation.created + else "Report output already exists" + ) + status_code = 201 if creation.created else 200 response_body = dict( status="ok", @@ -148,25 +131,10 @@ def get_report_output(country_id: str, report_id: int) -> Response: """ print(f"Getting report output {report_id} for country {country_id}") - with get_v1_session_factory().begin() as session: - requested_report = report_output_service.get_report_output( - session, country_id, report_id - ) - if requested_report is None: - raise NotFound(f"Report #{report_id} not found.") - if report_output_service.is_current_report_output(requested_report): - report_output = report_output_service.ensure_report_output_dual_write_state( - session, report_id, country_id - ) - response_id = None - else: - report_output = report_output_service.get_or_create_current_report_output( - session, requested_report - ) - response_id = report_id - result = _serialize_v1_report_output( - session, report_output, response_id=response_id - ) + view = report_output_service.resolve_report_output(country_id, report_id) + if view is None: + raise NotFound(f"Report #{report_id} not found.") + result = _serialize_v1_report_output(view) response_body = dict( status="ok", @@ -223,27 +191,16 @@ def update_report_output(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - with get_v1_session_factory().begin() as session: - # Do not synchronize before this mutation: doing so could overwrite - # the pending rerun that this PATCH is about to mark as running. - if not report_output_service.report_output_exists( - session, country_id, report_id - ): - raise NotFound(f"Report #{report_id} not found.") - success = report_output_service.update_report_output( - session, - country_id=country_id, - report_id=report_id, - status=status, - output=output, - error_message=error_message, - ) - if not success: - raise BadRequest("No fields to update") - updated_report = report_output_service.get_report_output( - session, country_id, report_id - ) - result = _serialize_v1_report_output(session, updated_report) + view = report_output_service.update_report_output( + country_id=country_id, + report_id=report_id, + status=status, + output=output, + error_message=error_message, + ) + if view is None: + raise BadRequest("No fields to update") + result = _serialize_v1_report_output(view) response_body = dict( status="ok", @@ -257,6 +214,8 @@ def update_report_output(country_id: str) -> Response: mimetype="application/json", ) + except LookupError: + raise NotFound(f"Report #{report_id} not found.") from None except HTTPException: # Let explicit client-error responses (BadRequest/NotFound/etc.) pass # through without being logged as "Unexpected error". diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index fc66a9a3b..28bdd3e9e 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,10 +1,12 @@ import json +from dataclasses import dataclass from datetime import datetime, timezone from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun, Simulation from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.report_spec_service import ( @@ -15,13 +17,34 @@ from policyengine_api.services.simulation_service import SimulationService +@dataclass(frozen=True) +class ReportOutputView: + report_output: ReportOutput + display_run: ReportOutputRun | None + response_id: int | None = None + + +@dataclass(frozen=True) +class ReportCreateResult: + view: ReportOutputView + created: bool + + class ReportOutputService: - """Report-output orchestration through one caller-owned ORM Session.""" + """Report-output orchestration with service-owned transactions.""" - def __init__(self): + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory self.report_spec_service = ReportSpecService() self.report_run_service = ReportRunService() - self.simulation_service = SimulationService() + self.simulation_service = SimulationService(session_factory) + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() @staticmethod def _utc_timestamp() -> datetime: @@ -326,7 +349,7 @@ def _sync_parent_pointers( latest_successful = runs_descending[0].id if runs_descending else None report_output.latest_successful_run_id = latest_successful - def ensure_report_output_dual_write_state( + def _ensure_report_output_dual_write_state( self, session: Session, report_output_id: int, @@ -390,7 +413,7 @@ def ensure_report_output_dual_write_state( session.flush() return report_output - def find_existing_report_output( + def _find_existing_report_output( self, session: Session, country_id: str, @@ -426,7 +449,7 @@ def _require_simulation( ) return simulation - def create_report_output( + def _create_report_output( self, session: Session, country_id: str, @@ -434,11 +457,11 @@ def create_report_output( simulation_2_id: int | None = None, year: str = "2025", ) -> ReportOutput: - existing = self.find_existing_report_output( + existing = self._find_existing_report_output( session, country_id, simulation_1_id, simulation_2_id, year ) if existing is not None: - return self.ensure_report_output_dual_write_state( + return self._ensure_report_output_dual_write_state( session, existing.id, country_id ) self._require_simulation(session, country_id, simulation_1_id) @@ -454,11 +477,11 @@ def create_report_output( ) session.add(report_output) session.flush() - return self.ensure_report_output_dual_write_state( + return self._ensure_report_output_dual_write_state( session, report_output.id, country_id ) - def get_report_output( + def _get_report_output( self, session: Session, country_id: str, report_output_id: int ) -> ReportOutput | None: if type(report_output_id) is not int or report_output_id < 0: @@ -469,15 +492,15 @@ def get_report_output( return self._select_report_output(session, report_output_id, country_id) @staticmethod - def is_current_report_output(report_output: ReportOutput) -> bool: + def _is_current_report_output(report_output: ReportOutput) -> bool: return report_output.api_version == get_report_output_cache_version( report_output.country_id ) - def get_or_create_current_report_output( + def _get_or_create_current_report_output( self, session: Session, report_output: ReportOutput ) -> ReportOutput: - existing = self.find_existing_report_output( + existing = self._find_existing_report_output( session, report_output.country_id, report_output.simulation_1_id, @@ -485,10 +508,10 @@ def get_or_create_current_report_output( report_output.year, ) if existing is not None: - return self.ensure_report_output_dual_write_state( + return self._ensure_report_output_dual_write_state( session, existing.id, report_output.country_id ) - return self.create_report_output( + return self._create_report_output( session, report_output.country_id, report_output.simulation_1_id, @@ -496,23 +519,93 @@ def get_or_create_current_report_output( report_output.year, ) - def report_output_exists( - self, session: Session, country_id: str, report_output_id: int - ) -> bool: - return ( - self._select_report_output(session, report_output_id, country_id) - is not None + def _build_view( + self, + session: Session, + report_output: ReportOutput, + *, + response_id: int | None = None, + ) -> ReportOutputView: + return ReportOutputView( + report_output=report_output, + display_run=self.report_run_service.select_display_run( + session, report_output + ), + response_id=response_id, ) + def create_or_reuse_report_output( + self, + country_id: str, + simulation_1_id: int, + simulation_2_id: int | None = None, + year: str = "2025", + ) -> ReportCreateResult: + with self._sessions.begin() as session: + existing = self._find_existing_report_output( + session, + country_id, + simulation_1_id, + simulation_2_id, + year, + ) + created = existing is None + report_output = ( + self._create_report_output( + session, + country_id, + simulation_1_id, + simulation_2_id, + year, + ) + if existing is None + else self._ensure_report_output_dual_write_state( + session, existing.id, country_id + ) + ) + return ReportCreateResult( + view=self._build_view(session, report_output), + created=created, + ) + + def resolve_report_output( + self, + country_id: str, + report_output_id: int, + ) -> ReportOutputView | None: + if type(report_output_id) is not int or report_output_id < 0: + raise Exception( + f"Invalid report output ID: {report_output_id}. " + "Must be a positive integer." + ) + with self._sessions.begin() as session: + requested = self._get_report_output(session, country_id, report_output_id) + if requested is None: + return None + if self._is_current_report_output(requested): + report_output = self._ensure_report_output_dual_write_state( + session, report_output_id, country_id + ) + response_id = None + else: + report_output = self._get_or_create_current_report_output( + session, requested + ) + response_id = report_output_id + return self._build_view( + session, + report_output, + response_id=response_id, + ) + def update_report_output( self, - session: Session, country_id: str, report_id: int, status: str | None = None, output: dict | list | str | None = None, error_message: str | None = None, - ) -> bool: + ) -> ReportOutputView | None: values = { key: value for key, value in { @@ -523,16 +616,33 @@ def update_report_output( if value is not None } if not values: - return False + return None if isinstance(values.get("output"), str): values["output"] = json.loads(values["output"]) - report_output = self._select_report_output( - session, report_id, country_id, for_update=True - ) - if report_output is None: - raise ValueError(f"Report output #{report_id} not found") - if status == "running": - runs = self._list_runs_descending(session, report_id) + with self._sessions.begin() as session: + report_output = self._select_report_output( + session, report_id, country_id, for_update=True + ) + if report_output is None: + raise LookupError(f"Report output #{report_id} not found") + self._update_report_output( + session, + report_output, + values, + requested_status=status, + ) + return self._build_view(session, report_output) + + def _update_report_output( + self, + session: Session, + report_output: ReportOutput, + values: dict, + *, + requested_status: str | None, + ) -> None: + if requested_status == "running": + runs = self._list_runs_descending(session, report_output.id) if not self._has_mutable_running_run(report_output, runs): raise ValueError( "Cannot mark report output running without an active pending " @@ -540,5 +650,8 @@ def update_report_output( ) for field, value in values.items(): setattr(report_output, field, value) - self.ensure_report_output_dual_write_state(session, report_id, country_id) - return True + self._ensure_report_output_dual_write_state( + session, + report_output.id, + report_output.country_id, + ) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index a2ccc4dc4..887fdcef8 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -10,11 +10,21 @@ from policyengine_api.constants import get_report_output_cache_version from policyengine_api.endpoints.household import get_calculate from policyengine_api.endpoints.policy import get_policy_search -from policyengine_api.data.v1_models import Household, Policy, ReportOutput, Simulation +from policyengine_api.data.v1_models import ( + Household, + Policy, + ReportOutput, + ReportOutputRun, + Simulation, +) from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.simulation_routes import simulation_bp +from policyengine_api.services.report_output_service import ( + ReportCreateResult, + ReportOutputView, +) from policyengine_api.services.simulation_service import SimulationCreateResult from tests.contract.clients import ( ASGIContractClient, @@ -340,49 +350,48 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", - return_value=None, - ) - ) - stack.enter_context( - patch( - "policyengine_api.routes.report_output_routes.report_output_service.create_report_output", - return_value=ReportOutput( - id=33, - country_id="us", - simulation_1_id=11, - simulation_2_id=None, - api_version=get_report_output_cache_version("us"), - status="pending", - year="2026", - ), - ) - ) - stack.enter_context( - patch( - "policyengine_api.routes.report_output_routes.report_output_service.get_report_output", - return_value=ReportOutput( - id=33, - country_id="us", - simulation_1_id=11, - simulation_2_id=None, - api_version=get_report_output_cache_version("us"), - status="pending", - year="2026", + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", + return_value=ReportCreateResult( + view=ReportOutputView( + report_output=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), + display_run=ReportOutputRun( + id="run-33", + report_output_id=33, + run_sequence=1, + status="pending", + ), + ), + created=True, ), ) ) stack.enter_context( patch( - "policyengine_api.routes.report_output_routes.report_output_service.ensure_report_output_dual_write_state", - return_value=ReportOutput( - id=33, - country_id="us", - simulation_1_id=11, - simulation_2_id=None, - api_version=get_report_output_cache_version("us"), - status="pending", - year="2026", + "policyengine_api.routes.report_output_routes.report_output_service.resolve_report_output", + return_value=ReportOutputView( + report_output=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), + display_run=ReportOutputRun( + id="run-33", + report_output_id=33, + run_sequence=1, + status="pending", + ), ), ) ) diff --git a/tests/unit/routes/test_route_exception_handling.py b/tests/unit/routes/test_route_exception_handling.py index fe4dd8522..b67e97274 100644 --- a/tests/unit/routes/test_route_exception_handling.py +++ b/tests/unit/routes/test_route_exception_handling.py @@ -63,7 +63,7 @@ def test_simulation_create_value_error_still_400(): def test_report_create_runtime_error_becomes_500(): client = _client_with(report_output_bp) with patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", side_effect=RuntimeError("db went away"), ): response = client.post( @@ -76,7 +76,7 @@ def test_report_create_runtime_error_becomes_500(): def test_report_create_value_error_still_400(): client = _client_with(report_output_bp) with patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", side_effect=ValueError("bad input"), ): response = client.post( diff --git a/tests/unit/services/test_report_output_service.py b/tests/unit/services/test_report_output_service.py index ac44cae5a..f4e6d67ea 100644 --- a/tests/unit/services/test_report_output_service.py +++ b/tests/unit/services/test_report_output_service.py @@ -2,214 +2,251 @@ from sqlalchemy import func, select from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.data.v1_models import ( - ReportOutput, - ReportOutputRun, -) +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun from policyengine_api.services.report_output_service import ReportOutputService from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.simulation_service import SimulationService -service = ReportOutputService() run_service = ReportRunService() -simulation_service = SimulationService() -def create_simulation(orm_session, *, policy_id=1, population_id="household-1"): - return simulation_service._create_simulation( - orm_session, - country_id="us", - population_id=population_id, - population_type="household", - policy_id=policy_id, +@pytest.fixture +def service(orm_session_factory): + return ReportOutputService(orm_session_factory) + + +def create_simulation( + orm_session_factory, + *, + policy_id=1, + population_id="household-1", +): + return ( + SimulationService(orm_session_factory) + .get_or_create_simulation( + country_id="us", + population_id=population_id, + population_type="household", + policy_id=policy_id, + ) + .simulation ) -def test_creates_mapped_report_with_spec_and_initial_run(orm_session): - simulation = create_simulation(orm_session) +def create_report(service, simulation_id, *, year="2025"): + return service.create_or_reuse_report_output( + "us", simulation_id, year=year + ).view.report_output + + +def test_creates_mapped_report_with_spec_and_initial_run( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) - report = service.create_report_output( - orm_session, + creation = service.create_or_reuse_report_output( country_id="us", simulation_1_id=simulation.id, year="2025", ) + report = creation.view.report_output + assert creation.created is True assert isinstance(report, ReportOutput) assert report.report_kind == "household_single" assert isinstance(report.report_spec_json, dict) - run = orm_session.get(ReportOutputRun, report.active_run_id) - assert isinstance(run, ReportOutputRun) - assert run.status == "pending" - assert run.report_spec_snapshot_json == report.report_spec_json + assert isinstance(creation.view.display_run, ReportOutputRun) + assert creation.view.display_run.status == "pending" + assert ( + creation.view.display_run.report_spec_snapshot_json == report.report_spec_json + ) -def test_create_reuses_current_report_and_repairs_dual_state(orm_session): - simulation = create_simulation(orm_session) - existing = ReportOutput( - country_id="us", - simulation_1_id=simulation.id, - simulation_2_id=None, - api_version=get_report_output_cache_version("us"), - status="pending", - year="2025", - ) - orm_session.add(existing) - orm_session.flush() +def test_create_reuses_current_report_and_repairs_dual_state( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + with orm_session_factory.begin() as session: + existing = ReportOutput( + country_id="us", + simulation_1_id=simulation.id, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2025", + ) + session.add(existing) + session.flush() + report_id = existing.id - result = service.create_report_output(orm_session, "us", simulation.id, year="2025") + result = service.create_or_reuse_report_output("us", simulation.id, year="2025") - assert result is existing - assert existing.active_run_id is not None - assert existing.report_spec_json is not None + assert result.created is False + assert result.view.report_output.id == report_id + assert result.view.report_output.active_run_id is not None + assert result.view.report_output.report_spec_json is not None @pytest.mark.parametrize("missing_secondary", [False, True]) -def test_create_rejects_missing_linked_simulation(orm_session, missing_secondary): - simulation = create_simulation(orm_session) if missing_secondary else None +def test_create_rejects_missing_linked_simulation( + service, + orm_session_factory, + missing_secondary, +): + simulation = create_simulation(orm_session_factory) if missing_secondary else None with pytest.raises(ValueError, match="references missing simulation"): - service.create_report_output( - orm_session, + service.create_or_reuse_report_output( "us", simulation.id if simulation else 999, 999 if missing_secondary else None, ) - assert orm_session.scalar(select(func.count()).select_from(ReportOutput)) == 0 + with orm_session_factory() as session: + assert session.scalar(select(func.count()).select_from(ReportOutput)) == 0 -def test_find_existing_report_uses_current_cache_version(orm_session): - simulation = create_simulation(orm_session) - stale = ReportOutput( - country_id="us", - simulation_1_id=simulation.id, - simulation_2_id=None, - api_version="stale", - status="pending", - year="2025", - ) - orm_session.add(stale) - orm_session.flush() - - assert ( - service.find_existing_report_output( - orm_session, "us", simulation.id, year="2025" +def test_create_ignores_stale_cache_version(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + with orm_session_factory.begin() as session: + stale = ReportOutput( + country_id="us", + simulation_1_id=simulation.id, + simulation_2_id=None, + api_version="stale", + status="pending", + year="2025", ) - is None - ) + session.add(stale) + session.flush() + stale_id = stale.id - current = service.create_report_output( - orm_session, "us", simulation.id, year="2025" - ) - assert ( - service.find_existing_report_output( - orm_session, "us", simulation.id, year="2025" - ) - is current + current = service.create_or_reuse_report_output("us", simulation.id, year="2025") + + assert current.created is True + assert current.view.report_output.id != stale_id + assert current.view.report_output.api_version == get_report_output_cache_version( + "us" ) -def test_get_report_is_scoped_to_country_and_validates_id(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_resolve_is_scoped_to_country_and_validates_id( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) - assert service.get_report_output(orm_session, "us", report.id) is report - assert service.get_report_output(orm_session, "uk", report.id) is None + assert service.resolve_report_output("us", report.id).report_output.id == report.id + assert service.resolve_report_output("uk", report.id) is None with pytest.raises(Exception, match="Invalid report output ID"): - service.get_report_output(orm_session, "us", -1) + service.resolve_report_output("us", -1) -def test_update_complete_stores_python_json_and_promotes_run(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_update_complete_stores_python_json_and_promotes_run( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) active_run_id = report.active_run_id - assert service.update_report_output( - orm_session, + view = service.update_report_output( "us", report.id, status="complete", output={"ok": True}, ) - assert report.output == {"ok": True} - assert report.active_run_id is None - assert report.latest_successful_run_id == active_run_id - run = orm_session.get(ReportOutputRun, active_run_id) - assert run.status == "complete" - assert run.output == {"ok": True} - assert run.started_at is not None - assert run.finished_at is not None + assert view.report_output.output == {"ok": True} + assert view.report_output.active_run_id is None + assert view.report_output.latest_successful_run_id == active_run_id + assert view.display_run.status == "complete" + assert view.display_run.output == {"ok": True} + assert view.display_run.started_at is not None + assert view.display_run.finished_at is not None -def test_update_accepts_existing_v1_json_string_boundary(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_update_accepts_existing_v1_json_string_boundary( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) - service.update_report_output(orm_session, "us", report.id, output='{"ok": true}') + view = service.update_report_output("us", report.id, output='{"ok": true}') - assert report.output == {"ok": True} + assert view.report_output.output == {"ok": True} -def test_update_running_requires_mutable_run(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_update_running_requires_mutable_run(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) service.update_report_output( - orm_session, "us", report.id, status="complete", output={"ok": True} + "us", report.id, status="complete", output={"ok": True} ) with pytest.raises(ValueError, match="without an active pending or running"): - service.update_report_output(orm_session, "us", report.id, status="running") + service.update_report_output("us", report.id, status="running") + + +def _add_rerun(orm_session_factory, report_id): + with orm_session_factory.begin() as session: + report = session.get(ReportOutput, report_id) + successful_run_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + session, report.id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + return rerun.id, successful_run_id -def test_update_running_targets_active_rerun(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_update_running_targets_active_rerun(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) service.update_report_output( - orm_session, "us", report.id, status="complete", output={"old": True} + "us", report.id, status="complete", output={"old": True} ) - successful_run_id = report.latest_successful_run_id - rerun = run_service.create_report_output_run( - orm_session, report.id, trigger_type="rerun" - ) - report.active_run_id = rerun.id + rerun_id, successful_run_id = _add_rerun(orm_session_factory, report.id) - service.update_report_output(orm_session, "us", report.id, status="running") + view = service.update_report_output("us", report.id, status="running") - assert rerun.status == "running" - assert rerun.started_at is not None - assert report.active_run_id == rerun.id - assert report.latest_successful_run_id == successful_run_id + assert view.display_run.id == rerun_id + assert view.display_run.status == "running" + assert view.display_run.started_at is not None + assert view.report_output.active_run_id == rerun_id + assert view.report_output.latest_successful_run_id == successful_run_id -def test_failed_rerun_preserves_latest_successful_pointer(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_failed_rerun_preserves_latest_successful_pointer( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) service.update_report_output( - orm_session, "us", report.id, status="complete", output={"old": True} - ) - successful_run_id = report.latest_successful_run_id - rerun = run_service.create_report_output_run( - orm_session, report.id, trigger_type="rerun" + "us", report.id, status="complete", output={"old": True} ) - report.active_run_id = rerun.id + rerun_id, successful_run_id = _add_rerun(orm_session_factory, report.id) - service.update_report_output( - orm_session, "us", report.id, status="error", error_message="failed" + view = service.update_report_output( + "us", report.id, status="error", error_message="failed" ) - assert report.active_run_id is None - assert report.latest_successful_run_id == successful_run_id - assert rerun.status == "error" - assert rerun.finished_at is not None + assert view.report_output.active_run_id is None + assert view.report_output.latest_successful_run_id == successful_run_id + assert view.display_run.id == rerun_id + assert view.display_run.status == "error" + assert view.display_run.finished_at is not None -def test_noop_and_missing_updates(orm_session): - simulation = create_simulation(orm_session) - report = service.create_report_output(orm_session, "us", simulation.id) +def test_noop_and_missing_updates(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) - assert service.update_report_output(orm_session, "us", report.id) is False - with pytest.raises(ValueError, match="Report output #999 not found"): - service.update_report_output(orm_session, "us", 999, status="pending") + assert service.update_report_output("us", report.id) is None + with pytest.raises(LookupError, match="Report output #999 not found"): + service.update_report_output("us", 999, status="pending") diff --git a/tests/unit/services/test_report_service_owned_sessions.py b/tests/unit/services/test_report_service_owned_sessions.py new file mode 100644 index 000000000..e9103a07b --- /dev/null +++ b/tests/unit/services/test_report_service_owned_sessions.py @@ -0,0 +1,110 @@ +import inspect +from pathlib import Path + +import pytest +from sqlalchemy import func, select + +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun +from policyengine_api.services.report_output_service import ( + ReportCreateResult, + ReportOutputService, + ReportOutputView, +) +from policyengine_api.services.simulation_service import SimulationService + + +ROUTE_PATH = ( + Path(__file__).parents[3] + / "policyengine_api" + / "routes" + / "report_output_routes.py" +) + + +@pytest.fixture +def simulation_id(orm_session_factory): + return ( + SimulationService(orm_session_factory) + .get_or_create_simulation("us", "household-1", "household", 1) + .simulation.id + ) + + +@pytest.fixture +def service(orm_session_factory): + return ReportOutputService(orm_session_factory) + + +def test_report_public_methods_do_not_accept_sessions(): + for method_name in ( + "create_or_reuse_report_output", + "resolve_report_output", + "update_report_output", + ): + parameters = inspect.signature( + getattr(ReportOutputService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_report_routes_do_not_manage_sessions_or_query_run_services(): + source = ROUTE_PATH.read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "report_run_service" not in source + assert "sqlalchemy" not in source + + +def test_create_or_reuse_returns_report_and_display_run(service, simulation_id): + created = service.create_or_reuse_report_output("us", simulation_id, year="2025") + reused = service.create_or_reuse_report_output("us", simulation_id, year="2025") + + assert isinstance(created, ReportCreateResult) + assert created.created is True + assert isinstance(created.view, ReportOutputView) + assert isinstance(created.view.report_output, ReportOutput) + assert isinstance(created.view.display_run, ReportOutputRun) + assert created.view.response_id is None + assert reused.created is False + assert reused.view.report_output.id == created.view.report_output.id + + +def test_create_rolls_back_all_rows_when_dual_write_fails( + service, + simulation_id, + orm_session_factory, + monkeypatch, +): + def fail_dual_write(*args, **kwargs): + raise RuntimeError("report dual write failed") + + monkeypatch.setattr( + service, "_ensure_report_output_dual_write_state", fail_dual_write + ) + + with pytest.raises(RuntimeError, match="report dual write failed"): + service.create_or_reuse_report_output("us", simulation_id, year="2026") + + with orm_session_factory() as session: + report_count = session.scalar(select(func.count()).select_from(ReportOutput)) + run_count = session.scalar(select(func.count()).select_from(ReportOutputRun)) + assert report_count == 0 + assert run_count == 0 + + +def test_update_returns_view_from_the_same_atomic_operation( + service, + simulation_id, +): + created = service.create_or_reuse_report_output("us", simulation_id) + + view = service.update_report_output( + "us", + created.view.report_output.id, + status="complete", + output={"result": 42}, + ) + + assert isinstance(view, ReportOutputView) + assert view.report_output.output == {"result": 42} + assert view.display_run.status == "complete" diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_stage5_routes.py index 9b3e80a49..7db0fa6fb 100644 --- a/tests/unit/test_stage5_routes.py +++ b/tests/unit/test_stage5_routes.py @@ -35,11 +35,12 @@ def create_simulation(factory, *, population_id="household-1", policy_id=1): def create_report(factory, simulation_id): - with factory.begin() as session: - report = report_service.create_report_output( - session, "us", simulation_id, year="2025" - ) - return report.id + report = ( + ReportOutputService(factory) + .create_or_reuse_report_output("us", simulation_id, year="2025") + .view.report_output + ) + return report.id def test_create_simulation_existing_row_repairs_dual_write_state( @@ -179,10 +180,10 @@ def test_report_routes_scope_reads_and_writes_to_country(orm_session_factory): def test_report_get_serializes_display_run_timestamps(orm_session_factory): simulation_id = create_simulation(orm_session_factory, policy_id=46) report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"ok": True} + ) with orm_session_factory.begin() as session: - report_service.update_report_output( - session, "us", report_id, status="complete", output={"ok": True} - ) report = session.get(ReportOutput, report_id) run = session.get(ReportOutputRun, report.latest_successful_run_id) run.requested_at = datetime(2026, 5, 4, 12, 0) @@ -201,10 +202,10 @@ def test_report_patch_updates_active_rerun_and_preserves_success( ): simulation_id = create_simulation(orm_session_factory, policy_id=47) report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"old": True} + ) with orm_session_factory.begin() as session: - report_service.update_report_output( - session, "us", report_id, status="complete", output={"old": True} - ) report = session.get(ReportOutput, report_id) successful_id = report.latest_successful_run_id rerun = run_service.create_report_output_run( @@ -235,10 +236,10 @@ def test_report_patch_updates_active_rerun_and_preserves_success( def test_report_patch_complete_promotes_active_rerun(orm_session_factory): simulation_id = create_simulation(orm_session_factory, policy_id=48) report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"old": True} + ) with orm_session_factory.begin() as session: - report_service.update_report_output( - session, "us", report_id, status="complete", output={"old": True} - ) report = session.get(ReportOutput, report_id) rerun = run_service.create_report_output_run( session, report_id, trigger_type="rerun" From 8824eb22128a87e7a38c270b2d6da0cd7a9e14e9 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:06:14 +0300 Subject: [PATCH 61/89] refactor: move saved-policy SQL behind services --- policyengine_api/endpoints/policy.py | 114 +++++------------- policyengine_api/services/policy_service.py | 28 +++++ .../services/user_policy_service.py | 101 ++++++++++++++++ tests/contract/test_v1_route_contracts.py | 59 +++------ .../endpoints/test_stage7_orm_endpoints.py | 14 +-- tests/unit/services/test_policy_service.py | 41 +++++++ .../services/test_service_owned_sessions.py | 5 +- .../unit/services/test_user_policy_service.py | 78 ++++++++++++ 8 files changed, 307 insertions(+), 133 deletions(-) create mode 100644 policyengine_api/services/user_policy_service.py create mode 100644 tests/unit/services/test_user_policy_service.py diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py index 275303592..023ea9abe 100644 --- a/policyengine_api/endpoints/policy.py +++ b/policyengine_api/endpoints/policy.py @@ -1,23 +1,14 @@ from policyengine_api.utils.payload_validators import validate_country import json from flask import Response, request -from sqlalchemy import select - -from policyengine_api.data.orm import get_v1_session_factory -from policyengine_api.data.v1_models import Policy, UserPolicy - - -USER_POLICY_IDENTITY_FIELDS = ( - "country_id", - "reform_id", - "baseline_id", - "user_id", - "year", - "geography", - "reform_label", - "baseline_label", - "dataset", -) + +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_policy_service import UserPolicyService + + +policy_service = PolicyService() +user_policy_service = UserPolicyService() def _serialize_user_policy(user_policy: UserPolicy) -> dict: @@ -54,14 +45,11 @@ def get_policy_search(country_id: str) -> dict: unique_only = request.args.get("unique_only", default=False, type=json.loads) try: - sessions = get_v1_session_factory() - with sessions() as session: - results = session.scalars( - select(Policy).where( - Policy.country_id == country_id, - Policy.label.contains(query, autoescape=True), - ) - ).all() + results = policy_service.search_policies( + country_id, + query, + unique_only=unique_only, + ) if not results: body = dict( @@ -70,24 +58,6 @@ def get_policy_search(country_id: str) -> dict: ) return Response(json.dumps(body), status=404, mimetype="application/json") - # If unique_only is true, filter results to only include - # items where everything except ID is unique - if unique_only: - processed_vals = set() - new_results = [] - - # Compare every label and hash to what's contained in processed_vals - # If a label-hash set aren't already in processed_vals, - # add them to new_results - for policy in results: - comparison_vals = policy.label, policy.policy_hash - if comparison_vals not in processed_vals: - new_results.append(policy) - processed_vals.add(comparison_vals) - - # Overwrite results with new_results - results = new_results - # Format into: [{ id: 1, label: "My policy" }, ...] policies = [dict(id=result.id, label=result.label) for result in results] body = dict( @@ -124,7 +94,7 @@ def set_user_policy(country_id: str) -> dict: added_date = payload.pop("added_date") updated_date = payload.pop("updated_date") budgetary_impact = payload.pop("budgetary_impact", None) - type = payload.pop("type", None) + policy_type = payload.pop("type", None) values = { "country_id": country_id, @@ -141,7 +111,7 @@ def set_user_policy(country_id: str) -> dict: "added_date": added_date, "updated_date": updated_date, "budgetary_impact": budgetary_impact, - "type": type, + "type": policy_type, } # When setting a user policy, "unique" records contain @@ -154,30 +124,19 @@ def set_user_policy(country_id: str) -> dict: # to be tested; type is not yet implemented try: - with get_v1_session_factory().begin() as session: - user_policy = session.scalar( - select(UserPolicy).where( - *( - getattr(UserPolicy, field) == values[field] - for field in USER_POLICY_IDENTITY_FIELDS - ) - ) + creation = user_policy_service.create_or_get_user_policy(values) + user_policy = creation.user_policy + if not creation.created: + response = dict( + status="ok", + message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", + result=dict(id=user_policy.id), + ) + return Response( + json.dumps(response), + status=200, + mimetype="application/json", ) - if user_policy is None: - user_policy = UserPolicy(**values) - session.add(user_policy) - session.flush() - else: - response = dict( - status="ok", - message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", - result=dict(id=user_policy.id), - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) except Exception as e: return Response( json.dumps( @@ -208,17 +167,8 @@ def get_user_policy(country_id: str, user_id: str) -> dict: Fetch all saved user policies by user id """ - # Get the policy record for a given policy ID. - sessions = get_v1_session_factory() - with sessions() as session: - user_policies = session.scalars( - select(UserPolicy).where( - UserPolicy.country_id == country_id, - UserPolicy.user_id == user_id, - ) - ).all() - - rows_parsed = [_serialize_user_policy(row) for row in user_policies] + user_policies = user_policy_service.list_user_policies(country_id, user_id) + rows_parsed = [_serialize_user_policy(row) for row in user_policies] if rows_parsed is None: response = dict( @@ -302,11 +252,7 @@ def update_user_policy(country_id: str) -> dict: ) try: - with get_v1_session_factory().begin() as session: - user_policy = session.get(UserPolicy, user_policy_id) - if user_policy is not None: - for key, value in payload.items(): - setattr(user_policy, key, value) + user_policy_service.update_user_policy(user_policy_id, payload) except Exception as e: return Response( json.dumps( diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index b2780b135..728565882 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -63,6 +63,34 @@ def get_policy_json( policy = self.get_policy(country_id, policy_id) return None if policy is None else policy.policy_json + def search_policies( + self, + country_id: str, + query: str = "", + *, + unique_only: bool = False, + ) -> list[Policy]: + with self._sessions() as session: + results = list( + session.scalars( + select(Policy).where( + Policy.country_id == country_id, + Policy.label.contains(query, autoescape=True), + ) + ) + ) + if not unique_only: + return results + + unique_results = [] + processed_values = set() + for policy in results: + identity = policy.label, policy.policy_hash + if identity not in processed_values: + unique_results.append(policy) + processed_values.add(identity) + return unique_results + def set_policy( self, country_id: str, diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py new file mode 100644 index 000000000..07ca93f5e --- /dev/null +++ b/policyengine_api/services/user_policy_service.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserPolicy + + +USER_POLICY_IDENTITY_FIELDS = ( + "country_id", + "reform_id", + "baseline_id", + "user_id", + "year", + "geography", + "reform_label", + "baseline_label", + "dataset", +) + + +@dataclass(frozen=True) +class UserPolicyCreateResult: + user_policy: UserPolicy + created: bool + + +class UserPolicyService: + """Saved-policy operations with service-owned ORM transactions.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() + + @staticmethod + def _find_matching_user_policy( + session: Session, + values: Mapping[str, Any], + ) -> UserPolicy | None: + return session.scalar( + select(UserPolicy).where( + *( + getattr(UserPolicy, field) == values[field] + for field in USER_POLICY_IDENTITY_FIELDS + ) + ) + ) + + def create_or_get_user_policy( + self, + values: Mapping[str, Any], + ) -> UserPolicyCreateResult: + with self._sessions.begin() as session: + user_policy = self._find_matching_user_policy(session, values) + created = user_policy is None + if user_policy is None: + user_policy = UserPolicy(**values) + session.add(user_policy) + session.flush() + return UserPolicyCreateResult( + user_policy=user_policy, + created=created, + ) + + def list_user_policies( + self, + country_id: str, + user_id: str, + ) -> list[UserPolicy]: + with self._sessions() as session: + return list( + session.scalars( + select(UserPolicy).where( + UserPolicy.country_id == country_id, + UserPolicy.user_id == user_id, + ) + ) + ) + + def update_user_policy( + self, + user_policy_id: int, + values: Mapping[str, Any], + ) -> UserPolicy | None: + with self._sessions.begin() as session: + user_policy = session.get(UserPolicy, user_policy_id) + if user_policy is None: + return None + for field, value in values.items(): + setattr(user_policy, field, value) + return user_policy diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 887fdcef8..5bd63df13 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -1,4 +1,4 @@ -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack import importlib import sys from types import SimpleNamespace @@ -202,42 +202,6 @@ def _json_payload(contract: ContractRequest) -> dict | None: return None -def _policy_search_session_factory(): - policies = [ - Policy( - id=123, - country_id="us", - label="Tax reform", - api_version="1", - policy_json={}, - policy_hash="hash-1", - ), - Policy( - id=124, - country_id="us", - label="Tax reform", - api_version="1", - policy_json={}, - policy_hash="hash-1", - ), - ] - - class Result: - def all(self): - return policies - - class Session: - def scalars(self, statement): - return Result() - - class Factory: - @contextmanager - def __call__(self): - yield Session() - - return Factory() - - def _fake_country(): return SimpleNamespace( metadata={}, @@ -271,8 +235,25 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.get_v1_session_factory", - return_value=_policy_search_session_factory(), + "policyengine_api.endpoints.policy.policy_service.search_policies", + return_value=[ + Policy( + id=123, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + Policy( + id=124, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + ], ) ) stack.enter_context( diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py index f1672adfb..f9450b380 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -129,15 +129,11 @@ def test_user_policy_endpoints_round_trip_through_orm_session_factory( "type": None, } - with patch( - "policyengine_api.endpoints.policy.get_v1_session_factory", - return_value=orm_session_factory, - ): - with app.test_request_context(json=payload): - created = set_user_policy("us") - listed = get_user_policy("us", "auth0|one") - with app.test_request_context(json={"id": 1, "reform_label": "Updated"}): - updated = update_user_policy("us") + with app.test_request_context(json=payload): + created = set_user_policy("us") + listed = get_user_policy("us", "auth0|one") + with app.test_request_context(json={"id": 1, "reform_label": "Updated"}): + updated = update_user_policy("us") assert created.status_code == 201 assert created.get_json()["result"]["dataset"] == "default" diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index ca7a24951..a2b1c8062 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -52,6 +52,47 @@ def test_get_policy_json_returns_python_object(service, existing_policy_record): } +def test_search_policies_filters_and_deduplicates(service, orm_session_factory): + with orm_session_factory.begin() as session: + session.add_all( + [ + Policy( + id=31, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="same-hash", + ), + Policy( + id=32, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={"different": True}, + policy_hash="same-hash", + ), + Policy( + id=33, + country_id="us", + label="Benefit reform", + api_version="1", + policy_json={}, + policy_hash="other-hash", + ), + ] + ) + + all_results = service.search_policies("us", "Tax", unique_only=False) + unique_results = service.search_policies("us", "Tax", unique_only=True) + escaped_wildcard_results = service.search_policies("us", "Tax%", unique_only=False) + + assert len(all_results) == 2 + assert len(unique_results) == 1 + assert unique_results[0].label == "Tax reform" + assert escaped_wildcard_results == [] + + def test_set_policy_adds_mapped_entity(service, monkeypatch): monkeypatch.setattr( "policyengine_api.services.policy_service.hash_object", diff --git a/tests/unit/services/test_service_owned_sessions.py b/tests/unit/services/test_service_owned_sessions.py index e1edc3b70..1976013ad 100644 --- a/tests/unit/services/test_service_owned_sessions.py +++ b/tests/unit/services/test_service_owned_sessions.py @@ -21,7 +21,10 @@ HouseholdService, ("get_household", "create_household", "update_household"), ), - (PolicyService, ("get_policy", "get_policy_json", "set_policy")), + ( + PolicyService, + ("get_policy", "get_policy_json", "search_policies", "set_policy"), + ), ( SimulationService, ( diff --git a/tests/unit/services/test_user_policy_service.py b/tests/unit/services/test_user_policy_service.py new file mode 100644 index 000000000..d271fc0b3 --- /dev/null +++ b/tests/unit/services/test_user_policy_service.py @@ -0,0 +1,78 @@ +import inspect +from pathlib import Path + +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.services.user_policy_service import ( + UserPolicyCreateResult, + UserPolicyService, +) + + +ENDPOINT_PATH = ( + Path(__file__).parents[3] / "policyengine_api" / "endpoints" / "policy.py" +) + + +def _values(**overrides): + values = { + "country_id": "us", + "reform_id": 2, + "reform_label": "Reform", + "baseline_id": 1, + "baseline_label": "Current law", + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": "enhanced_cps_2024", + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 2, + "budgetary_impact": None, + "type": None, + } + values.update(overrides) + return values + + +def test_user_policy_public_methods_do_not_accept_sessions(): + for method_name in ( + "create_or_get_user_policy", + "list_user_policies", + "update_user_policy", + ): + parameters = inspect.signature( + getattr(UserPolicyService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_legacy_policy_endpoints_do_not_manage_sessions_or_queries(): + source = ENDPOINT_PATH.read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "from sqlalchemy" not in source + assert "select(" not in source + + +def test_create_reuse_list_and_update_user_policy(orm_session_factory): + service = UserPolicyService(orm_session_factory) + + created = service.create_or_get_user_policy(_values()) + reused = service.create_or_get_user_policy( + _values(number_of_provisions=99, updated_date=99) + ) + listed = service.list_user_policies("us", "auth0|one") + updated = service.update_user_policy( + created.user_policy.id, + {"reform_label": "Updated", "updated_date": 3}, + ) + + assert isinstance(created, UserPolicyCreateResult) + assert created.created is True + assert reused.created is False + assert reused.user_policy.id == created.user_policy.id + assert len(listed) == 1 + assert isinstance(listed[0], UserPolicy) + assert updated.reform_label == "Updated" + assert updated.updated_date == 3 From 3eaea38d9b089cc55bea0aa2718165cef724da54 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:15:06 +0300 Subject: [PATCH 62/89] refactor: isolate household calculation persistence --- policyengine_api/country.py | 28 +- policyengine_api/endpoints/household.py | 184 +++--------- .../services/household_calculation_service.py | 261 ++++++++++++++++++ .../services/tracer_analysis_service.py | 2 +- tests/integration/test_simulations.py | 6 +- .../endpoints/test_stage7_orm_endpoints.py | 52 ++-- .../test_household_calculation_service.py | 189 +++++++++++++ 7 files changed, 512 insertions(+), 210 deletions(-) create mode 100644 policyengine_api/services/household_calculation_service.py create mode 100644 tests/unit/services/test_household_calculation_service.py diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 593a49cf6..42980dc1b 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -2,7 +2,7 @@ import inspect import json from policyengine_core.taxbenefitsystems import TaxBenefitSystem -from typing import Union, Optional +from typing import Union from policyengine_api.utils import get_safe_json from policyengine_core.parameters import ( ParameterNode, @@ -23,12 +23,10 @@ build_congressional_district_metadata, ) -from policyengine_api.data.orm import get_v1_session_factory -from policyengine_api.data.v1_models import Tracer from policyengine_api.constants import ( - COUNTRY_PACKAGE_VERSIONS, get_bundle_default_dataset_option, ) +from policyengine_api.services.household_calculation_service import CalculationResult class PolicyEngineCountry: @@ -362,9 +360,7 @@ def calculate( self, household: dict, reform: Union[dict, None], - household_id: Optional[int] = None, - policy_id: Optional[int] = None, - ): + ) -> CalculationResult: simulation, system = self._create_simulation(household, reform) household = json.loads(json.dumps(household)) @@ -432,20 +428,10 @@ def calculate( tracer_output = simulation.tracer.computation_log log_lines = tracer_output.lines(aggregate=False, max_depth=10) - if household_id is not None and policy_id is not None: - # write to local database - with get_v1_session_factory(local=True).begin() as session: - session.add( - Tracer( - household_id=household_id, - policy_id=policy_id, - country_id=self.country_id, - api_version=COUNTRY_PACKAGE_VERSIONS[self.country_id], - tracer_output=log_lines, - ) - ) - - return household + return CalculationResult( + household=household, + tracer_output=log_lines, + ) def _create_simulation( self, diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py index 78007242f..843436ceb 100644 --- a/policyengine_api/endpoints/household.py +++ b/policyengine_api/endpoints/household.py @@ -1,8 +1,6 @@ import json from flask import Response, request -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS import logging -from datetime import date from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs from policyengine_api.utils.input_validation import ( find_unrecognized_inputs, @@ -10,10 +8,17 @@ ) from policyengine_api.utils.payload_validators import validate_country from policyengine_core.errors import SituationParsingError -from sqlalchemy import select -from policyengine_api.data.orm import get_v1_session_factory -from policyengine_api.data.v1_models import ComputedHousehold, Household, Policy +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, + HouseholdNotFoundError, + InvalidHouseholdInputsError, + PolicyNotFoundError, + add_yearly_variables, +) + + +household_calculation_service = HouseholdCalculationService() def get_countries(): @@ -22,41 +27,6 @@ def get_countries(): return COUNTRIES -def add_yearly_variables(household, country_id, countries=None): - """ - Add yearly variables to a household dict before enqueueing calculation - """ - metadata = (countries or get_countries()).get(country_id).metadata - - variables = metadata["variables"] - entities = metadata["entities"] - household_year = get_household_year(household) - - for variable in variables: - if variables[variable]["definitionPeriod"] in ( - "year", - "month", - "eternity", - ): - entity_plural = entities[variables[variable]["entity"]]["plural"] - if entity_plural in household: - possible_entities = household[entity_plural].keys() - for entity in possible_entities: - if ( - variables[variable]["name"] - not in household[entity_plural][entity] - ): - if variables[variable]["isInputVariable"]: - household[entity_plural][entity][ - variables[variable]["name"] - ] = {household_year: variables[variable]["defaultValue"]} - else: - household[entity_plural][entity][ - variables[variable]["name"] - ] = {household_year: None} - return household - - def get_invalid_inputs_response(household_json, policy_json, country): invalid_inputs = find_unrecognized_inputs( household_json, @@ -79,27 +49,6 @@ def get_invalid_inputs_response(household_json, policy_json, country): ) -def get_household_year(household): - """Given a household dict, get the household's year - - Args: - household (dict): The household itself - """ - - # Set household_year based on current year - household_year = date.today().year - - # Determine if "age" variable present within household and return list of values at it - household_age_list = list( - household.get("people", {}).get("you", {}).get("age", {}).keys() - ) - # If it is, overwrite household_year with the value present - if len(household_age_list) > 0: - household_year = household_age_list[0] - - return household_year - - @validate_country def get_household_under_policy(country_id: str, household_id: str, policy_id: str): """Get a household's output data under a given policy. @@ -110,46 +59,13 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st policy_id (str): The policy ID. """ - api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - # Look in computed_households to see if already computed - - sessions = get_v1_session_factory(local=True) - with sessions() as session: - computed_household = session.scalar( - select(ComputedHousehold).where( - ComputedHousehold.household_id == int(household_id), - ComputedHousehold.policy_id == int(policy_id), - ComputedHousehold.country_id == country_id, - ComputedHousehold.api_version == api_version, - ) - ) - - if computed_household is not None: - return dict( - status="ok", - message=None, - result=computed_household.computed_household_json, - ) - - # Retrieve from the household table - - sessions = get_v1_session_factory() - with sessions() as session: - household = session.scalar( - select(Household).where( - Household.country_id == country_id, - Household.id == int(household_id), - ) - ) - policy = session.scalar( - select(Policy).where( - Policy.country_id == country_id, - Policy.id == int(policy_id), - ) + try: + calculation = household_calculation_service.calculate_stored_household( + country_id, + int(household_id), + int(policy_id), ) - - if household is None: + except HouseholdNotFoundError: response_body = dict( status="error", message=f"Household #{household_id} not found.", @@ -160,17 +76,7 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st mimetype="application/json", ) - # Add in any missing yearly variables - household_json = add_yearly_variables( - household.household_json, - country_id, - ) - deprecated_inputs = drop_deprecated_inputs(household_json) - household_json = deprecated_inputs.household - - # Retrieve from the policy table - - if policy is None: + except PolicyNotFoundError: response_body = dict( status="error", message=f"Policy #{policy_id} not found.", @@ -181,21 +87,17 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st mimetype="application/json", ) - country = get_countries().get(country_id) - invalid_inputs_response = get_invalid_inputs_response( - household_json, - policy.policy_json, - country, - ) - if invalid_inputs_response is not None: - return invalid_inputs_response - - try: - result = country.calculate( - household_json, - policy.policy_json, - household_id, - policy_id, + except InvalidHouseholdInputsError as error: + response_body = dict( + status="error", + message=format_unrecognized_inputs_message(error.invalid_inputs), + result=None, + errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], + ) + return Response( + json.dumps(response_body), + status=400, + mimetype="application/json", ) except Exception as e: logging.exception(e) @@ -209,34 +111,13 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st mimetype="application/json", ) - # Store the result in the computed_household table - - with get_v1_session_factory(local=True).begin() as session: - identity = (int(household_id), int(policy_id), country_id) - computed_household = session.get(ComputedHousehold, identity) - if computed_household is None: - computed_household = ComputedHousehold( - country_id=country_id, - household_id=int(household_id), - policy_id=int(policy_id), - computed_household_json=result, - api_version=api_version, - status="complete", - ) - session.add(computed_household) - else: - computed_household.computed_household_json = result - computed_household.api_version = api_version - computed_household.status = "complete" - response_body = dict( status="ok", message=None, - result=result, + result=calculation.household, ) - warning_messages = [w.message for w in deprecated_inputs.warnings] - if warning_messages: - response_body["warnings"] = warning_messages + if calculation.warnings: + response_body["warnings"] = list(calculation.warnings) return response_body @@ -273,7 +154,8 @@ def get_calculate(country_id: str, add_missing: bool = False) -> dict: return invalid_inputs_response try: - result = country.calculate(household_json, policy_json) + calculation = country.calculate(household_json, policy_json) + result = calculation if isinstance(calculation, dict) else calculation.household except SituationParsingError as e: # Malformed household payloads (e.g. a dict where a number belongs) # are client errors, not server errors — mostly bot traffic. diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py new file mode 100644 index 000000000..98c86ab5b --- /dev/null +++ b/policyengine_api/services/household_calculation_service.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from datetime import date +from typing import Any, Callable + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, + Tracer, +) +from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs +from policyengine_api.utils.input_validation import find_unrecognized_inputs + + +@dataclass(frozen=True) +class CalculationResult: + household: dict + tracer_output: list[str] + + +@dataclass(frozen=True) +class HouseholdCalculationResult: + household: dict + warnings: tuple[str, ...] = () + cached: bool = False + + +class HouseholdNotFoundError(LookupError): + pass + + +class PolicyNotFoundError(LookupError): + pass + + +class InvalidHouseholdInputsError(ValueError): + def __init__(self, invalid_inputs: list[Any]) -> None: + self.invalid_inputs = invalid_inputs + super().__init__("Household or policy contains unrecognized inputs") + + +def get_household_year(household: dict) -> int | str: + household_year: int | str = date.today().year + household_age_list = list( + household.get("people", {}).get("you", {}).get("age", {}).keys() + ) + if household_age_list: + household_year = household_age_list[0] + return household_year + + +def add_yearly_variables( + household: dict, + country_id: str, + countries: dict | None = None, +) -> dict: + if countries is None: + from policyengine_api.country import COUNTRIES + + countries = COUNTRIES + metadata = countries.get(country_id).metadata + variables = metadata["variables"] + entities = metadata["entities"] + household_year = get_household_year(household) + + for variable in variables.values(): + if variable["definitionPeriod"] not in ("year", "month", "eternity"): + continue + entity_plural = entities[variable["entity"]]["plural"] + for entity in household.get(entity_plural, {}).values(): + if variable["name"] not in entity: + entity[variable["name"]] = { + household_year: ( + variable["defaultValue"] + if variable["isInputVariable"] + else None + ) + } + return household + + +class HouseholdCalculationService: + """Orchestrate stored-household calculations with short DB scopes.""" + + def __init__( + self, + primary_session_factory: sessionmaker[Session] | None = None, + local_session_factory: sessionmaker[Session] | None = None, + country_provider: Callable[[], dict] | None = None, + ) -> None: + self._injected_primary_session_factory = primary_session_factory + self._injected_local_session_factory = local_session_factory + self._country_provider = country_provider + + @property + def _primary_sessions(self) -> sessionmaker[Session]: + return self._injected_primary_session_factory or get_v1_session_factory() + + @property + def _local_sessions(self) -> sessionmaker[Session]: + return self._injected_local_session_factory or get_v1_session_factory( + local=True + ) + + def _countries(self) -> dict: + if self._country_provider is not None: + return self._country_provider() + from policyengine_api.country import COUNTRIES + + return COUNTRIES + + def _get_cached_household( + self, + country_id: str, + household_id: int, + policy_id: int, + api_version: str, + ) -> ComputedHousehold | None: + with self._local_sessions() as session: + return session.scalar( + select(ComputedHousehold).where( + ComputedHousehold.household_id == household_id, + ComputedHousehold.policy_id == policy_id, + ComputedHousehold.country_id == country_id, + ComputedHousehold.api_version == api_version, + ) + ) + + def _get_inputs( + self, + country_id: str, + household_id: int, + policy_id: int, + ) -> tuple[Household | None, Policy | None]: + with self._primary_sessions() as session: + household = session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) + policy = session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, + ) + ) + return household, policy + + def _store_result( + self, + country_id: str, + household_id: int, + policy_id: int, + api_version: str, + calculation: CalculationResult, + ) -> None: + with self._local_sessions.begin() as session: + identity = (household_id, policy_id, country_id) + computed_household = session.get(ComputedHousehold, identity) + if computed_household is None: + computed_household = ComputedHousehold( + country_id=country_id, + household_id=household_id, + policy_id=policy_id, + computed_household_json=calculation.household, + api_version=api_version, + status="complete", + ) + session.add(computed_household) + else: + computed_household.computed_household_json = calculation.household + computed_household.api_version = api_version + computed_household.status = "complete" + if calculation.tracer_output: + session.add( + Tracer( + household_id=household_id, + policy_id=policy_id, + country_id=country_id, + api_version=api_version, + tracer_output=calculation.tracer_output, + ) + ) + + def calculate_stored_household( + self, + country_id: str, + household_id: int, + policy_id: int, + ) -> HouseholdCalculationResult: + api_version = COUNTRY_PACKAGE_VERSIONS[country_id] + cached = self._get_cached_household( + country_id, + household_id, + policy_id, + api_version, + ) + if cached is not None: + return HouseholdCalculationResult( + household=cached.computed_household_json, + cached=True, + ) + + household, policy = self._get_inputs(country_id, household_id, policy_id) + if household is None: + raise HouseholdNotFoundError(household_id) + if policy is None: + raise PolicyNotFoundError(policy_id) + + countries = self._countries() + country = countries.get(country_id) + household_json = add_yearly_variables( + deepcopy(household.household_json), + country_id, + countries, + ) + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( + household_json, + policy.policy_json, + country.metadata, + ) + if invalid_inputs: + raise InvalidHouseholdInputsError(invalid_inputs) + + raw_calculation = country.calculate(household_json, policy.policy_json) + if isinstance(raw_calculation, CalculationResult): + calculation = raw_calculation + elif hasattr(raw_calculation, "household"): + calculation = CalculationResult( + household=raw_calculation.household, + tracer_output=raw_calculation.tracer_output, + ) + else: + # Temporary compatibility for test doubles and country packages + # that have not yet adopted CalculationResult. + calculation = CalculationResult( + household=raw_calculation, + tracer_output=[], + ) + self._store_result( + country_id, + household_id, + policy_id, + api_version, + calculation, + ) + return HouseholdCalculationResult( + household=calculation.household, + warnings=tuple(warning.message for warning in deprecated_inputs.warnings), + ) diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index 307c54e65..ac4de09d2 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -1,4 +1,4 @@ -from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from typing import Generator, Literal import re import anthropic diff --git a/tests/integration/test_simulations.py b/tests/integration/test_simulations.py index 9be9ba7e2..77937a260 100644 --- a/tests/integration/test_simulations.py +++ b/tests/integration/test_simulations.py @@ -1,8 +1,4 @@ -import pytest -from unittest.mock import Mock, patch, MagicMock -import numpy as np from policyengine_api.country import COUNTRIES -from policyengine_api.endpoints.household import add_yearly_variables from tests.fixtures.integration.simulations import ( TEST_COUNTRY_ID, SMALL_AXES_COUNT, @@ -23,7 +19,7 @@ def test__given_any_number_of_axes__sim_returns_valid_arrays( base_household, small_axes_config ) country = COUNTRIES.get(TEST_COUNTRY_ID) - result = country.calculate(household_with_axes, {}) + result = country.calculate(household_with_axes, {}).household # This variable does not function like others; it is a list of member names and is not calculated FORBIDDEN_VARIABLES = ["members"] diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/endpoints/test_stage7_orm_endpoints.py index f9450b380..4f5bb80e5 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/endpoints/test_stage7_orm_endpoints.py @@ -17,6 +17,9 @@ set_user_policy, update_user_policy, ) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) def test_household_under_policy_returns_cached_json_object(orm_session_factory): @@ -33,11 +36,7 @@ def test_household_under_policy_returns_cached_json_object(orm_session_factory): ) ) - with patch( - "policyengine_api.endpoints.household.get_v1_session_factory", - return_value=orm_session_factory, - ): - response = get_household_under_policy("us", "1", "2") + response = get_household_under_policy("us", "1", "2") assert response["result"] == {"people": {"you": {"net_income": {"2026": 42}}}} @@ -67,32 +66,23 @@ def test_household_under_policy_calculates_and_caches_json_as_an_object( ] ) calculated = {"people": {"you": {"net_income": {"2026": 42}}}} - country = SimpleNamespace(calculate=Mock(return_value=calculated)) + country = SimpleNamespace( + calculate=Mock(return_value=calculated), + metadata={ + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + "parameters": {"gov.example.parameter": {}}, + }, + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) - with ( - patch( - "policyengine_api.endpoints.household.get_v1_session_factory", - return_value=orm_session_factory, - ), - patch( - "policyengine_api.endpoints.household.add_yearly_variables", - side_effect=lambda household, _: household, - ), - patch( - "policyengine_api.endpoints.household.drop_deprecated_inputs", - side_effect=lambda household: SimpleNamespace( - household=household, - warnings=[], - ), - ), - patch( - "policyengine_api.endpoints.household.get_invalid_inputs_response", - return_value=None, - ), - patch( - "policyengine_api.endpoints.household.get_countries", - return_value={"us": country}, - ), + with patch( + "policyengine_api.endpoints.household.household_calculation_service", + service, ): response = get_household_under_policy("us", "1", "2") @@ -100,8 +90,6 @@ def test_household_under_policy_calculates_and_caches_json_as_an_object( country.calculate.assert_called_once_with( {"people": {"you": {}}}, {"gov.example.parameter": 1}, - "1", - "2", ) with orm_session_factory() as session: cached = session.scalar(select(ComputedHousehold)) diff --git a/tests/unit/services/test_household_calculation_service.py b/tests/unit/services/test_household_calculation_service.py new file mode 100644 index 000000000..191b9aec5 --- /dev/null +++ b/tests/unit/services/test_household_calculation_service.py @@ -0,0 +1,189 @@ +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +from sqlalchemy import select + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, + Tracer, +) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +class TrackingSessionFactory: + def __init__(self, factory): + self.factory = factory + self.active_scopes = 0 + + @contextmanager + def __call__(self): + self.active_scopes += 1 + try: + with self.factory() as session: + yield session + finally: + self.active_scopes -= 1 + + @contextmanager + def begin(self): + self.active_scopes += 1 + try: + with self.factory.begin() as session: + yield session + finally: + self.active_scopes -= 1 + + +def _seed_inputs(factory): + with factory.begin() as session: + session.add_all( + [ + Household( + id=1, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + household_json={"people": {"you": {}}}, + household_hash="household-hash", + ), + Policy( + id=2, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + policy_json={}, + policy_hash="policy-hash", + ), + ] + ) + + +def test_household_endpoint_and_country_do_not_manage_persistence(): + endpoint_source = (PACKAGE_ROOT / "endpoints" / "household.py").read_text( + encoding="utf-8" + ) + country_source = (PACKAGE_ROOT / "country.py").read_text(encoding="utf-8") + assert "get_v1_session_factory" not in endpoint_source + assert "from sqlalchemy" not in endpoint_source + assert "select(" not in endpoint_source + assert "get_v1_session_factory" not in country_source + assert "Tracer(" not in country_source + + +def test_calculation_closes_reads_before_compute_and_persists_local_results( + orm_session_factory, +): + _seed_inputs(orm_session_factory) + primary = TrackingSessionFactory(orm_session_factory) + local = TrackingSessionFactory(orm_session_factory) + + class Country: + metadata = { + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + } + + def calculate(self, household, policy): + assert primary.active_scopes == 0 + assert local.active_scopes == 0 + return SimpleNamespace( + household={"people": {"you": {"net_income": {"2026": 42}}}}, + tracer_output=["net_income <2026>"], + ) + + service = HouseholdCalculationService( + primary_session_factory=primary, + local_session_factory=local, + country_provider=lambda: {"us": Country()}, + ) + + result = service.calculate_stored_household("us", 1, 2) + + assert result.household["people"]["you"]["net_income"]["2026"] == 42 + with orm_session_factory() as session: + cached = session.scalar(select(ComputedHousehold)) + tracer = session.scalar(select(Tracer)) + assert cached.computed_household_json == result.household + assert tracer.tracer_output == ["net_income <2026>"] + + +def test_calculation_uses_local_cache_without_recomputing(orm_session_factory): + _seed_inputs(orm_session_factory) + calculated = {"people": {"you": {"net_income": {"2026": 42}}}} + with orm_session_factory.begin() as session: + session.add( + ComputedHousehold( + household_id=1, + policy_id=2, + country_id="us", + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + computed_household_json=calculated, + status="complete", + ) + ) + country = SimpleNamespace( + metadata={"variables": {}, "entities": {}}, + calculate=lambda *_: (_ for _ in ()).throw( + AssertionError("cache hit should not calculate") + ), + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) + + result = service.calculate_stored_household("us", 1, 2) + + assert result.household == calculated + assert result.cached is True + + +def test_local_computed_household_and_tracer_write_roll_back_together( + orm_session_factory, + monkeypatch, +): + _seed_inputs(orm_session_factory) + country = SimpleNamespace( + metadata={ + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + }, + calculate=lambda *_: SimpleNamespace( + household={"people": {"you": {}}}, + tracer_output=["trace"], + ), + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) + session_type = orm_session_factory.class_ + original_flush = session_type.flush + + def fail_tracer_flush(session, *args, **kwargs): + has_tracer = any(isinstance(value, Tracer) for value in session.new) + original_flush(session, *args, **kwargs) + if has_tracer: + raise RuntimeError("tracer insert failed") + + monkeypatch.setattr(session_type, "flush", fail_tracer_flush) + + with pytest.raises(RuntimeError, match="tracer insert failed"): + service.calculate_stored_household("us", 1, 2) + + monkeypatch.setattr(session_type, "flush", original_flush) + with orm_session_factory() as session: + assert session.scalar(select(ComputedHousehold)) is None + assert session.scalar(select(Tracer)) is None From 63c3acd084e0e7428c7d1b0293423952f90ade84 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:26:36 +0300 Subject: [PATCH 63/89] refactor: inject persistence into economy and analysis services --- .../endpoints/economy/reform_impact.py | 30 --- policyengine_api/endpoints/simulation.py | 13 +- .../routes/simulation_analysis_routes.py | 2 - .../routes/tracer_analysis_routes.py | 2 - .../services/ai_analysis_service.py | 32 ++- policyengine_api/services/economy_service.py | 197 ++++++++++-------- .../services/reform_impacts_service.py | 195 ++++++++++++++++- .../services/simulation_analysis_service.py | 7 +- .../services/tracer_analysis_service.py | 41 ++-- tests/fixtures/services/economy_service.py | 16 +- .../services/tracer_fixture_service.py | 2 +- .../test_budget_window_in_flight_dedupe.py | 8 +- tests/unit/data/test_stage7_no_direct_sql.py | 1 - .../unit/services/test_ai_analysis_service.py | 71 ++++++- .../test_direct_orm_local_analysis.py | 23 +- tests/unit/services/test_economy_service.py | 16 +- tests/unit/services/test_execute_analysis.py | 13 +- .../services/test_reform_impacts_service.py | 79 ++++--- .../test_stage7_local_service_boundaries.py | 65 +++++- tests/unit/services/test_tracer_service.py | 19 +- 20 files changed, 580 insertions(+), 252 deletions(-) delete mode 100644 policyengine_api/endpoints/economy/reform_impact.py diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py deleted file mode 100644 index 5d8641d3c..000000000 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ /dev/null @@ -1,30 +0,0 @@ -from sqlalchemy import select - -from policyengine_api.data.orm import get_v1_session_factory -from policyengine_api.data.v1_models import ReformImpact - - -def set_comment_on_job( - comment: str, - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options_hash, -): - with get_v1_session_factory(local=True).begin() as session: - impacts = session.scalars( - select(ReformImpact).where( - ReformImpact.country_id == country_id, - ReformImpact.reform_policy_id == policy_id, - ReformImpact.baseline_policy_id == baseline_policy_id, - ReformImpact.region == region, - ReformImpact.time_period == time_period, - ReformImpact.options_hash == options_hash, - ReformImpact.dataset == dataset, - ) - ) - for impact in impacts: - impact.message = comment diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py index 08fa9c52d..7e83dee31 100644 --- a/policyengine_api/endpoints/simulation.py +++ b/policyengine_api/endpoints/simulation.py @@ -1,7 +1,5 @@ -from sqlalchemy import select - -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ReformImpact +from policyengine_api.services.reform_impacts_service import ReformImpactsService """ @@ -26,6 +24,7 @@ _MAX_SIMULATION_RESULTS = 1000 _DEFAULT_SIMULATION_RESULTS = 100 +reform_impacts_service = ReformImpactsService() def get_simulations( @@ -45,13 +44,7 @@ def get_simulations( max_results = _DEFAULT_SIMULATION_RESULTS max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - sessions = get_v1_session_factory(local=True) - with sessions() as session: - result = session.scalars( - select(ReformImpact) - .order_by(ReformImpact.start_time.desc()) - .limit(max_results) - ).all() + result = reform_impacts_service.get_recent_reform_impacts(max_results) # Format into [{}] diff --git a/policyengine_api/routes/simulation_analysis_routes.py b/policyengine_api/routes/simulation_analysis_routes.py index a3094292a..d9ddc0b97 100644 --- a/policyengine_api/routes/simulation_analysis_routes.py +++ b/policyengine_api/routes/simulation_analysis_routes.py @@ -10,7 +10,6 @@ validate_sim_analysis_payload, ) import json -from policyengine_api.data.orm import get_v1_session_factory simulation_analysis_bp = Blueprint("simulation_analysis", __name__) simulation_analysis_service = SimulationAnalysisService() @@ -44,7 +43,6 @@ def execute_simulation_analysis(country_id): audience = payload.get("audience", "") analysis, analysis_type = simulation_analysis_service.execute_analysis( - get_v1_session_factory(local=True), country_id, currency, dataset, diff --git a/policyengine_api/routes/tracer_analysis_routes.py b/policyengine_api/routes/tracer_analysis_routes.py index 9991e5514..43695d45c 100644 --- a/policyengine_api/routes/tracer_analysis_routes.py +++ b/policyengine_api/routes/tracer_analysis_routes.py @@ -8,7 +8,6 @@ TracerAnalysisService, ) import json -from policyengine_api.data.orm import get_v1_session_factory tracer_analysis_bp = Blueprint("tracer_analysis", __name__) tracer_analysis_service = TracerAnalysisService() @@ -30,7 +29,6 @@ def execute_tracer_analysis(country_id): raise BadRequest("variable must be a string") analysis, analysis_type = tracer_analysis_service.execute_analysis( - get_v1_session_factory(local=True), country_id, household_id, policy_id, diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index c34a9dbfd..6dd0bee93 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -1,12 +1,14 @@ import json import os from collections.abc import Generator +from typing import Callable import anthropic from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import Analysis @@ -25,10 +27,29 @@ class ErrorEvent(StreamEvent): class AIAnalysisService: - """AI analysis operations backed by caller-owned ORM sessions.""" + """AI analysis operations with short, service-owned ORM scopes.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + claude_client_factory: Callable[[], anthropic.Anthropic] | None = None, + ) -> None: + self._injected_session_factory = session_factory + self._claude_client_factory = claude_client_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory(local=True) def get_existing_analysis( self, + prompt: str, + ) -> Analysis | None: + with self._sessions() as session: + return self._get_existing_analysis(session, prompt) + + @staticmethod + def _get_existing_analysis( session: Session, prompt: str, ) -> Analysis | None: @@ -44,9 +65,12 @@ def get_existing_analysis( def trigger_ai_analysis( self, prompt: str, - session_factory: sessionmaker[Session], ) -> Generator[str, None, None]: - claude_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + claude_client = ( + self._claude_client_factory() + if self._claude_client_factory is not None + else anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + ) def generate(): response_text = "" @@ -71,7 +95,7 @@ def generate(): yield ( json.dumps(TextEvent(stream=event.text).model_dump()) + "\n" ) - with session_factory.begin() as session: + with self._sessions.begin() as session: session.add( Analysis( prompt=prompt, diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index d0b2af3c9..742d0f8bc 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -22,7 +22,6 @@ get_valid_state_codes, normalize_us_region, ) -from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ReformImpact from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger @@ -37,8 +36,6 @@ load_dotenv() -policy_service = PolicyService() -reform_impacts_service = ReformImpactsService() budget_window_cache = BudgetWindowCache() @@ -237,17 +234,56 @@ class EconomyService: with other services to access their respective tables """ + def __init__( + self, + *, + primary_session_factory=None, + local_session_factory=None, + policy_service_: PolicyService | None = None, + reform_impacts_service_: ReformImpactsService | None = None, + budget_window_cache_: BudgetWindowCache | None = None, + simulation_entrypoint_=None, + ) -> None: + self._primary_session_factory = primary_session_factory + self._local_session_factory = local_session_factory + self._injected_policy_service = policy_service_ + self._injected_reform_impacts_service = reform_impacts_service_ + self._injected_budget_window_cache = budget_window_cache_ + self._injected_simulation_entrypoint = simulation_entrypoint_ + + @property + def _policies(self) -> PolicyService: + if self._injected_policy_service is None: + self._injected_policy_service = PolicyService(self._primary_session_factory) + return self._injected_policy_service + + @property + def _reform_impacts(self) -> ReformImpactsService: + if self._injected_reform_impacts_service is None: + self._injected_reform_impacts_service = ReformImpactsService( + self._local_session_factory + ) + return self._injected_reform_impacts_service + + @property + def _budget_window_cache(self) -> BudgetWindowCache: + return self._injected_budget_window_cache or budget_window_cache + + @property + def _simulation_gateway(self): + return self._injected_simulation_entrypoint or simulation_entrypoint + def _get_policy_jsons( self, country_id: str, baseline_policy_id: int, reform_policy_id: int, ) -> tuple[dict | None, dict | None]: - baseline = policy_service.get_policy_json( + baseline = self._policies.get_policy_json( country_id, baseline_policy_id, ) - reform = policy_service.get_policy_json( + reform = self._policies.get_policy_json( country_id, reform_policy_id, ) @@ -347,14 +383,14 @@ def get_budget_window_economic_impact( ) cache_key = self._build_budget_window_cache_key(setup_options) - cached_result = budget_window_cache.get_completed_result(cache_key) + cached_result = self._budget_window_cache.get_completed_result(cache_key) if cached_result is not None: return BudgetWindowEconomicImpactResult.completed( cached_result, cache_status="result-hit", ) - batch_job_id = budget_window_cache.get_batch_job_id(cache_key) + batch_job_id = self._budget_window_cache.get_batch_job_id(cache_key) if batch_job_id: return self._get_budget_window_result_from_batch_job_id( batch_job_id=batch_job_id, @@ -366,7 +402,7 @@ def get_budget_window_economic_impact( claim_token = setup_options.process_id cache_status = "starting-claim-hit" - if budget_window_cache.claim_batch_start(cache_key, claim_token): + if self._budget_window_cache.claim_batch_start(cache_key, claim_token): cache_status = "miss" try: batch_execution = self._start_budget_window_batch( @@ -375,11 +411,13 @@ def get_budget_window_economic_impact( window_size=window_size, max_parallel=max_active_years, ) - budget_window_cache.store_batch_job_id( + self._budget_window_cache.store_batch_job_id( cache_key, batch_execution.batch_job_id ) except httpx.HTTPStatusError as error: - budget_window_cache.clear_starting_claim(cache_key, claim_token) + self._budget_window_cache.clear_starting_claim( + cache_key, claim_token + ) if ( error.response.status_code in BUDGET_WINDOW_SUBMISSION_VALIDATION_ERROR_STATUS_CODES @@ -391,7 +429,9 @@ def get_budget_window_economic_impact( ) raise except Exception: - budget_window_cache.clear_starting_claim(cache_key, claim_token) + self._budget_window_cache.clear_starting_claim( + cache_key, claim_token + ) raise return self._build_budget_window_computing_result( @@ -410,7 +450,7 @@ def _build_budget_window_cache_key( self, setup_options: EconomicImpactSetupOptions, ) -> str: - return budget_window_cache.build_key( + return self._budget_window_cache.build_key( country_id=setup_options.country_id, reform_policy_id=setup_options.reform_policy_id, baseline_policy_id=setup_options.baseline_policy_id, @@ -481,7 +521,7 @@ def _start_budget_window_batch( severity="INFO", ) - return simulation_entrypoint.run_budget_window_batch(sim_params) + return self._simulation_gateway.run_budget_window_batch(sim_params) def _build_budget_window_submission_error_message( self, error: httpx.HTTPStatusError @@ -512,14 +552,14 @@ def _get_budget_window_result_from_batch_job_id( queued_years_on_submit: list[str], cache_status: Optional[str] = None, ) -> BudgetWindowEconomicImpactResult: - batch_execution = simulation_entrypoint.get_budget_window_batch_by_id( + batch_execution = self._simulation_gateway.get_budget_window_batch_by_id( batch_job_id ) if batch_execution.status in EXECUTION_STATUSES_SUCCESS: result = batch_execution.result if not isinstance(result, dict) or not result: - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.failed( "Budget-window batch completed without a result", completed_years=batch_execution.completed_years, @@ -527,8 +567,8 @@ def _get_budget_window_result_from_batch_job_id( queued_years=batch_execution.queued_years or queued_years_on_submit, cache_status=cache_status, ) - budget_window_cache.set_completed_result(cache_key, result) - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.set_completed_result(cache_key, result) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.completed( result, cache_status=cache_status, @@ -536,7 +576,7 @@ def _get_budget_window_result_from_batch_job_id( if batch_execution.status in EXECUTION_STATUSES_FAILURE: error_message = batch_execution.error or "Budget-window batch failed" - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.failed( error_message, completed_years=batch_execution.completed_years, @@ -715,7 +755,7 @@ def _resolve_runtime_bundle_for_setup_options( ( setup_options.runtime_app_name, setup_options.model_version, - ) = simulation_entrypoint.resolve_app_name( + ) = self._simulation_gateway.resolve_app_name( setup_options.country_id, setup_options.model_version, policyengine_version=setup_options.policyengine_version, @@ -766,22 +806,19 @@ def _get_previous_impacts( """ previous_impacts: list[Any] = [] - sessions = get_v1_session_factory(local=True) - with sessions() as session: - previous_impacts = ( - reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix( - session, - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options_hash, - self._build_options_hash_lookup_pattern(options_hash), - api_version, - ) + previous_impacts = ( + self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + self._build_options_hash_lookup_pattern(options_hash), + api_version, ) + ) return previous_impacts def _get_most_recent_impact( @@ -842,7 +879,7 @@ def _handle_execution_state( """ if execution_state in EXECUTION_STATUSES_SUCCESS: result = self._with_policyengine_bundle( - result=simulation_entrypoint.get_execution_result(execution), + result=self._simulation_gateway.get_execution_result(execution), setup_options=setup_options, execution=execution, ) @@ -908,10 +945,10 @@ def _handle_computing_impact( setup_options: EconomicImpactSetupOptions, most_recent_impact: ReformImpact, ) -> EconomicImpactResult: - execution = simulation_entrypoint.get_execution_by_id( + execution = self._simulation_gateway.get_execution_by_id( most_recent_impact.execution_id ) - execution_state = simulation_entrypoint.get_execution_status(execution) + execution_state = self._simulation_gateway.get_execution_status(execution) return self._handle_execution_state( execution_state=execution_state, setup_options=setup_options, @@ -979,8 +1016,8 @@ def _handle_create_impact( if sim_params.get("time_period") is not None: sim_params["time_period"] = str(sim_params["time_period"]) - entrypoint_execution = simulation_entrypoint.run(sim_params) - execution_id = simulation_entrypoint.get_execution_id(entrypoint_execution) + entrypoint_execution = self._simulation_gateway.run(sim_params) + execution_id = self._simulation_gateway.get_execution_id(entrypoint_execution) run_id = getattr(entrypoint_execution, "run_id", None) or telemetry["run_id"] @@ -1090,7 +1127,7 @@ def _should_refresh_cached_impact( cached_resolved_app_name = cached_result.get("resolved_app_name") try: runtime_app_name, resolved_model_version = ( - simulation_entrypoint.resolve_app_name( + self._simulation_gateway.resolve_app_name( setup_options.country_id, setup_options.model_version, policyengine_version=setup_options.policyengine_version, @@ -1334,23 +1371,21 @@ def _set_reform_impact_computing( In the reform_impact table, set the status of the impact to "computing". """ try: - with get_v1_session_factory(local=True).begin() as session: - reform_impacts_service.set_reform_impact( - session, - country_id=setup_options.country_id, - policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options=setup_options.options, - options_hash=setup_options.options_hash, - status=ImpactStatus.COMPUTING.value, - api_version=setup_options.api_version, - reform_impact_json={}, - start_time=datetime.datetime.now(), - execution_id=execution_id, - ) + self._reform_impacts.set_reform_impact( + country_id=setup_options.country_id, + policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options=setup_options.options, + options_hash=setup_options.options_hash, + status=ImpactStatus.COMPUTING.value, + api_version=setup_options.api_version, + reform_impact_json={}, + start_time=datetime.datetime.now(), + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { @@ -1370,19 +1405,17 @@ def _set_reform_impact_complete( In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON. """ try: - with get_v1_session_factory(local=True).begin() as session: - reform_impacts_service.set_complete_reform_impact( - session, - country_id=setup_options.country_id, - reform_policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options_hash=setup_options.options_hash, - reform_impact_json=reform_impact_json, - execution_id=execution_id, - ) + self._reform_impacts.set_complete_reform_impact( + country_id=setup_options.country_id, + reform_policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options_hash=setup_options.options_hash, + reform_impact_json=reform_impact_json, + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { @@ -1402,19 +1435,17 @@ def _set_reform_impact_error( In the reform_impact table, set the status of the impact to "error" and store the error message. """ try: - with get_v1_session_factory(local=True).begin() as session: - reform_impacts_service.set_error_reform_impact( - session, - country_id=setup_options.country_id, - policy_id=setup_options.reform_policy_id, - baseline_policy_id=setup_options.baseline_policy_id, - region=setup_options.region, - dataset=setup_options.dataset, - time_period=setup_options.time_period, - options_hash=setup_options.options_hash, - message=message, - execution_id=execution_id, - ) + self._reform_impacts.set_error_reform_impact( + country_id=setup_options.country_id, + policy_id=setup_options.reform_policy_id, + baseline_policy_id=setup_options.baseline_policy_id, + region=setup_options.region, + dataset=setup_options.dataset, + time_period=setup_options.time_period, + options_hash=setup_options.options_hash, + message=message, + execution_id=execution_id, + ) except Exception as e: logger.log_struct( { diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 641afb3ed..d338e8687 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -2,13 +2,192 @@ from typing import Any from sqlalchemy import delete, or_, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker +from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ReformImpact class ReformImpactsService: - """Reform-impact operations performed through a caller-owned Session.""" + """Reform-impact operations with service-owned local transactions.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory(local=True) + + def get_recent_reform_impacts(self, max_results: int) -> list[ReformImpact]: + with self._sessions() as session: + return list( + session.scalars( + select(ReformImpact) + .order_by(ReformImpact.start_time.desc()) + .limit(max_results) + ) + ) + + def get_all_reform_impacts( + self, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + api_version, + ) -> list[ReformImpact]: + with self._sessions() as session: + return self._get_all_reform_impacts( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + api_version, + ) + + def get_all_reform_impacts_by_options_hash_prefix( + self, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + options_hash_prefix, + api_version, + ) -> list[ReformImpact]: + with self._sessions() as session: + return self._get_all_reform_impacts_by_options_hash_prefix( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + options_hash_prefix, + api_version, + ) + + def set_reform_impact( + self, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options: dict[str, Any], + options_hash, + status, + api_version, + reform_impact_json: dict[str, Any], + start_time, + execution_id: str, + ) -> ReformImpact: + with self._sessions.begin() as session: + return self._set_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options, + options_hash, + status, + api_version, + reform_impact_json, + start_time, + execution_id, + ) + + def delete_reform_impact( + self, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) -> None: + with self._sessions.begin() as session: + self._delete_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) + + def set_error_reform_impact( + self, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + message, + execution_id: str, + ) -> ReformImpact | None: + with self._sessions.begin() as session: + return self._set_error_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + message, + execution_id, + ) + + def set_complete_reform_impact( + self, + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + reform_impact_json: dict[str, Any], + execution_id, + ) -> ReformImpact | None: + with self._sessions.begin() as session: + return self._set_complete_reform_impact( + session, + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + reform_impact_json, + execution_id, + ) @staticmethod def _filters( @@ -38,7 +217,7 @@ def _scope(statement, **filters): *(getattr(ReformImpact, key) == value for key, value in filters.items()) ) - def get_all_reform_impacts( + def _get_all_reform_impacts( self, session: Session, country_id, @@ -64,7 +243,7 @@ def get_all_reform_impacts( ) return list(session.scalars(statement.order_by(ReformImpact.start_time.desc()))) - def get_all_reform_impacts_by_options_hash_prefix( + def _get_all_reform_impacts_by_options_hash_prefix( self, session: Session, country_id, @@ -101,7 +280,7 @@ def get_all_reform_impacts_by_options_hash_prefix( ) ) - def set_reform_impact( + def _set_reform_impact( self, session: Session, country_id, @@ -137,7 +316,7 @@ def set_reform_impact( session.flush() return impact - def delete_reform_impact( + def _delete_reform_impact( self, session: Session, country_id, @@ -163,7 +342,7 @@ def delete_reform_impact( ) ) - def set_error_reform_impact( + def _set_error_reform_impact( self, session: Session, country_id, @@ -197,7 +376,7 @@ def set_error_reform_impact( impact.end_time = self._now() return impact - def set_complete_reform_impact( + def _set_complete_reform_impact( self, session: Session, country_id, diff --git a/policyengine_api/services/simulation_analysis_service.py b/policyengine_api/services/simulation_analysis_service.py index 1204a7495..bb6382c87 100644 --- a/policyengine_api/services/simulation_analysis_service.py +++ b/policyengine_api/services/simulation_analysis_service.py @@ -1,7 +1,6 @@ from policyengine_api.services.ai_analysis_service import AIAnalysisService from policyengine_api.services.ai_prompt_service import AIPromptService from typing import Generator, Literal -from sqlalchemy.orm import Session, sessionmaker ai_prompt_service = AIPromptService() @@ -15,7 +14,6 @@ class SimulationAnalysisService(AIAnalysisService): def execute_analysis( self, - session_factory: sessionmaker[Session], country_id: str, currency: str, dataset: str | None, @@ -60,15 +58,14 @@ def execute_analysis( print("Checking if AI analysis already exists for this prompt") # If a calculated record exists for this prompt, return it as a # streaming response - with session_factory() as session: - existing_analysis = self.get_existing_analysis(session, prompt) + existing_analysis = self.get_existing_analysis(prompt) if existing_analysis is not None: return existing_analysis.analysis, "static" print("Found no existing AI analysis; triggering new analysis with Claude") # Otherwise, pass prompt to Claude, then return streaming function try: - analysis = self.trigger_ai_analysis(prompt, session_factory) + analysis = self.trigger_ai_analysis(prompt) return analysis, "streaming" except Exception as e: raise e diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index ac4de09d2..ff6573d89 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -5,7 +5,6 @@ from policyengine_api.services.ai_analysis_service import AIAnalysisService from werkzeug.exceptions import NotFound from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker from policyengine_api.data.v1_models import Tracer @@ -13,7 +12,6 @@ class TracerAnalysisService(AIAnalysisService): def execute_analysis( self, - session_factory: sessionmaker[Session], country_id: str, household_id: str, policy_id: str, @@ -31,14 +29,12 @@ def execute_analysis( # Retrieve tracer record from table try: - with session_factory() as session: - tracer: list[str] = self.get_tracer( - session, - country_id, - household_id, - policy_id, - api_version, - ) + tracer: list[str] = self.get_tracer( + country_id, + household_id, + policy_id, + api_version, + ) except Exception as e: raise e @@ -58,14 +54,13 @@ def execute_analysis( ) # If a calculated record exists for this prompt, return it as a string - with session_factory() as session: - existing_analysis = self.get_existing_analysis(session, prompt) + existing_analysis = self.get_existing_analysis(prompt) if existing_analysis is not None: return existing_analysis.analysis, "static" # Otherwise, pass prompt to Claude, then return streaming function try: - analysis: Generator = self.trigger_ai_analysis(prompt, session_factory) + analysis: Generator = self.trigger_ai_analysis(prompt) return analysis, "streaming" except Exception as e: print( @@ -75,23 +70,23 @@ def execute_analysis( def get_tracer( self, - session: Session, country_id: str, household_id: str, policy_id: str, api_version: str, ) -> list: try: - tracer = session.scalar( - select(Tracer) - .where( - Tracer.household_id == int(household_id), - Tracer.policy_id == int(policy_id), - Tracer.country_id == country_id, - Tracer.api_version == api_version, + with self._sessions() as session: + tracer = session.scalar( + select(Tracer) + .where( + Tracer.household_id == int(household_id), + Tracer.policy_id == int(policy_id), + Tracer.country_id == country_id, + Tracer.api_version == api_version, + ) + .order_by(Tracer.id.desc()) ) - .order_by(Tracer.id.desc()) - ) if tracer is None: raise NotFound("No household simulation tracer found") diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index e0a264644..6b16019b5 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -105,10 +105,10 @@ def mock_policy_service(): ) with patch( - "policyengine_api.services.economy_service.policy_service", - mock_service, - ) as mock: - yield mock + "policyengine_api.services.economy_service.PolicyService", + return_value=mock_service, + ): + yield mock_service @pytest.fixture @@ -122,10 +122,10 @@ def mock_reform_impacts_service(): mock_service.set_error_reform_impact.return_value = None with patch( - "policyengine_api.services.economy_service.reform_impacts_service", - mock_service, - ) as mock: - yield mock + "policyengine_api.services.economy_service.ReformImpactsService", + return_value=mock_service, + ): + yield mock_service @pytest.fixture diff --git a/tests/fixtures/services/tracer_fixture_service.py b/tests/fixtures/services/tracer_fixture_service.py index f22bc6ea4..ef34553cf 100644 --- a/tests/fixtures/services/tracer_fixture_service.py +++ b/tests/fixtures/services/tracer_fixture_service.py @@ -32,5 +32,5 @@ def test_tracer_data(orm_session): tracer_output=json.loads(valid_tracer_row["tracer_output"]), ) orm_session.add(tracer) - orm_session.flush() + orm_session.commit() return tracer diff --git a/tests/integration/test_budget_window_in_flight_dedupe.py b/tests/integration/test_budget_window_in_flight_dedupe.py index 14a273214..183483d5d 100644 --- a/tests/integration/test_budget_window_in_flight_dedupe.py +++ b/tests/integration/test_budget_window_in_flight_dedupe.py @@ -35,10 +35,12 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( from policyengine_api.libs.simulation_entrypoint import ( ModalBudgetWindowBatchExecution, ) - from policyengine_api.routes.economy_routes import economy_bp + from policyengine_api.routes import economy_routes from policyengine_api.services import economy_service as economy_service_module from policyengine_api.services.budget_window_cache import BudgetWindowCache + economy_bp = economy_routes.economy_bp + fake_cache = BudgetWindowCache(client=FakeRedis()) simulation_entrypoint = MagicMock() reform_impacts_service = MagicMock() @@ -65,8 +67,8 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( economy_service_module, "simulation_entrypoint", simulation_entrypoint ) monkeypatch.setattr( - economy_service_module, - "reform_impacts_service", + economy_routes.economy_service, + "_injected_reform_impacts_service", reform_impacts_service, ) monkeypatch.setattr( diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index cb9b6ca48..0086c37de 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -28,7 +28,6 @@ def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): "endpoints/household.py", "endpoints/policy.py", "endpoints/simulation.py", - "endpoints/economy/reform_impact.py", "country.py", "services/ai_analysis_service.py", "services/reform_impacts_service.py", diff --git a/tests/unit/services/test_ai_analysis_service.py b/tests/unit/services/test_ai_analysis_service.py index b39854d59..13867f6e3 100644 --- a/tests/unit/services/test_ai_analysis_service.py +++ b/tests/unit/services/test_ai_analysis_service.py @@ -1,18 +1,77 @@ import json +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest from sqlalchemy import select from policyengine_api.data.v1_models import Analysis from policyengine_api.services.ai_analysis_service import AIAnalysisService from tests.fixtures.services.ai_analysis_service import parse_to_chunks -import pytest pytest_plugins = ["tests.fixtures.services.ai_analysis_service"] -# Initialize the service -service = AIAnalysisService() - class TestTriggerAIAnalysis: + def test_claude_stream_runs_without_an_open_database_session( + self, + orm_session_factory, + ): + class TrackingSessions: + def __init__(self, delegate): + self.delegate = delegate + self.active = 0 + + @contextmanager + def __call__(self): + self.active += 1 + try: + with self.delegate() as session: + yield session + finally: + self.active -= 1 + + @contextmanager + def begin(self): + self.active += 1 + try: + with self.delegate.begin() as session: + yield session + finally: + self.active -= 1 + + sessions = TrackingSessions(orm_session_factory) + + class ClaudeStream: + def __enter__(self): + assert sessions.active == 0 + return self + + def __exit__(self, *args): + return None + + def __iter__(self): + assert sessions.active == 0 + yield SimpleNamespace(type="text", text="analysis") + + claude_client = SimpleNamespace( + messages=SimpleNamespace(stream=lambda **kwargs: ClaudeStream()) + ) + service = AIAnalysisService( + sessions, + claude_client_factory=lambda: claude_client, + ) + + assert list(service.trigger_ai_analysis("prompt")) == [ + json.dumps({"type": "text", "stream": "analysis"}) + "\n" + ] + assert sessions.active == 0 + + with orm_session_factory() as session: + stored = session.scalar(select(Analysis).where(Analysis.prompt == "prompt")) + assert stored is not None + assert stored.analysis == "analysis" + def test_trigger_ai_analysis_given_successful_streaming( self, mock_stream_text_events, orm_session_factory ): @@ -23,7 +82,7 @@ def test_trigger_ai_analysis_given_successful_streaming( # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote" - generator = service.trigger_ai_analysis(prompt, orm_session_factory) + generator = AIAnalysisService(orm_session_factory).trigger_ai_analysis(prompt) # THEN it should yield the expected chunks results = list(generator) @@ -62,7 +121,7 @@ def test_trigger_ai_analysis_given_error( # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote about erroneous systems" - generator = service.trigger_ai_analysis(prompt, orm_session_factory) + generator = AIAnalysisService(orm_session_factory).trigger_ai_analysis(prompt) # THEN it should yield the expected error message results = list(generator) diff --git a/tests/unit/services/test_direct_orm_local_analysis.py b/tests/unit/services/test_direct_orm_local_analysis.py index 2207f2c41..f60af4e31 100644 --- a/tests/unit/services/test_direct_orm_local_analysis.py +++ b/tests/unit/services/test_direct_orm_local_analysis.py @@ -6,7 +6,10 @@ from policyengine_api.services.tracer_analysis_service import TracerAnalysisService -def test_ai_analysis_service_returns_the_latest_mapped_analysis(orm_session): +def test_ai_analysis_service_returns_the_latest_mapped_analysis( + orm_session, + orm_session_factory, +): orm_session.add_all( [ Analysis(prompt="prompt", analysis="old", status="ok"), @@ -15,15 +18,15 @@ def test_ai_analysis_service_returns_the_latest_mapped_analysis(orm_session): ) orm_session.flush() - analysis = AIAnalysisService().get_existing_analysis(orm_session, "prompt") + orm_session.commit() + analysis = AIAnalysisService(orm_session_factory).get_existing_analysis("prompt") assert isinstance(analysis, Analysis) assert analysis.analysis == "new" -def test_reform_impact_service_writes_mapped_entity(orm_session): - impact = ReformImpactsService().set_reform_impact( - orm_session, +def test_reform_impact_service_writes_mapped_entity(orm_session_factory): + impact = ReformImpactsService(orm_session_factory).set_reform_impact( country_id="us", policy_id=2, baseline_policy_id=1, @@ -43,7 +46,10 @@ def test_reform_impact_service_writes_mapped_entity(orm_session): assert impact.options_json == {"dataset": "default"} -def test_tracer_service_reads_python_json_from_mapped_entity(orm_session): +def test_tracer_service_reads_python_json_from_mapped_entity( + orm_session, + orm_session_factory, +): orm_session.add( Tracer( household_id=1, @@ -53,10 +59,9 @@ def test_tracer_service_reads_python_json_from_mapped_entity(orm_session): tracer_output=["net_income <2026>", " dependency"], ) ) - orm_session.flush() + orm_session.commit() - tracer = TracerAnalysisService().get_tracer( - orm_session, + tracer = TracerAnalysisService(orm_session_factory).get_tracer( "us", "1", "2", diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 74092498d..4926ac737 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -1,6 +1,6 @@ import json from typing import Literal -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch import httpx import pytest @@ -406,7 +406,6 @@ def test__given_runtime_cache_version__uses_versioned_economy_cache_key( economy_service.get_economic_impact(**base_params) mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.assert_called_once_with( - ANY, MOCK_COUNTRY_ID, MOCK_POLICY_ID, MOCK_BASELINE_POLICY_ID, @@ -438,14 +437,14 @@ def test__given_alias_dataset__queries_previous_impacts_with_resolved_bundle( economy_service.get_economic_impact(**base_params) call_args = mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.call_args.args - assert call_args[5] == MOCK_RESOLVED_DATASET - assert call_args[7] == MOCK_LOOKUP_OPTIONS_HASH - assert call_args[8] == economy_service._build_options_hash_lookup_pattern( + assert call_args[4] == MOCK_RESOLVED_DATASET + assert call_args[6] == MOCK_LOOKUP_OPTIONS_HASH + assert call_args[7] == economy_service._build_options_hash_lookup_pattern( MOCK_LOOKUP_OPTIONS_HASH ) - assert "data\\_version=faux-populace-us-2099-test-release" in call_args[8] - assert "policyengine\\_version=3.4.0" in call_args[8] - assert "runtime_app_name" not in call_args[8] + assert "data\\_version=faux-populace-us-2099-test-release" in call_args[7] + assert "policyengine\\_version=3.4.0" in call_args[7] + assert "runtime_app_name" not in call_args[7] def test__given_completed_impact__uses_resolved_runtime_bundle_for_cache_lookup( self, @@ -1288,7 +1287,6 @@ def test_given_valid_parameters_calls_service_correctly( assert result == expected_impacts mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.assert_called_once_with( - ANY, MOCK_COUNTRY_ID, MOCK_POLICY_ID, MOCK_BASELINE_POLICY_ID, diff --git a/tests/unit/services/test_execute_analysis.py b/tests/unit/services/test_execute_analysis.py index 7963575b6..e4e2e7714 100644 --- a/tests/unit/services/test_execute_analysis.py +++ b/tests/unit/services/test_execute_analysis.py @@ -4,7 +4,6 @@ pytest_plugins = ["tests.fixtures.services.tracer_analysis_service"] -service = TracerAnalysisService() country_id = "us" household_id = "71424" policy_id = "2" @@ -26,9 +25,9 @@ def test_execute_analysis_static( THEN then a static analysis with the "static" flag should be returned. """ - analysis, analysis_type = service.execute_analysis( - orm_session_factory, country_id, household_id, policy_id, target_variable - ) + analysis, analysis_type = TracerAnalysisService( + orm_session_factory + ).execute_analysis(country_id, household_id, policy_id, target_variable) assert analysis == "Existing static analysis" assert analysis_type == "static" @@ -51,9 +50,9 @@ def test_execute_analysis_streaming( # When existing analysis value is None mock_get_existing_analysis.return_value = None - analysis, analysis_type = service.execute_analysis( - orm_session_factory, country_id, household_id, policy_id, target_variable - ) + analysis, analysis_type = TracerAnalysisService( + orm_session_factory + ).execute_analysis(country_id, household_id, policy_id, target_variable) expected_streaming_output = ["stream chunk 1", "stream chunk 2"] streaming_output = list(analysis) diff --git a/tests/unit/services/test_reform_impacts_service.py b/tests/unit/services/test_reform_impacts_service.py index 18fbec449..15b5029d5 100644 --- a/tests/unit/services/test_reform_impacts_service.py +++ b/tests/unit/services/test_reform_impacts_service.py @@ -1,17 +1,25 @@ from datetime import datetime +import pytest from sqlalchemy import select from policyengine_api.data.v1_models import ReformImpact from policyengine_api.services.reform_impacts_service import ReformImpactsService -service = ReformImpactsService() +@pytest.fixture +def service(orm_session_factory): + return ReformImpactsService(orm_session_factory) -def _create_impact(session, *, execution_id: str, options_hash: str, day: int): +def _create_impact( + service, + *, + execution_id: str, + options_hash: str, + day: int, +): return service.set_reform_impact( - session, country_id="us", policy_id=2, baseline_policy_id=1, @@ -28,22 +36,41 @@ def _create_impact(session, *, execution_id: str, options_hash: str, day: int): ) -def test_reform_impact_service_round_trips_models_and_transitions(orm_session): +def test_get_recent_reform_impacts_orders_and_limits_results(service): + older = _create_impact( + service, + execution_id="older-job", + options_hash="older", + day=1, + ) + newer = _create_impact( + service, + execution_id="newer-job", + options_hash="newer", + day=2, + ) + + assert [ + impact.reform_impact_id for impact in service.get_recent_reform_impacts(1) + ] == [newer.reform_impact_id] + assert older.reform_impact_id != newer.reform_impact_id + + +def test_reform_impact_service_round_trips_models_and_transitions(service): exact = _create_impact( - orm_session, + service, execution_id="exact-job", options_hash="hash-exact", day=1, ) compatible = _create_impact( - orm_session, + service, execution_id="compatible-job", options_hash="hash-compatible", day=2, ) exact_results = service.get_all_reform_impacts( - orm_session, "us", 2, 1, @@ -54,7 +81,6 @@ def test_reform_impact_service_round_trips_models_and_transitions(orm_session): "1", ) compatible_results = service.get_all_reform_impacts_by_options_hash_prefix( - orm_session, "us", 2, 1, @@ -66,10 +92,14 @@ def test_reform_impact_service_round_trips_models_and_transitions(orm_session): "1", ) - assert exact_results == [exact] - assert compatible_results == [exact, compatible] + assert [impact.reform_impact_id for impact in exact_results] == [ + exact.reform_impact_id + ] + assert [impact.reform_impact_id for impact in compatible_results] == [ + exact.reform_impact_id, + compatible.reform_impact_id, + ] completed = service.set_complete_reform_impact( - orm_session, "us", 2, 1, @@ -81,7 +111,6 @@ def test_reform_impact_service_round_trips_models_and_transitions(orm_session): "exact-job", ) failed = service.set_error_reform_impact( - orm_session, "us", 2, 1, @@ -99,22 +128,24 @@ def test_reform_impact_service_round_trips_models_and_transitions(orm_session): assert failed.message == "failed" -def test_reform_impact_service_deletes_only_matching_computing_rows(orm_session): +def test_reform_impact_service_deletes_only_matching_computing_rows( + service, + orm_session_factory, +): _create_impact( - orm_session, + service, execution_id="delete-job", options_hash="delete-hash", day=1, ) retained = _create_impact( - orm_session, + service, execution_id="retain-job", options_hash="retain-hash", day=2, ) service.delete_reform_impact( - orm_session, "us", 2, 1, @@ -124,19 +155,19 @@ def test_reform_impact_service_deletes_only_matching_computing_rows(orm_session) "delete-hash", ) - assert ( - orm_session.scalar( - select(ReformImpact).where(ReformImpact.execution_id == "delete-job") + with orm_session_factory() as session: + assert ( + session.scalar( + select(ReformImpact).where(ReformImpact.execution_id == "delete-job") + ) + is None ) - is None - ) - assert orm_session.get(ReformImpact, retained.reform_impact_id) is retained + assert session.get(ReformImpact, retained.reform_impact_id) is not None -def test_reform_impact_transitions_return_none_for_missing_execution(orm_session): +def test_reform_impact_transitions_return_none_for_missing_execution(service): assert ( service.set_error_reform_impact( - orm_session, "us", 2, 1, diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_stage7_local_service_boundaries.py index 2e1bd4224..2d972fb74 100644 --- a/tests/unit/services/test_stage7_local_service_boundaries.py +++ b/tests/unit/services/test_stage7_local_service_boundaries.py @@ -1,7 +1,15 @@ +import inspect from pathlib import Path import pytest +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.simulation_analysis_service import ( + SimulationAnalysisService, +) +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService + SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" @@ -22,12 +30,55 @@ def test_local_data_services_do_not_issue_queries_directly(module_name): assert "from policyengine_api.data import database" not in source -def test_migrated_local_services_accept_caller_owned_sessions(): +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + (AIAnalysisService, ("get_existing_analysis", "trigger_ai_analysis")), + (SimulationAnalysisService, ("execute_analysis",)), + (TracerAnalysisService, ("execute_analysis", "get_tracer")), + ( + ReformImpactsService, + ( + "get_recent_reform_impacts", + "get_all_reform_impacts", + "get_all_reform_impacts_by_options_hash_prefix", + "set_reform_impact", + "delete_reform_impact", + "set_error_reform_impact", + "set_complete_reform_impact", + ), + ), + ], +) +def test_local_service_public_methods_do_not_accept_persistence( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_analysis_routes_do_not_manage_sessions(): + route_root = SERVICE_ROOT.parent / "routes" for module_name in ( - "ai_analysis_service.py", - "reform_impacts_service.py", - "tracer_analysis_service.py", - "report_output_alias_service.py", + "simulation_analysis_routes.py", + "tracer_analysis_routes.py", ): - source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") - assert "from sqlalchemy.orm import Session" in source + source = (route_root / module_name).read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_legacy_simulation_endpoint_does_not_manage_sessions(): + source = (SERVICE_ROOT.parent / "endpoints" / "simulation.py").read_text( + encoding="utf-8" + ) + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_economy_service_delegates_persistence_without_opening_sessions(): + source = (SERVICE_ROOT / "economy_service.py").read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source diff --git a/tests/unit/services/test_tracer_service.py b/tests/unit/services/test_tracer_service.py index 04a35f636..8f5202fa7 100644 --- a/tests/unit/services/test_tracer_service.py +++ b/tests/unit/services/test_tracer_service.py @@ -8,14 +8,14 @@ pytest_plugins = ["tests.fixtures.services.tracer_fixture_service"] -tracer_service = TracerAnalysisService() - -def test_get_tracer_valid(test_tracer_data, orm_session): +def test_get_tracer_valid( + test_tracer_data, + orm_session_factory, +): # Test get_tracer successfully retrieves valid data from the database. - result = tracer_service.get_tracer( - orm_session, + result = TracerAnalysisService(orm_session_factory).get_tracer( test_tracer_data.country_id, test_tracer_data.household_id, test_tracer_data.policy_id, @@ -27,7 +27,7 @@ def test_get_tracer_valid(test_tracer_data, orm_session): assert result == valid_output -def test_get_tracer_not_found(orm_session): +def test_get_tracer_not_found(orm_session_factory): # Test get_tracer raises NotFound when no matching record exists. valid_country_val_in_db = "us" invalid_household_not_in_db = "9999999" @@ -40,10 +40,10 @@ def test_get_tracer_not_found(orm_session): invalid_api_version, ] with pytest.raises(NotFound): - tracer_service.get_tracer(orm_session, *data_not_in_db) + TracerAnalysisService(orm_session_factory).get_tracer(*data_not_in_db) -def test_get_tracer_database_error(orm_session): +def test_get_tracer_database_error(orm_session_factory): # Test get_tracer handles database errors properly. missing_country_id = "" valid_householdID = "71424" @@ -56,7 +56,6 @@ def test_get_tracer_database_error(orm_session): valid_api_version, ] with pytest.raises(Exception): - tracer_service.get_tracer( - orm_session, + TracerAnalysisService(orm_session_factory).get_tracer( *missing_parameter_causing_database_exception, ) From 7138b0335eca74288ec1a4e0dd8ce4dedd05cd1d Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 17:41:37 +0300 Subject: [PATCH 64/89] test: enforce service-owned persistence boundaries --- changelog.d/3788.changed.md | 2 +- .../test_budget_window_in_flight_dedupe.py | 18 +-- .../test_stage7_service_architecture.py | 126 ++++++++++++++++++ 3 files changed, 136 insertions(+), 10 deletions(-) create mode 100644 tests/unit/services/test_stage7_service_architecture.py diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index 48679b963..66844423b 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to caller-owned SQLAlchemy 2 sessions, mapped models, and Alembic while preserving the existing database schema and public API contracts. +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving the existing database schema and public API contracts. diff --git a/tests/integration/test_budget_window_in_flight_dedupe.py b/tests/integration/test_budget_window_in_flight_dedupe.py index 183483d5d..5d6cb82ea 100644 --- a/tests/integration/test_budget_window_in_flight_dedupe.py +++ b/tests/integration/test_budget_window_in_flight_dedupe.py @@ -36,8 +36,8 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( ModalBudgetWindowBatchExecution, ) from policyengine_api.routes import economy_routes - from policyengine_api.services import economy_service as economy_service_module from policyengine_api.services.budget_window_cache import BudgetWindowCache + from policyengine_api.services.economy_service import EconomyService economy_bp = economy_routes.economy_bp @@ -62,17 +62,17 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( ) ) - monkeypatch.setattr(economy_service_module, "budget_window_cache", fake_cache) monkeypatch.setattr( - economy_service_module, "simulation_entrypoint", simulation_entrypoint + economy_routes, + "economy_service", + EconomyService( + reform_impacts_service_=reform_impacts_service, + budget_window_cache_=fake_cache, + simulation_entrypoint_=simulation_entrypoint, + ), ) monkeypatch.setattr( - economy_routes.economy_service, - "_injected_reform_impacts_service", - reform_impacts_service, - ) - monkeypatch.setattr( - economy_service_module.EconomyService, + EconomyService, "_build_budget_window_batch_payload", lambda self, **kwargs: { "country_id": "us", diff --git a/tests/unit/services/test_stage7_service_architecture.py b/tests/unit/services/test_stage7_service_architecture.py new file mode 100644 index 000000000..48be69c26 --- /dev/null +++ b/tests/unit/services/test_stage7_service_architecture.py @@ -0,0 +1,126 @@ +import inspect +from pathlib import Path + +import pytest + +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.services.simulation_analysis_service import ( + SimulationAnalysisService, +) +from policyengine_api.services.simulation_service import SimulationService +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService +from policyengine_api.services.user_policy_service import UserPolicyService +from policyengine_api.services.user_service import UserService + + +PROJECT_ROOT = Path(__file__).parents[3] +PACKAGE_ROOT = PROJECT_ROOT / "policyengine_api" + + +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + ( + HouseholdService, + ("get_household", "create_household", "update_household"), + ), + (HouseholdCalculationService, ("calculate_stored_household",)), + ( + PolicyService, + ("get_policy", "get_policy_json", "search_policies", "set_policy"), + ), + ( + UserPolicyService, + ( + "create_or_get_user_policy", + "list_user_policies", + "update_user_policy", + ), + ), + (UserService, ("get_profile", "create_profile", "update_profile")), + ( + SimulationService, + ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", + ), + ), + ( + ReportOutputService, + ( + "create_or_reuse_report_output", + "resolve_report_output", + "update_report_output", + ), + ), + (AIAnalysisService, ("get_existing_analysis", "trigger_ai_analysis")), + (SimulationAnalysisService, ("execute_analysis",)), + (TracerAnalysisService, ("execute_analysis", "get_tracer")), + ( + ReformImpactsService, + ( + "get_recent_reform_impacts", + "get_all_reform_impacts", + "get_all_reform_impacts_by_options_hash_prefix", + "set_reform_impact", + "delete_reform_impact", + "set_error_reform_impact", + "set_complete_reform_impact", + ), + ), + ], +) +def test_route_facing_service_methods_hide_persistence_dependencies( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_presentation_layer_does_not_import_or_create_sqlalchemy_sessions(): + offenders = [] + presentation_paths = [ + *sorted((PACKAGE_ROOT / "routes").glob("*.py")), + *sorted((PACKAGE_ROOT / "endpoints").rglob("*.py")), + PACKAGE_ROOT / "country.py", + ] + banned_tokens = ( + "get_v1_session_factory", + "from sqlalchemy", + "import sqlalchemy", + "sessionmaker", + ) + for path in presentation_paths: + source = path.read_text(encoding="utf-8") + if any(token in source for token in banned_tokens): + offenders.append(str(path.relative_to(PACKAGE_ROOT))) + + assert offenders == [] + + +def test_removed_persistence_abstractions_stay_removed(): + removed_paths = ( + "data/v1_daos.py", + "data/data.py", + "endpoints/economy/reform_impact.py", + ) + assert [path for path in removed_paths if (PACKAGE_ROOT / path).exists()] == [] + + +def test_changelog_describes_service_owned_sessions(): + changelog = (PROJECT_ROOT / "changelog.d/3788.changed.md").read_text( + encoding="utf-8" + ) + assert "service-owned" in changelog + assert "caller-owned" not in changelog From 505a457387217829fee46e6000e0d29c084098c1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 20:45:07 +0300 Subject: [PATCH 65/89] style: format service persistence tests --- tests/unit/services/test_direct_orm_policy_household.py | 4 +--- tests/unit/services/test_user_service.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit/services/test_direct_orm_policy_household.py b/tests/unit/services/test_direct_orm_policy_household.py index c457b44c3..bcf480e10 100644 --- a/tests/unit/services/test_direct_orm_policy_household.py +++ b/tests/unit/services/test_direct_orm_policy_household.py @@ -45,9 +45,7 @@ def test_household_service_reads_updates_and_writes_mapped_models( "Direct ORM household", ) with orm_session_factory() as session: - stored = session.scalar( - select(Household).where(Household.id == household.id) - ) + stored = session.scalar(select(Household).where(Household.id == household.id)) assert household.id == stored.id assert stored.household_json == payload diff --git a/tests/unit/services/test_user_service.py b/tests/unit/services/test_user_service.py index 3b515b68e..f7804fba5 100644 --- a/tests/unit/services/test_user_service.py +++ b/tests/unit/services/test_user_service.py @@ -7,6 +7,7 @@ pytest_plugins = ["tests.fixtures.services.user_service"] + @pytest.fixture def service(orm_session_factory): return UserService(orm_session_factory) From 778dfc7e9e11e93a364400563dd1e9ae6de6d6ec Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 21:26:59 +0300 Subject: [PATCH 66/89] test: restore route modules after contract loading --- tests/contract/test_v1_route_contracts.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 5bd63df13..19da62062 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -104,13 +104,17 @@ def _load_blueprint_with_fake_service( blueprint_name: str, ): sentinel = object() + original_route_module = sys.modules.get(route_module_name, sentinel) original_service_module = sys.modules.get(service_module_name, sentinel) sys.modules.pop(route_module_name, None) sys.modules[service_module_name] = fake_service_module try: return getattr(importlib.import_module(route_module_name), blueprint_name) finally: - sys.modules.pop(route_module_name, None) + if original_route_module is sentinel: + sys.modules.pop(route_module_name, None) + else: + sys.modules[route_module_name] = original_route_module if original_service_module is sentinel: sys.modules.pop(service_module_name, None) else: From 9e8a6ded7672a299b75cd35ada41e81ddfacefe1 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 21:27:06 +0300 Subject: [PATCH 67/89] config: require explicit remote database instance --- .env.example | 3 + .github/scripts/cloud_run_env.sh | 2 - .github/scripts/deploy_cloud_run_candidate.sh | 4 +- .../scripts/validate_app_engine_deploy_env.sh | 1 + .../scripts/validate_cloud_run_deploy_env.sh | 2 +- .github/workflows/push.yml | 4 + README.md | 1 + gcp/export.py | 12 +++ gcp/policyengine_api/app.yaml | 2 + policyengine_api/data/orm.py | 10 +-- .../unit/data/test_remote_database_config.py | 14 +++- tests/unit/test_cloud_run_deploy_scripts.py | 84 ++++++++++++++++++- 12 files changed, 125 insertions(+), 14 deletions(-) diff --git a/.env.example b/.env.example index 562132a89..1c6ef9ce3 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,9 @@ # Password for connecting to the PolicyEngine database POLICYENGINE_DB_PASSWORD=policyengine_db_password +# Cloud SQL instance targeted by remote database connections +POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=policyengine-api:us-central1:policyengine-api-data + # Github Microdata Token POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN=policyengine_github_token diff --git a/.github/scripts/cloud_run_env.sh b/.github/scripts/cloud_run_env.sh index 893934f83..2263b3f23 100755 --- a/.github/scripts/cloud_run_env.sh +++ b/.github/scripts/cloud_run_env.sh @@ -9,7 +9,6 @@ cloud_run_set_defaults() { # image built by the staging track, so it must not embed the service name. CLOUD_RUN_IMAGE_NAME="${CLOUD_RUN_IMAGE_NAME:-policyengine-api}" CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT="${CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT:-policyengine-api-cr-runtime@policyengine-api.iam.gserviceaccount.com}" - CLOUD_RUN_CLOUD_SQL_INSTANCE="${CLOUD_RUN_CLOUD_SQL_INSTANCE:-policyengine-api:us-central1:policyengine-api-data}" CLOUD_RUN_CPU="${CLOUD_RUN_CPU:-4}" CLOUD_RUN_MEMORY="${CLOUD_RUN_MEMORY:-16Gi}" CLOUD_RUN_TIMEOUT="${CLOUD_RUN_TIMEOUT:-300}" @@ -57,7 +56,6 @@ cloud_run_set_defaults() { export CLOUD_RUN_ARTIFACT_REPOSITORY export CLOUD_RUN_IMAGE_NAME export CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT - export CLOUD_RUN_CLOUD_SQL_INSTANCE export CLOUD_RUN_CPU export CLOUD_RUN_MEMORY export CLOUD_RUN_TIMEOUT diff --git a/.github/scripts/deploy_cloud_run_candidate.sh b/.github/scripts/deploy_cloud_run_candidate.sh index 04d288397..ffbe6dd8d 100755 --- a/.github/scripts/deploy_cloud_run_candidate.sh +++ b/.github/scripts/deploy_cloud_run_candidate.sh @@ -8,7 +8,7 @@ cloud_run_set_defaults bash .github/scripts/validate_cloud_run_deploy_env.sh env_vars=( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=${CLOUD_RUN_CLOUD_SQL_INSTANCE}" + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" "POLICYENGINE_DB_USER=${POLICYENGINE_DB_USER:-policyengine}" "POLICYENGINE_DB_NAME=${POLICYENGINE_DB_NAME:-policyengine}" "GATEWAY_AUTH_REQUIRED=1" @@ -54,7 +54,7 @@ cloud_run_run gcloud run deploy "${CLOUD_RUN_SERVICE}" \ --allow-unauthenticated \ --execution-environment gen2 \ --service-account "${CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT}" \ - --add-cloudsql-instances "${CLOUD_RUN_CLOUD_SQL_INSTANCE}" \ + --add-cloudsql-instances "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" \ --port "${CLOUD_RUN_PORT}" \ --cpu "${CLOUD_RUN_CPU}" \ --cpu-boost \ diff --git a/.github/scripts/validate_app_engine_deploy_env.sh b/.github/scripts/validate_app_engine_deploy_env.sh index 4e081a4b0..95fdd9a2b 100644 --- a/.github/scripts/validate_app_engine_deploy_env.sh +++ b/.github/scripts/validate_app_engine_deploy_env.sh @@ -9,6 +9,7 @@ selected_url_env="$( )" required=( + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME SIM_ENTRYPOINT "${selected_url_env}" GATEWAY_AUTH_ISSUER diff --git a/.github/scripts/validate_cloud_run_deploy_env.sh b/.github/scripts/validate_cloud_run_deploy_env.sh index 389d18f7e..aed2784ed 100755 --- a/.github/scripts/validate_cloud_run_deploy_env.sh +++ b/.github/scripts/validate_cloud_run_deploy_env.sh @@ -23,7 +23,7 @@ cloud_run_require_env \ CLOUD_RUN_IMAGE_URI \ CLOUD_RUN_TAG \ CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT \ - CLOUD_RUN_CLOUD_SQL_INSTANCE \ + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME \ CLOUD_RUN_POLICYENGINE_DB_PASSWORD_SECRET \ CLOUD_RUN_GITHUB_MICRODATA_TOKEN_SECRET \ CLOUD_RUN_ANTHROPIC_API_KEY_SECRET \ diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index d913dc88d..7e5fb16f2 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -113,6 +113,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} permissions: contents: read id-token: write @@ -219,6 +220,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} ROUTE_IMPL_HEALTH: ${{ vars.ROUTE_IMPL_HEALTH }} ROUTE_IMPL_SPECIFICATION: ${{ vars.ROUTE_IMPL_SPECIFICATION }} ROUTE_IMPL_METADATA: ${{ vars.ROUTE_IMPL_METADATA }} @@ -453,6 +455,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} permissions: contents: read id-token: write @@ -584,6 +587,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} ROUTE_IMPL_HEALTH: ${{ vars.ROUTE_IMPL_HEALTH }} ROUTE_IMPL_SPECIFICATION: ${{ vars.ROUTE_IMPL_SPECIFICATION }} ROUTE_IMPL_METADATA: ${{ vars.ROUTE_IMPL_METADATA }} diff --git a/README.md b/README.md index 724a0e408..925405cd9 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,7 @@ make setup-env `make setup-env` creates a local `.env` from `.env.example`. At minimum, local development expects values for: - `POLICYENGINE_DB_PASSWORD` +- `POLICYENGINE_DB_INSTANCE_CONNECTION_NAME` - `POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN` - `ANTHROPIC_API_KEY` - `OPENAI_API_KEY` diff --git a/gcp/export.py b/gcp/export.py index a56bcbf39..5a726e7d5 100644 --- a/gcp/export.py +++ b/gcp/export.py @@ -1,6 +1,9 @@ import os DB_PD = os.environ["POLICYENGINE_DB_PASSWORD"] +POLICYENGINE_DB_INSTANCE_CONNECTION_NAME = os.environ[ + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" +] GITHUB_MICRODATA_TOKEN = os.environ["POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN"] ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"] OPENAI_API_KEY = os.environ["OPENAI_API_KEY"] @@ -35,6 +38,15 @@ with open(".dbpw", "w") as f: f.write(DB_PD) +app_config_location = "gcp/policyengine_api/app.yaml" +with open(app_config_location) as f: + app_config = f.read().replace( + ".policyengine_db_instance_connection_name", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME, + ) +with open(app_config_location, "w") as f: + f.write(app_config) + # in gcp/compute_api/Dockerfile, replace .github_microdata_token with the contents of the file for dockerfile_location in [ "gcp/policyengine_api/Dockerfile", diff --git a/gcp/policyengine_api/app.yaml b/gcp/policyengine_api/app.yaml index c3512b479..2db2d191a 100644 --- a/gcp/policyengine_api/app.yaml +++ b/gcp/policyengine_api/app.yaml @@ -20,6 +20,8 @@ liveness_check: runtime_config: operating_system: "ubuntu22" runtime_version: "22" +env_variables: + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ".policyengine_db_instance_connection_name" readiness_check: path: "/readiness-check" check_interval_sec: 30 diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 679092ba9..88f4c9426 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -20,9 +20,6 @@ load_dotenv() -DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME = ( - "policyengine-api:us-central1:policyengine-api-data" -) DEFAULT_REMOTE_DB_USER = "policyengine" DEFAULT_REMOTE_DB_NAME = "policyengine" CLOUD_SQL_IP_TYPE = IPTypes.PUBLIC @@ -39,10 +36,9 @@ def get_remote_database_config() -> dict[str, str]: return { - "instance_connection_name": os.environ.get( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", - DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME, - ), + "instance_connection_name": os.environ[ + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" + ], "db_user": os.environ.get("POLICYENGINE_DB_USER", DEFAULT_REMOTE_DB_USER), "db_name": os.environ.get("POLICYENGINE_DB_NAME", DEFAULT_REMOTE_DB_NAME), } diff --git a/tests/unit/data/test_remote_database_config.py b/tests/unit/data/test_remote_database_config.py index fbac992a4..118594f27 100644 --- a/tests/unit/data/test_remote_database_config.py +++ b/tests/unit/data/test_remote_database_config.py @@ -2,11 +2,23 @@ os.environ.setdefault("FLASK_DEBUG", "1") +import pytest + from policyengine_api.data.orm import get_remote_database_config -def test_remote_database_config_defaults_to_current_production_values(monkeypatch): +def test_remote_database_config_requires_instance_connection_name(monkeypatch): monkeypatch.delenv("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", raising=False) + + with pytest.raises(KeyError, match="POLICYENGINE_DB_INSTANCE_CONNECTION_NAME"): + get_remote_database_config() + + +def test_remote_database_config_defaults_non_target_settings(monkeypatch): + monkeypatch.setenv( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", + "policyengine-api:us-central1:policyengine-api-data", + ) monkeypatch.delenv("POLICYENGINE_DB_USER", raising=False) monkeypatch.delenv("POLICYENGINE_DB_NAME", raising=False) diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index e027ba947..d891c15bd 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -65,6 +65,7 @@ def _gateway_auth_env() -> dict[str, str]: def _required_runtime_env() -> dict[str, str]: return { + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": PRODUCTION_CLOUD_SQL_INSTANCE, "POLICYENGINE_DB_PASSWORD": "raw-db-secret-value", "POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN": ("raw-github-secret-value"), "ANTHROPIC_API_KEY": "raw-anthropic-secret-value", @@ -534,6 +535,7 @@ def test_validate_cloud_run_deploy_env_accepts_direct_mode_from_environment(): ROUTE_IMPL_HEALTH="fastapi_native", ROUTE_IMPL_SPECIFICATION="fastapi_native", ROUTE_IMPL_METADATA="fastapi_native", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ), ) @@ -613,6 +615,7 @@ def test_validate_cloud_run_deploy_env_requires_only_selected_url( ROUTE_IMPL_HEALTH="fastapi_native", ROUTE_IMPL_SPECIFICATION="fastapi_native", ROUTE_IMPL_METADATA="fastapi_native", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ) missing_result = _run_script( @@ -648,6 +651,7 @@ def test_validate_app_engine_deploy_env_accepts_direct_mode_from_environment(): _script_env( SIM_ENTRYPOINT="old_gateway_direct", OLD_SIMULATION_GATEWAY_URL="https://old-gateway.example.test", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ), ) @@ -677,6 +681,7 @@ def test_validate_app_engine_deploy_env_requires_only_selected_url( ): env = _script_env( SIM_ENTRYPOINT=entrypoint, + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ) missing_result = _run_script( @@ -693,15 +698,41 @@ def test_validate_app_engine_deploy_env_requires_only_selected_url( assert valid_result.returncode == 0, valid_result.stderr -def test_app_engine_image_contains_simulation_routing_environment_placeholders(): +def test_app_engine_bundle_contains_runtime_environment_placeholders(): dockerfile = (REPO / "gcp/policyengine_api/Dockerfile").read_text(encoding="utf-8") + app_config = (REPO / "gcp/policyengine_api/app.yaml").read_text(encoding="utf-8") export_script = (REPO / "gcp/export.py").read_text(encoding="utf-8") assert 'ENV SIMULATION_ENTRYPOINT_URL=".simulation_entrypoint_url"' in dockerfile assert 'ENV OLD_SIMULATION_GATEWAY_URL=".old_simulation_gateway_url"' in dockerfile assert 'ENV SIM_ENTRYPOINT=".sim_entrypoint"' in dockerfile + assert ( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + '".policyengine_db_instance_connection_name"' in app_config + ) assert '".old_simulation_gateway_url", OLD_SIMULATION_GATEWAY_URL' in export_script assert '".sim_entrypoint", SIM_ENTRYPOINT' in export_script + assert '".policyengine_db_instance_connection_name",' in export_script + assert "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME," in export_script + + +@pytest.mark.parametrize( + "validation_script", + [ + ".github/scripts/validate_app_engine_deploy_env.sh", + ".github/scripts/validate_cloud_run_deploy_env.sh", + ], +) +def test_deployment_validation_requires_database_instance_connection_name( + validation_script, +): + env = _script_env(**_required_runtime_env()) + env.pop("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME") + + result = _run_script(validation_script, env) + + assert result.returncode == 1 + assert "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" in result.stderr @pytest.mark.parametrize( @@ -730,6 +761,10 @@ def test_app_engine_export_requires_only_selected_url( ): (tmp_path / "gcp/policyengine_api").mkdir(parents=True) shutil.copy2(REPO / "gcp/export.py", tmp_path / "gcp/export.py") + shutil.copy2( + REPO / "gcp/policyengine_api/app.yaml", + tmp_path / "gcp/policyengine_api/app.yaml", + ) shutil.copy2( REPO / "gcp/policyengine_api/Dockerfile", tmp_path / "gcp/policyengine_api/Dockerfile", @@ -755,8 +790,15 @@ def test_app_engine_export_requires_only_selected_url( rendered = (tmp_path / "gcp/policyengine_api/Dockerfile").read_text( encoding="utf-8" ) + rendered_app_config = (tmp_path / "gcp/policyengine_api/app.yaml").read_text( + encoding="utf-8" + ) assert f'ENV {selected_url_env}="{selected_url}"' in rendered assert f'ENV {unselected_url_env}=""' in rendered + assert ( + f"POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + f'"{PRODUCTION_CLOUD_SQL_INSTANCE}"' in rendered_app_config + ) def test_build_cloud_run_image_dry_run_uses_cloud_run_dockerfile(): @@ -825,6 +867,30 @@ def test_deploy_cloud_run_candidate_dry_run_never_shifts_traffic(): assert result.stdout.count(f"{selector}=fastapi_native") == 1 +def test_deploy_cloud_run_candidate_uses_configured_database_instance(): + configured_instance = "project:region:configured-instance" + env = { + **_required_runtime_env(), + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": configured_instance, + } + result = _run_script( + ".github/scripts/deploy_cloud_run_candidate.sh", + _script_env( + **env, + CLOUD_RUN_IMAGE_URI="us-central1-docker.pkg.dev/project/repo/api:sha", + CLOUD_RUN_TAG="stage3-test", + ), + ) + + assert result.returncode == 0, result.stderr + assert f"--add-cloudsql-instances {configured_instance}" in result.stdout + assert ( + f"POLICYENGINE_DB_INSTANCE_CONNECTION_NAME={configured_instance}" + in result.stdout + ) + assert PRODUCTION_CLOUD_SQL_INSTANCE not in result.stdout + + @pytest.mark.parametrize( ("entrypoint", "selected_url_env", "selected_url", "unselected_url_env"), [ @@ -1501,6 +1567,22 @@ def test_cloud_run_deploy_jobs_use_environment_scoped_stage6_route_selectors(): assert f"{selector}: ${{{{ vars.{selector} }}}}" in job +def test_all_deploy_jobs_use_github_database_instance_variable(): + workflow = _push_workflow() + instance_env = ( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + "${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }}" + ) + + for job_name in ( + "deploy-staging", + "deploy-cloud-run-staging", + "deploy-production-candidate", + "deploy-cloud-run-candidate", + ): + assert instance_env in _workflow_job_block(workflow, job_name) + + def test_deployment_consumers_require_selector_from_environment(): consumers = ( ".github/request-simulation-model-versions.sh", From 2491e843add5581c630da4eca7d5e9084494903c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 22:32:39 +0300 Subject: [PATCH 68/89] refactor: consolidate Flask handlers into blueprints --- policyengine_api/api.py | 98 +------ policyengine_api/extensions.py | 6 + policyengine_api/routes/home_routes.py | 14 + policyengine_api/routes/household_routes.py | 167 +++++++++++- policyengine_api/routes/policy_routes.py | 247 +++++++++++++++++- .../routes/reform_impact_routes.py | 38 +++ policyengine_api/routes/system_routes.py | 30 +++ .../services/household_calculation_service.py | 40 +++ tests/contract/test_v1_route_contracts.py | 27 +- tests/fixtures/integration/simulations.py | 6 +- tests/unit/data/test_stage7_no_direct_sql.py | 6 +- .../test_calculate_deprecated_inputs.py | 45 ++-- .../test_calculate_error_statuses.py | 38 +-- .../test_reform_impact_routes.py} | 21 +- .../test_set_user_policy_dataset.py | 4 +- .../test_stage7_orm_routes.py} | 6 +- .../test_update_user_policy.py | 4 +- .../test_household_calculation_service.py | 10 +- .../test_stage7_local_service_boundaries.py | 4 +- .../unit/services/test_user_policy_service.py | 8 +- tests/unit/test_warmup.py | 4 +- tests/unit/test_yearly_var_removal.py | 4 +- 22 files changed, 638 insertions(+), 189 deletions(-) create mode 100644 policyengine_api/extensions.py create mode 100644 policyengine_api/routes/home_routes.py create mode 100644 policyengine_api/routes/reform_impact_routes.py create mode 100644 policyengine_api/routes/system_routes.py rename tests/unit/{endpoints => routes}/test_calculate_deprecated_inputs.py (90%) rename tests/unit/{endpoints => routes}/test_calculate_error_statuses.py (74%) rename tests/unit/{endpoints/test_get_simulations.py => routes/test_reform_impact_routes.py} (79%) rename tests/unit/{endpoints => routes}/test_set_user_policy_dataset.py (87%) rename tests/unit/{endpoints/test_stage7_orm_endpoints.py => routes/test_stage7_orm_routes.py} (95%) rename tests/unit/{endpoints => routes}/test_update_user_policy.py (95%) diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 9f052a93f..b0730e0ae 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -1,3 +1,4 @@ +# ruff: noqa: E402 """ This is the main Flask app for the PolicyEngine API. """ @@ -21,8 +22,7 @@ def log_timing(message): log_timing("Flask imports completed") -from flask_caching import Cache -from policyengine_api.utils import make_cache_key +from policyengine_api.extensions import cache from policyengine_api.migration_logging import register_migration_request_logging log_timing("Caching utilities import completed") @@ -33,6 +33,9 @@ def log_timing(message): from policyengine_api.routes.error_routes import error_bp log_timing("Error routes import completed") +from policyengine_api.routes.home_routes import home_bp + +log_timing("Home routes import completed") from policyengine_api.routes.economy_routes import economy_bp log_timing("Economy routes import completed") @@ -59,24 +62,11 @@ def log_timing(message): from policyengine_api.routes.ai_prompt_routes import ai_prompt_bp from policyengine_api.routes.simulation_routes import simulation_bp from policyengine_api.routes.report_output_routes import report_output_bp +from policyengine_api.routes.reform_impact_routes import reform_impact_bp +from policyengine_api.routes.system_routes import system_bp log_timing("Base AI routes import completed") -from .endpoints import ( - get_home, - get_policy_search, - get_household_under_policy, - get_calculate, - set_user_policy, - get_user_policy, - update_user_policy, - get_simulations, -) - -log_timing("Legacy endpoints import completed") - -from policyengine_api.readiness import is_ready - log_timing("Initialising API...") app = application = flask.Flask(__name__) @@ -92,7 +82,7 @@ def log_timing(message): "CACHE_DEFAULT_TIMEOUT": 300, } ) -cache = Cache(app) +cache.init_app(app) log_timing("Caching initialised") CORS(app) @@ -104,7 +94,7 @@ def log_timing(message): app.register_blueprint(error_bp) log_timing("Error routes registered") -app.route("/", methods=["GET"])(get_home) +app.register_blueprint(home_bp) log_timing("Home routes registered") app.register_blueprint(metadata_bp) @@ -117,27 +107,6 @@ def log_timing(message): app.register_blueprint(policy_bp) log_timing("Policy routes registered") -app.route("//policies", methods=["GET"])(get_policy_search) -log_timing("Policy search endpoint registered") - -app.route( - "//household//policy/", - methods=["GET"], -)(get_household_under_policy) -log_timing("Household under policy endpoint registered") - -app.route("//calculate", methods=["POST"])( - cache.cached(make_cache_key=make_cache_key)(get_calculate) -) -log_timing("Calculate endpoint registered") - -app.route("//calculate-full", methods=["POST"])( - cache.cached(make_cache_key=make_cache_key)( - lambda *args, **kwargs: get_calculate(*args, **kwargs, add_missing=True) - ) -) -log_timing("Calculate-full endpoint registered") - # Routes for economy microsimulation app.register_blueprint(economy_bp) log_timing("Economy routes registered") @@ -146,19 +115,10 @@ def log_timing(message): app.register_blueprint(simulation_analysis_bp) log_timing("Simulation analysis routes registered") -app.route("//user-policy", methods=["POST"])(set_user_policy) -log_timing("User policy set endpoint registered") - -app.route("//user-policy", methods=["PUT"])(update_user_policy) -log_timing("User policy update endpoint registered") - -app.route("//user-policy/", methods=["GET"])(get_user_policy) -log_timing("User policy get endpoint registered") - app.register_blueprint(user_profile_bp) log_timing("User profile routes registered") -app.route("/simulations", methods=["GET"])(get_simulations) +app.register_blueprint(reform_impact_bp) log_timing("Simulations endpoint registered") app.register_blueprint(tracer_analysis_bp) @@ -170,42 +130,8 @@ def log_timing(message): app.register_blueprint(simulation_bp) app.register_blueprint(report_output_bp) - - -@app.route("/liveness-check", methods=["GET"]) -def liveness_check(): - return flask.Response("OK", status=200, headers={"Content-Type": "text/plain"}) - - -log_timing("Liveness check endpoint registered") - - -@app.route("/readiness-check", methods=["GET"]) -def readiness_check(): - # 503 until the startup warmup has compiled the simulation machinery - # (policyengine_api.readiness); /liveness-check stays unconditional. - if not is_ready(): - return flask.Response( - "NOT READY", status=503, headers={"Content-Type": "text/plain"} - ) - return flask.Response("OK", status=200, headers={"Content-Type": "text/plain"}) - - -log_timing("Readiness check endpoint registered") - - -from policyengine_api.specification import OPENAPI_SPECIFICATION - -openapi_spec = OPENAPI_SPECIFICATION -log_timing("OpenAPI spec loaded") - - -@app.route("/specification", methods=["GET"]) -def get_specification(): - return flask.jsonify(openapi_spec) - - -log_timing("Specification endpoint registered") +app.register_blueprint(system_bp) +log_timing("System routes registered") log_timing("API initialised.") diff --git a/policyengine_api/extensions.py b/policyengine_api/extensions.py new file mode 100644 index 000000000..064b55d22 --- /dev/null +++ b/policyengine_api/extensions.py @@ -0,0 +1,6 @@ +"""Shared Flask extensions initialized by the application factory module.""" + +from flask_caching import Cache + + +cache = Cache() diff --git a/policyengine_api/routes/home_routes.py b/policyengine_api/routes/home_routes.py new file mode 100644 index 000000000..fdc230705 --- /dev/null +++ b/policyengine_api/routes/home_routes.py @@ -0,0 +1,14 @@ +from flask import Blueprint + + +home_bp = Blueprint("home", __name__) + + +@home_bp.route("/", methods=["GET"]) +def get_home() -> str: + """Get the home page of the PolicyEngine API.""" + return ( + "

PolicyEngine households API

" + "

Use this API to compute the impact of public policy on individual " + "households.

" + ) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 26e63dcc7..5a9222cd4 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -1,16 +1,29 @@ -from flask import Blueprint, Response, request -from werkzeug.exceptions import NotFound, BadRequest import json +import logging + +from flask import Blueprint, Response, request +from policyengine_core.errors import SituationParsingError +from werkzeug.exceptions import BadRequest, NotFound from policyengine_api.data.v1_models import Household +from policyengine_api.extensions import cache +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, + HouseholdNotFoundError, + InvalidHouseholdInputsError, + PolicyNotFoundError, +) from policyengine_api.services.household_service import HouseholdService +from policyengine_api.utils import make_cache_key +from policyengine_api.utils.input_validation import format_unrecognized_inputs_message from policyengine_api.utils.payload_validators import ( - validate_household_payload, validate_country, + validate_household_payload, ) household_bp = Blueprint("household", __name__) household_service = HouseholdService() +household_calculation_service = HouseholdCalculationService() def _serialize_household(household: Household) -> dict: @@ -144,3 +157,151 @@ def update_household(country_id: str, household_id: int) -> Response: status=200, mimetype="application/json", ) + + +@household_bp.route( + "//household//policy/", + methods=["GET"], +) +@validate_country +def get_household_under_policy(country_id: str, household_id: str, policy_id: str): + """Get a stored household's output under a stored policy.""" + try: + calculation = household_calculation_service.calculate_stored_household( + country_id, + int(household_id), + int(policy_id), + ) + except HouseholdNotFoundError: + return Response( + json.dumps( + dict( + status="error", + message=f"Household #{household_id} not found.", + ) + ), + status=404, + mimetype="application/json", + ) + except PolicyNotFoundError: + return Response( + json.dumps( + dict( + status="error", + message=f"Policy #{policy_id} not found.", + ) + ), + status=404, + mimetype="application/json", + ) + except InvalidHouseholdInputsError as error: + return Response( + json.dumps( + dict( + status="error", + message=format_unrecognized_inputs_message(error.invalid_inputs), + result=None, + errors=[ + invalid_input.to_dict() + for invalid_input in error.invalid_inputs + ], + ) + ), + status=400, + mimetype="application/json", + ) + except Exception as error: + logging.exception(error) + return Response( + json.dumps( + dict( + status="error", + message=( + f"Error calculating household #{household_id} under policy " + f"#{policy_id}: {error}" + ), + ) + ), + status=500, + mimetype="application/json", + ) + + response_body = dict(status="ok", message=None, result=calculation.household) + if calculation.warnings: + response_body["warnings"] = list(calculation.warnings) + return response_body + + +def _calculate(country_id: str, *, add_missing: bool) -> dict | Response: + payload = request.json + household_json = payload.get("household", {}) + policy_json = payload.get("policy", {}) + + try: + calculation = household_calculation_service.calculate_household( + country_id, + household_json, + policy_json, + add_missing=add_missing, + ) + except InvalidHouseholdInputsError as error: + return Response( + json.dumps( + { + "status": "error", + "message": format_unrecognized_inputs_message(error.invalid_inputs), + "result": None, + "errors": [ + invalid_input.to_dict() + for invalid_input in error.invalid_inputs + ], + } + ), + status=400, + mimetype="application/json", + ) + except SituationParsingError as error: + return Response( + json.dumps( + dict( + status="error", + message=f"Invalid household payload: {error}", + result=None, + ) + ), + status=400, + mimetype="application/json", + ) + except Exception as error: + logging.exception(error) + return Response( + json.dumps( + dict( + status="error", + message=f"Error calculating household under policy: {error}", + ) + ), + status=500, + mimetype="application/json", + ) + + response_body = dict(status="ok", message=None, result=calculation.household) + if calculation.warnings: + response_body["warnings"] = list(calculation.warnings) + return response_body + + +@household_bp.route("//calculate", methods=["POST"]) +@cache.cached(make_cache_key=make_cache_key) +@validate_country +def get_calculate(country_id: str) -> dict | Response: + """Calculate a household without adding omitted yearly variables.""" + return _calculate(country_id, add_missing=False) + + +@household_bp.route("//calculate-full", methods=["POST"]) +@cache.cached(make_cache_key=make_cache_key) +@validate_country +def get_calculate_full(country_id: str) -> dict | Response: + """Calculate a household after adding omitted yearly variables.""" + return _calculate(country_id, add_missing=True) diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index ee83c8f0b..2403cec1d 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,9 +1,11 @@ -from flask import Blueprint, Response, request import json -from policyengine_api.data.v1_models import Policy +from flask import Blueprint, Response, request +from werkzeug.exceptions import BadRequest, NotFound + +from policyengine_api.data.v1_models import Policy, UserPolicy from policyengine_api.services.policy_service import PolicyService -from werkzeug.exceptions import NotFound, BadRequest +from policyengine_api.services.user_policy_service import UserPolicyService from policyengine_api.utils.payload_validators import ( validate_country, validate_set_policy_payload, @@ -11,6 +13,7 @@ policy_bp = Blueprint("policy", __name__) policy_service = PolicyService() +user_policy_service = UserPolicyService() def _serialize_policy(policy: Policy) -> dict: @@ -90,3 +93,241 @@ def set_policy(country_id: str) -> Response: code = 200 if is_existing_policy else 201 return Response(json.dumps(response_body), status=code, mimetype="application/json") + + +def _serialize_user_policy(user_policy: UserPolicy) -> dict: + return { + column.name: getattr(user_policy, column.name) + for column in UserPolicy.__table__.columns + } + + +@policy_bp.route("//policies", methods=["GET"]) +@validate_country +def get_policy_search(country_id: str) -> Response: + """Search policies for a country.""" + query = request.args.get("query", "") + unique_only = request.args.get("unique_only", default=False, type=json.loads) + + try: + results = policy_service.search_policies( + country_id, + query, + unique_only=unique_only, + ) + if not results: + return Response( + json.dumps( + dict( + status="error", + message=( + f"No policies found for country {country_id} for query " + f"'{query}" + ), + ) + ), + status=404, + mimetype="application/json", + ) + + policies = [dict(id=result.id, label=result.label) for result in results] + return Response( + json.dumps( + dict( + status="ok", + message="Policies found", + result=policies, + ) + ), + status=200, + mimetype="application/json", + ) + except Exception as error: + return Response( + json.dumps( + dict( + status="error", + message=f"Internal server error: {error}", + ) + ), + status=500, + mimetype="application/json", + ) + + +@policy_bp.route("//user-policy", methods=["POST"]) +@validate_country +def set_user_policy(country_id: str) -> Response: + """Create a saved policy for a user, or return its existing ID.""" + payload = request.json + reform_label = payload.pop("reform_label", None) + reform_id = payload.pop("reform_id") + baseline_label = payload.pop("baseline_label", None) + baseline_id = payload.pop("baseline_id") + user_id = payload.pop("user_id") + year = payload.pop("year") + geography = payload.pop("geography") + dataset = payload.pop("dataset", None) + number_of_provisions = payload.pop("number_of_provisions") + api_version = payload.pop("api_version") + added_date = payload.pop("added_date") + updated_date = payload.pop("updated_date") + budgetary_impact = payload.pop("budgetary_impact", None) + policy_type = payload.pop("type", None) + + values = { + "country_id": country_id, + "reform_id": reform_id, + "reform_label": reform_label, + "baseline_id": baseline_id, + "baseline_label": baseline_label, + "user_id": user_id, + "year": year, + "geography": geography, + "dataset": dataset, + "number_of_provisions": number_of_provisions, + "api_version": api_version, + "added_date": added_date, + "updated_date": updated_date, + "budgetary_impact": budgetary_impact, + "type": policy_type, + } + + try: + creation = user_policy_service.create_or_get_user_policy(values) + user_policy = creation.user_policy + if not creation.created: + return Response( + json.dumps( + dict( + status="ok", + message=( + f"The reform #{reform_id} / baseline #{baseline_id} pair " + f"already exists for user {user_id}" + ), + result=dict(id=user_policy.id), + ) + ), + status=200, + mimetype="application/json", + ) + except Exception as error: + return Response( + json.dumps( + { + "message": ( + f"Internal database error: {error}; please try again later." + ) + } + ), + status=500, + mimetype="application/json", + ) + + return Response( + json.dumps( + dict( + status="ok", + message="Record created successfully", + result=_serialize_user_policy(user_policy), + ) + ), + status=201, + mimetype="application/json", + ) + + +@policy_bp.route("//user-policy/", methods=["GET"]) +@validate_country +def get_user_policy(country_id: str, user_id: str) -> dict: + """Fetch all saved policies for a user.""" + user_policies = user_policy_service.list_user_policies(country_id, user_id) + return dict( + status="ok", + message=None, + result=[_serialize_user_policy(row) for row in user_policies], + ) + + +UPDATE_USER_POLICY_ALLOWED_FIELDS = frozenset( + { + "reform_label", + "baseline_label", + "year", + "geography", + "dataset", + "number_of_provisions", + "api_version", + "added_date", + "updated_date", + "budgetary_impact", + "type", + } +) + + +@policy_bp.route("//user-policy", methods=["PUT"]) +@validate_country +def update_user_policy(country_id: str) -> Response: + """Update mutable fields on a saved user policy.""" + payload = request.json + if not isinstance(payload, dict) or "id" not in payload: + return Response( + json.dumps({"message": "Request body must include an 'id' field."}), + status=400, + mimetype="application/json", + ) + + user_policy_id = payload.pop("id") + unknown_keys = [ + key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS + ] + if unknown_keys: + return Response( + json.dumps( + { + "message": ( + "Request body contains unsupported fields: " + f"{sorted(unknown_keys)}" + ) + } + ), + status=400, + mimetype="application/json", + ) + + if not payload: + return Response( + json.dumps( + {"message": "Request body must include at least one field to update."} + ), + status=400, + mimetype="application/json", + ) + + try: + user_policy_service.update_user_policy(user_policy_id, payload) + except Exception as error: + return Response( + json.dumps( + { + "message": ( + f"Internal database error: {error}; please try again later." + ) + } + ), + status=500, + mimetype="application/json", + ) + + return Response( + json.dumps( + dict( + status="ok", + message="Record updated successfully", + result=dict(id=user_policy_id), + ) + ), + status=200, + mimetype="application/json", + ) diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py new file mode 100644 index 000000000..ac9866bd7 --- /dev/null +++ b/policyengine_api/routes/reform_impact_routes.py @@ -0,0 +1,38 @@ +from flask import Blueprint, request + +from policyengine_api.data.v1_models import ReformImpact +from policyengine_api.services.reform_impacts_service import ReformImpactsService + + +reform_impact_bp = Blueprint("reform_impact", __name__) +reform_impacts_service = ReformImpactsService() + +_MAX_SIMULATION_RESULTS = 1000 +_DEFAULT_SIMULATION_RESULTS = 100 + + +def _parse_result_limit(value: str | None) -> int: + if value is None: + return _DEFAULT_SIMULATION_RESULTS + try: + result_limit = int(value) + except (TypeError, ValueError): + return _DEFAULT_SIMULATION_RESULTS + return max(1, min(result_limit, _MAX_SIMULATION_RESULTS)) + + +@reform_impact_bp.route("/simulations", methods=["GET"]) +def get_simulations() -> dict: + """Return recent reform impacts, bounded to protect the database query.""" + result_limit = _parse_result_limit(request.args.get("max_results")) + impacts = reform_impacts_service.get_recent_reform_impacts(result_limit) + + return { + "result": [ + { + column.name: getattr(impact, column.name) + for column in ReformImpact.__table__.columns + } + for impact in impacts + ] + } diff --git a/policyengine_api/routes/system_routes.py b/policyengine_api/routes/system_routes.py new file mode 100644 index 000000000..f0e2748c0 --- /dev/null +++ b/policyengine_api/routes/system_routes.py @@ -0,0 +1,30 @@ +from flask import Blueprint, Response, jsonify + +from policyengine_api.readiness import is_ready +from policyengine_api.specification import OPENAPI_SPECIFICATION + + +system_bp = Blueprint("system", __name__) + + +@system_bp.route("/liveness-check", methods=["GET"]) +def liveness_check() -> Response: + return Response("OK", status=200, headers={"Content-Type": "text/plain"}) + + +@system_bp.route("/readiness-check", methods=["GET"]) +def readiness_check() -> Response: + # The service is not ready until startup warmup compiles the simulation + # machinery. Liveness remains unconditional so the worker is not restarted. + if not is_ready(): + return Response( + "NOT READY", + status=503, + headers={"Content-Type": "text/plain"}, + ) + return Response("OK", status=200, headers={"Content-Type": "text/plain"}) + + +@system_bp.route("/specification", methods=["GET"]) +def get_specification() -> Response: + return jsonify(OPENAPI_SPECIFICATION) diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 98c86ab5b..0e9d7d959 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -259,3 +259,43 @@ def calculate_stored_household( household=calculation.household, warnings=tuple(warning.message for warning in deprecated_inputs.warnings), ) + + def calculate_household( + self, + country_id: str, + household_json: dict, + policy_json: dict, + *, + add_missing: bool = False, + ) -> HouseholdCalculationResult: + """Validate and calculate request-provided household and policy data.""" + countries = self._countries() + country = countries.get(country_id) + household_json = deepcopy(household_json) + if add_missing: + household_json = add_yearly_variables( + household_json, + country_id, + countries, + ) + + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( + household_json, + policy_json, + country.metadata, + ) + if invalid_inputs: + raise InvalidHouseholdInputsError(invalid_inputs) + + raw_calculation = country.calculate(household_json, policy_json) + household = ( + raw_calculation + if isinstance(raw_calculation, dict) + else raw_calculation.household + ) + return HouseholdCalculationResult( + household=household, + warnings=tuple(warning.message for warning in deprecated_inputs.warnings), + ) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 19da62062..eae014cf1 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -8,8 +8,6 @@ from flask import Flask, Response from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.endpoints.household import get_calculate -from policyengine_api.endpoints.policy import get_policy_search from policyengine_api.data.v1_models import ( Household, Policy, @@ -17,6 +15,7 @@ ReportOutputRun, Simulation, ) +from policyengine_api.extensions import cache from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp @@ -25,6 +24,9 @@ ReportCreateResult, ReportOutputView, ) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationResult, +) from policyengine_api.services.simulation_service import SimulationCreateResult from tests.contract.clients import ( ASGIContractClient, @@ -145,15 +147,14 @@ def _load_contract_economy_blueprint(): def create_contract_flask_app() -> Flask: app = Flask(__name__) - app.config["TESTING"] = True + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) app.register_blueprint(_load_contract_metadata_blueprint()) app.register_blueprint(policy_bp) app.register_blueprint(household_bp) app.register_blueprint(_load_contract_economy_blueprint()) app.register_blueprint(simulation_bp) app.register_blueprint(report_output_bp) - app.route("//policies", methods=["GET"])(get_policy_search) - app.route("//calculate", methods=["POST"])(get_calculate) @app.route("/liveness-check") def liveness_check(): @@ -239,7 +240,7 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.policy_service.search_policies", + "policyengine_api.routes.policy_routes.policy_service.search_policies", return_value=[ Policy( id=123, @@ -301,14 +302,12 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.household.get_countries", - return_value={"us": _fake_country()}, - ) - ) - stack.enter_context( - patch( - "policyengine_api.endpoints.household.get_invalid_inputs_response", - return_value=None, + "policyengine_api.routes.household_routes.household_calculation_service.calculate_household", + return_value=HouseholdCalculationResult( + household={ + "people": {"you": {"age": {"2026": 40}}}, + } + ), ) ) stack.enter_context( diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index 741676047..b0fe44636 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -2,9 +2,9 @@ Test fixtures and constants for axes calculation tests """ -import pytest -from unittest.mock import Mock, MagicMock, patch -from policyengine_api.endpoints.household import add_yearly_variables +from policyengine_api.services.household_calculation_service import ( + add_yearly_variables, +) STANDARD_AXES_COUNT = ( 401 # Not formally defined anywhere, but this value is used throughout the API diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_stage7_no_direct_sql.py index 0086c37de..ae62e2d18 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_stage7_no_direct_sql.py @@ -25,9 +25,9 @@ def test_runtime_sql_is_confined_to_the_data_access_layer(): def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): relative_paths = ( - "endpoints/household.py", - "endpoints/policy.py", - "endpoints/simulation.py", + "routes/household_routes.py", + "routes/policy_routes.py", + "routes/reform_impact_routes.py", "country.py", "services/ai_analysis_service.py", "services/reform_impacts_service.py", diff --git a/tests/unit/endpoints/test_calculate_deprecated_inputs.py b/tests/unit/routes/test_calculate_deprecated_inputs.py similarity index 90% rename from tests/unit/endpoints/test_calculate_deprecated_inputs.py rename to tests/unit/routes/test_calculate_deprecated_inputs.py index 0d7799abd..aeb0af3e8 100644 --- a/tests/unit/endpoints/test_calculate_deprecated_inputs.py +++ b/tests/unit/routes/test_calculate_deprecated_inputs.py @@ -1,7 +1,11 @@ from flask import Flask import pytest -from policyengine_api.endpoints import household as household_endpoint +from policyengine_api.routes import household_routes +from policyengine_api.extensions import cache +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) class DummyCountry: @@ -43,18 +47,15 @@ def calculate(self, household, policy): def calculate_client(monkeypatch): country = DummyCountry() monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) app = Flask(__name__) - app.add_url_rule( - "//calculate", - "calculate", - household_endpoint.get_calculate, - methods=["POST"], - ) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) + app.register_blueprint(household_routes.household_bp) return app.test_client(), country @@ -261,14 +262,13 @@ def test__calculate__preserves_relationship_fields(calculate_client): def test__calculate_full__drops_deprecated_input_after_add_missing(monkeypatch): country = DummyCountry() monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) monkeypatch.setattr( - household_endpoint, - "add_yearly_variables", - lambda household, country_id: { + "policyengine_api.services.household_calculation_service.add_yearly_variables", + lambda household, country_id, countries=None: { **household, "people": { **household["people"], @@ -281,16 +281,9 @@ def test__calculate_full__drops_deprecated_input_after_add_missing(monkeypatch): ) app = Flask(__name__) - - def calculate_full(country_id): - return household_endpoint.get_calculate(country_id, add_missing=True) - - app.add_url_rule( - "//calculate-full", - "calculate_full", - calculate_full, - methods=["POST"], - ) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) + app.register_blueprint(household_routes.household_bp) client = app.test_client() household = { "people": { diff --git a/tests/unit/endpoints/test_calculate_error_statuses.py b/tests/unit/routes/test_calculate_error_statuses.py similarity index 74% rename from tests/unit/endpoints/test_calculate_error_statuses.py rename to tests/unit/routes/test_calculate_error_statuses.py index 286bab6a3..6f8f6cd19 100644 --- a/tests/unit/endpoints/test_calculate_error_statuses.py +++ b/tests/unit/routes/test_calculate_error_statuses.py @@ -1,8 +1,11 @@ from flask import Flask -import pytest from policyengine_core.errors import SituationParsingError -from policyengine_api.endpoints import household as household_endpoint +from policyengine_api.extensions import cache +from policyengine_api.routes import household_routes +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) from .test_calculate_deprecated_inputs import DummyCountry @@ -31,34 +34,19 @@ def calculate(self, household, policy): def make_client(monkeypatch, country, add_missing=False): monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) app = Flask(__name__) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) if add_missing: monkeypatch.setattr( - household_endpoint, - "add_yearly_variables", - lambda household, country_id: household, - ) - - def handler(country_id): - return household_endpoint.get_calculate(country_id, add_missing=True) - - app.add_url_rule( - "//calculate-full", - "calculate_full", - handler, - methods=["POST"], - ) - else: - app.add_url_rule( - "//calculate", - "calculate", - household_endpoint.get_calculate, - methods=["POST"], + "policyengine_api.services.household_calculation_service.add_yearly_variables", + lambda household, country_id, countries=None: household, ) + app.register_blueprint(household_routes.household_bp) return app.test_client() diff --git a/tests/unit/endpoints/test_get_simulations.py b/tests/unit/routes/test_reform_impact_routes.py similarity index 79% rename from tests/unit/endpoints/test_get_simulations.py rename to tests/unit/routes/test_reform_impact_routes.py index 0e85afaa6..78c9eafbe 100644 --- a/tests/unit/endpoints/test_get_simulations.py +++ b/tests/unit/routes/test_reform_impact_routes.py @@ -9,10 +9,19 @@ from datetime import datetime +from flask import Flask from sqlalchemy import func, select from policyengine_api.data.v1_models import ReformImpact -from policyengine_api.endpoints.simulation import get_simulations +from policyengine_api.routes.reform_impact_routes import reform_impact_bp + + +def _get_simulations(max_results=100): + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(reform_impact_bp) + query = "" if max_results is None else f"?max_results={max_results}" + return app.test_client().get(f"/simulations{query}").get_json() def _seed_reform_impacts(orm_session, n: int) -> None: @@ -39,7 +48,7 @@ def _seed_reform_impacts(orm_session, n: int) -> None: def test_get_simulations_default_limit_caps_at_100(orm_session): _seed_reform_impacts(orm_session, 150) - result = get_simulations() + result = _get_simulations() assert len(result["result"]) == 100 @@ -47,20 +56,20 @@ def test_get_simulations_clamps_huge_max_results(orm_session): _seed_reform_impacts(orm_session, 50) # A caller passing an absurdly large value must not crash and # must not cause a full scan; the value is clamped at 1000. - result = get_simulations(max_results=10**9) + result = _get_simulations(max_results=10**9) assert len(result["result"]) == 50 # only 50 seeded def test_get_simulations_clamps_negative_max_results(orm_session): _seed_reform_impacts(orm_session, 5) # max_results of 0 or negative must still return something sane. - result = get_simulations(max_results=0) + result = _get_simulations(max_results=0) assert 1 <= len(result["result"]) <= 5 def test_get_simulations_defaults_when_none(orm_session): _seed_reform_impacts(orm_session, 10) - result = get_simulations(max_results=None) + result = _get_simulations(max_results=None) assert len(result["result"]) == 10 # fewer than the default 100 @@ -68,7 +77,7 @@ def test_get_simulations_rejects_non_integer_gracefully(orm_session): _seed_reform_impacts(orm_session, 5) # A string like "100; DROP TABLE reform_impact" must not reach # the SQL statement; it falls back to the default. - result = get_simulations(max_results="100; DROP TABLE reform_impact") + result = _get_simulations(max_results="100; DROP TABLE reform_impact") assert len(result["result"]) == 5 # And the table must still exist. diff --git a/tests/unit/endpoints/test_set_user_policy_dataset.py b/tests/unit/routes/test_set_user_policy_dataset.py similarity index 87% rename from tests/unit/endpoints/test_set_user_policy_dataset.py rename to tests/unit/routes/test_set_user_policy_dataset.py index 4d7001071..887dd653d 100644 --- a/tests/unit/endpoints/test_set_user_policy_dataset.py +++ b/tests/unit/routes/test_set_user_policy_dataset.py @@ -5,13 +5,13 @@ from sqlalchemy import select from policyengine_api.data.v1_models import UserPolicy -from policyengine_api.endpoints.policy import set_user_policy +from policyengine_api.routes.policy_routes import policy_bp def create_client(): app = Flask(__name__) app.config["TESTING"] = True - app.route("//user-policy", methods=["POST"])(set_user_policy) + app.register_blueprint(policy_bp) return app.test_client() diff --git a/tests/unit/endpoints/test_stage7_orm_endpoints.py b/tests/unit/routes/test_stage7_orm_routes.py similarity index 95% rename from tests/unit/endpoints/test_stage7_orm_endpoints.py rename to tests/unit/routes/test_stage7_orm_routes.py index 4f5bb80e5..629e1e36b 100644 --- a/tests/unit/endpoints/test_stage7_orm_endpoints.py +++ b/tests/unit/routes/test_stage7_orm_routes.py @@ -11,8 +11,8 @@ Policy, UserPolicy, ) -from policyengine_api.endpoints.household import get_household_under_policy -from policyengine_api.endpoints.policy import ( +from policyengine_api.routes.household_routes import get_household_under_policy +from policyengine_api.routes.policy_routes import ( get_user_policy, set_user_policy, update_user_policy, @@ -81,7 +81,7 @@ def test_household_under_policy_calculates_and_caches_json_as_an_object( ) with patch( - "policyengine_api.endpoints.household.household_calculation_service", + "policyengine_api.routes.household_routes.household_calculation_service", service, ): response = get_household_under_policy("us", "1", "2") diff --git a/tests/unit/endpoints/test_update_user_policy.py b/tests/unit/routes/test_update_user_policy.py similarity index 95% rename from tests/unit/endpoints/test_update_user_policy.py rename to tests/unit/routes/test_update_user_policy.py index bb8b8a18f..0950ed0af 100644 --- a/tests/unit/endpoints/test_update_user_policy.py +++ b/tests/unit/routes/test_update_user_policy.py @@ -13,13 +13,13 @@ from flask import Flask from policyengine_api.data.v1_models import UserPolicy -from policyengine_api.endpoints import update_user_policy +from policyengine_api.routes.policy_routes import policy_bp def _create_test_client() -> Flask: app = Flask(__name__) app.config["TESTING"] = True - app.route("//user-policy", methods=["PUT"])(update_user_policy) + app.register_blueprint(policy_bp) return app.test_client() diff --git a/tests/unit/services/test_household_calculation_service.py b/tests/unit/services/test_household_calculation_service.py index 191b9aec5..c24bed5a1 100644 --- a/tests/unit/services/test_household_calculation_service.py +++ b/tests/unit/services/test_household_calculation_service.py @@ -68,14 +68,14 @@ def _seed_inputs(factory): ) -def test_household_endpoint_and_country_do_not_manage_persistence(): - endpoint_source = (PACKAGE_ROOT / "endpoints" / "household.py").read_text( +def test_household_route_and_country_do_not_manage_persistence(): + route_source = (PACKAGE_ROOT / "routes" / "household_routes.py").read_text( encoding="utf-8" ) country_source = (PACKAGE_ROOT / "country.py").read_text(encoding="utf-8") - assert "get_v1_session_factory" not in endpoint_source - assert "from sqlalchemy" not in endpoint_source - assert "select(" not in endpoint_source + assert "get_v1_session_factory" not in route_source + assert "from sqlalchemy" not in route_source + assert "select(" not in route_source assert "get_v1_session_factory" not in country_source assert "Tracer(" not in country_source diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_stage7_local_service_boundaries.py index 2d972fb74..45191bd56 100644 --- a/tests/unit/services/test_stage7_local_service_boundaries.py +++ b/tests/unit/services/test_stage7_local_service_boundaries.py @@ -71,8 +71,8 @@ def test_analysis_routes_do_not_manage_sessions(): assert "sqlalchemy" not in source -def test_legacy_simulation_endpoint_does_not_manage_sessions(): - source = (SERVICE_ROOT.parent / "endpoints" / "simulation.py").read_text( +def test_reform_impact_route_does_not_manage_sessions(): + source = (SERVICE_ROOT.parent / "routes" / "reform_impact_routes.py").read_text( encoding="utf-8" ) assert "get_v1_session_factory" not in source diff --git a/tests/unit/services/test_user_policy_service.py b/tests/unit/services/test_user_policy_service.py index d271fc0b3..e92951161 100644 --- a/tests/unit/services/test_user_policy_service.py +++ b/tests/unit/services/test_user_policy_service.py @@ -8,8 +8,8 @@ ) -ENDPOINT_PATH = ( - Path(__file__).parents[3] / "policyengine_api" / "endpoints" / "policy.py" +ROUTE_PATH = ( + Path(__file__).parents[3] / "policyengine_api" / "routes" / "policy_routes.py" ) @@ -48,8 +48,8 @@ def test_user_policy_public_methods_do_not_accept_sessions(): assert "session_factory" not in parameters -def test_legacy_policy_endpoints_do_not_manage_sessions_or_queries(): - source = ENDPOINT_PATH.read_text(encoding="utf-8") +def test_policy_routes_do_not_manage_sessions_or_queries(): + source = ROUTE_PATH.read_text(encoding="utf-8") assert "get_v1_session_factory" not in source assert "from sqlalchemy" not in source assert "select(" not in source diff --git a/tests/unit/test_warmup.py b/tests/unit/test_warmup.py index 32aea3579..8bf1f9c2f 100644 --- a/tests/unit/test_warmup.py +++ b/tests/unit/test_warmup.py @@ -101,6 +101,8 @@ def test_asgi_runs_warmup_and_marks_ready(): def test_readiness_check_gates_on_readiness(): - src = (REPO / "policyengine_api/api.py").read_text(encoding="utf-8") + src = (REPO / "policyengine_api/routes/system_routes.py").read_text( + encoding="utf-8" + ) assert "is_ready" in src assert "status=503" in src diff --git a/tests/unit/test_yearly_var_removal.py b/tests/unit/test_yearly_var_removal.py index c5d81009c..9c13fd378 100644 --- a/tests/unit/test_yearly_var_removal.py +++ b/tests/unit/test_yearly_var_removal.py @@ -1,7 +1,9 @@ import copy from types import SimpleNamespace -from policyengine_api.endpoints.household import add_yearly_variables +from policyengine_api.services.household_calculation_service import ( + add_yearly_variables, +) TEST_YEAR = "2023" From 42fc237a8695878b5412fbd8fe43f3e0402306c9 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 22:32:52 +0300 Subject: [PATCH 69/89] refactor: remove the legacy endpoints package --- .../data/congressional_districts.py | 2 +- policyengine_api/endpoints/__init__.py | 13 - policyengine_api/endpoints/home.py | 7 - policyengine_api/endpoints/household.py | 196 ------------- policyengine_api/endpoints/policy.py | 275 ------------------ policyengine_api/endpoints/simulation.py | 59 ---- .../economy_comparison.py} | 11 +- policyengine_api/setup_data.py | 2 - pyproject.toml | 3 - tests/unit/endpoints/__init__.py | 1 - tests/unit/endpoints/economy/__init__.py | 1 - .../routes/test_flask_route_boundaries.py | 81 ++++++ .../test_economy_comparison.py} | 74 ++--- .../test_stage7_service_architecture.py | 12 +- 14 files changed, 132 insertions(+), 605 deletions(-) delete mode 100644 policyengine_api/endpoints/__init__.py delete mode 100644 policyengine_api/endpoints/home.py delete mode 100644 policyengine_api/endpoints/household.py delete mode 100644 policyengine_api/endpoints/policy.py delete mode 100644 policyengine_api/endpoints/simulation.py rename policyengine_api/{endpoints/economy/compare.py => services/economy_comparison.py} (99%) delete mode 100644 policyengine_api/setup_data.py delete mode 100644 tests/unit/endpoints/__init__.py delete mode 100644 tests/unit/endpoints/economy/__init__.py create mode 100644 tests/unit/routes/test_flask_route_boundaries.py rename tests/unit/{endpoints/economy/test_compare.py => services/test_economy_comparison.py} (92%) diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index f5218fd36..6dbac4be1 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -74,7 +74,7 @@ class CongressionalDistrictMetadataItem(BaseModel): Uses Pydantic BaseModel for: - Runtime validation of data integrity - Automatic serialization/deserialization - - Consistency with existing codebase patterns (see policyengine_api/endpoints/economy/compare.py) + - Consistency with existing codebase patterns (see services/economy_comparison.py) - Self-documenting schema with type hints """ diff --git a/policyengine_api/endpoints/__init__.py b/policyengine_api/endpoints/__init__.py deleted file mode 100644 index 9639c2317..000000000 --- a/policyengine_api/endpoints/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .home import get_home -from .household import ( - get_household_under_policy, - get_calculate, -) -from .policy import ( - get_policy_search, - set_user_policy, - get_user_policy, - update_user_policy, -) - -from .simulation import get_simulations diff --git a/policyengine_api/endpoints/home.py b/policyengine_api/endpoints/home.py deleted file mode 100644 index 33017204c..000000000 --- a/policyengine_api/endpoints/home.py +++ /dev/null @@ -1,7 +0,0 @@ -def get_home() -> str: - """Get the home page of the PolicyEngine API. - - Returns: - str: The home page. - """ - return f"

PolicyEngine households API

Use this API to compute the impact of public policy on individual households.

" diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py deleted file mode 100644 index 843436ceb..000000000 --- a/policyengine_api/endpoints/household.py +++ /dev/null @@ -1,196 +0,0 @@ -import json -from flask import Response, request -import logging -from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs -from policyengine_api.utils.input_validation import ( - find_unrecognized_inputs, - format_unrecognized_inputs_message, -) -from policyengine_api.utils.payload_validators import validate_country -from policyengine_core.errors import SituationParsingError - -from policyengine_api.services.household_calculation_service import ( - HouseholdCalculationService, - HouseholdNotFoundError, - InvalidHouseholdInputsError, - PolicyNotFoundError, - add_yearly_variables, -) - - -household_calculation_service = HouseholdCalculationService() - - -def get_countries(): - from policyengine_api.country import COUNTRIES - - return COUNTRIES - - -def get_invalid_inputs_response(household_json, policy_json, country): - invalid_inputs = find_unrecognized_inputs( - household_json, - policy_json, - country.metadata, - ) - if not invalid_inputs: - return None - - response_body = dict( - status="error", - message=format_unrecognized_inputs_message(invalid_inputs), - result=None, - errors=[invalid_input.to_dict() for invalid_input in invalid_inputs], - ) - return Response( - json.dumps(response_body), - status=400, - mimetype="application/json", - ) - - -@validate_country -def get_household_under_policy(country_id: str, household_id: str, policy_id: str): - """Get a household's output data under a given policy. - - Args: - country_id (str): The country ID. - household_id (str): The household ID. - policy_id (str): The policy ID. - """ - - try: - calculation = household_calculation_service.calculate_stored_household( - country_id, - int(household_id), - int(policy_id), - ) - except HouseholdNotFoundError: - response_body = dict( - status="error", - message=f"Household #{household_id} not found.", - ) - return Response( - json.dumps(response_body), - status=404, - mimetype="application/json", - ) - - except PolicyNotFoundError: - response_body = dict( - status="error", - message=f"Policy #{policy_id} not found.", - ) - return Response( - json.dumps(response_body), - status=404, - mimetype="application/json", - ) - - except InvalidHouseholdInputsError as error: - response_body = dict( - status="error", - message=format_unrecognized_inputs_message(error.invalid_inputs), - result=None, - errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], - ) - return Response( - json.dumps(response_body), - status=400, - mimetype="application/json", - ) - except Exception as e: - logging.exception(e) - response_body = dict( - status="error", - message=f"Error calculating household #{household_id} under policy #{policy_id}: {e}", - ) - return Response( - json.dumps(response_body), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message=None, - result=calculation.household, - ) - if calculation.warnings: - response_body["warnings"] = list(calculation.warnings) - return response_body - - -@validate_country -def get_calculate(country_id: str, add_missing: bool = False) -> dict: - """Lightweight endpoint for passing in household and policy JSON objects and calculating without storing data. - - Args: - country_id (str): The country ID. - """ - - payload = request.json - household_json = payload.get("household", {}) - policy_json = payload.get("policy", {}) - - if add_missing: - # Add in any missing yearly variables to household_json - household_json = add_yearly_variables(household_json, country_id) - - # Strip deprecated inputs from a copy before the engine runs so - # partners who still pass removed/renamed variables get a warning + - # working response instead of a `VariableNotFoundError` HTTP 500. - deprecated_inputs = drop_deprecated_inputs(household_json) - household_json = deprecated_inputs.household - deprecation_warnings = deprecated_inputs.warnings - - country = get_countries().get(country_id) - invalid_inputs_response = get_invalid_inputs_response( - household_json, - policy_json, - country, - ) - if invalid_inputs_response is not None: - return invalid_inputs_response - - try: - calculation = country.calculate(household_json, policy_json) - result = calculation if isinstance(calculation, dict) else calculation.household - except SituationParsingError as e: - # Malformed household payloads (e.g. a dict where a number belongs) - # are client errors, not server errors — mostly bot traffic. - response_body = dict( - status="error", - message=f"Invalid household payload: {e}", - result=None, - ) - return Response( - json.dumps(response_body), - status=400, - mimetype="application/json", - ) - except Exception as e: - logging.exception(e) - response_body = dict( - status="error", - message=f"Error calculating household under policy: {e}", - ) - return Response( - json.dumps(response_body), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message=None, - result=result, - ) - - warning_messages = [w.message for w in deprecation_warnings] - if warning_messages: - # Serialize to strings on the wire; the structured dataclasses - # stay available for any future caller that wants the fields. - response_body["warnings"] = warning_messages - - return response_body diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py deleted file mode 100644 index 023ea9abe..000000000 --- a/policyengine_api/endpoints/policy.py +++ /dev/null @@ -1,275 +0,0 @@ -from policyengine_api.utils.payload_validators import validate_country -import json -from flask import Response, request - -from policyengine_api.data.v1_models import UserPolicy -from policyengine_api.services.policy_service import PolicyService -from policyengine_api.services.user_policy_service import UserPolicyService - - -policy_service = PolicyService() -user_policy_service = UserPolicyService() - - -def _serialize_user_policy(user_policy: UserPolicy) -> dict: - return { - column.name: getattr(user_policy, column.name) - for column in UserPolicy.__table__.columns - } - - -@validate_country -def get_policy_search(country_id: str) -> dict: - """ - Search for policies for a specified country - - Args: - country_id (str): The country ID. - - Query Parameters: - query (str): Optional search term to filter policies - unique_only (bool): If true, return only unique policy-label combinations - - Returns: - Response: Json response with: - - On success: list of policies with id and label - - On failure: error message and appropriate status code - - Example: - GET /api/policies/us?query=tax&unique_only=true - """ - - query = request.args.get("query", "") - # The "json.loads" default type is added to convert lowercase - # "true" and "false" to Python-friendly bool values - unique_only = request.args.get("unique_only", default=False, type=json.loads) - - try: - results = policy_service.search_policies( - country_id, - query, - unique_only=unique_only, - ) - - if not results: - body = dict( - status="error", - message=f"No policies found for country {country_id} for query '{query}", - ) - return Response(json.dumps(body), status=404, mimetype="application/json") - - # Format into: [{ id: 1, label: "My policy" }, ...] - policies = [dict(id=result.id, label=result.label) for result in results] - body = dict( - status="ok", - message="Policies found", - result=policies, - ) - return Response(json.dumps(body), status=200, mimetype="application/json") - except Exception as e: - body = dict(status="error", message=f"Internal server error: {e}") - return Response(json.dumps(body), status=500, mimetype="application/json") - - -@validate_country -def set_user_policy(country_id: str) -> dict: - """ - Adds a record (if unique, barring type) to the user_policy table - that defines a particular policy as saved by a user to "their - policies"; this table also contains an optional "type" column that - is currently unused - """ - - payload = request.json - reform_label = payload.pop("reform_label", None) - reform_id = payload.pop("reform_id") - baseline_label = payload.pop("baseline_label", None) - baseline_id = payload.pop("baseline_id") - user_id = payload.pop("user_id") - year = payload.pop("year") - geography = payload.pop("geography") - dataset = payload.pop("dataset", None) - number_of_provisions = payload.pop("number_of_provisions") - api_version = payload.pop("api_version") - added_date = payload.pop("added_date") - updated_date = payload.pop("updated_date") - budgetary_impact = payload.pop("budgetary_impact", None) - policy_type = payload.pop("type", None) - - values = { - "country_id": country_id, - "reform_id": reform_id, - "reform_label": reform_label, - "baseline_id": baseline_id, - "baseline_label": baseline_label, - "user_id": user_id, - "year": year, - "geography": geography, - "dataset": dataset, - "number_of_provisions": number_of_provisions, - "api_version": api_version, - "added_date": added_date, - "updated_date": updated_date, - "budgetary_impact": budgetary_impact, - "type": policy_type, - } - - # When setting a user policy, "unique" records contain - # a unique set of the following pieces of data: - # country_id, reform_id, baseline_id, user_id, year, - # geography, reform_label, baseline_label, dataset; - # added_date, budgetary_impact, updated_date, - # number_of_provisions, and api_version are - # all changeable, and thus do not need - # to be tested; type is not yet implemented - - try: - creation = user_policy_service.create_or_get_user_policy(values) - user_policy = creation.user_policy - if not creation.created: - response = dict( - status="ok", - message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", - result=dict(id=user_policy.id), - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message="Record created successfully", - result=dict( - **_serialize_user_policy(user_policy), - ), - ) - - return Response( - json.dumps(response_body), - status=201, - mimetype="application/json", - ) - - -@validate_country -def get_user_policy(country_id: str, user_id: str) -> dict: - """ - Fetch all saved user policies by user id - """ - - user_policies = user_policy_service.list_user_policies(country_id, user_id) - rows_parsed = [_serialize_user_policy(row) for row in user_policies] - - if rows_parsed is None: - response = dict( - status="ok", - message=f"No saved policies found for user {user_id}", - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) - return dict( - status="ok", - message=None, - result=rows_parsed, - ) - - -# Whitelist of attributes that callers may modify via update_user_policy. -# Identity attributes are intentionally excluded because they define the -# record and must not be reassigned through this endpoint. -UPDATE_USER_POLICY_ALLOWED_FIELDS = frozenset( - { - "reform_label", - "baseline_label", - "year", - "geography", - "dataset", - "number_of_provisions", - "api_version", - "added_date", - "updated_date", - "budgetary_impact", - "type", - } -) - - -@validate_country -def update_user_policy(country_id: str) -> dict: - """ - Update any parts of a user_policy, given a user_policy ID - """ - - payload = request.json - if not isinstance(payload, dict) or "id" not in payload: - return Response( - json.dumps({"message": "Request body must include an 'id' field."}), - status=400, - mimetype="application/json", - ) - - user_policy_id = payload.pop("id") - - # Reject unknown or identity attributes before applying payload values to - # the mapped entity. - unknown_keys = [ - key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS - ] - if unknown_keys: - return Response( - json.dumps( - { - "message": ( - "Request body contains unsupported fields: " - f"{sorted(unknown_keys)}" - ) - } - ), - status=400, - mimetype="application/json", - ) - - if not payload: - return Response( - json.dumps( - {"message": "Request body must include at least one field to update."} - ), - status=400, - mimetype="application/json", - ) - - try: - user_policy_service.update_user_policy(user_policy_id, payload) - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message="Record updated successfully", - result=dict(id=user_policy_id), - ) - - return Response( - json.dumps(response_body), - status=200, - mimetype="application/json", - ) diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py deleted file mode 100644 index 7e83dee31..000000000 --- a/policyengine_api/endpoints/simulation.py +++ /dev/null @@ -1,59 +0,0 @@ -from policyengine_api.data.v1_models import ReformImpact -from policyengine_api.services.reform_impacts_service import ReformImpactsService - -""" - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON, - options_hash VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME -); - -""" - - -_MAX_SIMULATION_RESULTS = 1000 -_DEFAULT_SIMULATION_RESULTS = 100 -reform_impacts_service = ReformImpactsService() - - -def get_simulations( - max_results: int | None = 100, -): - # Get the last N simulations ordered by start time. - # - # LIMIT is always applied (unbounded scans against reform_impact - # are expensive) and max_results is clamped to [1, - # _MAX_SIMULATION_RESULTS] before being bound as a parameter, so - # the value can never be interpolated into the SQL string. - if max_results is None: - max_results = _DEFAULT_SIMULATION_RESULTS - try: - max_results = int(max_results) - except (TypeError, ValueError): - max_results = _DEFAULT_SIMULATION_RESULTS - max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - - result = reform_impacts_service.get_recent_reform_impacts(max_results) - - # Format into [{}] - - return { - "result": [ - { - column.name: getattr(impact, column.name) - for column in ReformImpact.__table__.columns - } - for impact in result - ] - } diff --git a/policyengine_api/endpoints/economy/compare.py b/policyengine_api/services/economy_comparison.py similarity index 99% rename from policyengine_api/endpoints/economy/compare.py rename to policyengine_api/services/economy_comparison.py index bb10ed7f6..a50f909f8 100644 --- a/policyengine_api/endpoints/economy/compare.py +++ b/policyengine_api/services/economy_comparison.py @@ -1,12 +1,13 @@ +"""Reusable economy comparison calculations and response models.""" + import logging -from microdf import MicroDataFrame, MicroSeries + +import h5py import numpy as np -import sys -from policyengine_core.tools.hugging_face import download_huggingface_dataset import pandas as pd -import h5py +from microdf import MicroSeries +from policyengine_core.tools.hugging_face import download_huggingface_dataset from pydantic import BaseModel -from typing import Any logger = logging.getLogger(__name__) diff --git a/policyengine_api/setup_data.py b/policyengine_api/setup_data.py deleted file mode 100644 index 47be04faf..000000000 --- a/policyengine_api/setup_data.py +++ /dev/null @@ -1,2 +0,0 @@ -def setup_data(): - import policyengine_api.endpoints.search diff --git a/pyproject.toml b/pyproject.toml index efd3959d9..0c11a1518 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,9 +65,6 @@ dev = [ "towncrier>=24.8.0", ] -[project.scripts] -policyengine-api-setup = "policyengine_api.setup_data:setup_data" - [tool.hatch.build.targets.wheel] packages = ["policyengine_api"] diff --git a/tests/unit/endpoints/__init__.py b/tests/unit/endpoints/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/tests/unit/endpoints/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/unit/endpoints/economy/__init__.py b/tests/unit/endpoints/economy/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/tests/unit/endpoints/economy/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/unit/routes/test_flask_route_boundaries.py b/tests/unit/routes/test_flask_route_boundaries.py new file mode 100644 index 000000000..47bdbbf00 --- /dev/null +++ b/tests/unit/routes/test_flask_route_boundaries.py @@ -0,0 +1,81 @@ +from pathlib import Path + +from flask import Flask + +from policyengine_api.routes.home_routes import home_bp +from policyengine_api.routes.household_routes import household_bp +from policyengine_api.routes.policy_routes import policy_bp +from policyengine_api.routes.reform_impact_routes import reform_impact_bp +from policyengine_api.routes.system_routes import system_bp + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +def test_legacy_flask_urls_are_owned_by_blueprints(): + app = Flask(__name__) + app.register_blueprint(home_bp) + app.register_blueprint(household_bp) + app.register_blueprint(policy_bp) + app.register_blueprint(reform_impact_bp) + app.register_blueprint(system_bp) + + rules = { + (rule.rule, method) + for rule in app.url_map.iter_rules() + for method in rule.methods + if method not in {"HEAD", "OPTIONS"} + } + + assert { + ("/", "GET"), + ("//policies", "GET"), + ("//household//policy/", "GET"), + ("//calculate", "POST"), + ("//calculate-full", "POST"), + ("//user-policy", "POST"), + ("//user-policy", "PUT"), + ("//user-policy/", "GET"), + ("/simulations", "GET"), + ("/liveness-check", "GET"), + ("/readiness-check", "GET"), + ("/specification", "GET"), + } <= rules + + +def test_flask_app_assembles_blueprints_without_legacy_endpoint_wiring(): + source = (PACKAGE_ROOT / "api.py").read_text(encoding="utf-8") + + assert "policyengine_api.endpoints" not in source + assert "from .endpoints" not in source + assert "Legacy endpoints" not in source + assert "get_policy_search" not in source + assert "get_household_under_policy" not in source + assert "get_calculate" not in source + assert "set_user_policy" not in source + assert "get_user_policy" not in source + assert "update_user_policy" not in source + assert "get_simulations" not in source + assert "@app.route" not in source + assert "app.route(" not in source + + +def test_legacy_endpoints_package_is_removed(): + assert not any((PACKAGE_ROOT / "endpoints").rglob("*.py")) + + +def test_economy_comparison_logic_is_not_in_the_http_layer(): + comparison_module = PACKAGE_ROOT / "services/economy_comparison.py" + + assert comparison_module.exists() + source = comparison_module.read_text(encoding="utf-8") + assert "from flask" not in source + assert "import flask" not in source + + +def test_calculation_route_delegates_domain_processing_to_its_service(): + source = (PACKAGE_ROOT / "routes/household_routes.py").read_text(encoding="utf-8") + + assert "country.calculate(" not in source + assert "find_unrecognized_inputs(" not in source + assert "drop_deprecated_inputs(" not in source diff --git a/tests/unit/endpoints/economy/test_compare.py b/tests/unit/services/test_economy_comparison.py similarity index 92% rename from tests/unit/endpoints/economy/test_compare.py rename to tests/unit/services/test_economy_comparison.py index 86d2187dd..f0e451d45 100644 --- a/tests/unit/endpoints/economy/test_compare.py +++ b/tests/unit/services/test_economy_comparison.py @@ -4,7 +4,7 @@ import pandas as pd from pydantic import ValidationError -from policyengine_api.endpoints.economy.compare import ( +from policyengine_api.services.economy_comparison import ( UKConstituencyBreakdownByConstituency, UKConstituencyBreakdown, UKLocalAuthorityBreakdownByLA, @@ -121,9 +121,9 @@ def test__given_non_uk_country_canada__returns_none(self): result = uk_local_authority_breakdown({}, {}, "ca") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_country__returns_breakdown( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -228,9 +228,9 @@ def test__outcome_bucket_categorization_logic(self): f"Failed for {percent_change}: expected {expected_bucket}, got {bucket}" ) - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__outcome_buckets_are_correct( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -264,9 +264,9 @@ def test__outcome_buckets_are_correct( assert result.outcomes_by_region["uk"]["Gain more than 5%"] == 1 assert result.outcomes_by_region["uk"]["Gain less than 5%"] == 0 - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__downloads_from_correct_repos( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -313,9 +313,9 @@ def test__given_constituency_region_with_code__returns_none(self): result = uk_local_authority_breakdown({}, {}, "uk", "constituency/E12345678") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_specific_la_region__returns_only_that_la( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -354,9 +354,9 @@ def test__given_specific_la_region__returns_only_that_la( assert "Aberdeen City" not in result.by_local_authority assert "Isle of Anglesey" not in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_country_scotland_region__returns_only_scottish_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -395,9 +395,9 @@ def test__given_country_scotland_region__returns_only_scottish_las( assert "Hartlepool" not in result.by_local_authority assert "Isle of Anglesey" not in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_region__returns_all_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -434,9 +434,9 @@ def test__given_uk_region__returns_all_las( assert "Aberdeen City" in result.by_local_authority assert "Isle of Anglesey" in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_no_region__returns_all_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -520,9 +520,9 @@ def test__given_local_authority_region_with_code__returns_none(self): result = uk_constituency_breakdown({}, {}, "uk", "local_authority/E06000016") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_specific_constituency_region__returns_only_that_constituency( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -563,9 +563,9 @@ def test__given_specific_constituency_region__returns_only_that_constituency( assert "Edinburgh East" not in result.by_constituency assert "Cardiff South" not in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_country_scotland_region__returns_only_scottish_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -602,9 +602,9 @@ def test__given_country_scotland_region__returns_only_scottish_constituencies( assert "Aldershot" not in result.by_constituency assert "Cardiff South" not in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_region__returns_all_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -641,9 +641,9 @@ def test__given_uk_region__returns_all_constituencies( assert "Edinburgh East" in result.by_constituency assert "Cardiff South" in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_no_region__returns_all_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -677,9 +677,9 @@ def test__given_no_region__returns_all_constituencies( assert result is not None assert len(result.by_constituency) == 3 - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__country_filter_uses_prefix_not_substring( self, mock_read_csv, mock_h5py_file, mock_download ): diff --git a/tests/unit/services/test_stage7_service_architecture.py b/tests/unit/services/test_stage7_service_architecture.py index 48be69c26..e4986e3d2 100644 --- a/tests/unit/services/test_stage7_service_architecture.py +++ b/tests/unit/services/test_stage7_service_architecture.py @@ -31,7 +31,10 @@ HouseholdService, ("get_household", "create_household", "update_household"), ), - (HouseholdCalculationService, ("calculate_stored_household",)), + ( + HouseholdCalculationService, + ("calculate_stored_household", "calculate_household"), + ), ( PolicyService, ("get_policy", "get_policy_json", "search_policies", "set_policy"), @@ -92,7 +95,6 @@ def test_presentation_layer_does_not_import_or_create_sqlalchemy_sessions(): offenders = [] presentation_paths = [ *sorted((PACKAGE_ROOT / "routes").glob("*.py")), - *sorted((PACKAGE_ROOT / "endpoints").rglob("*.py")), PACKAGE_ROOT / "country.py", ] banned_tokens = ( @@ -110,12 +112,12 @@ def test_presentation_layer_does_not_import_or_create_sqlalchemy_sessions(): def test_removed_persistence_abstractions_stay_removed(): - removed_paths = ( + removed_files = ( "data/v1_daos.py", "data/data.py", - "endpoints/economy/reform_impact.py", ) - assert [path for path in removed_paths if (PACKAGE_ROOT / path).exists()] == [] + assert [path for path in removed_files if (PACKAGE_ROOT / path).exists()] == [] + assert not any((PACKAGE_ROOT / "endpoints").rglob("*.py")) def test_changelog_describes_service_owned_sessions(): From 30da81f93895de4eadbecaa7f56ffa0d6bc6e179 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 22:49:41 +0300 Subject: [PATCH 70/89] refactor: centralize Flask error responses --- policyengine_api/response_factory.py | 26 +++++ policyengine_api/routes/economy_routes.py | 10 +- policyengine_api/routes/error_routes.py | 37 ++---- policyengine_api/routes/household_routes.py | 105 +++++------------- policyengine_api/routes/policy_routes.py | 89 +++++---------- .../payload_validators/validate_country.py | 8 +- .../routes/test_flask_route_boundaries.py | 10 ++ tests/unit/test_response_factory.py | 44 ++++++++ 8 files changed, 153 insertions(+), 176 deletions(-) create mode 100644 policyengine_api/response_factory.py create mode 100644 tests/unit/test_response_factory.py diff --git a/policyengine_api/response_factory.py b/policyengine_api/response_factory.py new file mode 100644 index 000000000..84cdd4240 --- /dev/null +++ b/policyengine_api/response_factory.py @@ -0,0 +1,26 @@ +"""Factories for consistent Flask response construction.""" + +import json +from typing import Any + +from flask import Response + + +def _make_error_response( + message: object, + status_code: int, + *, + include_status: bool = True, + mimetype: str | None = "application/json", + **payload_fields: Any, +) -> Response: + """Build a JSON error response while preserving legacy payload variants.""" + payload = {"message": str(message), **payload_fields} + if include_status: + payload = {"status": "error", **payload} + + return Response( + json.dumps(payload), + status=status_code, + mimetype=mimetype, + ) diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 8507cfe3e..56bb35e12 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -4,6 +4,7 @@ EconomicImpactResult, BudgetWindowEconomicImpactResult, ) +from policyengine_api.response_factory import _make_error_response from policyengine_api.utils import get_current_law_policy_id from policyengine_api.utils.payload_validators import validate_country from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS @@ -24,14 +25,7 @@ def _json_response(payload: dict, status: int = 200) -> Response: def _bad_request_response(message: str) -> Response: - return _json_response( - { - "status": "error", - "message": message, - "result": None, - }, - status=400, - ) + return _make_error_response(message, 400, result=None) @economy_bp.route( diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py index e9fced1c0..a4801bb09 100644 --- a/policyengine_api/routes/error_routes.py +++ b/policyengine_api/routes/error_routes.py @@ -1,67 +1,50 @@ -import json -from flask import Response, Blueprint +from flask import Blueprint, Response from werkzeug.exceptions import ( HTTPException, ) +from policyengine_api.response_factory import _make_error_response + error_bp = Blueprint("error", __name__) @error_bp.app_errorhandler(404) def response_404(error) -> Response: """Specific handler for 404 Not Found errors""" - return make_error_response(error, 404) + return _make_error_response(error, 404, result=None) @error_bp.app_errorhandler(400) def response_400(error) -> Response: """Specific handler for 400 Bad Request errors""" - return make_error_response(error, 400) + return _make_error_response(error, 400, result=None) @error_bp.app_errorhandler(401) def response_401(error) -> Response: """Specific handler for 401 Unauthorized errors""" - return make_error_response(error, 401) + return _make_error_response(error, 401, result=None) @error_bp.app_errorhandler(403) def response_403(error) -> Response: """Specific handler for 403 Forbidden errors""" - return make_error_response(error, 403) + return _make_error_response(error, 403, result=None) @error_bp.app_errorhandler(500) def response_500(error) -> Response: """Specific handler for 500 Internal Server errors""" - return make_error_response(error, 500) + return _make_error_response(error, 500, result=None) @error_bp.app_errorhandler(HTTPException) def response_http_exception(error: HTTPException) -> Response: """Generic handler for HTTPException; should be raised if no specific handler is found""" - return make_error_response(str(error), error.code) + return _make_error_response(error, error.code, result=None) @error_bp.app_errorhandler(Exception) def response_generic_error(error: Exception) -> Response: """Handler for any unhandled exceptions""" - return make_error_response(str(error), 500) - - -def make_error_response( - error, - status_code: int, -) -> Response: - """Create a generic error response""" - return Response( - json.dumps( - { - "status": "error", - "message": str(error), - "result": None, - } - ), - status_code, - mimetype="application/json", - ) + return _make_error_response(error, 500, result=None) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 5a9222cd4..731895889 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -7,6 +7,7 @@ from policyengine_api.data.v1_models import Household from policyengine_api.extensions import cache +from policyengine_api.response_factory import _make_error_response from policyengine_api.services.household_calculation_service import ( HouseholdCalculationService, HouseholdNotFoundError, @@ -173,57 +174,28 @@ def get_household_under_policy(country_id: str, household_id: str, policy_id: st int(policy_id), ) except HouseholdNotFoundError: - return Response( - json.dumps( - dict( - status="error", - message=f"Household #{household_id} not found.", - ) - ), - status=404, - mimetype="application/json", + return _make_error_response( + f"Household #{household_id} not found.", + 404, ) except PolicyNotFoundError: - return Response( - json.dumps( - dict( - status="error", - message=f"Policy #{policy_id} not found.", - ) - ), - status=404, - mimetype="application/json", + return _make_error_response( + f"Policy #{policy_id} not found.", + 404, ) except InvalidHouseholdInputsError as error: - return Response( - json.dumps( - dict( - status="error", - message=format_unrecognized_inputs_message(error.invalid_inputs), - result=None, - errors=[ - invalid_input.to_dict() - for invalid_input in error.invalid_inputs - ], - ) - ), - status=400, - mimetype="application/json", + return _make_error_response( + format_unrecognized_inputs_message(error.invalid_inputs), + 400, + result=None, + errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], ) except Exception as error: logging.exception(error) - return Response( - json.dumps( - dict( - status="error", - message=( - f"Error calculating household #{household_id} under policy " - f"#{policy_id}: {error}" - ), - ) - ), - status=500, - mimetype="application/json", + return _make_error_response( + f"Error calculating household #{household_id} under policy " + f"#{policy_id}: {error}", + 500, ) response_body = dict(status="ok", message=None, result=calculation.household) @@ -245,44 +217,23 @@ def _calculate(country_id: str, *, add_missing: bool) -> dict | Response: add_missing=add_missing, ) except InvalidHouseholdInputsError as error: - return Response( - json.dumps( - { - "status": "error", - "message": format_unrecognized_inputs_message(error.invalid_inputs), - "result": None, - "errors": [ - invalid_input.to_dict() - for invalid_input in error.invalid_inputs - ], - } - ), - status=400, - mimetype="application/json", + return _make_error_response( + format_unrecognized_inputs_message(error.invalid_inputs), + 400, + result=None, + errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], ) except SituationParsingError as error: - return Response( - json.dumps( - dict( - status="error", - message=f"Invalid household payload: {error}", - result=None, - ) - ), - status=400, - mimetype="application/json", + return _make_error_response( + f"Invalid household payload: {error}", + 400, + result=None, ) except Exception as error: logging.exception(error) - return Response( - json.dumps( - dict( - status="error", - message=f"Error calculating household under policy: {error}", - ) - ), - status=500, - mimetype="application/json", + return _make_error_response( + f"Error calculating household under policy: {error}", + 500, ) response_body = dict(status="ok", message=None, result=calculation.household) diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 2403cec1d..ca35cc2b7 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -4,6 +4,7 @@ from werkzeug.exceptions import BadRequest, NotFound from policyengine_api.data.v1_models import Policy, UserPolicy +from policyengine_api.response_factory import _make_error_response from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.user_policy_service import UserPolicyService from policyengine_api.utils.payload_validators import ( @@ -116,18 +117,9 @@ def get_policy_search(country_id: str) -> Response: unique_only=unique_only, ) if not results: - return Response( - json.dumps( - dict( - status="error", - message=( - f"No policies found for country {country_id} for query " - f"'{query}" - ), - ) - ), - status=404, - mimetype="application/json", + return _make_error_response( + f"No policies found for country {country_id} for query '{query}", + 404, ) policies = [dict(id=result.id, label=result.label) for result in results] @@ -143,15 +135,9 @@ def get_policy_search(country_id: str) -> Response: mimetype="application/json", ) except Exception as error: - return Response( - json.dumps( - dict( - status="error", - message=f"Internal server error: {error}", - ) - ), - status=500, - mimetype="application/json", + return _make_error_response( + f"Internal server error: {error}", + 500, ) @@ -212,16 +198,10 @@ def set_user_policy(country_id: str) -> Response: mimetype="application/json", ) except Exception as error: - return Response( - json.dumps( - { - "message": ( - f"Internal database error: {error}; please try again later." - ) - } - ), - status=500, - mimetype="application/json", + return _make_error_response( + f"Internal database error: {error}; please try again later.", + 500, + include_status=False, ) return Response( @@ -272,10 +252,10 @@ def update_user_policy(country_id: str) -> Response: """Update mutable fields on a saved user policy.""" payload = request.json if not isinstance(payload, dict) or "id" not in payload: - return Response( - json.dumps({"message": "Request body must include an 'id' field."}), - status=400, - mimetype="application/json", + return _make_error_response( + "Request body must include an 'id' field.", + 400, + include_status=False, ) user_policy_id = payload.pop("id") @@ -283,41 +263,26 @@ def update_user_policy(country_id: str) -> Response: key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS ] if unknown_keys: - return Response( - json.dumps( - { - "message": ( - "Request body contains unsupported fields: " - f"{sorted(unknown_keys)}" - ) - } - ), - status=400, - mimetype="application/json", + return _make_error_response( + f"Request body contains unsupported fields: {sorted(unknown_keys)}", + 400, + include_status=False, ) if not payload: - return Response( - json.dumps( - {"message": "Request body must include at least one field to update."} - ), - status=400, - mimetype="application/json", + return _make_error_response( + "Request body must include at least one field to update.", + 400, + include_status=False, ) try: user_policy_service.update_user_policy(user_policy_id, payload) except Exception as error: - return Response( - json.dumps( - { - "message": ( - f"Internal database error: {error}; please try again later." - ) - } - ), - status=500, - mimetype="application/json", + return _make_error_response( + f"Internal database error: {error}; please try again later.", + 500, + include_status=False, ) return Response( diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index cfd6c61f0..346840d34 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -1,11 +1,13 @@ from functools import wraps from typing import Union + from flask import Response -import json + from policyengine_api.country_validation import ( InvalidCountryError, ensure_supported_country, ) +from policyengine_api.response_factory import _make_error_response def validate_country(func): @@ -26,7 +28,9 @@ def validate_country_wrapper( try: ensure_supported_country(country_id) except InvalidCountryError as error: - return Response(json.dumps(error.to_payload()), status=400) + # Preserve the legacy v1 content type while native FastAPI routes + # are contractually required to match the Flask response. + return _make_error_response(error, 400, mimetype=None) return func(country_id, *args, **kwargs) return validate_country_wrapper diff --git a/tests/unit/routes/test_flask_route_boundaries.py b/tests/unit/routes/test_flask_route_boundaries.py index 47bdbbf00..ae2f01b60 100644 --- a/tests/unit/routes/test_flask_route_boundaries.py +++ b/tests/unit/routes/test_flask_route_boundaries.py @@ -79,3 +79,13 @@ def test_calculation_route_delegates_domain_processing_to_its_service(): assert "country.calculate(" not in source assert "find_unrecognized_inputs(" not in source assert "drop_deprecated_inputs(" not in source + + +def test_routes_do_not_construct_json_error_responses_inline(): + offenders = [] + for path in sorted((PACKAGE_ROOT / "routes").glob("*.py")): + source = path.read_text(encoding="utf-8") + if '"status": "error"' in source or 'status="error"' in source: + offenders.append(path.name) + + assert offenders == [] diff --git a/tests/unit/test_response_factory.py b/tests/unit/test_response_factory.py new file mode 100644 index 000000000..b8a27c9e5 --- /dev/null +++ b/tests/unit/test_response_factory.py @@ -0,0 +1,44 @@ +from policyengine_api.response_factory import _make_error_response + + +def test_error_response_factory_builds_standard_json_response(): + response = _make_error_response("Policy #2 not found.", 404) + + assert response.status_code == 404 + assert response.mimetype == "application/json" + assert response.get_json() == { + "status": "error", + "message": "Policy #2 not found.", + } + + +def test_error_response_factory_includes_structured_payload_fields(): + response = _make_error_response( + "Invalid inputs", + 400, + result=None, + errors=[{"name": "unknown_variable"}], + ) + + assert response.get_json() == { + "status": "error", + "message": "Invalid inputs", + "result": None, + "errors": [{"name": "unknown_variable"}], + } + + +def test_error_response_factory_can_preserve_message_only_v1_payloads(): + response = _make_error_response( + "Database error", + 500, + include_status=False, + ) + + assert response.get_json() == {"message": "Database error"} + + +def test_error_response_factory_can_preserve_a_legacy_default_mimetype(): + response = _make_error_response("Invalid country", 400, mimetype=None) + + assert response.content_type == "text/html; charset=utf-8" From dfbc0f40cb740b8d4d130ea8b367c2e2ac5925e4 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 23:19:05 +0300 Subject: [PATCH 71/89] test: describe persistence coverage by behavior --- ...sting_schema.py => test_v1_schema_metadata_compatibility.py} | 2 +- ...t_stage7_no_direct_sql.py => test_runtime_sql_boundaries.py} | 2 ++ ...m_routes.py => test_household_and_user_policy_orm_routes.py} | 2 ++ ...vice_boundaries.py => test_local_data_service_boundaries.py} | 2 ++ ...ies.py => test_orchestration_service_database_boundaries.py} | 2 ++ ..._stage7_dao_boundaries.py => test_orm_service_boundaries.py} | 2 ++ ...ries.py => test_run_and_spec_service_database_boundaries.py} | 2 ++ ...chitecture.py => test_service_owned_session_architecture.py} | 2 ++ ...st_stage5_routes.py => test_simulation_and_report_routes.py} | 2 ++ 9 files changed, 17 insertions(+), 1 deletion(-) rename tests/integration/{test_stage7_existing_schema.py => test_v1_schema_metadata_compatibility.py} (93%) rename tests/unit/data/{test_stage7_no_direct_sql.py => test_runtime_sql_boundaries.py} (96%) rename tests/unit/routes/{test_stage7_orm_routes.py => test_household_and_user_policy_orm_routes.py} (98%) rename tests/unit/services/{test_stage7_local_service_boundaries.py => test_local_data_service_boundaries.py} (97%) rename tests/unit/services/{test_stage7_orchestration_boundaries.py => test_orchestration_service_database_boundaries.py} (88%) rename tests/unit/services/{test_stage7_dao_boundaries.py => test_orm_service_boundaries.py} (94%) rename tests/unit/services/{test_stage7_run_service_boundaries.py => test_run_and_spec_service_database_boundaries.py} (89%) rename tests/unit/services/{test_stage7_service_architecture.py => test_service_owned_session_architecture.py} (98%) rename tests/unit/{test_stage5_routes.py => test_simulation_and_report_routes.py} (99%) diff --git a/tests/integration/test_stage7_existing_schema.py b/tests/integration/test_v1_schema_metadata_compatibility.py similarity index 93% rename from tests/integration/test_stage7_existing_schema.py rename to tests/integration/test_v1_schema_metadata_compatibility.py index ba4260a0d..d5c6f26db 100644 --- a/tests/integration/test_stage7_existing_schema.py +++ b/tests/integration/test_v1_schema_metadata_compatibility.py @@ -1,4 +1,4 @@ -"""Optional read-only metadata comparison for an existing v1 schema.""" +"""Compare ORM metadata with an existing v1 schema without mutating it.""" import os diff --git a/tests/unit/data/test_stage7_no_direct_sql.py b/tests/unit/data/test_runtime_sql_boundaries.py similarity index 96% rename from tests/unit/data/test_stage7_no_direct_sql.py rename to tests/unit/data/test_runtime_sql_boundaries.py index ae62e2d18..6f42a5708 100644 --- a/tests/unit/data/test_stage7_no_direct_sql.py +++ b/tests/unit/data/test_runtime_sql_boundaries.py @@ -1,3 +1,5 @@ +"""Guards that confine runtime SQL access to approved persistence layers.""" + from pathlib import Path diff --git a/tests/unit/routes/test_stage7_orm_routes.py b/tests/unit/routes/test_household_and_user_policy_orm_routes.py similarity index 98% rename from tests/unit/routes/test_stage7_orm_routes.py rename to tests/unit/routes/test_household_and_user_policy_orm_routes.py index 629e1e36b..5850aac74 100644 --- a/tests/unit/routes/test_stage7_orm_routes.py +++ b/tests/unit/routes/test_household_and_user_policy_orm_routes.py @@ -1,3 +1,5 @@ +"""ORM integration behavior for household and saved-policy routes.""" + from types import SimpleNamespace from unittest.mock import Mock, patch diff --git a/tests/unit/services/test_stage7_local_service_boundaries.py b/tests/unit/services/test_local_data_service_boundaries.py similarity index 97% rename from tests/unit/services/test_stage7_local_service_boundaries.py rename to tests/unit/services/test_local_data_service_boundaries.py index 45191bd56..cdf0ab4f7 100644 --- a/tests/unit/services/test_stage7_local_service_boundaries.py +++ b/tests/unit/services/test_local_data_service_boundaries.py @@ -1,3 +1,5 @@ +"""Persistence boundaries for services backed by local data.""" + import inspect from pathlib import Path diff --git a/tests/unit/services/test_stage7_orchestration_boundaries.py b/tests/unit/services/test_orchestration_service_database_boundaries.py similarity index 88% rename from tests/unit/services/test_stage7_orchestration_boundaries.py rename to tests/unit/services/test_orchestration_service_database_boundaries.py index 83cf7d8dd..c16e42049 100644 --- a/tests/unit/services/test_stage7_orchestration_boundaries.py +++ b/tests/unit/services/test_orchestration_service_database_boundaries.py @@ -1,3 +1,5 @@ +"""Database-access boundaries for orchestration services.""" + from pathlib import Path import pytest diff --git a/tests/unit/services/test_stage7_dao_boundaries.py b/tests/unit/services/test_orm_service_boundaries.py similarity index 94% rename from tests/unit/services/test_stage7_dao_boundaries.py rename to tests/unit/services/test_orm_service_boundaries.py index 6bbe87133..1c48c8494 100644 --- a/tests/unit/services/test_stage7_dao_boundaries.py +++ b/tests/unit/services/test_orm_service_boundaries.py @@ -1,3 +1,5 @@ +"""ORM and removed-DAO boundaries for entity services.""" + from pathlib import Path import pytest diff --git a/tests/unit/services/test_stage7_run_service_boundaries.py b/tests/unit/services/test_run_and_spec_service_database_boundaries.py similarity index 89% rename from tests/unit/services/test_stage7_run_service_boundaries.py rename to tests/unit/services/test_run_and_spec_service_database_boundaries.py index febe5a852..fcf6cf398 100644 --- a/tests/unit/services/test_stage7_run_service_boundaries.py +++ b/tests/unit/services/test_run_and_spec_service_database_boundaries.py @@ -1,3 +1,5 @@ +"""Database-access boundaries for run and specification services.""" + from pathlib import Path import pytest diff --git a/tests/unit/services/test_stage7_service_architecture.py b/tests/unit/services/test_service_owned_session_architecture.py similarity index 98% rename from tests/unit/services/test_stage7_service_architecture.py rename to tests/unit/services/test_service_owned_session_architecture.py index e4986e3d2..4f41dd600 100644 --- a/tests/unit/services/test_stage7_service_architecture.py +++ b/tests/unit/services/test_service_owned_session_architecture.py @@ -1,3 +1,5 @@ +"""Architecture guards for service-owned SQLAlchemy sessions.""" + import inspect from pathlib import Path diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_simulation_and_report_routes.py similarity index 99% rename from tests/unit/test_stage5_routes.py rename to tests/unit/test_simulation_and_report_routes.py index 7db0fa6fb..2d2502f7c 100644 --- a/tests/unit/test_stage5_routes.py +++ b/tests/unit/test_simulation_and_report_routes.py @@ -1,3 +1,5 @@ +"""Route behavior for simulation and report lifecycle operations.""" + import json from datetime import datetime From f61394cb01cdce53fac5d644df933900348bc572 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 23:41:00 +0300 Subject: [PATCH 72/89] fix: preserve ORM JSON in economy submissions --- policyengine_api/services/economy_service.py | 21 ++++---- tests/fixtures/services/economy_service.py | 6 +-- tests/unit/services/test_economy_service.py | 57 ++++++++++++++++++++ 3 files changed, 72 insertions(+), 12 deletions(-) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 742d0f8bc..24b88b9a9 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -3,7 +3,7 @@ import json import uuid from enum import Enum -from typing import Annotated, Any, Literal, Optional +from typing import Any, Literal, Optional import httpx import numpy as np @@ -278,7 +278,7 @@ def _get_policy_jsons( country_id: str, baseline_policy_id: int, reform_policy_id: int, - ) -> tuple[dict | None, dict | None]: + ) -> tuple[dict[str, Any], dict[str, Any]]: baseline = self._policies.get_policy_json( country_id, baseline_policy_id, @@ -287,15 +287,18 @@ def _get_policy_jsons( country_id, reform_policy_id, ) - return baseline, reform + return ( + self._parse_json_object(baseline), + self._parse_json_object(reform), + ) @staticmethod def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: - """Accept ORM-decoded objects and legacy JSON text at the read boundary.""" + """Accept ORM-decoded objects and legacy JSON text at a read boundary.""" parsed = json.loads(value) if isinstance(value, str) else value if not isinstance(parsed, dict): - raise TypeError("Expected a JSON object for reform impact data") + raise TypeError("Expected a JSON object") return parsed def get_economic_impact( @@ -1039,8 +1042,8 @@ def _handle_create_impact( def _setup_sim_options( self, country_id: str, - reform_policy: Annotated[str, "String-formatted JSON"], - baseline_policy: Annotated[str, "String-formatted JSON"], + reform_policy: dict[str, Any] | str, + baseline_policy: dict[str, Any] | str, region: str, time_period: str, scope: Literal["macro", "household"] = "macro", @@ -1058,8 +1061,8 @@ def _setup_sim_options( { "country": country_id, "scope": scope, - "reform": json.loads(reform_policy), - "baseline": json.loads(baseline_policy), + "reform": self._parse_json_object(reform_policy), + "baseline": self._parse_json_object(baseline_policy), "time_period": time_period, "include_cliffs": include_cliffs, "region": self._setup_region(country_id=country_id, region=region), diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 6b16019b5..8e1bfe0e9 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -96,12 +96,12 @@ def mock_policyengine_version(): @pytest.fixture def mock_policy_service(): - """Mock PolicyService with get_policy_json method.""" + """Mock the ORM-facing PolicyService with decoded JSON objects.""" mock_service = MagicMock() mock_service.get_policy_json.side_effect = lambda country_id, policy_id: ( - MOCK_REFORM_POLICY_JSON + json.loads(MOCK_REFORM_POLICY_JSON) if policy_id == MOCK_POLICY_ID - else MOCK_BASELINE_POLICY_JSON + else json.loads(MOCK_BASELINE_POLICY_JSON) ) with patch( diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 4926ac737..af9f4b14b 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -13,6 +13,7 @@ ImpactAction, ImpactStatus, ) +from policyengine_api.services.policy_service import PolicyService from tests.fixtures.services.economy_service import ( MOCK_API_VERSION, MOCK_BASELINE_POLICY_ID, @@ -311,6 +312,62 @@ def test__given_no_previous_impact__creates_new_simulation( assert write_values["options"] == MOCK_OPTIONS assert write_values["reform_impact_json"] == {} + def test__given_policies_created_through_orm__submits_decoded_json( + self, + orm_session_factory, + monkeypatch, + ): + policy_service = PolicyService(orm_session_factory) + baseline_policy_id, _, _ = policy_service.set_policy( + "us", + "ORM baseline", + {}, + ) + reform = {"gov.example.parameter": {"2026": 1}} + reform_policy_id, _, _ = policy_service.set_policy( + "us", + "ORM reform", + reform, + ) + + reform_impacts = MagicMock() + reform_impacts.get_all_reform_impacts_by_options_hash_prefix.return_value = [] + simulation_gateway = MagicMock() + simulation_gateway.resolve_app_name.return_value = ( + "policyengine-simulation-test", + MOCK_MODEL_VERSION, + ) + simulation_gateway.get_execution_id.return_value = "execution-1" + simulation_gateway.run.return_value.run_id = "run-1" + monkeypatch.setattr( + "policyengine_api.services.economy_service.logger", + MagicMock(), + ) + + service = EconomyService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + policy_service_=policy_service, + reform_impacts_service_=reform_impacts, + simulation_entrypoint_=simulation_gateway, + ) + + result = service.get_economic_impact( + country_id="us", + policy_id=reform_policy_id, + baseline_policy_id=baseline_policy_id, + region="us", + dataset="default", + time_period="2026", + options={}, + api_version="test", + ) + + assert result.status is ImpactStatus.COMPUTING + submitted = simulation_gateway.run.call_args.args[0] + assert submitted["baseline"] == {} + assert submitted["reform"] == reform + def test__given_no_previous_impact__includes_metadata_in_simulation_params( self, economy_service, From 1e4391107503af1e5d33f232bdfe43bbd3fd1257 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 23:41:54 +0300 Subject: [PATCH 73/89] fix: scope saved policy updates by country --- policyengine_api/routes/policy_routes.py | 13 ++++++++- .../services/user_policy_service.py | 8 +++++- tests/unit/routes/test_update_user_policy.py | 27 +++++++++++++++++-- .../unit/services/test_user_policy_service.py | 17 ++++++++++++ 4 files changed, 61 insertions(+), 4 deletions(-) diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index ca35cc2b7..f857954a4 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -277,7 +277,11 @@ def update_user_policy(country_id: str) -> Response: ) try: - user_policy_service.update_user_policy(user_policy_id, payload) + user_policy = user_policy_service.update_user_policy( + country_id, + user_policy_id, + payload, + ) except Exception as error: return _make_error_response( f"Internal database error: {error}; please try again later.", @@ -285,6 +289,13 @@ def update_user_policy(country_id: str) -> Response: include_status=False, ) + if user_policy is None: + return _make_error_response( + f"User policy #{user_policy_id} not found.", + 404, + include_status=False, + ) + return Response( json.dumps( dict( diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py index 07ca93f5e..07339310f 100644 --- a/policyengine_api/services/user_policy_service.py +++ b/policyengine_api/services/user_policy_service.py @@ -89,11 +89,17 @@ def list_user_policies( def update_user_policy( self, + country_id: str, user_policy_id: int, values: Mapping[str, Any], ) -> UserPolicy | None: with self._sessions.begin() as session: - user_policy = session.get(UserPolicy, user_policy_id) + user_policy = session.scalar( + select(UserPolicy).where( + UserPolicy.id == user_policy_id, + UserPolicy.country_id == country_id, + ) + ) if user_policy is None: return None for field, value in values.items(): diff --git a/tests/unit/routes/test_update_user_policy.py b/tests/unit/routes/test_update_user_policy.py index 0950ed0af..262f6137e 100644 --- a/tests/unit/routes/test_update_user_policy.py +++ b/tests/unit/routes/test_update_user_policy.py @@ -23,10 +23,10 @@ def _create_test_client() -> Flask: return app.test_client() -def _insert_user_policy(orm_session) -> int: +def _insert_user_policy(orm_session, *, country_id: str = "us") -> int: now = int(time.time()) policy = UserPolicy( - country_id="us", + country_id=country_id, reform_label="old label", reform_id=2, baseline_label=None, @@ -108,3 +108,26 @@ def test_update_user_policy_requires_at_least_one_field(orm_session): client = _create_test_client() response = client.put("/us/user-policy", json={"id": policy_id}) assert response.status_code == 400 + + +def test_update_user_policy_returns_not_found_for_missing_id(): + response = _create_test_client().put( + "/us/user-policy", + json={"id": 999_999, "reform_label": "missing"}, + ) + + assert response.status_code == 404 + assert response.get_json()["message"] == "User policy #999999 not found." + + +def test_update_user_policy_does_not_update_another_country(orm_session): + policy_id = _insert_user_policy(orm_session, country_id="uk") + + response = _create_test_client().put( + "/us/user-policy", + json={"id": policy_id, "reform_label": "wrong country"}, + ) + + assert response.status_code == 404 + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "old label" diff --git a/tests/unit/services/test_user_policy_service.py b/tests/unit/services/test_user_policy_service.py index e92951161..44a96f363 100644 --- a/tests/unit/services/test_user_policy_service.py +++ b/tests/unit/services/test_user_policy_service.py @@ -64,6 +64,7 @@ def test_create_reuse_list_and_update_user_policy(orm_session_factory): ) listed = service.list_user_policies("us", "auth0|one") updated = service.update_user_policy( + "us", created.user_policy.id, {"reform_label": "Updated", "updated_date": 3}, ) @@ -76,3 +77,19 @@ def test_create_reuse_list_and_update_user_policy(orm_session_factory): assert isinstance(listed[0], UserPolicy) assert updated.reform_label == "Updated" assert updated.updated_date == 3 + + +def test_update_user_policy_requires_matching_country(orm_session_factory): + service = UserPolicyService(orm_session_factory) + created = service.create_or_get_user_policy(_values(country_id="uk")) + + assert ( + service.update_user_policy( + "us", + created.user_policy.id, + {"reform_label": "Wrong country"}, + ) + is None + ) + stored = service.list_user_policies("uk", "auth0|one")[0] + assert stored.reform_label == "Reform" From adbb676959cbf81bb3a954e3175fce70974e4b03 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 10 Aug 2026 23:43:57 +0300 Subject: [PATCH 74/89] fix: preserve v1 reform impact response shapes --- .../routes/reform_impact_routes.py | 41 ++++++++++++++----- .../unit/routes/test_reform_impact_routes.py | 35 +++++++++++++++- 2 files changed, 63 insertions(+), 13 deletions(-) diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py index ac9866bd7..e49076304 100644 --- a/policyengine_api/routes/reform_impact_routes.py +++ b/policyengine_api/routes/reform_impact_routes.py @@ -1,4 +1,7 @@ -from flask import Blueprint, request +from datetime import datetime +import json + +from flask import Blueprint, Response, request from policyengine_api.data.v1_models import ReformImpact from policyengine_api.services.reform_impacts_service import ReformImpactsService @@ -11,6 +14,24 @@ _DEFAULT_SIMULATION_RESULTS = 100 +def _serialize_v1_reform_impact(impact: ReformImpact) -> dict: + """Project canonical ORM values onto the historical v1 response shape.""" + + result = { + column.name: getattr(impact, column.name) + for column in ReformImpact.__table__.columns + } + for field in ("options_json", "reform_impact_json"): + value = result[field] + if value is not None and not isinstance(value, str): + result[field] = json.dumps(value) + for field in ("start_time", "end_time"): + value = result[field] + if isinstance(value, datetime): + result[field] = str(value) + return result + + def _parse_result_limit(value: str | None) -> int: if value is None: return _DEFAULT_SIMULATION_RESULTS @@ -22,17 +43,15 @@ def _parse_result_limit(value: str | None) -> int: @reform_impact_bp.route("/simulations", methods=["GET"]) -def get_simulations() -> dict: +def get_simulations() -> Response: """Return recent reform impacts, bounded to protect the database query.""" result_limit = _parse_result_limit(request.args.get("max_results")) impacts = reform_impacts_service.get_recent_reform_impacts(result_limit) - return { - "result": [ - { - column.name: getattr(impact, column.name) - for column in ReformImpact.__table__.columns - } - for impact in impacts - ] - } + return Response( + json.dumps( + {"result": [_serialize_v1_reform_impact(impact) for impact in impacts]} + ), + status=200, + mimetype="application/json", + ) diff --git a/tests/unit/routes/test_reform_impact_routes.py b/tests/unit/routes/test_reform_impact_routes.py index 78c9eafbe..4fe072ee4 100644 --- a/tests/unit/routes/test_reform_impact_routes.py +++ b/tests/unit/routes/test_reform_impact_routes.py @@ -8,20 +8,28 @@ """ from datetime import datetime +import json +from fastapi.testclient import TestClient from flask import Flask from sqlalchemy import func, select +from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.data.v1_models import ReformImpact from policyengine_api.routes.reform_impact_routes import reform_impact_bp def _get_simulations(max_results=100): + app = _create_app() + query = "" if max_results is None else f"?max_results={max_results}" + return app.test_client().get(f"/simulations{query}").get_json() + + +def _create_app() -> Flask: app = Flask(__name__) app.config["TESTING"] = True app.register_blueprint(reform_impact_bp) - query = "" if max_results is None else f"?max_results={max_results}" - return app.test_client().get(f"/simulations{query}").get_json() + return app def _seed_reform_impacts(orm_session, n: int) -> None: @@ -40,6 +48,7 @@ def _seed_reform_impacts(orm_session, n: int) -> None: reform_impact_json={}, status="complete", start_time=datetime(2026, 1, 1, 0, i // 60, i % 60), + end_time=datetime(2026, 1, 1, 1, i // 60, i % 60), execution_id=f"exec-{i}", ) ) @@ -82,3 +91,25 @@ def test_get_simulations_rejects_non_integer_gracefully(orm_session): # And the table must still exist. assert orm_session.scalar(select(func.count()).select_from(ReformImpact)) == 5 + + +def test_get_simulations_preserves_v1_json_and_timestamp_fields(orm_session): + _seed_reform_impacts(orm_session, 1) + + impact = _get_simulations(max_results=1)["result"][0] + + assert json.loads(impact["options_json"]) == {} + assert json.loads(impact["reform_impact_json"]) == {} + assert impact["start_time"] == "2026-01-01 00:00:00" + assert impact["end_time"] == "2026-01-01 01:00:00" + + +def test_get_simulations_matches_through_fastapi_fallback(orm_session): + _seed_reform_impacts(orm_session, 1) + app = _create_app() + + flask_response = app.test_client().get("/simulations?max_results=1") + asgi_response = TestClient(create_asgi_app(app)).get("/simulations?max_results=1") + + assert asgi_response.status_code == flask_response.status_code == 200 + assert asgi_response.json() == flask_response.get_json() From b76a080fa85e0f16f5ff388f2ec926ff57a7494c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 11 Aug 2026 00:04:44 +0300 Subject: [PATCH 75/89] test: qualify Alembic against ephemeral MySQL --- alembic.ini | 1 - docs/engineering/skills/alembic-migrations.md | 46 ++++++++++++++ migrations/env.py | 15 ++++- .../test_alembic_mysql_lifecycle.py | 63 +++++++++++++++++++ tests/unit/data/test_alembic_baseline.py | 17 +++++ 5 files changed, 138 insertions(+), 4 deletions(-) create mode 100644 tests/integration/test_alembic_mysql_lifecycle.py diff --git a/alembic.ini b/alembic.ini index 6d23a99ae..115fa521e 100644 --- a/alembic.ini +++ b/alembic.ini @@ -2,7 +2,6 @@ script_location = migrations prepend_sys_path = . path_separator = os -sqlalchemy.url = sqlite+pysqlite:///policyengine_api/data/policyengine.db [loggers] keys = root,sqlalchemy,alembic diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 10c1d55b3..cbb3a5b82 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -34,3 +34,49 @@ Before committing a migration: Never print database credentials or embed them in Alembic configuration. Never run a generated baseline's create operations against an existing database; verify schema equivalence and stamp it instead. + +## API v1 database targets + +Alembic manages the API v1 MySQL schema only. The local SQLite database is a +temporary application cache and is deliberately outside Alembic's lifecycle. +Every Alembic command must therefore receive an explicit MySQL URL: + +```bash +export ALEMBIC_DATABASE_URL="mysql+pymysql://..." +uv run alembic upgrade head +``` + +Do not put a database URL in `alembic.ini`. Keeping the configuration without a +default prevents an omitted environment variable from silently migrating the +wrong database. + +### Fresh database qualification + +Use a disposable schema named `policyengine_alembic_test` on local MySQL, then +run the lifecycle integration test: + +```bash +ALEMBIC_DATABASE_URL="mysql+pymysql://.../policyengine_alembic_test" \ + uv run pytest tests/integration/test_alembic_mysql_lifecycle.py -q +``` + +The test must upgrade from an empty schema, run `alembic check`, compare the +live schema with the ORM metadata, downgrade to `base`, and upgrade to `head` +again. Its host and schema-name checks prevent it from running against a shared +or deployed database. + +### Existing database adoption gate + +The baseline represents tables that already exist in deployed API v1 +databases. It must never be upgraded into one of those databases. Before the +first Alembic-managed release for each environment: + +1. Supply a read-only URL as `STAGE7_EXISTING_DATABASE_URL` and run + `uv run pytest tests/integration/test_v1_schema_metadata_compatibility.py -q`. +2. Review the result and stop if any metadata drift is reported. +3. With a separately authorized migration connection, run + `uv run alembic stamp head`. +4. Confirm the application starts without emitting schema DDL. + +Stamping is an explicit release operation after the read-only comparison; it +must not run automatically during application startup or ordinary CI. diff --git a/migrations/env.py b/migrations/env.py index ceafe5835..a7eb19f11 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -5,6 +5,7 @@ from alembic import context from sqlalchemy import engine_from_config, pool +from sqlalchemy.engine import make_url from policyengine_api.data.v1_models import V1Base @@ -13,9 +14,17 @@ if config.config_file_name is not None: fileConfig(config.config_file_name) -database_url = os.environ.get("ALEMBIC_DATABASE_URL") -if database_url: - config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) +database_url = os.environ.get("ALEMBIC_DATABASE_URL") or config.get_main_option( + "sqlalchemy.url" +) +if not database_url: + raise RuntimeError( + "ALEMBIC_DATABASE_URL is required; the local SQLite cache is not an " + "Alembic migration target" + ) +if make_url(database_url).get_backend_name() != "mysql": + raise RuntimeError("Alembic migrations must target a MySQL database") +config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) target_metadata = V1Base.metadata diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py new file mode 100644 index 000000000..c3782046d --- /dev/null +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -0,0 +1,63 @@ +"""Exercise the complete Alembic lifecycle against an ephemeral MySQL schema.""" + +import os + +from alembic import command +from alembic.autogenerate import compare_metadata +from alembic.config import Config +from alembic.migration import MigrationContext +import pytest +from sqlalchemy import create_engine, inspect +from sqlalchemy.engine import make_url + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base + + +def _ephemeral_mysql_url() -> str: + database_url = os.environ.get("ALEMBIC_DATABASE_URL", "") + if not database_url: + pytest.skip("ALEMBIC_DATABASE_URL is not set") + + url = make_url(database_url) + if url.get_backend_name() != "mysql": + pytest.fail("ALEMBIC_DATABASE_URL must use MySQL for this test") + if url.host not in {"127.0.0.1", "localhost"}: + pytest.fail("Alembic lifecycle tests may only target local MySQL") + if url.database != "policyengine_alembic_test": + pytest.fail( + "Alembic lifecycle tests require the policyengine_alembic_test schema" + ) + return database_url + + +def _alembic_config(database_url: str) -> Config: + config = Config(str(REPO / "alembic.ini")) + config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) + return config + + +def test_fresh_upgrade_check_downgrade_and_reupgrade(): + database_url = _ephemeral_mysql_url() + config = _alembic_config(database_url) + engine = create_engine(database_url) + + try: + command.upgrade(config, "head") + command.check(config) + + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() is not None + assert compare_metadata(context, V1Base.metadata) == [] + + command.downgrade(config, "base") + assert set(inspect(engine).get_table_names()) <= {"alembic_version"} + + command.upgrade(config, "head") + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() is not None + assert compare_metadata(context, V1Base.metadata) == [] + finally: + engine.dispose() diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index f16eccb4f..21e826388 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -3,6 +3,7 @@ from alembic import command from alembic.config import Config from alembic.script import ScriptDirectory +import pytest from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base @@ -18,6 +19,22 @@ def _mysql_offline_config() -> tuple[Config, StringIO]: return config, output +def test_default_configuration_requires_an_explicit_database_url(monkeypatch): + monkeypatch.delenv("ALEMBIC_DATABASE_URL", raising=False) + config = Config(str(REPO / "alembic.ini"), output_buffer=StringIO()) + + with pytest.raises(RuntimeError, match="ALEMBIC_DATABASE_URL"): + command.upgrade(config, "head", sql=True) + + +def test_sqlite_is_rejected_as_a_migration_target(monkeypatch): + monkeypatch.setenv("ALEMBIC_DATABASE_URL", "sqlite+pysqlite:///:memory:") + config = Config(str(REPO / "alembic.ini"), output_buffer=StringIO()) + + with pytest.raises(RuntimeError, match="MySQL"): + command.upgrade(config, "head", sql=True) + + def test_baseline_renders_the_v1_schema_for_mysql_without_connecting(): config, output = _mysql_offline_config() From e7f844f06697d7cec0082e6b0c2d54df90f9a8c8 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 11 Aug 2026 00:09:31 +0300 Subject: [PATCH 76/89] ci: qualify Alembic migrations on MySQL --- .github/workflows/pr.yml | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 2dd9b5235..ac3c60c3c 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -50,6 +50,38 @@ jobs: python-version: "3.12" - name: Run quality guards run: python scripts/run_quality_guards.py + + alembic-mysql: + name: Alembic MySQL lifecycle + runs-on: ubuntu-latest + services: + mysql: + image: mysql:8.4 + env: + MYSQL_ROOT_PASSWORD: policyengine_test + MYSQL_DATABASE: policyengine_alembic_test + ports: + - 3306:3306 + options: >- + --health-cmd="mysqladmin ping --silent" + --health-interval=5s + --health-timeout=5s + --health-retries=20 + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install dependencies + run: make install + - name: Test Alembic configuration and MySQL lifecycle + run: >- + pytest tests/unit/data/test_alembic_baseline.py + tests/integration/test_alembic_mysql_lifecycle.py -q + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test check-changelog: name: Check changelog fragment runs-on: ubuntu-latest From fe7b944ebe4091a8fe990657ff5a5b4feb8de15e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 11 Aug 2026 16:54:42 +0300 Subject: [PATCH 77/89] fix: remove orphaned question table via Alembic --- changelog.d/3788.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 29 +++++-- ...9b3a056e_remove_orphaned_question_table.py | 43 ++++++++++ .../test_alembic_mysql_lifecycle.py | 84 ++++++++++++++++++- tests/unit/data/test_alembic_baseline.py | 7 +- 5 files changed, 151 insertions(+), 14 deletions(-) create mode 100644 migrations/versions/01e49b3a056e_remove_orphaned_question_table.py diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index 66844423b..f9bcbdd45 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving the existing database schema and public API contracts. +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving public API contracts, including a generated migration that removes the orphaned `question` prototype table. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index cbb3a5b82..0f1f0a35a 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -61,9 +61,9 @@ ALEMBIC_DATABASE_URL="mysql+pymysql://.../policyengine_alembic_test" \ ``` The test must upgrade from an empty schema, run `alembic check`, compare the -live schema with the ORM metadata, downgrade to `base`, and upgrade to `head` -again. Its host and schema-name checks prevent it from running against a shared -or deployed database. +live schema with the ORM metadata, downgrade one revision, and upgrade to +`head` again. Its host and schema-name checks prevent it from running against a +shared or deployed database. ### Existing database adoption gate @@ -73,10 +73,25 @@ first Alembic-managed release for each environment: 1. Supply a read-only URL as `STAGE7_EXISTING_DATABASE_URL` and run `uv run pytest tests/integration/test_v1_schema_metadata_compatibility.py -q`. -2. Review the result and stop if any metadata drift is reported. -3. With a separately authorized migration connection, run - `uv run alembic stamp head`. -4. Confirm the application starts without emitting schema DDL. +2. Review every reported difference and stop if any difference is not exactly + accounted for by a reviewed post-baseline migration. +3. With a separately authorized migration connection, stamp the existing + database at the generated baseline revision: + `uv run alembic stamp eafc2a547a4e`. +4. Run `uv run alembic upgrade head` to apply every reviewed post-baseline + migration. +5. Repeat the metadata comparison and require zero remaining differences. +6. Confirm the application starts without emitting schema DDL. Stamping is an explicit release operation after the read-only comparison; it must not run automatically during application startup or ordinary CI. + +### Production-only table cleanup + +Revision `01e49b3a056e` was autogenerated after reflecting the orphaned +production `question` table. Its only post-generation correction is +`if_exists=True` on the generated drop: deployed databases contain the table, +while fresh databases built from the baseline do not. Integration tests cover +both paths and verify the generated downgrade recreates the table structure. +The nine removed prototype rows are not recreated by downgrade, so retain a +database backup if their contents may be needed later. diff --git a/migrations/versions/01e49b3a056e_remove_orphaned_question_table.py b/migrations/versions/01e49b3a056e_remove_orphaned_question_table.py new file mode 100644 index 000000000..e0869cff7 --- /dev/null +++ b/migrations/versions/01e49b3a056e_remove_orphaned_question_table.py @@ -0,0 +1,43 @@ +"""remove orphaned question table + +Revision ID: 01e49b3a056e +Revises: eafc2a547a4e +Create Date: 2026-08-11 16:46:20.938200 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "01e49b3a056e" +down_revision: Union[str, None] = "eafc2a547a4e" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + # Production has this orphaned prototype table; fresh baseline schemas do not. + op.drop_table("question", if_exists=True) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "question", + sa.Column("question_id", mysql.INTEGER(), autoincrement=True, nullable=False), + sa.Column("question", mysql.LONGTEXT(), nullable=False), + sa.Column("answer", mysql.LONGTEXT(), nullable=True), + sa.Column("policy_id", mysql.INTEGER(), autoincrement=False, nullable=True), + sa.Column("country_id", mysql.VARCHAR(length=3), nullable=False), + sa.Column("subtask", mysql.VARCHAR(length=32), nullable=False), + sa.Column("status", mysql.VARCHAR(length=32), nullable=False), + sa.PrimaryKeyConstraint("question_id"), + mysql_collate="utf8mb4_0900_ai_ci", + mysql_default_charset="utf8mb4", + mysql_engine="InnoDB", + ) + # ### end Alembic commands ### diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index c3782046d..abfa9a4be 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -7,13 +7,43 @@ from alembic.config import Config from alembic.migration import MigrationContext import pytest -from sqlalchemy import create_engine, inspect +from sqlalchemy import ( + Column, + Integer, + MetaData, + String, + Table, + Text, + create_engine, + inspect, +) +from sqlalchemy.dialects.mysql import LONGTEXT from sqlalchemy.engine import make_url from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base +BASELINE_REVISION = "eafc2a547a4e" + + +def _deployed_question_table() -> Table: + """Describe the orphaned production table without adding it to ORM metadata.""" + + metadata = MetaData() + return Table( + "question", + metadata, + Column("question_id", Integer, primary_key=True, autoincrement=True), + Column("question", Text().with_variant(LONGTEXT(), "mysql"), nullable=False), + Column("answer", Text().with_variant(LONGTEXT(), "mysql")), + Column("policy_id", Integer), + Column("country_id", String(3), nullable=False), + Column("subtask", String(32), nullable=False), + Column("status", String(32), nullable=False), + ) + + def _ephemeral_mysql_url() -> str: database_url = os.environ.get("ALEMBIC_DATABASE_URL", "") if not database_url: @@ -41,8 +71,11 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): database_url = _ephemeral_mysql_url() config = _alembic_config(database_url) engine = create_engine(database_url) + question = _deployed_question_table() try: + command.downgrade(config, "base") + question.drop(engine, checkfirst=True) command.upgrade(config, "head") command.check(config) @@ -51,8 +84,8 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): assert context.get_current_revision() is not None assert compare_metadata(context, V1Base.metadata) == [] - command.downgrade(config, "base") - assert set(inspect(engine).get_table_names()) <= {"alembic_version"} + command.downgrade(config, BASELINE_REVISION) + assert "question" in inspect(engine).get_table_names() command.upgrade(config, "head") with engine.connect() as connection: @@ -61,3 +94,48 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): assert compare_metadata(context, V1Base.metadata) == [] finally: engine.dispose() + + +def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema(): + database_url = _ephemeral_mysql_url() + config = _alembic_config(database_url) + engine = create_engine(database_url) + question = _deployed_question_table() + + try: + command.downgrade(config, "base") + question.drop(engine, checkfirst=True) + command.upgrade(config, BASELINE_REVISION) + question.create(engine) + with engine.begin() as connection: + connection.execute( + question.insert(), + { + "question": "Historical prototype", + "country_id": "uk", + "subtask": "complete", + "status": "ok", + }, + ) + + command.upgrade(config, "head") + + assert "question" not in inspect(engine).get_table_names() + + command.downgrade(config, BASELINE_REVISION) + + assert "question" in inspect(engine).get_table_names() + assert { + column["name"] for column in inspect(engine).get_columns("question") + } == { + "question_id", + "question", + "answer", + "policy_id", + "country_id", + "subtask", + "status", + } + finally: + command.upgrade(config, "head") + engine.dispose() diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index 21e826388..12c9530b3 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -49,10 +49,11 @@ def test_baseline_renders_the_v1_schema_for_mysql_without_connecting(): def test_baseline_is_the_single_root_revision(): config, _ = _mysql_offline_config() scripts = ScriptDirectory.from_config(config) - head = scripts.get_revision(scripts.get_current_head()) + baseline = scripts.get_revision("eafc2a547a4e") - assert head is not None - assert head.down_revision is None + assert scripts.get_bases() == ["eafc2a547a4e"] + assert baseline is not None + assert baseline.down_revision is None def test_baseline_renders_a_complete_mysql_downgrade_without_connecting(): From 228c025d340e1783504859e149d62d12c61b72a3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:04:55 +0300 Subject: [PATCH 78/89] fix: align v1 metadata with deployed schema --- changelog.d/3788.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 10 ++++ ...5f97_separate_local_tracer_and_require_.py | 54 +++++++++++++++++++ ..._require_explicit_reform_impact_dataset.py | 41 ++++++++++++++ policyengine_api/data/local_database.py | 2 + policyengine_api/data/local_models.py | 26 +++++++++ policyengine_api/data/v1_models.py | 10 ---- .../services/household_calculation_service.py | 2 +- .../services/tracer_analysis_service.py | 2 +- .../services/tracer_fixture_service.py | 2 +- .../test_alembic_mysql_lifecycle.py | 34 +++++++++++- tests/unit/conftest.py | 5 +- tests/unit/data/test_sqlalchemy_v2.py | 1 + tests/unit/data/test_v1_models.py | 7 ++- .../test_direct_orm_local_analysis.py | 3 +- .../test_household_calculation_service.py | 2 +- 16 files changed, 184 insertions(+), 19 deletions(-) create mode 100644 migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py create mode 100644 migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py create mode 100644 policyengine_api/data/local_models.py diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index f9bcbdd45..052617746 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving public API contracts, including a generated migration that removes the orphaned `question` prototype table. +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving public API contracts, including generated migrations that remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 0f1f0a35a..6e50c61b4 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -95,3 +95,13 @@ while fresh databases built from the baseline do not. Integration tests cover both paths and verify the generated downgrade recreates the table structure. The nine removed prototype rows are not recreated by downgrade, so retain a database backup if their contents may be needed later. + +### Local-only tracer cleanup + +Revision `17bb32415f97` was autogenerated after moving the local SQLite +`tracers` cache table out of production `V1Base` metadata and aligning +`reform_impact.execution_id` with its required application contract. Its only +post-generation correction is `if_exists=True` on the generated tracer drop: +fresh databases built from the original baseline contain the table, while +deployed MySQL databases do not. Integration tests cover both fresh and +production-shaped schemas and verify that `execution_id` becomes non-nullable. diff --git a/migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py b/migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py new file mode 100644 index 000000000..f75a69049 --- /dev/null +++ b/migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py @@ -0,0 +1,54 @@ +"""separate local tracer and require execution id + +Revision ID: 17bb32415f97 +Revises: 01e49b3a056e +Create Date: 2026-08-11 18:17:19.903942 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "17bb32415f97" +down_revision: Union[str, None] = "01e49b3a056e" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + # Production never had this local-only table; fresh baseline schemas do. + op.drop_table("tracers", if_exists=True) + op.alter_column( + "reform_impact", + "execution_id", + existing_type=mysql.VARCHAR(length=255), + nullable=False, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "execution_id", + existing_type=mysql.VARCHAR(length=255), + nullable=True, + ) + op.create_table( + "tracers", + sa.Column("id", mysql.INTEGER(), autoincrement=True, nullable=False), + sa.Column("household_id", mysql.INTEGER(), autoincrement=False, nullable=False), + sa.Column("policy_id", mysql.INTEGER(), autoincrement=False, nullable=False), + sa.Column("country_id", mysql.VARCHAR(length=3), nullable=False), + sa.Column("api_version", mysql.VARCHAR(length=10), nullable=False), + sa.Column("tracer_output", mysql.JSON(), nullable=False), + sa.PrimaryKeyConstraint("id"), + mysql_collate="utf8mb4_0900_ai_ci", + mysql_default_charset="utf8mb4", + mysql_engine="InnoDB", + ) + # ### end Alembic commands ### diff --git a/migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py b/migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py new file mode 100644 index 000000000..49f1e3f98 --- /dev/null +++ b/migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py @@ -0,0 +1,41 @@ +"""require explicit reform impact dataset + +Revision ID: 1914c0422236 +Revises: 17bb32415f97 +Create Date: 2026-08-11 19:01:40.778647 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "1914c0422236" +down_revision: Union[str, None] = "17bb32415f97" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "dataset", + existing_type=mysql.VARCHAR(length=255), + server_default=None, + existing_nullable=False, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "dataset", + existing_type=mysql.VARCHAR(length=255), + server_default=sa.text("'default'"), + existing_nullable=False, + ) + # ### end Alembic commands ### diff --git a/policyengine_api/data/local_database.py b/policyengine_api/data/local_database.py index 05e59ee74..2d29c0210 100644 --- a/policyengine_api/data/local_database.py +++ b/policyengine_api/data/local_database.py @@ -12,6 +12,7 @@ from sqlalchemy import Column, Engine, Integer, JSON, MetaData, String, Table +from policyengine_api.data.local_models import LocalV1Base from policyengine_api.data.v1_models import Policy, V1Base @@ -40,4 +41,5 @@ def create_local_v1_schema(engine: Engine) -> None: if table is not Policy.__table__ ] V1Base.metadata.create_all(engine, tables=production_tables) + LocalV1Base.metadata.create_all(engine) _sqlite_policy_metadata.create_all(engine) diff --git a/policyengine_api/data/local_models.py b/policyengine_api/data/local_models.py new file mode 100644 index 000000000..35ca26074 --- /dev/null +++ b/policyengine_api/data/local_models.py @@ -0,0 +1,26 @@ +"""Declarative mappings for the temporary local SQLite cache only. + +These tables are deliberately outside the production API v1 metadata and +Alembic lifecycle. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import Integer, JSON, String +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + +class LocalV1Base(DeclarativeBase): + pass + + +class Tracer(LocalV1Base): + __tablename__ = "tracers" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + household_id: Mapped[int] + policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + api_version: Mapped[str] = mapped_column(String(10)) + tracer_output: Mapped[Any] = mapped_column(JSON) diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 2529021d5..6e848d106 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -138,16 +138,6 @@ class UserProfile(V1Base): user_since: Mapped[int] = mapped_column(BigInteger) -class Tracer(V1Base): - __tablename__ = "tracers" - id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) - household_id: Mapped[int] - policy_id: Mapped[int] - country_id: Mapped[str] = mapped_column(String(3)) - api_version: Mapped[str] = mapped_column(String(10)) - tracer_output: Mapped[Any] = mapped_column(JSON) - - class Simulation(V1Base): __tablename__ = "simulations" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 0e9d7d959..0179dc2f9 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -9,12 +9,12 @@ from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.local_models import Tracer from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ( ComputedHousehold, Household, Policy, - Tracer, ) from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs from policyengine_api.utils.input_validation import find_unrecognized_inputs diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index ff6573d89..575cb6a4f 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -6,7 +6,7 @@ from werkzeug.exceptions import NotFound from sqlalchemy import select -from policyengine_api.data.v1_models import Tracer +from policyengine_api.data.local_models import Tracer class TracerAnalysisService(AIAnalysisService): diff --git a/tests/fixtures/services/tracer_fixture_service.py b/tests/fixtures/services/tracer_fixture_service.py index ef34553cf..2d3d015d0 100644 --- a/tests/fixtures/services/tracer_fixture_service.py +++ b/tests/fixtures/services/tracer_fixture_service.py @@ -1,6 +1,6 @@ import pytest import json -from policyengine_api.data.v1_models import Tracer +from policyengine_api.data.local_models import Tracer valid_tracer = { "tracer_output": [ diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index abfa9a4be..cff70731f 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -6,6 +6,7 @@ from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext +from alembic.operations import Operations import pytest from sqlalchemy import ( Column, @@ -16,6 +17,7 @@ Text, create_engine, inspect, + text, ) from sqlalchemy.dialects.mysql import LONGTEXT from sqlalchemy.engine import make_url @@ -84,6 +86,14 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): assert context.get_current_revision() is not None assert compare_metadata(context, V1Base.metadata) == [] + inspector = inspect(engine) + assert "tracers" not in inspector.get_table_names() + reform_impact_columns = { + column["name"]: column for column in inspector.get_columns("reform_impact") + } + assert reform_impact_columns["dataset"]["default"] is None + assert reform_impact_columns["execution_id"]["nullable"] is False + command.downgrade(config, BASELINE_REVISION) assert "question" in inspect(engine).get_table_names() @@ -108,6 +118,21 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema() command.upgrade(config, BASELINE_REVISION) question.create(engine) with engine.begin() as connection: + operations = Operations(MigrationContext.configure(connection)) + operations.drop_table("tracers") + operations.alter_column( + "reform_impact", + "execution_id", + existing_type=String(255), + nullable=True, + ) + operations.alter_column( + "reform_impact", + "dataset", + existing_type=String(255), + existing_nullable=False, + server_default=text("'default'"), + ) connection.execute( question.insert(), { @@ -120,7 +145,14 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema() command.upgrade(config, "head") - assert "question" not in inspect(engine).get_table_names() + inspector = inspect(engine) + assert "question" not in inspector.get_table_names() + assert "tracers" not in inspector.get_table_names() + reform_impact_columns = { + column["name"]: column for column in inspector.get_columns("reform_impact") + } + assert reform_impact_columns["dataset"]["default"] is None + assert reform_impact_columns["execution_id"]["nullable"] is False command.downgrade(config, BASELINE_REVISION) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 8da464ea5..b703a0b76 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -9,6 +9,7 @@ from policyengine_api.data import orm from policyengine_api.data.local_database import create_local_v1_schema +from policyengine_api.data.local_models import LocalV1Base from policyengine_api.data.v1_models import V1Base @@ -32,7 +33,9 @@ def isolated_orm_database(test_engine, monkeypatch): orm.clear_v1_session_factories() factory = orm.get_v1_session_factory() with factory.begin() as session: - for table in reversed(V1Base.metadata.sorted_tables): + local_tables = LocalV1Base.metadata.sorted_tables + production_tables = V1Base.metadata.sorted_tables + for table in reversed([*production_tables, *local_tables]): session.execute(table.delete()) try: yield diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 015e79663..89d9689e2 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -90,6 +90,7 @@ def test_local_schema_preserves_the_documented_sqlite_policy_key_exception(): policy_key = inspect(engine).get_pk_constraint("policy") assert policy_key["constrained_columns"] == ["id"] + assert "tracers" in inspect(engine).get_table_names() finally: engine.dispose() diff --git a/tests/unit/data/test_v1_models.py b/tests/unit/data/test_v1_models.py index 02b738531..24155d04d 100644 --- a/tests/unit/data/test_v1_models.py +++ b/tests/unit/data/test_v1_models.py @@ -1,6 +1,7 @@ from sqlalchemy.dialects import mysql from sqlalchemy.schema import CreateTable +from policyengine_api.data.local_models import LocalV1Base from policyengine_api.data.v1_models import V1Base @@ -16,7 +17,6 @@ "report_outputs", "simulation_runs", "simulations", - "tracers", "user_policies", "user_profiles", } @@ -26,6 +26,11 @@ def test_v1_metadata_contains_every_legacy_table(): assert set(V1Base.metadata.tables) == EXPECTED_TABLES +def test_tracer_metadata_is_local_only(): + assert "tracers" not in V1Base.metadata.tables + assert set(LocalV1Base.metadata.tables) == {"tracers"} + + def test_v1_metadata_compiles_for_the_production_mysql_dialect(): for table in V1Base.metadata.sorted_tables: assert str(CreateTable(table).compile(dialect=mysql.dialect())) diff --git a/tests/unit/services/test_direct_orm_local_analysis.py b/tests/unit/services/test_direct_orm_local_analysis.py index f60af4e31..c6a4b0ef6 100644 --- a/tests/unit/services/test_direct_orm_local_analysis.py +++ b/tests/unit/services/test_direct_orm_local_analysis.py @@ -1,6 +1,7 @@ from datetime import datetime -from policyengine_api.data.v1_models import Analysis, ReformImpact, Tracer +from policyengine_api.data.local_models import Tracer +from policyengine_api.data.v1_models import Analysis, ReformImpact from policyengine_api.services.ai_analysis_service import AIAnalysisService from policyengine_api.services.reform_impacts_service import ReformImpactsService from policyengine_api.services.tracer_analysis_service import TracerAnalysisService diff --git a/tests/unit/services/test_household_calculation_service.py b/tests/unit/services/test_household_calculation_service.py index c24bed5a1..e849861b8 100644 --- a/tests/unit/services/test_household_calculation_service.py +++ b/tests/unit/services/test_household_calculation_service.py @@ -6,11 +6,11 @@ from sqlalchemy import select from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.local_models import Tracer from policyengine_api.data.v1_models import ( ComputedHousehold, Household, Policy, - Tracer, ) from policyengine_api.services.household_calculation_service import ( HouseholdCalculationService, From 74addbfd29859960e81fa82e2fc9429a22aef540 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:43:44 +0300 Subject: [PATCH 79/89] feat: add safe v1 database migration orchestration --- alembic.ini => alembic-v1.ini | 2 +- migrations/{ => v1}/env.py | 25 +- migrations/{ => v1}/script.py.mako | 0 migrations/v1/versions/.gitkeep | 1 + ...9b3a056e_remove_orphaned_question_table.py | 0 ...5f97_separate_local_tracer_and_require_.py | 0 ..._require_explicit_reform_impact_dataset.py | 0 ...afc2a547a4e_baseline_existing_v1_schema.py | 0 migrations/versions/.gitkeep | 0 scripts/v1_alembic_changes.py | 64 ++++ scripts/v1_database_migration.py | 293 ++++++++++++++++++ scripts/write_v1_database_urls.py | 46 +++ .../test_alembic_mysql_lifecycle.py | 15 +- .../test_v1_schema_metadata_compatibility.py | 12 +- tests/unit/data/test_alembic_baseline.py | 6 +- tests/unit/data/test_v1_database_migration.py | 158 ++++++++++ 16 files changed, 599 insertions(+), 23 deletions(-) rename alembic.ini => alembic-v1.ini (94%) rename migrations/{ => v1}/env.py (74%) rename migrations/{ => v1}/script.py.mako (100%) create mode 100644 migrations/v1/versions/.gitkeep rename migrations/{ => v1}/versions/01e49b3a056e_remove_orphaned_question_table.py (100%) rename migrations/{ => v1}/versions/17bb32415f97_separate_local_tracer_and_require_.py (100%) rename migrations/{ => v1}/versions/1914c0422236_require_explicit_reform_impact_dataset.py (100%) rename migrations/{ => v1}/versions/eafc2a547a4e_baseline_existing_v1_schema.py (100%) delete mode 100644 migrations/versions/.gitkeep create mode 100644 scripts/v1_alembic_changes.py create mode 100644 scripts/v1_database_migration.py create mode 100644 scripts/write_v1_database_urls.py create mode 100644 tests/unit/data/test_v1_database_migration.py diff --git a/alembic.ini b/alembic-v1.ini similarity index 94% rename from alembic.ini rename to alembic-v1.ini index 115fa521e..3b9f297d1 100644 --- a/alembic.ini +++ b/alembic-v1.ini @@ -1,5 +1,5 @@ [alembic] -script_location = migrations +script_location = migrations/v1 prepend_sys_path = . path_separator = os diff --git a/migrations/env.py b/migrations/v1/env.py similarity index 74% rename from migrations/env.py rename to migrations/v1/env.py index a7eb19f11..e9d18e768 100644 --- a/migrations/env.py +++ b/migrations/v1/env.py @@ -42,20 +42,29 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: + provided_connection = config.attributes.get("connection") + if provided_connection is not None: + _run_migrations_with_connection(provided_connection) + return + connectable = engine_from_config( config.get_section(config.config_ini_section, {}), prefix="sqlalchemy.", poolclass=pool.NullPool, ) with connectable.connect() as connection: - context.configure( - connection=connection, - target_metadata=target_metadata, - compare_type=True, - compare_server_default=True, - ) - with context.begin_transaction(): - context.run_migrations() + _run_migrations_with_connection(connection) + + +def _run_migrations_with_connection(connection) -> None: + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + compare_server_default=True, + ) + with context.begin_transaction(): + context.run_migrations() if context.is_offline_mode(): diff --git a/migrations/script.py.mako b/migrations/v1/script.py.mako similarity index 100% rename from migrations/script.py.mako rename to migrations/v1/script.py.mako diff --git a/migrations/v1/versions/.gitkeep b/migrations/v1/versions/.gitkeep new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/migrations/v1/versions/.gitkeep @@ -0,0 +1 @@ + diff --git a/migrations/versions/01e49b3a056e_remove_orphaned_question_table.py b/migrations/v1/versions/01e49b3a056e_remove_orphaned_question_table.py similarity index 100% rename from migrations/versions/01e49b3a056e_remove_orphaned_question_table.py rename to migrations/v1/versions/01e49b3a056e_remove_orphaned_question_table.py diff --git a/migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py b/migrations/v1/versions/17bb32415f97_separate_local_tracer_and_require_.py similarity index 100% rename from migrations/versions/17bb32415f97_separate_local_tracer_and_require_.py rename to migrations/v1/versions/17bb32415f97_separate_local_tracer_and_require_.py diff --git a/migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py b/migrations/v1/versions/1914c0422236_require_explicit_reform_impact_dataset.py similarity index 100% rename from migrations/versions/1914c0422236_require_explicit_reform_impact_dataset.py rename to migrations/v1/versions/1914c0422236_require_explicit_reform_impact_dataset.py diff --git a/migrations/versions/eafc2a547a4e_baseline_existing_v1_schema.py b/migrations/v1/versions/eafc2a547a4e_baseline_existing_v1_schema.py similarity index 100% rename from migrations/versions/eafc2a547a4e_baseline_existing_v1_schema.py rename to migrations/v1/versions/eafc2a547a4e_baseline_existing_v1_schema.py diff --git a/migrations/versions/.gitkeep b/migrations/versions/.gitkeep deleted file mode 100644 index e69de29bb..000000000 diff --git a/scripts/v1_alembic_changes.py b/scripts/v1_alembic_changes.py new file mode 100644 index 000000000..335a46b63 --- /dev/null +++ b/scripts/v1_alembic_changes.py @@ -0,0 +1,64 @@ +"""Report whether a pull request changes the API v1 Alembic surface.""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import PurePosixPath + + +EXACT_PATHS = frozenset( + { + "alembic-v1.ini", + ".github/workflows/alembic-v1-check.yml", + "docs/engineering/skills/alembic-migrations.md", + "policyengine_api/data/v1_models.py", + "pyproject.toml", + "scripts/v1_alembic_changes.py", + "scripts/v1_database_migration.py", + "scripts/write_v1_database_urls.py", + "tests/integration/test_alembic_mysql_lifecycle.py", + "tests/integration/test_v1_schema_metadata_compatibility.py", + "tests/unit/test_alembic_workflows.py", + "uv.lock", + } +) +PATH_PREFIXES = ( + "migrations/v1/", + "tests/unit/data/test_alembic_", + "tests/unit/data/test_v1_database_migration", +) + + +def is_v1_alembic_path(path: str) -> bool: + """Return whether *path* can change v1 migration behavior.""" + + normalized = PurePosixPath(path).as_posix().removeprefix("./") + return normalized in EXACT_PATHS or normalized.startswith(PATH_PREFIXES) + + +def changed_paths(base: str, head: str) -> tuple[str, ...]: + """Return repository paths changed between two git revisions.""" + + result = subprocess.run( + ["git", "diff", "--name-only", base, head, "--"], + check=True, + capture_output=True, + text=True, + ) + return tuple(path for path in result.stdout.splitlines() if path) + + +def main(argv: list[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if len(args) != 2: + print("usage: v1_alembic_changes.py BASE HEAD", file=sys.stderr) + return 2 + + changed = any(is_v1_alembic_path(path) for path in changed_paths(*args)) + print(f"changed={str(changed).lower()}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py new file mode 100644 index 000000000..415e77fd6 --- /dev/null +++ b/scripts/v1_database_migration.py @@ -0,0 +1,293 @@ +"""Safely inspect, adopt, and upgrade the API v1 Cloud SQL schema.""" + +from __future__ import annotations + +import argparse +from enum import StrEnum +import os +from typing import Any +from urllib.parse import quote_plus + +from alembic import command +from alembic.autogenerate import compare_metadata +from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.script import ScriptDirectory +from sqlalchemy import Connection, create_engine, inspect, text +from sqlalchemy.pool import NullPool + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base + + +BASELINE_REVISION = "eafc2a547a4e" +ADOPTION_CONFIRMATION = f"ADOPT-{BASELINE_REVISION}" +ALEMBIC_CONFIG = REPO / "alembic-v1.ini" +MIGRATION_LOCK_NAME = "policyengine-api-v1-alembic" +EXPECTED_LEGACY_DIFFERENCES = frozenset( + { + "extra_table:question", + "nullable:reform_impact.execution_id:true->false", + "default:reform_impact.dataset:present->none", + } +) + + +class DatabaseState(StrEnum): + UNVERSIONED = "unversioned" + INVALID = "invalid" + PENDING = "pending" + HEAD = "head" + + +def build_database_url( + *, + username: str, + password: str, + host: str, + port: int, + database: str, +) -> str: + """Build an encoded PyMySQL URL without logging its credentials.""" + + return ( + f"mysql+pymysql://{quote_plus(username)}:{quote_plus(password)}@" + f"{host}:{port}/{quote_plus(database)}" + ) + + +def classify_database_state( + *, + version_table_exists: bool, + current_heads: set[str], + script_heads: set[str] | None = None, +) -> DatabaseState: + if not version_table_exists: + return DatabaseState.UNVERSIONED + if not current_heads: + return DatabaseState.INVALID + if script_heads is not None and current_heads == script_heads: + return DatabaseState.HEAD + return DatabaseState.PENDING + + +def require_adoption_confirmation(confirmation: str) -> None: + if confirmation != ADOPTION_CONFIRMATION: + raise ValueError( + "Explicit adoption confirmation is required; expected " + f"{ADOPTION_CONFIRMATION!r}" + ) + + +def describe_metadata_difference(difference: Any) -> str: + """Return a stable, data-free description of an Alembic metadata diff.""" + + item = difference + if isinstance(item, list) and len(item) == 1: + item = item[0] + if not isinstance(item, tuple) or not item: + raise ValueError("Unsupported metadata difference shape") + + operation = item[0] + if operation == "remove_table" and len(item) >= 2: + return f"extra_table:{item[1].name}" + if operation == "modify_nullable" and len(item) >= 7: + return ( + f"nullable:{item[2]}.{item[3]}:" + f"{str(item[5]).lower()}->{str(item[6]).lower()}" + ) + if operation == "modify_default" and len(item) >= 7: + existing = "none" if item[5] is None else "present" + target = "none" if item[6] is None else "present" + return f"default:{item[2]}.{item[3]}:{existing}->{target}" + raise ValueError(f"Unsupported metadata difference operation: {operation}") + + +def _config(connection: Connection) -> Config: + config = Config(str(ALEMBIC_CONFIG)) + config.attributes["connection"] = connection + return config + + +def _script_heads(config: Config) -> set[str]: + return set(ScriptDirectory.from_config(config).get_heads()) + + +def database_state(connection: Connection) -> DatabaseState: + inspector = inspect(connection) + version_table_exists = "alembic_version" in inspector.get_table_names() + context = MigrationContext.configure(connection) + config = _config(connection) + return classify_database_state( + version_table_exists=version_table_exists, + current_heads=set(context.get_current_heads()), + script_heads=_script_heads(config), + ) + + +def metadata_differences(connection: Connection) -> set[str]: + context = MigrationContext.configure( + connection, + opts={"compare_type": True, "compare_server_default": True}, + ) + return { + describe_metadata_difference(difference) + for difference in compare_metadata(context, V1Base.metadata) + } + + +def verify_legacy_schema( + connection: Connection, *, expected_question_rows: int +) -> None: + state = database_state(connection) + if state is not DatabaseState.UNVERSIONED: + raise RuntimeError( + f"Existing database must be unversioned before adoption; found {state}" + ) + + differences = metadata_differences(connection) + if differences != EXPECTED_LEGACY_DIFFERENCES: + unexpected = sorted(differences - EXPECTED_LEGACY_DIFFERENCES) + missing = sorted(EXPECTED_LEGACY_DIFFERENCES - differences) + raise RuntimeError( + "Existing schema does not match the reviewed legacy shape; " + f"unexpected={unexpected}, missing={missing}" + ) + + question_rows = connection.scalar(text("SELECT COUNT(*) FROM question")) + if question_rows != expected_question_rows: + raise RuntimeError( + "question row count changed; expected " + f"{expected_question_rows}, found {question_rows}" + ) + + for column in ("execution_id", "dataset"): + null_rows = connection.scalar( + text(f"SELECT COUNT(*) FROM reform_impact WHERE {column} IS NULL") + ) + if null_rows: + raise RuntimeError(f"reform_impact.{column} contains {null_rows} NULL rows") + + +def verify_head_schema(connection: Connection) -> None: + state = database_state(connection) + if state is not DatabaseState.HEAD: + raise RuntimeError(f"Database is not at the Alembic head; found {state}") + differences = metadata_differences(connection) + if differences: + raise RuntimeError( + f"Database metadata drift remains after migration: {sorted(differences)}" + ) + + +def _acquire_lock(connection: Connection) -> None: + acquired = connection.scalar( + text("SELECT GET_LOCK(:name, 60)"), {"name": MIGRATION_LOCK_NAME} + ) + if acquired != 1: + raise RuntimeError("Could not acquire the v1 Alembic migration lock") + + +def _release_lock(connection: Connection) -> None: + connection.execute( + text("SELECT RELEASE_LOCK(:name)"), {"name": MIGRATION_LOCK_NAME} + ) + + +def adopt_database( + connection: Connection, + *, + confirmation: str, + backup_id: str, + expected_question_rows: int, +) -> None: + require_adoption_confirmation(confirmation) + if not backup_id.strip(): + raise ValueError("A completed Cloud SQL backup ID is required") + + _acquire_lock(connection) + try: + verify_legacy_schema(connection, expected_question_rows=expected_question_rows) + config = _config(connection) + command.stamp(config, BASELINE_REVISION) + command.upgrade(config, "head") + verify_head_schema(connection) + finally: + _release_lock(connection) + + +def upgrade_database(connection: Connection, *, backup_id: str) -> None: + _acquire_lock(connection) + try: + state = database_state(connection) + if state is DatabaseState.UNVERSIONED: + raise RuntimeError( + "database is unversioned; run the explicit adoption workflow first" + ) + if state is DatabaseState.HEAD: + verify_head_schema(connection) + return + if not backup_id.strip(): + raise ValueError("A completed Cloud SQL backup ID is required") + + command.upgrade(_config(connection), "head") + verify_head_schema(connection) + finally: + _release_lock(connection) + + +def _database_url(mode: str) -> str: + env_name = ( + "STAGE7_EXISTING_DATABASE_URL" + if mode in {"verify-legacy", "verify-head", "state"} + else "ALEMBIC_DATABASE_URL" + ) + try: + return os.environ[env_name] + except KeyError as error: + raise RuntimeError(f"{env_name} is required") from error + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--mode", + required=True, + choices=("state", "verify-legacy", "verify-head", "adopt", "upgrade"), + ) + parser.add_argument("--confirmation", default="") + parser.add_argument("--backup-id", default="") + parser.add_argument("--expected-question-rows", type=int, default=9) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + engine = create_engine(_database_url(args.mode), poolclass=NullPool) + try: + with engine.connect() as connection: + if args.mode == "state": + print(database_state(connection).value) + elif args.mode == "verify-legacy": + verify_legacy_schema( + connection, + expected_question_rows=args.expected_question_rows, + ) + elif args.mode == "verify-head": + verify_head_schema(connection) + elif args.mode == "adopt": + adopt_database( + connection, + confirmation=args.confirmation, + backup_id=args.backup_id, + expected_question_rows=args.expected_question_rows, + ) + else: + upgrade_database(connection, backup_id=args.backup_id) + finally: + engine.dispose() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/write_v1_database_urls.py b/scripts/write_v1_database_urls.py new file mode 100644 index 000000000..752b06317 --- /dev/null +++ b/scripts/write_v1_database_urls.py @@ -0,0 +1,46 @@ +"""Write masked local-proxy database URLs into a GitHub Actions env file.""" + +from __future__ import annotations + +import os +from pathlib import Path + +from scripts.v1_database_migration import build_database_url + + +DATABASE_NAME = "policyengine" +READONLY_USER = "policyengine_schema_reader" +MIGRATION_USER = "policyengine_schema_migrator" +PROXY_HOST = "127.0.0.1" +PROXY_PORT = 3307 + + +def main() -> int: + output_path = Path(os.environ["GITHUB_ENV"]) + readonly_url = build_database_url( + username=READONLY_USER, + password=os.environ["POLICYENGINE_DB_READONLY_PASSWORD"], + host=PROXY_HOST, + port=PROXY_PORT, + database=DATABASE_NAME, + ) + migration_url = build_database_url( + username=MIGRATION_USER, + password=os.environ["POLICYENGINE_DB_MIGRATION_PASSWORD"], + host=PROXY_HOST, + port=PROXY_PORT, + database=DATABASE_NAME, + ) + + # URL-encoded passwords can differ from the exact GitHub secret value, so + # mask the complete derived URLs before any later command can mention them. + print(f"::add-mask::{readonly_url}") + print(f"::add-mask::{migration_url}") + with output_path.open("a", encoding="utf-8") as output: + output.write(f"STAGE7_EXISTING_DATABASE_URL={readonly_url}\n") + output.write(f"ALEMBIC_DATABASE_URL={migration_url}\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index cff70731f..857643c06 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -24,6 +24,10 @@ from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base +from scripts.v1_database_migration import ( + ADOPTION_CONFIRMATION, + adopt_database, +) BASELINE_REVISION = "eafc2a547a4e" @@ -64,7 +68,7 @@ def _ephemeral_mysql_url() -> str: def _alembic_config(database_url: str) -> Config: - config = Config(str(REPO / "alembic.ini")) + config = Config(str(REPO / "alembic-v1.ini")) config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) return config @@ -142,8 +146,15 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema() "status": "ok", }, ) + operations.drop_table("alembic_version") - command.upgrade(config, "head") + with engine.connect() as connection: + adopt_database( + connection, + confirmation=ADOPTION_CONFIRMATION, + backup_id="test-backup", + expected_question_rows=1, + ) inspector = inspect(engine) assert "question" not in inspector.get_table_names() diff --git a/tests/integration/test_v1_schema_metadata_compatibility.py b/tests/integration/test_v1_schema_metadata_compatibility.py index d5c6f26db..ed16a0247 100644 --- a/tests/integration/test_v1_schema_metadata_compatibility.py +++ b/tests/integration/test_v1_schema_metadata_compatibility.py @@ -2,25 +2,19 @@ import os -from alembic.autogenerate import compare_metadata -from alembic.migration import MigrationContext import pytest from sqlalchemy import create_engine -from policyengine_api.data.v1_models import V1Base +from scripts.v1_database_migration import metadata_differences -def compare_existing_schema(database_url: str) -> list: +def compare_existing_schema(database_url: str) -> list[str]: """Return metadata drift without mutating the target database.""" engine = create_engine(database_url) try: with engine.connect() as connection: - context = MigrationContext.configure( - connection, - opts={"compare_type": True, "compare_server_default": True}, - ) - return compare_metadata(context, V1Base.metadata) + return sorted(metadata_differences(connection)) finally: engine.dispose() diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index 12c9530b3..a6bf8aeb5 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -11,7 +11,7 @@ def _mysql_offline_config() -> tuple[Config, StringIO]: output = StringIO() - config = Config(str(REPO / "alembic.ini"), output_buffer=output) + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=output) config.set_main_option( "sqlalchemy.url", "mysql+pymysql://offline:offline@localhost/offline", @@ -21,7 +21,7 @@ def _mysql_offline_config() -> tuple[Config, StringIO]: def test_default_configuration_requires_an_explicit_database_url(monkeypatch): monkeypatch.delenv("ALEMBIC_DATABASE_URL", raising=False) - config = Config(str(REPO / "alembic.ini"), output_buffer=StringIO()) + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=StringIO()) with pytest.raises(RuntimeError, match="ALEMBIC_DATABASE_URL"): command.upgrade(config, "head", sql=True) @@ -29,7 +29,7 @@ def test_default_configuration_requires_an_explicit_database_url(monkeypatch): def test_sqlite_is_rejected_as_a_migration_target(monkeypatch): monkeypatch.setenv("ALEMBIC_DATABASE_URL", "sqlite+pysqlite:///:memory:") - config = Config(str(REPO / "alembic.ini"), output_buffer=StringIO()) + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=StringIO()) with pytest.raises(RuntimeError, match="MySQL"): command.upgrade(config, "head", sql=True) diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py new file mode 100644 index 000000000..de52c6fe3 --- /dev/null +++ b/tests/unit/data/test_v1_database_migration.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import pytest +from sqlalchemy import Column, Integer, MetaData, String, Table, text + +from scripts.v1_alembic_changes import is_v1_alembic_path +from scripts.v1_database_migration import ( + ADOPTION_CONFIRMATION, + BASELINE_REVISION, + DatabaseState, + build_database_url, + classify_database_state, + describe_metadata_difference, + require_adoption_confirmation, +) + + +@pytest.mark.parametrize( + "path", + [ + "alembic-v1.ini", + "migrations/v1/env.py", + "migrations/v1/versions/123_add_column.py", + "policyengine_api/data/v1_models.py", + "scripts/v1_database_migration.py", + "tests/integration/test_alembic_mysql_lifecycle.py", + "tests/integration/test_v1_schema_metadata_compatibility.py", + ".github/workflows/alembic-v1-check.yml", + "docs/engineering/skills/alembic-migrations.md", + "pyproject.toml", + "uv.lock", + ], +) +def test_v1_alembic_change_paths_trigger_qualification(path): + assert is_v1_alembic_path(path) + + +@pytest.mark.parametrize( + "path", + [ + "policyengine_api/routes/household_routes.py", + "tests/unit/routes/test_household_routes.py", + "docs/migration/cloud-run-operations.md", + ".github/workflows/push.yml", + ], +) +def test_unrelated_paths_do_not_trigger_v1_alembic_qualification(path): + assert not is_v1_alembic_path(path) + + +def test_database_url_percent_encodes_credentials_without_losing_driver(): + url = build_database_url( + username="schema reader", + password="p@ss:/word", + host="127.0.0.1", + port=3307, + database="policyengine", + ) + + assert url == ( + "mysql+pymysql://schema+reader:p%40ss%3A%2Fword@127.0.0.1:3307/policyengine" + ) + + +def test_database_state_is_unversioned_without_a_version_table(): + assert ( + classify_database_state(version_table_exists=False, current_heads=set()) + is DatabaseState.UNVERSIONED + ) + + +def test_database_state_is_invalid_when_version_table_has_no_revision(): + assert ( + classify_database_state(version_table_exists=True, current_heads=set()) + is DatabaseState.INVALID + ) + + +def test_database_state_is_head_only_when_all_script_heads_are_applied(): + assert ( + classify_database_state( + version_table_exists=True, + current_heads={"head-a"}, + script_heads={"head-a"}, + ) + is DatabaseState.HEAD + ) + assert ( + classify_database_state( + version_table_exists=True, + current_heads={BASELINE_REVISION}, + script_heads={"head-a"}, + ) + is DatabaseState.PENDING + ) + + +def test_adoption_requires_the_exact_explicit_confirmation(): + require_adoption_confirmation(ADOPTION_CONFIRMATION) + + with pytest.raises(ValueError, match="confirmation"): + require_adoption_confirmation("yes") + + +def test_metadata_difference_descriptions_are_stable_and_do_not_include_data(): + metadata = MetaData() + question = Table( + "question", + metadata, + Column("question_id", Integer, primary_key=True), + Column("question", String(255)), + ) + + assert describe_metadata_difference(("remove_table", question)) == ( + "extra_table:question" + ) + assert ( + describe_metadata_difference( + [ + ( + "modify_nullable", + None, + "reform_impact", + "execution_id", + {}, + True, + False, + ) + ] + ) + == "nullable:reform_impact.execution_id:true->false" + ) + assert ( + describe_metadata_difference( + [ + ( + "modify_default", + None, + "reform_impact", + "dataset", + {}, + text("'default'"), + None, + ) + ] + ) + == "default:reform_impact.dataset:present->none" + ) + + +def test_metadata_difference_rejects_unknown_shapes_without_repr_leakage(): + secret = "do-not-print-this-value" + difference = ("unknown", secret) + + with pytest.raises(ValueError, match="Unsupported metadata difference") as error: + describe_metadata_difference(difference) + + assert secret not in str(error.value) From bbd45403feb3ca0243516266698e6bd325a3e372 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:43:50 +0300 Subject: [PATCH 80/89] ci: gate releases on qualified v1 migrations --- .github/scripts/create_cloud_sql_backup.sh | 23 ++++ .github/scripts/start_cloud_sql_proxy.sh | 41 +++++++ .github/scripts/stop_cloud_sql_proxy.sh | 14 +++ .github/workflows/adopt-v1-cloud-sql.yml | 76 +++++++++++++ .github/workflows/alembic-v1-check.yml | 47 ++++++++ .github/workflows/pr.yml | 42 +++---- .github/workflows/push.yml | 85 +++++++++++++- changelog.d/3788.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 72 ++++++++---- tests/unit/test_alembic_workflows.py | 107 ++++++++++++++++++ 10 files changed, 458 insertions(+), 51 deletions(-) create mode 100644 .github/scripts/create_cloud_sql_backup.sh create mode 100644 .github/scripts/start_cloud_sql_proxy.sh create mode 100644 .github/scripts/stop_cloud_sql_proxy.sh create mode 100644 .github/workflows/adopt-v1-cloud-sql.yml create mode 100644 .github/workflows/alembic-v1-check.yml create mode 100644 tests/unit/test_alembic_workflows.py diff --git a/.github/scripts/create_cloud_sql_backup.sh b/.github/scripts/create_cloud_sql_backup.sh new file mode 100644 index 000000000..5f43e443b --- /dev/null +++ b/.github/scripts/create_cloud_sql_backup.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" +: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" + +instance_id="${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME##*:}" +description="policyengine-api-v1-alembic-${GITHUB_SHA:-manual}" +backup_id="$( + gcloud sql backups create \ + --project policyengine-api \ + --instance "${instance_id}" \ + --description "${description}" \ + --format 'value(id)' +)" + +if [[ -z "${backup_id}" ]]; then + echo "Cloud SQL did not return a completed backup ID." >&2 + exit 1 +fi + +printf 'backup_id=%s\n' "${backup_id}" >>"${GITHUB_OUTPUT}" diff --git a/.github/scripts/start_cloud_sql_proxy.sh b/.github/scripts/start_cloud_sql_proxy.sh new file mode 100644 index 000000000..38bf8d64c --- /dev/null +++ b/.github/scripts/start_cloud_sql_proxy.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" + +proxy_version="2.25.0" +proxy_sha256="091a9a12eddab6c028b6c563a4f2dacd067e8f7689c25a3fb4afce397e1f0c60" +proxy_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy" +pid_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.pid" +log_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.log" + +curl -fsSL \ + "https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v${proxy_version}/cloud-sql-proxy.linux.amd64" \ + --output "${proxy_path}" +printf '%s %s\n' "${proxy_sha256}" "${proxy_path}" | sha256sum --check --status +chmod +x "${proxy_path}" + +"${proxy_path}" \ + --quota-project policyengine-api \ + --address 127.0.0.1 \ + --port 3307 \ + "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" \ + >"${log_path}" 2>&1 & +proxy_pid="$!" +printf '%s\n' "${proxy_pid}" >"${pid_path}" + +for _ in $(seq 1 30); do + if ! kill -0 "${proxy_pid}" 2>/dev/null; then + echo "Cloud SQL Auth Proxy exited before becoming ready." >&2 + sed -n '1,120p' "${log_path}" >&2 + exit 1 + fi + if python -c 'import socket; socket.create_connection(("127.0.0.1", 3307), 1).close()' 2>/dev/null; then + exit 0 + fi + sleep 1 +done + +echo "Cloud SQL Auth Proxy did not become ready." >&2 +exit 1 diff --git a/.github/scripts/stop_cloud_sql_proxy.sh b/.github/scripts/stop_cloud_sql_proxy.sh new file mode 100644 index 000000000..89f24b26c --- /dev/null +++ b/.github/scripts/stop_cloud_sql_proxy.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash + +set -euo pipefail + +pid_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.pid" +if [[ ! -f "${pid_path}" ]]; then + exit 0 +fi + +proxy_pid="$(cat "${pid_path}")" +if kill -0 "${proxy_pid}" 2>/dev/null; then + kill "${proxy_pid}" + wait "${proxy_pid}" 2>/dev/null || true +fi diff --git a/.github/workflows/adopt-v1-cloud-sql.yml b/.github/workflows/adopt-v1-cloud-sql.yml new file mode 100644 index 000000000..e0d816064 --- /dev/null +++ b/.github/workflows/adopt-v1-cloud-sql.yml @@ -0,0 +1,76 @@ +name: Adopt existing v1 Cloud SQL schema + +on: + workflow_dispatch: + inputs: + confirmation: + description: Type ADOPT-eafc2a547a4e to stamp and upgrade the existing database + required: true + type: string + expected_question_rows: + description: Expected number of orphaned question rows that the migration removes + required: true + default: "9" + type: string + +concurrency: + group: policyengine-api-v1-cloud-sql-schema + cancel-in-progress: false + +jobs: + adopt: + name: Stamp and upgrade existing Cloud SQL database + if: github.ref == 'refs/heads/master' + runs-on: ubuntu-latest + environment: production-database + permissions: + contents: read + id-token: write + env: + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} + steps: + - name: Require explicit baseline confirmation + run: test "${{ inputs.confirmation }}" = "ADOPT-eafc2a547a4e" + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Authenticate to GCP + uses: google-github-actions/auth@v2 + with: + workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} + service_account: ${{ vars.GCP_DB_MIGRATION_SERVICE_ACCOUNT }} + - name: Set up GCloud + uses: google-github-actions/setup-gcloud@v2 + - name: Install dependencies + run: make install + - name: Start Cloud SQL Auth Proxy + run: bash .github/scripts/start_cloud_sql_proxy.sh + - name: Prepare masked database URLs + run: | + POLICYENGINE_DB_READONLY_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-readonly-password --project policyengine-api)" + POLICYENGINE_DB_MIGRATION_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-migration-password --project policyengine-api)" + export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD + uv run python scripts/write_v1_database_urls.py + - name: Verify legacy schema + run: >- + uv run python scripts/v1_database_migration.py + --mode verify-legacy + --expected-question-rows "${{ inputs.expected_question_rows }}" + - name: Create Cloud SQL backup + id: backup + run: bash .github/scripts/create_cloud_sql_backup.sh + - name: Stamp baseline and upgrade + run: >- + uv run python scripts/v1_database_migration.py + --mode adopt + --confirmation "${{ inputs.confirmation }}" + --backup-id "${{ steps.backup.outputs.backup_id }}" + --expected-question-rows "${{ inputs.expected_question_rows }}" + - name: Verify migrated schema + run: uv run python scripts/v1_database_migration.py --mode verify-head + - name: Stop Cloud SQL Auth Proxy + if: always() + run: bash .github/scripts/stop_cloud_sql_proxy.sh diff --git a/.github/workflows/alembic-v1-check.yml b/.github/workflows/alembic-v1-check.yml new file mode 100644 index 000000000..a26ab9268 --- /dev/null +++ b/.github/workflows/alembic-v1-check.yml @@ -0,0 +1,47 @@ +name: Alembic v1 checks + +on: + workflow_call: + workflow_dispatch: + +jobs: + mysql-lifecycle: + name: Alembic MySQL lifecycle + runs-on: ubuntu-latest + services: + mysql: + image: mysql:8.4 + env: + MYSQL_ROOT_PASSWORD: policyengine_test + MYSQL_DATABASE: policyengine_alembic_test + ports: + - 3306:3306 + options: >- + --health-cmd="mysqladmin ping --silent" + --health-interval=5s + --health-timeout=5s + --health-retries=20 + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install dependencies + run: make install + - name: Test Alembic configuration and MySQL lifecycle + run: >- + uv run pytest tests/unit/data/test_alembic_baseline.py + tests/unit/data/test_v1_database_migration.py + tests/integration/test_alembic_mysql_lifecycle.py -q + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test + - name: Require database at all v1 heads + run: uv run alembic -c alembic-v1.ini current --check-heads + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test + - name: Require no ungenerated v1 operations + run: uv run alembic -c alembic-v1.ini check + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index ac3c60c3c..5c42be20f 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -51,37 +51,29 @@ jobs: - name: Run quality guards run: python scripts/run_quality_guards.py - alembic-mysql: - name: Alembic MySQL lifecycle + detect-v1-alembic-changes: + name: Detect v1 Alembic changes runs-on: ubuntu-latest - services: - mysql: - image: mysql:8.4 - env: - MYSQL_ROOT_PASSWORD: policyengine_test - MYSQL_DATABASE: policyengine_alembic_test - ports: - - 3306:3306 - options: >- - --health-cmd="mysqladmin ping --silent" - --health-interval=5s - --health-timeout=5s - --health-retries=20 + outputs: + changed: ${{ steps.changes.outputs.changed }} steps: - name: Checkout repo uses: actions/checkout@v4 - - name: Setup Python - uses: actions/setup-python@v5 with: - python-version: "3.12" - - name: Install dependencies - run: make install - - name: Test Alembic configuration and MySQL lifecycle + fetch-depth: 0 + - name: Detect relevant changes + id: changes run: >- - pytest tests/unit/data/test_alembic_baseline.py - tests/integration/test_alembic_mysql_lifecycle.py -q - env: - ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test + python scripts/v1_alembic_changes.py + "${{ github.event.pull_request.base.sha }}" + "${{ github.event.pull_request.head.sha }}" + >> "${GITHUB_OUTPUT}" + + alembic-v1-check: + name: Alembic v1 qualification + needs: detect-v1-alembic-changes + if: needs.detect-v1-alembic-changes.outputs.changed == 'true' + uses: ./.github/workflows/alembic-v1-check.yml check-changelog: name: Check changelog fragment runs-on: ubuntu-latest diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 7e5fb16f2..842c8d11f 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -14,11 +14,10 @@ concurrency: cancel-in-progress: false jobs: - Lint: + lint: + name: Lint runs-on: ubuntu-latest - if: | - (github.repository == 'PolicyEngine/policyengine-uk') - && (github.event.head_commit.message == 'Update PolicyEngine API') + if: github.repository == 'PolicyEngine/policyengine-api' steps: - name: Checkout repo uses: actions/checkout@v4 @@ -29,6 +28,11 @@ jobs: - name: Format check with ruff run: ruff format --check . + alembic-v1-check: + name: Alembic v1 qualification + if: github.repository == 'PolicyEngine/policyengine-api' + uses: ./.github/workflows/alembic-v1-check.yml + ensure-staging-model-version-aligns-with-sim-api: name: Ensure staging model version aligns with simulation API runs-on: ubuntu-latest @@ -49,6 +53,7 @@ jobs: versioning: name: Update versioning + needs: [lint, alembic-v1-check] if: | (github.repository == 'PolicyEngine/policyengine-api') && !(github.event.head_commit.message == 'Update PolicyEngine API') @@ -86,7 +91,10 @@ jobs: publish-git-tag: name: Publish Git Tag runs-on: ubuntu-latest - needs: ensure-staging-model-version-aligns-with-sim-api + needs: + - ensure-staging-model-version-aligns-with-sim-api + - lint + - alembic-v1-check if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') @@ -100,12 +108,78 @@ jobs: - name: Publish Git Tag run: ".github/publish-git-tag.sh" + migrate-v1-cloud-sql: + name: Upgrade v1 Cloud SQL schema + runs-on: ubuntu-latest + needs: + - ensure-staging-model-version-aligns-with-sim-api + - publish-git-tag + if: | + (github.repository == 'PolicyEngine/policyengine-api') + && (github.event.head_commit.message == 'Update PolicyEngine API') + environment: production-database + permissions: + contents: read + id-token: write + env: + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Authenticate to GCP + uses: google-github-actions/auth@v2 + with: + workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} + service_account: ${{ vars.GCP_DB_MIGRATION_SERVICE_ACCOUNT }} + - name: Set up GCloud + uses: google-github-actions/setup-gcloud@v2 + - name: Install dependencies + run: make install + - name: Start Cloud SQL Auth Proxy + run: bash .github/scripts/start_cloud_sql_proxy.sh + - name: Prepare masked database URLs + run: | + POLICYENGINE_DB_READONLY_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-readonly-password --project policyengine-api)" + POLICYENGINE_DB_MIGRATION_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-migration-password --project policyengine-api)" + export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD + uv run python scripts/write_v1_database_urls.py + - name: Inspect current database revision + id: schema + run: | + state="$(uv run python scripts/v1_database_migration.py --mode state)" + echo "state=${state}" >> "${GITHUB_OUTPUT}" + - name: Refuse implicit baseline adoption + if: steps.schema.outputs.state == 'unversioned' || steps.schema.outputs.state == 'invalid' + run: | + echo "database is unversioned or has invalid Alembic state; run the explicit adoption workflow or repair the version table first" >&2 + exit 1 + - name: Create Cloud SQL backup + if: steps.schema.outputs.state == 'pending' + id: backup + run: bash .github/scripts/create_cloud_sql_backup.sh + - name: Upgrade pending v1 migrations + if: steps.schema.outputs.state == 'pending' + run: >- + uv run python scripts/v1_database_migration.py + --mode upgrade + --backup-id "${{ steps.backup.outputs.backup_id }}" + - name: Verify v1 database at head + run: uv run python scripts/v1_database_migration.py --mode verify-head + - name: Stop Cloud SQL Auth Proxy + if: always() + run: bash .github/scripts/stop_cloud_sql_proxy.sh + deploy-staging: name: Deploy staging App Engine version runs-on: ubuntu-latest needs: - ensure-staging-model-version-aligns-with-sim-api - publish-git-tag + - migrate-v1-cloud-sql if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') @@ -213,6 +287,7 @@ jobs: needs: - ensure-staging-model-version-aligns-with-sim-api - publish-git-tag + - migrate-v1-cloud-sql if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index 052617746..4f2629f98 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and Alembic while preserving public API contracts, including generated migrations that remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and a dedicated Cloud SQL/MySQL Alembic chain while preserving public API contracts. Add conditional pull-request migration qualification, mandatory release qualification, explicit existing-database adoption, backup-before-DDL, and a fail-closed pre-deployment migration gate. Generated migrations remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 6e50c61b4..9054a749a 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -9,9 +9,12 @@ AI **MUST NOT manually author Alembic revision scripts**. Generate every schema revision from reviewed SQLAlchemy metadata: ```bash -alembic revision --autogenerate -m "" +uv run alembic -c alembic-v1.ini revision --autogenerate -m "" ``` +The mandatory generation operation is `alembic revision --autogenerate`; the +explicit v1 configuration keeps these revisions in the Cloud SQL/MySQL chain. + If generated operations are wrong, first correct the model metadata and regenerate. AI may make minimal post-generation corrections only for dialect compatibility, reversibility, or an objectively incorrect autogenerated @@ -25,7 +28,8 @@ those narrow review corrections, stop and request a human migration decision. Before committing a migration: -1. Run `alembic check` and review the generated operations. +1. Run `uv run alembic -c alembic-v1.ini check` and review the generated + operations. 2. Upgrade a fresh database to `head`. 3. Compare any existing database to the ORM metadata before stamping it. 4. Downgrade one revision and upgrade to `head` again in an isolated database. @@ -43,12 +47,21 @@ Every Alembic command must therefore receive an explicit MySQL URL: ```bash export ALEMBIC_DATABASE_URL="mysql+pymysql://..." -uv run alembic upgrade head +uv run alembic -c alembic-v1.ini upgrade head ``` -Do not put a database URL in `alembic.ini`. Keeping the configuration without a -default prevents an omitted environment variable from silently migrating the -wrong database. +The v1 configuration is `alembic-v1.ini`; its revision tree is `migrations/v1`. +Do not put a database URL in the configuration. Keeping it without a default +prevents an omitted environment variable from silently migrating the wrong +database. + +## Separate v1 and v2 migration domains + +The current chain manages API v1 Cloud SQL/MySQL only. Stage 8 must introduce a +different configuration and a separate revision chain for the v2 +Supabase/Postgres schema. Never point `alembic-v1.ini` at Supabase, append v2 +Postgres revisions under `migrations/v1`, or use both Alembic and Supabase CLI +schema migrations as authorities for the same tables. ### Fresh database qualification @@ -69,22 +82,41 @@ shared or deployed database. The baseline represents tables that already exist in deployed API v1 databases. It must never be upgraded into one of those databases. Before the -first Alembic-managed release for each environment: - -1. Supply a read-only URL as `STAGE7_EXISTING_DATABASE_URL` and run - `uv run pytest tests/integration/test_v1_schema_metadata_compatibility.py -q`. -2. Review every reported difference and stop if any difference is not exactly - accounted for by a reviewed post-baseline migration. -3. With a separately authorized migration connection, stamp the existing - database at the generated baseline revision: - `uv run alembic stamp eafc2a547a4e`. -4. Run `uv run alembic upgrade head` to apply every reviewed post-baseline - migration. -5. Repeat the metadata comparison and require zero remaining differences. -6. Confirm the application starts without emitting schema DDL. +first Alembic-managed release: + +1. Run the manual `Adopt existing v1 Cloud SQL schema` workflow from `master`. +2. Supply the exact confirmation `ADOPT-eafc2a547a4e`; the workflow must reject + every other value. +3. Compare the database through the read-only schema account and require only + the explicitly reviewed legacy differences. +4. Require the known data invariants, including no nullable execution IDs or + datasets and the reviewed orphaned-question row count. +5. Create and complete an on-demand Cloud SQL backup. +6. With the separately authorized migration account, stamp the existing + database at `eafc2a547a4e` and upgrade to `head` while holding the migration + advisory lock. +7. Repeat the read-only metadata comparison, require zero drift, and require all + Alembic heads to be current. Stamping is an explicit release operation after the read-only comparison; it -must not run automatically during application startup or ordinary CI. +must not run automatically during application startup, PR CI, or an ordinary +release. The ordinary release migration job fails closed when it sees an +unversioned database. + +### CI/CD behavior + +- Pull requests first detect whether the v1 Alembic surface changed. Relevant + PRs invoke `.github/workflows/alembic-v1-check.yml`; unrelated PRs skip it. +- Every push to `master` runs the reusable disposable-MySQL qualification next + to lint, regardless of which paths changed. +- The release migration job runs before either staging deployment. It refuses + implicit adoption, takes a backup only when revisions are pending, upgrades + under the advisory lock, and requires zero post-migration drift. +- Database credentials and derived URLs must never be printed. The application + runtime account must not be granted schema-migration privileges. +- Production downgrades are never automatic. MySQL DDL is non-transactional, so + rollback means application compatibility plus a reviewed forward fix or a + database restore. ### Production-only table cleanup diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py new file mode 100644 index 000000000..5b250fcc2 --- /dev/null +++ b/tests/unit/test_alembic_workflows.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[2] + + +def _workflow(name: str) -> str: + return (REPO / ".github" / "workflows" / name).read_text(encoding="utf-8") + + +def test_pr_runs_reusable_alembic_check_only_for_relevant_changes(): + workflow = _workflow("pr.yml") + + assert "detect-v1-alembic-changes:" in workflow + assert "python scripts/v1_alembic_changes.py" in workflow + assert "alembic-v1-check:" in workflow + assert "needs.detect-v1-alembic-changes.outputs.changed == 'true'" in workflow + assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow + + +def test_push_always_runs_lint_and_alembic_qualification_before_versioning(): + workflow = _workflow("push.yml") + + assert "lint:" in workflow + assert "alembic-v1-check:" in workflow + assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow + assert "needs: [lint, alembic-v1-check]" in workflow + assert "github.repository == 'PolicyEngine/policyengine-uk'" not in workflow + + +def test_reusable_alembic_check_uses_only_disposable_mysql(): + workflow = _workflow("alembic-v1-check.yml") + + assert "workflow_call:" in workflow + assert "workflow_dispatch:" in workflow + assert "mysql:8.4" in workflow + assert "policyengine_alembic_test" in workflow + assert "alembic-v1.ini" in workflow + assert "STAGE7_EXISTING_DATABASE_URL" not in workflow + assert "POLICYENGINE_DB_MIGRATION_PASSWORD" not in workflow + + +def test_adoption_workflow_is_manual_explicit_and_backup_first(): + workflow = _workflow("adopt-v1-cloud-sql.yml") + + assert "workflow_dispatch:" in workflow + assert "ADOPT-eafc2a547a4e" in workflow + assert "environment: production-database" in workflow + assert "protection" not in workflow.lower() + assert workflow.index("Verify legacy schema") < workflow.index( + "Create Cloud SQL backup" + ) + assert workflow.index("Create Cloud SQL backup") < workflow.index( + "Stamp baseline and upgrade" + ) + assert "--mode adopt" in workflow + + +def test_release_migration_fails_closed_and_gates_both_staging_deploys(): + workflow = _workflow("push.yml") + + assert "migrate-v1-cloud-sql:" in workflow + assert "environment: production-database" in workflow + assert "--mode upgrade" in workflow + assert "--mode adopt" not in workflow + assert "database is unversioned" in workflow + + app_engine_job = workflow[workflow.index(" deploy-staging:") :] + app_engine_job = app_engine_job[ + : app_engine_job.index("\n deploy-cloud-run-staging:") + ] + assert "migrate-v1-cloud-sql" in app_engine_job + + cloud_run_job = workflow[workflow.index(" deploy-cloud-run-staging:") :] + cloud_run_job = cloud_run_job[ + : cloud_run_job.index("\n integration-tests-staging:") + ] + assert "migrate-v1-cloud-sql" in cloud_run_job + + +def test_cloud_sql_workflows_use_oidc_and_separate_database_credentials(): + workflows = _workflow("adopt-v1-cloud-sql.yml") + _workflow("push.yml") + + assert "google-github-actions/auth@v2" in workflows + assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in workflows + assert "policyengine-api-prod-db-readonly-password" in workflows + assert "policyengine-api-prod-db-migration-password" in workflows + assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in workflows + assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in workflows + assert ( + "POLICYENGINE_DB_PASSWORD: ${{ secrets.POLICYENGINE_DB_PASSWORD }}" + not in _workflow("adopt-v1-cloud-sql.yml") + ) + + +def test_v1_and_future_v2_alembic_domains_are_explicitly_separate(): + assert (REPO / "alembic-v1.ini").exists() + assert (REPO / "migrations" / "v1" / "env.py").exists() + + guidance = REPO / "docs" / "engineering" / "skills" / "alembic-migrations.md" + text = guidance.read_text(encoding="utf-8") + assert "alembic-v1.ini" in text + assert "migrations/v1" in text + assert "Supabase/Postgres" in text + assert "separate revision chain" in text From 42fd63002540de9eb274e542415728773c62ed39 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:54:07 +0300 Subject: [PATCH 81/89] fix: use installed Python in Alembic CI --- .github/workflows/alembic-v1-check.yml | 6 +++--- tests/unit/test_alembic_workflows.py | 8 ++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/.github/workflows/alembic-v1-check.yml b/.github/workflows/alembic-v1-check.yml index a26ab9268..529d8b6fd 100644 --- a/.github/workflows/alembic-v1-check.yml +++ b/.github/workflows/alembic-v1-check.yml @@ -32,16 +32,16 @@ jobs: run: make install - name: Test Alembic configuration and MySQL lifecycle run: >- - uv run pytest tests/unit/data/test_alembic_baseline.py + python -m pytest tests/unit/data/test_alembic_baseline.py tests/unit/data/test_v1_database_migration.py tests/integration/test_alembic_mysql_lifecycle.py -q env: ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test - name: Require database at all v1 heads - run: uv run alembic -c alembic-v1.ini current --check-heads + run: python -m alembic -c alembic-v1.ini current --check-heads env: ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test - name: Require no ungenerated v1 operations - run: uv run alembic -c alembic-v1.ini check + run: python -m alembic -c alembic-v1.ini check env: ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 5b250fcc2..0234cc862 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -42,6 +42,14 @@ def test_reusable_alembic_check_uses_only_disposable_mysql(): assert "POLICYENGINE_DB_MIGRATION_PASSWORD" not in workflow +def test_reusable_alembic_check_uses_the_installed_python_environment(): + workflow = _workflow("alembic-v1-check.yml") + + assert "python -m pytest" in workflow + assert "python -m alembic" in workflow + assert "uv run" not in workflow + + def test_adoption_workflow_is_manual_explicit_and_backup_first(): workflow = _workflow("adopt-v1-cloud-sql.yml") From ec947295baaacaa690b7f984339c41bff126fd11 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:30:12 +0300 Subject: [PATCH 82/89] fix: support connection-driven v1 adoption --- .github/scripts/create_cloud_sql_backup.sh | 19 +++++++++++--- migrations/v1/env.py | 26 +++++++++++-------- .../test_alembic_mysql_lifecycle.py | 5 +++- tests/unit/test_alembic_workflows.py | 14 ++++++++++ 4 files changed, 48 insertions(+), 16 deletions(-) diff --git a/.github/scripts/create_cloud_sql_backup.sh b/.github/scripts/create_cloud_sql_backup.sh index 5f43e443b..5c0e18ccd 100644 --- a/.github/scripts/create_cloud_sql_backup.sh +++ b/.github/scripts/create_cloud_sql_backup.sh @@ -6,13 +6,24 @@ set -euo pipefail : "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" instance_id="${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME##*:}" -description="policyengine-api-v1-alembic-${GITHUB_SHA:-manual}" +description="policyengine-api-v1-alembic-${GITHUB_SHA:-manual}-${GITHUB_RUN_ID:-local}-${GITHUB_RUN_ATTEMPT:-1}" +gcloud sql backups create \ + --project policyengine-api \ + --instance "${instance_id}" \ + --description "${description}" \ + --quiet + +# `gcloud sql backups create` waits for completion but does not consistently +# emit the created resource with value-format output. Recover the ID from the +# unique workflow description and require the service-reported successful state. backup_id="$( - gcloud sql backups create \ + gcloud sql backups list \ --project policyengine-api \ --instance "${instance_id}" \ - --description "${description}" \ - --format 'value(id)' + --filter="description=${description} AND status=SUCCESSFUL" \ + --sort-by='~startTime' \ + --limit=1 \ + --format='value(id)' )" if [[ -z "${backup_id}" ]]; then diff --git a/migrations/v1/env.py b/migrations/v1/env.py index e9d18e768..8fce659c9 100644 --- a/migrations/v1/env.py +++ b/migrations/v1/env.py @@ -14,17 +14,22 @@ if config.config_file_name is not None: fileConfig(config.config_file_name) -database_url = os.environ.get("ALEMBIC_DATABASE_URL") or config.get_main_option( - "sqlalchemy.url" -) -if not database_url: - raise RuntimeError( - "ALEMBIC_DATABASE_URL is required; the local SQLite cache is not an " - "Alembic migration target" +provided_connection = config.attributes.get("connection") +if provided_connection is not None: + if provided_connection.dialect.name != "mysql": + raise RuntimeError("Alembic migrations must target a MySQL database") +else: + database_url = os.environ.get("ALEMBIC_DATABASE_URL") or config.get_main_option( + "sqlalchemy.url" ) -if make_url(database_url).get_backend_name() != "mysql": - raise RuntimeError("Alembic migrations must target a MySQL database") -config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) + if not database_url: + raise RuntimeError( + "ALEMBIC_DATABASE_URL is required; the local SQLite cache is not an " + "Alembic migration target" + ) + if make_url(database_url).get_backend_name() != "mysql": + raise RuntimeError("Alembic migrations must target a MySQL database") + config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) target_metadata = V1Base.metadata @@ -42,7 +47,6 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: - provided_connection = config.attributes.get("connection") if provided_connection is not None: _run_migrations_with_connection(provided_connection) return diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index 857643c06..8a0ff9027 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -110,7 +110,9 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): engine.dispose() -def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema(): +def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( + monkeypatch, +): database_url = _ephemeral_mysql_url() config = _alembic_config(database_url) engine = create_engine(database_url) @@ -149,6 +151,7 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema() operations.drop_table("alembic_version") with engine.connect() as connection: + monkeypatch.delenv("ALEMBIC_DATABASE_URL") adopt_database( connection, confirmation=ADOPTION_CONFIRMATION, diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 0234cc862..407b6132c 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -103,6 +103,20 @@ def test_cloud_sql_workflows_use_oidc_and_separate_database_credentials(): ) +def test_backup_helper_recovers_and_verifies_the_created_backup_id(): + script = (REPO / ".github" / "scripts" / "create_cloud_sql_backup.sh").read_text( + encoding="utf-8" + ) + + assert "GITHUB_RUN_ID" in script + assert "gcloud sql backups create" in script + assert "gcloud sql backups list" in script + assert "status=SUCCESSFUL" in script + assert script.index("gcloud sql backups create") < script.index( + "gcloud sql backups list" + ) + + def test_v1_and_future_v2_alembic_domains_are_explicitly_separate(): assert (REPO / "alembic-v1.ini").exists() assert (REPO / "migrations" / "v1" / "env.py").exists() From 208b06ec3ac4f35a2e3f7f6578c885bea7986497 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:32:33 +0300 Subject: [PATCH 83/89] fix: commit Alembic revision updates --- scripts/v1_database_migration.py | 5 +- .../test_alembic_mysql_lifecycle.py | 7 ++- tests/unit/data/test_v1_database_migration.py | 54 +++++++++++++++++++ 3 files changed, 64 insertions(+), 2 deletions(-) diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index 415e77fd6..fa19f29c9 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -265,7 +265,10 @@ def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) engine = create_engine(_database_url(args.mode), poolclass=NullPool) try: - with engine.connect() as connection: + connection_context = ( + engine.begin() if args.mode in {"adopt", "upgrade"} else engine.connect() + ) + with connection_context as connection: if args.mode == "state": print(database_state(connection).value) elif args.mode == "verify-legacy": diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index 8a0ff9027..5fd573ee5 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -26,7 +26,9 @@ from policyengine_api.data.v1_models import V1Base from scripts.v1_database_migration import ( ADOPTION_CONFIRMATION, + DatabaseState, adopt_database, + database_state, ) @@ -150,7 +152,7 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( ) operations.drop_table("alembic_version") - with engine.connect() as connection: + with engine.begin() as connection: monkeypatch.delenv("ALEMBIC_DATABASE_URL") adopt_database( connection, @@ -159,6 +161,9 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( expected_question_rows=1, ) + with engine.connect() as connection: + assert database_state(connection) is DatabaseState.HEAD + inspector = inspect(engine) assert "question" not in inspector.get_table_names() assert "tracers" not in inspector.get_table_names() diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py index de52c6fe3..1e08b2f50 100644 --- a/tests/unit/data/test_v1_database_migration.py +++ b/tests/unit/data/test_v1_database_migration.py @@ -1,8 +1,11 @@ from __future__ import annotations +from contextlib import nullcontext + import pytest from sqlalchemy import Column, Integer, MetaData, String, Table, text +import scripts.v1_database_migration as migration from scripts.v1_alembic_changes import is_v1_alembic_path from scripts.v1_database_migration import ( ADOPTION_CONFIRMATION, @@ -156,3 +159,54 @@ def test_metadata_difference_rejects_unknown_shapes_without_repr_leakage(): describe_metadata_difference(difference) assert secret not in str(error.value) + + +def test_adoption_cli_commits_the_externally_supplied_connection(monkeypatch): + connection = object() + + class FakeEngine: + def begin(self): + return nullcontext(connection) + + def connect(self): + raise AssertionError("adoption must use a committing transaction") + + def dispose(self): + pass + + calls = [] + monkeypatch.setattr( + migration, "create_engine", lambda *args, **kwargs: FakeEngine() + ) + monkeypatch.setattr( + migration, + "adopt_database", + lambda supplied_connection, **kwargs: calls.append( + (supplied_connection, kwargs) + ), + ) + monkeypatch.setenv("ALEMBIC_DATABASE_URL", "mysql+pymysql://unused") + + assert ( + migration.main( + [ + "--mode", + "adopt", + "--confirmation", + ADOPTION_CONFIRMATION, + "--backup-id", + "verified-backup", + ] + ) + == 0 + ) + assert calls == [ + ( + connection, + { + "confirmation": ADOPTION_CONFIRMATION, + "backup_id": "verified-backup", + "expected_question_rows": 9, + }, + ) + ] From 3aea079e8b84456ae15d62d10fc83684421b7b85 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:08:34 +0300 Subject: [PATCH 84/89] chore: remove completed database adoption tooling --- .github/workflows/adopt-v1-cloud-sql.yml | 76 -------------- .github/workflows/push.yml | 4 +- changelog.d/3788.changed.md | 2 +- docs/engineering/skills/README.md | 2 +- docs/engineering/skills/alembic-migrations.md | 39 ++------ scripts/v1_database_migration.py | 99 ++----------------- .../test_alembic_mysql_lifecycle.py | 67 +++---------- tests/unit/data/test_v1_database_migration.py | 40 ++++---- tests/unit/test_alembic_workflows.py | 41 +++----- 9 files changed, 65 insertions(+), 305 deletions(-) delete mode 100644 .github/workflows/adopt-v1-cloud-sql.yml diff --git a/.github/workflows/adopt-v1-cloud-sql.yml b/.github/workflows/adopt-v1-cloud-sql.yml deleted file mode 100644 index e0d816064..000000000 --- a/.github/workflows/adopt-v1-cloud-sql.yml +++ /dev/null @@ -1,76 +0,0 @@ -name: Adopt existing v1 Cloud SQL schema - -on: - workflow_dispatch: - inputs: - confirmation: - description: Type ADOPT-eafc2a547a4e to stamp and upgrade the existing database - required: true - type: string - expected_question_rows: - description: Expected number of orphaned question rows that the migration removes - required: true - default: "9" - type: string - -concurrency: - group: policyengine-api-v1-cloud-sql-schema - cancel-in-progress: false - -jobs: - adopt: - name: Stamp and upgrade existing Cloud SQL database - if: github.ref == 'refs/heads/master' - runs-on: ubuntu-latest - environment: production-database - permissions: - contents: read - id-token: write - env: - POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} - steps: - - name: Require explicit baseline confirmation - run: test "${{ inputs.confirmation }}" = "ADOPT-eafc2a547a4e" - - name: Checkout repo - uses: actions/checkout@v4 - - name: Setup Python - uses: actions/setup-python@v5 - with: - python-version: "3.12" - - name: Authenticate to GCP - uses: google-github-actions/auth@v2 - with: - workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} - service_account: ${{ vars.GCP_DB_MIGRATION_SERVICE_ACCOUNT }} - - name: Set up GCloud - uses: google-github-actions/setup-gcloud@v2 - - name: Install dependencies - run: make install - - name: Start Cloud SQL Auth Proxy - run: bash .github/scripts/start_cloud_sql_proxy.sh - - name: Prepare masked database URLs - run: | - POLICYENGINE_DB_READONLY_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-readonly-password --project policyengine-api)" - POLICYENGINE_DB_MIGRATION_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-migration-password --project policyengine-api)" - export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD - uv run python scripts/write_v1_database_urls.py - - name: Verify legacy schema - run: >- - uv run python scripts/v1_database_migration.py - --mode verify-legacy - --expected-question-rows "${{ inputs.expected_question_rows }}" - - name: Create Cloud SQL backup - id: backup - run: bash .github/scripts/create_cloud_sql_backup.sh - - name: Stamp baseline and upgrade - run: >- - uv run python scripts/v1_database_migration.py - --mode adopt - --confirmation "${{ inputs.confirmation }}" - --backup-id "${{ steps.backup.outputs.backup_id }}" - --expected-question-rows "${{ inputs.expected_question_rows }}" - - name: Verify migrated schema - run: uv run python scripts/v1_database_migration.py --mode verify-head - - name: Stop Cloud SQL Auth Proxy - if: always() - run: bash .github/scripts/stop_cloud_sql_proxy.sh diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 842c8d11f..c5c06b61c 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -152,10 +152,10 @@ jobs: run: | state="$(uv run python scripts/v1_database_migration.py --mode state)" echo "state=${state}" >> "${GITHUB_OUTPUT}" - - name: Refuse implicit baseline adoption + - name: Refuse unversioned or invalid schema if: steps.schema.outputs.state == 'unversioned' || steps.schema.outputs.state == 'invalid' run: | - echo "database is unversioned or has invalid Alembic state; run the explicit adoption workflow or repair the version table first" >&2 + echo "database is unversioned or has invalid Alembic state; automatic baseline stamping is disabled and manual recovery is required" >&2 exit 1 - name: Create Cloud SQL backup if: steps.schema.outputs.state == 'pending' diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md index 4f2629f98..e86a42214 100644 --- a/changelog.d/3788.changed.md +++ b/changelog.d/3788.changed.md @@ -1 +1 @@ -Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and a dedicated Cloud SQL/MySQL Alembic chain while preserving public API contracts. Add conditional pull-request migration qualification, mandatory release qualification, explicit existing-database adoption, backup-before-DDL, and a fail-closed pre-deployment migration gate. Generated migrations remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and a dedicated Cloud SQL/MySQL Alembic chain while preserving public API contracts. Add conditional pull-request migration qualification, mandatory release qualification, backup-before-DDL, and a fail-closed pre-deployment migration gate. Generated migrations remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. diff --git a/docs/engineering/skills/README.md b/docs/engineering/skills/README.md index 1aa49f799..a3248f5b2 100644 --- a/docs/engineering/skills/README.md +++ b/docs/engineering/skills/README.md @@ -10,7 +10,7 @@ first, then keep adapters thin. Current skills: - `alembic-migrations.md`: mandatory autogenerated Alembic revision workflow, - adoption safeguards, and migration validation. + deployment safeguards, and migration validation. - `github-prs.md`: PR workflow and migration PR handoff expectations. - `migration_contracts.md`: API v2 migration route contracts, route-group metadata, generated migration artifacts, and quality guards. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 9054a749a..ea6891559 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -31,13 +31,14 @@ Before committing a migration: 1. Run `uv run alembic -c alembic-v1.ini check` and review the generated operations. 2. Upgrade a fresh database to `head`. -3. Compare any existing database to the ORM metadata before stamping it. +3. Compare the deployed database to the ORM metadata before release. 4. Downgrade one revision and upgrade to `head` again in an isolated database. 5. Confirm application startup performs no implicit DDL. -Never print database credentials or embed them in Alembic configuration. Never -run a generated baseline's create operations against an existing database; -verify schema equivalence and stamp it instead. +Never print database credentials or embed them in Alembic configuration. The +deployed v1 database is already versioned; never recreate or restamp its +baseline. Missing or invalid revision history must fail closed for manual +recovery. ## API v1 database targets @@ -78,31 +79,6 @@ live schema with the ORM metadata, downgrade one revision, and upgrade to `head` again. Its host and schema-name checks prevent it from running against a shared or deployed database. -### Existing database adoption gate - -The baseline represents tables that already exist in deployed API v1 -databases. It must never be upgraded into one of those databases. Before the -first Alembic-managed release: - -1. Run the manual `Adopt existing v1 Cloud SQL schema` workflow from `master`. -2. Supply the exact confirmation `ADOPT-eafc2a547a4e`; the workflow must reject - every other value. -3. Compare the database through the read-only schema account and require only - the explicitly reviewed legacy differences. -4. Require the known data invariants, including no nullable execution IDs or - datasets and the reviewed orphaned-question row count. -5. Create and complete an on-demand Cloud SQL backup. -6. With the separately authorized migration account, stamp the existing - database at `eafc2a547a4e` and upgrade to `head` while holding the migration - advisory lock. -7. Repeat the read-only metadata comparison, require zero drift, and require all - Alembic heads to be current. - -Stamping is an explicit release operation after the read-only comparison; it -must not run automatically during application startup, PR CI, or an ordinary -release. The ordinary release migration job fails closed when it sees an -unversioned database. - ### CI/CD behavior - Pull requests first detect whether the v1 Alembic surface changed. Relevant @@ -110,8 +86,9 @@ unversioned database. - Every push to `master` runs the reusable disposable-MySQL qualification next to lint, regardless of which paths changed. - The release migration job runs before either staging deployment. It refuses - implicit adoption, takes a backup only when revisions are pending, upgrades - under the advisory lock, and requires zero post-migration drift. + unversioned or invalid history, takes a backup only when revisions are + pending, upgrades under the advisory lock, and requires zero post-migration + drift. - Database credentials and derived URLs must never be printed. The application runtime account must not be granted schema-migration privileges. - Production downgrades are never automatic. MySQL DDL is non-transactional, so diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index fa19f29c9..bcd18a088 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -1,4 +1,4 @@ -"""Safely inspect, adopt, and upgrade the API v1 Cloud SQL schema.""" +"""Safely inspect and upgrade the API v1 Cloud SQL schema.""" from __future__ import annotations @@ -20,17 +20,8 @@ from policyengine_api.data.v1_models import V1Base -BASELINE_REVISION = "eafc2a547a4e" -ADOPTION_CONFIRMATION = f"ADOPT-{BASELINE_REVISION}" ALEMBIC_CONFIG = REPO / "alembic-v1.ini" MIGRATION_LOCK_NAME = "policyengine-api-v1-alembic" -EXPECTED_LEGACY_DIFFERENCES = frozenset( - { - "extra_table:question", - "nullable:reform_impact.execution_id:true->false", - "default:reform_impact.dataset:present->none", - } -) class DatabaseState(StrEnum): @@ -71,14 +62,6 @@ def classify_database_state( return DatabaseState.PENDING -def require_adoption_confirmation(confirmation: str) -> None: - if confirmation != ADOPTION_CONFIRMATION: - raise ValueError( - "Explicit adoption confirmation is required; expected " - f"{ADOPTION_CONFIRMATION!r}" - ) - - def describe_metadata_difference(difference: Any) -> str: """Return a stable, data-free description of an Alembic metadata diff.""" @@ -136,39 +119,6 @@ def metadata_differences(connection: Connection) -> set[str]: } -def verify_legacy_schema( - connection: Connection, *, expected_question_rows: int -) -> None: - state = database_state(connection) - if state is not DatabaseState.UNVERSIONED: - raise RuntimeError( - f"Existing database must be unversioned before adoption; found {state}" - ) - - differences = metadata_differences(connection) - if differences != EXPECTED_LEGACY_DIFFERENCES: - unexpected = sorted(differences - EXPECTED_LEGACY_DIFFERENCES) - missing = sorted(EXPECTED_LEGACY_DIFFERENCES - differences) - raise RuntimeError( - "Existing schema does not match the reviewed legacy shape; " - f"unexpected={unexpected}, missing={missing}" - ) - - question_rows = connection.scalar(text("SELECT COUNT(*) FROM question")) - if question_rows != expected_question_rows: - raise RuntimeError( - "question row count changed; expected " - f"{expected_question_rows}, found {question_rows}" - ) - - for column in ("execution_id", "dataset"): - null_rows = connection.scalar( - text(f"SELECT COUNT(*) FROM reform_impact WHERE {column} IS NULL") - ) - if null_rows: - raise RuntimeError(f"reform_impact.{column} contains {null_rows} NULL rows") - - def verify_head_schema(connection: Connection) -> None: state = database_state(connection) if state is not DatabaseState.HEAD: @@ -194,35 +144,14 @@ def _release_lock(connection: Connection) -> None: ) -def adopt_database( - connection: Connection, - *, - confirmation: str, - backup_id: str, - expected_question_rows: int, -) -> None: - require_adoption_confirmation(confirmation) - if not backup_id.strip(): - raise ValueError("A completed Cloud SQL backup ID is required") - - _acquire_lock(connection) - try: - verify_legacy_schema(connection, expected_question_rows=expected_question_rows) - config = _config(connection) - command.stamp(config, BASELINE_REVISION) - command.upgrade(config, "head") - verify_head_schema(connection) - finally: - _release_lock(connection) - - def upgrade_database(connection: Connection, *, backup_id: str) -> None: _acquire_lock(connection) try: state = database_state(connection) - if state is DatabaseState.UNVERSIONED: + if state in {DatabaseState.UNVERSIONED, DatabaseState.INVALID}: raise RuntimeError( - "database is unversioned; run the explicit adoption workflow first" + "database has no valid Alembic revision; automatic baseline " + "stamping is disabled" ) if state is DatabaseState.HEAD: verify_head_schema(connection) @@ -239,7 +168,7 @@ def upgrade_database(connection: Connection, *, backup_id: str) -> None: def _database_url(mode: str) -> str: env_name = ( "STAGE7_EXISTING_DATABASE_URL" - if mode in {"verify-legacy", "verify-head", "state"} + if mode in {"verify-head", "state"} else "ALEMBIC_DATABASE_URL" ) try: @@ -253,11 +182,9 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument( "--mode", required=True, - choices=("state", "verify-legacy", "verify-head", "adopt", "upgrade"), + choices=("state", "verify-head", "upgrade"), ) - parser.add_argument("--confirmation", default="") parser.add_argument("--backup-id", default="") - parser.add_argument("--expected-question-rows", type=int, default=9) return parser.parse_args(argv) @@ -266,25 +193,13 @@ def main(argv: list[str] | None = None) -> int: engine = create_engine(_database_url(args.mode), poolclass=NullPool) try: connection_context = ( - engine.begin() if args.mode in {"adopt", "upgrade"} else engine.connect() + engine.begin() if args.mode == "upgrade" else engine.connect() ) with connection_context as connection: if args.mode == "state": print(database_state(connection).value) - elif args.mode == "verify-legacy": - verify_legacy_schema( - connection, - expected_question_rows=args.expected_question_rows, - ) elif args.mode == "verify-head": verify_head_schema(connection) - elif args.mode == "adopt": - adopt_database( - connection, - confirmation=args.confirmation, - backup_id=args.backup_id, - expected_question_rows=args.expected_question_rows, - ) else: upgrade_database(connection, backup_id=args.backup_id) finally: diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index 5fd573ee5..9e4f0344f 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -6,7 +6,6 @@ from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext -from alembic.operations import Operations import pytest from sqlalchemy import ( Column, @@ -17,7 +16,6 @@ Text, create_engine, inspect, - text, ) from sqlalchemy.dialects.mysql import LONGTEXT from sqlalchemy.engine import make_url @@ -25,14 +23,14 @@ from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base from scripts.v1_database_migration import ( - ADOPTION_CONFIRMATION, DatabaseState, - adopt_database, database_state, + upgrade_database, ) BASELINE_REVISION = "eafc2a547a4e" +PREVIOUS_REVISION = "17bb32415f97" def _deployed_question_table() -> Table: @@ -112,9 +110,7 @@ def test_fresh_upgrade_check_downgrade_and_reupgrade(): engine.dispose() -def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( - monkeypatch, -): +def test_pending_upgrade_commits_head_revision(monkeypatch): database_url = _ephemeral_mysql_url() config = _alembic_config(database_url) engine = create_engine(database_url) @@ -123,46 +119,20 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( try: command.downgrade(config, "base") question.drop(engine, checkfirst=True) - command.upgrade(config, BASELINE_REVISION) - question.create(engine) - with engine.begin() as connection: - operations = Operations(MigrationContext.configure(connection)) - operations.drop_table("tracers") - operations.alter_column( - "reform_impact", - "execution_id", - existing_type=String(255), - nullable=True, - ) - operations.alter_column( - "reform_impact", - "dataset", - existing_type=String(255), - existing_nullable=False, - server_default=text("'default'"), - ) - connection.execute( - question.insert(), - { - "question": "Historical prototype", - "country_id": "uk", - "subtask": "complete", - "status": "ok", - }, - ) - operations.drop_table("alembic_version") + command.upgrade(config, "head") + command.downgrade(config, PREVIOUS_REVISION) + + with engine.connect() as connection: + assert database_state(connection) is DatabaseState.PENDING with engine.begin() as connection: monkeypatch.delenv("ALEMBIC_DATABASE_URL") - adopt_database( - connection, - confirmation=ADOPTION_CONFIRMATION, - backup_id="test-backup", - expected_question_rows=1, - ) + upgrade_database(connection, backup_id="test-backup") with engine.connect() as connection: assert database_state(connection) is DatabaseState.HEAD + context = MigrationContext.configure(connection) + assert compare_metadata(context, V1Base.metadata) == [] inspector = inspect(engine) assert "question" not in inspector.get_table_names() @@ -172,21 +142,6 @@ def test_upgrade_removes_orphaned_question_table_and_downgrade_restores_schema( } assert reform_impact_columns["dataset"]["default"] is None assert reform_impact_columns["execution_id"]["nullable"] is False - - command.downgrade(config, BASELINE_REVISION) - - assert "question" in inspect(engine).get_table_names() - assert { - column["name"] for column in inspect(engine).get_columns("question") - } == { - "question_id", - "question", - "answer", - "policy_id", - "country_id", - "subtask", - "status", - } finally: command.upgrade(config, "head") engine.dispose() diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py index 1e08b2f50..e4d654ed9 100644 --- a/tests/unit/data/test_v1_database_migration.py +++ b/tests/unit/data/test_v1_database_migration.py @@ -8,13 +8,11 @@ import scripts.v1_database_migration as migration from scripts.v1_alembic_changes import is_v1_alembic_path from scripts.v1_database_migration import ( - ADOPTION_CONFIRMATION, - BASELINE_REVISION, DatabaseState, build_database_url, classify_database_state, describe_metadata_difference, - require_adoption_confirmation, + upgrade_database, ) @@ -91,20 +89,13 @@ def test_database_state_is_head_only_when_all_script_heads_are_applied(): assert ( classify_database_state( version_table_exists=True, - current_heads={BASELINE_REVISION}, + current_heads={"previous-revision"}, script_heads={"head-a"}, ) is DatabaseState.PENDING ) -def test_adoption_requires_the_exact_explicit_confirmation(): - require_adoption_confirmation(ADOPTION_CONFIRMATION) - - with pytest.raises(ValueError, match="confirmation"): - require_adoption_confirmation("yes") - - def test_metadata_difference_descriptions_are_stable_and_do_not_include_data(): metadata = MetaData() question = Table( @@ -161,7 +152,7 @@ def test_metadata_difference_rejects_unknown_shapes_without_repr_leakage(): assert secret not in str(error.value) -def test_adoption_cli_commits_the_externally_supplied_connection(monkeypatch): +def test_upgrade_cli_commits_the_externally_supplied_connection(monkeypatch): connection = object() class FakeEngine: @@ -169,7 +160,7 @@ def begin(self): return nullcontext(connection) def connect(self): - raise AssertionError("adoption must use a committing transaction") + raise AssertionError("upgrade must use a committing transaction") def dispose(self): pass @@ -180,7 +171,7 @@ def dispose(self): ) monkeypatch.setattr( migration, - "adopt_database", + "upgrade_database", lambda supplied_connection, **kwargs: calls.append( (supplied_connection, kwargs) ), @@ -191,9 +182,7 @@ def dispose(self): migration.main( [ "--mode", - "adopt", - "--confirmation", - ADOPTION_CONFIRMATION, + "upgrade", "--backup-id", "verified-backup", ] @@ -204,9 +193,22 @@ def dispose(self): ( connection, { - "confirmation": ADOPTION_CONFIRMATION, "backup_id": "verified-backup", - "expected_question_rows": 9, }, ) ] + + +@pytest.mark.parametrize("state", [DatabaseState.UNVERSIONED, DatabaseState.INVALID]) +def test_upgrade_refuses_databases_without_valid_revision_history(monkeypatch, state): + class FakeConnection: + def scalar(self, *args, **kwargs): + return 1 + + def execute(self, *args, **kwargs): + pass + + monkeypatch.setattr(migration, "database_state", lambda connection: state) + + with pytest.raises(RuntimeError, match="no valid Alembic revision"): + upgrade_database(FakeConnection(), backup_id="verified-backup") diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 407b6132c..37d965d6c 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -50,29 +50,14 @@ def test_reusable_alembic_check_uses_the_installed_python_environment(): assert "uv run" not in workflow -def test_adoption_workflow_is_manual_explicit_and_backup_first(): - workflow = _workflow("adopt-v1-cloud-sql.yml") - - assert "workflow_dispatch:" in workflow - assert "ADOPT-eafc2a547a4e" in workflow - assert "environment: production-database" in workflow - assert "protection" not in workflow.lower() - assert workflow.index("Verify legacy schema") < workflow.index( - "Create Cloud SQL backup" - ) - assert workflow.index("Create Cloud SQL backup") < workflow.index( - "Stamp baseline and upgrade" - ) - assert "--mode adopt" in workflow - - def test_release_migration_fails_closed_and_gates_both_staging_deploys(): workflow = _workflow("push.yml") assert "migrate-v1-cloud-sql:" in workflow assert "environment: production-database" in workflow + assert "--mode state" in workflow assert "--mode upgrade" in workflow - assert "--mode adopt" not in workflow + assert "--mode verify-head" in workflow assert "database is unversioned" in workflow app_engine_job = workflow[workflow.index(" deploy-staging:") :] @@ -88,18 +73,20 @@ def test_release_migration_fails_closed_and_gates_both_staging_deploys(): assert "migrate-v1-cloud-sql" in cloud_run_job -def test_cloud_sql_workflows_use_oidc_and_separate_database_credentials(): - workflows = _workflow("adopt-v1-cloud-sql.yml") + _workflow("push.yml") - - assert "google-github-actions/auth@v2" in workflows - assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in workflows - assert "policyengine-api-prod-db-readonly-password" in workflows - assert "policyengine-api-prod-db-migration-password" in workflows - assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in workflows - assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in workflows +def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): + workflow = _workflow("push.yml") + migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] + migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + + assert "google-github-actions/auth@v2" in migration_job + assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in migration_job + assert "policyengine-api-prod-db-readonly-password" in migration_job + assert "policyengine-api-prod-db-migration-password" in migration_job + assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in migration_job + assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in migration_job assert ( "POLICYENGINE_DB_PASSWORD: ${{ secrets.POLICYENGINE_DB_PASSWORD }}" - not in _workflow("adopt-v1-cloud-sql.yml") + not in migration_job ) From 446ec19944a400402464460134c3600e3390db1e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:28:55 +0300 Subject: [PATCH 85/89] fix: make production migration job executable --- .github/workflows/push.yml | 8 ++++---- tests/unit/test_alembic_workflows.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index c5c06b61c..228cebd22 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -146,11 +146,11 @@ jobs: POLICYENGINE_DB_READONLY_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-readonly-password --project policyengine-api)" POLICYENGINE_DB_MIGRATION_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-migration-password --project policyengine-api)" export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD - uv run python scripts/write_v1_database_urls.py + python scripts/write_v1_database_urls.py - name: Inspect current database revision id: schema run: | - state="$(uv run python scripts/v1_database_migration.py --mode state)" + state="$(python scripts/v1_database_migration.py --mode state)" echo "state=${state}" >> "${GITHUB_OUTPUT}" - name: Refuse unversioned or invalid schema if: steps.schema.outputs.state == 'unversioned' || steps.schema.outputs.state == 'invalid' @@ -164,11 +164,11 @@ jobs: - name: Upgrade pending v1 migrations if: steps.schema.outputs.state == 'pending' run: >- - uv run python scripts/v1_database_migration.py + python scripts/v1_database_migration.py --mode upgrade --backup-id "${{ steps.backup.outputs.backup_id }}" - name: Verify v1 database at head - run: uv run python scripts/v1_database_migration.py --mode verify-head + run: python scripts/v1_database_migration.py --mode verify-head - name: Stop Cloud SQL Auth Proxy if: always() run: bash .github/scripts/stop_cloud_sql_proxy.sh diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 37d965d6c..0e348dc09 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -30,6 +30,16 @@ def test_push_always_runs_lint_and_alembic_qualification_before_versioning(): assert "github.repository == 'PolicyEngine/policyengine-uk'" not in workflow +def test_release_migration_uses_the_installed_python_environment(): + workflow = _workflow("push.yml") + migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] + migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + + assert "python scripts/write_v1_database_urls.py" in migration_job + assert "python scripts/v1_database_migration.py" in migration_job + assert "uv run" not in migration_job + + def test_reusable_alembic_check_uses_only_disposable_mysql(): workflow = _workflow("alembic-v1-check.yml") From c0edde0d29e773379ac19db0bd7af481c5ff43f0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:32:09 +0300 Subject: [PATCH 86/89] refactor: extract long CI shell blocks --- .github/scripts/detect_v1_alembic_changes.sh | 12 ++ .github/scripts/prepare_v1_database_urls.sh | 17 +++ .github/workflows/pr.yml | 6 +- .github/workflows/push.yml | 6 +- scripts/v1_alembic_changes.py | 2 + tests/unit/data/test_v1_database_migration.py | 2 + tests/unit/test_alembic_workflows.py | 128 +++++++++++++++++- 7 files changed, 159 insertions(+), 14 deletions(-) create mode 100644 .github/scripts/detect_v1_alembic_changes.sh create mode 100644 .github/scripts/prepare_v1_database_urls.sh diff --git a/.github/scripts/detect_v1_alembic_changes.sh b/.github/scripts/detect_v1_alembic_changes.sh new file mode 100644 index 000000000..7269c7fa6 --- /dev/null +++ b/.github/scripts/detect_v1_alembic_changes.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" + +if [[ "$#" -ne 2 ]]; then + echo "usage: detect_v1_alembic_changes.sh BASE HEAD" >&2 + exit 2 +fi + +python scripts/v1_alembic_changes.py "$1" "$2" >>"${GITHUB_OUTPUT}" diff --git a/.github/scripts/prepare_v1_database_urls.sh b/.github/scripts/prepare_v1_database_urls.sh new file mode 100644 index 000000000..cb7525c51 --- /dev/null +++ b/.github/scripts/prepare_v1_database_urls.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash + +set -euo pipefail + +POLICYENGINE_DB_READONLY_PASSWORD="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-readonly-password \ + --project policyengine-api +)" +POLICYENGINE_DB_MIGRATION_PASSWORD="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-migration-password \ + --project policyengine-api +)" +export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD + +python scripts/write_v1_database_urls.py diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 5c42be20f..51fbc310d 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -63,11 +63,7 @@ jobs: fetch-depth: 0 - name: Detect relevant changes id: changes - run: >- - python scripts/v1_alembic_changes.py - "${{ github.event.pull_request.base.sha }}" - "${{ github.event.pull_request.head.sha }}" - >> "${GITHUB_OUTPUT}" + run: bash .github/scripts/detect_v1_alembic_changes.sh "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" alembic-v1-check: name: Alembic v1 qualification diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 228cebd22..65cb54fd4 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -142,11 +142,7 @@ jobs: - name: Start Cloud SQL Auth Proxy run: bash .github/scripts/start_cloud_sql_proxy.sh - name: Prepare masked database URLs - run: | - POLICYENGINE_DB_READONLY_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-readonly-password --project policyengine-api)" - POLICYENGINE_DB_MIGRATION_PASSWORD="$(gcloud secrets versions access latest --secret policyengine-api-prod-db-migration-password --project policyengine-api)" - export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD - python scripts/write_v1_database_urls.py + run: bash .github/scripts/prepare_v1_database_urls.sh - name: Inspect current database revision id: schema run: | diff --git a/scripts/v1_alembic_changes.py b/scripts/v1_alembic_changes.py index 335a46b63..b73f120be 100644 --- a/scripts/v1_alembic_changes.py +++ b/scripts/v1_alembic_changes.py @@ -10,6 +10,8 @@ EXACT_PATHS = frozenset( { "alembic-v1.ini", + ".github/scripts/detect_v1_alembic_changes.sh", + ".github/scripts/prepare_v1_database_urls.sh", ".github/workflows/alembic-v1-check.yml", "docs/engineering/skills/alembic-migrations.md", "policyengine_api/data/v1_models.py", diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py index e4d654ed9..34dc41b8a 100644 --- a/tests/unit/data/test_v1_database_migration.py +++ b/tests/unit/data/test_v1_database_migration.py @@ -26,6 +26,8 @@ "scripts/v1_database_migration.py", "tests/integration/test_alembic_mysql_lifecycle.py", "tests/integration/test_v1_schema_metadata_compatibility.py", + ".github/scripts/detect_v1_alembic_changes.sh", + ".github/scripts/prepare_v1_database_urls.sh", ".github/workflows/alembic-v1-check.yml", "docs/engineering/skills/alembic-migrations.md", "pyproject.toml", diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 0e348dc09..a800b400b 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -1,6 +1,8 @@ from __future__ import annotations +import os from pathlib import Path +import subprocess REPO = Path(__file__).resolve().parents[2] @@ -10,11 +12,51 @@ def _workflow(name: str) -> str: return (REPO / ".github" / "workflows" / name).read_text(encoding="utf-8") +def _long_inline_run_blocks() -> list[str]: + offenders = [] + block_markers = {"|", "|-", "|+", ">", ">-", ">+"} + for path in sorted((REPO / ".github" / "workflows").glob("*.y*ml")): + lines = path.read_text(encoding="utf-8").splitlines() + line_index = 0 + while line_index < len(lines): + line = lines[line_index] + stripped = line.lstrip() + indentation = len(line) - len(stripped) + if not ( + stripped.startswith("run:") + and stripped.removeprefix("run:").strip() in block_markers + ): + line_index += 1 + continue + + body_start = line_index + 1 + line_index = body_start + while line_index < len(lines): + candidate = lines[line_index] + candidate_indentation = len(candidate) - len(candidate.lstrip()) + if candidate.strip() and candidate_indentation <= indentation: + break + line_index += 1 + + substantive_lines = [ + candidate + for candidate in lines[body_start:line_index] + if candidate.strip() and not candidate.lstrip().startswith("#") + ] + if len(substantive_lines) > 3: + offenders.append(f"{path.name}:{body_start}") + return offenders + + +def test_workflows_do_not_inline_long_shell_programs(): + assert _long_inline_run_blocks() == [] + + def test_pr_runs_reusable_alembic_check_only_for_relevant_changes(): workflow = _workflow("pr.yml") assert "detect-v1-alembic-changes:" in workflow - assert "python scripts/v1_alembic_changes.py" in workflow + assert "bash .github/scripts/detect_v1_alembic_changes.sh" in workflow assert "alembic-v1-check:" in workflow assert "needs.detect-v1-alembic-changes.outputs.changed == 'true'" in workflow assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow @@ -35,7 +77,7 @@ def test_release_migration_uses_the_installed_python_environment(): migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] migration_job = migration_job[: migration_job.index("\n deploy-staging:")] - assert "python scripts/write_v1_database_urls.py" in migration_job + assert "bash .github/scripts/prepare_v1_database_urls.sh" in migration_job assert "python scripts/v1_database_migration.py" in migration_job assert "uv run" not in migration_job @@ -90,8 +132,7 @@ def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): assert "google-github-actions/auth@v2" in migration_job assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in migration_job - assert "policyengine-api-prod-db-readonly-password" in migration_job - assert "policyengine-api-prod-db-migration-password" in migration_job + assert "prepare_v1_database_urls.sh" in migration_job assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in migration_job assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in migration_job assert ( @@ -100,6 +141,85 @@ def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): ) +def test_change_detector_script_appends_python_result_to_github_output(tmp_path): + bin_path = tmp_path / "bin" + bin_path.mkdir() + python_path = bin_path / "python" + python_path.write_text( + "#!/usr/bin/env bash\n" + 'printf "%s\\n" "$*" > "${ARGS_FILE}"\n' + 'printf "changed=true\\n"\n', + encoding="utf-8", + ) + python_path.chmod(0o755) + github_output = tmp_path / "github-output" + args_file = tmp_path / "args" + result = subprocess.run( + [ + "bash", + ".github/scripts/detect_v1_alembic_changes.sh", + "base-sha", + "head-sha", + ], + cwd=REPO, + env={ + **os.environ, + "ARGS_FILE": str(args_file), + "GITHUB_OUTPUT": str(github_output), + "PATH": f"{bin_path}:{os.environ['PATH']}", + }, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + assert github_output.read_text(encoding="utf-8") == "changed=true\n" + assert args_file.read_text(encoding="utf-8") == ( + "scripts/v1_alembic_changes.py base-sha head-sha\n" + ) + + +def test_database_url_script_fetches_both_gcp_secrets_and_writes_urls(tmp_path): + bin_path = tmp_path / "bin" + bin_path.mkdir() + gcloud_path = bin_path / "gcloud" + gcloud_path.write_text( + "#!/usr/bin/env bash\n" + 'if [[ "$*" == *"readonly-password"* ]]; then\n' + ' printf "reader-p@ss\\n"\n' + "else\n" + ' printf "migrator-p@ss\\n"\n' + "fi\n", + encoding="utf-8", + ) + gcloud_path.chmod(0o755) + github_env = tmp_path / "github-env" + result = subprocess.run( + ["bash", ".github/scripts/prepare_v1_database_urls.sh"], + cwd=REPO, + env={ + **os.environ, + "GITHUB_ENV": str(github_env), + "PATH": f"{bin_path}:{REPO / '.venv' / 'bin'}:{os.environ['PATH']}", + }, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + environment = github_env.read_text(encoding="utf-8") + assert ( + "STAGE7_EXISTING_DATABASE_URL=mysql+pymysql://" + "policyengine_schema_reader:reader-p%40ss@127.0.0.1:3307/policyengine" + ) in environment + assert ( + "ALEMBIC_DATABASE_URL=mysql+pymysql://" + "policyengine_schema_migrator:migrator-p%40ss@127.0.0.1:3307/policyengine" + ) in environment + + def test_backup_helper_recovers_and_verifies_the_created_backup_id(): script = (REPO / ".github" / "scripts" / "create_cloud_sql_backup.sh").read_text( encoding="utf-8" From 6de700729995b6a5205b01892f7d29f3ff00c8e0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:36:49 +0300 Subject: [PATCH 87/89] docs: clarify disposable Alembic credentials --- .github/workflows/alembic-v1-check.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/alembic-v1-check.yml b/.github/workflows/alembic-v1-check.yml index 529d8b6fd..6e9094b16 100644 --- a/.github/workflows/alembic-v1-check.yml +++ b/.github/workflows/alembic-v1-check.yml @@ -11,6 +11,8 @@ jobs: services: mysql: image: mysql:8.4 + # This database exists only for this job, so these test credentials are + # intentionally non-secret and safe to expose in the workflow. env: MYSQL_ROOT_PASSWORD: policyengine_test MYSQL_DATABASE: policyengine_alembic_test From 496451f0261c633116030d9dd2bf5fc2ae701983 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:42:38 +0300 Subject: [PATCH 88/89] ci: run Alembic checks on every PR --- .github/scripts/detect_v1_alembic_changes.sh | 12 ---- .github/workflows/pr.yml | 16 ----- scripts/v1_alembic_changes.py | 66 ------------------- tests/unit/data/test_v1_database_migration.py | 36 ---------- tests/unit/test_alembic_workflows.py | 46 +------------ 5 files changed, 3 insertions(+), 173 deletions(-) delete mode 100644 .github/scripts/detect_v1_alembic_changes.sh delete mode 100644 scripts/v1_alembic_changes.py diff --git a/.github/scripts/detect_v1_alembic_changes.sh b/.github/scripts/detect_v1_alembic_changes.sh deleted file mode 100644 index 7269c7fa6..000000000 --- a/.github/scripts/detect_v1_alembic_changes.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/usr/bin/env bash - -set -euo pipefail - -: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" - -if [[ "$#" -ne 2 ]]; then - echo "usage: detect_v1_alembic_changes.sh BASE HEAD" >&2 - exit 2 -fi - -python scripts/v1_alembic_changes.py "$1" "$2" >>"${GITHUB_OUTPUT}" diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 51fbc310d..83d23f875 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -51,24 +51,8 @@ jobs: - name: Run quality guards run: python scripts/run_quality_guards.py - detect-v1-alembic-changes: - name: Detect v1 Alembic changes - runs-on: ubuntu-latest - outputs: - changed: ${{ steps.changes.outputs.changed }} - steps: - - name: Checkout repo - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - name: Detect relevant changes - id: changes - run: bash .github/scripts/detect_v1_alembic_changes.sh "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" - alembic-v1-check: name: Alembic v1 qualification - needs: detect-v1-alembic-changes - if: needs.detect-v1-alembic-changes.outputs.changed == 'true' uses: ./.github/workflows/alembic-v1-check.yml check-changelog: name: Check changelog fragment diff --git a/scripts/v1_alembic_changes.py b/scripts/v1_alembic_changes.py deleted file mode 100644 index b73f120be..000000000 --- a/scripts/v1_alembic_changes.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Report whether a pull request changes the API v1 Alembic surface.""" - -from __future__ import annotations - -import subprocess -import sys -from pathlib import PurePosixPath - - -EXACT_PATHS = frozenset( - { - "alembic-v1.ini", - ".github/scripts/detect_v1_alembic_changes.sh", - ".github/scripts/prepare_v1_database_urls.sh", - ".github/workflows/alembic-v1-check.yml", - "docs/engineering/skills/alembic-migrations.md", - "policyengine_api/data/v1_models.py", - "pyproject.toml", - "scripts/v1_alembic_changes.py", - "scripts/v1_database_migration.py", - "scripts/write_v1_database_urls.py", - "tests/integration/test_alembic_mysql_lifecycle.py", - "tests/integration/test_v1_schema_metadata_compatibility.py", - "tests/unit/test_alembic_workflows.py", - "uv.lock", - } -) -PATH_PREFIXES = ( - "migrations/v1/", - "tests/unit/data/test_alembic_", - "tests/unit/data/test_v1_database_migration", -) - - -def is_v1_alembic_path(path: str) -> bool: - """Return whether *path* can change v1 migration behavior.""" - - normalized = PurePosixPath(path).as_posix().removeprefix("./") - return normalized in EXACT_PATHS or normalized.startswith(PATH_PREFIXES) - - -def changed_paths(base: str, head: str) -> tuple[str, ...]: - """Return repository paths changed between two git revisions.""" - - result = subprocess.run( - ["git", "diff", "--name-only", base, head, "--"], - check=True, - capture_output=True, - text=True, - ) - return tuple(path for path in result.stdout.splitlines() if path) - - -def main(argv: list[str] | None = None) -> int: - args = list(sys.argv[1:] if argv is None else argv) - if len(args) != 2: - print("usage: v1_alembic_changes.py BASE HEAD", file=sys.stderr) - return 2 - - changed = any(is_v1_alembic_path(path) for path in changed_paths(*args)) - print(f"changed={str(changed).lower()}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py index 34dc41b8a..e93bed60d 100644 --- a/tests/unit/data/test_v1_database_migration.py +++ b/tests/unit/data/test_v1_database_migration.py @@ -6,7 +6,6 @@ from sqlalchemy import Column, Integer, MetaData, String, Table, text import scripts.v1_database_migration as migration -from scripts.v1_alembic_changes import is_v1_alembic_path from scripts.v1_database_migration import ( DatabaseState, build_database_url, @@ -16,41 +15,6 @@ ) -@pytest.mark.parametrize( - "path", - [ - "alembic-v1.ini", - "migrations/v1/env.py", - "migrations/v1/versions/123_add_column.py", - "policyengine_api/data/v1_models.py", - "scripts/v1_database_migration.py", - "tests/integration/test_alembic_mysql_lifecycle.py", - "tests/integration/test_v1_schema_metadata_compatibility.py", - ".github/scripts/detect_v1_alembic_changes.sh", - ".github/scripts/prepare_v1_database_urls.sh", - ".github/workflows/alembic-v1-check.yml", - "docs/engineering/skills/alembic-migrations.md", - "pyproject.toml", - "uv.lock", - ], -) -def test_v1_alembic_change_paths_trigger_qualification(path): - assert is_v1_alembic_path(path) - - -@pytest.mark.parametrize( - "path", - [ - "policyengine_api/routes/household_routes.py", - "tests/unit/routes/test_household_routes.py", - "docs/migration/cloud-run-operations.md", - ".github/workflows/push.yml", - ], -) -def test_unrelated_paths_do_not_trigger_v1_alembic_qualification(path): - assert not is_v1_alembic_path(path) - - def test_database_url_percent_encodes_credentials_without_losing_driver(): url = build_database_url( username="schema reader", diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index a800b400b..9f3b8b008 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -52,14 +52,13 @@ def test_workflows_do_not_inline_long_shell_programs(): assert _long_inline_run_blocks() == [] -def test_pr_runs_reusable_alembic_check_only_for_relevant_changes(): +def test_pr_always_runs_reusable_alembic_check(): workflow = _workflow("pr.yml") - assert "detect-v1-alembic-changes:" in workflow - assert "bash .github/scripts/detect_v1_alembic_changes.sh" in workflow assert "alembic-v1-check:" in workflow - assert "needs.detect-v1-alembic-changes.outputs.changed == 'true'" in workflow assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow + assert "detect-v1-alembic-changes:" not in workflow + assert "needs.detect-v1-alembic-changes" not in workflow def test_push_always_runs_lint_and_alembic_qualification_before_versioning(): @@ -141,45 +140,6 @@ def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): ) -def test_change_detector_script_appends_python_result_to_github_output(tmp_path): - bin_path = tmp_path / "bin" - bin_path.mkdir() - python_path = bin_path / "python" - python_path.write_text( - "#!/usr/bin/env bash\n" - 'printf "%s\\n" "$*" > "${ARGS_FILE}"\n' - 'printf "changed=true\\n"\n', - encoding="utf-8", - ) - python_path.chmod(0o755) - github_output = tmp_path / "github-output" - args_file = tmp_path / "args" - result = subprocess.run( - [ - "bash", - ".github/scripts/detect_v1_alembic_changes.sh", - "base-sha", - "head-sha", - ], - cwd=REPO, - env={ - **os.environ, - "ARGS_FILE": str(args_file), - "GITHUB_OUTPUT": str(github_output), - "PATH": f"{bin_path}:{os.environ['PATH']}", - }, - capture_output=True, - text=True, - check=False, - ) - - assert result.returncode == 0, result.stderr - assert github_output.read_text(encoding="utf-8") == "changed=true\n" - assert args_file.read_text(encoding="utf-8") == ( - "scripts/v1_alembic_changes.py base-sha head-sha\n" - ) - - def test_database_url_script_fetches_both_gcp_secrets_and_writes_urls(tmp_path): bin_path = tmp_path / "bin" bin_path.mkdir() From 3b3daebd18cd70c83086a44498b535dfd2a277ef Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:06:48 +0300 Subject: [PATCH 89/89] ci: keep v1 migration credentials process-local --- .github/scripts/create_cloud_sql_backup.sh | 5 +- .github/scripts/migrate_v1_cloud_sql.sh | 53 ++++++++ .github/scripts/prepare_v1_database_urls.sh | 17 --- .github/workflows/push.yml | 26 +--- scripts/v1_database_migration.py | 52 ++++++-- scripts/write_v1_database_urls.py | 46 ------- tests/unit/data/test_v1_database_migration.py | 56 +++++++- tests/unit/test_alembic_workflows.py | 125 ++++++++++++++---- 8 files changed, 249 insertions(+), 131 deletions(-) create mode 100644 .github/scripts/migrate_v1_cloud_sql.sh delete mode 100644 .github/scripts/prepare_v1_database_urls.sh delete mode 100644 scripts/write_v1_database_urls.py diff --git a/.github/scripts/create_cloud_sql_backup.sh b/.github/scripts/create_cloud_sql_backup.sh index 5c0e18ccd..44407c234 100644 --- a/.github/scripts/create_cloud_sql_backup.sh +++ b/.github/scripts/create_cloud_sql_backup.sh @@ -3,7 +3,6 @@ set -euo pipefail : "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" -: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" instance_id="${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME##*:}" description="policyengine-api-v1-alembic-${GITHUB_SHA:-manual}-${GITHUB_RUN_ID:-local}-${GITHUB_RUN_ATTEMPT:-1}" @@ -11,7 +10,7 @@ gcloud sql backups create \ --project policyengine-api \ --instance "${instance_id}" \ --description "${description}" \ - --quiet + --quiet >&2 # `gcloud sql backups create` waits for completion but does not consistently # emit the created resource with value-format output. Recover the ID from the @@ -31,4 +30,4 @@ if [[ -z "${backup_id}" ]]; then exit 1 fi -printf 'backup_id=%s\n' "${backup_id}" >>"${GITHUB_OUTPUT}" +printf '%s\n' "${backup_id}" diff --git a/.github/scripts/migrate_v1_cloud_sql.sh b/.github/scripts/migrate_v1_cloud_sql.sh new file mode 100644 index 000000000..77d2bb23d --- /dev/null +++ b/.github/scripts/migrate_v1_cloud_sql.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" + +readonly_password="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-readonly-password \ + --project policyengine-api +)" +migration_password="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-migration-password \ + --project policyengine-api +)" + +if [[ -z "${readonly_password}" || -z "${migration_password}" ]]; then + echo "Cloud SQL database credentials must not be empty." >&2 + exit 1 +fi + +printf '::add-mask::%s\n' "${readonly_password}" +printf '::add-mask::%s\n' "${migration_password}" +export POLICYENGINE_DB_READONLY_PASSWORD="${readonly_password}" +export POLICYENGINE_DB_MIGRATION_PASSWORD="${migration_password}" + +# Never allow a job-level URL to bypass the credentials fetched for this release. +unset STAGE7_EXISTING_DATABASE_URL ALEMBIC_DATABASE_URL + +database_state="$(python scripts/v1_database_migration.py --mode state)" +echo "Detected v1 database state: ${database_state}" + +case "${database_state}" in + head) + ;; + pending) + backup_id="$(bash .github/scripts/create_cloud_sql_backup.sh)" + python scripts/v1_database_migration.py \ + --mode upgrade \ + --backup-id "${backup_id}" + ;; + unversioned | invalid) + echo "database is unversioned or has invalid Alembic state; automatic baseline stamping is disabled and manual recovery is required" >&2 + exit 1 + ;; + *) + echo "Unrecognized v1 database state: ${database_state}" >&2 + exit 1 + ;; +esac + +python scripts/v1_database_migration.py --mode verify-head diff --git a/.github/scripts/prepare_v1_database_urls.sh b/.github/scripts/prepare_v1_database_urls.sh deleted file mode 100644 index cb7525c51..000000000 --- a/.github/scripts/prepare_v1_database_urls.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/usr/bin/env bash - -set -euo pipefail - -POLICYENGINE_DB_READONLY_PASSWORD="$( - gcloud secrets versions access latest \ - --secret policyengine-api-prod-db-readonly-password \ - --project policyengine-api -)" -POLICYENGINE_DB_MIGRATION_PASSWORD="$( - gcloud secrets versions access latest \ - --secret policyengine-api-prod-db-migration-password \ - --project policyengine-api -)" -export POLICYENGINE_DB_READONLY_PASSWORD POLICYENGINE_DB_MIGRATION_PASSWORD - -python scripts/write_v1_database_urls.py diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 65cb54fd4..cab5ff0ff 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -141,30 +141,8 @@ jobs: run: make install - name: Start Cloud SQL Auth Proxy run: bash .github/scripts/start_cloud_sql_proxy.sh - - name: Prepare masked database URLs - run: bash .github/scripts/prepare_v1_database_urls.sh - - name: Inspect current database revision - id: schema - run: | - state="$(python scripts/v1_database_migration.py --mode state)" - echo "state=${state}" >> "${GITHUB_OUTPUT}" - - name: Refuse unversioned or invalid schema - if: steps.schema.outputs.state == 'unversioned' || steps.schema.outputs.state == 'invalid' - run: | - echo "database is unversioned or has invalid Alembic state; automatic baseline stamping is disabled and manual recovery is required" >&2 - exit 1 - - name: Create Cloud SQL backup - if: steps.schema.outputs.state == 'pending' - id: backup - run: bash .github/scripts/create_cloud_sql_backup.sh - - name: Upgrade pending v1 migrations - if: steps.schema.outputs.state == 'pending' - run: >- - python scripts/v1_database_migration.py - --mode upgrade - --backup-id "${{ steps.backup.outputs.backup_id }}" - - name: Verify v1 database at head - run: python scripts/v1_database_migration.py --mode verify-head + - name: Upgrade and verify v1 database + run: bash .github/scripts/migrate_v1_cloud_sql.sh - name: Stop Cloud SQL Auth Proxy if: always() run: bash .github/scripts/stop_cloud_sql_proxy.sh diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index bcd18a088..80611b9e8 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -6,7 +6,6 @@ from enum import StrEnum import os from typing import Any -from urllib.parse import quote_plus from alembic import command from alembic.autogenerate import compare_metadata @@ -14,6 +13,7 @@ from alembic.migration import MigrationContext from alembic.script import ScriptDirectory from sqlalchemy import Connection, create_engine, inspect, text +from sqlalchemy.engine import URL from sqlalchemy.pool import NullPool from policyengine_api.constants import REPO @@ -22,6 +22,11 @@ ALEMBIC_CONFIG = REPO / "alembic-v1.ini" MIGRATION_LOCK_NAME = "policyengine-api-v1-alembic" +DATABASE_NAME = "policyengine" +READONLY_USER = "policyengine_schema_reader" +MIGRATION_USER = "policyengine_schema_migrator" +PROXY_HOST = "127.0.0.1" +PROXY_PORT = 3307 class DatabaseState(StrEnum): @@ -38,12 +43,16 @@ def build_database_url( host: str, port: int, database: str, -) -> str: - """Build an encoded PyMySQL URL without logging its credentials.""" - - return ( - f"mysql+pymysql://{quote_plus(username)}:{quote_plus(password)}@" - f"{host}:{port}/{quote_plus(database)}" +) -> URL: + """Build a SQLAlchemy URL without stringifying its credentials.""" + + return URL.create( + drivername="mysql+pymysql", + username=username, + password=password, + host=host, + port=port, + database=database, ) @@ -165,16 +174,31 @@ def upgrade_database(connection: Connection, *, backup_id: str) -> None: _release_lock(connection) -def _database_url(mode: str) -> str: - env_name = ( +def _database_target(mode: str) -> str | URL: + url_env_name = ( "STAGE7_EXISTING_DATABASE_URL" if mode in {"verify-head", "state"} else "ALEMBIC_DATABASE_URL" ) - try: - return os.environ[env_name] - except KeyError as error: - raise RuntimeError(f"{env_name} is required") from error + if explicit_url := os.environ.get(url_env_name): + return explicit_url + + readonly = mode in {"verify-head", "state"} + password_env_name = ( + "POLICYENGINE_DB_READONLY_PASSWORD" + if readonly + else "POLICYENGINE_DB_MIGRATION_PASSWORD" + ) + password = os.environ.get(password_env_name) + if not password: + raise RuntimeError(f"{url_env_name} or {password_env_name} is required") + return build_database_url( + username=READONLY_USER if readonly else MIGRATION_USER, + password=password, + host=PROXY_HOST, + port=PROXY_PORT, + database=DATABASE_NAME, + ) def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: @@ -190,7 +214,7 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) - engine = create_engine(_database_url(args.mode), poolclass=NullPool) + engine = create_engine(_database_target(args.mode), poolclass=NullPool) try: connection_context = ( engine.begin() if args.mode == "upgrade" else engine.connect() diff --git a/scripts/write_v1_database_urls.py b/scripts/write_v1_database_urls.py deleted file mode 100644 index 752b06317..000000000 --- a/scripts/write_v1_database_urls.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Write masked local-proxy database URLs into a GitHub Actions env file.""" - -from __future__ import annotations - -import os -from pathlib import Path - -from scripts.v1_database_migration import build_database_url - - -DATABASE_NAME = "policyengine" -READONLY_USER = "policyengine_schema_reader" -MIGRATION_USER = "policyengine_schema_migrator" -PROXY_HOST = "127.0.0.1" -PROXY_PORT = 3307 - - -def main() -> int: - output_path = Path(os.environ["GITHUB_ENV"]) - readonly_url = build_database_url( - username=READONLY_USER, - password=os.environ["POLICYENGINE_DB_READONLY_PASSWORD"], - host=PROXY_HOST, - port=PROXY_PORT, - database=DATABASE_NAME, - ) - migration_url = build_database_url( - username=MIGRATION_USER, - password=os.environ["POLICYENGINE_DB_MIGRATION_PASSWORD"], - host=PROXY_HOST, - port=PROXY_PORT, - database=DATABASE_NAME, - ) - - # URL-encoded passwords can differ from the exact GitHub secret value, so - # mask the complete derived URLs before any later command can mention them. - print(f"::add-mask::{readonly_url}") - print(f"::add-mask::{migration_url}") - with output_path.open("a", encoding="utf-8") as output: - output.write(f"STAGE7_EXISTING_DATABASE_URL={readonly_url}\n") - output.write(f"ALEMBIC_DATABASE_URL={migration_url}\n") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py index e93bed60d..b4271cb16 100644 --- a/tests/unit/data/test_v1_database_migration.py +++ b/tests/unit/data/test_v1_database_migration.py @@ -4,6 +4,7 @@ import pytest from sqlalchemy import Column, Integer, MetaData, String, Table, text +from sqlalchemy.engine import URL import scripts.v1_database_migration as migration from scripts.v1_database_migration import ( @@ -15,7 +16,7 @@ ) -def test_database_url_percent_encodes_credentials_without_losing_driver(): +def test_database_url_uses_sqlalchemy_url_without_stringifying_credentials(): url = build_database_url( username="schema reader", password="p@ss:/word", @@ -24,9 +25,56 @@ def test_database_url_percent_encodes_credentials_without_losing_driver(): database="policyengine", ) - assert url == ( - "mysql+pymysql://schema+reader:p%40ss%3A%2Fword@127.0.0.1:3307/policyengine" - ) + assert isinstance(url, URL) + assert url.username == "schema reader" + assert url.password == "p@ss:/word" + assert url.host == "127.0.0.1" + assert url.port == 3307 + assert url.database == "policyengine" + + +@pytest.mark.parametrize( + ("mode", "password_name", "expected_user"), + [ + ("state", "POLICYENGINE_DB_READONLY_PASSWORD", "policyengine_schema_reader"), + ( + "verify-head", + "POLICYENGINE_DB_READONLY_PASSWORD", + "policyengine_schema_reader", + ), + ( + "upgrade", + "POLICYENGINE_DB_MIGRATION_PASSWORD", + "policyengine_schema_migrator", + ), + ], +) +def test_database_target_builds_in_memory_url_for_each_database_role( + monkeypatch, + mode, + password_name, + expected_user, +): + monkeypatch.delenv("STAGE7_EXISTING_DATABASE_URL", raising=False) + monkeypatch.delenv("ALEMBIC_DATABASE_URL", raising=False) + monkeypatch.setenv(password_name, "p@ssword") + + target = migration._database_target(mode) + + assert isinstance(target, URL) + assert target.username == expected_user + assert target.password == "p@ssword" + assert target.host == "127.0.0.1" + assert target.port == 3307 + assert target.database == "policyengine" + + +def test_database_target_preserves_explicit_url_override(monkeypatch): + explicit_url = "mysql+pymysql://root:test@127.0.0.1/test" + monkeypatch.setenv("ALEMBIC_DATABASE_URL", explicit_url) + monkeypatch.delenv("POLICYENGINE_DB_MIGRATION_PASSWORD", raising=False) + + assert migration._database_target("upgrade") == explicit_url def test_database_state_is_unversioned_without_a_version_table(): diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 9f3b8b008..deab1d7be 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -4,6 +4,8 @@ from pathlib import Path import subprocess +import pytest + REPO = Path(__file__).resolve().parents[2] @@ -75,10 +77,13 @@ def test_release_migration_uses_the_installed_python_environment(): workflow = _workflow("push.yml") migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") - assert "bash .github/scripts/prepare_v1_database_urls.sh" in migration_job - assert "python scripts/v1_database_migration.py" in migration_job - assert "uv run" not in migration_job + assert "bash .github/scripts/migrate_v1_cloud_sql.sh" in migration_job + assert "python scripts/v1_database_migration.py" in orchestration_script + assert "uv run" not in orchestration_script def test_reusable_alembic_check_uses_only_disposable_mysql(): @@ -103,13 +108,17 @@ def test_reusable_alembic_check_uses_the_installed_python_environment(): def test_release_migration_fails_closed_and_gates_both_staging_deploys(): workflow = _workflow("push.yml") + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") assert "migrate-v1-cloud-sql:" in workflow assert "environment: production-database" in workflow - assert "--mode state" in workflow - assert "--mode upgrade" in workflow - assert "--mode verify-head" in workflow - assert "database is unversioned" in workflow + assert "--mode state" in orchestration_script + assert "--mode upgrade" in orchestration_script + assert "--mode verify-head" in orchestration_script + assert "database is unversioned" in orchestration_script + assert "create_cloud_sql_backup.sh" in orchestration_script app_engine_job = workflow[workflow.index(" deploy-staging:") :] app_engine_job = app_engine_job[ @@ -128,10 +137,15 @@ def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): workflow = _workflow("push.yml") migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") assert "google-github-actions/auth@v2" in migration_job assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in migration_job - assert "prepare_v1_database_urls.sh" in migration_job + assert "migrate_v1_cloud_sql.sh" in migration_job + assert "prepare_v1_database_urls.sh" not in migration_job + assert "GITHUB_ENV" not in orchestration_script assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in migration_job assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in migration_job assert ( @@ -140,44 +154,108 @@ def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): ) -def test_database_url_script_fetches_both_gcp_secrets_and_writes_urls(tmp_path): +def _write_fake_migration_commands(tmp_path: Path) -> tuple[Path, Path]: bin_path = tmp_path / "bin" bin_path.mkdir() gcloud_path = bin_path / "gcloud" gcloud_path.write_text( "#!/usr/bin/env bash\n" + 'printf "%s\\n" "$*" >> "${GCLOUD_CALLS}"\n' 'if [[ "$*" == *"readonly-password"* ]]; then\n' ' printf "reader-p@ss\\n"\n' - "else\n" + 'elif [[ "$*" == *"migration-password"* ]]; then\n' ' printf "migrator-p@ss\\n"\n' + 'elif [[ "$*" == *"sql backups list"* ]]; then\n' + ' printf "backup-123\\n"\n' "fi\n", encoding="utf-8", ) gcloud_path.chmod(0o755) - github_env = tmp_path / "github-env" + + python_path = bin_path / "python" + python_path.write_text( + "#!/usr/bin/env bash\n" + ': "${POLICYENGINE_DB_READONLY_PASSWORD:?}"\n' + ': "${POLICYENGINE_DB_MIGRATION_PASSWORD:?}"\n' + 'printf "%s\\n" "$*" >> "${PYTHON_CALLS}"\n' + 'if [[ "$*" == *"--mode state"* ]]; then\n' + ' printf "%s\\n" "${DATABASE_STATE}"\n' + "fi\n", + encoding="utf-8", + ) + python_path.chmod(0o755) + return bin_path, gcloud_path + + +def _run_migration_orchestrator(tmp_path: Path, database_state: str): + bin_path, _ = _write_fake_migration_commands(tmp_path) + python_calls = tmp_path / "python-calls" + gcloud_calls = tmp_path / "gcloud-calls" result = subprocess.run( - ["bash", ".github/scripts/prepare_v1_database_urls.sh"], + ["bash", ".github/scripts/migrate_v1_cloud_sql.sh"], cwd=REPO, env={ **os.environ, - "GITHUB_ENV": str(github_env), - "PATH": f"{bin_path}:{REPO / '.venv' / 'bin'}:{os.environ['PATH']}", + "DATABASE_STATE": database_state, + "GCLOUD_CALLS": str(gcloud_calls), + "PYTHON_CALLS": str(python_calls), + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": "project:region:instance", + "PATH": f"{bin_path}:{os.environ['PATH']}", }, capture_output=True, text=True, check=False, ) + return result, python_calls, gcloud_calls + + +def test_migration_orchestrator_keeps_credentials_local_and_upgrades_pending_schema( + tmp_path, +): + result, python_calls, gcloud_calls = _run_migration_orchestrator( + tmp_path, "pending" + ) assert result.returncode == 0, result.stderr - environment = github_env.read_text(encoding="utf-8") - assert ( - "STAGE7_EXISTING_DATABASE_URL=mysql+pymysql://" - "policyengine_schema_reader:reader-p%40ss@127.0.0.1:3307/policyengine" - ) in environment - assert ( - "ALEMBIC_DATABASE_URL=mysql+pymysql://" - "policyengine_schema_migrator:migrator-p%40ss@127.0.0.1:3307/policyengine" - ) in environment + assert "::add-mask::reader-p@ss" in result.stdout + assert "::add-mask::migrator-p@ss" in result.stdout + assert "STAGE7_EXISTING_DATABASE_URL" not in result.stdout + assert "ALEMBIC_DATABASE_URL" not in result.stdout + calls = python_calls.read_text(encoding="utf-8").splitlines() + assert calls == [ + "scripts/v1_database_migration.py --mode state", + "scripts/v1_database_migration.py --mode upgrade --backup-id backup-123", + "scripts/v1_database_migration.py --mode verify-head", + ] + assert "sql backups create" in gcloud_calls.read_text(encoding="utf-8") + + +def test_migration_orchestrator_skips_backup_and_upgrade_at_head(tmp_path): + result, python_calls, gcloud_calls = _run_migration_orchestrator(tmp_path, "head") + + assert result.returncode == 0, result.stderr + assert python_calls.read_text(encoding="utf-8").splitlines() == [ + "scripts/v1_database_migration.py --mode state", + "scripts/v1_database_migration.py --mode verify-head", + ] + assert "sql backups create" not in gcloud_calls.read_text(encoding="utf-8") + + +@pytest.mark.parametrize("database_state", ["unversioned", "invalid"]) +def test_migration_orchestrator_refuses_unsafe_database_states( + tmp_path, + database_state, +): + result, python_calls, gcloud_calls = _run_migration_orchestrator( + tmp_path, database_state + ) + + assert result.returncode != 0 + assert "automatic baseline stamping is disabled" in result.stderr + assert python_calls.read_text(encoding="utf-8").splitlines() == [ + "scripts/v1_database_migration.py --mode state" + ] + assert "sql backups create" not in gcloud_calls.read_text(encoding="utf-8") def test_backup_helper_recovers_and_verifies_the_created_backup_id(): @@ -189,6 +267,7 @@ def test_backup_helper_recovers_and_verifies_the_created_backup_id(): assert "gcloud sql backups create" in script assert "gcloud sql backups list" in script assert "status=SUCCESSFUL" in script + assert "GITHUB_OUTPUT" not in script assert script.index("gcloud sql backups create") < script.index( "gcloud sql backups list" )