From c4009dccbdec2e5c16d7f3fdbb0dabe7142d9387 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:24:05 +0400 Subject: [PATCH 1/5] Add simulation.simulations: what to run, independent of model versions A simulation is one policy applied to one population: a household, or an economy named by a region code (state/ca) with an optional dataset name, where a null dataset name means the region's default dataset in whichever catalog version a run uses. A null policy means current law. The row holds no catalog version, so the same simulation can run again on a newer one; each run resolves the region and dataset in its own version (item 1.06). Revision 0002 creates the table with: - the population checked against the kind (a household simulation has a household and no region or dataset; an economy simulation has a region and no household); - foreign keys into the Public API's tables in the same database, all ON DELETE RESTRICT: (policy_id, country) to policies (id, country_id), (household_id, country) to households (id, country_id), and dynamic_id to dynamics, with an index on each referencing column; - a unique legacy_simulation_id for simulations migrated from the Public API's v1 database, and a unique (id, country, kind) for the reports' foreign keys; - no UPDATE privilege for the runtime role, since definitions never change. Creating the foreign keys needs the operator's REFERENCES grant on the Public API tables; a test shows the chain fails without it and leaves nothing applied. SimulationStore creates, reads and finds simulations by v1 ID, and create_migrated is idempotent by legacy_simulation_id for the one-time migration (item 2.05). Importing the package now registers every table, so the migration environment always compares the whole schema. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 16 +- .../public_api.py | 45 +++ .../schema.py | 2 +- .../simulations.py | 300 ++++++++++++++++++ .../tests/conftest.py | 40 ++- .../tests/public_api_support.py | 126 ++++++++ .../tests/simulation_support.py | 43 +++ .../tests/test_definitions.py | 31 ++ .../tests/test_migration_cycle.py | 33 ++ .../tests/test_schema.py | 13 + .../tests/test_simulations.py | 212 +++++++++++++ .../versions/0002_simulations_and_reports.py | 121 +++++++ 12 files changed, 978 insertions(+), 4 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py create mode 100644 libs/policyengine-simulation-persistence/tests/public_api_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/simulation_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_definitions.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_simulations.py create mode 100644 migrations/versions/0002_simulations_and_reports.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index fb098c8ed..000d62dae 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -1,4 +1,8 @@ -"""Tables and connections for the Simulation Entrypoint's ``simulation`` schema.""" +"""Tables and connections for the Simulation Entrypoint's ``simulation`` schema. + +Importing the package registers every table in ``simulation_metadata``, which +the migration environment compares with the database. +""" from policyengine_simulation_persistence.engine import ( create_simulation_engine, @@ -12,12 +16,22 @@ SIMULATION_SCHEMA, simulation_metadata, ) +from policyengine_simulation_persistence.simulations import ( + NewSimulation, + SimulationKind, + SimulationRecord, + SimulationStore, +) __all__ = [ "ALEMBIC_VERSION_TABLE", "HEAD_REVISION", "RUNTIME_ROLE", "SIMULATION_SCHEMA", + "NewSimulation", + "SimulationKind", + "SimulationRecord", + "SimulationStore", "create_simulation_engine", "normalized_database_url", "read_schema_revision", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py new file mode 100644 index 000000000..1da0d5b0d --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py @@ -0,0 +1,45 @@ +"""The Public API tables that rows in the simulation schema refer to. + +The Public API owns these tables and creates and changes them with its own +migration chain, in the same database as the ``simulation`` schema. They are +described here only so the simulation tables can declare foreign keys to them: +each lists just the key columns those foreign keys use, in a metadata object of +its own, so the simulation chain never compares, creates or changes them. + +The tables have no schema here because PostgreSQL reports a foreign key into +``public`` without a schema when ``public`` is on the search path; naming one +would make every schema comparison report the foreign keys as changed. The +migrations name ``public`` explicitly. + +Every foreign key into these tables is ``ON DELETE RESTRICT``: the Public API +cannot delete a row a simulation uses. Creating the foreign keys needs +``REFERENCES`` on each table for the migrator role; the runtime role needs no +privilege on them, because PostgreSQL runs a foreign key's check as the owner +of the referenced table. +""" + +from __future__ import annotations + +from sqlalchemy import Column, MetaData, String, Table, Uuid + +public_api_metadata = MetaData() + +policies = Table( + "policies", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("country_id", String(2), nullable=False), +) +households = Table( + "households", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("country_id", String(2), nullable=False), +) +dynamics = Table( + "dynamics", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), +) + +REFERENCED_PUBLIC_API_TABLES = tuple(sorted(public_api_metadata.tables)) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py index a7cd53c18..7e4080638 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py @@ -17,7 +17,7 @@ # The newest revision in migrations/versions. Each migration PR updates it, and # tests/test_schema.py fails when it differs from the chain's head. The # readiness check reports it next to the database's revision. -HEAD_REVISION = "0001" +HEAD_REVISION = "0002" NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py new file mode 100644 index 000000000..31db413ad --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py @@ -0,0 +1,300 @@ +"""``simulation.simulations``: what to run, independent of model versions. + +A simulation is one policy applied to one population: a household, or an +economy named by a region code (``state/ca``) and an optional dataset name. It +is a definition, not an execution, and never changes once created. A null +policy means current law, and a null dataset name means the region's default +dataset in whichever catalog version a run uses, so the same simulation can +run again on a newer catalog version. + +Foreign keys into the Public API's tables (``public_api``) tie the policy and +the household to the simulation's country, and name an optional dynamic. The +region and dataset are not foreign keys here, because their rows belong to one +catalog version; each run resolves them in its own. + +A simulation migrated from the Public API's v1 database keeps its v1 +simulation ID in ``legacy_simulation_id``. That database is separate, so the v1 +ID is a plain unique column, not a foreign key. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, fields +from enum import StrEnum +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from sqlalchemy import ( + BigInteger, + CheckConstraint, + Column, + DateTime, + ForeignKeyConstraint, + Index, + String, + Table, + Text, + UniqueConstraint, + Uuid, + func, + select, +) +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence import public_api +from policyengine_simulation_persistence.schema import simulation_metadata + +if TYPE_CHECKING: + from datetime import datetime + + from sqlalchemy import ColumnElement, Connection, Engine, RowMapping + +COUNTRY_PATTERN = re.compile(r"[a-z]{2}") +# Region codes and dataset names are the Public API's catalog keys: up to 255 +# characters, with no whitespace. +CATALOG_KEY_PATTERN = re.compile(r"\S{1,255}") +CATALOG_KEY_LENGTH = 255 + + +class SimulationKind(StrEnum): + """What a simulation runs on: one household or a whole economy.""" + + ECONOMY = "economy" + HOUSEHOLD = "household" + + +_KIND_VALUES = ", ".join(f"'{kind.value}'" for kind in SimulationKind) + +simulations = Table( + "simulations", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("kind", Text, nullable=False), + Column("country", Text, nullable=False), + Column("policy_id", Uuid(as_uuid=True), nullable=True), + Column("dynamic_id", Uuid(as_uuid=True), nullable=True), + Column("household_id", Uuid(as_uuid=True), nullable=True), + Column("region_code", String(CATALOG_KEY_LENGTH), nullable=True), + Column("dataset_name", String(CATALOG_KEY_LENGTH), nullable=True), + Column("legacy_simulation_id", BigInteger, nullable=True, unique=True), + Column( + "created_at", + DateTime(timezone=True), + nullable=False, + server_default=func.now(), + ), + CheckConstraint(f"kind IN ({_KIND_VALUES})", name="kind"), + CheckConstraint("country ~ '^[a-z]{2}$'", name="country"), + CheckConstraint( + "(kind = 'household' AND household_id IS NOT NULL " + "AND region_code IS NULL AND dataset_name IS NULL) " + "OR (kind = 'economy' AND household_id IS NULL " + "AND region_code IS NOT NULL)", + name="population_matches_kind", + ), + CheckConstraint( + "region_code IS NULL OR region_code ~ '^\\S+$'", name="region_code" + ), + CheckConstraint( + "dataset_name IS NULL OR dataset_name ~ '^\\S+$'", name="dataset_name" + ), + CheckConstraint( + "legacy_simulation_id IS NULL OR legacy_simulation_id > 0", + name="legacy_simulation_id", + ), + # Reports refer to a simulation together with its country and kind, so a + # report can only compare simulations of its own country and kind. + UniqueConstraint("id", "country", "kind"), + ForeignKeyConstraint( + ["policy_id", "country"], + [public_api.policies.c.id, public_api.policies.c.country_id], + name="fk_simulations_policy_country", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["dynamic_id"], + [public_api.dynamics.c.id], + name="fk_simulations_dynamic", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["household_id", "country"], + [public_api.households.c.id, public_api.households.c.country_id], + name="fk_simulations_household_country", + ondelete="RESTRICT", + ), +) +# A Public API delete of a policy, household or dynamic checks these for +# references. +Index("ix_simulations_policy_id", simulations.c.policy_id) +Index("ix_simulations_household_id", simulations.c.household_id) +Index("ix_simulations_dynamic_id", simulations.c.dynamic_id) + + +def _check_catalog_key(field_name: str, value: str | None) -> None: + if value is not None and not CATALOG_KEY_PATTERN.fullmatch(value): + msg = f"{field_name} must be 1 to 255 characters with no whitespace" + raise ValueError(msg) + + +def check_legacy_id(field_name: str, value: int | None) -> None: + """Reject a v1 ID that is not a positive integer.""" + if value is not None and value <= 0: + msg = f"{field_name} must be a positive v1 ID" + raise ValueError(msg) + + +def check_country(country: str) -> None: + """Reject a country that is not a lowercase two-letter code.""" + if not COUNTRY_PATTERN.fullmatch(country): + msg = "country must be a lowercase two-letter code" + raise ValueError(msg) + + +def insert_or_get_existing( + connection: Connection, + table: Table, + values: dict[str, Any], + unique_column: Column[Any], +) -> RowMapping: + """Insert a row, or return the row that already has its unique value.""" + created = ( + connection.execute( + postgres_insert(table) + .values(**values) + .on_conflict_do_nothing(index_elements=[unique_column]) + .returning(*table.c) + ) + .mappings() + .first() + ) + if created is not None: + return created + return ( + connection.execute( + select(table).where(unique_column == values[unique_column.name]) + ) + .mappings() + .one() + ) + + +@dataclass(frozen=True) +class NewSimulation: + """A simulation to create. + + The database checks that the population matches the kind and that the + policy, household and dynamic exist for the country; the fields are + checked here first. + """ + + kind: SimulationKind + country: str + policy_id: UUID | None = None + dynamic_id: UUID | None = None + household_id: UUID | None = None + region_code: str | None = None + dataset_name: str | None = None + legacy_simulation_id: int | None = None + + def __post_init__(self) -> None: + """Reject malformed fields without touching the database.""" + check_country(self.country) + _check_catalog_key("region_code", self.region_code) + _check_catalog_key("dataset_name", self.dataset_name) + check_legacy_id("legacy_simulation_id", self.legacy_simulation_id) + + +@dataclass(frozen=True) +class SimulationRecord: + """One ``simulation.simulations`` row.""" + + id: UUID + kind: SimulationKind + country: str + policy_id: UUID | None + dynamic_id: UUID | None + household_id: UUID | None + region_code: str | None + dataset_name: str | None + legacy_simulation_id: int | None + created_at: datetime + + @classmethod + def from_row(cls, row: RowMapping) -> SimulationRecord: + """Build the record from a selected ``simulations`` row.""" + values = {field.name: row[field.name] for field in fields(cls)} + values["kind"] = SimulationKind(values["kind"]) + return cls(**values) + + +def _new_simulation_values(new_simulation: NewSimulation) -> dict[str, Any]: + return { + "id": uuid4(), + **{ + field.name: getattr(new_simulation, field.name) + for field in fields(new_simulation) + }, + "kind": new_simulation.kind.value, + } + + +class SimulationStore: + """Typed SQLAlchemy Core reads and inserts for simulation rows.""" + + def __init__(self, engine: Engine) -> None: + """Use ``engine``, which connects as the runtime role.""" + self.engine = engine + + def create(self, new_simulation: NewSimulation) -> SimulationRecord: + """Insert one simulation and return it.""" + with self.engine.begin() as connection: + row = ( + connection.execute( + postgres_insert(simulations) + .values(**_new_simulation_values(new_simulation)) + .returning(*simulations.c) + ) + .mappings() + .one() + ) + return SimulationRecord.from_row(row) + + def create_migrated(self, new_simulation: NewSimulation) -> SimulationRecord: + """Create a simulation for a v1 simulation, or return the existing one. + + The v1 simulation ID must be set. An existing row is returned as it is, + even if ``new_simulation`` describes it differently, so running the + one-time migration again creates nothing new. + """ + if new_simulation.legacy_simulation_id is None: + msg = "a migrated simulation needs its legacy_simulation_id" + raise ValueError(msg) + with self.engine.begin() as connection: + row = insert_or_get_existing( + connection, + simulations, + _new_simulation_values(new_simulation), + simulations.c.legacy_simulation_id, + ) + return SimulationRecord.from_row(row) + + def get(self, simulation_id: UUID) -> SimulationRecord | None: + """Return the simulation with this ID, or None.""" + return self._one(simulations.c.id == simulation_id) + + def get_by_legacy_simulation_id( + self, legacy_simulation_id: int + ) -> SimulationRecord | None: + """Return the simulation migrated from this v1 simulation, or None.""" + return self._one(simulations.c.legacy_simulation_id == legacy_simulation_id) + + def _one(self, condition: ColumnElement[bool]) -> SimulationRecord | None: + with self.engine.connect() as connection: + row = ( + connection.execute(select(simulations).where(condition)) + .mappings() + .first() + ) + return None if row is None else SimulationRecord.from_row(row) diff --git a/libs/policyengine-simulation-persistence/tests/conftest.py b/libs/policyengine-simulation-persistence/tests/conftest.py index 1726049dc..12a4821b6 100644 --- a/libs/policyengine-simulation-persistence/tests/conftest.py +++ b/libs/policyengine-simulation-persistence/tests/conftest.py @@ -15,12 +15,19 @@ from uuid import uuid4 import pytest +from alembic import command from migration_support import ( MIGRATION_DATABASE_URL_VARIABLE, MIGRATOR_ROLE, TEST_ROLE_PASSWORD, + alembic_config, ) from policyengine_simulation_persistence.schema import RUNTIME_ROLE +from public_api_support import ( + PublicApiRows, + create_public_api_tables, + insert_public_api_rows, +) from sqlalchemy import create_engine, text from sqlalchemy.engine import URL, make_url @@ -70,14 +77,27 @@ def postgres_server_url() -> URL: return make_url(configured).set(drivername="postgresql+psycopg") +@pytest.fixture +def references_granted() -> bool: + """Whether the migrator gets REFERENCES on the Public API tables. + + The operator grants it in production; a test overrides this fixture to + leave it out. + """ + return True + + @pytest.fixture def disposable_database( postgres_server_url: URL, monkeypatch: pytest.MonkeyPatch, + references_granted, ) -> Iterator[DisposableDatabase]: - """Create an empty database owned by the migrator, and drop it afterwards. + """Create a database owned by the migrator, and drop it afterwards. - ``SIMULATION_MIGRATION_DATABASE_URL`` points at it, as in the deploy job. + It holds only the Public API tables the simulation schema refers to + (``public_api_support``). ``SIMULATION_MIGRATION_DATABASE_URL`` points at + it, as in the deploy job. """ database_name = f"simulation_test_{uuid4().hex[:12]}" server_engine = create_engine(postgres_server_url, isolation_level="AUTOCOMMIT") @@ -106,6 +126,9 @@ def disposable_database( ), superuser_url=_canonical(postgres_server_url.set(database=database_name)), ) + create_public_api_tables( + database.superuser_url, grant_references=references_granted + ) monkeypatch.setenv(MIGRATION_DATABASE_URL_VARIABLE, database.migrator_url) yield database finally: @@ -114,3 +137,16 @@ def disposable_database( text(f'DROP DATABASE IF EXISTS "{database_name}" WITH (FORCE)') ) server_engine.dispose() + + +@pytest.fixture +def migrated_database(disposable_database: DisposableDatabase) -> DisposableDatabase: + """Return a disposable database with the whole chain applied.""" + command.upgrade(alembic_config(), "head") + return disposable_database + + +@pytest.fixture +def public_api_rows(migrated_database: DisposableDatabase) -> PublicApiRows: + """Insert one row of each kind the simulation tables refer to.""" + return insert_public_api_rows(migrated_database.superuser_url) diff --git a/libs/policyengine-simulation-persistence/tests/public_api_support.py b/libs/policyengine-simulation-persistence/tests/public_api_support.py new file mode 100644 index 000000000..cf144c68d --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/public_api_support.py @@ -0,0 +1,126 @@ +"""The Public API tables the simulation schema refers to, reduced to their keys. + +In production these tables belong to the Public API's own migration chain, in +the same database as the simulation schema. The tests create only the columns +and unique constraints the simulation schema's foreign keys use, with the +Public API's constraint names, owned by the superuser, and grant the migrator +``REFERENCES`` on them, as the operator does in production. The runtime role +gets no privilege on them. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +from migration_support import MIGRATOR_ROLE +from sqlalchemy import create_engine, text + +if TYPE_CHECKING: + from sqlalchemy import Engine + +PUBLIC_API_TABLE_STATEMENTS = ( + "CREATE TABLE public.policies (" + " id uuid PRIMARY KEY," + " country_id varchar(2) NOT NULL," + " CONSTRAINT uq_policies_id_country UNIQUE (id, country_id))", + "CREATE TABLE public.households (" + " id uuid PRIMARY KEY," + " country_id varchar(2) NOT NULL," + " CONSTRAINT uq_households_id_country UNIQUE (id, country_id))", + "CREATE TABLE public.dynamics (id uuid PRIMARY KEY)", +) +REFERENCED_TABLES = ("policies", "households", "dynamics") + + +def superuser_engine(superuser_url: str, **options: object) -> Engine: + """Return an engine that connects to the test database as the superuser.""" + return create_engine( + superuser_url.replace("postgresql://", "postgresql+psycopg://", 1), + **options, + ) + + +def create_public_api_tables(superuser_url: str, *, grant_references: bool) -> None: + """Create the referenced Public API tables in a new test database.""" + engine = superuser_engine(superuser_url, isolation_level="AUTOCOMMIT") + grant = ( + "GRANT REFERENCES ON " + + ", ".join(f"public.{table}" for table in REFERENCED_TABLES) + + f" TO {MIGRATOR_ROLE}" + ) + statements = ( + (*PUBLIC_API_TABLE_STATEMENTS, grant) + if grant_references + else PUBLIC_API_TABLE_STATEMENTS + ) + try: + with engine.connect() as connection: + for statement in statements: + connection.execute(text(statement)) + finally: + engine.dispose() + + +@dataclass(frozen=True) +class PublicApiRows: + """Public API rows a simulation can refer to.""" + + us_policy_id: UUID + uk_policy_id: UUID + us_household_id: UUID + uk_household_id: UUID + dynamic_id: UUID + + +def insert_public_api_rows(superuser_url: str) -> PublicApiRows: + """Insert rows into the referenced Public API tables, as their owner.""" + rows = PublicApiRows( + us_policy_id=uuid4(), + uk_policy_id=uuid4(), + us_household_id=uuid4(), + uk_household_id=uuid4(), + dynamic_id=uuid4(), + ) + statements = ( + ( + "INSERT INTO public.policies (id, country_id) VALUES " + "(:us_policy, 'us'), (:uk_policy, 'uk')", + {"us_policy": rows.us_policy_id, "uk_policy": rows.uk_policy_id}, + ), + ( + "INSERT INTO public.households (id, country_id) VALUES " + "(:us_household, 'us'), (:uk_household, 'uk')", + { + "us_household": rows.us_household_id, + "uk_household": rows.uk_household_id, + }, + ), + ( + "INSERT INTO public.dynamics (id) VALUES (:dynamic)", + {"dynamic": rows.dynamic_id}, + ), + ) + engine = superuser_engine(superuser_url) + try: + with engine.begin() as connection: + for statement, parameters in statements: + connection.execute(text(statement), parameters) + finally: + engine.dispose() + return rows + + +def delete_public_api_row(superuser_url: str, table_name: str, row_id: UUID) -> None: + """Delete one Public API row as its owner, as the Public API would.""" + assert table_name in REFERENCED_TABLES + engine = superuser_engine(superuser_url) + try: + with engine.begin() as connection: + connection.execute( + text(f"DELETE FROM public.{table_name} WHERE id = :row_id"), # noqa: S608 # table names come from REFERENCED_TABLES + {"row_id": row_id}, + ) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/simulation_support.py b/libs/policyengine-simulation-persistence/tests/simulation_support.py new file mode 100644 index 000000000..030ee2e93 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/simulation_support.py @@ -0,0 +1,43 @@ +"""Builders shared by the simulation and report tests.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence import ( + NewSimulation, + SimulationKind, + SimulationStore, + create_simulation_engine, +) + +if TYPE_CHECKING: + from public_api_support import PublicApiRows + + +def economy_simulation(rows: PublicApiRows, **overrides: object) -> NewSimulation: + """Return a California economy simulation of the US reform.""" + values: dict[str, object] = { + "kind": SimulationKind.ECONOMY, + "country": "us", + "policy_id": rows.us_policy_id, + "region_code": "state/ca", + } + values.update(overrides) + return NewSimulation(**values) # type: ignore[arg-type] + + +def household_simulation(rows: PublicApiRows, **overrides: object) -> NewSimulation: + """Return a current-law simulation of the US household.""" + values: dict[str, object] = { + "kind": SimulationKind.HOUSEHOLD, + "country": "us", + "household_id": rows.us_household_id, + } + values.update(overrides) + return NewSimulation(**values) # type: ignore[arg-type] + + +def simulation_store(database) -> SimulationStore: + """Return a store that connects as the runtime role.""" + return SimulationStore(create_simulation_engine(database.runtime_url)) diff --git a/libs/policyengine-simulation-persistence/tests/test_definitions.py b/libs/policyengine-simulation-persistence/tests/test_definitions.py new file mode 100644 index 000000000..f9b843a5c --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_definitions.py @@ -0,0 +1,31 @@ +"""New simulation and report definitions are checked before any write.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_persistence import NewSimulation, SimulationKind + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("country", "USA"), + ("country", "Us"), + ("region_code", ""), + ("region_code", " state/ca"), + ("region_code", "x" * 256), + ("dataset_name", ""), + ("dataset_name", "populace us"), + ("legacy_simulation_id", 0), + ], +) +def test_malformed_simulation_definitions_are_rejected(field_name, value) -> None: + values = { + "kind": SimulationKind.ECONOMY, + "country": "us", + "region_code": "state/ca", + field_name: value, + } + + with pytest.raises(ValueError, match=field_name): + NewSimulation(**values) diff --git a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py index d033ec4aa..828c1a386 100644 --- a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py +++ b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py @@ -19,6 +19,7 @@ SIMULATION_SCHEMA, create_simulation_engine, ) +from simulation_support import economy_simulation, simulation_store from sqlalchemy import create_engine, inspect, text from sqlalchemy.exc import OperationalError, ProgrammingError @@ -196,3 +197,35 @@ def test_a_migration_waiting_for_a_lock_gives_up_instead_of_blocking( finally: holder.dispose() assert waited < 30 + + +@pytest.mark.parametrize("references_granted", [False]) +def test_the_chain_needs_references_on_the_public_api_tables( + disposable_database, +) -> None: + """Without the operator's REFERENCES grant, no foreign key can be created.""" + with pytest.raises(ProgrammingError, match="permission denied for table"): + command.upgrade(alembic_config(), "head") + + assert applied_revision(disposable_database.migrator_url) is None + + +@pytest.mark.parametrize("table_name", ["simulations"]) +def test_the_runtime_role_cannot_change_or_delete_a_definition( + migrated_database, public_api_rows, table_name +) -> None: + store = simulation_store(migrated_database) + store.create(economy_simulation(public_api_rows)) + engine = create_simulation_engine(migrated_database.runtime_url) + try: + for statement in ( + f"UPDATE simulation.{table_name} SET created_at = now()", # noqa: S608 # fixed table names + f"DELETE FROM simulation.{table_name}", # noqa: S608 # fixed table names + ): + with ( + pytest.raises(ProgrammingError, match="permission denied"), + engine.begin() as connection, + ): + connection.execute(text(statement)) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index d42b8bf4d..2125d6c14 100644 --- a/libs/policyengine-simulation-persistence/tests/test_schema.py +++ b/libs/policyengine-simulation-persistence/tests/test_schema.py @@ -2,6 +2,10 @@ from __future__ import annotations +import importlib +import pkgutil + +import policyengine_simulation_persistence import pytest from alembic import command from alembic.script import ScriptDirectory @@ -104,3 +108,12 @@ def test_migrations_refuse_offline_sql_generation(monkeypatch) -> None: with pytest.raises(RuntimeError, match="online"): command.upgrade(alembic_config(), "head", sql=True) + + +def test_importing_the_package_registers_every_table() -> None: + package_tables = set(simulation_metadata.tables) + for module in pkgutil.iter_modules(policyengine_simulation_persistence.__path__): + importlib.import_module(f"policyengine_simulation_persistence.{module.name}") + + assert set(simulation_metadata.tables) == package_tables + assert "simulation.simulations" in package_tables diff --git a/libs/policyengine-simulation-persistence/tests/test_simulations.py b/libs/policyengine-simulation-persistence/tests/test_simulations.py new file mode 100644 index 000000000..4b15d4d97 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_simulations.py @@ -0,0 +1,212 @@ +"""simulation.simulations: what to run, independent of model versions.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from policyengine_simulation_persistence import ( + NewSimulation, + SimulationKind, + SimulationRecord, +) +from public_api_support import delete_public_api_row +from simulation_support import ( + economy_simulation, + household_simulation, + simulation_store, +) +from sqlalchemy.exc import IntegrityError + +pytestmark = pytest.mark.postgres + + +def test_an_economy_simulation_names_a_region_by_code( + migrated_database, public_api_rows +) -> None: + store = simulation_store(migrated_database) + + created = store.create( + economy_simulation(public_api_rows, dataset_name="populace_us_2024") + ) + + assert isinstance(created, SimulationRecord) + assert created.kind is SimulationKind.ECONOMY + assert created.country == "us" + assert created.policy_id == public_api_rows.us_policy_id + assert created.region_code == "state/ca" + assert created.dataset_name == "populace_us_2024" + assert created.household_id is None + assert created.created_at.tzinfo is not None + assert store.get(created.id) == created + + +def test_a_null_dataset_name_means_the_region_default( + migrated_database, public_api_rows +) -> None: + created = simulation_store(migrated_database).create( + economy_simulation(public_api_rows) + ) + + assert created.dataset_name is None + + +def test_current_law_has_no_policy(migrated_database, public_api_rows) -> None: + created = simulation_store(migrated_database).create( + economy_simulation(public_api_rows, policy_id=None) + ) + + assert created.policy_id is None + + +def test_a_household_simulation_names_its_household_and_dynamic( + migrated_database, public_api_rows +) -> None: + created = simulation_store(migrated_database).create( + household_simulation( + public_api_rows, + policy_id=public_api_rows.us_policy_id, + dynamic_id=public_api_rows.dynamic_id, + ) + ) + + assert created.kind is SimulationKind.HOUSEHOLD + assert created.household_id == public_api_rows.us_household_id + assert created.dynamic_id == public_api_rows.dynamic_id + assert created.region_code is None + + +def _mismatched_population(rows, case: str) -> NewSimulation: + household = {"kind": SimulationKind.HOUSEHOLD, "household_id": rows.us_household_id} + economy = {"kind": SimulationKind.ECONOMY, "region_code": "state/ca"} + values = { + "household-without-household": {**household, "household_id": None}, + "household-with-region": {**household, "region_code": "state/ca"}, + "household-with-dataset": {**household, "dataset_name": "populace_us_2024"}, + "economy-without-region": {**economy, "region_code": None}, + "economy-with-household": {**economy, "household_id": rows.us_household_id}, + }[case] + return NewSimulation(country="us", **values) + + +@pytest.mark.parametrize( + "case", + [ + "household-without-household", + "household-with-region", + "household-with-dataset", + "economy-without-region", + "economy-with-household", + ], +) +def test_the_population_must_match_the_kind( + migrated_database, public_api_rows, case +) -> None: + with pytest.raises(IntegrityError, match="population_matches_kind"): + simulation_store(migrated_database).create( + _mismatched_population(public_api_rows, case) + ) + + +@pytest.mark.parametrize( + ("kind", "reference"), + [ + (SimulationKind.ECONOMY, "policy"), + (SimulationKind.HOUSEHOLD, "household"), + ], +) +def test_a_policy_or_household_must_be_of_the_simulation_country( + migrated_database, public_api_rows, kind, reference +) -> None: + if reference == "policy": + new_simulation = economy_simulation( + public_api_rows, country="uk", region_code="country/wales" + ) + constraint = "fk_simulations_policy_country" + else: + new_simulation = household_simulation(public_api_rows, country="uk") + constraint = "fk_simulations_household_country" + assert new_simulation.kind is kind + + with pytest.raises(IntegrityError, match=constraint): + simulation_store(migrated_database).create(new_simulation) + + +def test_an_unknown_dynamic_is_rejected(migrated_database, public_api_rows) -> None: + with pytest.raises(IntegrityError, match="fk_simulations_dynamic"): + simulation_store(migrated_database).create( + economy_simulation(public_api_rows, dynamic_id=uuid4()) + ) + + +def _simulation_using(rows, table_name: str) -> NewSimulation: + if table_name == "households": + return household_simulation(rows) + if table_name == "dynamics": + return economy_simulation(rows, dynamic_id=rows.dynamic_id) + return economy_simulation(rows) + + +@pytest.mark.parametrize( + ("table_name", "row_attribute"), + [ + ("policies", "us_policy_id"), + ("households", "us_household_id"), + ("dynamics", "dynamic_id"), + ], +) +def test_the_public_api_cannot_delete_a_row_a_simulation_uses( + migrated_database, public_api_rows, table_name, row_attribute +) -> None: + simulation_store(migrated_database).create( + _simulation_using(public_api_rows, table_name) + ) + + with pytest.raises(IntegrityError, match="violates foreign key constraint"): + delete_public_api_row( + migrated_database.superuser_url, + table_name, + getattr(public_api_rows, row_attribute), + ) + + +def test_a_migrated_simulation_is_created_once( + migrated_database, public_api_rows +) -> None: + store = simulation_store(migrated_database) + first = store.create_migrated( + economy_simulation(public_api_rows, legacy_simulation_id=4021) + ) + + again = store.create_migrated( + economy_simulation( + public_api_rows, legacy_simulation_id=4021, region_code="state/ny" + ) + ) + + assert again == first + assert store.get_by_legacy_simulation_id(4021) == first + assert store.get_by_legacy_simulation_id(4022) is None + + +def test_a_migrated_simulation_needs_its_legacy_id( + migrated_database, public_api_rows +) -> None: + with pytest.raises(ValueError, match="legacy_simulation_id"): + simulation_store(migrated_database).create_migrated( + economy_simulation(public_api_rows) + ) + + +def test_two_simulations_cannot_share_a_legacy_id( + migrated_database, public_api_rows +) -> None: + store = simulation_store(migrated_database) + store.create(economy_simulation(public_api_rows, legacy_simulation_id=7)) + + with pytest.raises(IntegrityError, match="legacy_simulation_id"): + store.create(economy_simulation(public_api_rows, legacy_simulation_id=7)) + + +def test_an_unknown_simulation_reads_as_none(migrated_database) -> None: + assert simulation_store(migrated_database).get(uuid4()) is None diff --git a/migrations/versions/0002_simulations_and_reports.py b/migrations/versions/0002_simulations_and_reports.py new file mode 100644 index 000000000..d5231da6e --- /dev/null +++ b/migrations/versions/0002_simulations_and_reports.py @@ -0,0 +1,121 @@ +"""Add the simulation definitions table. + +``simulations`` records what each simulation runs, independent of model +versions: its kind, country, policy (null for current law), dynamic, and either +a household or a region code with an optional dataset name. Foreign keys into +the Public API's policies, households and dynamics in the same database are +``ON DELETE RESTRICT`` and need ``REFERENCES`` on those tables for the migrator +role. Definitions never change, so the runtime role may read and insert them +but not update them. Nothing writes the table yet. + +Revision ID: 0002 +Revises: 0001 +Create Date: 2026-10-22 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "0002" +down_revision = "0001" +branch_labels = None +depends_on = None + +SCHEMA = "simulation" +PUBLIC_API_SCHEMA = "public" + + +def _public_api(table_and_column: str) -> str: + return f"{PUBLIC_API_SCHEMA}.{table_and_column}" + + +def _create_simulations() -> None: + op.create_table( + "simulations", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("country", sa.Text(), nullable=False), + sa.Column("policy_id", sa.Uuid(), nullable=True), + sa.Column("dynamic_id", sa.Uuid(), nullable=True), + sa.Column("household_id", sa.Uuid(), nullable=True), + sa.Column("region_code", sa.String(length=255), nullable=True), + sa.Column("dataset_name", sa.String(length=255), nullable=True), + sa.Column("legacy_simulation_id", sa.BigInteger(), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "kind IN ('economy', 'household')", + name="ck_simulations_kind", + ), + sa.CheckConstraint("country ~ '^[a-z]{2}$'", name="ck_simulations_country"), + sa.CheckConstraint( + "(kind = 'household' AND household_id IS NOT NULL " + "AND region_code IS NULL AND dataset_name IS NULL) " + "OR (kind = 'economy' AND household_id IS NULL " + "AND region_code IS NOT NULL)", + name="ck_simulations_population_matches_kind", + ), + sa.CheckConstraint( + "region_code IS NULL OR region_code ~ '^\\S+$'", + name="ck_simulations_region_code", + ), + sa.CheckConstraint( + "dataset_name IS NULL OR dataset_name ~ '^\\S+$'", + name="ck_simulations_dataset_name", + ), + sa.CheckConstraint( + "legacy_simulation_id IS NULL OR legacy_simulation_id > 0", + name="ck_simulations_legacy_simulation_id", + ), + sa.ForeignKeyConstraint( + ["policy_id", "country"], + [_public_api("policies.id"), _public_api("policies.country_id")], + name="fk_simulations_policy_country", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["dynamic_id"], + [_public_api("dynamics.id")], + name="fk_simulations_dynamic", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["household_id", "country"], + [_public_api("households.id"), _public_api("households.country_id")], + name="fk_simulations_household_country", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_simulations"), + sa.UniqueConstraint( + "legacy_simulation_id", name="uq_simulations_legacy_simulation_id" + ), + sa.UniqueConstraint( + "id", "country", "kind", name="uq_simulations_id_country_kind" + ), + schema=SCHEMA, + ) + for column in ("policy_id", "household_id", "dynamic_id"): + op.create_index( + f"ix_simulations_{column}", "simulations", [column], schema=SCHEMA + ) + + +def upgrade() -> None: + """Create the simulations table; the runtime role may not update it.""" + _create_simulations() + op.execute("REVOKE UPDATE ON simulation.simulations FROM simulation_runtime") + + +def downgrade() -> None: + """Drop the table ``upgrade`` created, with its indexes and privileges.""" + for column in ("dynamic_id", "household_id", "policy_id"): + op.drop_index( + f"ix_simulations_{column}", table_name="simulations", schema=SCHEMA + ) + op.drop_table("simulations", schema=SCHEMA) From 3a2dd8f6f1e77bf5930cfc5fa5b6a04112c850c3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:26:53 +0400 Subject: [PATCH 2/5] Add simulation.reports: which two simulations a report compares A report is a definition: its baseline simulation, its reform simulation (required for an economy report; a household report may have only a baseline), the year (1900 to 2200) and its options as a JSON object. Composite foreign keys (simulation id, country, kind) into simulation.simulations make both simulations share the report's country and kind; they are ON DELETE RESTRICT, so a simulation a report uses cannot be deleted. A report migrated from the Public API's v1 database keeps its v1 report ID in the unique legacy_report_id. Revision 0002 now creates the reports table too, and revokes UPDATE on both definition tables from the runtime role. ReportStore creates, reads and finds reports by v1 ID; create_migrated is idempotent by legacy_report_id, for the one-time migration (item 2.05). NewReport checks the country, year, options and v1 ID before any write; the database checks them again for any other writer. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 18 ++ .../__init__.py | 10 + .../reports.py | 226 ++++++++++++++++++ .../tests/report_support.py | 56 +++++ .../tests/test_definitions.py | 43 +++- .../tests/test_migration_cycle.py | 8 +- .../tests/test_reports.py | 190 +++++++++++++++ .../tests/test_schema.py | 2 +- .../versions/0002_simulations_and_reports.py | 83 ++++++- 9 files changed, 624 insertions(+), 12 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py create mode 100644 libs/policyengine-simulation-persistence/tests/report_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_reports.py diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index ab82daf90..aa689f9a2 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -8,6 +8,24 @@ The temporary Stage 12 comparison tables are not part of this schema: they are owned and migrated by the Public API, and `libs/policyengine-stage12-persistence` writes them with DML only. +## Tables + +| Table | Holds | +| --- | --- | +| `simulations` | What to run, independent of model versions: kind (`economy` or `household`), country, policy (null for current law), dynamic, and either a household or a region code (`state/ca`) with an optional dataset name (null means the region's default dataset in each catalog version) | +| `reports` | Which two simulations a report compares (an economy report needs a reform; a household report may have only a baseline), both of the report's own country and kind, with its year and options | + +Both are definitions: the runtime role may read and insert them but not update +or delete them. Rows migrated from the Public API's v1 database keep their v1 +IDs in the unique `legacy_simulation_id` and `legacy_report_id` columns, and +the stores' `create_migrated` methods return the existing row when the same v1 +ID is migrated again. + +Foreign keys into the Public API's tables in the same database (`policies` and +`households`, each together with its country, and `dynamics`) are +`ON DELETE RESTRICT`. Creating them needs `REFERENCES` on those tables for the +migrator role; the runtime role needs no privilege on them. + ## Ownership rules These rules are registered in `rules/registry.toml` and checked by diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 000d62dae..5e87ac336 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -9,6 +9,12 @@ normalized_database_url, read_schema_revision, ) +from policyengine_simulation_persistence.reports import ( + NewReport, + ReportKind, + ReportRecord, + ReportStore, +) from policyengine_simulation_persistence.schema import ( ALEMBIC_VERSION_TABLE, HEAD_REVISION, @@ -28,7 +34,11 @@ "HEAD_REVISION", "RUNTIME_ROLE", "SIMULATION_SCHEMA", + "NewReport", "NewSimulation", + "ReportKind", + "ReportRecord", + "ReportStore", "SimulationKind", "SimulationRecord", "SimulationStore", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py new file mode 100644 index 000000000..f68a4caa8 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -0,0 +1,226 @@ +"""``simulation.reports``: which two simulations a report compares. + +A report is an economy or household comparison a user can save: its baseline +simulation, its reform simulation (required for an economy report; a household +report may have only a baseline), the year and further options. It is a +definition and never changes. Both simulations must have the report's country +and kind, which composite foreign keys into ``simulation.simulations`` check. + +A report's calculations are its report runs (item 1.06); a report with none has +not been calculated yet. A report migrated from the Public API's v1 database +keeps its v1 report ID in ``legacy_report_id``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from dataclasses import field as dataclass_field +from enum import StrEnum +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from sqlalchemy import ( + BigInteger, + CheckConstraint, + Column, + DateTime, + ForeignKeyConstraint, + Index, + Integer, + Table, + Text, + Uuid, + func, + select, + text, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence.schema import simulation_metadata +from policyengine_simulation_persistence.simulations import ( + check_country, + check_legacy_id, + insert_or_get_existing, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + + from sqlalchemy import ColumnElement, Engine, RowMapping + +FIRST_YEAR = 1900 +LAST_YEAR = 2200 + + +class ReportKind(StrEnum): + """What a report compares; its simulations have the same kind.""" + + ECONOMY = "economy" + HOUSEHOLD = "household" + + +_KIND_VALUES = ", ".join(f"'{kind.value}'" for kind in ReportKind) +_SIMULATION_KEY = [simulations.c.id, simulations.c.country, simulations.c.kind] + +reports = Table( + "reports", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("kind", Text, nullable=False), + Column("country", Text, nullable=False), + Column("baseline_simulation_id", Uuid(as_uuid=True), nullable=False), + Column("reform_simulation_id", Uuid(as_uuid=True), nullable=True), + Column("year", Integer, nullable=False), + Column("options", JSONB, nullable=False, server_default=text("'{}'::jsonb")), + Column("legacy_report_id", BigInteger, nullable=True, unique=True), + Column( + "created_at", + DateTime(timezone=True), + nullable=False, + server_default=func.now(), + ), + CheckConstraint(f"kind IN ({_KIND_VALUES})", name="kind"), + CheckConstraint("country ~ '^[a-z]{2}$'", name="country"), + CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), + CheckConstraint("jsonb_typeof(options) = 'object'", name="options"), + CheckConstraint( + "kind = 'household' OR reform_simulation_id IS NOT NULL", + name="economy_has_reform", + ), + CheckConstraint( + "legacy_report_id IS NULL OR legacy_report_id > 0", + name="legacy_report_id", + ), + ForeignKeyConstraint( + ["baseline_simulation_id", "country", "kind"], + _SIMULATION_KEY, + name="fk_reports_baseline_simulation", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["reform_simulation_id", "country", "kind"], + _SIMULATION_KEY, + name="fk_reports_reform_simulation", + ondelete="RESTRICT", + ), +) +Index("ix_reports_baseline_simulation_id", reports.c.baseline_simulation_id) +Index("ix_reports_reform_simulation_id", reports.c.reform_simulation_id) + + +@dataclass(frozen=True) +class NewReport: + """A report to create. + + The database checks that both simulations exist with the report's country + and kind, and that an economy report has a reform; the fields are checked + here first. + """ + + kind: ReportKind + country: str + baseline_simulation_id: UUID + year: int + reform_simulation_id: UUID | None = None + options: dict[str, Any] = dataclass_field(default_factory=dict) + legacy_report_id: int | None = None + + def __post_init__(self) -> None: + """Reject malformed fields without touching the database.""" + check_country(self.country) + if not FIRST_YEAR <= self.year <= LAST_YEAR: + msg = f"year must be between {FIRST_YEAR} and {LAST_YEAR}" + raise ValueError(msg) + if not isinstance(self.options, dict): + msg = "options must be a JSON object" + raise TypeError(msg) + check_legacy_id("legacy_report_id", self.legacy_report_id) + + +@dataclass(frozen=True) +class ReportRecord: + """One ``simulation.reports`` row.""" + + id: UUID + kind: ReportKind + country: str + baseline_simulation_id: UUID + reform_simulation_id: UUID | None + year: int + options: dict[str, Any] + legacy_report_id: int | None + created_at: datetime + + @classmethod + def from_row(cls, row: RowMapping) -> ReportRecord: + """Build the record from a selected ``reports`` row.""" + values = {field.name: row[field.name] for field in fields(cls)} + values["kind"] = ReportKind(values["kind"]) + return cls(**values) + + +def _new_report_values(new_report: NewReport) -> dict[str, Any]: + return { + "id": uuid4(), + **{field.name: getattr(new_report, field.name) for field in fields(new_report)}, + "kind": new_report.kind.value, + } + + +class ReportStore: + """Typed SQLAlchemy Core reads and inserts for report rows.""" + + def __init__(self, engine: Engine) -> None: + """Use ``engine``, which connects as the runtime role.""" + self.engine = engine + + def create(self, new_report: NewReport) -> ReportRecord: + """Insert one report and return it.""" + with self.engine.begin() as connection: + row = ( + connection.execute( + postgres_insert(reports) + .values(**_new_report_values(new_report)) + .returning(*reports.c) + ) + .mappings() + .one() + ) + return ReportRecord.from_row(row) + + def create_migrated(self, new_report: NewReport) -> ReportRecord: + """Create a report for a v1 report, or return the existing one. + + The v1 report ID must be set. An existing row is returned as it is, + even if ``new_report`` describes it differently, so running the + one-time migration again creates nothing new. The report has no run + until one is added. + """ + if new_report.legacy_report_id is None: + msg = "a migrated report needs its legacy_report_id" + raise ValueError(msg) + with self.engine.begin() as connection: + row = insert_or_get_existing( + connection, + reports, + _new_report_values(new_report), + reports.c.legacy_report_id, + ) + return ReportRecord.from_row(row) + + def get(self, report_id: UUID) -> ReportRecord | None: + """Return the report with this ID, or None.""" + return self._one(reports.c.id == report_id) + + def get_by_legacy_report_id(self, legacy_report_id: int) -> ReportRecord | None: + """Return the report migrated from this v1 report, or None.""" + return self._one(reports.c.legacy_report_id == legacy_report_id) + + def _one(self, condition: ColumnElement[bool]) -> ReportRecord | None: + with self.engine.connect() as connection: + row = ( + connection.execute(select(reports).where(condition)).mappings().first() + ) + return None if row is None else ReportRecord.from_row(row) diff --git a/libs/policyengine-simulation-persistence/tests/report_support.py b/libs/policyengine-simulation-persistence/tests/report_support.py new file mode 100644 index 000000000..2ce21677e --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/report_support.py @@ -0,0 +1,56 @@ +"""Builders for the report tests.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence import ( + NewReport, + ReportKind, + ReportStore, + create_simulation_engine, +) +from simulation_support import ( + economy_simulation, + household_simulation, + simulation_store, +) + +if TYPE_CHECKING: + from policyengine_simulation_persistence import SimulationRecord + from public_api_support import PublicApiRows + + +def report_store(database) -> ReportStore: + """Return a store that connects as the runtime role.""" + return ReportStore(create_simulation_engine(database.runtime_url)) + + +def economy_pair( + database, rows: PublicApiRows +) -> tuple[SimulationRecord, SimulationRecord]: + """Create a current-law baseline and a reform of California's economy.""" + store = simulation_store(database) + baseline = store.create(economy_simulation(rows, policy_id=None)) + reform = store.create(economy_simulation(rows)) + return baseline, reform + + +def economy_report( + baseline: SimulationRecord, reform: SimulationRecord, **overrides: object +) -> NewReport: + """Return a 2026 US economy report comparing the two simulations.""" + values: dict[str, object] = { + "kind": ReportKind.ECONOMY, + "country": "us", + "baseline_simulation_id": baseline.id, + "reform_simulation_id": reform.id, + "year": 2026, + } + values.update(overrides) + return NewReport(**values) # type: ignore[arg-type] + + +def household_baseline(database, rows: PublicApiRows) -> SimulationRecord: + """Create a current-law simulation of the US household.""" + return simulation_store(database).create(household_simulation(rows)) diff --git a/libs/policyengine-simulation-persistence/tests/test_definitions.py b/libs/policyengine-simulation-persistence/tests/test_definitions.py index f9b843a5c..614b53ea1 100644 --- a/libs/policyengine-simulation-persistence/tests/test_definitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_definitions.py @@ -2,8 +2,15 @@ from __future__ import annotations +from uuid import uuid4 + import pytest -from policyengine_simulation_persistence import NewSimulation, SimulationKind +from policyengine_simulation_persistence import ( + NewReport, + NewSimulation, + ReportKind, + SimulationKind, +) @pytest.mark.parametrize( @@ -29,3 +36,37 @@ def test_malformed_simulation_definitions_are_rejected(field_name, value) -> Non with pytest.raises(ValueError, match=field_name): NewSimulation(**values) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("country", "gb "), + ("year", 1899), + ("year", 2201), + ("legacy_report_id", -4), + ], +) +def test_malformed_report_definitions_are_rejected(field_name, value) -> None: + values = { + "kind": ReportKind.ECONOMY, + "country": "us", + "baseline_simulation_id": uuid4(), + "reform_simulation_id": uuid4(), + "year": 2026, + field_name: value, + } + + with pytest.raises(ValueError, match=field_name): + NewReport(**values) + + +def test_report_options_must_be_a_json_object() -> None: + with pytest.raises(TypeError, match="options"): + NewReport( + kind=ReportKind.ECONOMY, + country="us", + baseline_simulation_id=uuid4(), + year=2026, + options=["not", "an", "object"], # type: ignore[arg-type] + ) diff --git a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py index 828c1a386..fcf332e5e 100644 --- a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py +++ b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py @@ -19,7 +19,7 @@ SIMULATION_SCHEMA, create_simulation_engine, ) -from simulation_support import economy_simulation, simulation_store +from report_support import economy_pair, economy_report, report_store from sqlalchemy import create_engine, inspect, text from sqlalchemy.exc import OperationalError, ProgrammingError @@ -210,12 +210,12 @@ def test_the_chain_needs_references_on_the_public_api_tables( assert applied_revision(disposable_database.migrator_url) is None -@pytest.mark.parametrize("table_name", ["simulations"]) +@pytest.mark.parametrize("table_name", ["simulations", "reports"]) def test_the_runtime_role_cannot_change_or_delete_a_definition( migrated_database, public_api_rows, table_name ) -> None: - store = simulation_store(migrated_database) - store.create(economy_simulation(public_api_rows)) + baseline, reform = economy_pair(migrated_database, public_api_rows) + report_store(migrated_database).create(economy_report(baseline, reform)) engine = create_simulation_engine(migrated_database.runtime_url) try: for statement in ( diff --git a/libs/policyengine-simulation-persistence/tests/test_reports.py b/libs/policyengine-simulation-persistence/tests/test_reports.py new file mode 100644 index 000000000..9d40c597a --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_reports.py @@ -0,0 +1,190 @@ +"""simulation.reports: which simulations a report compares, its year and options.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from policyengine_simulation_persistence import ReportKind, ReportRecord +from public_api_support import superuser_engine +from report_support import ( + economy_pair, + economy_report, + household_baseline, + report_store, +) +from simulation_support import economy_simulation, simulation_store +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError + +pytestmark = pytest.mark.postgres + + +def test_an_economy_report_compares_a_baseline_and_a_reform( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + store = report_store(migrated_database) + + created = store.create( + economy_report(baseline, reform, options={"target": "cliffs"}) + ) + + assert isinstance(created, ReportRecord) + assert created.kind is ReportKind.ECONOMY + assert created.baseline_simulation_id == baseline.id + assert created.reform_simulation_id == reform.id + assert created.year == 2026 + assert created.options == {"target": "cliffs"} + assert created.legacy_report_id is None + assert store.get(created.id) == created + + +def test_a_household_report_may_have_only_a_baseline( + migrated_database, public_api_rows +) -> None: + baseline = household_baseline(migrated_database, public_api_rows) + + created = report_store(migrated_database).create( + economy_report( + baseline, + baseline, + kind=ReportKind.HOUSEHOLD, + reform_simulation_id=None, + ) + ) + + assert created.kind is ReportKind.HOUSEHOLD + assert created.reform_simulation_id is None + assert created.options == {} + + +def test_an_economy_report_needs_a_reform(migrated_database, public_api_rows) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(IntegrityError, match="economy_has_reform"): + report_store(migrated_database).create( + economy_report(baseline, reform, reform_simulation_id=None) + ) + + +def test_simulations_must_have_the_report_country( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + uk_reform = simulation_store(migrated_database).create( + economy_simulation( + public_api_rows, + country="uk", + policy_id=public_api_rows.uk_policy_id, + region_code="country/wales", + ) + ) + + with pytest.raises(IntegrityError, match="fk_reports_reform_simulation"): + report_store(migrated_database).create(economy_report(baseline, uk_reform)) + + +def test_simulations_must_have_the_report_kind( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(IntegrityError, match="fk_reports_baseline_simulation"): + report_store(migrated_database).create( + economy_report( + baseline, + reform, + kind=ReportKind.HOUSEHOLD, + reform_simulation_id=None, + ) + ) + + +def test_an_unknown_simulation_is_rejected(migrated_database, public_api_rows) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(IntegrityError, match="fk_reports_reform_simulation"): + report_store(migrated_database).create( + economy_report(baseline, reform, reform_simulation_id=uuid4()) + ) + + +@pytest.mark.parametrize( + ("column", "value", "constraint"), + [ + ("year", "1899", "ck_reports_year"), + ("options", "'[]'::jsonb", "ck_reports_options"), + ], +) +def test_the_database_checks_the_year_and_options( + migrated_database, public_api_rows, column, value, constraint +) -> None: + """The database checks what NewReport also checks, for any other writer.""" + baseline, reform = economy_pair(migrated_database, public_api_rows) + values = {"year": "2026", "options": "'{}'::jsonb", column: value} + engine = superuser_engine(migrated_database.superuser_url) + try: + with ( + pytest.raises(IntegrityError, match=constraint), + engine.begin() as connection, + ): + connection.execute( + text( + "INSERT INTO simulation.reports (id, kind, country, " # noqa: S608 # fixed test values + "baseline_simulation_id, reform_simulation_id, year, options) " + f"VALUES (:id, 'economy', 'us', :baseline, :reform, " + f"{values['year']}, {values['options']})" + ), + {"id": uuid4(), "baseline": baseline.id, "reform": reform.id}, + ) + finally: + engine.dispose() + + +def test_the_public_api_cannot_delete_a_simulation_a_report_uses( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + report_store(migrated_database).create(economy_report(baseline, reform)) + engine = superuser_engine(migrated_database.superuser_url) + try: + with ( + pytest.raises(IntegrityError, match="violates foreign key constraint"), + engine.begin() as connection, + ): + connection.execute( + text("DELETE FROM simulation.simulations WHERE id = :id"), + {"id": baseline.id}, + ) + finally: + engine.dispose() + + +def test_a_migrated_report_is_created_once(migrated_database, public_api_rows) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + store = report_store(migrated_database) + first = store.create_migrated(economy_report(baseline, reform, legacy_report_id=91)) + + again = store.create_migrated( + economy_report(baseline, reform, legacy_report_id=91, year=2027) + ) + + assert again == first + assert store.get_by_legacy_report_id(91) == first + assert store.get_by_legacy_report_id(92) is None + + +def test_a_migrated_report_needs_its_legacy_id( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(ValueError, match="legacy_report_id"): + report_store(migrated_database).create_migrated( + economy_report(baseline, reform) + ) + + +def test_an_unknown_report_reads_as_none(migrated_database) -> None: + assert report_store(migrated_database).get(uuid4()) is None diff --git a/libs/policyengine-simulation-persistence/tests/test_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index 2125d6c14..8ac5b7b05 100644 --- a/libs/policyengine-simulation-persistence/tests/test_schema.py +++ b/libs/policyengine-simulation-persistence/tests/test_schema.py @@ -116,4 +116,4 @@ def test_importing_the_package_registers_every_table() -> None: importlib.import_module(f"policyengine_simulation_persistence.{module.name}") assert set(simulation_metadata.tables) == package_tables - assert "simulation.simulations" in package_tables + assert {"simulation.simulations", "simulation.reports"} <= package_tables diff --git a/migrations/versions/0002_simulations_and_reports.py b/migrations/versions/0002_simulations_and_reports.py index d5231da6e..e6351e1ac 100644 --- a/migrations/versions/0002_simulations_and_reports.py +++ b/migrations/versions/0002_simulations_and_reports.py @@ -1,12 +1,14 @@ -"""Add the simulation definitions table. +"""Add the simulation and report definition tables. ``simulations`` records what each simulation runs, independent of model versions: its kind, country, policy (null for current law), dynamic, and either a household or a region code with an optional dataset name. Foreign keys into the Public API's policies, households and dynamics in the same database are ``ON DELETE RESTRICT`` and need ``REFERENCES`` on those tables for the migrator -role. Definitions never change, so the runtime role may read and insert them -but not update them. Nothing writes the table yet. +role. ``reports`` records which simulations each report compares, of the +report's own country and kind, with its year and options. Definitions never +change, so the runtime role may read and insert them but not update them. +Nothing writes the tables yet. Revision ID: 0002 Revises: 0001 @@ -17,6 +19,7 @@ import sqlalchemy as sa from alembic import op +from sqlalchemy.dialects import postgresql revision = "0002" down_revision = "0001" @@ -106,14 +109,82 @@ def _create_simulations() -> None: ) +def _create_reports() -> None: + simulation_key = [ + f"{SCHEMA}.simulations.id", + f"{SCHEMA}.simulations.country", + f"{SCHEMA}.simulations.kind", + ] + op.create_table( + "reports", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("country", sa.Text(), nullable=False), + sa.Column("baseline_simulation_id", sa.Uuid(), nullable=False), + sa.Column("reform_simulation_id", sa.Uuid(), nullable=True), + sa.Column("year", sa.Integer(), nullable=False), + sa.Column( + "options", + postgresql.JSONB(), + server_default=sa.text("'{}'::jsonb"), + nullable=False, + ), + sa.Column("legacy_report_id", sa.BigInteger(), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint("kind IN ('economy', 'household')", name="ck_reports_kind"), + sa.CheckConstraint("country ~ '^[a-z]{2}$'", name="ck_reports_country"), + sa.CheckConstraint("year BETWEEN 1900 AND 2200", name="ck_reports_year"), + sa.CheckConstraint( + "jsonb_typeof(options) = 'object'", name="ck_reports_options" + ), + sa.CheckConstraint( + "kind = 'household' OR reform_simulation_id IS NOT NULL", + name="ck_reports_economy_has_reform", + ), + sa.CheckConstraint( + "legacy_report_id IS NULL OR legacy_report_id > 0", + name="ck_reports_legacy_report_id", + ), + sa.ForeignKeyConstraint( + ["baseline_simulation_id", "country", "kind"], + simulation_key, + name="fk_reports_baseline_simulation", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["reform_simulation_id", "country", "kind"], + simulation_key, + name="fk_reports_reform_simulation", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_reports"), + sa.UniqueConstraint("legacy_report_id", name="uq_reports_legacy_report_id"), + schema=SCHEMA, + ) + for column in ("baseline_simulation_id", "reform_simulation_id"): + op.create_index(f"ix_reports_{column}", "reports", [column], schema=SCHEMA) + + def upgrade() -> None: - """Create the simulations table; the runtime role may not update it.""" + """Create both tables; the runtime role may not update either.""" _create_simulations() - op.execute("REVOKE UPDATE ON simulation.simulations FROM simulation_runtime") + _create_reports() + op.execute( + "REVOKE UPDATE ON simulation.simulations, simulation.reports " + "FROM simulation_runtime" + ) def downgrade() -> None: - """Drop the table ``upgrade`` created, with its indexes and privileges.""" + """Drop the tables ``upgrade`` created, with their indexes and privileges.""" + for column in ("reform_simulation_id", "baseline_simulation_id"): + op.drop_index(f"ix_reports_{column}", table_name="reports", schema=SCHEMA) + op.drop_table("reports", schema=SCHEMA) for column in ("dynamic_id", "household_id", "policy_id"): op.drop_index( f"ix_simulations_{column}", table_name="simulations", schema=SCHEMA From 3a15fb6ed0bb195d5e5b8695ac9b54c8e28ccfe2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:53:04 +0400 Subject: [PATCH 3/5] Accept a simulation or report kind given as text NewSimulation and NewReport typed kind as an enum but neither checked nor converted it, so a kind given as plain text (as a mapping of v1 records easily does) passed their checks and then failed with AttributeError when the store read kind.value. Both now convert the kind to its enum member when they are created, and reject an unknown kind with a ValueError that names the field. Co-Authored-By: Claude Opus 5.5 --- .../reports.py | 3 ++ .../simulations.py | 14 ++++++++++ .../tests/test_definitions.py | 28 +++++++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py index f68a4caa8..99faccf40 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -42,6 +42,7 @@ check_country, check_legacy_id, insert_or_get_existing, + parse_kind, simulations, ) @@ -129,6 +130,8 @@ class NewReport: def __post_init__(self) -> None: """Reject malformed fields without touching the database.""" + # A frozen dataclass sets its own field through object.__setattr__. + object.__setattr__(self, "kind", parse_kind(self.kind, ReportKind)) check_country(self.country) if not FIRST_YEAR <= self.year <= LAST_YEAR: msg = f"year must be between {FIRST_YEAR} and {LAST_YEAR}" diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py index 31db413ad..0ad0e7629 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py @@ -145,6 +145,18 @@ def check_legacy_id(field_name: str, value: int | None) -> None: raise ValueError(msg) +def parse_kind[KindT: StrEnum](value: object, kind_type: type[KindT]) -> KindT: + """Return the ``kind_type`` member for ``value``, which may be its text. + + Raise ValueError naming ``kind`` when the value is not one of its members. + """ + try: + return kind_type(value) + except ValueError: + msg = f"kind must be one of {', '.join(kind_type)}" + raise ValueError(msg) from None + + def check_country(country: str) -> None: """Reject a country that is not a lowercase two-letter code.""" if not COUNTRY_PATTERN.fullmatch(country): @@ -200,6 +212,8 @@ class NewSimulation: def __post_init__(self) -> None: """Reject malformed fields without touching the database.""" + # A frozen dataclass sets its own field through object.__setattr__. + object.__setattr__(self, "kind", parse_kind(self.kind, SimulationKind)) check_country(self.country) _check_catalog_key("region_code", self.region_code) _check_catalog_key("dataset_name", self.dataset_name) diff --git a/libs/policyengine-simulation-persistence/tests/test_definitions.py b/libs/policyengine-simulation-persistence/tests/test_definitions.py index 614b53ea1..d15f35b35 100644 --- a/libs/policyengine-simulation-persistence/tests/test_definitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_definitions.py @@ -70,3 +70,31 @@ def test_report_options_must_be_a_json_object() -> None: year=2026, options=["not", "an", "object"], # type: ignore[arg-type] ) + + +def test_a_kind_given_as_text_becomes_the_enum() -> None: + simulation = NewSimulation(kind="economy", country="us", region_code="state/ca") # type: ignore[arg-type] + report = NewReport( + kind="household", # type: ignore[arg-type] + country="us", + baseline_simulation_id=uuid4(), + year=2026, + ) + + assert simulation.kind is SimulationKind.ECONOMY + assert report.kind is ReportKind.HOUSEHOLD + + +def test_an_unknown_simulation_kind_is_rejected() -> None: + with pytest.raises(ValueError, match="kind"): + NewSimulation(kind="microsimulation", country="us", region_code="state/ca") # type: ignore[arg-type] + + +def test_an_unknown_report_kind_is_rejected() -> None: + with pytest.raises(ValueError, match="kind"): + NewReport( + kind="microsimulation", # type: ignore[arg-type] + country="us", + baseline_simulation_id=uuid4(), + year=2026, + ) From e2dd3625dc5654cfc75555d2a533bceb994b8e1d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:54:33 +0400 Subject: [PATCH 4/5] Test every definition check in the database and correct the roles table The database repeats each check NewSimulation and NewReport make, for any other writer, but only the report year and options checks had a test, so a mistake in another check's expression would have gone unnoticed. New tests insert rows directly as the runtime role and expect each of the other checks to refuse them: the kinds, countries, region code, dataset name, both v1 IDs, and an economy report without a reform. The README's roles table still gave the runtime role UPDATE on every table, and did not say the migrator needs REFERENCES on the Public API tables before revision 0002 deploys; both are corrected. The unused REFERENCED_PUBLIC_API_TABLES constant is gone. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 14 ++++-- .../public_api.py | 2 - .../tests/test_reports.py | 43 ++++++++++++++++++- .../tests/test_simulations.py | 35 +++++++++++++++ 4 files changed, 86 insertions(+), 8 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index aa689f9a2..1436930e8 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -57,12 +57,18 @@ These rules are registered in `rules/registry.toml` and checked by | Role | Used by | Privileges | | --- | --- | --- | -| `simulation_migrator` | The deploy workflow's migration job, through `SIMULATION_MIGRATION_DATABASE_URL` | Creates the `simulation` schema if it is missing, and owns every object in it | -| `simulation_runtime` | The Simulation Entrypoint, through `SIMULATION_DATABASE_URL` | `USAGE` on the schema, `SELECT` on the version table, and `SELECT`, `INSERT` and `UPDATE` on the tables the migrations create | +| `simulation_migrator` | The deploy workflow's migration job, through `SIMULATION_MIGRATION_DATABASE_URL` | Creates the `simulation` schema if it is missing, and owns every object in it; `REFERENCES` on the Public API tables its foreign keys point to | +| `simulation_runtime` | The Simulation Entrypoint, through `SIMULATION_DATABASE_URL` | `USAGE` on the schema, `SELECT` on the version table, and `SELECT`, `INSERT` and `UPDATE` on the tables the migrations create, except `UPDATE` on the definition tables (`simulations`, `reports`); never `DELETE`, and nothing on the Public API tables | The initial revision grants the runtime role its privileges, including default -privileges on tables the migrator creates later. Both roles must exist before -the first migration runs; an operator creates them. The migrator needs either +privileges on tables the migrator creates later; revision 0002 revokes +`UPDATE` on the definition tables, which never change. Both roles must exist +before the first migration runs; an operator creates them. Before revision +0002 is deployed, the owner of the Public API's tables must also grant the +migrator `REFERENCES` on `policies`, `households` and `dynamics` (the +operator step grants it on `tax_benefit_model_versions`, `datasets` and +`regions` too, for later revisions); without it the migration fails and +applies nothing. The migrator needs either `CREATE` on the database or a `simulation` schema created in advance and owned by it: the environment creates the schema only when it is missing, because PostgreSQL checks `CREATE` on the database even for diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py index 1da0d5b0d..65f709a1f 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py @@ -41,5 +41,3 @@ public_api_metadata, Column("id", Uuid(as_uuid=True), primary_key=True), ) - -REFERENCED_PUBLIC_API_TABLES = tuple(sorted(public_api_metadata.tables)) diff --git a/libs/policyengine-simulation-persistence/tests/test_reports.py b/libs/policyengine-simulation-persistence/tests/test_reports.py index 9d40c597a..d1061a8b7 100644 --- a/libs/policyengine-simulation-persistence/tests/test_reports.py +++ b/libs/policyengine-simulation-persistence/tests/test_reports.py @@ -5,7 +5,12 @@ from uuid import uuid4 import pytest -from policyengine_simulation_persistence import ReportKind, ReportRecord +from policyengine_simulation_persistence import ( + ReportKind, + ReportRecord, + create_simulation_engine, +) +from policyengine_simulation_persistence.reports import reports from public_api_support import superuser_engine from report_support import ( economy_pair, @@ -14,7 +19,7 @@ report_store, ) from simulation_support import economy_simulation, simulation_store -from sqlalchemy import text +from sqlalchemy import insert, text from sqlalchemy.exc import IntegrityError pytestmark = pytest.mark.postgres @@ -188,3 +193,37 @@ def test_a_migrated_report_needs_its_legacy_id( def test_an_unknown_report_reads_as_none(migrated_database) -> None: assert report_store(migrated_database).get(uuid4()) is None + + +@pytest.mark.parametrize( + ("column", "value", "constraint"), + [ + ("kind", "microsimulation", "ck_reports_kind"), + ("country", "GB", "ck_reports_country"), + ("legacy_report_id", -4, "ck_reports_legacy_report_id"), + ("reform_simulation_id", None, "ck_reports_economy_has_reform"), + ], +) +def test_the_database_checks_each_report_field( + migrated_database, public_api_rows, column, value, constraint +) -> None: + """The database checks what NewReport also checks, for any other writer.""" + baseline, reform = economy_pair(migrated_database, public_api_rows) + values = { + "id": uuid4(), + "kind": "economy", + "country": "us", + "baseline_simulation_id": baseline.id, + "reform_simulation_id": reform.id, + "year": 2026, + column: value, + } + engine = create_simulation_engine(migrated_database.runtime_url) + try: + with ( + pytest.raises(IntegrityError, match=constraint), + engine.begin() as connection, + ): + connection.execute(insert(reports).values(**values)) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_simulations.py b/libs/policyengine-simulation-persistence/tests/test_simulations.py index 4b15d4d97..f574b8f26 100644 --- a/libs/policyengine-simulation-persistence/tests/test_simulations.py +++ b/libs/policyengine-simulation-persistence/tests/test_simulations.py @@ -9,13 +9,16 @@ NewSimulation, SimulationKind, SimulationRecord, + create_simulation_engine, ) +from policyengine_simulation_persistence.simulations import simulations from public_api_support import delete_public_api_row from simulation_support import ( economy_simulation, household_simulation, simulation_store, ) +from sqlalchemy import insert from sqlalchemy.exc import IntegrityError pytestmark = pytest.mark.postgres @@ -210,3 +213,35 @@ def test_two_simulations_cannot_share_a_legacy_id( def test_an_unknown_simulation_reads_as_none(migrated_database) -> None: assert simulation_store(migrated_database).get(uuid4()) is None + + +@pytest.mark.parametrize( + ("column", "value", "constraint"), + [ + ("kind", "microsimulation", "ck_simulations_kind"), + ("country", "USA", "ck_simulations_country"), + ("region_code", "state ca", "ck_simulations_region_code"), + ("dataset_name", "populace us", "ck_simulations_dataset_name"), + ("legacy_simulation_id", 0, "ck_simulations_legacy_simulation_id"), + ], +) +def test_the_database_checks_each_simulation_field( + migrated_database, public_api_rows, column, value, constraint +) -> None: + """The database checks what NewSimulation also checks, for any other writer.""" + values = { + "id": uuid4(), + "kind": "economy", + "country": "us", + "region_code": "state/ca", + column: value, + } + engine = create_simulation_engine(migrated_database.runtime_url) + try: + with ( + pytest.raises(IntegrityError, match=constraint), + engine.begin() as connection, + ): + connection.execute(insert(simulations).values(**values)) + finally: + engine.dispose() From c7404a186361be6d72b0b0380d897e30ad82f253 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:01:53 +0400 Subject: [PATCH 5/5] Name the check constraints as their table definitions do Revision 0002 passed full check constraint names, such as ck_simulations_kind, which the naming convention prefixed a second time, so the database held ck_simulations_ck_simulations_kind while the table definitions name it ck_simulations_kind. alembic check does not compare check constraints, so nothing noticed; a later revision dropping one by its defined name would have failed. The revision now marks the names as final (op.f), and a test compares every constraint name in the migrated database with the names the table definitions give. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_migration_cycle.py | 37 +++++++++++++++++++ .../versions/0002_simulations_and_reports.py | 28 ++++++++------ 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py index fcf332e5e..b889e61a0 100644 --- a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py +++ b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re import time import pytest @@ -18,10 +19,13 @@ HEAD_REVISION, SIMULATION_SCHEMA, create_simulation_engine, + simulation_metadata, ) from report_support import economy_pair, economy_report, report_store from sqlalchemy import create_engine, inspect, text +from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError, ProgrammingError +from sqlalchemy.schema import CreateTable pytestmark = pytest.mark.postgres @@ -229,3 +233,36 @@ def test_the_runtime_role_cannot_change_or_delete_a_definition( connection.execute(text(statement)) finally: engine.dispose() + + +def _declared_constraint_names() -> set[str]: + """Return the names the table definitions give their constraints.""" + names: set[str] = set() + for table in simulation_metadata.sorted_tables: + ddl = str(CreateTable(table).compile(dialect=postgresql.dialect())) + names.update(re.findall(r"CONSTRAINT (\w+)", ddl)) + return names + + +def test_every_constraint_has_the_name_its_definition_gives_it( + migrated_database, +) -> None: + """Compare names here, since alembic check does not compare check constraints.""" + engine = _superuser_engine(migrated_database) + try: + with engine.connect() as connection: + database_names = set( + connection.execute( + text( + "SELECT conname FROM pg_constraint " + "JOIN pg_namespace ON pg_namespace.oid = connamespace " + "WHERE nspname = 'simulation' " + "AND contype IN ('c', 'f', 'p', 'u') " + "AND conrelid <> 'simulation.alembic_version'::regclass" + ) + ).scalars() + ) + finally: + engine.dispose() + + assert database_names == _declared_constraint_names() diff --git a/migrations/versions/0002_simulations_and_reports.py b/migrations/versions/0002_simulations_and_reports.py index e6351e1ac..b3fd790de 100644 --- a/migrations/versions/0002_simulations_and_reports.py +++ b/migrations/versions/0002_simulations_and_reports.py @@ -54,27 +54,29 @@ def _create_simulations() -> None: ), sa.CheckConstraint( "kind IN ('economy', 'household')", - name="ck_simulations_kind", + name=op.f("ck_simulations_kind"), + ), + sa.CheckConstraint( + "country ~ '^[a-z]{2}$'", name=op.f("ck_simulations_country") ), - sa.CheckConstraint("country ~ '^[a-z]{2}$'", name="ck_simulations_country"), sa.CheckConstraint( "(kind = 'household' AND household_id IS NOT NULL " "AND region_code IS NULL AND dataset_name IS NULL) " "OR (kind = 'economy' AND household_id IS NULL " "AND region_code IS NOT NULL)", - name="ck_simulations_population_matches_kind", + name=op.f("ck_simulations_population_matches_kind"), ), sa.CheckConstraint( "region_code IS NULL OR region_code ~ '^\\S+$'", - name="ck_simulations_region_code", + name=op.f("ck_simulations_region_code"), ), sa.CheckConstraint( "dataset_name IS NULL OR dataset_name ~ '^\\S+$'", - name="ck_simulations_dataset_name", + name=op.f("ck_simulations_dataset_name"), ), sa.CheckConstraint( "legacy_simulation_id IS NULL OR legacy_simulation_id > 0", - name="ck_simulations_legacy_simulation_id", + name=op.f("ck_simulations_legacy_simulation_id"), ), sa.ForeignKeyConstraint( ["policy_id", "country"], @@ -136,19 +138,21 @@ def _create_reports() -> None: server_default=sa.text("now()"), nullable=False, ), - sa.CheckConstraint("kind IN ('economy', 'household')", name="ck_reports_kind"), - sa.CheckConstraint("country ~ '^[a-z]{2}$'", name="ck_reports_country"), - sa.CheckConstraint("year BETWEEN 1900 AND 2200", name="ck_reports_year"), sa.CheckConstraint( - "jsonb_typeof(options) = 'object'", name="ck_reports_options" + "kind IN ('economy', 'household')", name=op.f("ck_reports_kind") + ), + sa.CheckConstraint("country ~ '^[a-z]{2}$'", name=op.f("ck_reports_country")), + sa.CheckConstraint("year BETWEEN 1900 AND 2200", name=op.f("ck_reports_year")), + sa.CheckConstraint( + "jsonb_typeof(options) = 'object'", name=op.f("ck_reports_options") ), sa.CheckConstraint( "kind = 'household' OR reform_simulation_id IS NOT NULL", - name="ck_reports_economy_has_reform", + name=op.f("ck_reports_economy_has_reform"), ), sa.CheckConstraint( "legacy_report_id IS NULL OR legacy_report_id > 0", - name="ck_reports_legacy_report_id", + name=op.f("ck_reports_legacy_report_id"), ), sa.ForeignKeyConstraint( ["baseline_simulation_id", "country", "kind"],