From 52b01a1ff3f6ca65ea4dc8ab3967f2bc23bb675d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:16:01 +0400 Subject: [PATCH 01/22] Record each simulation run and key v2 results by run Add simulation_runs, one row per run of a simulation with the catalog, country package and PolicyEngine.py versions it used, its status, Modal call and output file; and report_run_simulation_runs, the simulation runs each report run combines. Poverty and inequality belong to a simulation run, the other result tables name the two simulation runs they compare, report runs record their Modal call and the two results without a table, and household jobs calculate one household for one year under one policy. Revision 2ffd50f47a2b was autogenerated, with two documented enum-type corrections. Co-Authored-By: Claude Opus 5.5 --- changelog.d/v2-simulation-runs.changed.md | 1 + docs/engineering/skills/alembic-migrations.md | 48 + ...add_simulation_runs_and_key_results_by_.py | 1019 +++++++++++++++++ .../data/v2/metadata_validation.py | 6 +- policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/__init__.py | 6 + policyengine_api/data/v2/models/households.py | 34 +- policyengine_api/data/v2/models/metadata.py | 4 - policyengine_api/data/v2/models/reports.py | 157 ++- .../data/v2/models/simulations.py | 123 +- policyengine_api/data/v2/privileges.py | 2 + .../integration/test_alembic_v2_lifecycle.py | 20 + .../test_alembic_v2_simulation_runs.py | 374 ++++++ .../test_v2_catalog_publication.py | 7 +- tests/unit/v2/test_alembic_v2.py | 56 +- tests/unit/v2/test_metadata_validation.py | 8 + tests/unit/v2/test_models.py | 11 +- tests/unit/v2/test_simulation_run_models.py | 286 +++++ 18 files changed, 2056 insertions(+), 108 deletions(-) create mode 100644 changelog.d/v2-simulation-runs.changed.md create mode 100644 migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py create mode 100644 tests/integration/test_alembic_v2_simulation_runs.py create mode 100644 tests/unit/v2/test_simulation_run_models.py diff --git a/changelog.d/v2-simulation-runs.changed.md b/changelog.d/v2-simulation-runs.changed.md new file mode 100644 index 000000000..863a41eb3 --- /dev/null +++ b/changelog.d/v2-simulation-runs.changed.md @@ -0,0 +1 @@ +API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index bd2719994..c3a529ff4 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -453,3 +453,51 @@ schema. The dropped values are not recoverable. Adding a required column fails on a table that has rows, so the downgrade works only while both tables are empty. + +### API v2 simulation runs + +Revision `2ffd50f47a2b` (parent `c614b6bcccfc`) was autogenerated after these +model changes: + +- `simulations` keeps only what to calculate. Its execution columns + (`status`, `error_message`, `started_at`, `completed_at`, + `household_result`, `output_dataset_id`) move to the new `simulation_runs`, + one row per run with the catalog version, country package version and + PolicyEngine.py version it used, its Modal call ID and its output file. +- The new `report_run_simulation_runs` names the simulation runs each report + run combines, as its baseline or its reform. +- `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); + the other result tables belong to a report run and name the two simulation + runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`). + `aggregates` and `change_aggregates` are unchanged. +- `report_runs` gains its Modal call ID and two JSON results that have no + table (`labor_supply_result`, `cliff_impact_result`). +- `household_jobs` calculates one household (`household_id`) for one `year` + under one policy, and records its catalog and package versions and its + Modal call ID. + +It has two post-generation corrections, each commented in the revision: + +1. **Dialect:** `simulation_runs.status` reuses the `v2_simulation_status` + type the baseline created for `simulations.status`. The generated + `create_table` would create the type again and fail, so the column uses + `postgresql.ENUM(..., create_type=False)`. +2. **Reversibility:** the downgrade drops `v2_simulation_role` after dropping + `report_run_simulation_runs`, because Alembic does not drop a table's enum + type with it. + +The upgrade adds required columns without server defaults to `household_jobs` +and the result tables, so it needs them empty, as they are while no managed +role holds a write privilege on them; on a table with rows it fails, and the +deploy stops before anything changes. The downgrade restores required columns +the same way, so it needs `simulations`, `household_jobs` and the result +tables empty. The previous Public API revision keeps working at this revision: +none of its roles can read or write a table the revision changes, except the +catalog writer, which only counts the rows of `simulations` and `report_runs`. +`tests/integration/test_alembic_v2_simulation_runs.py` checks that the +revision changes no other table, and the v2 lifecycle test covers the +downgrade and re-upgrade. + +The name `simulation_runs` is also a v1 Cloud SQL table. The v2 metadata +validation (`policyengine_api/data/v2/metadata_validation.py`) now treats it as +a name the two databases share, as it does `simulations`. diff --git a/migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py b/migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py new file mode 100644 index 000000000..3c2c663f3 --- /dev/null +++ b/migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py @@ -0,0 +1,1019 @@ +"""add simulation runs and key results by run + +Revision ID: 2ffd50f47a2b +Revises: c614b6bcccfc +Create Date: 2026-11-04 00:09:14.105644 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +Two post-generation corrections, each commented where it is made: +simulation_runs.status reuses the existing v2_simulation_status type, and the +downgrade drops the v2_simulation_role type it no longer uses. + +The upgrade adds required columns without server defaults to household_jobs +and to the result tables, so it needs those tables empty, as they are while no +managed role can write them. The downgrade restores required columns the same +way, so it needs simulations, household_jobs and the result tables empty. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel +from sqlalchemy.dialects import postgresql + +revision: str = "2ffd50f47a2b" +down_revision: Union[str, None] = "c614b6bcccfc" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "simulation_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("simulation_id", sa.Uuid(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column( + "country_package_version", + sqlmodel.sql.sqltypes.AutoString(length=128), + nullable=False, + ), + sa.Column( + "policyengine_version", + sqlmodel.sql.sqltypes.AutoString(length=128), + nullable=False, + ), + sa.Column("dataset_id", sa.Uuid(), nullable=True), + sa.Column("region_id", sa.Uuid(), nullable=True), + sa.Column( + "status", + # Dialect correction: the baseline created v2_simulation_status for + # simulations.status, which this revision moves here, so the column + # reuses the existing type instead of creating it again. + postgresql.ENUM( + "pending", + "running", + "succeeded", + "failed", + name="v2_simulation_status", + create_type=False, + ), + nullable=False, + ), + sa.Column( + "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True + ), + sa.Column("error_message", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("started_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("household_result", sa.JSON(), nullable=True), + sa.Column( + "output_uri", sqlmodel.sql.sqltypes.AutoString(length=2048), nullable=True + ), + sa.Column( + "output_sha256", sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True + ), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name=op.f("ck_simulation_runs_terminal_completion"), + ), + sa.CheckConstraint( + "(output_uri IS NULL) = (output_sha256 IS NULL)", + name=op.f("ck_simulation_runs_output_pair"), + ), + sa.CheckConstraint( + "output_sha256 IS NULL OR length(output_sha256) = 64", + name=op.f("ck_simulation_runs_output_sha256_length"), + ), + sa.ForeignKeyConstraint( + ["simulation_id"], + ["simulations.id"], + name=op.f("fk_simulation_runs_simulation_id_simulations"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_simulation_runs")), + ) + op.create_index( + "ix_simulation_runs_current_output", + "simulation_runs", + [ + "simulation_id", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ], + unique=False, + ) + op.create_index( + "ix_simulation_runs_status_started_at", + "simulation_runs", + ["status", "started_at"], + unique=False, + ) + op.create_table( + "report_run_simulation_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_run_id", sa.Uuid(), nullable=False), + sa.Column("simulation_run_id", sa.Uuid(), nullable=False), + sa.Column( + "role", + sa.Enum("baseline", "reform", name="v2_simulation_role"), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_report_run_simulation_runs_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["simulation_run_id"], + ["simulation_runs.id"], + name=op.f( + "fk_report_run_simulation_runs_simulation_run_id_simulation_runs" + ), + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_simulation_runs")), + sa.UniqueConstraint( + "report_run_id", + "role", + "simulation_run_id", + name="uq_report_run_simulation_runs_member", + ), + ) + op.create_index( + op.f("ix_report_run_simulation_runs_simulation_run_id"), + "report_run_simulation_runs", + ["simulation_run_id"], + unique=False, + ) + op.add_column( + "budget_summary", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "budget_summary", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_budget_summary_baseline_simulation_id_simulations"), + "budget_summary", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_budget_summary_reform_simulation_id_simulations"), + "budget_summary", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + "budget_summary", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + "budget_summary", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("budget_summary", "reform_simulation_id") + op.drop_column("budget_summary", "baseline_simulation_id") + op.add_column( + "congressional_district_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + "congressional_district_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + "congressional_district_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f( + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + ), + "congressional_district_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f( + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + ), + "congressional_district_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("congressional_district_impacts", "reform_simulation_id") + op.drop_column("congressional_district_impacts", "baseline_simulation_id") + op.add_column( + "constituency_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "constituency_impacts", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + "constituency_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + "constituency_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + "constituency_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + "constituency_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("constituency_impacts", "reform_simulation_id") + op.drop_column("constituency_impacts", "baseline_simulation_id") + op.add_column( + "decile_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "decile_impacts", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + "decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_decile_impacts_reform_simulation_id_simulations"), + "decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + "decile_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + "decile_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("decile_impacts", "reform_simulation_id") + op.drop_column("decile_impacts", "baseline_simulation_id") + op.add_column( + "household_jobs", + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + ) + op.add_column( + "household_jobs", sa.Column("household_id", sa.Uuid(), nullable=False) + ) + op.add_column("household_jobs", sa.Column("year", sa.Integer(), nullable=False)) + op.add_column( + "household_jobs", + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "household_jobs", + sa.Column( + "country_package_version", + sqlmodel.sql.sqltypes.AutoString(length=128), + nullable=False, + ), + ) + op.add_column( + "household_jobs", + sa.Column( + "policyengine_version", + sqlmodel.sql.sqltypes.AutoString(length=128), + nullable=False, + ), + ) + op.add_column( + "household_jobs", + sa.Column( + "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True + ), + ) + op.create_index( + op.f("ix_household_jobs_household_id"), + "household_jobs", + ["household_id"], + unique=False, + ) + op.create_check_constraint( + op.f("ck_household_jobs_terminal_completion"), + "household_jobs", + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + ) + op.create_check_constraint( + op.f("ck_household_jobs_year"), "household_jobs", "year BETWEEN 1900 AND 2200" + ) + op.add_column( + "inequality", sa.Column("simulation_run_id", sa.Uuid(), nullable=False) + ) + op.drop_constraint( + op.f("uq_inequality_run_simulation_measure"), "inequality", type_="unique" + ) + op.create_unique_constraint( + "uq_inequality_simulation_run_measure", + "inequality", + ["simulation_run_id", "income_variable", "entity"], + ) + op.drop_constraint( + op.f("fk_inequality_simulation_id_simulations"), + "inequality", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_inequality_simulation_run_id_simulation_runs"), + "inequality", + "simulation_runs", + ["simulation_run_id"], + ["id"], + ondelete="CASCADE", + ) + op.drop_column("inequality", "simulation_id") + op.drop_column("inequality", "report_run_id") + op.add_column( + "intra_decile_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "intra_decile_impacts", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("intra_decile_impacts", "reform_simulation_id") + op.drop_column("intra_decile_impacts", "baseline_simulation_id") + op.add_column( + "local_authority_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "local_authority_impacts", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + "local_authority_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + "local_authority_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + "local_authority_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + "local_authority_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("local_authority_impacts", "reform_simulation_id") + op.drop_column("local_authority_impacts", "baseline_simulation_id") + op.add_column("poverty", sa.Column("simulation_run_id", sa.Uuid(), nullable=False)) + op.drop_index(op.f("ix_poverty_run_simulation_type"), table_name="poverty") + op.create_unique_constraint( + "uq_poverty_simulation_run_measure", + "poverty", + ["simulation_run_id", "poverty_type", "entity", "filter_variable"], + postgresql_nulls_not_distinct=True, + ) + op.drop_constraint( + op.f("fk_poverty_report_run_id_report_runs"), "poverty", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_poverty_simulation_id_simulations"), "poverty", type_="foreignkey" + ) + op.create_foreign_key( + op.f("fk_poverty_simulation_run_id_simulation_runs"), + "poverty", + "simulation_runs", + ["simulation_run_id"], + ["id"], + ondelete="CASCADE", + ) + op.drop_column("poverty", "simulation_id") + op.drop_column("poverty", "report_run_id") + op.add_column( + "program_statistics", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "program_statistics", + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.drop_constraint( + op.f("fk_program_statistics_baseline_simulation_id_simulations"), + "program_statistics", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_program_statistics_reform_simulation_id_simulations"), + "program_statistics", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + "program_statistics", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + "program_statistics", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("program_statistics", "reform_simulation_id") + op.drop_column("program_statistics", "baseline_simulation_id") + op.add_column( + "report_runs", + sa.Column( + "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True + ), + ) + op.add_column( + "report_runs", sa.Column("labor_supply_result", sa.JSON(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("cliff_impact_result", sa.JSON(), nullable=True) + ) + op.create_index( + "ix_report_runs_status_started_at", + "report_runs", + ["status", "started_at"], + unique=False, + ) + op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") + op.drop_constraint( + op.f("fk_simulations_output_dataset_id_datasets"), + "simulations", + type_="foreignkey", + ) + op.drop_column("simulations", "household_result") + op.drop_column("simulations", "completed_at") + op.drop_column("simulations", "output_dataset_id") + op.drop_column("simulations", "error_message") + op.drop_column("simulations", "started_at") + op.drop_column("simulations", "status") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "simulations", + sa.Column( + "status", + postgresql.ENUM( + "pending", "running", "succeeded", "failed", name="v2_simulation_status" + ), + autoincrement=False, + nullable=False, + ), + ) + op.add_column( + "simulations", + sa.Column( + "started_at", + postgresql.TIMESTAMP(timezone=True), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + ) + op.add_column( + "simulations", + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + ) + op.add_column( + "simulations", + sa.Column( + "completed_at", + postgresql.TIMESTAMP(timezone=True), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "simulations", + sa.Column( + "household_result", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), + ) + op.create_foreign_key( + op.f("fk_simulations_output_dataset_id_datasets"), + "simulations", + "datasets", + ["output_dataset_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_index( + op.f("ix_simulations_status_created_at"), + "simulations", + ["status", "created_at"], + unique=False, + ) + op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") + op.drop_column("report_runs", "cliff_impact_result") + op.drop_column("report_runs", "labor_supply_result") + op.drop_column("report_runs", "modal_call_id") + op.add_column( + "program_statistics", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "program_statistics", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + "program_statistics", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + "program_statistics", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_program_statistics_reform_simulation_id_simulations"), + "program_statistics", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_program_statistics_baseline_simulation_id_simulations"), + "program_statistics", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("program_statistics", "reform_simulation_run_id") + op.drop_column("program_statistics", "baseline_simulation_run_id") + op.add_column( + "poverty", + sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.add_column( + "poverty", + sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.drop_constraint( + op.f("fk_poverty_simulation_run_id_simulation_runs"), + "poverty", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_poverty_simulation_id_simulations"), + "poverty", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_poverty_report_run_id_report_runs"), + "poverty", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) + op.drop_constraint("uq_poverty_simulation_run_measure", "poverty", type_="unique") + op.create_index( + op.f("ix_poverty_run_simulation_type"), + "poverty", + ["report_run_id", "simulation_id", "poverty_type"], + unique=False, + ) + op.drop_column("poverty", "simulation_run_id") + op.add_column( + "local_authority_impacts", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "local_authority_impacts", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + "local_authority_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + "local_authority_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + "local_authority_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + "local_authority_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("local_authority_impacts", "reform_simulation_run_id") + op.drop_column("local_authority_impacts", "baseline_simulation_run_id") + op.add_column( + "intra_decile_impacts", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "intra_decile_impacts", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + "intra_decile_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + "intra_decile_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("intra_decile_impacts", "reform_simulation_run_id") + op.drop_column("intra_decile_impacts", "baseline_simulation_run_id") + op.add_column( + "inequality", + sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.add_column( + "inequality", + sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.drop_constraint( + op.f("fk_inequality_simulation_run_id_simulation_runs"), + "inequality", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) + op.create_foreign_key( + op.f("fk_inequality_simulation_id_simulations"), + "inequality", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + "uq_inequality_simulation_run_measure", "inequality", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_inequality_run_simulation_measure"), + "inequality", + ["report_run_id", "simulation_id", "income_variable", "entity"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("inequality", "simulation_run_id") + op.drop_constraint(op.f("ck_household_jobs_year"), "household_jobs", type_="check") + op.drop_constraint( + op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" + ) + op.drop_index(op.f("ix_household_jobs_household_id"), table_name="household_jobs") + op.drop_column("household_jobs", "modal_call_id") + op.drop_column("household_jobs", "policyengine_version") + op.drop_column("household_jobs", "country_package_version") + op.drop_column("household_jobs", "tax_benefit_model_version_id") + op.drop_column("household_jobs", "year") + op.drop_column("household_jobs", "household_id") + op.drop_column("household_jobs", "updated_at") + op.add_column( + "decile_impacts", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "decile_impacts", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + "decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + "decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_decile_impacts_reform_simulation_id_simulations"), + "decile_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + "decile_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("decile_impacts", "reform_simulation_run_id") + op.drop_column("decile_impacts", "baseline_simulation_run_id") + op.add_column( + "constituency_impacts", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "constituency_impacts", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + "constituency_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + "constituency_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + "constituency_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + "constituency_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("constituency_impacts", "reform_simulation_run_id") + op.drop_column("constituency_impacts", "baseline_simulation_run_id") + op.add_column( + "congressional_district_impacts", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "congressional_district_impacts", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f( + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + ), + "congressional_district_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f( + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + ), + "congressional_district_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + "congressional_district_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + "congressional_district_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("congressional_district_impacts", "reform_simulation_run_id") + op.drop_column("congressional_district_impacts", "baseline_simulation_run_id") + op.add_column( + "budget_summary", + sa.Column( + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "budget_summary", + sa.Column( + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + ), + ) + op.drop_constraint( + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + "budget_summary", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + "budget_summary", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_budget_summary_reform_simulation_id_simulations"), + "budget_summary", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_budget_summary_baseline_simulation_id_simulations"), + "budget_summary", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("budget_summary", "reform_simulation_run_id") + op.drop_column("budget_summary", "baseline_simulation_run_id") + op.drop_index( + op.f("ix_report_run_simulation_runs_simulation_run_id"), + table_name="report_run_simulation_runs", + ) + op.drop_table("report_run_simulation_runs") + op.drop_index("ix_simulation_runs_status_started_at", table_name="simulation_runs") + op.drop_index("ix_simulation_runs_current_output", table_name="simulation_runs") + op.drop_table("simulation_runs") + # Post-generation reversibility correction: Alembic creates the + # v2_simulation_role type with its table but does not drop it with it. + sa.Enum(name="v2_simulation_role").drop(op.get_bind()) + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/metadata_validation.py b/policyengine_api/data/v2/metadata_validation.py index 7445269f1..40a1ef286 100644 --- a/policyengine_api/data/v2/metadata_validation.py +++ b/policyengine_api/data/v2/metadata_validation.py @@ -3,8 +3,9 @@ from collections.abc import Iterable # These names belong only to the existing v1 database domain. Names that are -# intentionally shared between v1 and v2, such as `simulations` and -# `user_policies`, are excluded from this targeted rejection set. +# intentionally shared between v1 and v2, such as `simulations`, +# `simulation_runs` and `user_policies`, are excluded from this targeted +# rejection set. V1_ONLY_TABLES = frozenset( { "analysis", @@ -16,7 +17,6 @@ "reform_impact", "report_output_runs", "report_outputs", - "simulation_runs", "tracers", "user_profiles", } diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 420601582..5ca67210e 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "c614b6bcccfc" +V2_ALEMBIC_HEAD_REVISION = "2ffd50f47a2b" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index 8498cfbb6..d9b4a3c04 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -82,11 +82,14 @@ ProgramStatistics, Report, ReportRun, + ReportRunSimulationRun, ReportRunStatus, ReportRunTrigger, + SimulationRole, ) from policyengine_api.data.v2.models.simulations import ( # noqa: E402 Simulation, + SimulationRun, SimulationStatus, SimulationType, ) @@ -138,9 +141,12 @@ "RegionType", "Report", "ReportRun", + "ReportRunSimulationRun", "ReportRunStatus", "ReportRunTrigger", "Simulation", + "SimulationRole", + "SimulationRun", "SimulationStatus", "SimulationType", "Stage12AggregationStatus", diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 181266d9d..8ab02d5f0 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -8,11 +8,7 @@ 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.base import TimestampedModel, enum_type from policyengine_api.data.v2.models.policies import Dynamic, Policy if TYPE_CHECKING: @@ -25,7 +21,7 @@ class HouseholdJobStatus(StrEnum): - """Lifecycle state of a queued household calculation job.""" + """Lifecycle state of a household calculation job.""" PENDING = "pending" RUNNING = "running" @@ -96,18 +92,34 @@ class Household(TimestampedModel, table=True): ) -class HouseholdJob(IdentifiedModel, table=True): - """A queued household calculation request, with its status and result or error. +class HouseholdJob(TimestampedModel, table=True): + """One calculation of one household for one year under one policy. - Deleting the job's policy or dynamic keeps the job but clears the reference. + ``policy_id`` is null for current law. ``request_data`` holds the + calculation's options, such as axes. The job records the catalog version + and .py version it ran on, its Modal call, and its status and result or + error. ``household_id`` and ``tax_benefit_model_version_id`` are logical + references to the Public API's rows, which the simulation service checks + before it writes the job. Deleting the job's policy or dynamic keeps the + job but clears the reference. """ __tablename__ = "household_jobs" __table_args__ = ( + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", + name="ck_household_jobs_year", + ), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name="ck_household_jobs_terminal_completion", + ), sa.Index("ix_household_jobs_status_created_at", "status", "created_at"), ) country: str = Field(max_length=16) + household_id: UUID = Field(index=True) + year: int request_data: dict[str, Any] = Field(sa_type=sa.JSON) policy_id: UUID | None = Field( default=None, @@ -119,10 +131,14 @@ class HouseholdJob(IdentifiedModel, table=True): foreign_key="dynamics.id", ondelete="SET NULL", ) + tax_benefit_model_version_id: UUID + country_package_version: str = Field(max_length=128) + policyengine_version: str = Field(max_length=128) status: HouseholdJobStatus = Field( default=HouseholdJobStatus.PENDING, sa_type=enum_type(HouseholdJobStatus, "v2_household_job_status"), ) + modal_call_id: str | None = Field(default=None, max_length=255) error_message: str | None = None result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) started_at: datetime | None = Field( diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index b53e42a7b..e736010cc 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -248,10 +248,6 @@ class Dataset(TimestampedModel, table=True): 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") diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 5c73947c1..c379a56ed 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,4 +1,9 @@ -"""Canonical SQLModel tables for stable reports, runs, and run outputs.""" +"""Canonical SQLModel tables for stable reports, runs, and run outputs. + +Results that compare a baseline with a reform belong to a report run and name +the two simulation runs they compare; poverty and inequality, which describe +one simulation, belong to a simulation run. +""" from datetime import datetime from enum import StrEnum @@ -20,7 +25,7 @@ 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.simulations import Simulation, SimulationRun from policyengine_api.data.v2.models.users import User @@ -65,6 +70,13 @@ class DecileType(StrEnum): WEALTH = "wealth" +class SimulationRole(StrEnum): + """The part a simulation run plays in the report run that combines it.""" + + BASELINE = "baseline" + REFORM = "reform" + + class Report(TimestampedModel, table=True): """Stable report definition; execution state belongs to ReportRun.""" @@ -145,7 +157,12 @@ class Report(TimestampedModel, table=True): class ReportRun(TimestampedModel, table=True): - """One immutable-version execution attempt for a stable report.""" + """One immutable-version execution attempt for a stable report. + + Its ``simulation_runs`` rows name the simulation runs it combines. + ``labor_supply_result`` and ``cliff_impact_result`` hold the two parts of an + economy comparison that have no result table. + """ __tablename__ = "report_runs" __table_args__ = ( @@ -167,6 +184,7 @@ class ReportRun(TimestampedModel, table=True): "completed_at", "id", ), + sa.Index("ix_report_runs_status_started_at", "status", "started_at"), ) report_id: UUID = Field( @@ -184,6 +202,7 @@ class ReportRun(TimestampedModel, table=True): sa_type=enum_type(ReportRunTrigger, "v2_report_run_trigger"), ) idempotency_key: UUID | None = Field(default=None) + modal_call_id: str | None = Field(default=None, max_length=255) started_at: datetime | None = Field( default=None, sa_type=sa.DateTime(timezone=True), @@ -194,8 +213,14 @@ class ReportRun(TimestampedModel, table=True): ) error_message: str | None = None markdown: str | None = Field(default=None, sa_type=sa.Text) + labor_supply_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + cliff_impact_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) report: Report = Relationship(back_populates="runs") + simulation_runs: list["ReportRunSimulationRun"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) aggregates: list["AggregateOutput"] = Relationship( back_populates="report_run", cascade_delete=True, @@ -219,10 +244,6 @@ class ReportRun(TimestampedModel, table=True): 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, @@ -231,16 +252,46 @@ class ReportRun(TimestampedModel, table=True): 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 ReportRunSimulationRun(IdentifiedModel, table=True): + """One simulation run a report run combines, as its baseline or its reform. + + A single-year report run names one of each (a household report may name + only a baseline); a budget-window run names one of each per year. + """ + + __tablename__ = "report_run_simulation_runs" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "role", + "simulation_run_id", + name="uq_report_run_simulation_runs_member", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + index=True, + ) + role: SimulationRole = Field( + sa_type=enum_type(SimulationRole, "v2_simulation_role"), + ) + + report_run: ReportRun = Relationship(back_populates="simulation_runs") + simulation_run: SimulationRun = Relationship(back_populates="report_runs") + + class AggregateOutput(IdentifiedModel, table=True): """A sum, mean or count of one variable in one simulation of a report run. @@ -342,12 +393,12 @@ class BudgetSummary(IdentifiedModel, table=True): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) variable_name: str = Field(max_length=512) @@ -387,12 +438,12 @@ class DecileImpact(IdentifiedModel, table=True): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) income_variable: str = Field(max_length=512) @@ -434,12 +485,12 @@ class IntraDecileImpact(IdentifiedModel, table=True): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) decile_type: DecileType = Field( @@ -457,27 +508,22 @@ class IntraDecileImpact(IdentifiedModel, table=True): class Inequality(IdentifiedModel, table=True): - """Gini and top/bottom income shares for one simulation of a report run.""" + """Gini and top/bottom income shares for one simulation run.""" __tablename__ = "inequality" __table_args__ = ( sa.UniqueConstraint( - "report_run_id", - "simulation_id", + "simulation_run_id", "income_variable", "entity", - name="uq_inequality_run_simulation_measure", + name="uq_inequality_simulation_run_measure", ), ) - report_run_id: UUID = Field( - foreign_key="report_runs.id", + simulation_run_id: UUID = Field( + foreign_key="simulation_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 @@ -485,30 +531,33 @@ class Inequality(IdentifiedModel, table=True): top_1_share: float | None = None bottom_50_share: float | None = None - report_run: ReportRun = Relationship(back_populates="inequalities") + simulation_run: SimulationRun = Relationship(back_populates="inequalities") class Poverty(IdentifiedModel, table=True): - """Poverty headcount and rate of one type for one simulation of a report run.""" + """Poverty headcount and rate of one type for one simulation run. + + One row per run, poverty type, entity and filter variable; rows without a + filter variable count as equal, so each run has one unfiltered row per type + and entity. + """ __tablename__ = "poverty" __table_args__ = ( - sa.Index( - "ix_poverty_run_simulation_type", - "report_run_id", - "simulation_id", + sa.UniqueConstraint( + "simulation_run_id", "poverty_type", + "entity", + "filter_variable", + name="uq_poverty_simulation_run_measure", + postgresql_nulls_not_distinct=True, ), ) - report_run_id: UUID = Field( - foreign_key="report_runs.id", + simulation_run_id: UUID = Field( + foreign_key="simulation_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) @@ -516,7 +565,7 @@ class Poverty(IdentifiedModel, table=True): total_population: float | None = None rate: float | None = None - report_run: ReportRun = Relationship(back_populates="poverty_results") + simulation_run: SimulationRun = Relationship(back_populates="poverty_results") class ProgramStatistics(IdentifiedModel, table=True): @@ -539,12 +588,12 @@ class ProgramStatistics(IdentifiedModel, table=True): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) program_name: str = Field(max_length=512) @@ -568,12 +617,12 @@ class GeographicImpactBase(IdentifiedModel): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", ondelete="RESTRICT", ) average_household_income_change: float diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index a6821eadc..f59038b00 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -1,4 +1,4 @@ -"""Canonical SQLModel table for v2 simulations.""" +"""Canonical SQLModel tables for v2 simulations and their runs.""" from datetime import datetime from enum import StrEnum @@ -18,11 +18,16 @@ from policyengine_api.data.v2.models.policies import Dynamic, Policy if TYPE_CHECKING: - from policyengine_api.data.v2.models.reports import Report + from policyengine_api.data.v2.models.reports import ( + Inequality, + Poverty, + Report, + ReportRunSimulationRun, + ) class SimulationStatus(StrEnum): - """Lifecycle state of a v2 simulation.""" + """Lifecycle state of one run of a v2 simulation.""" PENDING = "pending" RUNNING = "running" @@ -40,8 +45,9 @@ class SimulationType(StrEnum): 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. + It defines what to calculate; each calculation of it is a ``SimulationRun``. + Deleting its policy, dynamic or region clears that reference; its household + and input dataset cannot be deleted while it exists. """ __tablename__ = "simulations" @@ -61,7 +67,6 @@ class Simulation(TimestampedModel, table=True): "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( @@ -93,43 +98,20 @@ class Simulation(TimestampedModel, table=True): 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") @@ -145,3 +127,86 @@ class Simulation(TimestampedModel, table=True): back_populates="reform_simulation", sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, ) + runs: list["SimulationRun"] = Relationship( + back_populates="simulation", + cascade_delete=True, + ) + + +class SimulationRun(TimestampedModel, table=True): + """One run of a simulation on one catalog version and .py version. + + ``tax_benefit_model_version_id``, ``dataset_id`` and ``region_id`` name the + catalog version, dataset and region the run used, which a run on a newer + version resolves again. They are logical references: the catalog belongs to + the Public API, and the simulation service checks them before it writes the + run. A household simulation's run stores its household in + ``household_result``; an economy run's output file is ``output_uri``, with + its SHA-256. Each poverty and inequality result belongs to one run. + """ + + __tablename__ = "simulation_runs" + __table_args__ = ( + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name="ck_simulation_runs_terminal_completion", + ), + sa.CheckConstraint( + "(output_uri IS NULL) = (output_sha256 IS NULL)", + name="ck_simulation_runs_output_pair", + ), + sa.CheckConstraint( + "output_sha256 IS NULL OR length(output_sha256) = 64", + name="ck_simulation_runs_output_sha256_length", + ), + sa.Index( + "ix_simulation_runs_current_output", + "simulation_id", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ), + sa.Index("ix_simulation_runs_status_started_at", "status", "started_at"), + ) + + simulation_id: UUID = Field( + foreign_key="simulations.id", + ondelete="CASCADE", + ) + tax_benefit_model_version_id: UUID + country_package_version: str = Field(max_length=128) + policyengine_version: str = Field(max_length=128) + dataset_id: UUID | None = None + region_id: UUID | None = None + status: SimulationStatus = Field( + default=SimulationStatus.PENDING, + sa_type=enum_type(SimulationStatus, "v2_simulation_status"), + ) + modal_call_id: str | None = Field(default=None, max_length=255) + error_message: str | 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) + output_uri: str | None = Field(default=None, max_length=2048) + output_sha256: str | None = Field(default=None, max_length=64) + + simulation: Simulation = Relationship(back_populates="runs") + report_runs: list["ReportRunSimulationRun"] = Relationship( + back_populates="simulation_run" + ) + poverty_results: list["Poverty"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, + ) + inequalities: list["Inequality"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, + ) diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 7bf2897c8..814aa07fa 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -107,8 +107,10 @@ def _access( "local_authority_impacts": _access(), "poverty": _access(), "program_statistics": _access(), + "report_run_simulation_runs": _access(), "report_runs": _access(catalog_writer=READ), "reports": _access(catalog_writer=READ), + "simulation_runs": _access(), "simulations": _access(catalog_writer=READ), # Metadata catalog: the catalog writer publishes it, the Public API reads it. "dataset_versions": _access(catalog_writer=READ), diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 3c2cc8161..3c01b0a84 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -37,6 +37,8 @@ # The catalog format 2 revision, the parent of the area impact coordinate removal. CATALOG_FORMAT_2_REVISION = "f5c1d9b4685b" AREA_IMPACT_COORDINATE_REVISION = "e86f19a6c4da" +# The parameter search file revision, the parent of the simulation runs revision. +SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) AREA_IMPACT_TABLES = ("constituency_impacts", "local_authority_impacts") MAP_COORDINATE_COLUMNS = ("x", "y") @@ -409,6 +411,20 @@ def _assert_countries_shape(engine) -> None: assert check_name not in _check_names(engine, table_name) +def _assert_pre_simulation_runs_shape(engine) -> None: + _assert_revision(engine, SIMULATION_RUNS_PREVIOUS_REVISION) + table_names = set(inspect(engine).get_table_names(schema="public")) + assert {"simulation_runs", "report_run_simulation_runs"}.isdisjoint(table_names) + simulation_columns = { + column["name"] for column in inspect(engine).get_columns("simulations") + } + poverty_columns = { + column["name"] for column in inspect(engine).get_columns("poverty") + } + assert {"status", "household_result", "output_dataset_id"} <= simulation_columns + assert {"report_run_id", "simulation_id"} <= poverty_columns + + def _assert_pre_countries_shape(engine) -> None: assert "countries" not in inspect(engine).get_table_names(schema="public") for (table_name, _column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): @@ -430,6 +446,10 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: _upgrade_and_assert_head(config, engine) _assert_countries_shape(engine) + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + _assert_pre_simulation_runs_shape(engine) + _upgrade_and_assert_head(config, engine) + command.downgrade(config, CATALOG_FORMAT_2_REVISION) _assert_area_impact_coordinates_restored(engine) _upgrade_and_assert_head(config, engine) diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py new file mode 100644 index 000000000..83ef41ade --- /dev/null +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -0,0 +1,374 @@ +"""Check the simulation runs revision against disposable Postgres. + +The revision adds simulation runs and the table naming the simulation runs a +report run combines, keys every result table by run, and gives household jobs +one household, one year and one policy. It changes no other table, so the +previous Public API revision, whose roles can't read or write any table it +changes, keeps working while it is deployed. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, inspect, text +from sqlalchemy.exc import IntegrityError + +from policyengine_api.constants import REPO +from policyengine_api.data.v2.migration_target import ( + V2_ALEMBIC_DISPOSABLE_TEST, + load_v2_alembic_settings, +) +from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection, Engine + +SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" +NEW_TABLES = frozenset({"simulation_runs", "report_run_simulation_runs"}) +CHANGED_TABLES = frozenset( + { + "budget_summary", + "congressional_district_impacts", + "constituency_impacts", + "decile_impacts", + "household_jobs", + "inequality", + "intra_decile_impacts", + "local_authority_impacts", + "poverty", + "program_statistics", + "report_runs", + "simulations", + } +) +SHA256_HEX = "a" * 64 + + +def _disposable_url() -> str: + database_url = os.environ.get(V2_MIGRATION_DATABASE_URL, "") + if not database_url: + pytest.skip(f"{V2_MIGRATION_DATABASE_URL} is not set") + settings = load_v2_alembic_settings( + { + V2_MIGRATION_DATABASE_URL: database_url, + V2_ALEMBIC_DISPOSABLE_TEST: os.environ.get(V2_ALEMBIC_DISPOSABLE_TEST, ""), + } + ) + if not settings.disposable_test: + pytest.fail("v2 lifecycle tests require disposable-test mode") + return settings.url.render_as_string(hide_password=False) + + +@pytest.fixture +def config() -> Config: + return Config(str(REPO / "alembic-v2.ini")) + + +@pytest.fixture +def engine(config: Config) -> Iterator[Engine]: + database_engine = create_engine(_disposable_url()) + try: + yield database_engine + finally: + command.upgrade(config, "head") + database_engine.dispose() + + +def _table_shape(engine: Engine, table_name: str) -> dict[str, object]: + """Describe one table's columns, keys, indexes and checks for comparison.""" + inspector = inspect(engine) + return { + "columns": [ + (column["name"], str(column["type"]), column["nullable"], column["default"]) + for column in inspector.get_columns(table_name) + ], + "primary_key": inspector.get_pk_constraint(table_name), + "foreign_keys": sorted( + inspector.get_foreign_keys(table_name), key=lambda key: key["name"] + ), + "indexes": sorted( + inspector.get_indexes(table_name), key=lambda index: index["name"] + ), + "unique_constraints": sorted( + inspector.get_unique_constraints(table_name), + key=lambda constraint: constraint["name"], + ), + "check_constraints": sorted( + inspector.get_check_constraints(table_name), + key=lambda constraint: constraint["name"], + ), + } + + +def _shapes(engine: Engine) -> dict[str, dict[str, object]]: + return { + table_name: _table_shape(engine, table_name) + for table_name in inspect(engine).get_table_names(schema="public") + } + + +def test_the_revision_adds_two_tables_and_changes_only_simulation_side_tables( + config: Config, engine: Engine +) -> None: + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + before = _shapes(engine) + command.upgrade(config, "head") + after = _shapes(engine) + + assert set(after) - set(before) == NEW_TABLES + assert set(before) <= set(after) + changed = { + table_name for table_name in before if before[table_name] != after[table_name] + } + assert changed == CHANGED_TABLES + + +@contextmanager +def _refused(connection: Connection) -> Iterator[None]: + """Expect the statement inside to violate a constraint, then roll it back.""" + savepoint = connection.begin_nested() + with pytest.raises(IntegrityError): + yield + savepoint.rollback() + + +def _insert(connection: Connection, table: str, **values: object) -> None: + columns = ", ".join(values) + placeholders = ", ".join(f":{column}" for column in values) + connection.execute( + text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"), values + ) + + +def _insert_simulation_run( + connection: Connection, + simulation_id: object, + version_id: object, + *, + status: str = "pending", + **values: object, +) -> object: + run_id = uuid4() + _insert( + connection, + "simulation_runs", + id=run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=version_id, + country_package_version="1.0.0", + policyengine_version="4.0.0", + status=status, + **values, + ) + return run_id + + +def test_run_tables_enforce_completion_output_and_uniqueness( + config: Config, engine: Engine +) -> None: + command.upgrade(config, "head") + model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() + simulation_id, report_id, report_run_id = uuid4(), uuid4(), uuid4() + + with engine.connect() as connection: + transaction = connection.begin() + try: + _insert( + connection, + "tax_benefit_models", + id=model_id, + name=f"runs-{model_id.hex[:8]}", + ) + connection.execute( + text( + "INSERT INTO tax_benefit_model_versions " + "(id, model_id, version, current_law_id, metadata_time_periods) " + "VALUES (:id, :model_id, '1.0.0', 1, '[2026]')" + ), + {"id": version_id, "model_id": model_id}, + ) + _insert( + connection, + "datasets", + id=dataset_id, + name="populace", + year=2026, + is_output_dataset=False, + tax_benefit_model_version_id=version_id, + ) + _insert( + connection, + "simulations", + id=simulation_id, + simulation_type="economy", + dataset_id=dataset_id, + tax_benefit_model_version_id=version_id, + year=2026, + ) + baseline_run_id = _insert_simulation_run( + connection, simulation_id, version_id + ) + reform_run_id = _insert_simulation_run( + connection, + simulation_id, + version_id, + status="succeeded", + completed_at="2026-11-04T00:00:00Z", + output_uri="gs://bucket/run.parquet", + output_sha256=SHA256_HEX, + ) + + with _refused(connection): + connection.execute( + text( + "UPDATE simulation_runs SET status = 'succeeded' WHERE id = :id" + ), + {"id": baseline_run_id}, + ) + with _refused(connection): + _insert_simulation_run( + connection, + simulation_id, + version_id, + output_uri="gs://bucket/run.parquet", + ) + with _refused(connection): + _insert_simulation_run( + connection, + simulation_id, + version_id, + output_uri="gs://bucket/run.parquet", + output_sha256="short", + ) + + _insert( + connection, + "reports", + id=report_id, + label="Report", + country="us", + tax_benefit_model_id=model_id, + inputs="{}", + ) + _insert( + connection, + "report_runs", + id=report_run_id, + report_id=report_id, + country_package_version="1.0.0", + policyengine_version="4.0.0", + status="running", + trigger="initial", + ) + _insert( + connection, + "report_run_simulation_runs", + id=uuid4(), + report_run_id=report_run_id, + simulation_run_id=baseline_run_id, + role="baseline", + ) + with _refused(connection): + _insert( + connection, + "report_run_simulation_runs", + id=uuid4(), + report_run_id=report_run_id, + simulation_run_id=baseline_run_id, + role="baseline", + ) + with _refused(connection): + _insert( + connection, + "budget_summary", + id=uuid4(), + report_run_id=report_run_id, + baseline_simulation_run_id=baseline_run_id, + reform_simulation_run_id=uuid4(), + variable_name="household_net_income", + entity="household", + ) + + _insert( + connection, + "poverty", + id=uuid4(), + simulation_run_id=reform_run_id, + poverty_type="poverty", + entity="person", + ) + with _refused(connection): + _insert( + connection, + "poverty", + id=uuid4(), + simulation_run_id=reform_run_id, + poverty_type="poverty", + entity="person", + ) + + household_job = { + "country": "us", + "household_id": uuid4(), + "request_data": "{}", + "tax_benefit_model_version_id": version_id, + "country_package_version": "1.0.0", + "policyengine_version": "4.0.0", + } + _insert( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="pending", + **household_job, + ) + with _refused(connection): + _insert( + connection, + "household_jobs", + id=uuid4(), + year=1800, + status="pending", + **household_job, + ) + with _refused(connection): + _insert( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="failed", + **household_job, + ) + + connection.execute( + text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} + ) + assert ( + connection.execute( + text("SELECT count(*) FROM report_run_simulation_runs") + ).scalar_one() + == 0 + ) + connection.execute( + text("DELETE FROM simulations WHERE id = :id"), {"id": simulation_id} + ) + remaining_runs = connection.execute( + text("SELECT count(*) FROM simulation_runs") + ).scalar_one() + remaining_poverty = connection.execute( + text("SELECT count(*) FROM poverty") + ).scalar_one() + assert (remaining_runs, remaining_poverty) == (0, 0) + finally: + transaction.rollback() diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index 0875d2864..e0bbbfc67 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -290,7 +290,6 @@ def test_republish_preserves_existing_run_and_dataset_version_records( ) 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, ) @@ -316,8 +315,8 @@ def test_republish_preserves_existing_run_and_dataset_version_records( report.id, report_run.id, simulation.dataset_id, - simulation.output_dataset_id, report.dataset_id, + dataset_version.dataset_id, ) publish_catalog(publication_engine, catalog) @@ -332,8 +331,8 @@ def test_republish_preserves_existing_run_and_dataset_version_records( 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_report.dataset_id == protected_ids[5] + assert persisted_dataset_version.dataset_id == protected_ids[6] def test_new_version_is_additive_and_concurrent_retries_serialize( diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 2099bb908..48147fdf5 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -47,6 +47,8 @@ "legacy_user_policy_mappings", "parameter_search_files", "parameter_tree_children", + "report_run_simulation_runs", + "simulation_runs", "stage12_evaluation_reports", "stage12_evaluation_simulations", "user_identities", @@ -62,6 +64,12 @@ CATALOG_FORMAT_2_REVISION = "f5c1d9b4685b" AREA_IMPACT_COORDINATE_REVISION = "e86f19a6c4da" PARAMETER_TREE_REVISION = "3869eb814a90" +# The parameter search file revision, parent of the simulation runs revision. +SIMULATION_RUNS_PARENT_REVISION = "c614b6bcccfc" +SIMULATION_RUNS_REVISION_PATH = ( + REPO + / "migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py" +) DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -243,6 +251,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, AREA_IMPACT_COORDINATE_REVISION, CATALOG_FORMAT_2_REVISION, @@ -989,7 +998,7 @@ def test_parameter_search_file_revision_creates_only_the_files_table() -> None: if revision.down_revision == PARAMETER_TREE_REVISION ] assert len(children) == 1 - assert children[0].revision == V2_ALEMBIC_HEAD_REVISION + assert children[0].revision == SIMULATION_RUNS_PARENT_REVISION revision = Path(children[0].path).read_text(encoding="utf-8") upgrade = revision[ revision.index("def upgrade()") : revision.index("def downgrade()") @@ -1009,3 +1018,48 @@ def test_parameter_search_file_revision_creates_only_the_files_table() -> None: assert "op.execute(" not in body assert "op.bulk_insert(" not in body assert "op.add_column(" not in body + + +def test_simulation_runs_revision_has_only_its_two_documented_corrections() -> None: + """The revision is generated; its two corrections handle enum types. + + ``simulation_runs.status`` reuses the type the baseline created for + ``simulations.status``, and the downgrade drops the new role type. + """ + revision = SIMULATION_RUNS_REVISION_PATH.read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert ( + f'down_revision: Union[str, None] = "{SIMULATION_RUNS_PARENT_REVISION}"' + in revision + ) + assert revision.count("Dialect correction:") == 1 + assert revision.count("Post-generation reversibility correction:") == 1 + assert 'name="v2_simulation_status",\n create_type=False,' in ( + upgrade + ) + assert 'sa.Enum(name="v2_simulation_role").drop(op.get_bind())' in downgrade + assert 'op.create_table(\n "simulation_runs"' in upgrade + assert 'op.create_table(\n "report_run_simulation_runs"' in upgrade + assert 'op.drop_table("report_run_simulation_runs")' in downgrade + assert 'op.drop_table("simulation_runs")' in downgrade + assert upgrade.count("op.drop_table(") == 0 + for removed_column in ( + "status", + "error_message", + "started_at", + "completed_at", + "household_result", + "output_dataset_id", + ): + assert f'op.drop_column("simulations", "{removed_column}")' in upgrade + for body in (upgrade, downgrade): + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body diff --git a/tests/unit/v2/test_metadata_validation.py b/tests/unit/v2/test_metadata_validation.py index 935e4ed9b..ad9e0724d 100644 --- a/tests/unit/v2/test_metadata_validation.py +++ b/tests/unit/v2/test_metadata_validation.py @@ -33,3 +33,11 @@ def test_rejection_sets_do_not_block_current_metadata() -> None: def test_new_table_name_requires_no_allowlist_update() -> None: validate_v2_metadata_table_names({*V2_METADATA.tables, "future_domain_table"}) + + +@pytest.mark.parametrize("table_name", ["simulations", "simulation_runs"]) +def test_names_shared_with_v1_are_accepted(table_name: str) -> None: + """v2 reuses these v1 table names for its own tables, in its own database.""" + assert table_name in V2_METADATA.tables + assert table_name not in V1_ONLY_TABLES + validate_v2_metadata_table_names({*V2_METADATA.tables, table_name}) diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 10ff02660..db60cd8cf 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -23,7 +23,9 @@ ParameterTreeChild, ParameterValue, Policy, + ReportRunSimulationRun, Simulation, + SimulationRun, TaxBenefitModelVersion, User, UserHouseholdAssociation, @@ -46,6 +48,8 @@ ("user_report_associations", "country_id"), ) +# Outputs of a report run. Poverty and inequality describe one simulation, so +# they belong to a simulation run instead (test_simulation_run_models.py). RUN_OUTPUT_TABLES = frozenset( { "aggregates", @@ -54,10 +58,8 @@ "congressional_district_impacts", "constituency_impacts", "decile_impacts", - "inequality", "intra_decile_impacts", "local_authority_impacts", - "poverty", "program_statistics", } ) @@ -74,6 +76,8 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: HouseholdJob: "households", LegacyHouseholdMapping: "household_mappings", Simulation: "simulations", + SimulationRun: "simulations", + ReportRunSimulationRun: "reports", UserHouseholdAssociation: "associations", UserPolicy: "associations", LegacyPolicyMapping: "policy_mappings", @@ -500,11 +504,12 @@ def test_named_checks_and_required_indexes_cover_core_invariants() -> None: "ck_users_standard_primary_country", "ck_countries_society_wide_requires_household", "ck_report_runs_terminal_completion", + "ck_simulation_runs_terminal_completion", "ck_parameter_values_single_owner", }.issubset(constraint_names) assert { "ix_users_email", - "ix_simulations_status_created_at", + "ix_simulation_runs_current_output", "ix_report_runs_current_output", "uq_parameter_values_canonical_parameter_start_date", }.issubset(index_names) diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py new file mode 100644 index 000000000..64cea61f4 --- /dev/null +++ b/tests/unit/v2/test_simulation_run_models.py @@ -0,0 +1,286 @@ +"""Structural tests for simulation runs, report runs, household jobs and results. + +A simulation defines what to calculate and each calculation of it is a run that +records the catalog and .py versions it used. A report run names the simulation +runs it combines, results are keyed by run, and a household job calculates one +household for one year under one policy. +""" + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql +from sqlalchemy.schema import CreateTable + +from policyengine_api.data.v2.models import ( + V2_METADATA, + ReportRunSimulationRun, + SimulationRole, + SimulationRun, +) + +RUN_STATUSES = ["pending", "running", "succeeded", "failed"] +PER_SIMULATION_RESULT_TABLES = ("inequality", "poverty") +COMPARISON_RESULT_TABLES = ( + "budget_summary", + "congressional_district_impacts", + "constituency_impacts", + "decile_impacts", + "intra_decile_impacts", + "local_authority_impacts", + "program_statistics", +) + + +def _foreign_key(table: sa.Table, column_name: str) -> sa.ForeignKey: + (foreign_key,) = table.c[column_name].foreign_keys + return foreign_key + + +def _check_names(table: sa.Table) -> set[str]: + return { + constraint.name + for constraint in table.constraints + if isinstance(constraint, sa.CheckConstraint) + } + + +def _unique_column_sets(table: sa.Table) -> set[tuple[str, ...]]: + return { + tuple(column.name for column in constraint.columns) + for constraint in table.constraints + if isinstance(constraint, sa.UniqueConstraint) + } + + +def _indexes(table: sa.Table) -> dict[str, tuple[str, ...]]: + return { + index.name: tuple(column.name for column in index.columns) + for index in table.indexes + } + + +def test_simulations_keep_only_what_to_calculate() -> None: + simulations = V2_METADATA.tables["simulations"] + + assert { + "status", + "error_message", + "started_at", + "completed_at", + "household_result", + "output_dataset_id", + }.isdisjoint(simulations.c.keys()) + assert { + "simulation_type", + "dataset_id", + "household_id", + "policy_id", + "dynamic_id", + "tax_benefit_model_version_id", + "region_id", + "filter_field", + "filter_value", + "filter_strategy", + "year", + } <= set(simulations.c.keys()) + assert "ix_simulations_status_created_at" not in _indexes(simulations) + + +def test_each_simulation_run_records_its_versions_status_and_output() -> None: + simulation_runs = V2_METADATA.tables["simulation_runs"] + + required = { + "simulation_id", + "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", + "status", + "created_at", + "updated_at", + } + optional = { + "dataset_id", + "region_id", + "modal_call_id", + "error_message", + "started_at", + "completed_at", + "household_result", + "output_uri", + "output_sha256", + } + assert set(simulation_runs.c.keys()) == {"id"} | required | optional + assert all(not simulation_runs.c[name].nullable for name in required) + assert all(simulation_runs.c[name].nullable for name in optional) + assert simulation_runs.c.status.type.name == "v2_simulation_status" + assert list(simulation_runs.c.status.type.enums) == RUN_STATUSES + assert simulation_runs.c.modal_call_id.type.length == 255 + assert simulation_runs.c.output_uri.type.length == 2048 + assert simulation_runs.c.output_sha256.type.length == 64 + assert simulation_runs.c.country_package_version.type.length == 128 + assert simulation_runs.c.policyengine_version.type.length == 128 + + simulation_key = _foreign_key(simulation_runs, "simulation_id") + assert simulation_key.target_fullname == "simulations.id" + assert simulation_key.ondelete == "CASCADE" + # The catalog version, dataset and region are the Public API's rows, so a + # run refers to them without a foreign key. + assert {key.target_fullname for key in simulation_runs.foreign_keys} == { + "simulations.id" + } + assert SimulationRun.__module__.endswith(".simulations") + + +def test_run_tables_require_a_completion_time_and_index_running_runs() -> None: + simulation_runs = V2_METADATA.tables["simulation_runs"] + report_runs = V2_METADATA.tables["report_runs"] + household_jobs = V2_METADATA.tables["household_jobs"] + + assert { + "ck_simulation_runs_terminal_completion", + "ck_simulation_runs_output_pair", + "ck_simulation_runs_output_sha256_length", + } <= _check_names(simulation_runs) + assert "ck_household_jobs_terminal_completion" in _check_names(household_jobs) + assert _indexes(simulation_runs) == { + "ix_simulation_runs_current_output": ( + "simulation_id", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ), + "ix_simulation_runs_status_started_at": ("status", "started_at"), + } + assert _indexes(report_runs)["ix_report_runs_status_started_at"] == ( + "status", + "started_at", + ) + + +def test_report_runs_record_their_modal_call_and_untabled_results() -> None: + report_runs = V2_METADATA.tables["report_runs"] + + assert report_runs.c.modal_call_id.nullable + assert report_runs.c.modal_call_id.type.length == 255 + for column_name in ("labor_supply_result", "cliff_impact_result"): + assert report_runs.c[column_name].nullable + assert isinstance(report_runs.c[column_name].type, sa.JSON) + + +def test_report_runs_name_the_simulation_runs_they_combine() -> None: + members = V2_METADATA.tables["report_run_simulation_runs"] + + assert set(members.c.keys()) == { + "id", + "created_at", + "report_run_id", + "simulation_run_id", + "role", + } + assert all(not column.nullable for column in members.columns) + report_run_key = _foreign_key(members, "report_run_id") + simulation_run_key = _foreign_key(members, "simulation_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + assert simulation_run_key.target_fullname == "simulation_runs.id" + assert simulation_run_key.ondelete == "RESTRICT" + assert members.c.role.type.name == "v2_simulation_role" + assert list(members.c.role.type.enums) == ["baseline", "reform"] + assert [role.value for role in SimulationRole] == ["baseline", "reform"] + assert _unique_column_sets(members) == { + ("report_run_id", "role", "simulation_run_id") + } + assert "ix_report_run_simulation_runs_simulation_run_id" in _indexes(members) + assert ReportRunSimulationRun.__module__.endswith(".reports") + + +def test_household_jobs_calculate_one_household_year_and_policy() -> None: + household_jobs = V2_METADATA.tables["household_jobs"] + + for column_name in ( + "household_id", + "year", + "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", + "request_data", + "updated_at", + ): + assert not household_jobs.c[column_name].nullable, column_name + for column_name in ("policy_id", "dynamic_id", "modal_call_id"): + assert household_jobs.c[column_name].nullable, column_name + assert household_jobs.c.modal_call_id.type.length == 255 + assert "ck_household_jobs_year" in _check_names(household_jobs) + assert "ix_household_jobs_household_id" in _indexes(household_jobs) + assert "household_id" not in { + key.parent.name for key in household_jobs.foreign_keys + } + assert "tax_benefit_model_version_id" not in { + key.parent.name for key in household_jobs.foreign_keys + } + + +def test_per_simulation_results_belong_to_one_simulation_run() -> None: + for table_name in PER_SIMULATION_RESULT_TABLES: + table = V2_METADATA.tables[table_name] + assert {"report_run_id", "simulation_id"}.isdisjoint(table.c.keys()) + assert not table.c.simulation_run_id.nullable + simulation_run_key = _foreign_key(table, "simulation_run_id") + assert simulation_run_key.target_fullname == "simulation_runs.id" + assert simulation_run_key.ondelete == "CASCADE" + + assert _unique_column_sets(V2_METADATA.tables["inequality"]) == { + ("simulation_run_id", "income_variable", "entity") + } + assert _unique_column_sets(V2_METADATA.tables["poverty"]) == { + ("simulation_run_id", "poverty_type", "entity", "filter_variable") + } + + +def test_poverty_rows_without_a_filter_variable_are_unique_per_run() -> None: + statement = str( + CreateTable(V2_METADATA.tables["poverty"]).compile(dialect=postgresql.dialect()) + ) + + assert ( + "CONSTRAINT uq_poverty_simulation_run_measure UNIQUE NULLS NOT DISTINCT " + "(simulation_run_id, poverty_type, entity, filter_variable)" + ) in statement + + +def test_comparison_results_name_the_two_simulation_runs_they_compare() -> None: + for table_name in COMPARISON_RESULT_TABLES: + table = V2_METADATA.tables[table_name] + assert {"baseline_simulation_id", "reform_simulation_id"}.isdisjoint( + table.c.keys() + ), table_name + report_run_key = _foreign_key(table, "report_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + for column_name in ("baseline_simulation_run_id", "reform_simulation_run_id"): + assert not table.c[column_name].nullable + run_key = _foreign_key(table, column_name) + assert run_key.target_fullname == "simulation_runs.id", table_name + assert run_key.ondelete == "RESTRICT" + assert any( + columns[0] == "report_run_id" for columns in _unique_column_sets(table) + ), table_name + + +def test_custom_totals_keep_their_simulation_references() -> None: + """``aggregates`` and ``change_aggregates`` stay unchanged until they return.""" + aggregates = V2_METADATA.tables["aggregates"] + change_aggregates = V2_METADATA.tables["change_aggregates"] + + assert _foreign_key(aggregates, "simulation_id").target_fullname == ( + "simulations.id" + ) + for column_name in ("baseline_simulation_id", "reform_simulation_id"): + assert _foreign_key(change_aggregates, column_name).target_fullname == ( + "simulations.id" + ) + for table in (aggregates, change_aggregates): + assert _foreign_key(table, "report_run_id").target_fullname == ( + "report_runs.id" + ) From 019907f7e7267b24ccad93d119e23bb85afb90ff Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:21:05 +0400 Subject: [PATCH 02/22] Drop foreign keys between the two services' v2 tables The simulation service and its Modal workers write simulations, reports, household jobs, runs and results; the Public API writes users, saved items, policies, households and the catalog. Revision 7a83a87d029a drops the 14 foreign keys from the first group to the second, keeping each column as a plain UUID that the writing service checks, and the models drop the ORM relationships across the two groups. Keys within a group stay. The revision is autogenerated with no corrections. Co-Authored-By: Claude Opus 5.5 --- changelog.d/v2-simulation-runs.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 14 ++ ...029a_drop_foreign_keys_between_the_two_.py | 196 ++++++++++++++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/households.py | 35 +--- policyengine_api/data/v2/models/metadata.py | 13 -- policyengine_api/data/v2/models/policies.py | 8 - policyengine_api/data/v2/models/reports.py | 57 ++--- .../data/v2/models/simulations.py | 59 +----- policyengine_api/data/v2/models/users.py | 6 +- .../integration/test_alembic_v2_lifecycle.py | 15 ++ .../test_alembic_v2_simulation_runs.py | 38 +++- .../test_v2_catalog_publication.py | 9 +- tests/unit/v2/test_alembic_v2.py | 79 +++++++ tests/unit/v2/test_models.py | 9 +- tests/unit/v2/test_report_run_creation.py | 5 +- tests/unit/v2/test_report_runs.py | 17 +- tests/unit/v2/test_simulation_run_models.py | 82 ++++++++ 18 files changed, 479 insertions(+), 167 deletions(-) create mode 100644 migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py diff --git a/changelog.d/v2-simulation-runs.changed.md b/changelog.d/v2-simulation-runs.changed.md index 863a41eb3..eef5d02ad 100644 --- a/changelog.d/v2-simulation-runs.changed.md +++ b/changelog.d/v2-simulation-runs.changed.md @@ -1 +1 @@ -API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. +API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. Its simulation tables no longer have foreign keys to the tables the Public API writes. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index c3a529ff4..21a1e8a01 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -501,3 +501,17 @@ downgrade and re-upgrade. The name `simulation_runs` is also a v1 Cloud SQL table. The v2 metadata validation (`policyengine_api/data/v2/metadata_validation.py`) now treats it as a name the two databases share, as it does `simulations`. + +### API v2 foreign keys between the two services' tables + +Revision `7a83a87d029a` (parent `2ffd50f47a2b`) was autogenerated after the +models stopped declaring foreign keys from the tables the simulation service +and its Modal workers write (`simulations`, `reports`, `household_jobs`) to the +tables the Public API writes (`users`, `policies`, `dynamics`, `households`, +and the catalog's `tax_benefit_models`, `tax_benefit_model_versions`, +`datasets` and `regions`). The 14 columns stay as plain UUIDs: the service that +writes a row checks the rows it names. Foreign keys between tables of the same +service stay. It has no post-generation corrections. Its downgrade restores the +keys, so it fails while a column names a row the referenced table lacks. +`tests/unit/v2/test_simulation_run_models.py` checks that no foreign key or ORM +relationship joins the two services' tables. diff --git a/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py b/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py new file mode 100644 index 000000000..d132e8560 --- /dev/null +++ b/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py @@ -0,0 +1,196 @@ +"""drop foreign keys between the two services tables + +Revision ID: 7a83a87d029a +Revises: 2ffd50f47a2b +Create Date: 2026-11-04 00:19:04.175069 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +No post-generation corrections. The tables the simulation service writes keep +the columns naming the Public API's rows, without foreign keys. The downgrade +restores the keys, so it fails while any of those columns names a row the +Public API's table lacks. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "7a83a87d029a" +down_revision: Union[str, None] = "2ffd50f47a2b" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint( + op.f("fk_household_jobs_dynamic_id_dynamics"), + "household_jobs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_household_jobs_policy_id_policies"), + "household_jobs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_reports_household_id_households"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), + "reports", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_reports_region_id_regions"), "reports", type_="foreignkey" + ) + op.drop_constraint(op.f("fk_reports_user_id_users"), "reports", type_="foreignkey") + op.drop_constraint( + op.f("fk_reports_policy_id_policies"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_reports_dataset_id_datasets"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_dynamic_id_dynamics"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_household_id_households"), + "simulations", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_simulations_region_id_regions"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_dataset_id_datasets"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), + "simulations", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_simulations_policy_id_policies"), "simulations", type_="foreignkey" + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_foreign_key( + op.f("fk_simulations_policy_id_policies"), + "simulations", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), + "simulations", + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_simulations_dataset_id_datasets"), + "simulations", + "datasets", + ["dataset_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_simulations_region_id_regions"), + "simulations", + "regions", + ["region_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_simulations_household_id_households"), + "simulations", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_simulations_dynamic_id_dynamics"), + "simulations", + "dynamics", + ["dynamic_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_dataset_id_datasets"), + "reports", + "datasets", + ["dataset_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_policy_id_policies"), + "reports", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_user_id_users"), + "reports", + "users", + ["user_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_region_id_regions"), + "reports", + "regions", + ["region_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), + "reports", + "tax_benefit_models", + ["tax_benefit_model_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_reports_household_id_households"), + "reports", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_household_jobs_policy_id_policies"), + "household_jobs", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_household_jobs_dynamic_id_dynamics"), + "household_jobs", + "dynamics", + ["dynamic_id"], + ["id"], + ondelete="SET NULL", + ) + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 5ca67210e..1c7bc1d5e 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "2ffd50f47a2b" +V2_ALEMBIC_HEAD_REVISION = "7a83a87d029a" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 8ab02d5f0..85bb805d2 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -6,18 +6,16 @@ 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 TimestampedModel, enum_type -from policyengine_api.data.v2.models.policies import Dynamic, Policy 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): @@ -33,8 +31,9 @@ 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; simulations, reports and household jobs name a household + without a foreign key. """ __tablename__ = "households" @@ -77,13 +76,11 @@ class Household(TimestampedModel, table=True): ) 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", ) @@ -98,10 +95,9 @@ class HouseholdJob(TimestampedModel, table=True): ``policy_id`` is null for current law. ``request_data`` holds the calculation's options, such as axes. The job records the catalog version and .py version it ran on, its Modal call, and its status and result or - error. ``household_id`` and ``tax_benefit_model_version_id`` are logical - references to the Public API's rows, which the simulation service checks - before it writes the job. Deleting the job's policy or dynamic keeps the - job but clears the reference. + error. Its household, policy, dynamic and catalog version are logical + references to the Public API's rows, without foreign keys: the simulation + service checks them before it writes the job. """ __tablename__ = "household_jobs" @@ -121,16 +117,8 @@ class HouseholdJob(TimestampedModel, table=True): household_id: UUID = Field(index=True) year: int 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", - ) + policy_id: UUID | None = None + dynamic_id: UUID | None = None tax_benefit_model_version_id: UUID country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) @@ -149,6 +137,3 @@ class HouseholdJob(TimestampedModel, table=True): 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 e736010cc..6613a8b0e 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -18,8 +18,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): @@ -80,7 +78,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): @@ -128,9 +125,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") @@ -190,8 +184,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): @@ -244,11 +236,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"}, - ) - 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 f0da6e8cc..229a5f81a 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 index c379a56ed..8f25d35ec 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -18,15 +18,7 @@ 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, SimulationRun -from policyengine_api.data.v2.models.users import User class ReportRunStatus(StrEnum): @@ -78,7 +70,12 @@ class SimulationRole(StrEnum): class Report(TimestampedModel, table=True): - """Stable report definition; execution state belongs to ReportRun.""" + """Stable report definition; execution state belongs to ReportRun. + + Its user, model, policy, household, dataset and region are logical + references to the Public API's rows, without foreign keys: the simulation + service checks them before it writes the report. + """ __tablename__ = "reports" __table_args__ = ( @@ -93,21 +90,9 @@ class Report(TimestampedModel, table=True): 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", - ) + user_id: UUID | None = None + tax_benefit_model_id: UUID = Field(index=True) + policy_id: UUID | None = None baseline_simulation_id: UUID | None = Field( default=None, foreign_key="simulations.id", @@ -118,27 +103,12 @@ class Report(TimestampedModel, table=True): 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", - ) + household_id: UUID | None = None + dataset_id: UUID | None = None + region_id: UUID | None = None 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"}, @@ -147,9 +117,6 @@ class Report(TimestampedModel, table=True): 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, diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index f59038b00..ff24aae72 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -9,13 +9,6 @@ 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 ( @@ -46,8 +39,9 @@ class Simulation(TimestampedModel, table=True): """A simulation of exactly one household or one dataset, matching its type. It defines what to calculate; each calculation of it is a ``SimulationRun``. - Deleting its policy, dynamic or region clears that reference; its household - and input dataset cannot be deleted while it exists. + Its dataset, household, policy, dynamic, catalog version and region are + logical references to the Public API's rows, without foreign keys: the + simulation service checks them before it writes the simulation. """ __tablename__ = "simulations" @@ -73,52 +67,17 @@ class Simulation(TimestampedModel, table=True): 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, - ) - region_id: UUID | None = Field( - default=None, - foreign_key="regions.id", - ondelete="SET NULL", - ) + dataset_id: UUID | None = None + household_id: UUID | None = None + policy_id: UUID | None = None + dynamic_id: UUID | None = None + tax_benefit_model_version_id: UUID = Field(index=True) + region_id: UUID | 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 - dataset: Dataset | None = Relationship( - back_populates="input_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.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"}, diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index f3b987123..06c2b17f5 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,8 @@ 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; a report naming the user + as its owner keeps that ID, because it has no foreign key to this table. """ __tablename__ = "users" @@ -71,7 +70,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/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 3c01b0a84..769a0101e 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -39,6 +39,9 @@ AREA_IMPACT_COORDINATE_REVISION = "e86f19a6c4da" # The parameter search file revision, the parent of the simulation runs revision. SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" +# The simulation runs revision, parent of the revision that drops the foreign +# keys from the simulation service's tables to the Public API's tables. +CROSS_SERVICE_KEYS_PREVIOUS_REVISION = "2ffd50f47a2b" V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) AREA_IMPACT_TABLES = ("constituency_impacts", "local_authority_impacts") MAP_COORDINATE_COLUMNS = ("x", "y") @@ -411,6 +414,14 @@ def _assert_countries_shape(engine) -> None: assert check_name not in _check_names(engine, table_name) +def _assert_cross_service_keys_shape(engine) -> None: + _assert_revision(engine, CROSS_SERVICE_KEYS_PREVIOUS_REVISION) + referred_tables = { + key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") + } + assert {"datasets", "households", "policies"} <= referred_tables + + def _assert_pre_simulation_runs_shape(engine) -> None: _assert_revision(engine, SIMULATION_RUNS_PREVIOUS_REVISION) table_names = set(inspect(engine).get_table_names(schema="public")) @@ -446,6 +457,10 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: _upgrade_and_assert_head(config, engine) _assert_countries_shape(engine) + command.downgrade(config, CROSS_SERVICE_KEYS_PREVIOUS_REVISION) + _assert_cross_service_keys_shape(engine) + _upgrade_and_assert_head(config, engine) + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) _assert_pre_simulation_runs_shape(engine) _upgrade_and_assert_head(config, engine) diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 83ef41ade..02de26f6b 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,10 +1,11 @@ -"""Check the simulation runs revision against disposable Postgres. - -The revision adds simulation runs and the table naming the simulation runs a -report run combines, keys every result table by run, and gives household jobs -one household, one year and one policy. It changes no other table, so the -previous Public API revision, whose roles can't read or write any table it -changes, keeps working while it is deployed. +"""Check the simulation runs revisions against disposable Postgres. + +The first revision adds simulation runs and the table naming the simulation +runs a report run combines, keys every result table by run, and gives +household jobs one household, one year and one policy. The second drops the +foreign keys from those tables to the Public API's tables. Neither changes any +other table, so the previous Public API revision, whose roles can't read or +write any table they change, keeps working while they are deployed. """ from __future__ import annotations @@ -47,9 +48,15 @@ "poverty", "program_statistics", "report_runs", + "reports", "simulations", } ) +# The tables the simulation service and its Modal workers write; the Public API +# writes every other table. +SIMULATION_SIDE_TABLES = ( + CHANGED_TABLES | NEW_TABLES | {"aggregates", "change_aggregates"} +) SHA256_HEX = "a" * 64 @@ -132,6 +139,23 @@ def test_the_revision_adds_two_tables_and_changes_only_simulation_side_tables( assert changed == CHANGED_TABLES +def test_head_has_no_foreign_key_between_the_two_services_tables( + config: Config, engine: Engine +) -> None: + command.upgrade(config, "head") + inspector = inspect(engine) + + cross_service_keys = sorted( + f"{table_name}.{key['name']} -> {key['referred_table']}" + for table_name in inspector.get_table_names(schema="public") + for key in inspector.get_foreign_keys(table_name) + if (table_name in SIMULATION_SIDE_TABLES) + != (key["referred_table"] in SIMULATION_SIDE_TABLES) + ) + + assert cross_service_keys == [] + + @contextmanager def _refused(connection: Connection) -> Iterator[None]: """Expect the statement inside to violate a constraint, then roll it back.""" diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index e0bbbfc67..c539f0faf 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -80,6 +80,11 @@ def _disposable_url() -> str: return database_url +# Simulations and reports have no foreign key to the catalog, so truncating the +# catalog no longer reaches them; the tests that store them name them here. +TRUNCATE_PUBLICATION_ROWS = "TRUNCATE tax_benefit_models, simulations, reports CASCADE" + + @pytest.fixture def publication_engine() -> Engine: database_url = _disposable_url() @@ -87,10 +92,10 @@ def publication_engine() -> Engine: command.upgrade(config, "head") engine = create_engine(database_url, poolclass=NullPool) with engine.begin() as connection: - connection.execute(text("TRUNCATE tax_benefit_models CASCADE")) + connection.execute(text(TRUNCATE_PUBLICATION_ROWS)) yield engine with engine.begin() as connection: - connection.execute(text("TRUNCATE tax_benefit_models CASCADE")) + connection.execute(text(TRUNCATE_PUBLICATION_ROWS)) engine.dispose() diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 48147fdf5..080b6e329 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -70,6 +70,28 @@ REPO / "migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py" ) +# The simulation runs revision, parent of the cross-service foreign key revision. +CROSS_SERVICE_KEYS_PARENT_REVISION = "2ffd50f47a2b" +# The foreign keys from the simulation service's tables to the Public API's. +CROSS_SERVICE_KEYS = { + "household_jobs": {"dynamics", "policies"}, + "reports": { + "datasets", + "households", + "policies", + "regions", + "tax_benefit_models", + "users", + }, + "simulations": { + "datasets", + "dynamics", + "households", + "policies", + "regions", + "tax_benefit_model_versions", + }, +} DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -251,6 +273,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + CROSS_SERVICE_KEYS_PARENT_REVISION, SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, AREA_IMPACT_COORDINATE_REVISION, @@ -1063,3 +1086,59 @@ def test_simulation_runs_revision_has_only_its_two_documented_corrections() -> N for body in (upgrade, downgrade): assert "op.execute(" not in body assert "op.bulk_insert(" not in body + + +def test_cross_service_key_revision_only_drops_and_restores_foreign_keys() -> None: + """The revision after the simulation runs revision is generated. + + Its upgrade drops each foreign key from a table the simulation service + writes to a table the Public API writes; its downgrade restores them. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == CROSS_SERVICE_KEYS_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + dropped_tables = re.findall( + r'op\.drop_constraint\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*' + r'type_="foreignkey"', + upgrade, + ) + restored_keys = re.findall( + r'op\.create_foreign_key\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*"([a-z_]+)"', + downgrade, + ) + expected_count = sum(len(targets) for targets in CROSS_SERVICE_KEYS.values()) + assert len(dropped_tables) == upgrade.count("op.drop_constraint(") + assert len(dropped_tables) == expected_count + assert {table: dropped_tables.count(table) for table in CROSS_SERVICE_KEYS} == { + table: len(targets) for table, targets in CROSS_SERVICE_KEYS.items() + } + assert len(restored_keys) == downgrade.count("op.create_foreign_key(") + restored_targets: dict[str, set[str]] = {} + for table, target in restored_keys: + restored_targets.setdefault(table, set()).add(target) + assert restored_targets == CROSS_SERVICE_KEYS + for body in (upgrade, downgrade): + for operation in ( + "op.add_column(", + "op.drop_column(", + "op.create_table(", + "op.drop_table(", + "op.execute(", + "op.bulk_insert(", + ): + assert operation not in body diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index db60cd8cf..976ebe7e6 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -436,14 +436,15 @@ def test_regions_have_one_same_model_version_default_logical_dataset() -> None: } -def test_reports_and_simulations_snapshot_selected_datasets() -> None: +def test_reports_and_simulations_name_their_selected_dataset() -> None: + """The dataset is the Public API's catalog row, named without a foreign key.""" 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"} + assert isinstance(table.c.dataset_id.type, sa.Uuid) + assert table.c.dataset_id.nullable + assert not table.c.dataset_id.foreign_keys def test_stage9_adds_no_dataset_version_relationship_to_run_tables() -> None: diff --git a/tests/unit/v2/test_report_run_creation.py b/tests/unit/v2/test_report_run_creation.py index 3a1b1b82c..7aa06a392 100644 --- a/tests/unit/v2/test_report_run_creation.py +++ b/tests/unit/v2/test_report_run_creation.py @@ -38,10 +38,13 @@ def engine(): def _create_report(session: Session) -> Report: + model = TaxBenefitModel(name=f"model-{uuid4()}") + session.add(model) + session.flush() report = Report( label="Distributional report", country="us", - tax_benefit_model=TaxBenefitModel(name=f"model-{uuid4()}"), + tax_benefit_model_id=model.id, inputs={}, ) session.add(report) diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py index 4383a8110..0221624e5 100644 --- a/tests/unit/v2/test_report_runs.py +++ b/tests/unit/v2/test_report_runs.py @@ -56,11 +56,13 @@ def engine(): def _create_report(session: Session, *, report_type: str | None = None) -> Report: model = TaxBenefitModel(name=f"model-{uuid4()}") + session.add(model) + session.flush() report = Report( label="Distributional report", country="us", type=report_type, - tax_benefit_model=model, + tax_benefit_model_id=model.id, inputs={"reform": {"gov.example": 1}}, ) session.add(report) @@ -249,23 +251,26 @@ def test_worker_retry_resumes_the_same_run_and_terminal_runs_stay_terminal( 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, + model_id=report.tax_benefit_model_id, version="1.2.3", current_law_id=1, metadata_time_periods=[2026], ) + session.add(version) + session.flush() dataset = Dataset( name="dataset", storage_path="datasets/test.h5", year=2026, - tax_benefit_model_version=version, + tax_benefit_model_version_id=version.id, ) + session.add(dataset) + session.flush() simulation = Simulation( simulation_type=SimulationType.ECONOMY, - dataset=dataset, - tax_benefit_model_version=version, + dataset_id=dataset.id, + tax_benefit_model_version_id=version.id, ) session.add(simulation) session.flush() diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 64cea61f4..ffc60870c 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -8,6 +8,7 @@ import sqlalchemy as sa from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import configure_mappers from sqlalchemy.schema import CreateTable from policyengine_api.data.v2.models import ( @@ -15,6 +16,7 @@ ReportRunSimulationRun, SimulationRole, SimulationRun, + User, ) RUN_STATUSES = ["pending", "running", "succeeded", "failed"] @@ -28,6 +30,49 @@ "local_authority_impacts", "program_statistics", ) +# The tables the simulation service and its Modal workers write. The Public API +# writes every other table. +SIMULATION_SIDE_TABLES = frozenset( + { + "simulations", + "simulation_runs", + "reports", + "report_runs", + "report_run_simulation_runs", + "household_jobs", + "aggregates", + "change_aggregates", + *PER_SIMULATION_RESULT_TABLES, + *COMPARISON_RESULT_TABLES, + } +) +# Columns naming the Public API's rows, which the simulation service checks +# before it writes the row that names them. +LOGICAL_REFERENCES = { + "simulations": ( + "dataset_id", + "household_id", + "policy_id", + "dynamic_id", + "tax_benefit_model_version_id", + "region_id", + ), + "reports": ( + "user_id", + "tax_benefit_model_id", + "policy_id", + "household_id", + "dataset_id", + "region_id", + ), + "simulation_runs": ("tax_benefit_model_version_id", "dataset_id", "region_id"), + "household_jobs": ( + "household_id", + "policy_id", + "dynamic_id", + "tax_benefit_model_version_id", + ), +} def _foreign_key(table: sa.Table, column_name: str) -> sa.ForeignKey: @@ -284,3 +329,40 @@ def test_custom_totals_keep_their_simulation_references() -> None: assert _foreign_key(table, "report_run_id").target_fullname == ( "report_runs.id" ) + + +def test_no_foreign_key_joins_tables_that_different_services_write() -> None: + """Each service checks the other service's rows its own rows name.""" + cross_service_keys = sorted( + f"{table.name}.{key.parent.name} -> {key.column.table.name}" + for table in V2_METADATA.tables.values() + for key in table.foreign_keys + if (table.name in SIMULATION_SIDE_TABLES) + != (key.column.table.name in SIMULATION_SIDE_TABLES) + ) + + assert cross_service_keys == [] + assert set(V2_METADATA.tables) >= SIMULATION_SIDE_TABLES + + +def test_references_to_public_api_rows_stay_as_uuid_columns() -> None: + for table_name, column_names in LOGICAL_REFERENCES.items(): + table = V2_METADATA.tables[table_name] + for column_name in column_names: + column = table.c[column_name] + assert isinstance(column.type, sa.Uuid), f"{table_name}.{column_name}" + assert not column.foreign_keys, f"{table_name}.{column_name}" + + +def test_no_relationship_joins_models_that_different_services_write() -> None: + configure_mappers() + cross_service_relationships = sorted( + f"{mapper.class_.__name__}.{relationship.key}" + for mapper in sa.inspect(User).registry.mappers + if mapper.local_table is not None and mapper.local_table.metadata is V2_METADATA + for relationship in mapper.relationships + if (mapper.local_table.name in SIMULATION_SIDE_TABLES) + != (relationship.mapper.local_table.name in SIMULATION_SIDE_TABLES) + ) + + assert cross_service_relationships == [] From 0ec365a4b92239539ce91780f03e1d7707f8ffc7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:34:52 +0400 Subject: [PATCH 03/22] Store the v2 economy output fields that had no column The simulation API's economy output carries fields with no column in the v2 result tables. Revision 7f43d2c7c3c6, autogenerated with no corrections, adds report_runs.spm_provenance and data_version, the congressional districts' winner, loser and no-change shares, and the UK areas' disclosure tier and interval bounds. The detailed budget names neither a program's entity nor whether it is a tax, so both become nullable in program_statistics, whose unique key treats a missing entity as equal. The revision-by-revision downgrade checks move from the v2 lifecycle test, which reached the test-file length limit, to the simulation runs integration tests. Co-Authored-By: Claude Opus 5.5 --- changelog.d/v2-simulation-runs.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 17 + ...tore_the_economy_output_fields_without_.py | 174 ++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/reports.py | 50 +- .../integration/test_alembic_v2_lifecycle.py | 35 -- .../test_alembic_v2_simulation_runs.py | 508 +++++++++++------- tests/unit/v2/test_alembic_v2.py | 65 +++ tests/unit/v2/test_simulation_run_models.py | 48 +- 9 files changed, 646 insertions(+), 255 deletions(-) create mode 100644 migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py diff --git a/changelog.d/v2-simulation-runs.changed.md b/changelog.d/v2-simulation-runs.changed.md index eef5d02ad..9e459c853 100644 --- a/changelog.d/v2-simulation-runs.changed.md +++ b/changelog.d/v2-simulation-runs.changed.md @@ -1 +1 @@ -API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. Its simulation tables no longer have foreign keys to the tables the Public API writes. +API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. Its simulation tables no longer have foreign keys to the tables the Public API writes. Result tables gain the economy output's fields that had no column. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 21a1e8a01..786565d15 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -515,3 +515,20 @@ service stay. It has no post-generation corrections. Its downgrade restores the keys, so it fails while a column names a row the referenced table lacks. `tests/unit/v2/test_simulation_run_models.py` checks that no foreign key or ORM relationship joins the two services' tables. + +### API v2 economy output fields + +Revision `7f43d2c7c3c6` (parent `7a83a87d029a`) was autogenerated after the +models gained the fields of the simulation API's economy output that had no +column: `report_runs.spm_provenance` (JSON) and `data_version`; the winner, +loser and no-change shares of `congressional_district_impacts`; and the UK +disclosure control fields (`disclosure_tier`, checked to be `reported` or +`approximate`, and the four `_lower`/`_upper` interval bounds) of +`constituency_impacts` and `local_authority_impacts`. The output's detailed +budget names neither a program's entity nor whether it is a tax, so +`program_statistics.entity` and `is_tax` become nullable, and its unique key +treats a missing entity as equal (`NULLS NOT DISTINCT`). It has no +post-generation corrections. The three share columns are required, so the +upgrade needs `congressional_district_impacts` empty; the downgrade needs +every `program_statistics` row to have an entity and a tax flag. + diff --git a/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py b/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py new file mode 100644 index 000000000..2ff0acd21 --- /dev/null +++ b/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py @@ -0,0 +1,174 @@ +"""store the economy output fields without a column + +Revision ID: 7f43d2c7c3c6 +Revises: 7a83a87d029a +Create Date: 2026-11-04 00:30:48.602215 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +No post-generation corrections. The upgrade adds required columns without +server defaults to congressional_district_impacts, so it needs that table +empty. The downgrade makes program_statistics.entity and is_tax required +again, so it needs every row of that table to have both. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "7f43d2c7c3c6" +down_revision: Union[str, None] = "7a83a87d029a" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "congressional_district_impacts", + sa.Column("winner_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("loser_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("no_change_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "constituency_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) + op.create_check_constraint( + op.f("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) + op.add_column( + "local_authority_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) + op.create_check_constraint( + op.f("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) + op.alter_column( + "program_statistics", + "entity", + existing_type=sa.VARCHAR(length=128), + nullable=True, + ) + op.alter_column( + "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=True + ) + op.drop_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + type_="unique", + ) + op.create_unique_constraint( + "uq_program_statistics_run_program_entity", + "program_statistics", + ["report_run_id", "program_name", "entity"], + postgresql_nulls_not_distinct=True, + ) + op.add_column("report_runs", sa.Column("spm_provenance", sa.JSON(), nullable=True)) + op.add_column( + "report_runs", + sa.Column( + "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True + ), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("report_runs", "data_version") + op.drop_column("report_runs", "spm_provenance") + op.drop_constraint( + "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + ["report_run_id", "program_name", "entity"], + postgresql_nulls_not_distinct=False, + ) + op.alter_column( + "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=False + ) + op.alter_column( + "program_statistics", + "entity", + existing_type=sa.VARCHAR(length=128), + nullable=False, + ) + op.drop_constraint( + op.f("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + type_="check", + ) + op.drop_column("local_authority_impacts", "relative_household_income_change_upper") + op.drop_column("local_authority_impacts", "relative_household_income_change_lower") + op.drop_column("local_authority_impacts", "average_household_income_change_upper") + op.drop_column("local_authority_impacts", "average_household_income_change_lower") + op.drop_column("local_authority_impacts", "disclosure_tier") + op.drop_constraint( + op.f("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + type_="check", + ) + op.drop_column("constituency_impacts", "relative_household_income_change_upper") + op.drop_column("constituency_impacts", "relative_household_income_change_lower") + op.drop_column("constituency_impacts", "average_household_income_change_upper") + op.drop_column("constituency_impacts", "average_household_income_change_lower") + op.drop_column("constituency_impacts", "disclosure_tier") + op.drop_column("congressional_district_impacts", "no_change_percentage") + op.drop_column("congressional_district_impacts", "loser_percentage") + op.drop_column("congressional_district_impacts", "winner_percentage") + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 1c7bc1d5e..b94e05ae3 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "7a83a87d029a" +V2_ALEMBIC_HEAD_REVISION = "7f43d2c7c3c6" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 8f25d35ec..238901192 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -127,8 +127,9 @@ class ReportRun(TimestampedModel, table=True): """One immutable-version execution attempt for a stable report. Its ``simulation_runs`` rows name the simulation runs it combines. - ``labor_supply_result`` and ``cliff_impact_result`` hold the two parts of an - economy comparison that have no result table. + ``labor_supply_result``, ``cliff_impact_result`` and ``spm_provenance`` hold + the parts of an economy comparison that have no result table, and + ``data_version`` the release of the data it used. """ __tablename__ = "report_runs" @@ -182,6 +183,8 @@ class ReportRun(TimestampedModel, table=True): markdown: str | None = Field(default=None, sa_type=sa.Text) labor_supply_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) cliff_impact_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + spm_provenance: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + data_version: str | None = Field(default=None, max_length=128) report: Report = Relationship(back_populates="runs") simulation_runs: list["ReportRunSimulationRun"] = Relationship( @@ -538,7 +541,9 @@ class Poverty(IdentifiedModel, table=True): 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. + One row per run, program and entity; ``is_tax`` marks taxes. The economy + output's detailed budget names neither the entity nor whether a program is + a tax, so both may be null, and rows without an entity count as equal. """ __tablename__ = "program_statistics" @@ -548,6 +553,7 @@ class ProgramStatistics(IdentifiedModel, table=True): "program_name", "entity", name="uq_program_statistics_run_program_entity", + postgresql_nulls_not_distinct=True, ), ) @@ -564,8 +570,8 @@ class ProgramStatistics(IdentifiedModel, table=True): ondelete="RESTRICT", ) program_name: str = Field(max_length=512) - entity: str = Field(max_length=128) - is_tax: bool = False + entity: str | None = Field(default=None, max_length=128) + is_tax: bool | None = None baseline_total: float | None = None reform_total: float | None = None change: float | None = None @@ -597,10 +603,33 @@ class GeographicImpactBase(IdentifiedModel): population: float +class UKAreaImpactBase(GeographicImpactBase): + """Shared non-table disclosure fields of a UK area's impact. + + ``disclosure_tier`` is ``reported``, or ``approximate`` when the area's + sample is thin; the ``_lower`` and ``_upper`` fields bound each income + change's 90% interval. All are null for results without disclosure control. + """ + + disclosure_tier: str | None = Field(default=None, max_length=16) + average_household_income_change_lower: float | None = None + average_household_income_change_upper: float | None = None + relative_household_income_change_lower: float | None = None + relative_household_income_change_upper: float | None = None + + +def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: + return sa.CheckConstraint( + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + name=f"ck_{table_name}_disclosure_tier", + ) + + class CongressionalDistrictImpact(GeographicImpactBase, table=True): """Household income change and population in one US congressional district. - One row per run and district GEOID. + One row per run and district GEOID, with the shares of people who gain, + lose or see no change. """ __tablename__ = "congressional_district_impacts" @@ -615,13 +644,16 @@ class CongressionalDistrictImpact(GeographicImpactBase, table=True): district_geoid: int state_fips: int district_number: int + winner_percentage: float + loser_percentage: float + no_change_percentage: float report_run: ReportRun = Relationship( back_populates="congressional_district_impacts" ) -class ConstituencyImpact(GeographicImpactBase, table=True): +class ConstituencyImpact(UKAreaImpactBase, table=True): """Household income change in one UK constituency. One row per run and constituency code. @@ -634,6 +666,7 @@ class ConstituencyImpact(GeographicImpactBase, table=True): "constituency_code", name="uq_constituency_impacts_run_code", ), + _disclosure_tier_check("constituency_impacts"), ) constituency_code: str = Field(max_length=64) @@ -642,7 +675,7 @@ class ConstituencyImpact(GeographicImpactBase, table=True): report_run: ReportRun = Relationship(back_populates="constituency_impacts") -class LocalAuthorityImpact(GeographicImpactBase, table=True): +class LocalAuthorityImpact(UKAreaImpactBase, table=True): """Household income change in one UK local authority. One row per run and local authority code. @@ -655,6 +688,7 @@ class LocalAuthorityImpact(GeographicImpactBase, table=True): "local_authority_code", name="uq_local_authority_impacts_run_code", ), + _disclosure_tier_check("local_authority_impacts"), ) local_authority_code: str = Field(max_length=64) diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 769a0101e..3c2cc8161 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -37,11 +37,6 @@ # The catalog format 2 revision, the parent of the area impact coordinate removal. CATALOG_FORMAT_2_REVISION = "f5c1d9b4685b" AREA_IMPACT_COORDINATE_REVISION = "e86f19a6c4da" -# The parameter search file revision, the parent of the simulation runs revision. -SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" -# The simulation runs revision, parent of the revision that drops the foreign -# keys from the simulation service's tables to the Public API's tables. -CROSS_SERVICE_KEYS_PREVIOUS_REVISION = "2ffd50f47a2b" V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) AREA_IMPACT_TABLES = ("constituency_impacts", "local_authority_impacts") MAP_COORDINATE_COLUMNS = ("x", "y") @@ -414,28 +409,6 @@ def _assert_countries_shape(engine) -> None: assert check_name not in _check_names(engine, table_name) -def _assert_cross_service_keys_shape(engine) -> None: - _assert_revision(engine, CROSS_SERVICE_KEYS_PREVIOUS_REVISION) - referred_tables = { - key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") - } - assert {"datasets", "households", "policies"} <= referred_tables - - -def _assert_pre_simulation_runs_shape(engine) -> None: - _assert_revision(engine, SIMULATION_RUNS_PREVIOUS_REVISION) - table_names = set(inspect(engine).get_table_names(schema="public")) - assert {"simulation_runs", "report_run_simulation_runs"}.isdisjoint(table_names) - simulation_columns = { - column["name"] for column in inspect(engine).get_columns("simulations") - } - poverty_columns = { - column["name"] for column in inspect(engine).get_columns("poverty") - } - assert {"status", "household_result", "output_dataset_id"} <= simulation_columns - assert {"report_run_id", "simulation_id"} <= poverty_columns - - def _assert_pre_countries_shape(engine) -> None: assert "countries" not in inspect(engine).get_table_names(schema="public") for (table_name, _column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): @@ -457,14 +430,6 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: _upgrade_and_assert_head(config, engine) _assert_countries_shape(engine) - command.downgrade(config, CROSS_SERVICE_KEYS_PREVIOUS_REVISION) - _assert_cross_service_keys_shape(engine) - _upgrade_and_assert_head(config, engine) - - command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) - _assert_pre_simulation_runs_shape(engine) - _upgrade_and_assert_head(config, engine) - command.downgrade(config, CATALOG_FORMAT_2_REVISION) _assert_area_impact_coordinates_restored(engine) _upgrade_and_assert_head(config, engine) diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 02de26f6b..1f1f18f6e 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,11 +1,13 @@ """Check the simulation runs revisions against disposable Postgres. -The first revision adds simulation runs and the table naming the simulation -runs a report run combines, keys every result table by run, and gives -household jobs one household, one year and one policy. The second drops the -foreign keys from those tables to the Public API's tables. Neither changes any -other table, so the previous Public API revision, whose roles can't read or -write any table they change, keeps working while they are deployed. +Three revisions follow the area impact coordinate removal. The first adds +simulation runs and the table naming the simulation runs a report run +combines, keys every result table by run, and gives household jobs one +household, one year and one policy. The second drops the foreign keys from +those tables to the Public API's tables. The third adds the economy output's +fields that had no column. None changes any other table, so the previous +Public API revision, whose roles can't read or write any table they change, +keeps working while they are deployed. """ from __future__ import annotations @@ -34,6 +36,8 @@ from sqlalchemy import Connection, Engine SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" +SIMULATION_RUNS_REVISION = "2ffd50f47a2b" +CROSS_SERVICE_KEYS_REVISION = "7a83a87d029a" NEW_TABLES = frozenset({"simulation_runs", "report_run_simulation_runs"}) CHANGED_TABLES = frozenset( { @@ -57,7 +61,9 @@ SIMULATION_SIDE_TABLES = ( CHANGED_TABLES | NEW_TABLES | {"aggregates", "change_aggregates"} ) +RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} SHA256_HEX = "a" * 64 +COMPLETED_AT = "2026-11-04T00:00:00Z" def _disposable_url() -> str: @@ -90,6 +96,18 @@ def engine(config: Config) -> Iterator[Engine]: database_engine.dispose() +@pytest.fixture +def connection(config: Config, engine: Engine) -> Iterator[Connection]: + """Yield a connection at head whose writes are rolled back after the test.""" + command.upgrade(config, "head") + with engine.connect() as database_connection: + transaction = database_connection.begin() + try: + yield database_connection + finally: + transaction.rollback() + + def _table_shape(engine: Engine, table_name: str) -> dict[str, object]: """Describe one table's columns, keys, indexes and checks for comparison.""" inspector = inspect(engine) @@ -123,7 +141,20 @@ def _shapes(engine: Engine) -> dict[str, dict[str, object]]: } -def test_the_revision_adds_two_tables_and_changes_only_simulation_side_tables( +def _column_names(engine: Engine, table_name: str) -> set[str]: + return {column["name"] for column in inspect(engine).get_columns(table_name)} + + +def _enum_type_names(engine: Engine) -> set[str]: + with engine.connect() as type_connection: + return set( + type_connection.execute( + text("SELECT typname FROM pg_type WHERE typtype = 'e'") + ).scalars() + ) + + +def test_the_revisions_add_two_tables_and_change_only_simulation_side_tables( config: Config, engine: Engine ) -> None: command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) @@ -139,6 +170,42 @@ def test_the_revision_adds_two_tables_and_changes_only_simulation_side_tables( assert changed == CHANGED_TABLES +def test_each_revision_downgrades_to_its_parent_and_upgrades_again( + config: Config, engine: Engine +) -> None: + command.upgrade(config, "head") + + command.downgrade(config, CROSS_SERVICE_KEYS_REVISION) + assert {"spm_provenance", "data_version"}.isdisjoint( + _column_names(engine, "report_runs") + ) + program_entity = next( + column + for column in inspect(engine).get_columns("program_statistics") + if column["name"] == "entity" + ) + assert not program_entity["nullable"] + + command.downgrade(config, SIMULATION_RUNS_REVISION) + assert {"datasets", "households", "policies"} <= { + key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") + } + + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + assert NEW_TABLES.isdisjoint(inspect(engine).get_table_names(schema="public")) + assert {"status", "household_result", "output_dataset_id"} <= _column_names( + engine, "simulations" + ) + assert {"report_run_id", "simulation_id"} <= _column_names(engine, "poverty") + enum_types = _enum_type_names(engine) + assert "v2_simulation_status" in enum_types + assert "v2_simulation_role" not in enum_types + + command.upgrade(config, "head") + command.check(config) + assert "v2_simulation_role" in _enum_type_names(engine) + + def test_head_has_no_foreign_key_between_the_two_services_tables( config: Config, engine: Engine ) -> None: @@ -173,10 +240,31 @@ def _insert(connection: Connection, table: str, **values: object) -> None: ) +def _count(connection: Connection, table: str, column: str, value: object) -> int: + return connection.execute( + text(f"SELECT count(*) FROM {table} WHERE {column} = :value"), + {"value": value}, + ).scalar_one() + + +def _insert_simulation(connection: Connection) -> object: + """Store an economy simulation; its catalog references need no catalog rows.""" + simulation_id = uuid4() + _insert( + connection, + "simulations", + id=simulation_id, + simulation_type="economy", + dataset_id=uuid4(), + tax_benefit_model_version_id=uuid4(), + year=2026, + ) + return simulation_id + + def _insert_simulation_run( connection: Connection, simulation_id: object, - version_id: object, *, status: str = "pending", **values: object, @@ -187,212 +275,214 @@ def _insert_simulation_run( "simulation_runs", id=run_id, simulation_id=simulation_id, - tax_benefit_model_version_id=version_id, - country_package_version="1.0.0", - policyengine_version="4.0.0", + tax_benefit_model_version_id=uuid4(), status=status, + **RUN_VERSIONS, **values, ) return run_id -def test_run_tables_enforce_completion_output_and_uniqueness( - config: Config, engine: Engine +def _insert_report_run(connection: Connection) -> object: + report_id, report_run_id = uuid4(), uuid4() + _insert( + connection, + "reports", + id=report_id, + label="Report", + country="us", + tax_benefit_model_id=uuid4(), + inputs="{}", + ) + _insert( + connection, + "report_runs", + id=report_run_id, + report_id=report_id, + status="running", + trigger="initial", + **RUN_VERSIONS, + ) + return report_run_id + + +def test_simulation_runs_need_a_completion_time_and_a_whole_output( + connection: Connection, ) -> None: - command.upgrade(config, "head") - model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() - simulation_id, report_id, report_run_id = uuid4(), uuid4(), uuid4() + simulation_id = _insert_simulation(connection) + run_id = _insert_simulation_run(connection, simulation_id) + + with _refused(connection): + connection.execute( + text("UPDATE simulation_runs SET status = 'succeeded' WHERE id = :id"), + {"id": run_id}, + ) + with _refused(connection): + _insert_simulation_run( + connection, simulation_id, output_uri="gs://bucket/run.parquet" + ) + with _refused(connection): + _insert_simulation_run( + connection, + simulation_id, + output_uri="gs://bucket/run.parquet", + output_sha256="short", + ) + finished_run_id = _insert_simulation_run( + connection, + simulation_id, + status="succeeded", + completed_at=COMPLETED_AT, + output_uri="gs://bucket/run.parquet", + output_sha256=SHA256_HEX, + ) + _insert( + connection, + "poverty", + id=uuid4(), + simulation_run_id=finished_run_id, + poverty_type="spm", + entity="person", + ) - with engine.connect() as connection: - transaction = connection.begin() - try: - _insert( - connection, - "tax_benefit_models", - id=model_id, - name=f"runs-{model_id.hex[:8]}", - ) - connection.execute( - text( - "INSERT INTO tax_benefit_model_versions " - "(id, model_id, version, current_law_id, metadata_time_periods) " - "VALUES (:id, :model_id, '1.0.0', 1, '[2026]')" - ), - {"id": version_id, "model_id": model_id}, - ) - _insert( - connection, - "datasets", - id=dataset_id, - name="populace", - year=2026, - is_output_dataset=False, - tax_benefit_model_version_id=version_id, - ) - _insert( - connection, - "simulations", - id=simulation_id, - simulation_type="economy", - dataset_id=dataset_id, - tax_benefit_model_version_id=version_id, - year=2026, - ) - baseline_run_id = _insert_simulation_run( - connection, simulation_id, version_id - ) - reform_run_id = _insert_simulation_run( - connection, - simulation_id, - version_id, - status="succeeded", - completed_at="2026-11-04T00:00:00Z", - output_uri="gs://bucket/run.parquet", - output_sha256=SHA256_HEX, - ) - - with _refused(connection): - connection.execute( - text( - "UPDATE simulation_runs SET status = 'succeeded' WHERE id = :id" - ), - {"id": baseline_run_id}, - ) - with _refused(connection): - _insert_simulation_run( - connection, - simulation_id, - version_id, - output_uri="gs://bucket/run.parquet", - ) - with _refused(connection): - _insert_simulation_run( - connection, - simulation_id, - version_id, - output_uri="gs://bucket/run.parquet", - output_sha256="short", - ) - - _insert( - connection, - "reports", - id=report_id, - label="Report", - country="us", - tax_benefit_model_id=model_id, - inputs="{}", - ) - _insert( - connection, - "report_runs", - id=report_run_id, - report_id=report_id, - country_package_version="1.0.0", - policyengine_version="4.0.0", - status="running", - trigger="initial", - ) - _insert( - connection, - "report_run_simulation_runs", - id=uuid4(), - report_run_id=report_run_id, - simulation_run_id=baseline_run_id, - role="baseline", - ) - with _refused(connection): - _insert( - connection, - "report_run_simulation_runs", - id=uuid4(), - report_run_id=report_run_id, - simulation_run_id=baseline_run_id, - role="baseline", - ) - with _refused(connection): - _insert( - connection, - "budget_summary", - id=uuid4(), - report_run_id=report_run_id, - baseline_simulation_run_id=baseline_run_id, - reform_simulation_run_id=uuid4(), - variable_name="household_net_income", - entity="household", - ) - - _insert( - connection, - "poverty", - id=uuid4(), - simulation_run_id=reform_run_id, - poverty_type="poverty", - entity="person", - ) - with _refused(connection): - _insert( - connection, - "poverty", - id=uuid4(), - simulation_run_id=reform_run_id, - poverty_type="poverty", - entity="person", - ) - - household_job = { - "country": "us", - "household_id": uuid4(), - "request_data": "{}", - "tax_benefit_model_version_id": version_id, - "country_package_version": "1.0.0", - "policyengine_version": "4.0.0", - } - _insert( - connection, - "household_jobs", - id=uuid4(), - year=2026, - status="pending", - **household_job, - ) - with _refused(connection): - _insert( - connection, - "household_jobs", - id=uuid4(), - year=1800, - status="pending", - **household_job, - ) - with _refused(connection): - _insert( - connection, - "household_jobs", - id=uuid4(), - year=2026, - status="failed", - **household_job, - ) - - connection.execute( - text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} - ) - assert ( - connection.execute( - text("SELECT count(*) FROM report_run_simulation_runs") - ).scalar_one() - == 0 - ) - connection.execute( - text("DELETE FROM simulations WHERE id = :id"), {"id": simulation_id} - ) - remaining_runs = connection.execute( - text("SELECT count(*) FROM simulation_runs") - ).scalar_one() - remaining_poverty = connection.execute( - text("SELECT count(*) FROM poverty") - ).scalar_one() - assert (remaining_runs, remaining_poverty) == (0, 0) - finally: - transaction.rollback() + connection.execute( + text("DELETE FROM simulations WHERE id = :id"), {"id": simulation_id} + ) + assert _count(connection, "simulation_runs", "simulation_id", simulation_id) == 0 + assert _count(connection, "poverty", "simulation_run_id", finished_run_id) == 0 + + +def test_report_run_members_are_unique_and_keep_their_simulation_runs( + connection: Connection, +) -> None: + run_id = _insert_simulation_run(connection, _insert_simulation(connection)) + report_run_id = _insert_report_run(connection) + member = { + "report_run_id": report_run_id, + "simulation_run_id": run_id, + "role": "baseline", + } + _insert(connection, "report_run_simulation_runs", id=uuid4(), **member) + + with _refused(connection): + _insert(connection, "report_run_simulation_runs", id=uuid4(), **member) + with _refused(connection): + connection.execute( + text("DELETE FROM simulation_runs WHERE id = :id"), {"id": run_id} + ) + connection.execute( + text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} + ) + assert ( + _count(connection, "report_run_simulation_runs", "report_run_id", report_run_id) + == 0 + ) + + +def test_results_are_unique_per_run_and_name_runs_that_exist( + connection: Connection, +) -> None: + simulation_id = _insert_simulation(connection) + baseline_run_id = _insert_simulation_run(connection, simulation_id) + reform_run_id = _insert_simulation_run(connection, simulation_id) + comparison_runs = { + "report_run_id": _insert_report_run(connection), + "baseline_simulation_run_id": baseline_run_id, + "reform_simulation_run_id": reform_run_id, + } + + with _refused(connection): + _insert( + connection, + "budget_summary", + id=uuid4(), + variable_name="household_net_income", + entity="household", + **{**comparison_runs, "reform_simulation_run_id": uuid4()}, + ) + poverty = { + "simulation_run_id": reform_run_id, + "poverty_type": "spm", + "entity": "person", + } + _insert(connection, "poverty", id=uuid4(), **poverty) + with _refused(connection): + _insert(connection, "poverty", id=uuid4(), **poverty) + _insert( + connection, + "program_statistics", + id=uuid4(), + program_name="income_tax", + **comparison_runs, + ) + with _refused(connection): + _insert( + connection, + "program_statistics", + id=uuid4(), + program_name="income_tax", + **comparison_runs, + ) + constituency = { + "constituency_name": "Aldershot", + "average_household_income_change": 1.0, + "relative_household_income_change": 0.01, + "population": 100.0, + **comparison_runs, + } + _insert( + connection, + "constituency_impacts", + id=uuid4(), + constituency_code="E14000530", + disclosure_tier="approximate", + **constituency, + ) + with _refused(connection): + _insert( + connection, + "constituency_impacts", + id=uuid4(), + constituency_code="E14000531", + disclosure_tier="withheld", + **constituency, + ) + + +def test_household_jobs_take_a_valid_year_and_record_when_they_end( + connection: Connection, +) -> None: + household_job = { + "country": "us", + "household_id": uuid4(), + "request_data": "{}", + "tax_benefit_model_version_id": uuid4(), + **RUN_VERSIONS, + } + _insert( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="pending", + **household_job, + ) + + with _refused(connection): + _insert( + connection, + "household_jobs", + id=uuid4(), + year=1800, + status="pending", + **household_job, + ) + with _refused(connection): + _insert( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="failed", + **household_job, + ) diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 080b6e329..e90ec3a27 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -72,6 +72,8 @@ ) # The simulation runs revision, parent of the cross-service foreign key revision. CROSS_SERVICE_KEYS_PARENT_REVISION = "2ffd50f47a2b" +# The cross-service foreign key revision, parent of the economy output revision. +ECONOMY_OUTPUT_PARENT_REVISION = "7a83a87d029a" # The foreign keys from the simulation service's tables to the Public API's. CROSS_SERVICE_KEYS = { "household_jobs": {"dynamics", "policies"}, @@ -273,6 +275,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + ECONOMY_OUTPUT_PARENT_REVISION, CROSS_SERVICE_KEYS_PARENT_REVISION, SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, @@ -1142,3 +1145,65 @@ def test_cross_service_key_revision_only_drops_and_restores_foreign_keys() -> No "op.bulk_insert(", ): assert operation not in body + + +def test_economy_output_revision_only_adds_columns_and_relaxes_two() -> None: + """The revision after the cross-service key revision is generated. + + It adds the economy output's fields that had no column and lets a detailed + budget row omit its entity and tax flag; it removes nothing. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == ECONOMY_OUTPUT_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + added_columns = set( + re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) + ) + uk_disclosure_columns = { + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", + } + assert added_columns == { + ("report_runs", "spm_provenance"), + ("report_runs", "data_version"), + ("congressional_district_impacts", "winner_percentage"), + ("congressional_district_impacts", "loser_percentage"), + ("congressional_district_impacts", "no_change_percentage"), + *(("constituency_impacts", column) for column in uk_disclosure_columns), + *(("local_authority_impacts", column) for column in uk_disclosure_columns), + } + assert upgrade.count("op.add_column(") == len(added_columns) + relaxed_columns = re.findall( + r'op\.alter_column\(\s*"(\w+)",\s*"(\w+)",\s*' + r"existing_type=sa\.\w+\([^)]*\),\s*nullable=True", + upgrade, + ) + assert upgrade.count("op.alter_column(") == 2 + assert set(relaxed_columns) == { + ("program_statistics", "entity"), + ("program_statistics", "is_tax"), + } + assert downgrade.count("op.drop_column(") == len(added_columns) + for body in (upgrade, downgrade): + assert "op.drop_table(" not in body + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body + assert "op.drop_column(" not in upgrade diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index ffc60870c..d321048a4 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -208,9 +208,55 @@ def test_report_runs_record_their_modal_call_and_untabled_results() -> None: assert report_runs.c.modal_call_id.nullable assert report_runs.c.modal_call_id.type.length == 255 - for column_name in ("labor_supply_result", "cliff_impact_result"): + for column_name in ( + "labor_supply_result", + "cliff_impact_result", + "spm_provenance", + ): assert report_runs.c[column_name].nullable assert isinstance(report_runs.c[column_name].type, sa.JSON) + assert report_runs.c.data_version.nullable + assert report_runs.c.data_version.type.length == 128 + + +def test_uk_area_impacts_record_disclosure_control() -> None: + for table_name in ("constituency_impacts", "local_authority_impacts"): + table = V2_METADATA.tables[table_name] + assert table.c.disclosure_tier.nullable + assert table.c.disclosure_tier.type.length == 16 + for change in ("average", "relative"): + for bound in ("lower", "upper"): + column = table.c[f"{change}_household_income_change_{bound}"] + assert column.nullable + assert isinstance(column.type, sa.Float) + assert f"ck_{table_name}_disclosure_tier" in _check_names(table) + + +def test_congressional_district_impacts_record_winner_and_loser_shares() -> None: + districts = V2_METADATA.tables["congressional_district_impacts"] + + for column_name in ( + "winner_percentage", + "loser_percentage", + "no_change_percentage", + ): + assert not districts.c[column_name].nullable + assert isinstance(districts.c[column_name].type, sa.Float) + + +def test_detailed_budget_rows_need_neither_an_entity_nor_a_tax_flag() -> None: + """The economy output names each program's totals, not its entity or kind.""" + program_statistics = V2_METADATA.tables["program_statistics"] + + assert program_statistics.c.entity.nullable + assert program_statistics.c.is_tax.nullable + statement = str( + CreateTable(program_statistics).compile(dialect=postgresql.dialect()) + ) + assert ( + "CONSTRAINT uq_program_statistics_run_program_entity UNIQUE NULLS NOT " + "DISTINCT (report_run_id, program_name, entity)" + ) in statement def test_report_runs_name_the_simulation_runs_they_combine() -> None: From 13669e0e9d14f1dc78ded6f6fa1a0e4434f383de Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:37:02 +0400 Subject: [PATCH 04/22] Let the v2 privilege steps skip a role not created yet A pull request can now declare a database role before an operator creates it, by listing the role in ROLES_AWAITING_PROVISIONING. While a listed role is missing, reconcile_v2_privileges.py apply and verify skip it and name it in their output and job summary instead of failing the deploy; once the role exists, they grant and check its privileges. A missing role that is not listed still blocks every grant and fails verification. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 7 ++ .../data/v2/privilege_reconciliation.py | 72 ++++++++++++++++--- policyengine_api/data/v2/privileges.py | 6 ++ .../test_v2_privilege_reconciliation.py | 55 ++++++++++++++ .../unit/v2/test_privilege_reconciliation.py | 64 +++++++++++++++++ 5 files changed, 196 insertions(+), 8 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 786565d15..308ccc65c 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -248,6 +248,13 @@ table is missing. The roles themselves are provisioned outside the repository; `--create-missing-roles` is accepted only for the disposable test database. +A pull request may declare a role before an operator creates it by listing it +in `ROLES_AWAITING_PROVISIONING` in `privileges.py`. While a listed role is +missing, `apply` and `verify` skip it and name it in their output and the job +summary instead of failing; once it exists, the next deploy grants and checks +its privileges. Remove the role from the set once it exists in staging and +production, so that its absence fails the deploy again. + The declarations include write privileges for code paths production doesn't turn on: the dual-write copies into `legacy_policy_mappings` and `legacy_household_mappings`, and the writes behind the native `/v2/policies` diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index d9d0cacec..9e89c075e 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -11,6 +11,11 @@ - ``verify`` checks every declared privilege with ``has_table_privilege`` (rule ``api.v2.deployed-privileges-match``) and fails on any that is missing. +Both skip a missing role the manifest lists in ``ROLES_AWAITING_PROVISIONING``, +and name it in their output, so a pull request can declare a role before an +operator creates it. Once the role exists, ``apply`` grants its privileges and +``verify`` checks them. + ``--create-missing-roles`` creates the managed roles as ``NOLOGIN`` roles. It is refused unless the target is the disposable test database (``V2_ALEMBIC_DISPOSABLE_TEST=1``); in staging and production the roles are @@ -42,6 +47,7 @@ ) from policyengine_api.data.v2.privileges import ( PERSISTENT_ENVIRONMENTS, + ROLES_AWAITING_PROVISIONING, TablePrivilege, declared_privileges, ) @@ -88,12 +94,17 @@ class DefaultPrivilege: @dataclass(frozen=True) class ReconciliationPlan: - """What ``apply`` grants and reports for one database.""" + """What ``apply`` grants and reports for one database. + + ``skipped_roles`` are missing roles awaiting creation by an operator; they + don't block the plan, and nothing is granted to them. + """ grants: tuple[PrivilegeEntry, ...] undeclared: tuple[PrivilegeEntry, ...] missing_roles: tuple[str, ...] missing_tables: tuple[str, ...] + skipped_roles: tuple[str, ...] = () @property def can_apply(self) -> bool: @@ -107,18 +118,22 @@ def plan_reconciliation( *, existing_roles: Collection[str], existing_tables: Collection[str], + roles_awaiting_creation: Collection[str] = (), ) -> ReconciliationPlan: """Compare declared privileges with what the database grants. - ``granted`` holds the managed roles' direct grants, by role then table. + ``granted`` holds the managed roles' direct grants, by role then table. A + missing role in ``roles_awaiting_creation`` is skipped instead of missing. """ - missing_roles = tuple(sorted(set(declared) - set(existing_roles))) + absent_roles = set(declared) - set(existing_roles) + skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) + missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) declared_tables = {table for tables in declared.values() for table in tables} missing_tables = tuple(sorted(declared_tables - set(existing_tables))) grants = sorted( PrivilegeEntry(role, table, privilege.value) for role, tables in declared.items() - if role not in missing_roles + if role not in absent_roles for table, privileges in tables.items() if table not in missing_tables for privilege in privileges @@ -140,6 +155,7 @@ def plan_reconciliation( undeclared=tuple(undeclared), missing_roles=missing_roles, missing_tables=missing_tables, + skipped_roles=skipped_roles, ) @@ -193,6 +209,13 @@ def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: ] +def _skipped_roles_line(roles: Iterable[str]) -> str: + return ( + "Declared roles not created yet, skipped until an operator creates them: " + + ", ".join(f"`{role}`" for role in roles) + ) + + def render_summary( plan: ReconciliationPlan, default_privileges: Sequence[DefaultPrivilege], @@ -202,6 +225,8 @@ def render_summary( ) -> str: """Return the Markdown report for the job summary and the log.""" lines = [f"## v2 table privileges ({deployment_environment})", ""] + if plan.skipped_roles: + lines += [_skipped_roles_line(plan.skipped_roles), ""] if plan.missing_roles: lines += [ "Declared roles missing from the database (nothing was granted): " @@ -363,10 +388,13 @@ def create_missing_roles(connection: Connection, roles: Iterable[str]) -> list[s def apply_declared_privileges( connection: Connection, declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + *, + roles_awaiting_creation: Collection[str] = (), ) -> tuple[ReconciliationPlan, list[DefaultPrivilege]]: """Grant every declared privilege a role lacks; never revoke. - Grants nothing when a declared role or table is missing. + Grants nothing when a declared role or table is missing, unless the missing + role awaits creation, which is skipped. """ roles = list(declared) plan = plan_reconciliation( @@ -374,6 +402,7 @@ def apply_declared_privileges( read_granted_privileges(connection, roles), existing_roles=read_existing_roles(connection, roles), existing_tables=read_public_tables(connection), + roles_awaiting_creation=roles_awaiting_creation, ) if plan.can_apply: for statement in grant_statements( @@ -383,14 +412,33 @@ def apply_declared_privileges( return plan, read_default_privileges(connection, roles) +def skipped_roles( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + roles_awaiting_creation: Collection[str], +) -> tuple[str, ...]: + """Return the declared roles awaiting creation that don't exist yet.""" + existing_roles = read_existing_roles(connection, list(declared)) + return tuple( + sorted((set(declared) & set(roles_awaiting_creation)) - existing_roles) + ) + + def check_declared_privileges( connection: Connection, declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + *, + roles_awaiting_creation: Collection[str] = (), ) -> list[PrivilegeEntry]: """Return every declared privilege a role doesn't effectively hold. - A missing role or table counts as missing every privilege declared for it. + A missing role or table counts as missing every privilege declared for it, + except a missing role awaiting creation, which is skipped. """ + skipped = set(skipped_roles(connection, declared, roles_awaiting_creation)) + declared = { + role: tables for role, tables in declared.items() if role not in skipped + } existing_roles = read_existing_roles(connection, list(declared)) existing_tables = read_public_tables(connection) @@ -457,12 +505,18 @@ def run( msg = "--create-missing-roles is only allowed for the disposable test database" raise V2ConfigurationError(msg) declared = declared_privileges(environment) + awaiting = ROLES_AWAITING_PROVISIONING engine = engine_builder(settings) try: with engine.begin() as connection: qualify_v2_connection(connection, settings) if mode == "verify": - failures = check_declared_privileges(connection, declared) + skipped = skipped_roles(connection, declared, awaiting) + if skipped: + print(_skipped_roles_line(skipped)) + failures = check_declared_privileges( + connection, declared, roles_awaiting_creation=awaiting + ) if failures: rows = "\n".join( f"- `{failure.role}` lacks {failure.privilege} on " @@ -481,7 +535,9 @@ def run( created = create_missing_roles(connection, declared) if created: print("Created roles: " + ", ".join(created)) - plan, defaults = apply_declared_privileges(connection, declared) + plan, defaults = apply_declared_privileges( + connection, declared, roles_awaiting_creation=awaiting + ) _write_summary( values, render_summary( diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 814aa07fa..b7c463f95 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -44,6 +44,12 @@ CATALOG_WRITER_ROLE: Final = "policyengine_v2_catalog_writer" STAGE12_ROLE_PREFIX: Final = "policyengine_stage12_" PERSISTENT_ENVIRONMENTS: Final = frozenset({"staging", "production"}) +# Declared roles an operator hasn't created in staging and production yet. While +# one is missing, the deploy's ``apply`` and ``verify`` skip it and name it, +# instead of failing; once it exists they grant and check its privileges. +# Remove a role from this set when it exists in both environments, so that its +# absence fails the deploy again. +ROLES_AWAITING_PROVISIONING: Final[frozenset[str]] = frozenset() ENVIRONMENT_PATTERN: Final = re.compile(r"^[a-z][a-z0-9_]{0,30}$") diff --git a/tests/integration/test_v2_privilege_reconciliation.py b/tests/integration/test_v2_privilege_reconciliation.py index f720c3fb9..3fa2e3376 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -15,6 +15,7 @@ from sqlalchemy import create_engine, text from sqlalchemy.pool import NullPool +from policyengine_api.data.v2 import privilege_reconciliation from policyengine_api.data.v2.migration_target import V2_ALEMBIC_DISPOSABLE_TEST from policyengine_api.data.v2.privilege_reconciliation import ( DEPLOYMENT_ENVIRONMENT, @@ -137,3 +138,57 @@ def test_a_missing_role_blocks_every_grant_and_fails_verification( PrivilegeEntry(role, "regions", "SELECT"), ] ) + + +def test_a_role_awaiting_creation_is_skipped_until_it_exists( + superuser_engine, temporary_role +) -> None: + role = temporary_role() + absent_role = f"policyengine_privilege_test_absent_{uuid4().hex[:12]}" + declared = {role: {"regions": READ}, absent_role: {"regions": READ}} + with superuser_engine.begin() as connection: + plan, _defaults = apply_declared_privileges( + connection, declared, roles_awaiting_creation={absent_role} + ) + missing = check_declared_privileges( + connection, declared, roles_awaiting_creation={absent_role} + ) + + assert plan.can_apply + assert plan.skipped_roles == (absent_role,) + assert plan.grants == (PrivilegeEntry(role, "regions", "SELECT"),) + assert missing == [] + assert _has_privilege(superuser_engine, role, "regions", "SELECT") + + +def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( + disposable_v2_database_url, temporary_role, monkeypatch, capsys +) -> None: + """``apply`` and ``verify`` succeed while a role awaiting creation is absent.""" + role = temporary_role() + absent_role = f"policyengine_privilege_test_absent_{uuid4().hex[:12]}" + monkeypatch.setattr( + privilege_reconciliation, + "declared_privileges", + lambda _environment: {role: {"regions": READ}, absent_role: {"regions": READ}}, + ) + monkeypatch.setattr( + privilege_reconciliation, + "ROLES_AWAITING_PROVISIONING", + frozenset({absent_role}), + ) + environ = { + V2_MIGRATION_DATABASE_URL: disposable_v2_database_url, + V2_ALEMBIC_DISPOSABLE_TEST: "1", + DEPLOYMENT_ENVIRONMENT: "test", + } + + apply_status = run("apply", environ=environ) + apply_report = capsys.readouterr().out + verify_status = run("verify", environ=environ) + verify_output = capsys.readouterr().out + + assert (apply_status, verify_status) == (0, 0) + assert f"`{absent_role}`" in apply_report + assert absent_role in verify_output + assert "Every declared v2 table privilege is held" in verify_output diff --git a/tests/unit/v2/test_privilege_reconciliation.py b/tests/unit/v2/test_privilege_reconciliation.py index 73ca30e19..4c829f037 100644 --- a/tests/unit/v2/test_privilege_reconciliation.py +++ b/tests/unit/v2/test_privilege_reconciliation.py @@ -120,6 +120,70 @@ def test_missing_roles_and_tables_are_recorded_and_block_the_plan() -> None: ) +def test_a_missing_role_awaiting_creation_is_skipped_without_blocking() -> None: + plan = plan_reconciliation( + _declared(), + granted={}, + existing_roles={RUNTIME}, + existing_tables={"users", "legacy_policy_mappings", "reports"}, + roles_awaiting_creation={STAGE12}, + ) + + assert plan.can_apply + assert plan.missing_roles == () + assert plan.skipped_roles == (STAGE12,) + assert {entry.role for entry in plan.grants} == {RUNTIME} + + +def test_a_role_awaiting_creation_gets_its_privileges_once_it_exists() -> None: + declared = {RUNTIME: {"users": READ}, STAGE12: {"users": READ}} + + plan = plan_reconciliation( + declared, + granted={}, + existing_roles={RUNTIME, STAGE12}, + existing_tables={"users"}, + roles_awaiting_creation={STAGE12}, + ) + + assert plan.skipped_roles == () + assert plan.grants == ( + PrivilegeEntry(STAGE12, "users", "SELECT"), + PrivilegeEntry(RUNTIME, "users", "SELECT"), + ) + + +def test_a_missing_role_not_awaiting_creation_still_blocks_the_plan() -> None: + plan = plan_reconciliation( + _declared(), + granted={}, + existing_roles={RUNTIME}, + existing_tables={"users", "legacy_policy_mappings", "reports"}, + roles_awaiting_creation={"policyengine_v2_some_other_role"}, + ) + + assert plan.missing_roles == (STAGE12,) + assert plan.skipped_roles == () + assert not plan.can_apply + + +def test_summary_names_the_roles_skipped_until_an_operator_creates_them() -> None: + plan = ReconciliationPlan( + grants=(), + undeclared=(), + missing_roles=(), + missing_tables=(), + skipped_roles=(STAGE12,), + ) + + summary = render_summary(plan, (), deployment_environment="staging", applied=True) + + assert ( + "Declared roles not created yet, skipped until an operator creates them: " + f"`{STAGE12}`" + ) in summary + + def test_grant_statements_group_privileges_per_role_and_table() -> None: plan = ReconciliationPlan( grants=( From b874e2ff14a69853a624db1c5094e6aaff0159d9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:42:02 +0400 Subject: [PATCH 05/22] Declare the v2 table privileges of the simulation API's roles The privilege manifest declares two roles an operator will create, both awaiting provisioning so the deploy skips them until they exist: - policyengine_v2_simulation_service, the simulation API on Cloud Run, creates simulations, reports, their runs, report run members and household jobs (SELECT, INSERT), and reads the policy, household and catalog rows a run needs; - policyengine_v2_modal_worker, its Modal workers, updates simulation runs, report runs and household jobs, inserts results, and reads the simulation definitions and the rows a run needs. The Public API's runtime role reads the simulation tables and results, and writes none of them. New integration tests run these writes as each role, and CI runs them and the simulation runs migration tests. Co-Authored-By: Claude Opus 5.5 --- .github/scripts/test_alembic_v2_lifecycle.sh | 3 +- .github/workflows/v2-integration-check.yml | 4 +- docs/engineering/skills/alembic-migrations.md | 5 + policyengine_api/data/v2/privileges.py | 111 ++++++--- tests/integration/conftest.py | 6 + .../test_v2_simulation_role_privileges.py | 228 ++++++++++++++++++ tests/unit/test_alembic_workflows.py | 2 + tests/unit/v2/test_privilege_manifest.py | 98 +++++++- 8 files changed, 426 insertions(+), 31 deletions(-) create mode 100644 tests/integration/test_v2_simulation_role_privileges.py diff --git a/.github/scripts/test_alembic_v2_lifecycle.sh b/.github/scripts/test_alembic_v2_lifecycle.sh index 0ed37eba0..091ab8326 100755 --- a/.github/scripts/test_alembic_v2_lifecycle.sh +++ b/.github/scripts/test_alembic_v2_lifecycle.sh @@ -5,4 +5,5 @@ set -euo pipefail uv run --frozen --extra dev pytest -q \ tests/unit/v2/test_alembic_v2.py \ tests/integration/test_alembic_v2_lifecycle.py \ - tests/integration/test_alembic_v2_catalog_tables.py + tests/integration/test_alembic_v2_catalog_tables.py \ + tests/integration/test_alembic_v2_simulation_runs.py diff --git a/.github/workflows/v2-integration-check.yml b/.github/workflows/v2-integration-check.yml index 14da930ea..a0b8a698c 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.github/workflows/v2-integration-check.yml @@ -82,8 +82,8 @@ jobs: run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_policy_persistence.py tests/integration/test_v1_policy_dual_write.py tests/integration/test_mysql_policy_dual_write.py tests/integration/test_v2_user_policy_mirroring.py tests/integration/test_v1_user_policy_dual_write.py - name: Test signed-in identity provisioning and user associations run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_user_identity_persistence.py - - name: Test comparison-run persistence and privilege reconciliation as the managed roles - run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_stage12_comparison_persistence.py tests/integration/test_v2_privilege_reconciliation.py + - name: Test comparison-run persistence, the simulation API's roles and privilege reconciliation as the managed roles + run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_stage12_comparison_persistence.py tests/integration/test_v2_simulation_role_privileges.py tests/integration/test_v2_privilege_reconciliation.py - name: Test v2 household persistence and immediate v1 copying run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_household_persistence.py tests/integration/test_v1_household_dual_write.py tests/integration/test_mysql_household_spm_dual_write.py - name: Test the row counts v1 ORM updates report on MySQL diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 308ccc65c..92c80400e 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -225,6 +225,8 @@ and a revision runs once per database, so it can't repair drift. | `policyengine_v2_runtime` | The Public API on Cloud Run | | `policyengine_v2_catalog_writer` | `scripts/publish_v2_metadata_catalog.py` | | `policyengine_stage12_` | The simulation API and its Modal worker, on the temporary comparison-run tables | +| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs and household jobs, and reads the policy, household and catalog rows a run needs | +| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs | Every v2 table has an entry for every role, so "no access" is a decision, not an omission. The migration role owns every table and isn't listed. Three rules @@ -237,6 +239,9 @@ check the declarations: `policyengine_v2_runtime`, catalog publication as `policyengine_v2_catalog_writer`, and the comparison-run services as `policyengine_stage12_test`, holding only their declared privileges. + `tests/integration/test_v2_simulation_role_privileges.py` runs the + simulation API's writes as its two roles; their code lives in the simulation + API's repository. - `api.v2.deployed-privileges-match`: every deploy runs `scripts/reconcile_v2_privileges.py apply` right after the schema upgrade, then `verify`, before the catalog is seeded. diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index b7c463f95..7ef0bbc6f 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -24,6 +24,15 @@ (``scripts/publish_v2_metadata_catalog.py``). Its privileges follow what publication writes and reads; its code isn't organized as connector modules, so the deploy check, not the connector analysis, covers it. +- ``policyengine_v2_simulation_service``: the simulation API on Cloud Run. It + creates simulations, reports, their runs and household jobs, and reads the + policy, household and catalog rows a run needs. +- ``policyengine_v2_modal_worker``: the simulation API's Modal workers. They + write the status of the runs and household jobs they were given, and their + results, and read the rows a run needs. Both simulation API roles' code lives + in that repository, so the deploy check covers them here. + +The Public API reads the simulation tables it serves, and writes none of them. The migration role (``policyengine_v2_migrator``) owns every table, and Supabase's ``postgres`` role administers the project; neither is managed here. @@ -36,12 +45,14 @@ from __future__ import annotations import re -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import StrEnum from typing import Final RUNTIME_ROLE: Final = "policyengine_v2_runtime" CATALOG_WRITER_ROLE: Final = "policyengine_v2_catalog_writer" +SIMULATION_SERVICE_ROLE: Final = "policyengine_v2_simulation_service" +MODAL_WORKER_ROLE: Final = "policyengine_v2_modal_worker" STAGE12_ROLE_PREFIX: Final = "policyengine_stage12_" PERSISTENT_ENVIRONMENTS: Final = frozenset({"staging", "production"}) # Declared roles an operator hasn't created in staging and production yet. While @@ -49,7 +60,9 @@ # instead of failing; once it exists they grant and check its privileges. # Remove a role from this set when it exists in both environments, so that its # absence fails the deploy again. -ROLES_AWAITING_PROVISIONING: Final[frozenset[str]] = frozenset() +ROLES_AWAITING_PROVISIONING: Final[frozenset[str]] = frozenset( + {SIMULATION_SERVICE_ROLE, MODAL_WORKER_ROLE} +) ENVIRONMENT_PATTERN: Final = re.compile(r"^[a-z][a-z0-9_]{0,30}$") @@ -65,6 +78,7 @@ class TablePrivilege(StrEnum): NO_ACCESS: Final[frozenset[TablePrivilege]] = frozenset() READ: Final = frozenset({TablePrivilege.SELECT}) READ_INSERT: Final = frozenset({TablePrivilege.SELECT, TablePrivilege.INSERT}) +READ_UPDATE: Final = frozenset({TablePrivilege.SELECT, TablePrivilege.UPDATE}) READ_INSERT_UPDATE: Final = frozenset( {TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE} ) @@ -85,6 +99,8 @@ class TableAccess: runtime: frozenset[TablePrivilege] catalog_writer: frozenset[TablePrivilege] stage12: frozenset[TablePrivilege] + simulation_service: frozenset[TablePrivilege] + modal_worker: frozenset[TablePrivilege] def _access( @@ -92,35 +108,62 @@ def _access( runtime: frozenset[TablePrivilege] = NO_ACCESS, catalog_writer: frozenset[TablePrivilege] = NO_ACCESS, stage12: frozenset[TablePrivilege] = NO_ACCESS, + simulation_service: frozenset[TablePrivilege] = NO_ACCESS, + modal_worker: frozenset[TablePrivilege] = NO_ACCESS, ) -> TableAccess: - return TableAccess(runtime=runtime, catalog_writer=catalog_writer, stage12=stage12) + return TableAccess( + runtime=runtime, + catalog_writer=catalog_writer, + stage12=stage12, + simulation_service=simulation_service, + modal_worker=modal_worker, + ) + + +# The simulation API creates these rows; its Modal workers read them, and write +# the status of the runs and household jobs among them. +_SIMULATION_DEFINITION = _access( + runtime=READ, simulation_service=READ_INSERT, modal_worker=READ +) +_RUN_STATUS = _access( + runtime=READ, simulation_service=READ_INSERT, modal_worker=READ_UPDATE +) +# The Modal workers write results, keyed by run, once each. +_RESULT = _access(runtime=READ, modal_worker=READ_INSERT) +# What a run needs from the Public API's policies, households and catalog. +_RUN_INPUT_READS: Final = {"simulation_service": READ, "modal_worker": READ} 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. + # Custom totals, unused until they return. "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_run_simulation_runs": _access(), - "report_runs": _access(catalog_writer=READ), - "reports": _access(catalog_writer=READ), - "simulation_runs": _access(), - "simulations": _access(catalog_writer=READ), + # Simulations, reports, their runs and household jobs. The catalog writer + # checks that publication leaves the rows of three of them untouched. + "household_jobs": _RUN_STATUS, + "report_run_simulation_runs": _SIMULATION_DEFINITION, + "report_runs": replace(_RUN_STATUS, catalog_writer=READ), + "reports": replace(_SIMULATION_DEFINITION, catalog_writer=READ), + "simulation_runs": _RUN_STATUS, + "simulations": replace(_SIMULATION_DEFINITION, catalog_writer=READ), + # Results of simulation runs and report runs. + "budget_summary": _RESULT, + "congressional_district_impacts": _RESULT, + "constituency_impacts": _RESULT, + "decile_impacts": _RESULT, + "inequality": _RESULT, + "intra_decile_impacts": _RESULT, + "local_authority_impacts": _RESULT, + "poverty": _RESULT, + "program_statistics": _RESULT, # 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), + "datasets": _access( + runtime=READ, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS + ), "parameter_nodes": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), # Written by catalog publication for each version it publishes, from the # installed packages' catalog; never rewritten. @@ -128,20 +171,28 @@ def _access( # Written with each catalog version that publication publishes. "parameter_tree_children": _access(runtime=READ, catalog_writer=READ_INSERT), # The Public API also inserts a v2 policy's own parameter values. - "parameter_values": _access(runtime=READ_INSERT, catalog_writer=READ_INSERT_UPDATE), - "parameters": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), - "regions": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), + "parameter_values": _access( + runtime=READ_INSERT, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS + ), + "parameters": _access( + runtime=READ, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS + ), + "regions": _access( + runtime=READ, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS + ), "tax_benefit_model_versions": _access( - runtime=READ, catalog_writer=READ_INSERT_UPDATE + runtime=READ, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS + ), + "tax_benefit_models": _access( + runtime=READ, catalog_writer=READ_INSERT_UPDATE, **_RUN_INPUT_READS ), - "tax_benefit_models": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), "variables": _access(runtime=READ, catalog_writer=READ_INSERT_UPDATE), # v2 policies and households, created by the native routes and by Stage 10 # dual-write. A copy locks its mapping row (FOR UPDATE), which needs UPDATE. - "households": _access(runtime=READ_INSERT), + "households": _access(runtime=READ_INSERT, **_RUN_INPUT_READS), "legacy_household_mappings": _access(runtime=READ_INSERT_UPDATE), "legacy_policy_mappings": _access(runtime=READ_INSERT_UPDATE), - "policies": _access(runtime=READ_INSERT), + "policies": _access(runtime=READ_INSERT, **_RUN_INPUT_READS), # Stage 10 saved-policy mirror. "legacy_user_mappings": _access(runtime=READ_WRITE), "legacy_user_policy_mappings": _access(runtime=READ_WRITE), @@ -191,4 +242,10 @@ def declared_privileges( table: access.catalog_writer for table, access in TABLE_ACCESS.items() }, stage12: {table: access.stage12 for table, access in TABLE_ACCESS.items()}, + SIMULATION_SERVICE_ROLE: { + table: access.simulation_service for table, access in TABLE_ACCESS.items() + }, + MODAL_WORKER_ROLE: { + table: access.modal_worker for table, access in TABLE_ACCESS.items() + }, } diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index b90efebd2..ecdbec122 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -69,6 +69,8 @@ class ManagedV2Roles: deployment_environment: str runtime: str stage12: str + simulation_service: str + modal_worker: str @pytest.fixture(scope="session") @@ -87,7 +89,9 @@ def managed_v2_roles(disposable_v2_database_url: str) -> ManagedV2Roles: create_missing_roles, ) from policyengine_api.data.v2.privileges import ( + MODAL_WORKER_ROLE, RUNTIME_ROLE, + SIMULATION_SERVICE_ROLE, declared_privileges, stage12_role, ) @@ -112,6 +116,8 @@ def managed_v2_roles(disposable_v2_database_url: str) -> ManagedV2Roles: deployment_environment=deployment_environment, runtime=RUNTIME_ROLE, stage12=stage12_role(deployment_environment), + simulation_service=SIMULATION_SERVICE_ROLE, + modal_worker=MODAL_WORKER_ROLE, ) diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py new file mode 100644 index 000000000..bffc9679e --- /dev/null +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -0,0 +1,228 @@ +"""Run the simulation tables' writes as the roles that do them in production. + +The simulation API creates simulations, reports, their runs and household +jobs; its Modal workers write their runs' status and results; the Public API +reads them. Each role holds only its declared privileges +(``policyengine_api/data/v2/privileges.py``), so a statement outside them +fails here as it would in staging and production. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, text +from sqlalchemy.exc import ProgrammingError +from sqlalchemy.pool import NullPool + +from policyengine_api.constants import REPO + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} + + +@pytest.fixture +def connection(disposable_v2_database_url, managed_v2_roles) -> Iterator[Connection]: + """Yield a connection at head whose writes are rolled back after the test. + + It connects as the disposable database's superuser; ``_act_as`` switches it + to one managed role for the rest of the transaction. + """ + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + try: + with engine.connect() as database_connection: + transaction = database_connection.begin() + try: + yield database_connection + finally: + transaction.rollback() + finally: + engine.dispose() + + +def _act_as(connection: Connection, role: str) -> None: + quoted_role = connection.dialect.identifier_preparer.quote(role) + connection.execute(text(f"SET LOCAL ROLE {quoted_role}")) + + +def _insert(connection: Connection, table: str, **values: object) -> None: + columns = ", ".join(values) + placeholders = ", ".join(f":{column}" for column in values) + connection.execute( + text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"), values + ) + + +def _refused(connection: Connection, statement: str, **values: object) -> None: + """Expect a permission error, then return the connection to a usable state.""" + savepoint = connection.begin_nested() + with pytest.raises(ProgrammingError, match="permission denied"): + connection.execute(text(statement), values) + savepoint.rollback() + + +def _create_simulation_rows(connection: Connection) -> dict[str, object]: + """Create what the simulation API creates for one economy report run.""" + ids = {name: uuid4() for name in ("simulation", "run", "report", "report_run")} + _insert( + connection, + "simulations", + id=ids["simulation"], + simulation_type="economy", + dataset_id=uuid4(), + tax_benefit_model_version_id=uuid4(), + year=2026, + ) + _insert( + connection, + "simulation_runs", + id=ids["run"], + simulation_id=ids["simulation"], + tax_benefit_model_version_id=uuid4(), + status="pending", + **RUN_VERSIONS, + ) + _insert( + connection, + "reports", + id=ids["report"], + label="Report", + country="us", + tax_benefit_model_id=uuid4(), + inputs="{}", + ) + _insert( + connection, + "report_runs", + id=ids["report_run"], + report_id=ids["report"], + status="pending", + trigger="initial", + **RUN_VERSIONS, + ) + _insert( + connection, + "report_run_simulation_runs", + id=uuid4(), + report_run_id=ids["report_run"], + simulation_run_id=ids["run"], + role="baseline", + ) + _insert( + connection, + "household_jobs", + id=uuid4(), + country="us", + household_id=uuid4(), + year=2026, + request_data="{}", + tax_benefit_model_version_id=uuid4(), + status="pending", + **RUN_VERSIONS, + ) + return ids + + +def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( + connection: Connection, managed_v2_roles +) -> None: + _act_as(connection, managed_v2_roles.simulation_service) + + ids = _create_simulation_rows(connection) + run_count = connection.execute( + text("SELECT count(*) FROM simulation_runs WHERE id = :id"), {"id": ids["run"]} + ).scalar_one() + connection.execute(text("SELECT count(*) FROM policies")).scalar_one() + + assert run_count == 1 + _refused( + connection, + "UPDATE simulation_runs SET status = 'running' WHERE id = :id", + id=ids["run"], + ) + _refused( + connection, + "INSERT INTO poverty (id, simulation_run_id, poverty_type, entity) " + "VALUES (:id, :run_id, 'spm', 'person')", + id=uuid4(), + run_id=ids["run"], + ) + _refused(connection, "DELETE FROM simulations WHERE id = :id", id=ids["simulation"]) + _refused(connection, "UPDATE policies SET content_hash = content_hash") + + +def test_modal_workers_write_their_run_status_and_results_only( + connection: Connection, managed_v2_roles +) -> None: + _act_as(connection, managed_v2_roles.simulation_service) + ids = _create_simulation_rows(connection) + _act_as(connection, managed_v2_roles.modal_worker) + + updated = connection.execute( + text( + "UPDATE simulation_runs SET status = 'succeeded', completed_at = now() " + "WHERE id = :id AND status = 'pending' RETURNING id" + ), + {"id": ids["run"]}, + ).scalar_one() + _insert( + connection, + "poverty", + id=uuid4(), + simulation_run_id=ids["run"], + poverty_type="spm", + entity="person", + rate=0.1, + ) + connection.execute( + text("UPDATE report_runs SET status = 'running' WHERE id = :id"), + {"id": ids["report_run"]}, + ) + connection.execute(text("SELECT count(*) FROM parameter_values")).scalar_one() + + assert updated == ids["run"] + _refused( + connection, + "INSERT INTO simulations (id, simulation_type, dataset_id, " + "tax_benefit_model_version_id) " + "VALUES (:id, 'economy', :dataset_id, :version_id)", + id=uuid4(), + dataset_id=uuid4(), + version_id=uuid4(), + ) + _refused( + connection, "DELETE FROM poverty WHERE simulation_run_id = :id", id=ids["run"] + ) + _refused( + connection, + "UPDATE poverty SET rate = 0 WHERE simulation_run_id = :id", + id=ids["run"], + ) + _refused(connection, "UPDATE households SET content_hash = content_hash") + + +def test_the_public_api_reads_runs_and_results_but_writes_neither( + connection: Connection, managed_v2_roles +) -> None: + _act_as(connection, managed_v2_roles.runtime) + + for table in ("simulation_runs", "report_run_simulation_runs", "poverty"): + connection.execute(text(f"SELECT count(*) FROM {table}")).scalar_one() + _refused( + connection, + "INSERT INTO report_run_simulation_runs " + "(id, report_run_id, simulation_run_id, role) " + "VALUES (:id, :report_run_id, :run_id, 'baseline')", + id=uuid4(), + report_run_id=uuid4(), + run_id=uuid4(), + ) diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 86310e018..d798d2036 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -204,6 +204,7 @@ def test_reusable_v2_alembic_check_uses_only_disposable_postgres(): assert "test_alembic_v2.py" in lifecycle_script assert "test_alembic_v2_lifecycle.py" in lifecycle_script assert "test_alembic_v2_catalog_tables.py" in lifecycle_script + assert "test_alembic_v2_simulation_runs.py" in lifecycle_script assert "redis:7.2-alpine" not in workflow assert "RUNTIME_CACHE_TEST_URL" not in workflow assert "test_v2_catalog_installed.py" not in workflow @@ -241,6 +242,7 @@ def test_reusable_v2_integration_check_uses_disposable_databases_and_coverage(): assert "test_v2_user_identity_persistence.py" in workflow assert "test_v2_stage12_comparison_persistence.py" in workflow assert "test_v2_privilege_reconciliation.py" in workflow + assert "test_v2_simulation_role_privileges.py" in workflow assert "test_v2_household_persistence.py" in workflow assert "test_v1_household_dual_write.py" in workflow assert "test_mysql_household_spm_dual_write.py" in workflow diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 7a3a76b44..5c635e22e 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -18,9 +18,15 @@ ) from policyengine_api.data.v2.privileges import ( CATALOG_WRITER_ROLE, + MODAL_WORKER_ROLE, NO_ACCESS, + READ, + READ_INSERT, + READ_UPDATE, READ_WRITE, + ROLES_AWAITING_PROVISIONING, RUNTIME_ROLE, + SIMULATION_SERVICE_ROLE, TABLE_ACCESS, InvalidDeploymentEnvironmentError, TablePrivilege, @@ -33,6 +39,47 @@ PROOF_FIXTURES = PROJECT_ROOT / "rules/fixtures/api.v2.runtime-privileges-declared" STAGE12_PACKAGES = frozenset({"comparison_runs"}) TEST_ENVIRONMENT = "staging" +# The rows the simulation service creates, and the run tables among them. +SIMULATION_SERVICE_CREATES = frozenset( + { + "simulations", + "reports", + "simulation_runs", + "report_runs", + "report_run_simulation_runs", + "household_jobs", + } +) +RUN_STATUS_TABLES = frozenset({"simulation_runs", "report_runs", "household_jobs"}) +# The policy, household and catalog rows a run needs. +RUN_INPUT_TABLES = frozenset( + { + "policies", + "households", + "tax_benefit_models", + "tax_benefit_model_versions", + "regions", + "datasets", + "parameters", + "parameter_values", + } +) +RESULT_TABLES = frozenset( + { + "budget_summary", + "congressional_district_impacts", + "constituency_impacts", + "decile_impacts", + "inequality", + "intra_decile_impacts", + "local_authority_impacts", + "poverty", + "program_statistics", + } +) +WRITE_PRIVILEGES = frozenset( + {TablePrivilege.INSERT, TablePrivilege.UPDATE, TablePrivilege.DELETE} +) def _table_by_model() -> dict[str, str]: @@ -118,14 +165,63 @@ def test_stage12_role_name_follows_the_deployment_environment() -> None: stage12_role(invalid) -def test_declared_privileges_name_the_three_managed_roles() -> None: +def test_declared_privileges_name_the_five_managed_roles() -> None: assert set(declared_privileges("production")) == { RUNTIME_ROLE, CATALOG_WRITER_ROLE, "policyengine_stage12_production", + "policyengine_v2_simulation_service", + "policyengine_v2_modal_worker", } +def test_only_the_new_simulation_roles_await_provisioning() -> None: + declared_roles = set(declared_privileges(TEST_ENVIRONMENT)) + + assert {SIMULATION_SERVICE_ROLE, MODAL_WORKER_ROLE} == ROLES_AWAITING_PROVISIONING + assert declared_roles >= ROLES_AWAITING_PROVISIONING + + +def test_simulation_service_creates_simulation_rows_and_reads_run_inputs() -> None: + declared = declared_privileges(TEST_ENVIRONMENT)[SIMULATION_SERVICE_ROLE] + + expected = dict.fromkeys(TABLE_ACCESS, NO_ACCESS) + expected |= dict.fromkeys(SIMULATION_SERVICE_CREATES, READ_INSERT) + expected |= dict.fromkeys(RUN_INPUT_TABLES, READ) + assert declared == expected + + +def test_modal_workers_write_only_run_status_and_results() -> None: + declared = declared_privileges(TEST_ENVIRONMENT)[MODAL_WORKER_ROLE] + + expected = dict.fromkeys(TABLE_ACCESS, NO_ACCESS) + expected |= dict.fromkeys(SIMULATION_SERVICE_CREATES, READ) + expected |= dict.fromkeys(RUN_STATUS_TABLES, READ_UPDATE) + expected |= dict.fromkeys(RESULT_TABLES, READ_INSERT) + expected |= dict.fromkeys(RUN_INPUT_TABLES, READ) + assert declared == expected + + +def test_public_api_reads_simulations_runs_and_results_without_writing_them() -> None: + declared = declared_privileges(TEST_ENVIRONMENT)[RUNTIME_ROLE] + + for table in SIMULATION_SERVICE_CREATES | RESULT_TABLES: + assert declared[table] == READ, table + for table in ("aggregates", "change_aggregates"): + assert declared[table] == NO_ACCESS, table + + +def test_no_new_role_writes_the_public_api_tables() -> None: + simulation_tables = SIMULATION_SERVICE_CREATES | RESULT_TABLES + for role in (SIMULATION_SERVICE_ROLE, MODAL_WORKER_ROLE): + writes = { + table + for table, privileges in declared_privileges(TEST_ENVIRONMENT)[role].items() + if privileges & WRITE_PRIVILEGES + } + assert writes <= simulation_tables, role + + def test_rule_proof_runtime_privileges_declared() -> None: """Proof for rule api.v2.runtime-privileges-declared (rules/registry.toml). From 205b2a6cc0c5fdea9e37a79008c1f85394f4b813 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:09:55 +0400 Subject: [PATCH 06/22] Tell apart v2 result rows of each run pair, decile grouping and group A budget-window report run compares one pair of simulation runs per year, so each comparison result table's unique key now includes the pair after the report run. decile_impacts.decile_variable names the variable that groups households (UK wealth deciles share the income variable and entity), and poverty.filter_group names the group one filter variable selects (age: child, adult, senior); both join their table's unique key. Revision 2e3de7c572ed is autogenerated with no corrections. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 20 ++ ...ell_apart_result_rows_of_each_run_pair_.py | 248 ++++++++++++++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/reports.py | 55 +++- .../test_alembic_v2_simulation_runs.py | 47 ++++ tests/unit/v2/test_alembic_v2.py | 66 +++++ tests/unit/v2/test_simulation_run_models.py | 44 +++- 7 files changed, 462 insertions(+), 20 deletions(-) create mode 100644 migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 92c80400e..d6eb1fbe3 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -544,3 +544,23 @@ post-generation corrections. The three share columns are required, so the upgrade needs `congressional_district_impacts` empty; the downgrade needs every `program_statistics` row to have an entity and a tax flag. +### API v2 result row identity + +Revision `2e3de7c572ed` (parent `7f43d2c7c3c6`) was autogenerated after the +models made each result row tell apart what it measures: + +- the comparison result tables' unique keys include the pair of simulation + runs (`baseline_simulation_run_id`, `reform_simulation_run_id`) after + `report_run_id`, because a budget-window report run compares one pair per + year; +- `decile_impacts.decile_variable` names the variable that groups households + when it is not the income variable's own deciles (for example + `household_wealth_decile`), and joins the table's unique key; +- `poverty.filter_group` names the group a filter selects (`child` for + `age`), and joins the table's unique key. + +Missing values in those keys count as equal (`NULLS NOT DISTINCT`). It has no +post-generation corrections and needs no table empty; its downgrade fails +while a report run holds comparison rows of more than one pair, or a run holds +rows that differ only in a removed column. + diff --git a/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py b/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py new file mode 100644 index 000000000..68bdef772 --- /dev/null +++ b/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py @@ -0,0 +1,248 @@ +"""tell apart result rows of each run pair and group + +Revision ID: 2e3de7c572ed +Revises: 7f43d2c7c3c6 +Create Date: 2026-11-04 01:07:49.171506 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +No post-generation corrections. The upgrade adds two nullable columns and +widens unique keys, so it needs no table empty. The downgrade narrows them +again, so it fails while a report run holds comparison rows of more than one +pair of simulation runs, or a run holds rows that differ only in a removed +column. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "2e3de7c572ed" +down_revision: Union[str, None] = "7f43d2c7c3c6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint( + op.f("uq_budget_summary_run_variable_entity"), "budget_summary", type_="unique" + ) + op.create_unique_constraint( + "uq_budget_summary_run_variable_entity", + "budget_summary", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "variable_name", + "entity", + ], + ) + op.drop_constraint( + op.f("uq_congressional_district_impacts_run_geoid"), + "congressional_district_impacts", + type_="unique", + ) + op.create_unique_constraint( + "uq_congressional_district_impacts_run_geoid", + "congressional_district_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "district_geoid", + ], + ) + op.drop_constraint( + op.f("uq_constituency_impacts_run_code"), "constituency_impacts", type_="unique" + ) + op.create_unique_constraint( + "uq_constituency_impacts_run_code", + "constituency_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "constituency_code", + ], + ) + op.add_column( + "decile_impacts", + sa.Column( + "decile_variable", + sqlmodel.sql.sqltypes.AutoString(length=512), + nullable=True, + ), + ) + op.drop_constraint( + op.f("uq_decile_impacts_run_measure"), "decile_impacts", type_="unique" + ) + op.create_unique_constraint( + "uq_decile_impacts_run_measure", + "decile_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "income_variable", + "decile_variable", + "entity", + "decile", + "quantiles", + ], + postgresql_nulls_not_distinct=True, + ) + op.drop_constraint( + op.f("uq_intra_decile_impacts_run_type_decile"), + "intra_decile_impacts", + type_="unique", + ) + op.create_unique_constraint( + "uq_intra_decile_impacts_run_type_decile", + "intra_decile_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "decile_type", + "decile", + ], + ) + op.drop_constraint( + op.f("uq_local_authority_impacts_run_code"), + "local_authority_impacts", + type_="unique", + ) + op.create_unique_constraint( + "uq_local_authority_impacts_run_code", + "local_authority_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "local_authority_code", + ], + ) + op.add_column( + "poverty", + sa.Column( + "filter_group", sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True + ), + ) + op.drop_constraint( + op.f("uq_poverty_simulation_run_measure"), "poverty", type_="unique" + ) + op.create_unique_constraint( + "uq_poverty_simulation_run_measure", + "poverty", + [ + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + ], + postgresql_nulls_not_distinct=True, + ) + op.drop_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + type_="unique", + ) + op.create_unique_constraint( + "uq_program_statistics_run_program_entity", + "program_statistics", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "program_name", + "entity", + ], + postgresql_nulls_not_distinct=True, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint( + "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + ["report_run_id", "program_name", "entity"], + postgresql_nulls_not_distinct=True, + ) + op.drop_constraint("uq_poverty_simulation_run_measure", "poverty", type_="unique") + op.create_unique_constraint( + op.f("uq_poverty_simulation_run_measure"), + "poverty", + ["simulation_run_id", "poverty_type", "entity", "filter_variable"], + postgresql_nulls_not_distinct=True, + ) + op.drop_column("poverty", "filter_group") + op.drop_constraint( + "uq_local_authority_impacts_run_code", "local_authority_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_local_authority_impacts_run_code"), + "local_authority_impacts", + ["report_run_id", "local_authority_code"], + postgresql_nulls_not_distinct=False, + ) + op.drop_constraint( + "uq_intra_decile_impacts_run_type_decile", + "intra_decile_impacts", + type_="unique", + ) + op.create_unique_constraint( + op.f("uq_intra_decile_impacts_run_type_decile"), + "intra_decile_impacts", + ["report_run_id", "decile_type", "decile"], + postgresql_nulls_not_distinct=False, + ) + op.drop_constraint( + "uq_decile_impacts_run_measure", "decile_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_decile_impacts_run_measure"), + "decile_impacts", + ["report_run_id", "income_variable", "entity", "decile", "quantiles"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("decile_impacts", "decile_variable") + op.drop_constraint( + "uq_constituency_impacts_run_code", "constituency_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_constituency_impacts_run_code"), + "constituency_impacts", + ["report_run_id", "constituency_code"], + postgresql_nulls_not_distinct=False, + ) + op.drop_constraint( + "uq_congressional_district_impacts_run_geoid", + "congressional_district_impacts", + type_="unique", + ) + op.create_unique_constraint( + op.f("uq_congressional_district_impacts_run_geoid"), + "congressional_district_impacts", + ["report_run_id", "district_geoid"], + postgresql_nulls_not_distinct=False, + ) + op.drop_constraint( + "uq_budget_summary_run_variable_entity", "budget_summary", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_budget_summary_run_variable_entity"), + "budget_summary", + ["report_run_id", "variable_name", "entity"], + postgresql_nulls_not_distinct=False, + ) + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index b94e05ae3..03f0fb38e 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "7f43d2c7c3c6" +V2_ALEMBIC_HEAD_REVISION = "2e3de7c572ed" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 238901192..9c2d06ad1 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,8 +1,10 @@ """Canonical SQLModel tables for stable reports, runs, and run outputs. Results that compare a baseline with a reform belong to a report run and name -the two simulation runs they compare; poverty and inequality, which describe -one simulation, belong to a simulation run. +the two simulation runs they compare; each such row is unique within its report +run and pair of simulation runs, because a budget-window run compares one pair +per year. Poverty and inequality, which describe one simulation, belong to a +simulation run. """ from datetime import datetime @@ -346,13 +348,15 @@ class ChangeAggregate(IdentifiedModel, table=True): class BudgetSummary(IdentifiedModel, table=True): """Baseline and reform totals of one budget variable for a report run. - One row per run, variable and entity. + One row per report run, pair of simulation runs, variable and entity. """ __tablename__ = "budget_summary" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "variable_name", "entity", name="uq_budget_summary_run_variable_entity", @@ -383,19 +387,26 @@ class BudgetSummary(IdentifiedModel, table=True): 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. + ``decile`` runs from 1 to ``quantiles``; there is one row per report run, + pair of simulation runs, income variable, decile variable, entity, quantile + count and decile. ``decile_variable`` names a variable that groups the + households instead of the income variable's own deciles, such as + ``household_wealth_decile``; rows without one count as equal. """ __tablename__ = "decile_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "income_variable", + "decile_variable", "entity", "decile", "quantiles", name="uq_decile_impacts_run_measure", + postgresql_nulls_not_distinct=True, ), sa.CheckConstraint("quantiles > 0", name="ck_decile_impacts_quantiles"), sa.CheckConstraint( @@ -417,6 +428,7 @@ class DecileImpact(IdentifiedModel, table=True): ondelete="RESTRICT", ) income_variable: str = Field(max_length=512) + decile_variable: str | None = Field(default=None, max_length=512) entity: str = Field(max_length=128) decile: int quantiles: int = 10 @@ -434,13 +446,16 @@ class DecileImpact(IdentifiedModel, table=True): 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. + One row per report run, pair of simulation runs, decile type and decile, + where ``decile`` is 0 to 10. """ __tablename__ = "intra_decile_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "decile_type", "decile", name="uq_intra_decile_impacts_run_type_decile", @@ -507,9 +522,10 @@ class Inequality(IdentifiedModel, table=True): class Poverty(IdentifiedModel, table=True): """Poverty headcount and rate of one type for one simulation run. - One row per run, poverty type, entity and filter variable; rows without a - filter variable count as equal, so each run has one unfiltered row per type - and entity. + One row per run, poverty type, entity, filter variable and filter group: + ``filter_variable`` selects people (``age``) and ``filter_group`` names the + group (``child``). Missing values count as equal, so each run has one + unfiltered row per type and entity. """ __tablename__ = "poverty" @@ -519,6 +535,7 @@ class Poverty(IdentifiedModel, table=True): "poverty_type", "entity", "filter_variable", + "filter_group", name="uq_poverty_simulation_run_measure", postgresql_nulls_not_distinct=True, ), @@ -531,6 +548,7 @@ class Poverty(IdentifiedModel, table=True): 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) + filter_group: str | None = Field(default=None, max_length=64) headcount: float | None = None total_population: float | None = None rate: float | None = None @@ -541,7 +559,8 @@ class Poverty(IdentifiedModel, table=True): 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. The economy + One row per report run, pair of simulation runs, program and entity; + ``is_tax`` marks taxes. The economy output's detailed budget names neither the entity nor whether a program is a tax, so both may be null, and rows without an entity count as equal. """ @@ -550,6 +569,8 @@ class ProgramStatistics(IdentifiedModel, table=True): __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "program_name", "entity", name="uq_program_statistics_run_program_entity", @@ -628,14 +649,16 @@ def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: class CongressionalDistrictImpact(GeographicImpactBase, table=True): """Household income change and population in one US congressional district. - One row per run and district GEOID, with the shares of people who gain, - lose or see no change. + One row per report run, pair of simulation runs and district GEOID, with the + shares of people who gain, lose or see no change. """ __tablename__ = "congressional_district_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "district_geoid", name="uq_congressional_district_impacts_run_geoid", ), @@ -656,13 +679,15 @@ class CongressionalDistrictImpact(GeographicImpactBase, table=True): class ConstituencyImpact(UKAreaImpactBase, table=True): """Household income change in one UK constituency. - One row per run and constituency code. + One row per report run, pair of simulation runs and constituency code. """ __tablename__ = "constituency_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "constituency_code", name="uq_constituency_impacts_run_code", ), @@ -678,13 +703,15 @@ class ConstituencyImpact(UKAreaImpactBase, table=True): class LocalAuthorityImpact(UKAreaImpactBase, table=True): """Household income change in one UK local authority. - One row per run and local authority code. + One row per report run, pair of simulation runs and local authority code. """ __tablename__ = "local_authority_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "local_authority_code", name="uq_local_authority_impacts_run_code", ), diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 1f1f18f6e..5ec391ecf 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -38,6 +38,7 @@ SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" SIMULATION_RUNS_REVISION = "2ffd50f47a2b" CROSS_SERVICE_KEYS_REVISION = "7a83a87d029a" +ECONOMY_OUTPUT_REVISION = "7f43d2c7c3c6" NEW_TABLES = frozenset({"simulation_runs", "report_run_simulation_runs"}) CHANGED_TABLES = frozenset( { @@ -175,6 +176,16 @@ def test_each_revision_downgrades_to_its_parent_and_upgrades_again( ) -> None: command.upgrade(config, "head") + command.downgrade(config, ECONOMY_OUTPUT_REVISION) + assert "decile_variable" not in _column_names(engine, "decile_impacts") + assert "filter_group" not in _column_names(engine, "poverty") + budget_key = next( + constraint + for constraint in inspect(engine).get_unique_constraints("budget_summary") + if constraint["name"] == "uq_budget_summary_run_variable_entity" + ) + assert budget_key["column_names"] == ["report_run_id", "variable_name", "entity"] + command.downgrade(config, CROSS_SERVICE_KEYS_REVISION) assert {"spm_provenance", "data_version"}.isdisjoint( _column_names(engine, "report_runs") @@ -408,6 +419,42 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( _insert(connection, "poverty", id=uuid4(), **poverty) with _refused(connection): _insert(connection, "poverty", id=uuid4(), **poverty) + for group in ("child", "adult"): + _insert( + connection, + "poverty", + id=uuid4(), + filter_variable="age", + filter_group=group, + **poverty, + ) + decile = { + "income_variable": "household_net_income", + "entity": "household", + "decile": 1, + "quantiles": 10, + **comparison_runs, + } + _insert(connection, "decile_impacts", id=uuid4(), **decile) + _insert( + connection, + "decile_impacts", + id=uuid4(), + decile_variable="household_wealth_decile", + **decile, + ) + with _refused(connection): + _insert(connection, "decile_impacts", id=uuid4(), **decile) + budget = {"variable_name": "household_tax", "entity": "household"} + _insert(connection, "budget_summary", id=uuid4(), **budget, **comparison_runs) + next_year_runs = { + **comparison_runs, + "baseline_simulation_run_id": _insert_simulation_run(connection, simulation_id), + "reform_simulation_run_id": _insert_simulation_run(connection, simulation_id), + } + _insert(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) + with _refused(connection): + _insert(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) _insert( connection, "program_statistics", diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index e90ec3a27..bcf1dc956 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -74,6 +74,18 @@ CROSS_SERVICE_KEYS_PARENT_REVISION = "2ffd50f47a2b" # The cross-service foreign key revision, parent of the economy output revision. ECONOMY_OUTPUT_PARENT_REVISION = "7a83a87d029a" +# The economy output revision, parent of the result row identity revision. +RESULT_IDENTITY_PARENT_REVISION = "7f43d2c7c3c6" +# Each comparison result table's unique key, which gains the pair of runs. +COMPARISON_UNIQUE_KEYS = { + "budget_summary": "uq_budget_summary_run_variable_entity", + "congressional_district_impacts": "uq_congressional_district_impacts_run_geoid", + "constituency_impacts": "uq_constituency_impacts_run_code", + "decile_impacts": "uq_decile_impacts_run_measure", + "intra_decile_impacts": "uq_intra_decile_impacts_run_type_decile", + "local_authority_impacts": "uq_local_authority_impacts_run_code", + "program_statistics": "uq_program_statistics_run_program_entity", +} # The foreign keys from the simulation service's tables to the Public API's. CROSS_SERVICE_KEYS = { "household_jobs": {"dynamics", "policies"}, @@ -275,6 +287,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + RESULT_IDENTITY_PARENT_REVISION, ECONOMY_OUTPUT_PARENT_REVISION, CROSS_SERVICE_KEYS_PARENT_REVISION, SIMULATION_RUNS_PARENT_REVISION, @@ -1207,3 +1220,56 @@ def test_economy_output_revision_only_adds_columns_and_relaxes_two() -> None: assert "op.execute(" not in body assert "op.bulk_insert(" not in body assert "op.drop_column(" not in upgrade + + +def test_result_identity_revision_adds_two_columns_and_widens_unique_keys() -> None: + """The revision after the economy output revision is generated. + + It adds ``decile_impacts.decile_variable`` and ``poverty.filter_group`` and + rebuilds each comparison result table's unique key with the pair of + simulation runs, and ``poverty``'s with its group. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == RESULT_IDENTITY_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + added_columns = set( + re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) + ) + assert added_columns == { + ("decile_impacts", "decile_variable"), + ("poverty", "filter_group"), + } + rebuilt_keys = { + table: name + for name, table in re.findall( + r'op\.create_unique_constraint\(\s*"(\w+)",\s*"(\w+)"', upgrade + ) + } + assert rebuilt_keys == { + **COMPARISON_UNIQUE_KEYS, + "poverty": "uq_poverty_simulation_run_measure", + } + assert upgrade.count("op.drop_constraint(") == len(rebuilt_keys) + assert upgrade.count('"baseline_simulation_run_id",') == len(COMPARISON_UNIQUE_KEYS) + assert downgrade.count("op.create_unique_constraint(") == len(rebuilt_keys) + assert downgrade.count("op.drop_column(") == len(added_columns) + for body in (upgrade, downgrade): + assert "op.drop_table(" not in body + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body + assert "op.drop_column(" not in upgrade diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index d321048a4..6f352eb31 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -255,7 +255,8 @@ def test_detailed_budget_rows_need_neither_an_entity_nor_a_tax_flag() -> None: ) assert ( "CONSTRAINT uq_program_statistics_run_program_entity UNIQUE NULLS NOT " - "DISTINCT (report_run_id, program_name, entity)" + "DISTINCT (report_run_id, baseline_simulation_run_id, " + "reform_simulation_run_id, program_name, entity)" ) in statement @@ -325,7 +326,13 @@ def test_per_simulation_results_belong_to_one_simulation_run() -> None: ("simulation_run_id", "income_variable", "entity") } assert _unique_column_sets(V2_METADATA.tables["poverty"]) == { - ("simulation_run_id", "poverty_type", "entity", "filter_variable") + ( + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + ) } @@ -336,7 +343,7 @@ def test_poverty_rows_without_a_filter_variable_are_unique_per_run() -> None: assert ( "CONSTRAINT uq_poverty_simulation_run_measure UNIQUE NULLS NOT DISTINCT " - "(simulation_run_id, poverty_type, entity, filter_variable)" + "(simulation_run_id, poverty_type, entity, filter_variable, filter_group)" ) in statement @@ -354,8 +361,13 @@ def test_comparison_results_name_the_two_simulation_runs_they_compare() -> None: run_key = _foreign_key(table, column_name) assert run_key.target_fullname == "simulation_runs.id", table_name assert run_key.ondelete == "RESTRICT" - assert any( - columns[0] == "report_run_id" for columns in _unique_column_sets(table) + # A budget-window run compares one pair of simulation runs per year, so + # each row is unique within its report run and pair. + (unique_columns,) = _unique_column_sets(table) + assert unique_columns[:3] == ( + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", ), table_name @@ -412,3 +424,25 @@ def test_no_relationship_joins_models_that_different_services_write() -> None: ) assert cross_service_relationships == [] + + +def test_poverty_rows_name_the_group_their_filter_selects() -> None: + """One filter variable selects several groups (``age``: child, adult, senior).""" + poverty = V2_METADATA.tables["poverty"] + + assert poverty.c.filter_group.nullable + assert poverty.c.filter_group.type.length == 64 + + +def test_decile_rows_name_the_variable_that_groups_them() -> None: + """Wealth deciles of income changes share the income variable and entity.""" + decile_impacts = V2_METADATA.tables["decile_impacts"] + + assert decile_impacts.c.decile_variable.nullable + assert decile_impacts.c.decile_variable.type.length == 512 + statement = str(CreateTable(decile_impacts).compile(dialect=postgresql.dialect())) + assert ( + "CONSTRAINT uq_decile_impacts_run_measure UNIQUE NULLS NOT DISTINCT " + "(report_run_id, baseline_simulation_run_id, reform_simulation_run_id, " + "income_variable, decile_variable, entity, decile, quantiles)" + ) in statement From b5cd8d7423633685b1d296dbd14b141f49c57d15 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:11:17 +0400 Subject: [PATCH 07/22] Let the simulation API's roles read dynamics, and name skipped roles A simulation and a household job may name a dynamic, which a run reads, so the simulation service and Modal worker roles get SELECT on dynamics. reports.user_id is the owner the Public API names and checks, since the simulation service can't read users. The manifest and the Alembic guide state that the workers' UPDATE is table-level. verify now names a skipped role in the job summary as apply does, the catalog publication test checks that the catalog rows a simulation and a report name survive republishing, and the guide's descriptions of the new revisions are corrected. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 13 ++++++++----- policyengine_api/data/v2/models/reports.py | 4 +++- .../data/v2/privilege_reconciliation.py | 6 +++++- policyengine_api/data/v2/privileges.py | 15 ++++++++++----- tests/integration/test_v2_catalog_publication.py | 16 ++++++++++++++++ .../test_v2_privilege_reconciliation.py | 6 +++++- tests/unit/v2/test_privilege_manifest.py | 3 ++- 7 files changed, 49 insertions(+), 14 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d6eb1fbe3..c3e64628d 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -226,7 +226,7 @@ and a revision runs once per database, so it can't repair drift. | `policyengine_v2_catalog_writer` | `scripts/publish_v2_metadata_catalog.py` | | `policyengine_stage12_` | The simulation API and its Modal worker, on the temporary comparison-run tables | | `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs and household jobs, and reads the policy, household and catalog rows a run needs | -| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs | +| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs. Its `UPDATE` is table-level: the workers' code, not the database, limits it to their own run's status and result columns | Every v2 table has an entry for every role, so "no access" is a decision, not an omission. The migration role owns every table and isn't listed. Three rules @@ -282,6 +282,7 @@ When a pull request adds or changes a v2 table, or a connector's use of one: tests/integration/test_v2_stage12_comparison_persistence.py \ tests/integration/test_v2_metadata_routes.py \ tests/integration/test_v2_catalog_publication.py \ + tests/integration/test_v2_simulation_role_privileges.py \ tests/integration/test_v2_privilege_reconciliation.py ``` @@ -473,9 +474,10 @@ model changes: - `simulations` keeps only what to calculate. Its execution columns (`status`, `error_message`, `started_at`, `completed_at`, - `household_result`, `output_dataset_id`) move to the new `simulation_runs`, - one row per run with the catalog version, country package version and - PolicyEngine.py version it used, its Modal call ID and its output file. + `household_result`) move to the new `simulation_runs`, one row per run with + the catalog version, country package version and PolicyEngine.py version it + used, and its Modal call ID; `output_dataset_id` is replaced by the run's + output file (`output_uri`, `output_sha256`). - The new `report_run_simulation_runs` names the simulation runs each report run combines, as its baseline or its reform. - `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); @@ -505,7 +507,8 @@ deploy stops before anything changes. The downgrade restores required columns the same way, so it needs `simulations`, `household_jobs` and the result tables empty. The previous Public API revision keeps working at this revision: none of its roles can read or write a table the revision changes, except the -catalog writer, which only counts the rows of `simulations` and `report_runs`. +catalog writer, which only counts the rows of `simulations`, `reports` and +`report_runs`. `tests/integration/test_alembic_v2_simulation_runs.py` checks that the revision changes no other table, and the v2 lifecycle test covers the downgrade and re-upgrade. diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 9c2d06ad1..bfc209713 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -76,7 +76,9 @@ class Report(TimestampedModel, table=True): Its user, model, policy, household, dataset and region are logical references to the Public API's rows, without foreign keys: the simulation - service checks them before it writes the report. + service checks them before it writes the report, except ``user_id``, the + owner the Public API names in its request, which the Public API checks + because the simulation service can't read users. """ __tablename__ = "reports" diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index 9e89c075e..8213a23a4 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -513,7 +513,11 @@ def run( if mode == "verify": skipped = skipped_roles(connection, declared, awaiting) if skipped: - print(_skipped_roles_line(skipped)) + _write_summary( + values, + f"## v2 table privileges ({environment})\n\n" + f"{_skipped_roles_line(skipped)}\n", + ) failures = check_declared_privileges( connection, declared, roles_awaiting_creation=awaiting ) diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 7ef0bbc6f..f300ff0d3 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -26,11 +26,14 @@ so the deploy check, not the connector analysis, covers it. - ``policyengine_v2_simulation_service``: the simulation API on Cloud Run. It creates simulations, reports, their runs and household jobs, and reads the - policy, household and catalog rows a run needs. + policy, dynamic, household and catalog rows a run needs. - ``policyengine_v2_modal_worker``: the simulation API's Modal workers. They write the status of the runs and household jobs they were given, and their - results, and read the rows a run needs. Both simulation API roles' code lives - in that repository, so the deploy check covers them here. + results, and read the rows a run needs. Their ``UPDATE`` is table-level, like + every privilege here: the database doesn't limit it to the status and result + columns or to the worker's own run, so the workers' code does. Both + simulation API roles' code lives in that repository, so the deploy check + covers them here. The Public API reads the simulation tables it serves, and writes none of them. @@ -130,7 +133,8 @@ def _access( ) # The Modal workers write results, keyed by run, once each. _RESULT = _access(runtime=READ, modal_worker=READ_INSERT) -# What a run needs from the Public API's policies, households and catalog. +# What a run needs from the Public API's policies, dynamics, households and +# catalog. _RUN_INPUT_READS: Final = {"simulation_service": READ, "modal_worker": READ} @@ -140,7 +144,8 @@ def _access( # Custom totals, unused until they return. "aggregates": _access(), "change_aggregates": _access(), - "dynamics": _access(), + # A simulation's or household job's dynamic, which a run reads. + "dynamics": _access(**_RUN_INPUT_READS), # Simulations, reports, their runs and household jobs. The catalog writer # checks that publication leaves the rows of three of them untouched. "household_jobs": _RUN_STATUS, diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index c539f0faf..aad982354 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -338,6 +338,22 @@ def test_republish_preserves_existing_run_and_dataset_version_records( assert persisted_simulation.dataset_id == protected_ids[4] assert persisted_report.dataset_id == protected_ids[5] assert persisted_dataset_version.dataset_id == protected_ids[6] + # The simulation and report name catalog rows without foreign keys, + # so republishing must keep every row they name. + for dataset_id in ( + persisted_simulation.dataset_id, + persisted_report.dataset_id, + ): + assert session.get(Dataset, dataset_id) is not None + for region_id in (persisted_simulation.region_id, persisted_report.region_id): + assert session.get(Region, region_id) is not None + assert ( + session.get( + TaxBenefitModelVersion, + persisted_simulation.tax_benefit_model_version_id, + ) + is not None + ) def test_new_version_is_additive_and_concurrent_retries_serialize( diff --git a/tests/integration/test_v2_privilege_reconciliation.py b/tests/integration/test_v2_privilege_reconciliation.py index 3fa2e3376..0a764e0ce 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -162,7 +162,7 @@ def test_a_role_awaiting_creation_is_skipped_until_it_exists( def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( - disposable_v2_database_url, temporary_role, monkeypatch, capsys + disposable_v2_database_url, temporary_role, monkeypatch, capsys, tmp_path ) -> None: """``apply`` and ``verify`` succeed while a role awaiting creation is absent.""" role = temporary_role() @@ -177,10 +177,12 @@ def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( "ROLES_AWAITING_PROVISIONING", frozenset({absent_role}), ) + summary_path = tmp_path / "summary.md" environ = { V2_MIGRATION_DATABASE_URL: disposable_v2_database_url, V2_ALEMBIC_DISPOSABLE_TEST: "1", DEPLOYMENT_ENVIRONMENT: "test", + "GITHUB_STEP_SUMMARY": str(summary_path), } apply_status = run("apply", environ=environ) @@ -192,3 +194,5 @@ def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( assert f"`{absent_role}`" in apply_report assert absent_role in verify_output assert "Every declared v2 table privilege is held" in verify_output + # Each step names the skipped role in the job summary. + assert summary_path.read_text(encoding="utf-8").count(f"`{absent_role}`") == 2 diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 5c635e22e..749262f12 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -51,10 +51,11 @@ } ) RUN_STATUS_TABLES = frozenset({"simulation_runs", "report_runs", "household_jobs"}) -# The policy, household and catalog rows a run needs. +# The policy, dynamic, household and catalog rows a run needs. RUN_INPUT_TABLES = frozenset( { "policies", + "dynamics", "households", "tax_benefit_models", "tax_benefit_model_versions", From 601561f2ed7882d38cd67567a28312c3c0776e29 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:35:45 +0400 Subject: [PATCH 08/22] Keep a v2 report run's untabled results per compared pair A budget-window report run compares one pair of simulation runs per year, and each pair's economy output has its own labour supply response, cliff impact, SPM receipts and data release. Revision dce0caa329cc, autogenerated with no corrections, moves those four values from report_runs to the new report_run_comparisons, one row per report run and pair, which the Modal workers insert and the Public API reads. The report run relationships are named for the rows they hold (simulation_run_members, comparisons, report_run_memberships), the simulation-side integration tests share their SQL helpers, and the Alembic guide says what deleting a user now leaves in reports.user_id. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 18 ++- ...c_keep_untabled_comparison_results_per_.py | 120 ++++++++++++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/__init__.py | 2 + policyengine_api/data/v2/models/reports.py | 62 +++++++-- .../data/v2/models/simulations.py | 2 +- policyengine_api/data/v2/privileges.py | 1 + tests/integration/simulation_rows.py | 45 ++++++ .../test_alembic_v2_simulation_runs.py | 129 +++++++++--------- .../test_v2_simulation_role_privileges.py | 49 ++++--- tests/unit/v2/test_alembic_v2.py | 47 +++++++ tests/unit/v2/test_privilege_manifest.py | 1 + tests/unit/v2/test_simulation_run_models.py | 51 ++++++- 13 files changed, 424 insertions(+), 105 deletions(-) create mode 100644 migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py create mode 100644 tests/integration/simulation_rows.py diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index c3e64628d..e148d4d5a 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -225,7 +225,7 @@ and a revision runs once per database, so it can't repair drift. | `policyengine_v2_runtime` | The Public API on Cloud Run | | `policyengine_v2_catalog_writer` | `scripts/publish_v2_metadata_catalog.py` | | `policyengine_stage12_` | The simulation API and its Modal worker, on the temporary comparison-run tables | -| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs and household jobs, and reads the policy, household and catalog rows a run needs | +| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs and household jobs, and reads the policy, dynamic, household and catalog rows a run needs | | `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs. Its `UPDATE` is table-level: the workers' code, not the database, limits it to their own run's status and result columns | Every v2 table has an entry for every role, so "no access" is a decision, not @@ -526,7 +526,10 @@ tables the Public API writes (`users`, `policies`, `dynamics`, `households`, and the catalog's `tax_benefit_models`, `tax_benefit_model_versions`, `datasets` and `regions`). The 14 columns stay as plain UUIDs: the service that writes a row checks the rows it names. Foreign keys between tables of the same -service stay. It has no post-generation corrections. Its downgrade restores the +service stay. Deleting a Public API row no longer clears or blocks the rows +that name it: in particular, deleting a user leaves that user's ID in +`reports.user_id`, which no managed role may update; the owner of a saved +report is the user association, which the deletion removes. It has no post-generation corrections. Its downgrade restores the keys, so it fails while a column names a row the referenced table lacks. `tests/unit/v2/test_simulation_run_models.py` checks that no foreign key or ORM relationship joins the two services' tables. @@ -567,3 +570,14 @@ post-generation corrections and needs no table empty; its downgrade fails while a report run holds comparison rows of more than one pair, or a run holds rows that differ only in a removed column. +### API v2 comparisons + +Revision `dce0caa329cc` (parent `2e3de7c572ed`) was autogenerated after the +models moved the parts of a comparison that have no result table +(`labor_supply_result`, `cliff_impact_result`, `spm_provenance`, +`data_version`) from `report_runs` to the new `report_run_comparisons`, one row +per report run and compared pair of simulation runs, because a budget-window +run compares one pair per year. It has no post-generation corrections. The +four `report_runs` columns it drops were added by the earlier revisions of +this change and never written; its downgrade drops the new table and its +rows. diff --git a/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py b/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py new file mode 100644 index 000000000..853202f49 --- /dev/null +++ b/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py @@ -0,0 +1,120 @@ +"""keep untabled comparison results per compared pair + +Revision ID: dce0caa329cc +Revises: 2e3de7c572ed +Create Date: 2026-11-04 01:33:30.494861 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +No post-generation corrections. The upgrade drops four nullable report_runs +columns, which nothing has written: no managed role could write report_runs +before the previous revisions. The downgrade drops report_run_comparisons and +its rows. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel +from sqlalchemy.dialects import postgresql + +revision: str = "dce0caa329cc" +down_revision: Union[str, None] = "2e3de7c572ed" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "report_run_comparisons", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_run_id", sa.Uuid(), nullable=False), + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("labor_supply_result", sa.JSON(), nullable=True), + sa.Column("cliff_impact_result", sa.JSON(), nullable=True), + sa.Column("spm_provenance", sa.JSON(), nullable=True), + sa.Column( + "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True + ), + sa.ForeignKeyConstraint( + ["baseline_simulation_run_id"], + ["simulation_runs.id"], + name=op.f( + "fk_report_run_comparisons_baseline_simulation_run_id_simulation_runs" + ), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["reform_simulation_run_id"], + ["simulation_runs.id"], + name=op.f( + "fk_report_run_comparisons_reform_simulation_run_id_simulation_runs" + ), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_report_run_comparisons_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_comparisons")), + sa.UniqueConstraint( + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + name="uq_report_run_comparisons_run_pair", + ), + ) + op.drop_column("report_runs", "cliff_impact_result") + op.drop_column("report_runs", "labor_supply_result") + op.drop_column("report_runs", "data_version") + op.drop_column("report_runs", "spm_provenance") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "report_runs", + sa.Column( + "spm_provenance", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "report_runs", + sa.Column( + "data_version", sa.VARCHAR(length=128), autoincrement=False, nullable=True + ), + ) + op.add_column( + "report_runs", + sa.Column( + "labor_supply_result", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "report_runs", + sa.Column( + "cliff_impact_result", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), + ) + op.drop_table("report_run_comparisons") + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 03f0fb38e..da6e3f5e5 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "2e3de7c572ed" +V2_ALEMBIC_HEAD_REVISION = "dce0caa329cc" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index d9b4a3c04..cc182268d 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -82,6 +82,7 @@ ProgramStatistics, Report, ReportRun, + ReportRunComparison, ReportRunSimulationRun, ReportRunStatus, ReportRunTrigger, @@ -141,6 +142,7 @@ "RegionType", "Report", "ReportRun", + "ReportRunComparison", "ReportRunSimulationRun", "ReportRunStatus", "ReportRunTrigger", diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index bfc209713..71bfbc931 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -130,10 +130,9 @@ class Report(TimestampedModel, table=True): class ReportRun(TimestampedModel, table=True): """One immutable-version execution attempt for a stable report. - Its ``simulation_runs`` rows name the simulation runs it combines. - ``labor_supply_result``, ``cliff_impact_result`` and ``spm_provenance`` hold - the parts of an economy comparison that have no result table, and - ``data_version`` the release of the data it used. + Its ``simulation_run_members`` name the simulation runs it combines, and + its ``comparisons`` hold, per compared pair, the parts of an economy + comparison that have no result table. """ __tablename__ = "report_runs" @@ -185,13 +184,13 @@ class ReportRun(TimestampedModel, table=True): ) error_message: str | None = None markdown: str | None = Field(default=None, sa_type=sa.Text) - labor_supply_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - cliff_impact_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - spm_provenance: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - data_version: str | None = Field(default=None, max_length=128) report: Report = Relationship(back_populates="runs") - simulation_runs: list["ReportRunSimulationRun"] = Relationship( + simulation_run_members: list["ReportRunSimulationRun"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) + comparisons: list["ReportRunComparison"] = Relationship( back_populates="report_run", cascade_delete=True, ) @@ -262,8 +261,49 @@ class ReportRunSimulationRun(IdentifiedModel, table=True): sa_type=enum_type(SimulationRole, "v2_simulation_role"), ) - report_run: ReportRun = Relationship(back_populates="simulation_runs") - simulation_run: SimulationRun = Relationship(back_populates="report_runs") + report_run: ReportRun = Relationship(back_populates="simulation_run_members") + simulation_run: SimulationRun = Relationship( + back_populates="report_run_memberships" + ) + + +class ReportRunComparison(IdentifiedModel, table=True): + """The parts of one compared pair of a report run that have no result table. + + A report run compares one baseline and one reform simulation run, or one + pair per year for a budget window. Each pair's labour supply response, + cliff impact, SPM receipts and data release are one row here, as the + simulation API's economy output gives them. + """ + + __tablename__ = "report_run_comparisons" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + name="uq_report_run_comparisons_run_pair", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + labor_supply_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + cliff_impact_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + spm_provenance: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + data_version: str | None = Field(default=None, max_length=128) + + report_run: ReportRun = Relationship(back_populates="comparisons") class AggregateOutput(IdentifiedModel, table=True): diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index ff24aae72..4a2689a3a 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -158,7 +158,7 @@ class SimulationRun(TimestampedModel, table=True): output_sha256: str | None = Field(default=None, max_length=64) simulation: Simulation = Relationship(back_populates="runs") - report_runs: list["ReportRunSimulationRun"] = Relationship( + report_run_memberships: list["ReportRunSimulationRun"] = Relationship( back_populates="simulation_run" ) poverty_results: list["Poverty"] = Relationship( diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index f300ff0d3..2dcef8ff8 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -155,6 +155,7 @@ def _access( "simulation_runs": _RUN_STATUS, "simulations": replace(_SIMULATION_DEFINITION, catalog_writer=READ), # Results of simulation runs and report runs. + "report_run_comparisons": _RESULT, "budget_summary": _RESULT, "congressional_district_impacts": _RESULT, "constituency_impacts": _RESULT, diff --git a/tests/integration/simulation_rows.py b/tests/integration/simulation_rows.py new file mode 100644 index 000000000..8a44d0a71 --- /dev/null +++ b/tests/integration/simulation_rows.py @@ -0,0 +1,45 @@ +"""Write simulation-side v2 rows with plain SQL in the integration tests. + +The statements run as whichever role the connection acts as, so a test can +check what that role may do. Tests write inside a transaction they roll back. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +# The package versions a test run records. +RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} + + +def insert_row(connection: Connection, table: str, **values: object) -> None: + """Insert one row of the named columns into ``table``.""" + columns = ", ".join(values) + placeholders = ", ".join(f":{column}" for column in values) + connection.execute( + text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"), values + ) + + +@contextmanager +def refused( + connection: Connection, + error: type[Exception] = IntegrityError, + *, + match: str | None = None, +) -> Iterator[None]: + """Expect the statements inside to fail with ``error``, then roll them back.""" + savepoint = connection.begin_nested() + with pytest.raises(error, match=match): + yield + savepoint.rollback() diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 5ec391ecf..337fd2ae2 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,19 +1,21 @@ """Check the simulation runs revisions against disposable Postgres. -Three revisions follow the area impact coordinate removal. The first adds +Five revisions follow the area impact coordinate removal. The first adds simulation runs and the table naming the simulation runs a report run combines, keys every result table by run, and gives household jobs one household, one year and one policy. The second drops the foreign keys from those tables to the Public API's tables. The third adds the economy output's -fields that had no column. None changes any other table, so the previous -Public API revision, whose roles can't read or write any table they change, -keeps working while they are deployed. +fields that had no column. The fourth makes each result row tell apart its +pair of simulation runs, decile grouping and poverty group. The fifth moves +the parts of a comparison without a result table from the report run to one +row per compared pair. None changes any other table, so the previous Public +API revision, whose roles can't read or write any table they change, keeps +working while they are deployed. """ from __future__ import annotations import os -from contextlib import contextmanager from typing import TYPE_CHECKING from uuid import uuid4 @@ -21,7 +23,6 @@ from alembic import command from alembic.config import Config from sqlalchemy import create_engine, inspect, text -from sqlalchemy.exc import IntegrityError from policyengine_api.constants import REPO from policyengine_api.data.v2.migration_target import ( @@ -29,6 +30,7 @@ load_v2_alembic_settings, ) from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL +from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused if TYPE_CHECKING: from collections.abc import Iterator @@ -39,7 +41,10 @@ SIMULATION_RUNS_REVISION = "2ffd50f47a2b" CROSS_SERVICE_KEYS_REVISION = "7a83a87d029a" ECONOMY_OUTPUT_REVISION = "7f43d2c7c3c6" -NEW_TABLES = frozenset({"simulation_runs", "report_run_simulation_runs"}) +RESULT_IDENTITY_REVISION = "2e3de7c572ed" +NEW_TABLES = frozenset( + {"simulation_runs", "report_run_simulation_runs", "report_run_comparisons"} +) CHANGED_TABLES = frozenset( { "budget_summary", @@ -62,7 +67,6 @@ SIMULATION_SIDE_TABLES = ( CHANGED_TABLES | NEW_TABLES | {"aggregates", "change_aggregates"} ) -RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} SHA256_HEX = "a" * 64 COMPLETED_AT = "2026-11-04T00:00:00Z" @@ -176,6 +180,12 @@ def test_each_revision_downgrades_to_its_parent_and_upgrades_again( ) -> None: command.upgrade(config, "head") + command.downgrade(config, RESULT_IDENTITY_REVISION) + assert "report_run_comparisons" not in inspect(engine).get_table_names() + assert {"labor_supply_result", "data_version"} <= _column_names( + engine, "report_runs" + ) + command.downgrade(config, ECONOMY_OUTPUT_REVISION) assert "decile_variable" not in _column_names(engine, "decile_impacts") assert "filter_group" not in _column_names(engine, "poverty") @@ -234,23 +244,6 @@ def test_head_has_no_foreign_key_between_the_two_services_tables( assert cross_service_keys == [] -@contextmanager -def _refused(connection: Connection) -> Iterator[None]: - """Expect the statement inside to violate a constraint, then roll it back.""" - savepoint = connection.begin_nested() - with pytest.raises(IntegrityError): - yield - savepoint.rollback() - - -def _insert(connection: Connection, table: str, **values: object) -> None: - columns = ", ".join(values) - placeholders = ", ".join(f":{column}" for column in values) - connection.execute( - text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"), values - ) - - def _count(connection: Connection, table: str, column: str, value: object) -> int: return connection.execute( text(f"SELECT count(*) FROM {table} WHERE {column} = :value"), @@ -261,7 +254,7 @@ def _count(connection: Connection, table: str, column: str, value: object) -> in def _insert_simulation(connection: Connection) -> object: """Store an economy simulation; its catalog references need no catalog rows.""" simulation_id = uuid4() - _insert( + insert_row( connection, "simulations", id=simulation_id, @@ -281,7 +274,7 @@ def _insert_simulation_run( **values: object, ) -> object: run_id = uuid4() - _insert( + insert_row( connection, "simulation_runs", id=run_id, @@ -296,7 +289,7 @@ def _insert_simulation_run( def _insert_report_run(connection: Connection) -> object: report_id, report_run_id = uuid4(), uuid4() - _insert( + insert_row( connection, "reports", id=report_id, @@ -305,7 +298,7 @@ def _insert_report_run(connection: Connection) -> object: tax_benefit_model_id=uuid4(), inputs="{}", ) - _insert( + insert_row( connection, "report_runs", id=report_run_id, @@ -323,16 +316,16 @@ def test_simulation_runs_need_a_completion_time_and_a_whole_output( simulation_id = _insert_simulation(connection) run_id = _insert_simulation_run(connection, simulation_id) - with _refused(connection): + with refused(connection): connection.execute( text("UPDATE simulation_runs SET status = 'succeeded' WHERE id = :id"), {"id": run_id}, ) - with _refused(connection): + with refused(connection): _insert_simulation_run( connection, simulation_id, output_uri="gs://bucket/run.parquet" ) - with _refused(connection): + with refused(connection): _insert_simulation_run( connection, simulation_id, @@ -347,7 +340,7 @@ def test_simulation_runs_need_a_completion_time_and_a_whole_output( output_uri="gs://bucket/run.parquet", output_sha256=SHA256_HEX, ) - _insert( + insert_row( connection, "poverty", id=uuid4(), @@ -373,11 +366,11 @@ def test_report_run_members_are_unique_and_keep_their_simulation_runs( "simulation_run_id": run_id, "role": "baseline", } - _insert(connection, "report_run_simulation_runs", id=uuid4(), **member) + insert_row(connection, "report_run_simulation_runs", id=uuid4(), **member) - with _refused(connection): - _insert(connection, "report_run_simulation_runs", id=uuid4(), **member) - with _refused(connection): + with refused(connection): + insert_row(connection, "report_run_simulation_runs", id=uuid4(), **member) + with refused(connection): connection.execute( text("DELETE FROM simulation_runs WHERE id = :id"), {"id": run_id} ) @@ -402,8 +395,8 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( "reform_simulation_run_id": reform_run_id, } - with _refused(connection): - _insert( + with refused(connection): + insert_row( connection, "budget_summary", id=uuid4(), @@ -416,11 +409,11 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( "poverty_type": "spm", "entity": "person", } - _insert(connection, "poverty", id=uuid4(), **poverty) - with _refused(connection): - _insert(connection, "poverty", id=uuid4(), **poverty) + insert_row(connection, "poverty", id=uuid4(), **poverty) + with refused(connection): + insert_row(connection, "poverty", id=uuid4(), **poverty) for group in ("child", "adult"): - _insert( + insert_row( connection, "poverty", id=uuid4(), @@ -435,35 +428,47 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( "quantiles": 10, **comparison_runs, } - _insert(connection, "decile_impacts", id=uuid4(), **decile) - _insert( + insert_row(connection, "decile_impacts", id=uuid4(), **decile) + insert_row( connection, "decile_impacts", id=uuid4(), decile_variable="household_wealth_decile", **decile, ) - with _refused(connection): - _insert(connection, "decile_impacts", id=uuid4(), **decile) + with refused(connection): + insert_row(connection, "decile_impacts", id=uuid4(), **decile) budget = {"variable_name": "household_tax", "entity": "household"} - _insert(connection, "budget_summary", id=uuid4(), **budget, **comparison_runs) + insert_row(connection, "budget_summary", id=uuid4(), **budget, **comparison_runs) next_year_runs = { **comparison_runs, "baseline_simulation_run_id": _insert_simulation_run(connection, simulation_id), "reform_simulation_run_id": _insert_simulation_run(connection, simulation_id), } - _insert(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) - with _refused(connection): - _insert(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) - _insert( + insert_row(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) + with refused(connection): + insert_row(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) + # One comparison row per compared pair holds the untabled parts. + for runs in (comparison_runs, next_year_runs): + insert_row( + connection, + "report_run_comparisons", + id=uuid4(), + labor_supply_result="{}", + data_version="1.0.0", + **runs, + ) + with refused(connection): + insert_row(connection, "report_run_comparisons", id=uuid4(), **next_year_runs) + insert_row( connection, "program_statistics", id=uuid4(), program_name="income_tax", **comparison_runs, ) - with _refused(connection): - _insert( + with refused(connection): + insert_row( connection, "program_statistics", id=uuid4(), @@ -477,7 +482,7 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( "population": 100.0, **comparison_runs, } - _insert( + insert_row( connection, "constituency_impacts", id=uuid4(), @@ -485,8 +490,8 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( disclosure_tier="approximate", **constituency, ) - with _refused(connection): - _insert( + with refused(connection): + insert_row( connection, "constituency_impacts", id=uuid4(), @@ -506,7 +511,7 @@ def test_household_jobs_take_a_valid_year_and_record_when_they_end( "tax_benefit_model_version_id": uuid4(), **RUN_VERSIONS, } - _insert( + insert_row( connection, "household_jobs", id=uuid4(), @@ -515,8 +520,8 @@ def test_household_jobs_take_a_valid_year_and_record_when_they_end( **household_job, ) - with _refused(connection): - _insert( + with refused(connection): + insert_row( connection, "household_jobs", id=uuid4(), @@ -524,8 +529,8 @@ def test_household_jobs_take_a_valid_year_and_record_when_they_end( status="pending", **household_job, ) - with _refused(connection): - _insert( + with refused(connection): + insert_row( connection, "household_jobs", id=uuid4(), diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index bffc9679e..7b2239667 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -20,14 +20,13 @@ from sqlalchemy.pool import NullPool from policyengine_api.constants import REPO +from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused if TYPE_CHECKING: from collections.abc import Iterator from sqlalchemy import Connection -RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} - @pytest.fixture def connection(disposable_v2_database_url, managed_v2_roles) -> Iterator[Connection]: @@ -54,26 +53,16 @@ def _act_as(connection: Connection, role: str) -> None: connection.execute(text(f"SET LOCAL ROLE {quoted_role}")) -def _insert(connection: Connection, table: str, **values: object) -> None: - columns = ", ".join(values) - placeholders = ", ".join(f":{column}" for column in values) - connection.execute( - text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"), values - ) - - def _refused(connection: Connection, statement: str, **values: object) -> None: - """Expect a permission error, then return the connection to a usable state.""" - savepoint = connection.begin_nested() - with pytest.raises(ProgrammingError, match="permission denied"): + """Expect a permission error from ``statement``.""" + with refused(connection, ProgrammingError, match="permission denied"): connection.execute(text(statement), values) - savepoint.rollback() def _create_simulation_rows(connection: Connection) -> dict[str, object]: """Create what the simulation API creates for one economy report run.""" ids = {name: uuid4() for name in ("simulation", "run", "report", "report_run")} - _insert( + insert_row( connection, "simulations", id=ids["simulation"], @@ -82,7 +71,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: tax_benefit_model_version_id=uuid4(), year=2026, ) - _insert( + insert_row( connection, "simulation_runs", id=ids["run"], @@ -91,7 +80,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: status="pending", **RUN_VERSIONS, ) - _insert( + insert_row( connection, "reports", id=ids["report"], @@ -100,7 +89,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: tax_benefit_model_id=uuid4(), inputs="{}", ) - _insert( + insert_row( connection, "report_runs", id=ids["report_run"], @@ -109,7 +98,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: trigger="initial", **RUN_VERSIONS, ) - _insert( + insert_row( connection, "report_run_simulation_runs", id=uuid4(), @@ -117,7 +106,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: simulation_run_id=ids["run"], role="baseline", ) - _insert( + insert_row( connection, "household_jobs", id=uuid4(), @@ -157,6 +146,15 @@ def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( run_id=ids["run"], ) _refused(connection, "DELETE FROM simulations WHERE id = :id", id=ids["simulation"]) + _refused( + connection, + "INSERT INTO report_run_comparisons " + "(id, report_run_id, baseline_simulation_run_id, reform_simulation_run_id) " + "VALUES (:id, :report_run_id, :run_id, :run_id)", + id=uuid4(), + report_run_id=ids["report_run"], + run_id=ids["run"], + ) _refused(connection, "UPDATE policies SET content_hash = content_hash") @@ -174,7 +172,7 @@ def test_modal_workers_write_their_run_status_and_results_only( ), {"id": ids["run"]}, ).scalar_one() - _insert( + insert_row( connection, "poverty", id=uuid4(), @@ -187,6 +185,15 @@ def test_modal_workers_write_their_run_status_and_results_only( text("UPDATE report_runs SET status = 'running' WHERE id = :id"), {"id": ids["report_run"]}, ) + insert_row( + connection, + "report_run_comparisons", + id=uuid4(), + report_run_id=ids["report_run"], + baseline_simulation_run_id=ids["run"], + reform_simulation_run_id=ids["run"], + cliff_impact_result="{}", + ) connection.execute(text("SELECT count(*) FROM parameter_values")).scalar_one() assert updated == ids["run"] diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index bcf1dc956..1171a9cb3 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -47,6 +47,7 @@ "legacy_user_policy_mappings", "parameter_search_files", "parameter_tree_children", + "report_run_comparisons", "report_run_simulation_runs", "simulation_runs", "stage12_evaluation_reports", @@ -76,6 +77,14 @@ ECONOMY_OUTPUT_PARENT_REVISION = "7a83a87d029a" # The economy output revision, parent of the result row identity revision. RESULT_IDENTITY_PARENT_REVISION = "7f43d2c7c3c6" +# The result row identity revision, parent of the comparisons revision. +COMPARISONS_PARENT_REVISION = "2e3de7c572ed" +UNTABLED_COMPARISON_COLUMNS = ( + "labor_supply_result", + "cliff_impact_result", + "spm_provenance", + "data_version", +) # Each comparison result table's unique key, which gains the pair of runs. COMPARISON_UNIQUE_KEYS = { "budget_summary": "uq_budget_summary_run_variable_entity", @@ -287,6 +296,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + COMPARISONS_PARENT_REVISION, RESULT_IDENTITY_PARENT_REVISION, ECONOMY_OUTPUT_PARENT_REVISION, CROSS_SERVICE_KEYS_PARENT_REVISION, @@ -1273,3 +1283,40 @@ def test_result_identity_revision_adds_two_columns_and_widens_unique_keys() -> N assert "op.execute(" not in body assert "op.bulk_insert(" not in body assert "op.drop_column(" not in upgrade + + +def test_comparisons_revision_moves_untabled_results_to_one_row_per_pair() -> None: + """The revision after the result row identity revision is generated. + + It creates ``report_run_comparisons`` and drops the four columns it takes + over from ``report_runs``; the downgrade puts them back. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == COMPARISONS_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + assert 'op.create_table(\n "report_run_comparisons"' in upgrade + assert "uq_report_run_comparisons_run_pair" in upgrade + assert upgrade.count("op.drop_column(") == len(UNTABLED_COMPARISON_COLUMNS) + assert downgrade.count("op.add_column(") == len(UNTABLED_COMPARISON_COLUMNS) + for column_name in UNTABLED_COMPARISON_COLUMNS: + assert f'op.drop_column("report_runs", "{column_name}")' in upgrade + assert f'"{column_name}"' in downgrade + assert 'op.drop_table("report_run_comparisons")' in downgrade + for body in (upgrade, downgrade): + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 749262f12..f38a659d4 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -67,6 +67,7 @@ ) RESULT_TABLES = frozenset( { + "report_run_comparisons", "budget_summary", "congressional_district_impacts", "constituency_impacts", diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 6f352eb31..b6eb45749 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -13,6 +13,8 @@ from policyengine_api.data.v2.models import ( V2_METADATA, + ReportRun, + ReportRunComparison, ReportRunSimulationRun, SimulationRole, SimulationRun, @@ -39,6 +41,7 @@ "reports", "report_runs", "report_run_simulation_runs", + "report_run_comparisons", "household_jobs", "aggregates", "change_aggregates", @@ -203,20 +206,54 @@ def test_run_tables_require_a_completion_time_and_index_running_runs() -> None: ) -def test_report_runs_record_their_modal_call_and_untabled_results() -> None: +def test_report_runs_record_their_modal_call() -> None: report_runs = V2_METADATA.tables["report_runs"] assert report_runs.c.modal_call_id.nullable assert report_runs.c.modal_call_id.type.length == 255 - for column_name in ( + # A budget-window run compares one pair per year, so the parts of a + # comparison without a result table belong to its comparison row. + assert { "labor_supply_result", "cliff_impact_result", "spm_provenance", - ): - assert report_runs.c[column_name].nullable - assert isinstance(report_runs.c[column_name].type, sa.JSON) - assert report_runs.c.data_version.nullable - assert report_runs.c.data_version.type.length == 128 + "data_version", + }.isdisjoint(report_runs.c.keys()) + + +def test_each_comparison_keeps_the_parts_without_a_result_table() -> None: + comparisons = V2_METADATA.tables["report_run_comparisons"] + + for column_name in ("labor_supply_result", "cliff_impact_result", "spm_provenance"): + assert comparisons.c[column_name].nullable + assert isinstance(comparisons.c[column_name].type, sa.JSON) + assert comparisons.c.data_version.nullable + assert comparisons.c.data_version.type.length == 128 + report_run_key = _foreign_key(comparisons, "report_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + for column_name in ("baseline_simulation_run_id", "reform_simulation_run_id"): + assert not comparisons.c[column_name].nullable + run_key = _foreign_key(comparisons, column_name) + assert run_key.target_fullname == "simulation_runs.id" + assert run_key.ondelete == "RESTRICT" + assert _unique_column_sets(comparisons) == { + ("report_run_id", "baseline_simulation_run_id", "reform_simulation_run_id") + } + assert ReportRunComparison.__module__.endswith(".reports") + + +def test_report_run_relationships_name_member_and_comparison_rows() -> None: + relationships = sa.inspect(ReportRun).relationships + + assert relationships["simulation_run_members"].mapper.class_ is ( + ReportRunSimulationRun + ) + assert relationships["comparisons"].mapper.class_ is ReportRunComparison + assert ( + sa.inspect(SimulationRun).relationships["report_run_memberships"].mapper.class_ + is ReportRunSimulationRun + ) def test_uk_area_impacts_record_disclosure_control() -> None: From 1af1cf4728b0637324c310b2027cd2ca6e51b2ce Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 02:13:57 +0400 Subject: [PATCH 09/22] Tidy the v2 privilege check's report and the simulation runs docs verify now writes one job summary section, names the roles it skipped, and says "Every other declared v2 table privilege is held" while it skips any, instead of claiming every privilege is held. New tests cover verify failing once a role awaiting creation exists without its privileges, and a Modal worker finishing a household job. The rules registry describes the skipped roles and the simulation API's role tests, and the Alembic guide notes that NULLS NOT DISTINCT needs PostgreSQL 15. Docstrings name PolicyEngine.py in full, and the report-run module says no managed role can run its row-locking operations. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 4 +- docs/generated/rules.md | 4 +- policyengine_api/data/v2/models/households.py | 4 +- .../data/v2/models/simulations.py | 2 +- .../data/v2/privilege_reconciliation.py | 64 ++++++++++++------- policyengine_api/data/v2/report_runs.py | 7 +- rules/registry.toml | 6 +- .../test_alembic_v2_simulation_runs.py | 11 ++-- .../test_v2_privilege_reconciliation.py | 37 ++++++++++- .../test_v2_simulation_role_privileges.py | 14 +++- tests/unit/v2/test_simulation_run_models.py | 6 +- 11 files changed, 114 insertions(+), 45 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index e148d4d5a..d69b74ba2 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -565,7 +565,9 @@ models made each result row tell apart what it measures: - `poverty.filter_group` names the group a filter selects (`child` for `age`), and joins the table's unique key. -Missing values in those keys count as equal (`NULLS NOT DISTINCT`). It has no +Missing values in those keys count as equal (`NULLS NOT DISTINCT`, which +PostgreSQL supports from version 15; `poverty` and `program_statistics` use it +from the earlier revisions too). It has no post-generation corrections and needs no table empty; its downgrade fails while a report run holds comparison rows of more than one pair, or a run holds rows that differ only in a removed column. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 98a079179..6e2a75655 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -244,9 +244,9 @@ flowchart LR | `api.migration.bundle-supported-by-simulation-api` | The installed PolicyEngine bundle is the latest one the simulation API has published, checked on every pull request and before each staging and production deploy. | mechanical | enforced | yes | CI job `ensure-policyengine-bundle-supported-by-simulation-api` | `pr.yml:ensure-policyengine-bundle-supported-by-simulation-api` | — | | `api.migration.contracts-consistent` | Migration route groups and workflow contracts are internally consistent, and their generated docs are current. | mechanical | enforced | yes | quality guard `migration-contracts` | `pr.yml:quality-guards` | **unproven** | | `api.migration.cross-database-integration` | Cross-database integration tests pass against real MySQL and Postgres. | mechanical | enforced | yes | CI job `v2-integration-check` | `pr.yml:v2-integration-check` | — | -| `api.v2.deployed-privileges-match` | Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing. | mechanical | enforced | yes | CI job `seed`, `catalog-runtime-integration` | `seed-v2-database.yml:seed`, `v2-integration-check.yml:catalog-runtime-integration` | yes | +| `api.v2.deployed-privileges-match` | Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead. | mechanical | enforced | yes | CI job `seed`, `catalog-runtime-integration` | `seed-v2-database.yml:seed`, `v2-integration-check.yml:catalog-runtime-integration` | yes | | `api.v2.runtime-privileges-declared` | Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py, and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it. | mechanical | enforced | yes | pytest `tests/unit/v2/test_privilege_manifest.py` | `pr.yml:test` | yes | -| `api.v2.runtime-privileges-tested-as-runtime-role` | The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, and the comparison-run services as policyengine_stage12_, holding only their declared privileges, so a statement those privileges don't allow fails CI. | mechanical | enforced | yes | CI job `v2-integration-check`, `catalog-runtime-integration` | `pr.yml:v2-integration-check`, `v2-integration-check.yml:catalog-runtime-integration` | — | +| `api.v2.runtime-privileges-tested-as-runtime-role` | The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, the comparison-run services as policyengine_stage12_, and the simulation API's table writes as policyengine_v2_simulation_service and policyengine_v2_modal_worker, holding only their declared privileges, so a statement those privileges don't allow fails CI. | mechanical | enforced | yes | CI job `v2-integration-check`, `catalog-runtime-integration` | `pr.yml:v2-integration-check`, `v2-integration-check.yml:catalog-runtime-integration` | — | ```mermaid flowchart LR diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 85bb805d2..351dff625 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -94,8 +94,8 @@ class HouseholdJob(TimestampedModel, table=True): ``policy_id`` is null for current law. ``request_data`` holds the calculation's options, such as axes. The job records the catalog version - and .py version it ran on, its Modal call, and its status and result or - error. Its household, policy, dynamic and catalog version are logical + and PolicyEngine.py version it ran on, its Modal call, and its status and + result or error. Its household, policy, dynamic and catalog version are logical references to the Public API's rows, without foreign keys: the simulation service checks them before it writes the job. """ diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 4a2689a3a..5a2e31322 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -93,7 +93,7 @@ class Simulation(TimestampedModel, table=True): class SimulationRun(TimestampedModel, table=True): - """One run of a simulation on one catalog version and .py version. + """One run of a simulation on one catalog version and PolicyEngine.py version. ``tax_benefit_model_version_id``, ``dataset_id`` and ``region_id`` name the catalog version, dataset and region the run used, which a run on a newer diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index 8213a23a4..3a6f7f04a 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -486,6 +486,45 @@ def _write_summary(environ: Mapping[str, str], summary: str) -> None: summary_file.write(summary) +def _verify( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + roles_awaiting_creation: Collection[str], + environment: str, + environ: Mapping[str, str], +) -> int: + """Check every declared privilege, skipping roles not created yet. + + The roles skipped and the privileges missing are written to stdout and the + job summary under one heading. + """ + skipped = skipped_roles(connection, declared, roles_awaiting_creation) + checked = {role: tables for role, tables in declared.items() if role not in skipped} + failures = check_declared_privileges(connection, checked) + heading = f"## v2 table privileges ({environment})\n\n" + skipped_note = f"{_skipped_roles_line(skipped)}\n\n" if skipped else "" + if failures: + rows = "\n".join( + f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" + for failure in failures + ) + _write_summary( + environ, + f"{heading}{skipped_note}" + f"Declared privileges missing after reconciling:\n\n{rows}\n", + ) + return 1 + if skipped: + _write_summary( + environ, + f"{heading}{skipped_note}" + "Every other declared v2 table privilege is held.\n", + ) + return 0 + print(f"Every declared v2 table privilege is held ({environment}).") + return 0 + + def _build_engine(settings: V2AlembicSettings) -> Engine: return create_engine(settings.url, poolclass=NullPool) @@ -511,30 +550,7 @@ def run( with engine.begin() as connection: qualify_v2_connection(connection, settings) if mode == "verify": - skipped = skipped_roles(connection, declared, awaiting) - if skipped: - _write_summary( - values, - f"## v2 table privileges ({environment})\n\n" - f"{_skipped_roles_line(skipped)}\n", - ) - failures = check_declared_privileges( - connection, declared, roles_awaiting_creation=awaiting - ) - if failures: - rows = "\n".join( - f"- `{failure.role}` lacks {failure.privilege} on " - f"`{failure.table}`" - for failure in failures - ) - _write_summary( - values, - f"## v2 table privileges ({environment})\n\n" - f"Declared privileges missing after reconciling:\n\n{rows}\n", - ) - return 1 - print(f"Every declared v2 table privilege is held ({environment}).") - return 0 + return _verify(connection, declared, awaiting, environment, values) if create_roles: created = create_missing_roles(connection, declared) if created: diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 14685f150..423904534 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -1,4 +1,9 @@ -"""Durable report-definition and report-run operations for API v2-alpha.""" +"""Durable report-definition and report-run operations for API v2-alpha. + +No managed role can run these operations: they lock ``reports`` rows with +``SELECT ... FOR UPDATE``, which needs ``UPDATE``, and only the simulation API's +roles write report runs, with their own code. +""" from __future__ import annotations diff --git a/rules/registry.toml b/rules/registry.toml index 04f5be4bf..03573dcca 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -711,9 +711,9 @@ test = "tests/unit/v2/test_privilege_manifest.py::test_rule_proof_runtime_privil [[rule]] id = "api.v2.runtime-privileges-tested-as-runtime-role" category = "migration" -statement = "The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, and the comparison-run services as policyengine_stage12_, holding only their declared privileges, so a statement those privileges don't allow fails CI." +statement = "The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, the comparison-run services as policyengine_stage12_, and the simulation API's table writes as policyengine_v2_simulation_service and policyengine_v2_modal_worker, holding only their declared privileges, so a statement those privileges don't allow fails CI." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" -scope = ["tests/integration/conftest.py", "tests/integration/test_v2_user_identity_persistence.py", "tests/integration/test_v2_stage12_comparison_persistence.py", "tests/integration/test_v2_metadata_routes.py", "tests/integration/test_v2_catalog_publication.py"] +scope = ["tests/integration/conftest.py", "tests/integration/test_v2_user_identity_persistence.py", "tests/integration/test_v2_stage12_comparison_persistence.py", "tests/integration/test_v2_metadata_routes.py", "tests/integration/test_v2_catalog_publication.py", "tests/integration/test_v2_simulation_role_privileges.py"] tier = "mechanical" status = "enforced" autofix = false @@ -731,7 +731,7 @@ invocation = "reconcile_v2_privileges.py apply --create-missing-roles" [[rule]] id = "api.v2.deployed-privileges-match" category = "migration" -statement = "Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing." +statement = "Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] tier = "mechanical" diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 337fd2ae2..d941d43ee 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,6 +1,6 @@ """Check the simulation runs revisions against disposable Postgres. -Five revisions follow the area impact coordinate removal. The first adds +Five revisions follow the parameter tree revision. The first adds simulation runs and the table naming the simulation runs a report run combines, keys every result table by run, and gives household jobs one household, one year and one policy. The second drops the foreign keys from @@ -8,9 +8,10 @@ fields that had no column. The fourth makes each result row tell apart its pair of simulation runs, decile grouping and poverty group. The fifth moves the parts of a comparison without a result table from the report run to one -row per compared pair. None changes any other table, so the previous Public -API revision, whose roles can't read or write any table they change, keeps -working while they are deployed. +row per compared pair. None changes any other table. The previous Public API +revision keeps working while they are deployed: its roles can't read or write +any table they change, except the catalog writer, which counts the rows of +simulations, reports and report_runs. """ from __future__ import annotations @@ -159,7 +160,7 @@ def _enum_type_names(engine: Engine) -> set[str]: ) -def test_the_revisions_add_two_tables_and_change_only_simulation_side_tables( +def test_the_revisions_add_three_tables_and_change_only_simulation_side_tables( config: Config, engine: Engine ) -> None: command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) diff --git a/tests/integration/test_v2_privilege_reconciliation.py b/tests/integration/test_v2_privilege_reconciliation.py index 0a764e0ce..7916907c1 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -193,6 +193,41 @@ def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( assert (apply_status, verify_status) == (0, 0) assert f"`{absent_role}`" in apply_report assert absent_role in verify_output - assert "Every declared v2 table privilege is held" in verify_output + assert "Every other declared v2 table privilege is held." in verify_output + assert "Every declared v2 table privilege is held" not in verify_output # Each step names the skipped role in the job summary. assert summary_path.read_text(encoding="utf-8").count(f"`{absent_role}`") == 2 + + +def test_verify_fails_once_a_role_awaiting_creation_exists_without_its_privileges( + disposable_v2_database_url, temporary_role, monkeypatch, capsys, tmp_path +) -> None: + """A role awaiting creation is checked as soon as it exists.""" + role = temporary_role() + absent_role = f"policyengine_privilege_test_absent_{uuid4().hex[:12]}" + monkeypatch.setattr( + privilege_reconciliation, + "declared_privileges", + lambda _environment: {role: {"regions": READ}, absent_role: {"regions": READ}}, + ) + monkeypatch.setattr( + privilege_reconciliation, + "ROLES_AWAITING_PROVISIONING", + frozenset({role, absent_role}), + ) + summary_path = tmp_path / "summary.md" + environ = { + V2_MIGRATION_DATABASE_URL: disposable_v2_database_url, + V2_ALEMBIC_DISPOSABLE_TEST: "1", + DEPLOYMENT_ENVIRONMENT: "test", + "GITHUB_STEP_SUMMARY": str(summary_path), + } + + verify_status = run("verify", environ=environ) + capsys.readouterr() + summary = summary_path.read_text(encoding="utf-8") + + assert verify_status == 1 + assert f"`{role}` lacks SELECT on `regions`" in summary + assert f"`{absent_role}`" in summary + assert summary.count("## v2 table privileges") == 1 diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index 7b2239667..3f0de7734 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -61,7 +61,10 @@ def _refused(connection: Connection, statement: str, **values: object) -> None: def _create_simulation_rows(connection: Connection) -> dict[str, object]: """Create what the simulation API creates for one economy report run.""" - ids = {name: uuid4() for name in ("simulation", "run", "report", "report_run")} + ids = { + name: uuid4() + for name in ("simulation", "run", "report", "report_run", "household_job") + } insert_row( connection, "simulations", @@ -109,7 +112,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: insert_row( connection, "household_jobs", - id=uuid4(), + id=ids["household_job"], country="us", household_id=uuid4(), year=2026, @@ -194,6 +197,13 @@ def test_modal_workers_write_their_run_status_and_results_only( reform_simulation_run_id=ids["run"], cliff_impact_result="{}", ) + connection.execute( + text( + "UPDATE household_jobs SET status = 'succeeded', result = '{}', " + "completed_at = now() WHERE id = :id AND status = 'pending'" + ), + {"id": ids["household_job"]}, + ) connection.execute(text("SELECT count(*) FROM parameter_values")).scalar_one() assert updated == ids["run"] diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index b6eb45749..42027a4c1 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -1,9 +1,9 @@ """Structural tests for simulation runs, report runs, household jobs and results. A simulation defines what to calculate and each calculation of it is a run that -records the catalog and .py versions it used. A report run names the simulation -runs it combines, results are keyed by run, and a household job calculates one -household for one year under one policy. +records the catalog and PolicyEngine.py versions it used. A report run names +the simulation runs it combines, results are keyed by run, and a household job +calculates one household for one year under one policy. """ import sqlalchemy as sa From 23bfc02e549ae696434c4f01dfb449900d314868 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 02:22:41 +0400 Subject: [PATCH 10/22] Decide the v2 privilege check's skipped roles in one place check_declared_privileges no longer takes the roles awaiting creation: verify asks skipped_roles once and checks the other declared roles. The report-run module's note names the two operations that lock reports rows, and the tests check that the worker's household job update changed its row and that verify names the skipped role as skipped. Co-Authored-By: Claude Opus 5.5 --- .../data/v2/privilege_reconciliation.py | 18 +++++++----------- policyengine_api/data/v2/report_runs.py | 7 ++++--- .../test_v2_privilege_reconciliation.py | 11 +++++++++-- .../test_v2_simulation_role_privileges.py | 8 +++++--- 4 files changed, 25 insertions(+), 19 deletions(-) diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index 3a6f7f04a..c449c93cf 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -427,18 +427,12 @@ def skipped_roles( def check_declared_privileges( connection: Connection, declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], - *, - roles_awaiting_creation: Collection[str] = (), ) -> list[PrivilegeEntry]: """Return every declared privilege a role doesn't effectively hold. - A missing role or table counts as missing every privilege declared for it, - except a missing role awaiting creation, which is skipped. + A missing role or table counts as missing every privilege declared for it; + leave out the roles ``skipped_roles`` returns to skip them. """ - skipped = set(skipped_roles(connection, declared, roles_awaiting_creation)) - declared = { - role: tables for role, tables in declared.items() if role not in skipped - } existing_roles = read_existing_roles(connection, list(declared)) existing_tables = read_public_tables(connection) @@ -493,10 +487,12 @@ def _verify( environment: str, environ: Mapping[str, str], ) -> int: - """Check every declared privilege, skipping roles not created yet. + """Check every declared privilege, skipping missing roles awaiting creation. - The roles skipped and the privileges missing are written to stdout and the - job summary under one heading. + A missing role in ``roles_awaiting_creation`` is skipped; any other missing + role counts as missing every privilege declared for it. The roles skipped + and the privileges missing are written to stdout and the job summary under + one heading. """ skipped = skipped_roles(connection, declared, roles_awaiting_creation) checked = {role: tables for role, tables in declared.items() if role not in skipped} diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 423904534..1b1debd61 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -1,8 +1,9 @@ """Durable report-definition and report-run operations for API v2-alpha. -No managed role can run these operations: they lock ``reports`` rows with -``SELECT ... FOR UPDATE``, which needs ``UPDATE``, and only the simulation API's -roles write report runs, with their own code. +``set_report_type`` and ``create_report_run`` lock ``reports`` rows with +``SELECT ... FOR UPDATE``, which needs ``UPDATE`` on ``reports``; no managed +role holds it, so none can run them. The simulation API creates reports and +report runs with its own code. """ from __future__ import annotations diff --git a/tests/integration/test_v2_privilege_reconciliation.py b/tests/integration/test_v2_privilege_reconciliation.py index 7916907c1..3edc1cdd1 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -23,6 +23,7 @@ apply_declared_privileges, check_declared_privileges, run, + skipped_roles, ) from policyengine_api.data.v2.privileges import READ, TablePrivilege from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL @@ -150,13 +151,16 @@ def test_a_role_awaiting_creation_is_skipped_until_it_exists( plan, _defaults = apply_declared_privileges( connection, declared, roles_awaiting_creation={absent_role} ) + skipped = skipped_roles(connection, declared, {absent_role}) missing = check_declared_privileges( - connection, declared, roles_awaiting_creation={absent_role} + connection, + {name: tables for name, tables in declared.items() if name not in skipped}, ) assert plan.can_apply assert plan.skipped_roles == (absent_role,) assert plan.grants == (PrivilegeEntry(role, "regions", "SELECT"),) + assert skipped == (absent_role,) assert missing == [] assert _has_privilege(superuser_engine, role, "regions", "SELECT") @@ -229,5 +233,8 @@ def test_verify_fails_once_a_role_awaiting_creation_exists_without_its_privilege assert verify_status == 1 assert f"`{role}` lacks SELECT on `regions`" in summary - assert f"`{absent_role}`" in summary + assert ( + "Declared roles not created yet, skipped until an operator creates them: " + f"`{absent_role}`" + ) in summary assert summary.count("## v2 table privileges") == 1 diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index 3f0de7734..c47d5c0f6 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -197,16 +197,18 @@ def test_modal_workers_write_their_run_status_and_results_only( reform_simulation_run_id=ids["run"], cliff_impact_result="{}", ) - connection.execute( + finished_job = connection.execute( text( "UPDATE household_jobs SET status = 'succeeded', result = '{}', " - "completed_at = now() WHERE id = :id AND status = 'pending'" + "completed_at = now() WHERE id = :id AND status = 'pending' " + "RETURNING id" ), {"id": ids["household_job"]}, - ) + ).scalar_one() connection.execute(text("SELECT count(*) FROM parameter_values")).scalar_one() assert updated == ids["run"] + assert finished_job == ids["household_job"] _refused( connection, "INSERT INTO simulations (id, simulation_type, dataset_id, " From d674032a25b4f764ce8bd73095de39813562591f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 02:31:03 +0400 Subject: [PATCH 11/22] Count the three new v2 tables in the import side-effect test The test pins the number of tables the v2 metadata registers; this change adds simulation_runs, report_run_simulation_runs and report_run_comparisons, so it is 45. Co-Authored-By: Claude Opus 5.5 --- tests/unit/v2/test_import_side_effects.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index a3010eff6..df7828892 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 42 +assert len(V2_METADATA.tables) == 45 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules From aa2875f407e57456969ac3fcb4ba6bb142b7fd35 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:45:52 +0400 Subject: [PATCH 12/22] Replace the five v2 simulation runs revisions with one The five revisions after c614b6bcccfc are deleted and revision 1356f658347d is autogenerated from that revision against the same models, so the schema at head is unchanged. Its two corrections are the ones the first of the five had: simulation_runs.status reuses v2_simulation_status, and the downgrade drops v2_simulation_role. The per-revision tests become one test of the revision's source and one downgrade-and-upgrade test. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 103 +-- ...ecord_simulation_runs_and_key_results_.py} | 697 +++++++++++++++--- ...ell_apart_result_rows_of_each_run_pair_.py | 248 ------- ...029a_drop_foreign_keys_between_the_two_.py | 196 ----- ...tore_the_economy_output_fields_without_.py | 174 ----- ...c_keep_untabled_comparison_results_per_.py | 120 --- policyengine_api/data/v2/migration_target.py | 2 +- .../test_alembic_v2_simulation_runs.py | 72 +- tests/unit/v2/test_alembic_v2.py | 290 +------- 9 files changed, 675 insertions(+), 1227 deletions(-) rename migrations/v2/versions/{2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py => 1356f658347d_record_simulation_runs_and_key_results_.py} (66%) delete mode 100644 migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py delete mode 100644 migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py delete mode 100644 migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py delete mode 100644 migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d3a3d0660..94ac98f65 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -499,7 +499,7 @@ table that has rows, so the downgrade works only while both tables are empty. ### API v2 simulation runs -Revision `2ffd50f47a2b` (parent `c614b6bcccfc`) was autogenerated after these +Revision `1356f658347d` (parent `c614b6bcccfc`) was autogenerated after these model changes: - `simulations` keeps only what to calculate. Its execution columns @@ -509,16 +509,37 @@ model changes: used, and its Modal call ID; `output_dataset_id` is replaced by the run's output file (`output_uri`, `output_sha256`). - The new `report_run_simulation_runs` names the simulation runs each report - run combines, as its baseline or its reform. + run combines, as its baseline or its reform, and the new + `report_run_comparisons` keeps, per compared pair, the parts of an economy + comparison that have no result table (`labor_supply_result`, + `cliff_impact_result`, `spm_provenance`, `data_version`). - `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); the other result tables belong to a report run and name the two simulation - runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`). - `aggregates` and `change_aggregates` are unchanged. -- `report_runs` gains its Modal call ID and two JSON results that have no - table (`labor_supply_result`, `cliff_impact_result`). + runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`), + which join each table's unique key. `aggregates` and `change_aggregates` are + unchanged. +- The result tables gain the economy output's fields that had no column: the + winner, loser and no-change shares of `congressional_district_impacts`, and + the UK disclosure control fields (`disclosure_tier`, checked to be + `reported` or `approximate`, and the four `_lower`/`_upper` interval bounds) + of `constituency_impacts` and `local_authority_impacts`. The detailed budget + names neither a program's entity nor whether it is a tax, so + `program_statistics.entity` and `is_tax` become nullable. + `decile_impacts.decile_variable` names the variable that groups households + when it is not the income variable's own deciles (for example + `household_wealth_decile`), and `poverty.filter_group` names the group a + filter selects (`child` for `age`); both join their table's unique key. +- `report_runs` gains its Modal call ID. - `household_jobs` calculates one household (`household_id`) for one `year` under one policy, and records its catalog and package versions and its Modal call ID. +- The foreign keys from the tables the simulation service and its Modal workers + write to the tables the Public API writes are dropped; the service that + writes a row checks the rows it names. + +Missing values in the unique keys of `poverty`, `decile_impacts` and +`program_statistics` count as equal (`NULLS NOT DISTINCT`, which PostgreSQL +supports from version 15). It has two post-generation corrections, each commented in the revision: @@ -540,76 +561,8 @@ none of its roles can read or write a table the revision changes, except the catalog writer, which only counts the rows of `simulations`, `reports` and `report_runs`. `tests/integration/test_alembic_v2_simulation_runs.py` checks that the -revision changes no other table, and the v2 lifecycle test covers the -downgrade and re-upgrade. +revision changes no other table, and downgrades and upgrades it again. The name `simulation_runs` is also a v1 Cloud SQL table. The v2 metadata validation (`policyengine_api/data/v2/metadata_validation.py`) now treats it as a name the two databases share, as it does `simulations`. - -### API v2 foreign keys between the two services' tables - -Revision `7a83a87d029a` (parent `2ffd50f47a2b`) was autogenerated after the -models stopped declaring foreign keys from the tables the simulation service -and its Modal workers write (`simulations`, `reports`, `household_jobs`) to the -tables the Public API writes (`users`, `policies`, `dynamics`, `households`, -and the catalog's `tax_benefit_models`, `tax_benefit_model_versions`, -`datasets` and `regions`). The 14 columns stay as plain UUIDs: the service that -writes a row checks the rows it names. Foreign keys between tables of the same -service stay. Deleting a Public API row no longer clears or blocks the rows -that name it: in particular, deleting a user leaves that user's ID in -`reports.user_id`, which no managed role may update; the owner of a saved -report is the user association, which the deletion removes. It has no post-generation corrections. Its downgrade restores the -keys, so it fails while a column names a row the referenced table lacks. -`tests/unit/v2/test_simulation_run_models.py` checks that no foreign key or ORM -relationship joins the two services' tables. - -### API v2 economy output fields - -Revision `7f43d2c7c3c6` (parent `7a83a87d029a`) was autogenerated after the -models gained the fields of the simulation API's economy output that had no -column: `report_runs.spm_provenance` (JSON) and `data_version`; the winner, -loser and no-change shares of `congressional_district_impacts`; and the UK -disclosure control fields (`disclosure_tier`, checked to be `reported` or -`approximate`, and the four `_lower`/`_upper` interval bounds) of -`constituency_impacts` and `local_authority_impacts`. The output's detailed -budget names neither a program's entity nor whether it is a tax, so -`program_statistics.entity` and `is_tax` become nullable, and its unique key -treats a missing entity as equal (`NULLS NOT DISTINCT`). It has no -post-generation corrections. The three share columns are required, so the -upgrade needs `congressional_district_impacts` empty; the downgrade needs -every `program_statistics` row to have an entity and a tax flag. - -### API v2 result row identity - -Revision `2e3de7c572ed` (parent `7f43d2c7c3c6`) was autogenerated after the -models made each result row tell apart what it measures: - -- the comparison result tables' unique keys include the pair of simulation - runs (`baseline_simulation_run_id`, `reform_simulation_run_id`) after - `report_run_id`, because a budget-window report run compares one pair per - year; -- `decile_impacts.decile_variable` names the variable that groups households - when it is not the income variable's own deciles (for example - `household_wealth_decile`), and joins the table's unique key; -- `poverty.filter_group` names the group a filter selects (`child` for - `age`), and joins the table's unique key. - -Missing values in those keys count as equal (`NULLS NOT DISTINCT`, which -PostgreSQL supports from version 15; `poverty` and `program_statistics` use it -from the earlier revisions too). It has no -post-generation corrections and needs no table empty; its downgrade fails -while a report run holds comparison rows of more than one pair, or a run holds -rows that differ only in a removed column. - -### API v2 comparisons - -Revision `dce0caa329cc` (parent `2e3de7c572ed`) was autogenerated after the -models moved the parts of a comparison that have no result table -(`labor_supply_result`, `cliff_impact_result`, `spm_provenance`, -`data_version`) from `report_runs` to the new `report_run_comparisons`, one row -per report run and compared pair of simulation runs, because a budget-window -run compares one pair per year. It has no post-generation corrections. The -four `report_runs` columns it drops were added by the earlier revisions of -this change and never written; its downgrade drops the new table and its -rows. diff --git a/migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py similarity index 66% rename from migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py rename to migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 3c2c663f3..2f1ebefca 100644 --- a/migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -1,13 +1,13 @@ -"""add simulation runs and key results by run +"""record simulation runs and key results by run -Revision ID: 2ffd50f47a2b +Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-04 00:09:14.105644 +Create Date: 2026-11-05 23:44:44.321856 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate Two post-generation corrections, each commented where it is made: simulation_runs.status reuses the existing v2_simulation_status type, and the -downgrade drops the v2_simulation_role type it no longer uses. +downgrade drops the v2_simulation_role type the upgrade creates. The upgrade adds required columns without server defaults to household_jobs and to the result tables, so it needs those tables empty, as they are while no @@ -22,7 +22,7 @@ import sqlmodel from sqlalchemy.dialects import postgresql -revision: str = "2ffd50f47a2b" +revision: str = "1356f658347d" down_revision: Union[str, None] = "c614b6bcccfc" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -126,6 +126,54 @@ def upgrade() -> None: ["status", "started_at"], unique=False, ) + op.create_table( + "report_run_comparisons", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_run_id", sa.Uuid(), nullable=False), + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("labor_supply_result", sa.JSON(), nullable=True), + sa.Column("cliff_impact_result", sa.JSON(), nullable=True), + sa.Column("spm_provenance", sa.JSON(), nullable=True), + sa.Column( + "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True + ), + sa.ForeignKeyConstraint( + ["baseline_simulation_run_id"], + ["simulation_runs.id"], + name=op.f( + "fk_report_run_comparisons_baseline_simulation_run_id_simulation_runs" + ), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["reform_simulation_run_id"], + ["simulation_runs.id"], + name=op.f( + "fk_report_run_comparisons_reform_simulation_run_id_simulation_runs" + ), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_report_run_comparisons_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_comparisons")), + sa.UniqueConstraint( + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + name="uq_report_run_comparisons_run_pair", + ), + ) op.create_table( "report_run_simulation_runs", sa.Column("id", sa.Uuid(), nullable=False), @@ -179,15 +227,29 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("uq_budget_summary_run_variable_entity"), "budget_summary", type_="unique" + ) + op.create_unique_constraint( + "uq_budget_summary_run_variable_entity", "budget_summary", - type_="foreignkey", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "variable_name", + "entity", + ], ) op.drop_constraint( op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) + op.drop_constraint( + op.f("fk_budget_summary_baseline_simulation_id_simulations"), + "budget_summary", + type_="foreignkey", + ) op.create_foreign_key( op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", @@ -214,6 +276,33 @@ def upgrade() -> None: "congressional_district_impacts", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) + op.add_column( + "congressional_district_impacts", + sa.Column("winner_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("loser_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("no_change_percentage", sa.Float(), nullable=False), + ) + op.drop_constraint( + op.f("uq_congressional_district_impacts_run_geoid"), + "congressional_district_impacts", + type_="unique", + ) + op.create_unique_constraint( + "uq_congressional_district_impacts_run_geoid", + "congressional_district_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "district_geoid", + ], + ) op.drop_constraint( op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", @@ -226,21 +315,21 @@ def upgrade() -> None: ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -254,6 +343,43 @@ def upgrade() -> None: "constituency_impacts", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) + op.add_column( + "constituency_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) + op.drop_constraint( + op.f("uq_constituency_impacts_run_code"), "constituency_impacts", type_="unique" + ) + op.create_unique_constraint( + "uq_constituency_impacts_run_code", + "constituency_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "constituency_code", + ], + ) op.drop_constraint( op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", @@ -280,6 +406,11 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_check_constraint( + op.f("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) op.drop_column("constituency_impacts", "reform_simulation_id") op.drop_column("constituency_impacts", "baseline_simulation_id") op.add_column( @@ -290,29 +421,55 @@ def upgrade() -> None: "decile_impacts", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) + op.add_column( + "decile_impacts", + sa.Column( + "decile_variable", + sqlmodel.sql.sqltypes.AutoString(length=512), + nullable=True, + ), + ) op.drop_constraint( - op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + op.f("uq_decile_impacts_run_measure"), "decile_impacts", type_="unique" + ) + op.create_unique_constraint( + "uq_decile_impacts_run_measure", "decile_impacts", - type_="foreignkey", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "income_variable", + "decile_variable", + "entity", + "decile", + "quantiles", + ], + postgresql_nulls_not_distinct=True, ) op.drop_constraint( op.f("fk_decile_impacts_reform_simulation_id_simulations"), "decile_impacts", type_="foreignkey", ) + op.drop_constraint( + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + "decile_impacts", + type_="foreignkey", + ) op.create_foreign_key( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -363,6 +520,16 @@ def upgrade() -> None: ["household_id"], unique=False, ) + op.drop_constraint( + op.f("fk_household_jobs_policy_id_policies"), + "household_jobs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_household_jobs_dynamic_id_dynamics"), + "household_jobs", + type_="foreignkey", + ) op.create_check_constraint( op.f("ck_household_jobs_terminal_completion"), "household_jobs", @@ -383,12 +550,12 @@ def upgrade() -> None: ["simulation_run_id", "income_variable", "entity"], ) op.drop_constraint( - op.f("fk_inequality_simulation_id_simulations"), + op.f("fk_inequality_report_run_id_report_runs"), "inequality", type_="foreignkey", ) op.drop_constraint( - op.f("fk_inequality_report_run_id_report_runs"), + op.f("fk_inequality_simulation_id_simulations"), "inequality", type_="foreignkey", ) @@ -411,15 +578,31 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("uq_intra_decile_impacts_run_type_decile"), "intra_decile_impacts", - type_="foreignkey", + type_="unique", + ) + op.create_unique_constraint( + "uq_intra_decile_impacts_run_type_decile", + "intra_decile_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "decile_type", + "decile", + ], ) op.drop_constraint( op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + "intra_decile_impacts", + type_="foreignkey", + ) op.create_foreign_key( op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", @@ -446,16 +629,55 @@ def upgrade() -> None: "local_authority_impacts", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) + op.add_column( + "local_authority_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("uq_local_authority_impacts_run_code"), "local_authority_impacts", - type_="foreignkey", + type_="unique", + ) + op.create_unique_constraint( + "uq_local_authority_impacts_run_code", + "local_authority_impacts", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "local_authority_code", + ], ) op.drop_constraint( op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) + op.drop_constraint( + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + "local_authority_impacts", + type_="foreignkey", + ) op.create_foreign_key( op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", @@ -472,21 +694,38 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_check_constraint( + op.f("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) op.drop_column("local_authority_impacts", "reform_simulation_id") op.drop_column("local_authority_impacts", "baseline_simulation_id") op.add_column("poverty", sa.Column("simulation_run_id", sa.Uuid(), nullable=False)) + op.add_column( + "poverty", + sa.Column( + "filter_group", sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True + ), + ) op.drop_index(op.f("ix_poverty_run_simulation_type"), table_name="poverty") op.create_unique_constraint( "uq_poverty_simulation_run_measure", "poverty", - ["simulation_run_id", "poverty_type", "entity", "filter_variable"], + [ + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + ], postgresql_nulls_not_distinct=True, ) op.drop_constraint( - op.f("fk_poverty_report_run_id_report_runs"), "poverty", type_="foreignkey" + op.f("fk_poverty_simulation_id_simulations"), "poverty", type_="foreignkey" ) op.drop_constraint( - op.f("fk_poverty_simulation_id_simulations"), "poverty", type_="foreignkey" + op.f("fk_poverty_report_run_id_report_runs"), "poverty", type_="foreignkey" ) op.create_foreign_key( op.f("fk_poverty_simulation_run_id_simulation_runs"), @@ -506,16 +745,42 @@ def upgrade() -> None: "program_statistics", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) + op.alter_column( + "program_statistics", + "entity", + existing_type=sa.VARCHAR(length=128), + nullable=True, + ) + op.alter_column( + "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=True + ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("uq_program_statistics_run_program_entity"), "program_statistics", - type_="foreignkey", + type_="unique", + ) + op.create_unique_constraint( + "uq_program_statistics_run_program_entity", + "program_statistics", + [ + "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "program_name", + "entity", + ], + postgresql_nulls_not_distinct=True, ) op.drop_constraint( op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", type_="foreignkey", ) + op.drop_constraint( + op.f("fk_program_statistics_baseline_simulation_id_simulations"), + "program_statistics", + type_="foreignkey", + ) op.create_foreign_key( op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", @@ -540,28 +805,62 @@ def upgrade() -> None: "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True ), ) - op.add_column( - "report_runs", sa.Column("labor_supply_result", sa.JSON(), nullable=True) - ) - op.add_column( - "report_runs", sa.Column("cliff_impact_result", sa.JSON(), nullable=True) - ) op.create_index( "ix_report_runs_status_started_at", "report_runs", ["status", "started_at"], unique=False, ) + op.drop_constraint( + op.f("fk_reports_region_id_regions"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_reports_policy_id_policies"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_reports_dataset_id_datasets"), "reports", type_="foreignkey" + ) + op.drop_constraint(op.f("fk_reports_user_id_users"), "reports", type_="foreignkey") + op.drop_constraint( + op.f("fk_reports_household_id_households"), "reports", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), + "reports", + type_="foreignkey", + ) op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") + op.drop_constraint( + op.f("fk_simulations_dataset_id_datasets"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_dynamic_id_dynamics"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_policy_id_policies"), "simulations", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), + "simulations", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_simulations_household_id_households"), + "simulations", + type_="foreignkey", + ) op.drop_constraint( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", type_="foreignkey", ) - op.drop_column("simulations", "household_result") - op.drop_column("simulations", "completed_at") - op.drop_column("simulations", "output_dataset_id") + op.drop_constraint( + op.f("fk_simulations_region_id_regions"), "simulations", type_="foreignkey" + ) op.drop_column("simulations", "error_message") + op.drop_column("simulations", "output_dataset_id") + op.drop_column("simulations", "completed_at") + op.drop_column("simulations", "household_result") op.drop_column("simulations", "started_at") op.drop_column("simulations", "status") # ### end Alembic commands ### @@ -591,11 +890,12 @@ def downgrade() -> None: ) op.add_column( "simulations", - sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), - ) - op.add_column( - "simulations", - sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + sa.Column( + "household_result", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), ) op.add_column( "simulations", @@ -608,12 +908,19 @@ def downgrade() -> None: ) op.add_column( "simulations", - sa.Column( - "household_result", - postgresql.JSON(astext_type=sa.Text()), - autoincrement=False, - nullable=True, - ), + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + ) + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + ) + op.create_foreign_key( + op.f("fk_simulations_region_id_regions"), + "simulations", + "regions", + ["region_id"], + ["id"], + ondelete="SET NULL", ) op.create_foreign_key( op.f("fk_simulations_output_dataset_id_datasets"), @@ -623,15 +930,101 @@ def downgrade() -> None: ["id"], ondelete="SET NULL", ) + op.create_foreign_key( + op.f("fk_simulations_household_id_households"), + "simulations", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), + "simulations", + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_simulations_policy_id_policies"), + "simulations", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_simulations_dynamic_id_dynamics"), + "simulations", + "dynamics", + ["dynamic_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_simulations_dataset_id_datasets"), + "simulations", + "datasets", + ["dataset_id"], + ["id"], + ondelete="RESTRICT", + ) op.create_index( op.f("ix_simulations_status_created_at"), "simulations", ["status", "created_at"], unique=False, ) + op.create_foreign_key( + op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), + "reports", + "tax_benefit_models", + ["tax_benefit_model_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_reports_household_id_households"), + "reports", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_reports_user_id_users"), + "reports", + "users", + ["user_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_dataset_id_datasets"), + "reports", + "datasets", + ["dataset_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_policy_id_policies"), + "reports", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_reports_region_id_regions"), + "reports", + "regions", + ["region_id"], + ["id"], + ondelete="SET NULL", + ) op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") - op.drop_column("report_runs", "cliff_impact_result") - op.drop_column("report_runs", "labor_supply_result") op.drop_column("report_runs", "modal_call_id") op.add_column( "program_statistics", @@ -656,21 +1049,39 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_id_simulations"), + op.f("fk_program_statistics_baseline_simulation_id_simulations"), "program_statistics", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + ["report_run_id", "program_name", "entity"], + postgresql_nulls_not_distinct=False, + ) + op.alter_column( + "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=False + ) + op.alter_column( + "program_statistics", + "entity", + existing_type=sa.VARCHAR(length=128), + nullable=False, + ) op.drop_column("program_statistics", "reform_simulation_run_id") op.drop_column("program_statistics", "baseline_simulation_run_id") op.add_column( @@ -686,14 +1097,6 @@ def downgrade() -> None: "poverty", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_poverty_simulation_id_simulations"), - "poverty", - "simulations", - ["simulation_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_foreign_key( op.f("fk_poverty_report_run_id_report_runs"), "poverty", @@ -702,6 +1105,14 @@ def downgrade() -> None: ["id"], ondelete="CASCADE", ) + op.create_foreign_key( + op.f("fk_poverty_simulation_id_simulations"), + "poverty", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) op.drop_constraint("uq_poverty_simulation_run_measure", "poverty", type_="unique") op.create_index( op.f("ix_poverty_run_simulation_type"), @@ -709,6 +1120,7 @@ def downgrade() -> None: ["report_run_id", "simulation_id", "poverty_type"], unique=False, ) + op.drop_column("poverty", "filter_group") op.drop_column("poverty", "simulation_run_id") op.add_column( "local_authority_impacts", @@ -722,6 +1134,11 @@ def downgrade() -> None: "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) + op.drop_constraint( + op.f("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + type_="check", + ) op.drop_constraint( op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", @@ -733,21 +1150,35 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_local_authority_impacts_run_code", "local_authority_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_local_authority_impacts_run_code"), + "local_authority_impacts", + ["report_run_id", "local_authority_code"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("local_authority_impacts", "relative_household_income_change_upper") + op.drop_column("local_authority_impacts", "relative_household_income_change_lower") + op.drop_column("local_authority_impacts", "average_household_income_change_upper") + op.drop_column("local_authority_impacts", "average_household_income_change_lower") + op.drop_column("local_authority_impacts", "disclosure_tier") op.drop_column("local_authority_impacts", "reform_simulation_run_id") op.drop_column("local_authority_impacts", "baseline_simulation_run_id") op.add_column( @@ -773,21 +1204,32 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_intra_decile_impacts_run_type_decile", + "intra_decile_impacts", + type_="unique", + ) + op.create_unique_constraint( + op.f("uq_intra_decile_impacts_run_type_decile"), + "intra_decile_impacts", + ["report_run_id", "decile_type", "decile"], + postgresql_nulls_not_distinct=False, + ) op.drop_column("intra_decile_impacts", "reform_simulation_run_id") op.drop_column("intra_decile_impacts", "baseline_simulation_run_id") op.add_column( @@ -803,14 +1245,6 @@ def downgrade() -> None: "inequality", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_inequality_report_run_id_report_runs"), - "inequality", - "report_runs", - ["report_run_id"], - ["id"], - ondelete="CASCADE", - ) op.create_foreign_key( op.f("fk_inequality_simulation_id_simulations"), "inequality", @@ -819,6 +1253,14 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) op.drop_constraint( "uq_inequality_simulation_run_measure", "inequality", type_="unique" ) @@ -833,6 +1275,22 @@ def downgrade() -> None: op.drop_constraint( op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" ) + op.create_foreign_key( + op.f("fk_household_jobs_dynamic_id_dynamics"), + "household_jobs", + "dynamics", + ["dynamic_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_foreign_key( + op.f("fk_household_jobs_policy_id_policies"), + "household_jobs", + "policies", + ["policy_id"], + ["id"], + ondelete="SET NULL", + ) op.drop_index(op.f("ix_household_jobs_household_id"), table_name="household_jobs") op.drop_column("household_jobs", "modal_call_id") op.drop_column("household_jobs", "policyengine_version") @@ -854,31 +1312,41 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), "decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_decile_impacts_reform_simulation_id_simulations"), "decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_decile_impacts_run_measure", "decile_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_decile_impacts_run_measure"), + "decile_impacts", + ["report_run_id", "income_variable", "entity", "decile", "quantiles"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("decile_impacts", "decile_variable") op.drop_column("decile_impacts", "reform_simulation_run_id") op.drop_column("decile_impacts", "baseline_simulation_run_id") op.add_column( @@ -893,6 +1361,11 @@ def downgrade() -> None: "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) + op.drop_constraint( + op.f("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + type_="check", + ) op.drop_constraint( op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), "constituency_impacts", @@ -919,6 +1392,20 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_constituency_impacts_run_code", "constituency_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_constituency_impacts_run_code"), + "constituency_impacts", + ["report_run_id", "constituency_code"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("constituency_impacts", "relative_household_income_change_upper") + op.drop_column("constituency_impacts", "relative_household_income_change_lower") + op.drop_column("constituency_impacts", "average_household_income_change_upper") + op.drop_column("constituency_impacts", "average_household_income_change_lower") + op.drop_column("constituency_impacts", "disclosure_tier") op.drop_column("constituency_impacts", "reform_simulation_run_id") op.drop_column("constituency_impacts", "baseline_simulation_run_id") op.add_column( @@ -935,14 +1422,14 @@ def downgrade() -> None: ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", @@ -963,6 +1450,20 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_congressional_district_impacts_run_geoid", + "congressional_district_impacts", + type_="unique", + ) + op.create_unique_constraint( + op.f("uq_congressional_district_impacts_run_geoid"), + "congressional_district_impacts", + ["report_run_id", "district_geoid"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("congressional_district_impacts", "no_change_percentage") + op.drop_column("congressional_district_impacts", "loser_percentage") + op.drop_column("congressional_district_impacts", "winner_percentage") op.drop_column("congressional_district_impacts", "reform_simulation_run_id") op.drop_column("congressional_district_impacts", "baseline_simulation_run_id") op.add_column( @@ -988,21 +1489,30 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) + op.drop_constraint( + "uq_budget_summary_run_variable_entity", "budget_summary", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_budget_summary_run_variable_entity"), + "budget_summary", + ["report_run_id", "variable_name", "entity"], + postgresql_nulls_not_distinct=False, + ) op.drop_column("budget_summary", "reform_simulation_run_id") op.drop_column("budget_summary", "baseline_simulation_run_id") op.drop_index( @@ -1010,10 +1520,11 @@ def downgrade() -> None: table_name="report_run_simulation_runs", ) op.drop_table("report_run_simulation_runs") - op.drop_index("ix_simulation_runs_status_started_at", table_name="simulation_runs") - op.drop_index("ix_simulation_runs_current_output", table_name="simulation_runs") - op.drop_table("simulation_runs") # Post-generation reversibility correction: Alembic creates the # v2_simulation_role type with its table but does not drop it with it. sa.Enum(name="v2_simulation_role").drop(op.get_bind()) + op.drop_table("report_run_comparisons") + op.drop_index("ix_simulation_runs_status_started_at", table_name="simulation_runs") + op.drop_index("ix_simulation_runs_current_output", table_name="simulation_runs") + op.drop_table("simulation_runs") # ### end Alembic commands ### diff --git a/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py b/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py deleted file mode 100644 index 68bdef772..000000000 --- a/migrations/v2/versions/2e3de7c572ed_tell_apart_result_rows_of_each_run_pair_.py +++ /dev/null @@ -1,248 +0,0 @@ -"""tell apart result rows of each run pair and group - -Revision ID: 2e3de7c572ed -Revises: 7f43d2c7c3c6 -Create Date: 2026-11-04 01:07:49.171506 -Generation: uv run alembic -c alembic-v2.ini revision --autogenerate - -No post-generation corrections. The upgrade adds two nullable columns and -widens unique keys, so it needs no table empty. The downgrade narrows them -again, so it fails while a report run holds comparison rows of more than one -pair of simulation runs, or a run holds rows that differ only in a removed -column. -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel - - -revision: str = "2e3de7c572ed" -down_revision: Union[str, None] = "7f43d2c7c3c6" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint( - op.f("uq_budget_summary_run_variable_entity"), "budget_summary", type_="unique" - ) - op.create_unique_constraint( - "uq_budget_summary_run_variable_entity", - "budget_summary", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "variable_name", - "entity", - ], - ) - op.drop_constraint( - op.f("uq_congressional_district_impacts_run_geoid"), - "congressional_district_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_congressional_district_impacts_run_geoid", - "congressional_district_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "district_geoid", - ], - ) - op.drop_constraint( - op.f("uq_constituency_impacts_run_code"), "constituency_impacts", type_="unique" - ) - op.create_unique_constraint( - "uq_constituency_impacts_run_code", - "constituency_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "constituency_code", - ], - ) - op.add_column( - "decile_impacts", - sa.Column( - "decile_variable", - sqlmodel.sql.sqltypes.AutoString(length=512), - nullable=True, - ), - ) - op.drop_constraint( - op.f("uq_decile_impacts_run_measure"), "decile_impacts", type_="unique" - ) - op.create_unique_constraint( - "uq_decile_impacts_run_measure", - "decile_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "income_variable", - "decile_variable", - "entity", - "decile", - "quantiles", - ], - postgresql_nulls_not_distinct=True, - ) - op.drop_constraint( - op.f("uq_intra_decile_impacts_run_type_decile"), - "intra_decile_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_intra_decile_impacts_run_type_decile", - "intra_decile_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "decile_type", - "decile", - ], - ) - op.drop_constraint( - op.f("uq_local_authority_impacts_run_code"), - "local_authority_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_local_authority_impacts_run_code", - "local_authority_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "local_authority_code", - ], - ) - op.add_column( - "poverty", - sa.Column( - "filter_group", sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True - ), - ) - op.drop_constraint( - op.f("uq_poverty_simulation_run_measure"), "poverty", type_="unique" - ) - op.create_unique_constraint( - "uq_poverty_simulation_run_measure", - "poverty", - [ - "simulation_run_id", - "poverty_type", - "entity", - "filter_variable", - "filter_group", - ], - postgresql_nulls_not_distinct=True, - ) - op.drop_constraint( - op.f("uq_program_statistics_run_program_entity"), - "program_statistics", - type_="unique", - ) - op.create_unique_constraint( - "uq_program_statistics_run_program_entity", - "program_statistics", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "program_name", - "entity", - ], - postgresql_nulls_not_distinct=True, - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint( - "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_program_statistics_run_program_entity"), - "program_statistics", - ["report_run_id", "program_name", "entity"], - postgresql_nulls_not_distinct=True, - ) - op.drop_constraint("uq_poverty_simulation_run_measure", "poverty", type_="unique") - op.create_unique_constraint( - op.f("uq_poverty_simulation_run_measure"), - "poverty", - ["simulation_run_id", "poverty_type", "entity", "filter_variable"], - postgresql_nulls_not_distinct=True, - ) - op.drop_column("poverty", "filter_group") - op.drop_constraint( - "uq_local_authority_impacts_run_code", "local_authority_impacts", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_local_authority_impacts_run_code"), - "local_authority_impacts", - ["report_run_id", "local_authority_code"], - postgresql_nulls_not_distinct=False, - ) - op.drop_constraint( - "uq_intra_decile_impacts_run_type_decile", - "intra_decile_impacts", - type_="unique", - ) - op.create_unique_constraint( - op.f("uq_intra_decile_impacts_run_type_decile"), - "intra_decile_impacts", - ["report_run_id", "decile_type", "decile"], - postgresql_nulls_not_distinct=False, - ) - op.drop_constraint( - "uq_decile_impacts_run_measure", "decile_impacts", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_decile_impacts_run_measure"), - "decile_impacts", - ["report_run_id", "income_variable", "entity", "decile", "quantiles"], - postgresql_nulls_not_distinct=False, - ) - op.drop_column("decile_impacts", "decile_variable") - op.drop_constraint( - "uq_constituency_impacts_run_code", "constituency_impacts", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_constituency_impacts_run_code"), - "constituency_impacts", - ["report_run_id", "constituency_code"], - postgresql_nulls_not_distinct=False, - ) - op.drop_constraint( - "uq_congressional_district_impacts_run_geoid", - "congressional_district_impacts", - type_="unique", - ) - op.create_unique_constraint( - op.f("uq_congressional_district_impacts_run_geoid"), - "congressional_district_impacts", - ["report_run_id", "district_geoid"], - postgresql_nulls_not_distinct=False, - ) - op.drop_constraint( - "uq_budget_summary_run_variable_entity", "budget_summary", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_budget_summary_run_variable_entity"), - "budget_summary", - ["report_run_id", "variable_name", "entity"], - postgresql_nulls_not_distinct=False, - ) - # ### end Alembic commands ### diff --git a/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py b/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py deleted file mode 100644 index d132e8560..000000000 --- a/migrations/v2/versions/7a83a87d029a_drop_foreign_keys_between_the_two_.py +++ /dev/null @@ -1,196 +0,0 @@ -"""drop foreign keys between the two services tables - -Revision ID: 7a83a87d029a -Revises: 2ffd50f47a2b -Create Date: 2026-11-04 00:19:04.175069 -Generation: uv run alembic -c alembic-v2.ini revision --autogenerate - -No post-generation corrections. The tables the simulation service writes keep -the columns naming the Public API's rows, without foreign keys. The downgrade -restores the keys, so it fails while any of those columns names a row the -Public API's table lacks. -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel - - -revision: str = "7a83a87d029a" -down_revision: Union[str, None] = "2ffd50f47a2b" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_constraint( - op.f("fk_household_jobs_dynamic_id_dynamics"), - "household_jobs", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_household_jobs_policy_id_policies"), - "household_jobs", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_reports_household_id_households"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), - "reports", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_reports_region_id_regions"), "reports", type_="foreignkey" - ) - op.drop_constraint(op.f("fk_reports_user_id_users"), "reports", type_="foreignkey") - op.drop_constraint( - op.f("fk_reports_policy_id_policies"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_reports_dataset_id_datasets"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_dynamic_id_dynamics"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_household_id_households"), - "simulations", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_simulations_region_id_regions"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_dataset_id_datasets"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), - "simulations", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_simulations_policy_id_policies"), "simulations", type_="foreignkey" - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.create_foreign_key( - op.f("fk_simulations_policy_id_policies"), - "simulations", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), - "simulations", - "tax_benefit_model_versions", - ["tax_benefit_model_version_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_simulations_dataset_id_datasets"), - "simulations", - "datasets", - ["dataset_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_simulations_region_id_regions"), - "simulations", - "regions", - ["region_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_simulations_household_id_households"), - "simulations", - "households", - ["household_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_simulations_dynamic_id_dynamics"), - "simulations", - "dynamics", - ["dynamic_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_dataset_id_datasets"), - "reports", - "datasets", - ["dataset_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_policy_id_policies"), - "reports", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_user_id_users"), - "reports", - "users", - ["user_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_region_id_regions"), - "reports", - "regions", - ["region_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), - "reports", - "tax_benefit_models", - ["tax_benefit_model_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_reports_household_id_households"), - "reports", - "households", - ["household_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_household_jobs_policy_id_policies"), - "household_jobs", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_household_jobs_dynamic_id_dynamics"), - "household_jobs", - "dynamics", - ["dynamic_id"], - ["id"], - ondelete="SET NULL", - ) - # ### end Alembic commands ### diff --git a/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py b/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py deleted file mode 100644 index 2ff0acd21..000000000 --- a/migrations/v2/versions/7f43d2c7c3c6_store_the_economy_output_fields_without_.py +++ /dev/null @@ -1,174 +0,0 @@ -"""store the economy output fields without a column - -Revision ID: 7f43d2c7c3c6 -Revises: 7a83a87d029a -Create Date: 2026-11-04 00:30:48.602215 -Generation: uv run alembic -c alembic-v2.ini revision --autogenerate - -No post-generation corrections. The upgrade adds required columns without -server defaults to congressional_district_impacts, so it needs that table -empty. The downgrade makes program_statistics.entity and is_tax required -again, so it needs every row of that table to have both. -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel - - -revision: str = "7f43d2c7c3c6" -down_revision: Union[str, None] = "7a83a87d029a" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.add_column( - "congressional_district_impacts", - sa.Column("winner_percentage", sa.Float(), nullable=False), - ) - op.add_column( - "congressional_district_impacts", - sa.Column("loser_percentage", sa.Float(), nullable=False), - ) - op.add_column( - "congressional_district_impacts", - sa.Column("no_change_percentage", sa.Float(), nullable=False), - ) - op.add_column( - "constituency_impacts", - sa.Column( - "disclosure_tier", - sqlmodel.sql.sqltypes.AutoString(length=16), - nullable=True, - ), - ) - op.add_column( - "constituency_impacts", - sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), - ) - op.add_column( - "constituency_impacts", - sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), - ) - op.add_column( - "constituency_impacts", - sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), - ) - op.add_column( - "constituency_impacts", - sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), - ) - op.create_check_constraint( - op.f("ck_constituency_impacts_disclosure_tier"), - "constituency_impacts", - "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", - ) - op.add_column( - "local_authority_impacts", - sa.Column( - "disclosure_tier", - sqlmodel.sql.sqltypes.AutoString(length=16), - nullable=True, - ), - ) - op.add_column( - "local_authority_impacts", - sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), - ) - op.add_column( - "local_authority_impacts", - sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), - ) - op.add_column( - "local_authority_impacts", - sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), - ) - op.add_column( - "local_authority_impacts", - sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), - ) - op.create_check_constraint( - op.f("ck_local_authority_impacts_disclosure_tier"), - "local_authority_impacts", - "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", - ) - op.alter_column( - "program_statistics", - "entity", - existing_type=sa.VARCHAR(length=128), - nullable=True, - ) - op.alter_column( - "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=True - ) - op.drop_constraint( - op.f("uq_program_statistics_run_program_entity"), - "program_statistics", - type_="unique", - ) - op.create_unique_constraint( - "uq_program_statistics_run_program_entity", - "program_statistics", - ["report_run_id", "program_name", "entity"], - postgresql_nulls_not_distinct=True, - ) - op.add_column("report_runs", sa.Column("spm_provenance", sa.JSON(), nullable=True)) - op.add_column( - "report_runs", - sa.Column( - "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True - ), - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_column("report_runs", "data_version") - op.drop_column("report_runs", "spm_provenance") - op.drop_constraint( - "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_program_statistics_run_program_entity"), - "program_statistics", - ["report_run_id", "program_name", "entity"], - postgresql_nulls_not_distinct=False, - ) - op.alter_column( - "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=False - ) - op.alter_column( - "program_statistics", - "entity", - existing_type=sa.VARCHAR(length=128), - nullable=False, - ) - op.drop_constraint( - op.f("ck_local_authority_impacts_disclosure_tier"), - "local_authority_impacts", - type_="check", - ) - op.drop_column("local_authority_impacts", "relative_household_income_change_upper") - op.drop_column("local_authority_impacts", "relative_household_income_change_lower") - op.drop_column("local_authority_impacts", "average_household_income_change_upper") - op.drop_column("local_authority_impacts", "average_household_income_change_lower") - op.drop_column("local_authority_impacts", "disclosure_tier") - op.drop_constraint( - op.f("ck_constituency_impacts_disclosure_tier"), - "constituency_impacts", - type_="check", - ) - op.drop_column("constituency_impacts", "relative_household_income_change_upper") - op.drop_column("constituency_impacts", "relative_household_income_change_lower") - op.drop_column("constituency_impacts", "average_household_income_change_upper") - op.drop_column("constituency_impacts", "average_household_income_change_lower") - op.drop_column("constituency_impacts", "disclosure_tier") - op.drop_column("congressional_district_impacts", "no_change_percentage") - op.drop_column("congressional_district_impacts", "loser_percentage") - op.drop_column("congressional_district_impacts", "winner_percentage") - # ### end Alembic commands ### diff --git a/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py b/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py deleted file mode 100644 index 853202f49..000000000 --- a/migrations/v2/versions/dce0caa329cc_keep_untabled_comparison_results_per_.py +++ /dev/null @@ -1,120 +0,0 @@ -"""keep untabled comparison results per compared pair - -Revision ID: dce0caa329cc -Revises: 2e3de7c572ed -Create Date: 2026-11-04 01:33:30.494861 -Generation: uv run alembic -c alembic-v2.ini revision --autogenerate - -No post-generation corrections. The upgrade drops four nullable report_runs -columns, which nothing has written: no managed role could write report_runs -before the previous revisions. The downgrade drops report_run_comparisons and -its rows. -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel -from sqlalchemy.dialects import postgresql - -revision: str = "dce0caa329cc" -down_revision: Union[str, None] = "2e3de7c572ed" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.create_table( - "report_run_comparisons", - sa.Column("id", sa.Uuid(), nullable=False), - sa.Column( - "created_at", - sa.DateTime(timezone=True), - server_default=sa.text("now()"), - nullable=False, - ), - sa.Column("report_run_id", sa.Uuid(), nullable=False), - sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), - sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), - sa.Column("labor_supply_result", sa.JSON(), nullable=True), - sa.Column("cliff_impact_result", sa.JSON(), nullable=True), - sa.Column("spm_provenance", sa.JSON(), nullable=True), - sa.Column( - "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True - ), - sa.ForeignKeyConstraint( - ["baseline_simulation_run_id"], - ["simulation_runs.id"], - name=op.f( - "fk_report_run_comparisons_baseline_simulation_run_id_simulation_runs" - ), - ondelete="RESTRICT", - ), - sa.ForeignKeyConstraint( - ["reform_simulation_run_id"], - ["simulation_runs.id"], - name=op.f( - "fk_report_run_comparisons_reform_simulation_run_id_simulation_runs" - ), - ondelete="RESTRICT", - ), - sa.ForeignKeyConstraint( - ["report_run_id"], - ["report_runs.id"], - name=op.f("fk_report_run_comparisons_report_run_id_report_runs"), - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_comparisons")), - sa.UniqueConstraint( - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - name="uq_report_run_comparisons_run_pair", - ), - ) - op.drop_column("report_runs", "cliff_impact_result") - op.drop_column("report_runs", "labor_supply_result") - op.drop_column("report_runs", "data_version") - op.drop_column("report_runs", "spm_provenance") - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.add_column( - "report_runs", - sa.Column( - "spm_provenance", - postgresql.JSON(astext_type=sa.Text()), - autoincrement=False, - nullable=True, - ), - ) - op.add_column( - "report_runs", - sa.Column( - "data_version", sa.VARCHAR(length=128), autoincrement=False, nullable=True - ), - ) - op.add_column( - "report_runs", - sa.Column( - "labor_supply_result", - postgresql.JSON(astext_type=sa.Text()), - autoincrement=False, - nullable=True, - ), - ) - op.add_column( - "report_runs", - sa.Column( - "cliff_impact_result", - postgresql.JSON(astext_type=sa.Text()), - autoincrement=False, - nullable=True, - ), - ) - op.drop_table("report_run_comparisons") - # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index da6e3f5e5..6c50fae72 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "dce0caa329cc" +V2_ALEMBIC_HEAD_REVISION = "1356f658347d" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index d941d43ee..4679f3793 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,17 +1,14 @@ -"""Check the simulation runs revisions against disposable Postgres. - -Five revisions follow the parameter tree revision. The first adds -simulation runs and the table naming the simulation runs a report run -combines, keys every result table by run, and gives household jobs one -household, one year and one policy. The second drops the foreign keys from -those tables to the Public API's tables. The third adds the economy output's -fields that had no column. The fourth makes each result row tell apart its -pair of simulation runs, decile grouping and poverty group. The fifth moves -the parts of a comparison without a result table from the report run to one -row per compared pair. None changes any other table. The previous Public API -revision keeps working while they are deployed: its roles can't read or write -any table they change, except the catalog writer, which counts the rows of -simulations, reports and report_runs. +"""Check the simulation runs revision against disposable Postgres. + +One revision follows the parameter search file revision. It adds simulation +runs and the table naming the simulation runs a report run combines, keys +every result table by run, gives household jobs one household, one year and +one policy, drops the foreign keys from those tables to the Public API's +tables, adds the economy output's fields that had no column, and keeps the +parts of a comparison without a result table in one row per compared pair. It +changes no other table. The previous Public API revision keeps working while +it is deployed: its roles can't read or write any table it changes, except the +catalog writer, which counts the rows of simulations, reports and report_runs. """ from __future__ import annotations @@ -39,10 +36,6 @@ from sqlalchemy import Connection, Engine SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" -SIMULATION_RUNS_REVISION = "2ffd50f47a2b" -CROSS_SERVICE_KEYS_REVISION = "7a83a87d029a" -ECONOMY_OUTPUT_REVISION = "7f43d2c7c3c6" -RESULT_IDENTITY_REVISION = "2e3de7c572ed" NEW_TABLES = frozenset( {"simulation_runs", "report_run_simulation_runs", "report_run_comparisons"} ) @@ -160,7 +153,7 @@ def _enum_type_names(engine: Engine) -> set[str]: ) -def test_the_revisions_add_three_tables_and_change_only_simulation_side_tables( +def test_the_revision_adds_three_tables_and_changes_only_simulation_side_tables( config: Config, engine: Engine ) -> None: command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) @@ -176,49 +169,26 @@ def test_the_revisions_add_three_tables_and_change_only_simulation_side_tables( assert changed == CHANGED_TABLES -def test_each_revision_downgrades_to_its_parent_and_upgrades_again( +def test_the_revision_downgrades_to_its_parent_and_upgrades_again( config: Config, engine: Engine ) -> None: command.upgrade(config, "head") - command.downgrade(config, RESULT_IDENTITY_REVISION) - assert "report_run_comparisons" not in inspect(engine).get_table_names() - assert {"labor_supply_result", "data_version"} <= _column_names( - engine, "report_runs" - ) - - command.downgrade(config, ECONOMY_OUTPUT_REVISION) - assert "decile_variable" not in _column_names(engine, "decile_impacts") - assert "filter_group" not in _column_names(engine, "poverty") - budget_key = next( - constraint - for constraint in inspect(engine).get_unique_constraints("budget_summary") - if constraint["name"] == "uq_budget_summary_run_variable_entity" - ) - assert budget_key["column_names"] == ["report_run_id", "variable_name", "entity"] - - command.downgrade(config, CROSS_SERVICE_KEYS_REVISION) - assert {"spm_provenance", "data_version"}.isdisjoint( - _column_names(engine, "report_runs") + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + assert NEW_TABLES.isdisjoint(inspect(engine).get_table_names(schema="public")) + assert {"status", "household_result", "output_dataset_id"} <= _column_names( + engine, "simulations" ) + assert {"report_run_id", "simulation_id"} <= _column_names(engine, "poverty") + assert {"datasets", "households", "policies"} <= { + key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") + } program_entity = next( column for column in inspect(engine).get_columns("program_statistics") if column["name"] == "entity" ) assert not program_entity["nullable"] - - command.downgrade(config, SIMULATION_RUNS_REVISION) - assert {"datasets", "households", "policies"} <= { - key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") - } - - command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) - assert NEW_TABLES.isdisjoint(inspect(engine).get_table_names(schema="public")) - assert {"status", "household_result", "output_dataset_id"} <= _column_names( - engine, "simulations" - ) - assert {"report_run_id", "simulation_id"} <= _column_names(engine, "poverty") enum_types = _enum_type_names(engine) assert "v2_simulation_status" in enum_types assert "v2_simulation_role" not in enum_types diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 1171a9cb3..f8df401c4 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -67,54 +67,6 @@ PARAMETER_TREE_REVISION = "3869eb814a90" # The parameter search file revision, parent of the simulation runs revision. SIMULATION_RUNS_PARENT_REVISION = "c614b6bcccfc" -SIMULATION_RUNS_REVISION_PATH = ( - REPO - / "migrations/v2/versions/2ffd50f47a2b_add_simulation_runs_and_key_results_by_.py" -) -# The simulation runs revision, parent of the cross-service foreign key revision. -CROSS_SERVICE_KEYS_PARENT_REVISION = "2ffd50f47a2b" -# The cross-service foreign key revision, parent of the economy output revision. -ECONOMY_OUTPUT_PARENT_REVISION = "7a83a87d029a" -# The economy output revision, parent of the result row identity revision. -RESULT_IDENTITY_PARENT_REVISION = "7f43d2c7c3c6" -# The result row identity revision, parent of the comparisons revision. -COMPARISONS_PARENT_REVISION = "2e3de7c572ed" -UNTABLED_COMPARISON_COLUMNS = ( - "labor_supply_result", - "cliff_impact_result", - "spm_provenance", - "data_version", -) -# Each comparison result table's unique key, which gains the pair of runs. -COMPARISON_UNIQUE_KEYS = { - "budget_summary": "uq_budget_summary_run_variable_entity", - "congressional_district_impacts": "uq_congressional_district_impacts_run_geoid", - "constituency_impacts": "uq_constituency_impacts_run_code", - "decile_impacts": "uq_decile_impacts_run_measure", - "intra_decile_impacts": "uq_intra_decile_impacts_run_type_decile", - "local_authority_impacts": "uq_local_authority_impacts_run_code", - "program_statistics": "uq_program_statistics_run_program_entity", -} -# The foreign keys from the simulation service's tables to the Public API's. -CROSS_SERVICE_KEYS = { - "household_jobs": {"dynamics", "policies"}, - "reports": { - "datasets", - "households", - "policies", - "regions", - "tax_benefit_models", - "users", - }, - "simulations": { - "datasets", - "dynamics", - "households", - "policies", - "regions", - "tax_benefit_model_versions", - }, -} DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -296,10 +248,6 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, - COMPARISONS_PARENT_REVISION, - RESULT_IDENTITY_PARENT_REVISION, - ECONOMY_OUTPUT_PARENT_REVISION, - CROSS_SERVICE_KEYS_PARENT_REVISION, SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, AREA_IMPACT_COORDINATE_REVISION, @@ -1069,13 +1017,22 @@ def test_parameter_search_file_revision_creates_only_the_files_table() -> None: assert "op.add_column(" not in body -def test_simulation_runs_revision_has_only_its_two_documented_corrections() -> None: - """The revision is generated; its two corrections handle enum types. +def _simulation_runs_revision_source() -> str: + """Return the source of the one revision after the parameter search files.""" + (path,) = (REPO / "migrations/v2/versions").glob(f"{V2_ALEMBIC_HEAD_REVISION}_*.py") + return path.read_text(encoding="utf-8") - ``simulation_runs.status`` reuses the type the baseline created for - ``simulations.status``, and the downgrade drops the new role type. + +def test_simulation_runs_revision_is_generated_with_two_documented_corrections() -> ( + None +): + """One generated revision follows the parameter search file revision. + + Its two corrections handle enum types: ``simulation_runs.status`` reuses the + type the baseline created for ``simulations.status``, and the downgrade + drops the role type the upgrade creates with ``report_run_simulation_runs``. """ - revision = SIMULATION_RUNS_REVISION_PATH.read_text(encoding="utf-8") + revision = _simulation_runs_revision_source() upgrade = revision[ revision.index("def upgrade()") : revision.index("def downgrade()") ] @@ -1095,10 +1052,13 @@ def test_simulation_runs_revision_has_only_its_two_documented_corrections() -> N upgrade ) assert 'sa.Enum(name="v2_simulation_role").drop(op.get_bind())' in downgrade - assert 'op.create_table(\n "simulation_runs"' in upgrade - assert 'op.create_table(\n "report_run_simulation_runs"' in upgrade - assert 'op.drop_table("report_run_simulation_runs")' in downgrade - assert 'op.drop_table("simulation_runs")' in downgrade + for table_name in ( + "simulation_runs", + "report_run_simulation_runs", + "report_run_comparisons", + ): + assert f'op.create_table(\n "{table_name}"' in upgrade + assert f'op.drop_table("{table_name}")' in downgrade assert upgrade.count("op.drop_table(") == 0 for removed_column in ( "status", @@ -1112,211 +1072,3 @@ def test_simulation_runs_revision_has_only_its_two_documented_corrections() -> N for body in (upgrade, downgrade): assert "op.execute(" not in body assert "op.bulk_insert(" not in body - - -def test_cross_service_key_revision_only_drops_and_restores_foreign_keys() -> None: - """The revision after the simulation runs revision is generated. - - Its upgrade drops each foreign key from a table the simulation service - writes to a table the Public API writes; its downgrade restores them. - """ - script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) - children = [ - revision - for revision in script.walk_revisions() - if revision.down_revision == CROSS_SERVICE_KEYS_PARENT_REVISION - ] - assert len(children) == 1 - revision = Path(children[0].path).read_text(encoding="utf-8") - upgrade = revision[ - revision.index("def upgrade()") : revision.index("def downgrade()") - ] - downgrade = revision[revision.index("def downgrade()") :] - - assert ( - "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" - in revision - ) - assert "No post-generation corrections." in revision - dropped_tables = re.findall( - r'op\.drop_constraint\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*' - r'type_="foreignkey"', - upgrade, - ) - restored_keys = re.findall( - r'op\.create_foreign_key\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*"([a-z_]+)"', - downgrade, - ) - expected_count = sum(len(targets) for targets in CROSS_SERVICE_KEYS.values()) - assert len(dropped_tables) == upgrade.count("op.drop_constraint(") - assert len(dropped_tables) == expected_count - assert {table: dropped_tables.count(table) for table in CROSS_SERVICE_KEYS} == { - table: len(targets) for table, targets in CROSS_SERVICE_KEYS.items() - } - assert len(restored_keys) == downgrade.count("op.create_foreign_key(") - restored_targets: dict[str, set[str]] = {} - for table, target in restored_keys: - restored_targets.setdefault(table, set()).add(target) - assert restored_targets == CROSS_SERVICE_KEYS - for body in (upgrade, downgrade): - for operation in ( - "op.add_column(", - "op.drop_column(", - "op.create_table(", - "op.drop_table(", - "op.execute(", - "op.bulk_insert(", - ): - assert operation not in body - - -def test_economy_output_revision_only_adds_columns_and_relaxes_two() -> None: - """The revision after the cross-service key revision is generated. - - It adds the economy output's fields that had no column and lets a detailed - budget row omit its entity and tax flag; it removes nothing. - """ - script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) - children = [ - revision - for revision in script.walk_revisions() - if revision.down_revision == ECONOMY_OUTPUT_PARENT_REVISION - ] - assert len(children) == 1 - revision = Path(children[0].path).read_text(encoding="utf-8") - upgrade = revision[ - revision.index("def upgrade()") : revision.index("def downgrade()") - ] - downgrade = revision[revision.index("def downgrade()") :] - - assert ( - "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" - in revision - ) - assert "No post-generation corrections." in revision - added_columns = set( - re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) - ) - uk_disclosure_columns = { - "disclosure_tier", - "average_household_income_change_lower", - "average_household_income_change_upper", - "relative_household_income_change_lower", - "relative_household_income_change_upper", - } - assert added_columns == { - ("report_runs", "spm_provenance"), - ("report_runs", "data_version"), - ("congressional_district_impacts", "winner_percentage"), - ("congressional_district_impacts", "loser_percentage"), - ("congressional_district_impacts", "no_change_percentage"), - *(("constituency_impacts", column) for column in uk_disclosure_columns), - *(("local_authority_impacts", column) for column in uk_disclosure_columns), - } - assert upgrade.count("op.add_column(") == len(added_columns) - relaxed_columns = re.findall( - r'op\.alter_column\(\s*"(\w+)",\s*"(\w+)",\s*' - r"existing_type=sa\.\w+\([^)]*\),\s*nullable=True", - upgrade, - ) - assert upgrade.count("op.alter_column(") == 2 - assert set(relaxed_columns) == { - ("program_statistics", "entity"), - ("program_statistics", "is_tax"), - } - assert downgrade.count("op.drop_column(") == len(added_columns) - for body in (upgrade, downgrade): - assert "op.drop_table(" not in body - assert "op.execute(" not in body - assert "op.bulk_insert(" not in body - assert "op.drop_column(" not in upgrade - - -def test_result_identity_revision_adds_two_columns_and_widens_unique_keys() -> None: - """The revision after the economy output revision is generated. - - It adds ``decile_impacts.decile_variable`` and ``poverty.filter_group`` and - rebuilds each comparison result table's unique key with the pair of - simulation runs, and ``poverty``'s with its group. - """ - script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) - children = [ - revision - for revision in script.walk_revisions() - if revision.down_revision == RESULT_IDENTITY_PARENT_REVISION - ] - assert len(children) == 1 - revision = Path(children[0].path).read_text(encoding="utf-8") - upgrade = revision[ - revision.index("def upgrade()") : revision.index("def downgrade()") - ] - downgrade = revision[revision.index("def downgrade()") :] - - assert ( - "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" - in revision - ) - assert "No post-generation corrections." in revision - added_columns = set( - re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) - ) - assert added_columns == { - ("decile_impacts", "decile_variable"), - ("poverty", "filter_group"), - } - rebuilt_keys = { - table: name - for name, table in re.findall( - r'op\.create_unique_constraint\(\s*"(\w+)",\s*"(\w+)"', upgrade - ) - } - assert rebuilt_keys == { - **COMPARISON_UNIQUE_KEYS, - "poverty": "uq_poverty_simulation_run_measure", - } - assert upgrade.count("op.drop_constraint(") == len(rebuilt_keys) - assert upgrade.count('"baseline_simulation_run_id",') == len(COMPARISON_UNIQUE_KEYS) - assert downgrade.count("op.create_unique_constraint(") == len(rebuilt_keys) - assert downgrade.count("op.drop_column(") == len(added_columns) - for body in (upgrade, downgrade): - assert "op.drop_table(" not in body - assert "op.execute(" not in body - assert "op.bulk_insert(" not in body - assert "op.drop_column(" not in upgrade - - -def test_comparisons_revision_moves_untabled_results_to_one_row_per_pair() -> None: - """The revision after the result row identity revision is generated. - - It creates ``report_run_comparisons`` and drops the four columns it takes - over from ``report_runs``; the downgrade puts them back. - """ - script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) - children = [ - revision - for revision in script.walk_revisions() - if revision.down_revision == COMPARISONS_PARENT_REVISION - ] - assert len(children) == 1 - revision = Path(children[0].path).read_text(encoding="utf-8") - upgrade = revision[ - revision.index("def upgrade()") : revision.index("def downgrade()") - ] - downgrade = revision[revision.index("def downgrade()") :] - - assert ( - "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" - in revision - ) - assert "No post-generation corrections." in revision - assert 'op.create_table(\n "report_run_comparisons"' in upgrade - assert "uq_report_run_comparisons_run_pair" in upgrade - assert upgrade.count("op.drop_column(") == len(UNTABLED_COMPARISON_COLUMNS) - assert downgrade.count("op.add_column(") == len(UNTABLED_COMPARISON_COLUMNS) - for column_name in UNTABLED_COMPARISON_COLUMNS: - assert f'op.drop_column("report_runs", "{column_name}")' in upgrade - assert f'"{column_name}"' in downgrade - assert 'op.drop_table("report_run_comparisons")' in downgrade - for body in (upgrade, downgrade): - assert "op.execute(" not in body - assert "op.bulk_insert(" not in body From 4375609400dd17f943d4d19118addcf99b1976f2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:51:08 +0400 Subject: [PATCH 13/22] Keep the foreign keys between the two services' v2 tables Each service writes only its own tables, and the foreign keys between them stay, so PostgreSQL checks every row a simulation-side row names when it is written. The 14 keys from simulations, reports and household jobs to the Public API's tables are kept, with their ORM relationships, and the new columns that name the Public API's rows get keys with the delete rules the same columns have on simulations: a simulation run's catalog version and dataset (RESTRICT) and region (SET NULL), and a household job's household and catalog version (RESTRICT). reports.user_id gains an index for account deletion's SET NULL. Revision 1356f658347d is regenerated; its corrections are unchanged. The integration tests now store the catalog rows and household a row names. Co-Authored-By: Claude Opus 5.5 --- changelog.d/v2-simulation-runs.changed.md | 2 +- docs/engineering/skills/alembic-migrations.md | 12 +- ...record_simulation_runs_and_key_results_.py | 380 +++++++----------- policyengine_api/data/v2/models/households.py | 48 ++- policyengine_api/data/v2/models/metadata.py | 13 + policyengine_api/data/v2/models/policies.py | 8 + policyengine_api/data/v2/models/reports.py | 62 ++- .../data/v2/models/simulations.py | 87 +++- policyengine_api/data/v2/models/users.py | 6 +- tests/integration/simulation_rows.py | 58 +++ .../test_alembic_v2_simulation_runs.py | 143 +++++-- .../test_v2_catalog_publication.py | 12 +- .../test_v2_simulation_role_privileges.py | 38 +- tests/unit/v2/test_models.py | 9 +- tests/unit/v2/test_simulation_run_models.py | 100 +++-- 15 files changed, 593 insertions(+), 385 deletions(-) diff --git a/changelog.d/v2-simulation-runs.changed.md b/changelog.d/v2-simulation-runs.changed.md index 9e459c853..88dcb1697 100644 --- a/changelog.d/v2-simulation-runs.changed.md +++ b/changelog.d/v2-simulation-runs.changed.md @@ -1 +1 @@ -API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. Its simulation tables no longer have foreign keys to the tables the Public API writes. Result tables gain the economy output's fields that had no column. +API v2's database records each run of a simulation and of a report with the versions it used, keys every result to the run that produced it, and stores household calculation jobs for one household, one year and one policy. Result tables gain the economy output's fields that had no column. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 94ac98f65..8292516af 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -533,9 +533,15 @@ model changes: - `household_jobs` calculates one household (`household_id`) for one `year` under one policy, and records its catalog and package versions and its Modal call ID. -- The foreign keys from the tables the simulation service and its Modal workers - write to the tables the Public API writes are dropped; the service that - writes a row checks the rows it names. +- The new columns that name the Public API's rows are foreign keys with the + delete rules the same columns have on `simulations`: + `simulation_runs.tax_benefit_model_version_id` and `dataset_id` (`RESTRICT`) + and `region_id` (`SET NULL`), and `household_jobs.household_id` and + `tax_benefit_model_version_id` (`RESTRICT`). `reports.user_id` gains an + index, which account deletion's `SET NULL` searches. Each service writes only + its own tables, and the foreign keys between them stay: PostgreSQL checks + each named row when a row is written, with the privileges of the named + table's owner, so a writing role needs no privilege on the named table. Missing values in the unique keys of `poverty`, `decile_impacts` and `program_statistics` count as equal (`NULLS NOT DISTINCT`, which PostgreSQL diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 2f1ebefca..b758bd240 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -2,10 +2,10 @@ Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-05 23:44:44.321856 +Create Date: 2026-11-05 23:50:12.713355 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate -Two post-generation corrections, each commented where it is made: +Post-generation corrections, each commented where it is made: simulation_runs.status reuses the existing v2_simulation_status type, and the downgrade drops the v2_simulation_role type the upgrade creates. @@ -99,12 +99,32 @@ def upgrade() -> None: "output_sha256 IS NULL OR length(output_sha256) = 64", name=op.f("ck_simulation_runs_output_sha256_length"), ), + sa.ForeignKeyConstraint( + ["dataset_id"], + ["datasets.id"], + name=op.f("fk_simulation_runs_dataset_id_datasets"), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["region_id"], + ["regions.id"], + name=op.f("fk_simulation_runs_region_id_regions"), + ondelete="SET NULL", + ), sa.ForeignKeyConstraint( ["simulation_id"], ["simulations.id"], name=op.f("fk_simulation_runs_simulation_id_simulations"), ondelete="CASCADE", ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + ["tax_benefit_model_versions.id"], + name=op.f( + "fk_simulation_runs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + ondelete="RESTRICT", + ), sa.PrimaryKeyConstraint("id", name=op.f("pk_simulation_runs")), ) op.create_index( @@ -126,6 +146,12 @@ def upgrade() -> None: ["status", "started_at"], unique=False, ) + op.create_index( + op.f("ix_simulation_runs_tax_benefit_model_version_id"), + "simulation_runs", + ["tax_benefit_model_version_id"], + unique=False, + ) op.create_table( "report_run_comparisons", sa.Column("id", sa.Uuid(), nullable=False), @@ -251,18 +277,18 @@ def upgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -315,21 +341,21 @@ def upgrade() -> None: ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -381,12 +407,12 @@ def upgrade() -> None: ], ) op.drop_constraint( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) @@ -473,8 +499,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("decile_impacts", "reform_simulation_id") op.drop_column("decile_impacts", "baseline_simulation_id") + op.drop_column("decile_impacts", "reform_simulation_id") op.add_column( "household_jobs", sa.Column( @@ -520,15 +546,29 @@ def upgrade() -> None: ["household_id"], unique=False, ) - op.drop_constraint( - op.f("fk_household_jobs_policy_id_policies"), + op.create_index( + op.f("ix_household_jobs_tax_benefit_model_version_id"), "household_jobs", - type_="foreignkey", + ["tax_benefit_model_version_id"], + unique=False, ) - op.drop_constraint( - op.f("fk_household_jobs_dynamic_id_dynamics"), + op.create_foreign_key( + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", - type_="foreignkey", + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_household_jobs_household_id_households"), + "household_jobs", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", ) op.create_check_constraint( op.f("ck_household_jobs_terminal_completion"), @@ -550,12 +590,12 @@ def upgrade() -> None: ["simulation_run_id", "income_variable", "entity"], ) op.drop_constraint( - op.f("fk_inequality_report_run_id_report_runs"), + op.f("fk_inequality_simulation_id_simulations"), "inequality", type_="foreignkey", ) op.drop_constraint( - op.f("fk_inequality_simulation_id_simulations"), + op.f("fk_inequality_report_run_id_report_runs"), "inequality", type_="foreignkey", ) @@ -567,8 +607,8 @@ def upgrade() -> None: ["id"], ondelete="CASCADE", ) - op.drop_column("inequality", "simulation_id") op.drop_column("inequality", "report_run_id") + op.drop_column("inequality", "simulation_id") op.add_column( "intra_decile_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -594,28 +634,28 @@ def upgrade() -> None: ], ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -679,18 +719,18 @@ def upgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -735,8 +775,8 @@ def upgrade() -> None: ["id"], ondelete="CASCADE", ) - op.drop_column("poverty", "simulation_id") op.drop_column("poverty", "report_run_id") + op.drop_column("poverty", "simulation_id") op.add_column( "program_statistics", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -782,18 +822,18 @@ def upgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -811,58 +851,19 @@ def upgrade() -> None: ["status", "started_at"], unique=False, ) - op.drop_constraint( - op.f("fk_reports_region_id_regions"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_reports_policy_id_policies"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_reports_dataset_id_datasets"), "reports", type_="foreignkey" - ) - op.drop_constraint(op.f("fk_reports_user_id_users"), "reports", type_="foreignkey") - op.drop_constraint( - op.f("fk_reports_household_id_households"), "reports", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), - "reports", - type_="foreignkey", - ) + op.create_index(op.f("ix_reports_user_id"), "reports", ["user_id"], unique=False) op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") - op.drop_constraint( - op.f("fk_simulations_dataset_id_datasets"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_dynamic_id_dynamics"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_policy_id_policies"), "simulations", type_="foreignkey" - ) - op.drop_constraint( - op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), - "simulations", - type_="foreignkey", - ) - op.drop_constraint( - op.f("fk_simulations_household_id_households"), - "simulations", - type_="foreignkey", - ) op.drop_constraint( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", type_="foreignkey", ) - op.drop_constraint( - op.f("fk_simulations_region_id_regions"), "simulations", type_="foreignkey" - ) - op.drop_column("simulations", "error_message") + op.drop_column("simulations", "status") op.drop_column("simulations", "output_dataset_id") - op.drop_column("simulations", "completed_at") - op.drop_column("simulations", "household_result") op.drop_column("simulations", "started_at") - op.drop_column("simulations", "status") + op.drop_column("simulations", "error_message") + op.drop_column("simulations", "household_result") + op.drop_column("simulations", "completed_at") # ### end Alembic commands ### @@ -871,18 +872,7 @@ def downgrade() -> None: op.add_column( "simulations", sa.Column( - "status", - postgresql.ENUM( - "pending", "running", "succeeded", "failed", name="v2_simulation_status" - ), - autoincrement=False, - nullable=False, - ), - ) - op.add_column( - "simulations", - sa.Column( - "started_at", + "completed_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, @@ -897,10 +887,14 @@ def downgrade() -> None: nullable=True, ), ) + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + ) op.add_column( "simulations", sa.Column( - "completed_at", + "started_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, @@ -912,15 +906,14 @@ def downgrade() -> None: ) op.add_column( "simulations", - sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), - ) - op.create_foreign_key( - op.f("fk_simulations_region_id_regions"), - "simulations", - "regions", - ["region_id"], - ["id"], - ondelete="SET NULL", + sa.Column( + "status", + postgresql.ENUM( + "pending", "running", "succeeded", "failed", name="v2_simulation_status" + ), + autoincrement=False, + nullable=False, + ), ) op.create_foreign_key( op.f("fk_simulations_output_dataset_id_datasets"), @@ -930,100 +923,13 @@ def downgrade() -> None: ["id"], ondelete="SET NULL", ) - op.create_foreign_key( - op.f("fk_simulations_household_id_households"), - "simulations", - "households", - ["household_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_simulations_tax_benefit_model_version_id_tax_benefit_84f8"), - "simulations", - "tax_benefit_model_versions", - ["tax_benefit_model_version_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_simulations_policy_id_policies"), - "simulations", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_simulations_dynamic_id_dynamics"), - "simulations", - "dynamics", - ["dynamic_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_simulations_dataset_id_datasets"), - "simulations", - "datasets", - ["dataset_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_index( op.f("ix_simulations_status_created_at"), "simulations", ["status", "created_at"], unique=False, ) - op.create_foreign_key( - op.f("fk_reports_tax_benefit_model_id_tax_benefit_models"), - "reports", - "tax_benefit_models", - ["tax_benefit_model_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_reports_household_id_households"), - "reports", - "households", - ["household_id"], - ["id"], - ondelete="RESTRICT", - ) - op.create_foreign_key( - op.f("fk_reports_user_id_users"), - "reports", - "users", - ["user_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_dataset_id_datasets"), - "reports", - "datasets", - ["dataset_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_policy_id_policies"), - "reports", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - op.f("fk_reports_region_id_regions"), - "reports", - "regions", - ["region_id"], - ["id"], - ondelete="SET NULL", - ) + op.drop_index(op.f("ix_reports_user_id"), table_name="reports") op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") op.drop_column("report_runs", "modal_call_id") op.add_column( @@ -1039,12 +945,12 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) op.drop_constraint( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) @@ -1086,11 +992,11 @@ def downgrade() -> None: op.drop_column("program_statistics", "baseline_simulation_run_id") op.add_column( "poverty", - sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), + sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), ) op.add_column( "poverty", - sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), ) op.drop_constraint( op.f("fk_poverty_simulation_run_id_simulation_runs"), @@ -1140,12 +1046,12 @@ def downgrade() -> None: type_="check", ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) @@ -1194,28 +1100,28 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1234,25 +1140,17 @@ def downgrade() -> None: op.drop_column("intra_decile_impacts", "baseline_simulation_run_id") op.add_column( "inequality", - sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), + sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), ) op.add_column( "inequality", - sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), ) op.drop_constraint( op.f("fk_inequality_simulation_run_id_simulation_runs"), "inequality", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_inequality_simulation_id_simulations"), - "inequality", - "simulations", - ["simulation_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_foreign_key( op.f("fk_inequality_report_run_id_report_runs"), "inequality", @@ -1261,6 +1159,14 @@ def downgrade() -> None: ["id"], ondelete="CASCADE", ) + op.create_foreign_key( + op.f("fk_inequality_simulation_id_simulations"), + "inequality", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) op.drop_constraint( "uq_inequality_simulation_run_measure", "inequality", type_="unique" ) @@ -1275,21 +1181,21 @@ def downgrade() -> None: op.drop_constraint( op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" ) - op.create_foreign_key( - op.f("fk_household_jobs_dynamic_id_dynamics"), + op.drop_constraint( + op.f("fk_household_jobs_household_id_households"), "household_jobs", - "dynamics", - ["dynamic_id"], - ["id"], - ondelete="SET NULL", + type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_household_jobs_policy_id_policies"), + op.drop_constraint( + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", - "policies", - ["policy_id"], - ["id"], - ondelete="SET NULL", + type_="foreignkey", + ) + op.drop_index( + op.f("ix_household_jobs_tax_benefit_model_version_id"), + table_name="household_jobs", ) op.drop_index(op.f("ix_household_jobs_household_id"), table_name="household_jobs") op.drop_column("household_jobs", "modal_call_id") @@ -1302,13 +1208,13 @@ def downgrade() -> None: op.add_column( "decile_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "decile_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1377,18 +1283,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1422,14 +1328,14 @@ def downgrade() -> None: ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", @@ -1479,12 +1385,12 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) @@ -1524,6 +1430,10 @@ def downgrade() -> None: # v2_simulation_role type with its table but does not drop it with it. sa.Enum(name="v2_simulation_role").drop(op.get_bind()) op.drop_table("report_run_comparisons") + op.drop_index( + op.f("ix_simulation_runs_tax_benefit_model_version_id"), + table_name="simulation_runs", + ) op.drop_index("ix_simulation_runs_status_started_at", table_name="simulation_runs") op.drop_index("ix_simulation_runs_current_output", table_name="simulation_runs") op.drop_table("simulation_runs") diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 351dff625..7a7fd4ca0 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -6,16 +6,18 @@ 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 TimestampedModel, enum_type +from policyengine_api.data.v2.models.policies import Dynamic, Policy 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): @@ -31,9 +33,8 @@ 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 saved links use - cannot be deleted; simulations, reports and household jobs name a household - without a foreign key. + household again finds the existing row. Households that simulations, reports + or saved links use cannot be deleted. """ __tablename__ = "households" @@ -76,11 +77,13 @@ class Household(TimestampedModel, table=True): ) default_year: int | None = None household_data: dict[str, Any] = Field( - sa_type=sa.JSON().with_variant(postgresql.JSONB(), "postgresql") + sa_type=sa.JSON().with_variant(sa.dialects.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", ) @@ -95,9 +98,11 @@ class HouseholdJob(TimestampedModel, table=True): ``policy_id`` is null for current law. ``request_data`` holds the calculation's options, such as axes. The job records the catalog version and PolicyEngine.py version it ran on, its Modal call, and its status and - result or error. Its household, policy, dynamic and catalog version are logical - references to the Public API's rows, without foreign keys: the simulation - service checks them before it writes the job. + result or error. The simulation service writes the job and the Public API + the rows it names; each service writes only its own tables, and the foreign + keys between them stay. The job's household and catalog version can't be + deleted while it exists; deleting its policy or dynamic keeps the job but + clears the reference. """ __tablename__ = "household_jobs" @@ -114,12 +119,28 @@ class HouseholdJob(TimestampedModel, table=True): ) country: str = Field(max_length=16) - household_id: UUID = Field(index=True) + household_id: UUID = Field( + foreign_key="households.id", + ondelete="RESTRICT", + index=True, + ) year: int request_data: dict[str, Any] = Field(sa_type=sa.JSON) - policy_id: UUID | None = None - dynamic_id: UUID | None = None - tax_benefit_model_version_id: UUID + 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, + ) country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) status: HouseholdJobStatus = Field( @@ -137,3 +158,6 @@ class HouseholdJob(TimestampedModel, table=True): 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 1fa1622e7..06081e25c 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -18,6 +18,8 @@ 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): @@ -78,6 +80,7 @@ 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): @@ -125,6 +128,9 @@ 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") @@ -185,6 +191,8 @@ 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): @@ -237,6 +245,11 @@ 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"}, + ) + 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 229a5f81a..f0da6e8cc 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -11,11 +11,14 @@ 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): @@ -83,6 +86,9 @@ 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", ) @@ -104,3 +110,5 @@ 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 index 71bfbc931..9427ffbb5 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -20,7 +20,15 @@ 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, SimulationRun +from policyengine_api.data.v2.models.users import User class ReportRunStatus(StrEnum): @@ -74,11 +82,12 @@ class SimulationRole(StrEnum): class Report(TimestampedModel, table=True): """Stable report definition; execution state belongs to ReportRun. - Its user, model, policy, household, dataset and region are logical - references to the Public API's rows, without foreign keys: the simulation - service checks them before it writes the report, except ``user_id``, the - owner the Public API names in its request, which the Public API checks - because the simulation service can't read users. + The simulation service writes reports, and the Public API the user, model, + policy, household, dataset and region a report names; each service writes + only its own tables, and the foreign keys between them stay. Deleting the + user, policy, dataset or region clears the reference, which for the user + is found through ``user_id``'s index; the model and household can't be + deleted while a report names them. """ __tablename__ = "reports" @@ -94,9 +103,22 @@ class Report(TimestampedModel, table=True): description: str | None = None country: str = Field(max_length=16) type: str | None = Field(default=None, max_length=128) - user_id: UUID | None = None - tax_benefit_model_id: UUID = Field(index=True) - policy_id: UUID | None = None + user_id: UUID | None = Field( + default=None, + foreign_key="users.id", + ondelete="SET NULL", + index=True, + ) + 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", @@ -107,12 +129,27 @@ class Report(TimestampedModel, table=True): foreign_key="simulations.id", ondelete="SET NULL", ) - household_id: UUID | None = None - dataset_id: UUID | None = None - region_id: UUID | None = None + 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"}, @@ -121,6 +158,9 @@ class Report(TimestampedModel, table=True): 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, diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 5a2e31322..ca84d91fb 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -9,6 +9,13 @@ 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 ( @@ -39,9 +46,11 @@ class Simulation(TimestampedModel, table=True): """A simulation of exactly one household or one dataset, matching its type. It defines what to calculate; each calculation of it is a ``SimulationRun``. - Its dataset, household, policy, dynamic, catalog version and region are - logical references to the Public API's rows, without foreign keys: the - simulation service checks them before it writes the simulation. + The simulation service writes it, and the Public API the rows it names; + each service writes only its own tables, and the foreign keys between them + stay. Deleting its policy, dynamic or region clears that reference; its + household, input dataset and catalog version cannot be deleted while it + exists. """ __tablename__ = "simulations" @@ -67,17 +76,52 @@ class Simulation(TimestampedModel, table=True): default=SimulationType.ECONOMY, sa_type=enum_type(SimulationType, "v2_simulation_type"), ) - dataset_id: UUID | None = None - household_id: UUID | None = None - policy_id: UUID | None = None - dynamic_id: UUID | None = None - tax_benefit_model_version_id: UUID = Field(index=True) - region_id: UUID | None = None + 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, + ) + region_id: UUID | None = Field( + default=None, + foreign_key="regions.id", + ondelete="SET NULL", + ) 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 + dataset: Dataset | None = Relationship( + back_populates="input_simulations", + sa_relationship_kwargs={"foreign_keys": "Simulation.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"}, @@ -97,9 +141,10 @@ class SimulationRun(TimestampedModel, table=True): ``tax_benefit_model_version_id``, ``dataset_id`` and ``region_id`` name the catalog version, dataset and region the run used, which a run on a newer - version resolves again. They are logical references: the catalog belongs to - the Public API, and the simulation service checks them before it writes the - run. A household simulation's run stores its household in + version resolves again. They are foreign keys to the Public API's catalog + with the same delete rules as the simulation's: the catalog version and + dataset can't be deleted while a run names them, and deleting the region + clears the reference. A household simulation's run stores its household in ``household_result``; an economy run's output file is ``output_uri``, with its SHA-256. Each poverty and inequality result belongs to one run. """ @@ -134,11 +179,23 @@ class SimulationRun(TimestampedModel, table=True): foreign_key="simulations.id", ondelete="CASCADE", ) - tax_benefit_model_version_id: UUID + tax_benefit_model_version_id: UUID = Field( + foreign_key="tax_benefit_model_versions.id", + ondelete="RESTRICT", + index=True, + ) country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) - dataset_id: UUID | None = None - region_id: UUID | None = None + dataset_id: UUID | None = Field( + default=None, + foreign_key="datasets.id", + ondelete="RESTRICT", + ) + 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"), diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 06c2b17f5..f3b987123 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -23,6 +23,7 @@ UserReportAssociation, UserSimulationAssociation, ) + from policyengine_api.data.v2.models.reports import Report class UserType(StrEnum): @@ -36,8 +37,8 @@ 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; a report naming the user - as its owner keeps that ID, because it has no foreign key to this table. + their sign-in identities and saved associations; reports they own are kept + without an owner. """ __tablename__ = "users" @@ -70,6 +71,7 @@ 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/tests/integration/simulation_rows.py b/tests/integration/simulation_rows.py index 8a44d0a71..4a6767dc1 100644 --- a/tests/integration/simulation_rows.py +++ b/tests/integration/simulation_rows.py @@ -7,7 +7,9 @@ from __future__ import annotations from contextlib import contextmanager +from dataclasses import dataclass from typing import TYPE_CHECKING +from uuid import UUID, uuid4 import pytest from sqlalchemy import text @@ -22,6 +24,62 @@ RUN_VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "4.0.0"} +@dataclass(frozen=True) +class PublicApiRows: + """The Public API's rows a simulation, report or household job names.""" + + model_id: UUID + version_id: UUID + dataset_id: UUID + household_id: UUID + + +def insert_public_api_rows(connection: Connection) -> PublicApiRows: + """Store a model, one of its versions, a dataset and a US household. + + The simulation tables' foreign keys need them. Run it as the disposable + database's superuser: the simulation API's roles can't write these tables. + """ + rows = PublicApiRows( + model_id=uuid4(), version_id=uuid4(), dataset_id=uuid4(), household_id=uuid4() + ) + # Model names are unique and at most 32 characters. + insert_row( + connection, + "tax_benefit_models", + id=rows.model_id, + name=f"runs-{rows.model_id.hex[:24]}", + ) + insert_row( + connection, + "tax_benefit_model_versions", + id=rows.version_id, + model_id=rows.model_id, + version="1.0.0", + current_law_id=1, + metadata_time_periods="[2026]", + ) + insert_row( + connection, + "datasets", + id=rows.dataset_id, + name="populace", + year=2026, + is_output_dataset=False, + tax_benefit_model_version_id=rows.version_id, + ) + insert_row( + connection, + "households", + id=rows.household_id, + country_id="us", + household_data="{}", + canonicalization_version=1, + content_hash=uuid4().hex * 2, + ) + return rows + + def insert_row(connection: Connection, table: str, **values: object) -> None: """Insert one row of the named columns into ``table``.""" columns = ", ".join(values) diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 4679f3793..5862861aa 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -28,7 +28,13 @@ load_v2_alembic_settings, ) from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL -from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused +from tests.integration.simulation_rows import ( + RUN_VERSIONS, + PublicApiRows, + insert_public_api_rows, + insert_row, + refused, +) if TYPE_CHECKING: from collections.abc import Iterator @@ -61,6 +67,38 @@ SIMULATION_SIDE_TABLES = ( CHANGED_TABLES | NEW_TABLES | {"aggregates", "change_aggregates"} ) +# The foreign keys from the simulation service's tables to the Public API's +# rows, with their delete rules. +PUBLIC_API_FOREIGN_KEYS = { + ("simulations", "dataset_id"): ("datasets", "RESTRICT"), + ("simulations", "household_id"): ("households", "RESTRICT"), + ("simulations", "policy_id"): ("policies", "SET NULL"), + ("simulations", "dynamic_id"): ("dynamics", "SET NULL"), + ("simulations", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulations", "region_id"): ("regions", "SET NULL"), + ("reports", "user_id"): ("users", "SET NULL"), + ("reports", "tax_benefit_model_id"): ("tax_benefit_models", "RESTRICT"), + ("reports", "policy_id"): ("policies", "SET NULL"), + ("reports", "household_id"): ("households", "RESTRICT"), + ("reports", "dataset_id"): ("datasets", "SET NULL"), + ("reports", "region_id"): ("regions", "SET NULL"), + ("simulation_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulation_runs", "dataset_id"): ("datasets", "RESTRICT"), + ("simulation_runs", "region_id"): ("regions", "SET NULL"), + ("household_jobs", "household_id"): ("households", "RESTRICT"), + ("household_jobs", "policy_id"): ("policies", "SET NULL"), + ("household_jobs", "dynamic_id"): ("dynamics", "SET NULL"), + ("household_jobs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), +} SHA256_HEX = "a" * 64 COMPLETED_AT = "2026-11-04T00:00:00Z" @@ -198,21 +236,25 @@ def test_the_revision_downgrades_to_its_parent_and_upgrades_again( assert "v2_simulation_role" in _enum_type_names(engine) -def test_head_has_no_foreign_key_between_the_two_services_tables( +def test_head_keeps_the_foreign_keys_to_the_public_api_tables( config: Config, engine: Engine ) -> None: + """Each service writes only its own tables; the keys between them stay.""" command.upgrade(config, "head") inspector = inspect(engine) - cross_service_keys = sorted( - f"{table_name}.{key['name']} -> {key['referred_table']}" - for table_name in inspector.get_table_names(schema="public") + public_api_keys = { + (table_name, column_name): (key["referred_table"], key["options"]["ondelete"]) + for table_name in SIMULATION_SIDE_TABLES for key in inspector.get_foreign_keys(table_name) - if (table_name in SIMULATION_SIDE_TABLES) - != (key["referred_table"] in SIMULATION_SIDE_TABLES) - ) + if key["referred_table"] not in SIMULATION_SIDE_TABLES + for column_name in key["constrained_columns"] + } - assert cross_service_keys == [] + assert public_api_keys == PUBLIC_API_FOREIGN_KEYS + assert {"ix_reports_user_id"} <= { + index["name"] for index in inspector.get_indexes("reports") + } def _count(connection: Connection, table: str, column: str, value: object) -> int: @@ -222,16 +264,23 @@ def _count(connection: Connection, table: str, column: str, value: object) -> in ).scalar_one() -def _insert_simulation(connection: Connection) -> object: - """Store an economy simulation; its catalog references need no catalog rows.""" +@pytest.fixture +def public_api_rows(connection: Connection) -> PublicApiRows: + return insert_public_api_rows(connection) + + +def _insert_simulation( + connection: Connection, public_api_rows: PublicApiRows +) -> object: + """Store an economy simulation of the stored dataset and catalog version.""" simulation_id = uuid4() insert_row( connection, "simulations", id=simulation_id, simulation_type="economy", - dataset_id=uuid4(), - tax_benefit_model_version_id=uuid4(), + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, year=2026, ) return simulation_id @@ -244,13 +293,18 @@ def _insert_simulation_run( status: str = "pending", **values: object, ) -> object: + """Store a run of the simulation on the simulation's catalog version.""" run_id = uuid4() + version_id = connection.execute( + text("SELECT tax_benefit_model_version_id FROM simulations WHERE id = :id"), + {"id": simulation_id}, + ).scalar_one() insert_row( connection, "simulation_runs", id=run_id, simulation_id=simulation_id, - tax_benefit_model_version_id=uuid4(), + tax_benefit_model_version_id=version_id, status=status, **RUN_VERSIONS, **values, @@ -258,7 +312,9 @@ def _insert_simulation_run( return run_id -def _insert_report_run(connection: Connection) -> object: +def _insert_report_run( + connection: Connection, public_api_rows: PublicApiRows +) -> object: report_id, report_run_id = uuid4(), uuid4() insert_row( connection, @@ -266,7 +322,7 @@ def _insert_report_run(connection: Connection) -> object: id=report_id, label="Report", country="us", - tax_benefit_model_id=uuid4(), + tax_benefit_model_id=public_api_rows.model_id, inputs="{}", ) insert_row( @@ -281,10 +337,41 @@ def _insert_report_run(connection: Connection) -> object: return report_run_id +def test_rows_naming_a_missing_public_api_row_are_refused( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + """PostgreSQL checks each row a simulation-side row names when it is written.""" + simulation_id = _insert_simulation(connection, public_api_rows) + + with refused(connection, match="foreign key constraint"): + insert_row( + connection, + "simulation_runs", + id=uuid4(), + simulation_id=simulation_id, + tax_benefit_model_version_id=uuid4(), + status="pending", + **RUN_VERSIONS, + ) + with refused(connection, match="foreign key constraint"): + insert_row( + connection, + "household_jobs", + id=uuid4(), + country="us", + household_id=uuid4(), + year=2026, + request_data="{}", + tax_benefit_model_version_id=public_api_rows.version_id, + status="pending", + **RUN_VERSIONS, + ) + + def test_simulation_runs_need_a_completion_time_and_a_whole_output( - connection: Connection, + connection: Connection, public_api_rows: PublicApiRows ) -> None: - simulation_id = _insert_simulation(connection) + simulation_id = _insert_simulation(connection, public_api_rows) run_id = _insert_simulation_run(connection, simulation_id) with refused(connection): @@ -328,10 +415,12 @@ def test_simulation_runs_need_a_completion_time_and_a_whole_output( def test_report_run_members_are_unique_and_keep_their_simulation_runs( - connection: Connection, + connection: Connection, public_api_rows: PublicApiRows ) -> None: - run_id = _insert_simulation_run(connection, _insert_simulation(connection)) - report_run_id = _insert_report_run(connection) + run_id = _insert_simulation_run( + connection, _insert_simulation(connection, public_api_rows) + ) + report_run_id = _insert_report_run(connection, public_api_rows) member = { "report_run_id": report_run_id, "simulation_run_id": run_id, @@ -355,13 +444,13 @@ def test_report_run_members_are_unique_and_keep_their_simulation_runs( def test_results_are_unique_per_run_and_name_runs_that_exist( - connection: Connection, + connection: Connection, public_api_rows: PublicApiRows ) -> None: - simulation_id = _insert_simulation(connection) + simulation_id = _insert_simulation(connection, public_api_rows) baseline_run_id = _insert_simulation_run(connection, simulation_id) reform_run_id = _insert_simulation_run(connection, simulation_id) comparison_runs = { - "report_run_id": _insert_report_run(connection), + "report_run_id": _insert_report_run(connection, public_api_rows), "baseline_simulation_run_id": baseline_run_id, "reform_simulation_run_id": reform_run_id, } @@ -473,13 +562,13 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( def test_household_jobs_take_a_valid_year_and_record_when_they_end( - connection: Connection, + connection: Connection, public_api_rows: PublicApiRows ) -> None: household_job = { "country": "us", - "household_id": uuid4(), + "household_id": public_api_rows.household_id, "request_data": "{}", - "tax_benefit_model_version_id": uuid4(), + "tax_benefit_model_version_id": public_api_rows.version_id, **RUN_VERSIONS, } insert_row( diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index cb98c0b5a..17078dceb 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -82,11 +82,6 @@ def _disposable_url() -> str: return database_url -# Simulations and reports have no foreign key to the catalog, so truncating the -# catalog no longer reaches them; the tests that store them name them here. -TRUNCATE_PUBLICATION_ROWS = "TRUNCATE tax_benefit_models, simulations, reports CASCADE" - - @pytest.fixture def publication_engine() -> Engine: database_url = _disposable_url() @@ -94,10 +89,10 @@ def publication_engine() -> Engine: command.upgrade(config, "head") engine = create_engine(database_url, poolclass=NullPool) with engine.begin() as connection: - connection.execute(text(TRUNCATE_PUBLICATION_ROWS)) + connection.execute(text("TRUNCATE tax_benefit_models CASCADE")) yield engine with engine.begin() as connection: - connection.execute(text(TRUNCATE_PUBLICATION_ROWS)) + connection.execute(text("TRUNCATE tax_benefit_models CASCADE")) engine.dispose() @@ -548,8 +543,7 @@ def test_republish_preserves_existing_run_and_dataset_version_records( assert persisted_simulation.dataset_id == protected_ids[4] assert persisted_report.dataset_id == protected_ids[5] assert persisted_dataset_version.dataset_id == protected_ids[6] - # The simulation and report name catalog rows without foreign keys, - # so republishing must keep every row they name. + # Republishing keeps every catalog row the simulation and report name. for dataset_id in ( persisted_simulation.dataset_id, persisted_report.dataset_id, diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index c47d5c0f6..6152dfb37 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -20,7 +20,13 @@ from sqlalchemy.pool import NullPool from policyengine_api.constants import REPO -from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused +from tests.integration.simulation_rows import ( + RUN_VERSIONS, + PublicApiRows, + insert_public_api_rows, + insert_row, + refused, +) if TYPE_CHECKING: from collections.abc import Iterator @@ -48,6 +54,12 @@ def connection(disposable_v2_database_url, managed_v2_roles) -> Iterator[Connect engine.dispose() +@pytest.fixture +def public_api_rows(connection: Connection) -> PublicApiRows: + """Store, as the superuser, the Public API's rows the simulation rows name.""" + return insert_public_api_rows(connection) + + def _act_as(connection: Connection, role: str) -> None: quoted_role = connection.dialect.identifier_preparer.quote(role) connection.execute(text(f"SET LOCAL ROLE {quoted_role}")) @@ -59,7 +71,9 @@ def _refused(connection: Connection, statement: str, **values: object) -> None: connection.execute(text(statement), values) -def _create_simulation_rows(connection: Connection) -> dict[str, object]: +def _create_simulation_rows( + connection: Connection, public_api_rows: PublicApiRows +) -> dict[str, object]: """Create what the simulation API creates for one economy report run.""" ids = { name: uuid4() @@ -70,8 +84,8 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: "simulations", id=ids["simulation"], simulation_type="economy", - dataset_id=uuid4(), - tax_benefit_model_version_id=uuid4(), + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, year=2026, ) insert_row( @@ -79,7 +93,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: "simulation_runs", id=ids["run"], simulation_id=ids["simulation"], - tax_benefit_model_version_id=uuid4(), + tax_benefit_model_version_id=public_api_rows.version_id, status="pending", **RUN_VERSIONS, ) @@ -89,7 +103,7 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: id=ids["report"], label="Report", country="us", - tax_benefit_model_id=uuid4(), + tax_benefit_model_id=public_api_rows.model_id, inputs="{}", ) insert_row( @@ -114,10 +128,10 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: "household_jobs", id=ids["household_job"], country="us", - household_id=uuid4(), + household_id=public_api_rows.household_id, year=2026, request_data="{}", - tax_benefit_model_version_id=uuid4(), + tax_benefit_model_version_id=public_api_rows.version_id, status="pending", **RUN_VERSIONS, ) @@ -125,11 +139,11 @@ def _create_simulation_rows(connection: Connection) -> dict[str, object]: def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( - connection: Connection, managed_v2_roles + connection: Connection, managed_v2_roles, public_api_rows: PublicApiRows ) -> None: _act_as(connection, managed_v2_roles.simulation_service) - ids = _create_simulation_rows(connection) + ids = _create_simulation_rows(connection, public_api_rows) run_count = connection.execute( text("SELECT count(*) FROM simulation_runs WHERE id = :id"), {"id": ids["run"]} ).scalar_one() @@ -162,10 +176,10 @@ def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( def test_modal_workers_write_their_run_status_and_results_only( - connection: Connection, managed_v2_roles + connection: Connection, managed_v2_roles, public_api_rows: PublicApiRows ) -> None: _act_as(connection, managed_v2_roles.simulation_service) - ids = _create_simulation_rows(connection) + ids = _create_simulation_rows(connection, public_api_rows) _act_as(connection, managed_v2_roles.modal_worker) updated = connection.execute( diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 976ebe7e6..db60cd8cf 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -436,15 +436,14 @@ def test_regions_have_one_same_model_version_default_logical_dataset() -> None: } -def test_reports_and_simulations_name_their_selected_dataset() -> None: - """The dataset is the Public API's catalog row, named without a foreign key.""" +def test_reports_and_simulations_snapshot_selected_datasets() -> None: reports = V2_METADATA.tables["reports"] simulations = V2_METADATA.tables["simulations"] for table in (reports, simulations): - assert isinstance(table.c.dataset_id.type, sa.Uuid) - assert table.c.dataset_id.nullable - assert not table.c.dataset_id.foreign_keys + 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: diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 42027a4c1..33618b73e 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -8,7 +8,6 @@ import sqlalchemy as sa from sqlalchemy.dialects import postgresql -from sqlalchemy.orm import configure_mappers from sqlalchemy.schema import CreateTable from policyengine_api.data.v2.models import ( @@ -18,7 +17,6 @@ ReportRunSimulationRun, SimulationRole, SimulationRun, - User, ) RUN_STATUSES = ["pending", "running", "succeeded", "failed"] @@ -49,31 +47,37 @@ *COMPARISON_RESULT_TABLES, } ) -# Columns naming the Public API's rows, which the simulation service checks -# before it writes the row that names them. -LOGICAL_REFERENCES = { - "simulations": ( - "dataset_id", - "household_id", - "policy_id", - "dynamic_id", - "tax_benefit_model_version_id", - "region_id", +# The foreign keys from the simulation service's tables to the Public API's rows, +# with their delete rules. Each service writes only its own tables, and +# PostgreSQL checks every row named on insert. +PUBLIC_API_FOREIGN_KEYS = { + ("simulations", "dataset_id"): ("datasets", "RESTRICT"), + ("simulations", "household_id"): ("households", "RESTRICT"), + ("simulations", "policy_id"): ("policies", "SET NULL"), + ("simulations", "dynamic_id"): ("dynamics", "SET NULL"), + ("simulations", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", ), - "reports": ( - "user_id", - "tax_benefit_model_id", - "policy_id", - "household_id", - "dataset_id", - "region_id", + ("simulations", "region_id"): ("regions", "SET NULL"), + ("reports", "user_id"): ("users", "SET NULL"), + ("reports", "tax_benefit_model_id"): ("tax_benefit_models", "RESTRICT"), + ("reports", "policy_id"): ("policies", "SET NULL"), + ("reports", "household_id"): ("households", "RESTRICT"), + ("reports", "dataset_id"): ("datasets", "SET NULL"), + ("reports", "region_id"): ("regions", "SET NULL"), + ("simulation_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", ), - "simulation_runs": ("tax_benefit_model_version_id", "dataset_id", "region_id"), - "household_jobs": ( - "household_id", - "policy_id", - "dynamic_id", - "tax_benefit_model_version_id", + ("simulation_runs", "dataset_id"): ("datasets", "RESTRICT"), + ("simulation_runs", "region_id"): ("regions", "SET NULL"), + ("household_jobs", "household_id"): ("households", "RESTRICT"), + ("household_jobs", "policy_id"): ("policies", "SET NULL"), + ("household_jobs", "dynamic_id"): ("dynamics", "SET NULL"), + ("household_jobs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", ), } @@ -426,41 +430,31 @@ def test_custom_totals_keep_their_simulation_references() -> None: ) -def test_no_foreign_key_joins_tables_that_different_services_write() -> None: - """Each service checks the other service's rows its own rows name.""" - cross_service_keys = sorted( - f"{table.name}.{key.parent.name} -> {key.column.table.name}" +def test_simulation_tables_keep_foreign_keys_to_the_rows_they_name() -> None: + """Each service writes only its own tables; the keys between them stay.""" + public_api_keys = { + (table.name, key.parent.name): (key.column.table.name, key.ondelete) for table in V2_METADATA.tables.values() + if table.name in SIMULATION_SIDE_TABLES for key in table.foreign_keys - if (table.name in SIMULATION_SIDE_TABLES) - != (key.column.table.name in SIMULATION_SIDE_TABLES) - ) + if key.column.table.name not in SIMULATION_SIDE_TABLES + } - assert cross_service_keys == [] + assert public_api_keys == PUBLIC_API_FOREIGN_KEYS assert set(V2_METADATA.tables) >= SIMULATION_SIDE_TABLES -def test_references_to_public_api_rows_stay_as_uuid_columns() -> None: - for table_name, column_names in LOGICAL_REFERENCES.items(): - table = V2_METADATA.tables[table_name] - for column_name in column_names: - column = table.c[column_name] - assert isinstance(column.type, sa.Uuid), f"{table_name}.{column_name}" - assert not column.foreign_keys, f"{table_name}.{column_name}" - - -def test_no_relationship_joins_models_that_different_services_write() -> None: - configure_mappers() - cross_service_relationships = sorted( - f"{mapper.class_.__name__}.{relationship.key}" - for mapper in sa.inspect(User).registry.mappers - if mapper.local_table is not None and mapper.local_table.metadata is V2_METADATA - for relationship in mapper.relationships - if (mapper.local_table.name in SIMULATION_SIDE_TABLES) - != (relationship.mapper.local_table.name in SIMULATION_SIDE_TABLES) - ) +def test_reports_index_their_owner_for_account_deletion() -> None: + """Deleting a user clears ``reports.user_id``, which an index finds.""" + reports = V2_METADATA.tables["reports"] + + assert ("user_id",) in _indexes(reports).values() - assert cross_service_relationships == [] + +def test_the_catalog_versions_runs_and_jobs_use_are_indexed() -> None: + for table_name in ("simulation_runs", "household_jobs"): + table = V2_METADATA.tables[table_name] + assert ("tax_benefit_model_version_id",) in _indexes(table).values(), table_name def test_poverty_rows_name_the_group_their_filter_selects() -> None: From d6693ad36f61bf8d1b9ab5d9e70ea7251110e45c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:56:37 +0400 Subject: [PATCH 14/22] Name a v2 report run's baseline and reform simulation runs A report run compares one pair of simulation runs, so report_runs gains baseline_simulation_run_id (required) and reform_simulation_run_id (empty for a household report with only a baseline), each a foreign key to simulation_runs with RESTRICT. report_run_simulation_runs and its v2_simulation_role type are gone, and the comparison result tables' unique keys are per report run again. create_report_run takes the simulation runs a new run compares. Revision 1356f658347d is regenerated; its one correction left is simulation_runs.status reusing v2_simulation_status. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 36 +- ...record_simulation_runs_and_key_results_.py | 387 ++++++------------ policyengine_api/data/v2/models/__init__.py | 4 - policyengine_api/data/v2/models/reports.py | 118 ++---- .../data/v2/models/simulations.py | 4 - policyengine_api/data/v2/privileges.py | 1 - policyengine_api/data/v2/report_runs.py | 20 +- tests/integration/simulation_rows.py | 11 + .../integration/test_alembic_v2_lifecycle.py | 55 ++- .../test_alembic_v2_simulation_runs.py | 91 ++-- .../test_v2_catalog_publication.py | 10 + .../test_v2_simulation_role_privileges.py | 19 +- tests/unit/v2/test_alembic_v2.py | 19 +- tests/unit/v2/test_import_side_effects.py | 2 +- tests/unit/v2/test_models.py | 2 - tests/unit/v2/test_privilege_manifest.py | 1 - tests/unit/v2/test_report_run_creation.py | 16 +- tests/unit/v2/test_report_runs.py | 13 + tests/unit/v2/test_simulation_run_models.py | 72 ++-- 19 files changed, 376 insertions(+), 505 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 8292516af..bff133779 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -508,16 +508,18 @@ model changes: the catalog version, country package version and PolicyEngine.py version it used, and its Modal call ID; `output_dataset_id` is replaced by the run's output file (`output_uri`, `output_sha256`). -- The new `report_run_simulation_runs` names the simulation runs each report - run combines, as its baseline or its reform, and the new - `report_run_comparisons` keeps, per compared pair, the parts of an economy - comparison that have no result table (`labor_supply_result`, - `cliff_impact_result`, `spm_provenance`, `data_version`). +- A report run compares one pair of simulation runs: `report_runs` names its + `baseline_simulation_run_id` (required) and `reform_simulation_run_id` + (empty for a household report with only a baseline), each a foreign key to + `simulation_runs` with `RESTRICT`. The new `report_run_comparisons` keeps + the parts of an economy comparison that have no result table + (`labor_supply_result`, `cliff_impact_result`, `spm_provenance`, + `data_version`). - `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); the other result tables belong to a report run and name the two simulation - runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`), - which join each table's unique key. `aggregates` and `change_aggregates` are - unchanged. + runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`); + each of their rows is unique within its report run. `aggregates` and + `change_aggregates` are unchanged. - The result tables gain the economy output's fields that had no column: the winner, loser and no-change shares of `congressional_district_impacts`, and the UK disclosure control fields (`disclosure_tier`, checked to be @@ -547,18 +549,14 @@ Missing values in the unique keys of `poverty`, `decile_impacts` and `program_statistics` count as equal (`NULLS NOT DISTINCT`, which PostgreSQL supports from version 15). -It has two post-generation corrections, each commented in the revision: +It has one post-generation correction, commented in the revision: +`simulation_runs.status` reuses the `v2_simulation_status` type the baseline +created for `simulations.status`. The generated `create_table` would create +the type again and fail, so the column uses +`postgresql.ENUM(..., create_type=False)`. -1. **Dialect:** `simulation_runs.status` reuses the `v2_simulation_status` - type the baseline created for `simulations.status`. The generated - `create_table` would create the type again and fail, so the column uses - `postgresql.ENUM(..., create_type=False)`. -2. **Reversibility:** the downgrade drops `v2_simulation_role` after dropping - `report_run_simulation_runs`, because Alembic does not drop a table's enum - type with it. - -The upgrade adds required columns without server defaults to `household_jobs` -and the result tables, so it needs them empty, as they are while no managed +The upgrade adds required columns without server defaults to `household_jobs`, +`report_runs` and the result tables, so it needs them empty, as they are while no managed role holds a write privilege on them; on a table with rows it fails, and the deploy stops before anything changes. The downgrade restores required columns the same way, so it needs `simulations`, `household_jobs` and the result diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index b758bd240..83aa7a97d 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -2,17 +2,17 @@ Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-05 23:50:12.713355 +Create Date: 2026-11-05 23:56:01.811595 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate -Post-generation corrections, each commented where it is made: -simulation_runs.status reuses the existing v2_simulation_status type, and the -downgrade drops the v2_simulation_role type the upgrade creates. +One post-generation correction, commented where it is made: +simulation_runs.status reuses the existing v2_simulation_status type. -The upgrade adds required columns without server defaults to household_jobs -and to the result tables, so it needs those tables empty, as they are while no -managed role can write them. The downgrade restores required columns the same -way, so it needs simulations, household_jobs and the result tables empty. +The upgrade adds required columns without server defaults to household_jobs, +report_runs and the result tables, so it needs those tables empty, as they are +while no managed role can write them. The downgrade restores required columns +the same way, so it needs simulations, household_jobs and the result tables +empty. """ from typing import Sequence, Union @@ -200,50 +200,6 @@ def upgrade() -> None: name="uq_report_run_comparisons_run_pair", ), ) - op.create_table( - "report_run_simulation_runs", - sa.Column("id", sa.Uuid(), nullable=False), - sa.Column( - "created_at", - sa.DateTime(timezone=True), - server_default=sa.text("now()"), - nullable=False, - ), - sa.Column("report_run_id", sa.Uuid(), nullable=False), - sa.Column("simulation_run_id", sa.Uuid(), nullable=False), - sa.Column( - "role", - sa.Enum("baseline", "reform", name="v2_simulation_role"), - nullable=False, - ), - sa.ForeignKeyConstraint( - ["report_run_id"], - ["report_runs.id"], - name=op.f("fk_report_run_simulation_runs_report_run_id_report_runs"), - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["simulation_run_id"], - ["simulation_runs.id"], - name=op.f( - "fk_report_run_simulation_runs_simulation_run_id_simulation_runs" - ), - ondelete="RESTRICT", - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_simulation_runs")), - sa.UniqueConstraint( - "report_run_id", - "role", - "simulation_run_id", - name="uq_report_run_simulation_runs_member", - ), - ) - op.create_index( - op.f("ix_report_run_simulation_runs_simulation_run_id"), - "report_run_simulation_runs", - ["simulation_run_id"], - unique=False, - ) op.add_column( "budget_summary", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -252,20 +208,6 @@ def upgrade() -> None: "budget_summary", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) - op.drop_constraint( - op.f("uq_budget_summary_run_variable_entity"), "budget_summary", type_="unique" - ) - op.create_unique_constraint( - "uq_budget_summary_run_variable_entity", - "budget_summary", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "variable_name", - "entity", - ], - ) op.drop_constraint( op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", @@ -292,8 +234,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("budget_summary", "reform_simulation_id") op.drop_column("budget_summary", "baseline_simulation_id") + op.drop_column("budget_summary", "reform_simulation_id") op.add_column( "congressional_district_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -315,27 +257,12 @@ def upgrade() -> None: sa.Column("no_change_percentage", sa.Float(), nullable=False), ) op.drop_constraint( - op.f("uq_congressional_district_impacts_run_geoid"), - "congressional_district_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_congressional_district_impacts_run_geoid", - "congressional_district_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "district_geoid", - ], - ) - op.drop_constraint( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", type_="foreignkey", ) @@ -359,8 +286,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("congressional_district_impacts", "reform_simulation_id") op.drop_column("congressional_district_impacts", "baseline_simulation_id") + op.drop_column("congressional_district_impacts", "reform_simulation_id") op.add_column( "constituency_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -393,19 +320,6 @@ def upgrade() -> None: "constituency_impacts", sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) - op.drop_constraint( - op.f("uq_constituency_impacts_run_code"), "constituency_impacts", type_="unique" - ) - op.create_unique_constraint( - "uq_constituency_impacts_run_code", - "constituency_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "constituency_code", - ], - ) op.drop_constraint( op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", @@ -437,8 +351,8 @@ def upgrade() -> None: "constituency_impacts", "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", ) - op.drop_column("constituency_impacts", "reform_simulation_id") op.drop_column("constituency_impacts", "baseline_simulation_id") + op.drop_column("constituency_impacts", "reform_simulation_id") op.add_column( "decile_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -463,8 +377,6 @@ def upgrade() -> None: "decile_impacts", [ "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "income_variable", "decile_variable", "entity", @@ -484,18 +396,18 @@ def upgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -590,12 +502,12 @@ def upgrade() -> None: ["simulation_run_id", "income_variable", "entity"], ) op.drop_constraint( - op.f("fk_inequality_simulation_id_simulations"), + op.f("fk_inequality_report_run_id_report_runs"), "inequality", type_="foreignkey", ) op.drop_constraint( - op.f("fk_inequality_report_run_id_report_runs"), + op.f("fk_inequality_simulation_id_simulations"), "inequality", type_="foreignkey", ) @@ -617,22 +529,6 @@ def upgrade() -> None: "intra_decile_impacts", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) - op.drop_constraint( - op.f("uq_intra_decile_impacts_run_type_decile"), - "intra_decile_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_intra_decile_impacts_run_type_decile", - "intra_decile_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "decile_type", - "decile", - ], - ) op.drop_constraint( op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", @@ -659,8 +555,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("intra_decile_impacts", "reform_simulation_id") op.drop_column("intra_decile_impacts", "baseline_simulation_id") + op.drop_column("intra_decile_impacts", "reform_simulation_id") op.add_column( "local_authority_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -694,43 +590,28 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("uq_local_authority_impacts_run_code"), - "local_authority_impacts", - type_="unique", - ) - op.create_unique_constraint( - "uq_local_authority_impacts_run_code", - "local_authority_impacts", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "local_authority_code", - ], - ) - op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -739,8 +620,8 @@ def upgrade() -> None: "local_authority_impacts", "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", ) - op.drop_column("local_authority_impacts", "reform_simulation_id") op.drop_column("local_authority_impacts", "baseline_simulation_id") + op.drop_column("local_authority_impacts", "reform_simulation_id") op.add_column("poverty", sa.Column("simulation_run_id", sa.Uuid(), nullable=False)) op.add_column( "poverty", @@ -802,13 +683,7 @@ def upgrade() -> None: op.create_unique_constraint( "uq_program_statistics_run_program_entity", "program_statistics", - [ - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - "program_name", - "entity", - ], + ["report_run_id", "program_name", "entity"], postgresql_nulls_not_distinct=True, ) op.drop_constraint( @@ -822,35 +697,70 @@ def upgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) - op.drop_column("program_statistics", "reform_simulation_id") op.drop_column("program_statistics", "baseline_simulation_id") + op.drop_column("program_statistics", "reform_simulation_id") + op.add_column( + "report_runs", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "report_runs", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=True) + ) op.add_column( "report_runs", sa.Column( "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True ), ) + op.create_index( + op.f("ix_report_runs_baseline_simulation_run_id"), + "report_runs", + ["baseline_simulation_run_id"], + unique=False, + ) + op.create_index( + op.f("ix_report_runs_reform_simulation_run_id"), + "report_runs", + ["reform_simulation_run_id"], + unique=False, + ) op.create_index( "ix_report_runs_status_started_at", "report_runs", ["status", "started_at"], unique=False, ) + op.create_foreign_key( + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + "report_runs", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "report_runs", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) op.create_index(op.f("ix_reports_user_id"), "reports", ["user_id"], unique=False) op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") op.drop_constraint( @@ -861,9 +771,9 @@ def upgrade() -> None: op.drop_column("simulations", "status") op.drop_column("simulations", "output_dataset_id") op.drop_column("simulations", "started_at") + op.drop_column("simulations", "completed_at") op.drop_column("simulations", "error_message") op.drop_column("simulations", "household_result") - op.drop_column("simulations", "completed_at") # ### end Alembic commands ### @@ -872,25 +782,25 @@ def downgrade() -> None: op.add_column( "simulations", sa.Column( - "completed_at", - postgresql.TIMESTAMP(timezone=True), + "household_result", + postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, ), ) + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + ) op.add_column( "simulations", sa.Column( - "household_result", - postgresql.JSON(astext_type=sa.Text()), + "completed_at", + postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, ), ) - op.add_column( - "simulations", - sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), - ) op.add_column( "simulations", sa.Column( @@ -930,27 +840,45 @@ def downgrade() -> None: unique=False, ) op.drop_index(op.f("ix_reports_user_id"), table_name="reports") + op.drop_constraint( + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "report_runs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + "report_runs", + type_="foreignkey", + ) op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") + op.drop_index( + op.f("ix_report_runs_reform_simulation_run_id"), table_name="report_runs" + ) + op.drop_index( + op.f("ix_report_runs_baseline_simulation_run_id"), table_name="report_runs" + ) op.drop_column("report_runs", "modal_call_id") + op.drop_column("report_runs", "reform_simulation_run_id") + op.drop_column("report_runs", "baseline_simulation_run_id") op.add_column( "program_statistics", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "program_statistics", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) @@ -1031,13 +959,13 @@ def downgrade() -> None: op.add_column( "local_authority_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "local_authority_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1046,40 +974,31 @@ def downgrade() -> None: type_="check", ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) - op.drop_constraint( - "uq_local_authority_impacts_run_code", "local_authority_impacts", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_local_authority_impacts_run_code"), - "local_authority_impacts", - ["report_run_id", "local_authority_code"], - postgresql_nulls_not_distinct=False, - ) op.drop_column("local_authority_impacts", "relative_household_income_change_upper") op.drop_column("local_authority_impacts", "relative_household_income_change_lower") op.drop_column("local_authority_impacts", "average_household_income_change_upper") @@ -1090,13 +1009,13 @@ def downgrade() -> None: op.add_column( "intra_decile_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "intra_decile_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1125,17 +1044,6 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_constraint( - "uq_intra_decile_impacts_run_type_decile", - "intra_decile_impacts", - type_="unique", - ) - op.create_unique_constraint( - op.f("uq_intra_decile_impacts_run_type_decile"), - "intra_decile_impacts", - ["report_run_id", "decile_type", "decile"], - postgresql_nulls_not_distinct=False, - ) op.drop_column("intra_decile_impacts", "reform_simulation_run_id") op.drop_column("intra_decile_impacts", "baseline_simulation_run_id") op.add_column( @@ -1151,14 +1059,6 @@ def downgrade() -> None: "inequality", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_inequality_report_run_id_report_runs"), - "inequality", - "report_runs", - ["report_run_id"], - ["id"], - ondelete="CASCADE", - ) op.create_foreign_key( op.f("fk_inequality_simulation_id_simulations"), "inequality", @@ -1167,6 +1067,14 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) op.drop_constraint( "uq_inequality_simulation_run_measure", "inequality", type_="unique" ) @@ -1218,12 +1126,12 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) @@ -1258,13 +1166,13 @@ def downgrade() -> None: op.add_column( "constituency_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "constituency_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1298,15 +1206,6 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_constraint( - "uq_constituency_impacts_run_code", "constituency_impacts", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_constituency_impacts_run_code"), - "constituency_impacts", - ["report_run_id", "constituency_code"], - postgresql_nulls_not_distinct=False, - ) op.drop_column("constituency_impacts", "relative_household_income_change_upper") op.drop_column("constituency_impacts", "relative_household_income_change_lower") op.drop_column("constituency_impacts", "average_household_income_change_upper") @@ -1317,13 +1216,13 @@ def downgrade() -> None: op.add_column( "congressional_district_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "congressional_district_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1341,32 +1240,21 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) - op.drop_constraint( - "uq_congressional_district_impacts_run_geoid", - "congressional_district_impacts", - type_="unique", - ) - op.create_unique_constraint( - op.f("uq_congressional_district_impacts_run_geoid"), - "congressional_district_impacts", - ["report_run_id", "district_geoid"], - postgresql_nulls_not_distinct=False, - ) op.drop_column("congressional_district_impacts", "no_change_percentage") op.drop_column("congressional_district_impacts", "loser_percentage") op.drop_column("congressional_district_impacts", "winner_percentage") @@ -1375,13 +1263,13 @@ def downgrade() -> None: op.add_column( "budget_summary", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "budget_summary", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1410,25 +1298,8 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_constraint( - "uq_budget_summary_run_variable_entity", "budget_summary", type_="unique" - ) - op.create_unique_constraint( - op.f("uq_budget_summary_run_variable_entity"), - "budget_summary", - ["report_run_id", "variable_name", "entity"], - postgresql_nulls_not_distinct=False, - ) op.drop_column("budget_summary", "reform_simulation_run_id") op.drop_column("budget_summary", "baseline_simulation_run_id") - op.drop_index( - op.f("ix_report_run_simulation_runs_simulation_run_id"), - table_name="report_run_simulation_runs", - ) - op.drop_table("report_run_simulation_runs") - # Post-generation reversibility correction: Alembic creates the - # v2_simulation_role type with its table but does not drop it with it. - sa.Enum(name="v2_simulation_role").drop(op.get_bind()) op.drop_table("report_run_comparisons") op.drop_index( op.f("ix_simulation_runs_tax_benefit_model_version_id"), diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index cc182268d..87f73fb6e 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -83,10 +83,8 @@ Report, ReportRun, ReportRunComparison, - ReportRunSimulationRun, ReportRunStatus, ReportRunTrigger, - SimulationRole, ) from policyengine_api.data.v2.models.simulations import ( # noqa: E402 Simulation, @@ -143,11 +141,9 @@ "Report", "ReportRun", "ReportRunComparison", - "ReportRunSimulationRun", "ReportRunStatus", "ReportRunTrigger", "Simulation", - "SimulationRole", "SimulationRun", "SimulationStatus", "SimulationType", diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 9427ffbb5..be2a9e2ac 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,10 +1,9 @@ """Canonical SQLModel tables for stable reports, runs, and run outputs. -Results that compare a baseline with a reform belong to a report run and name -the two simulation runs they compare; each such row is unique within its report -run and pair of simulation runs, because a budget-window run compares one pair -per year. Poverty and inequality, which describe one simulation, belong to a -simulation run. +A report run compares one baseline simulation run with one reform simulation +run. Results that compare them belong to the report run and name the two +simulation runs; each such row is unique within its report run. Poverty and +inequality, which describe one simulation, belong to a simulation run. """ from datetime import datetime @@ -72,13 +71,6 @@ class DecileType(StrEnum): WEALTH = "wealth" -class SimulationRole(StrEnum): - """The part a simulation run plays in the report run that combines it.""" - - BASELINE = "baseline" - REFORM = "reform" - - class Report(TimestampedModel, table=True): """Stable report definition; execution state belongs to ReportRun. @@ -170,9 +162,10 @@ class Report(TimestampedModel, table=True): class ReportRun(TimestampedModel, table=True): """One immutable-version execution attempt for a stable report. - Its ``simulation_run_members`` name the simulation runs it combines, and - its ``comparisons`` hold, per compared pair, the parts of an economy - comparison that have no result table. + It compares ``baseline_simulation_run_id`` with + ``reform_simulation_run_id``; a household report with only a baseline has + no reform run. Its ``comparisons`` hold the parts of an economy comparison + that have no result table. """ __tablename__ = "report_runs" @@ -202,6 +195,17 @@ class ReportRun(TimestampedModel, table=True): foreign_key="reports.id", ondelete="CASCADE", ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + index=True, + ) + reform_simulation_run_id: UUID | None = Field( + default=None, + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + index=True, + ) country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) status: ReportRunStatus = Field( @@ -226,10 +230,6 @@ class ReportRun(TimestampedModel, table=True): markdown: str | None = Field(default=None, sa_type=sa.Text) report: Report = Relationship(back_populates="runs") - simulation_run_members: list["ReportRunSimulationRun"] = Relationship( - back_populates="report_run", - cascade_delete=True, - ) comparisons: list["ReportRunComparison"] = Relationship( back_populates="report_run", cascade_delete=True, @@ -271,42 +271,6 @@ class ReportRun(TimestampedModel, table=True): ) -class ReportRunSimulationRun(IdentifiedModel, table=True): - """One simulation run a report run combines, as its baseline or its reform. - - A single-year report run names one of each (a household report may name - only a baseline); a budget-window run names one of each per year. - """ - - __tablename__ = "report_run_simulation_runs" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "role", - "simulation_run_id", - name="uq_report_run_simulation_runs_member", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - index=True, - ) - role: SimulationRole = Field( - sa_type=enum_type(SimulationRole, "v2_simulation_role"), - ) - - report_run: ReportRun = Relationship(back_populates="simulation_run_members") - simulation_run: SimulationRun = Relationship( - back_populates="report_run_memberships" - ) - - class ReportRunComparison(IdentifiedModel, table=True): """The parts of one compared pair of a report run that have no result table. @@ -430,15 +394,13 @@ class ChangeAggregate(IdentifiedModel, table=True): class BudgetSummary(IdentifiedModel, table=True): """Baseline and reform totals of one budget variable for a report run. - One row per report run, pair of simulation runs, variable and entity. + One row per report run, variable and entity. """ __tablename__ = "budget_summary" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "variable_name", "entity", name="uq_budget_summary_run_variable_entity", @@ -470,18 +432,16 @@ 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 report run, - pair of simulation runs, income variable, decile variable, entity, quantile - count and decile. ``decile_variable`` names a variable that groups the - households instead of the income variable's own deciles, such as - ``household_wealth_decile``; rows without one count as equal. + income variable, decile variable, entity, quantile count and decile. + ``decile_variable`` names a variable that groups the households instead of + the income variable's own deciles, such as ``household_wealth_decile``; rows + without one count as equal. """ __tablename__ = "decile_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "income_variable", "decile_variable", "entity", @@ -528,16 +488,14 @@ class DecileImpact(IdentifiedModel, table=True): class IntraDecileImpact(IdentifiedModel, table=True): """Shares of one decile whose income rises or falls by more or less than 5%. - One row per report run, pair of simulation runs, decile type and decile, - where ``decile`` is 0 to 10. + One row per report run, decile type and decile, where ``decile`` is 0 to + 10. """ __tablename__ = "intra_decile_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "decile_type", "decile", name="uq_intra_decile_impacts_run_type_decile", @@ -641,18 +599,16 @@ class Poverty(IdentifiedModel, table=True): class ProgramStatistics(IdentifiedModel, table=True): """Baseline and reform totals, counts, winners and losers for one program. - One row per report run, pair of simulation runs, program and entity; - ``is_tax`` marks taxes. The economy - output's detailed budget names neither the entity nor whether a program is - a tax, so both may be null, and rows without an entity count as equal. + One row per report run, program and entity; ``is_tax`` marks taxes. The + economy output's detailed budget names neither the entity nor whether a + program is a tax, so both may be null, and rows without an entity count as + equal. """ __tablename__ = "program_statistics" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "program_name", "entity", name="uq_program_statistics_run_program_entity", @@ -731,16 +687,14 @@ def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: class CongressionalDistrictImpact(GeographicImpactBase, table=True): """Household income change and population in one US congressional district. - One row per report run, pair of simulation runs and district GEOID, with the - shares of people who gain, lose or see no change. + One row per report run and district GEOID, with the shares of people who + gain, lose or see no change. """ __tablename__ = "congressional_district_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "district_geoid", name="uq_congressional_district_impacts_run_geoid", ), @@ -761,15 +715,13 @@ class CongressionalDistrictImpact(GeographicImpactBase, table=True): class ConstituencyImpact(UKAreaImpactBase, table=True): """Household income change in one UK constituency. - One row per report run, pair of simulation runs and constituency code. + One row per report run and constituency code. """ __tablename__ = "constituency_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "constituency_code", name="uq_constituency_impacts_run_code", ), @@ -785,15 +737,13 @@ class ConstituencyImpact(UKAreaImpactBase, table=True): class LocalAuthorityImpact(UKAreaImpactBase, table=True): """Household income change in one UK local authority. - One row per report run, pair of simulation runs and local authority code. + One row per report run and local authority code. """ __tablename__ = "local_authority_impacts" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", "local_authority_code", name="uq_local_authority_impacts_run_code", ), diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index ca84d91fb..51dd99771 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -22,7 +22,6 @@ Inequality, Poverty, Report, - ReportRunSimulationRun, ) @@ -215,9 +214,6 @@ class SimulationRun(TimestampedModel, table=True): output_sha256: str | None = Field(default=None, max_length=64) simulation: Simulation = Relationship(back_populates="runs") - report_run_memberships: list["ReportRunSimulationRun"] = Relationship( - back_populates="simulation_run" - ) poverty_results: list["Poverty"] = Relationship( back_populates="simulation_run", cascade_delete=True, diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 2dcef8ff8..98cefb580 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -149,7 +149,6 @@ def _access( # Simulations, reports, their runs and household jobs. The catalog writer # checks that publication leaves the rows of three of them untouched. "household_jobs": _RUN_STATUS, - "report_run_simulation_runs": _SIMULATION_DEFINITION, "report_runs": replace(_RUN_STATUS, catalog_writer=READ), "reports": replace(_SIMULATION_DEFINITION, catalog_writer=READ), "simulation_runs": _RUN_STATUS, diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 1b1debd61..38af95c2e 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -53,6 +53,17 @@ class ReportRunPackageVersions: policyengine_version: str +@dataclass(frozen=True) +class ReportRunSimulationRuns: + """The simulation runs one report run compares. + + A household report with only a baseline has no reform run. + """ + + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + 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. @@ -94,9 +105,14 @@ def create_report_run( report_id: UUID, versions: ReportRunPackageVersions, trigger: ReportRunTrigger, + simulation_runs: ReportRunSimulationRuns, idempotency_key: UUID | None = None, ) -> ReportRun: - """Create one run, or return the run for a retried idempotent request.""" + """Create one run, or return the run for a retried idempotent request. + + A retried request gets the run its first attempt created, with the + simulation runs that attempt named. + """ _locked_report(session, report_id) if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: msg = "manual report reruns require an idempotency key" @@ -107,6 +123,8 @@ def create_report_run( run = ReportRun( report_id=report_id, + baseline_simulation_run_id=simulation_runs.baseline_simulation_run_id, + reform_simulation_run_id=simulation_runs.reform_simulation_run_id, country_package_version=versions.country_package_version, policyengine_version=versions.policyengine_version, trigger=trigger, diff --git a/tests/integration/simulation_rows.py b/tests/integration/simulation_rows.py index 4a6767dc1..214914f8e 100644 --- a/tests/integration/simulation_rows.py +++ b/tests/integration/simulation_rows.py @@ -101,3 +101,14 @@ def refused( with pytest.raises(error, match=match): yield savepoint.rollback() + + +def delete_public_api_rows(connection: Connection, rows: PublicApiRows) -> None: + """Delete what ``insert_public_api_rows`` stored, once nothing names it.""" + for table, row_id in ( + ("households", rows.household_id), + ("datasets", rows.dataset_id), + ("tax_benefit_model_versions", rows.version_id), + ("tax_benefit_models", rows.model_id), + ): + connection.execute(text(f"DELETE FROM {table} WHERE id = :id"), {"id": row_id}) diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 3c2cc8161..18feb946b 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -25,6 +25,13 @@ disposable_v2_url, table_columns, ) +from tests.integration.simulation_rows import ( + RUN_VERSIONS, + PublicApiRows, + delete_public_api_rows, + insert_public_api_rows, + insert_row, +) BASELINE_REVISION = "f5ef4347cb2a" STAGE_11_PREVIOUS_REVISION = "af34023a728f" @@ -886,19 +893,35 @@ def report_run_checks() -> set[str]: for constraint in inspect(engine).get_check_constraints("report_runs") } + report_id = uuid4() + report_run_id = uuid4() + simulation_id = uuid4() + request_key = uuid4() + public_api_rows: PublicApiRows | None = None try: command.upgrade(config, "head") assert isinstance(idempotency_column_type(), postgresql.UUID) assert "ck_report_runs_idempotency_key_nonblank" not in report_run_checks() - model_id = uuid4() - report_id = uuid4() - report_run_id = uuid4() - request_key = uuid4() with engine.begin() as connection: - connection.execute( - text("INSERT INTO tax_benefit_models (id, name) VALUES (:id, :name)"), - {"id": model_id, "name": f"uuid-cast-{model_id.hex[:8]}"}, + public_api_rows = insert_public_api_rows(connection) + insert_row( + connection, + "simulations", + id=simulation_id, + simulation_type="economy", + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, + ) + simulation_run_id = uuid4() + insert_row( + connection, + "simulation_runs", + id=simulation_run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=public_api_rows.version_id, + status="pending", + **RUN_VERSIONS, ) connection.execute( text( @@ -906,19 +929,21 @@ def report_run_checks() -> set[str]: "(id, label, country, tax_benefit_model_id, inputs) " "VALUES (:id, 'UUID cast report', 'us', :model_id, '{}')" ), - {"id": report_id, "model_id": model_id}, + {"id": report_id, "model_id": public_api_rows.model_id}, ) connection.execute( text( "INSERT INTO report_runs " - "(id, report_id, country_package_version, " - "policyengine_version, status, trigger, idempotency_key) " - "VALUES (:id, :report_id, '1.0', '1.0', 'pending', " - "'manual', :request_key)" + "(id, report_id, baseline_simulation_run_id, " + "country_package_version, policyengine_version, status, " + "trigger, idempotency_key) " + "VALUES (:id, :report_id, :simulation_run_id, '1.0', '1.0', " + "'pending', 'manual', :request_key)" ), { "id": report_run_id, "report_id": report_id, + "simulation_run_id": simulation_run_id, "request_key": str(request_key), }, ) @@ -941,9 +966,11 @@ def report_run_checks() -> set[str]: {"report_id": report_id}, ) connection.execute( - text("DELETE FROM tax_benefit_models WHERE id = :model_id"), - {"model_id": model_id}, + text("DELETE FROM simulations WHERE id = :simulation_id"), + {"simulation_id": simulation_id}, ) + if public_api_rows is not None: + delete_public_api_rows(connection, public_api_rows) command.upgrade(config, "head") engine.dispose() diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 5862861aa..6040077c5 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -42,9 +42,7 @@ from sqlalchemy import Connection, Engine SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" -NEW_TABLES = frozenset( - {"simulation_runs", "report_run_simulation_runs", "report_run_comparisons"} -) +NEW_TABLES = frozenset({"simulation_runs", "report_run_comparisons"}) CHANGED_TABLES = frozenset( { "budget_summary", @@ -227,13 +225,13 @@ def test_the_revision_downgrades_to_its_parent_and_upgrades_again( if column["name"] == "entity" ) assert not program_entity["nullable"] - enum_types = _enum_type_names(engine) - assert "v2_simulation_status" in enum_types - assert "v2_simulation_role" not in enum_types + assert "v2_simulation_status" in _enum_type_names(engine) command.upgrade(config, "head") command.check(config) - assert "v2_simulation_role" in _enum_type_names(engine) + assert {"baseline_simulation_run_id", "reform_simulation_run_id"} <= ( + _column_names(engine, "report_runs") + ) def test_head_keeps_the_foreign_keys_to_the_public_api_tables( @@ -313,8 +311,12 @@ def _insert_simulation_run( def _insert_report_run( - connection: Connection, public_api_rows: PublicApiRows + connection: Connection, + public_api_rows: PublicApiRows, + baseline_run_id: object, + reform_run_id: object | None = None, ) -> object: + """Store a report and one run of it comparing the two simulation runs.""" report_id, report_run_id = uuid4(), uuid4() insert_row( connection, @@ -330,6 +332,8 @@ def _insert_report_run( "report_runs", id=report_run_id, report_id=report_id, + baseline_simulation_run_id=baseline_run_id, + reform_simulation_run_id=reform_run_id, status="running", trigger="initial", **RUN_VERSIONS, @@ -414,33 +418,31 @@ def test_simulation_runs_need_a_completion_time_and_a_whole_output( assert _count(connection, "poverty", "simulation_run_id", finished_run_id) == 0 -def test_report_run_members_are_unique_and_keep_their_simulation_runs( +def test_report_runs_name_a_baseline_run_and_keep_their_simulation_runs( connection: Connection, public_api_rows: PublicApiRows ) -> None: - run_id = _insert_simulation_run( - connection, _insert_simulation(connection, public_api_rows) - ) - report_run_id = _insert_report_run(connection, public_api_rows) - member = { - "report_run_id": report_run_id, - "simulation_run_id": run_id, - "role": "baseline", - } - insert_row(connection, "report_run_simulation_runs", id=uuid4(), **member) + simulation_id = _insert_simulation(connection, public_api_rows) + baseline_run_id = _insert_simulation_run(connection, simulation_id) + reform_run_id = _insert_simulation_run(connection, simulation_id) - with refused(connection): - insert_row(connection, "report_run_simulation_runs", id=uuid4(), **member) - with refused(connection): - connection.execute( - text("DELETE FROM simulation_runs WHERE id = :id"), {"id": run_id} - ) + with refused(connection, match="baseline_simulation_run_id"): + _insert_report_run(connection, public_api_rows, None) + with refused(connection, match="foreign key constraint"): + _insert_report_run(connection, public_api_rows, uuid4()) + # A household report with only a baseline names no reform run. + _insert_report_run(connection, public_api_rows, baseline_run_id) + report_run_id = _insert_report_run( + connection, public_api_rows, baseline_run_id, reform_run_id + ) + for run_id in (baseline_run_id, reform_run_id): + with refused(connection, match="foreign key constraint"): + connection.execute( + text("DELETE FROM simulation_runs WHERE id = :id"), {"id": run_id} + ) connection.execute( text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} ) - assert ( - _count(connection, "report_run_simulation_runs", "report_run_id", report_run_id) - == 0 - ) + assert _count(connection, "simulation_runs", "id", reform_run_id) == 1 def test_results_are_unique_per_run_and_name_runs_that_exist( @@ -450,7 +452,9 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( baseline_run_id = _insert_simulation_run(connection, simulation_id) reform_run_id = _insert_simulation_run(connection, simulation_id) comparison_runs = { - "report_run_id": _insert_report_run(connection, public_api_rows), + "report_run_id": _insert_report_run( + connection, public_api_rows, baseline_run_id, reform_run_id + ), "baseline_simulation_run_id": baseline_run_id, "reform_simulation_run_id": reform_run_id, } @@ -500,26 +504,25 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( insert_row(connection, "decile_impacts", id=uuid4(), **decile) budget = {"variable_name": "household_tax", "entity": "household"} insert_row(connection, "budget_summary", id=uuid4(), **budget, **comparison_runs) - next_year_runs = { - **comparison_runs, - "baseline_simulation_run_id": _insert_simulation_run(connection, simulation_id), - "reform_simulation_run_id": _insert_simulation_run(connection, simulation_id), - } - insert_row(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) + # A report run compares one pair, so its rows are unique within it. with refused(connection): - insert_row(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) - # One comparison row per compared pair holds the untabled parts. - for runs in (comparison_runs, next_year_runs): insert_row( connection, - "report_run_comparisons", + "budget_summary", id=uuid4(), - labor_supply_result="{}", - data_version="1.0.0", - **runs, + **budget, + **{**comparison_runs, "reform_simulation_run_id": baseline_run_id}, ) + insert_row( + connection, + "report_run_comparisons", + id=uuid4(), + labor_supply_result="{}", + data_version="1.0.0", + **comparison_runs, + ) with refused(connection): - insert_row(connection, "report_run_comparisons", id=uuid4(), **next_year_runs) + insert_row(connection, "report_run_comparisons", id=uuid4(), **comparison_runs) insert_row( connection, "program_statistics", diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index 17078dceb..10bd85d74 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -36,6 +36,7 @@ Report, ReportRun, Simulation, + SimulationRun, TaxBenefitModel, TaxBenefitModelVersion, ) @@ -512,8 +513,17 @@ def test_republish_preserves_existing_run_and_dataset_version_records( ) session.add_all((dataset_version, simulation, report)) session.flush() + simulation_run = SimulationRun( + simulation_id=simulation.id, + tax_benefit_model_version_id=model_version.id, + country_package_version="existing-country-version", + policyengine_version="existing-policyengine-version", + ) + session.add(simulation_run) + session.flush() report_run = ReportRun( report_id=report.id, + baseline_simulation_run_id=simulation_run.id, country_package_version="existing-country-version", policyengine_version="existing-policyengine-version", ) diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index 6152dfb37..503bf73e0 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -111,18 +111,11 @@ def _create_simulation_rows( "report_runs", id=ids["report_run"], report_id=ids["report"], + baseline_simulation_run_id=ids["run"], status="pending", trigger="initial", **RUN_VERSIONS, ) - insert_row( - connection, - "report_run_simulation_runs", - id=uuid4(), - report_run_id=ids["report_run"], - simulation_run_id=ids["run"], - role="baseline", - ) insert_row( connection, "household_jobs", @@ -248,14 +241,14 @@ def test_the_public_api_reads_runs_and_results_but_writes_neither( ) -> None: _act_as(connection, managed_v2_roles.runtime) - for table in ("simulation_runs", "report_run_simulation_runs", "poverty"): + for table in ("simulation_runs", "report_runs", "poverty"): connection.execute(text(f"SELECT count(*) FROM {table}")).scalar_one() _refused( connection, - "INSERT INTO report_run_simulation_runs " - "(id, report_run_id, simulation_run_id, role) " - "VALUES (:id, :report_run_id, :run_id, 'baseline')", + "INSERT INTO report_runs (id, report_id, baseline_simulation_run_id, " + "country_package_version, policyengine_version, status, trigger) " + "VALUES (:id, :report_id, :run_id, '1.0.0', '4.0.0', 'pending', 'initial')", id=uuid4(), - report_run_id=uuid4(), + report_id=uuid4(), run_id=uuid4(), ) diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index f8df401c4..54f215b76 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -48,7 +48,6 @@ "parameter_search_files", "parameter_tree_children", "report_run_comparisons", - "report_run_simulation_runs", "simulation_runs", "stage12_evaluation_reports", "stage12_evaluation_simulations", @@ -1023,14 +1022,11 @@ def _simulation_runs_revision_source() -> str: return path.read_text(encoding="utf-8") -def test_simulation_runs_revision_is_generated_with_two_documented_corrections() -> ( - None -): +def test_simulation_runs_revision_is_generated_with_one_documented_correction() -> None: """One generated revision follows the parameter search file revision. - Its two corrections handle enum types: ``simulation_runs.status`` reuses the - type the baseline created for ``simulations.status``, and the downgrade - drops the role type the upgrade creates with ``report_run_simulation_runs``. + Its one correction makes ``simulation_runs.status`` reuse the type the + baseline created for ``simulations.status``. """ revision = _simulation_runs_revision_source() upgrade = revision[ @@ -1047,16 +1043,11 @@ def test_simulation_runs_revision_is_generated_with_two_documented_corrections() in revision ) assert revision.count("Dialect correction:") == 1 - assert revision.count("Post-generation reversibility correction:") == 1 + assert "reversibility correction" not in revision assert 'name="v2_simulation_status",\n create_type=False,' in ( upgrade ) - assert 'sa.Enum(name="v2_simulation_role").drop(op.get_bind())' in downgrade - for table_name in ( - "simulation_runs", - "report_run_simulation_runs", - "report_run_comparisons", - ): + for table_name in ("simulation_runs", "report_run_comparisons"): assert f'op.create_table(\n "{table_name}"' in upgrade assert f'op.drop_table("{table_name}")' in downgrade assert upgrade.count("op.drop_table(") == 0 diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index df7828892..13325fd74 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 45 +assert len(V2_METADATA.tables) == 44 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index db60cd8cf..6b83ff44e 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -23,7 +23,6 @@ ParameterTreeChild, ParameterValue, Policy, - ReportRunSimulationRun, Simulation, SimulationRun, TaxBenefitModelVersion, @@ -77,7 +76,6 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: LegacyHouseholdMapping: "household_mappings", Simulation: "simulations", SimulationRun: "simulations", - ReportRunSimulationRun: "reports", UserHouseholdAssociation: "associations", UserPolicy: "associations", LegacyPolicyMapping: "policy_mappings", diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index f38a659d4..94dc35993 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -46,7 +46,6 @@ "reports", "simulation_runs", "report_runs", - "report_run_simulation_runs", "household_jobs", } ) diff --git a/tests/unit/v2/test_report_run_creation.py b/tests/unit/v2/test_report_run_creation.py index 7aa06a392..b1e1d8f0b 100644 --- a/tests/unit/v2/test_report_run_creation.py +++ b/tests/unit/v2/test_report_run_creation.py @@ -3,7 +3,7 @@ from __future__ import annotations from contextlib import contextmanager -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.exc import IntegrityError @@ -21,9 +21,16 @@ from policyengine_api.data.v2.report_runs import ( ReportNotFoundError, ReportRunPackageVersions, + ReportRunSimulationRuns, create_report_run, ) +# The simulation runs each report run compares. SQLite doesn't enforce foreign +# keys, so these tests store no simulation run rows for them. +SIMULATION_RUNS = ReportRunSimulationRuns( + baseline_simulation_run_id=UUID(int=101), reform_simulation_run_id=UUID(int=102) +) + @pytest.fixture def engine(): @@ -65,6 +72,7 @@ def test_unknown_report_is_rejected_before_any_validation(engine) -> None: policyengine_version="", ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, ) @@ -83,6 +91,7 @@ def test_manual_run_without_a_key_is_rejected_before_version_checks(engine) -> N policyengine_version="", ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, ) @@ -109,6 +118,7 @@ def test_empty_package_versions_are_rejected( policyengine_version=policyengine_version, ), trigger=ReportRunTrigger.SYSTEM, + simulation_runs=SIMULATION_RUNS, ) @@ -124,6 +134,7 @@ def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: policyengine_version="4.5.6", ), trigger=ReportRunTrigger.INITIAL, + simulation_runs=SIMULATION_RUNS, ) stored = session.exec(select(ReportRun)).one() @@ -134,6 +145,8 @@ def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: assert run.trigger is ReportRunTrigger.INITIAL assert run.idempotency_key is None assert run.status is ReportRunStatus.PENDING + assert run.baseline_simulation_run_id == UUID(int=101) + assert run.reform_simulation_run_id == UUID(int=102) def test_conflict_without_a_matching_keyed_run_reraises( @@ -160,6 +173,7 @@ def failing_savepoint(): policyengine_version="4.5.6", ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, idempotency_key=uuid4(), ) diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py index 0221624e5..89f83282d 100644 --- a/tests/unit/v2/test_report_runs.py +++ b/tests/unit/v2/test_report_runs.py @@ -27,6 +27,7 @@ ) from policyengine_api.data.v2.report_runs import ( ReportRunPackageVersions, + ReportRunSimulationRuns, ReportRunStateError, ReportTypeImmutableError, begin_report_run, @@ -40,6 +41,11 @@ DEPLOYED_COUNTRY_VERSIONS = {"us": "1.2.3"} DEPLOYED_POLICYENGINE_VERSION = "4.5.6" NOW = datetime(2026, 8, 14, 12, 0, tzinfo=UTC) +# The simulation runs each report run compares. SQLite doesn't enforce foreign +# keys, so these tests store no simulation run rows for them. +SIMULATION_RUNS = ReportRunSimulationRuns( + baseline_simulation_run_id=UUID(int=101), reform_simulation_run_id=UUID(int=102) +) @pytest.fixture @@ -86,6 +92,7 @@ def _create_run( policyengine_version=policyengine_version, ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, idempotency_key=key if key is not None else uuid4(), ) @@ -152,6 +159,7 @@ def test_manual_rerun_requires_an_idempotency_key(engine) -> None: policyengine_version=DEPLOYED_POLICYENGINE_VERSION, ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, ) @@ -214,6 +222,7 @@ def request_rerun() -> UUID: policyengine_version=DEPLOYED_POLICYENGINE_VERSION, ), trigger=ReportRunTrigger.MANUAL, + simulation_runs=SIMULATION_RUNS, idempotency_key=request_key, ) run_id = run.id @@ -310,6 +319,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=1), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -318,6 +328,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=2), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -326,6 +337,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=3), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="9.9.9", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -334,6 +346,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=4), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.RUNNING, diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 33618b73e..4b73b57a7 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -2,8 +2,8 @@ A simulation defines what to calculate and each calculation of it is a run that records the catalog and PolicyEngine.py versions it used. A report run names -the simulation runs it combines, results are keyed by run, and a household job -calculates one household for one year under one policy. +its baseline and reform simulation runs, results are keyed by run, and a +household job calculates one household for one year under one policy. """ import sqlalchemy as sa @@ -14,8 +14,6 @@ V2_METADATA, ReportRun, ReportRunComparison, - ReportRunSimulationRun, - SimulationRole, SimulationRun, ) @@ -38,7 +36,6 @@ "simulation_runs", "reports", "report_runs", - "report_run_simulation_runs", "report_run_comparisons", "household_jobs", "aggregates", @@ -247,17 +244,12 @@ def test_each_comparison_keeps_the_parts_without_a_result_table() -> None: assert ReportRunComparison.__module__.endswith(".reports") -def test_report_run_relationships_name_member_and_comparison_rows() -> None: +def test_report_run_relationships_name_comparison_rows() -> None: relationships = sa.inspect(ReportRun).relationships - assert relationships["simulation_run_members"].mapper.class_ is ( - ReportRunSimulationRun - ) assert relationships["comparisons"].mapper.class_ is ReportRunComparison - assert ( - sa.inspect(SimulationRun).relationships["report_run_memberships"].mapper.class_ - is ReportRunSimulationRun - ) + assert "simulation_run_members" not in relationships + assert "report_run_memberships" not in sa.inspect(SimulationRun).relationships def test_uk_area_impacts_record_disclosure_control() -> None: @@ -296,36 +288,29 @@ def test_detailed_budget_rows_need_neither_an_entity_nor_a_tax_flag() -> None: ) assert ( "CONSTRAINT uq_program_statistics_run_program_entity UNIQUE NULLS NOT " - "DISTINCT (report_run_id, baseline_simulation_run_id, " - "reform_simulation_run_id, program_name, entity)" + "DISTINCT (report_run_id, program_name, entity)" ) in statement -def test_report_runs_name_the_simulation_runs_they_combine() -> None: - members = V2_METADATA.tables["report_run_simulation_runs"] +def test_report_runs_name_their_baseline_and_reform_simulation_runs() -> None: + """A report run compares one pair; a household report may have no reform.""" + report_runs = V2_METADATA.tables["report_runs"] - assert set(members.c.keys()) == { - "id", - "created_at", - "report_run_id", - "simulation_run_id", - "role", - } - assert all(not column.nullable for column in members.columns) - report_run_key = _foreign_key(members, "report_run_id") - simulation_run_key = _foreign_key(members, "simulation_run_id") - assert report_run_key.target_fullname == "report_runs.id" - assert report_run_key.ondelete == "CASCADE" - assert simulation_run_key.target_fullname == "simulation_runs.id" - assert simulation_run_key.ondelete == "RESTRICT" - assert members.c.role.type.name == "v2_simulation_role" - assert list(members.c.role.type.enums) == ["baseline", "reform"] - assert [role.value for role in SimulationRole] == ["baseline", "reform"] - assert _unique_column_sets(members) == { - ("report_run_id", "role", "simulation_run_id") + assert not report_runs.c.baseline_simulation_run_id.nullable + assert report_runs.c.reform_simulation_run_id.nullable + for column_name in ("baseline_simulation_run_id", "reform_simulation_run_id"): + run_key = _foreign_key(report_runs, column_name) + assert run_key.target_fullname == "simulation_runs.id" + assert run_key.ondelete == "RESTRICT" + assert (column_name,) in _indexes(report_runs).values() + assert "report_run_simulation_runs" not in V2_METADATA.tables + enum_type_names = { + column.type.name + for table in V2_METADATA.tables.values() + for column in table.columns + if isinstance(column.type, sa.Enum) } - assert "ix_report_run_simulation_runs_simulation_run_id" in _indexes(members) - assert ReportRunSimulationRun.__module__.endswith(".reports") + assert "v2_simulation_role" not in enum_type_names def test_household_jobs_calculate_one_household_year_and_policy() -> None: @@ -402,13 +387,12 @@ def test_comparison_results_name_the_two_simulation_runs_they_compare() -> None: run_key = _foreign_key(table, column_name) assert run_key.target_fullname == "simulation_runs.id", table_name assert run_key.ondelete == "RESTRICT" - # A budget-window run compares one pair of simulation runs per year, so - # each row is unique within its report run and pair. + # A report run compares one pair of simulation runs, so each row is + # unique within its report run. (unique_columns,) = _unique_column_sets(table) - assert unique_columns[:3] == ( - "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", + assert unique_columns[0] == "report_run_id", table_name + assert {"baseline_simulation_run_id", "reform_simulation_run_id"}.isdisjoint( + unique_columns ), table_name From 41b22d179fc1aa01687e7f71a39013136afc4e52 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:59:16 +0400 Subject: [PATCH 15/22] Type a v2 report run's labour supply and cliff impact results report_run_comparisons is gone: a report run compares one pair, so the parts of an economy comparison that had no result table belong to it again, typed. report_runs gains the baseline's and reform's cliff gap and share (checked to be set together), the SPM receipt as jsonb and the data release. The new labor_supply_responses holds one row per report run with the earnings, revenue and weekly hours responses, and labor_supply_decile_responses one row per report run and income decile. The Modal worker inserts into both and the Public API reads them. Revision 1356f658347d is regenerated. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 13 +- ...record_simulation_runs_and_key_results_.py | 287 +++++++++++------- policyengine_api/data/v2/models/__init__.py | 6 +- policyengine_api/data/v2/models/reports.py | 101 ++++-- policyengine_api/data/v2/privileges.py | 3 +- .../test_alembic_v2_simulation_runs.py | 110 ++++++- .../test_v2_simulation_role_privileges.py | 19 +- tests/unit/v2/test_alembic_v2.py | 9 +- tests/unit/v2/test_import_side_effects.py | 2 +- tests/unit/v2/test_privilege_manifest.py | 3 +- tests/unit/v2/test_simulation_run_models.py | 127 ++++++-- 11 files changed, 473 insertions(+), 207 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index bff133779..9ffa4d8ea 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -511,10 +511,15 @@ model changes: - A report run compares one pair of simulation runs: `report_runs` names its `baseline_simulation_run_id` (required) and `reform_simulation_run_id` (empty for a household report with only a baseline), each a foreign key to - `simulation_runs` with `RESTRICT`. The new `report_run_comparisons` keeps - the parts of an economy comparison that have no result table - (`labor_supply_result`, `cliff_impact_result`, `spm_provenance`, - `data_version`). + `simulation_runs` with `RESTRICT`. The parts of an economy comparison that + had no result table are typed: `report_runs` gains the baseline's and + reform's cliff gap and share (four nullable numbers, checked to be set + together), the SPM receipt (`spm_provenance`, `jsonb`) and the data release + (`data_version`). The new `labor_supply_responses` holds one row per report + run (the earnings changes from the income and substitution effects and their + total, the revenue change, the two relative responses, and baseline, reform + and changed weekly hours split into the two effects), and the new + `labor_supply_decile_responses` one row per report run and income decile. - `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); the other result tables belong to a report run and name the two simulation runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`); diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 83aa7a97d..6d599d4f2 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -2,7 +2,7 @@ Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-05 23:56:01.811595 +Create Date: 2026-11-05 23:58:34.990506 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate One post-generation correction, commented where it is made: @@ -153,7 +153,7 @@ def upgrade() -> None: unique=False, ) op.create_table( - "report_run_comparisons", + "labor_supply_decile_responses", sa.Column("id", sa.Uuid(), nullable=False), sa.Column( "created_at", @@ -162,44 +162,54 @@ def upgrade() -> None: nullable=False, ), sa.Column("report_run_id", sa.Uuid(), nullable=False), - sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), - sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), - sa.Column("labor_supply_result", sa.JSON(), nullable=True), - sa.Column("cliff_impact_result", sa.JSON(), nullable=True), - sa.Column("spm_provenance", sa.JSON(), nullable=True), - sa.Column( - "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True - ), - sa.ForeignKeyConstraint( - ["baseline_simulation_run_id"], - ["simulation_runs.id"], - name=op.f( - "fk_report_run_comparisons_baseline_simulation_run_id_simulation_runs" - ), - ondelete="RESTRICT", - ), - sa.ForeignKeyConstraint( - ["reform_simulation_run_id"], - ["simulation_runs.id"], - name=op.f( - "fk_report_run_comparisons_reform_simulation_run_id_simulation_runs" - ), - ondelete="RESTRICT", - ), + sa.Column("decile", sa.Integer(), nullable=False), + sa.Column("average_income_effect", sa.Float(), nullable=True), + sa.Column("average_substitution_effect", sa.Float(), nullable=True), + sa.Column("relative_income_effect", sa.Float(), nullable=True), + sa.Column("relative_substitution_effect", sa.Float(), nullable=True), sa.ForeignKeyConstraint( ["report_run_id"], ["report_runs.id"], - name=op.f("fk_report_run_comparisons_report_run_id_report_runs"), + name=op.f("fk_labor_supply_decile_responses_report_run_id_report_runs"), ondelete="CASCADE", ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_report_run_comparisons")), + sa.PrimaryKeyConstraint("id", name=op.f("pk_labor_supply_decile_responses")), sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - name="uq_report_run_comparisons_run_pair", + "decile", + name="uq_labor_supply_decile_responses_run_decile", ), ) + op.create_table( + "labor_supply_responses", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_run_id", sa.Uuid(), nullable=False), + sa.Column("income_effect_earnings_change", sa.Float(), nullable=False), + sa.Column("substitution_effect_earnings_change", sa.Float(), nullable=False), + sa.Column("total_earnings_change", sa.Float(), nullable=False), + sa.Column("revenue_change", sa.Float(), nullable=False), + sa.Column("relative_income_effect", sa.Float(), nullable=False), + sa.Column("relative_substitution_effect", sa.Float(), nullable=False), + sa.Column("baseline_weekly_hours", sa.Float(), nullable=False), + sa.Column("reform_weekly_hours", sa.Float(), nullable=False), + sa.Column("weekly_hours_change", sa.Float(), nullable=False), + sa.Column("weekly_hours_income_effect", sa.Float(), nullable=False), + sa.Column("weekly_hours_substitution_effect", sa.Float(), nullable=False), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_labor_supply_responses_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_labor_supply_responses")), + sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), + ) op.add_column( "budget_summary", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -209,28 +219,28 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -268,21 +278,21 @@ def upgrade() -> None: ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -386,28 +396,28 @@ def upgrade() -> None: postgresql_nulls_not_distinct=True, ) op.drop_constraint( - op.f("fk_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), "decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_decile_impacts_reform_simulation_id_simulations"), "decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -465,20 +475,20 @@ def upgrade() -> None: unique=False, ) op.create_foreign_key( - op.f( - "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" - ), + op.f("fk_household_jobs_household_id_households"), "household_jobs", - "tax_benefit_model_versions", - ["tax_benefit_model_version_id"], + "households", + ["household_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_household_jobs_household_id_households"), + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", - "households", - ["household_id"], + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], ["id"], ondelete="RESTRICT", ) @@ -530,28 +540,28 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -590,12 +600,12 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) @@ -727,6 +737,34 @@ def upgrade() -> None: "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True ), ) + op.add_column( + "report_runs", sa.Column("baseline_cliff_gap", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("baseline_cliff_share", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("reform_cliff_gap", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("reform_cliff_share", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", + sa.Column( + "spm_provenance", + sa.JSON().with_variant( + postgresql.JSONB(astext_type=sa.Text()), "postgresql" + ), + nullable=True, + ), + ) + op.add_column( + "report_runs", + sa.Column( + "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True + ), + ) op.create_index( op.f("ix_report_runs_baseline_simulation_run_id"), "report_runs", @@ -746,21 +784,26 @@ def upgrade() -> None: unique=False, ) op.create_foreign_key( - op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), "report_runs", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), "report_runs", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) + op.create_check_constraint( + op.f("ck_report_runs_cliff_impact_complete"), + "report_runs", + "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_share IS NULL)", + ) op.create_index(op.f("ix_reports_user_id"), "reports", ["user_id"], unique=False) op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") op.drop_constraint( @@ -768,12 +811,12 @@ def upgrade() -> None: "simulations", type_="foreignkey", ) - op.drop_column("simulations", "status") op.drop_column("simulations", "output_dataset_id") - op.drop_column("simulations", "started_at") op.drop_column("simulations", "completed_at") - op.drop_column("simulations", "error_message") + op.drop_column("simulations", "status") op.drop_column("simulations", "household_result") + op.drop_column("simulations", "error_message") + op.drop_column("simulations", "started_at") # ### end Alembic commands ### @@ -782,8 +825,8 @@ def downgrade() -> None: op.add_column( "simulations", sa.Column( - "household_result", - postgresql.JSON(astext_type=sa.Text()), + "started_at", + postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, ), @@ -795,8 +838,8 @@ def downgrade() -> None: op.add_column( "simulations", sa.Column( - "completed_at", - postgresql.TIMESTAMP(timezone=True), + "household_result", + postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, ), @@ -804,27 +847,27 @@ def downgrade() -> None: op.add_column( "simulations", sa.Column( - "started_at", - postgresql.TIMESTAMP(timezone=True), + "status", + postgresql.ENUM( + "pending", "running", "succeeded", "failed", name="v2_simulation_status" + ), autoincrement=False, - nullable=True, + nullable=False, ), ) - op.add_column( - "simulations", - sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), - ) op.add_column( "simulations", sa.Column( - "status", - postgresql.ENUM( - "pending", "running", "succeeded", "failed", name="v2_simulation_status" - ), + "completed_at", + postgresql.TIMESTAMP(timezone=True), autoincrement=False, - nullable=False, + nullable=True, ), ) + op.add_column( + "simulations", + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + ) op.create_foreign_key( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", @@ -841,12 +884,15 @@ def downgrade() -> None: ) op.drop_index(op.f("ix_reports_user_id"), table_name="reports") op.drop_constraint( - op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + op.f("ck_report_runs_cliff_impact_complete"), "report_runs", type_="check" + ) + op.drop_constraint( + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) op.drop_constraint( - op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) @@ -857,6 +903,12 @@ def downgrade() -> None: op.drop_index( op.f("ix_report_runs_baseline_simulation_run_id"), table_name="report_runs" ) + op.drop_column("report_runs", "data_version") + op.drop_column("report_runs", "spm_provenance") + op.drop_column("report_runs", "reform_cliff_share") + op.drop_column("report_runs", "reform_cliff_gap") + op.drop_column("report_runs", "baseline_cliff_share") + op.drop_column("report_runs", "baseline_cliff_gap") op.drop_column("report_runs", "modal_call_id") op.drop_column("report_runs", "reform_simulation_run_id") op.drop_column("report_runs", "baseline_simulation_run_id") @@ -984,18 +1036,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1019,28 +1071,28 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1090,14 +1142,14 @@ def downgrade() -> None: op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" ) op.drop_constraint( - op.f("fk_household_jobs_household_id_households"), + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", type_="foreignkey", ) op.drop_constraint( - op.f( - "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" - ), + op.f("fk_household_jobs_household_id_households"), "household_jobs", type_="foreignkey", ) @@ -1126,28 +1178,28 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), "decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_decile_impacts_reform_simulation_id_simulations"), "decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_decile_impacts_baseline_simulation_id_simulations"), "decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1227,14 +1279,14 @@ def downgrade() -> None: ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", @@ -1273,34 +1325,35 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.drop_column("budget_summary", "reform_simulation_run_id") op.drop_column("budget_summary", "baseline_simulation_run_id") - op.drop_table("report_run_comparisons") + op.drop_table("labor_supply_responses") + op.drop_table("labor_supply_decile_responses") op.drop_index( op.f("ix_simulation_runs_tax_benefit_model_version_id"), table_name="simulation_runs", diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index 87f73fb6e..bb8469d0c 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -76,13 +76,14 @@ DecileType, Inequality, IntraDecileImpact, + LaborSupplyDecileResponse, + LaborSupplyResponse, LocalAuthorityImpact, OutputStatus, Poverty, ProgramStatistics, Report, ReportRun, - ReportRunComparison, ReportRunStatus, ReportRunTrigger, ) @@ -122,6 +123,8 @@ "HouseholdJobStatus", "Inequality", "IntraDecileImpact", + "LaborSupplyDecileResponse", + "LaborSupplyResponse", "LegacyHouseholdMapping", "LegacyPolicyMapping", "LegacyUserMapping", @@ -140,7 +143,6 @@ "RegionType", "Report", "ReportRun", - "ReportRunComparison", "ReportRunStatus", "ReportRunTrigger", "Simulation", diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index be2a9e2ac..2d7a8ff72 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -12,6 +12,7 @@ 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 ( @@ -164,8 +165,10 @@ class ReportRun(TimestampedModel, table=True): It compares ``baseline_simulation_run_id`` with ``reform_simulation_run_id``; a household report with only a baseline has - no reform run. Its ``comparisons`` hold the parts of an economy comparison - that have no result table. + no reform run. An economy comparison's cliff impact is the baseline's and + reform's cliff gap and share, stored together or not at all; + ``spm_provenance`` is the SPM receipt and ``data_version`` the data release + it used. Its labour supply response has its own two tables. """ __tablename__ = "report_runs" @@ -179,6 +182,12 @@ class ReportRun(TimestampedModel, table=True): "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", name="ck_report_runs_terminal_completion", ), + sa.CheckConstraint( + "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) " + "AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) " + "AND (baseline_cliff_gap IS NULL) = (reform_cliff_share IS NULL)", + name="ck_report_runs_cliff_impact_complete", + ), sa.Index( "ix_report_runs_current_output", "report_id", @@ -228,9 +237,22 @@ class ReportRun(TimestampedModel, table=True): ) error_message: str | None = None markdown: str | None = Field(default=None, sa_type=sa.Text) + baseline_cliff_gap: float | None = None + baseline_cliff_share: float | None = None + reform_cliff_gap: float | None = None + reform_cliff_share: float | None = None + spm_provenance: dict[str, Any] | None = Field( + default=None, + sa_type=sa.JSON().with_variant(postgresql.JSONB(), "postgresql"), + ) + data_version: str | None = Field(default=None, max_length=128) report: Report = Relationship(back_populates="runs") - comparisons: list["ReportRunComparison"] = Relationship( + labor_supply_responses: list["LaborSupplyResponse"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) + labor_supply_decile_responses: list["LaborSupplyDecileResponse"] = Relationship( back_populates="report_run", cascade_delete=True, ) @@ -271,22 +293,56 @@ class ReportRun(TimestampedModel, table=True): ) -class ReportRunComparison(IdentifiedModel, table=True): - """The parts of one compared pair of a report run that have no result table. +class LaborSupplyResponse(IdentifiedModel, table=True): + """The labour supply response to the reform in one report run. + + One row per report run. The earnings changes come from the income and + substitution effects (.py's ``income_lsr`` and ``substitution_lsr``) and + total ``total_earnings_change``; ``relative_income_effect`` and + ``relative_substitution_effect`` are relative to baseline earnings. The + weekly hours are the total hours worked under the baseline and the reform, + and their change split into the two effects. + """ + + __tablename__ = "labor_supply_responses" + __table_args__ = ( + sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + income_effect_earnings_change: float + substitution_effect_earnings_change: float + total_earnings_change: float + revenue_change: float + relative_income_effect: float + relative_substitution_effect: float + baseline_weekly_hours: float + reform_weekly_hours: float + weekly_hours_change: float + weekly_hours_income_effect: float + weekly_hours_substitution_effect: float + + report_run: ReportRun = Relationship(back_populates="labor_supply_responses") + - A report run compares one baseline and one reform simulation run, or one - pair per year for a budget window. Each pair's labour supply response, - cliff impact, SPM receipts and data release are one row here, as the - simulation API's economy output gives them. +class LaborSupplyDecileResponse(IdentifiedModel, table=True): + """The labour supply response in one income decile of one report run. + + One row per report run and decile, as .py numbers them. The average effects + are per household; the relative effects are relative to the decile's + baseline earnings, which .py reports only for deciles above zero, so they + may be null. """ - __tablename__ = "report_run_comparisons" + __tablename__ = "labor_supply_decile_responses" __table_args__ = ( sa.UniqueConstraint( "report_run_id", - "baseline_simulation_run_id", - "reform_simulation_run_id", - name="uq_report_run_comparisons_run_pair", + "decile", + name="uq_labor_supply_decile_responses_run_decile", ), ) @@ -294,20 +350,13 @@ class ReportRunComparison(IdentifiedModel, table=True): foreign_key="report_runs.id", ondelete="CASCADE", ) - baseline_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - labor_supply_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - cliff_impact_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - spm_provenance: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) - data_version: str | None = Field(default=None, max_length=128) + decile: int + average_income_effect: float | None = None + average_substitution_effect: float | None = None + relative_income_effect: float | None = None + relative_substitution_effect: float | None = None - report_run: ReportRun = Relationship(back_populates="comparisons") + report_run: ReportRun = Relationship(back_populates="labor_supply_decile_responses") class AggregateOutput(IdentifiedModel, table=True): diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 98cefb580..57d8d443d 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -154,7 +154,8 @@ def _access( "simulation_runs": _RUN_STATUS, "simulations": replace(_SIMULATION_DEFINITION, catalog_writer=READ), # Results of simulation runs and report runs. - "report_run_comparisons": _RESULT, + "labor_supply_decile_responses": _RESULT, + "labor_supply_responses": _RESULT, "budget_summary": _RESULT, "congressional_district_impacts": _RESULT, "constituency_impacts": _RESULT, diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 6040077c5..2ed3ad317 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -42,7 +42,9 @@ from sqlalchemy import Connection, Engine SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" -NEW_TABLES = frozenset({"simulation_runs", "report_run_comparisons"}) +NEW_TABLES = frozenset( + {"simulation_runs", "labor_supply_responses", "labor_supply_decile_responses"} +) CHANGED_TABLES = frozenset( { "budget_summary", @@ -97,6 +99,19 @@ "RESTRICT", ), } +LABOR_SUPPLY_RESPONSE_COLUMNS = ( + "income_effect_earnings_change", + "substitution_effect_earnings_change", + "total_earnings_change", + "revenue_change", + "relative_income_effect", + "relative_substitution_effect", + "baseline_weekly_hours", + "reform_weekly_hours", + "weekly_hours_change", + "weekly_hours_income_effect", + "weekly_hours_substitution_effect", +) SHA256_HEX = "a" * 64 COMPLETED_AT = "2026-11-04T00:00:00Z" @@ -513,16 +528,6 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( **budget, **{**comparison_runs, "reform_simulation_run_id": baseline_run_id}, ) - insert_row( - connection, - "report_run_comparisons", - id=uuid4(), - labor_supply_result="{}", - data_version="1.0.0", - **comparison_runs, - ) - with refused(connection): - insert_row(connection, "report_run_comparisons", id=uuid4(), **comparison_runs) insert_row( connection, "program_statistics", @@ -564,6 +569,89 @@ def test_results_are_unique_per_run_and_name_runs_that_exist( ) +def test_a_report_runs_labor_supply_and_cliff_impact_are_typed( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + simulation_id = _insert_simulation(connection, public_api_rows) + report_run_id = _insert_report_run( + connection, + public_api_rows, + _insert_simulation_run(connection, simulation_id), + _insert_simulation_run(connection, simulation_id), + ) + response = dict.fromkeys(LABOR_SUPPLY_RESPONSE_COLUMNS, 1.0) + + insert_row( + connection, + "labor_supply_responses", + id=uuid4(), + report_run_id=report_run_id, + **response, + ) + with refused(connection): + insert_row( + connection, + "labor_supply_responses", + id=uuid4(), + report_run_id=report_run_id, + **response, + ) + # .py reports a relative response only for deciles above zero. + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=0, + average_income_effect=-1.0, + average_substitution_effect=2.0, + ) + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=1, + average_income_effect=-1.0, + average_substitution_effect=2.0, + relative_income_effect=-0.01, + relative_substitution_effect=0.02, + ) + with refused(connection): + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=1, + ) + with refused(connection, match="ck_report_runs_cliff_impact_complete"): + connection.execute( + text("UPDATE report_runs SET baseline_cliff_gap = 1.0 WHERE id = :id"), + {"id": report_run_id}, + ) + connection.execute( + text( + "UPDATE report_runs SET baseline_cliff_gap = 1.0, " + "baseline_cliff_share = 0.1, reform_cliff_gap = 2.0, " + "reform_cliff_share = 0.2, data_version = '1.0.0', " + "spm_provenance = CAST(:receipt AS jsonb) WHERE id = :id" + ), + {"id": report_run_id, "receipt": '{"source": "census"}'}, + ) + receipt_type = connection.execute( + text("SELECT jsonb_typeof(spm_provenance) FROM report_runs WHERE id = :id"), + {"id": report_run_id}, + ).scalar_one() + assert receipt_type == "object" + + connection.execute( + text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} + ) + for table in ("labor_supply_responses", "labor_supply_decile_responses"): + assert _count(connection, table, "report_run_id", report_run_id) == 0 + + def test_household_jobs_take_a_valid_year_and_record_when_they_end( connection: Connection, public_api_rows: PublicApiRows ) -> None: diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index 503bf73e0..a995979e5 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -158,12 +158,10 @@ def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( _refused(connection, "DELETE FROM simulations WHERE id = :id", id=ids["simulation"]) _refused( connection, - "INSERT INTO report_run_comparisons " - "(id, report_run_id, baseline_simulation_run_id, reform_simulation_run_id) " - "VALUES (:id, :report_run_id, :run_id, :run_id)", + "INSERT INTO labor_supply_decile_responses (id, report_run_id, decile) " + "VALUES (:id, :report_run_id, 1)", id=uuid4(), report_run_id=ids["report_run"], - run_id=ids["run"], ) _refused(connection, "UPDATE policies SET content_hash = content_hash") @@ -192,17 +190,20 @@ def test_modal_workers_write_their_run_status_and_results_only( rate=0.1, ) connection.execute( - text("UPDATE report_runs SET status = 'running' WHERE id = :id"), + text( + "UPDATE report_runs SET status = 'running', baseline_cliff_gap = 1.0, " + "baseline_cliff_share = 0.1, reform_cliff_gap = 2.0, " + "reform_cliff_share = 0.2 WHERE id = :id" + ), {"id": ids["report_run"]}, ) insert_row( connection, - "report_run_comparisons", + "labor_supply_decile_responses", id=uuid4(), report_run_id=ids["report_run"], - baseline_simulation_run_id=ids["run"], - reform_simulation_run_id=ids["run"], - cliff_impact_result="{}", + decile=1, + average_income_effect=-1.0, ) finished_job = connection.execute( text( diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 54f215b76..e004b6cf1 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -47,7 +47,8 @@ "legacy_user_policy_mappings", "parameter_search_files", "parameter_tree_children", - "report_run_comparisons", + "labor_supply_decile_responses", + "labor_supply_responses", "simulation_runs", "stage12_evaluation_reports", "stage12_evaluation_simulations", @@ -1047,7 +1048,11 @@ def test_simulation_runs_revision_is_generated_with_one_documented_correction() assert 'name="v2_simulation_status",\n create_type=False,' in ( upgrade ) - for table_name in ("simulation_runs", "report_run_comparisons"): + for table_name in ( + "simulation_runs", + "labor_supply_responses", + "labor_supply_decile_responses", + ): assert f'op.create_table(\n "{table_name}"' in upgrade assert f'op.drop_table("{table_name}")' in downgrade assert upgrade.count("op.drop_table(") == 0 diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index 13325fd74..df7828892 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 44 +assert len(V2_METADATA.tables) == 45 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 94dc35993..cb9363e95 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -66,7 +66,8 @@ ) RESULT_TABLES = frozenset( { - "report_run_comparisons", + "labor_supply_responses", + "labor_supply_decile_responses", "budget_summary", "congressional_district_impacts", "constituency_impacts", diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 4b73b57a7..84255db5b 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -12,12 +12,39 @@ from policyengine_api.data.v2.models import ( V2_METADATA, + LaborSupplyDecileResponse, + LaborSupplyResponse, ReportRun, - ReportRunComparison, SimulationRun, ) RUN_STATUSES = ["pending", "running", "succeeded", "failed"] +# The baseline's and reform's cliff gap and share, stored together or not at all. +CLIFF_IMPACT_COLUMNS = ( + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", +) +LABOR_SUPPLY_RESPONSE_COLUMNS = ( + "income_effect_earnings_change", + "substitution_effect_earnings_change", + "total_earnings_change", + "revenue_change", + "relative_income_effect", + "relative_substitution_effect", + "baseline_weekly_hours", + "reform_weekly_hours", + "weekly_hours_change", + "weekly_hours_income_effect", + "weekly_hours_substitution_effect", +) +LABOR_SUPPLY_DECILE_COLUMNS = ( + "average_income_effect", + "average_substitution_effect", + "relative_income_effect", + "relative_substitution_effect", +) PER_SIMULATION_RESULT_TABLES = ("inequality", "poverty") COMPARISON_RESULT_TABLES = ( "budget_summary", @@ -36,7 +63,8 @@ "simulation_runs", "reports", "report_runs", - "report_run_comparisons", + "labor_supply_responses", + "labor_supply_decile_responses", "household_jobs", "aggregates", "change_aggregates", @@ -212,44 +240,77 @@ def test_report_runs_record_their_modal_call() -> None: assert report_runs.c.modal_call_id.nullable assert report_runs.c.modal_call_id.type.length == 255 - # A budget-window run compares one pair per year, so the parts of a - # comparison without a result table belong to its comparison row. - assert { - "labor_supply_result", - "cliff_impact_result", - "spm_provenance", - "data_version", - }.isdisjoint(report_runs.c.keys()) - - -def test_each_comparison_keeps_the_parts_without_a_result_table() -> None: - comparisons = V2_METADATA.tables["report_run_comparisons"] - - for column_name in ("labor_supply_result", "cliff_impact_result", "spm_provenance"): - assert comparisons.c[column_name].nullable - assert isinstance(comparisons.c[column_name].type, sa.JSON) - assert comparisons.c.data_version.nullable - assert comparisons.c.data_version.type.length == 128 - report_run_key = _foreign_key(comparisons, "report_run_id") + assert "report_run_comparisons" not in V2_METADATA.tables + + +def test_report_runs_record_cliff_impact_spm_receipt_and_data_version() -> None: + """The comparison parts with no result table are typed, except the receipt.""" + report_runs = V2_METADATA.tables["report_runs"] + + for column_name in CLIFF_IMPACT_COLUMNS: + assert report_runs.c[column_name].nullable, column_name + assert isinstance(report_runs.c[column_name].type, sa.Float), column_name + assert "ck_report_runs_cliff_impact_complete" in _check_names(report_runs) + assert report_runs.c.spm_provenance.nullable + spm_type = report_runs.c.spm_provenance.type + assert isinstance(spm_type.dialect_impl(postgresql.dialect()), postgresql.JSONB) + assert report_runs.c.data_version.nullable + assert report_runs.c.data_version.type.length == 128 + assert {"labor_supply_result", "cliff_impact_result"}.isdisjoint( + report_runs.c.keys() + ) + + +def test_labor_supply_responses_have_one_typed_row_per_report_run() -> None: + responses = V2_METADATA.tables["labor_supply_responses"] + + assert set(responses.c.keys()) == { + "id", + "created_at", + "report_run_id", + *LABOR_SUPPLY_RESPONSE_COLUMNS, + } + for column_name in LABOR_SUPPLY_RESPONSE_COLUMNS: + assert not responses.c[column_name].nullable, column_name + assert isinstance(responses.c[column_name].type, sa.Float), column_name + report_run_key = _foreign_key(responses, "report_run_id") assert report_run_key.target_fullname == "report_runs.id" assert report_run_key.ondelete == "CASCADE" - for column_name in ("baseline_simulation_run_id", "reform_simulation_run_id"): - assert not comparisons.c[column_name].nullable - run_key = _foreign_key(comparisons, column_name) - assert run_key.target_fullname == "simulation_runs.id" - assert run_key.ondelete == "RESTRICT" - assert _unique_column_sets(comparisons) == { - ("report_run_id", "baseline_simulation_run_id", "reform_simulation_run_id") + assert _unique_column_sets(responses) == {("report_run_id",)} + + +def test_labor_supply_decile_responses_have_one_row_per_decile() -> None: + deciles = V2_METADATA.tables["labor_supply_decile_responses"] + + assert set(deciles.c.keys()) == { + "id", + "created_at", + "report_run_id", + "decile", + *LABOR_SUPPLY_DECILE_COLUMNS, } - assert ReportRunComparison.__module__.endswith(".reports") + assert not deciles.c.decile.nullable + # .py reports a relative response only for deciles above zero. + for column_name in LABOR_SUPPLY_DECILE_COLUMNS: + assert deciles.c[column_name].nullable, column_name + assert isinstance(deciles.c[column_name].type, sa.Float), column_name + report_run_key = _foreign_key(deciles, "report_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + assert _unique_column_sets(deciles) == {("report_run_id", "decile")} -def test_report_run_relationships_name_comparison_rows() -> None: +def test_report_run_relationships_name_labor_supply_rows() -> None: relationships = sa.inspect(ReportRun).relationships - assert relationships["comparisons"].mapper.class_ is ReportRunComparison - assert "simulation_run_members" not in relationships - assert "report_run_memberships" not in sa.inspect(SimulationRun).relationships + assert relationships["labor_supply_responses"].mapper.class_ is ( + LaborSupplyResponse + ) + assert relationships["labor_supply_decile_responses"].mapper.class_ is ( + LaborSupplyDecileResponse + ) + assert "comparisons" not in relationships + assert LaborSupplyResponse.__module__.endswith(".reports") def test_uk_area_impacts_record_disclosure_control() -> None: From c4742f225ef6eab1a902e13b843f6da102e89ff8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:03:34 +0400 Subject: [PATCH 16/22] Store v2 budget windows: a run per year and one row per calculation Following the simulation API's budget-window design, without its scheduling columns: simulation_runs gains a required year, which joins the index for a simulation's newest successful run, so a budget window compares a report's same two simulations in every year. reports.budget_window_size (1 to 75 years) marks a budget-window report. budget_window_runs holds one row per calculation (report, catalog version, trigger, whether to recalculate, status, error and times), reusing the report run's status and trigger types, and report_runs.budget_window_run_id and year name the window run and year a report run calculates, at most one per year, kept to the same report by a composite foreign key. The simulation service creates window runs; the Modal worker and the Public API read them. create_report_run now takes its package versions and simulation runs as one ReportRunInputs, keeping it within ruff's argument limit. Three model tests that still asserted the dropped foreign keys and index set are updated. Revision 1356f658347d is regenerated with three corrections of one kind: new columns reusing existing enum types. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 23 +- ...record_simulation_runs_and_key_results_.py | 366 ++++++++++++------ policyengine_api/data/v2/models/__init__.py | 2 + policyengine_api/data/v2/models/reports.py | 100 ++++- .../data/v2/models/simulations.py | 11 +- policyengine_api/data/v2/privileges.py | 1 + policyengine_api/data/v2/report_runs.py | 16 +- .../integration/test_alembic_v2_lifecycle.py | 1 + .../test_alembic_v2_simulation_runs.py | 111 +++++- .../test_v2_catalog_publication.py | 1 + .../test_v2_simulation_role_privileges.py | 33 +- tests/unit/v2/test_alembic_v2.py | 20 +- tests/unit/v2/test_import_side_effects.py | 2 +- tests/unit/v2/test_privilege_manifest.py | 1 + tests/unit/v2/test_report_run_creation.py | 51 ++- tests/unit/v2/test_report_runs.py | 31 +- tests/unit/v2/test_simulation_run_models.py | 102 ++++- 17 files changed, 696 insertions(+), 176 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 9ffa4d8ea..d56cde7f9 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -537,6 +537,18 @@ model changes: `household_wealth_decile`), and `poverty.filter_group` names the group a filter selects (`child` for `age`); both join their table's unique key. - `report_runs` gains its Modal call ID. +- Budget windows: a simulation run calculates one `year`, which joins the + index for a simulation's newest successful run, so a budget window compares + a report's same two simulations in every year. `reports.budget_window_size` + (1 to 75 years, an economy report with a year, ending by 2200) marks a + budget-window report. The new `budget_window_runs` holds one row per + calculation of one: its report, catalog version, trigger, whether its years + calculate their simulations again (`recalculate`), status, error and times. + `report_runs.budget_window_run_id` and `year` name the window run and year a + report run calculates, at most one report run per year of a window run; a + composite foreign key keeps it to a window run of the same report. That key + is `NO ACTION`, which PostgreSQL checks at the end of the statement, so + deleting a report deletes its window runs and their report runs together. - `household_jobs` calculates one household (`household_id`) for one `year` under one policy, and records its catalog and package versions and its Modal call ID. @@ -554,11 +566,12 @@ Missing values in the unique keys of `poverty`, `decile_impacts` and `program_statistics` count as equal (`NULLS NOT DISTINCT`, which PostgreSQL supports from version 15). -It has one post-generation correction, commented in the revision: -`simulation_runs.status` reuses the `v2_simulation_status` type the baseline -created for `simulations.status`. The generated `create_table` would create -the type again and fail, so the column uses -`postgresql.ENUM(..., create_type=False)`. +It has three post-generation corrections of one kind, each commented in the +revision: `simulation_runs.status` reuses the `v2_simulation_status` type the +baseline created for `simulations.status`, and `budget_window_runs.status` and +`trigger` reuse the baseline's `v2_report_run_status` and +`v2_report_run_trigger`. The generated `create_table` would create each type +again and fail, so each column uses `postgresql.ENUM(..., create_type=False)`. The upgrade adds required columns without server defaults to `household_jobs`, `report_runs` and the result tables, so it needs them empty, as they are while no managed diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 6d599d4f2..48021b291 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -2,11 +2,13 @@ Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-05 23:58:34.990506 +Create Date: 2026-11-06 00:03:10.204134 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate -One post-generation correction, commented where it is made: -simulation_runs.status reuses the existing v2_simulation_status type. +Three post-generation corrections of one kind, each commented where it is +made: simulation_runs.status reuses the existing v2_simulation_status type, +and budget_window_runs.status and trigger reuse the existing +v2_report_run_status and v2_report_run_trigger types. The upgrade adds required columns without server defaults to household_jobs, report_runs and the result tables, so it needs those tables empty, as they are @@ -46,6 +48,7 @@ def upgrade() -> None: nullable=False, ), sa.Column("simulation_id", sa.Uuid(), nullable=False), + sa.Column("year", sa.Integer(), nullable=False), sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), sa.Column( "country_package_version", @@ -99,6 +102,9 @@ def upgrade() -> None: "output_sha256 IS NULL OR length(output_sha256) = 64", name=op.f("ck_simulation_runs_output_sha256_length"), ), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_simulation_runs_year") + ), sa.ForeignKeyConstraint( ["dataset_id"], ["datasets.id"], @@ -132,6 +138,7 @@ def upgrade() -> None: "simulation_runs", [ "simulation_id", + "year", "status", "country_package_version", "policyengine_version", @@ -152,6 +159,89 @@ def upgrade() -> None: ["tax_benefit_model_version_id"], unique=False, ) + op.create_table( + "budget_window_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_id", sa.Uuid(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column( + "trigger", + # Dialect correction: a budget-window run starts for the same reasons + # as a report run, so the column reuses the baseline's + # v2_report_run_trigger type instead of creating it again. + postgresql.ENUM( + "initial", + "manual", + "system", + name="v2_report_run_trigger", + create_type=False, + ), + nullable=False, + ), + sa.Column("recalculate", sa.Boolean(), nullable=False), + sa.Column( + "status", + # Dialect correction: a budget-window run has the same states as a + # report run, so the column reuses the baseline's v2_report_run_status + # type instead of creating it again. + postgresql.ENUM( + "pending", + "running", + "succeeded", + "failed", + name="v2_report_run_status", + create_type=False, + ), + nullable=False, + ), + sa.Column("error_message", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("started_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name=op.f("ck_budget_window_runs_terminal_completion"), + ), + sa.ForeignKeyConstraint( + ["report_id"], + ["reports.id"], + name=op.f("fk_budget_window_runs_report_id_reports"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + ["tax_benefit_model_versions.id"], + name=op.f( + "fk_budget_window_runs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_budget_window_runs")), + sa.UniqueConstraint("id", "report_id", name="uq_budget_window_runs_id_report"), + ) + op.create_index( + "ix_budget_window_runs_report_id_created_at", + "budget_window_runs", + ["report_id", "created_at"], + unique=False, + ) + op.create_index( + op.f("ix_budget_window_runs_tax_benefit_model_version_id"), + "budget_window_runs", + ["tax_benefit_model_version_id"], + unique=False, + ) op.create_table( "labor_supply_decile_responses", sa.Column("id", sa.Uuid(), nullable=False), @@ -219,28 +309,28 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -267,12 +357,12 @@ def upgrade() -> None: sa.Column("no_change_percentage", sa.Float(), nullable=False), ) op.drop_constraint( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", type_="foreignkey", ) @@ -331,28 +421,28 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), "constituency_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), "constituency_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -474,14 +564,6 @@ def upgrade() -> None: ["tax_benefit_model_version_id"], unique=False, ) - op.create_foreign_key( - op.f("fk_household_jobs_household_id_households"), - "household_jobs", - "households", - ["household_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_foreign_key( op.f( "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" @@ -492,6 +574,14 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_household_jobs_household_id_households"), + "household_jobs", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) op.create_check_constraint( op.f("ck_household_jobs_terminal_completion"), "household_jobs", @@ -512,12 +602,12 @@ def upgrade() -> None: ["simulation_run_id", "income_variable", "entity"], ) op.drop_constraint( - op.f("fk_inequality_report_run_id_report_runs"), + op.f("fk_inequality_simulation_id_simulations"), "inequality", type_="foreignkey", ) op.drop_constraint( - op.f("fk_inequality_simulation_id_simulations"), + op.f("fk_inequality_report_run_id_report_runs"), "inequality", type_="foreignkey", ) @@ -540,12 +630,12 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) @@ -600,28 +690,28 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -653,10 +743,10 @@ def upgrade() -> None: postgresql_nulls_not_distinct=True, ) op.drop_constraint( - op.f("fk_poverty_simulation_id_simulations"), "poverty", type_="foreignkey" + op.f("fk_poverty_report_run_id_report_runs"), "poverty", type_="foreignkey" ) op.drop_constraint( - op.f("fk_poverty_report_run_id_report_runs"), "poverty", type_="foreignkey" + op.f("fk_poverty_simulation_id_simulations"), "poverty", type_="foreignkey" ) op.create_foreign_key( op.f("fk_poverty_simulation_run_id_simulation_runs"), @@ -697,28 +787,28 @@ def upgrade() -> None: postgresql_nulls_not_distinct=True, ) op.drop_constraint( - op.f("fk_program_statistics_reform_simulation_id_simulations"), + op.f("fk_program_statistics_baseline_simulation_id_simulations"), "program_statistics", type_="foreignkey", ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -731,6 +821,10 @@ def upgrade() -> None: op.add_column( "report_runs", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=True) ) + op.add_column( + "report_runs", sa.Column("budget_window_run_id", sa.Uuid(), nullable=True) + ) + op.add_column("report_runs", sa.Column("year", sa.Integer(), nullable=True)) op.add_column( "report_runs", sa.Column( @@ -783,6 +877,11 @@ def upgrade() -> None: ["status", "started_at"], unique=False, ) + op.create_unique_constraint( + "uq_report_runs_budget_window_run_id_year", + "report_runs", + ["budget_window_run_id", "year"], + ) op.create_foreign_key( op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), "report_runs", @@ -791,6 +890,13 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + "fk_report_runs_budget_window_run_id_budget_window_runs", + "report_runs", + "budget_window_runs", + ["budget_window_run_id", "report_id"], + ["id", "report_id"], + ) op.create_foreign_key( op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), "report_runs", @@ -799,42 +905,60 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_check_constraint( + op.f("ck_report_runs_budget_window_year"), + "report_runs", + "budget_window_run_id IS NULL OR year IS NOT NULL", + ) op.create_check_constraint( op.f("ck_report_runs_cliff_impact_complete"), "report_runs", "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_share IS NULL)", ) + op.create_check_constraint( + op.f("ck_report_runs_year"), + "report_runs", + "year IS NULL OR year BETWEEN 1900 AND 2200", + ) + op.add_column( + "reports", sa.Column("budget_window_size", sa.Integer(), nullable=True) + ) op.create_index(op.f("ix_reports_user_id"), "reports", ["user_id"], unique=False) + op.create_check_constraint( + op.f("ck_reports_budget_window_size"), + "reports", + "budget_window_size IS NULL OR (household_id IS NULL AND budget_window_size BETWEEN 1 AND 75 AND year IS NOT NULL AND year + budget_window_size - 1 <= 2200)", + ) op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") op.drop_constraint( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", type_="foreignkey", ) - op.drop_column("simulations", "output_dataset_id") - op.drop_column("simulations", "completed_at") + op.drop_column("simulations", "started_at") + op.drop_column("simulations", "error_message") op.drop_column("simulations", "status") op.drop_column("simulations", "household_result") - op.drop_column("simulations", "error_message") - op.drop_column("simulations", "started_at") + op.drop_column("simulations", "completed_at") + op.drop_column("simulations", "output_dataset_id") # ### end Alembic commands ### def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "simulations", + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + ) op.add_column( "simulations", sa.Column( - "started_at", + "completed_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, ), ) - op.add_column( - "simulations", - sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), - ) op.add_column( "simulations", sa.Column( @@ -855,19 +979,19 @@ def downgrade() -> None: nullable=False, ), ) + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + ) op.add_column( "simulations", sa.Column( - "completed_at", + "started_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=True, ), ) - op.add_column( - "simulations", - sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), - ) op.create_foreign_key( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", @@ -882,20 +1006,34 @@ def downgrade() -> None: ["status", "created_at"], unique=False, ) + op.drop_constraint(op.f("ck_reports_budget_window_size"), "reports", type_="check") op.drop_index(op.f("ix_reports_user_id"), table_name="reports") + op.drop_column("reports", "budget_window_size") + op.drop_constraint(op.f("ck_report_runs_year"), "report_runs", type_="check") op.drop_constraint( op.f("ck_report_runs_cliff_impact_complete"), "report_runs", type_="check" ) + op.drop_constraint( + op.f("ck_report_runs_budget_window_year"), "report_runs", type_="check" + ) op.drop_constraint( op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) + op.drop_constraint( + "fk_report_runs_budget_window_run_id_budget_window_runs", + "report_runs", + type_="foreignkey", + ) op.drop_constraint( op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) + op.drop_constraint( + "uq_report_runs_budget_window_run_id_year", "report_runs", type_="unique" + ) op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") op.drop_index( op.f("ix_report_runs_reform_simulation_run_id"), table_name="report_runs" @@ -910,6 +1048,8 @@ def downgrade() -> None: op.drop_column("report_runs", "baseline_cliff_share") op.drop_column("report_runs", "baseline_cliff_gap") op.drop_column("report_runs", "modal_call_id") + op.drop_column("report_runs", "year") + op.drop_column("report_runs", "budget_window_run_id") op.drop_column("report_runs", "reform_simulation_run_id") op.drop_column("report_runs", "baseline_simulation_run_id") op.add_column( @@ -925,28 +1065,28 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) op.drop_constraint( - op.f("fk_program_statistics_reform_simulation_run_id_simulation_runs"), + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), "program_statistics", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_id_simulations"), + op.f("fk_program_statistics_baseline_simulation_id_simulations"), "program_statistics", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -983,14 +1123,6 @@ def downgrade() -> None: "poverty", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_poverty_report_run_id_report_runs"), - "poverty", - "report_runs", - ["report_run_id"], - ["id"], - ondelete="CASCADE", - ) op.create_foreign_key( op.f("fk_poverty_simulation_id_simulations"), "poverty", @@ -999,6 +1131,14 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_poverty_report_run_id_report_runs"), + "poverty", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) op.drop_constraint("uq_poverty_simulation_run_measure", "poverty", type_="unique") op.create_index( op.f("ix_poverty_run_simulation_type"), @@ -1026,28 +1166,28 @@ def downgrade() -> None: type_="check", ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), "local_authority_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1081,18 +1221,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1111,14 +1251,6 @@ def downgrade() -> None: "inequality", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_inequality_simulation_id_simulations"), - "inequality", - "simulations", - ["simulation_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_foreign_key( op.f("fk_inequality_report_run_id_report_runs"), "inequality", @@ -1127,6 +1259,14 @@ def downgrade() -> None: ["id"], ondelete="CASCADE", ) + op.create_foreign_key( + op.f("fk_inequality_simulation_id_simulations"), + "inequality", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) op.drop_constraint( "uq_inequality_simulation_run_measure", "inequality", type_="unique" ) @@ -1142,14 +1282,14 @@ def downgrade() -> None: op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" ) op.drop_constraint( - op.f( - "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" - ), + op.f("fk_household_jobs_household_id_households"), "household_jobs", type_="foreignkey", ) op.drop_constraint( - op.f("fk_household_jobs_household_id_households"), + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", type_="foreignkey", ) @@ -1233,28 +1373,28 @@ def downgrade() -> None: type_="check", ) op.drop_constraint( - op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), "constituency_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), "constituency_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1292,18 +1432,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1325,28 +1465,28 @@ def downgrade() -> None: ), ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1354,6 +1494,14 @@ def downgrade() -> None: op.drop_column("budget_summary", "baseline_simulation_run_id") op.drop_table("labor_supply_responses") op.drop_table("labor_supply_decile_responses") + op.drop_index( + op.f("ix_budget_window_runs_tax_benefit_model_version_id"), + table_name="budget_window_runs", + ) + op.drop_index( + "ix_budget_window_runs_report_id_created_at", table_name="budget_window_runs" + ) + op.drop_table("budget_window_runs") op.drop_index( op.f("ix_simulation_runs_tax_benefit_model_version_id"), table_name="simulation_runs", diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index bb8469d0c..4ae843094 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -69,6 +69,7 @@ AggregateOutput, AggregateType, BudgetSummary, + BudgetWindowRun, ChangeAggregate, CongressionalDistrictImpact, ConstituencyImpact, @@ -107,6 +108,7 @@ "AggregateOutput", "AggregateType", "BudgetSummary", + "BudgetWindowRun", "CatalogEntity", "CatalogEnumValue", "ChangeAggregate", diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 2d7a8ff72..215bc4e4a 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -30,6 +30,10 @@ from policyengine_api.data.v2.models.simulations import Simulation, SimulationRun from policyengine_api.data.v2.models.users import User +# A budget window compares at most this many years, none after LAST_YEAR. +MAX_BUDGET_WINDOW_YEARS = 75 +LAST_YEAR = 2200 + class ReportRunStatus(StrEnum): """Lifecycle state of a report run; finished runs record a completion time.""" @@ -89,6 +93,12 @@ class Report(TimestampedModel, table=True): "year IS NULL OR year BETWEEN 1900 AND 2200", name="ck_reports_year", ), + sa.CheckConstraint( + "budget_window_size IS NULL OR (household_id IS NULL " + f"AND budget_window_size BETWEEN 1 AND {MAX_BUDGET_WINDOW_YEARS} " + f"AND year IS NOT NULL AND year + budget_window_size - 1 <= {LAST_YEAR})", + name="ck_reports_budget_window_size", + ), sa.Index("ix_reports_country_type_created_at", "country", "type", "created_at"), ) @@ -138,6 +148,7 @@ class Report(TimestampedModel, table=True): ondelete="SET NULL", ) year: int | None = None + budget_window_size: int | None = None inputs: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) user: User | None = Relationship(back_populates="reports") @@ -158,6 +169,62 @@ class Report(TimestampedModel, table=True): back_populates="report", cascade_delete=True, ) + budget_window_runs: list["BudgetWindowRun"] = Relationship( + back_populates="report", + cascade_delete=True, + ) + + +class BudgetWindowRun(TimestampedModel, table=True): + """One calculation of a budget-window report on one catalog version. + + Each year of the window is a report run of the same report that names this + row and its year. ``recalculate`` says whether those years calculate their + simulations again instead of reusing successful runs. What starts the later + years is decided where the simulation service schedules them. + """ + + __tablename__ = "budget_window_runs" + __table_args__ = ( + sa.UniqueConstraint("id", "report_id", name="uq_budget_window_runs_id_report"), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name="ck_budget_window_runs_terminal_completion", + ), + sa.Index( + "ix_budget_window_runs_report_id_created_at", "report_id", "created_at" + ), + ) + + report_id: UUID = Field( + foreign_key="reports.id", + ondelete="CASCADE", + ) + tax_benefit_model_version_id: UUID = Field( + foreign_key="tax_benefit_model_versions.id", + ondelete="RESTRICT", + index=True, + ) + trigger: ReportRunTrigger = Field( + default=ReportRunTrigger.INITIAL, + sa_type=enum_type(ReportRunTrigger, "v2_report_run_trigger"), + ) + recalculate: bool = False + status: ReportRunStatus = Field( + default=ReportRunStatus.PENDING, + sa_type=enum_type(ReportRunStatus, "v2_report_run_status"), + ) + error_message: str | 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), + ) + + report: Report = Relationship(back_populates="budget_window_runs") class ReportRun(TimestampedModel, table=True): @@ -165,10 +232,12 @@ class ReportRun(TimestampedModel, table=True): It compares ``baseline_simulation_run_id`` with ``reform_simulation_run_id``; a household report with only a baseline has - no reform run. An economy comparison's cliff impact is the baseline's and - reform's cliff gap and share, stored together or not at all; - ``spm_provenance`` is the SPM receipt and ``data_version`` the data release - it used. Its labour supply response has its own two tables. + no reform run. A year of a budget window names its ``budget_window_run_id`` + and ``year``, at most one report run per year of a window run. An economy + comparison's cliff impact is the baseline's and reform's cliff gap and + share, stored together or not at all; ``spm_provenance`` is the SPM receipt + and ``data_version`` the data release it used. Its labour supply response + has its own two tables. """ __tablename__ = "report_runs" @@ -182,6 +251,27 @@ class ReportRun(TimestampedModel, table=True): "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", name="ck_report_runs_terminal_completion", ), + sa.CheckConstraint( + "year IS NULL OR year BETWEEN 1900 AND 2200", + name="ck_report_runs_year", + ), + sa.CheckConstraint( + "budget_window_run_id IS NULL OR year IS NOT NULL", + name="ck_report_runs_budget_window_year", + ), + sa.UniqueConstraint( + "budget_window_run_id", + "year", + name="uq_report_runs_budget_window_run_id_year", + ), + # A year's report run belongs to a budget-window run of its own report. + # NO ACTION waits for the end of the statement, so deleting the report + # deletes the window run and its years together. + sa.ForeignKeyConstraint( + ["budget_window_run_id", "report_id"], + ["budget_window_runs.id", "budget_window_runs.report_id"], + name="fk_report_runs_budget_window_run_id_budget_window_runs", + ), sa.CheckConstraint( "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) " "AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) " @@ -215,6 +305,8 @@ class ReportRun(TimestampedModel, table=True): ondelete="RESTRICT", index=True, ) + budget_window_run_id: UUID | None = None + year: int | None = None country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) status: ReportRunStatus = Field( diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 51dd99771..075ced605 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -136,7 +136,10 @@ class Simulation(TimestampedModel, table=True): class SimulationRun(TimestampedModel, table=True): - """One run of a simulation on one catalog version and PolicyEngine.py version. + """One run of a simulation for one year on one catalog and .py version. + + A simulation has a run per year it is calculated for, so a budget window + compares a report's same two simulations in every year of the window. ``tax_benefit_model_version_id``, ``dataset_id`` and ``region_id`` name the catalog version, dataset and region the run used, which a run on a newer @@ -162,9 +165,14 @@ class SimulationRun(TimestampedModel, table=True): "output_sha256 IS NULL OR length(output_sha256) = 64", name="ck_simulation_runs_output_sha256_length", ), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", + name="ck_simulation_runs_year", + ), sa.Index( "ix_simulation_runs_current_output", "simulation_id", + "year", "status", "country_package_version", "policyengine_version", @@ -178,6 +186,7 @@ class SimulationRun(TimestampedModel, table=True): foreign_key="simulations.id", ondelete="CASCADE", ) + year: int tax_benefit_model_version_id: UUID = Field( foreign_key="tax_benefit_model_versions.id", ondelete="RESTRICT", diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 57d8d443d..0d26f0357 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -148,6 +148,7 @@ def _access( "dynamics": _access(**_RUN_INPUT_READS), # Simulations, reports, their runs and household jobs. The catalog writer # checks that publication leaves the rows of three of them untouched. + "budget_window_runs": _SIMULATION_DEFINITION, "household_jobs": _RUN_STATUS, "report_runs": replace(_RUN_STATUS, catalog_writer=READ), "reports": replace(_SIMULATION_DEFINITION, catalog_writer=READ), diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 38af95c2e..1429bf379 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -64,6 +64,14 @@ class ReportRunSimulationRuns: reform_simulation_run_id: UUID | None +@dataclass(frozen=True) +class ReportRunInputs: + """The package versions and simulation runs one report run uses.""" + + versions: ReportRunPackageVersions + simulation_runs: ReportRunSimulationRuns + + 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. @@ -103,9 +111,8 @@ def create_report_run( session: Session, *, report_id: UUID, - versions: ReportRunPackageVersions, + inputs: ReportRunInputs, trigger: ReportRunTrigger, - simulation_runs: ReportRunSimulationRuns, idempotency_key: UUID | None = None, ) -> ReportRun: """Create one run, or return the run for a retried idempotent request. @@ -117,14 +124,15 @@ def create_report_run( if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: msg = "manual report reruns require an idempotency key" raise ValueError(msg) + versions = inputs.versions 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, - baseline_simulation_run_id=simulation_runs.baseline_simulation_run_id, - reform_simulation_run_id=simulation_runs.reform_simulation_run_id, + baseline_simulation_run_id=inputs.simulation_runs.baseline_simulation_run_id, + reform_simulation_run_id=inputs.simulation_runs.reform_simulation_run_id, country_package_version=versions.country_package_version, policyengine_version=versions.policyengine_version, trigger=trigger, diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 18feb946b..8e3086f15 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -920,6 +920,7 @@ def report_run_checks() -> set[str]: id=simulation_run_id, simulation_id=simulation_id, tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, status="pending", **RUN_VERSIONS, ) diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 2ed3ad317..7b34b1154 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -43,7 +43,12 @@ SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" NEW_TABLES = frozenset( - {"simulation_runs", "labor_supply_responses", "labor_supply_decile_responses"} + { + "simulation_runs", + "budget_window_runs", + "labor_supply_responses", + "labor_supply_decile_responses", + } ) CHANGED_TABLES = frozenset( { @@ -98,6 +103,10 @@ "tax_benefit_model_versions", "RESTRICT", ), + ("budget_window_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), } LABOR_SUPPLY_RESPONSE_COLUMNS = ( "income_effect_earnings_change", @@ -204,7 +213,7 @@ def _enum_type_names(engine: Engine) -> set[str]: ) -def test_the_revision_adds_three_tables_and_changes_only_simulation_side_tables( +def test_the_revision_adds_its_tables_and_changes_only_simulation_side_tables( config: Config, engine: Engine ) -> None: command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) @@ -304,9 +313,10 @@ def _insert_simulation_run( simulation_id: object, *, status: str = "pending", + year: int = 2026, **values: object, ) -> object: - """Store a run of the simulation on the simulation's catalog version.""" + """Store a run of the simulation for ``year`` on its catalog version.""" run_id = uuid4() version_id = connection.execute( text("SELECT tax_benefit_model_version_id FROM simulations WHERE id = :id"), @@ -319,6 +329,7 @@ def _insert_simulation_run( simulation_id=simulation_id, tax_benefit_model_version_id=version_id, status=status, + year=year, **RUN_VERSIONS, **values, ) @@ -370,6 +381,7 @@ def test_rows_naming_a_missing_public_api_row_are_refused( simulation_id=simulation_id, tax_benefit_model_version_id=uuid4(), status="pending", + year=2026, **RUN_VERSIONS, ) with refused(connection, match="foreign key constraint"): @@ -652,6 +664,99 @@ def test_a_report_runs_labor_supply_and_cliff_impact_are_typed( assert _count(connection, table, "report_run_id", report_run_id) == 0 +def _insert_report( + connection: Connection, public_api_rows: PublicApiRows, **values: object +) -> object: + report_id = uuid4() + insert_row( + connection, + "reports", + id=report_id, + label="Budget window", + country="us", + tax_benefit_model_id=public_api_rows.model_id, + inputs="{}", + **values, + ) + return report_id + + +def _insert_budget_window_run( + connection: Connection, public_api_rows: PublicApiRows, report_id: object +) -> object: + window_run_id = uuid4() + insert_row( + connection, + "budget_window_runs", + id=window_run_id, + report_id=report_id, + tax_benefit_model_version_id=public_api_rows.version_id, + trigger="initial", + recalculate=False, + status="running", + ) + return window_run_id + + +def test_budget_windows_hold_one_report_run_per_year( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + for refused_values in ( + {"year": 2026, "budget_window_size": 0}, + {"year": 2026, "budget_window_size": 76}, + {"year": 2190, "budget_window_size": 12}, + {"budget_window_size": 3}, + {"year": 2026, "budget_window_size": 3, "household_id": None}, + ): + values = dict(refused_values) + if "household_id" in values: + values["household_id"] = public_api_rows.household_id + with refused(connection, match="ck_reports_budget_window_size"): + _insert_report(connection, public_api_rows, **values) + report_id = _insert_report( + connection, public_api_rows, year=2026, budget_window_size=3 + ) + other_report_id = _insert_report(connection, public_api_rows, year=2026) + window_run_id = _insert_budget_window_run(connection, public_api_rows, report_id) + simulation_id = _insert_simulation(connection, public_api_rows) + + def insert_year(year: int | None, *, of_report: object = report_id) -> None: + insert_row( + connection, + "report_runs", + id=uuid4(), + report_id=of_report, + budget_window_run_id=window_run_id, + year=year, + baseline_simulation_run_id=_insert_simulation_run( + connection, simulation_id, year=year or 2026 + ), + status="pending", + trigger="initial", + **RUN_VERSIONS, + ) + + insert_year(2026) + insert_year(2027) + with refused(connection, match="uq_report_runs_budget_window_run_id_year"): + insert_year(2027) + with refused(connection, match="ck_report_runs_budget_window_year"): + insert_year(None) + with refused(connection, match="foreign key constraint"): + insert_year(2028, of_report=other_report_id) + with refused(connection, match="ck_simulation_runs_year"): + _insert_simulation_run(connection, simulation_id, year=1800) + with refused(connection, match="foreign key constraint"): + connection.execute( + text("DELETE FROM budget_window_runs WHERE id = :id"), + {"id": window_run_id}, + ) + + connection.execute(text("DELETE FROM reports WHERE id = :id"), {"id": report_id}) + assert _count(connection, "budget_window_runs", "report_id", report_id) == 0 + assert _count(connection, "report_runs", "report_id", report_id) == 0 + + def test_household_jobs_take_a_valid_year_and_record_when_they_end( connection: Connection, public_api_rows: PublicApiRows ) -> None: diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index 10bd85d74..c46fa0f6e 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -515,6 +515,7 @@ def test_republish_preserves_existing_run_and_dataset_version_records( session.flush() simulation_run = SimulationRun( simulation_id=simulation.id, + year=2025, tax_benefit_model_version_id=model_version.id, country_package_version="existing-country-version", policyengine_version="existing-policyengine-version", diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index a995979e5..fba0221d7 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -77,7 +77,14 @@ def _create_simulation_rows( """Create what the simulation API creates for one economy report run.""" ids = { name: uuid4() - for name in ("simulation", "run", "report", "report_run", "household_job") + for name in ( + "simulation", + "run", + "report", + "budget_window_run", + "report_run", + "household_job", + ) } insert_row( connection, @@ -94,6 +101,7 @@ def _create_simulation_rows( id=ids["run"], simulation_id=ids["simulation"], tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, status="pending", **RUN_VERSIONS, ) @@ -104,13 +112,27 @@ def _create_simulation_rows( label="Report", country="us", tax_benefit_model_id=public_api_rows.model_id, + year=2026, + budget_window_size=1, inputs="{}", ) + insert_row( + connection, + "budget_window_runs", + id=ids["budget_window_run"], + report_id=ids["report"], + tax_benefit_model_version_id=public_api_rows.version_id, + trigger="initial", + recalculate=False, + status="pending", + ) insert_row( connection, "report_runs", id=ids["report_run"], report_id=ids["report"], + budget_window_run_id=ids["budget_window_run"], + year=2026, baseline_simulation_run_id=ids["run"], status="pending", trigger="initial", @@ -229,6 +251,15 @@ def test_modal_workers_write_their_run_status_and_results_only( _refused( connection, "DELETE FROM poverty WHERE simulation_run_id = :id", id=ids["run"] ) + _refused( + connection, + "INSERT INTO budget_window_runs (id, report_id, " + "tax_benefit_model_version_id, trigger, recalculate, status) " + "VALUES (:id, :report_id, :version_id, 'manual', true, 'pending')", + id=uuid4(), + report_id=ids["report"], + version_id=public_api_rows.version_id, + ) _refused( connection, "UPDATE poverty SET rate = 0 WHERE simulation_run_id = :id", diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index e004b6cf1..e6430f6a4 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -40,6 +40,7 @@ V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) POST_BASELINE_TABLE_NAMES = frozenset( { + "budget_window_runs", "countries", "legacy_household_mappings", "legacy_policy_mappings", @@ -1023,11 +1024,12 @@ def _simulation_runs_revision_source() -> str: return path.read_text(encoding="utf-8") -def test_simulation_runs_revision_is_generated_with_one_documented_correction() -> None: +def test_simulation_runs_revision_is_generated_with_documented_corrections() -> None: """One generated revision follows the parameter search file revision. - Its one correction makes ``simulation_runs.status`` reuse the type the - baseline created for ``simulations.status``. + Its corrections make new columns reuse enum types the baseline created: + ``simulation_runs.status`` reuses ``simulations.status``'s, and + ``budget_window_runs.status`` and ``trigger`` reuse the report run's. """ revision = _simulation_runs_revision_source() upgrade = revision[ @@ -1043,13 +1045,17 @@ def test_simulation_runs_revision_is_generated_with_one_documented_correction() f'down_revision: Union[str, None] = "{SIMULATION_RUNS_PARENT_REVISION}"' in revision ) - assert revision.count("Dialect correction:") == 1 + assert revision.count("Dialect correction:") == 3 assert "reversibility correction" not in revision - assert 'name="v2_simulation_status",\n create_type=False,' in ( - upgrade - ) + for type_name in ( + "v2_simulation_status", + "v2_report_run_status", + "v2_report_run_trigger", + ): + assert f'name="{type_name}",\n create_type=False,' in upgrade for table_name in ( "simulation_runs", + "budget_window_runs", "labor_supply_responses", "labor_supply_decile_responses", ): diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index df7828892..45fc7d21a 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 45 +assert len(V2_METADATA.tables) == 46 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index cb9363e95..306bbac38 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -46,6 +46,7 @@ "reports", "simulation_runs", "report_runs", + "budget_window_runs", "household_jobs", } ) diff --git a/tests/unit/v2/test_report_run_creation.py b/tests/unit/v2/test_report_run_creation.py index b1e1d8f0b..ac0959f9b 100644 --- a/tests/unit/v2/test_report_run_creation.py +++ b/tests/unit/v2/test_report_run_creation.py @@ -20,6 +20,7 @@ ) from policyengine_api.data.v2.report_runs import ( ReportNotFoundError, + ReportRunInputs, ReportRunPackageVersions, ReportRunSimulationRuns, create_report_run, @@ -67,12 +68,14 @@ def test_unknown_report_is_rejected_before_any_validation(engine) -> None: create_report_run( session, report_id=uuid4(), - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="", + policyengine_version="", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, ) @@ -86,12 +89,14 @@ def test_manual_run_without_a_key_is_rejected_before_version_checks(engine) -> N create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="", + policyengine_version="", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, ) @@ -113,12 +118,14 @@ def test_empty_package_versions_are_rejected( create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_package_version, - policyengine_version=policyengine_version, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version=country_package_version, + policyengine_version=policyengine_version, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.SYSTEM, - simulation_runs=SIMULATION_RUNS, ) @@ -129,12 +136,14 @@ def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: run = create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version="4.5.6", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.INITIAL, - simulation_runs=SIMULATION_RUNS, ) stored = session.exec(select(ReportRun)).one() @@ -168,12 +177,14 @@ def failing_savepoint(): create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version="4.5.6", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, idempotency_key=uuid4(), ) diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py index 89f83282d..ecbf58ff4 100644 --- a/tests/unit/v2/test_report_runs.py +++ b/tests/unit/v2/test_report_runs.py @@ -26,6 +26,7 @@ TaxBenefitModelVersion, ) from policyengine_api.data.v2.report_runs import ( + ReportRunInputs, ReportRunPackageVersions, ReportRunSimulationRuns, ReportRunStateError, @@ -87,12 +88,14 @@ def _create_run( return create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_version, - policyengine_version=policyengine_version, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version=country_version, + policyengine_version=policyengine_version, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, idempotency_key=key if key is not None else uuid4(), ) @@ -154,12 +157,14 @@ def test_manual_rerun_requires_an_idempotency_key(engine) -> None: create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, ) @@ -217,12 +222,14 @@ def request_rerun() -> UUID: run = create_report_run( session, report_id=report_id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, - simulation_runs=SIMULATION_RUNS, idempotency_key=request_key, ) run_id = run.id diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 84255db5b..67c3d7fda 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -12,6 +12,7 @@ from policyengine_api.data.v2.models import ( V2_METADATA, + BudgetWindowRun, LaborSupplyDecileResponse, LaborSupplyResponse, ReportRun, @@ -63,6 +64,7 @@ "simulation_runs", "reports", "report_runs", + "budget_window_runs", "labor_supply_responses", "labor_supply_decile_responses", "household_jobs", @@ -104,6 +106,10 @@ "tax_benefit_model_versions", "RESTRICT", ), + ("budget_window_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), } @@ -167,6 +173,7 @@ def test_each_simulation_run_records_its_versions_status_and_output() -> None: required = { "simulation_id", + "year", "tax_benefit_model_version_id", "country_package_version", "policyengine_version", @@ -199,11 +206,13 @@ def test_each_simulation_run_records_its_versions_status_and_output() -> None: simulation_key = _foreign_key(simulation_runs, "simulation_id") assert simulation_key.target_fullname == "simulations.id" assert simulation_key.ondelete == "CASCADE" - # The catalog version, dataset and region are the Public API's rows, so a - # run refers to them without a foreign key. assert {key.target_fullname for key in simulation_runs.foreign_keys} == { - "simulations.id" + "simulations.id", + "tax_benefit_model_versions.id", + "datasets.id", + "regions.id", } + assert "ck_simulation_runs_year" in _check_names(simulation_runs) assert SimulationRun.__module__.endswith(".simulations") @@ -218,9 +227,12 @@ def test_run_tables_require_a_completion_time_and_index_running_runs() -> None: "ck_simulation_runs_output_sha256_length", } <= _check_names(simulation_runs) assert "ck_household_jobs_terminal_completion" in _check_names(household_jobs) + # A simulation has a run per year, so its newest successful run is found + # per year. assert _indexes(simulation_runs) == { "ix_simulation_runs_current_output": ( "simulation_id", + "year", "status", "country_package_version", "policyengine_version", @@ -228,6 +240,9 @@ def test_run_tables_require_a_completion_time_and_index_running_runs() -> None: "id", ), "ix_simulation_runs_status_started_at": ("status", "started_at"), + "ix_simulation_runs_tax_benefit_model_version_id": ( + "tax_benefit_model_version_id", + ), } assert _indexes(report_runs)["ix_report_runs_status_started_at"] == ( "status", @@ -374,6 +389,75 @@ def test_report_runs_name_their_baseline_and_reform_simulation_runs() -> None: assert "v2_simulation_role" not in enum_type_names +def test_budget_window_reports_name_how_many_years_they_compare() -> None: + reports = V2_METADATA.tables["reports"] + + assert reports.c.budget_window_size.nullable + assert isinstance(reports.c.budget_window_size.type, sa.Integer) + assert "ck_reports_budget_window_size" in _check_names(reports) + + +def test_budget_window_runs_record_each_calculation_of_a_window() -> None: + windows = V2_METADATA.tables["budget_window_runs"] + + required = { + "report_id", + "tax_benefit_model_version_id", + "trigger", + "recalculate", + "status", + "created_at", + "updated_at", + } + optional = {"error_message", "started_at", "completed_at"} + assert set(windows.c.keys()) == {"id"} | required | optional + assert all(not windows.c[name].nullable for name in required) + assert all(windows.c[name].nullable for name in optional) + # A window run has a report run's states and starts for a report run's reasons. + assert windows.c.status.type.name == "v2_report_run_status" + assert windows.c.trigger.type.name == "v2_report_run_trigger" + report_key = _foreign_key(windows, "report_id") + assert report_key.target_fullname == "reports.id" + assert report_key.ondelete == "CASCADE" + version_key = _foreign_key(windows, "tax_benefit_model_version_id") + assert version_key.target_fullname == "tax_benefit_model_versions.id" + assert version_key.ondelete == "RESTRICT" + assert ("id", "report_id") in _unique_column_sets(windows) + assert "ck_budget_window_runs_terminal_completion" in _check_names(windows) + assert _indexes(windows)["ix_budget_window_runs_report_id_created_at"] == ( + "report_id", + "created_at", + ) + assert BudgetWindowRun.__module__.endswith(".reports") + + +def test_report_runs_hold_one_year_of_a_budget_window_run() -> None: + report_runs = V2_METADATA.tables["report_runs"] + + assert report_runs.c.budget_window_run_id.nullable + assert report_runs.c.year.nullable + (window_key,) = ( + constraint + for constraint in report_runs.foreign_key_constraints + if constraint.referred_table.name == "budget_window_runs" + ) + # A year's report run belongs to a window run of the same report. The check + # waits for the end of the statement, so deleting the report deletes both. + assert [column.name for column in window_key.columns] == [ + "budget_window_run_id", + "report_id", + ] + assert [element.target_fullname for element in window_key.elements] == [ + "budget_window_runs.id", + "budget_window_runs.report_id", + ] + assert window_key.ondelete is None + assert ("budget_window_run_id", "year") in _unique_column_sets(report_runs) + assert {"ck_report_runs_year", "ck_report_runs_budget_window_year"} <= ( + _check_names(report_runs) + ) + + def test_household_jobs_calculate_one_household_year_and_policy() -> None: household_jobs = V2_METADATA.tables["household_jobs"] @@ -392,12 +476,12 @@ def test_household_jobs_calculate_one_household_year_and_policy() -> None: assert household_jobs.c.modal_call_id.type.length == 255 assert "ck_household_jobs_year" in _check_names(household_jobs) assert "ix_household_jobs_household_id" in _indexes(household_jobs) - assert "household_id" not in { - key.parent.name for key in household_jobs.foreign_keys - } - assert "tax_benefit_model_version_id" not in { - key.parent.name for key in household_jobs.foreign_keys - } + assert _foreign_key(household_jobs, "household_id").target_fullname == ( + "households.id" + ) + assert _foreign_key( + household_jobs, "tax_benefit_model_version_id" + ).target_fullname == ("tax_benefit_model_versions.id") def test_per_simulation_results_belong_to_one_simulation_run() -> None: From c6bb73a879b853eb1175e4060f299a74ed9ee89a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:11:06 +0400 Subject: [PATCH 17/22] Grant the simulation API's roles UPDATE on run columns only The privilege manifest can now declare a privilege on some of a table's columns (update_columns), beside whole-table privileges. The Modal worker updates only a run's status, Modal call ID, error, times and result columns (a simulation run's output file, digest and household result, a report run's cliff impact, SPM receipt and data version, a household job's result), and the simulation service only the status, call ID, error and completion columns of the runs and household jobs whose Modal calls it starts. PostgreSQL then refuses any change to what a run is. apply reads column grants from pg_attribute, grants each declared column privilege a role lacks with one GRANT per column privilege, and reports undeclared column grants. verify checks each declared column with has_column_privilege and fails when a role also holds the privilege on the whole table, which would override the column limits. The two privilege rules and the role tests change to match. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 19 +- docs/generated/rules.md | 4 +- .../data/v2/privilege_reconciliation.py | 341 ++++++++++++++---- .../data/v2/privilege_requirements.py | 10 +- policyengine_api/data/v2/privileges.py | 117 ++++-- rules/registry.toml | 5 +- .../test_v2_privilege_reconciliation.py | 79 +++- .../test_v2_simulation_role_privileges.py | 57 ++- tests/unit/v2/test_privilege_manifest.py | 109 +++++- .../unit/v2/test_privilege_reconciliation.py | 139 +++++++ 10 files changed, 765 insertions(+), 115 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d56cde7f9..d1737580e 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -255,11 +255,15 @@ and a revision runs once per database, so it can't repair drift. | `policyengine_v2_runtime` | The Public API on Cloud Run | | `policyengine_v2_catalog_writer` | `scripts/publish_v2_metadata_catalog.py` | | `policyengine_stage12_` | The simulation API and its Modal worker, on the temporary comparison-run tables | -| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs and household jobs, and reads the policy, dynamic, household and catalog rows a run needs | -| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs. Its `UPDATE` is table-level: the workers' code, not the database, limits it to their own run's status and result columns | +| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs, budget-window runs and household jobs, and reads the policy, dynamic, household and catalog rows a run needs. It updates only the status, Modal call ID, error and completion columns of the runs and household jobs whose Modal calls it starts | +| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs. It updates only a run's status, Modal call ID, error, time and result columns; the workers' code, not the database, keeps a worker to its own run | Every v2 table has an entry for every role, so "no access" is a decision, not -an omission. The migration role owns every table and isn't listed. Three rules +an omission. A privilege is declared on the whole table, or with +`update_columns(...)` on only the columns it names; PostgreSQL then refuses a +statement that sets any other column, so neither simulation API role can +change what a run is. A role never holds both kinds of `UPDATE` on one table: +the whole-table one would let it set every column. The migration role owns every table and isn't listed. Three rules check the declarations: - `api.v2.runtime-privileges-declared`: a unit test reads each connector @@ -276,9 +280,12 @@ check the declarations: `scripts/reconcile_v2_privileges.py apply` right after the schema upgrade, then `verify`, before the catalog is seeded. -`apply` grants each declared privilege a role lacks. It never revokes: it -reports privileges the file doesn't declare, and the projects' default -privileges, in the job summary. It grants nothing when a declared role or +`apply` grants each declared privilege a role lacks, as one `GRANT` per role +and table and one per column privilege. It never revokes: it reports +privileges the file doesn't declare, on whole tables or on columns, and the +projects' default privileges, in the job summary. `verify` checks each +column of a column privilege with `has_column_privilege`, and fails when a +role holds the same privilege on the whole table. It grants nothing when a declared role or table is missing. The roles themselves are provisioned outside the repository; `--create-missing-roles` is accepted only for the disposable test database. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index d68a06194..ebe284c0a 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -285,8 +285,8 @@ flowchart LR | `api.migration.bundle-supported-by-simulation-api` | The installed PolicyEngine bundle is the latest one the simulation API has published, checked on every pull request and before each staging and production deploy. | mechanical | enforced | yes | CI job `ensure-policyengine-bundle-supported-by-simulation-api` | `pr.yml:ensure-policyengine-bundle-supported-by-simulation-api` | — | | `api.migration.contracts-consistent` | Migration route groups and workflow contracts are internally consistent, and their generated docs are current. | mechanical | enforced | yes | quality guard `migration-contracts` | `pr.yml:quality-guards` | **unproven** | | `api.migration.cross-database-integration` | Cross-database integration tests pass against real MySQL and Postgres. | mechanical | enforced | yes | CI job `v2-integration-check` | `pr.yml:v2-integration-check` | — | -| `api.v2.deployed-privileges-match` | Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead. | mechanical | enforced | yes | CI job `seed`, `catalog-runtime-integration` | `seed-v2-database.yml:seed`, `v2-integration-check.yml:catalog-runtime-integration` | yes | -| `api.v2.runtime-privileges-declared` | Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py, and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it. | mechanical | enforced | yes | pytest `tests/unit/v2/test_privilege_manifest.py` | `pr.yml:test` | yes | +| `api.v2.deployed-privileges-match` | Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares, on the whole table or on the columns it names, right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing or a role holds on a whole table a privilege declared for only some of its columns. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead. | mechanical | enforced | yes | CI job `seed`, `catalog-runtime-integration` | `seed-v2-database.yml:seed`, `v2-integration-check.yml:catalog-runtime-integration` | yes | +| `api.v2.runtime-privileges-declared` | Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py; a privilege declared on only some of a table's columns names columns of that table and is never also declared on the whole table; and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it, on the whole table. | mechanical | enforced | yes | pytest `tests/unit/v2/test_privilege_manifest.py` | `pr.yml:test` | yes | | `api.v2.runtime-privileges-tested-as-runtime-role` | The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, the comparison-run services as policyengine_stage12_, and the simulation API's table writes as policyengine_v2_simulation_service and policyengine_v2_modal_worker, holding only their declared privileges, so a statement those privileges don't allow fails CI. | mechanical | enforced | yes | CI job `v2-integration-check`, `catalog-runtime-integration` | `pr.yml:v2-integration-check`, `v2-integration-check.yml:catalog-runtime-integration` | — | ```mermaid diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index c449c93cf..faea4818f 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -4,12 +4,16 @@ connected as the migration role, which owns every table and so may grant on it. Two modes: -- ``apply`` reads the managed roles' current grants, grants every declared - privilege a role lacks, and reports privileges the manifest doesn't declare. - It never revokes, and it changes nothing when a declared role or table is - missing. Running it again grants nothing more. -- ``verify`` checks every declared privilege with ``has_table_privilege`` - (rule ``api.v2.deployed-privileges-match``) and fails on any that is missing. +- ``apply`` reads the managed roles' current grants, on whole tables and on + columns, grants every declared privilege a role lacks, and reports privileges + the manifest doesn't declare. It never revokes, and it changes nothing when a + declared role or table is missing. Running it again grants nothing more. +- ``verify`` checks every declared privilege with ``has_table_privilege``, or + ``has_column_privilege`` for each column of a column privilege (rule + ``api.v2.deployed-privileges-match``), and fails on any that is missing. It + also fails when a role holds on the whole table a privilege the manifest + declares for only some of its columns, since that lets the role use it on + every column. Both skip a missing role the manifest lists in ``ROLES_AWAITING_PROVISIONING``, and name it in their output, so a pull request can declare a role before an @@ -48,6 +52,8 @@ from policyengine_api.data.v2.privileges import ( PERSISTENT_ENVIRONMENTS, ROLES_AWAITING_PROVISIONING, + ColumnPrivilege, + DeclaredPrivilege, TablePrivilege, declared_privileges, ) @@ -74,11 +80,16 @@ class UnsafeIdentifierError(PrivilegeReconciliationError): @dataclass(frozen=True, order=True) class PrivilegeEntry: - """One privilege one role holds, or should hold, on one table.""" + """One privilege one role holds, or should hold, on one table. + + ``column`` names the one column a column privilege is on; it is empty for a + privilege on the whole table. + """ role: str table: str privilege: str + column: str = "" @dataclass(frozen=True) @@ -112,44 +123,69 @@ def can_apply(self) -> bool: return not self.missing_roles and not self.missing_tables -def plan_reconciliation( - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], +def _declared_entries( + role: str, table: str, privileges: Collection[DeclaredPrivilege] +) -> list[PrivilegeEntry]: + """Return one entry per whole-table privilege and per column of a column one.""" + entries = [] + for privilege in privileges: + if isinstance(privilege, ColumnPrivilege): + entries += [ + PrivilegeEntry(role, table, privilege.privilege.value, column) + for column in privilege.columns + ] + else: + entries.append(PrivilegeEntry(role, table, privilege.value)) + return entries + + +def plan_reconciliation( # noqa: PLR0913 # the database facts it compares + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], granted: Mapping[str, Mapping[str, Collection[str]]], *, existing_roles: Collection[str], existing_tables: Collection[str], roles_awaiting_creation: Collection[str] = (), + granted_columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] + | None = None, ) -> ReconciliationPlan: """Compare declared privileges with what the database grants. - ``granted`` holds the managed roles' direct grants, by role then table. A - missing role in ``roles_awaiting_creation`` is skipped instead of missing. + ``granted`` holds the managed roles' direct grants on whole tables, by role + then table, and ``granted_columns`` their direct grants on columns, by role, + table, then column. A missing role in ``roles_awaiting_creation`` is skipped + instead of missing. """ + granted_columns = granted_columns or {} absent_roles = set(declared) - set(existing_roles) skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) declared_tables = {table for tables in declared.values() for table in tables} missing_tables = tuple(sorted(declared_tables - set(existing_tables))) - grants = sorted( - PrivilegeEntry(role, table, privilege.value) + declared_entries = { + entry for role, tables in declared.items() - if role not in absent_roles for table, privileges in tables.items() - if table not in missing_tables - for privilege in privileges - if privilege.value not in granted.get(role, {}).get(table, ()) - ) - undeclared = sorted( + for entry in _declared_entries(role, table, privileges) + } + granted_entries = { PrivilegeEntry(role, table, privilege) for role, tables in granted.items() for table, privileges in tables.items() for privilege in privileges - if privilege - not in { - declared_privilege.value - for declared_privilege in declared.get(role, {}).get(table, frozenset()) - } + } | { + PrivilegeEntry(role, table, privilege, column) + for role, tables in granted_columns.items() + for table, columns in tables.items() + for column, privileges in columns.items() + for privilege in privileges + } + grants = sorted( + entry + for entry in declared_entries - granted_entries + if entry.role not in absent_roles and entry.table not in missing_tables ) + undeclared = sorted(granted_entries - declared_entries) return ReconciliationPlan( grants=tuple(grants), undeclared=tuple(undeclared), @@ -161,7 +197,7 @@ def plan_reconciliation( def _safe(identifier: str) -> str: if SAFE_IDENTIFIER.fullmatch(identifier) is None: - msg = f"refusing to use {identifier!r} as a role or table name" + msg = f"refusing to use {identifier!r} as a role, table or column name" raise UnsafeIdentifierError(msg) return identifier @@ -169,44 +205,111 @@ def _safe(identifier: str) -> str: def grant_statements( plan: ReconciliationPlan, quote_identifier: Callable[[str], str] ) -> list[str]: - """Return one GRANT per role and table, with its privileges in name order.""" - privileges_by_target: dict[tuple[str, str], set[str]] = defaultdict(set) + """Return the GRANTs for a plan, in role and table order. + + Each role and table gets one GRANT of its whole-table privileges, then one + per column privilege, naming its columns in name order. + """ + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str, str], set[str]] = defaultdict(set) for entry in plan.grants: - privileges_by_target[(_safe(entry.role), _safe(entry.table))].add( - TablePrivilege(entry.privilege).value + target = (_safe(entry.role), _safe(entry.table)) + privilege = TablePrivilege(entry.privilege).value + if entry.column: + columns_by_privilege[(*target, privilege)].add(_safe(entry.column)) + else: + table_privileges[target].add(privilege) + statements: list[tuple[tuple[str, str, str], str]] = [ + ( + (role, table, ""), + f"GRANT {', '.join(sorted(privileges))} ON TABLE " + f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", ) - return [ - f"GRANT {', '.join(sorted(privileges))} ON TABLE " - f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " - f"TO {quote_identifier(role)}" - for (role, table), privileges in sorted( - privileges_by_target.items(), key=lambda item: (item[0][0], item[0][1]) + for (role, table), privileges in table_privileges.items() + ] + statements += [ + ( + (role, table, privilege), + f"GRANT {privilege} " + f"({', '.join(quote_identifier(column) for column in sorted(columns))}) " + f"ON TABLE {quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", ) + for (role, table, privilege), columns in columns_by_privilege.items() ] + return [statement for _target, statement in sorted(statements)] def verification_failures( - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], has_privilege: Callable[[str, str, str], bool], + has_column_privilege: Callable[[str, str, str, str], bool] | None = None, ) -> list[PrivilegeEntry]: - """Return every declared privilege the role doesn't effectively hold.""" + """Return every declared privilege the role doesn't effectively hold. + + A column privilege is checked on each of its columns with + ``has_column_privilege``, which is required when one is declared. + """ + failures = [] + for role, tables in declared.items(): + for table, privileges in tables.items(): + for entry in _declared_entries(role, table, privileges): + if not entry.column: + held = has_privilege(role, table, entry.privilege) + elif has_column_privilege is None: + msg = "checking a column privilege needs has_column_privilege" + raise ValueError(msg) + else: + held = has_column_privilege( + role, table, entry.column, entry.privilege + ) + if not held: + failures.append(entry) + return sorted(failures) + + +def table_wide_overrides( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + has_privilege: Callable[[str, str, str], bool], +) -> list[PrivilegeEntry]: + """Return each column privilege a role also holds on the whole table. + + Holding it on the whole table lets the role use it on every column, past + the columns the manifest declares. + """ return sorted( - PrivilegeEntry(role, table, privilege.value) + PrivilegeEntry(role, table, privilege.privilege.value) for role, tables in declared.items() for table, privileges in tables.items() for privilege in privileges - if not has_privilege(role, table, privilege.value) + if isinstance(privilege, ColumnPrivilege) + and has_privilege(role, table, privilege.privilege.value) ) def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: - rows_by_target: dict[tuple[str, str], list[str]] = defaultdict(list) + """Return one row per role and table, naming each column privilege's columns.""" + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str], dict[str, set[str]]] = defaultdict( + lambda: defaultdict(set) + ) for entry in entries: - rows_by_target[(entry.role, entry.table)].append(entry.privilege) - return [ - f"| `{role}` | `{table}` | {', '.join(sorted(privileges))} |" - for (role, table), privileges in sorted(rows_by_target.items()) - ] + target = (entry.role, entry.table) + if entry.column: + columns_by_privilege[target][entry.privilege].add(entry.column) + else: + table_privileges[target].add(entry.privilege) + rows = [] + for role, table in sorted(set(table_privileges) | set(columns_by_privilege)): + described = sorted(table_privileges.get((role, table), ())) + [ + f"{privilege} ({', '.join(sorted(columns))})" + for privilege, columns in sorted( + columns_by_privilege.get((role, table), {}).items() + ) + ] + rows.append(f"| `{role}` | `{table}` | {', '.join(described)} |") + return rows def _skipped_roles_line(roles: Iterable[str]) -> str: @@ -326,10 +429,34 @@ def render_summary( """ ).bindparams(bindparam("roles", expanding=True)) +_GRANTED_COLUMN_PRIVILEGES = text( + """ + SELECT grantee_role.rolname AS role_name, + table_class.relname AS table_name, + table_column.attname AS column_name, + acl_entry.privilege_type AS privilege + FROM pg_attribute AS table_column + JOIN pg_class AS table_class ON table_class.oid = table_column.attrelid + JOIN pg_namespace AS table_namespace + ON table_namespace.oid = table_class.relnamespace + CROSS JOIN LATERAL aclexplode(table_column.attacl) AS acl_entry + JOIN pg_roles AS grantee_role ON grantee_role.oid = acl_entry.grantee + WHERE table_namespace.nspname = :schema + AND table_class.relkind IN ('r', 'p', 'v', 'm') + AND table_column.attnum > 0 + AND NOT table_column.attisdropped + AND grantee_role.rolname IN :roles + """ +).bindparams(bindparam("roles", expanding=True)) + _HAS_TABLE_PRIVILEGE = text( "SELECT has_table_privilege(:role, :qualified_table, :privilege)" ) +_HAS_COLUMN_PRIVILEGE = text( + "SELECT has_column_privilege(:role, :qualified_table, :column, :privilege)" +) + def read_existing_roles(connection: Connection, roles: Collection[str]) -> set[str]: """Return which of ``roles`` exist.""" @@ -353,6 +480,23 @@ def read_granted_privileges( return {role: dict(tables) for role, tables in granted.items()} +def read_granted_column_privileges( + connection: Connection, roles: Collection[str] +) -> dict[str, dict[str, dict[str, set[str]]]]: + """Return the roles' direct grants on columns of public tables and views.""" + granted: dict[str, dict[str, dict[str, set[str]]]] = defaultdict( + lambda: defaultdict(lambda: defaultdict(set)) + ) + for row in connection.execute( + _GRANTED_COLUMN_PRIVILEGES, {"schema": SCHEMA, "roles": list(roles)} + ): + granted[row.role_name][row.table_name][row.column_name].add(row.privilege) + return { + role: {table: dict(columns) for table, columns in tables.items()} + for role, tables in granted.items() + } + + def read_default_privileges( connection: Connection, roles: Collection[str] ) -> list[DefaultPrivilege]: @@ -387,7 +531,7 @@ def create_missing_roles(connection: Connection, roles: Iterable[str]) -> list[s def apply_declared_privileges( connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], *, roles_awaiting_creation: Collection[str] = (), ) -> tuple[ReconciliationPlan, list[DefaultPrivilege]]: @@ -403,6 +547,7 @@ def apply_declared_privileges( existing_roles=read_existing_roles(connection, roles), existing_tables=read_public_tables(connection), roles_awaiting_creation=roles_awaiting_creation, + granted_columns=read_granted_column_privileges(connection, roles), ) if plan.can_apply: for statement in grant_statements( @@ -414,7 +559,7 @@ def apply_declared_privileges( def skipped_roles( connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], roles_awaiting_creation: Collection[str], ) -> tuple[str, ...]: """Return the declared roles awaiting creation that don't exist yet.""" @@ -424,16 +569,14 @@ def skipped_roles( ) -def check_declared_privileges( - connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], -) -> list[PrivilegeEntry]: - """Return every declared privilege a role doesn't effectively hold. +def _privilege_checks( + connection: Connection, roles: Collection[str] +) -> tuple[Callable[[str, str, str], bool], Callable[[str, str, str, str], bool]]: + """Return checks of a role's effective table and column privileges. - A missing role or table counts as missing every privilege declared for it; - leave out the roles ``skipped_roles`` returns to skip them. + A missing role or table holds nothing. """ - existing_roles = read_existing_roles(connection, list(declared)) + existing_roles = read_existing_roles(connection, roles) existing_tables = read_public_tables(connection) def has_privilege(role: str, table: str, privilege: str) -> bool: @@ -450,7 +593,46 @@ def has_privilege(role: str, table: str, privilege: str) -> bool: ).scalar_one() ) - return verification_failures(declared, has_privilege) + def has_column_privilege( + role: str, table: str, column: str, privilege: str + ) -> bool: + if role not in existing_roles or table not in existing_tables: + return False + return bool( + connection.execute( + _HAS_COLUMN_PRIVILEGE, + { + "role": role, + "qualified_table": f"{SCHEMA}.{_safe(table)}", + "column": _safe(column), + "privilege": privilege, + }, + ).scalar_one() + ) + + return has_privilege, has_column_privilege + + +def check_declared_privileges( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], +) -> list[PrivilegeEntry]: + """Return every declared privilege a role doesn't effectively hold. + + A missing role or table counts as missing every privilege declared for it; + leave out the roles ``skipped_roles`` returns to skip them. + """ + has_privilege, has_column_privilege = _privilege_checks(connection, list(declared)) + return verification_failures(declared, has_privilege, has_column_privilege) + + +def check_table_wide_overrides( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], +) -> list[PrivilegeEntry]: + """Return each column privilege a role also holds on the whole table.""" + has_privilege, _has_column_privilege = _privilege_checks(connection, list(declared)) + return table_wide_overrides(declared, has_privilege) # --- Command-line interface ----------------------------------------------- @@ -480,9 +662,37 @@ def _write_summary(environ: Mapping[str, str], summary: str) -> None: summary_file.write(summary) +def _failure_row(failure: PrivilegeEntry) -> str: + if failure.column: + return ( + f"- `{failure.role}` lacks {failure.privilege} on column " + f"`{failure.column}` of `{failure.table}`" + ) + return f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" + + +def _verification_report( + failures: Sequence[PrivilegeEntry], overrides: Sequence[PrivilegeEntry] +) -> str: + sections = [] + if failures: + rows = "\n".join(_failure_row(failure) for failure in failures) + sections.append(f"Declared privileges missing after reconciling:\n\n{rows}\n") + if overrides: + rows = "\n".join( + f"- `{override.role}` holds {override.privilege} on the whole of " + f"`{override.table}`, which overrides its declared columns" + for override in overrides + ) + sections.append( + f"Privileges held on whole tables, past their declared columns:\n\n{rows}\n" + ) + return "\n".join(sections) + + def _verify( connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], roles_awaiting_creation: Collection[str], environment: str, environ: Mapping[str, str], @@ -490,24 +700,21 @@ def _verify( """Check every declared privilege, skipping missing roles awaiting creation. A missing role in ``roles_awaiting_creation`` is skipped; any other missing - role counts as missing every privilege declared for it. The roles skipped - and the privileges missing are written to stdout and the job summary under - one heading. + role counts as missing every privilege declared for it. A role holding on a + whole table a privilege declared for only some of its columns fails too. + The roles skipped, the privileges missing and the whole-table privileges + are written to stdout and the job summary under one heading. """ skipped = skipped_roles(connection, declared, roles_awaiting_creation) checked = {role: tables for role, tables in declared.items() if role not in skipped} failures = check_declared_privileges(connection, checked) + overrides = check_table_wide_overrides(connection, checked) heading = f"## v2 table privileges ({environment})\n\n" skipped_note = f"{_skipped_roles_line(skipped)}\n\n" if skipped else "" - if failures: - rows = "\n".join( - f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" - for failure in failures - ) + if failures or overrides: _write_summary( environ, - f"{heading}{skipped_note}" - f"Declared privileges missing after reconciling:\n\n{rows}\n", + f"{heading}{skipped_note}{_verification_report(failures, overrides)}", ) return 1 if skipped: diff --git a/policyengine_api/data/v2/privilege_requirements.py b/policyengine_api/data/v2/privilege_requirements.py index f159f747f..ea860adfc 100644 --- a/policyengine_api/data/v2/privilege_requirements.py +++ b/policyengine_api/data/v2/privilege_requirements.py @@ -25,7 +25,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from policyengine_api.data.v2.privileges import TablePrivilege +from policyengine_api.data.v2.privileges import DeclaredPrivilege, TablePrivilege if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping @@ -226,9 +226,13 @@ def required_privileges( def undeclared_requirements( requirements: Iterable[RequiredPrivilege], - declared: Mapping[str, frozenset[TablePrivilege]], + declared: Mapping[str, frozenset[DeclaredPrivilege]], ) -> list[RequiredPrivilege]: - """Return the requirements the declared table privileges don't cover, sorted.""" + """Return the requirements the declared table privileges don't cover, sorted. + + A privilege declared on only some columns doesn't cover a connector's need + for it on the whole table. + """ return sorted( requirement for requirement in requirements diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 0d26f0357..18f3ab8f1 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -10,7 +10,10 @@ ``api.v2.runtime-privileges-tested-as-runtime-role``). Every table in the v2 SQLModel metadata, plus ``alembic_version``, has an entry -for every managed role, so "no access" is a decision, not an omission. +for every managed role, so "no access" is a decision, not an omission. A +privilege is on the whole table, or a ``ColumnPrivilege`` on the columns it +names; a role never holds both kinds of the same privilege on one table, since +the whole-table one would let it change every column. Roles and what decides their privileges: @@ -25,15 +28,20 @@ publication writes and reads; its code isn't organized as connector modules, so the deploy check, not the connector analysis, covers it. - ``policyengine_v2_simulation_service``: the simulation API on Cloud Run. It - creates simulations, reports, their runs and household jobs, and reads the - policy, dynamic, household and catalog rows a run needs. + creates simulations, reports, their runs, budget-window runs and household + jobs, and reads the policy, dynamic, household and catalog rows a run needs. + It starts the Modal calls of runs and household jobs, so it may update their + status, Modal call ID, error and completion columns: it records a call's ID + once Modal accepts the call, and fails the run when Modal refuses it. - ``policyengine_v2_modal_worker``: the simulation API's Modal workers. They write the status of the runs and household jobs they were given, and their - results, and read the rows a run needs. Their ``UPDATE`` is table-level, like - every privilege here: the database doesn't limit it to the status and result - columns or to the worker's own run, so the workers' code does. Both - simulation API roles' code lives in that repository, so the deploy check - covers them here. + results, and read the rows a run needs. + +Both simulation API roles update only the columns ``update_columns`` names, so +the database refuses any change to what a run is (its simulation or report, +versions, dataset, region, year, simulation runs or budget-window run). It +doesn't limit a worker to its own run; the workers' code does. Their code lives +in the simulation API's repository, so the deploy check covers them here. The Public API reads the simulation tables it serves, and writes none of them. @@ -78,10 +86,28 @@ class TablePrivilege(StrEnum): DELETE = "DELETE" +@dataclass(frozen=True) +class ColumnPrivilege: + """A privilege on some of a table's columns instead of the whole table. + + PostgreSQL refuses a statement that uses the privilege on any other column. + """ + + privilege: TablePrivilege + columns: frozenset[str] + + +DeclaredPrivilege = TablePrivilege | ColumnPrivilege + + +def update_columns(*columns: str) -> ColumnPrivilege: + """Declare ``UPDATE`` on only the named columns.""" + return ColumnPrivilege(TablePrivilege.UPDATE, frozenset(columns)) + + NO_ACCESS: Final[frozenset[TablePrivilege]] = frozenset() READ: Final = frozenset({TablePrivilege.SELECT}) READ_INSERT: Final = frozenset({TablePrivilege.SELECT, TablePrivilege.INSERT}) -READ_UPDATE: Final = frozenset({TablePrivilege.SELECT, TablePrivilege.UPDATE}) READ_INSERT_UPDATE: Final = frozenset( {TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE} ) @@ -99,20 +125,20 @@ class TablePrivilege(StrEnum): class TableAccess: """What each managed role may do on one table.""" - runtime: frozenset[TablePrivilege] - catalog_writer: frozenset[TablePrivilege] - stage12: frozenset[TablePrivilege] - simulation_service: frozenset[TablePrivilege] - modal_worker: frozenset[TablePrivilege] + runtime: frozenset[DeclaredPrivilege] + catalog_writer: frozenset[DeclaredPrivilege] + stage12: frozenset[DeclaredPrivilege] + simulation_service: frozenset[DeclaredPrivilege] + modal_worker: frozenset[DeclaredPrivilege] def _access( *, - runtime: frozenset[TablePrivilege] = NO_ACCESS, - catalog_writer: frozenset[TablePrivilege] = NO_ACCESS, - stage12: frozenset[TablePrivilege] = NO_ACCESS, - simulation_service: frozenset[TablePrivilege] = NO_ACCESS, - modal_worker: frozenset[TablePrivilege] = NO_ACCESS, + runtime: frozenset[DeclaredPrivilege] = NO_ACCESS, + catalog_writer: frozenset[DeclaredPrivilege] = NO_ACCESS, + stage12: frozenset[DeclaredPrivilege] = NO_ACCESS, + simulation_service: frozenset[DeclaredPrivilege] = NO_ACCESS, + modal_worker: frozenset[DeclaredPrivilege] = NO_ACCESS, ) -> TableAccess: return TableAccess( runtime=runtime, @@ -123,14 +149,43 @@ def _access( ) -# The simulation API creates these rows; its Modal workers read them, and write -# the status of the runs and household jobs among them. +# The simulation API creates these rows and its Modal workers read them. _SIMULATION_DEFINITION = _access( runtime=READ, simulation_service=READ_INSERT, modal_worker=READ ) -_RUN_STATUS = _access( - runtime=READ, simulation_service=READ_INSERT, modal_worker=READ_UPDATE +# A run's state, which the Modal worker running it writes. +_RUN_STATE_COLUMNS: Final = ( + "status", + "modal_call_id", + "error_message", + "started_at", + "completed_at", + "updated_at", +) +# The state the simulation service writes for a run whose Modal call it starts: +# the call's ID once Modal accepts the call, or the failure when Modal refuses +# it. +_CALL_START_COLUMNS: Final = ( + "status", + "modal_call_id", + "error_message", + "completed_at", + "updated_at", ) + + +def _run_table( + *result_columns: str, catalog_writer: frozenset[DeclaredPrivilege] = NO_ACCESS +) -> TableAccess: + """Access to a run table whose Modal worker writes the named result columns.""" + return _access( + runtime=READ, + catalog_writer=catalog_writer, + simulation_service=READ_INSERT | {update_columns(*_CALL_START_COLUMNS)}, + modal_worker=READ | {update_columns(*_RUN_STATE_COLUMNS, *result_columns)}, + ) + + # The Modal workers write results, keyed by run, once each. _RESULT = _access(runtime=READ, modal_worker=READ_INSERT) # What a run needs from the Public API's policies, dynamics, households and @@ -149,10 +204,18 @@ def _access( # Simulations, reports, their runs and household jobs. The catalog writer # checks that publication leaves the rows of three of them untouched. "budget_window_runs": _SIMULATION_DEFINITION, - "household_jobs": _RUN_STATUS, - "report_runs": replace(_RUN_STATUS, catalog_writer=READ), + "household_jobs": _run_table("result"), + "report_runs": _run_table( + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", + "spm_provenance", + "data_version", + catalog_writer=READ, + ), "reports": replace(_SIMULATION_DEFINITION, catalog_writer=READ), - "simulation_runs": _RUN_STATUS, + "simulation_runs": _run_table("output_uri", "output_sha256", "household_result"), "simulations": replace(_SIMULATION_DEFINITION, catalog_writer=READ), # Results of simulation runs and report runs. "labor_supply_decile_responses": _RESULT, @@ -240,7 +303,7 @@ def stage12_role(deployment_environment: str) -> str: def declared_privileges( deployment_environment: str, -) -> dict[str, dict[str, frozenset[TablePrivilege]]]: +) -> dict[str, dict[str, frozenset[DeclaredPrivilege]]]: """Return each managed role's declared privileges, keyed by role then table.""" stage12 = stage12_role(deployment_environment) return { diff --git a/rules/registry.toml b/rules/registry.toml index 1258bde72..613e03313 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1018,7 +1018,7 @@ invocation = "v2-integration-check.yml" [[rule]] id = "api.v2.runtime-privileges-declared" category = "migration" -statement = "Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py, and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it." +statement = "Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py; a privilege declared on only some of a table's columns names columns of that table and is never also declared on the whole table; and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it, on the whole table." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/models/*.py", "policyengine_api/services/v2/*/database_connectors/*.py"] tier = "mechanical" @@ -1028,6 +1028,7 @@ autofix = false kind = "pytest" nodes = [ "tests/unit/v2/test_privilege_manifest.py::test_every_v2_table_has_an_entry_for_every_managed_role", + "tests/unit/v2/test_privilege_manifest.py::test_column_privileges_name_their_tables_columns_alone", "tests/unit/v2/test_privilege_manifest.py::test_v2_schema_has_no_sequences_the_manifest_would_need_to_cover", "tests/unit/v2/test_privilege_manifest.py::test_runtime_connectors_need_only_declared_runtime_privileges", "tests/unit/v2/test_privilege_manifest.py::test_stage12_connectors_need_only_declared_stage12_privileges", @@ -1064,7 +1065,7 @@ invocation = "reconcile_v2_privileges.py apply --create-missing-roles" [[rule]] id = "api.v2.deployed-privileges-match" category = "migration" -statement = "Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead." +statement = "Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares, on the whole table or on the columns it names, right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing or a role holds on a whole table a privilege declared for only some of its columns. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] tier = "mechanical" diff --git a/tests/integration/test_v2_privilege_reconciliation.py b/tests/integration/test_v2_privilege_reconciliation.py index 3edc1cdd1..97fdb3f34 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -22,10 +22,11 @@ PrivilegeEntry, apply_declared_privileges, check_declared_privileges, + check_table_wide_overrides, run, skipped_roles, ) -from policyengine_api.data.v2.privileges import READ, TablePrivilege +from policyengine_api.data.v2.privileges import READ, TablePrivilege, update_columns from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL if TYPE_CHECKING: @@ -74,6 +75,82 @@ def _has_privilege(engine: Engine, role: str, table: str, privilege: str) -> boo ) +def _has_column_privilege( + engine: Engine, role: str, table: str, column: str, privilege: str +) -> bool: + with engine.connect() as connection: + return bool( + connection.execute( + text("SELECT has_column_privilege(:role, :table, :column, :privilege)"), + { + "role": role, + "table": f"public.{table}", + "column": column, + "privilege": privilege, + }, + ).scalar_one() + ) + + +def test_apply_grants_declared_columns_alone_and_verify_checks_them( + superuser_engine, temporary_role +) -> None: + role = temporary_role() + declared = {role: {"regions": READ | {update_columns("label")}}} + with superuser_engine.begin() as connection: + first_plan, _defaults = apply_declared_privileges(connection, declared) + with superuser_engine.begin() as connection: + second_plan, _defaults = apply_declared_privileges(connection, declared) + missing = check_declared_privileges(connection, declared) + overrides = check_table_wide_overrides(connection, declared) + + assert first_plan.grants == ( + PrivilegeEntry(role, "regions", "SELECT"), + PrivilegeEntry(role, "regions", "UPDATE", column="label"), + ) + assert (second_plan.grants, second_plan.undeclared) == ((), ()) + assert (missing, overrides) == ([], []) + assert _has_column_privilege(superuser_engine, role, "regions", "label", "UPDATE") + assert not _has_column_privilege( + superuser_engine, role, "regions", "code", "UPDATE" + ) + assert not _has_privilege(superuser_engine, role, "regions", "UPDATE") + + +def test_verify_fails_when_a_role_can_update_the_whole_table( + disposable_v2_database_url, superuser_engine, temporary_role, monkeypatch, tmp_path +) -> None: + """``UPDATE`` on the whole table would let the role change every column.""" + role = temporary_role() + monkeypatch.setattr( + privilege_reconciliation, + "declared_privileges", + lambda _environment: {role: {"regions": READ | {update_columns("label")}}}, + ) + monkeypatch.setattr( + privilege_reconciliation, "ROLES_AWAITING_PROVISIONING", frozenset() + ) + summary_path = tmp_path / "summary.md" + environ = { + V2_MIGRATION_DATABASE_URL: disposable_v2_database_url, + V2_ALEMBIC_DISPOSABLE_TEST: "1", + DEPLOYMENT_ENVIRONMENT: "test", + "GITHUB_STEP_SUMMARY": str(summary_path), + } + + apply_status = run("apply", environ=environ) + with superuser_engine.begin() as connection: + connection.execute(text(f'GRANT UPDATE ON TABLE public.regions TO "{role}"')) + verify_status = run("verify", environ=environ) + summary = summary_path.read_text(encoding="utf-8") + + assert (apply_status, verify_status) == (0, 1) + assert ( + f"`{role}` holds UPDATE on the whole of `regions`, which overrides its " + "declared columns" + ) in summary + + def test_apply_grants_nothing_the_second_time_and_verify_passes( disposable_v2_database_url, managed_v2_roles, capsys ) -> None: diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index fba0221d7..a71dc7633 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -153,7 +153,7 @@ def _create_simulation_rows( return ids -def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( +def test_the_simulation_service_creates_rows_and_records_only_its_calls( connection: Connection, managed_v2_roles, public_api_rows: PublicApiRows ) -> None: _act_as(connection, managed_v2_roles.simulation_service) @@ -165,10 +165,35 @@ def test_the_simulation_service_creates_rows_but_writes_no_status_or_result( connection.execute(text("SELECT count(*) FROM policies")).scalar_one() assert run_count == 1 + # The service records a Modal call's ID once Modal accepts the call, and + # fails the run when Modal refuses it, while the run is still pending. + recorded = connection.execute( + text( + "UPDATE report_runs SET modal_call_id = 'fc-accepted', " + "updated_at = now() WHERE id = :id AND status = 'pending' " + "AND modal_call_id IS NULL RETURNING id" + ), + {"id": ids["report_run"]}, + ).scalar_one() + failed = connection.execute( + text( + "UPDATE household_jobs SET status = 'failed', error_message = 'refused', " + "completed_at = now(), updated_at = now() " + "WHERE id = :id AND status = 'pending' RETURNING id" + ), + {"id": ids["household_job"]}, + ).scalar_one() + assert (recorded, failed) == (ids["report_run"], ids["household_job"]) + for statement in ( + "UPDATE simulation_runs SET started_at = now() WHERE id = :id", + "UPDATE simulation_runs SET output_uri = 'gs://bucket/run' WHERE id = :id", + "UPDATE simulation_runs SET year = 2027 WHERE id = :id", + ): + _refused(connection, statement, id=ids["run"]) _refused( connection, - "UPDATE simulation_runs SET status = 'running' WHERE id = :id", - id=ids["run"], + "UPDATE report_runs SET baseline_cliff_gap = 1.0 WHERE id = :id", + id=ids["report_run"], ) _refused( connection, @@ -213,9 +238,12 @@ def test_modal_workers_write_their_run_status_and_results_only( ) connection.execute( text( - "UPDATE report_runs SET status = 'running', baseline_cliff_gap = 1.0, " + "UPDATE report_runs SET status = 'running', started_at = now(), " + "modal_call_id = 'fc-coordinator', baseline_cliff_gap = 1.0, " "baseline_cliff_share = 0.1, reform_cliff_gap = 2.0, " - "reform_cliff_share = 0.2 WHERE id = :id" + "reform_cliff_share = 0.2, data_version = '1.0.0', " + "spm_provenance = CAST('{}' AS jsonb), updated_at = now() " + "WHERE id = :id" ), {"id": ids["report_run"]}, ) @@ -266,6 +294,25 @@ def test_modal_workers_write_their_run_status_and_results_only( id=ids["run"], ) _refused(connection, "UPDATE households SET content_hash = content_hash") + # What a run is stays as the simulation service wrote it. + for statement, run_key in ( + ("UPDATE simulation_runs SET year = 2027 WHERE id = :id", "run"), + ( + "UPDATE simulation_runs SET simulation_id = simulation_id WHERE id = :id", + "run", + ), + ( + "UPDATE report_runs SET baseline_simulation_run_id = " + "baseline_simulation_run_id WHERE id = :id", + "report_run", + ), + ("UPDATE report_runs SET markdown = 'edited' WHERE id = :id", "report_run"), + ( + "UPDATE household_jobs SET household_id = household_id WHERE id = :id", + "household_job", + ), + ): + _refused(connection, statement, id=ids[run_key]) def test_the_public_api_reads_runs_and_results_but_writes_neither( diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 306bbac38..09f4420e5 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -22,16 +22,17 @@ NO_ACCESS, READ, READ_INSERT, - READ_UPDATE, READ_WRITE, ROLES_AWAITING_PROVISIONING, RUNTIME_ROLE, SIMULATION_SERVICE_ROLE, TABLE_ACCESS, + ColumnPrivilege, InvalidDeploymentEnvironmentError, TablePrivilege, declared_privileges, stage12_role, + update_columns, ) PROJECT_ROOT = Path(__file__).parents[3] @@ -51,6 +52,72 @@ } ) RUN_STATUS_TABLES = frozenset({"simulation_runs", "report_runs", "household_jobs"}) +# The run columns the Modal workers update: a run's state and its results. +WORKER_RUN_STATE_COLUMNS = frozenset( + { + "status", + "modal_call_id", + "error_message", + "started_at", + "completed_at", + "updated_at", + } +) +WORKER_RESULT_COLUMNS = { + "simulation_runs": frozenset({"output_uri", "output_sha256", "household_result"}), + "report_runs": frozenset( + { + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", + "spm_provenance", + "data_version", + } + ), + "household_jobs": frozenset({"result"}), +} +# The run columns the simulation service updates: it records a Modal call's ID +# when Modal accepts the call, and fails the run when Modal refuses it. +SERVICE_RUN_STATE_COLUMNS = frozenset( + {"status", "modal_call_id", "error_message", "completed_at", "updated_at"} +) +# What a run is: no role may change these after the simulation service writes +# them. +RUN_DEFINITION_COLUMNS = { + "simulation_runs": frozenset( + { + "simulation_id", + "year", + "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", + "dataset_id", + "region_id", + } + ), + "report_runs": frozenset( + { + "report_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "budget_window_run_id", + "year", + "country_package_version", + "policyengine_version", + } + ), + "household_jobs": frozenset( + { + "household_id", + "year", + "policy_id", + "dynamic_id", + "request_data", + "tax_benefit_model_version_id", + } + ), +} # The policy, dynamic, household and catalog rows a run needs. RUN_INPUT_TABLES = frozenset( { @@ -191,6 +258,8 @@ def test_simulation_service_creates_simulation_rows_and_reads_run_inputs() -> No expected = dict.fromkeys(TABLE_ACCESS, NO_ACCESS) expected |= dict.fromkeys(SIMULATION_SERVICE_CREATES, READ_INSERT) expected |= dict.fromkeys(RUN_INPUT_TABLES, READ) + for table in RUN_STATUS_TABLES: + expected[table] = READ_INSERT | {update_columns(*SERVICE_RUN_STATE_COLUMNS)} assert declared == expected @@ -199,12 +268,47 @@ def test_modal_workers_write_only_run_status_and_results() -> None: expected = dict.fromkeys(TABLE_ACCESS, NO_ACCESS) expected |= dict.fromkeys(SIMULATION_SERVICE_CREATES, READ) - expected |= dict.fromkeys(RUN_STATUS_TABLES, READ_UPDATE) expected |= dict.fromkeys(RESULT_TABLES, READ_INSERT) expected |= dict.fromkeys(RUN_INPUT_TABLES, READ) + for table in RUN_STATUS_TABLES: + worker_columns = WORKER_RUN_STATE_COLUMNS | WORKER_RESULT_COLUMNS[table] + expected[table] = READ | {update_columns(*worker_columns)} assert declared == expected +def test_column_privileges_name_their_tables_columns_alone() -> None: + """A whole-table privilege would override the same privilege's column list.""" + for role, privileges_by_table in declared_privileges(TEST_ENVIRONMENT).items(): + for table, privileges in privileges_by_table.items(): + column_privileges = [ + privilege + for privilege in privileges + if isinstance(privilege, ColumnPrivilege) + ] + kinds = [privilege.privilege for privilege in column_privileges] + assert len(kinds) == len(set(kinds)), (role, table) + for privilege in column_privileges: + assert privilege.columns, (role, table) + assert privilege.columns <= set(V2_METADATA.tables[table].c.keys()), ( + role, + table, + ) + assert privilege.privilege not in privileges, (role, table) + + +def test_no_role_can_change_what_a_run_is() -> None: + for role, privileges_by_table in declared_privileges(TEST_ENVIRONMENT).items(): + for table, definition_columns in RUN_DEFINITION_COLUMNS.items(): + privileges = privileges_by_table[table] + assert TablePrivilege.UPDATE not in privileges, (role, table) + for privilege in privileges: + if isinstance(privilege, ColumnPrivilege): + assert privilege.columns.isdisjoint(definition_columns), ( + role, + table, + ) + + def test_public_api_reads_simulations_runs_and_results_without_writing_them() -> None: declared = declared_privileges(TEST_ENVIRONMENT)[RUNTIME_ROLE] @@ -221,6 +325,7 @@ def test_no_new_role_writes_the_public_api_tables() -> None: table for table, privileges in declared_privileges(TEST_ENVIRONMENT)[role].items() if privileges & WRITE_PRIVILEGES + or any(isinstance(privilege, ColumnPrivilege) for privilege in privileges) } assert writes <= simulation_tables, role diff --git a/tests/unit/v2/test_privilege_reconciliation.py b/tests/unit/v2/test_privilege_reconciliation.py index 4c829f037..1977d4391 100644 --- a/tests/unit/v2/test_privilege_reconciliation.py +++ b/tests/unit/v2/test_privilege_reconciliation.py @@ -15,18 +15,25 @@ grant_statements, plan_reconciliation, render_summary, + table_wide_overrides, verification_failures, ) from policyengine_api.data.v2.privileges import ( READ, READ_WRITE, TablePrivilege, + update_columns, ) PROJECT_ROOT = Path(__file__).parents[3] PROOF_FIXTURES = PROJECT_ROOT / "rules/fixtures/api.v2.deployed-privileges-match" RUNTIME = "policyengine_v2_runtime" STAGE12 = "policyengine_stage12_staging" +WORKER = "policyengine_v2_modal_worker" +# The Modal worker may read a run and update two of its columns. +WORKER_DECLARED = { + WORKER: {"simulation_runs": READ | {update_columns("status", "completed_at")}} +} def _quote(identifier: str) -> str: @@ -265,6 +272,138 @@ def test_summary_lists_grants_extras_and_default_privileges_without_secrets() -> assert "postgresql" not in summary.lower().replace("postgresql's", "") +def test_missing_declared_column_privileges_become_column_grants() -> None: + plan = plan_reconciliation( + WORKER_DECLARED, + granted={WORKER: {"simulation_runs": {"SELECT"}}}, + existing_roles={WORKER}, + existing_tables={"simulation_runs"}, + granted_columns={WORKER: {"simulation_runs": {"status": {"UPDATE"}}}}, + ) + + assert plan.grants == ( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ) + assert plan.undeclared == () + + +def test_column_and_whole_table_privileges_beyond_the_manifest_are_reported() -> None: + plan = plan_reconciliation( + WORKER_DECLARED, + granted={WORKER: {"simulation_runs": {"SELECT", "UPDATE"}}}, + existing_roles={WORKER}, + existing_tables={"simulation_runs"}, + granted_columns={ + WORKER: { + "simulation_runs": { + "status": {"UPDATE"}, + "completed_at": {"UPDATE"}, + "simulation_id": {"UPDATE"}, + } + } + }, + ) + + assert plan.grants == () + assert plan.undeclared == ( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="simulation_id"), + ) + + +def test_grant_statements_grant_column_privileges_per_role_and_table() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "SELECT"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="status"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ), + undeclared=(), + missing_roles=(), + missing_tables=(), + ) + + assert grant_statements(plan, _quote) == [ + 'GRANT SELECT ON TABLE "public"."simulation_runs" ' + 'TO "policyengine_v2_modal_worker"', + 'GRANT UPDATE ("completed_at", "status") ON TABLE "public"."simulation_runs" ' + 'TO "policyengine_v2_modal_worker"', + ] + + +def test_grant_statements_refuse_an_unsafe_column_name() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column='status"; --'), + ), + undeclared=(), + missing_roles=(), + missing_tables=(), + ) + + with pytest.raises(privilege_reconciliation.UnsafeIdentifierError): + grant_statements(plan, _quote) + + +def test_verification_reports_each_declared_column_the_role_lacks() -> None: + failures = verification_failures( + WORKER_DECLARED, + lambda role, table, privilege: privilege == "SELECT", + has_column_privilege=lambda role, table, column, privilege: column == "status", + ) + + assert failures == [ + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at") + ] + + +def test_column_privileges_need_a_column_check() -> None: + with pytest.raises(ValueError, match="column"): + verification_failures(WORKER_DECLARED, lambda role, table, privilege: True) + + +def test_a_whole_table_privilege_overrides_its_declared_columns() -> None: + """``UPDATE`` on the whole table lets a role change every column.""" + + def has_privilege(role: str, table: str, privilege: str) -> bool: + return privilege in {"SELECT", "UPDATE"} + + assert table_wide_overrides(WORKER_DECLARED, has_privilege) == [ + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE") + ] + assert ( + table_wide_overrides( + WORKER_DECLARED, lambda role, table, privilege: privilege == "SELECT" + ) + == [] + ) + + +def test_summary_lists_column_grants_with_their_columns() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "SELECT"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="status"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ), + undeclared=( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="year"), + ), + missing_roles=(), + missing_tables=(), + ) + + summary = render_summary(plan, (), deployment_environment="staging", applied=True) + + assert ( + "| `policyengine_v2_modal_worker` | `simulation_runs` | " + "SELECT, UPDATE (completed_at, status) |" + ) in summary + assert ( + "| `policyengine_v2_modal_worker` | `simulation_runs` | UPDATE (year) |" + ) in summary + + def _fixture_privileges(name: str) -> set[tuple[str, str, str]]: entries = json.loads((PROOF_FIXTURES / name).read_text(encoding="utf-8")) return {(entry["role"], entry["table"], entry["privilege"]) for entry in entries} From b990f6cb7660f665b8d837836b3857e665077fd8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:14:01 +0400 Subject: [PATCH 18/22] Pass the v2 privilege plan the database's grants as one value plan_reconciliation takes GrantedPrivileges, the managed roles' grants on whole tables and on columns, instead of a separate argument for each, so it stays within ruff's argument limit without a suppression. Co-Authored-By: Claude Opus 5.5 --- .../data/v2/privilege_reconciliation.py | 38 +++++++----- .../unit/v2/test_privilege_reconciliation.py | 59 +++++++++++-------- 2 files changed, 57 insertions(+), 40 deletions(-) diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index faea4818f..96581f2e6 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -36,7 +36,7 @@ import re import sys from collections import defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING @@ -92,6 +92,20 @@ class PrivilegeEntry: column: str = "" +@dataclass(frozen=True) +class GrantedPrivileges: + """The managed roles' direct grants in one database. + + ``tables`` holds grants on whole tables, by role then table, and + ``columns`` grants on columns, by role, table, then column. + """ + + tables: Mapping[str, Mapping[str, Collection[str]]] + columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] = field( + default_factory=dict + ) + + @dataclass(frozen=True) class DefaultPrivilege: """A default privilege PostgreSQL gives on objects a role creates later.""" @@ -139,24 +153,18 @@ def _declared_entries( return entries -def plan_reconciliation( # noqa: PLR0913 # the database facts it compares +def plan_reconciliation( declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], - granted: Mapping[str, Mapping[str, Collection[str]]], + granted: GrantedPrivileges, *, existing_roles: Collection[str], existing_tables: Collection[str], roles_awaiting_creation: Collection[str] = (), - granted_columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] - | None = None, ) -> ReconciliationPlan: """Compare declared privileges with what the database grants. - ``granted`` holds the managed roles' direct grants on whole tables, by role - then table, and ``granted_columns`` their direct grants on columns, by role, - table, then column. A missing role in ``roles_awaiting_creation`` is skipped - instead of missing. + A missing role in ``roles_awaiting_creation`` is skipped instead of missing. """ - granted_columns = granted_columns or {} absent_roles = set(declared) - set(existing_roles) skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) @@ -170,12 +178,12 @@ def plan_reconciliation( # noqa: PLR0913 # the database facts it compares } granted_entries = { PrivilegeEntry(role, table, privilege) - for role, tables in granted.items() + for role, tables in granted.tables.items() for table, privileges in tables.items() for privilege in privileges } | { PrivilegeEntry(role, table, privilege, column) - for role, tables in granted_columns.items() + for role, tables in granted.columns.items() for table, columns in tables.items() for column, privileges in columns.items() for privilege in privileges @@ -543,11 +551,13 @@ def apply_declared_privileges( roles = list(declared) plan = plan_reconciliation( declared, - read_granted_privileges(connection, roles), + GrantedPrivileges( + tables=read_granted_privileges(connection, roles), + columns=read_granted_column_privileges(connection, roles), + ), existing_roles=read_existing_roles(connection, roles), existing_tables=read_public_tables(connection), roles_awaiting_creation=roles_awaiting_creation, - granted_columns=read_granted_column_privileges(connection, roles), ) if plan.can_apply: for statement in grant_statements( diff --git a/tests/unit/v2/test_privilege_reconciliation.py b/tests/unit/v2/test_privilege_reconciliation.py index 1977d4391..208a5a1e5 100644 --- a/tests/unit/v2/test_privilege_reconciliation.py +++ b/tests/unit/v2/test_privilege_reconciliation.py @@ -10,6 +10,7 @@ from policyengine_api.data.v2 import privilege_reconciliation from policyengine_api.data.v2.privilege_reconciliation import ( DefaultPrivilege, + GrantedPrivileges, PrivilegeEntry, ReconciliationPlan, grant_statements, @@ -58,7 +59,7 @@ def _declared() -> dict[str, dict[str, frozenset[TablePrivilege]]]: def test_missing_declared_privileges_become_grants() -> None: plan = plan_reconciliation( _declared(), - granted={RUNTIME: {"users": {"SELECT", "INSERT"}}}, + granted=GrantedPrivileges({RUNTIME: {"users": {"SELECT", "INSERT"}}}), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports"}, ) @@ -76,14 +77,16 @@ def test_missing_declared_privileges_become_grants() -> None: def test_privileges_beyond_the_manifest_are_reported_not_granted_or_revoked() -> None: plan = plan_reconciliation( _declared(), - granted={ - RUNTIME: { - "users": {"SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE"}, - "legacy_policy_mappings": {"SELECT", "INSERT"}, - "reports": {"SELECT"}, - }, - STAGE12: {"users": {"SELECT"}}, - }, + granted=GrantedPrivileges( + { + RUNTIME: { + "users": {"SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE"}, + "legacy_policy_mappings": {"SELECT", "INSERT"}, + "reports": {"SELECT"}, + }, + STAGE12: {"users": {"SELECT"}}, + } + ), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports"}, ) @@ -101,7 +104,7 @@ def test_privileges_beyond_the_manifest_are_reported_not_granted_or_revoked() -> def test_privileges_on_tables_outside_the_manifest_are_reported() -> None: plan = plan_reconciliation( _declared(), - granted={RUNTIME: {"scratch_table": {"SELECT"}}}, + granted=GrantedPrivileges({RUNTIME: {"scratch_table": {"SELECT"}}}), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports", "scratch_table"}, ) @@ -112,7 +115,7 @@ def test_privileges_on_tables_outside_the_manifest_are_reported() -> None: def test_missing_roles_and_tables_are_recorded_and_block_the_plan() -> None: plan = plan_reconciliation( _declared(), - granted={}, + granted=GrantedPrivileges({}), existing_roles={RUNTIME}, existing_tables={"users", "reports"}, ) @@ -130,7 +133,7 @@ def test_missing_roles_and_tables_are_recorded_and_block_the_plan() -> None: def test_a_missing_role_awaiting_creation_is_skipped_without_blocking() -> None: plan = plan_reconciliation( _declared(), - granted={}, + granted=GrantedPrivileges({}), existing_roles={RUNTIME}, existing_tables={"users", "legacy_policy_mappings", "reports"}, roles_awaiting_creation={STAGE12}, @@ -147,7 +150,7 @@ def test_a_role_awaiting_creation_gets_its_privileges_once_it_exists() -> None: plan = plan_reconciliation( declared, - granted={}, + granted=GrantedPrivileges({}), existing_roles={RUNTIME, STAGE12}, existing_tables={"users"}, roles_awaiting_creation={STAGE12}, @@ -163,7 +166,7 @@ def test_a_role_awaiting_creation_gets_its_privileges_once_it_exists() -> None: def test_a_missing_role_not_awaiting_creation_still_blocks_the_plan() -> None: plan = plan_reconciliation( _declared(), - granted={}, + granted=GrantedPrivileges({}), existing_roles={RUNTIME}, existing_tables={"users", "legacy_policy_mappings", "reports"}, roles_awaiting_creation={"policyengine_v2_some_other_role"}, @@ -275,10 +278,12 @@ def test_summary_lists_grants_extras_and_default_privileges_without_secrets() -> def test_missing_declared_column_privileges_become_column_grants() -> None: plan = plan_reconciliation( WORKER_DECLARED, - granted={WORKER: {"simulation_runs": {"SELECT"}}}, + granted=GrantedPrivileges( + tables={WORKER: {"simulation_runs": {"SELECT"}}}, + columns={WORKER: {"simulation_runs": {"status": {"UPDATE"}}}}, + ), existing_roles={WORKER}, existing_tables={"simulation_runs"}, - granted_columns={WORKER: {"simulation_runs": {"status": {"UPDATE"}}}}, ) assert plan.grants == ( @@ -290,18 +295,20 @@ def test_missing_declared_column_privileges_become_column_grants() -> None: def test_column_and_whole_table_privileges_beyond_the_manifest_are_reported() -> None: plan = plan_reconciliation( WORKER_DECLARED, - granted={WORKER: {"simulation_runs": {"SELECT", "UPDATE"}}}, + granted=GrantedPrivileges( + tables={WORKER: {"simulation_runs": {"SELECT", "UPDATE"}}}, + columns={ + WORKER: { + "simulation_runs": { + "status": {"UPDATE"}, + "completed_at": {"UPDATE"}, + "simulation_id": {"UPDATE"}, + } + } + }, + ), existing_roles={WORKER}, existing_tables={"simulation_runs"}, - granted_columns={ - WORKER: { - "simulation_runs": { - "status": {"UPDATE"}, - "completed_at": {"UPDATE"}, - "simulation_id": {"UPDATE"}, - } - } - }, ) assert plan.grants == () From 357719fad5816db0c45c8ed9b6981bdc9d0f39f0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:33:34 +0400 Subject: [PATCH 19/22] Fix v2 tests that still asserted the old schema From the review round and a search for stale tests: - the decile unique key no longer includes the run pair; - a budget-window run's year runs are deleted with it (CASCADE) instead of blocking its deletion, so every foreign key names a delete rule, as test_every_table_has_named_primary_foreign_and_relational_constraints requires; deleting a report still removes both; - the run-definition columns no role may update include a report run's trigger and idempotency key and a household job's country and versions; - the two revision lifecycle tests are one downgrade-and-upgrade test, and the module docstring describes the final schema. Revision 1356f658347d is regenerated. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 12 +- ...record_simulation_runs_and_key_results_.py | 259 +++++++++--------- policyengine_api/data/v2/models/reports.py | 6 +- .../test_alembic_v2_simulation_runs.py | 61 ++--- tests/unit/v2/test_privilege_manifest.py | 5 + tests/unit/v2/test_simulation_run_models.py | 9 +- 6 files changed, 179 insertions(+), 173 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d1737580e..1fe4951a3 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -553,17 +553,19 @@ model changes: calculate their simulations again (`recalculate`), status, error and times. `report_runs.budget_window_run_id` and `year` name the window run and year a report run calculates, at most one report run per year of a window run; a - composite foreign key keeps it to a window run of the same report. That key - is `NO ACTION`, which PostgreSQL checks at the end of the statement, so - deleting a report deletes its window runs and their report runs together. + composite foreign key keeps it to a window run of the same report, and + deletes it with the window run (`CASCADE`), as deleting a report deletes its + runs. - `household_jobs` calculates one household (`household_id`) for one `year` under one policy, and records its catalog and package versions and its Modal call ID. - The new columns that name the Public API's rows are foreign keys with the delete rules the same columns have on `simulations`: `simulation_runs.tax_benefit_model_version_id` and `dataset_id` (`RESTRICT`) - and `region_id` (`SET NULL`), and `household_jobs.household_id` and - `tax_benefit_model_version_id` (`RESTRICT`). `reports.user_id` gains an + and `region_id` (`SET NULL`), `household_jobs.household_id` and + `tax_benefit_model_version_id` (`RESTRICT`), and + `budget_window_runs.tax_benefit_model_version_id` (`RESTRICT`), 20 keys + between the two services' tables in all. `reports.user_id` gains an index, which account deletion's `SET NULL` searches. Each service writes only its own tables, and the foreign keys between them stay: PostgreSQL checks each named row when a row is written, with the privileges of the named diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 48021b291..69a386c3b 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -2,7 +2,7 @@ Revision ID: 1356f658347d Revises: c614b6bcccfc -Create Date: 2026-11-06 00:03:10.204134 +Create Date: 2026-11-06 00:33:17.526471 Generation: uv run alembic -c alembic-v2.ini revision --autogenerate Three post-generation corrections of one kind, each commented where it is @@ -309,33 +309,33 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) - op.drop_column("budget_summary", "baseline_simulation_id") op.drop_column("budget_summary", "reform_simulation_id") + op.drop_column("budget_summary", "baseline_simulation_id") op.add_column( "congressional_district_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -357,37 +357,37 @@ def upgrade() -> None: sa.Column("no_change_percentage", sa.Float(), nullable=False), ) op.drop_constraint( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", type_="foreignkey", ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) - op.drop_column("congressional_district_impacts", "baseline_simulation_id") op.drop_column("congressional_district_impacts", "reform_simulation_id") + op.drop_column("congressional_district_impacts", "baseline_simulation_id") op.add_column( "constituency_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -421,28 +421,28 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), "constituency_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), "constituency_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) @@ -451,8 +451,8 @@ def upgrade() -> None: "constituency_impacts", "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", ) - op.drop_column("constituency_impacts", "baseline_simulation_id") op.drop_column("constituency_impacts", "reform_simulation_id") + op.drop_column("constituency_impacts", "baseline_simulation_id") op.add_column( "decile_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -511,8 +511,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("decile_impacts", "baseline_simulation_id") op.drop_column("decile_impacts", "reform_simulation_id") + op.drop_column("decile_impacts", "baseline_simulation_id") op.add_column( "household_jobs", sa.Column( @@ -565,20 +565,20 @@ def upgrade() -> None: unique=False, ) op.create_foreign_key( - op.f( - "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" - ), + op.f("fk_household_jobs_household_id_households"), "household_jobs", - "tax_benefit_model_versions", - ["tax_benefit_model_version_id"], + "households", + ["household_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_household_jobs_household_id_households"), + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", - "households", - ["household_id"], + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], ["id"], ondelete="RESTRICT", ) @@ -602,12 +602,12 @@ def upgrade() -> None: ["simulation_run_id", "income_variable", "entity"], ) op.drop_constraint( - op.f("fk_inequality_simulation_id_simulations"), + op.f("fk_inequality_report_run_id_report_runs"), "inequality", type_="foreignkey", ) op.drop_constraint( - op.f("fk_inequality_report_run_id_report_runs"), + op.f("fk_inequality_simulation_id_simulations"), "inequality", type_="foreignkey", ) @@ -630,33 +630,33 @@ def upgrade() -> None: sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=False), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["baseline_simulation_run_id"], + ["reform_simulation_run_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", "simulation_runs", - ["reform_simulation_run_id"], + ["baseline_simulation_run_id"], ["id"], ondelete="RESTRICT", ) - op.drop_column("intra_decile_impacts", "baseline_simulation_id") op.drop_column("intra_decile_impacts", "reform_simulation_id") + op.drop_column("intra_decile_impacts", "baseline_simulation_id") op.add_column( "local_authority_impacts", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -690,12 +690,12 @@ def upgrade() -> None: sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), ) op.drop_constraint( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", type_="foreignkey", ) @@ -720,8 +720,8 @@ def upgrade() -> None: "local_authority_impacts", "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", ) - op.drop_column("local_authority_impacts", "baseline_simulation_id") op.drop_column("local_authority_impacts", "reform_simulation_id") + op.drop_column("local_authority_impacts", "baseline_simulation_id") op.add_column("poverty", sa.Column("simulation_run_id", sa.Uuid(), nullable=False)) op.add_column( "poverty", @@ -787,12 +787,12 @@ def upgrade() -> None: postgresql_nulls_not_distinct=True, ) op.drop_constraint( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", type_="foreignkey", ) op.drop_constraint( - op.f("fk_program_statistics_reform_simulation_id_simulations"), + op.f("fk_program_statistics_baseline_simulation_id_simulations"), "program_statistics", type_="foreignkey", ) @@ -812,8 +812,8 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) - op.drop_column("program_statistics", "baseline_simulation_id") op.drop_column("program_statistics", "reform_simulation_id") + op.drop_column("program_statistics", "baseline_simulation_id") op.add_column( "report_runs", sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), @@ -882,20 +882,13 @@ def upgrade() -> None: "report_runs", ["budget_window_run_id", "year"], ) - op.create_foreign_key( - op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), - "report_runs", - "simulation_runs", - ["baseline_simulation_run_id"], - ["id"], - ondelete="RESTRICT", - ) op.create_foreign_key( "fk_report_runs_budget_window_run_id_budget_window_runs", "report_runs", "budget_window_runs", ["budget_window_run_id", "report_id"], ["id", "report_id"], + ondelete="CASCADE", ) op.create_foreign_key( op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), @@ -905,6 +898,14 @@ def upgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "report_runs", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) op.create_check_constraint( op.f("ck_report_runs_budget_window_year"), "report_runs", @@ -935,12 +936,12 @@ def upgrade() -> None: "simulations", type_="foreignkey", ) + op.drop_column("simulations", "household_result") op.drop_column("simulations", "started_at") - op.drop_column("simulations", "error_message") + op.drop_column("simulations", "output_dataset_id") op.drop_column("simulations", "status") - op.drop_column("simulations", "household_result") op.drop_column("simulations", "completed_at") - op.drop_column("simulations", "output_dataset_id") + op.drop_column("simulations", "error_message") # ### end Alembic commands ### @@ -948,7 +949,7 @@ def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### op.add_column( "simulations", - sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), ) op.add_column( "simulations", @@ -959,15 +960,6 @@ def downgrade() -> None: nullable=True, ), ) - op.add_column( - "simulations", - sa.Column( - "household_result", - postgresql.JSON(astext_type=sa.Text()), - autoincrement=False, - nullable=True, - ), - ) op.add_column( "simulations", sa.Column( @@ -981,7 +973,7 @@ def downgrade() -> None: ) op.add_column( "simulations", - sa.Column("error_message", sa.VARCHAR(), autoincrement=False, nullable=True), + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), ) op.add_column( "simulations", @@ -992,6 +984,15 @@ def downgrade() -> None: nullable=True, ), ) + op.add_column( + "simulations", + sa.Column( + "household_result", + postgresql.JSON(astext_type=sa.Text()), + autoincrement=False, + nullable=True, + ), + ) op.create_foreign_key( op.f("fk_simulations_output_dataset_id_datasets"), "simulations", @@ -1017,17 +1018,17 @@ def downgrade() -> None: op.f("ck_report_runs_budget_window_year"), "report_runs", type_="check" ) op.drop_constraint( - op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) op.drop_constraint( - "fk_report_runs_budget_window_run_id_budget_window_runs", + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), "report_runs", type_="foreignkey", ) op.drop_constraint( - op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "fk_report_runs_budget_window_run_id_budget_window_runs", "report_runs", type_="foreignkey", ) @@ -1055,13 +1056,13 @@ def downgrade() -> None: op.add_column( "program_statistics", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "program_statistics", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1075,18 +1076,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_program_statistics_reform_simulation_id_simulations"), + op.f("fk_program_statistics_baseline_simulation_id_simulations"), "program_statistics", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_program_statistics_baseline_simulation_id_simulations"), + op.f("fk_program_statistics_reform_simulation_id_simulations"), "program_statistics", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1151,13 +1152,13 @@ def downgrade() -> None: op.add_column( "local_authority_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "local_authority_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1176,18 +1177,18 @@ def downgrade() -> None: type_="foreignkey", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), "local_authority_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1201,38 +1202,38 @@ def downgrade() -> None: op.add_column( "intra_decile_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "intra_decile_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( - op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_intra_decile_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), "intra_decile_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), "intra_decile_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1251,14 +1252,6 @@ def downgrade() -> None: "inequality", type_="foreignkey", ) - op.create_foreign_key( - op.f("fk_inequality_report_run_id_report_runs"), - "inequality", - "report_runs", - ["report_run_id"], - ["id"], - ondelete="CASCADE", - ) op.create_foreign_key( op.f("fk_inequality_simulation_id_simulations"), "inequality", @@ -1267,6 +1260,14 @@ def downgrade() -> None: ["id"], ondelete="RESTRICT", ) + op.create_foreign_key( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) op.drop_constraint( "uq_inequality_simulation_run_measure", "inequality", type_="unique" ) @@ -1282,14 +1283,14 @@ def downgrade() -> None: op.f("ck_household_jobs_terminal_completion"), "household_jobs", type_="check" ) op.drop_constraint( - op.f("fk_household_jobs_household_id_households"), + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), "household_jobs", type_="foreignkey", ) op.drop_constraint( - op.f( - "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" - ), + op.f("fk_household_jobs_household_id_households"), "household_jobs", type_="foreignkey", ) @@ -1308,13 +1309,13 @@ def downgrade() -> None: op.add_column( "decile_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "decile_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1358,13 +1359,13 @@ def downgrade() -> None: op.add_column( "constituency_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "constituency_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( @@ -1373,28 +1374,28 @@ def downgrade() -> None: type_="check", ) op.drop_constraint( - op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), "constituency_impacts", type_="foreignkey", ) op.drop_constraint( - op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + op.f("fk_constituency_impacts_baseline_simulation_run_id_simulation_runs"), "constituency_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), "constituency_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_constituency_impacts_baseline_simulation_id_simulations"), + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), "constituency_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1408,42 +1409,42 @@ def downgrade() -> None: op.add_column( "congressional_district_impacts", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "congressional_district_impacts", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", ) op.drop_constraint( op.f( - "fk_congressional_district_impacts_reform_simulation_run_id_simulation_runs" + "fk_congressional_district_impacts_baseline_simulation_run_id_simulation_runs" ), "congressional_district_impacts", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), "congressional_district_impacts", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), "congressional_district_impacts", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) @@ -1455,38 +1456,38 @@ def downgrade() -> None: op.add_column( "budget_summary", sa.Column( - "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.add_column( "budget_summary", sa.Column( - "baseline_simulation_id", sa.UUID(), autoincrement=False, nullable=False + "reform_simulation_id", sa.UUID(), autoincrement=False, nullable=False ), ) op.drop_constraint( - op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.drop_constraint( - op.f("fk_budget_summary_reform_simulation_run_id_simulation_runs"), + op.f("fk_budget_summary_baseline_simulation_run_id_simulation_runs"), "budget_summary", type_="foreignkey", ) op.create_foreign_key( - op.f("fk_budget_summary_baseline_simulation_id_simulations"), + op.f("fk_budget_summary_reform_simulation_id_simulations"), "budget_summary", "simulations", - ["baseline_simulation_id"], + ["reform_simulation_id"], ["id"], ondelete="RESTRICT", ) op.create_foreign_key( - op.f("fk_budget_summary_reform_simulation_id_simulations"), + op.f("fk_budget_summary_baseline_simulation_id_simulations"), "budget_summary", "simulations", - ["reform_simulation_id"], + ["baseline_simulation_id"], ["id"], ondelete="RESTRICT", ) diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 215bc4e4a..152b867b4 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -264,13 +264,13 @@ class ReportRun(TimestampedModel, table=True): "year", name="uq_report_runs_budget_window_run_id_year", ), - # A year's report run belongs to a budget-window run of its own report. - # NO ACTION waits for the end of the statement, so deleting the report - # deletes the window run and its years together. + # A year's report run belongs to a budget-window run of its own report, + # which owns it as the report owns its runs. sa.ForeignKeyConstraint( ["budget_window_run_id", "report_id"], ["budget_window_runs.id", "budget_window_runs.report_id"], name="fk_report_runs_budget_window_run_id_budget_window_runs", + ondelete="CASCADE", ), sa.CheckConstraint( "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) " diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py index 7b34b1154..7657f16b9 100644 --- a/tests/integration/test_alembic_v2_simulation_runs.py +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -1,14 +1,16 @@ """Check the simulation runs revision against disposable Postgres. One revision follows the parameter search file revision. It adds simulation -runs and the table naming the simulation runs a report run combines, keys -every result table by run, gives household jobs one household, one year and -one policy, drops the foreign keys from those tables to the Public API's -tables, adds the economy output's fields that had no column, and keeps the -parts of a comparison without a result table in one row per compared pair. It -changes no other table. The previous Public API revision keeps working while -it is deployed: its roles can't read or write any table it changes, except the -catalog writer, which counts the rows of simulations, reports and report_runs. +runs, one per simulation and year; has each report run name its baseline and +reform simulation runs; keys every result table by run; types a report run's +cliff impact and labour supply; adds the budget-window tables; gives household +jobs one household, one year and one policy; and adds the economy output's +fields that had no column. The foreign keys from those tables to the Public +API's tables stay, and the new columns that name the Public API's rows get +keys too. It changes no other table. The previous Public API revision keeps +working while it is deployed: its roles can't read or write any table it +changes, except the catalog writer, which counts the rows of simulations, +reports and report_runs. """ from __future__ import annotations @@ -213,29 +215,15 @@ def _enum_type_names(engine: Engine) -> set[str]: ) -def test_the_revision_adds_its_tables_and_changes_only_simulation_side_tables( - config: Config, engine: Engine -) -> None: - command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) - before = _shapes(engine) - command.upgrade(config, "head") - after = _shapes(engine) - - assert set(after) - set(before) == NEW_TABLES - assert set(before) <= set(after) - changed = { - table_name for table_name in before if before[table_name] != after[table_name] - } - assert changed == CHANGED_TABLES - - def test_the_revision_downgrades_to_its_parent_and_upgrades_again( config: Config, engine: Engine ) -> None: + """The revision adds its tables and changes only simulation-side tables.""" command.upgrade(config, "head") command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) - assert NEW_TABLES.isdisjoint(inspect(engine).get_table_names(schema="public")) + before = _shapes(engine) + assert NEW_TABLES.isdisjoint(before) assert {"status", "household_result", "output_dataset_id"} <= _column_names( engine, "simulations" ) @@ -253,6 +241,13 @@ def test_the_revision_downgrades_to_its_parent_and_upgrades_again( command.upgrade(config, "head") command.check(config) + after = _shapes(engine) + assert set(after) - set(before) == NEW_TABLES + assert set(before) <= set(after) + changed = { + table_name for table_name in before if before[table_name] != after[table_name] + } + assert changed == CHANGED_TABLES assert {"baseline_simulation_run_id", "reform_simulation_run_id"} <= ( _column_names(engine, "report_runs") ) @@ -746,14 +741,18 @@ def insert_year(year: int | None, *, of_report: object = report_id) -> None: insert_year(2028, of_report=other_report_id) with refused(connection, match="ck_simulation_runs_year"): _insert_simulation_run(connection, simulation_id, year=1800) - with refused(connection, match="foreign key constraint"): - connection.execute( - text("DELETE FROM budget_window_runs WHERE id = :id"), - {"id": window_run_id}, - ) + # A window run owns its years, as the report owns its runs. + second_window_run_id = _insert_budget_window_run( + connection, public_api_rows, report_id + ) + connection.execute( + text("DELETE FROM budget_window_runs WHERE id = :id"), + {"id": window_run_id}, + ) + assert _count(connection, "report_runs", "budget_window_run_id", window_run_id) == 0 connection.execute(text("DELETE FROM reports WHERE id = :id"), {"id": report_id}) - assert _count(connection, "budget_window_runs", "report_id", report_id) == 0 + assert _count(connection, "budget_window_runs", "id", second_window_run_id) == 0 assert _count(connection, "report_runs", "report_id", report_id) == 0 diff --git a/tests/unit/v2/test_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 09f4420e5..f66b69d79 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -105,16 +105,21 @@ "year", "country_package_version", "policyengine_version", + "trigger", + "idempotency_key", } ), "household_jobs": frozenset( { + "country", "household_id", "year", "policy_id", "dynamic_id", "request_data", "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", } ), } diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 67c3d7fda..281fbb36b 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -441,8 +441,8 @@ def test_report_runs_hold_one_year_of_a_budget_window_run() -> None: for constraint in report_runs.foreign_key_constraints if constraint.referred_table.name == "budget_window_runs" ) - # A year's report run belongs to a window run of the same report. The check - # waits for the end of the statement, so deleting the report deletes both. + # A year's report run belongs to a window run of the same report, which + # owns it as the report owns its runs. assert [column.name for column in window_key.columns] == [ "budget_window_run_id", "report_id", @@ -451,7 +451,7 @@ def test_report_runs_hold_one_year_of_a_budget_window_run() -> None: "budget_window_runs.id", "budget_window_runs.report_id", ] - assert window_key.ondelete is None + assert window_key.ondelete == "CASCADE" assert ("budget_window_run_id", "year") in _unique_column_sets(report_runs) assert {"ck_report_runs_year", "ck_report_runs_budget_window_year"} <= ( _check_names(report_runs) @@ -603,6 +603,5 @@ def test_decile_rows_name_the_variable_that_groups_them() -> None: statement = str(CreateTable(decile_impacts).compile(dialect=postgresql.dialect())) assert ( "CONSTRAINT uq_decile_impacts_run_measure UNIQUE NULLS NOT DISTINCT " - "(report_run_id, baseline_simulation_run_id, reform_simulation_run_id, " - "income_variable, decile_variable, entity, decile, quantiles)" + "(report_run_id, income_variable, decile_variable, entity, decile, quantiles)" ) in statement From 56212cc0067d09b8f338a53e5b4c7bfa862d513a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:36:58 +0400 Subject: [PATCH 20/22] Split the v2 reports and privilege modules under the file-length limit The models' reports module had 897 lines and privilege_reconciliation 810, over api.quality.file-length's 800. Code moves unchanged: - The result tables and their enums move from models.reports to the new models.run_results; reports keeps reports, budget-window runs and report runs. Both are imported from the models package as before. The package's new import line carries the same E402 suppression as its neighbours, with its reason, because the tables must be declared after the naming convention. Import Linter's model exceptions and ruff's runtime base classes name the new module. - The planning and reporting half of privilege_reconciliation moves to the new privilege_plans, which runs no SQL; privilege_reconciliation keeps the database access and the command line and re-exports the moved names. _safe, _skipped_roles_line and _verification_report, now used from a second module, become safe_identifier, skipped_roles_line and verification_report. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/data/v2/models/__init__.py | 12 +- policyengine_api/data/v2/models/reports.py | 563 +----------------- .../data/v2/models/run_results.py | 554 +++++++++++++++++ .../data/v2/models/simulations.py | 7 +- policyengine_api/data/v2/privilege_plans.py | 373 ++++++++++++ .../data/v2/privilege_reconciliation.py | 411 ++----------- pyproject.toml | 5 +- rules/registry.toml | 2 + tests/unit/v2/test_simulation_run_models.py | 2 +- 9 files changed, 1020 insertions(+), 909 deletions(-) create mode 100644 policyengine_api/data/v2/models/run_results.py create mode 100644 policyengine_api/data/v2/privilege_plans.py diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index 4ae843094..5cd21269d 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -66,10 +66,16 @@ LegacyUserPolicyMapping, ) from policyengine_api.data.v2.models.reports import ( # noqa: E402 + BudgetWindowRun, + Report, + ReportRun, + ReportRunStatus, + ReportRunTrigger, +) +from policyengine_api.data.v2.models.run_results import ( # noqa: E402 # tables declared after the naming convention above AggregateOutput, AggregateType, BudgetSummary, - BudgetWindowRun, ChangeAggregate, CongressionalDistrictImpact, ConstituencyImpact, @@ -83,10 +89,6 @@ OutputStatus, Poverty, ProgramStatistics, - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, ) from policyengine_api.data.v2.models.simulations import ( # noqa: E402 Simulation, diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 152b867b4..11f820ab4 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,14 +1,13 @@ -"""Canonical SQLModel tables for stable reports, runs, and run outputs. +"""Canonical SQLModel tables for stable reports and their runs. A report run compares one baseline simulation run with one reform simulation -run. Results that compare them belong to the report run and name the two -simulation runs; each such row is unique within its report run. Poverty and -inequality, which describe one simulation, belong to a simulation run. +run; a budget-window report has a run per year of each calculation. Their +results are in ``run_results``. """ from datetime import datetime from enum import StrEnum -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import UUID import sqlalchemy as sa @@ -16,7 +15,6 @@ from sqlmodel import Field, Relationship from policyengine_api.data.v2.models.base import ( - IdentifiedModel, TimestampedModel, enum_type, ) @@ -27,9 +25,24 @@ TaxBenefitModel, ) from policyengine_api.data.v2.models.policies import Policy -from policyengine_api.data.v2.models.simulations import Simulation, SimulationRun +from policyengine_api.data.v2.models.simulations import Simulation from policyengine_api.data.v2.models.users import User +if TYPE_CHECKING: + from policyengine_api.data.v2.models.run_results import ( + AggregateOutput, + BudgetSummary, + ChangeAggregate, + CongressionalDistrictImpact, + ConstituencyImpact, + DecileImpact, + IntraDecileImpact, + LaborSupplyDecileResponse, + LaborSupplyResponse, + LocalAuthorityImpact, + ProgramStatistics, + ) + # A budget window compares at most this many years, none after LAST_YEAR. MAX_BUDGET_WINDOW_YEARS = 75 LAST_YEAR = 2200 @@ -52,30 +65,6 @@ class ReportRunTrigger(StrEnum): 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. @@ -383,515 +372,3 @@ class ReportRun(TimestampedModel, table=True): back_populates="report_run", cascade_delete=True, ) - - -class LaborSupplyResponse(IdentifiedModel, table=True): - """The labour supply response to the reform in one report run. - - One row per report run. The earnings changes come from the income and - substitution effects (.py's ``income_lsr`` and ``substitution_lsr``) and - total ``total_earnings_change``; ``relative_income_effect`` and - ``relative_substitution_effect`` are relative to baseline earnings. The - weekly hours are the total hours worked under the baseline and the reform, - and their change split into the two effects. - """ - - __tablename__ = "labor_supply_responses" - __table_args__ = ( - sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - income_effect_earnings_change: float - substitution_effect_earnings_change: float - total_earnings_change: float - revenue_change: float - relative_income_effect: float - relative_substitution_effect: float - baseline_weekly_hours: float - reform_weekly_hours: float - weekly_hours_change: float - weekly_hours_income_effect: float - weekly_hours_substitution_effect: float - - report_run: ReportRun = Relationship(back_populates="labor_supply_responses") - - -class LaborSupplyDecileResponse(IdentifiedModel, table=True): - """The labour supply response in one income decile of one report run. - - One row per report run and decile, as .py numbers them. The average effects - are per household; the relative effects are relative to the decile's - baseline earnings, which .py reports only for deciles above zero, so they - may be null. - """ - - __tablename__ = "labor_supply_decile_responses" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "decile", - name="uq_labor_supply_decile_responses_run_decile", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - decile: int - average_income_effect: float | None = None - average_substitution_effect: float | None = None - relative_income_effect: float | None = None - relative_substitution_effect: float | None = None - - report_run: ReportRun = Relationship(back_populates="labor_supply_decile_responses") - - -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 report 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_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.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 report run, - income variable, decile variable, entity, quantile count and decile. - ``decile_variable`` names a variable that groups the households instead of - the income variable's own deciles, such as ``household_wealth_decile``; rows - without one count as equal. - """ - - __tablename__ = "decile_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "income_variable", - "decile_variable", - "entity", - "decile", - "quantiles", - name="uq_decile_impacts_run_measure", - postgresql_nulls_not_distinct=True, - ), - 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_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - income_variable: str = Field(max_length=512) - decile_variable: str | None = Field(default=None, 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 report 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_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.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 run.""" - - __tablename__ = "inequality" - __table_args__ = ( - sa.UniqueConstraint( - "simulation_run_id", - "income_variable", - "entity", - name="uq_inequality_simulation_run_measure", - ), - ) - - simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="CASCADE", - ) - 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 - - simulation_run: SimulationRun = Relationship(back_populates="inequalities") - - -class Poverty(IdentifiedModel, table=True): - """Poverty headcount and rate of one type for one simulation run. - - One row per run, poverty type, entity, filter variable and filter group: - ``filter_variable`` selects people (``age``) and ``filter_group`` names the - group (``child``). Missing values count as equal, so each run has one - unfiltered row per type and entity. - """ - - __tablename__ = "poverty" - __table_args__ = ( - sa.UniqueConstraint( - "simulation_run_id", - "poverty_type", - "entity", - "filter_variable", - "filter_group", - name="uq_poverty_simulation_run_measure", - postgresql_nulls_not_distinct=True, - ), - ) - - simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="CASCADE", - ) - 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) - filter_group: str | None = Field(default=None, max_length=64) - headcount: float | None = None - total_population: float | None = None - rate: float | None = None - - simulation_run: SimulationRun = Relationship(back_populates="poverty_results") - - -class ProgramStatistics(IdentifiedModel, table=True): - """Baseline and reform totals, counts, winners and losers for one program. - - One row per report run, program and entity; ``is_tax`` marks taxes. The - economy output's detailed budget names neither the entity nor whether a - program is a tax, so both may be null, and rows without an entity count as - equal. - """ - - __tablename__ = "program_statistics" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "program_name", - "entity", - name="uq_program_statistics_run_program_entity", - postgresql_nulls_not_distinct=True, - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - program_name: str = Field(max_length=512) - entity: str | None = Field(default=None, max_length=128) - is_tax: bool | None = None - 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_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - reform_simulation_run_id: UUID = Field( - foreign_key="simulation_runs.id", - ondelete="RESTRICT", - ) - average_household_income_change: float - relative_household_income_change: float - population: float - - -class UKAreaImpactBase(GeographicImpactBase): - """Shared non-table disclosure fields of a UK area's impact. - - ``disclosure_tier`` is ``reported``, or ``approximate`` when the area's - sample is thin; the ``_lower`` and ``_upper`` fields bound each income - change's 90% interval. All are null for results without disclosure control. - """ - - disclosure_tier: str | None = Field(default=None, max_length=16) - average_household_income_change_lower: float | None = None - average_household_income_change_upper: float | None = None - relative_household_income_change_lower: float | None = None - relative_household_income_change_upper: float | None = None - - -def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: - return sa.CheckConstraint( - "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", - name=f"ck_{table_name}_disclosure_tier", - ) - - -class CongressionalDistrictImpact(GeographicImpactBase, table=True): - """Household income change and population in one US congressional district. - - One row per report run and district GEOID, with the shares of people who - gain, lose or see no change. - """ - - __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 - winner_percentage: float - loser_percentage: float - no_change_percentage: float - - report_run: ReportRun = Relationship( - back_populates="congressional_district_impacts" - ) - - -class ConstituencyImpact(UKAreaImpactBase, table=True): - """Household income change in one UK constituency. - - One row per report run and constituency code. - """ - - __tablename__ = "constituency_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "constituency_code", - name="uq_constituency_impacts_run_code", - ), - _disclosure_tier_check("constituency_impacts"), - ) - - constituency_code: str = Field(max_length=64) - constituency_name: str = Field(max_length=255) - - report_run: ReportRun = Relationship(back_populates="constituency_impacts") - - -class LocalAuthorityImpact(UKAreaImpactBase, table=True): - """Household income change in one UK local authority. - - One row per report 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", - ), - _disclosure_tier_check("local_authority_impacts"), - ) - - local_authority_code: str = Field(max_length=64) - local_authority_name: str = Field(max_length=255) - - report_run: ReportRun = Relationship(back_populates="local_authority_impacts") diff --git a/policyengine_api/data/v2/models/run_results.py b/policyengine_api/data/v2/models/run_results.py new file mode 100644 index 000000000..1f6943783 --- /dev/null +++ b/policyengine_api/data/v2/models/run_results.py @@ -0,0 +1,554 @@ +"""Canonical SQLModel tables for the results of simulation runs and report runs. + +Results that compare a report run's baseline and reform belong to the report +run and name its two simulation runs; each such row is unique within its report +run. Poverty and inequality, which describe one simulation, belong to a +simulation run. +""" + +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, enum_type +from policyengine_api.data.v2.models.reports import ReportRun +from policyengine_api.data.v2.models.simulations import SimulationRun + + +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 LaborSupplyResponse(IdentifiedModel, table=True): + """The labour supply response to the reform in one report run. + + One row per report run. The earnings changes come from the income and + substitution effects (.py's ``income_lsr`` and ``substitution_lsr``) and + total ``total_earnings_change``; ``relative_income_effect`` and + ``relative_substitution_effect`` are relative to baseline earnings. The + weekly hours are the total hours worked under the baseline and the reform, + and their change split into the two effects. + """ + + __tablename__ = "labor_supply_responses" + __table_args__ = ( + sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + income_effect_earnings_change: float + substitution_effect_earnings_change: float + total_earnings_change: float + revenue_change: float + relative_income_effect: float + relative_substitution_effect: float + baseline_weekly_hours: float + reform_weekly_hours: float + weekly_hours_change: float + weekly_hours_income_effect: float + weekly_hours_substitution_effect: float + + report_run: ReportRun = Relationship(back_populates="labor_supply_responses") + + +class LaborSupplyDecileResponse(IdentifiedModel, table=True): + """The labour supply response in one income decile of one report run. + + One row per report run and decile, as .py numbers them. The average effects + are per household; the relative effects are relative to the decile's + baseline earnings, which .py reports only for deciles above zero, so they + may be null. + """ + + __tablename__ = "labor_supply_decile_responses" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "decile", + name="uq_labor_supply_decile_responses_run_decile", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + decile: int + average_income_effect: float | None = None + average_substitution_effect: float | None = None + relative_income_effect: float | None = None + relative_substitution_effect: float | None = None + + report_run: ReportRun = Relationship(back_populates="labor_supply_decile_responses") + + +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 report 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_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.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 report run, + income variable, decile variable, entity, quantile count and decile. + ``decile_variable`` names a variable that groups the households instead of + the income variable's own deciles, such as ``household_wealth_decile``; rows + without one count as equal. + """ + + __tablename__ = "decile_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "income_variable", + "decile_variable", + "entity", + "decile", + "quantiles", + name="uq_decile_impacts_run_measure", + postgresql_nulls_not_distinct=True, + ), + 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_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + income_variable: str = Field(max_length=512) + decile_variable: str | None = Field(default=None, 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 report 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_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.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 run.""" + + __tablename__ = "inequality" + __table_args__ = ( + sa.UniqueConstraint( + "simulation_run_id", + "income_variable", + "entity", + name="uq_inequality_simulation_run_measure", + ), + ) + + simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="CASCADE", + ) + 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 + + simulation_run: SimulationRun = Relationship(back_populates="inequalities") + + +class Poverty(IdentifiedModel, table=True): + """Poverty headcount and rate of one type for one simulation run. + + One row per run, poverty type, entity, filter variable and filter group: + ``filter_variable`` selects people (``age``) and ``filter_group`` names the + group (``child``). Missing values count as equal, so each run has one + unfiltered row per type and entity. + """ + + __tablename__ = "poverty" + __table_args__ = ( + sa.UniqueConstraint( + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + name="uq_poverty_simulation_run_measure", + postgresql_nulls_not_distinct=True, + ), + ) + + simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="CASCADE", + ) + 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) + filter_group: str | None = Field(default=None, max_length=64) + headcount: float | None = None + total_population: float | None = None + rate: float | None = None + + simulation_run: SimulationRun = Relationship(back_populates="poverty_results") + + +class ProgramStatistics(IdentifiedModel, table=True): + """Baseline and reform totals, counts, winners and losers for one program. + + One row per report run, program and entity; ``is_tax`` marks taxes. The + economy output's detailed budget names neither the entity nor whether a + program is a tax, so both may be null, and rows without an entity count as + equal. + """ + + __tablename__ = "program_statistics" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "program_name", + "entity", + name="uq_program_statistics_run_program_entity", + postgresql_nulls_not_distinct=True, + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + program_name: str = Field(max_length=512) + entity: str | None = Field(default=None, max_length=128) + is_tax: bool | None = None + 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_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + average_household_income_change: float + relative_household_income_change: float + population: float + + +class UKAreaImpactBase(GeographicImpactBase): + """Shared non-table disclosure fields of a UK area's impact. + + ``disclosure_tier`` is ``reported``, or ``approximate`` when the area's + sample is thin; the ``_lower`` and ``_upper`` fields bound each income + change's 90% interval. All are null for results without disclosure control. + """ + + disclosure_tier: str | None = Field(default=None, max_length=16) + average_household_income_change_lower: float | None = None + average_household_income_change_upper: float | None = None + relative_household_income_change_lower: float | None = None + relative_household_income_change_upper: float | None = None + + +def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: + return sa.CheckConstraint( + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + name=f"ck_{table_name}_disclosure_tier", + ) + + +class CongressionalDistrictImpact(GeographicImpactBase, table=True): + """Household income change and population in one US congressional district. + + One row per report run and district GEOID, with the shares of people who + gain, lose or see no change. + """ + + __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 + winner_percentage: float + loser_percentage: float + no_change_percentage: float + + report_run: ReportRun = Relationship( + back_populates="congressional_district_impacts" + ) + + +class ConstituencyImpact(UKAreaImpactBase, table=True): + """Household income change in one UK constituency. + + One row per report run and constituency code. + """ + + __tablename__ = "constituency_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "constituency_code", + name="uq_constituency_impacts_run_code", + ), + _disclosure_tier_check("constituency_impacts"), + ) + + constituency_code: str = Field(max_length=64) + constituency_name: str = Field(max_length=255) + + report_run: ReportRun = Relationship(back_populates="constituency_impacts") + + +class LocalAuthorityImpact(UKAreaImpactBase, table=True): + """Household income change in one UK local authority. + + One row per report 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", + ), + _disclosure_tier_check("local_authority_impacts"), + ) + + local_authority_code: str = Field(max_length=64) + local_authority_name: str = Field(max_length=255) + + report_run: ReportRun = Relationship(back_populates="local_authority_impacts") diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 075ced605..dafb7b41d 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -18,11 +18,8 @@ from policyengine_api.data.v2.models.policies import Dynamic, Policy if TYPE_CHECKING: - from policyengine_api.data.v2.models.reports import ( - Inequality, - Poverty, - Report, - ) + from policyengine_api.data.v2.models.reports import Report + from policyengine_api.data.v2.models.run_results import Inequality, Poverty class SimulationStatus(StrEnum): diff --git a/policyengine_api/data/v2/privilege_plans.py b/policyengine_api/data/v2/privilege_plans.py new file mode 100644 index 000000000..671243cc8 --- /dev/null +++ b/policyengine_api/data/v2/privilege_plans.py @@ -0,0 +1,373 @@ +"""Plan, state and report the table privileges the v2 manifest declares. + +The pure half of ``privilege_reconciliation``: it compares the declared +privileges with what a database grants, writes the GRANT statements, finds the +privileges a role lacks or holds too widely, and renders the reports. It runs +no SQL; ``privilege_reconciliation`` reads the database and runs the result. +""" + +from __future__ import annotations + +import re +from collections import defaultdict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from policyengine_api.data.v2.privileges import ( + ColumnPrivilege, + DeclaredPrivilege, + TablePrivilege, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Collection, Iterable, Mapping, Sequence + +SCHEMA = "public" +SAFE_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]{0,62}$") + + +class PrivilegeReconciliationError(RuntimeError): + """Raised when privileges can't be reconciled or don't match the manifest.""" + + +class UnsafeIdentifierError(PrivilegeReconciliationError): + """Raised for a role or table name that isn't a plain lowercase identifier.""" + + +@dataclass(frozen=True, order=True) +class PrivilegeEntry: + """One privilege one role holds, or should hold, on one table. + + ``column`` names the one column a column privilege is on; it is empty for a + privilege on the whole table. + """ + + role: str + table: str + privilege: str + column: str = "" + + +@dataclass(frozen=True) +class GrantedPrivileges: + """The managed roles' direct grants in one database. + + ``tables`` holds grants on whole tables, by role then table, and + ``columns`` grants on columns, by role, table, then column. + """ + + tables: Mapping[str, Mapping[str, Collection[str]]] + columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] = field( + default_factory=dict + ) + + +@dataclass(frozen=True) +class DefaultPrivilege: + """A default privilege PostgreSQL gives on objects a role creates later.""" + + creating_role: str + schema: str + object_type: str + grantee: str + privileges: tuple[str, ...] + + +@dataclass(frozen=True) +class ReconciliationPlan: + """What ``apply`` grants and reports for one database. + + ``skipped_roles`` are missing roles awaiting creation by an operator; they + don't block the plan, and nothing is granted to them. + """ + + grants: tuple[PrivilegeEntry, ...] + undeclared: tuple[PrivilegeEntry, ...] + missing_roles: tuple[str, ...] + missing_tables: tuple[str, ...] + skipped_roles: tuple[str, ...] = () + + @property + def can_apply(self) -> bool: + """Whether every declared role and table exists, so grants may run.""" + return not self.missing_roles and not self.missing_tables + + +def _declared_entries( + role: str, table: str, privileges: Collection[DeclaredPrivilege] +) -> list[PrivilegeEntry]: + """Return one entry per whole-table privilege and per column of a column one.""" + entries = [] + for privilege in privileges: + if isinstance(privilege, ColumnPrivilege): + entries += [ + PrivilegeEntry(role, table, privilege.privilege.value, column) + for column in privilege.columns + ] + else: + entries.append(PrivilegeEntry(role, table, privilege.value)) + return entries + + +def plan_reconciliation( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + granted: GrantedPrivileges, + *, + existing_roles: Collection[str], + existing_tables: Collection[str], + roles_awaiting_creation: Collection[str] = (), +) -> ReconciliationPlan: + """Compare declared privileges with what the database grants. + + A missing role in ``roles_awaiting_creation`` is skipped instead of missing. + """ + absent_roles = set(declared) - set(existing_roles) + skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) + missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) + declared_tables = {table for tables in declared.values() for table in tables} + missing_tables = tuple(sorted(declared_tables - set(existing_tables))) + declared_entries = { + entry + for role, tables in declared.items() + for table, privileges in tables.items() + for entry in _declared_entries(role, table, privileges) + } + granted_entries = { + PrivilegeEntry(role, table, privilege) + for role, tables in granted.tables.items() + for table, privileges in tables.items() + for privilege in privileges + } | { + PrivilegeEntry(role, table, privilege, column) + for role, tables in granted.columns.items() + for table, columns in tables.items() + for column, privileges in columns.items() + for privilege in privileges + } + grants = sorted( + entry + for entry in declared_entries - granted_entries + if entry.role not in absent_roles and entry.table not in missing_tables + ) + undeclared = sorted(granted_entries - declared_entries) + return ReconciliationPlan( + grants=tuple(grants), + undeclared=tuple(undeclared), + missing_roles=missing_roles, + missing_tables=missing_tables, + skipped_roles=skipped_roles, + ) + + +def safe_identifier(identifier: str) -> str: + """Return ``identifier`` if it is a plain lowercase SQL name, else raise.""" + if SAFE_IDENTIFIER.fullmatch(identifier) is None: + msg = f"refusing to use {identifier!r} as a role, table or column name" + raise UnsafeIdentifierError(msg) + return identifier + + +def grant_statements( + plan: ReconciliationPlan, quote_identifier: Callable[[str], str] +) -> list[str]: + """Return the GRANTs for a plan, in role and table order. + + Each role and table gets one GRANT of its whole-table privileges, then one + per column privilege, naming its columns in name order. + """ + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str, str], set[str]] = defaultdict(set) + for entry in plan.grants: + target = (safe_identifier(entry.role), safe_identifier(entry.table)) + privilege = TablePrivilege(entry.privilege).value + if entry.column: + columns_by_privilege[(*target, privilege)].add( + safe_identifier(entry.column) + ) + else: + table_privileges[target].add(privilege) + statements: list[tuple[tuple[str, str, str], str]] = [ + ( + (role, table, ""), + f"GRANT {', '.join(sorted(privileges))} ON TABLE " + f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", + ) + for (role, table), privileges in table_privileges.items() + ] + statements += [ + ( + (role, table, privilege), + f"GRANT {privilege} " + f"({', '.join(quote_identifier(column) for column in sorted(columns))}) " + f"ON TABLE {quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", + ) + for (role, table, privilege), columns in columns_by_privilege.items() + ] + return [statement for _target, statement in sorted(statements)] + + +def verification_failures( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + has_privilege: Callable[[str, str, str], bool], + has_column_privilege: Callable[[str, str, str, str], bool] | None = None, +) -> list[PrivilegeEntry]: + """Return every declared privilege the role doesn't effectively hold. + + A column privilege is checked on each of its columns with + ``has_column_privilege``, which is required when one is declared. + """ + failures = [] + for role, tables in declared.items(): + for table, privileges in tables.items(): + for entry in _declared_entries(role, table, privileges): + if not entry.column: + held = has_privilege(role, table, entry.privilege) + elif has_column_privilege is None: + msg = "checking a column privilege needs has_column_privilege" + raise ValueError(msg) + else: + held = has_column_privilege( + role, table, entry.column, entry.privilege + ) + if not held: + failures.append(entry) + return sorted(failures) + + +def table_wide_overrides( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + has_privilege: Callable[[str, str, str], bool], +) -> list[PrivilegeEntry]: + """Return each column privilege a role also holds on the whole table. + + Holding it on the whole table lets the role use it on every column, past + the columns the manifest declares. + """ + return sorted( + PrivilegeEntry(role, table, privilege.privilege.value) + for role, tables in declared.items() + for table, privileges in tables.items() + for privilege in privileges + if isinstance(privilege, ColumnPrivilege) + and has_privilege(role, table, privilege.privilege.value) + ) + + +def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: + """Return one row per role and table, naming each column privilege's columns.""" + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str], dict[str, set[str]]] = defaultdict( + lambda: defaultdict(set) + ) + for entry in entries: + target = (entry.role, entry.table) + if entry.column: + columns_by_privilege[target][entry.privilege].add(entry.column) + else: + table_privileges[target].add(entry.privilege) + rows = [] + for role, table in sorted(set(table_privileges) | set(columns_by_privilege)): + described = sorted(table_privileges.get((role, table), ())) + [ + f"{privilege} ({', '.join(sorted(columns))})" + for privilege, columns in sorted( + columns_by_privilege.get((role, table), {}).items() + ) + ] + rows.append(f"| `{role}` | `{table}` | {', '.join(described)} |") + return rows + + +def skipped_roles_line(roles: Iterable[str]) -> str: + """Return the report sentence naming roles skipped until they are created.""" + return ( + "Declared roles not created yet, skipped until an operator creates them: " + + ", ".join(f"`{role}`" for role in roles) + ) + + +def render_summary( + plan: ReconciliationPlan, + default_privileges: Sequence[DefaultPrivilege], + *, + deployment_environment: str, + applied: bool, +) -> str: + """Return the Markdown report for the job summary and the log.""" + lines = [f"## v2 table privileges ({deployment_environment})", ""] + if plan.skipped_roles: + lines += [skipped_roles_line(plan.skipped_roles), ""] + if plan.missing_roles: + lines += [ + "Declared roles missing from the database (nothing was granted): " + + ", ".join(f"`{role}`" for role in plan.missing_roles), + "", + ] + if plan.missing_tables: + lines += [ + "Declared tables missing from the database (nothing was granted): " + + ", ".join(f"`{table}`" for table in plan.missing_tables), + "", + ] + grant_heading = "Granted" if applied else "Would grant" + lines += [f"### {grant_heading}", ""] + if plan.grants: + lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] + lines += _privilege_rows(plan.grants) + else: + lines.append("Nothing: every declared privilege was already granted.") + lines += [ + "", + "### Privileges the manifest doesn't declare (reported, not revoked)", + "", + ] + if plan.undeclared: + lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] + lines += _privilege_rows(plan.undeclared) + else: + lines.append("None.") + lines += ["", "### Default privileges for new tables", ""] + if default_privileges: + lines += [ + "| Creating role | Schema | Objects | Grantee | Privileges |", + "| --- | --- | --- | --- | --- |", + ] + lines += [ + f"| `{default.creating_role}` | `{default.schema}` | " + f"{default.object_type} | `{default.grantee}` | " + f"{', '.join(default.privileges)} |" + for default in default_privileges + ] + else: + lines.append("None.") + return "\n".join(lines) + "\n" + + +def _failure_row(failure: PrivilegeEntry) -> str: + if failure.column: + return ( + f"- `{failure.role}` lacks {failure.privilege} on column " + f"`{failure.column}` of `{failure.table}`" + ) + return f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" + + +def verification_report( + failures: Sequence[PrivilegeEntry], overrides: Sequence[PrivilegeEntry] +) -> str: + """Return verify's report of missing privileges and whole-table overrides.""" + sections = [] + if failures: + rows = "\n".join(_failure_row(failure) for failure in failures) + sections.append(f"Declared privileges missing after reconciling:\n\n{rows}\n") + if overrides: + rows = "\n".join( + f"- `{override.role}` holds {override.privilege} on the whole of " + f"`{override.table}`, which overrides its declared columns" + for override in overrides + ) + sections.append( + f"Privileges held on whole tables, past their declared columns:\n\n{rows}\n" + ) + return "\n".join(sections) diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index 96581f2e6..008e33e87 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -27,16 +27,17 @@ The report goes to stdout and, when set, ``GITHUB_STEP_SUMMARY``. It names roles, tables and privileges only, never a connection URL or credential. + +This module reads the database and runs the commands; ``privilege_plans`` +compares, states and reports the privileges, and its names are re-exported here. """ from __future__ import annotations import argparse import os -import re import sys from collections import defaultdict -from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING @@ -49,12 +50,27 @@ load_v2_alembic_settings, qualify_v2_connection, ) +from policyengine_api.data.v2.privilege_plans import ( + SCHEMA, + DefaultPrivilege, + GrantedPrivileges, + PrivilegeEntry, + PrivilegeReconciliationError, + ReconciliationPlan, + UnsafeIdentifierError, + grant_statements, + plan_reconciliation, + render_summary, + safe_identifier, + skipped_roles_line, + table_wide_overrides, + verification_failures, + verification_report, +) from policyengine_api.data.v2.privileges import ( PERSISTENT_ENVIRONMENTS, ROLES_AWAITING_PROVISIONING, - ColumnPrivilege, DeclaredPrivilege, - TablePrivilege, declared_privileges, ) from policyengine_api.data.v2.settings import V2ConfigurationError @@ -64,324 +80,39 @@ from sqlalchemy import Connection, Engine +# Names other modules import from here; the planning and reporting half lives +# in privilege_plans. +__all__ = [ + "DEPLOYMENT_ENVIRONMENT", + "GITHUB_STEP_SUMMARY", + "SCHEMA", + "DefaultPrivilege", + "GrantedPrivileges", + "PrivilegeEntry", + "PrivilegeReconciliationError", + "ReconciliationPlan", + "UnsafeIdentifierError", + "apply_declared_privileges", + "check_declared_privileges", + "check_table_wide_overrides", + "create_missing_roles", + "grant_statements", + "main", + "plan_reconciliation", + "read_default_privileges", + "read_existing_roles", + "read_granted_column_privileges", + "read_granted_privileges", + "read_public_tables", + "render_summary", + "run", + "skipped_roles", + "table_wide_overrides", + "verification_failures", +] + DEPLOYMENT_ENVIRONMENT = "DEPLOYMENT_ENVIRONMENT" GITHUB_STEP_SUMMARY = "GITHUB_STEP_SUMMARY" -SCHEMA = "public" -SAFE_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]{0,62}$") - - -class PrivilegeReconciliationError(RuntimeError): - """Raised when privileges can't be reconciled or don't match the manifest.""" - - -class UnsafeIdentifierError(PrivilegeReconciliationError): - """Raised for a role or table name that isn't a plain lowercase identifier.""" - - -@dataclass(frozen=True, order=True) -class PrivilegeEntry: - """One privilege one role holds, or should hold, on one table. - - ``column`` names the one column a column privilege is on; it is empty for a - privilege on the whole table. - """ - - role: str - table: str - privilege: str - column: str = "" - - -@dataclass(frozen=True) -class GrantedPrivileges: - """The managed roles' direct grants in one database. - - ``tables`` holds grants on whole tables, by role then table, and - ``columns`` grants on columns, by role, table, then column. - """ - - tables: Mapping[str, Mapping[str, Collection[str]]] - columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] = field( - default_factory=dict - ) - - -@dataclass(frozen=True) -class DefaultPrivilege: - """A default privilege PostgreSQL gives on objects a role creates later.""" - - creating_role: str - schema: str - object_type: str - grantee: str - privileges: tuple[str, ...] - - -@dataclass(frozen=True) -class ReconciliationPlan: - """What ``apply`` grants and reports for one database. - - ``skipped_roles`` are missing roles awaiting creation by an operator; they - don't block the plan, and nothing is granted to them. - """ - - grants: tuple[PrivilegeEntry, ...] - undeclared: tuple[PrivilegeEntry, ...] - missing_roles: tuple[str, ...] - missing_tables: tuple[str, ...] - skipped_roles: tuple[str, ...] = () - - @property - def can_apply(self) -> bool: - """Whether every declared role and table exists, so grants may run.""" - return not self.missing_roles and not self.missing_tables - - -def _declared_entries( - role: str, table: str, privileges: Collection[DeclaredPrivilege] -) -> list[PrivilegeEntry]: - """Return one entry per whole-table privilege and per column of a column one.""" - entries = [] - for privilege in privileges: - if isinstance(privilege, ColumnPrivilege): - entries += [ - PrivilegeEntry(role, table, privilege.privilege.value, column) - for column in privilege.columns - ] - else: - entries.append(PrivilegeEntry(role, table, privilege.value)) - return entries - - -def plan_reconciliation( - declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], - granted: GrantedPrivileges, - *, - existing_roles: Collection[str], - existing_tables: Collection[str], - roles_awaiting_creation: Collection[str] = (), -) -> ReconciliationPlan: - """Compare declared privileges with what the database grants. - - A missing role in ``roles_awaiting_creation`` is skipped instead of missing. - """ - absent_roles = set(declared) - set(existing_roles) - skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) - missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) - declared_tables = {table for tables in declared.values() for table in tables} - missing_tables = tuple(sorted(declared_tables - set(existing_tables))) - declared_entries = { - entry - for role, tables in declared.items() - for table, privileges in tables.items() - for entry in _declared_entries(role, table, privileges) - } - granted_entries = { - PrivilegeEntry(role, table, privilege) - for role, tables in granted.tables.items() - for table, privileges in tables.items() - for privilege in privileges - } | { - PrivilegeEntry(role, table, privilege, column) - for role, tables in granted.columns.items() - for table, columns in tables.items() - for column, privileges in columns.items() - for privilege in privileges - } - grants = sorted( - entry - for entry in declared_entries - granted_entries - if entry.role not in absent_roles and entry.table not in missing_tables - ) - undeclared = sorted(granted_entries - declared_entries) - return ReconciliationPlan( - grants=tuple(grants), - undeclared=tuple(undeclared), - missing_roles=missing_roles, - missing_tables=missing_tables, - skipped_roles=skipped_roles, - ) - - -def _safe(identifier: str) -> str: - if SAFE_IDENTIFIER.fullmatch(identifier) is None: - msg = f"refusing to use {identifier!r} as a role, table or column name" - raise UnsafeIdentifierError(msg) - return identifier - - -def grant_statements( - plan: ReconciliationPlan, quote_identifier: Callable[[str], str] -) -> list[str]: - """Return the GRANTs for a plan, in role and table order. - - Each role and table gets one GRANT of its whole-table privileges, then one - per column privilege, naming its columns in name order. - """ - table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) - columns_by_privilege: dict[tuple[str, str, str], set[str]] = defaultdict(set) - for entry in plan.grants: - target = (_safe(entry.role), _safe(entry.table)) - privilege = TablePrivilege(entry.privilege).value - if entry.column: - columns_by_privilege[(*target, privilege)].add(_safe(entry.column)) - else: - table_privileges[target].add(privilege) - statements: list[tuple[tuple[str, str, str], str]] = [ - ( - (role, table, ""), - f"GRANT {', '.join(sorted(privileges))} ON TABLE " - f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " - f"TO {quote_identifier(role)}", - ) - for (role, table), privileges in table_privileges.items() - ] - statements += [ - ( - (role, table, privilege), - f"GRANT {privilege} " - f"({', '.join(quote_identifier(column) for column in sorted(columns))}) " - f"ON TABLE {quote_identifier(SCHEMA)}.{quote_identifier(table)} " - f"TO {quote_identifier(role)}", - ) - for (role, table, privilege), columns in columns_by_privilege.items() - ] - return [statement for _target, statement in sorted(statements)] - - -def verification_failures( - declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], - has_privilege: Callable[[str, str, str], bool], - has_column_privilege: Callable[[str, str, str, str], bool] | None = None, -) -> list[PrivilegeEntry]: - """Return every declared privilege the role doesn't effectively hold. - - A column privilege is checked on each of its columns with - ``has_column_privilege``, which is required when one is declared. - """ - failures = [] - for role, tables in declared.items(): - for table, privileges in tables.items(): - for entry in _declared_entries(role, table, privileges): - if not entry.column: - held = has_privilege(role, table, entry.privilege) - elif has_column_privilege is None: - msg = "checking a column privilege needs has_column_privilege" - raise ValueError(msg) - else: - held = has_column_privilege( - role, table, entry.column, entry.privilege - ) - if not held: - failures.append(entry) - return sorted(failures) - - -def table_wide_overrides( - declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], - has_privilege: Callable[[str, str, str], bool], -) -> list[PrivilegeEntry]: - """Return each column privilege a role also holds on the whole table. - - Holding it on the whole table lets the role use it on every column, past - the columns the manifest declares. - """ - return sorted( - PrivilegeEntry(role, table, privilege.privilege.value) - for role, tables in declared.items() - for table, privileges in tables.items() - for privilege in privileges - if isinstance(privilege, ColumnPrivilege) - and has_privilege(role, table, privilege.privilege.value) - ) - - -def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: - """Return one row per role and table, naming each column privilege's columns.""" - table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) - columns_by_privilege: dict[tuple[str, str], dict[str, set[str]]] = defaultdict( - lambda: defaultdict(set) - ) - for entry in entries: - target = (entry.role, entry.table) - if entry.column: - columns_by_privilege[target][entry.privilege].add(entry.column) - else: - table_privileges[target].add(entry.privilege) - rows = [] - for role, table in sorted(set(table_privileges) | set(columns_by_privilege)): - described = sorted(table_privileges.get((role, table), ())) + [ - f"{privilege} ({', '.join(sorted(columns))})" - for privilege, columns in sorted( - columns_by_privilege.get((role, table), {}).items() - ) - ] - rows.append(f"| `{role}` | `{table}` | {', '.join(described)} |") - return rows - - -def _skipped_roles_line(roles: Iterable[str]) -> str: - return ( - "Declared roles not created yet, skipped until an operator creates them: " - + ", ".join(f"`{role}`" for role in roles) - ) - - -def render_summary( - plan: ReconciliationPlan, - default_privileges: Sequence[DefaultPrivilege], - *, - deployment_environment: str, - applied: bool, -) -> str: - """Return the Markdown report for the job summary and the log.""" - lines = [f"## v2 table privileges ({deployment_environment})", ""] - if plan.skipped_roles: - lines += [_skipped_roles_line(plan.skipped_roles), ""] - if plan.missing_roles: - lines += [ - "Declared roles missing from the database (nothing was granted): " - + ", ".join(f"`{role}`" for role in plan.missing_roles), - "", - ] - if plan.missing_tables: - lines += [ - "Declared tables missing from the database (nothing was granted): " - + ", ".join(f"`{table}`" for table in plan.missing_tables), - "", - ] - grant_heading = "Granted" if applied else "Would grant" - lines += [f"### {grant_heading}", ""] - if plan.grants: - lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] - lines += _privilege_rows(plan.grants) - else: - lines.append("Nothing: every declared privilege was already granted.") - lines += [ - "", - "### Privileges the manifest doesn't declare (reported, not revoked)", - "", - ] - if plan.undeclared: - lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] - lines += _privilege_rows(plan.undeclared) - else: - lines.append("None.") - lines += ["", "### Default privileges for new tables", ""] - if default_privileges: - lines += [ - "| Creating role | Schema | Objects | Grantee | Privileges |", - "| --- | --- | --- | --- | --- |", - ] - lines += [ - f"| `{default.creating_role}` | `{default.schema}` | " - f"{default.object_type} | `{default.grantee}` | " - f"{', '.join(default.privileges)} |" - for default in default_privileges - ] - else: - lines.append("None.") - return "\n".join(lines) + "\n" # --- Database access ------------------------------------------------------- @@ -533,7 +264,7 @@ def create_missing_roles(connection: Connection, roles: Iterable[str]) -> list[s missing = sorted(set(roles) - read_existing_roles(connection, list(roles))) quote = connection.dialect.identifier_preparer.quote for role in missing: - connection.execute(text(f"CREATE ROLE {quote(_safe(role))} NOLOGIN")) + connection.execute(text(f"CREATE ROLE {quote(safe_identifier(role))} NOLOGIN")) return missing @@ -597,7 +328,7 @@ def has_privilege(role: str, table: str, privilege: str) -> bool: _HAS_TABLE_PRIVILEGE, { "role": role, - "qualified_table": f"{SCHEMA}.{_safe(table)}", + "qualified_table": f"{SCHEMA}.{safe_identifier(table)}", "privilege": privilege, }, ).scalar_one() @@ -613,8 +344,8 @@ def has_column_privilege( _HAS_COLUMN_PRIVILEGE, { "role": role, - "qualified_table": f"{SCHEMA}.{_safe(table)}", - "column": _safe(column), + "qualified_table": f"{SCHEMA}.{safe_identifier(table)}", + "column": safe_identifier(column), "privilege": privilege, }, ).scalar_one() @@ -672,34 +403,6 @@ def _write_summary(environ: Mapping[str, str], summary: str) -> None: summary_file.write(summary) -def _failure_row(failure: PrivilegeEntry) -> str: - if failure.column: - return ( - f"- `{failure.role}` lacks {failure.privilege} on column " - f"`{failure.column}` of `{failure.table}`" - ) - return f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" - - -def _verification_report( - failures: Sequence[PrivilegeEntry], overrides: Sequence[PrivilegeEntry] -) -> str: - sections = [] - if failures: - rows = "\n".join(_failure_row(failure) for failure in failures) - sections.append(f"Declared privileges missing after reconciling:\n\n{rows}\n") - if overrides: - rows = "\n".join( - f"- `{override.role}` holds {override.privilege} on the whole of " - f"`{override.table}`, which overrides its declared columns" - for override in overrides - ) - sections.append( - f"Privileges held on whole tables, past their declared columns:\n\n{rows}\n" - ) - return "\n".join(sections) - - def _verify( connection: Connection, declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], @@ -720,11 +423,11 @@ def _verify( failures = check_declared_privileges(connection, checked) overrides = check_table_wide_overrides(connection, checked) heading = f"## v2 table privileges ({environment})\n\n" - skipped_note = f"{_skipped_roles_line(skipped)}\n\n" if skipped else "" + skipped_note = f"{skipped_roles_line(skipped)}\n\n" if skipped else "" if failures or overrides: _write_summary( environ, - f"{heading}{skipped_note}{_verification_report(failures, overrides)}", + f"{heading}{skipped_note}{verification_report(failures, overrides)}", ) return 1 if skipped: diff --git a/pyproject.toml b/pyproject.toml index 024e94cb9..a1e0e213f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -275,6 +275,8 @@ ignore_imports = [ "policyengine_api.data.v2.models.policy_mappings -> sqlmodel", "policyengine_api.data.v2.models.reports -> sqlalchemy", "policyengine_api.data.v2.models.reports -> sqlmodel", + "policyengine_api.data.v2.models.run_results -> sqlalchemy", + "policyengine_api.data.v2.models.run_results -> sqlmodel", "policyengine_api.data.v2.models.simulations -> sqlalchemy", "policyengine_api.data.v2.models.simulations -> sqlmodel", "policyengine_api.data.v2.models.users -> sqlalchemy", @@ -421,7 +423,8 @@ 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.data.v2.models.run_results.GeographicImpactBase", + "policyengine_api.data.v2.models.run_results.UKAreaImpactBase", "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/rules/registry.toml b/rules/registry.toml index 613e03313..d3cfb4f05 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -440,6 +440,8 @@ allowed_ignore_imports = [ "policyengine_api.data.v2.models.policy_mappings -> sqlmodel", "policyengine_api.data.v2.models.reports -> sqlalchemy", "policyengine_api.data.v2.models.reports -> sqlmodel", + "policyengine_api.data.v2.models.run_results -> sqlalchemy", + "policyengine_api.data.v2.models.run_results -> sqlmodel", "policyengine_api.data.v2.models.simulations -> sqlalchemy", "policyengine_api.data.v2.models.simulations -> sqlmodel", "policyengine_api.data.v2.models.users -> sqlalchemy", diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py index 281fbb36b..24298ab60 100644 --- a/tests/unit/v2/test_simulation_run_models.py +++ b/tests/unit/v2/test_simulation_run_models.py @@ -325,7 +325,7 @@ def test_report_run_relationships_name_labor_supply_rows() -> None: LaborSupplyDecileResponse ) assert "comparisons" not in relationships - assert LaborSupplyResponse.__module__.endswith(".reports") + assert LaborSupplyResponse.__module__.endswith(".run_results") def test_uk_area_impacts_record_disclosure_control() -> None: From 7642bd430b0a780f24ba320b611fa3bb5c30a230 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:33:21 +0400 Subject: [PATCH 21/22] Test every run column the simulation API's roles may update From the second review round: - A role test checks, for the simulation service and the Modal worker, every column of simulation_runs, report_runs, household_jobs and budget_window_runs: an UPDATE of it succeeds exactly when the manifest declares it for that role. Seen failing with an undeclared column grant. - api.v2.deployed-privileges-match's scope names privilege_plans, which now holds the planning and verification code. - The Alembic guide's privilege paragraph keeps apply's sentences together. - The report runs module says complete_report_run sets markdown, which no managed role may update. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 12 +++--- policyengine_api/data/v2/report_runs.py | 5 ++- rules/registry.toml | 2 +- .../test_v2_simulation_role_privileges.py | 38 +++++++++++++++++++ 4 files changed, 48 insertions(+), 9 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 1fe4951a3..62630e6f8 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -283,12 +283,12 @@ check the declarations: `apply` grants each declared privilege a role lacks, as one `GRANT` per role and table and one per column privilege. It never revokes: it reports privileges the file doesn't declare, on whole tables or on columns, and the -projects' default privileges, in the job summary. `verify` checks each -column of a column privilege with `has_column_privilege`, and fails when a -role holds the same privilege on the whole table. It grants nothing when a declared role or -table is missing. The roles themselves are provisioned outside the -repository; `--create-missing-roles` is accepted only for the disposable test -database. +projects' default privileges, in the job summary. It grants nothing when a +declared role or table is missing. `verify` checks each column of a column +privilege with `has_column_privilege`, and fails when a role holds the same +privilege on the whole table. The roles themselves are provisioned outside +the repository; `--create-missing-roles` is accepted only for the disposable +test database. A pull request may declare a role before an operator creates it by listing it in `ROLES_AWAITING_PROVISIONING` in `privileges.py`. While a listed role is diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 1429bf379..f1980aadd 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -2,8 +2,9 @@ ``set_report_type`` and ``create_report_run`` lock ``reports`` rows with ``SELECT ... FOR UPDATE``, which needs ``UPDATE`` on ``reports``; no managed -role holds it, so none can run them. The simulation API creates reports and -report runs with its own code. +role holds it, so none can run them. ``complete_report_run`` sets +``markdown``, which no managed role may update either. The simulation API +creates and completes reports and report runs with its own code. """ from __future__ import annotations diff --git a/rules/registry.toml b/rules/registry.toml index d3cfb4f05..458e46d69 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1069,7 +1069,7 @@ id = "api.v2.deployed-privileges-match" category = "migration" statement = "Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares, on the whole table or on the columns it names, right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing or a role holds on a whole table a privilege declared for only some of its columns. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" -scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] +scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_plans.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] tier = "mechanical" status = "enforced" autofix = false diff --git a/tests/integration/test_v2_simulation_role_privileges.py b/tests/integration/test_v2_simulation_role_privileges.py index a71dc7633..fbef49bf9 100644 --- a/tests/integration/test_v2_simulation_role_privileges.py +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -20,6 +20,12 @@ from sqlalchemy.pool import NullPool from policyengine_api.constants import REPO +from policyengine_api.data.v2.models import V2_METADATA +from policyengine_api.data.v2.privileges import ( + ColumnPrivilege, + TablePrivilege, + declared_privileges, +) from tests.integration.simulation_rows import ( RUN_VERSIONS, PublicApiRows, @@ -331,3 +337,35 @@ def test_the_public_api_reads_runs_and_results_but_writes_neither( report_id=uuid4(), run_id=uuid4(), ) + + +@pytest.mark.parametrize("role_kind", ["simulation_service", "modal_worker"]) +def test_each_simulation_role_updates_exactly_its_declared_columns( + connection: Connection, managed_v2_roles, role_kind: str +) -> None: + """Every column of the run tables is updatable exactly when it is declared.""" + role = getattr(managed_v2_roles, role_kind) + declared = declared_privileges(managed_v2_roles.deployment_environment)[role] + _act_as(connection, role) + + for table in ( + "simulation_runs", + "report_runs", + "household_jobs", + "budget_window_runs", + ): + updatable = { + column + for privilege in declared[table] + if isinstance(privilege, ColumnPrivilege) + and privilege.privilege is TablePrivilege.UPDATE + for column in privilege.columns + } + for column in V2_METADATA.tables[table].c: + # The privilege is checked when the statement starts, so no row is + # needed. + statement = f"UPDATE {table} SET {column.name} = {column.name} WHERE false" + if column.name in updatable: + connection.execute(text(statement)) + else: + _refused(connection, statement) From 772ca6b14d06f1b1dfdaf19f7353cd6d8ee356a7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:55:52 +0400 Subject: [PATCH 22/22] Correct the v2 simulation runs downgrade's stated precondition The third review round found that the downgrade doesn't need household_jobs empty: it drops that table's new columns with their values. It needs simulations empty, which leaves no run or result row either. The revision's docstring and the Alembic guide say so; no operation changes. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/alembic-migrations.md | 5 +++-- .../1356f658347d_record_simulation_runs_and_key_results_.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 62630e6f8..79d780554 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -586,8 +586,9 @@ The upgrade adds required columns without server defaults to `household_jobs`, `report_runs` and the result tables, so it needs them empty, as they are while no managed role holds a write privilege on them; on a table with rows it fails, and the deploy stops before anything changes. The downgrade restores required columns -the same way, so it needs `simulations`, `household_jobs` and the result -tables empty. The previous Public API revision keeps working at this revision: +of `simulations` and the result tables the same way, so it needs +`simulations` empty, which leaves no run or result row either; it drops +`household_jobs`' new columns with their values. The previous Public API revision keeps working at this revision: none of its roles can read or write a table the revision changes, except the catalog writer, which only counts the rows of `simulations`, `reports` and `report_runs`. diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py index 69a386c3b..c118ea5d0 100644 --- a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -13,8 +13,9 @@ The upgrade adds required columns without server defaults to household_jobs, report_runs and the result tables, so it needs those tables empty, as they are while no managed role can write them. The downgrade restores required columns -the same way, so it needs simulations, household_jobs and the result tables -empty. +of simulations and the result tables the same way, so it needs simulations +empty, which leaves no run or result row either; it drops household_jobs' new +columns with their values. """ from typing import Sequence, Union