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 37b94289a..a45a660a2 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.github/workflows/v2-integration-check.yml @@ -79,8 +79,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..88dcb1697 --- /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. 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 9d92a76d9..79d780554 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -255,9 +255,15 @@ and a revision runs once per database, so it can't repair drift. | `policyengine_v2_runtime` | The Public API on Cloud Run | | `policyengine_v2_catalog_writer` | `scripts/publish_v2_metadata_catalog.py` | | `policyengine_stage12_` | The simulation API and its Modal worker, on the temporary comparison-run tables | +| `policyengine_v2_simulation_service` | The simulation API on Cloud Run: creates simulations, reports, their runs, budget-window runs and household jobs, and reads the policy, dynamic, household and catalog rows a run needs. It updates only the status, Modal call ID, error and completion columns of the runs and household jobs whose Modal calls it starts | +| `policyengine_v2_modal_worker` | The simulation API's Modal workers: write their runs' and household jobs' status and results, and read the rows a run needs. It updates only a run's status, Modal call ID, error, time and result columns; the workers' code, not the database, keeps a worker to its own run | Every v2 table has an entry for every role, so "no access" is a decision, not -an omission. The migration role owns every table and isn't listed. Three rules +an omission. A privilege is declared on the whole table, or with +`update_columns(...)` on only the columns it names; PostgreSQL then refuses a +statement that sets any other column, so neither simulation API role can +change what a run is. A role never holds both kinds of `UPDATE` on one table: +the whole-table one would let it set every column. The migration role owns every table and isn't listed. Three rules check the declarations: - `api.v2.runtime-privileges-declared`: a unit test reads each connector @@ -267,16 +273,29 @@ 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. -`apply` grants each declared privilege a role lacks. It never revokes: it -reports privileges the file doesn't declare, and the projects' default -privileges, in the job summary. It grants nothing when a declared role or -table is missing. The roles themselves are provisioned outside the -repository; `--create-missing-roles` is accepted only for the disposable test -database. +`apply` grants each declared privilege a role lacks, as one `GRANT` per role +and table and one per column privilege. It never revokes: it reports +privileges the file doesn't declare, on whole tables or on columns, and the +projects' default privileges, in the job summary. It grants nothing when a +declared role or table is missing. `verify` checks each column of a column +privilege with `has_column_privilege`, and fails when a role holds the same +privilege on the whole table. The roles themselves are provisioned outside +the repository; `--create-missing-roles` is accepted only for the disposable +test database. + +A pull request may declare a role before an operator creates it by listing it +in `ROLES_AWAITING_PROVISIONING` in `privileges.py`. While a listed role is +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 @@ -300,6 +319,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 ``` @@ -483,3 +503,98 @@ 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 `1356f658347d` (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`). +- A report run compares one pair of simulation runs: `report_runs` names its + `baseline_simulation_run_id` (required) and `reform_simulation_run_id` + (empty for a household report with only a baseline), each a foreign key to + `simulation_runs` with `RESTRICT`. The parts of an economy comparison that + had no result table are typed: `report_runs` gains the baseline's and + reform's cliff gap and share (four nullable numbers, checked to be set + together), the SPM receipt (`spm_provenance`, `jsonb`) and the data release + (`data_version`). The new `labor_supply_responses` holds one row per report + run (the earnings changes from the income and substitution effects and their + total, the revenue change, the two relative responses, and baseline, reform + and changed weekly hours split into the two effects), and the new + `labor_supply_decile_responses` one row per report run and income decile. +- `poverty` and `inequality` belong to a simulation run (`simulation_run_id`); + the other result tables belong to a report run and name the two simulation + runs they compare (`baseline_simulation_run_id`, `reform_simulation_run_id`); + each of their rows is unique within its report run. `aggregates` and + `change_aggregates` are unchanged. +- The result tables gain the economy output's fields that had no column: the + winner, loser and no-change shares of `congressional_district_impacts`, and + the UK disclosure control fields (`disclosure_tier`, checked to be + `reported` or `approximate`, and the four `_lower`/`_upper` interval bounds) + of `constituency_impacts` and `local_authority_impacts`. The detailed budget + names neither a program's entity nor whether it is a tax, so + `program_statistics.entity` and `is_tax` become nullable. + `decile_impacts.decile_variable` names the variable that groups households + when it is not the income variable's own deciles (for example + `household_wealth_decile`), and `poverty.filter_group` names the group a + filter selects (`child` for `age`); both join their table's unique key. +- `report_runs` gains its Modal call ID. +- Budget windows: a simulation run calculates one `year`, which joins the + index for a simulation's newest successful run, so a budget window compares + a report's same two simulations in every year. `reports.budget_window_size` + (1 to 75 years, an economy report with a year, ending by 2200) marks a + budget-window report. The new `budget_window_runs` holds one row per + calculation of one: its report, catalog version, trigger, whether its years + calculate their simulations again (`recalculate`), status, error and times. + `report_runs.budget_window_run_id` and `year` name the window run and year a + report run calculates, at most one report run per year of a window run; a + composite foreign key keeps it to a window run of the same report, and + deletes it with the window run (`CASCADE`), as deleting a report deletes its + runs. +- `household_jobs` calculates one household (`household_id`) for one `year` + under one policy, and records its catalog and package versions and its + Modal call ID. +- The new columns that name the Public API's rows are foreign keys with the + delete rules the same columns have on `simulations`: + `simulation_runs.tax_benefit_model_version_id` and `dataset_id` (`RESTRICT`) + and `region_id` (`SET NULL`), `household_jobs.household_id` and + `tax_benefit_model_version_id` (`RESTRICT`), and + `budget_window_runs.tax_benefit_model_version_id` (`RESTRICT`), 20 keys + between the two services' tables in all. `reports.user_id` gains an + index, which account deletion's `SET NULL` searches. Each service writes only + its own tables, and the foreign keys between them stay: PostgreSQL checks + each named row when a row is written, with the privileges of the named + table's owner, so a writing role needs no privilege on the named table. + +Missing values in the unique keys of `poverty`, `decile_impacts` and +`program_statistics` count as equal (`NULLS NOT DISTINCT`, which PostgreSQL +supports from version 15). + +It has three post-generation corrections of one kind, each commented in the +revision: `simulation_runs.status` reuses the `v2_simulation_status` type the +baseline created for `simulations.status`, and `budget_window_runs.status` and +`trigger` reuse the baseline's `v2_report_run_status` and +`v2_report_run_trigger`. The generated `create_table` would create each type +again and fail, so each column uses `postgresql.ENUM(..., create_type=False)`. + +The upgrade adds required columns without server defaults to `household_jobs`, +`report_runs` and the result tables, so it needs them empty, as they are while no managed +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 +of `simulations` and the result tables the same way, so it needs +`simulations` empty, which leaves no run or result row either; it drops +`household_jobs`' new columns with their values. The previous Public API revision keeps working at this revision: +none of its roles can read or write a table the revision changes, except the +catalog writer, which only counts the rows of `simulations`, `reports` and +`report_runs`. +`tests/integration/test_alembic_v2_simulation_runs.py` checks that the +revision changes no other table, and downgrades and upgrades it again. + +The name `simulation_runs` is also a v1 Cloud SQL table. The v2 metadata +validation (`policyengine_api/data/v2/metadata_validation.py`) now treats it as +a name the two databases share, as it does `simulations`. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 3a9edb3b9..ebe284c0a 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -285,9 +285,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.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.deployed-privileges-match` | Every v2 deploy grants the managed roles each privilege policyengine_api/data/v2/privileges.py declares, on the whole table or on the columns it names, right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing or a role holds on a whole table a privilege declared for only some of its columns. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead. | mechanical | enforced | yes | CI job `seed`, `catalog-runtime-integration` | `seed-v2-database.yml:seed`, `v2-integration-check.yml:catalog-runtime-integration` | yes | +| `api.v2.runtime-privileges-declared` | Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py; a privilege declared on only some of a table's columns names columns of that table and is never also declared on the whole table; and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it, on the whole table. | mechanical | enforced | yes | pytest `tests/unit/v2/test_privilege_manifest.py` | `pr.yml:test` | yes | +| `api.v2.runtime-privileges-tested-as-runtime-role` | The cross-database integration tests run the v2 user services and the parameter tree read as policyengine_v2_runtime, catalog publication as policyengine_v2_catalog_writer, the comparison-run services as policyengine_stage12_, and the simulation API's table writes as policyengine_v2_simulation_service and policyengine_v2_modal_worker, holding only their declared privileges, so a statement those privileges don't allow fails CI. | mechanical | enforced | yes | CI job `v2-integration-check`, `catalog-runtime-integration` | `pr.yml:v2-integration-check`, `v2-integration-check.yml:catalog-runtime-integration` | — | ```mermaid flowchart LR diff --git a/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py new file mode 100644 index 000000000..c118ea5d0 --- /dev/null +++ b/migrations/v2/versions/1356f658347d_record_simulation_runs_and_key_results_.py @@ -0,0 +1,1514 @@ +"""record simulation runs and key results by run + +Revision ID: 1356f658347d +Revises: c614b6bcccfc +Create Date: 2026-11-06 00:33:17.526471 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate + +Three post-generation corrections of one kind, each commented where it is +made: simulation_runs.status reuses the existing v2_simulation_status type, +and budget_window_runs.status and trigger reuse the existing +v2_report_run_status and v2_report_run_trigger types. + +The upgrade adds required columns without server defaults to household_jobs, +report_runs and the result tables, so it needs those tables empty, as they are +while no managed role can write them. The downgrade restores required columns +of simulations and the result tables the same way, so it needs simulations +empty, which leaves no run or result row either; it drops household_jobs' new +columns with their values. +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel +from sqlalchemy.dialects import postgresql + +revision: str = "1356f658347d" +down_revision: Union[str, None] = "c614b6bcccfc" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +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("year", sa.Integer(), 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.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_simulation_runs_year") + ), + sa.ForeignKeyConstraint( + ["dataset_id"], + ["datasets.id"], + name=op.f("fk_simulation_runs_dataset_id_datasets"), + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["region_id"], + ["regions.id"], + name=op.f("fk_simulation_runs_region_id_regions"), + ondelete="SET NULL", + ), + sa.ForeignKeyConstraint( + ["simulation_id"], + ["simulations.id"], + name=op.f("fk_simulation_runs_simulation_id_simulations"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + ["tax_benefit_model_versions.id"], + name=op.f( + "fk_simulation_runs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_simulation_runs")), + ) + op.create_index( + "ix_simulation_runs_current_output", + "simulation_runs", + [ + "simulation_id", + "year", + "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_index( + op.f("ix_simulation_runs_tax_benefit_model_version_id"), + "simulation_runs", + ["tax_benefit_model_version_id"], + unique=False, + ) + op.create_table( + "budget_window_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_id", sa.Uuid(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column( + "trigger", + # Dialect correction: a budget-window run starts for the same reasons + # as a report run, so the column reuses the baseline's + # v2_report_run_trigger type instead of creating it again. + postgresql.ENUM( + "initial", + "manual", + "system", + name="v2_report_run_trigger", + create_type=False, + ), + nullable=False, + ), + sa.Column("recalculate", sa.Boolean(), nullable=False), + sa.Column( + "status", + # Dialect correction: a budget-window run has the same states as a + # report run, so the column reuses the baseline's v2_report_run_status + # type instead of creating it again. + postgresql.ENUM( + "pending", + "running", + "succeeded", + "failed", + name="v2_report_run_status", + create_type=False, + ), + nullable=False, + ), + sa.Column("error_message", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("started_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name=op.f("ck_budget_window_runs_terminal_completion"), + ), + sa.ForeignKeyConstraint( + ["report_id"], + ["reports.id"], + name=op.f("fk_budget_window_runs_report_id_reports"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + ["tax_benefit_model_versions.id"], + name=op.f( + "fk_budget_window_runs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_budget_window_runs")), + sa.UniqueConstraint("id", "report_id", name="uq_budget_window_runs_id_report"), + ) + op.create_index( + "ix_budget_window_runs_report_id_created_at", + "budget_window_runs", + ["report_id", "created_at"], + unique=False, + ) + op.create_index( + op.f("ix_budget_window_runs_tax_benefit_model_version_id"), + "budget_window_runs", + ["tax_benefit_model_version_id"], + unique=False, + ) + op.create_table( + "labor_supply_decile_responses", + sa.Column("id", sa.Uuid(), nullable=False), + 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("decile", sa.Integer(), nullable=False), + sa.Column("average_income_effect", sa.Float(), nullable=True), + sa.Column("average_substitution_effect", sa.Float(), nullable=True), + sa.Column("relative_income_effect", sa.Float(), nullable=True), + sa.Column("relative_substitution_effect", sa.Float(), nullable=True), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_labor_supply_decile_responses_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_labor_supply_decile_responses")), + sa.UniqueConstraint( + "report_run_id", + "decile", + name="uq_labor_supply_decile_responses_run_decile", + ), + ) + op.create_table( + "labor_supply_responses", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column("report_run_id", sa.Uuid(), nullable=False), + sa.Column("income_effect_earnings_change", sa.Float(), nullable=False), + sa.Column("substitution_effect_earnings_change", sa.Float(), nullable=False), + sa.Column("total_earnings_change", sa.Float(), nullable=False), + sa.Column("revenue_change", sa.Float(), nullable=False), + sa.Column("relative_income_effect", sa.Float(), nullable=False), + sa.Column("relative_substitution_effect", sa.Float(), nullable=False), + sa.Column("baseline_weekly_hours", sa.Float(), nullable=False), + sa.Column("reform_weekly_hours", sa.Float(), nullable=False), + sa.Column("weekly_hours_change", sa.Float(), nullable=False), + sa.Column("weekly_hours_income_effect", sa.Float(), nullable=False), + sa.Column("weekly_hours_substitution_effect", sa.Float(), nullable=False), + sa.ForeignKeyConstraint( + ["report_run_id"], + ["report_runs.id"], + name=op.f("fk_labor_supply_responses_report_run_id_report_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_labor_supply_responses")), + sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), + ) + op.add_column( + "budget_summary", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + 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.add_column( + "congressional_district_impacts", + sa.Column("winner_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("loser_percentage", sa.Float(), nullable=False), + ) + op.add_column( + "congressional_district_impacts", + sa.Column("no_change_percentage", sa.Float(), nullable=False), + ) + op.drop_constraint( + op.f("fk_congressional_district_impacts_reform_simulation_id__ca94"), + "congressional_district_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_congressional_district_impacts_baseline_simulation_i_d68d"), + "congressional_district_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f( + "fk_congressional_district_impacts_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.add_column( + "constituency_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "constituency_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) + op.drop_constraint( + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + "constituency_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_constituency_impacts_baseline_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.create_check_constraint( + op.f("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) + op.drop_column("constituency_impacts", "reform_simulation_id") + op.drop_column("constituency_impacts", "baseline_simulation_id") + op.add_column( + "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.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", + "income_variable", + "decile_variable", + "entity", + "decile", + "quantiles", + ], + postgresql_nulls_not_distinct=True, + ) + 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_baseline_simulation_run_id_simulation_runs"), + "decile_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + 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.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_index( + op.f("ix_household_jobs_tax_benefit_model_version_id"), + "household_jobs", + ["tax_benefit_model_version_id"], + unique=False, + ) + op.create_foreign_key( + op.f("fk_household_jobs_household_id_households"), + "household_jobs", + "households", + ["household_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + "household_jobs", + "tax_benefit_model_versions", + ["tax_benefit_model_version_id"], + ["id"], + ondelete="RESTRICT", + ) + 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_report_run_id_report_runs"), + "inequality", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_inequality_simulation_id_simulations"), + "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", "report_run_id") + op.drop_column("inequality", "simulation_id") + op.add_column( + "intra_decile_impacts", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + 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_reform_simulation_id_simulations"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_baseline_simulation_id_simulations"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + 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.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.add_column( + "local_authority_impacts", + sa.Column( + "disclosure_tier", + sqlmodel.sql.sqltypes.AutoString(length=16), + nullable=True, + ), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("average_household_income_change_upper", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_lower", sa.Float(), nullable=True), + ) + op.add_column( + "local_authority_impacts", + sa.Column("relative_household_income_change_upper", sa.Float(), nullable=True), + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_baseline_simulation_id_simulations"), + "local_authority_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + "local_authority_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + "local_authority_impacts", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + 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_check_constraint( + op.f("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + ) + op.drop_column("local_authority_impacts", "reform_simulation_id") + op.drop_column("local_authority_impacts", "baseline_simulation_id") + op.add_column("poverty", sa.Column("simulation_run_id", sa.Uuid(), nullable=False)) + op.add_column( + "poverty", + sa.Column( + "filter_group", sqlmodel.sql.sqltypes.AutoString(length=64), nullable=True + ), + ) + op.drop_index(op.f("ix_poverty_run_simulation_type"), table_name="poverty") + op.create_unique_constraint( + "uq_poverty_simulation_run_measure", + "poverty", + [ + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + ], + 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", "report_run_id") + op.drop_column("poverty", "simulation_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.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.drop_constraint( + op.f("fk_program_statistics_reform_simulation_id_simulations"), + "program_statistics", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_program_statistics_baseline_simulation_id_simulations"), + "program_statistics", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + "program_statistics", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + 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.drop_column("program_statistics", "reform_simulation_id") + op.drop_column("program_statistics", "baseline_simulation_id") + op.add_column( + "report_runs", + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + ) + op.add_column( + "report_runs", sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("budget_window_run_id", sa.Uuid(), nullable=True) + ) + op.add_column("report_runs", sa.Column("year", sa.Integer(), nullable=True)) + op.add_column( + "report_runs", + sa.Column( + "modal_call_id", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True + ), + ) + op.add_column( + "report_runs", sa.Column("baseline_cliff_gap", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("baseline_cliff_share", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("reform_cliff_gap", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", sa.Column("reform_cliff_share", sa.Float(), nullable=True) + ) + op.add_column( + "report_runs", + sa.Column( + "spm_provenance", + sa.JSON().with_variant( + postgresql.JSONB(astext_type=sa.Text()), "postgresql" + ), + nullable=True, + ), + ) + op.add_column( + "report_runs", + sa.Column( + "data_version", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True + ), + ) + op.create_index( + op.f("ix_report_runs_baseline_simulation_run_id"), + "report_runs", + ["baseline_simulation_run_id"], + unique=False, + ) + op.create_index( + op.f("ix_report_runs_reform_simulation_run_id"), + "report_runs", + ["reform_simulation_run_id"], + unique=False, + ) + op.create_index( + "ix_report_runs_status_started_at", + "report_runs", + ["status", "started_at"], + unique=False, + ) + op.create_unique_constraint( + "uq_report_runs_budget_window_run_id_year", + "report_runs", + ["budget_window_run_id", "year"], + ) + op.create_foreign_key( + "fk_report_runs_budget_window_run_id_budget_window_runs", + "report_runs", + "budget_window_runs", + ["budget_window_run_id", "report_id"], + ["id", "report_id"], + ondelete="CASCADE", + ) + op.create_foreign_key( + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + "report_runs", + "simulation_runs", + ["reform_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "report_runs", + "simulation_runs", + ["baseline_simulation_run_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_check_constraint( + op.f("ck_report_runs_budget_window_year"), + "report_runs", + "budget_window_run_id IS NULL OR year IS NOT NULL", + ) + op.create_check_constraint( + op.f("ck_report_runs_cliff_impact_complete"), + "report_runs", + "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) AND (baseline_cliff_gap IS NULL) = (reform_cliff_share IS NULL)", + ) + op.create_check_constraint( + op.f("ck_report_runs_year"), + "report_runs", + "year IS NULL OR year BETWEEN 1900 AND 2200", + ) + op.add_column( + "reports", sa.Column("budget_window_size", sa.Integer(), nullable=True) + ) + op.create_index(op.f("ix_reports_user_id"), "reports", ["user_id"], unique=False) + op.create_check_constraint( + op.f("ck_reports_budget_window_size"), + "reports", + "budget_window_size IS NULL OR (household_id IS NULL AND budget_window_size BETWEEN 1 AND 75 AND year IS NOT NULL AND year + budget_window_size - 1 <= 2200)", + ) + op.drop_index(op.f("ix_simulations_status_created_at"), table_name="simulations") + op.drop_constraint( + op.f("fk_simulations_output_dataset_id_datasets"), + "simulations", + type_="foreignkey", + ) + op.drop_column("simulations", "household_result") + op.drop_column("simulations", "started_at") + op.drop_column("simulations", "output_dataset_id") + op.drop_column("simulations", "status") + op.drop_column("simulations", "completed_at") + op.drop_column("simulations", "error_message") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "simulations", + sa.Column("error_message", sa.VARCHAR(), 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( + "status", + postgresql.ENUM( + "pending", "running", "succeeded", "failed", name="v2_simulation_status" + ), + autoincrement=False, + nullable=False, + ), + ) + op.add_column( + "simulations", + sa.Column("output_dataset_id", sa.UUID(), autoincrement=False, nullable=True), + ) + op.add_column( + "simulations", + sa.Column( + "started_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_constraint(op.f("ck_reports_budget_window_size"), "reports", type_="check") + op.drop_index(op.f("ix_reports_user_id"), table_name="reports") + op.drop_column("reports", "budget_window_size") + op.drop_constraint(op.f("ck_report_runs_year"), "report_runs", type_="check") + op.drop_constraint( + op.f("ck_report_runs_cliff_impact_complete"), "report_runs", type_="check" + ) + op.drop_constraint( + op.f("ck_report_runs_budget_window_year"), "report_runs", type_="check" + ) + op.drop_constraint( + op.f("fk_report_runs_baseline_simulation_run_id_simulation_runs"), + "report_runs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_report_runs_reform_simulation_run_id_simulation_runs"), + "report_runs", + type_="foreignkey", + ) + op.drop_constraint( + "fk_report_runs_budget_window_run_id_budget_window_runs", + "report_runs", + type_="foreignkey", + ) + op.drop_constraint( + "uq_report_runs_budget_window_run_id_year", "report_runs", type_="unique" + ) + op.drop_index("ix_report_runs_status_started_at", table_name="report_runs") + op.drop_index( + op.f("ix_report_runs_reform_simulation_run_id"), table_name="report_runs" + ) + op.drop_index( + op.f("ix_report_runs_baseline_simulation_run_id"), table_name="report_runs" + ) + op.drop_column("report_runs", "data_version") + op.drop_column("report_runs", "spm_provenance") + op.drop_column("report_runs", "reform_cliff_share") + op.drop_column("report_runs", "reform_cliff_gap") + op.drop_column("report_runs", "baseline_cliff_share") + op.drop_column("report_runs", "baseline_cliff_gap") + op.drop_column("report_runs", "modal_call_id") + op.drop_column("report_runs", "year") + op.drop_column("report_runs", "budget_window_run_id") + op.drop_column("report_runs", "reform_simulation_run_id") + op.drop_column("report_runs", "baseline_simulation_run_id") + op.add_column( + "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_reform_simulation_run_id_simulation_runs"), + "program_statistics", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_program_statistics_baseline_simulation_run_id_simulation_runs"), + "program_statistics", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_program_statistics_baseline_simulation_id_simulations"), + "program_statistics", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_program_statistics_reform_simulation_id_simulations"), + "program_statistics", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + "uq_program_statistics_run_program_entity", "program_statistics", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_program_statistics_run_program_entity"), + "program_statistics", + ["report_run_id", "program_name", "entity"], + postgresql_nulls_not_distinct=False, + ) + op.alter_column( + "program_statistics", "is_tax", existing_type=sa.BOOLEAN(), nullable=False + ) + op.alter_column( + "program_statistics", + "entity", + existing_type=sa.VARCHAR(length=128), + nullable=False, + ) + op.drop_column("program_statistics", "reform_simulation_run_id") + op.drop_column("program_statistics", "baseline_simulation_run_id") + op.add_column( + "poverty", + sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.add_column( + "poverty", + sa.Column("report_run_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", "filter_group") + 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("ck_local_authority_impacts_disclosure_tier"), + "local_authority_impacts", + type_="check", + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_reform_simulation_run_id_simulation_runs"), + "local_authority_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_local_authority_impacts_baseline_simulation_run_id_simulation_runs"), + "local_authority_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_local_authority_impacts_reform_simulation_id_simulations"), + "local_authority_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + 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.drop_column("local_authority_impacts", "relative_household_income_change_upper") + op.drop_column("local_authority_impacts", "relative_household_income_change_lower") + op.drop_column("local_authority_impacts", "average_household_income_change_upper") + op.drop_column("local_authority_impacts", "average_household_income_change_lower") + op.drop_column("local_authority_impacts", "disclosure_tier") + op.drop_column("local_authority_impacts", "reform_simulation_run_id") + op.drop_column("local_authority_impacts", "baseline_simulation_run_id") + op.add_column( + "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_baseline_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_intra_decile_impacts_reform_simulation_run_id_simulation_runs"), + "intra_decile_impacts", + type_="foreignkey", + ) + 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.create_foreign_key( + op.f("fk_intra_decile_impacts_reform_simulation_id_simulations"), + "intra_decile_impacts", + "simulations", + ["reform_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("simulation_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.add_column( + "inequality", + sa.Column("report_run_id", sa.UUID(), autoincrement=False, nullable=False), + ) + op.drop_constraint( + op.f("fk_inequality_simulation_run_id_simulation_runs"), + "inequality", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_inequality_simulation_id_simulations"), + "inequality", + "simulations", + ["simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_inequality_report_run_id_report_runs"), + "inequality", + "report_runs", + ["report_run_id"], + ["id"], + ondelete="CASCADE", + ) + 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_constraint( + op.f( + "fk_household_jobs_tax_benefit_model_version_id_tax_benefit_model_versions" + ), + "household_jobs", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_household_jobs_household_id_households"), + "household_jobs", + type_="foreignkey", + ) + op.drop_index( + op.f("ix_household_jobs_tax_benefit_model_version_id"), + table_name="household_jobs", + ) + op.drop_index(op.f("ix_household_jobs_household_id"), table_name="household_jobs") + op.drop_column("household_jobs", "modal_call_id") + 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_reform_simulation_run_id_simulation_runs"), + "decile_impacts", + type_="foreignkey", + ) + op.drop_constraint( + op.f("fk_decile_impacts_baseline_simulation_run_id_simulation_runs"), + "decile_impacts", + type_="foreignkey", + ) + op.create_foreign_key( + op.f("fk_decile_impacts_reform_simulation_id_simulations"), + "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_constraint( + "uq_decile_impacts_run_measure", "decile_impacts", type_="unique" + ) + op.create_unique_constraint( + op.f("uq_decile_impacts_run_measure"), + "decile_impacts", + ["report_run_id", "income_variable", "entity", "decile", "quantiles"], + postgresql_nulls_not_distinct=False, + ) + op.drop_column("decile_impacts", "decile_variable") + op.drop_column("decile_impacts", "reform_simulation_run_id") + op.drop_column("decile_impacts", "baseline_simulation_run_id") + op.add_column( + "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("ck_constituency_impacts_disclosure_tier"), + "constituency_impacts", + type_="check", + ) + op.drop_constraint( + op.f("fk_constituency_impacts_reform_simulation_run_id_simulation_runs"), + "constituency_impacts", + 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_baseline_simulation_id_simulations"), + "constituency_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.create_foreign_key( + op.f("fk_constituency_impacts_reform_simulation_id_simulations"), + "constituency_impacts", + "simulations", + ["reform_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_column("constituency_impacts", "relative_household_income_change_upper") + op.drop_column("constituency_impacts", "relative_household_income_change_lower") + op.drop_column("constituency_impacts", "average_household_income_change_upper") + op.drop_column("constituency_impacts", "average_household_income_change_lower") + op.drop_column("constituency_impacts", "disclosure_tier") + op.drop_column("constituency_impacts", "reform_simulation_run_id") + op.drop_column("constituency_impacts", "baseline_simulation_run_id") + op.add_column( + "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_baseline_simulation_i_d68d"), + "congressional_district_impacts", + "simulations", + ["baseline_simulation_id"], + ["id"], + ondelete="RESTRICT", + ) + 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.drop_column("congressional_district_impacts", "no_change_percentage") + op.drop_column("congressional_district_impacts", "loser_percentage") + op.drop_column("congressional_district_impacts", "winner_percentage") + op.drop_column("congressional_district_impacts", "reform_simulation_run_id") + op.drop_column("congressional_district_impacts", "baseline_simulation_run_id") + op.add_column( + "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_table("labor_supply_responses") + op.drop_table("labor_supply_decile_responses") + op.drop_index( + op.f("ix_budget_window_runs_tax_benefit_model_version_id"), + table_name="budget_window_runs", + ) + op.drop_index( + "ix_budget_window_runs_report_id_created_at", table_name="budget_window_runs" + ) + op.drop_table("budget_window_runs") + op.drop_index( + op.f("ix_simulation_runs_tax_benefit_model_version_id"), + table_name="simulation_runs", + ) + op.drop_index("ix_simulation_runs_status_started_at", table_name="simulation_runs") + op.drop_index("ix_simulation_runs_current_output", table_name="simulation_runs") + op.drop_table("simulation_runs") + # ### end Alembic commands ### diff --git a/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..6c50fae72 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "c614b6bcccfc" +V2_ALEMBIC_HEAD_REVISION = "1356f658347d" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index 8498cfbb6..5cd21269d 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -66,6 +66,13 @@ LegacyUserPolicyMapping, ) from policyengine_api.data.v2.models.reports import ( # noqa: E402 + BudgetWindowRun, + Report, + ReportRun, + ReportRunStatus, + ReportRunTrigger, +) +from policyengine_api.data.v2.models.run_results import ( # noqa: E402 # tables declared after the naming convention above AggregateOutput, AggregateType, BudgetSummary, @@ -76,17 +83,16 @@ DecileType, Inequality, IntraDecileImpact, + LaborSupplyDecileResponse, + LaborSupplyResponse, LocalAuthorityImpact, OutputStatus, Poverty, ProgramStatistics, - Report, - ReportRun, - ReportRunStatus, - ReportRunTrigger, ) from policyengine_api.data.v2.models.simulations import ( # noqa: E402 Simulation, + SimulationRun, SimulationStatus, SimulationType, ) @@ -104,6 +110,7 @@ "AggregateOutput", "AggregateType", "BudgetSummary", + "BudgetWindowRun", "CatalogEntity", "CatalogEnumValue", "ChangeAggregate", @@ -120,6 +127,8 @@ "HouseholdJobStatus", "Inequality", "IntraDecileImpact", + "LaborSupplyDecileResponse", + "LaborSupplyResponse", "LegacyHouseholdMapping", "LegacyPolicyMapping", "LegacyUserMapping", @@ -141,6 +150,7 @@ "ReportRunStatus", "ReportRunTrigger", "Simulation", + "SimulationRun", "SimulationStatus", "SimulationType", "Stage12AggregationStatus", diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 181266d9d..7a7fd4ca0 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -8,11 +8,7 @@ import sqlalchemy as sa from sqlmodel import Field, Relationship -from policyengine_api.data.v2.models.base import ( - IdentifiedModel, - TimestampedModel, - enum_type, -) +from policyengine_api.data.v2.models.base import TimestampedModel, enum_type from policyengine_api.data.v2.models.policies import Dynamic, Policy if TYPE_CHECKING: @@ -25,7 +21,7 @@ class HouseholdJobStatus(StrEnum): - """Lifecycle state of a queued household calculation job.""" + """Lifecycle state of a household calculation job.""" PENDING = "pending" RUNNING = "running" @@ -96,18 +92,39 @@ 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. The simulation service writes the job and the Public API + the rows it names; each service writes only its own tables, and the foreign + keys between them stay. The job's household and catalog version can't be + deleted while it exists; deleting its policy or dynamic keeps the job but + clears the reference. """ __tablename__ = "household_jobs" __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( + foreign_key="households.id", + ondelete="RESTRICT", + index=True, + ) + year: int request_data: dict[str, Any] = Field(sa_type=sa.JSON) policy_id: UUID | None = Field( default=None, @@ -119,10 +136,18 @@ class HouseholdJob(IdentifiedModel, table=True): foreign_key="dynamics.id", ondelete="SET NULL", ) + tax_benefit_model_version_id: UUID = Field( + foreign_key="tax_benefit_model_versions.id", + ondelete="RESTRICT", + index=True, + ) + country_package_version: str = Field(max_length=128) + policyengine_version: str = Field(max_length=128) status: HouseholdJobStatus = Field( default=HouseholdJobStatus.PENDING, sa_type=enum_type(HouseholdJobStatus, "v2_household_job_status"), ) + modal_call_id: str | None = Field(default=None, max_length=255) error_message: str | None = None result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) started_at: datetime | None = Field( diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 1eec3e005..06081e25c 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -249,10 +249,6 @@ class Dataset(TimestampedModel, table=True): back_populates="dataset", sa_relationship_kwargs={"foreign_keys": "Simulation.dataset_id"}, ) - output_simulations: list["Simulation"] = Relationship( - back_populates="output_dataset", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, - ) reports: list["Report"] = Relationship(back_populates="dataset") diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 5c73947c1..11f820ab4 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,15 +1,20 @@ -"""Canonical SQLModel tables for stable reports, runs, and run outputs.""" +"""Canonical SQLModel tables for stable reports and their runs. + +A report run compares one baseline simulation run with one reform simulation +run; a budget-window report has a run per year of each calculation. Their +results are in ``run_results``. +""" from datetime import datetime from enum import StrEnum -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import UUID import sqlalchemy as sa +from sqlalchemy.dialects import postgresql from sqlmodel import Field, Relationship from policyengine_api.data.v2.models.base import ( - IdentifiedModel, TimestampedModel, enum_type, ) @@ -23,6 +28,25 @@ from policyengine_api.data.v2.models.simulations import Simulation from policyengine_api.data.v2.models.users import User +if TYPE_CHECKING: + from policyengine_api.data.v2.models.run_results import ( + AggregateOutput, + BudgetSummary, + ChangeAggregate, + CongressionalDistrictImpact, + ConstituencyImpact, + DecileImpact, + IntraDecileImpact, + LaborSupplyDecileResponse, + LaborSupplyResponse, + LocalAuthorityImpact, + ProgramStatistics, + ) + +# A budget window compares at most this many years, none after LAST_YEAR. +MAX_BUDGET_WINDOW_YEARS = 75 +LAST_YEAR = 2200 + class ReportRunStatus(StrEnum): """Lifecycle state of a report run; finished runs record a completion time.""" @@ -41,32 +65,16 @@ class ReportRunTrigger(StrEnum): SYSTEM = "system" -class OutputStatus(StrEnum): - """Lifecycle state of one computed output row of a report run.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - - -class AggregateType(StrEnum): - """How an aggregate combines a variable's values.""" - - SUM = "sum" - MEAN = "mean" - COUNT = "count" - - -class DecileType(StrEnum): - """Whether deciles rank by income or by wealth.""" - - INCOME = "income" - WEALTH = "wealth" - - class Report(TimestampedModel, table=True): - """Stable report definition; execution state belongs to ReportRun.""" + """Stable report definition; execution state belongs to ReportRun. + + The simulation service writes reports, and the Public API the user, model, + policy, household, dataset and region a report names; each service writes + only its own tables, and the foreign keys between them stay. Deleting the + user, policy, dataset or region clears the reference, which for the user + is found through ``user_id``'s index; the model and household can't be + deleted while a report names them. + """ __tablename__ = "reports" __table_args__ = ( @@ -74,6 +82,12 @@ class Report(TimestampedModel, table=True): "year IS NULL OR year BETWEEN 1900 AND 2200", name="ck_reports_year", ), + sa.CheckConstraint( + "budget_window_size IS NULL OR (household_id IS NULL " + f"AND budget_window_size BETWEEN 1 AND {MAX_BUDGET_WINDOW_YEARS} " + f"AND year IS NOT NULL AND year + budget_window_size - 1 <= {LAST_YEAR})", + name="ck_reports_budget_window_size", + ), sa.Index("ix_reports_country_type_created_at", "country", "type", "created_at"), ) @@ -85,6 +99,7 @@ class Report(TimestampedModel, table=True): default=None, foreign_key="users.id", ondelete="SET NULL", + index=True, ) tax_benefit_model_id: UUID = Field( foreign_key="tax_benefit_models.id", @@ -122,6 +137,7 @@ class Report(TimestampedModel, table=True): ondelete="SET NULL", ) year: int | None = None + budget_window_size: int | None = None inputs: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) user: User | None = Relationship(back_populates="reports") @@ -142,10 +158,76 @@ class Report(TimestampedModel, table=True): back_populates="report", cascade_delete=True, ) + budget_window_runs: list["BudgetWindowRun"] = Relationship( + back_populates="report", + cascade_delete=True, + ) + + +class BudgetWindowRun(TimestampedModel, table=True): + """One calculation of a budget-window report on one catalog version. + + Each year of the window is a report run of the same report that names this + row and its year. ``recalculate`` says whether those years calculate their + simulations again instead of reusing successful runs. What starts the later + years is decided where the simulation service schedules them. + """ + + __tablename__ = "budget_window_runs" + __table_args__ = ( + sa.UniqueConstraint("id", "report_id", name="uq_budget_window_runs_id_report"), + sa.CheckConstraint( + "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", + name="ck_budget_window_runs_terminal_completion", + ), + sa.Index( + "ix_budget_window_runs_report_id_created_at", "report_id", "created_at" + ), + ) + + report_id: UUID = Field( + foreign_key="reports.id", + ondelete="CASCADE", + ) + tax_benefit_model_version_id: UUID = Field( + foreign_key="tax_benefit_model_versions.id", + ondelete="RESTRICT", + index=True, + ) + trigger: ReportRunTrigger = Field( + default=ReportRunTrigger.INITIAL, + sa_type=enum_type(ReportRunTrigger, "v2_report_run_trigger"), + ) + recalculate: bool = False + status: ReportRunStatus = Field( + default=ReportRunStatus.PENDING, + sa_type=enum_type(ReportRunStatus, "v2_report_run_status"), + ) + error_message: str | None = None + started_at: datetime | None = Field( + default=None, + sa_type=sa.DateTime(timezone=True), + ) + completed_at: datetime | None = Field( + default=None, + sa_type=sa.DateTime(timezone=True), + ) + + report: Report = Relationship(back_populates="budget_window_runs") class ReportRun(TimestampedModel, table=True): - """One immutable-version execution attempt for a stable report.""" + """One immutable-version execution attempt for a stable report. + + It compares ``baseline_simulation_run_id`` with + ``reform_simulation_run_id``; a household report with only a baseline has + no reform run. A year of a budget window names its ``budget_window_run_id`` + and ``year``, at most one report run per year of a window run. An economy + comparison's cliff impact is the baseline's and reform's cliff gap and + share, stored together or not at all; ``spm_provenance`` is the SPM receipt + and ``data_version`` the data release it used. Its labour supply response + has its own two tables. + """ __tablename__ = "report_runs" __table_args__ = ( @@ -158,6 +240,33 @@ class ReportRun(TimestampedModel, table=True): "status NOT IN ('succeeded', 'failed') OR completed_at IS NOT NULL", name="ck_report_runs_terminal_completion", ), + sa.CheckConstraint( + "year IS NULL OR year BETWEEN 1900 AND 2200", + name="ck_report_runs_year", + ), + sa.CheckConstraint( + "budget_window_run_id IS NULL OR year IS NOT NULL", + name="ck_report_runs_budget_window_year", + ), + sa.UniqueConstraint( + "budget_window_run_id", + "year", + name="uq_report_runs_budget_window_run_id_year", + ), + # A year's report run belongs to a budget-window run of its own report, + # which owns it as the report owns its runs. + sa.ForeignKeyConstraint( + ["budget_window_run_id", "report_id"], + ["budget_window_runs.id", "budget_window_runs.report_id"], + name="fk_report_runs_budget_window_run_id_budget_window_runs", + ondelete="CASCADE", + ), + sa.CheckConstraint( + "(baseline_cliff_gap IS NULL) = (baseline_cliff_share IS NULL) " + "AND (baseline_cliff_gap IS NULL) = (reform_cliff_gap IS NULL) " + "AND (baseline_cliff_gap IS NULL) = (reform_cliff_share IS NULL)", + name="ck_report_runs_cliff_impact_complete", + ), sa.Index( "ix_report_runs_current_output", "report_id", @@ -167,12 +276,26 @@ class ReportRun(TimestampedModel, table=True): "completed_at", "id", ), + sa.Index("ix_report_runs_status_started_at", "status", "started_at"), ) report_id: UUID = Field( foreign_key="reports.id", ondelete="CASCADE", ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + index=True, + ) + reform_simulation_run_id: UUID | None = Field( + default=None, + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + index=True, + ) + budget_window_run_id: UUID | None = None + year: int | None = None country_package_version: str = Field(max_length=128) policyengine_version: str = Field(max_length=128) status: ReportRunStatus = Field( @@ -184,6 +307,7 @@ class ReportRun(TimestampedModel, table=True): sa_type=enum_type(ReportRunTrigger, "v2_report_run_trigger"), ) idempotency_key: UUID | None = Field(default=None) + modal_call_id: str | None = Field(default=None, max_length=255) started_at: datetime | None = Field( default=None, sa_type=sa.DateTime(timezone=True), @@ -194,8 +318,25 @@ class ReportRun(TimestampedModel, table=True): ) error_message: str | None = None markdown: str | None = Field(default=None, sa_type=sa.Text) + baseline_cliff_gap: float | None = None + baseline_cliff_share: float | None = None + reform_cliff_gap: float | None = None + reform_cliff_share: float | None = None + spm_provenance: dict[str, Any] | None = Field( + default=None, + sa_type=sa.JSON().with_variant(postgresql.JSONB(), "postgresql"), + ) + data_version: str | None = Field(default=None, max_length=128) report: Report = Relationship(back_populates="runs") + labor_supply_responses: list["LaborSupplyResponse"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) + labor_supply_decile_responses: list["LaborSupplyDecileResponse"] = Relationship( + back_populates="report_run", + cascade_delete=True, + ) aggregates: list["AggregateOutput"] = Relationship( back_populates="report_run", cascade_delete=True, @@ -219,10 +360,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,417 +368,7 @@ 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 AggregateOutput(IdentifiedModel, table=True): - """A sum, mean or count of one variable in one simulation of a report run. - - ``filter_config`` optionally restricts which entities are included. - """ - - __tablename__ = "aggregates" - __table_args__ = ( - sa.Index("ix_aggregates_report_run_status", "report_run_id", "status"), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable: str = Field(max_length=512) - aggregate_type: AggregateType = Field( - sa_type=enum_type(AggregateType, "v2_aggregate_type") - ) - entity: str | None = Field(default=None, max_length=128) - filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) - status: OutputStatus = Field( - default=OutputStatus.PENDING, - sa_type=enum_type(OutputStatus, "v2_output_status"), - ) - error_message: str | None = None - result: float | None = None - - report_run: ReportRun = Relationship(back_populates="aggregates") - - -class ChangeAggregate(IdentifiedModel, table=True): - """An aggregate of a variable's change from baseline to reform in a report run. - - ``change_geq`` and ``change_leq`` optionally restrict it to entities whose - change falls within those bounds. - """ - - __tablename__ = "change_aggregates" - __table_args__ = ( - sa.Index( - "ix_change_aggregates_report_run_status", - "report_run_id", - "status", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable: str = Field(max_length=512) - aggregate_type: AggregateType = Field( - sa_type=enum_type(AggregateType, "v2_aggregate_type"), - ) - entity: str | None = Field(default=None, max_length=128) - filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) - change_geq: float | None = None - change_leq: float | None = None - status: OutputStatus = Field( - default=OutputStatus.PENDING, - sa_type=enum_type(OutputStatus, "v2_output_status"), - ) - error_message: str | None = None - result: float | None = None - - report_run: ReportRun = Relationship(back_populates="change_aggregates") - - -class BudgetSummary(IdentifiedModel, table=True): - """Baseline and reform totals of one budget variable for a report run. - - One row per run, variable and entity. - """ - - __tablename__ = "budget_summary" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "variable_name", - "entity", - name="uq_budget_summary_run_variable_entity", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - variable_name: str = Field(max_length=512) - entity: str = Field(max_length=128) - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - - report_run: ReportRun = Relationship(back_populates="budget_summaries") - - -class DecileImpact(IdentifiedModel, table=True): - """Mean income change and better/worse-off counts in one income quantile. - - ``decile`` runs from 1 to ``quantiles``; there is one row per run, income - variable, entity, quantile count and decile. - """ - - __tablename__ = "decile_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "income_variable", - "entity", - "decile", - "quantiles", - name="uq_decile_impacts_run_measure", - ), - sa.CheckConstraint("quantiles > 0", name="ck_decile_impacts_quantiles"), - sa.CheckConstraint( - "decile BETWEEN 1 AND quantiles", - name="ck_decile_impacts_decile", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - income_variable: str = Field(max_length=512) - entity: str = Field(max_length=128) - decile: int - quantiles: int = 10 - baseline_mean: float | None = None - reform_mean: float | None = None - absolute_change: float | None = None - relative_change: float | None = None - count_better_off: float | None = None - count_worse_off: float | None = None - count_no_change: float | None = None - - report_run: ReportRun = Relationship(back_populates="decile_impacts") - - -class IntraDecileImpact(IdentifiedModel, table=True): - """Shares of one decile whose income rises or falls by more or less than 5%. - - One row per run, decile type and decile, where ``decile`` is 0 to 10. - """ - - __tablename__ = "intra_decile_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "decile_type", - "decile", - name="uq_intra_decile_impacts_run_type_decile", - ), - sa.CheckConstraint( - "decile BETWEEN 0 AND 10", - name="ck_intra_decile_impacts_decile", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - decile_type: DecileType = Field( - default=DecileType.INCOME, - sa_type=enum_type(DecileType, "v2_decile_type"), - ) - decile: int - lose_more_than_5pct: float | None = None - lose_less_than_5pct: float | None = None - no_change: float | None = None - gain_less_than_5pct: float | None = None - gain_more_than_5pct: float | None = None - - report_run: ReportRun = Relationship(back_populates="intra_decile_impacts") - - -class Inequality(IdentifiedModel, table=True): - """Gini and top/bottom income shares for one simulation of a report run.""" - - __tablename__ = "inequality" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "simulation_id", - "income_variable", - "entity", - name="uq_inequality_run_simulation_measure", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - income_variable: str = Field(max_length=512) - entity: str = Field(default="household", max_length=128) - gini: float | None = None - top_10_share: float | None = None - top_1_share: float | None = None - bottom_50_share: float | None = None - - report_run: ReportRun = Relationship(back_populates="inequalities") - - -class Poverty(IdentifiedModel, table=True): - """Poverty headcount and rate of one type for one simulation of a report run.""" - - __tablename__ = "poverty" - __table_args__ = ( - sa.Index( - "ix_poverty_run_simulation_type", - "report_run_id", - "simulation_id", - "poverty_type", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - poverty_type: str = Field(max_length=128) - entity: str = Field(default="person", max_length=128) - filter_variable: str | None = Field(default=None, max_length=512) - headcount: float | None = None - total_population: float | None = None - rate: float | None = None - - report_run: ReportRun = Relationship(back_populates="poverty_results") - - -class ProgramStatistics(IdentifiedModel, table=True): - """Baseline and reform totals, counts, winners and losers for one program. - - One row per run, program and entity; ``is_tax`` marks taxes. - """ - - __tablename__ = "program_statistics" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "program_name", - "entity", - name="uq_program_statistics_run_program_entity", - ), - ) - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - program_name: str = Field(max_length=512) - entity: str = Field(max_length=128) - is_tax: bool = False - baseline_total: float | None = None - reform_total: float | None = None - change: float | None = None - baseline_count: float | None = None - reform_count: float | None = None - winners: float | None = None - losers: float | None = None - - report_run: ReportRun = Relationship(back_populates="program_statistics") - - -class GeographicImpactBase(IdentifiedModel): - """Shared non-table fields for geographic report-run impacts.""" - - report_run_id: UUID = Field( - foreign_key="report_runs.id", - ondelete="CASCADE", - ) - baseline_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - reform_simulation_id: UUID = Field( - foreign_key="simulations.id", - ondelete="RESTRICT", - ) - average_household_income_change: float - relative_household_income_change: float - population: float - - -class CongressionalDistrictImpact(GeographicImpactBase, table=True): - """Household income change and population in one US congressional district. - - One row per run and district GEOID. - """ - - __tablename__ = "congressional_district_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "district_geoid", - name="uq_congressional_district_impacts_run_geoid", - ), - ) - - district_geoid: int - state_fips: int - district_number: int - - report_run: ReportRun = Relationship( - back_populates="congressional_district_impacts" - ) - - -class ConstituencyImpact(GeographicImpactBase, table=True): - """Household income change in one UK constituency. - - One row per run and constituency code. - """ - - __tablename__ = "constituency_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "constituency_code", - name="uq_constituency_impacts_run_code", - ), - ) - - constituency_code: str = Field(max_length=64) - constituency_name: str = Field(max_length=255) - - report_run: ReportRun = Relationship(back_populates="constituency_impacts") - - -class LocalAuthorityImpact(GeographicImpactBase, table=True): - """Household income change in one UK local authority. - - One row per run and local authority code. - """ - - __tablename__ = "local_authority_impacts" - __table_args__ = ( - sa.UniqueConstraint( - "report_run_id", - "local_authority_code", - name="uq_local_authority_impacts_run_code", - ), - ) - - local_authority_code: str = Field(max_length=64) - local_authority_name: str = Field(max_length=255) - - report_run: ReportRun = Relationship(back_populates="local_authority_impacts") diff --git a/policyengine_api/data/v2/models/run_results.py b/policyengine_api/data/v2/models/run_results.py new file mode 100644 index 000000000..1f6943783 --- /dev/null +++ b/policyengine_api/data/v2/models/run_results.py @@ -0,0 +1,554 @@ +"""Canonical SQLModel tables for the results of simulation runs and report runs. + +Results that compare a report run's baseline and reform belong to the report +run and name its two simulation runs; each such row is unique within its report +run. Poverty and inequality, which describe one simulation, belong to a +simulation run. +""" + +from enum import StrEnum +from typing import Any +from uuid import UUID + +import sqlalchemy as sa +from sqlmodel import Field, Relationship + +from policyengine_api.data.v2.models.base import IdentifiedModel, enum_type +from policyengine_api.data.v2.models.reports import ReportRun +from policyengine_api.data.v2.models.simulations import SimulationRun + + +class OutputStatus(StrEnum): + """Lifecycle state of one computed output row of a report run.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class AggregateType(StrEnum): + """How an aggregate combines a variable's values.""" + + SUM = "sum" + MEAN = "mean" + COUNT = "count" + + +class DecileType(StrEnum): + """Whether deciles rank by income or by wealth.""" + + INCOME = "income" + WEALTH = "wealth" + + +class LaborSupplyResponse(IdentifiedModel, table=True): + """The labour supply response to the reform in one report run. + + One row per report run. The earnings changes come from the income and + substitution effects (.py's ``income_lsr`` and ``substitution_lsr``) and + total ``total_earnings_change``; ``relative_income_effect`` and + ``relative_substitution_effect`` are relative to baseline earnings. The + weekly hours are the total hours worked under the baseline and the reform, + and their change split into the two effects. + """ + + __tablename__ = "labor_supply_responses" + __table_args__ = ( + sa.UniqueConstraint("report_run_id", name="uq_labor_supply_responses_run"), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + income_effect_earnings_change: float + substitution_effect_earnings_change: float + total_earnings_change: float + revenue_change: float + relative_income_effect: float + relative_substitution_effect: float + baseline_weekly_hours: float + reform_weekly_hours: float + weekly_hours_change: float + weekly_hours_income_effect: float + weekly_hours_substitution_effect: float + + report_run: ReportRun = Relationship(back_populates="labor_supply_responses") + + +class LaborSupplyDecileResponse(IdentifiedModel, table=True): + """The labour supply response in one income decile of one report run. + + One row per report run and decile, as .py numbers them. The average effects + are per household; the relative effects are relative to the decile's + baseline earnings, which .py reports only for deciles above zero, so they + may be null. + """ + + __tablename__ = "labor_supply_decile_responses" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "decile", + name="uq_labor_supply_decile_responses_run_decile", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + decile: int + average_income_effect: float | None = None + average_substitution_effect: float | None = None + relative_income_effect: float | None = None + relative_substitution_effect: float | None = None + + report_run: ReportRun = Relationship(back_populates="labor_supply_decile_responses") + + +class AggregateOutput(IdentifiedModel, table=True): + """A sum, mean or count of one variable in one simulation of a report run. + + ``filter_config`` optionally restricts which entities are included. + """ + + __tablename__ = "aggregates" + __table_args__ = ( + sa.Index("ix_aggregates_report_run_status", "report_run_id", "status"), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + simulation_id: UUID = Field( + foreign_key="simulations.id", + ondelete="RESTRICT", + ) + variable: str = Field(max_length=512) + aggregate_type: AggregateType = Field( + sa_type=enum_type(AggregateType, "v2_aggregate_type") + ) + entity: str | None = Field(default=None, max_length=128) + filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) + status: OutputStatus = Field( + default=OutputStatus.PENDING, + sa_type=enum_type(OutputStatus, "v2_output_status"), + ) + error_message: str | None = None + result: float | None = None + + report_run: ReportRun = Relationship(back_populates="aggregates") + + +class ChangeAggregate(IdentifiedModel, table=True): + """An aggregate of a variable's change from baseline to reform in a report run. + + ``change_geq`` and ``change_leq`` optionally restrict it to entities whose + change falls within those bounds. + """ + + __tablename__ = "change_aggregates" + __table_args__ = ( + sa.Index( + "ix_change_aggregates_report_run_status", + "report_run_id", + "status", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_id: UUID = Field( + foreign_key="simulations.id", + ondelete="RESTRICT", + ) + reform_simulation_id: UUID = Field( + foreign_key="simulations.id", + ondelete="RESTRICT", + ) + variable: str = Field(max_length=512) + aggregate_type: AggregateType = Field( + sa_type=enum_type(AggregateType, "v2_aggregate_type"), + ) + entity: str | None = Field(default=None, max_length=128) + filter_config: dict[str, Any] = Field(default_factory=dict, sa_type=sa.JSON) + change_geq: float | None = None + change_leq: float | None = None + status: OutputStatus = Field( + default=OutputStatus.PENDING, + sa_type=enum_type(OutputStatus, "v2_output_status"), + ) + error_message: str | None = None + result: float | None = None + + report_run: ReportRun = Relationship(back_populates="change_aggregates") + + +class BudgetSummary(IdentifiedModel, table=True): + """Baseline and reform totals of one budget variable for a report run. + + One row per report run, variable and entity. + """ + + __tablename__ = "budget_summary" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "variable_name", + "entity", + name="uq_budget_summary_run_variable_entity", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + variable_name: str = Field(max_length=512) + entity: str = Field(max_length=128) + baseline_total: float | None = None + reform_total: float | None = None + change: float | None = None + + report_run: ReportRun = Relationship(back_populates="budget_summaries") + + +class DecileImpact(IdentifiedModel, table=True): + """Mean income change and better/worse-off counts in one income quantile. + + ``decile`` runs from 1 to ``quantiles``; there is one row per report run, + income variable, decile variable, entity, quantile count and decile. + ``decile_variable`` names a variable that groups the households instead of + the income variable's own deciles, such as ``household_wealth_decile``; rows + without one count as equal. + """ + + __tablename__ = "decile_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "income_variable", + "decile_variable", + "entity", + "decile", + "quantiles", + name="uq_decile_impacts_run_measure", + postgresql_nulls_not_distinct=True, + ), + sa.CheckConstraint("quantiles > 0", name="ck_decile_impacts_quantiles"), + sa.CheckConstraint( + "decile BETWEEN 1 AND quantiles", + name="ck_decile_impacts_decile", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + income_variable: str = Field(max_length=512) + decile_variable: str | None = Field(default=None, max_length=512) + entity: str = Field(max_length=128) + decile: int + quantiles: int = 10 + baseline_mean: float | None = None + reform_mean: float | None = None + absolute_change: float | None = None + relative_change: float | None = None + count_better_off: float | None = None + count_worse_off: float | None = None + count_no_change: float | None = None + + report_run: ReportRun = Relationship(back_populates="decile_impacts") + + +class IntraDecileImpact(IdentifiedModel, table=True): + """Shares of one decile whose income rises or falls by more or less than 5%. + + One row per report run, decile type and decile, where ``decile`` is 0 to + 10. + """ + + __tablename__ = "intra_decile_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "decile_type", + "decile", + name="uq_intra_decile_impacts_run_type_decile", + ), + sa.CheckConstraint( + "decile BETWEEN 0 AND 10", + name="ck_intra_decile_impacts_decile", + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + decile_type: DecileType = Field( + default=DecileType.INCOME, + sa_type=enum_type(DecileType, "v2_decile_type"), + ) + decile: int + lose_more_than_5pct: float | None = None + lose_less_than_5pct: float | None = None + no_change: float | None = None + gain_less_than_5pct: float | None = None + gain_more_than_5pct: float | None = None + + report_run: ReportRun = Relationship(back_populates="intra_decile_impacts") + + +class Inequality(IdentifiedModel, table=True): + """Gini and top/bottom income shares for one simulation run.""" + + __tablename__ = "inequality" + __table_args__ = ( + sa.UniqueConstraint( + "simulation_run_id", + "income_variable", + "entity", + name="uq_inequality_simulation_run_measure", + ), + ) + + simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="CASCADE", + ) + income_variable: str = Field(max_length=512) + entity: str = Field(default="household", max_length=128) + gini: float | None = None + top_10_share: float | None = None + top_1_share: float | None = None + bottom_50_share: float | None = None + + simulation_run: SimulationRun = Relationship(back_populates="inequalities") + + +class Poverty(IdentifiedModel, table=True): + """Poverty headcount and rate of one type for one simulation run. + + One row per run, poverty type, entity, filter variable and filter group: + ``filter_variable`` selects people (``age``) and ``filter_group`` names the + group (``child``). Missing values count as equal, so each run has one + unfiltered row per type and entity. + """ + + __tablename__ = "poverty" + __table_args__ = ( + sa.UniqueConstraint( + "simulation_run_id", + "poverty_type", + "entity", + "filter_variable", + "filter_group", + name="uq_poverty_simulation_run_measure", + postgresql_nulls_not_distinct=True, + ), + ) + + simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="CASCADE", + ) + poverty_type: str = Field(max_length=128) + entity: str = Field(default="person", max_length=128) + filter_variable: str | None = Field(default=None, max_length=512) + filter_group: str | None = Field(default=None, max_length=64) + headcount: float | None = None + total_population: float | None = None + rate: float | None = None + + simulation_run: SimulationRun = Relationship(back_populates="poverty_results") + + +class ProgramStatistics(IdentifiedModel, table=True): + """Baseline and reform totals, counts, winners and losers for one program. + + One row per report run, program and entity; ``is_tax`` marks taxes. The + economy output's detailed budget names neither the entity nor whether a + program is a tax, so both may be null, and rows without an entity count as + equal. + """ + + __tablename__ = "program_statistics" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "program_name", + "entity", + name="uq_program_statistics_run_program_entity", + postgresql_nulls_not_distinct=True, + ), + ) + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + program_name: str = Field(max_length=512) + entity: str | None = Field(default=None, max_length=128) + is_tax: bool | None = None + baseline_total: float | None = None + reform_total: float | None = None + change: float | None = None + baseline_count: float | None = None + reform_count: float | None = None + winners: float | None = None + losers: float | None = None + + report_run: ReportRun = Relationship(back_populates="program_statistics") + + +class GeographicImpactBase(IdentifiedModel): + """Shared non-table fields for geographic report-run impacts.""" + + report_run_id: UUID = Field( + foreign_key="report_runs.id", + ondelete="CASCADE", + ) + baseline_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + reform_simulation_run_id: UUID = Field( + foreign_key="simulation_runs.id", + ondelete="RESTRICT", + ) + average_household_income_change: float + relative_household_income_change: float + population: float + + +class UKAreaImpactBase(GeographicImpactBase): + """Shared non-table disclosure fields of a UK area's impact. + + ``disclosure_tier`` is ``reported``, or ``approximate`` when the area's + sample is thin; the ``_lower`` and ``_upper`` fields bound each income + change's 90% interval. All are null for results without disclosure control. + """ + + disclosure_tier: str | None = Field(default=None, max_length=16) + average_household_income_change_lower: float | None = None + average_household_income_change_upper: float | None = None + relative_household_income_change_lower: float | None = None + relative_household_income_change_upper: float | None = None + + +def _disclosure_tier_check(table_name: str) -> sa.CheckConstraint: + return sa.CheckConstraint( + "disclosure_tier IS NULL OR disclosure_tier IN ('reported', 'approximate')", + name=f"ck_{table_name}_disclosure_tier", + ) + + +class CongressionalDistrictImpact(GeographicImpactBase, table=True): + """Household income change and population in one US congressional district. + + One row per report run and district GEOID, with the shares of people who + gain, lose or see no change. + """ + + __tablename__ = "congressional_district_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "district_geoid", + name="uq_congressional_district_impacts_run_geoid", + ), + ) + + district_geoid: int + state_fips: int + district_number: int + winner_percentage: float + loser_percentage: float + no_change_percentage: float + + report_run: ReportRun = Relationship( + back_populates="congressional_district_impacts" + ) + + +class ConstituencyImpact(UKAreaImpactBase, table=True): + """Household income change in one UK constituency. + + One row per report run and constituency code. + """ + + __tablename__ = "constituency_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "constituency_code", + name="uq_constituency_impacts_run_code", + ), + _disclosure_tier_check("constituency_impacts"), + ) + + constituency_code: str = Field(max_length=64) + constituency_name: str = Field(max_length=255) + + report_run: ReportRun = Relationship(back_populates="constituency_impacts") + + +class LocalAuthorityImpact(UKAreaImpactBase, table=True): + """Household income change in one UK local authority. + + One row per report run and local authority code. + """ + + __tablename__ = "local_authority_impacts" + __table_args__ = ( + sa.UniqueConstraint( + "report_run_id", + "local_authority_code", + name="uq_local_authority_impacts_run_code", + ), + _disclosure_tier_check("local_authority_impacts"), + ) + + local_authority_code: str = Field(max_length=64) + local_authority_name: str = Field(max_length=255) + + report_run: ReportRun = Relationship(back_populates="local_authority_impacts") diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index a6821eadc..dafb7b41d 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 @@ -19,10 +19,11 @@ if TYPE_CHECKING: from policyengine_api.data.v2.models.reports import Report + from policyengine_api.data.v2.models.run_results import Inequality, Poverty class SimulationStatus(StrEnum): - """Lifecycle state of a v2 simulation.""" + """Lifecycle state of one run of a v2 simulation.""" PENDING = "pending" RUNNING = "running" @@ -40,8 +41,12 @@ 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``. + The simulation service writes it, and the Public API the rows it names; + each service writes only its own tables, and the foreign keys between them + stay. Deleting its policy, dynamic or region clears that reference; its + household, input dataset and catalog version cannot be deleted while it + exists. """ __tablename__ = "simulations" @@ -61,7 +66,6 @@ class Simulation(TimestampedModel, table=True): "year IS NULL OR year BETWEEN 1900 AND 2200", name="ck_simulations_year", ), - sa.Index("ix_simulations_status_created_at", "status", "created_at"), ) simulation_type: SimulationType = Field( @@ -93,43 +97,20 @@ class Simulation(TimestampedModel, table=True): ondelete="RESTRICT", index=True, ) - output_dataset_id: UUID | None = Field( - default=None, - foreign_key="datasets.id", - ondelete="SET NULL", - ) region_id: UUID | None = Field( default=None, foreign_key="regions.id", ondelete="SET NULL", ) - status: SimulationStatus = Field( - default=SimulationStatus.PENDING, - sa_type=enum_type(SimulationStatus, "v2_simulation_status"), - ) - error_message: str | None = None filter_field: str | None = Field(default=None, max_length=128) filter_value: str | None = Field(default=None, max_length=255) filter_strategy: str | None = Field(default=None, max_length=64) year: int | None = None - started_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - completed_at: datetime | None = Field( - default=None, - sa_type=sa.DateTime(timezone=True), - ) - household_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) dataset: Dataset | None = Relationship( back_populates="input_simulations", sa_relationship_kwargs={"foreign_keys": "Simulation.dataset_id"}, ) - output_dataset: Dataset | None = Relationship( - back_populates="output_simulations", - sa_relationship_kwargs={"foreign_keys": "Simulation.output_dataset_id"}, - ) household: Household | None = Relationship(back_populates="simulations") policy: Policy | None = Relationship(back_populates="simulations") dynamic: Dynamic | None = Relationship(back_populates="simulations") @@ -145,3 +126,105 @@ class Simulation(TimestampedModel, table=True): back_populates="reform_simulation", sa_relationship_kwargs={"foreign_keys": "Report.reform_simulation_id"}, ) + runs: list["SimulationRun"] = Relationship( + back_populates="simulation", + cascade_delete=True, + ) + + +class SimulationRun(TimestampedModel, table=True): + """One run of a simulation for one year on one catalog and .py version. + + A simulation has a run per year it is calculated for, so a budget window + compares a report's same two simulations in every year of the window. + + ``tax_benefit_model_version_id``, ``dataset_id`` and ``region_id`` name the + catalog version, dataset and region the run used, which a run on a newer + version resolves again. They are foreign keys to the Public API's catalog + with the same delete rules as the simulation's: the catalog version and + dataset can't be deleted while a run names them, and deleting the region + clears the reference. A household simulation's run stores its household in + ``household_result``; an economy run's output file is ``output_uri``, with + its SHA-256. Each poverty and inequality result belongs to one run. + """ + + __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.CheckConstraint( + "year BETWEEN 1900 AND 2200", + name="ck_simulation_runs_year", + ), + sa.Index( + "ix_simulation_runs_current_output", + "simulation_id", + "year", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ), + sa.Index("ix_simulation_runs_status_started_at", "status", "started_at"), + ) + + simulation_id: UUID = Field( + foreign_key="simulations.id", + ondelete="CASCADE", + ) + year: int + tax_benefit_model_version_id: UUID = Field( + foreign_key="tax_benefit_model_versions.id", + ondelete="RESTRICT", + index=True, + ) + country_package_version: str = Field(max_length=128) + policyengine_version: str = Field(max_length=128) + dataset_id: UUID | None = Field( + default=None, + foreign_key="datasets.id", + ondelete="RESTRICT", + ) + region_id: UUID | None = Field( + default=None, + foreign_key="regions.id", + ondelete="SET NULL", + ) + status: SimulationStatus = Field( + default=SimulationStatus.PENDING, + sa_type=enum_type(SimulationStatus, "v2_simulation_status"), + ) + modal_call_id: str | None = Field(default=None, max_length=255) + error_message: str | None = None + started_at: datetime | None = Field( + default=None, + sa_type=sa.DateTime(timezone=True), + ) + completed_at: datetime | None = Field( + default=None, + sa_type=sa.DateTime(timezone=True), + ) + household_result: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) + output_uri: str | None = Field(default=None, max_length=2048) + output_sha256: str | None = Field(default=None, max_length=64) + + simulation: Simulation = Relationship(back_populates="runs") + poverty_results: list["Poverty"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, + ) + inequalities: list["Inequality"] = Relationship( + back_populates="simulation_run", + cascade_delete=True, + ) diff --git a/policyengine_api/data/v2/privilege_plans.py b/policyengine_api/data/v2/privilege_plans.py new file mode 100644 index 000000000..671243cc8 --- /dev/null +++ b/policyengine_api/data/v2/privilege_plans.py @@ -0,0 +1,373 @@ +"""Plan, state and report the table privileges the v2 manifest declares. + +The pure half of ``privilege_reconciliation``: it compares the declared +privileges with what a database grants, writes the GRANT statements, finds the +privileges a role lacks or holds too widely, and renders the reports. It runs +no SQL; ``privilege_reconciliation`` reads the database and runs the result. +""" + +from __future__ import annotations + +import re +from collections import defaultdict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from policyengine_api.data.v2.privileges import ( + ColumnPrivilege, + DeclaredPrivilege, + TablePrivilege, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Collection, Iterable, Mapping, Sequence + +SCHEMA = "public" +SAFE_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]{0,62}$") + + +class PrivilegeReconciliationError(RuntimeError): + """Raised when privileges can't be reconciled or don't match the manifest.""" + + +class UnsafeIdentifierError(PrivilegeReconciliationError): + """Raised for a role or table name that isn't a plain lowercase identifier.""" + + +@dataclass(frozen=True, order=True) +class PrivilegeEntry: + """One privilege one role holds, or should hold, on one table. + + ``column`` names the one column a column privilege is on; it is empty for a + privilege on the whole table. + """ + + role: str + table: str + privilege: str + column: str = "" + + +@dataclass(frozen=True) +class GrantedPrivileges: + """The managed roles' direct grants in one database. + + ``tables`` holds grants on whole tables, by role then table, and + ``columns`` grants on columns, by role, table, then column. + """ + + tables: Mapping[str, Mapping[str, Collection[str]]] + columns: Mapping[str, Mapping[str, Mapping[str, Collection[str]]]] = field( + default_factory=dict + ) + + +@dataclass(frozen=True) +class DefaultPrivilege: + """A default privilege PostgreSQL gives on objects a role creates later.""" + + creating_role: str + schema: str + object_type: str + grantee: str + privileges: tuple[str, ...] + + +@dataclass(frozen=True) +class ReconciliationPlan: + """What ``apply`` grants and reports for one database. + + ``skipped_roles`` are missing roles awaiting creation by an operator; they + don't block the plan, and nothing is granted to them. + """ + + grants: tuple[PrivilegeEntry, ...] + undeclared: tuple[PrivilegeEntry, ...] + missing_roles: tuple[str, ...] + missing_tables: tuple[str, ...] + skipped_roles: tuple[str, ...] = () + + @property + def can_apply(self) -> bool: + """Whether every declared role and table exists, so grants may run.""" + return not self.missing_roles and not self.missing_tables + + +def _declared_entries( + role: str, table: str, privileges: Collection[DeclaredPrivilege] +) -> list[PrivilegeEntry]: + """Return one entry per whole-table privilege and per column of a column one.""" + entries = [] + for privilege in privileges: + if isinstance(privilege, ColumnPrivilege): + entries += [ + PrivilegeEntry(role, table, privilege.privilege.value, column) + for column in privilege.columns + ] + else: + entries.append(PrivilegeEntry(role, table, privilege.value)) + return entries + + +def plan_reconciliation( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + granted: GrantedPrivileges, + *, + existing_roles: Collection[str], + existing_tables: Collection[str], + roles_awaiting_creation: Collection[str] = (), +) -> ReconciliationPlan: + """Compare declared privileges with what the database grants. + + A missing role in ``roles_awaiting_creation`` is skipped instead of missing. + """ + absent_roles = set(declared) - set(existing_roles) + skipped_roles = tuple(sorted(absent_roles & set(roles_awaiting_creation))) + missing_roles = tuple(sorted(absent_roles - set(skipped_roles))) + declared_tables = {table for tables in declared.values() for table in tables} + missing_tables = tuple(sorted(declared_tables - set(existing_tables))) + declared_entries = { + entry + for role, tables in declared.items() + for table, privileges in tables.items() + for entry in _declared_entries(role, table, privileges) + } + granted_entries = { + PrivilegeEntry(role, table, privilege) + for role, tables in granted.tables.items() + for table, privileges in tables.items() + for privilege in privileges + } | { + PrivilegeEntry(role, table, privilege, column) + for role, tables in granted.columns.items() + for table, columns in tables.items() + for column, privileges in columns.items() + for privilege in privileges + } + grants = sorted( + entry + for entry in declared_entries - granted_entries + if entry.role not in absent_roles and entry.table not in missing_tables + ) + undeclared = sorted(granted_entries - declared_entries) + return ReconciliationPlan( + grants=tuple(grants), + undeclared=tuple(undeclared), + missing_roles=missing_roles, + missing_tables=missing_tables, + skipped_roles=skipped_roles, + ) + + +def safe_identifier(identifier: str) -> str: + """Return ``identifier`` if it is a plain lowercase SQL name, else raise.""" + if SAFE_IDENTIFIER.fullmatch(identifier) is None: + msg = f"refusing to use {identifier!r} as a role, table or column name" + raise UnsafeIdentifierError(msg) + return identifier + + +def grant_statements( + plan: ReconciliationPlan, quote_identifier: Callable[[str], str] +) -> list[str]: + """Return the GRANTs for a plan, in role and table order. + + Each role and table gets one GRANT of its whole-table privileges, then one + per column privilege, naming its columns in name order. + """ + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str, str], set[str]] = defaultdict(set) + for entry in plan.grants: + target = (safe_identifier(entry.role), safe_identifier(entry.table)) + privilege = TablePrivilege(entry.privilege).value + if entry.column: + columns_by_privilege[(*target, privilege)].add( + safe_identifier(entry.column) + ) + else: + table_privileges[target].add(privilege) + statements: list[tuple[tuple[str, str, str], str]] = [ + ( + (role, table, ""), + f"GRANT {', '.join(sorted(privileges))} ON TABLE " + f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", + ) + for (role, table), privileges in table_privileges.items() + ] + statements += [ + ( + (role, table, privilege), + f"GRANT {privilege} " + f"({', '.join(quote_identifier(column) for column in sorted(columns))}) " + f"ON TABLE {quote_identifier(SCHEMA)}.{quote_identifier(table)} " + f"TO {quote_identifier(role)}", + ) + for (role, table, privilege), columns in columns_by_privilege.items() + ] + return [statement for _target, statement in sorted(statements)] + + +def verification_failures( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + has_privilege: Callable[[str, str, str], bool], + has_column_privilege: Callable[[str, str, str, str], bool] | None = None, +) -> list[PrivilegeEntry]: + """Return every declared privilege the role doesn't effectively hold. + + A column privilege is checked on each of its columns with + ``has_column_privilege``, which is required when one is declared. + """ + failures = [] + for role, tables in declared.items(): + for table, privileges in tables.items(): + for entry in _declared_entries(role, table, privileges): + if not entry.column: + held = has_privilege(role, table, entry.privilege) + elif has_column_privilege is None: + msg = "checking a column privilege needs has_column_privilege" + raise ValueError(msg) + else: + held = has_column_privilege( + role, table, entry.column, entry.privilege + ) + if not held: + failures.append(entry) + return sorted(failures) + + +def table_wide_overrides( + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + has_privilege: Callable[[str, str, str], bool], +) -> list[PrivilegeEntry]: + """Return each column privilege a role also holds on the whole table. + + Holding it on the whole table lets the role use it on every column, past + the columns the manifest declares. + """ + return sorted( + PrivilegeEntry(role, table, privilege.privilege.value) + for role, tables in declared.items() + for table, privileges in tables.items() + for privilege in privileges + if isinstance(privilege, ColumnPrivilege) + and has_privilege(role, table, privilege.privilege.value) + ) + + +def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: + """Return one row per role and table, naming each column privilege's columns.""" + table_privileges: dict[tuple[str, str], set[str]] = defaultdict(set) + columns_by_privilege: dict[tuple[str, str], dict[str, set[str]]] = defaultdict( + lambda: defaultdict(set) + ) + for entry in entries: + target = (entry.role, entry.table) + if entry.column: + columns_by_privilege[target][entry.privilege].add(entry.column) + else: + table_privileges[target].add(entry.privilege) + rows = [] + for role, table in sorted(set(table_privileges) | set(columns_by_privilege)): + described = sorted(table_privileges.get((role, table), ())) + [ + f"{privilege} ({', '.join(sorted(columns))})" + for privilege, columns in sorted( + columns_by_privilege.get((role, table), {}).items() + ) + ] + rows.append(f"| `{role}` | `{table}` | {', '.join(described)} |") + return rows + + +def skipped_roles_line(roles: Iterable[str]) -> str: + """Return the report sentence naming roles skipped until they are created.""" + return ( + "Declared roles not created yet, skipped until an operator creates them: " + + ", ".join(f"`{role}`" for role in roles) + ) + + +def render_summary( + plan: ReconciliationPlan, + default_privileges: Sequence[DefaultPrivilege], + *, + deployment_environment: str, + applied: bool, +) -> str: + """Return the Markdown report for the job summary and the log.""" + lines = [f"## v2 table privileges ({deployment_environment})", ""] + if plan.skipped_roles: + lines += [skipped_roles_line(plan.skipped_roles), ""] + if plan.missing_roles: + lines += [ + "Declared roles missing from the database (nothing was granted): " + + ", ".join(f"`{role}`" for role in plan.missing_roles), + "", + ] + if plan.missing_tables: + lines += [ + "Declared tables missing from the database (nothing was granted): " + + ", ".join(f"`{table}`" for table in plan.missing_tables), + "", + ] + grant_heading = "Granted" if applied else "Would grant" + lines += [f"### {grant_heading}", ""] + if plan.grants: + lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] + lines += _privilege_rows(plan.grants) + else: + lines.append("Nothing: every declared privilege was already granted.") + lines += [ + "", + "### Privileges the manifest doesn't declare (reported, not revoked)", + "", + ] + if plan.undeclared: + lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] + lines += _privilege_rows(plan.undeclared) + else: + lines.append("None.") + lines += ["", "### Default privileges for new tables", ""] + if default_privileges: + lines += [ + "| Creating role | Schema | Objects | Grantee | Privileges |", + "| --- | --- | --- | --- | --- |", + ] + lines += [ + f"| `{default.creating_role}` | `{default.schema}` | " + f"{default.object_type} | `{default.grantee}` | " + f"{', '.join(default.privileges)} |" + for default in default_privileges + ] + else: + lines.append("None.") + return "\n".join(lines) + "\n" + + +def _failure_row(failure: PrivilegeEntry) -> str: + if failure.column: + return ( + f"- `{failure.role}` lacks {failure.privilege} on column " + f"`{failure.column}` of `{failure.table}`" + ) + return f"- `{failure.role}` lacks {failure.privilege} on `{failure.table}`" + + +def verification_report( + failures: Sequence[PrivilegeEntry], overrides: Sequence[PrivilegeEntry] +) -> str: + """Return verify's report of missing privileges and whole-table overrides.""" + sections = [] + if failures: + rows = "\n".join(_failure_row(failure) for failure in failures) + sections.append(f"Declared privileges missing after reconciling:\n\n{rows}\n") + if overrides: + rows = "\n".join( + f"- `{override.role}` holds {override.privilege} on the whole of " + f"`{override.table}`, which overrides its declared columns" + for override in overrides + ) + sections.append( + f"Privileges held on whole tables, past their declared columns:\n\n{rows}\n" + ) + return "\n".join(sections) diff --git a/policyengine_api/data/v2/privilege_reconciliation.py b/policyengine_api/data/v2/privilege_reconciliation.py index d9d0cacec..008e33e87 100644 --- a/policyengine_api/data/v2/privilege_reconciliation.py +++ b/policyengine_api/data/v2/privilege_reconciliation.py @@ -4,12 +4,21 @@ connected as the migration role, which owns every table and so may grant on it. Two modes: -- ``apply`` reads the managed roles' current grants, grants every declared - privilege a role lacks, and reports privileges the manifest doesn't declare. - It never revokes, and it changes nothing when a declared role or table is - missing. Running it again grants nothing more. -- ``verify`` checks every declared privilege with ``has_table_privilege`` - (rule ``api.v2.deployed-privileges-match``) and fails on any that is missing. +- ``apply`` reads the managed roles' current grants, on whole tables and on + columns, grants every declared privilege a role lacks, and reports privileges + the manifest doesn't declare. It never revokes, and it changes nothing when a + declared role or table is missing. Running it again grants nothing more. +- ``verify`` checks every declared privilege with ``has_table_privilege``, or + ``has_column_privilege`` for each column of a column privilege (rule + ``api.v2.deployed-privileges-match``), and fails on any that is missing. It + also fails when a role holds on the whole table a privilege the manifest + declares for only some of its columns, since that lets the role use it on + every column. + +Both skip a missing role the manifest lists in ``ROLES_AWAITING_PROVISIONING``, +and name it in their output, so a pull request can declare a role before an +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 @@ -18,16 +27,17 @@ The report goes to stdout and, when set, ``GITHUB_STEP_SUMMARY``. It names roles, tables and privileges only, never a connection URL or credential. + +This module reads the database and runs the commands; ``privilege_plans`` +compares, states and reports the privileges, and its names are re-exported here. """ from __future__ import annotations import argparse import os -import re import sys from collections import defaultdict -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -40,9 +50,27 @@ load_v2_alembic_settings, qualify_v2_connection, ) +from policyengine_api.data.v2.privilege_plans import ( + SCHEMA, + DefaultPrivilege, + GrantedPrivileges, + PrivilegeEntry, + PrivilegeReconciliationError, + ReconciliationPlan, + UnsafeIdentifierError, + grant_statements, + plan_reconciliation, + render_summary, + safe_identifier, + skipped_roles_line, + table_wide_overrides, + verification_failures, + verification_report, +) from policyengine_api.data.v2.privileges import ( PERSISTENT_ENVIRONMENTS, - TablePrivilege, + ROLES_AWAITING_PROVISIONING, + DeclaredPrivilege, declared_privileges, ) from policyengine_api.data.v2.settings import V2ConfigurationError @@ -52,200 +80,39 @@ from sqlalchemy import Connection, Engine +# Names other modules import from here; the planning and reporting half lives +# in privilege_plans. +__all__ = [ + "DEPLOYMENT_ENVIRONMENT", + "GITHUB_STEP_SUMMARY", + "SCHEMA", + "DefaultPrivilege", + "GrantedPrivileges", + "PrivilegeEntry", + "PrivilegeReconciliationError", + "ReconciliationPlan", + "UnsafeIdentifierError", + "apply_declared_privileges", + "check_declared_privileges", + "check_table_wide_overrides", + "create_missing_roles", + "grant_statements", + "main", + "plan_reconciliation", + "read_default_privileges", + "read_existing_roles", + "read_granted_column_privileges", + "read_granted_privileges", + "read_public_tables", + "render_summary", + "run", + "skipped_roles", + "table_wide_overrides", + "verification_failures", +] + DEPLOYMENT_ENVIRONMENT = "DEPLOYMENT_ENVIRONMENT" GITHUB_STEP_SUMMARY = "GITHUB_STEP_SUMMARY" -SCHEMA = "public" -SAFE_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]{0,62}$") - - -class PrivilegeReconciliationError(RuntimeError): - """Raised when privileges can't be reconciled or don't match the manifest.""" - - -class UnsafeIdentifierError(PrivilegeReconciliationError): - """Raised for a role or table name that isn't a plain lowercase identifier.""" - - -@dataclass(frozen=True, order=True) -class PrivilegeEntry: - """One privilege one role holds, or should hold, on one table.""" - - role: str - table: str - privilege: str - - -@dataclass(frozen=True) -class DefaultPrivilege: - """A default privilege PostgreSQL gives on objects a role creates later.""" - - creating_role: str - schema: str - object_type: str - grantee: str - privileges: tuple[str, ...] - - -@dataclass(frozen=True) -class ReconciliationPlan: - """What ``apply`` grants and reports for one database.""" - - grants: tuple[PrivilegeEntry, ...] - undeclared: tuple[PrivilegeEntry, ...] - missing_roles: tuple[str, ...] - missing_tables: tuple[str, ...] - - @property - def can_apply(self) -> bool: - """Whether every declared role and table exists, so grants may run.""" - return not self.missing_roles and not self.missing_tables - - -def plan_reconciliation( - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], - granted: Mapping[str, Mapping[str, Collection[str]]], - *, - existing_roles: Collection[str], - existing_tables: Collection[str], -) -> ReconciliationPlan: - """Compare declared privileges with what the database grants. - - ``granted`` holds the managed roles' direct grants, by role then table. - """ - missing_roles = tuple(sorted(set(declared) - set(existing_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 - for table, privileges in tables.items() - if table not in missing_tables - for privilege in privileges - if privilege.value not in granted.get(role, {}).get(table, ()) - ) - undeclared = sorted( - PrivilegeEntry(role, table, privilege) - for role, tables in granted.items() - for table, privileges in tables.items() - for privilege in privileges - if privilege - not in { - declared_privilege.value - for declared_privilege in declared.get(role, {}).get(table, frozenset()) - } - ) - return ReconciliationPlan( - grants=tuple(grants), - undeclared=tuple(undeclared), - missing_roles=missing_roles, - missing_tables=missing_tables, - ) - - -def _safe(identifier: str) -> str: - if SAFE_IDENTIFIER.fullmatch(identifier) is None: - msg = f"refusing to use {identifier!r} as a role or table name" - raise UnsafeIdentifierError(msg) - return identifier - - -def grant_statements( - plan: ReconciliationPlan, quote_identifier: Callable[[str], str] -) -> list[str]: - """Return one GRANT per role and table, with its privileges in name order.""" - privileges_by_target: dict[tuple[str, str], set[str]] = defaultdict(set) - for entry in plan.grants: - privileges_by_target[(_safe(entry.role), _safe(entry.table))].add( - TablePrivilege(entry.privilege).value - ) - return [ - f"GRANT {', '.join(sorted(privileges))} ON TABLE " - f"{quote_identifier(SCHEMA)}.{quote_identifier(table)} " - f"TO {quote_identifier(role)}" - for (role, table), privileges in sorted( - privileges_by_target.items(), key=lambda item: (item[0][0], item[0][1]) - ) - ] - - -def verification_failures( - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], - has_privilege: Callable[[str, str, str], bool], -) -> list[PrivilegeEntry]: - """Return every declared privilege the role doesn't effectively hold.""" - return sorted( - PrivilegeEntry(role, table, privilege.value) - for role, tables in declared.items() - for table, privileges in tables.items() - for privilege in privileges - if not has_privilege(role, table, privilege.value) - ) - - -def _privilege_rows(entries: Iterable[PrivilegeEntry]) -> list[str]: - rows_by_target: dict[tuple[str, str], list[str]] = defaultdict(list) - for entry in entries: - rows_by_target[(entry.role, entry.table)].append(entry.privilege) - return [ - f"| `{role}` | `{table}` | {', '.join(sorted(privileges))} |" - for (role, table), privileges in sorted(rows_by_target.items()) - ] - - -def render_summary( - plan: ReconciliationPlan, - default_privileges: Sequence[DefaultPrivilege], - *, - deployment_environment: str, - applied: bool, -) -> str: - """Return the Markdown report for the job summary and the log.""" - lines = [f"## v2 table privileges ({deployment_environment})", ""] - if plan.missing_roles: - lines += [ - "Declared roles missing from the database (nothing was granted): " - + ", ".join(f"`{role}`" for role in plan.missing_roles), - "", - ] - if plan.missing_tables: - lines += [ - "Declared tables missing from the database (nothing was granted): " - + ", ".join(f"`{table}`" for table in plan.missing_tables), - "", - ] - grant_heading = "Granted" if applied else "Would grant" - lines += [f"### {grant_heading}", ""] - if plan.grants: - lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] - lines += _privilege_rows(plan.grants) - else: - lines.append("Nothing: every declared privilege was already granted.") - lines += [ - "", - "### Privileges the manifest doesn't declare (reported, not revoked)", - "", - ] - if plan.undeclared: - lines += ["| Role | Table | Privileges |", "| --- | --- | --- |"] - lines += _privilege_rows(plan.undeclared) - else: - lines.append("None.") - lines += ["", "### Default privileges for new tables", ""] - if default_privileges: - lines += [ - "| Creating role | Schema | Objects | Grantee | Privileges |", - "| --- | --- | --- | --- | --- |", - ] - lines += [ - f"| `{default.creating_role}` | `{default.schema}` | " - f"{default.object_type} | `{default.grantee}` | " - f"{', '.join(default.privileges)} |" - for default in default_privileges - ] - else: - lines.append("None.") - return "\n".join(lines) + "\n" # --- Database access ------------------------------------------------------- @@ -301,10 +168,34 @@ def render_summary( """ ).bindparams(bindparam("roles", expanding=True)) +_GRANTED_COLUMN_PRIVILEGES = text( + """ + SELECT grantee_role.rolname AS role_name, + table_class.relname AS table_name, + table_column.attname AS column_name, + acl_entry.privilege_type AS privilege + FROM pg_attribute AS table_column + JOIN pg_class AS table_class ON table_class.oid = table_column.attrelid + JOIN pg_namespace AS table_namespace + ON table_namespace.oid = table_class.relnamespace + CROSS JOIN LATERAL aclexplode(table_column.attacl) AS acl_entry + JOIN pg_roles AS grantee_role ON grantee_role.oid = acl_entry.grantee + WHERE table_namespace.nspname = :schema + AND table_class.relkind IN ('r', 'p', 'v', 'm') + AND table_column.attnum > 0 + AND NOT table_column.attisdropped + AND grantee_role.rolname IN :roles + """ +).bindparams(bindparam("roles", expanding=True)) + _HAS_TABLE_PRIVILEGE = text( "SELECT has_table_privilege(:role, :qualified_table, :privilege)" ) +_HAS_COLUMN_PRIVILEGE = text( + "SELECT has_column_privilege(:role, :qualified_table, :column, :privilege)" +) + def read_existing_roles(connection: Connection, roles: Collection[str]) -> set[str]: """Return which of ``roles`` exist.""" @@ -328,6 +219,23 @@ def read_granted_privileges( return {role: dict(tables) for role, tables in granted.items()} +def read_granted_column_privileges( + connection: Connection, roles: Collection[str] +) -> dict[str, dict[str, dict[str, set[str]]]]: + """Return the roles' direct grants on columns of public tables and views.""" + granted: dict[str, dict[str, dict[str, set[str]]]] = defaultdict( + lambda: defaultdict(lambda: defaultdict(set)) + ) + for row in connection.execute( + _GRANTED_COLUMN_PRIVILEGES, {"schema": SCHEMA, "roles": list(roles)} + ): + granted[row.role_name][row.table_name][row.column_name].add(row.privilege) + return { + role: {table: dict(columns) for table, columns in tables.items()} + for role, tables in granted.items() + } + + def read_default_privileges( connection: Connection, roles: Collection[str] ) -> list[DefaultPrivilege]: @@ -356,24 +264,31 @@ def create_missing_roles(connection: Connection, roles: Iterable[str]) -> list[s missing = sorted(set(roles) - read_existing_roles(connection, list(roles))) quote = connection.dialect.identifier_preparer.quote for role in missing: - connection.execute(text(f"CREATE ROLE {quote(_safe(role))} NOLOGIN")) + connection.execute(text(f"CREATE ROLE {quote(safe_identifier(role))} NOLOGIN")) return missing def apply_declared_privileges( connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + *, + roles_awaiting_creation: Collection[str] = (), ) -> tuple[ReconciliationPlan, list[DefaultPrivilege]]: """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( declared, - read_granted_privileges(connection, roles), + GrantedPrivileges( + tables=read_granted_privileges(connection, roles), + columns=read_granted_column_privileges(connection, roles), + ), existing_roles=read_existing_roles(connection, roles), existing_tables=read_public_tables(connection), + roles_awaiting_creation=roles_awaiting_creation, ) if plan.can_apply: for statement in grant_statements( @@ -383,15 +298,26 @@ def apply_declared_privileges( return plan, read_default_privileges(connection, roles) -def check_declared_privileges( +def skipped_roles( connection: Connection, - declared: Mapping[str, Mapping[str, frozenset[TablePrivilege]]], -) -> list[PrivilegeEntry]: - """Return every declared privilege a role doesn't effectively hold. + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + roles_awaiting_creation: Collection[str], +) -> tuple[str, ...]: + """Return the declared roles awaiting creation that don't exist yet.""" + existing_roles = read_existing_roles(connection, list(declared)) + return tuple( + sorted((set(declared) & set(roles_awaiting_creation)) - existing_roles) + ) + + +def _privilege_checks( + connection: Connection, roles: Collection[str] +) -> tuple[Callable[[str, str, str], bool], Callable[[str, str, str, str], bool]]: + """Return checks of a role's effective table and column privileges. - A missing role or table counts as missing every privilege declared for it. + A missing role or table holds nothing. """ - existing_roles = read_existing_roles(connection, list(declared)) + existing_roles = read_existing_roles(connection, roles) existing_tables = read_public_tables(connection) def has_privilege(role: str, table: str, privilege: str) -> bool: @@ -402,13 +328,52 @@ def has_privilege(role: str, table: str, privilege: str) -> bool: _HAS_TABLE_PRIVILEGE, { "role": role, - "qualified_table": f"{SCHEMA}.{_safe(table)}", + "qualified_table": f"{SCHEMA}.{safe_identifier(table)}", "privilege": privilege, }, ).scalar_one() ) - return verification_failures(declared, has_privilege) + def has_column_privilege( + role: str, table: str, column: str, privilege: str + ) -> bool: + if role not in existing_roles or table not in existing_tables: + return False + return bool( + connection.execute( + _HAS_COLUMN_PRIVILEGE, + { + "role": role, + "qualified_table": f"{SCHEMA}.{safe_identifier(table)}", + "column": safe_identifier(column), + "privilege": privilege, + }, + ).scalar_one() + ) + + return has_privilege, has_column_privilege + + +def check_declared_privileges( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], +) -> list[PrivilegeEntry]: + """Return every declared privilege a role doesn't effectively hold. + + A missing role or table counts as missing every privilege declared for it; + leave out the roles ``skipped_roles`` returns to skip them. + """ + has_privilege, has_column_privilege = _privilege_checks(connection, list(declared)) + return verification_failures(declared, has_privilege, has_column_privilege) + + +def check_table_wide_overrides( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], +) -> list[PrivilegeEntry]: + """Return each column privilege a role also holds on the whole table.""" + has_privilege, _has_column_privilege = _privilege_checks(connection, list(declared)) + return table_wide_overrides(declared, has_privilege) # --- Command-line interface ----------------------------------------------- @@ -438,6 +403,44 @@ def _write_summary(environ: Mapping[str, str], summary: str) -> None: summary_file.write(summary) +def _verify( + connection: Connection, + declared: Mapping[str, Mapping[str, frozenset[DeclaredPrivilege]]], + 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. A role holding on a + whole table a privilege declared for only some of its columns fails too. + The roles skipped, the privileges missing and the whole-table privileges + are written to stdout and the job summary under one heading. + """ + skipped = skipped_roles(connection, declared, roles_awaiting_creation) + checked = {role: tables for role, tables in declared.items() if role not in skipped} + failures = check_declared_privileges(connection, checked) + overrides = check_table_wide_overrides(connection, checked) + heading = f"## v2 table privileges ({environment})\n\n" + skipped_note = f"{skipped_roles_line(skipped)}\n\n" if skipped else "" + if failures or overrides: + _write_summary( + environ, + f"{heading}{skipped_note}{verification_report(failures, overrides)}", + ) + 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 +460,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/privilege_requirements.py b/policyengine_api/data/v2/privilege_requirements.py index f159f747f..ea860adfc 100644 --- a/policyengine_api/data/v2/privilege_requirements.py +++ b/policyengine_api/data/v2/privilege_requirements.py @@ -25,7 +25,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from policyengine_api.data.v2.privileges import TablePrivilege +from policyengine_api.data.v2.privileges import DeclaredPrivilege, TablePrivilege if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping @@ -226,9 +226,13 @@ def required_privileges( def undeclared_requirements( requirements: Iterable[RequiredPrivilege], - declared: Mapping[str, frozenset[TablePrivilege]], + declared: Mapping[str, frozenset[DeclaredPrivilege]], ) -> list[RequiredPrivilege]: - """Return the requirements the declared table privileges don't cover, sorted.""" + """Return the requirements the declared table privileges don't cover, sorted. + + A privilege declared on only some columns doesn't cover a connector's need + for it on the whole table. + """ return sorted( requirement for requirement in requirements diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 7bf2897c8..18f3ab8f1 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -10,7 +10,10 @@ ``api.v2.runtime-privileges-tested-as-runtime-role``). Every table in the v2 SQLModel metadata, plus ``alembic_version``, has an entry -for every managed role, so "no access" is a decision, not an omission. +for every managed role, so "no access" is a decision, not an omission. A +privilege is on the whole table, or a ``ColumnPrivilege`` on the columns it +names; a role never holds both kinds of the same privilege on one table, since +the whole-table one would let it change every column. Roles and what decides their privileges: @@ -24,6 +27,23 @@ (``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, budget-window runs and household + jobs, and reads the policy, dynamic, household and catalog rows a run needs. + It starts the Modal calls of runs and household jobs, so it may update their + status, Modal call ID, error and completion columns: it records a call's ID + once Modal accepts the call, and fails the run when Modal refuses it. +- ``policyengine_v2_modal_worker``: the simulation API's Modal workers. They + write the status of the runs and household jobs they were given, and their + results, and read the rows a run needs. + +Both simulation API roles update only the columns ``update_columns`` names, so +the database refuses any change to what a run is (its simulation or report, +versions, dataset, region, year, simulation runs or budget-window run). It +doesn't limit a worker to its own run; the workers' code does. Their code lives +in the simulation API's repository, so the deploy check covers them here. + +The Public API reads the simulation tables it serves, and writes none of them. The migration role (``policyengine_v2_migrator``) owns every table, and Supabase's ``postgres`` role administers the project; neither is managed here. @@ -36,14 +56,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}$") @@ -56,6 +86,25 @@ class TablePrivilege(StrEnum): DELETE = "DELETE" +@dataclass(frozen=True) +class ColumnPrivilege: + """A privilege on some of a table's columns instead of the whole table. + + PostgreSQL refuses a statement that uses the privilege on any other column. + """ + + privilege: TablePrivilege + columns: frozenset[str] + + +DeclaredPrivilege = TablePrivilege | ColumnPrivilege + + +def update_columns(*columns: str) -> ColumnPrivilege: + """Declare ``UPDATE`` on only the named columns.""" + return ColumnPrivilege(TablePrivilege.UPDATE, frozenset(columns)) + + NO_ACCESS: Final[frozenset[TablePrivilege]] = frozenset() READ: Final = frozenset({TablePrivilege.SELECT}) READ_INSERT: Final = frozenset({TablePrivilege.SELECT, TablePrivilege.INSERT}) @@ -76,43 +125,115 @@ class TablePrivilege(StrEnum): class TableAccess: """What each managed role may do on one table.""" - runtime: frozenset[TablePrivilege] - catalog_writer: frozenset[TablePrivilege] - stage12: frozenset[TablePrivilege] + runtime: frozenset[DeclaredPrivilege] + catalog_writer: frozenset[DeclaredPrivilege] + stage12: frozenset[DeclaredPrivilege] + simulation_service: frozenset[DeclaredPrivilege] + modal_worker: frozenset[DeclaredPrivilege] def _access( *, - runtime: frozenset[TablePrivilege] = NO_ACCESS, - catalog_writer: frozenset[TablePrivilege] = NO_ACCESS, - stage12: frozenset[TablePrivilege] = NO_ACCESS, + runtime: frozenset[DeclaredPrivilege] = NO_ACCESS, + catalog_writer: frozenset[DeclaredPrivilege] = NO_ACCESS, + stage12: frozenset[DeclaredPrivilege] = NO_ACCESS, + simulation_service: frozenset[DeclaredPrivilege] = NO_ACCESS, + modal_worker: frozenset[DeclaredPrivilege] = NO_ACCESS, ) -> TableAccess: - return TableAccess(runtime=runtime, 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 and its Modal workers read them. +_SIMULATION_DEFINITION = _access( + runtime=READ, simulation_service=READ_INSERT, modal_worker=READ +) +# A run's state, which the Modal worker running it writes. +_RUN_STATE_COLUMNS: Final = ( + "status", + "modal_call_id", + "error_message", + "started_at", + "completed_at", + "updated_at", +) +# The state the simulation service writes for a run whose Modal call it starts: +# the call's ID once Modal accepts the call, or the failure when Modal refuses +# it. +_CALL_START_COLUMNS: Final = ( + "status", + "modal_call_id", + "error_message", + "completed_at", + "updated_at", +) + + +def _run_table( + *result_columns: str, catalog_writer: frozenset[DeclaredPrivilege] = NO_ACCESS +) -> TableAccess: + """Access to a run table whose Modal worker writes the named result columns.""" + return _access( + runtime=READ, + catalog_writer=catalog_writer, + simulation_service=READ_INSERT | {update_columns(*_CALL_START_COLUMNS)}, + modal_worker=READ | {update_columns(*_RUN_STATE_COLUMNS, *result_columns)}, + ) + + +# The Modal workers write results, keyed by run, once each. +_RESULT = _access(runtime=READ, modal_worker=READ_INSERT) +# What a run needs from the Public API's policies, dynamics, households and +# 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. + "budget_window_runs": _SIMULATION_DEFINITION, + "household_jobs": _run_table("result"), + "report_runs": _run_table( + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", + "spm_provenance", + "data_version", + catalog_writer=READ, + ), + "reports": replace(_SIMULATION_DEFINITION, catalog_writer=READ), + "simulation_runs": _run_table("output_uri", "output_sha256", "household_result"), + "simulations": replace(_SIMULATION_DEFINITION, catalog_writer=READ), + # Results of simulation runs and report runs. + "labor_supply_decile_responses": _RESULT, + "labor_supply_responses": _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 +241,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), @@ -174,7 +303,7 @@ def stage12_role(deployment_environment: str) -> str: def declared_privileges( deployment_environment: str, -) -> dict[str, dict[str, frozenset[TablePrivilege]]]: +) -> dict[str, dict[str, frozenset[DeclaredPrivilege]]]: """Return each managed role's declared privileges, keyed by role then table.""" stage12 = stage12_role(deployment_environment) return { @@ -183,4 +312,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..f1980aadd 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -1,4 +1,11 @@ -"""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. ``complete_report_run`` sets +``markdown``, which no managed role may update either. The simulation API +creates and completes reports and report runs with its own code. +""" from __future__ import annotations @@ -47,6 +54,25 @@ class ReportRunPackageVersions: policyengine_version: str +@dataclass(frozen=True) +class ReportRunSimulationRuns: + """The simulation runs one report run compares. + + A household report with only a baseline has no reform run. + """ + + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + +@dataclass(frozen=True) +class ReportRunInputs: + """The package versions and simulation runs one report run uses.""" + + versions: ReportRunPackageVersions + simulation_runs: ReportRunSimulationRuns + + def _locked_report(session: Session, report_id: UUID) -> Report: # SQLModel's select remains the query surface. Row locking is the bounded # SQLAlchemy capability needed to serialize first-run/type updates. @@ -86,21 +112,28 @@ def create_report_run( session: Session, *, report_id: UUID, - versions: ReportRunPackageVersions, + inputs: ReportRunInputs, trigger: ReportRunTrigger, idempotency_key: UUID | None = None, ) -> ReportRun: - """Create one run, or return the run for a retried idempotent request.""" + """Create one run, or return the run for a retried idempotent request. + + A retried request gets the run its first attempt created, with the + simulation runs that attempt named. + """ _locked_report(session, report_id) if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: msg = "manual report reruns require an idempotency key" raise ValueError(msg) + versions = inputs.versions if not versions.country_package_version or not versions.policyengine_version: msg = "report run package versions must be non-empty" raise ValueError(msg) run = ReportRun( report_id=report_id, + baseline_simulation_run_id=inputs.simulation_runs.baseline_simulation_run_id, + reform_simulation_run_id=inputs.simulation_runs.reform_simulation_run_id, country_package_version=versions.country_package_version, policyengine_version=versions.policyengine_version, trigger=trigger, diff --git a/pyproject.toml b/pyproject.toml index 024e94cb9..a1e0e213f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -275,6 +275,8 @@ ignore_imports = [ "policyengine_api.data.v2.models.policy_mappings -> sqlmodel", "policyengine_api.data.v2.models.reports -> sqlalchemy", "policyengine_api.data.v2.models.reports -> sqlmodel", + "policyengine_api.data.v2.models.run_results -> sqlalchemy", + "policyengine_api.data.v2.models.run_results -> sqlmodel", "policyengine_api.data.v2.models.simulations -> sqlalchemy", "policyengine_api.data.v2.models.simulations -> sqlmodel", "policyengine_api.data.v2.models.users -> sqlalchemy", @@ -421,7 +423,8 @@ runtime-evaluated-base-classes = [ "policyengine_api.data.v2.models.base.IdentifiedModel", "policyengine_api.data.v2.models.base.TimestampedModel", "policyengine_api.data.v2.models.comparison_runs.Stage12RunTimestamps", - "policyengine_api.data.v2.models.reports.GeographicImpactBase", + "policyengine_api.data.v2.models.run_results.GeographicImpactBase", + "policyengine_api.data.v2.models.run_results.UKAreaImpactBase", "policyengine_api.fastapi_routes.v2.households.document_models.HouseholdEntityRecord", "policyengine_api.fastapi_routes.v2.households.document_models.StrictHouseholdDocumentModel", "policyengine_api.fastapi_routes.v2.households.response_models.HouseholdSuccessResponse", diff --git a/rules/registry.toml b/rules/registry.toml index bd7cba33f..458e46d69 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -440,6 +440,8 @@ allowed_ignore_imports = [ "policyengine_api.data.v2.models.policy_mappings -> sqlmodel", "policyengine_api.data.v2.models.reports -> sqlalchemy", "policyengine_api.data.v2.models.reports -> sqlmodel", + "policyengine_api.data.v2.models.run_results -> sqlalchemy", + "policyengine_api.data.v2.models.run_results -> sqlmodel", "policyengine_api.data.v2.models.simulations -> sqlalchemy", "policyengine_api.data.v2.models.simulations -> sqlmodel", "policyengine_api.data.v2.models.users -> sqlalchemy", @@ -1018,7 +1020,7 @@ invocation = "v2-integration-check.yml" [[rule]] id = "api.v2.runtime-privileges-declared" category = "migration" -statement = "Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py, and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it." +statement = "Every v2 table has an entry for each managed database role in policyengine_api/data/v2/privileges.py; a privilege declared on only some of a table's columns names columns of that table and is never also declared on the whole table; and every privilege a v2 connector module needs (from its reads, creates, updates or deletes module name and its row locks) is declared for the role that runs it, on the whole table." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/models/*.py", "policyengine_api/services/v2/*/database_connectors/*.py"] tier = "mechanical" @@ -1028,6 +1030,7 @@ autofix = false kind = "pytest" nodes = [ "tests/unit/v2/test_privilege_manifest.py::test_every_v2_table_has_an_entry_for_every_managed_role", + "tests/unit/v2/test_privilege_manifest.py::test_column_privileges_name_their_tables_columns_alone", "tests/unit/v2/test_privilege_manifest.py::test_v2_schema_has_no_sequences_the_manifest_would_need_to_cover", "tests/unit/v2/test_privilege_manifest.py::test_runtime_connectors_need_only_declared_runtime_privileges", "tests/unit/v2/test_privilege_manifest.py::test_stage12_connectors_need_only_declared_stage12_privileges", @@ -1044,9 +1047,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 @@ -1064,9 +1067,9 @@ 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, on the whole table or on the columns it names, right after the schema upgrade, reports privileges it doesn't declare without revoking them, and stops before the catalog is seeded if a declared privilege is still missing or a role holds on a whole table a privilege declared for only some of its columns. A declared role listed in ROLES_AWAITING_PROVISIONING that doesn't exist yet is named in the job summary and skipped instead." rationale = "docs/engineering/skills/alembic-migrations.md#table-privileges" -scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] +scope = ["policyengine_api/data/v2/privileges.py", "policyengine_api/data/v2/privilege_plans.py", "policyengine_api/data/v2/privilege_reconciliation.py", "scripts/reconcile_v2_privileges.py", ".github/workflows/seed-v2-database.yml"] tier = "mechanical" status = "enforced" autofix = false diff --git a/tests/integration/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..214914f8e --- /dev/null +++ b/tests/integration/simulation_rows.py @@ -0,0 +1,114 @@ +"""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 dataclasses import dataclass +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +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"} + + +@dataclass(frozen=True) +class PublicApiRows: + """The Public API's rows a simulation, report or household job names.""" + + model_id: UUID + version_id: UUID + dataset_id: UUID + household_id: UUID + + +def insert_public_api_rows(connection: Connection) -> PublicApiRows: + """Store a model, one of its versions, a dataset and a US household. + + The simulation tables' foreign keys need them. Run it as the disposable + database's superuser: the simulation API's roles can't write these tables. + """ + rows = PublicApiRows( + model_id=uuid4(), version_id=uuid4(), dataset_id=uuid4(), household_id=uuid4() + ) + # Model names are unique and at most 32 characters. + insert_row( + connection, + "tax_benefit_models", + id=rows.model_id, + name=f"runs-{rows.model_id.hex[:24]}", + ) + insert_row( + connection, + "tax_benefit_model_versions", + id=rows.version_id, + model_id=rows.model_id, + version="1.0.0", + current_law_id=1, + metadata_time_periods="[2026]", + ) + insert_row( + connection, + "datasets", + id=rows.dataset_id, + name="populace", + year=2026, + is_output_dataset=False, + tax_benefit_model_version_id=rows.version_id, + ) + insert_row( + connection, + "households", + id=rows.household_id, + country_id="us", + household_data="{}", + canonicalization_version=1, + content_hash=uuid4().hex * 2, + ) + return rows + + +def insert_row(connection: Connection, table: str, **values: object) -> None: + """Insert one row of the named columns into ``table``.""" + columns = ", ".join(values) + 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() + + +def delete_public_api_rows(connection: Connection, rows: PublicApiRows) -> None: + """Delete what ``insert_public_api_rows`` stored, once nothing names it.""" + for table, row_id in ( + ("households", rows.household_id), + ("datasets", rows.dataset_id), + ("tax_benefit_model_versions", rows.version_id), + ("tax_benefit_models", rows.model_id), + ): + connection.execute(text(f"DELETE FROM {table} WHERE id = :id"), {"id": row_id}) diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 3c2cc8161..8e3086f15 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -25,6 +25,13 @@ disposable_v2_url, table_columns, ) +from tests.integration.simulation_rows import ( + RUN_VERSIONS, + PublicApiRows, + delete_public_api_rows, + insert_public_api_rows, + insert_row, +) BASELINE_REVISION = "f5ef4347cb2a" STAGE_11_PREVIOUS_REVISION = "af34023a728f" @@ -886,19 +893,36 @@ def report_run_checks() -> set[str]: for constraint in inspect(engine).get_check_constraints("report_runs") } + report_id = uuid4() + report_run_id = uuid4() + simulation_id = uuid4() + request_key = uuid4() + public_api_rows: PublicApiRows | None = None try: command.upgrade(config, "head") assert isinstance(idempotency_column_type(), postgresql.UUID) assert "ck_report_runs_idempotency_key_nonblank" not in report_run_checks() - model_id = uuid4() - report_id = uuid4() - report_run_id = uuid4() - request_key = uuid4() with engine.begin() as connection: - connection.execute( - text("INSERT INTO tax_benefit_models (id, name) VALUES (:id, :name)"), - {"id": model_id, "name": f"uuid-cast-{model_id.hex[:8]}"}, + public_api_rows = insert_public_api_rows(connection) + insert_row( + connection, + "simulations", + id=simulation_id, + simulation_type="economy", + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, + ) + simulation_run_id = uuid4() + insert_row( + connection, + "simulation_runs", + id=simulation_run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, + status="pending", + **RUN_VERSIONS, ) connection.execute( text( @@ -906,19 +930,21 @@ def report_run_checks() -> set[str]: "(id, label, country, tax_benefit_model_id, inputs) " "VALUES (:id, 'UUID cast report', 'us', :model_id, '{}')" ), - {"id": report_id, "model_id": model_id}, + {"id": report_id, "model_id": public_api_rows.model_id}, ) connection.execute( text( "INSERT INTO report_runs " - "(id, report_id, country_package_version, " - "policyengine_version, status, trigger, idempotency_key) " - "VALUES (:id, :report_id, '1.0', '1.0', 'pending', " - "'manual', :request_key)" + "(id, report_id, baseline_simulation_run_id, " + "country_package_version, policyengine_version, status, " + "trigger, idempotency_key) " + "VALUES (:id, :report_id, :simulation_run_id, '1.0', '1.0', " + "'pending', 'manual', :request_key)" ), { "id": report_run_id, "report_id": report_id, + "simulation_run_id": simulation_run_id, "request_key": str(request_key), }, ) @@ -941,9 +967,11 @@ def report_run_checks() -> set[str]: {"report_id": report_id}, ) connection.execute( - text("DELETE FROM tax_benefit_models WHERE id = :model_id"), - {"model_id": model_id}, + text("DELETE FROM simulations WHERE id = :simulation_id"), + {"simulation_id": simulation_id}, ) + if public_api_rows is not None: + delete_public_api_rows(connection, public_api_rows) command.upgrade(config, "head") engine.dispose() diff --git a/tests/integration/test_alembic_v2_simulation_runs.py b/tests/integration/test_alembic_v2_simulation_runs.py new file mode 100644 index 000000000..7657f16b9 --- /dev/null +++ b/tests/integration/test_alembic_v2_simulation_runs.py @@ -0,0 +1,795 @@ +"""Check the simulation runs revision against disposable Postgres. + +One revision follows the parameter search file revision. It adds simulation +runs, one per simulation and year; has each report run name its baseline and +reform simulation runs; keys every result table by run; types a report run's +cliff impact and labour supply; adds the budget-window tables; gives household +jobs one household, one year and one policy; and adds the economy output's +fields that had no column. The foreign keys from those tables to the Public +API's tables stay, and the new columns that name the Public API's rows get +keys too. It changes no other table. The previous Public API revision keeps +working while it is deployed: its roles can't read or write any table it +changes, except the catalog writer, which counts the rows of simulations, +reports and report_runs. +""" + +from __future__ import annotations + +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, + PublicApiRows, + insert_public_api_rows, + insert_row, + refused, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection, Engine + +SIMULATION_RUNS_PREVIOUS_REVISION = "c614b6bcccfc" +NEW_TABLES = frozenset( + { + "simulation_runs", + "budget_window_runs", + "labor_supply_responses", + "labor_supply_decile_responses", + } +) +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"} +) +# The foreign keys from the simulation service's tables to the Public API's +# rows, with their delete rules. +PUBLIC_API_FOREIGN_KEYS = { + ("simulations", "dataset_id"): ("datasets", "RESTRICT"), + ("simulations", "household_id"): ("households", "RESTRICT"), + ("simulations", "policy_id"): ("policies", "SET NULL"), + ("simulations", "dynamic_id"): ("dynamics", "SET NULL"), + ("simulations", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulations", "region_id"): ("regions", "SET NULL"), + ("reports", "user_id"): ("users", "SET NULL"), + ("reports", "tax_benefit_model_id"): ("tax_benefit_models", "RESTRICT"), + ("reports", "policy_id"): ("policies", "SET NULL"), + ("reports", "household_id"): ("households", "RESTRICT"), + ("reports", "dataset_id"): ("datasets", "SET NULL"), + ("reports", "region_id"): ("regions", "SET NULL"), + ("simulation_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulation_runs", "dataset_id"): ("datasets", "RESTRICT"), + ("simulation_runs", "region_id"): ("regions", "SET NULL"), + ("household_jobs", "household_id"): ("households", "RESTRICT"), + ("household_jobs", "policy_id"): ("policies", "SET NULL"), + ("household_jobs", "dynamic_id"): ("dynamics", "SET NULL"), + ("household_jobs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("budget_window_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), +} +LABOR_SUPPLY_RESPONSE_COLUMNS = ( + "income_effect_earnings_change", + "substitution_effect_earnings_change", + "total_earnings_change", + "revenue_change", + "relative_income_effect", + "relative_substitution_effect", + "baseline_weekly_hours", + "reform_weekly_hours", + "weekly_hours_change", + "weekly_hours_income_effect", + "weekly_hours_substitution_effect", +) +SHA256_HEX = "a" * 64 +COMPLETED_AT = "2026-11-04T00:00:00Z" + + +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_revision_downgrades_to_its_parent_and_upgrades_again( + config: Config, engine: Engine +) -> None: + """The revision adds its tables and changes only simulation-side tables.""" + command.upgrade(config, "head") + + command.downgrade(config, SIMULATION_RUNS_PREVIOUS_REVISION) + before = _shapes(engine) + assert NEW_TABLES.isdisjoint(before) + assert {"status", "household_result", "output_dataset_id"} <= _column_names( + engine, "simulations" + ) + assert {"report_run_id", "simulation_id"} <= _column_names(engine, "poverty") + assert {"datasets", "households", "policies"} <= { + key["referred_table"] for key in inspect(engine).get_foreign_keys("simulations") + } + program_entity = next( + column + for column in inspect(engine).get_columns("program_statistics") + if column["name"] == "entity" + ) + assert not program_entity["nullable"] + assert "v2_simulation_status" in _enum_type_names(engine) + + command.upgrade(config, "head") + command.check(config) + after = _shapes(engine) + assert set(after) - set(before) == NEW_TABLES + assert set(before) <= set(after) + changed = { + table_name for table_name in before if before[table_name] != after[table_name] + } + assert changed == CHANGED_TABLES + assert {"baseline_simulation_run_id", "reform_simulation_run_id"} <= ( + _column_names(engine, "report_runs") + ) + + +def test_head_keeps_the_foreign_keys_to_the_public_api_tables( + config: Config, engine: Engine +) -> None: + """Each service writes only its own tables; the keys between them stay.""" + command.upgrade(config, "head") + inspector = inspect(engine) + + public_api_keys = { + (table_name, column_name): (key["referred_table"], key["options"]["ondelete"]) + for table_name in SIMULATION_SIDE_TABLES + for key in inspector.get_foreign_keys(table_name) + if key["referred_table"] not in SIMULATION_SIDE_TABLES + for column_name in key["constrained_columns"] + } + + assert public_api_keys == PUBLIC_API_FOREIGN_KEYS + assert {"ix_reports_user_id"} <= { + index["name"] for index in inspector.get_indexes("reports") + } + + +def _count(connection: Connection, table: str, column: str, value: object) -> int: + return connection.execute( + text(f"SELECT count(*) FROM {table} WHERE {column} = :value"), + {"value": value}, + ).scalar_one() + + +@pytest.fixture +def public_api_rows(connection: Connection) -> PublicApiRows: + return insert_public_api_rows(connection) + + +def _insert_simulation( + connection: Connection, public_api_rows: PublicApiRows +) -> object: + """Store an economy simulation of the stored dataset and catalog version.""" + simulation_id = uuid4() + insert_row( + connection, + "simulations", + id=simulation_id, + simulation_type="economy", + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, + ) + return simulation_id + + +def _insert_simulation_run( + connection: Connection, + simulation_id: object, + *, + status: str = "pending", + year: int = 2026, + **values: object, +) -> object: + """Store a run of the simulation for ``year`` on its catalog version.""" + run_id = uuid4() + version_id = connection.execute( + text("SELECT tax_benefit_model_version_id FROM simulations WHERE id = :id"), + {"id": simulation_id}, + ).scalar_one() + insert_row( + connection, + "simulation_runs", + id=run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=version_id, + status=status, + year=year, + **RUN_VERSIONS, + **values, + ) + return run_id + + +def _insert_report_run( + connection: Connection, + public_api_rows: PublicApiRows, + baseline_run_id: object, + reform_run_id: object | None = None, +) -> object: + """Store a report and one run of it comparing the two simulation runs.""" + report_id, report_run_id = uuid4(), uuid4() + insert_row( + connection, + "reports", + id=report_id, + label="Report", + country="us", + tax_benefit_model_id=public_api_rows.model_id, + inputs="{}", + ) + insert_row( + connection, + "report_runs", + id=report_run_id, + report_id=report_id, + baseline_simulation_run_id=baseline_run_id, + reform_simulation_run_id=reform_run_id, + status="running", + trigger="initial", + **RUN_VERSIONS, + ) + return report_run_id + + +def test_rows_naming_a_missing_public_api_row_are_refused( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + """PostgreSQL checks each row a simulation-side row names when it is written.""" + simulation_id = _insert_simulation(connection, public_api_rows) + + with refused(connection, match="foreign key constraint"): + insert_row( + connection, + "simulation_runs", + id=uuid4(), + simulation_id=simulation_id, + tax_benefit_model_version_id=uuid4(), + status="pending", + year=2026, + **RUN_VERSIONS, + ) + with refused(connection, match="foreign key constraint"): + insert_row( + connection, + "household_jobs", + id=uuid4(), + country="us", + household_id=uuid4(), + year=2026, + request_data="{}", + tax_benefit_model_version_id=public_api_rows.version_id, + status="pending", + **RUN_VERSIONS, + ) + + +def test_simulation_runs_need_a_completion_time_and_a_whole_output( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + simulation_id = _insert_simulation(connection, public_api_rows) + 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_runs_name_a_baseline_run_and_keep_their_simulation_runs( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + simulation_id = _insert_simulation(connection, public_api_rows) + baseline_run_id = _insert_simulation_run(connection, simulation_id) + reform_run_id = _insert_simulation_run(connection, simulation_id) + + with refused(connection, match="baseline_simulation_run_id"): + _insert_report_run(connection, public_api_rows, None) + with refused(connection, match="foreign key constraint"): + _insert_report_run(connection, public_api_rows, uuid4()) + # A household report with only a baseline names no reform run. + _insert_report_run(connection, public_api_rows, baseline_run_id) + report_run_id = _insert_report_run( + connection, public_api_rows, baseline_run_id, reform_run_id + ) + for run_id in (baseline_run_id, reform_run_id): + with refused(connection, match="foreign key constraint"): + connection.execute( + text("DELETE FROM simulation_runs WHERE id = :id"), {"id": run_id} + ) + connection.execute( + text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} + ) + assert _count(connection, "simulation_runs", "id", reform_run_id) == 1 + + +def test_results_are_unique_per_run_and_name_runs_that_exist( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + simulation_id = _insert_simulation(connection, public_api_rows) + baseline_run_id = _insert_simulation_run(connection, simulation_id) + reform_run_id = _insert_simulation_run(connection, simulation_id) + comparison_runs = { + "report_run_id": _insert_report_run( + connection, public_api_rows, baseline_run_id, reform_run_id + ), + "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) + # A report run compares one pair, so its rows are unique within it. + with refused(connection): + insert_row( + connection, + "budget_summary", + id=uuid4(), + **budget, + **{**comparison_runs, "reform_simulation_run_id": baseline_run_id}, + ) + 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_a_report_runs_labor_supply_and_cliff_impact_are_typed( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + simulation_id = _insert_simulation(connection, public_api_rows) + report_run_id = _insert_report_run( + connection, + public_api_rows, + _insert_simulation_run(connection, simulation_id), + _insert_simulation_run(connection, simulation_id), + ) + response = dict.fromkeys(LABOR_SUPPLY_RESPONSE_COLUMNS, 1.0) + + insert_row( + connection, + "labor_supply_responses", + id=uuid4(), + report_run_id=report_run_id, + **response, + ) + with refused(connection): + insert_row( + connection, + "labor_supply_responses", + id=uuid4(), + report_run_id=report_run_id, + **response, + ) + # .py reports a relative response only for deciles above zero. + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=0, + average_income_effect=-1.0, + average_substitution_effect=2.0, + ) + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=1, + average_income_effect=-1.0, + average_substitution_effect=2.0, + relative_income_effect=-0.01, + relative_substitution_effect=0.02, + ) + with refused(connection): + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=report_run_id, + decile=1, + ) + with refused(connection, match="ck_report_runs_cliff_impact_complete"): + connection.execute( + text("UPDATE report_runs SET baseline_cliff_gap = 1.0 WHERE id = :id"), + {"id": report_run_id}, + ) + connection.execute( + text( + "UPDATE report_runs SET baseline_cliff_gap = 1.0, " + "baseline_cliff_share = 0.1, reform_cliff_gap = 2.0, " + "reform_cliff_share = 0.2, data_version = '1.0.0', " + "spm_provenance = CAST(:receipt AS jsonb) WHERE id = :id" + ), + {"id": report_run_id, "receipt": '{"source": "census"}'}, + ) + receipt_type = connection.execute( + text("SELECT jsonb_typeof(spm_provenance) FROM report_runs WHERE id = :id"), + {"id": report_run_id}, + ).scalar_one() + assert receipt_type == "object" + + connection.execute( + text("DELETE FROM report_runs WHERE id = :id"), {"id": report_run_id} + ) + for table in ("labor_supply_responses", "labor_supply_decile_responses"): + assert _count(connection, table, "report_run_id", report_run_id) == 0 + + +def _insert_report( + connection: Connection, public_api_rows: PublicApiRows, **values: object +) -> object: + report_id = uuid4() + insert_row( + connection, + "reports", + id=report_id, + label="Budget window", + country="us", + tax_benefit_model_id=public_api_rows.model_id, + inputs="{}", + **values, + ) + return report_id + + +def _insert_budget_window_run( + connection: Connection, public_api_rows: PublicApiRows, report_id: object +) -> object: + window_run_id = uuid4() + insert_row( + connection, + "budget_window_runs", + id=window_run_id, + report_id=report_id, + tax_benefit_model_version_id=public_api_rows.version_id, + trigger="initial", + recalculate=False, + status="running", + ) + return window_run_id + + +def test_budget_windows_hold_one_report_run_per_year( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + for refused_values in ( + {"year": 2026, "budget_window_size": 0}, + {"year": 2026, "budget_window_size": 76}, + {"year": 2190, "budget_window_size": 12}, + {"budget_window_size": 3}, + {"year": 2026, "budget_window_size": 3, "household_id": None}, + ): + values = dict(refused_values) + if "household_id" in values: + values["household_id"] = public_api_rows.household_id + with refused(connection, match="ck_reports_budget_window_size"): + _insert_report(connection, public_api_rows, **values) + report_id = _insert_report( + connection, public_api_rows, year=2026, budget_window_size=3 + ) + other_report_id = _insert_report(connection, public_api_rows, year=2026) + window_run_id = _insert_budget_window_run(connection, public_api_rows, report_id) + simulation_id = _insert_simulation(connection, public_api_rows) + + def insert_year(year: int | None, *, of_report: object = report_id) -> None: + insert_row( + connection, + "report_runs", + id=uuid4(), + report_id=of_report, + budget_window_run_id=window_run_id, + year=year, + baseline_simulation_run_id=_insert_simulation_run( + connection, simulation_id, year=year or 2026 + ), + status="pending", + trigger="initial", + **RUN_VERSIONS, + ) + + insert_year(2026) + insert_year(2027) + with refused(connection, match="uq_report_runs_budget_window_run_id_year"): + insert_year(2027) + with refused(connection, match="ck_report_runs_budget_window_year"): + insert_year(None) + with refused(connection, match="foreign key constraint"): + insert_year(2028, of_report=other_report_id) + with refused(connection, match="ck_simulation_runs_year"): + _insert_simulation_run(connection, simulation_id, year=1800) + # A window run owns its years, as the report owns its runs. + second_window_run_id = _insert_budget_window_run( + connection, public_api_rows, report_id + ) + connection.execute( + text("DELETE FROM budget_window_runs WHERE id = :id"), + {"id": window_run_id}, + ) + assert _count(connection, "report_runs", "budget_window_run_id", window_run_id) == 0 + + connection.execute(text("DELETE FROM reports WHERE id = :id"), {"id": report_id}) + assert _count(connection, "budget_window_runs", "id", second_window_run_id) == 0 + assert _count(connection, "report_runs", "report_id", report_id) == 0 + + +def test_household_jobs_take_a_valid_year_and_record_when_they_end( + connection: Connection, public_api_rows: PublicApiRows +) -> None: + household_job = { + "country": "us", + "household_id": public_api_rows.household_id, + "request_data": "{}", + "tax_benefit_model_version_id": public_api_rows.version_id, + **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 e5daf84d4..c46fa0f6e 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -36,6 +36,7 @@ Report, ReportRun, Simulation, + SimulationRun, TaxBenefitModel, TaxBenefitModelVersion, ) @@ -500,7 +501,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, ) @@ -513,8 +513,18 @@ def test_republish_preserves_existing_run_and_dataset_version_records( ) session.add_all((dataset_version, simulation, report)) session.flush() + simulation_run = SimulationRun( + simulation_id=simulation.id, + year=2025, + tax_benefit_model_version_id=model_version.id, + country_package_version="existing-country-version", + policyengine_version="existing-policyengine-version", + ) + session.add(simulation_run) + session.flush() report_run = ReportRun( report_id=report.id, + baseline_simulation_run_id=simulation_run.id, country_package_version="existing-country-version", policyengine_version="existing-policyengine-version", ) @@ -526,8 +536,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) @@ -542,8 +552,23 @@ 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] + # Republishing keeps every catalog row the simulation and report 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..97fdb3f34 100644 --- a/tests/integration/test_v2_privilege_reconciliation.py +++ b/tests/integration/test_v2_privilege_reconciliation.py @@ -15,15 +15,18 @@ 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, PrivilegeEntry, apply_declared_privileges, check_declared_privileges, + check_table_wide_overrides, run, + skipped_roles, ) -from policyengine_api.data.v2.privileges import READ, TablePrivilege +from policyengine_api.data.v2.privileges import READ, TablePrivilege, update_columns from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL if TYPE_CHECKING: @@ -72,6 +75,82 @@ def _has_privilege(engine: Engine, role: str, table: str, privilege: str) -> boo ) +def _has_column_privilege( + engine: Engine, role: str, table: str, column: str, privilege: str +) -> bool: + with engine.connect() as connection: + return bool( + connection.execute( + text("SELECT has_column_privilege(:role, :table, :column, :privilege)"), + { + "role": role, + "table": f"public.{table}", + "column": column, + "privilege": privilege, + }, + ).scalar_one() + ) + + +def test_apply_grants_declared_columns_alone_and_verify_checks_them( + superuser_engine, temporary_role +) -> None: + role = temporary_role() + declared = {role: {"regions": READ | {update_columns("label")}}} + with superuser_engine.begin() as connection: + first_plan, _defaults = apply_declared_privileges(connection, declared) + with superuser_engine.begin() as connection: + second_plan, _defaults = apply_declared_privileges(connection, declared) + missing = check_declared_privileges(connection, declared) + overrides = check_table_wide_overrides(connection, declared) + + assert first_plan.grants == ( + PrivilegeEntry(role, "regions", "SELECT"), + PrivilegeEntry(role, "regions", "UPDATE", column="label"), + ) + assert (second_plan.grants, second_plan.undeclared) == ((), ()) + assert (missing, overrides) == ([], []) + assert _has_column_privilege(superuser_engine, role, "regions", "label", "UPDATE") + assert not _has_column_privilege( + superuser_engine, role, "regions", "code", "UPDATE" + ) + assert not _has_privilege(superuser_engine, role, "regions", "UPDATE") + + +def test_verify_fails_when_a_role_can_update_the_whole_table( + disposable_v2_database_url, superuser_engine, temporary_role, monkeypatch, tmp_path +) -> None: + """``UPDATE`` on the whole table would let the role change every column.""" + role = temporary_role() + monkeypatch.setattr( + privilege_reconciliation, + "declared_privileges", + lambda _environment: {role: {"regions": READ | {update_columns("label")}}}, + ) + monkeypatch.setattr( + privilege_reconciliation, "ROLES_AWAITING_PROVISIONING", frozenset() + ) + summary_path = tmp_path / "summary.md" + environ = { + V2_MIGRATION_DATABASE_URL: disposable_v2_database_url, + V2_ALEMBIC_DISPOSABLE_TEST: "1", + DEPLOYMENT_ENVIRONMENT: "test", + "GITHUB_STEP_SUMMARY": str(summary_path), + } + + apply_status = run("apply", environ=environ) + with superuser_engine.begin() as connection: + connection.execute(text(f'GRANT UPDATE ON TABLE public.regions TO "{role}"')) + verify_status = run("verify", environ=environ) + summary = summary_path.read_text(encoding="utf-8") + + assert (apply_status, verify_status) == (0, 1) + assert ( + f"`{role}` holds UPDATE on the whole of `regions`, which overrides its " + "declared columns" + ) in summary + + def test_apply_grants_nothing_the_second_time_and_verify_passes( disposable_v2_database_url, managed_v2_roles, capsys ) -> None: @@ -137,3 +216,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..fbef49bf9 --- /dev/null +++ b/tests/integration/test_v2_simulation_role_privileges.py @@ -0,0 +1,371 @@ +"""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 policyengine_api.data.v2.models import V2_METADATA +from policyengine_api.data.v2.privileges import ( + ColumnPrivilege, + TablePrivilege, + declared_privileges, +) +from tests.integration.simulation_rows import ( + RUN_VERSIONS, + PublicApiRows, + insert_public_api_rows, + 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() + + +@pytest.fixture +def public_api_rows(connection: Connection) -> PublicApiRows: + """Store, as the superuser, the Public API's rows the simulation rows name.""" + return insert_public_api_rows(connection) + + +def _act_as(connection: Connection, role: str) -> None: + quoted_role = connection.dialect.identifier_preparer.quote(role) + connection.execute(text(f"SET LOCAL ROLE {quoted_role}")) + + +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, public_api_rows: PublicApiRows +) -> dict[str, object]: + """Create what the simulation API creates for one economy report run.""" + ids = { + name: uuid4() + for name in ( + "simulation", + "run", + "report", + "budget_window_run", + "report_run", + "household_job", + ) + } + insert_row( + connection, + "simulations", + id=ids["simulation"], + simulation_type="economy", + dataset_id=public_api_rows.dataset_id, + tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, + ) + insert_row( + connection, + "simulation_runs", + id=ids["run"], + simulation_id=ids["simulation"], + tax_benefit_model_version_id=public_api_rows.version_id, + year=2026, + status="pending", + **RUN_VERSIONS, + ) + insert_row( + connection, + "reports", + id=ids["report"], + label="Report", + country="us", + tax_benefit_model_id=public_api_rows.model_id, + year=2026, + budget_window_size=1, + inputs="{}", + ) + insert_row( + connection, + "budget_window_runs", + id=ids["budget_window_run"], + report_id=ids["report"], + tax_benefit_model_version_id=public_api_rows.version_id, + trigger="initial", + recalculate=False, + status="pending", + ) + insert_row( + connection, + "report_runs", + id=ids["report_run"], + report_id=ids["report"], + budget_window_run_id=ids["budget_window_run"], + year=2026, + baseline_simulation_run_id=ids["run"], + status="pending", + trigger="initial", + **RUN_VERSIONS, + ) + insert_row( + connection, + "household_jobs", + id=ids["household_job"], + country="us", + household_id=public_api_rows.household_id, + year=2026, + request_data="{}", + tax_benefit_model_version_id=public_api_rows.version_id, + status="pending", + **RUN_VERSIONS, + ) + return ids + + +def test_the_simulation_service_creates_rows_and_records_only_its_calls( + connection: Connection, managed_v2_roles, public_api_rows: PublicApiRows +) -> None: + _act_as(connection, managed_v2_roles.simulation_service) + + ids = _create_simulation_rows(connection, public_api_rows) + run_count = connection.execute( + text("SELECT count(*) FROM simulation_runs WHERE id = :id"), {"id": ids["run"]} + ).scalar_one() + connection.execute(text("SELECT count(*) FROM policies")).scalar_one() + + assert run_count == 1 + # The service records a Modal call's ID once Modal accepts the call, and + # fails the run when Modal refuses it, while the run is still pending. + recorded = connection.execute( + text( + "UPDATE report_runs SET modal_call_id = 'fc-accepted', " + "updated_at = now() WHERE id = :id AND status = 'pending' " + "AND modal_call_id IS NULL RETURNING id" + ), + {"id": ids["report_run"]}, + ).scalar_one() + failed = connection.execute( + text( + "UPDATE household_jobs SET status = 'failed', error_message = 'refused', " + "completed_at = now(), updated_at = now() " + "WHERE id = :id AND status = 'pending' RETURNING id" + ), + {"id": ids["household_job"]}, + ).scalar_one() + assert (recorded, failed) == (ids["report_run"], ids["household_job"]) + for statement in ( + "UPDATE simulation_runs SET started_at = now() WHERE id = :id", + "UPDATE simulation_runs SET output_uri = 'gs://bucket/run' WHERE id = :id", + "UPDATE simulation_runs SET year = 2027 WHERE id = :id", + ): + _refused(connection, statement, id=ids["run"]) + _refused( + connection, + "UPDATE report_runs SET baseline_cliff_gap = 1.0 WHERE id = :id", + id=ids["report_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 labor_supply_decile_responses (id, report_run_id, decile) " + "VALUES (:id, :report_run_id, 1)", + id=uuid4(), + report_run_id=ids["report_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, public_api_rows: PublicApiRows +) -> None: + _act_as(connection, managed_v2_roles.simulation_service) + ids = _create_simulation_rows(connection, public_api_rows) + _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', started_at = now(), " + "modal_call_id = 'fc-coordinator', baseline_cliff_gap = 1.0, " + "baseline_cliff_share = 0.1, reform_cliff_gap = 2.0, " + "reform_cliff_share = 0.2, data_version = '1.0.0', " + "spm_provenance = CAST('{}' AS jsonb), updated_at = now() " + "WHERE id = :id" + ), + {"id": ids["report_run"]}, + ) + insert_row( + connection, + "labor_supply_decile_responses", + id=uuid4(), + report_run_id=ids["report_run"], + decile=1, + average_income_effect=-1.0, + ) + 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, + "INSERT INTO budget_window_runs (id, report_id, " + "tax_benefit_model_version_id, trigger, recalculate, status) " + "VALUES (:id, :report_id, :version_id, 'manual', true, 'pending')", + id=uuid4(), + report_id=ids["report"], + version_id=public_api_rows.version_id, + ) + _refused( + connection, + "UPDATE poverty SET rate = 0 WHERE simulation_run_id = :id", + id=ids["run"], + ) + _refused(connection, "UPDATE households SET content_hash = content_hash") + # What a run is stays as the simulation service wrote it. + for statement, run_key in ( + ("UPDATE simulation_runs SET year = 2027 WHERE id = :id", "run"), + ( + "UPDATE simulation_runs SET simulation_id = simulation_id WHERE id = :id", + "run", + ), + ( + "UPDATE report_runs SET baseline_simulation_run_id = " + "baseline_simulation_run_id WHERE id = :id", + "report_run", + ), + ("UPDATE report_runs SET markdown = 'edited' WHERE id = :id", "report_run"), + ( + "UPDATE household_jobs SET household_id = household_id WHERE id = :id", + "household_job", + ), + ): + _refused(connection, statement, id=ids[run_key]) + + +def test_the_public_api_reads_runs_and_results_but_writes_neither( + connection: Connection, managed_v2_roles +) -> None: + _act_as(connection, managed_v2_roles.runtime) + + for table in ("simulation_runs", "report_runs", "poverty"): + connection.execute(text(f"SELECT count(*) FROM {table}")).scalar_one() + _refused( + connection, + "INSERT INTO report_runs (id, report_id, baseline_simulation_run_id, " + "country_package_version, policyengine_version, status, trigger) " + "VALUES (:id, :report_id, :run_id, '1.0.0', '4.0.0', 'pending', 'initial')", + id=uuid4(), + report_id=uuid4(), + run_id=uuid4(), + ) + + +@pytest.mark.parametrize("role_kind", ["simulation_service", "modal_worker"]) +def test_each_simulation_role_updates_exactly_its_declared_columns( + connection: Connection, managed_v2_roles, role_kind: str +) -> None: + """Every column of the run tables is updatable exactly when it is declared.""" + role = getattr(managed_v2_roles, role_kind) + declared = declared_privileges(managed_v2_roles.deployment_environment)[role] + _act_as(connection, role) + + for table in ( + "simulation_runs", + "report_runs", + "household_jobs", + "budget_window_runs", + ): + updatable = { + column + for privilege in declared[table] + if isinstance(privilege, ColumnPrivilege) + and privilege.privilege is TablePrivilege.UPDATE + for column in privilege.columns + } + for column in V2_METADATA.tables[table].c: + # The privilege is checked when the statement starts, so no row is + # needed. + statement = f"UPDATE {table} SET {column.name} = {column.name} WHERE false" + if column.name in updatable: + connection.execute(text(statement)) + else: + _refused(connection, statement) diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index efe4e06b1..79a3b1780 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -208,6 +208,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 @@ -245,6 +246,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..e6430f6a4 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -40,6 +40,7 @@ V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) POST_BASELINE_TABLE_NAMES = frozenset( { + "budget_window_runs", "countries", "legacy_household_mappings", "legacy_policy_mappings", @@ -47,6 +48,9 @@ "legacy_user_policy_mappings", "parameter_search_files", "parameter_tree_children", + "labor_supply_decile_responses", + "labor_supply_responses", + "simulation_runs", "stage12_evaluation_reports", "stage12_evaluation_simulations", "user_identities", @@ -62,6 +66,8 @@ 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" DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -243,6 +249,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + SIMULATION_RUNS_PARENT_REVISION, PARAMETER_TREE_REVISION, AREA_IMPACT_COORDINATE_REVISION, CATALOG_FORMAT_2_REVISION, @@ -989,7 +996,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 +1016,61 @@ 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 _simulation_runs_revision_source() -> str: + """Return the source of the one revision after the parameter search files.""" + (path,) = (REPO / "migrations/v2/versions").glob(f"{V2_ALEMBIC_HEAD_REVISION}_*.py") + return path.read_text(encoding="utf-8") + + +def test_simulation_runs_revision_is_generated_with_documented_corrections() -> None: + """One generated revision follows the parameter search file revision. + + Its corrections make new columns reuse enum types the baseline created: + ``simulation_runs.status`` reuses ``simulations.status``'s, and + ``budget_window_runs.status`` and ``trigger`` reuse the report run's. + """ + revision = _simulation_runs_revision_source() + upgrade = revision[ + 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:") == 3 + assert "reversibility correction" not in revision + for type_name in ( + "v2_simulation_status", + "v2_report_run_status", + "v2_report_run_trigger", + ): + assert f'name="{type_name}",\n create_type=False,' in upgrade + for table_name in ( + "simulation_runs", + "budget_window_runs", + "labor_supply_responses", + "labor_supply_decile_responses", + ): + assert f'op.create_table(\n "{table_name}"' in upgrade + assert f'op.drop_table("{table_name}")' in downgrade + assert upgrade.count("op.drop_table(") == 0 + for removed_column in ( + "status", + "error_message", + "started_at", + "completed_at", + "household_result", + "output_dataset_id", + ): + assert f'op.drop_column("simulations", "{removed_column}")' in upgrade + for body in (upgrade, downgrade): + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index a3010eff6..45fc7d21a 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 42 +assert len(V2_METADATA.tables) == 46 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules diff --git a/tests/unit/v2/test_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..6b83ff44e 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -24,6 +24,7 @@ ParameterValue, Policy, Simulation, + SimulationRun, TaxBenefitModelVersion, User, UserHouseholdAssociation, @@ -46,6 +47,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 +57,8 @@ "congressional_district_impacts", "constituency_impacts", "decile_impacts", - "inequality", "intra_decile_impacts", "local_authority_impacts", - "poverty", "program_statistics", } ) @@ -74,6 +75,7 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: HouseholdJob: "households", LegacyHouseholdMapping: "household_mappings", Simulation: "simulations", + SimulationRun: "simulations", UserHouseholdAssociation: "associations", UserPolicy: "associations", LegacyPolicyMapping: "policy_mappings", @@ -500,11 +502,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..f66b69d79 100644 --- a/tests/unit/v2/test_privilege_manifest.py +++ b/tests/unit/v2/test_privilege_manifest.py @@ -18,14 +18,21 @@ ) from policyengine_api.data.v2.privileges import ( CATALOG_WRITER_ROLE, + MODAL_WORKER_ROLE, NO_ACCESS, + READ, + READ_INSERT, READ_WRITE, + ROLES_AWAITING_PROVISIONING, RUNTIME_ROLE, + SIMULATION_SERVICE_ROLE, TABLE_ACCESS, + ColumnPrivilege, InvalidDeploymentEnvironmentError, TablePrivilege, declared_privileges, stage12_role, + update_columns, ) PROJECT_ROOT = Path(__file__).parents[3] @@ -33,6 +40,121 @@ 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", + "budget_window_runs", + "household_jobs", + } +) +RUN_STATUS_TABLES = frozenset({"simulation_runs", "report_runs", "household_jobs"}) +# The run columns the Modal workers update: a run's state and its results. +WORKER_RUN_STATE_COLUMNS = frozenset( + { + "status", + "modal_call_id", + "error_message", + "started_at", + "completed_at", + "updated_at", + } +) +WORKER_RESULT_COLUMNS = { + "simulation_runs": frozenset({"output_uri", "output_sha256", "household_result"}), + "report_runs": frozenset( + { + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", + "spm_provenance", + "data_version", + } + ), + "household_jobs": frozenset({"result"}), +} +# The run columns the simulation service updates: it records a Modal call's ID +# when Modal accepts the call, and fails the run when Modal refuses it. +SERVICE_RUN_STATE_COLUMNS = frozenset( + {"status", "modal_call_id", "error_message", "completed_at", "updated_at"} +) +# What a run is: no role may change these after the simulation service writes +# them. +RUN_DEFINITION_COLUMNS = { + "simulation_runs": frozenset( + { + "simulation_id", + "year", + "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", + "dataset_id", + "region_id", + } + ), + "report_runs": frozenset( + { + "report_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "budget_window_run_id", + "year", + "country_package_version", + "policyengine_version", + "trigger", + "idempotency_key", + } + ), + "household_jobs": frozenset( + { + "country", + "household_id", + "year", + "policy_id", + "dynamic_id", + "request_data", + "tax_benefit_model_version_id", + "country_package_version", + "policyengine_version", + } + ), +} +# 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( + { + "labor_supply_responses", + "labor_supply_decile_responses", + "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 +240,101 @@ 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) + for table in RUN_STATUS_TABLES: + expected[table] = READ_INSERT | {update_columns(*SERVICE_RUN_STATE_COLUMNS)} + 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(RESULT_TABLES, READ_INSERT) + expected |= dict.fromkeys(RUN_INPUT_TABLES, READ) + for table in RUN_STATUS_TABLES: + worker_columns = WORKER_RUN_STATE_COLUMNS | WORKER_RESULT_COLUMNS[table] + expected[table] = READ | {update_columns(*worker_columns)} + assert declared == expected + + +def test_column_privileges_name_their_tables_columns_alone() -> None: + """A whole-table privilege would override the same privilege's column list.""" + for role, privileges_by_table in declared_privileges(TEST_ENVIRONMENT).items(): + for table, privileges in privileges_by_table.items(): + column_privileges = [ + privilege + for privilege in privileges + if isinstance(privilege, ColumnPrivilege) + ] + kinds = [privilege.privilege for privilege in column_privileges] + assert len(kinds) == len(set(kinds)), (role, table) + for privilege in column_privileges: + assert privilege.columns, (role, table) + assert privilege.columns <= set(V2_METADATA.tables[table].c.keys()), ( + role, + table, + ) + assert privilege.privilege not in privileges, (role, table) + + +def test_no_role_can_change_what_a_run_is() -> None: + for role, privileges_by_table in declared_privileges(TEST_ENVIRONMENT).items(): + for table, definition_columns in RUN_DEFINITION_COLUMNS.items(): + privileges = privileges_by_table[table] + assert TablePrivilege.UPDATE not in privileges, (role, table) + for privilege in privileges: + if isinstance(privilege, ColumnPrivilege): + assert privilege.columns.isdisjoint(definition_columns), ( + role, + table, + ) + + +def test_public_api_reads_simulations_runs_and_results_without_writing_them() -> None: + declared = declared_privileges(TEST_ENVIRONMENT)[RUNTIME_ROLE] + + 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 + or any(isinstance(privilege, ColumnPrivilege) for privilege in 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..208a5a1e5 100644 --- a/tests/unit/v2/test_privilege_reconciliation.py +++ b/tests/unit/v2/test_privilege_reconciliation.py @@ -10,23 +10,31 @@ from policyengine_api.data.v2 import privilege_reconciliation from policyengine_api.data.v2.privilege_reconciliation import ( DefaultPrivilege, + GrantedPrivileges, PrivilegeEntry, ReconciliationPlan, grant_statements, plan_reconciliation, render_summary, + table_wide_overrides, verification_failures, ) from policyengine_api.data.v2.privileges import ( READ, READ_WRITE, TablePrivilege, + update_columns, ) PROJECT_ROOT = Path(__file__).parents[3] PROOF_FIXTURES = PROJECT_ROOT / "rules/fixtures/api.v2.deployed-privileges-match" RUNTIME = "policyengine_v2_runtime" STAGE12 = "policyengine_stage12_staging" +WORKER = "policyengine_v2_modal_worker" +# The Modal worker may read a run and update two of its columns. +WORKER_DECLARED = { + WORKER: {"simulation_runs": READ | {update_columns("status", "completed_at")}} +} def _quote(identifier: str) -> str: @@ -51,7 +59,7 @@ def _declared() -> dict[str, dict[str, frozenset[TablePrivilege]]]: def test_missing_declared_privileges_become_grants() -> None: plan = plan_reconciliation( _declared(), - granted={RUNTIME: {"users": {"SELECT", "INSERT"}}}, + granted=GrantedPrivileges({RUNTIME: {"users": {"SELECT", "INSERT"}}}), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports"}, ) @@ -69,14 +77,16 @@ def test_missing_declared_privileges_become_grants() -> None: def test_privileges_beyond_the_manifest_are_reported_not_granted_or_revoked() -> None: plan = plan_reconciliation( _declared(), - granted={ - RUNTIME: { - "users": {"SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE"}, - "legacy_policy_mappings": {"SELECT", "INSERT"}, - "reports": {"SELECT"}, - }, - STAGE12: {"users": {"SELECT"}}, - }, + granted=GrantedPrivileges( + { + RUNTIME: { + "users": {"SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE"}, + "legacy_policy_mappings": {"SELECT", "INSERT"}, + "reports": {"SELECT"}, + }, + STAGE12: {"users": {"SELECT"}}, + } + ), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports"}, ) @@ -94,7 +104,7 @@ def test_privileges_beyond_the_manifest_are_reported_not_granted_or_revoked() -> def test_privileges_on_tables_outside_the_manifest_are_reported() -> None: plan = plan_reconciliation( _declared(), - granted={RUNTIME: {"scratch_table": {"SELECT"}}}, + granted=GrantedPrivileges({RUNTIME: {"scratch_table": {"SELECT"}}}), existing_roles={RUNTIME, STAGE12}, existing_tables={"users", "legacy_policy_mappings", "reports", "scratch_table"}, ) @@ -105,7 +115,7 @@ def test_privileges_on_tables_outside_the_manifest_are_reported() -> None: def test_missing_roles_and_tables_are_recorded_and_block_the_plan() -> None: plan = plan_reconciliation( _declared(), - granted={}, + granted=GrantedPrivileges({}), existing_roles={RUNTIME}, existing_tables={"users", "reports"}, ) @@ -120,6 +130,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=GrantedPrivileges({}), + 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=GrantedPrivileges({}), + 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=GrantedPrivileges({}), + 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=( @@ -201,6 +275,142 @@ def test_summary_lists_grants_extras_and_default_privileges_without_secrets() -> assert "postgresql" not in summary.lower().replace("postgresql's", "") +def test_missing_declared_column_privileges_become_column_grants() -> None: + plan = plan_reconciliation( + WORKER_DECLARED, + granted=GrantedPrivileges( + tables={WORKER: {"simulation_runs": {"SELECT"}}}, + columns={WORKER: {"simulation_runs": {"status": {"UPDATE"}}}}, + ), + existing_roles={WORKER}, + existing_tables={"simulation_runs"}, + ) + + assert plan.grants == ( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ) + assert plan.undeclared == () + + +def test_column_and_whole_table_privileges_beyond_the_manifest_are_reported() -> None: + plan = plan_reconciliation( + WORKER_DECLARED, + granted=GrantedPrivileges( + tables={WORKER: {"simulation_runs": {"SELECT", "UPDATE"}}}, + columns={ + WORKER: { + "simulation_runs": { + "status": {"UPDATE"}, + "completed_at": {"UPDATE"}, + "simulation_id": {"UPDATE"}, + } + } + }, + ), + existing_roles={WORKER}, + existing_tables={"simulation_runs"}, + ) + + assert plan.grants == () + assert plan.undeclared == ( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="simulation_id"), + ) + + +def test_grant_statements_grant_column_privileges_per_role_and_table() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "SELECT"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="status"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ), + undeclared=(), + missing_roles=(), + missing_tables=(), + ) + + assert grant_statements(plan, _quote) == [ + 'GRANT SELECT ON TABLE "public"."simulation_runs" ' + 'TO "policyengine_v2_modal_worker"', + 'GRANT UPDATE ("completed_at", "status") ON TABLE "public"."simulation_runs" ' + 'TO "policyengine_v2_modal_worker"', + ] + + +def test_grant_statements_refuse_an_unsafe_column_name() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column='status"; --'), + ), + undeclared=(), + missing_roles=(), + missing_tables=(), + ) + + with pytest.raises(privilege_reconciliation.UnsafeIdentifierError): + grant_statements(plan, _quote) + + +def test_verification_reports_each_declared_column_the_role_lacks() -> None: + failures = verification_failures( + WORKER_DECLARED, + lambda role, table, privilege: privilege == "SELECT", + has_column_privilege=lambda role, table, column, privilege: column == "status", + ) + + assert failures == [ + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at") + ] + + +def test_column_privileges_need_a_column_check() -> None: + with pytest.raises(ValueError, match="column"): + verification_failures(WORKER_DECLARED, lambda role, table, privilege: True) + + +def test_a_whole_table_privilege_overrides_its_declared_columns() -> None: + """``UPDATE`` on the whole table lets a role change every column.""" + + def has_privilege(role: str, table: str, privilege: str) -> bool: + return privilege in {"SELECT", "UPDATE"} + + assert table_wide_overrides(WORKER_DECLARED, has_privilege) == [ + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE") + ] + assert ( + table_wide_overrides( + WORKER_DECLARED, lambda role, table, privilege: privilege == "SELECT" + ) + == [] + ) + + +def test_summary_lists_column_grants_with_their_columns() -> None: + plan = ReconciliationPlan( + grants=( + PrivilegeEntry(WORKER, "simulation_runs", "SELECT"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="status"), + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="completed_at"), + ), + undeclared=( + PrivilegeEntry(WORKER, "simulation_runs", "UPDATE", column="year"), + ), + missing_roles=(), + missing_tables=(), + ) + + summary = render_summary(plan, (), deployment_environment="staging", applied=True) + + assert ( + "| `policyengine_v2_modal_worker` | `simulation_runs` | " + "SELECT, UPDATE (completed_at, status) |" + ) in summary + assert ( + "| `policyengine_v2_modal_worker` | `simulation_runs` | UPDATE (year) |" + ) in summary + + def _fixture_privileges(name: str) -> set[tuple[str, str, str]]: entries = json.loads((PROOF_FIXTURES / name).read_text(encoding="utf-8")) return {(entry["role"], entry["table"], entry["privilege"]) for entry in entries} diff --git a/tests/unit/v2/test_report_run_creation.py b/tests/unit/v2/test_report_run_creation.py index 3a1b1b82c..ac0959f9b 100644 --- a/tests/unit/v2/test_report_run_creation.py +++ b/tests/unit/v2/test_report_run_creation.py @@ -3,7 +3,7 @@ from __future__ import annotations from contextlib import contextmanager -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.exc import IntegrityError @@ -20,10 +20,18 @@ ) from policyengine_api.data.v2.report_runs import ( ReportNotFoundError, + ReportRunInputs, ReportRunPackageVersions, + ReportRunSimulationRuns, create_report_run, ) +# The simulation runs each report run compares. SQLite doesn't enforce foreign +# keys, so these tests store no simulation run rows for them. +SIMULATION_RUNS = ReportRunSimulationRuns( + baseline_simulation_run_id=UUID(int=101), reform_simulation_run_id=UUID(int=102) +) + @pytest.fixture def engine(): @@ -38,10 +46,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) @@ -57,9 +68,12 @@ def test_unknown_report_is_rejected_before_any_validation(engine) -> None: create_report_run( session, report_id=uuid4(), - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="", + policyengine_version="", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, ) @@ -75,9 +89,12 @@ def test_manual_run_without_a_key_is_rejected_before_version_checks(engine) -> N create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="", - policyengine_version="", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="", + policyengine_version="", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, ) @@ -101,9 +118,12 @@ def test_empty_package_versions_are_rejected( create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_package_version, - policyengine_version=policyengine_version, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version=country_package_version, + policyengine_version=policyengine_version, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.SYSTEM, ) @@ -116,9 +136,12 @@ def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: run = create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version="4.5.6", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.INITIAL, ) @@ -131,6 +154,8 @@ def test_run_without_a_key_is_inserted_without_a_savepoint(engine) -> None: assert run.trigger is ReportRunTrigger.INITIAL assert run.idempotency_key is None assert run.status is ReportRunStatus.PENDING + assert run.baseline_simulation_run_id == UUID(int=101) + assert run.reform_simulation_run_id == UUID(int=102) def test_conflict_without_a_matching_keyed_run_reraises( @@ -152,9 +177,12 @@ def failing_savepoint(): create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version="4.5.6", + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version="4.5.6", + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, idempotency_key=uuid4(), diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py index 4383a8110..ecbf58ff4 100644 --- a/tests/unit/v2/test_report_runs.py +++ b/tests/unit/v2/test_report_runs.py @@ -26,7 +26,9 @@ TaxBenefitModelVersion, ) from policyengine_api.data.v2.report_runs import ( + ReportRunInputs, ReportRunPackageVersions, + ReportRunSimulationRuns, ReportRunStateError, ReportTypeImmutableError, begin_report_run, @@ -40,6 +42,11 @@ DEPLOYED_COUNTRY_VERSIONS = {"us": "1.2.3"} DEPLOYED_POLICYENGINE_VERSION = "4.5.6" NOW = datetime(2026, 8, 14, 12, 0, tzinfo=UTC) +# The simulation runs each report run compares. SQLite doesn't enforce foreign +# keys, so these tests store no simulation run rows for them. +SIMULATION_RUNS = ReportRunSimulationRuns( + baseline_simulation_run_id=UUID(int=101), reform_simulation_run_id=UUID(int=102) +) @pytest.fixture @@ -56,11 +63,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) @@ -79,9 +88,12 @@ def _create_run( return create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version=country_version, - policyengine_version=policyengine_version, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version=country_version, + policyengine_version=policyengine_version, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, idempotency_key=key if key is not None else uuid4(), @@ -145,9 +157,12 @@ def test_manual_rerun_requires_an_idempotency_key(engine) -> None: create_report_run( session, report_id=report.id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, ) @@ -207,9 +222,12 @@ def request_rerun() -> UUID: run = create_report_run( session, report_id=report_id, - versions=ReportRunPackageVersions( - country_package_version="1.2.3", - policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + inputs=ReportRunInputs( + versions=ReportRunPackageVersions( + country_package_version="1.2.3", + policyengine_version=DEPLOYED_POLICYENGINE_VERSION, + ), + simulation_runs=SIMULATION_RUNS, ), trigger=ReportRunTrigger.MANUAL, idempotency_key=request_key, @@ -249,23 +267,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() @@ -305,6 +326,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=1), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -313,6 +335,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=2), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -321,6 +344,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=3), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="9.9.9", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.SUCCEEDED, @@ -329,6 +353,7 @@ def test_selector_uses_versions_success_completion_time_and_stable_id(engine) -> ReportRun( id=UUID(int=4), report=report, + baseline_simulation_run_id=UUID(int=101), country_package_version="1.2.3", policyengine_version=DEPLOYED_POLICYENGINE_VERSION, status=ReportRunStatus.RUNNING, diff --git a/tests/unit/v2/test_simulation_run_models.py b/tests/unit/v2/test_simulation_run_models.py new file mode 100644 index 000000000..24298ab60 --- /dev/null +++ b/tests/unit/v2/test_simulation_run_models.py @@ -0,0 +1,607 @@ +"""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 +its baseline and reform simulation runs, results are keyed by run, and a +household job calculates one household for one year under one policy. +""" + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql +from sqlalchemy.schema import CreateTable + +from policyengine_api.data.v2.models import ( + V2_METADATA, + BudgetWindowRun, + LaborSupplyDecileResponse, + LaborSupplyResponse, + ReportRun, + SimulationRun, +) + +RUN_STATUSES = ["pending", "running", "succeeded", "failed"] +# The baseline's and reform's cliff gap and share, stored together or not at all. +CLIFF_IMPACT_COLUMNS = ( + "baseline_cliff_gap", + "baseline_cliff_share", + "reform_cliff_gap", + "reform_cliff_share", +) +LABOR_SUPPLY_RESPONSE_COLUMNS = ( + "income_effect_earnings_change", + "substitution_effect_earnings_change", + "total_earnings_change", + "revenue_change", + "relative_income_effect", + "relative_substitution_effect", + "baseline_weekly_hours", + "reform_weekly_hours", + "weekly_hours_change", + "weekly_hours_income_effect", + "weekly_hours_substitution_effect", +) +LABOR_SUPPLY_DECILE_COLUMNS = ( + "average_income_effect", + "average_substitution_effect", + "relative_income_effect", + "relative_substitution_effect", +) +PER_SIMULATION_RESULT_TABLES = ("inequality", "poverty") +COMPARISON_RESULT_TABLES = ( + "budget_summary", + "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", + "budget_window_runs", + "labor_supply_responses", + "labor_supply_decile_responses", + "household_jobs", + "aggregates", + "change_aggregates", + *PER_SIMULATION_RESULT_TABLES, + *COMPARISON_RESULT_TABLES, + } +) +# The foreign keys from the simulation service's tables to the Public API's rows, +# with their delete rules. Each service writes only its own tables, and +# PostgreSQL checks every row named on insert. +PUBLIC_API_FOREIGN_KEYS = { + ("simulations", "dataset_id"): ("datasets", "RESTRICT"), + ("simulations", "household_id"): ("households", "RESTRICT"), + ("simulations", "policy_id"): ("policies", "SET NULL"), + ("simulations", "dynamic_id"): ("dynamics", "SET NULL"), + ("simulations", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulations", "region_id"): ("regions", "SET NULL"), + ("reports", "user_id"): ("users", "SET NULL"), + ("reports", "tax_benefit_model_id"): ("tax_benefit_models", "RESTRICT"), + ("reports", "policy_id"): ("policies", "SET NULL"), + ("reports", "household_id"): ("households", "RESTRICT"), + ("reports", "dataset_id"): ("datasets", "SET NULL"), + ("reports", "region_id"): ("regions", "SET NULL"), + ("simulation_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("simulation_runs", "dataset_id"): ("datasets", "RESTRICT"), + ("simulation_runs", "region_id"): ("regions", "SET NULL"), + ("household_jobs", "household_id"): ("households", "RESTRICT"), + ("household_jobs", "policy_id"): ("policies", "SET NULL"), + ("household_jobs", "dynamic_id"): ("dynamics", "SET NULL"), + ("household_jobs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), + ("budget_window_runs", "tax_benefit_model_version_id"): ( + "tax_benefit_model_versions", + "RESTRICT", + ), +} + + +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", + "year", + "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" + assert {key.target_fullname for key in simulation_runs.foreign_keys} == { + "simulations.id", + "tax_benefit_model_versions.id", + "datasets.id", + "regions.id", + } + assert "ck_simulation_runs_year" in _check_names(simulation_runs) + assert SimulationRun.__module__.endswith(".simulations") + + +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) + # A simulation has a run per year, so its newest successful run is found + # per year. + assert _indexes(simulation_runs) == { + "ix_simulation_runs_current_output": ( + "simulation_id", + "year", + "status", + "country_package_version", + "policyengine_version", + "completed_at", + "id", + ), + "ix_simulation_runs_status_started_at": ("status", "started_at"), + "ix_simulation_runs_tax_benefit_model_version_id": ( + "tax_benefit_model_version_id", + ), + } + assert _indexes(report_runs)["ix_report_runs_status_started_at"] == ( + "status", + "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 + assert "report_run_comparisons" not in V2_METADATA.tables + + +def test_report_runs_record_cliff_impact_spm_receipt_and_data_version() -> None: + """The comparison parts with no result table are typed, except the receipt.""" + report_runs = V2_METADATA.tables["report_runs"] + + for column_name in CLIFF_IMPACT_COLUMNS: + assert report_runs.c[column_name].nullable, column_name + assert isinstance(report_runs.c[column_name].type, sa.Float), column_name + assert "ck_report_runs_cliff_impact_complete" in _check_names(report_runs) + assert report_runs.c.spm_provenance.nullable + spm_type = report_runs.c.spm_provenance.type + assert isinstance(spm_type.dialect_impl(postgresql.dialect()), postgresql.JSONB) + assert report_runs.c.data_version.nullable + assert report_runs.c.data_version.type.length == 128 + assert {"labor_supply_result", "cliff_impact_result"}.isdisjoint( + report_runs.c.keys() + ) + + +def test_labor_supply_responses_have_one_typed_row_per_report_run() -> None: + responses = V2_METADATA.tables["labor_supply_responses"] + + assert set(responses.c.keys()) == { + "id", + "created_at", + "report_run_id", + *LABOR_SUPPLY_RESPONSE_COLUMNS, + } + for column_name in LABOR_SUPPLY_RESPONSE_COLUMNS: + assert not responses.c[column_name].nullable, column_name + assert isinstance(responses.c[column_name].type, sa.Float), column_name + report_run_key = _foreign_key(responses, "report_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + assert _unique_column_sets(responses) == {("report_run_id",)} + + +def test_labor_supply_decile_responses_have_one_row_per_decile() -> None: + deciles = V2_METADATA.tables["labor_supply_decile_responses"] + + assert set(deciles.c.keys()) == { + "id", + "created_at", + "report_run_id", + "decile", + *LABOR_SUPPLY_DECILE_COLUMNS, + } + assert not deciles.c.decile.nullable + # .py reports a relative response only for deciles above zero. + for column_name in LABOR_SUPPLY_DECILE_COLUMNS: + assert deciles.c[column_name].nullable, column_name + assert isinstance(deciles.c[column_name].type, sa.Float), column_name + report_run_key = _foreign_key(deciles, "report_run_id") + assert report_run_key.target_fullname == "report_runs.id" + assert report_run_key.ondelete == "CASCADE" + assert _unique_column_sets(deciles) == {("report_run_id", "decile")} + + +def test_report_run_relationships_name_labor_supply_rows() -> None: + relationships = sa.inspect(ReportRun).relationships + + assert relationships["labor_supply_responses"].mapper.class_ is ( + LaborSupplyResponse + ) + assert relationships["labor_supply_decile_responses"].mapper.class_ is ( + LaborSupplyDecileResponse + ) + assert "comparisons" not in relationships + assert LaborSupplyResponse.__module__.endswith(".run_results") + + +def test_uk_area_impacts_record_disclosure_control() -> None: + for table_name in ("constituency_impacts", "local_authority_impacts"): + table = V2_METADATA.tables[table_name] + assert table.c.disclosure_tier.nullable + assert table.c.disclosure_tier.type.length == 16 + for change in ("average", "relative"): + for bound in ("lower", "upper"): + column = table.c[f"{change}_household_income_change_{bound}"] + assert column.nullable + assert isinstance(column.type, sa.Float) + assert f"ck_{table_name}_disclosure_tier" in _check_names(table) + + +def test_congressional_district_impacts_record_winner_and_loser_shares() -> None: + districts = V2_METADATA.tables["congressional_district_impacts"] + + for column_name in ( + "winner_percentage", + "loser_percentage", + "no_change_percentage", + ): + assert not districts.c[column_name].nullable + assert isinstance(districts.c[column_name].type, sa.Float) + + +def test_detailed_budget_rows_need_neither_an_entity_nor_a_tax_flag() -> None: + """The economy output names each program's totals, not its entity or kind.""" + program_statistics = V2_METADATA.tables["program_statistics"] + + assert program_statistics.c.entity.nullable + assert program_statistics.c.is_tax.nullable + statement = str( + CreateTable(program_statistics).compile(dialect=postgresql.dialect()) + ) + assert ( + "CONSTRAINT uq_program_statistics_run_program_entity UNIQUE NULLS NOT " + "DISTINCT (report_run_id, program_name, entity)" + ) in statement + + +def test_report_runs_name_their_baseline_and_reform_simulation_runs() -> None: + """A report run compares one pair; a household report may have no reform.""" + report_runs = V2_METADATA.tables["report_runs"] + + assert not report_runs.c.baseline_simulation_run_id.nullable + assert report_runs.c.reform_simulation_run_id.nullable + for column_name in ("baseline_simulation_run_id", "reform_simulation_run_id"): + run_key = _foreign_key(report_runs, column_name) + assert run_key.target_fullname == "simulation_runs.id" + assert run_key.ondelete == "RESTRICT" + assert (column_name,) in _indexes(report_runs).values() + assert "report_run_simulation_runs" not in V2_METADATA.tables + enum_type_names = { + column.type.name + for table in V2_METADATA.tables.values() + for column in table.columns + if isinstance(column.type, sa.Enum) + } + assert "v2_simulation_role" not in enum_type_names + + +def test_budget_window_reports_name_how_many_years_they_compare() -> None: + reports = V2_METADATA.tables["reports"] + + assert reports.c.budget_window_size.nullable + assert isinstance(reports.c.budget_window_size.type, sa.Integer) + assert "ck_reports_budget_window_size" in _check_names(reports) + + +def test_budget_window_runs_record_each_calculation_of_a_window() -> None: + windows = V2_METADATA.tables["budget_window_runs"] + + required = { + "report_id", + "tax_benefit_model_version_id", + "trigger", + "recalculate", + "status", + "created_at", + "updated_at", + } + optional = {"error_message", "started_at", "completed_at"} + assert set(windows.c.keys()) == {"id"} | required | optional + assert all(not windows.c[name].nullable for name in required) + assert all(windows.c[name].nullable for name in optional) + # A window run has a report run's states and starts for a report run's reasons. + assert windows.c.status.type.name == "v2_report_run_status" + assert windows.c.trigger.type.name == "v2_report_run_trigger" + report_key = _foreign_key(windows, "report_id") + assert report_key.target_fullname == "reports.id" + assert report_key.ondelete == "CASCADE" + version_key = _foreign_key(windows, "tax_benefit_model_version_id") + assert version_key.target_fullname == "tax_benefit_model_versions.id" + assert version_key.ondelete == "RESTRICT" + assert ("id", "report_id") in _unique_column_sets(windows) + assert "ck_budget_window_runs_terminal_completion" in _check_names(windows) + assert _indexes(windows)["ix_budget_window_runs_report_id_created_at"] == ( + "report_id", + "created_at", + ) + assert BudgetWindowRun.__module__.endswith(".reports") + + +def test_report_runs_hold_one_year_of_a_budget_window_run() -> None: + report_runs = V2_METADATA.tables["report_runs"] + + assert report_runs.c.budget_window_run_id.nullable + assert report_runs.c.year.nullable + (window_key,) = ( + constraint + for constraint in report_runs.foreign_key_constraints + if constraint.referred_table.name == "budget_window_runs" + ) + # A year's report run belongs to a window run of the same report, which + # owns it as the report owns its runs. + assert [column.name for column in window_key.columns] == [ + "budget_window_run_id", + "report_id", + ] + assert [element.target_fullname for element in window_key.elements] == [ + "budget_window_runs.id", + "budget_window_runs.report_id", + ] + assert window_key.ondelete == "CASCADE" + assert ("budget_window_run_id", "year") in _unique_column_sets(report_runs) + assert {"ck_report_runs_year", "ck_report_runs_budget_window_year"} <= ( + _check_names(report_runs) + ) + + +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 _foreign_key(household_jobs, "household_id").target_fullname == ( + "households.id" + ) + assert _foreign_key( + household_jobs, "tax_benefit_model_version_id" + ).target_fullname == ("tax_benefit_model_versions.id") + + +def test_per_simulation_results_belong_to_one_simulation_run() -> None: + 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 report run compares one pair of simulation runs, so each row is + # unique within its report run. + (unique_columns,) = _unique_column_sets(table) + assert unique_columns[0] == "report_run_id", table_name + assert {"baseline_simulation_run_id", "reform_simulation_run_id"}.isdisjoint( + unique_columns + ), 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_simulation_tables_keep_foreign_keys_to_the_rows_they_name() -> None: + """Each service writes only its own tables; the keys between them stay.""" + public_api_keys = { + (table.name, key.parent.name): (key.column.table.name, key.ondelete) + for table in V2_METADATA.tables.values() + if table.name in SIMULATION_SIDE_TABLES + for key in table.foreign_keys + if key.column.table.name not in SIMULATION_SIDE_TABLES + } + + assert public_api_keys == PUBLIC_API_FOREIGN_KEYS + assert set(V2_METADATA.tables) >= SIMULATION_SIDE_TABLES + + +def test_reports_index_their_owner_for_account_deletion() -> None: + """Deleting a user clears ``reports.user_id``, which an index finds.""" + reports = V2_METADATA.tables["reports"] + + assert ("user_id",) in _indexes(reports).values() + + +def test_the_catalog_versions_runs_and_jobs_use_are_indexed() -> None: + for table_name in ("simulation_runs", "household_jobs"): + table = V2_METADATA.tables[table_name] + assert ("tax_benefit_model_version_id",) in _indexes(table).values(), table_name + + +def test_poverty_rows_name_the_group_their_filter_selects() -> None: + """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, income_variable, decile_variable, entity, decile, quantiles)" + ) in statement