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/changelog.d/v2-simulation-runs.changed.md b/changelog.d/v2-simulation-runs.changed.md new file mode 100644 index 000000000..9e459c853 --- /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. 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 bd2719994..d69b74ba2 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, 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 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. @@ -248,6 +253,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` @@ -270,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 ``` @@ -453,3 +466,120 @@ 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`) 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`); + 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`, `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. + +The name `simulation_runs` is also a v1 Cloud SQL table. The v2 metadata +validation (`policyengine_api/data/v2/metadata_validation.py`) now treats it as +a name the two databases share, as it does `simulations`. + +### API v2 foreign keys between the two services' tables + +Revision `7a83a87d029a` (parent `2ffd50f47a2b`) was autogenerated after the +models stopped declaring foreign keys from the tables the simulation service +and its Modal workers write (`simulations`, `reports`, `household_jobs`) to the +tables the Public API writes (`users`, `policies`, `dynamics`, `households`, +and the catalog's `tax_benefit_models`, `tax_benefit_model_versions`, +`datasets` and `regions`). The 14 columns stay as plain UUIDs: the service that +writes a row checks the rows it names. Foreign keys between tables of the same +service stay. Deleting a Public API row no longer clears or blocks the rows +that name it: in particular, deleting a user leaves that user's ID in +`reports.user_id`, which no managed role may update; the owner of a saved +report is the user association, which the deletion removes. It has no post-generation corrections. Its downgrade restores the +keys, so it fails while a column names a row the referenced table lacks. +`tests/unit/v2/test_simulation_run_models.py` checks that no foreign key or ORM +relationship joins the two services' tables. + +### API v2 economy output fields + +Revision `7f43d2c7c3c6` (parent `7a83a87d029a`) was autogenerated after the +models gained the fields of the simulation API's economy output that had no +column: `report_runs.spm_provenance` (JSON) and `data_version`; the winner, +loser and no-change shares of `congressional_district_impacts`; and the UK +disclosure control fields (`disclosure_tier`, checked to be `reported` or +`approximate`, and the four `_lower`/`_upper` interval bounds) of +`constituency_impacts` and `local_authority_impacts`. The output's detailed +budget names neither a program's entity nor whether it is a tax, so +`program_statistics.entity` and `is_tax` become nullable, and its unique key +treats a missing entity as equal (`NULLS NOT DISTINCT`). It has no +post-generation corrections. The three share columns are required, so the +upgrade needs `congressional_district_impacts` empty; the downgrade needs +every `program_statistics` row to have an entity and a tax flag. + +### API v2 result row identity + +Revision `2e3de7c572ed` (parent `7f43d2c7c3c6`) was autogenerated after the +models made each result row tell apart what it measures: + +- the comparison result tables' unique keys include the pair of simulation + runs (`baseline_simulation_run_id`, `reform_simulation_run_id`) after + `report_run_id`, because a budget-window report run compares one pair per + year; +- `decile_impacts.decile_variable` names the variable that groups households + when it is not the income variable's own deciles (for example + `household_wealth_decile`), and joins the table's unique key; +- `poverty.filter_group` names the group a filter selects (`child` for + `age`), and joins the table's unique key. + +Missing values in those keys count as equal (`NULLS NOT DISTINCT`, which +PostgreSQL supports from version 15; `poverty` and `program_statistics` use it +from the earlier revisions too). It has no +post-generation corrections and needs no table empty; its downgrade fails +while a report run holds comparison rows of more than one pair, or a run holds +rows that differ only in a removed column. + +### API v2 comparisons + +Revision `dce0caa329cc` (parent `2e3de7c572ed`) was autogenerated after the +models moved the parts of a comparison that have no result table +(`labor_supply_result`, `cliff_impact_result`, `spm_provenance`, +`data_version`) from `report_runs` to the new `report_run_comparisons`, one row +per report run and compared pair of simulation runs, because a budget-window +run compares one pair per year. It has no post-generation corrections. The +four `report_runs` columns it drops were added by the earlier revisions of +this change and never written; its downgrade drops the new table and its +rows. diff --git a/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/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/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/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/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/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/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..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 = "c614b6bcccfc" +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 8498cfbb6..cc182268d 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -82,11 +82,15 @@ ProgramStatistics, Report, ReportRun, + ReportRunComparison, + ReportRunSimulationRun, ReportRunStatus, ReportRunTrigger, + SimulationRole, ) from policyengine_api.data.v2.models.simulations import ( # noqa: E402 Simulation, + SimulationRun, SimulationStatus, SimulationType, ) @@ -138,9 +142,13 @@ "RegionType", "Report", "ReportRun", + "ReportRunComparison", + "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..351dff625 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -6,26 +6,20 @@ from uuid import UUID import sqlalchemy as sa +from sqlalchemy.dialects import postgresql from sqlmodel import Field, Relationship -from policyengine_api.data.v2.models.base import ( - IdentifiedModel, - TimestampedModel, - enum_type, -) -from policyengine_api.data.v2.models.policies import Dynamic, Policy +from policyengine_api.data.v2.models.base import TimestampedModel, enum_type if TYPE_CHECKING: from policyengine_api.data.v2.models.associations import UserHouseholdAssociation from policyengine_api.data.v2.models.household_mappings import ( LegacyHouseholdMapping, ) - from policyengine_api.data.v2.models.reports import Report - from policyengine_api.data.v2.models.simulations import Simulation class HouseholdJobStatus(StrEnum): - """Lifecycle state of a queued household calculation job.""" + """Lifecycle state of a household calculation job.""" PENDING = "pending" RUNNING = "running" @@ -37,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" @@ -81,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", ) @@ -96,33 +89,44 @@ 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 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. """ __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, - 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) 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( @@ -133,6 +137,3 @@ class HouseholdJob(IdentifiedModel, 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 b53e42a7b..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,15 +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"}, - ) - output_simulations: list["Simulation"] = Relationship( - back_populates="output_dataset", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, - ) - reports: list["Report"] = Relationship(back_populates="dataset") class DatasetVersion(IdentifiedModel, table=True): diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 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 5c73947c1..71bfbc931 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,4 +1,11 @@ -"""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; 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 from enum import StrEnum @@ -13,15 +20,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 -from policyengine_api.data.v2.models.users import User +from policyengine_api.data.v2.models.simulations import Simulation, SimulationRun class ReportRunStatus(StrEnum): @@ -65,8 +64,22 @@ 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.""" + """Stable report definition; execution state belongs to ReportRun. + + Its user, model, policy, household, dataset and region are logical + references to the Public API's rows, without foreign keys: the simulation + service checks them before it writes the report, except ``user_id``, the + owner the Public API names in its request, which the Public API checks + because the simulation service can't read users. + """ __tablename__ = "reports" __table_args__ = ( @@ -81,21 +94,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", @@ -106,27 +107,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"}, @@ -135,9 +121,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, @@ -145,7 +128,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_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" __table_args__ = ( @@ -167,6 +155,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 +173,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), @@ -196,6 +186,14 @@ class ReportRun(TimestampedModel, table=True): markdown: str | None = Field(default=None, sa_type=sa.Text) report: Report = Relationship(back_populates="runs") + simulation_run_members: list["ReportRunSimulationRun"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) + comparisons: list["ReportRunComparison"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) aggregates: list["AggregateOutput"] = Relationship( back_populates="report_run", cascade_delete=True, @@ -219,10 +217,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 +225,87 @@ 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_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): """A sum, mean or count of one variable in one simulation of a report run. @@ -325,13 +390,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", @@ -342,12 +409,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) @@ -362,19 +429,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( @@ -387,15 +461,16 @@ 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) + decile_variable: str | None = Field(default=None, max_length=512) entity: str = Field(max_length=128) decile: int quantiles: int = 10 @@ -413,13 +488,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", @@ -434,12 +512,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 +535,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,53 +558,65 @@ 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, filter variable and filter group: + ``filter_variable`` selects people (``age``) and ``filter_group`` names the + group (``child``). Missing values count as equal, so each run has one + unfiltered row per type and entity. + """ __tablename__ = "poverty" __table_args__ = ( - sa.Index( - "ix_poverty_run_simulation_type", - "report_run_id", - "simulation_id", + sa.UniqueConstraint( + "simulation_run_id", "poverty_type", + "entity", + "filter_variable", + "filter_group", + 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) + filter_group: str | None = Field(default=None, max_length=64) headcount: float | None = None 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): """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 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. """ __tablename__ = "program_statistics" __table_args__ = ( sa.UniqueConstraint( "report_run_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", "program_name", "entity", name="uq_program_statistics_run_program_entity", + postgresql_nulls_not_distinct=True, ), ) @@ -539,17 +624,17 @@ 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) - 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 @@ -568,12 +653,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 @@ -581,16 +666,41 @@ 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 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", ), @@ -599,25 +709,31 @@ 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. + 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", ), + _disclosure_tier_check("constituency_impacts"), ) constituency_code: str = Field(max_length=64) @@ -626,19 +742,22 @@ 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. + 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", ), + _disclosure_tier_check("local_authority_impacts"), ) local_authority_code: str = Field(max_length=64) diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index a6821eadc..5a2e31322 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 @@ -9,20 +9,18 @@ from sqlmodel import Field, Relationship from policyengine_api.data.v2.models.base import TimestampedModel, enum_type -from policyengine_api.data.v2.models.households import Household -from policyengine_api.data.v2.models.metadata import ( - Dataset, - Region, - TaxBenefitModelVersion, -) -from policyengine_api.data.v2.models.policies import Dynamic, Policy if TYPE_CHECKING: - from policyengine_api.data.v2.models.reports import Report + 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 +38,10 @@ 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``. + 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" @@ -61,57 +61,90 @@ 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( 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", + 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 + + baseline_reports: list["Report"] = Relationship( + back_populates="baseline_simulation", + sa_relationship_kwargs={"foreign_keys": "Report.baseline_simulation_id"}, ) - dynamic_id: UUID | None = Field( - default=None, - foreign_key="dynamics.id", - ondelete="SET NULL", + reform_reports: list["Report"] = Relationship( + back_populates="reform_simulation", + sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, ) - tax_benefit_model_version_id: UUID = Field( - foreign_key="tax_benefit_model_versions.id", - ondelete="RESTRICT", - index=True, + runs: list["SimulationRun"] = Relationship( + back_populates="simulation", + cascade_delete=True, ) - output_dataset_id: UUID | None = Field( - default=None, - foreign_key="datasets.id", - ondelete="SET NULL", + + +class SimulationRun(TimestampedModel, table=True): + """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 + 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"), ) - region_id: UUID | None = Field( - default=None, - foreign_key="regions.id", - ondelete="SET NULL", + + 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 - 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), @@ -121,27 +154,18 @@ class Simulation(TimestampedModel, table=True): 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) - dataset: Dataset | None = Relationship( - back_populates="input_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.dataset_id"}, + simulation: Simulation = Relationship(back_populates="runs") + report_run_memberships: list["ReportRunSimulationRun"] = Relationship( + back_populates="simulation_run" ) - output_dataset: Dataset | None = Relationship( - back_populates="output_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, + poverty_results: list["Poverty"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, ) - household: Household | None = Relationship(back_populates="simulations") - policy: Policy | None = Relationship(back_populates="simulations") - dynamic: Dynamic | None = Relationship(back_populates="simulations") - tax_benefit_model_version: TaxBenefitModelVersion = Relationship( - back_populates="simulations" - ) - region: Region | None = Relationship(back_populates="simulations") - baseline_reports: list["Report"] = Relationship( - back_populates="baseline_simulation", - sa_relationship_kwargs={"foreign_keys": "Report.baseline_simulation_id"}, - ) - reform_reports: list["Report"] = Relationship( - back_populates="reform_simulation", - sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, + inequalities: list["Inequality"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, ) 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/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index d9d0cacec..c449c93cf 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,13 +412,26 @@ 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]]], ) -> 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; + leave out the roles ``skipped_roles`` returns to skip them. """ existing_roles = read_existing_roles(connection, list(declared)) existing_tables = read_public_tables(connection) @@ -438,6 +480,47 @@ 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 missing roles awaiting creation. + + A missing role in ``roles_awaiting_creation`` is skipped; any other missing + role counts as missing every privilege declared for it. The roles skipped + and the privileges missing are written to stdout and the job summary under + one heading. + """ + 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) @@ -457,31 +540,20 @@ 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) - 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: 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 7bf2897c8..2dcef8ff8 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -24,6 +24,18 @@ (``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, 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. 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. The migration role (``policyengine_v2_migrator``) owns every table, and Supabase's ``postgres`` role administers the project; neither is managed here. @@ -36,14 +48,24 @@ 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 +# 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( + {SIMULATION_SERVICE_ROLE, MODAL_WORKER_ROLE} +) ENVIRONMENT_PATTERN: Final = re.compile(r"^[a-z][a-z0-9_]{0,30}$") @@ -59,6 +81,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} ) @@ -79,6 +102,8 @@ class TableAccess: runtime: frozenset[TablePrivilege] catalog_writer: frozenset[TablePrivilege] stage12: frozenset[TablePrivilege] + simulation_service: frozenset[TablePrivilege] + modal_worker: frozenset[TablePrivilege] def _access( @@ -86,33 +111,65 @@ 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, dynamics, 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_runs": _access(catalog_writer=READ), - "reports": _access(catalog_writer=READ), - "simulations": _access(catalog_writer=READ), + # 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, + "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. + "report_run_comparisons": _RESULT, + "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. @@ -120,20 +177,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), @@ -183,4 +248,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/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index 14685f150..1b1debd61 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -1,4 +1,10 @@ -"""Durable report-definition and report-run operations for API v2-alpha.""" +"""Durable report-definition and report-run operations for API v2-alpha. + +``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/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/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/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 new file mode 100644 index 000000000..d941d43ee --- /dev/null +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -0,0 +1,541 @@ +"""Check the simulation runs revisions against disposable Postgres. + +Five revisions follow the parameter tree revision. The first adds +simulation runs and the table naming the simulation runs a report run +combines, keys every result table by run, and gives household jobs one +household, one year and one policy. The second drops the foreign keys from +those tables to the Public API's tables. The third adds the economy output's +fields that had no column. The fourth makes each result row tell apart its +pair of simulation runs, decile grouping and poverty group. The fifth moves +the parts of a comparison without a result table from the report run to one +row per compared pair. None changes any other table. The previous Public API +revision keeps working while they are deployed: its roles can't read or write +any table they change, except the catalog writer, which counts the rows of +simulations, reports and report_runs. +""" + +from __future__ import annotations + +import os +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 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 +from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection, Engine + +SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" +SIMULATION_RUNS_REVISION = "2ffd50f47a2b" +CROSS_SERVICE_KEYS_REVISION = "7a83a87d029a" +ECONOMY_OUTPUT_REVISION = "7f43d2c7c3c6" +RESULT_IDENTITY_REVISION = "2e3de7c572ed" +NEW_TABLES = frozenset( + {"simulation_runs", "report_run_simulation_runs", "report_run_comparisons"} +) +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", + "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 +COMPLETED_AT = "2026-11-04T00:00:00Z" + + +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() + + +@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) + 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 _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_three_tables_and_change_only_simulation_side_tables( + config: Config, engine: Engine +) -> None: + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + before = _shapes(engine) + command.upgrade(config, "head") + after = _shapes(engine) + + assert set(after) - set(before) == NEW_TABLES + assert set(before) <= set(after) + changed = { + table_name for table_name in before if before[table_name] != after[table_name] + } + assert changed == CHANGED_TABLES + + +def test_each_revision_downgrades_to_its_parent_and_upgrades_again( + config: Config, engine: Engine +) -> None: + command.upgrade(config, "head") + + command.downgrade(config, RESULT_IDENTITY_REVISION) + assert "report_run_comparisons" not in inspect(engine).get_table_names() + assert {"labor_supply_result", "data_version"} <= _column_names( + engine, "report_runs" + ) + + command.downgrade(config, ECONOMY_OUTPUT_REVISION) + assert "decile_variable" not in _column_names(engine, "decile_impacts") + assert "filter_group" not in _column_names(engine, "poverty") + budget_key = next( + constraint + for constraint in inspect(engine).get_unique_constraints("budget_summary") + if constraint["name"] == "uq_budget_summary_run_variable_entity" + ) + assert budget_key["column_names"] == ["report_run_id", "variable_name", "entity"] + + command.downgrade(config, CROSS_SERVICE_KEYS_REVISION) + assert {"spm_provenance", "data_version"}.isdisjoint( + _column_names(engine, "report_runs") + ) + 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: + 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 == [] + + +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_row( + 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, + *, + status: str = "pending", + **values: object, +) -> object: + run_id = uuid4() + insert_row( + connection, + "simulation_runs", + id=run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=uuid4(), + status=status, + **RUN_VERSIONS, + **values, + ) + return run_id + + +def _insert_report_run(connection: Connection) -> object: + report_id, report_run_id = uuid4(), uuid4() + insert_row( + connection, + "reports", + id=report_id, + label="Report", + country="us", + tax_benefit_model_id=uuid4(), + inputs="{}", + ) + insert_row( + 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: + 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_row( + connection, + "poverty", + id=uuid4(), + simulation_run_id=finished_run_id, + poverty_type="spm", + entity="person", + ) + + 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_row(connection, "report_run_simulation_runs", id=uuid4(), **member) + + 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} + ) + 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_row( + 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_row(connection, "poverty", id=uuid4(), **poverty) + with refused(connection): + insert_row(connection, "poverty", id=uuid4(), **poverty) + for group in ("child", "adult"): + insert_row( + 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_row(connection, "decile_impacts", id=uuid4(), **decile) + insert_row( + connection, + "decile_impacts", + id=uuid4(), + decile_variable="household_wealth_decile", + **decile, + ) + with refused(connection): + insert_row(connection, "decile_impacts", id=uuid4(), **decile) + budget = {"variable_name": "household_tax", "entity": "household"} + insert_row(connection, "budget_summary", id=uuid4(), **budget, **comparison_runs) + next_year_runs = { + **comparison_runs, + "baseline_simulation_run_id": _insert_simulation_run(connection, simulation_id), + "reform_simulation_run_id": _insert_simulation_run(connection, simulation_id), + } + insert_row(connection, "budget_summary", id=uuid4(), **budget, **next_year_runs) + 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_row( + 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_row( + connection, + "constituency_impacts", + id=uuid4(), + constituency_code="E14000530", + disclosure_tier="approximate", + **constituency, + ) + with refused(connection): + insert_row( + 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_row( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="pending", + **household_job, + ) + + with refused(connection): + insert_row( + connection, + "household_jobs", + id=uuid4(), + year=1800, + status="pending", + **household_job, + ) + with refused(connection): + insert_row( + connection, + "household_jobs", + id=uuid4(), + year=2026, + status="failed", + **household_job, + ) diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index 0875d2864..aad982354 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() @@ -290,7 +295,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 +320,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 +336,24 @@ 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] + # 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 f720c3fb9..3edc1cdd1 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, @@ -22,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 @@ -137,3 +139,102 @@ 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} + ) + skipped = skipped_roles(connection, declared, {absent_role}) + missing = check_declared_privileges( + 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") + + +def test_the_deploy_steps_pass_and_name_a_role_awaiting_creation( + 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() + 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}), + ) + 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) + 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 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 ( + "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 new file mode 100644 index 000000000..c47d5c0f6 --- /dev/null +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -0,0 +1,247 @@ +"""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 +from tests.integration.simulation_rows import RUN_VERSIONS, insert_row, refused + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + + +@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 _refused(connection: Connection, statement: str, **values: object) -> None: + """Expect a permission error from ``statement``.""" + with refused(connection, ProgrammingError, match="permission denied"): + connection.execute(text(statement), values) + + +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", "household_job") + } + insert_row( + connection, + "simulations", + id=ids["simulation"], + simulation_type="economy", + dataset_id=uuid4(), + tax_benefit_model_version_id=uuid4(), + year=2026, + ) + insert_row( + connection, + "simulation_runs", + id=ids["run"], + simulation_id=ids["simulation"], + tax_benefit_model_version_id=uuid4(), + status="pending", + **RUN_VERSIONS, + ) + insert_row( + connection, + "reports", + id=ids["report"], + label="Report", + country="us", + tax_benefit_model_id=uuid4(), + inputs="{}", + ) + insert_row( + connection, + "report_runs", + id=ids["report_run"], + report_id=ids["report"], + status="pending", + trigger="initial", + **RUN_VERSIONS, + ) + insert_row( + connection, + "report_run_simulation_runs", + id=uuid4(), + report_run_id=ids["report_run"], + simulation_run_id=ids["run"], + role="baseline", + ) + insert_row( + connection, + "household_jobs", + id=ids["household_job"], + 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, + "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") + + +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_row( + 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"]}, + ) + 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="{}", + ) + finished_job = connection.execute( + text( + "UPDATE household_jobs SET status = 'succeeded', result = '{}', " + "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, " + "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_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 2099bb908..1171a9cb3 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -47,6 +47,9 @@ "legacy_user_policy_mappings", "parameter_search_files", "parameter_tree_children", + "report_run_comparisons", + "report_run_simulation_runs", + "simulation_runs", "stage12_evaluation_reports", "stage12_evaluation_simulations", "user_identities", @@ -62,6 +65,56 @@ 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" +) +# The simulation runs revision, parent of the cross-service foreign key revision. +CROSS_SERVICE_KEYS_PARENT_REVISION = "2ffd50f47a2b" +# The cross-service foreign key revision, parent of the economy output revision. +ECONOMY_OUTPUT_PARENT_REVISION = "7a83a87d029a" +# The economy output revision, parent of the result row identity revision. +RESULT_IDENTITY_PARENT_REVISION = "7f43d2c7c3c6" +# The result row identity revision, parent of the comparisons revision. +COMPARISONS_PARENT_REVISION = "2e3de7c572ed" +UNTABLED_COMPARISON_COLUMNS = ( + "labor_supply_result", + "cliff_impact_result", + "spm_provenance", + "data_version", +) +# Each comparison result table's unique key, which gains the pair of runs. +COMPARISON_UNIQUE_KEYS = { + "budget_summary": "uq_budget_summary_run_variable_entity", + "congressional_district_impacts": "uq_congressional_district_impacts_run_geoid", + "constituency_impacts": "uq_constituency_impacts_run_code", + "decile_impacts": "uq_decile_impacts_run_measure", + "intra_decile_impacts": "uq_intra_decile_impacts_run_type_decile", + "local_authority_impacts": "uq_local_authority_impacts_run_code", + "program_statistics": "uq_program_statistics_run_program_entity", +} +# The foreign keys from the simulation service's tables to the Public API's. +CROSS_SERVICE_KEYS = { + "household_jobs": {"dynamics", "policies"}, + "reports": { + "datasets", + "households", + "policies", + "regions", + "tax_benefit_models", + "users", + }, + "simulations": { + "datasets", + "dynamics", + "households", + "policies", + "regions", + "tax_benefit_model_versions", + }, +} DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -243,6 +296,11 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + COMPARISONS_PARENT_REVISION, + RESULT_IDENTITY_PARENT_REVISION, + ECONOMY_OUTPUT_PARENT_REVISION, + CROSS_SERVICE_KEYS_PARENT_REVISION, + SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, AREA_IMPACT_COORDINATE_REVISION, CATALOG_FORMAT_2_REVISION, @@ -989,7 +1047,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 +1067,256 @@ 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 + + +def test_cross_service_key_revision_only_drops_and_restores_foreign_keys() -> None: + """The revision after the simulation runs revision is generated. + + Its upgrade drops each foreign key from a table the simulation service + writes to a table the Public API writes; its downgrade restores them. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == CROSS_SERVICE_KEYS_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + dropped_tables = re.findall( + r'op\.drop_constraint\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*' + r'type_="foreignkey"', + upgrade, + ) + restored_keys = re.findall( + r'op\.create_foreign_key\(\s*op\.f\("[^"]+"\),\s*"([a-z_]+)",\s*"([a-z_]+)"', + downgrade, + ) + expected_count = sum(len(targets) for targets in CROSS_SERVICE_KEYS.values()) + assert len(dropped_tables) == upgrade.count("op.drop_constraint(") + assert len(dropped_tables) == expected_count + assert {table: dropped_tables.count(table) for table in CROSS_SERVICE_KEYS} == { + table: len(targets) for table, targets in CROSS_SERVICE_KEYS.items() + } + assert len(restored_keys) == downgrade.count("op.create_foreign_key(") + restored_targets: dict[str, set[str]] = {} + for table, target in restored_keys: + restored_targets.setdefault(table, set()).add(target) + assert restored_targets == CROSS_SERVICE_KEYS + for body in (upgrade, downgrade): + for operation in ( + "op.add_column(", + "op.drop_column(", + "op.create_table(", + "op.drop_table(", + "op.execute(", + "op.bulk_insert(", + ): + assert operation not in body + + +def test_economy_output_revision_only_adds_columns_and_relaxes_two() -> None: + """The revision after the cross-service key revision is generated. + + It adds the economy output's fields that had no column and lets a detailed + budget row omit its entity and tax flag; it removes nothing. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == ECONOMY_OUTPUT_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + added_columns = set( + re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) + ) + uk_disclosure_columns = { + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", + } + assert added_columns == { + ("report_runs", "spm_provenance"), + ("report_runs", "data_version"), + ("congressional_district_impacts", "winner_percentage"), + ("congressional_district_impacts", "loser_percentage"), + ("congressional_district_impacts", "no_change_percentage"), + *(("constituency_impacts", column) for column in uk_disclosure_columns), + *(("local_authority_impacts", column) for column in uk_disclosure_columns), + } + assert upgrade.count("op.add_column(") == len(added_columns) + relaxed_columns = re.findall( + r'op\.alter_column\(\s*"(\w+)",\s*"(\w+)",\s*' + r"existing_type=sa\.\w+\([^)]*\),\s*nullable=True", + upgrade, + ) + assert upgrade.count("op.alter_column(") == 2 + assert set(relaxed_columns) == { + ("program_statistics", "entity"), + ("program_statistics", "is_tax"), + } + assert downgrade.count("op.drop_column(") == len(added_columns) + for body in (upgrade, downgrade): + assert "op.drop_table(" not in body + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body + assert "op.drop_column(" not in upgrade + + +def test_result_identity_revision_adds_two_columns_and_widens_unique_keys() -> None: + """The revision after the economy output revision is generated. + + It adds ``decile_impacts.decile_variable`` and ``poverty.filter_group`` and + rebuilds each comparison result table's unique key with the pair of + simulation runs, and ``poverty``'s with its group. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == RESULT_IDENTITY_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + added_columns = set( + re.findall(r'op\.add_column\(\s*"(\w+)",\s*sa\.Column\(\s*"(\w+)"', upgrade) + ) + assert added_columns == { + ("decile_impacts", "decile_variable"), + ("poverty", "filter_group"), + } + rebuilt_keys = { + table: name + for name, table in re.findall( + r'op\.create_unique_constraint\(\s*"(\w+)",\s*"(\w+)"', upgrade + ) + } + assert rebuilt_keys == { + **COMPARISON_UNIQUE_KEYS, + "poverty": "uq_poverty_simulation_run_measure", + } + assert upgrade.count("op.drop_constraint(") == len(rebuilt_keys) + assert upgrade.count('"baseline_simulation_run_id",') == len(COMPARISON_UNIQUE_KEYS) + assert downgrade.count("op.create_unique_constraint(") == len(rebuilt_keys) + assert downgrade.count("op.drop_column(") == len(added_columns) + for body in (upgrade, downgrade): + assert "op.drop_table(" not in body + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body + assert "op.drop_column(" not in upgrade + + +def test_comparisons_revision_moves_untabled_results_to_one_row_per_pair() -> None: + """The revision after the result row identity revision is generated. + + It creates ``report_run_comparisons`` and drops the four columns it takes + over from ``report_runs``; the downgrade puts them back. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == COMPARISONS_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert "No post-generation corrections." in revision + assert 'op.create_table(\n "report_run_comparisons"' in upgrade + assert "uq_report_run_comparisons_run_pair" in upgrade + assert upgrade.count("op.drop_column(") == len(UNTABLED_COMPARISON_COLUMNS) + assert downgrade.count("op.add_column(") == len(UNTABLED_COMPARISON_COLUMNS) + for column_name in UNTABLED_COMPARISON_COLUMNS: + assert f'op.drop_column("report_runs", "{column_name}")' in upgrade + assert f'"{column_name}"' in downgrade + assert 'op.drop_table("report_run_comparisons")' in downgrade + for body in (upgrade, downgrade): + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body 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 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..976ebe7e6 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", @@ -432,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: @@ -500,11 +505,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_privilege_manifest.py b/tests/unit/v2/test_privilege_manifest.py index 7a3a76b44..f38a659d4 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,49 @@ 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, dynamic, household and catalog rows a run needs. +RUN_INPUT_TABLES = frozenset( + { + "policies", + "dynamics", + "households", + "tax_benefit_models", + "tax_benefit_model_versions", + "regions", + "datasets", + "parameters", + "parameter_values", + } +) +RESULT_TABLES = frozenset( + { + "report_run_comparisons", + "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 +167,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). 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=( 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 new file mode 100644 index 000000000..42027a4c1 --- /dev/null +++ b/tests/unit/v2/test_simulation_run_models.py @@ -0,0 +1,485 @@ +"""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 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 sqlalchemy.dialects import postgresql +from sqlalchemy.orm import configure_mappers +from sqlalchemy.schema import CreateTable + +from policyengine_api.data.v2.models import ( + V2_METADATA, + ReportRun, + ReportRunComparison, + ReportRunSimulationRun, + SimulationRole, + SimulationRun, + User, +) + +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", +) +# 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", + "report_run_comparisons", + "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: + (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() -> 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 + # A budget-window run compares one pair per year, so the parts of a + # comparison without a result table belong to its comparison row. + assert { + "labor_supply_result", + "cliff_impact_result", + "spm_provenance", + "data_version", + }.isdisjoint(report_runs.c.keys()) + + +def test_each_comparison_keeps_the_parts_without_a_result_table() -> None: + comparisons = V2_METADATA.tables["report_run_comparisons"] + + for column_name in ("labor_supply_result", "cliff_impact_result", "spm_provenance"): + assert comparisons.c[column_name].nullable + assert isinstance(comparisons.c[column_name].type, sa.JSON) + assert comparisons.c.data_version.nullable + assert comparisons.c.data_version.type.length == 128 + report_run_key = _foreign_key(comparisons, "report_run_id") + assert report_run_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: + 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, baseline_simulation_run_id, " + "reform_simulation_run_id, program_name, entity)" + ) in statement + + +def test_report_runs_name_the_simulation_runs_they_combine() -> None: + members = V2_METADATA.tables["report_run_simulation_runs"] + + 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", + "filter_group", + ) + } + + +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, filter_group)" + ) 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" + # 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 + + +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" + ) + + +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 == [] + + +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