Calculate budget windows on the v2 economy path

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/economy_runs.py

@@ -51,6 +51,9 @@
     "uk": "run_economy_simulation_uk",
 }
 ECONOMY_COMPARISON_FUNCTION = "compare_economy_simulations"
+# How many years of one budget window the Simulation Entrypoint calculates at
+# once. The functions above leave room for several windows at this cap.
+BUDGET_WINDOW_MAX_CONCURRENT_YEARS = 5
 
 # The parameters whose change makes policyengine.py calculate a labour supply
 # response, by country: a parameter at or below one of these paths. They are

libs/policyengine-simulation-persistence/README.md

@@ -13,9 +13,10 @@ writes them with DML only.
 | 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 |
+| `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, and, for a budget-window report, its window size (1 to 75 years from its year) |
 | `simulation_runs` | One execution of a simulation for a catalog version and year: the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (`initial`, `manual` or `system`) and attempt number |
-| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call |
+| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call; a budget-window report's report run names the budget-window run it is a year of |
+| `budget_window_runs` | One calculation of a budget-window report on a catalog version: its trigger, whether its years calculate their simulations again, how many years it calculates at once, its status and error, when the safety check next looks at it, and the lease of the caller starting its next years |
 
 `simulations` and `reports` 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
@@ -70,6 +71,20 @@ steps whose simulation run failed) are failed so the work can be submitted
 again. A comparison step can start only once both of its simulation runs
 have succeeded; starting it is left to the stage 2 report path.
 
+### Budget-window runs
+
+A budget-window run groups one report run per year of its report's window,
+each of the same report and catalog version (a composite foreign key) and at
+most one per year (`uq_report_runs_budget_window_run_id_year`).
+`BudgetWindowRunStore` creates runs (`create_first_run` adds one only while
+the report has none, under a transaction-level advisory lock, as first report
+runs are), reads each year's report run with both of its simulation runs, and
+reads each year's `budget_impacts` row; nothing stores a total, so readers sum
+the years. A caller starting a run's next years first claims it
+(`claim_scheduling`, a lease on `scheduling_claimed_until`), and releasing the
+claim sets the run's next safety-check time. A run is finished once, with a
+conditional write, as `succeeded` or `failed`.
+
 ### Results
 
 Every economy result is stored as rows (`result_tables`), one table per
@@ -179,7 +194,8 @@ These rules are registered in `rules/registry.toml` and checked by
 The initial revision grants the runtime role its privileges, including default
 privileges on tables the migrator creates later; revision 0002 revokes
 `UPDATE` on the definition tables, which never change, and revision 0004 on
-the result tables, whose rows are written once. Both roles must exist
+the result tables, whose rows are written once; revision 0005's
+`budget_window_runs` takes the default privileges. 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`, and

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

@@ -4,6 +4,18 @@
 the migration environment compares with the database.
 """
 
+from policyengine_simulation_persistence.budget_window_runs import (
+    ONE_REPORT_RUN_PER_YEAR,
+    BudgetWindowRunRecord,
+    BudgetWindowStatus,
+    BudgetWindowYear,
+    NewBudgetWindowRun,
+    StartedBudgetWindowRun,
+    YearBudget,
+)
+from policyengine_simulation_persistence.budget_window_store import (
+    BudgetWindowRunStore,
+)
 from policyengine_simulation_persistence.catalog_reads import (
     CATALOG_READ_TABLES,
     CatalogError,
@@ -41,6 +53,7 @@
     record_output_files,
 )
 from policyengine_simulation_persistence.reports import (
+    MAX_BUDGET_WINDOW_YEARS,
     NewReport,
     ReportKind,
     ReportRecord,
@@ -85,13 +98,19 @@
     "ALEMBIC_VERSION_TABLE",
     "CATALOG_READ_TABLES",
     "HEAD_REVISION",
+    "MAX_BUDGET_WINDOW_YEARS",
+    "ONE_REPORT_RUN_PER_YEAR",
     "OPTIONAL_SECTIONS",
     "OUTPUT_FILE_RETENTION",
     "RESULT_FIELD_STORAGE",
     "RUNTIME_ROLE",
     "SIMULATION_SCHEMA",
     "SUCCEEDED_RUN_REUSE_WINDOW",
     "AbandonedDispatches",
+    "BudgetWindowRunRecord",
+    "BudgetWindowRunStore",
+    "BudgetWindowStatus",
+    "BudgetWindowYear",
     "CatalogError",
     "CatalogReader",
     "CatalogVersion",
@@ -100,6 +119,7 @@
     "EconomyResultLabels",
     "GeographyLabel",
     "HouseholdOutput",
+    "NewBudgetWindowRun",
     "NewReport",
     "NewReportRun",
     "NewSimulation",
@@ -123,8 +143,10 @@
     "SimulationRecord",
     "SimulationRunRecord",
     "SimulationStore",
+    "StartedBudgetWindowRun",
     "StartedReportRun",
     "StartedSimulationRun",
+    "YearBudget",
     "available_output_files",
     "check_storable",
     "create_simulation_engine",

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

@@ -0,0 +1,226 @@
+"""``budget_window_runs``: calculations of budget-window reports.
+
+A budget-window report (``reports.budget_window_size``) compares its two
+simulations in each year of its window. One budget-window run is one
+calculation of it on a catalog version: it records what started it (its
+trigger), whether its years calculate their simulations again (``force``), how
+many years it calculates at once, its status and error, and when the safety
+check next looks at it. Each year it calculates is one of its report runs
+(``report_runs.budget_window_run_id``), at most one per year, of its report and
+catalog version. The run stores no figures of its own: its totals are the sums
+of its years' budget rows, added up when it is read.
+
+A run is ``running`` until every year of its window has succeeded, when it is
+``succeeded``, or until one year has failed, when it is ``failed`` with that
+year's error. ``scheduling_claimed_until`` is the lease of the one caller
+starting its next years at a time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from enum import StrEnum
+from typing import TYPE_CHECKING, Self
+
+from sqlalchemy import (
+    Boolean,
+    CheckConstraint,
+    Column,
+    DateTime,
+    ForeignKeyConstraint,
+    Index,
+    Integer,
+    Table,
+    Text,
+    UniqueConstraint,
+    Uuid,
+)
+
+from policyengine_simulation_persistence import public_api
+from policyengine_simulation_persistence.reports import reports
+from policyengine_simulation_persistence.runs import (
+    ReportRunRecord,
+    RunTrigger,
+    SimulationRunRecord,
+)
+from policyengine_simulation_persistence.schema import simulation_metadata
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from uuid import UUID
+
+    from sqlalchemy import RowMapping
+
+
+class BudgetWindowStatus(StrEnum):
+    """Where a budget-window run stands."""
+
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+
+
+# The unique constraint that allows one report run per year of a
+# budget-window run; a caller that loses a year's insert reads its name.
+ONE_REPORT_RUN_PER_YEAR = "uq_report_runs_budget_window_run_id_year"
+
+FINISHED_BUDGET_WINDOW_STATUSES = (
+    BudgetWindowStatus.SUCCEEDED,
+    BudgetWindowStatus.FAILED,
+)
+
+
+def _values(enum: type[StrEnum]) -> str:
+    return ", ".join(f"'{member.value}'" for member in enum)
+
+
+budget_window_runs = Table(
+    "budget_window_runs",
+    simulation_metadata,
+    Column("id", Uuid(as_uuid=True), primary_key=True),
+    Column("report_id", Uuid(as_uuid=True), nullable=False),
+    Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False),
+    Column("trigger", Text, nullable=False),
+    Column("force", Boolean, nullable=False),
+    Column("max_concurrent_years", Integer, nullable=False),
+    Column("status", Text, nullable=False),
+    Column("error_code", Text, nullable=True),
+    Column("error_message", Text, nullable=True),
+    Column("next_check_at", DateTime(timezone=True), nullable=True),
+    Column("scheduling_claimed_until", DateTime(timezone=True), nullable=True),
+    Column("completed_at", DateTime(timezone=True), nullable=True),
+    Column("created_at", DateTime(timezone=True), nullable=False),
+    Column("updated_at", DateTime(timezone=True), nullable=False),
+    CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"),
+    CheckConstraint(f"status IN ({_values(BudgetWindowStatus)})", name="status"),
+    CheckConstraint("max_concurrent_years >= 1", name="max_concurrent_years"),
+    # Only a running run is checked, and every running run is due a check.
+    CheckConstraint(
+        "(status = 'running') = (next_check_at IS NOT NULL)", name="next_check_at"
+    ),
+    CheckConstraint(
+        "(status IN ('succeeded', 'failed')) = (completed_at IS NOT NULL)",
+        name="completed_at",
+    ),
+    CheckConstraint("(status = 'failed') = (error_code IS NOT NULL)", name="error"),
+    ForeignKeyConstraint(
+        ["report_id"],
+        [reports.c.id],
+        name="fk_budget_window_runs_report",
+        ondelete="RESTRICT",
+    ),
+    ForeignKeyConstraint(
+        ["tax_benefit_model_version_id"],
+        [public_api.tax_benefit_model_versions.c.id],
+        name="fk_budget_window_runs_catalog_version",
+        ondelete="RESTRICT",
+    ),
+    # Each year's report run refers to the run with its report and catalog
+    # version, so it is a run of the same report on the same version.
+    UniqueConstraint(
+        "id",
+        "report_id",
+        "tax_benefit_model_version_id",
+        name="uq_budget_window_runs_id_report_catalog_version",
+    ),
+)
+# A report's newest budget-window run is its current one.
+Index(
+    "ix_budget_window_runs_report_id_created_at",
+    budget_window_runs.c.report_id,
+    budget_window_runs.c.created_at,
+)
+# The safety check reads only running runs that are due a check.
+Index(
+    "ix_budget_window_runs_due_checks",
+    budget_window_runs.c.next_check_at,
+    postgresql_where=budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value,
+)
+# A Public API delete of a catalog version checks this for references.
+Index(
+    "ix_budget_window_runs_tax_benefit_model_version_id",
+    budget_window_runs.c.tax_benefit_model_version_id,
+)
+
+
+@dataclass(frozen=True)
+class NewBudgetWindowRun:
+    """A calculation of a budget-window report to create.
+
+    ``force`` calculates every year's simulations again instead of reusing
+    runs that exist; ``max_concurrent_years`` caps how many of its years are
+    calculated at once.
+    """
+
+    report_id: UUID
+    tax_benefit_model_version_id: UUID
+    trigger: RunTrigger
+    max_concurrent_years: int
+    force: bool = False
+
+    def __post_init__(self) -> None:
+        """Reject a cap below one year without touching the database."""
+        if self.max_concurrent_years < 1:
+            msg = "max_concurrent_years must be at least 1"
+            raise ValueError(msg)
+
+
+@dataclass(frozen=True)
+class BudgetWindowRunRecord:
+    """One ``simulation.budget_window_runs`` row."""
+
+    id: UUID
+    report_id: UUID
+    tax_benefit_model_version_id: UUID
+    trigger: RunTrigger
+    force: bool
+    max_concurrent_years: int
+    status: BudgetWindowStatus
+    error_code: str | None
+    error_message: str | None
+    next_check_at: datetime | None
+    scheduling_claimed_until: datetime | None
+    completed_at: datetime | None
+    created_at: datetime
+    updated_at: datetime
+
+    @property
+    def finished(self) -> bool:
+        """Whether the run has succeeded or failed."""
+        return self.status in FINISHED_BUDGET_WINDOW_STATUSES
+
+    @classmethod
+    def from_row(cls, row: RowMapping) -> Self:
+        """Build the record from a selected ``budget_window_runs`` row."""
+        values = {field.name: row[field.name] for field in fields(cls)}
+        values["trigger"] = RunTrigger(values["trigger"])
+        values["status"] = BudgetWindowStatus(values["status"])
+        return cls(**values)
+
+
+@dataclass(frozen=True)
+class StartedBudgetWindowRun:
+    """The run a first start resolved to, and whether this call created it."""
+
+    run: BudgetWindowRunRecord
+    created: bool
+
+
+@dataclass(frozen=True)
+class BudgetWindowYear:
+    """One year of a budget-window run: its report run and both simulation runs."""
+
+    report_run: ReportRunRecord
+    baseline_run: SimulationRunRecord
+    reform_run: SimulationRunRecord
+
+
+@dataclass(frozen=True)
+class YearBudget:
+    """A year's budget figures, from its comparison's ``budget_impacts`` row."""
+
+    year: int
+    budgetary_impact: float
+    tax_revenue_impact: float
+    state_tax_revenue_impact: float
+    benefit_spending_impact: float