Add simulation and report definitions to the simulation schema

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

libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py

@@ -1,23 +1,47 @@
-"""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,
     RUNTIME_ROLE,
     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",

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),
+)

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)

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",