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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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