diff --git a/changelog.d/899.added.md b/changelog.d/899.added.md new file mode 100644 index 000000000..803a44a5f --- /dev/null +++ b/changelog.d/899.added.md @@ -0,0 +1 @@ +Added the `simulation.simulations` and `simulation.reports` tables, in migration revision `0002` of the `simulation` schema. A simulation records what to run independently of model versions: its kind, country, policy, dynamic, and either a household or a region code with an optional dataset name. A report records which baseline and reform simulations it compares, its year and its options. Both refer to the Public API's policies, households and dynamics with `ON DELETE RESTRICT` foreign keys and keep a unique v1 ID. Nothing writes either table yet. diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index ab82daf90..1436930e8 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 @@ -39,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/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index fb098c8ed..5e87ac336 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -1,10 +1,20 @@ -"""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, 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, @@ -12,12 +22,26 @@ SIMULATION_SCHEMA, simulation_metadata, ) +from policyengine_simulation_persistence.simulations import ( + NewSimulation, + SimulationKind, + SimulationRecord, + SimulationStore, +) __all__ = [ "ALEMBIC_VERSION_TABLE", "HEAD_REVISION", "RUNTIME_ROLE", "SIMULATION_SCHEMA", + "NewReport", + "NewSimulation", + "ReportKind", + "ReportRecord", + "ReportStore", + "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..65f709a1f --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py @@ -0,0 +1,43 @@ +"""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), +) 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..99faccf40 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -0,0 +1,229 @@ +"""``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, + parse_kind, + 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.""" + # 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}" + 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/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..0ad0e7629 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py @@ -0,0 +1,314 @@ +"""``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 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): + 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.""" + # 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) + 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/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/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..d15f35b35 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_definitions.py @@ -0,0 +1,100 @@ +"""New simulation and report definitions are checked before any write.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from policyengine_simulation_persistence import ( + NewReport, + NewSimulation, + ReportKind, + 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) + + +@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] + ) + + +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, + ) diff --git a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py index d033ec4aa..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,9 +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 @@ -196,3 +201,68 @@ 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", "reports"]) +def test_the_runtime_role_cannot_change_or_delete_a_definition( + migrated_database, public_api_rows, table_name +) -> None: + 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 ( + 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() + + +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/libs/policyengine-simulation-persistence/tests/test_reports.py b/libs/policyengine-simulation-persistence/tests/test_reports.py new file mode 100644 index 000000000..d1061a8b7 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_reports.py @@ -0,0 +1,229 @@ +"""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, + create_simulation_engine, +) +from policyengine_simulation_persistence.reports import reports +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 insert, 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 + + +@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_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index d42b8bf4d..8ac5b7b05 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", "simulation.reports"} <= 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..f574b8f26 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_simulations.py @@ -0,0 +1,247 @@ +"""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, + 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 + + +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 + + +@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() diff --git a/migrations/versions/0002_simulations_and_reports.py b/migrations/versions/0002_simulations_and_reports.py new file mode 100644 index 000000000..b3fd790de --- /dev/null +++ b/migrations/versions/0002_simulations_and_reports.py @@ -0,0 +1,196 @@ +"""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. ``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 +Create Date: 2026-10-22 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +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=op.f("ck_simulations_kind"), + ), + sa.CheckConstraint( + "country ~ '^[a-z]{2}$'", name=op.f("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=op.f("ck_simulations_population_matches_kind"), + ), + sa.CheckConstraint( + "region_code IS NULL OR region_code ~ '^\\S+$'", + name=op.f("ck_simulations_region_code"), + ), + sa.CheckConstraint( + "dataset_name IS NULL OR dataset_name ~ '^\\S+$'", + name=op.f("ck_simulations_dataset_name"), + ), + sa.CheckConstraint( + "legacy_simulation_id IS NULL OR legacy_simulation_id > 0", + name=op.f("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 _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=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=op.f("ck_reports_economy_has_reform"), + ), + sa.CheckConstraint( + "legacy_report_id IS NULL OR legacy_report_id > 0", + name=op.f("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 both tables; the runtime role may not update either.""" + _create_simulations() + _create_reports() + op.execute( + "REVOKE UPDATE ON simulation.simulations, simulation.reports " + "FROM simulation_runtime" + ) + + +def downgrade() -> None: + """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 + ) + op.drop_table("simulations", schema=SCHEMA)