diff --git a/changelog.d/stop-using-simulation-tables.changed.md b/changelog.d/stop-using-simulation-tables.changed.md new file mode 100644 index 000000000..0c87812e9 --- /dev/null +++ b/changelog.d/stop-using-simulation-tables.changed.md @@ -0,0 +1 @@ +Stopped the Public API from reading or writing its copies of the simulation-domain v2 tables and the retired `country` and `label` association columns, so a later revision can drop them while this release serves traffic. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 51c36817d..4c32f7d33 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -55,10 +55,10 @@ python -m pytest tests/unit/test_cloud_run_deploy_scripts.py tests/unit/test_cap For v2 platform foundation work, run the smallest applicable group while iterating, then all groups before qualification. -SQLModel schema, report/run behavior, lazy configuration, and import effects: +SQLModel schema, lazy configuration, and import effects: ```bash -uv run pytest tests/unit/v2/test_models.py tests/unit/v2/test_model_persistence.py tests/unit/v2/test_report_runs.py tests/unit/v2/test_settings.py tests/unit/v2/test_import_side_effects.py -q +uv run pytest tests/unit/v2/test_models.py tests/unit/v2/test_model_persistence.py tests/unit/v2/test_settings.py tests/unit/v2/test_import_side_effects.py -q ``` The generated-only v2 Alembic chain and disposable Postgres lifecycle: diff --git a/docs/migration/stage-12a-user-login.md b/docs/migration/stage-12a-user-login.md index c01657bc3..d73b3efd0 100644 --- a/docs/migration/stage-12a-user-login.md +++ b/docs/migration/stage-12a-user-login.md @@ -311,7 +311,8 @@ from the final SQLModel metadata. It has no post-generation corrections. kept: `country` becomes nullable and nothing reads or writes either column. The previous Public API revision still maps them, so dropping them would break it during the deploy window or an application rollback (see - "Rollback"). + "Rollback"). A later release stopped mapping them: they stay declared on the + tables, unmapped, only until a revision drops them. The models drop the `user_associations` relationships from v2 simulations and reports. `V2_ALEMBIC_HEAD_REVISION` is `fbb7488ebe8a`. diff --git a/docs/migration/stage-9-v2-metadata.md b/docs/migration/stage-9-v2-metadata.md index 4ce7d2c8e..893685987 100644 --- a/docs/migration/stage-9-v2-metadata.md +++ b/docs/migration/stage-9-v2-metadata.md @@ -23,9 +23,9 @@ Existing clients must not be redirected to the v2 resource routes. The initializer creates only reusable logical input `Dataset` rows. Each row has `is_output_dataset=false` and a null `storage_path`. It creates no -package-derived `DatasetVersion` rows or simulation/report output datasets and -does not modify `Simulation`, `Report`, or `ReportRun` rows or their dataset -references. +package-derived `DatasetVersion` rows or simulation/report output datasets. +The simulation service owns simulations, reports and report runs in its own +schema; the Public API no longer has models for them. ## Configuration and credential boundaries diff --git a/policyengine_api/data/v2/catalog/publication.py b/policyengine_api/data/v2/catalog/publication.py index 5e629fb05..a858a517d 100644 --- a/policyengine_api/data/v2/catalog/publication.py +++ b/policyengine_api/data/v2/catalog/publication.py @@ -24,12 +24,7 @@ PublicationEvidence, ) from policyengine_api.data.v2.migration_target import V2_ALEMBIC_HEAD_REVISION -from policyengine_api.data.v2.models import ( - DatasetVersion, - Report, - ReportRun, - Simulation, -) +from policyengine_api.data.v2.models import DatasetVersion if TYPE_CHECKING: from collections.abc import Callable @@ -43,12 +38,10 @@ "alembic_version", sa.column("version_num", sa.String), ) -PROTECTED_TABLES = ( - DatasetVersion.__table__, - Simulation.__table__, - Report.__table__, - ReportRun.__table__, -) +# User-created rows that publication must never change. Simulations, reports +# and report runs belong to the simulation service's own schema, and a later +# v2 revision drops the Public API's copies, so publication doesn't read them. +PROTECTED_TABLES = (DatasetVersion.__table__,) LOGGER = logging.getLogger(__name__) @@ -80,7 +73,7 @@ def _acquire_publication_lock(connection: Connection) -> None: ).scalar_one() -def _protected_row_counts(connection: Connection) -> tuple[int, int, int, int]: +def _protected_row_counts(connection: Connection) -> tuple[int, ...]: return tuple( connection.execute(sa.select(sa.func.count()).select_from(table)).scalar_one() for table in PROTECTED_TABLES @@ -120,7 +113,7 @@ def publish_catalog( assert_country_matches(connection, country) assert_canonical_value_uniqueness(connection) if _protected_row_counts(connection) != before: - msg = "publication changed simulation, report, or dataset-version rows" + msg = "publication changed dataset-version rows" raise CatalogPublicationError(msg) if checkpoint is not None: checkpoint("after_validation", connection) diff --git a/policyengine_api/data/v2/household_migration_qualification.py b/policyengine_api/data/v2/household_migration_qualification.py index ce679323e..646f64d84 100644 --- a/policyengine_api/data/v2/household_migration_qualification.py +++ b/policyengine_api/data/v2/household_migration_qualification.py @@ -10,18 +10,12 @@ from alembic.config import Config from alembic.migration import MigrationContext from alembic.script import ScriptDirectory -from sqlalchemy import Connection, Engine, func, inspect, select +from sqlalchemy import Connection, Engine, column, func, inspect, select, table from sqlalchemy.exc import NoInspectionAvailable from sqlalchemy.pool import NullPool from sqlmodel import create_engine -from policyengine_api.data.v2.models import ( - Household, - HouseholdJob, - Report, - Simulation, - UserHouseholdAssociation, -) +from policyengine_api.data.v2.models import Household, UserHouseholdAssociation from policyengine_api.data.v2.settings import ( V2ConfigurationError, V2DatabaseSettings, @@ -33,6 +27,12 @@ HOUSEHOLD_MIGRATION_REVISION = "724b1b11a33e" V2_ALEMBIC_CONFIG = "alembic-v2.ini" +# A target before the household revision has these tables, but no model maps +# them any more and a later revision drops them, so the counts name them +# directly. The existing-table check skips them on a target without them. +SIMULATIONS = table("simulations", column("household_id")) +REPORTS = table("reports", column("household_id")) +HOUSEHOLD_JOBS = table("household_jobs") class ScalarExecutor(Protocol): @@ -130,19 +130,19 @@ def count(statement: object, table_name: str) -> int: ), simulation_household_references=count( select(func.count()) - .select_from(Simulation) - .where(Simulation.household_id.is_not(None)), - "simulations", + .select_from(SIMULATIONS) + .where(SIMULATIONS.c.household_id.is_not(None)), + SIMULATIONS.name, ), report_household_references=count( select(func.count()) - .select_from(Report) - .where(Report.household_id.is_not(None)), - "reports", + .select_from(REPORTS) + .where(REPORTS.c.household_id.is_not(None)), + REPORTS.name, ), household_jobs=count( - select(func.count()).select_from(HouseholdJob), - "household_jobs", + select(func.count()).select_from(HOUSEHOLD_JOBS), + HOUSEHOLD_JOBS.name, ), ) diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index c63655bb6..255af0d6f 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -17,6 +17,11 @@ } # Import in dependency order so every reviewed table registers exactly once. +# retired_tables declares, without models, the tables and columns a later +# revision drops; see its docstring. +from policyengine_api.data.v2.models import ( # noqa: E402 + retired_tables as retired_tables, +) from policyengine_api.data.v2.models.associations import ( # noqa: E402 UserHouseholdAssociation, UserPolicy, @@ -36,8 +41,6 @@ ) from policyengine_api.data.v2.models.households import ( # noqa: E402 Household, - HouseholdJob, - HouseholdJobStatus, ) from policyengine_api.data.v2.models.metadata import ( # noqa: E402 Dataset, @@ -60,31 +63,6 @@ LegacyUserMapping, LegacyUserPolicyMapping, ) -from policyengine_api.data.v2.models.reports import ( # noqa: E402 - AggregateOutput, - AggregateType, - BudgetSummary, - ChangeAggregate, - CongressionalDistrictImpact, - ConstituencyImpact, - DecileImpact, - DecileType, - Inequality, - IntraDecileImpact, - LocalAuthorityImpact, - OutputStatus, - Poverty, - ProgramStatistics, - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, -) -from policyengine_api.data.v2.models.simulations import ( # noqa: E402 - Simulation, - SimulationStatus, - SimulationType, -) from policyengine_api.data.v2.models.users import ( # noqa: E402 User, UserIdentity, @@ -96,43 +74,20 @@ __all__ = [ "V2_METADATA", - "AggregateOutput", - "AggregateType", - "BudgetSummary", - "ChangeAggregate", - "CongressionalDistrictImpact", - "ConstituencyImpact", "Dataset", "DatasetVersion", - "DecileImpact", - "DecileType", "Dynamic", "Household", - "HouseholdJob", - "HouseholdJobStatus", - "Inequality", - "IntraDecileImpact", "LegacyHouseholdMapping", "LegacyPolicyMapping", "LegacyUserMapping", "LegacyUserPolicyMapping", - "LocalAuthorityImpact", - "OutputStatus", "Parameter", "ParameterNode", "ParameterValue", "Policy", - "Poverty", - "ProgramStatistics", "Region", "RegionType", - "Report", - "ReportRun", - "ReportRunStatus", - "ReportRunTrigger", - "Simulation", - "SimulationStatus", - "SimulationType", "Stage12AggregationStatus", "Stage12ComparisonReport", "Stage12ComparisonSimulation", diff --git a/policyengine_api/data/v2/models/associations.py b/policyengine_api/data/v2/models/associations.py index 2c5cd21b1..acc95d433 100644 --- a/policyengine_api/data/v2/models/associations.py +++ b/policyengine_api/data/v2/models/associations.py @@ -4,11 +4,10 @@ nullable country-scoped v1 (legacy) reference, or both; a database check requires at least one. -The simulation and report associations still carry the pre-Stage 12A -``country`` and ``label`` columns. Nothing reads or writes them; they stay -nullable only so the previous Public API revision, whose models still map -them, can load these rows during a deploy or rollback. A later revision drops -them once that revision can no longer be deployed. +The simulation and report association tables still have the pre-Stage 12A +``country`` and ``label`` columns, which a later revision drops. These models +don't declare them, so nothing reads or writes them; ``retired_tables`` +declares them on the tables only so the metadata matches the schema. """ from datetime import datetime @@ -217,9 +216,6 @@ class UserSimulationAssociation(TimestampedModel, table=True): country_id: str = Field(max_length=2) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) - # Retired; see the module docstring. - country: str | None = Field(default=None, max_length=16) - label: str | None = Field(default=None, max_length=255) user: User = Relationship(back_populates="simulation_associations") @@ -299,8 +295,5 @@ class UserReportAssociation(TimestampedModel, table=True): default=None, sa_type=sa.DateTime(timezone=True), ) - # Retired; see the module docstring. - country: str | None = Field(default=None, max_length=16) - label: str | None = Field(default=None, max_length=255) user: User = Relationship(back_populates="report_associations") diff --git a/policyengine_api/data/v2/models/base.py b/policyengine_api/data/v2/models/base.py index fae3a25e2..27bb5317b 100644 --- a/policyengine_api/data/v2/models/base.py +++ b/policyengine_api/data/v2/models/base.py @@ -31,13 +31,13 @@ "multi-column indexes are not expressible by one SQLModel Field." ), "ambiguous_foreign_key_relationships": ( - "Dataset input/output, baseline/reform, and overlapping composite " - "region-default joins need SQLAlchemy relationship hints exposed " - "by SQLModel Relationship." + "Overlapping composite region-default joins need SQLAlchemy " + "relationship hints exposed by SQLModel Relationship." ), - "transaction_conflict_recovery": ( - "Concurrent report idempotency requires a savepoint and bounded " - "IntegrityError recovery around the database uniqueness constraint." + "retired_table_declarations": ( + "Tables and columns a later revision drops stay declared as plain " + "SQLAlchemy tables and columns that no model maps, so the metadata " + "matches the schema while nothing reads or writes them." ), } ) diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 78e76b2ab..83da7d076 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -1,44 +1,26 @@ -"""Canonical SQLModel tables for v2 households and household jobs.""" +"""Canonical SQLModel table for v2 households.""" -from datetime import datetime -from enum import StrEnum from typing import TYPE_CHECKING, Any -from uuid import UUID import sqlalchemy as sa +from sqlalchemy.dialects import postgresql from sqlmodel import Field, Relationship -from policyengine_api.data.v2.models.base import ( - IdentifiedModel, - TimestampedModel, - enum_type, -) -from policyengine_api.data.v2.models.policies import Dynamic, Policy +from policyengine_api.data.v2.models.base import TimestampedModel if TYPE_CHECKING: from policyengine_api.data.v2.models.associations import UserHouseholdAssociation from policyengine_api.data.v2.models.household_mappings import ( LegacyHouseholdMapping, ) - from policyengine_api.data.v2.models.reports import Report - from policyengine_api.data.v2.models.simulations import Simulation - - -class HouseholdJobStatus(StrEnum): - """Lifecycle state of a queued household calculation job.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" class Household(TimestampedModel, table=True): """Immutable household input, stored once per canonicalization version and hash. ``content_hash`` is the SHA-256 of the canonical document, so saving the same - household again finds the existing row. Households that simulations, reports - or saved links use cannot be deleted. + household again finds the existing row. Households that saved links use + cannot be deleted. """ __tablename__ = "households" @@ -81,58 +63,14 @@ class Household(TimestampedModel, table=True): country_id: str = Field(max_length=2) default_year: int | None = None household_data: dict[str, Any] = Field( - sa_type=sa.JSON().with_variant(sa.dialects.postgresql.JSONB(), "postgresql") + sa_type=sa.JSON().with_variant(postgresql.JSONB(), "postgresql") ) canonicalization_version: int content_hash: str = Field(max_length=64) - simulations: list["Simulation"] = Relationship(back_populates="household") - reports: list["Report"] = Relationship(back_populates="household") user_associations: list["UserHouseholdAssociation"] = Relationship( back_populates="household", ) legacy_mappings: list["LegacyHouseholdMapping"] = Relationship( back_populates="household" ) - - -class HouseholdJob(IdentifiedModel, table=True): - """A queued household calculation request, with its status and result or error. - - Deleting the job's policy or dynamic keeps the job but clears the reference. - """ - - __tablename__ = "household_jobs" - __table_args__ = ( - sa.Index("ix_household_jobs_status_created_at", "status", "created_at"), - ) - - country: str = Field(max_length=16) - request_data: dict[str, Any] = Field(sa_type=sa.JSON) - policy_id: UUID | None = Field( - default=None, - foreign_key="policies.id", - ondelete="SET NULL", - ) - dynamic_id: UUID | None = Field( - default=None, - foreign_key="dynamics.id", - ondelete="SET NULL", - ) - status: HouseholdJobStatus = Field( - default=HouseholdJobStatus.PENDING, - sa_type=enum_type(HouseholdJobStatus, "v2_household_job_status"), - ) - error_message: str | None = None - result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - started_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - completed_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - - policy: Policy | None = Relationship(back_populates="household_jobs") - dynamic: Dynamic | None = Relationship(back_populates="household_jobs") diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 8d91864df..faa52b11f 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -16,8 +16,6 @@ if TYPE_CHECKING: from policyengine_api.data.v2.models.policies import Dynamic, Policy - from policyengine_api.data.v2.models.reports import Report - from policyengine_api.data.v2.models.simulations import Simulation class RegionType(StrEnum): @@ -56,7 +54,6 @@ class TaxBenefitModel(TimestampedModel, table=True): back_populates="tax_benefit_model" ) policies: list["Policy"] = Relationship(back_populates="tax_benefit_model") - reports: list["Report"] = Relationship(back_populates="tax_benefit_model") class TaxBenefitModelVersion(IdentifiedModel, table=True): @@ -99,9 +96,6 @@ class TaxBenefitModelVersion(IdentifiedModel, table=True): back_populates="tax_benefit_model_version", cascade_delete=True, ) - simulations: list["Simulation"] = Relationship( - back_populates="tax_benefit_model_version" - ) datasets: list["Dataset"] = Relationship(back_populates="tax_benefit_model_version") regions: list["Region"] = Relationship(back_populates="tax_benefit_model_version") policies: list["Policy"] = Relationship(back_populates="tax_benefit_model_version") @@ -161,8 +155,6 @@ class Region(TimestampedModel, table=True): back_populates="default_for_regions", sa_relationship_kwargs={"viewonly": True}, ) - simulations: list["Simulation"] = Relationship(back_populates="region") - reports: list["Report"] = Relationship(back_populates="region") class Dataset(TimestampedModel, table=True): @@ -215,15 +207,6 @@ class Dataset(TimestampedModel, table=True): back_populates="default_dataset", sa_relationship_kwargs={"viewonly": True}, ) - input_simulations: list["Simulation"] = Relationship( - back_populates="dataset", - sa_relationship_kwargs={"foreign_keys": "Simulation.dataset_id"}, - ) - output_simulations: list["Simulation"] = Relationship( - back_populates="output_dataset", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, - ) - reports: list["Report"] = Relationship(back_populates="dataset") class DatasetVersion(IdentifiedModel, table=True): diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 97e4a619e..6eaf97162 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -11,14 +11,11 @@ if TYPE_CHECKING: from policyengine_api.data.v2.models.associations import UserPolicy - from policyengine_api.data.v2.models.households import HouseholdJob from policyengine_api.data.v2.models.metadata import ( ParameterValue, TaxBenefitModelVersion, ) from policyengine_api.data.v2.models.policy_mappings import LegacyPolicyMapping - from policyengine_api.data.v2.models.reports import Report - from policyengine_api.data.v2.models.simulations import Simulation class Policy(TimestampedModel, table=True): @@ -86,9 +83,6 @@ class Policy(TimestampedModel, table=True): back_populates="policy", cascade_delete=True, ) - simulations: list["Simulation"] = Relationship(back_populates="policy") - household_jobs: list["HouseholdJob"] = Relationship(back_populates="policy") - reports: list["Report"] = Relationship(back_populates="policy") user_associations: list["UserPolicy"] = Relationship( back_populates="policy", ) @@ -110,5 +104,3 @@ class Dynamic(TimestampedModel, table=True): back_populates="dynamic", cascade_delete=True, ) - simulations: list["Simulation"] = Relationship(back_populates="dynamic") - household_jobs: list["HouseholdJob"] = Relationship(back_populates="dynamic") diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py deleted file mode 100644 index d10c36947..000000000 --- a/policyengine_api/data/v2/models/reports.py +++ /dev/null @@ -1,651 +0,0 @@ -"""Canonical SQLModel tables for stable reports, runs, and run outputs.""" - -from datetime import datetime -from enum import StrEnum -from typing import Any -from uuid import UUID - -import sqlalchemy as sa -from sqlmodel import Field, Relationship - -from policyengine_api.data.v2.models.base import ( - IdentifiedModel, - TimestampedModel, - enum_type, -) -from policyengine_api.data.v2.models.households import Household -from policyengine_api.data.v2.models.metadata import ( - Dataset, - Region, - TaxBenefitModel, -) -from policyengine_api.data.v2.models.policies import Policy -from policyengine_api.data.v2.models.simulations import Simulation -from policyengine_api.data.v2.models.users import User - - -class ReportRunStatus(StrEnum): - """Lifecycle state of a report run; finished runs record a completion time.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - - -class ReportRunTrigger(StrEnum): - """Why a report run started: the report's first run, a manual or a system rerun.""" - - INITIAL = "initial" - MANUAL = "manual" - SYSTEM = "system" - - -class OutputStatus(StrEnum): - """Lifecycle state of one computed output row of a report run.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - - -class AggregateType(StrEnum): - """How an aggregate combines a variable's values.""" - - SUM = "sum" - MEAN = "mean" - COUNT = "count" - - -class DecileType(StrEnum): - """Whether deciles rank by income or by wealth.""" - - INCOME = "income" - WEALTH = "wealth" - - -class Report(TimestampedModel, table=True): - """Stable report definition; execution state belongs to ReportRun.""" - - __tablename__ = "reports" - __table_args__ = ( - sa.CheckConstraint( - "year IS NULL OR year BETWEEN 1900 AND 2200", - name="ck_reports_year", - ), - sa.Index("ix_reports_country_type_created_at", "country", "type", "created_at"), - ) - - label: str = Field(max_length=255) - description: str | None = None - country: str = Field(max_length=16) - type: str | None = Field(default=None, max_length=128) - user_id: UUID | None = Field( - default=None, - foreign_key="users.id", - ondelete="SET NULL", - ) - tax_benefit_model_id: UUID = Field( - foreign_key="tax_benefit_models.id", - ondelete="RESTRICT", - index=True, - ) - policy_id: UUID | None = Field( - default=None, - foreign_key="policies.id", - ondelete="SET NULL", - ) - baseline_simulation_id: UUID | None = Field( - default=None, - foreign_key="simulations.id", - ondelete="SET NULL", - ) - reform_simulation_id: UUID | None = Field( - default=None, - foreign_key="simulations.id", - ondelete="SET NULL", - ) - household_id: UUID | None = Field( - default=None, - foreign_key="households.id", - ondelete="RESTRICT", - ) - dataset_id: UUID | None = Field( - default=None, - foreign_key="datasets.id", - ondelete="SET NULL", - ) - region_id: UUID | None = Field( - default=None, - foreign_key="regions.id", - ondelete="SET NULL", - ) - year: int | None = None - inputs: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) - - user: User | None = Relationship(back_populates="reports") - tax_benefit_model: TaxBenefitModel = Relationship(back_populates="reports") - policy: Policy | None = Relationship(back_populates="reports") - baseline_simulation: Simulation | None = Relationship( - back_populates="baseline_reports", - sa_relationship_kwargs={"foreign_keys": "Report.baseline_simulation_id"}, - ) - reform_simulation: Simulation | None = Relationship( - back_populates="reform_reports", - sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, - ) - household: Household | None = Relationship(back_populates="reports") - dataset: Dataset | None = Relationship(back_populates="reports") - region: Region | None = Relationship(back_populates="reports") - runs: list["ReportRun"] = Relationship( - back_populates="report", - cascade_delete=True, - ) - - -class ReportRun(TimestampedModel, table=True): - """One immutable-version execution attempt for a stable report.""" - - __tablename__ = "report_runs" - __table_args__ = ( - sa.UniqueConstraint( - "report_id", - "idempotency_key", - name="uq_report_runs_report_idempotency_key", - ), - sa.CheckConstraint( - "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", - name="ck_report_runs_terminal_completion", - ), - sa.Index( - "ix_report_runs_current_output", - "report_id", - "status", - "country_package_version", - "policyengine_version", - "completed_at", - "id", - ), - ) - - report_id: UUID = Field( - foreign_key="reports.id", - ondelete="CASCADE", - ) - country_package_version: str = Field(max_length=128) - policyengine_version: str = Field(max_length=128) - status: ReportRunStatus = Field( - default=ReportRunStatus.PENDING, - sa_type=enum_type(ReportRunStatus, "v2_report_run_status"), - ) - trigger: ReportRunTrigger = Field( - default=ReportRunTrigger.INITIAL, - sa_type=enum_type(ReportRunTrigger, "v2_report_run_trigger"), - ) - idempotency_key: UUID | None = Field(default=None) - started_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - completed_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - error_message: str | None = None - markdown: str | None = Field(default=None, sa_type=sa.Text) - - report: Report = Relationship(back_populates="runs") - aggregates: list["AggregateOutput"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - change_aggregates: list["ChangeAggregate"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - budget_summaries: list["BudgetSummary"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - congressional_district_impacts: list["CongressionalDistrictImpact"] = Relationship( - back_populates="report_run", cascade_delete=True - ) - constituency_impacts: list["ConstituencyImpact"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - decile_impacts: list["DecileImpact"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - inequalities: list["Inequality"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - intra_decile_impacts: list["IntraDecileImpact"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - local_authority_impacts: list["LocalAuthorityImpact"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - poverty_results: list["Poverty"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - program_statistics: list["ProgramStatistics"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) - - -class AggregateOutput(IdentifiedModel, table=True): - """A sum, mean or count of one variable in one simulation of a report run. - - ``filter_config`` optionally restricts which entities are included. - """ - - __tablename__ = "aggregates" - __table_args__ = ( - sa.Index("ix_aggregates_report_run_status", "report_run_id", "status"), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable: str = Field(max_length=512) - aggregate_type: AggregateType = Field( - sa_type=enum_type(AggregateType, "v2_aggregate_type") - ) - entity: str | None = Field(default=None, max_length=128) - filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) - status: OutputStatus = Field( - default=OutputStatus.PENDING, - sa_type=enum_type(OutputStatus, "v2_output_status"), - ) - error_message: str | None = None - result: float | None = None - - report_run: ReportRun = Relationship(back_populates="aggregates") - - -class ChangeAggregate(IdentifiedModel, table=True): - """An aggregate of a variable's change from baseline to reform in a report run. - - ``change_geq`` and ``change_leq`` optionally restrict it to entities whose - change falls within those bounds. - """ - - __tablename__ = "change_aggregates" - __table_args__ = ( - sa.Index( - "ix_change_aggregates_report_run_status", - "report_run_id", - "status", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable: str = Field(max_length=512) - aggregate_type: AggregateType = Field( - sa_type=enum_type(AggregateType, "v2_aggregate_type"), - ) - entity: str | None = Field(default=None, max_length=128) - filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) - change_geq: float | None = None - change_leq: float | None = None - status: OutputStatus = Field( - default=OutputStatus.PENDING, - sa_type=enum_type(OutputStatus, "v2_output_status"), - ) - error_message: str | None = None - result: float | None = None - - report_run: ReportRun = Relationship(back_populates="change_aggregates") - - -class BudgetSummary(IdentifiedModel, table=True): - """Baseline and reform totals of one budget variable for a report run. - - One row per run, variable and entity. - """ - - __tablename__ = "budget_summary" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "variable_name", - "entity", - name="uq_budget_summary_run_variable_entity", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable_name: str = Field(max_length=512) - entity: str = Field(max_length=128) - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - - report_run: ReportRun = Relationship(back_populates="budget_summaries") - - -class DecileImpact(IdentifiedModel, table=True): - """Mean income change and better/worse-off counts in one income quantile. - - ``decile`` runs from 1 to ``quantiles``; there is one row per run, income - variable, entity, quantile count and decile. - """ - - __tablename__ = "decile_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "income_variable", - "entity", - "decile", - "quantiles", - name="uq_decile_impacts_run_measure", - ), - sa.CheckConstraint("quantiles > 0", name="ck_decile_impacts_quantiles"), - sa.CheckConstraint( - "decile BETWEEN 1 AND quantiles", - name="ck_decile_impacts_decile", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - income_variable: str = Field(max_length=512) - entity: str = Field(max_length=128) - decile: int - quantiles: int = 10 - baseline_mean: float | None = None - reform_mean: float | None = None - absolute_change: float | None = None - relative_change: float | None = None - count_better_off: float | None = None - count_worse_off: float | None = None - count_no_change: float | None = None - - report_run: ReportRun = Relationship(back_populates="decile_impacts") - - -class IntraDecileImpact(IdentifiedModel, table=True): - """Shares of one decile whose income rises or falls by more or less than 5%. - - One row per run, decile type and decile, where ``decile`` is 0 to 10. - """ - - __tablename__ = "intra_decile_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "decile_type", - "decile", - name="uq_intra_decile_impacts_run_type_decile", - ), - sa.CheckConstraint( - "decile BETWEEN 0 AND 10", - name="ck_intra_decile_impacts_decile", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - decile_type: DecileType = Field( - default=DecileType.INCOME, - sa_type=enum_type(DecileType, "v2_decile_type"), - ) - decile: int - lose_more_than_5pct: float | None = None - lose_less_than_5pct: float | None = None - no_change: float | None = None - gain_less_than_5pct: float | None = None - gain_more_than_5pct: float | None = None - - report_run: ReportRun = Relationship(back_populates="intra_decile_impacts") - - -class Inequality(IdentifiedModel, table=True): - """Gini and top/bottom income shares for one simulation of a report run.""" - - __tablename__ = "inequality" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "simulation_id", - "income_variable", - "entity", - name="uq_inequality_run_simulation_measure", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - income_variable: str = Field(max_length=512) - entity: str = Field(default="household", max_length=128) - gini: float | None = None - top_10_share: float | None = None - top_1_share: float | None = None - bottom_50_share: float | None = None - - report_run: ReportRun = Relationship(back_populates="inequalities") - - -class Poverty(IdentifiedModel, table=True): - """Poverty headcount and rate of one type for one simulation of a report run.""" - - __tablename__ = "poverty" - __table_args__ = ( - sa.Index( - "ix_poverty_run_simulation_type", - "report_run_id", - "simulation_id", - "poverty_type", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - poverty_type: str = Field(max_length=128) - entity: str = Field(default="person", max_length=128) - filter_variable: str | None = Field(default=None, max_length=512) - headcount: float | None = None - total_population: float | None = None - rate: float | None = None - - report_run: ReportRun = Relationship(back_populates="poverty_results") - - -class ProgramStatistics(IdentifiedModel, table=True): - """Baseline and reform totals, counts, winners and losers for one program. - - One row per run, program and entity; ``is_tax`` marks taxes. - """ - - __tablename__ = "program_statistics" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "program_name", - "entity", - name="uq_program_statistics_run_program_entity", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - program_name: str = Field(max_length=512) - entity: str = Field(max_length=128) - is_tax: bool = False - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - baseline_count: float | None = None - reform_count: float | None = None - winners: float | None = None - losers: float | None = None - - report_run: ReportRun = Relationship(back_populates="program_statistics") - - -class GeographicImpactBase(IdentifiedModel): - """Shared non-table fields for geographic report-run impacts.""" - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - average_household_income_change: float - relative_household_income_change: float - population: float - - -class CongressionalDistrictImpact(GeographicImpactBase, table=True): - """Household income change and population in one US congressional district. - - One row per run and district GEOID. - """ - - __tablename__ = "congressional_district_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "district_geoid", - name="uq_congressional_district_impacts_run_geoid", - ), - ) - - district_geoid: int - state_fips: int - district_number: int - - report_run: ReportRun = Relationship( - back_populates="congressional_district_impacts" - ) - - -class ConstituencyImpact(GeographicImpactBase, table=True): - """Household income change in one UK constituency, with its map coordinates. - - One row per run and constituency code. - """ - - __tablename__ = "constituency_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "constituency_code", - name="uq_constituency_impacts_run_code", - ), - ) - - constituency_code: str = Field(max_length=64) - constituency_name: str = Field(max_length=255) - x: int - y: int - - report_run: ReportRun = Relationship(back_populates="constituency_impacts") - - -class LocalAuthorityImpact(GeographicImpactBase, table=True): - """Household income change in one UK local authority, with its map coordinates. - - One row per run and local authority code. - """ - - __tablename__ = "local_authority_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "local_authority_code", - name="uq_local_authority_impacts_run_code", - ), - ) - - local_authority_code: str = Field(max_length=64) - local_authority_name: str = Field(max_length=255) - x: int - y: int - - report_run: ReportRun = Relationship(back_populates="local_authority_impacts") diff --git a/policyengine_api/data/v2/models/retired_tables.py b/policyengine_api/data/v2/models/retired_tables.py new file mode 100644 index 000000000..9959fcdd1 --- /dev/null +++ b/policyengine_api/data/v2/models/retired_tables.py @@ -0,0 +1,519 @@ +"""Schema-only declarations of the v2 tables and columns a later revision drops. + +The simulation service owns household jobs, simulations, reports, report runs +and their outputs in its own ``simulation`` schema. The Public API's copies of +those 15 tables, and the retired ``country`` and ``label`` columns of the +simulation and report associations, stay in the database until a later v2 +revision drops them. That revision runs while this release serves traffic, so +this release must never read or write them. + +They stay declared here only so the v2 metadata keeps matching the schema: +``alembic check`` and the migration environment's head-schema validation +compare the two. They are plain SQLAlchemy tables and columns that no model +maps, so no ORM query, insert, update, relationship or cascade reaches them. +Nothing imports this module except the models package, which registers it. +Deleting it is what makes autogeneration produce the revision that drops them. +""" + +from sqlalchemy import ( + JSON, + Boolean, + CheckConstraint, + Column, + DateTime, + Enum, + Float, + ForeignKey, + Index, + Integer, + String, + Table, + Text, + UniqueConstraint, + Uuid, + func, +) +from sqlmodel import SQLModel + +from policyengine_api.data.v2.models.associations import ( + UserReportAssociation, + UserSimulationAssociation, +) + +METADATA = SQLModel.metadata + +# Each lifecycle enum has the same four values. +LIFECYCLE_STATES = ("pending", "running", "succeeded", "failed") +REPORT_RUN_TRIGGERS = ("initial", "manual", "system") +AGGREGATE_TYPES = ("sum", "mean", "count") +DECILE_TYPES = ("income", "wealth") + +# The retired association columns, by table, with their string lengths. +RETIRED_ASSOCIATION_COLUMNS = { + UserSimulationAssociation.__tablename__: {"country": 16, "label": 255}, + UserReportAssociation.__tablename__: {"country": 16, "label": 255}, +} + + +def _enum(values: tuple[str, ...], name: str) -> Enum: + return Enum(*values, name=name) + + +def _identified_columns() -> list[Column]: + return [ + Column("id", Uuid(), primary_key=True), + Column( + "created_at", + DateTime(timezone=True), + server_default=func.now(), + nullable=False, + ), + ] + + +def _timestamped_columns() -> list[Column]: + return [ + *_identified_columns(), + Column( + "updated_at", + DateTime(timezone=True), + server_default=func.now(), + nullable=False, + ), + ] + + +def _reference( + name: str, + target: str, + ondelete: str, + *, + nullable: bool, + index: bool = False, +) -> Column: + return Column( + name, + Uuid(), + ForeignKey(target, ondelete=ondelete), + nullable=nullable, + index=index, + ) + + +def _report_run_reference() -> Column: + return _reference("report_run_id", "report_runs.id", "CASCADE", nullable=False) + + +def _simulation_reference(name: str) -> Column: + return _reference(name, "simulations.id", "RESTRICT", nullable=False) + + +def _comparison_references() -> list[Column]: + return [ + _report_run_reference(), + _simulation_reference("baseline_simulation_id"), + _simulation_reference("reform_simulation_id"), + ] + + +def _geographic_impact_columns() -> list[Column]: + return [ + *_identified_columns(), + *_comparison_references(), + Column("average_household_income_change", Float(), nullable=False), + Column("relative_household_income_change", Float(), nullable=False), + Column("population", Float(), nullable=False), + ] + + +HOUSEHOLD_JOBS = Table( + "household_jobs", + METADATA, + *_identified_columns(), + Column("country", String(16), nullable=False), + Column("request_data", JSON(), nullable=False), + _reference("policy_id", "policies.id", "SET NULL", nullable=True), + _reference("dynamic_id", "dynamics.id", "SET NULL", nullable=True), + Column( + "status", + _enum(LIFECYCLE_STATES, "v2_household_job_status"), + nullable=False, + ), + Column("error_message", String(), nullable=True), + Column("result", JSON(), nullable=True), + Column("started_at", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + Index("ix_household_jobs_status_created_at", "status", "created_at"), +) + +SIMULATIONS = Table( + "simulations", + METADATA, + *_timestamped_columns(), + Column( + "simulation_type", + _enum(("household", "economy"), "v2_simulation_type"), + nullable=False, + ), + _reference("dataset_id", "datasets.id", "RESTRICT", nullable=True), + _reference("household_id", "households.id", "RESTRICT", nullable=True), + _reference("policy_id", "policies.id", "SET NULL", nullable=True), + _reference("dynamic_id", "dynamics.id", "SET NULL", nullable=True), + _reference( + "tax_benefit_model_version_id", + "tax_benefit_model_versions.id", + "RESTRICT", + nullable=False, + index=True, + ), + _reference("output_dataset_id", "datasets.id", "SET NULL", nullable=True), + _reference("region_id", "regions.id", "SET NULL", nullable=True), + Column("status", _enum(LIFECYCLE_STATES, "v2_simulation_status"), nullable=False), + Column("error_message", String(), nullable=True), + Column("filter_field", String(128), nullable=True), + Column("filter_value", String(255), nullable=True), + Column("filter_strategy", String(64), nullable=True), + Column("year", Integer(), nullable=True), + Column("started_at", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + Column("household_result", JSON(), nullable=True), + CheckConstraint( + "(simulation_type = 'household' AND household_id IS NOT NULL " + "AND dataset_id IS NULL) OR " + "(simulation_type = 'economy' AND dataset_id IS NOT NULL " + "AND household_id IS NULL)", + name="ck_simulations_type_input", + ), + CheckConstraint( + "(filter_field IS NULL) = (filter_value IS NULL)", + name="ck_simulations_filter_pair", + ), + CheckConstraint( + "year IS NULL OR year BETWEEN 1900 AND 2200", + name="ck_simulations_year", + ), + Index("ix_simulations_status_created_at", "status", "created_at"), +) + +REPORTS = Table( + "reports", + METADATA, + *_timestamped_columns(), + Column("label", String(255), nullable=False), + Column("description", String(), nullable=True), + Column("country", String(16), nullable=False), + Column("type", String(128), nullable=True), + _reference("user_id", "users.id", "SET NULL", nullable=True), + _reference( + "tax_benefit_model_id", + "tax_benefit_models.id", + "RESTRICT", + nullable=False, + index=True, + ), + _reference("policy_id", "policies.id", "SET NULL", nullable=True), + _reference("baseline_simulation_id", "simulations.id", "SET NULL", nullable=True), + _reference("reform_simulation_id", "simulations.id", "SET NULL", nullable=True), + _reference("household_id", "households.id", "RESTRICT", nullable=True), + _reference("dataset_id", "datasets.id", "SET NULL", nullable=True), + _reference("region_id", "regions.id", "SET NULL", nullable=True), + Column("year", Integer(), nullable=True), + Column("inputs", JSON(), nullable=False), + CheckConstraint( + "year IS NULL OR year BETWEEN 1900 AND 2200", + name="ck_reports_year", + ), + Index("ix_reports_country_type_created_at", "country", "type", "created_at"), +) + +REPORT_RUNS = Table( + "report_runs", + METADATA, + *_timestamped_columns(), + _reference("report_id", "reports.id", "CASCADE", nullable=False), + Column("country_package_version", String(128), nullable=False), + Column("policyengine_version", String(128), nullable=False), + Column("status", _enum(LIFECYCLE_STATES, "v2_report_run_status"), nullable=False), + Column( + "trigger", + _enum(REPORT_RUN_TRIGGERS, "v2_report_run_trigger"), + nullable=False, + ), + Column("idempotency_key", Uuid(), nullable=True), + Column("started_at", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + Column("error_message", String(), nullable=True), + Column("markdown", Text(), nullable=True), + UniqueConstraint( + "report_id", + "idempotency_key", + name="uq_report_runs_report_idempotency_key", + ), + CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name="ck_report_runs_terminal_completion", + ), + Index( + "ix_report_runs_current_output", + "report_id", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ), +) + +AGGREGATES = Table( + "aggregates", + METADATA, + *_identified_columns(), + _report_run_reference(), + _simulation_reference("simulation_id"), + Column("variable", String(512), nullable=False), + Column( + "aggregate_type", + _enum(AGGREGATE_TYPES, "v2_aggregate_type"), + nullable=False, + ), + Column("entity", String(128), nullable=True), + Column("filter_config", JSON(), nullable=False), + Column("status", _enum(LIFECYCLE_STATES, "v2_output_status"), nullable=False), + Column("error_message", String(), nullable=True), + Column("result", Float(), nullable=True), + Index("ix_aggregates_report_run_status", "report_run_id", "status"), +) + +CHANGE_AGGREGATES = Table( + "change_aggregates", + METADATA, + *_identified_columns(), + *_comparison_references(), + Column("variable", String(512), nullable=False), + Column( + "aggregate_type", + _enum(AGGREGATE_TYPES, "v2_aggregate_type"), + nullable=False, + ), + Column("entity", String(128), nullable=True), + Column("filter_config", JSON(), nullable=False), + Column("change_geq", Float(), nullable=True), + Column("change_leq", Float(), nullable=True), + Column("status", _enum(LIFECYCLE_STATES, "v2_output_status"), nullable=False), + Column("error_message", String(), nullable=True), + Column("result", Float(), nullable=True), + Index("ix_change_aggregates_report_run_status", "report_run_id", "status"), +) + +BUDGET_SUMMARY = Table( + "budget_summary", + METADATA, + *_identified_columns(), + *_comparison_references(), + Column("variable_name", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("baseline_total", Float(), nullable=True), + Column("reform_total", Float(), nullable=True), + Column("change", Float(), nullable=True), + UniqueConstraint( + "report_run_id", + "variable_name", + "entity", + name="uq_budget_summary_run_variable_entity", + ), +) + +DECILE_IMPACTS = Table( + "decile_impacts", + METADATA, + *_identified_columns(), + *_comparison_references(), + Column("income_variable", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("decile", Integer(), nullable=False), + Column("quantiles", Integer(), nullable=False), + Column("baseline_mean", Float(), nullable=True), + Column("reform_mean", Float(), nullable=True), + Column("absolute_change", Float(), nullable=True), + Column("relative_change", Float(), nullable=True), + Column("count_better_off", Float(), nullable=True), + Column("count_worse_off", Float(), nullable=True), + Column("count_no_change", Float(), nullable=True), + UniqueConstraint( + "report_run_id", + "income_variable", + "entity", + "decile", + "quantiles", + name="uq_decile_impacts_run_measure", + ), + CheckConstraint("quantiles > 0", name="ck_decile_impacts_quantiles"), + CheckConstraint("decile BETWEEN 1 AND quantiles", name="ck_decile_impacts_decile"), +) + +INTRA_DECILE_IMPACTS = Table( + "intra_decile_impacts", + METADATA, + *_identified_columns(), + *_comparison_references(), + Column("decile_type", _enum(DECILE_TYPES, "v2_decile_type"), nullable=False), + Column("decile", Integer(), nullable=False), + Column("lose_more_than_5pct", Float(), nullable=True), + Column("lose_less_than_5pct", Float(), nullable=True), + Column("no_change", Float(), nullable=True), + Column("gain_less_than_5pct", Float(), nullable=True), + Column("gain_more_than_5pct", Float(), nullable=True), + UniqueConstraint( + "report_run_id", + "decile_type", + "decile", + name="uq_intra_decile_impacts_run_type_decile", + ), + CheckConstraint( + "decile BETWEEN 0 AND 10", + name="ck_intra_decile_impacts_decile", + ), +) + +INEQUALITY = Table( + "inequality", + METADATA, + *_identified_columns(), + _report_run_reference(), + _simulation_reference("simulation_id"), + Column("income_variable", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("gini", Float(), nullable=True), + Column("top_10_share", Float(), nullable=True), + Column("top_1_share", Float(), nullable=True), + Column("bottom_50_share", Float(), nullable=True), + UniqueConstraint( + "report_run_id", + "simulation_id", + "income_variable", + "entity", + name="uq_inequality_run_simulation_measure", + ), +) + +POVERTY = Table( + "poverty", + METADATA, + *_identified_columns(), + _report_run_reference(), + _simulation_reference("simulation_id"), + Column("poverty_type", String(128), nullable=False), + Column("entity", String(128), nullable=False), + Column("filter_variable", String(512), nullable=True), + Column("headcount", Float(), nullable=True), + Column("total_population", Float(), nullable=True), + Column("rate", Float(), nullable=True), + Index( + "ix_poverty_run_simulation_type", + "report_run_id", + "simulation_id", + "poverty_type", + ), +) + +PROGRAM_STATISTICS = Table( + "program_statistics", + METADATA, + *_identified_columns(), + *_comparison_references(), + Column("program_name", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("is_tax", Boolean(), nullable=False), + Column("baseline_total", Float(), nullable=True), + Column("reform_total", Float(), nullable=True), + Column("change", Float(), nullable=True), + Column("baseline_count", Float(), nullable=True), + Column("reform_count", Float(), nullable=True), + Column("winners", Float(), nullable=True), + Column("losers", Float(), nullable=True), + UniqueConstraint( + "report_run_id", + "program_name", + "entity", + name="uq_program_statistics_run_program_entity", + ), +) + +CONGRESSIONAL_DISTRICT_IMPACTS = Table( + "congressional_district_impacts", + METADATA, + *_geographic_impact_columns(), + Column("district_geoid", Integer(), nullable=False), + Column("state_fips", Integer(), nullable=False), + Column("district_number", Integer(), nullable=False), + UniqueConstraint( + "report_run_id", + "district_geoid", + name="uq_congressional_district_impacts_run_geoid", + ), +) + +CONSTITUENCY_IMPACTS = Table( + "constituency_impacts", + METADATA, + *_geographic_impact_columns(), + Column("constituency_code", String(64), nullable=False), + Column("constituency_name", String(255), nullable=False), + Column("x", Integer(), nullable=False), + Column("y", Integer(), nullable=False), + UniqueConstraint( + "report_run_id", + "constituency_code", + name="uq_constituency_impacts_run_code", + ), +) + +LOCAL_AUTHORITY_IMPACTS = Table( + "local_authority_impacts", + METADATA, + *_geographic_impact_columns(), + Column("local_authority_code", String(64), nullable=False), + Column("local_authority_name", String(255), nullable=False), + Column("x", Integer(), nullable=False), + Column("y", Integer(), nullable=False), + UniqueConstraint( + "report_run_id", + "local_authority_code", + name="uq_local_authority_impacts_run_code", + ), +) + +RETIRED_TABLES = ( + HOUSEHOLD_JOBS, + SIMULATIONS, + REPORTS, + REPORT_RUNS, + AGGREGATES, + CHANGE_AGGREGATES, + BUDGET_SUMMARY, + DECILE_IMPACTS, + INTRA_DECILE_IMPACTS, + INEQUALITY, + POVERTY, + PROGRAM_STATISTICS, + CONGRESSIONAL_DISTRICT_IMPACTS, + CONSTITUENCY_IMPACTS, + LOCAL_AUTHORITY_IMPACTS, +) + + +def _declare_retired_association_columns() -> None: + """Add the retired columns to the association tables, but not their models. + + The association models no longer declare these columns, so their mappers + never select, insert or update them: a column appended to a table after + its mapper exists stays unmapped. + """ + for table_name, columns in RETIRED_ASSOCIATION_COLUMNS.items(): + for column_name, length in columns.items(): + METADATA.tables[table_name].append_column( + Column(column_name, String(length), nullable=True) + ) + + +_declare_retired_association_columns() diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py deleted file mode 100644 index a6821eadc..000000000 --- a/policyengine_api/data/v2/models/simulations.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Canonical SQLModel table for v2 simulations.""" - -from datetime import datetime -from enum import StrEnum -from typing import TYPE_CHECKING, Any -from uuid import UUID - -import sqlalchemy as sa -from sqlmodel import Field, Relationship - -from policyengine_api.data.v2.models.base import TimestampedModel, enum_type -from policyengine_api.data.v2.models.households import Household -from policyengine_api.data.v2.models.metadata import ( - Dataset, - Region, - TaxBenefitModelVersion, -) -from policyengine_api.data.v2.models.policies import Dynamic, Policy - -if TYPE_CHECKING: - from policyengine_api.data.v2.models.reports import Report - - -class SimulationStatus(StrEnum): - """Lifecycle state of a v2 simulation.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - - -class SimulationType(StrEnum): - """Whether a simulation runs on one household or on a population dataset.""" - - HOUSEHOLD = "household" - ECONOMY = "economy" - - -class Simulation(TimestampedModel, table=True): - """A simulation of exactly one household or one dataset, matching its type. - - Deleting its policy, dynamic, region or output dataset clears that reference; - its household and input dataset cannot be deleted while it exists. - """ - - __tablename__ = "simulations" - __table_args__ = ( - sa.CheckConstraint( - "(simulation_type = 'household' AND household_id IS NOT NULL " - "AND dataset_id IS NULL) OR " - "(simulation_type = 'economy' AND dataset_id IS NOT NULL " - "AND household_id IS NULL)", - name="ck_simulations_type_input", - ), - sa.CheckConstraint( - "(filter_field IS NULL) = (filter_value IS NULL)", - name="ck_simulations_filter_pair", - ), - sa.CheckConstraint( - "year IS NULL OR year BETWEEN 1900 AND 2200", - name="ck_simulations_year", - ), - sa.Index("ix_simulations_status_created_at", "status", "created_at"), - ) - - simulation_type: SimulationType = Field( - default=SimulationType.ECONOMY, - sa_type=enum_type(SimulationType, "v2_simulation_type"), - ) - dataset_id: UUID | None = Field( - default=None, - foreign_key="datasets.id", - ondelete="RESTRICT", - ) - household_id: UUID | None = Field( - default=None, - foreign_key="households.id", - ondelete="RESTRICT", - ) - policy_id: UUID | None = Field( - default=None, - foreign_key="policies.id", - ondelete="SET NULL", - ) - dynamic_id: UUID | None = Field( - default=None, - foreign_key="dynamics.id", - ondelete="SET NULL", - ) - tax_benefit_model_version_id: UUID = Field( - foreign_key="tax_benefit_model_versions.id", - ondelete="RESTRICT", - index=True, - ) - output_dataset_id: UUID | None = Field( - default=None, - foreign_key="datasets.id", - ondelete="SET NULL", - ) - region_id: UUID | None = Field( - default=None, - foreign_key="regions.id", - ondelete="SET NULL", - ) - status: SimulationStatus = Field( - default=SimulationStatus.PENDING, - sa_type=enum_type(SimulationStatus, "v2_simulation_status"), - ) - error_message: str | None = None - filter_field: str | None = Field(default=None, max_length=128) - filter_value: str | None = Field(default=None, max_length=255) - filter_strategy: str | None = Field(default=None, max_length=64) - year: int | None = None - started_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - completed_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - household_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - - dataset: Dataset | None = Relationship( - back_populates="input_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.dataset_id"}, - ) - output_dataset: Dataset | None = Relationship( - back_populates="output_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, - ) - household: Household | None = Relationship(back_populates="simulations") - policy: Policy | None = Relationship(back_populates="simulations") - dynamic: Dynamic | None = Relationship(back_populates="simulations") - tax_benefit_model_version: TaxBenefitModelVersion = Relationship( - back_populates="simulations" - ) - region: Region | None = Relationship(back_populates="simulations") - baseline_reports: list["Report"] = Relationship( - back_populates="baseline_simulation", - sa_relationship_kwargs={"foreign_keys": "Report.baseline_simulation_id"}, - ) - reform_reports: list["Report"] = Relationship( - back_populates="reform_simulation", - sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, - ) diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 85c0ea19a..660e8bd0f 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -23,7 +23,6 @@ UserReportAssociation, UserSimulationAssociation, ) - from policyengine_api.data.v2.models.reports import Report class UserType(StrEnum): @@ -37,8 +36,7 @@ class User(IdentifiedModel, table=True): """A v2 user account, unique by email when one is given. A ``standard`` user always has a primary country. Deleting a user deletes - their sign-in identities and saved associations; reports they own are kept - without an owner. + their sign-in identities and saved associations. """ __tablename__ = "users" @@ -70,7 +68,6 @@ class User(IdentifiedModel, table=True): sa_column_kwargs={"server_default": UserType.STANDARD.value}, ) - reports: list["Report"] = Relationship(back_populates="user") identities: list["UserIdentity"] = Relationship( back_populates="user", cascade_delete=True, diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 8f8c34515..a43e0c00d 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -93,23 +93,23 @@ def _access( TABLE_ACCESS: Final[dict[str, TableAccess]] = { # Alembic's revision record; the catalog writer checks the schema revision. "alembic_version": _access(catalog_writer=READ), - # Simulation-domain result tables, unused by the managed roles. + # Simulation-domain tables, which the simulation service owns in its own + # schema. No managed role uses them, and a later revision drops them. "aggregates": _access(), "budget_summary": _access(), "change_aggregates": _access(), "congressional_district_impacts": _access(), "constituency_impacts": _access(), "decile_impacts": _access(), - "dynamics": _access(), "household_jobs": _access(), "inequality": _access(), "intra_decile_impacts": _access(), "local_authority_impacts": _access(), "poverty": _access(), "program_statistics": _access(), - "report_runs": _access(catalog_writer=READ), - "reports": _access(catalog_writer=READ), - "simulations": _access(catalog_writer=READ), + "report_runs": _access(), + "reports": _access(), + "simulations": _access(), # Metadata catalog: the catalog writer publishes it, the Public API reads it. "dataset_versions": _access(catalog_writer=READ), "datasets": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), @@ -129,6 +129,8 @@ def _access( "legacy_household_mappings": _access(runtime=READ_INSERT_UPDATE), "legacy_policy_mappings": _access(runtime=READ_INSERT_UPDATE), "policies": _access(runtime=READ_INSERT), + # Named parameter values a simulation can apply; no managed role uses them. + "dynamics": _access(), # Stage 10 saved-policy mirror. "legacy_user_mappings": _access(runtime=READ_WRITE), "legacy_user_policy_mappings": _access(runtime=READ_WRITE), diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py deleted file mode 100644 index 14685f150..000000000 --- a/policyengine_api/data/v2/report_runs.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Durable report-definition and report-run operations for API v2-alpha.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -from sqlalchemy.exc import IntegrityError -from sqlmodel import Session, select - -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, POLICYENGINE_VERSION -from policyengine_api.data.v2.models import ( - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, -) -from policyengine_api.data.v2.models.base import utc_now - -if TYPE_CHECKING: - from collections.abc import Mapping - from datetime import datetime - from uuid import UUID - - -class ReportNotFoundError(LookupError): - """Raised when a requested durable report does not exist.""" - - -class ReportRunNotFoundError(LookupError): - """Raised when a requested durable report run does not exist.""" - - -class ReportTypeImmutableError(ValueError): - """Raised when a report type change would reinterpret existing runs.""" - - -class ReportRunStateError(ValueError): - """Raised when a worker transition is invalid for the durable run.""" - - -@dataclass(frozen=True) -class ReportRunPackageVersions: - """The country package and PolicyEngine.py versions one report run uses.""" - - country_package_version: str - policyengine_version: str - - -def _locked_report(session: Session, report_id: UUID) -> Report: - # SQLModel's select remains the query surface. Row locking is the bounded - # SQLAlchemy capability needed to serialize first-run/type updates. - statement = select(Report).where(Report.id == report_id).with_for_update() - report = session.exec(statement).one_or_none() - if report is None: - msg = f"report {report_id} does not exist" - raise ReportNotFoundError(msg) - return report - - -def set_report_type( - session: Session, - *, - report_id: UUID, - report_type: str | None, -) -> Report: - """Set a report's optional type only while it has no execution history.""" - report = _locked_report(session, report_id) - if report.type == report_type: - return report - - existing_run_id = session.exec( - select(ReportRun.id).where(ReportRun.report_id == report_id).limit(1) - ).first() - if existing_run_id is not None: - msg = "report type is immutable after the first report run" - raise ReportTypeImmutableError(msg) - - report.type = report_type - session.add(report) - session.flush() - return report - - -def create_report_run( - session: Session, - *, - report_id: UUID, - versions: ReportRunPackageVersions, - trigger: ReportRunTrigger, - idempotency_key: UUID | None = None, -) -> ReportRun: - """Create one run, or return the run for a retried idempotent request.""" - _locked_report(session, report_id) - if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: - msg = "manual report reruns require an idempotency key" - raise ValueError(msg) - if not versions.country_package_version or not versions.policyengine_version: - msg = "report run package versions must be non-empty" - raise ValueError(msg) - - run = ReportRun( - report_id=report_id, - country_package_version=versions.country_package_version, - policyengine_version=versions.policyengine_version, - trigger=trigger, - idempotency_key=idempotency_key, - ) - if idempotency_key is None: - session.add(run) - session.flush() - return run - - try: - # The report-scoped unique constraint is authoritative under races. - # A savepoint contains the expected conflict without rolling back - # unrelated caller work in the surrounding transaction. - with session.begin_nested(): - session.add(run) - session.flush() - except IntegrityError: - existing = session.exec( - select(ReportRun).where( - ReportRun.report_id == report_id, - ReportRun.idempotency_key == idempotency_key, - ) - ).one_or_none() - if existing is None: - raise - return existing - return run - - -def begin_report_run( - session: Session, - *, - report_run_id: UUID, - started_at: datetime | None = None, -) -> ReportRun: - """Start or resume the same pending/running durable run after worker retry.""" - run = session.exec( - select(ReportRun).where(ReportRun.id == report_run_id).with_for_update() - ).one_or_none() - if run is None: - msg = f"report run {report_run_id} does not exist" - raise ReportRunNotFoundError(msg) - if run.status in {ReportRunStatus.SUCCEEDED, ReportRunStatus.FAILED}: - msg = "a terminal report run cannot be resumed" - raise ReportRunStateError(msg) - if run.status is ReportRunStatus.PENDING: - run.status = ReportRunStatus.RUNNING - run.started_at = started_at or utc_now() - session.add(run) - session.flush() - return run - - -def complete_report_run( - session: Session, - *, - report_run_id: UUID, - completed_at: datetime | None = None, - markdown: str | None = None, -) -> ReportRun: - """Mark the selected durable run successful without replacing history.""" - run = begin_report_run(session, report_run_id=report_run_id) - run.status = ReportRunStatus.SUCCEEDED - run.completed_at = completed_at or utc_now() - run.markdown = markdown - run.error_message = None - session.add(run) - session.flush() - return run - - -def fail_report_run( - session: Session, - *, - report_run_id: UUID, - error_message: str, - completed_at: datetime | None = None, -) -> ReportRun: - """Mark the selected durable run failed without deleting older success.""" - run = begin_report_run(session, report_run_id=report_run_id) - run.status = ReportRunStatus.FAILED - run.completed_at = completed_at or utc_now() - run.error_message = error_message - session.add(run) - session.flush() - return run - - -def select_current_report_run( - session: Session, - *, - report_id: UUID, - country_package_versions: Mapping[str, str] = COUNTRY_PACKAGE_VERSIONS, - policyengine_version: str = POLICYENGINE_VERSION, -) -> ReportRun | None: - """Select the deterministic current successful run for deployed versions.""" - report = session.get(Report, report_id) - if report is None: - msg = f"report {report_id} does not exist" - raise ReportNotFoundError(msg) - country_package_version = country_package_versions.get(report.country) - if country_package_version is None: - return None - - statement = ( - select(ReportRun) - .where( - ReportRun.report_id == report_id, - ReportRun.status == ReportRunStatus.SUCCEEDED, - ReportRun.country_package_version == country_package_version, - ReportRun.policyengine_version == policyengine_version, - ReportRun.completed_at.is_not(None), - ) - .order_by(ReportRun.completed_at.desc(), ReportRun.id.desc()) - .limit(1) - ) - return session.exec(statement).first() diff --git a/pyproject.toml b/pyproject.toml index ed101271e..7342f72d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -241,7 +241,6 @@ runtime-evaluated-base-classes = [ "policyengine_api.data.v2.models.base.IdentifiedModel", "policyengine_api.data.v2.models.base.TimestampedModel", "policyengine_api.data.v2.models.comparison_runs.Stage12RunTimestamps", - "policyengine_api.data.v2.models.reports.GeographicImpactBase", "policyengine_api.fastapi_routes.v2.households.document_models.HouseholdEntityRecord", "policyengine_api.fastapi_routes.v2.households.document_models.StrictHouseholdDocumentModel", "policyengine_api.fastapi_routes.v2.households.response_models.HouseholdSuccessResponse", diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index 06e48f64c..43263f4ad 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -22,13 +22,10 @@ from policyengine_api.data.v2.models import ( Dataset, DatasetVersion, - Region, - Report, - ReportRun, - Simulation, TaxBenefitModel, TaxBenefitModelVersion, ) +from policyengine_api.data.v2.privileges import CATALOG_WRITER_ROLE from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL from tests.fixtures.v2_catalog import normalized_catalog @@ -216,7 +213,7 @@ def test_same_version_with_changed_content_fails_without_mutation( assert _counts(publication_engine) == before_counts -def test_republish_preserves_existing_run_and_dataset_version_records( +def test_republish_preserves_existing_dataset_version_records( publication_engine: Engine, ) -> None: catalog = normalized_catalog() @@ -231,18 +228,6 @@ def test_republish_preserves_existing_run_and_dataset_version_records( TaxBenefitModelVersion.model_id == model.id ) ).one() - input_dataset = session.exec( - select(Dataset).where( - Dataset.tax_benefit_model_version_id == model_version.id, - Dataset.name == "populace_us_2024", - ) - ).one() - region = session.exec( - select(Region).where( - Region.tax_benefit_model_version_id == model_version.id, - Region.code == "us", - ) - ).one() output_dataset = Dataset( name="existing-run-output", description="Existing generated output", @@ -259,52 +244,40 @@ def test_republish_preserves_existing_run_and_dataset_version_records( dataset_id=output_dataset.id, tax_benefit_model_id=model.id, ) - simulation = Simulation( - dataset_id=input_dataset.id, - output_dataset_id=output_dataset.id, - tax_benefit_model_version_id=model_version.id, - region_id=region.id, - ) - report = Report( - label="Existing report", - country="us", - tax_benefit_model_id=model.id, - dataset_id=input_dataset.id, - region_id=region.id, - ) - session.add_all((dataset_version, simulation, report)) - session.flush() - report_run = ReportRun( - report_id=report.id, - country_package_version="existing-country-version", - policyengine_version="existing-policyengine-version", - ) - session.add(report_run) + session.add(dataset_version) session.commit() - protected_ids = ( - dataset_version.id, - simulation.id, - report.id, - report_run.id, - simulation.dataset_id, - simulation.output_dataset_id, - report.dataset_id, - ) + protected_ids = (dataset_version.id, output_dataset.id) publish_catalog(publication_engine, catalog) with Session(publication_engine) as session: persisted_dataset_version = session.get(DatasetVersion, protected_ids[0]) - persisted_simulation = session.get(Simulation, protected_ids[1]) - persisted_report = session.get(Report, protected_ids[2]) - persisted_report_run = session.get(ReportRun, protected_ids[3]) assert persisted_dataset_version is not None - assert persisted_simulation is not None - assert persisted_report is not None - assert persisted_report_run is not None - assert persisted_simulation.dataset_id == protected_ids[4] - assert persisted_simulation.output_dataset_id == protected_ids[5] - assert persisted_report.dataset_id == protected_ids[6] + assert persisted_dataset_version.dataset_id == protected_ids[1] + assert session.get(Dataset, protected_ids[1]) is not None + + +def test_catalog_writer_publishes_without_access_to_simulation_tables( + publication_engine: Engine, + v2_engine_as_role, +) -> None: + # A later v2 revision drops simulations, reports and report_runs while + # this release serves traffic, so publication must work without them. + with publication_engine.begin() as connection: + connection.execute( + text( + "REVOKE ALL ON simulations, reports, report_runs " + f"FROM {CATALOG_WRITER_ROLE}" + ) + ) + catalog_writer_engine = v2_engine_as_role(CATALOG_WRITER_ROLE, poolclass=NullPool) + try: + evidence = publish_catalog(catalog_writer_engine, normalized_catalog()) + finally: + catalog_writer_engine.dispose() + + assert evidence.as_dict()["outcome"] == "ok" + assert _counts(publication_engine)["tax_benefit_models"] == 2 def test_new_version_is_additive_and_concurrent_retries_serialize( diff --git a/tests/unit/v2/test_catalog_publication.py b/tests/unit/v2/test_catalog_publication.py index a3cbe0d65..6f7ba6583 100644 --- a/tests/unit/v2/test_catalog_publication.py +++ b/tests/unit/v2/test_catalog_publication.py @@ -145,6 +145,31 @@ def test_revision_check_rejects_non_postgres_missing_and_wrong_revisions() -> No ) +class _CountRecordingConnection: + """Records the tables each row-count statement selects from.""" + + def __init__(self) -> None: + self.counted_tables: list[str] = [] + + def execute(self, statement): + self.counted_tables.extend( + from_clause.name for from_clause in statement.get_final_froms() + ) + return _ScalarResult(value=0) + + +def test_protected_row_counts_read_only_dataset_versions() -> None: + # The simulation service owns simulations, reports and report runs, and a + # later v2 revision drops the Public API's copies of those tables, so + # publication must not read them. + connection = _CountRecordingConnection() + + counts = publication._protected_row_counts(connection) + + assert counts == (0,) + assert connection.counted_tables == ["dataset_versions"] + + def test_copy_streams_rows_through_psycopg_without_an_orm_write() -> None: connection = FakeConnection() rows = ((index, f"row-{index}") for index in range(3)) diff --git a/tests/unit/v2/test_household_migration_qualification.py b/tests/unit/v2/test_household_migration_qualification.py index 38199cdc9..8757498b6 100644 --- a/tests/unit/v2/test_household_migration_qualification.py +++ b/tests/unit/v2/test_household_migration_qualification.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING import pytest +from sqlalchemy.dialects import postgresql from policyengine_api.data.v2 import household_migration_qualification as qualification from policyengine_api.data.v2.settings import ( @@ -17,6 +18,8 @@ if TYPE_CHECKING: from collections.abc import Iterator + from sqlalchemy.sql import ClauseElement + ENVIRONMENT = { V2_MIGRATION_DATABASE_URL: ( "postgresql+psycopg://migration:test-password@db." @@ -123,6 +126,51 @@ def test_pending_household_revision_runs_row_qualification() -> None: assert engine.disposed +class RecordingExecutor: + """Records each count statement, compiled for PostgreSQL, and returns zero.""" + + def __init__(self) -> None: + self.statements: list[str] = [] + + def scalar(self, statement: ClauseElement) -> int: + compiled = statement.compile(dialect=postgresql.dialect()) + self.statements.append(" ".join(str(compiled).split())) + return 0 + + +def test_counts_query_every_table_the_household_revision_replaced() -> None: + # The simulation, report and household job tables have no models any more; + # the counts still name them for a target before revision 724b1b11a33e. + executor = RecordingExecutor() + + counts = qualification.read_household_data_counts(executor) + + assert counts == qualification.HouseholdDataCounts(0, 0, 0, 0, 0) + assert executor.statements == [ + "SELECT count(*) AS count_1 FROM households", + "SELECT count(*) AS count_1 FROM user_household_associations", + "SELECT count(*) AS count_1 FROM simulations " + "WHERE simulations.household_id IS NOT NULL", + "SELECT count(*) AS count_1 FROM reports " + "WHERE reports.household_id IS NOT NULL", + "SELECT count(*) AS count_1 FROM household_jobs", + ] + + +def test_counts_skip_tables_the_target_does_not_have() -> None: + executor = RecordingExecutor() + + counts = qualification.read_household_data_counts( + executor, existing_tables={"households", "user_household_associations"} + ) + + assert counts == qualification.HouseholdDataCounts(0, 0, 0, 0, 0) + assert executor.statements == [ + "SELECT count(*) AS count_1 FROM households", + "SELECT count(*) AS count_1 FROM user_household_associations", + ] + + @pytest.mark.parametrize( "counts", [ diff --git a/tests/unit/v2/test_model_persistence.py b/tests/unit/v2/test_model_persistence.py index be7879424..d14a774c7 100644 --- a/tests/unit/v2/test_model_persistence.py +++ b/tests/unit/v2/test_model_persistence.py @@ -1,5 +1,6 @@ """Canonical SQLModel persistence and bounded SQLAlchemy escape-hatch tests.""" +import re from datetime import UTC, datetime, timedelta from pathlib import Path @@ -21,6 +22,8 @@ TaxBenefitModelVersion, User, UserPolicy, + UserReportAssociation, + UserSimulationAssociation, ) from policyengine_api.data.v2.models.base import DIRECT_SQLALCHEMY_EXCEPTIONS @@ -38,6 +41,38 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: return engine +# A later revision drops these tables and columns while this release serves +# traffic, so no statement this release runs may name them. +RETIRED_SIMULATION_TABLE_REFERENCE = re.compile( + r"\b(household_jobs|simulations|reports|report_runs|aggregates|" + r"change_aggregates|budget_summary|decile_impacts|intra_decile_impacts|" + r"inequality|poverty|program_statistics|congressional_district_impacts|" + r"constituency_impacts|local_authority_impacts)\b" +) +RETIRED_COLUMN_TABLE_REFERENCE = re.compile( + r"\buser_(simulation|report)_associations\b" +) +RETIRED_COLUMN_REFERENCE = re.compile(r"\b(country|label)\b") + + +def _record_statements(engine) -> list[str]: + statements: list[str] = [] + + @sa.event.listens_for(engine, "before_cursor_execute") + def record(_connection, _cursor, statement, *_arguments) -> None: + statements.append(statement) + + return statements + + +def _assert_no_retired_schema_references(statements: list[str]) -> None: + assert statements + for statement in statements: + assert not RETIRED_SIMULATION_TABLE_REFERENCE.search(statement), statement + if RETIRED_COLUMN_TABLE_REFERENCE.search(statement): + assert not RETIRED_COLUMN_REFERENCE.search(statement), statement + + def _policy_graph(*, content_hash: str = "a" * 64): model = TaxBenefitModel(name=f"policy-model-{content_hash[:8]}") version = TaxBenefitModelVersion( @@ -94,7 +129,7 @@ def test_direct_sqlalchemy_categories_are_complete_and_documented() -> None: "typed_json_and_text", "named_constraints_and_indexes", "ambiguous_foreign_key_relationships", - "transaction_conflict_recovery", + "retired_table_declarations", } assert all(DIRECT_SQLALCHEMY_EXCEPTIONS.values()) @@ -445,3 +480,57 @@ def test_v2_models_do_not_create_a_parallel_sqlalchemy_orm_layer() -> None: assert "mapped_column" not in model_sources assert "sqlalchemy.orm.Session" not in model_sources assert "class_=sqlalchemy" not in model_sources + + +def test_orm_user_deletion_reads_no_retired_table_or_column() -> None: + engine = _relational_sqlite_engine() + with Session(engine) as session: + user = User(primary_country="us") + session.add(user) + session.add( + UserReportAssociation(user=user, country_id="us", legacy_report_id=615) + ) + session.flush() + statements = _record_statements(engine) + + session.delete(user) + session.flush() + + _assert_no_retired_schema_references(statements) + engine.dispose() + + +@pytest.mark.parametrize( + ("association_model", "legacy_reference"), + [ + (UserSimulationAssociation, "legacy_simulation_id"), + (UserReportAssociation, "legacy_report_id"), + ], +) +def test_simulation_and_report_associations_omit_retired_columns( + association_model: type[UserSimulationAssociation | UserReportAssociation], + legacy_reference: str, +) -> None: + engine = _relational_sqlite_engine() + with Session(engine) as session: + user = User(primary_country="us") + session.add(user) + session.flush() + statements = _record_statements(engine) + + association = association_model( + user_id=user.id, country_id="us", **{legacy_reference: 615} + ) + session.add(association) + session.flush() + session.expire_all() + stored = session.exec( + select(association_model).where(association_model.id == association.id) + ).one() + stored.name = "Renamed" + session.flush() + session.delete(stored) + session.flush() + + _assert_no_retired_schema_references(statements) + engine.dispose() diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 20947ffcb..ddde6b525 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -6,19 +6,18 @@ from sqlalchemy.schema import CreateTable from policyengine_api.data.v1_models import V1Base +from policyengine_api.data.v2 import models as v2_models from policyengine_api.data.v2.models import ( V2_METADATA, DatasetVersion, Dynamic, Household, - HouseholdJob, LegacyHouseholdMapping, LegacyPolicyMapping, LegacyUserMapping, LegacyUserPolicyMapping, ParameterValue, Policy, - Simulation, TaxBenefitModelVersion, User, UserHouseholdAssociation, @@ -45,6 +44,42 @@ } ) +# The simulation service owns these tables in its own schema. The v2 metadata +# still declares them, so it matches the schema until a later revision drops +# them, but no model maps them: nothing this release runs reads or writes them. +RETIRED_SIMULATION_TABLES = frozenset( + {"household_jobs", "simulations", "reports", "report_runs", *RUN_OUTPUT_TABLES} +) +# Retired at Stage 12A; dropped by the same later revision. +RETIRED_ASSOCIATION_COLUMNS = frozenset({"country", "label"}) +REMOVED_SIMULATION_MODEL_NAMES = frozenset( + { + "AggregateOutput", + "AggregateType", + "BudgetSummary", + "ChangeAggregate", + "CongressionalDistrictImpact", + "ConstituencyImpact", + "DecileImpact", + "DecileType", + "HouseholdJob", + "HouseholdJobStatus", + "Inequality", + "IntraDecileImpact", + "LocalAuthorityImpact", + "OutputStatus", + "Poverty", + "ProgramStatistics", + "Report", + "ReportRun", + "ReportRunStatus", + "ReportRunTrigger", + "Simulation", + "SimulationStatus", + "SimulationType", + } +) + def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: expected_modules = { @@ -53,9 +88,7 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: Policy: "policies", Dynamic: "policies", Household: "households", - HouseholdJob: "households", LegacyHouseholdMapping: "household_mappings", - Simulation: "simulations", UserHouseholdAssociation: "associations", UserPolicy: "associations", LegacyPolicyMapping: "policy_mappings", @@ -80,13 +113,47 @@ def _v2_mappers(): ) -def test_every_metadata_table_has_one_mapped_model() -> None: +def test_every_kept_metadata_table_has_one_mapped_model() -> None: metadata_table_names = set(V2_METADATA.tables) model_table_names = {mapper.local_table.name for mapper in _v2_mappers()} assert V2_METADATA is not V1Base.metadata - assert model_table_names == metadata_table_names - assert len(_v2_mappers()) == len(metadata_table_names) + assert metadata_table_names >= RETIRED_SIMULATION_TABLES + assert model_table_names == metadata_table_names - RETIRED_SIMULATION_TABLES + assert len(_v2_mappers()) == len(model_table_names) + + +def test_models_package_exports_no_simulation_domain_model() -> None: + exported_names = set(v2_models.__all__) | set(vars(v2_models)) + + assert REMOVED_SIMULATION_MODEL_NAMES.isdisjoint(exported_names) + + +def test_kept_models_have_no_relationship_to_a_retired_table() -> None: + related_tables = { + f"{mapper.class_.__name__}.{relationship.key}": ( + relationship.mapper.local_table.name + ) + for mapper in _v2_mappers() + for relationship in mapper.relationships + } + + assert not { + relationship: table_name + for relationship, table_name in related_tables.items() + if table_name in RETIRED_SIMULATION_TABLES + } + + +def test_association_models_do_not_map_the_retired_columns() -> None: + for model in (UserSimulationAssociation, UserReportAssociation): + table = V2_METADATA.tables[model.__tablename__] + # Declared on the table, so the metadata matches the schema... + assert set(table.c.keys()) >= RETIRED_ASSOCIATION_COLUMNS + assert all(table.c[name].nullable for name in RETIRED_ASSOCIATION_COLUMNS) + # ...but neither mapped nor a model field. + assert RETIRED_ASSOCIATION_COLUMNS.isdisjoint(sa.inspect(model).columns.keys()) + assert RETIRED_ASSOCIATION_COLUMNS.isdisjoint(model.model_fields) def test_every_table_has_named_primary_foreign_and_relational_constraints() -> None: @@ -298,48 +365,6 @@ def test_all_datetime_columns_are_timezone_aware() -> None: assert all(column.type.timezone for column in datetime_columns) -def test_report_definition_and_run_columns_are_separated() -> None: - report = V2_METADATA.tables["reports"] - report_run = V2_METADATA.tables["report_runs"] - - assert report.c.type.nullable - assert { - "country", - "type", - "tax_benefit_model_id", - "policy_id", - "baseline_simulation_id", - "reform_simulation_id", - "household_id", - "dataset_id", - "region_id", - "year", - "inputs", - }.issubset(report.c.keys()) - assert { - "country_package_version", - "policyengine_version", - "status", - "trigger", - "idempotency_key", - "started_at", - "completed_at", - "error_message", - "markdown", - }.issubset(report_run.c.keys()) - assert { - "country_package_version", - "policyengine_version", - "status", - "error_message", - "markdown", - }.isdisjoint(report.c.keys()) - assert not report_run.c.country_package_version.nullable - assert not report_run.c.policyengine_version.nullable - assert isinstance(report_run.c.idempotency_key.type, sa.Uuid) - assert report_run.c.idempotency_key.type.as_uuid - - def test_standard_users_require_a_supported_primary_country() -> None: users = V2_METADATA.tables["users"] @@ -414,53 +439,6 @@ def test_regions_have_one_same_model_version_default_logical_dataset() -> None: } -def test_reports_and_simulations_snapshot_selected_datasets() -> None: - reports = V2_METADATA.tables["reports"] - simulations = V2_METADATA.tables["simulations"] - - for table in (reports, simulations): - dataset_foreign_key = next(iter(table.c.dataset_id.foreign_keys)) - assert dataset_foreign_key.target_fullname == "datasets.id" - assert dataset_foreign_key.ondelete in {"RESTRICT", "SET NULL"} - - -def test_stage9_adds_no_dataset_version_relationship_to_run_tables() -> None: - for table_name in ("simulations", "reports", "report_runs"): - assert "dataset_version_id" not in V2_METADATA.tables[table_name].c - - -def test_run_outputs_reference_report_runs_not_base_reports() -> None: - for table_name in RUN_OUTPUT_TABLES: - table = V2_METADATA.tables[table_name] - assert "report_run_id" in table.c - assert not table.c.report_run_id.nullable - assert "report_id" not in table.c - foreign_key = next(iter(table.c.report_run_id.foreign_keys)) - assert foreign_key.target_fullname == "report_runs.id" - assert foreign_key.ondelete == "CASCADE" - - -def test_report_rerun_constraints_allow_same_versions_but_deduplicate_requests() -> ( - None -): - report_runs = V2_METADATA.tables["report_runs"] - unique_column_sets = { - tuple(column.name for column in constraint.columns) - for constraint in report_runs.constraints - if isinstance(constraint, sa.UniqueConstraint) - } - - assert ("report_id", "idempotency_key") in unique_column_sets - assert ( - "report_id", - "country_package_version", - "policyengine_version", - ) not in unique_column_sets - assert "ix_report_runs_current_output" in { - index.name for index in report_runs.indexes - } - - def test_named_checks_and_required_indexes_cover_core_invariants() -> None: constraint_names = { constraint.name @@ -473,15 +451,11 @@ def test_named_checks_and_required_indexes_cover_core_invariants() -> None: assert { "ck_regions_required_filter_values", - "ck_simulations_type_input", "ck_users_primary_country", - "ck_report_runs_terminal_completion", "ck_parameter_values_single_owner", }.issubset(constraint_names) assert { "ix_users_email", - "ix_simulations_status_created_at", - "ix_report_runs_current_output", "uq_parameter_values_canonical_parameter_start_date", }.issubset(index_names) @@ -601,10 +575,6 @@ def test_simulation_and_report_associations_have_no_cross_service_keys() -> None simulations = V2_METADATA.tables["user_simulation_associations"] reports = V2_METADATA.tables["user_report_associations"] - # The retired columns stay, nullable, for the previous Public API revision. - for table in (simulations, reports): - assert table.c.country.nullable - assert table.c.label.nullable assert { foreign_key.target_fullname for foreign_key in simulations.foreign_keys } == {"users.id"} @@ -619,7 +589,6 @@ def test_simulation_and_report_associations_have_no_cross_service_keys() -> None assert isinstance(postgres_type, postgresql.JSONB) assert set(sa.inspect(UserSimulationAssociation).relationships.keys()) == {"user"} assert set(sa.inspect(UserReportAssociation).relationships.keys()) == {"user"} - assert "user_associations" not in sa.inspect(Simulation).relationships def test_user_policy_records_v3_reform_fields() -> None: diff --git a/tests/unit/v2/test_report_run_creation.py b/tests/unit/v2/test_report_run_creation.py deleted file mode 100644 index 3a1b1b82c..000000000 --- a/tests/unit/v2/test_report_run_creation.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Pin report-run creation's validation, key-less insert and conflict paths.""" - -from __future__ import annotations - -from contextlib import contextmanager -from uuid import uuid4 - -import pytest -from sqlalchemy.exc import IntegrityError -from sqlalchemy.pool import StaticPool -from sqlmodel import Session, create_engine, select - -from policyengine_api.data.v2.models import ( - V2_METADATA, - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, - TaxBenefitModel, -) -from policyengine_api.data.v2.report_runs import ( - ReportNotFoundError, - ReportRunPackageVersions, - create_report_run, -) - - -@pytest.fixture -def engine(): - test_engine = create_engine( - "sqlite://", - connect_args={"check_same_thread": False}, - poolclass=StaticPool, - ) - V2_METADATA.create_all(test_engine) - yield test_engine - test_engine.dispose() - - -def _create_report(session: Session) -> Report: - report = Report( - label="Distributional report", - country="us", - tax_benefit_model=TaxBenefitModel(name=f"model-{uuid4()}"), - inputs={}, - ) - session.add(report) - session.flush() - return report - - -def test_unknown_report_is_rejected_before_any_validation(engine) -> None: - with ( - Session(engine) as session, - pytest.raises(ReportNotFoundError, match="does not exist"), - ): - create_report_run( - session, - report_id=uuid4(), - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", - ), - trigger=ReportRunTrigger.MANUAL, - ) - - -def test_manual_run_without_a_key_is_rejected_before_version_checks(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - - with pytest.raises( - ValueError, match=r"^manual report reruns require an idempotency key$" - ): - create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", - ), - trigger=ReportRunTrigger.MANUAL, - ) - - -@pytest.mark.parametrize( - ("country_package_version", "policyengine_version"), - [("", "4.5.6"), ("1.2.3", "")], -) -def test_empty_package_versions_are_rejected( - engine, - country_package_version: str, - policyengine_version: str, -) -> None: - with Session(engine) as session: - report = _create_report(session) - - with pytest.raises( - ValueError, match=r"^report run package versions must be non-empty$" - ): - create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_package_version, - policyengine_version=policyengine_version, - ), - trigger=ReportRunTrigger.SYSTEM, - ) - - -def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - - run = create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", - ), - trigger=ReportRunTrigger.INITIAL, - ) - - stored = session.exec(select(ReportRun)).one() - assert stored is run - assert run.report_id == report.id - assert run.country_package_version == "1.2.3" - assert run.policyengine_version == "4.5.6" - assert run.trigger is ReportRunTrigger.INITIAL - assert run.idempotency_key is None - assert run.status is ReportRunStatus.PENDING - - -def test_conflict_without_a_matching_keyed_run_reraises( - engine, - monkeypatch, -) -> None: - conflict = IntegrityError("insert", {}, Exception("unrelated constraint")) - - @contextmanager - def failing_savepoint(): - raise conflict - yield - - with Session(engine) as session: - report = _create_report(session) - monkeypatch.setattr(session, "begin_nested", failing_savepoint) - - with pytest.raises(IntegrityError) as raised: - create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", - ), - trigger=ReportRunTrigger.MANUAL, - idempotency_key=uuid4(), - ) - - assert raised.value is conflict diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py deleted file mode 100644 index 4383a8110..000000000 --- a/tests/unit/v2/test_report_runs.py +++ /dev/null @@ -1,421 +0,0 @@ -"""Report definition, rerun, idempotency, worker, and selector tests.""" - -import threading -from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime, timedelta -from pathlib import Path -from uuid import UUID, uuid4 - -import pytest -import sqlalchemy as sa -from sqlalchemy.pool import StaticPool -from sqlmodel import Session, create_engine, select - -from policyengine_api.data.v2.models import ( - V2_METADATA, - AggregateOutput, - AggregateType, - Dataset, - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, - Simulation, - SimulationType, - TaxBenefitModel, - TaxBenefitModelVersion, -) -from policyengine_api.data.v2.report_runs import ( - ReportRunPackageVersions, - ReportRunStateError, - ReportTypeImmutableError, - begin_report_run, - complete_report_run, - create_report_run, - fail_report_run, - select_current_report_run, - set_report_type, -) - -DEPLOYED_COUNTRY_VERSIONS = {"us": "1.2.3"} -DEPLOYED_POLICYENGINE_VERSION = "4.5.6" -NOW = datetime(2026, 8, 14, 12, 0, tzinfo=UTC) - - -@pytest.fixture -def engine(): - test_engine = create_engine( - "sqlite://", - connect_args={"check_same_thread": False}, - poolclass=StaticPool, - ) - V2_METADATA.create_all(test_engine) - yield test_engine - test_engine.dispose() - - -def _create_report(session: Session, *, report_type: str | None = None) -> Report: - model = TaxBenefitModel(name=f"model-{uuid4()}") - report = Report( - label="Distributional report", - country="us", - type=report_type, - tax_benefit_model=model, - inputs={"reform": {"gov.example": 1}}, - ) - session.add(report) - session.flush() - return report - - -def _create_run( - session: Session, - report: Report, - *, - key: UUID | None = None, - country_version: str = "1.2.3", - policyengine_version: str = DEPLOYED_POLICYENGINE_VERSION, -) -> ReportRun: - return create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_version, - policyengine_version=policyengine_version, - ), - trigger=ReportRunTrigger.MANUAL, - idempotency_key=key if key is not None else uuid4(), - ) - - -def test_untyped_and_typed_report_definitions_store_no_execution_versions( - engine, -) -> None: - with Session(engine) as session: - untyped = _create_report(session) - typed = _create_report(session, report_type="marginal_tax_rate") - - assert untyped.type is None - assert typed.type == "marginal_tax_rate" - assert "country_package_version" not in Report.__table__.c - assert "policyengine_version" not in Report.__table__.c - - -def test_report_type_can_change_before_first_run_but_not_after(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - set_report_type( - session, - report_id=report.id, - report_type="marginal_tax_rate", - ) - _create_run(session, report) - - unchanged = set_report_type( - session, - report_id=report.id, - report_type="marginal_tax_rate", - ) - assert unchanged.id == report.id - with pytest.raises(ReportTypeImmutableError): - set_report_type( - session, - report_id=report.id, - report_type="economy_comparison", - ) - - -def test_new_manual_keys_create_distinct_same_version_runs(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - first = _create_run(session, report) - second = _create_run(session, report) - - assert first.id != second.id - assert first.country_package_version == second.country_package_version - assert first.policyengine_version == second.policyengine_version - assert len(session.exec(select(ReportRun)).all()) == 2 - - -def test_manual_rerun_requires_an_idempotency_key(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - - with pytest.raises(ValueError, match="require an idempotency key"): - create_report_run( - session, - report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ), - trigger=ReportRunTrigger.MANUAL, - ) - - -def test_idempotency_key_round_trips_as_a_python_uuid(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - request_key = uuid4() - run = _create_run(session, report, key=request_key) - session.flush() - session.expire(run) - - assert run.idempotency_key == request_key - assert isinstance(run.idempotency_key, UUID) - - -def test_transport_retry_returns_the_existing_report_scoped_run(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - request_key = uuid4() - first = _create_run(session, report, key=request_key) - retried = _create_run(session, report, key=request_key) - - assert retried.id == first.id - assert len(session.exec(select(ReportRun)).all()) == 1 - - -def test_concurrent_idempotent_requests_resolve_to_one_run(tmp_path: Path) -> None: - sqlite_path = tmp_path / "report-rerun-concurrency.db" - test_engine = create_engine( - f"sqlite:///{sqlite_path}", - connect_args={"check_same_thread": False, "timeout": 10}, - ) - - # BEGIN IMMEDIATE gives this explicit SQLite-only test fixture the same - # serialization point that SELECT FOR UPDATE provides in Postgres. - @sa.event.listens_for(test_engine, "connect") - def _set_sqlite_transaction_mode(dbapi_connection, _record) -> None: - dbapi_connection.isolation_level = None - - @sa.event.listens_for(test_engine, "begin") - def _begin_immediate(connection) -> None: - connection.exec_driver_sql("BEGIN IMMEDIATE") - - V2_METADATA.create_all(test_engine) - with Session(test_engine) as session: - report_id = _create_report(session).id - session.commit() - - barrier = threading.Barrier(2) - request_key = uuid4() - - def request_rerun() -> UUID: - with Session(test_engine) as session: - barrier.wait() - run = create_report_run( - session, - report_id=report_id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ), - trigger=ReportRunTrigger.MANUAL, - idempotency_key=request_key, - ) - run_id = run.id - session.commit() - return run_id - - with ThreadPoolExecutor(max_workers=2) as executor: - run_ids = set(executor.map(lambda _: request_rerun(), range(2))) - - with Session(test_engine) as session: - assert len(run_ids) == 1 - assert len(session.exec(select(ReportRun)).all()) == 1 - test_engine.dispose() - - -def test_worker_retry_resumes_the_same_run_and_terminal_runs_stay_terminal( - engine, -) -> None: - with Session(engine) as session: - report = _create_report(session) - run = _create_run(session, report) - started = begin_report_run(session, report_run_id=run.id, started_at=NOW) - resumed = begin_report_run(session, report_run_id=run.id) - - assert resumed.id == run.id - assert resumed.status is ReportRunStatus.RUNNING - assert resumed.started_at == started.started_at - assert len(session.exec(select(ReportRun)).all()) == 1 - - complete_report_run(session, report_run_id=run.id, completed_at=NOW) - with pytest.raises(ReportRunStateError): - begin_report_run(session, report_run_id=run.id) - - -def test_outputs_from_repeated_runs_are_preserved(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - model = report.tax_benefit_model - version = TaxBenefitModelVersion( - model=model, - version="1.2.3", - current_law_id=1, - metadata_time_periods=[2026], - ) - dataset = Dataset( - name="dataset", - storage_path="datasets/test.h5", - year=2026, - tax_benefit_model_version=version, - ) - simulation = Simulation( - simulation_type=SimulationType.ECONOMY, - dataset=dataset, - tax_benefit_model_version=version, - ) - session.add(simulation) - session.flush() - first = _create_run(session, report) - second = _create_run(session, report) - session.add_all( - [ - AggregateOutput( - report_run=first, - simulation_id=simulation.id, - variable="household_net_income", - aggregate_type=AggregateType.SUM, - result=1.0, - ), - AggregateOutput( - report_run=second, - simulation_id=simulation.id, - variable="household_net_income", - aggregate_type=AggregateType.SUM, - result=2.0, - ), - ] - ) - session.flush() - - outputs = session.exec( - select(AggregateOutput).order_by(AggregateOutput.result) - ).all() - assert [output.result for output in outputs] == [1.0, 2.0] - assert outputs[0].report_run_id != outputs[1].report_run_id - - -def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - candidates = [ - ReportRun( - id=UUID(int=1), - report=report, - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - status=ReportRunStatus.SUCCEEDED, - completed_at=NOW, - ), - ReportRun( - id=UUID(int=2), - report=report, - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - status=ReportRunStatus.SUCCEEDED, - completed_at=NOW, - ), - ReportRun( - id=UUID(int=3), - report=report, - country_package_version="9.9.9", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - status=ReportRunStatus.SUCCEEDED, - completed_at=NOW + timedelta(days=1), - ), - ReportRun( - id=UUID(int=4), - report=report, - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - status=ReportRunStatus.RUNNING, - ), - ] - session.add_all(candidates) - session.flush() - - current = select_current_report_run( - session, - report_id=report.id, - country_package_versions=DEPLOYED_COUNTRY_VERSIONS, - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ) - - assert current is not None - assert current.id == UUID(int=2) - - -def test_pending_and_failed_reruns_do_not_displace_success(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - successful = _create_run(session, report) - complete_report_run(session, report_run_id=successful.id, completed_at=NOW) - pending = _create_run(session, report) - failed = _create_run(session, report) - fail_report_run( - session, - report_run_id=failed.id, - error_message="transient failure", - completed_at=NOW + timedelta(hours=1), - ) - - current = select_current_report_run( - session, - report_id=report.id, - country_package_versions=DEPLOYED_COUNTRY_VERSIONS, - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ) - - assert pending.status is ReportRunStatus.PENDING - assert current is not None - assert current.id == successful.id - - -def test_new_successful_rerun_becomes_current_without_cache_state(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - first = _create_run(session, report) - complete_report_run(session, report_run_id=first.id, completed_at=NOW) - rerun = _create_run(session, report) - complete_report_run( - session, - report_run_id=rerun.id, - completed_at=NOW + timedelta(minutes=1), - ) - - # No Redis/cache collaborator participates in durable selection; a - # flush or expiry therefore has no state to invalidate here. - current = select_current_report_run( - session, - report_id=report.id, - country_package_versions=DEPLOYED_COUNTRY_VERSIONS, - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ) - - assert current is not None - assert current.id == rerun.id - assert len(session.exec(select(ReportRun)).all()) == 2 - - -def test_selector_returns_none_without_a_matching_success(engine) -> None: - with Session(engine) as session: - report = _create_report(session) - mismatched = _create_run( - session, - report, - country_version="0.0.1", - ) - complete_report_run(session, report_run_id=mismatched.id, completed_at=NOW) - - assert ( - select_current_report_run( - session, - report_id=report.id, - country_package_versions=DEPLOYED_COUNTRY_VERSIONS, - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, - ) - is None - )