diff --git a/.github/workflows/v2-integration-check.yml b/.github/workflows/v2-integration-check.yml index 3f402ae80..e6beed211 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.github/workflows/v2-integration-check.yml @@ -77,12 +77,12 @@ jobs: run: uv run coverage run --branch -m pytest -q tests/integration/test_v2_catalog_installed.py env: RUN_V2_CATALOG_COMPATIBILITY: "1" - - name: Test v2 metadata publication and resource routes - run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_metadata_routes.py + - name: Test v2 metadata publication, resource routes, the countries list and economy reports + run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_metadata_routes.py tests/integration/test_v2_country_routes.py tests/integration/test_v2_economy_reports.py - name: Test v2 policy persistence and immediate v1 mirroring run: uv run 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 coverage run -a --branch -m pytest -q tests/integration/test_v2_user_identity_persistence.py + run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_user_identity_persistence.py tests/integration/test_v2_saved_reference_updates.py - name: Test comparison-run persistence and privilege reconciliation as the managed roles run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_stage12_comparison_persistence.py tests/integration/test_v2_privilege_reconciliation.py - name: Test v2 household persistence and immediate v1 copying diff --git a/changelog.d/v2-countries-route.added.md b/changelog.d/v2-countries-route.added.md new file mode 100644 index 000000000..66b8d8e46 --- /dev/null +++ b/changelog.d/v2-countries-route.added.md @@ -0,0 +1 @@ +Added `GET /v2/countries`, which lists every row of the countries table: each served country's ID and whether it supports household and society-wide calculations. diff --git a/changelog.d/v2-countries-startup-check.added.md b/changelog.d/v2-countries-startup-check.added.md new file mode 100644 index 000000000..8e30ae192 --- /dev/null +++ b/changelog.d/v2-countries-startup-check.added.md @@ -0,0 +1 @@ +The Public API refuses to start when the v2 `countries` table and its country list in code name different countries, or when the table or the privilege to read it is missing. The check runs on Cloud Run and wherever a v2 runtime database is configured; when it cannot query the table, it logs why and startup continues. diff --git a/changelog.d/v2-countries-table.added.md b/changelog.d/v2-countries-table.added.md new file mode 100644 index 000000000..5c7e480ac --- /dev/null +++ b/changelog.d/v2-countries-table.added.md @@ -0,0 +1 @@ +Added a v2 `countries` table that records, for each country the Public API serves, whether it supports household and society-wide calculations. Its revision lists the US and the UK, both with both. Users' primary country and the country of every v2 household, policy and user association now reference this table instead of a fixed US-and-UK check. diff --git a/changelog.d/v2-economy-reports.added.md b/changelog.d/v2-economy-reports.added.md new file mode 100644 index 000000000..acb1fe5c1 --- /dev/null +++ b/changelog.d/v2-economy-reports.added.md @@ -0,0 +1 @@ +Added `/v2/reports` for economy reports, which the simulation service calculates and stores. `POST /v2/reports` creates a report and its first run on the country's current catalog version. `GET /v2/reports/{report_id}` and `GET /v2/reports/by-legacy-number/{number}` read one, `GET /v2/reports?report_id=...` lists up to 25 chosen reports without their results, and reading a report that has never run, such as one migrated from v1, starts its first run. `POST /v2/reports/{report_id}/runs` runs a report again on the current catalog version, reusing existing simulation runs unless `force` is set, and `GET /v2/reports/{report_id}/runs` lists its runs with their versions. A country whose `countries` row does not allow society-wide calculations gets a 422 naming it when it submits or re-runs a report, before anything is sent to the simulation service. A report whose stored result the simulation service cannot read back is a 502. Nothing calls these routes yet. diff --git a/changelog.d/v2-saved-records-v2-ids.changed.md b/changelog.d/v2-saved-records-v2-ids.changed.md new file mode 100644 index 000000000..f7f44f31e --- /dev/null +++ b/changelog.d/v2-saved-records-v2-ids.changed.md @@ -0,0 +1 @@ +`/v2/user-reports` accepts a v2 `report_id` and `/v2/user-simulations` a v2 `simulation_id`, alone or with the v1 ID, when a record is saved or repointed; the simulation service is asked first whether the report or simulation exists in the country, and an unknown one is a 422. A saved report or simulation holding both IDs must name the v2 record migrated from that v1 record, or the request is a 400. A shared report (`/v2/shared-reports/{id}`) now shows its `report_id`. diff --git a/docs/engineering/migration-contracts.md b/docs/engineering/migration-contracts.md index f39a5e369..7f56f8d13 100644 --- a/docs/engineering/migration-contracts.md +++ b/docs/engineering/migration-contracts.md @@ -7,8 +7,8 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | Metric | Count | | --- | ---: | | route group count | 9 | -| workflow count | 16 | -| request count | 89 | +| workflow count | 18 | +| request count | 96 | | db entity count | 6 | | sim flow count | 3 | @@ -18,12 +18,12 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | --- | --- | --- | --- | | `health` | `health`, `simulation-gateway-check`, `liveness-check`, `readiness-check` | `none` | `none` | | `specification` | `specification` | `none` | `none` | -| `metadata` | `metadata`, `datasets`, `economy-options`, `parameter-values`, `parameters`, `regions`, `tax-benefit-model-versions`, `tax-benefit-models`, `variables` | `metadata` | `none` | +| `metadata` | `metadata`, `countries`, `datasets`, `economy-options`, `parameter-values`, `parameters`, `regions`, `tax-benefit-model-versions`, `tax-benefit-models`, `variables` | `metadata` | `none` | | `policy` | `policy`, `policies`, `user-policy`, `user-policies` | `policy` | `none` | | `household` | `household`, `households`, `user-household`, `user-households`, `calculate`, `calculate-full` | `household` | `household` | | `economy` | `economy` | `simulation` | `economy` | | `simulation` | `simulation`, `simulations`, `user-simulations` | `simulation` | `economy` | -| `report` | `report`, `user-reports`, `shared-reports` | `report` | `report` | +| `report` | `report`, `reports`, `user-reports`, `shared-reports` | `report` | `report` | | `user_profile` | `user-profile`, `me` | `user` | `none` | ## App V2 Workflow Contracts @@ -128,12 +128,12 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `POST` | `/v2/user-simulations?country_id=us` | 201 | `simulation` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-simulations?country_id=us` | 200 | `simulation` | `Authorization` | `status`, `message`, `result.items`, `result.offset`, `result.limit`, `result.has_more` | -| `PATCH` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.updated_at` | +| `PATCH` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.updated_at` | | `DELETE` | `/v2/user-simulations/{association_id}?country_id=us` | 204 | `simulation` | `Authorization` | | | `POST` | `/v2/user-simulations?country_id=us` | 201 | `simulation` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-simulations?country_id=us` | 200 | `simulation` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.items`, `result.offset`, `result.limit`, `result.has_more` | -| `PATCH` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.updated_at` | +| `PATCH` | `/v2/user-simulations/{association_id}?country_id=us` | 200 | `simulation` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.simulation_id`, `result.item.legacy_simulation_id`, `result.item.name`, `result.item.description`, `result.item.updated_at` | | `DELETE` | `/v2/user-simulations/{association_id}?country_id=us` | 204 | `simulation` | `PolicyEngine-Unauthenticated-User` | | ### `user_report_associations_v2` @@ -146,14 +146,14 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `POST` | `/v2/user-reports?country_id=us` | 201 | `report` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-reports?country_id=us` | 200 | `report` | `Authorization` | `status`, `message`, `result.items`, `result.offset`, `result.limit`, `result.has_more` | -| `PATCH` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.updated_at` | +| `PATCH` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `Authorization` | `status`, `message`, `result.item.id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.updated_at` | | `DELETE` | `/v2/user-reports/{association_id}?country_id=us` | 204 | `report` | `Authorization` | | | `POST` | `/v2/user-reports?country_id=us` | 201 | `report` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.user_id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/user-reports?country_id=us` | 200 | `report` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.items`, `result.offset`, `result.limit`, `result.has_more` | -| `PATCH` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.updated_at` | +| `PATCH` | `/v2/user-reports/{association_id}?country_id=us` | 200 | `report` | `PolicyEngine-Unauthenticated-User` | `status`, `message`, `result.item.id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.description`, `result.item.user_policy_id`, `result.item.baseline_user_policy_id`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.last_run_at`, `result.item.updated_at` | | `DELETE` | `/v2/user-reports/{association_id}?country_id=us` | 204 | `report` | `PolicyEngine-Unauthenticated-User` | | -| `GET` | `/v2/shared-reports/{association_id}?country_id=us` | 200 | `report` | none | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.baseline_reform`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | +| `GET` | `/v2/shared-reports/{association_id}?country_id=us` | 200 | `report` | none | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.report_id`, `result.item.legacy_report_id`, `result.item.name`, `result.item.source_note`, `result.item.provisions`, `result.item.validation`, `result.item.budget_window_years`, `result.item.baseline_reform`, `result.item.last_run_at`, `result.item.created_at`, `result.item.updated_at` | ### `user_account_v2` @@ -190,6 +190,29 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `GET` | `/us/metadata` | 200 | `metadata` | none | `status`, `result.current_law_id`, `result.economy_options.region`, `result.economy_options.time_period` | | `GET` | `/uk/metadata` | 200 | `metadata` | none | `status`, `result.current_law_id`, `result.economy_options.region`, `result.economy_options.time_period` | +### `country_selection_v2` + +- Current contract: `typed_v2_resources` +- Future owner: Countries table and one country list in code + +| Method | Path | Status | Route group | Required headers | Stable response fields | +| --- | --- | ---: | --- | --- | --- | +| `GET` | `/v2/countries` | 200 | `metadata` | none | `status`, `message`, `result.items` | + +### `economy_reports_v2` + +- Current contract: `typed_v2_resources` +- Future owner: Economy reports on v2 + +| Method | Path | Status | Route group | Required headers | Stable response fields | +| --- | --- | ---: | --- | --- | --- | +| `POST` | `/v2/reports?country_id=us` | 201 | `report` | none | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.kind`, `result.item.year`, `result.item.baseline_simulation`, `result.item.reform_simulation`, `result.item.status`, `result.item.run`, `result.item.result` | +| `GET` | `/v2/reports?country_id=us&report_id={report_id}` | 200 | `report` | none | `status`, `message`, `result.items` | +| `GET` | `/v2/reports/{report_id}?country_id=us` | 200 | `report` | none | `status`, `message`, `result.item.id`, `result.item.legacy_report_id`, `result.item.status`, `result.item.run`, `result.item.result` | +| `GET` | `/v2/reports/by-legacy-number/{legacy_report_id}?country_id=us` | 200 | `report` | none | `status`, `message`, `result.item.id`, `result.item.legacy_report_id`, `result.item.status` | +| `POST` | `/v2/reports/{report_id}/runs?country_id=us` | 201 | `report` | none | `status`, `message`, `result.item.id`, `result.item.trigger`, `result.item.tax_benefit_model_version_id`, `result.item.policyengine_version`, `result.item.status` | +| `GET` | `/v2/reports/{report_id}/runs?country_id=us` | 200 | `report` | none | `status`, `message`, `result.items` | + ### `metadata_resources_v2_preview` - Current contract: `typed_v2_resources` diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d4a043d91..0b5c165b4 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -294,3 +294,33 @@ last dependent table is dropped. Its only post-generation correction drops the nine generated `v2_*` enum types at the end of the baseline downgrade. The v2 Postgres lifecycle test covers empty upgrade, downgrade to base, and re-upgrade so stale enum types cannot make the generated baseline non-reversible. + +### API v2 countries rows + +Revision `fb1a3e40d3d6` was autogenerated after adding the `countries` table +and replacing the US-and-UK checks on `users.primary_country` and on the +`country_id` of households, policies and the four user association tables +with foreign keys to it. It contains one hand-written operation: it inserts +the `uk` and `us` rows, both household and society-wide capable, between the +generated table creation and the generated foreign keys, which otherwise fail +on any database whose rows already name either country. The rows are exactly +the values the replaced checks allowed. The generated downgrade removes them +with the table, and its restored checks fail if any record names another +country. + +Inserting rows is a data operation, which this skill otherwise forbids. The +maintainers approved this insert as a one-off exception, because the new +foreign keys need the rows in the same revision. Unit tests pin the +operation, and the v2 Postgres lifecycle test covers an upgrade that keeps +existing rows, the downgrade, a downgrade refused by another country's +records, and the re-upgrade. + +The table must list exactly the countries in +`policyengine_api/supported_countries.py`: the Public API refuses to start +when they differ (`policyengine_api/startup_country_check.py`). Listing +another country later needs a revision that inserts its row, and the same +human migration decision. The deploy runs that revision before it starts +the new release, so instances of the previous release that start afterwards, +and a rollback to it, refuse to start until the revision is downgraded. The +module docstring of `supported_countries.py` lists the other checks a new +country must widen. diff --git a/docs/generated/migration_contracts.json b/docs/generated/migration_contracts.json index 7a0f1b67d..a81086a9c 100644 --- a/docs/generated/migration_contracts.json +++ b/docs/generated/migration_contracts.json @@ -1,10 +1,10 @@ { "metadata": { "db_entity_count": 6, - "request_count": 89, + "request_count": 96, "route_group_count": 9, "sim_flow_count": 3, - "workflow_count": 16 + "workflow_count": 18 }, "route_groups": [ { @@ -31,6 +31,7 @@ "name": "metadata", "path_segments": [ "metadata", + "countries", "datasets", "economy-options", "parameter-values", @@ -89,6 +90,7 @@ "name": "report", "path_segments": [ "report", + "reports", "user-reports", "shared-reports" ], @@ -801,6 +803,7 @@ "status", "message", "result.item.id", + "result.item.simulation_id", "result.item.legacy_simulation_id", "result.item.name", "result.item.description", @@ -890,6 +893,7 @@ "status", "message", "result.item.id", + "result.item.simulation_id", "result.item.legacy_simulation_id", "result.item.name", "result.item.description", @@ -1000,6 +1004,7 @@ "status", "message", "result.item.id", + "result.item.report_id", "result.item.legacy_report_id", "result.item.name", "result.item.description", @@ -1110,6 +1115,7 @@ "status", "message", "result.item.id", + "result.item.report_id", "result.item.legacy_report_id", "result.item.name", "result.item.description", @@ -1144,6 +1150,7 @@ "message", "result.item.id", "result.item.country_id", + "result.item.report_id", "result.item.legacy_report_id", "result.item.name", "result.item.source_note", @@ -1348,6 +1355,122 @@ } ] }, + { + "current_contract": "typed_v2_resources", + "future_owner_pr": "Countries table and one country list in code", + "name": "country_selection_v2", + "requests": [ + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/countries", + "route_group": "metadata", + "stable_response_fields": [ + "status", + "message", + "result.items" + ] + } + ] + }, + { + "current_contract": "typed_v2_resources", + "future_owner_pr": "Economy reports on v2", + "name": "economy_reports_v2", + "requests": [ + { + "expected_status": 201, + "headers": [], + "method": "POST", + "path": "/v2/reports?country_id=us", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.item.id", + "result.item.country_id", + "result.item.kind", + "result.item.year", + "result.item.baseline_simulation", + "result.item.reform_simulation", + "result.item.status", + "result.item.run", + "result.item.result" + ] + }, + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/reports?country_id=us&report_id={report_id}", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.items" + ] + }, + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/reports/{report_id}?country_id=us", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.item.id", + "result.item.legacy_report_id", + "result.item.status", + "result.item.run", + "result.item.result" + ] + }, + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/reports/by-legacy-number/{legacy_report_id}?country_id=us", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.item.id", + "result.item.legacy_report_id", + "result.item.status" + ] + }, + { + "expected_status": 201, + "headers": [], + "method": "POST", + "path": "/v2/reports/{report_id}/runs?country_id=us", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.item.id", + "result.item.trigger", + "result.item.tax_benefit_model_version_id", + "result.item.policyengine_version", + "result.item.status" + ] + }, + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/reports/{report_id}/runs?country_id=us", + "route_group": "report", + "stable_response_fields": [ + "status", + "message", + "result.items" + ] + } + ] + }, { "current_contract": "typed_v2_resources", "future_owner_pr": "Later metadata read cutover and v2 route-prefix removal", diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 8c3cc5f5d..e7f4bfedb 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -164,7 +164,7 @@ flowchart LR | --- | --- | --- | --- | --- | --- | --- | --- | | `api.catalog.initializer-not-at-startup` | Application startup modules never import the metadata catalog initializer. | mechanical | enforced | yes | pytest `tests/unit/v2/test_catalog_initialization.py` | `pr.yml:test` | **unproven** | | `api.startup.fail-closed-configuration` | Default Cloud SQL startup needs no Supabase configuration, and a selected v2 runtime without its database URL fails closed. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | -| `api.startup.no-import-side-effects` | Importing v2 modules, starting up and serving a request open no network connections and create no files or runtime SQLite databases. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | +| `api.startup.no-import-side-effects` | Importing v2 modules opens no network connections, starting up and serving a request open none unless the process runs on Cloud Run or has a v2 runtime database configured, and none of them creates files or runtime SQLite databases. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | | `api.startup.no-runtime-sqlite-modules` | Runtime SQLite modules are absent from the production package. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | ```mermaid @@ -294,7 +294,7 @@ flowchart LR | `api.quality.command-line-modules-ruff-exemptions` | Application modules that also run as command-line programs (if __name__ == "__main__") may print their output (T201); every other ruff rule applies to them. | mechanical | enforced | yes | ruff per-file ignore of T201 for `policyengine_api/{data/v2/catalog/initialization,data/v2/household_migration_qualification,data/v2/policy_migration_qualification,data/v2/privilege_reconciliation,data/v2/user_association_migration_qualification,worker_spm_release}.py` | `pr.yml:lint` | — | | `api.quality.command-line-scripts-ruff-exemptions` | Command-line scripts under scripts/ may print their output, run other programs, fetch fixed URLs, stamp today's date and pick random samples (T201, S603, S607, S310, DTZ011, S311); every other ruff rule applies to them. | mechanical | enforced | yes | ruff per-file ignore of DTZ011, S310, S311, S603, S607, T201 for `scripts/**` | `pr.yml:lint` | — | | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — | -| `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | +| `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, v2 routes and services, the user sign-in and account-deletion libraries, the country list and its startup check pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | | `api.quality.ruff-lint` | All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in pyproject.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212. | mechanical | enforced | yes | ruff check (ALL) | `pr.yml:lint` | — | | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** | diff --git a/docs/migration/stage-12a-user-login.md b/docs/migration/stage-12a-user-login.md index c01657bc3..6b18d20b2 100644 --- a/docs/migration/stage-12a-user-login.md +++ b/docs/migration/stage-12a-user-login.md @@ -241,9 +241,9 @@ contract shared with `policyengine-app-v2`. The generated route inventory is in A saved report's ID is its shareable link, so `GET /v2/shared-reports/{id}` takes no credential and sits outside the authenticated user-report router. It -never provisions a user. Its item is `{id, country_id, legacy_report_id, name, -source_note, provisions, validation, budget_window_years, baseline_reform, -last_run_at, created_at, updated_at}`: +never provisions a user. Its item is `{id, country_id, report_id, +legacy_report_id, name, source_note, provisions, validation, +budget_window_years, baseline_reform, last_run_at, created_at, updated_at}`: - `user_id`, `user_policy_id`, and `baseline_user_policy_id` are left out; the owner and the owner's saved reforms stay private. @@ -583,5 +583,22 @@ The disposable-Postgres suite proves the following: - A retention policy for unauthenticated users whose browsers never return. Nothing expires them today. - An explicit claim flow for legacy-app users (`POST /v2/me/legacy-claims`). -- Accepting v2 `simulation_id` / `report_id` once the Simulation Entrypoint owns - those records. + +## v2 report and simulation IDs + +Saved reports and saved simulations also accept the v2 IDs of the +Simulation Entrypoint's reports and simulations: `report_id` on +`/v2/user-reports` and `simulation_id` on `/v2/user-simulations`, on create +(alone or with the v1 ID) and on `PATCH` (to repoint; neither ID can be +cleared). The Simulation Entrypoint owns those records, so the columns keep +no foreign key; before a v2 ID is saved, the Public API asks the entrypoint +for that record, outside the v2 transaction, as it checks a v1 ID in the v1 +database: + +- an unknown record, or one of another country, is a 422; +- a saved report or simulation holding both IDs records a migration, so its + v2 record must be the one the entrypoint records as migrated from that v1 + record, or the request is a 400. A `PATCH` that sets one ID is checked + against the other one the saved record keeps, and is a 409 if the saved + record changes during that check; +- an entrypoint that cannot be used is a 503, and nothing is saved. diff --git a/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py b/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py new file mode 100644 index 000000000..1b02ec819 --- /dev/null +++ b/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py @@ -0,0 +1,195 @@ +"""add countries table + +Revision ID: fb1a3e40d3d6 +Revises: fbb7488ebe8a +Create Date: 2026-10-22 02:16:50.126715 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "fb1a3e40d3d6" +down_revision: Union[str, None] = "fbb7488ebe8a" +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! ### + countries = op.create_table( + "countries", + sa.Column("id", sqlmodel.sql.sqltypes.AutoString(length=2), nullable=False), + sa.Column("household_enabled", sa.Boolean(), nullable=False), + sa.Column("society_wide_enabled", sa.Boolean(), nullable=False), + sa.CheckConstraint( + "household_enabled OR NOT society_wide_enabled", + name=op.f("ck_countries_society_wide_requires_household"), + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_countries")), + ) + # Approved hand-written data operation: the one exception to generated-only + # revisions, decided by the maintainers for this table alone. The foreign + # keys below fail on any database whose rows already name the US or the UK + # unless both countries are listed first, so their rows are inserted here, + # in the same revision. They are the countries the replaced checks + # allowed, and both support household and society-wide calculations. The + # generated downgrade removes them with the table. Listing another + # country later needs its own revision and the same kind of decision. + op.bulk_insert( + countries, + [ + {"id": "uk", "household_enabled": True, "society_wide_enabled": True}, + {"id": "us", "household_enabled": True, "society_wide_enabled": True}, + ], + ) + op.create_foreign_key( + op.f("fk_households_country_id_countries"), + "households", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_households_country"), "households", type_="check") + op.create_foreign_key( + op.f("fk_policies_country_id_countries"), + "policies", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_policies_country"), "policies", type_="check") + op.create_foreign_key( + op.f("fk_user_household_associations_country_id_countries"), + "user_household_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_household_associations_country"), + "user_household_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_user_policies_country_id_countries"), + "user_policies", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_user_policies_country"), "user_policies", type_="check") + op.create_foreign_key( + op.f("fk_user_report_associations_country_id_countries"), + "user_report_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_report_associations_country"), + "user_report_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_user_simulation_associations_country_id_countries"), + "user_simulation_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_simulation_associations_country"), + "user_simulation_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_users_primary_country_countries"), + "users", + "countries", + ["primary_country"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_users_primary_country"), "users", type_="check") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_check_constraint( + op.f("ck_users_primary_country"), + "users", + "primary_country::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_users_primary_country_countries"), "users", type_="foreignkey" + ) + op.create_check_constraint( + op.f("ck_user_simulation_associations_country"), + "user_simulation_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_simulation_associations_country_id_countries"), + "user_simulation_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_report_associations_country"), + "user_report_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_report_associations_country_id_countries"), + "user_report_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_policies_country"), + "user_policies", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_policies_country_id_countries"), + "user_policies", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_household_associations_country"), + "user_household_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_household_associations_country_id_countries"), + "user_household_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_policies_country"), + "policies", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_policies_country_id_countries"), "policies", type_="foreignkey" + ) + op.create_check_constraint( + op.f("ck_households_country"), + "households", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_households_country_id_countries"), "households", type_="foreignkey" + ) + op.drop_table("countries") + # ### end Alembic commands ### diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index 1f981c56a..fc5809cf1 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -10,6 +10,7 @@ from policyengine_api.observability import get_runtime from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.runtime_cache.client import close_runtime_cache_clients +from policyengine_api.startup_country_check import check_countries_table_at_startup from policyengine_api.warmup import run_startup_warmup @@ -24,6 +25,10 @@ def _close_runtime_resources() -> None: shutdown_callback=_close_runtime_resources, ) +# Refuse to start when the v2 countries table and the country list in code +# disagree (see policyengine_api.startup_country_check). +check_countries_table_at_startup() + # Warm the simulation machinery before serving (see policyengine_api.warmup). # POLICYENGINE_API_STARTUP_WARMUP=0 skips it. if os.environ.get("POLICYENGINE_API_STARTUP_WARMUP", "1") != "0": diff --git a/policyengine_api/data/v2/catalog/catalog_selection.py b/policyengine_api/data/v2/catalog/catalog_selection.py index d9b50fc76..1e97b6e3c 100644 --- a/policyengine_api/data/v2/catalog/catalog_selection.py +++ b/policyengine_api/data/v2/catalog/catalog_selection.py @@ -12,8 +12,9 @@ TaxBenefitModel, TaxBenefitModelVersion, ) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS -SUPPORTED_V2_COUNTRY_IDS = frozenset({"us", "uk"}) +SUPPORTED_V2_COUNTRY_IDS = frozenset(SUPPORTED_COUNTRY_IDS) MAXIMUM_POLICYENGINE_VERSION_LENGTH = 128 diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index 0872eea09..0817e41ae 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -27,15 +27,15 @@ RegionRecord, VariableRecord, ) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS if TYPE_CHECKING: from policyengine_api.json_types import JSONValue -SUPPORTED_COUNTRIES = ("us", "uk") +SUPPORTED_COUNTRIES = SUPPORTED_COUNTRY_IDS REQUIRED_DEPENDENCIES = ( "policyengine-core", - "policyengine-us", - "policyengine-uk", + *(f"policyengine-{country_id}" for country_id in SUPPORTED_COUNTRY_IDS), ) REVIEWED_DEFAULT_DATASETS = { "us": "populace_us_2024", @@ -928,7 +928,7 @@ def extract_installed_catalog() -> NormalizedCatalog: bundle=get_current_bundle(), policyengine_version=policyengine_version, models={ - "us": policyengine.us.model, - "uk": policyengine.uk.model, + country_id: getattr(policyengine, country_id).model + for country_id in SUPPORTED_COUNTRIES }, ) diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index fd4c05fdf..2c7b6986d 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 = "fbb7488ebe8a" +V2_ALEMBIC_HEAD_REVISION = "fb1a3e40d3d6" 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 c63655bb6..a3aedae09 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -31,6 +31,7 @@ Stage12RunStatus, Stage12SimulationRole, ) +from policyengine_api.data.v2.models.countries import Country # noqa: E402 from policyengine_api.data.v2.models.household_mappings import ( # noqa: E402 LegacyHouseholdMapping, ) @@ -102,6 +103,7 @@ "ChangeAggregate", "CongressionalDistrictImpact", "ConstituencyImpact", + "Country", "Dataset", "DatasetVersion", "DecileImpact", diff --git a/policyengine_api/data/v2/models/associations.py b/policyengine_api/data/v2/models/associations.py index 2c5cd21b1..dd2beb337 100644 --- a/policyengine_api/data/v2/models/associations.py +++ b/policyengine_api/data/v2/models/associations.py @@ -51,10 +51,6 @@ class UserHouseholdAssociation(TimestampedModel, table=True): "country_id", name="uq_user_household_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_household_associations_country", - ), sa.CheckConstraint( "household_id IS NOT NULL OR legacy_household_id IS NOT NULL", name="ck_user_household_associations_household_ref", @@ -91,7 +87,11 @@ class UserHouseholdAssociation(TimestampedModel, table=True): ) household_id: UUID | None = Field(default=None, index=True) legacy_household_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) @@ -113,10 +113,6 @@ class UserPolicy(TimestampedModel, table=True): "country_id", name="uq_user_policies_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_policies_country", - ), sa.CheckConstraint( "policy_id IS NOT NULL OR legacy_policy_id IS NOT NULL", name="ck_user_policies_policy_ref", @@ -157,7 +153,11 @@ class UserPolicy(TimestampedModel, table=True): ) policy_id: UUID | None = Field(default=None, index=True) legacy_policy_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) baseline: str | None = Field(default=None, max_length=16) @@ -172,10 +172,13 @@ class UserPolicy(TimestampedModel, table=True): class UserSimulationAssociation(TimestampedModel, table=True): - """A user's saved v1 simulation; ``simulation_id`` is a logical reference. + """A user's saved simulation: a v2 simulation, a v1 simulation, or both. - Simulations are owned by the Simulation Entrypoint, so the v2 reference - deliberately has no cross-service foreign key. + Simulations are owned by the Simulation Entrypoint, so ``simulation_id`` is + a logical reference with no cross-service foreign key; the Public API asks + the entrypoint whether it exists before saving it. + ``legacy_simulation_id`` is the v1 simulation. A row holding both records + that the v2 simulation was migrated from that v1 simulation. """ __tablename__ = "user_simulation_associations" @@ -185,10 +188,6 @@ class UserSimulationAssociation(TimestampedModel, table=True): "country_id", name="uq_user_simulation_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_simulation_associations_country", - ), sa.CheckConstraint( "simulation_id IS NOT NULL OR legacy_simulation_id IS NOT NULL", name="ck_user_simulation_associations_simulation_ref", @@ -214,7 +213,11 @@ class UserSimulationAssociation(TimestampedModel, table=True): ) simulation_id: UUID | None = Field(default=None, index=True) legacy_simulation_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) # Retired; see the module docstring. @@ -225,11 +228,15 @@ class UserSimulationAssociation(TimestampedModel, table=True): class UserReportAssociation(TimestampedModel, table=True): - """A user's saved v1 report; ``report_id`` is a logical reference. - - ``legacy_report_id`` is the requested v1 report output identifier. The - optional ``user_policy_id`` links a v3 report to the same owner's saved - reform and is cleared when that reform is deleted. + """A user's saved report: a v2 report, a v1 report output, or both. + + Reports are owned by the Simulation Entrypoint, so ``report_id`` is a + logical reference with no cross-service foreign key; the Public API asks + the entrypoint whether it exists before saving it. ``legacy_report_id`` is + the requested v1 report output identifier. A row holding both records that + the v2 report was migrated from that v1 report output. The optional + ``user_policy_id`` links a v3 report to the same owner's saved reform and + is cleared when that reform is deleted. """ __tablename__ = "user_report_associations" @@ -239,10 +246,6 @@ class UserReportAssociation(TimestampedModel, table=True): "country_id", name="uq_user_report_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_report_associations_country", - ), sa.CheckConstraint( "report_id IS NOT NULL OR legacy_report_id IS NOT NULL", name="ck_user_report_associations_report_ref", @@ -273,7 +276,11 @@ class UserReportAssociation(TimestampedModel, table=True): ) report_id: UUID | None = Field(default=None, index=True) legacy_report_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) user_policy_id: UUID | None = Field( diff --git a/policyengine_api/data/v2/models/countries.py b/policyengine_api/data/v2/models/countries.py new file mode 100644 index 000000000..fc3ae8677 --- /dev/null +++ b/policyengine_api/data/v2/models/countries.py @@ -0,0 +1,37 @@ +"""Canonical SQLModel table recording what each served country supports. + +Which countries exist is decided in code: ``policyengine_api.supported_countries`` +lists them, and this table must hold exactly one row per listed country. A +row records the calculations its country supports: a household-only country +has ``household_enabled`` set and ``society_wide_enabled`` unset. + +``users.primary_country`` and the ``country_id`` of households, policies and +the four user association tables reference ``countries.id``. Other v2 country +columns do not: the mapping tables that pair v1 records with their v2 copies +keep their US-and-UK checks, and the simulation-domain tables, the comparison +run tables and the retired association ``country`` columns store unconstrained +text. +""" + +import sqlalchemy as sa +from sqlmodel import Field, SQLModel + + +class Country(SQLModel, table=True): + """One served country and the calculations it supports. + + A society-wide country must also be household capable. A country that a + user, household, policy or user association names cannot be deleted. + """ + + __tablename__ = "countries" + __table_args__ = ( + sa.CheckConstraint( + "household_enabled OR NOT society_wide_enabled", + name="ck_countries_society_wide_requires_household", + ), + ) + + id: str = Field(primary_key=True, max_length=2) + household_enabled: bool + society_wide_enabled: bool diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 78e76b2ab..181266d9d 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -53,10 +53,6 @@ class Household(TimestampedModel, table=True): "content_hash", name="uq_households_canonicalization_content_hash", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_households_country", - ), sa.CheckConstraint( "default_year IS NULL OR default_year BETWEEN 1900 AND 2200", name="ck_households_default_year", @@ -78,7 +74,11 @@ class Household(TimestampedModel, table=True): ), ) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) default_year: int | None = None household_data: dict[str, Any] = Field( sa_type=sa.JSON().with_variant(sa.dialects.postgresql.JSONB(), "postgresql") diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 97e4a619e..f0da6e8cc 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -40,10 +40,6 @@ class Policy(TimestampedModel, table=True): "content_hash", name="uq_policies_canonicalization_content_hash", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_policies_country", - ), sa.CheckConstraint( "canonicalization_version > 0", name="ck_policies_canonicalization_version", @@ -64,7 +60,11 @@ class Policy(TimestampedModel, table=True): ), ) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) tax_benefit_model_id: UUID = Field( foreign_key="tax_benefit_models.id", ondelete="RESTRICT", diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 85c0ea19a..f3b987123 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -44,10 +44,6 @@ class User(IdentifiedModel, table=True): __tablename__ = "users" __table_args__ = ( sa.UniqueConstraint("email", name="uq_users_email"), - sa.CheckConstraint( - "primary_country IN ('us', 'uk')", - name="ck_users_primary_country", - ), sa.CheckConstraint( "user_type IN ('unauthenticated', 'standard')", name="ck_users_user_type", @@ -61,7 +57,12 @@ class User(IdentifiedModel, table=True): first_name: str | None = Field(default=None, max_length=255) last_name: str | None = Field(default=None, max_length=255) email: str | None = Field(default=None, max_length=320, index=True) - primary_country: str | None = Field(default=None, max_length=2) + primary_country: str | None = Field( + default=None, + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) # The server default keeps inserts by earlier Public API revisions, which # do not map this column, creating standard users. user_type: str = Field( diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 8f8c34515..cec2fa66e 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -132,6 +132,9 @@ def _access( # Stage 10 saved-policy mirror. "legacy_user_mappings": _access(runtime=READ_WRITE), "legacy_user_policy_mappings": _access(runtime=READ_WRITE), + # Served countries; the migration that creates the table inserts its rows, + # and the Public API reads them. + "countries": _access(runtime=READ), # Users and their saved work (Stage 12A). "user_household_associations": _access(runtime=READ_WRITE), "user_identities": _access(runtime=READ_WRITE), diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 5d5a1a38d..9657599ef 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -13,15 +13,24 @@ from policyengine_api.fastapi_routes.v2.user_auth import authenticate_account_request from policyengine_api.libs.user_auth import UNAUTHENTICATED_USER_HEADER from policyengine_api.readiness import is_ready +from policyengine_api.services.v2.simulation_entrypoint.services import ( + DeferredEconomyReportSimulation, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationServiceUnavailableError, +) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Sequence from uuid import UUID from policyengine_api.json_types import JSONObject from policyengine_api.services.v2.associations.types import ( LegacyReferenceChecker, + SimulationReferenceChecker, ) + from policyengine_api.services.v2.countries.types import CountryRead + from policyengine_api.services.v2.economy_reports.types import EconomyReportFields from policyengine_api.services.v2.households.types import ( HouseholdPage, HouseholdRead, @@ -53,6 +62,11 @@ PolicyPage, PolicyRead, ) + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportRunRecord, + EconomyReportSimulation, + ) from policyengine_api.services.v2.user_accounts.types import ( AccountDeletion, AccountExport, @@ -312,6 +326,72 @@ def get_economy_options( ... +class V2CountryResourceService(Protocol): + """Route-independent countries list read for one request.""" + + def list_countries(self) -> list[CountryRead]: + """Return every served country and its capability, ordered by ID.""" + ... + + +class V2EconomyReportResourceService(Protocol): + """Router-facing economy report operations, run by the Simulation Entrypoint.""" + + def submit_report( + self, *, country_id: str, report_input: EconomyReportFields + ) -> EconomyReportRecord: + """Create a report and its first run on the current catalog version. + + Raise ``EconomyReportsUnsupportedError`` for a country without economy + reports and ``SimulationRequestRefusedError`` when the entrypoint + cannot run it. + """ + ... + + def get_report(self, *, country_id: str, report_id: UUID) -> EconomyReportRecord: + """Return a report of the country, starting its first run if it has none. + + Raise ``EconomyReportNotFoundError`` when the country has no such report. + """ + ... + + def get_report_by_legacy_number( + self, *, country_id: str, legacy_report_id: int + ) -> EconomyReportRecord: + """Return the report migrated from a v1 report output number, as ``get_report``. + + Raise ``EconomyReportNotFoundError`` when none of the country's was. + """ + ... + + def run_report( + self, *, country_id: str, report_id: UUID, force: bool + ) -> EconomyReportRunRecord: + """Run a report again on the current catalog version (trigger ``manual``). + + Existing simulation runs are reused unless ``force`` is set. Raise + ``EconomyReportsUnsupportedError`` for a country without economy + reports, ``EconomyReportNotFoundError`` for an unknown report and + ``SimulationRequestRefusedError`` when the entrypoint cannot run it. + """ + ... + + def list_reports( + self, *, country_id: str, report_ids: Sequence[UUID] + ) -> tuple[EconomyReportRecord, ...]: + """Return the country's reports among ``report_ids``, in order, once each.""" + ... + + def list_report_runs( + self, *, country_id: str, report_id: UUID + ) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, newest first. + + Raise ``EconomyReportNotFoundError`` when the country has no such report. + """ + ... + + class V2PolicyResourceService(Protocol): """Route-independent native policy operations for one request.""" @@ -510,7 +590,12 @@ def create_user_simulation( principal: AccountPrincipal, association_input: UserSimulationCreateFields, ) -> UserSimulationRead: - """Save a v1 simulation for the caller, provisioning their user if needed.""" + """Save a simulation for the caller, provisioning their user if needed. + + It names a v2 simulation, checked with the simulation service, a v1 + simulation, or both; with both, the v2 simulation must be the one + migrated from that v1 simulation. + """ ... def get_user_simulation( @@ -545,7 +630,10 @@ def patch_user_simulation( ) -> UserSimulationRead: """Update the supplied fields of the caller's association and return it. - A new ``legacy_simulation_id`` repoints it to another v1 simulation. + A new ``simulation_id`` or ``legacy_simulation_id`` repoints it; it is + checked as on create, together with the reference the association + keeps, and AssociationConflictError is raised if the association + changes during that check. """ ... @@ -569,9 +657,11 @@ def create_user_report( principal: AccountPrincipal, association_input: UserReportCreateFields, ) -> UserReportRead: - """Save a v1 report for the caller, provisioning their user if needed. + """Save a report for the caller, provisioning their user if needed. - Linked saved reforms must be the caller's own. + It names a v2 report, checked with the simulation service, a v1 report + output, or both; with both, the v2 report must be the one migrated from + that v1 report output. Linked saved reforms must be the caller's own. """ ... @@ -617,7 +707,13 @@ def patch_user_report( association_id: UUID, association_input: UserReportUpdateInput, ) -> UserReportRead: - """Update the supplied fields of the caller's association and return it.""" + """Update the supplied fields of the caller's association and return it. + + A new ``report_id`` or ``legacy_report_id`` repoints it; it is checked + as on create, together with the reference the association keeps, and + AssociationConflictError is raised if the association changes during + that check. + """ ... def delete_user_report( @@ -742,6 +838,69 @@ def _default_v2_metadata_reader_factory() -> V2MetadataResourceReader: ) +def _default_v2_country_service_factory() -> V2CountryResourceService: + from policyengine_api.data.v2.database import ( + get_v2_session_factory, + ) + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.services import V2CountryService + + return V2CountryService(CountryDatabaseSession(get_v2_session_factory())) + + +SIMULATION_CLIENT_PACKAGE = "policyengine_api_simulation_client" + + +def _default_economy_report_simulation() -> EconomyReportSimulation: + """Return the process's adapter for the entrypoint's economy report routes. + + Raise ``SimulationServiceUnavailableError`` when the generated client is + not installed, or the entrypoint URL or gateway auth is unusable. + """ + try: + from policyengine_api.libs.simulation_economy_reports import ( + shared_economy_reports, + ) + except ModuleNotFoundError as error: + if (error.name or "").split(".")[0] != SIMULATION_CLIENT_PACKAGE: + raise + msg = ( + "the generated simulation client (policyengine-api-simulation-client) " + "is not installed" + ) + raise SimulationServiceUnavailableError(msg) from error + return shared_economy_reports() + + +def _deferred_economy_report_simulation() -> EconomyReportSimulation: + """Build the adapter only when a request reaches the simulation service. + + Checks the Public API makes itself, such as the household-only refusal, + then answer without the simulation client or its settings. + """ + return DeferredEconomyReportSimulation(_default_economy_report_simulation) + + +def _default_v2_economy_report_service_factory() -> V2EconomyReportResourceService: + from policyengine_api.data.v2.database import ( + get_v2_session_factory, + ) + from policyengine_api.services.v2.economy_reports.database_session import ( + EconomyReportDatabaseSession, + ) + from policyengine_api.services.v2.economy_reports.services import ( + V2EconomyReportService, + ) + + return V2EconomyReportService( + EconomyReportDatabaseSession(get_v2_session_factory()), + _deferred_economy_report_simulation(), + current_policyengine_version=_running_policyengine_version(), + ) + + def _default_v2_policy_service_factory() -> V2PolicyResourceService: from policyengine_api.data.v2.database import ( get_v2_session_factory, @@ -826,6 +985,15 @@ def _default_legacy_reference_checker() -> LegacyReferenceChecker: return DeferredLegacyReferenceService(_build_v1_legacy_reference_service) +def _default_simulation_reference_checker() -> SimulationReferenceChecker: + """Defer the simulation client until a request checks a v2 reference.""" + from policyengine_api.services.v2.simulation_refs.services import ( + SimulationReferenceService, + ) + + return SimulationReferenceService(_deferred_economy_report_simulation()) + + def _default_v2_user_policy_service_factory() -> V2UserPolicyResourceService: from policyengine_api.data.v2.database import ( get_v2_session_factory, @@ -872,6 +1040,7 @@ def _default_v2_user_simulation_service_factory() -> V2UserSimulationResourceSer return V2UserSimulationService( UserSimulationDatabaseSession(get_v2_session_factory()), _default_legacy_reference_checker(), + _default_simulation_reference_checker(), ) @@ -887,6 +1056,7 @@ def _default_v2_user_report_service_factory() -> V2UserReportResourceService: return V2UserReportService( UserReportDatabaseSession(get_v2_session_factory()), _default_legacy_reference_checker(), + _default_simulation_reference_checker(), ) @@ -917,6 +1087,10 @@ class NativeRouteDependencies: metadata_reader_factory: Callable[[], MetadataReader] specification_provider: Callable[[], JSONObject] v2_metadata_reader_factory: Callable[[], V2MetadataResourceReader] | None = None + v2_country_service_factory: Callable[[], V2CountryResourceService] | None = None + v2_economy_report_service_factory: ( + Callable[[], V2EconomyReportResourceService] | None + ) = None v2_policy_service_factory: Callable[[], V2PolicyResourceService] | None = None v2_household_service_factory: Callable[[], V2HouseholdResourceService] | None = None v2_user_policy_service_factory: Callable[[], V2UserPolicyResourceService] | None = ( @@ -947,6 +1121,10 @@ def defaults(cls) -> NativeRouteDependencies: metadata_reader_factory=_default_metadata_reader_factory, specification_provider=_default_specification_provider, v2_metadata_reader_factory=_default_v2_metadata_reader_factory, + v2_country_service_factory=_default_v2_country_service_factory, + v2_economy_report_service_factory=( + _default_v2_economy_report_service_factory + ), v2_policy_service_factory=_default_v2_policy_service_factory, v2_household_service_factory=_default_v2_household_service_factory, v2_user_policy_service_factory=(_default_v2_user_policy_service_factory), diff --git a/policyengine_api/fastapi_routes/v2/countries/__init__.py b/policyengine_api/fastapi_routes/v2/countries/__init__.py new file mode 100644 index 000000000..3d7f8eaf0 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/__init__.py @@ -0,0 +1 @@ +"""FastAPI adapters for the v2 countries list.""" diff --git a/policyengine_api/fastapi_routes/v2/countries/response_models.py b/policyengine_api/fastapi_routes/v2/countries/response_models.py new file mode 100644 index 000000000..3ea3c7ee8 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/response_models.py @@ -0,0 +1,71 @@ +"""Strict HTTP response models for the native v2 countries list.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from pydantic import BaseModel, ConfigDict + +from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse + +if TYPE_CHECKING: + from policyengine_api.services.v2.countries.types import CountryRead + + +class StrictCountryAPIModel(BaseModel): + """Strict response base with dataclass conversion support.""" + + model_config = ConfigDict(extra="forbid", from_attributes=True) + + +class CountryItem(StrictCountryAPIModel): + """One served country and the calculations it supports. + + ``id`` is the value other v2 resources take as ``country_id``. + ``household_enabled`` means household calculations and saved households + are available; ``society_wide_enabled`` means society-wide reports are, + and is only set for a country that is also household enabled. + """ + + id: str + household_enabled: bool + society_wide_enabled: bool + + @classmethod + def from_read(cls, country: CountryRead) -> CountryItem: + """Build the response item from the service read record.""" + return cls.model_validate(country) + + +class CountryListResult(StrictCountryAPIModel): + """Every served country, ordered by ``id``.""" + + items: list[CountryItem] + + +class CountryListResponse(StrictCountryAPIModel): + """Success envelope: ``status`` "ok", a null ``message`` and the countries.""" + + status: Literal["ok"] = "ok" + message: None = None + result: CountryListResult + + +COUNTRY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { + 405: { + "model": V2ErrorResponse, + "description": "The countries list supports GET only.", + }, + 422: { + "model": V2ErrorResponse, + "description": "The request has query parameters; the list takes none.", + }, + 500: { + "model": V2ErrorResponse, + "description": "The countries query failed internally.", + }, + 503: { + "model": V2ErrorResponse, + "description": "The v2 database is not configured or cannot be reached.", + }, +} diff --git a/policyengine_api/fastapi_routes/v2/countries/routes.py b/policyengine_api/fastapi_routes/v2/countries/routes.py new file mode 100644 index 000000000..279aaf54f --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/routes.py @@ -0,0 +1,84 @@ +"""Native FastAPI route for the v2 countries list.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +from fastapi import APIRouter, Depends +from sqlalchemy.exc import SQLAlchemyError +from starlette.responses import JSONResponse + +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + _default_v2_country_service_factory, +) +from policyengine_api.fastapi_routes.query_parameters import query_dependency +from policyengine_api.fastapi_routes.v2.countries.response_models import ( + COUNTRY_ERROR_RESPONSES, + CountryItem, + CountryListResponse, + CountryListResult, +) +from policyengine_api.fastapi_routes.v2.errors import V2ErrorMapping, run_v2_operation +from policyengine_api.query_parameters import CountryListQuery + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2CountryResourceService, + ) + + +CountryListQueryDependency = Annotated[ + CountryListQuery, Depends(query_dependency(CountryListQuery)) +] + +COUNTRY_ERROR_MAPPINGS = ( + V2ErrorMapping( + (V2ConfigurationError, SQLAlchemyError), + 503, + "V2 countries are unavailable", + ), +) + + +def _service_factory( + dependencies: NativeRouteDependencies, +) -> Callable[[], V2CountryResourceService]: + if dependencies.v2_country_service_factory is not None: + return dependencies.v2_country_service_factory + return _default_v2_country_service_factory + + +def build_v2_country_router( + dependencies: NativeRouteDependencies, +) -> APIRouter: + """Build the countries route without opening a database connection.""" + router = APIRouter(prefix="/v2", responses=COUNTRY_ERROR_RESPONSES) + service_factory = _service_factory(dependencies) + + @router.get( + "/countries", + response_model=CountryListResponse, + summary="List the served countries and the calculations each supports", + ) + def list_countries( + _query: CountryListQueryDependency, + ) -> CountryListResponse | JSONResponse: + def read() -> CountryListResponse: + countries = service_factory().list_countries() + return CountryListResponse( + result=CountryListResult( + items=[CountryItem.from_read(country) for country in countries] + ) + ) + + return run_v2_operation( + read, + COUNTRY_ERROR_MAPPINGS, + fallback_message="V2 countries query failed", + ) + + return router diff --git a/policyengine_api/fastapi_routes/v2/errors.py b/policyengine_api/fastapi_routes/v2/errors.py index 0c1f38525..f180a2bf0 100644 --- a/policyengine_api/fastapi_routes/v2/errors.py +++ b/policyengine_api/fastapi_routes/v2/errors.py @@ -1,5 +1,6 @@ """Shared error response contract for all native API v2 routes.""" +import traceback from collections.abc import Callable, Mapping, Sequence from typing import Annotated, Literal, NamedTuple, TypeVar @@ -94,3 +95,19 @@ def run_v2_operation( message = str(error) if mapping.message is None else mapping.message return v2_error_response(mapping.status_code, message) return v2_error_response(500, fallback_message) + + +def failure_log_attributes(error: BaseException) -> dict[str, str]: + """Describe a failure for a log record without its message. + + Exception messages can carry bound SQL parameters or token claims, so the + record holds the error's type and the locations of its innermost frames. + """ + frames = traceback.extract_tb(error.__traceback__)[-8:] + return { + "error_type": type(error).__name__, + "error_frames": "; ".join( + f"{frame.filename.rsplit('/', 1)[-1]}:{frame.lineno} in {frame.name}" + for frame in frames + ), + } diff --git a/policyengine_api/fastapi_routes/v2/reports/__init__.py b/policyengine_api/fastapi_routes/v2/reports/__init__.py new file mode 100644 index 000000000..4e77012f8 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/reports/__init__.py @@ -0,0 +1 @@ +"""Native API v2 economy report routes.""" diff --git a/policyengine_api/fastapi_routes/v2/reports/request_models.py b/policyengine_api/fastapi_routes/v2/reports/request_models.py new file mode 100644 index 000000000..3fe27d938 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/reports/request_models.py @@ -0,0 +1,24 @@ +"""Strict HTTP request models for v2 economy reports.""" + +from policyengine_api.query_parameters import CountryId +from policyengine_api.services.v2.economy_reports.types import ( + EconomyReportFields, + EconomyReportRunFields, +) + + +class EconomyReportCreateRequest(EconomyReportFields): + """A new economy report; ``country_id`` must match the query's. + + ``baseline`` and ``reform`` each name a v2 policy and dynamic, and an + empty object means current law. Options are not accepted yet. + """ + + country_id: CountryId + + +class EconomyReportRunRequest(EconomyReportRunFields): + """A re-run on the current catalog version; ``{}`` reuses existing runs. + + With ``force``, both simulations are calculated again as new attempts. + """ diff --git a/policyengine_api/fastapi_routes/v2/reports/response_models.py b/policyengine_api/fastapi_routes/v2/reports/response_models.py new file mode 100644 index 000000000..398f76a05 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/reports/response_models.py @@ -0,0 +1,245 @@ +"""Strict HTTP response models for v2 economy reports.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, JsonValue + +from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse +from policyengine_api.services.v2.simulation_entrypoint.types import ReportStatus + +if TYPE_CHECKING: + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportRunRecord, + ) + + +class StrictReportAPIModel(BaseModel): + """Strict response base that converts the service's dataclass records.""" + + model_config = ConfigDict( + extra="forbid", + allow_inf_nan=False, + from_attributes=True, + ) + + +class RunErrorItem(StrictReportAPIModel): + """Why a run failed, without user data.""" + + code: str + message: str | None + + +class SimulationItem(StrictReportAPIModel): + """A simulation: one v2 policy and dynamic applied to one region. + + A null ``policy_id`` and ``dynamic_id`` mean current law. A null + ``dataset_name`` means the region's default dataset in each catalog + version. + """ + + id: UUID + kind: str + country_id: str + policy_id: UUID | None + dynamic_id: UUID | None + region_code: str | None + dataset_name: str | None + legacy_simulation_id: int | None + created_at: datetime + + +class SimulationRunItem(StrictReportAPIModel): + """One execution of a simulation for a catalog version and year. + + ``status`` is ``pending``, ``dispatching``, ``running``, ``succeeded`` or + ``failed``; ``attempt`` counts forced recalculations. + """ + + id: UUID + simulation_id: UUID + status: str + attempt: int + trigger: str + error: RunErrorItem | None + dispatched_at: datetime | None + completed_at: datetime | None + created_at: datetime + + +class ComparisonItem(StrictReportAPIModel): + """A report run's comparison step, which computes the report's result.""" + + status: str + error: RunErrorItem | None + dispatched_at: datetime | None + completed_at: datetime | None + + +class ReportRunItem(StrictReportAPIModel): + """One calculation of a report: two simulation runs and their comparison. + + ``trigger`` is ``initial`` for the run a submission starts, ``manual`` for + a re-run and ``system`` for the run the first read of a report with no run + starts. ``tax_benefit_model_version_id`` and ``policyengine_version`` name + the catalog version it ran on. While it runs, ``stage`` is ``queued``, + ``simulating`` or ``outputs``. + """ + + id: UUID + report_id: UUID + trigger: str + tax_benefit_model_version_id: UUID + policyengine_version: str | None + year: int + status: ReportStatus + stage: str | None + error: RunErrorItem | None + baseline_run: SimulationRunItem + reform_run: SimulationRunItem + comparison: ComparisonItem + created_at: datetime + + @classmethod + def from_record(cls, record: EconomyReportRunRecord) -> ReportRunItem: + """Build the response item from the service's report run record.""" + return cls.model_validate(record) + + +class ReportSummaryItem(StrictReportAPIModel): + """An economy report comparing a baseline and a reform simulation. + + ``status`` is the newest run's, or ``not_run`` for a report with no run. + ``legacy_report_id`` is the v1 report output it was migrated from. + """ + + id: UUID + country_id: str + kind: Literal["economy"] = "economy" + year: int + options: dict[str, JsonValue] + legacy_report_id: int | None + created_at: datetime + baseline_simulation: SimulationItem + reform_simulation: SimulationItem + status: ReportStatus + run: ReportRunItem | None + + @classmethod + def from_record(cls, record: EconomyReportRecord) -> ReportSummaryItem: + """Build the summary from the service's report record, without its result.""" + return cls.model_validate(record) + + +class ReportItem(ReportSummaryItem): + """An economy report with its result. + + ``result`` holds today's economy result format (``SingleYearMacroOutput``) + once the newest run has succeeded, and is null otherwise. + """ + + result: dict[str, JsonValue] | None + + @classmethod + def from_record(cls, record: EconomyReportRecord) -> ReportItem: + """Build the response item from the service's report record.""" + return cls.model_validate(record) + + +class ReportDetailResult(StrictReportAPIModel): + """Result holding one report in ``item``.""" + + item: ReportItem + + +class ReportListResult(StrictReportAPIModel): + """The listed reports in ``items``, in the order asked, without results.""" + + items: list[ReportSummaryItem] + + +class ReportRunDetailResult(StrictReportAPIModel): + """Result holding one report run in ``item``.""" + + item: ReportRunItem + + +class ReportRunListResult(StrictReportAPIModel): + """A report's runs in ``items``, newest first.""" + + items: list[ReportRunItem] + + +ResultT = TypeVar("ResultT") + + +class ReportSuccessResponse(StrictReportAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + + status: Literal["ok"] = "ok" + message: None = None + result: ResultT + + +class ReportDetailResponse(ReportSuccessResponse[ReportDetailResult]): + """Success response holding one report.""" + + +class ReportListResponse(ReportSuccessResponse[ReportListResult]): + """Success response holding listed reports.""" + + +class ReportRunDetailResponse(ReportSuccessResponse[ReportRunDetailResult]): + """Success response holding one report run.""" + + +class ReportRunListResponse(ReportSuccessResponse[ReportRunListResult]): + """Success response holding a report's runs.""" + + +def _error(description: str) -> dict[str, Any]: + return {"model": V2ErrorResponse, "description": description} + + +# Every report route can answer these. +REPORT_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { + 422: _error("The request does not match the schema."), + 500: _error("The operation could not be completed."), + 503: _error("The simulation service or the v2 database is unavailable."), +} +# Submitting or re-running a report can also be refused for the report itself. +REPORT_REFUSAL_RESPONSE: dict[int | str, dict[str, Any]] = { + 422: _error( + "The request does not match the schema, the country does not have " + "economy reports (such as a household-only country), or the simulation " + "service cannot run the report (for example, no such region or " + "dataset in the current catalog version); the message says which." + ), +} +REPORT_NOT_FOUND_RESPONSE: dict[int | str, dict[str, Any]] = { + 404: _error("No report has this ID or v1 number in this country."), +} +REPORT_BODY_RESPONSE: dict[int | str, dict[str, Any]] = { + 413: _error("The request body exceeds 16 KiB."), +} +SUBMIT_RESPONSES: dict[int | str, dict[str, Any]] = { + 400: _error("The body and query name different countries."), + **REPORT_BODY_RESPONSE, + **REPORT_REFUSAL_RESPONSE, +} +READ_RESPONSES: dict[int | str, dict[str, Any]] = { + **REPORT_NOT_FOUND_RESPONSE, + 502: _error( + "The simulation service has the report but cannot read back its stored result." + ), +} +RUN_RESPONSES: dict[int | str, dict[str, Any]] = { + **REPORT_NOT_FOUND_RESPONSE, + **REPORT_BODY_RESPONSE, + **REPORT_REFUSAL_RESPONSE, +} diff --git a/policyengine_api/fastapi_routes/v2/reports/routes.py b/policyengine_api/fastapi_routes/v2/reports/routes.py new file mode 100644 index 000000000..639b1c239 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/reports/routes.py @@ -0,0 +1,368 @@ +"""Native FastAPI routes for v2 economy reports. + +Reports are calculated and stored by the Simulation Entrypoint; these routes +need no user credential. A user keeps a report through ``/v2/user-reports``. +""" + +from __future__ import annotations + +import contextlib +import logging +from typing import TYPE_CHECKING, Annotated, TypeVar +from uuid import UUID + +from fastapi import APIRouter, Depends, Path +from sqlalchemy.exc import SQLAlchemyError +from starlette.responses import JSONResponse + +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + _default_v2_economy_report_service_factory, +) +from policyengine_api.fastapi_routes.query_parameters import query_dependency +from policyengine_api.fastapi_routes.v2.errors import ( + V2ErrorMapping, + failure_log_attributes, + run_v2_operation, + v2_error_response, +) +from policyengine_api.fastapi_routes.v2.reports.request_models import ( + EconomyReportCreateRequest, + EconomyReportRunRequest, +) +from policyengine_api.fastapi_routes.v2.reports.response_models import ( + READ_RESPONSES, + REPORT_ERROR_RESPONSES, + REPORT_NOT_FOUND_RESPONSE, + RUN_RESPONSES, + SUBMIT_RESPONSES, + ReportDetailResponse, + ReportDetailResult, + ReportItem, + ReportListResponse, + ReportListResult, + ReportRunDetailResponse, + ReportRunDetailResult, + ReportRunItem, + ReportRunListResponse, + ReportRunListResult, + ReportSummaryItem, +) +from policyengine_api.fastapi_routes.v2.user_resources import request_size_limit +from policyengine_api.query_parameters import ( + MAXIMUM_LEGACY_RESOURCE_ID, + ReportCollectionQuery, + ReportCreateQuery, + ReportDetailQuery, +) +from policyengine_api.services.v2.economy_reports.validators import ( + EconomyReportNotFoundError, + EconomyReportsUnsupportedError, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationRecordUnreadableError, + SimulationRequestRefusedError, + SimulationServiceUnavailableError, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2EconomyReportResourceService, + ) + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportRunRecord, + ) + +logger = logging.getLogger(__name__) + +MAXIMUM_REPORT_REQUEST_BYTES = 16 * 1024 +enforce_report_request_size = request_size_limit( + MAXIMUM_REPORT_REQUEST_BYTES, label="16 KiB" +) + +ReportCreateQueryDependency = Annotated[ + ReportCreateQuery, Depends(query_dependency(ReportCreateQuery)) +] +ReportCollectionQueryDependency = Annotated[ + ReportCollectionQuery, Depends(query_dependency(ReportCollectionQuery)) +] +ReportDetailQueryDependency = Annotated[ + ReportDetailQuery, Depends(query_dependency(ReportDetailQuery)) +] +LegacyReportNumber = Annotated[ + int, + Path( + ge=1, + le=MAXIMUM_LEGACY_RESOURCE_ID, + description="The v1 report output number the report was migrated from", + ), +] + +# Refusals a caller can act on carry their own message. +REPORT_REFUSALS: tuple[type[Exception], ...] = ( + EconomyReportsUnsupportedError, + SimulationRequestRefusedError, + EconomyReportNotFoundError, +) +REPORT_ERROR_MAPPINGS = ( + V2ErrorMapping((EconomyReportNotFoundError,), 404), + V2ErrorMapping( + (SimulationRecordUnreadableError,), + 502, + "The simulation service cannot read back this report's result", + ), + V2ErrorMapping( + (EconomyReportsUnsupportedError, SimulationRequestRefusedError), 422 + ), + V2ErrorMapping( + (SimulationServiceUnavailableError,), + 503, + "The simulation service is unavailable", + ), + V2ErrorMapping( + (V2ConfigurationError, SQLAlchemyError), + 503, + "V2 report persistence is unavailable", + ), +) + +OperationT = TypeVar("OperationT") + + +def _log_failure(error: Exception) -> None: + """Record why an operation failed, unless it was a refusal the caller sees. + + Errors are recorded by type and the locations of their innermost frames, + not their messages, which can carry bound SQL parameters; only the + simulation service's unavailability adds its reason. Logging never + changes the response. + """ + with contextlib.suppress(Exception): + attributes = failure_log_attributes(error) + if isinstance(error, SimulationServiceUnavailableError): + attributes["reason"] = str(error)[:500] + logger.error( + "v2_report_operation_failed", + extra={"policyengine_attributes": attributes}, + ) + + +def _report_operation( + operation: Callable[[], OperationT], +) -> OperationT | JSONResponse: + def logged() -> OperationT: + try: + return operation() + except Exception as error: + if not isinstance(error, REPORT_REFUSALS): + _log_failure(error) + raise + + return run_v2_operation( + logged, + REPORT_ERROR_MAPPINGS, + fallback_message="V2 report operation failed", + ) + + +def _service_factory( + dependencies: NativeRouteDependencies, +) -> Callable[[], V2EconomyReportResourceService]: + if dependencies.v2_economy_report_service_factory is not None: + return dependencies.v2_economy_report_service_factory + return _default_v2_economy_report_service_factory + + +def _report_detail(report: EconomyReportRecord) -> ReportDetailResponse: + return ReportDetailResponse( + result=ReportDetailResult(item=ReportItem.from_record(report)) + ) + + +def _report_list(reports: tuple[EconomyReportRecord, ...]) -> ReportListResponse: + return ReportListResponse( + result=ReportListResult( + items=[ReportSummaryItem.from_record(report) for report in reports] + ) + ) + + +def _report_run_detail(report_run: EconomyReportRunRecord) -> ReportRunDetailResponse: + return ReportRunDetailResponse( + result=ReportRunDetailResult(item=ReportRunItem.from_record(report_run)) + ) + + +def _report_run_list( + runs: tuple[EconomyReportRunRecord, ...], +) -> ReportRunListResponse: + return ReportRunListResponse( + result=ReportRunListResult( + items=[ReportRunItem.from_record(run) for run in runs] + ) + ) + + +def build_v2_report_router( + dependencies: NativeRouteDependencies, +) -> APIRouter: + """Build the economy report routes without opening a connection.""" + router = APIRouter(prefix="/v2", responses=REPORT_ERROR_RESPONSES) + service_factory = _service_factory(dependencies) + + @router.post( + "/reports", + response_model=ReportDetailResponse, + status_code=201, + responses=SUBMIT_RESPONSES, + summary="Submit an economy report", + description=( + "Creates the baseline and reform simulations, the report and its " + "first run in the simulation service, on the country's current " + "catalog version, and starts the run. Read the report to follow it. " + "A country without economy reports, such as a household-only one, " + "is refused with 422 before anything is created." + ), + ) + def submit_report( + body: EconomyReportCreateRequest, + query: ReportCreateQueryDependency, + _size: Annotated[None, Depends(enforce_report_request_size)], + ) -> ReportDetailResponse | JSONResponse: + if body.country_id != query.country_id: + return v2_error_response(400, "Body country_id must match query country_id") + return _report_operation( + lambda: _report_detail( + service_factory().submit_report( + country_id=query.country_id, report_input=body + ) + ) + ) + + @router.get( + "/reports", + response_model=ReportListResponse, + summary="List chosen economy reports", + description=( + "Lists the country's reports among the report_id values (at most " + "25), in that order and once each, with their status and newest run " + "but without results. Unknown reports and other countries' are left " + "out. Listing does not start the first run of a report that has " + "none; reading it does." + ), + ) + def list_reports( + query: ReportCollectionQueryDependency, + ) -> ReportListResponse | JSONResponse: + return _report_operation( + lambda: _report_list( + service_factory().list_reports( + country_id=query.country_id, report_ids=query.report_id + ) + ) + ) + + @router.get( + "/reports/by-legacy-number/{legacy_report_id}", + response_model=ReportDetailResponse, + responses=READ_RESPONSES, + summary="Find the report migrated from a v1 report number", + description=( + "Answers as reading the report by its ID, including starting the " + "first run of a report that has none. 404 means no report of this " + "country was migrated from that number." + ), + ) + def get_report_by_legacy_number( + legacy_report_id: LegacyReportNumber, + query: ReportDetailQueryDependency, + ) -> ReportDetailResponse | JSONResponse: + return _report_operation( + lambda: _report_detail( + service_factory().get_report_by_legacy_number( + country_id=query.country_id, legacy_report_id=legacy_report_id + ) + ) + ) + + @router.get( + "/reports/{report_id}", + response_model=ReportDetailResponse, + responses=READ_RESPONSES, + summary="Read an economy report", + description=( + "Returns the report, its newest run and, once that run succeeded, its " + "result. A report with no run, such as one migrated from v1, is " + "calculated the first time it is read: the read starts a run with " + "the system trigger on the current catalog version and answers with " + "it. If the simulation service cannot run it, the report is " + "returned not run." + ), + ) + def get_report( + report_id: UUID, + query: ReportDetailQueryDependency, + ) -> ReportDetailResponse | JSONResponse: + return _report_operation( + lambda: _report_detail( + service_factory().get_report( + country_id=query.country_id, report_id=report_id + ) + ) + ) + + @router.post( + "/reports/{report_id}/runs", + response_model=ReportRunDetailResponse, + status_code=201, + responses=RUN_RESPONSES, + summary="Run an economy report again", + description=( + "Runs the report on the country's current catalog version with the " + "manual trigger. Simulation runs that already exist for that version " + "and year are reused unless force is set, which calculates both " + "simulations again. A country without economy reports is refused " + "with 422 before the simulation service is asked." + ), + ) + def run_report( + report_id: UUID, + body: EconomyReportRunRequest, + query: ReportDetailQueryDependency, + _size: Annotated[None, Depends(enforce_report_request_size)], + ) -> ReportRunDetailResponse | JSONResponse: + return _report_operation( + lambda: _report_run_detail( + service_factory().run_report( + country_id=query.country_id, report_id=report_id, force=body.force + ) + ) + ) + + @router.get( + "/reports/{report_id}/runs", + response_model=ReportRunListResponse, + responses=REPORT_NOT_FOUND_RESPONSE, + summary="List an economy report's runs", + description=( + "Every run of the report, newest first, each with the catalog version " + "and PolicyEngine.py version it ran on." + ), + ) + def list_report_runs( + report_id: UUID, + query: ReportDetailQueryDependency, + ) -> ReportRunListResponse | JSONResponse: + return _report_operation( + lambda: _report_run_list( + service_factory().list_report_runs( + country_id=query.country_id, report_id=report_id + ) + ) + ) + + return router diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index 0521719be..81f1bce8f 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -7,6 +7,9 @@ from fastapi import APIRouter, Request from starlette.responses import JSONResponse +from policyengine_api.fastapi_routes.v2.countries.routes import ( + build_v2_country_router, +) from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse from policyengine_api.fastapi_routes.v2.households.routes import ( build_v2_household_router, @@ -22,6 +25,7 @@ build_v2_metadata_parameter_router, ) from policyengine_api.fastapi_routes.v2.policies.routes import build_v2_policy_router +from policyengine_api.fastapi_routes.v2.reports.routes import build_v2_report_router from policyengine_api.fastapi_routes.v2.user_households.routes import ( build_v2_user_household_router, ) @@ -53,6 +57,8 @@ def build_v2_router( router.include_router(build_v2_user_policy_router(dependencies)) router.include_router(build_v2_household_router(dependencies)) router.include_router(build_v2_policy_router(dependencies)) + router.include_router(build_v2_report_router(dependencies)) + router.include_router(build_v2_country_router(dependencies)) router.include_router(build_v2_metadata_model_router(dependencies)) router.include_router(build_v2_metadata_parameter_router(dependencies)) router.include_router(build_v2_metadata_geography_router(dependencies)) diff --git a/policyengine_api/fastapi_routes/v2/user_reports/response_models.py b/policyengine_api/fastapi_routes/v2/user_reports/response_models.py index de4be2d19..06dc6bc04 100644 --- a/policyengine_api/fastapi_routes/v2/user_reports/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_reports/response_models.py @@ -127,7 +127,8 @@ class ReportValidationItem(StoredSnapshotItem): class UserReportItem(StrictUserReportAPIModel): """A report saved by a user, with the user's name, notes and saved presentation. - ``legacy_report_id`` is the v1 report output it shows. ``user_policy_id`` is + ``report_id`` is the v2 report it shows and ``legacy_report_id`` the v1 + report output; a report migrated from v1 has both. ``user_policy_id`` is the user's saved reform the report was run from, and ``baseline_user_policy_id`` the saved reform it was compared with instead of current law; each is null when not set or once that reform is deleted. @@ -164,10 +165,15 @@ class SharedBaselineReformItem(StrictUserReportAPIModel): class SharedUserReportItem(StrictUserReportAPIModel): - """A saved report as anyone holding its link reads it.""" + """A saved report as anyone holding its link reads it. + + ``report_id`` is the v2 report it shows and ``legacy_report_id`` the v1 + report output. + """ id: UUID country_id: CountryId + report_id: UUID | None legacy_report_id: int | None name: str | None source_note: str | None diff --git a/policyengine_api/fastapi_routes/v2/user_reports/routes.py b/policyengine_api/fastapi_routes/v2/user_reports/routes.py index 241f19882..a39d7b6d2 100644 --- a/policyengine_api/fastapi_routes/v2/user_reports/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_reports/routes.py @@ -46,6 +46,16 @@ ) from policyengine_api.services.v2.identity.types import AccountPrincipal +# Saving or repointing to a v2 report reads it from the simulation service. +UNREADABLE_REPORT_RESPONSES: dict[int | str, dict[str, Any]] = { + 502: { + "model": V2ErrorResponse, + "description": ( + "The simulation service has the v2 report but cannot read it back." + ), + }, +} + SHARED_REPORT_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { 404: { "model": V2ErrorResponse, @@ -121,13 +131,17 @@ def build_v2_user_report_router( "/user-reports", response_model=UserReportDetailResponse, status_code=201, - summary="Save a v1 report for the current user", + responses=UNREADABLE_REPORT_RESPONSES, + summary="Save a report for the current user", description=( "Requires a bearer access token or the PolicyEngine-Unauthenticated-User " "header; the owner is taken from that credential. " - "Stage 12A accepts legacy_report_id only (the requested v1 report " - "output ID), checked in the v1 database. An optional user_policy_id " - "must name the same user's saved policy in the same country." + "The body names a v2 report (report_id), checked with the simulation " + "service, a v1 report output (legacy_report_id), checked in the v1 " + "database, or both; with both, the v2 report must be the one migrated " + "from that v1 report output (400 otherwise). An optional " + "user_policy_id must name the same user's saved policy in the same " + "country." ), ) def create_user_report( @@ -183,6 +197,7 @@ def read() -> UserReportPageResponse: @router.patch( "/user-reports/{association_id}", response_model=UserReportDetailResponse, + responses=UNREADABLE_REPORT_RESPONSES, summary="Update or repoint one of the current user's saved reports", ) def patch_user_report_route( diff --git a/policyengine_api/fastapi_routes/v2/user_resources.py b/policyengine_api/fastapi_routes/v2/user_resources.py index 581065a57..4386df9d2 100644 --- a/policyengine_api/fastapi_routes/v2/user_resources.py +++ b/policyengine_api/fastapi_routes/v2/user_resources.py @@ -4,7 +4,6 @@ import contextlib import logging -import traceback from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar # FastAPI reads the request-size dependency's annotations when routes are @@ -17,6 +16,7 @@ BEARER_CHALLENGE, V2ErrorResponse, V2RequestTooLargeError, + failure_log_attributes, v2_error_response, ) from policyengine_api.services.v2.associations.validators import ( @@ -31,6 +31,13 @@ from policyengine_api.services.v2.legacy_refs.validators import ( LegacyReferenceNotFoundError, ) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationRecordUnreadableError, + SimulationServiceUnavailableError, +) +from policyengine_api.services.v2.simulation_refs.validators import ( + SimulationReferenceNotFoundError, +) from policyengine_api.services.v2.user_accounts.validators import ( AccountDeletionUnavailableError, AccountEditRequiresSignInError, @@ -123,8 +130,8 @@ async def enforce(request: Request) -> None: 422: { "model": V2ErrorResponse, "description": ( - "The request does not match the schema, or a referenced v1 artifact " - "or linked record does not exist." + "The request does not match the schema, or a referenced v1 artifact, " + "v2 report or simulation, or linked record does not exist." ), }, 500: { @@ -133,7 +140,9 @@ async def enforce(request: Request) -> None: }, 503: { "model": V2ErrorResponse, - "description": "User authentication or persistence is unavailable.", + "description": ( + "User authentication, persistence or the simulation service is unavailable." + ), }, } @@ -151,19 +160,9 @@ def _log_failure(event: str, error: Exception) -> None: """ # Observability must not alter the response. with contextlib.suppress(Exception): - frames = traceback.extract_tb(error.__traceback__)[-8:] logger.error( event, - extra={ - "policyengine_attributes": { - "error_type": type(error).__name__, - "error_frames": "; ".join( - f"{frame.filename.rsplit('/', 1)[-1]}:{frame.lineno} " - f"in {frame.name}" - for frame in frames - ), - } - }, + extra={"policyengine_attributes": failure_log_attributes(error)}, ) @@ -192,7 +191,12 @@ class UserResourceFailure(NamedTuple): UserResourceFailure((AccountEditRequiresSignInError,), 403), UserResourceFailure((AssociationReferenceConflictError,), 400), UserResourceFailure( - (LegacyReferenceNotFoundError, AssociationLinkNotFoundError), 422 + ( + LegacyReferenceNotFoundError, + SimulationReferenceNotFoundError, + AssociationLinkNotFoundError, + ), + 422, ), UserResourceFailure( ( @@ -208,7 +212,19 @@ class UserResourceFailure(NamedTuple): (IntegrityError,), 409, "The record conflicts with stored state" ), UserResourceFailure((AccountIdentityDeletionError,), 502), + UserResourceFailure( + (SimulationRecordUnreadableError,), + 502, + "The simulation service cannot read back this report", + logged_event="v2_user_simulation_record_unreadable", + ), UserResourceFailure((AccountDeletionUnavailableError,), 503), + UserResourceFailure( + (SimulationServiceUnavailableError,), + 503, + "The simulation service is unavailable", + logged_event="v2_user_simulation_service_unavailable", + ), # Schema drift and rejected data surface here too, so keep a record. UserResourceFailure( (V2ConfigurationError, SQLAlchemyError), diff --git a/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py b/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py index 70df6f78a..ad964eb5c 100644 --- a/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py @@ -26,8 +26,9 @@ class StrictUserSimulationAPIModel(BaseModel): class UserSimulationItem(StrictUserSimulationAPIModel): """A simulation saved by a user, with the user's name and description for it. - ``legacy_simulation_id`` is the v1 simulation it saves. Associations can be - renamed, pointed at another v1 simulation of the same country, or deleted. + ``simulation_id`` is the v2 simulation it saves and ``legacy_simulation_id`` + the v1 simulation. Associations can be renamed, pointed at another + simulation of the same country, or deleted. """ id: UUID diff --git a/policyengine_api/fastapi_routes/v2/user_simulations/routes.py b/policyengine_api/fastapi_routes/v2/user_simulations/routes.py index 0cac726a4..13b69ed3d 100644 --- a/policyengine_api/fastapi_routes/v2/user_simulations/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_simulations/routes.py @@ -98,12 +98,14 @@ def build_v2_user_simulation_router( "/user-simulations", response_model=UserSimulationDetailResponse, status_code=201, - summary="Save a v1 simulation for the current user", + summary="Save a simulation for the current user", description=( "Requires a bearer access token or the PolicyEngine-Unauthenticated-User " "header; the owner is taken from that credential. " - "Stage 12A accepts legacy_simulation_id only, checked in the v1 " - "database." + "The body names a v2 simulation (simulation_id), checked with the " + "simulation service, a v1 simulation (legacy_simulation_id), checked " + "in the v1 database, or both; with both, the v2 simulation must be " + "the one migrated from that v1 simulation (400 otherwise)." ), ) def create_user_simulation( diff --git a/policyengine_api/libs/gateway_auth.py b/policyengine_api/libs/gateway_auth.py index 817c98c6c..dafaf3125 100644 --- a/policyengine_api/libs/gateway_auth.py +++ b/policyengine_api/libs/gateway_auth.py @@ -296,3 +296,25 @@ def auth_flow( self._token_provider.invalidate() request.headers["Authorization"] = f"Bearer {self._token_provider.get_token()}" yield request + + +def gateway_bearer_auth_from_environment() -> GatewayBearerAuth | None: + """Return bearer auth from the ``GATEWAY_AUTH_*`` variables, or None. + + None means no variable is set and ``GATEWAY_AUTH_REQUIRED`` is not + enabled. Raise GatewayAuthError when the variables are only partly set, or + when auth is required but not configured. + """ + _require_all_or_none_gateway_auth_env() + token_provider = GatewayAuthTokenProvider() + if token_provider.configured: + return GatewayBearerAuth(token_provider) + if gateway_auth_required(): + msg = ( + "Gateway auth is required in this runtime: set " + f"{GATEWAY_AUTH_ISSUER_ENV}, {GATEWAY_AUTH_AUDIENCE_ENV}, " + f"{GATEWAY_AUTH_CLIENT_ID_ENV}, and " + f"{GATEWAY_AUTH_CLIENT_SECRET_ENV}." + ) + raise GatewayAuthError(msg) + return None diff --git a/policyengine_api/libs/simulation_economy_reports.py b/policyengine_api/libs/simulation_economy_reports.py new file mode 100644 index 000000000..193ff68da --- /dev/null +++ b/policyengine_api/libs/simulation_economy_reports.py @@ -0,0 +1,397 @@ +"""The Simulation Entrypoint's economy report routes, through its generated client. + +This is the only module that imports ``policyengine_api_simulation_client``, +the client generated from the Simulation Entrypoint's OpenAPI document. +Everything else uses ``EconomyReportSimulation``, which this module +implements, so the rest of the code and its tests do not need the client. +Nothing imports this module at module level, so the application starts +without the client installed and the economy report routes then answer 503. + +The economy report routes exist only on the Cloud Run Simulation Entrypoint, +so requests always go to ``SIMULATION_ENTRYPOINT_URL``, whichever entrypoint +``SIM_ENTRYPOINT`` selects for the v1 routes. They carry the gateway bearer +token, which needs ``simulation:submit`` to submit and ``simulation:read`` +to read. + +The entrypoint's OpenAPI document declares its error answers, so the +generated client parses them too: 404 (no such record), 422 (it cannot run +the request, with its reason or the refused fields), 503 (it is not +configured), and, for report reads, 500 (a succeeded run's stored result +cannot be read back). Each becomes one of the protocol's errors. A 404 whose +``detail`` is exactly "Not Found" is FastAPI's answer for a route the +entrypoint does not serve, such as a revision without these routes, so it +means the entrypoint cannot be used rather than that the record is missing; +the entrypoint's own 404s name the missing record. +""" + +from __future__ import annotations + +from functools import cache +from typing import TYPE_CHECKING, Any, NoReturn, TypeVar + +import httpx +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import ( + get_economy_report, + get_economy_report_by_legacy_id, + get_simulation, + list_economy_report_runs, + run_economy_report, + submit_economy_report, +) +from policyengine_api_simulation_client.models import ( + EconomyErrorResponse, + EconomyRefusalResponse, + EconomyReportResponse, + EconomyReportRunRequest, + EconomyReportRunRequestTrigger, + EconomyReportRunResponse, + EconomyReportRunsResponse, + EconomyReportSubmissionRequest, + EconomyReportSubmissionRequestCountry, + SimulationChoiceRequest, + SimulationResponse, +) +from policyengine_api_simulation_client.types import UNSET +from policyengine_observability import instrument_httpx + +from policyengine_api.libs.gateway_auth import ( + GatewayAuthError, + gateway_bearer_auth_from_environment, +) +from policyengine_api.libs.simulation_http import ( + attach_request_identifiers, + required_base_url, +) +from policyengine_api.observability import get_runtime +from policyengine_api.services.v2.simulation_entrypoint.types import ( + ComparisonRecord, + EconomyReportRecord, + EconomyReportRunRecord, + ReportStatus, + RunError, + SimulationRecord, + SimulationRecordNotFoundError, + SimulationRecordUnreadableError, + SimulationRequestRefusedError, + SimulationRunRecord, + SimulationServiceUnavailableError, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from uuid import UUID + + from policyengine_api_simulation_client.models import ( + ComparisonResponse, + RunErrorResponse, + SimulationRunResponse, + ) + from policyengine_api_simulation_client.types import Response + + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportSubmission, + RunTrigger, + SimulationChoice, + ) + +SIMULATION_ENTRYPOINT_URL_ENV = "SIMULATION_ENTRYPOINT_URL" +# FastAPI's 404 detail for a route the application does not serve. +UNSERVED_ROUTE_DETAIL = "Not Found" +REQUEST_TIMEOUT_SECONDS = 30.0 +ParsedT = TypeVar("ParsedT") + + +def _refusal_reason(refusal: EconomyRefusalResponse) -> str: + """Return a 422's reason: the entrypoint's text, or each refused field.""" + if isinstance(refusal.detail, str): + return refusal.detail + return "; ".join( + ".".join(str(part) for part in field_error.loc) + f": {field_error.msg}" + for field_error in refusal.detail + ) + + +def _raise_declared_error(status_code: int, answer: object) -> None: + """Raise the protocol's error for one of the entrypoint's declared answers.""" + if isinstance(answer, EconomyRefusalResponse): + raise SimulationRequestRefusedError(_refusal_reason(answer)) + if not isinstance(answer, EconomyErrorResponse): + return + if status_code == httpx.codes.NOT_FOUND and answer.detail != UNSERVED_ROUTE_DETAIL: + raise SimulationRecordNotFoundError(answer.detail) + if status_code == httpx.codes.INTERNAL_SERVER_ERROR: + raise SimulationRecordUnreadableError(answer.detail) + msg = ( + f"the Simulation Entrypoint answered HTTP {status_code}: {answer.detail[:200]}" + ) + raise SimulationServiceUnavailableError(msg) + + +def _parsed(call: Callable[[], Response[Any]], expected: type[ParsedT]) -> ParsedT: + """Run one generated route call and return its parsed success answer. + + A declared error answer raises its protocol error. Network failures, a + gateway token that cannot be fetched, undeclared answers and answers that + do not parse mean the entrypoint cannot be used now. + """ + try: + response = call() + except httpx.HTTPError as error: + msg = f"the Simulation Entrypoint could not be reached: {type(error).__name__}" + raise SimulationServiceUnavailableError(msg) from error + except GatewayAuthError as error: + msg = f"the gateway token for the Simulation Entrypoint is unavailable: {error}" + raise SimulationServiceUnavailableError(msg) from error + except (KeyError, TypeError, ValueError, AttributeError) as error: + _raise_unreadable(error) + if isinstance(response.parsed, expected): + return response.parsed + _raise_declared_error(int(response.status_code), response.parsed) + msg = ( + "the Simulation Entrypoint answered HTTP " + f"{int(response.status_code)} without the expected body" + ) + raise SimulationServiceUnavailableError(msg) + + +def _raise_unreadable(error: Exception) -> NoReturn: + msg = ( + f"the Simulation Entrypoint's answer could not be read: {type(error).__name__}" + ) + raise SimulationServiceUnavailableError(msg) from error + + +def _run_error(error: RunErrorResponse | None) -> RunError | None: + return None if error is None else RunError(code=error.code, message=error.message) + + +def _simulation(simulation: SimulationResponse) -> SimulationRecord: + return SimulationRecord( + id=simulation.id, + kind=simulation.kind, + country_id=simulation.country, + policy_id=simulation.policy_id, + dynamic_id=simulation.dynamic_id, + region_code=simulation.region_code, + dataset_name=simulation.dataset_name, + legacy_simulation_id=simulation.legacy_simulation_id, + created_at=simulation.created_at, + ) + + +def _simulation_run(run: SimulationRunResponse) -> SimulationRunRecord: + return SimulationRunRecord( + id=run.id, + simulation_id=run.simulation_id, + status=run.status, + attempt=run.attempt, + trigger=run.trigger, + error=_run_error(run.error), + dispatched_at=run.dispatched_at, + completed_at=run.completed_at, + created_at=run.created_at, + ) + + +def _comparison(comparison: ComparisonResponse) -> ComparisonRecord: + return ComparisonRecord( + status=comparison.status, + error=_run_error(comparison.error), + dispatched_at=comparison.dispatched_at, + completed_at=comparison.completed_at, + ) + + +def _report_run(run: EconomyReportRunResponse) -> EconomyReportRunRecord: + return EconomyReportRunRecord( + id=run.id, + report_id=run.report_id, + trigger=run.trigger, + tax_benefit_model_version_id=run.catalog_version_id, + policyengine_version=run.policyengine_version, + year=run.year, + status=ReportStatus(run.status.value), + stage=None if run.stage is None else run.stage.value, + error=_run_error(run.error), + baseline_run=_simulation_run(run.baseline_run), + reform_run=_simulation_run(run.reform_run), + comparison=_comparison(run.comparison), + created_at=run.created_at, + ) + + +def _report(report: EconomyReportResponse) -> EconomyReportRecord: + return EconomyReportRecord( + id=report.id, + country_id=report.country, + year=report.year, + options=report.options.to_dict(), + legacy_report_id=report.legacy_report_id, + created_at=report.created_at, + baseline_simulation=_simulation(report.baseline_simulation), + reform_simulation=_simulation(report.reform_simulation), + status=ReportStatus(report.status.value), + run=None if report.run is None else _report_run(report.run), + result=None if report.result is None else report.result.to_dict(), + ) + + +def _simulation_choice(choice: SimulationChoice) -> SimulationChoiceRequest: + return SimulationChoiceRequest( + policy_id=choice.policy_id, dynamic_id=choice.dynamic_id + ) + + +class SimulationEntrypointEconomyReports: + """``EconomyReportSimulation`` over the generated simulation client.""" + + def __init__(self, client: AuthenticatedClient, *, base_url: str) -> None: + """Send every request through ``client``, whose base URL is ``base_url``.""" + self._client = client + self.base_url = base_url + + def submit_report(self, submission: EconomyReportSubmission) -> EconomyReportRecord: + """Create a report and its first run in the entrypoint, and start it.""" + body = EconomyReportSubmissionRequest( + country=EconomyReportSubmissionRequestCountry(submission.country_id), + policyengine_version=submission.policyengine_version, + year=submission.year, + region_code=submission.region_code, + baseline=_simulation_choice(submission.baseline), + reform=_simulation_choice(submission.reform), + dataset_name=( + UNSET if submission.dataset_name is None else submission.dataset_name + ), + ) + return _report( + _parsed( + lambda: submit_economy_report.sync_detailed( + client=self._client, body=body + ), + EconomyReportResponse, + ) + ) + + def get_report(self, report_id: UUID) -> EconomyReportRecord: + """Return a report with its newest run and, once succeeded, its result.""" + return _report( + _parsed( + lambda: get_economy_report.sync_detailed( + report_id, client=self._client + ), + EconomyReportResponse, + ) + ) + + def get_report_by_legacy_id(self, legacy_report_id: int) -> EconomyReportRecord: + """Return the report migrated from a v1 report output number.""" + return _report( + _parsed( + lambda: get_economy_report_by_legacy_id.sync_detailed( + legacy_report_id, client=self._client + ), + EconomyReportResponse, + ) + ) + + def run_report( + self, + report_id: UUID, + *, + policyengine_version: str, + trigger: RunTrigger, + force: bool, + ) -> EconomyReportRunRecord: + """Run a report again on the catalog version of ``policyengine_version``.""" + body = EconomyReportRunRequest( + policyengine_version=policyengine_version, + trigger=EconomyReportRunRequestTrigger(trigger.value), + force=force, + ) + return _report_run( + _parsed( + lambda: run_economy_report.sync_detailed( + report_id, client=self._client, body=body + ), + EconomyReportRunResponse, + ) + ) + + def list_report_runs(self, report_id: UUID) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, newest first.""" + runs = _parsed( + lambda: list_economy_report_runs.sync_detailed( + report_id, client=self._client + ), + EconomyReportRunsResponse, + ) + return tuple(_report_run(run) for run in runs.runs) + + def get_simulation(self, simulation_id: UUID) -> SimulationRecord: + """Return a simulation.""" + return _simulation( + _parsed( + lambda: get_simulation.sync_detailed( + simulation_id, client=self._client + ), + SimulationResponse, + ) + ) + + +def simulation_entrypoint_economy_reports( + base_url: str, + *, + auth: httpx.Auth | None, + transport: httpx.BaseTransport | None = None, +) -> SimulationEntrypointEconomyReports: + """Build the adapter for the entrypoint at ``base_url``. + + ``auth`` adds the bearer token; ``transport`` replaces the network in tests. + """ + http_client = httpx.Client( + base_url=base_url, + timeout=REQUEST_TIMEOUT_SECONDS, + auth=auth, + transport=transport, + event_hooks={"request": [attach_request_identifiers]}, + ) + instrument_httpx(http_client, get_runtime()) + # The httpx client carries the refreshing gateway token, so the generated + # client's own fixed token is never sent. + client = AuthenticatedClient(base_url=base_url, token="").set_httpx_client( + http_client + ) + return SimulationEntrypointEconomyReports(client, base_url=base_url) + + +def _entrypoint_url() -> str: + try: + return required_base_url(SIMULATION_ENTRYPOINT_URL_ENV) + except ValueError as error: + raise SimulationServiceUnavailableError(str(error)) from error + + +def _gateway_auth() -> httpx.Auth | None: + """Return the gateway bearer auth, or None where it may be left unconfigured.""" + try: + return gateway_bearer_auth_from_environment() + except GatewayAuthError as error: + raise SimulationServiceUnavailableError(str(error)) from error + + +def economy_reports_from_environment() -> SimulationEntrypointEconomyReports: + """Build the adapter from ``SIMULATION_ENTRYPOINT_URL`` and the gateway auth. + + Raise ``SimulationServiceUnavailableError`` when either is unusable. + """ + return simulation_entrypoint_economy_reports( + _entrypoint_url(), auth=_gateway_auth() + ) + + +@cache +def shared_economy_reports() -> SimulationEntrypointEconomyReports: + """Return one adapter per process, so its gateway token is fetched once.""" + return economy_reports_from_environment() diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 91ed85f9c..584a73416 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -1,79 +1,39 @@ """HTTP client for the configured PolicyEngine Simulation Entrypoint.""" -import os import sys from dataclasses import dataclass, field from http import HTTPStatus -from urllib.parse import urlparse import httpx from policyengine_observability import instrument_httpx from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError from policyengine_api.gcp_logging import logger -from policyengine_api.libs.gateway_auth import ( - GatewayAuthError, - GatewayAuthTokenProvider, - GatewayBearerAuth, - _require_all_or_none_gateway_auth_env, - gateway_auth_required, +from policyengine_api.libs.gateway_auth import gateway_bearer_auth_from_environment +from policyengine_api.libs.simulation_http import ( + attach_request_identifiers, + required_base_url, ) from policyengine_api.migration_flags import get_sim_entrypoint from policyengine_api.observability import get_runtime -from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.request_context import ( - REQUEST_ID_HEADER, current_observability_id, - current_request_id, ) from policyengine_api.spm import SPMValidationError from policyengine_api.worker_spm import raise_worker_spm_error, validate_worker_spm -def _required_base_url(env_name: str) -> str: - value = os.environ.get(env_name, "").strip() - if not value: - msg = f"{env_name} is required" - raise ValueError(msg) - - parsed = urlparse(value) - if ( - parsed.scheme not in {"http", "https"} - or not parsed.hostname - or parsed.username - or parsed.password - or parsed.query - or parsed.fragment - ): - msg = ( - f"{env_name} must be an absolute HTTP(S) base URL without " - "credentials, a query, or a fragment" - ) - raise ValueError(msg) - return value.rstrip("/") - - def resolve_simulation_entrypoint_url(entrypoint: str | None = None) -> str: """Resolve old/new upstream URLs without conflating their configuration.""" selected_entrypoint = entrypoint or get_sim_entrypoint() if selected_entrypoint == "old_gateway_direct": - return _required_base_url("OLD_SIMULATION_GATEWAY_URL") + return required_base_url("OLD_SIMULATION_GATEWAY_URL") if selected_entrypoint == "cloud_run_simulation_entrypoint": - return _required_base_url("SIMULATION_ENTRYPOINT_URL") + return required_base_url("SIMULATION_ENTRYPOINT_URL") msg = f"Unsupported simulation entrypoint: {selected_entrypoint!r}" raise ValueError(msg) -def _attach_current_request_id(request: httpx.Request) -> None: - """Attach the active API request ID to an outbound simulation request.""" - request_id = current_request_id() - if request_id is not None: - request.headers[REQUEST_ID_HEADER] = request_id - observability_id = current_observability_id() - if observability_id is not None: - request.headers[OBSERVABILITY_ID_HEADER] = observability_id - - class SimulationRunSegments(BaseModel): """How many of a running job's segments have finished, out of its total.""" @@ -223,22 +183,8 @@ def __init__(self, entrypoint: str | None = None) -> None: """ self.entrypoint = entrypoint or get_sim_entrypoint() self.base_url = resolve_simulation_entrypoint_url(self.entrypoint) - self._token_provider = GatewayAuthTokenProvider() - _require_all_or_none_gateway_auth_env() - auth = ( - GatewayBearerAuth(self._token_provider) - if self._token_provider.configured - else None - ) + auth = gateway_bearer_auth_from_environment() if auth is None: - if gateway_auth_required(): - msg = ( - "Gateway auth is required in this runtime: set " - "GATEWAY_AUTH_ISSUER, GATEWAY_AUTH_AUDIENCE, " - "GATEWAY_AUTH_CLIENT_ID, and " - "GATEWAY_AUTH_CLIENT_SECRET." - ) - raise GatewayAuthError(msg) print( # noqa: T201 # the unauthenticated-client warning goes to stderr whatever the logging setup "SimulationEntrypointClient initialized without gateway auth; " "all GATEWAY_AUTH_* env vars are unset and " @@ -249,7 +195,7 @@ def __init__(self, entrypoint: str | None = None) -> None: self.client = httpx.Client( timeout=30.0, auth=auth, - event_hooks={"request": [_attach_current_request_id]}, + event_hooks={"request": [attach_request_identifiers]}, ) instrument_httpx(self.client, get_runtime()) diff --git a/policyengine_api/libs/simulation_http.py b/policyengine_api/libs/simulation_http.py new file mode 100644 index 000000000..d4bdb0345 --- /dev/null +++ b/policyengine_api/libs/simulation_http.py @@ -0,0 +1,61 @@ +"""HTTP conventions shared by the Public API's clients of the simulation service. + +Both the Simulation Entrypoint client of the v1 routes and the economy report +adapter of the v2 routes read the service's base URL from the environment and +send the current API request's correlation identifiers. Importing this module +has no side effects. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_api.request_context import ( + REQUEST_ID_HEADER, + current_observability_id, + current_request_id, +) + +if TYPE_CHECKING: + import httpx + + +def required_base_url(env_name: str) -> str: + """Return the base URL in ``env_name`` without a trailing slash. + + Raise ValueError when it is missing, or is not an absolute HTTP(S) URL free + of credentials, a query and a fragment. + """ + value = os.environ.get(env_name, "").strip() + if not value: + msg = f"{env_name} is required" + raise ValueError(msg) + + parsed = urlparse(value) + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username + or parsed.password + or parsed.query + or parsed.fragment + ): + msg = ( + f"{env_name} must be an absolute HTTP(S) base URL without " + "credentials, a query, or a fragment" + ) + raise ValueError(msg) + return value.rstrip("/") + + +def attach_request_identifiers(request: httpx.Request) -> None: + """Send the current API request's correlation identifiers with ``request``.""" + request_id = current_request_id() + if request_id is not None: + request.headers[REQUEST_ID_HEADER] = request_id + observability_id = current_observability_id() + if observability_id is not None: + request.headers[OBSERVABILITY_ID_HEADER] = observability_id diff --git a/policyengine_api/migration_registry.py b/policyengine_api/migration_registry.py index e49de6986..51f92a8ad 100644 --- a/policyengine_api/migration_registry.py +++ b/policyengine_api/migration_registry.py @@ -37,6 +37,7 @@ class RouteGroupConfig: name="metadata", path_segments=( "metadata", + "countries", "datasets", "economy-options", "parameter-values", @@ -80,7 +81,7 @@ class RouteGroupConfig: ), RouteGroupConfig( name="report", - path_segments=("report", "user-reports", "shared-reports"), + path_segments=("report", "reports", "user-reports", "shared-reports"), db_entity="report", sim_flow="report", ), diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 06e38cf01..60edab0be 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -29,6 +29,7 @@ validate_policyengine_version, ) from policyengine_api.spm import SPMSelection +from policyengine_api.supported_countries import SupportedCountryId from policyengine_api.utils.budget_window import ( BUDGET_WINDOW_MAX_END_YEAR, BUDGET_WINDOW_MAX_YEARS, @@ -38,9 +39,10 @@ from collections.abc import Iterable, Mapping DEFAULT_QUERY_LIMIT = 100 +# Each listed report is one read from the simulation service. +MAXIMUM_LISTED_REPORTS = 25 MAXIMUM_QUERY_LIMIT = 500 MAXIMUM_USER_ID_LENGTH = 255 -SUPPORTED_COUNTRY_IDS = frozenset({"us", "uk"}) def normalize_country_id(value: object) -> object: @@ -57,7 +59,7 @@ def validate_legacy_user_id(value: str) -> str: CountryId = Annotated[ - Literal["us", "uk"], + SupportedCountryId, BeforeValidator(normalize_country_id), Field(description="Supported PolicyEngine country ID"), ] @@ -214,6 +216,10 @@ class PaginationQuery(StrictQueryParameters): limit: Limit = DEFAULT_QUERY_LIMIT +class CountryListQuery(StrictQueryParameters): + """Query contract for the countries list, which takes no parameters.""" + + class PolicyCreateQuery(CatalogQuery): """Query contract for native immutable policy creation.""" @@ -278,6 +284,30 @@ class UserSimulationCollectionQuery(CountryQuery, PaginationQuery): legacy_simulation_id: LegacyResourceId | None = None +class ReportCreateQuery(CountryQuery): + """Query contract for submitting an economy report.""" + + +class ReportDetailQuery(CountryQuery): + """Query contract for reading one report, or acting on it, in a country.""" + + +class ReportCollectionQuery(CountryQuery): + """Query contract for listing chosen reports of a country. + + ``report_id`` is repeated once per report, at most 25 times. + """ + + report_id: Annotated[ + list[ResourceId], + Field( + min_length=1, + max_length=MAXIMUM_LISTED_REPORTS, + description="A report to list; repeat for each report", + ), + ] + + class UserReportCollectionQuery(CountryQuery, PaginationQuery): """Query contract for the signed-in user's report associations.""" diff --git a/policyengine_api/services/v2/associations/transformations.py b/policyengine_api/services/v2/associations/transformations.py new file mode 100644 index 000000000..9b85570ab --- /dev/null +++ b/policyengine_api/services/v2/associations/transformations.py @@ -0,0 +1,24 @@ +"""Pure transformations shared by the saved-record association services.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.associations.types import RecordReferences + +if TYPE_CHECKING: + from uuid import UUID + + +def repointed_references( + stored: RecordReferences, *, v2_id: UUID | None, legacy_id: int | None +) -> RecordReferences: + """Return the references an association holds once an update applies. + + ``v2_id`` and ``legacy_id`` are the update's; None leaves the stored one, + since neither reference can be cleared. + """ + return RecordReferences( + v2_id=stored.v2_id if v2_id is None else v2_id, + legacy_id=stored.legacy_id if legacy_id is None else legacy_id, + ) diff --git a/policyengine_api/services/v2/associations/types.py b/policyengine_api/services/v2/associations/types.py index e3da0b36c..232be9f11 100644 --- a/policyengine_api/services/v2/associations/types.py +++ b/policyengine_api/services/v2/associations/types.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Protocol, TypeAlias from pydantic import BaseModel, ConfigDict, FiniteFloat, StringConstraints @@ -15,6 +16,7 @@ # does); Pydantic needs TypeAliasType to build the recursive validator. if TYPE_CHECKING: from collections.abc import Collection, Mapping + from uuid import UUID from policyengine_api.services.v2.legacy_refs.types import LegacyRecordKind @@ -74,3 +76,53 @@ def require_existing( ) -> None: """Raise ``LegacyReferenceNotFoundError`` if any requested ID is missing.""" ... + + +class SimulationReferenceChecker(Protocol): + """Existence checks of v2 reports and simulations, run before a v2 transaction. + + Both raise ``SimulationReferenceNotFoundError`` when the simulation service + has no such record in the country, and ``SimulationServiceUnavailableError`` + when it cannot be asked. + """ + + def require_report( + self, + *, + country_id: str, + report_id: UUID, + legacy_report_id: int | None = None, + ) -> None: + """Raise unless the country has the report. + + With ``legacy_report_id``, also raise ``AssociationReferenceConflictError`` + unless the report was migrated from that v1 report output. + """ + ... + + def require_simulation( + self, + *, + country_id: str, + simulation_id: UUID, + legacy_simulation_id: int | None = None, + ) -> None: + """Raise unless the country has the simulation. + + With ``legacy_simulation_id``, also raise + ``AssociationReferenceConflictError`` unless the simulation was + migrated from that v1 simulation. + """ + ... + + +@dataclass(frozen=True) +class RecordReferences: + """The v2 record and the v1 record a saved association names. + + Either may be None, not both. When both are set, the v2 record is the one + migrated from the v1 record. + """ + + v2_id: UUID | None + legacy_id: int | None diff --git a/policyengine_api/services/v2/associations/validators.py b/policyengine_api/services/v2/associations/validators.py index a5cff660e..c2a9f8aca 100644 --- a/policyengine_api/services/v2/associations/validators.py +++ b/policyengine_api/services/v2/associations/validators.py @@ -4,10 +4,18 @@ from typing import TYPE_CHECKING +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationServiceUnavailableError, +) + if TYPE_CHECKING: from uuid import UUID - from policyengine_api.services.v2.associations.types import LegacyReferenceChecker + from policyengine_api.services.v2.associations.types import ( + LegacyReferenceChecker, + RecordReferences, + SimulationReferenceChecker, + ) class AssociationNotFoundError(LookupError): @@ -44,6 +52,20 @@ def require_legacy_reference_checker( return legacy_references +class SimulationReferenceCheckUnavailableError(SimulationServiceUnavailableError): + """The service was built without the v2 report and simulation checker.""" + + +def require_simulation_reference_checker( + simulation_references: SimulationReferenceChecker | None, +) -> SimulationReferenceChecker: + """Return the checker, or raise SimulationReferenceCheckUnavailableError.""" + if simulation_references is None: + msg = "v2 report and simulation references cannot be checked here" + raise SimulationReferenceCheckUnavailableError(msg) + return simulation_references + + def require_not_mirror_owned(*, mirror_owned: bool, association_id: UUID) -> None: """Raise MirrorOwnedAssociationError if the v1 saved-policy mirror owns the row.""" if mirror_owned: @@ -52,3 +74,19 @@ def require_not_mirror_owned(*, mirror_owned: bool, association_id: UUID) -> Non "and cannot be changed here" ) raise MirrorOwnedAssociationError(msg) + + +def require_unchanged_references( + current: RecordReferences, *, checked: RecordReferences, association_id: UUID +) -> None: + """Raise ``AssociationConflictError`` if the references changed since the check. + + A saved record's references are checked before its transaction opens, so + another update in between makes that check stale. + """ + if current != checked: + msg = ( + f"association {association_id} changed while its references were " + "checked; try again" + ) + raise AssociationConflictError(msg) diff --git a/policyengine_api/services/v2/countries/__init__.py b/policyengine_api/services/v2/countries/__init__.py new file mode 100644 index 000000000..f0f7b972b --- /dev/null +++ b/policyengine_api/services/v2/countries/__init__.py @@ -0,0 +1 @@ +"""Application services for the v2 countries list.""" diff --git a/policyengine_api/services/v2/countries/database_connectors/__init__.py b/policyengine_api/services/v2/countries/database_connectors/__init__.py new file mode 100644 index 000000000..01c1fa018 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_connectors/__init__.py @@ -0,0 +1 @@ +"""SQL-facing connector functions for the v2 countries service.""" diff --git a/policyengine_api/services/v2/countries/database_connectors/reads.py b/policyengine_api/services/v2/countries/database_connectors/reads.py new file mode 100644 index 000000000..a13a8f3e4 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_connectors/reads.py @@ -0,0 +1,17 @@ +"""Database selections used by the v2 countries list and the startup check.""" + +from __future__ import annotations + +from sqlmodel import Session, col, select + +from policyengine_api.data.v2.models import Country + + +def read_countries(session: Session) -> list[Country]: + """Return every countries row, ordered by country ID.""" + return list(session.exec(select(Country).order_by(col(Country.id))).all()) + + +def read_country_ids(session: Session) -> list[str]: + """Return the ID of every countries row, ordered.""" + return list(session.exec(select(Country.id).order_by(col(Country.id))).all()) diff --git a/policyengine_api/services/v2/countries/database_session.py b/policyengine_api/services/v2/countries/database_session.py new file mode 100644 index 000000000..455d4aa39 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_session.py @@ -0,0 +1,26 @@ +"""Database session lifetime management for the v2 countries service.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy.orm import sessionmaker + from sqlmodel import Session + + +class CountryDatabaseSession: + """Open the read sessions the countries service uses.""" + + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" + self._session_factory = session_factory + + @contextmanager + def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" + with self._session_factory() as session: + yield session diff --git a/policyengine_api/services/v2/countries/services.py b/policyengine_api/services/v2/countries/services.py new file mode 100644 index 000000000..96e0ad601 --- /dev/null +++ b/policyengine_api/services/v2/countries/services.py @@ -0,0 +1,49 @@ +"""Router-facing and startup service for the v2 countries table.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.countries.database_connectors.reads import ( + read_countries, + read_country_ids, +) +from policyengine_api.services.v2.countries.transformations import country_read +from policyengine_api.services.v2.countries.validators import ( + validate_countries_match_code_list, +) + +if TYPE_CHECKING: + from collections.abc import Iterable + + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.types import CountryRead + + +class V2CountryService: + """List the served countries, and check them against the list in code.""" + + def __init__(self, database_session: CountryDatabaseSession) -> None: + """Read through ``database_session``, which opens one session per call.""" + self._database_session = database_session + + def list_countries(self) -> list[CountryRead]: + """Return every countries row, ordered by country ID.""" + with self._database_session.read() as session: + return [country_read(country) for country in read_countries(session)] + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + """Raise ``CountryListMismatchError`` unless the rows name exactly these. + + The application runs this at startup with the country list in code. + """ + with self._database_session.read() as session: + stored_country_ids = read_country_ids(session) + validate_countries_match_code_list( + stored_country_ids=stored_country_ids, + listed_country_ids=listed_country_ids, + ) diff --git a/policyengine_api/services/v2/countries/transformations.py b/policyengine_api/services/v2/countries/transformations.py new file mode 100644 index 000000000..fd1d65fe2 --- /dev/null +++ b/policyengine_api/services/v2/countries/transformations.py @@ -0,0 +1,19 @@ +"""Conversions between stored countries rows and service read records.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.countries.types import CountryRead + +if TYPE_CHECKING: + from policyengine_api.data.v2.models import Country + + +def country_read(country: Country) -> CountryRead: + """Convert one stored countries row to the service's read record.""" + return CountryRead( + id=country.id, + household_enabled=country.household_enabled, + society_wide_enabled=country.society_wide_enabled, + ) diff --git a/policyengine_api/services/v2/countries/types.py b/policyengine_api/services/v2/countries/types.py new file mode 100644 index 000000000..b4a1e74fc --- /dev/null +++ b/policyengine_api/services/v2/countries/types.py @@ -0,0 +1,19 @@ +"""Framework-independent data exchanged by the v2 countries layers.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CountryRead: + """One served country and the calculations it supports. + + ``household_enabled`` allows household calculations and saved households; + ``society_wide_enabled`` allows society-wide reports, and implies + ``household_enabled``. + """ + + id: str + household_enabled: bool + society_wide_enabled: bool diff --git a/policyengine_api/services/v2/countries/validators.py b/policyengine_api/services/v2/countries/validators.py new file mode 100644 index 000000000..fbb30ca5b --- /dev/null +++ b/policyengine_api/services/v2/countries/validators.py @@ -0,0 +1,48 @@ +"""Checks of the stored countries rows against the country list in code.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + +class CountryListMismatchError(RuntimeError): + """The countries table and the country list in code name different countries. + + ``listed_without_row`` holds the listed countries that have no row, and + ``rows_not_listed`` the rows for countries the list does not name, each + sorted by ID. + """ + + def __init__( + self, + *, + listed_without_row: tuple[str, ...], + rows_not_listed: tuple[str, ...], + ) -> None: + """Record both differences and describe them in the message.""" + self.listed_without_row = listed_without_row + self.rows_not_listed = rows_not_listed + super().__init__( + "the countries table does not match the country list in code: " + f"listed without a row: {', '.join(listed_without_row) or 'none'}; " + f"rows not listed: {', '.join(rows_not_listed) or 'none'}" + ) + + +def validate_countries_match_code_list( + *, + stored_country_ids: Iterable[str], + listed_country_ids: Iterable[str], +) -> None: + """Raise ``CountryListMismatchError`` unless both name the same countries.""" + stored = set(stored_country_ids) + listed = set(listed_country_ids) + if stored == listed: + return + raise CountryListMismatchError( + listed_without_row=tuple(sorted(listed - stored)), + rows_not_listed=tuple(sorted(stored - listed)), + ) diff --git a/policyengine_api/services/v2/economy_reports/__init__.py b/policyengine_api/services/v2/economy_reports/__init__.py new file mode 100644 index 000000000..922209028 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/__init__.py @@ -0,0 +1 @@ +"""Economy reports, calculated by the Simulation Entrypoint and served by API v2.""" diff --git a/policyengine_api/services/v2/economy_reports/database_connectors/__init__.py b/policyengine_api/services/v2/economy_reports/database_connectors/__init__.py new file mode 100644 index 000000000..2155996a5 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/database_connectors/__init__.py @@ -0,0 +1 @@ +"""Database reads used by the v2 economy report service.""" diff --git a/policyengine_api/services/v2/economy_reports/database_connectors/reads.py b/policyengine_api/services/v2/economy_reports/database_connectors/reads.py new file mode 100644 index 000000000..645e61df3 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/database_connectors/reads.py @@ -0,0 +1,15 @@ +"""Database selections used by the v2 economy report service.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.data.v2.models import Country + +if TYPE_CHECKING: + from sqlmodel import Session + + +def read_country(session: Session, country_id: str) -> Country | None: + """Return the country's countries row, or None when it has none.""" + return session.get(Country, country_id) diff --git a/policyengine_api/services/v2/economy_reports/database_session.py b/policyengine_api/services/v2/economy_reports/database_session.py new file mode 100644 index 000000000..c3bea8262 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/database_session.py @@ -0,0 +1,30 @@ +"""Database session lifetime management for the v2 economy report service.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy.orm import sessionmaker + from sqlmodel import Session + + +class EconomyReportDatabaseSession: + """Open the read sessions the economy report service uses. + + Reports themselves live in the Simulation Entrypoint; the Public API reads + only what each country supports. + """ + + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" + self._session_factory = session_factory + + @contextmanager + def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" + with self._session_factory() as session: + yield session diff --git a/policyengine_api/services/v2/economy_reports/services.py b/policyengine_api/services/v2/economy_reports/services.py new file mode 100644 index 000000000..14cb7b8a9 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/services.py @@ -0,0 +1,280 @@ +"""Application services for v2 economy reports. + +The Simulation Entrypoint owns economy reports. This service checks what the +Public API decides itself, then asks the entrypoint through +``EconomyReportSimulation``: + +- A country's ``countries`` row must allow economy (society-wide) reports. A + household-only country is refused before anything is sent to the + entrypoint, and the Public API writes nothing for a report. +- Reports run on the current catalog version: the one published for the + PolicyEngine.py version the Public API runs, which is also the catalog its + v2 metadata routes serve by default. +- Every report is read within a country: a report of another country is not + found. +- A report with no run (one migrated from v1) is calculated the first time + it is read: reading it starts a run with the ``system`` trigger. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.economy_reports.database_connectors.reads import ( + read_country, +) +from policyengine_api.services.v2.economy_reports.transformations import ( + economy_report_submission, + reports_of_country, +) +from policyengine_api.services.v2.economy_reports.validators import ( + economy_reports_supported, + report_not_found, + require_economy_reports_supported, + require_report_in_country, +) +from policyengine_api.services.v2.simulation_entrypoint.services import ( + read_report_country, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + RunTrigger, + SimulationRecordNotFoundError, + SimulationRecordUnreadableError, + SimulationRequestRefusedError, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from uuid import UUID + + from policyengine_api.services.v2.economy_reports.database_session import ( + EconomyReportDatabaseSession, + ) + from policyengine_api.services.v2.economy_reports.types import EconomyReportFields + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportRunRecord, + EconomyReportSimulation, + ) + +logger = logging.getLogger(__name__) + + +class V2EconomyReportService: + """Sequence economy report operations between the countries table and sim.""" + + def __init__( + self, + database_session: EconomyReportDatabaseSession, + simulation: EconomyReportSimulation, + *, + current_policyengine_version: str, + ) -> None: + """Read countries through ``database_session``, reports through ``simulation``. + + ``current_policyengine_version`` selects the catalog version new runs use. + """ + self._database_session = database_session + self._simulation = simulation + self._current_policyengine_version = current_policyengine_version + + def _require_economy_reports(self, country_id: str) -> None: + with self._database_session.read() as session: + country = read_country(session, country_id) + require_economy_reports_supported(country, country_id=country_id) + + def submit_report( + self, *, country_id: str, report_input: EconomyReportFields + ) -> EconomyReportRecord: + """Create a report, its simulations and its first run in the entrypoint. + + Raise ``EconomyReportsUnsupportedError`` for a country without economy + reports, before the entrypoint is asked, and + ``SimulationRequestRefusedError`` when the entrypoint cannot run it. + """ + self._require_economy_reports(country_id) + return self._simulation.submit_report( + economy_report_submission( + report_input, + country_id=country_id, + policyengine_version=self._current_policyengine_version, + ) + ) + + def _report_in_country( + self, + read: Callable[[], EconomyReportRecord], + *, + country_id: str, + reference: str, + ) -> EconomyReportRecord: + try: + report: EconomyReportRecord | None = read() + except SimulationRecordNotFoundError: + report = None + return require_report_in_country( + report, country_id=country_id, reference=reference + ) + + def _start_first_run(self, report: EconomyReportRecord) -> EconomyReportRecord: + """Start the first run of a report that has none, and read it again. + + A migrated report has no run until it is first read. Its country must + still have economy reports, and a run the entrypoint refuses (for + example, a region the current catalog version lacks) leaves it not run. + """ + if report.run is not None: + return report + with self._database_session.read() as session: + country = read_country(session, report.country_id) + if not economy_reports_supported(country): + return report + try: + self._simulation.run_report( + report.id, + policyengine_version=self._current_policyengine_version, + trigger=RunTrigger.SYSTEM, + force=False, + ) + except SimulationRecordNotFoundError as error: + raise report_not_found( + country_id=report.country_id, reference=str(report.id) + ) from error + except SimulationRequestRefusedError as error: + logger.warning( + "v2_economy_report_first_run_refused", + extra={ + "policyengine_attributes": { + "report_id": str(report.id), + "reason": str(error)[:500], + } + }, + ) + return report + return self._report_in_country( + lambda: self._simulation.get_report(report.id), + country_id=report.country_id, + reference=str(report.id), + ) + + def get_report(self, *, country_id: str, report_id: UUID) -> EconomyReportRecord: + """Return a report of this country, starting its first run if it has none. + + Raise ``EconomyReportNotFoundError`` when the entrypoint has no such + report in this country. + """ + return self._start_first_run( + self._report_in_country( + lambda: self._simulation.get_report(report_id), + country_id=country_id, + reference=str(report_id), + ) + ) + + def get_report_by_legacy_number( + self, *, country_id: str, legacy_report_id: int + ) -> EconomyReportRecord: + """Return the report migrated from a v1 report output number. + + Like ``get_report``, its first run starts if it has none. Raise + ``EconomyReportNotFoundError`` when no report of this country was + migrated from that number. + """ + return self._start_first_run( + self._report_in_country( + lambda: self._simulation.get_report_by_legacy_id(legacy_report_id), + country_id=country_id, + reference=f"migrated from v1 report {legacy_report_id}", + ) + ) + + def _runs_in_country( + self, *, country_id: str, report_id: UUID + ) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, once it is known to be the country's. + + The country is read from the report's simulations, so a report whose + stored result cannot be read back can still run again or list its + runs. Raise ``EconomyReportNotFoundError`` when the country has no + such report. + """ + reference = str(report_id) + try: + runs = self._simulation.list_report_runs(report_id) + found_country_id = read_report_country(self._simulation, report_id, runs) + except SimulationRecordNotFoundError as error: + raise report_not_found( + country_id=country_id, reference=reference + ) from error + if found_country_id != country_id: + raise report_not_found(country_id=country_id, reference=reference) + return runs + + def run_report( + self, *, country_id: str, report_id: UUID, force: bool + ) -> EconomyReportRunRecord: + """Run a report of the country again on the current catalog version. + + The run has the ``manual`` trigger and reuses existing simulation runs + unless ``force`` is set. Raise ``EconomyReportsUnsupportedError`` for a + country without economy reports, before the entrypoint is asked, + ``EconomyReportNotFoundError`` when the country has no such report, and + ``SimulationRequestRefusedError`` when the entrypoint cannot run it. + """ + self._require_economy_reports(country_id) + self._runs_in_country(country_id=country_id, report_id=report_id) + try: + return self._simulation.run_report( + report_id, + policyengine_version=self._current_policyengine_version, + trigger=RunTrigger.MANUAL, + force=force, + ) + except SimulationRecordNotFoundError as error: + raise report_not_found( + country_id=country_id, reference=str(report_id) + ) from error + + def list_report_runs( + self, *, country_id: str, report_id: UUID + ) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, newest first, each with its catalog version. + + Raise ``EconomyReportNotFoundError`` when the country has no such report. + """ + return self._runs_in_country(country_id=country_id, report_id=report_id) + + def _report_or_none(self, report_id: UUID) -> EconomyReportRecord | None: + """Read a listed report; leave out an unknown one or an unreadable one. + + A report whose stored result cannot be read back is logged and left out + of listings, as the caller can tell from the IDs it asked for. + """ + try: + return self._simulation.get_report(report_id) + except SimulationRecordNotFoundError: + return None + except SimulationRecordUnreadableError: + logger.warning( + "v2_economy_report_unreadable_in_listing", + extra={"policyengine_attributes": {"report_id": str(report_id)}}, + ) + return None + + def list_reports( + self, *, country_id: str, report_ids: Sequence[UUID] + ) -> tuple[EconomyReportRecord, ...]: + """Return the country's reports among ``report_ids``, in that order, once each. + + Unknown reports, other countries' and reports whose stored result + cannot be read back are left out. Listing a report with no run does not + start one; reading it does. + """ + return reports_of_country( + ( + self._report_or_none(report_id) + for report_id in dict.fromkeys(report_ids) + ), + country_id=country_id, + ) diff --git a/policyengine_api/services/v2/economy_reports/transformations.py b/policyengine_api/services/v2/economy_reports/transformations.py new file mode 100644 index 000000000..901a4fae3 --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/transformations.py @@ -0,0 +1,46 @@ +"""Pure representation transformations for v2 economy report operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportSubmission, +) + +if TYPE_CHECKING: + from collections.abc import Iterable + + from policyengine_api.services.v2.economy_reports.types import EconomyReportFields + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + ) + + +def economy_report_submission( + report_input: EconomyReportFields, + *, + country_id: str, + policyengine_version: str, +) -> EconomyReportSubmission: + """Build the entrypoint request for a new report on one catalog version.""" + return EconomyReportSubmission( + country_id=country_id, + policyengine_version=policyengine_version, + year=report_input.year, + region_code=report_input.region_code, + dataset_name=report_input.dataset_name, + baseline=report_input.baseline, + reform=report_input.reform, + ) + + +def reports_of_country( + reports: Iterable[EconomyReportRecord | None], *, country_id: str +) -> tuple[EconomyReportRecord, ...]: + """Keep the found reports of the country, in order, each once.""" + kept: dict[object, EconomyReportRecord] = {} + for report in reports: + if report is not None and report.country_id == country_id: + kept.setdefault(report.id, report) + return tuple(kept.values()) diff --git a/policyengine_api/services/v2/economy_reports/types.py b/policyengine_api/services/v2/economy_reports/types.py new file mode 100644 index 000000000..5ac6108fc --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/types.py @@ -0,0 +1,47 @@ +"""Framework-independent request data of the v2 economy report routes. + +What the Simulation Entrypoint answers, and the protocol through which the +Public API asks it, are in ``policyengine_api.services.v2.simulation_entrypoint``. +""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, StrictBool + +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationChoice, +) + +# The bounds and form the Simulation Entrypoint accepts. +ReportYear = Annotated[int, Field(ge=1900, le=2200, description="The report's year")] +CatalogKey = Annotated[str, Field(min_length=1, max_length=255, pattern=r"^\S+$")] + + +class StrictEconomyReportInput(BaseModel): + """Reject fields outside the reviewed economy report contract.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class EconomyReportFields(StrictEconomyReportInput): + """What an economy report compares, apart from its country. + + ``region_code`` is a region of the country's current catalog version + (``state/ca``, ``country/wales``), and ``dataset_name`` one of its + datasets; without it the region's default dataset is used. + """ + + kind: Literal["economy"] + year: ReportYear + region_code: CatalogKey + dataset_name: CatalogKey | None = None + baseline: SimulationChoice + reform: SimulationChoice + + +class EconomyReportRunFields(StrictEconomyReportInput): + """A re-run; ``force`` calculates both simulations again instead of reusing runs.""" + + force: StrictBool = False diff --git a/policyengine_api/services/v2/economy_reports/validators.py b/policyengine_api/services/v2/economy_reports/validators.py new file mode 100644 index 000000000..4c551c6af --- /dev/null +++ b/policyengine_api/services/v2/economy_reports/validators.py @@ -0,0 +1,57 @@ +"""Database-independent validation for v2 economy report operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from policyengine_api.data.v2.models import Country + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + ) + + +class EconomyReportsUnsupportedError(ValueError): + """The country does not have economy (society-wide) reports. + + Its countries row leaves ``society_wide_enabled`` unset, as for a + household-only country, or it has no row. + """ + + +class EconomyReportNotFoundError(LookupError): + """The simulation service has no such report in the requested country.""" + + +def economy_reports_supported(country: Country | None) -> bool: + """Return whether a countries row allows economy (society-wide) reports.""" + return country is not None and country.society_wide_enabled + + +def require_economy_reports_supported( + country: Country | None, *, country_id: str +) -> None: + """Raise ``EconomyReportsUnsupportedError`` unless the row allows them.""" + if not economy_reports_supported(country): + msg = f"Economy reports are not available for country {country_id!r}" + raise EconomyReportsUnsupportedError(msg) + + +def require_report_in_country( + report: EconomyReportRecord | None, *, country_id: str, reference: str +) -> EconomyReportRecord: + """Return the report, or raise ``EconomyReportNotFoundError``. + + A report of another country is not found, as it would be in a country's + own list. ``reference`` names the requested report in the message. + """ + if report is None or report.country_id != country_id: + raise report_not_found(country_id=country_id, reference=reference) + return report + + +def report_not_found(*, country_id: str, reference: str) -> EconomyReportNotFoundError: + """Build the error for a report the country does not have.""" + return EconomyReportNotFoundError( + f"Economy report {reference} was not found in country {country_id!r}" + ) diff --git a/policyengine_api/services/v2/simulation_entrypoint/__init__.py b/policyengine_api/services/v2/simulation_entrypoint/__init__.py new file mode 100644 index 000000000..c1b92310d --- /dev/null +++ b/policyengine_api/services/v2/simulation_entrypoint/__init__.py @@ -0,0 +1 @@ +"""The Simulation Entrypoint's routes as the Public API uses them.""" diff --git a/policyengine_api/services/v2/simulation_entrypoint/services.py b/policyengine_api/services/v2/simulation_entrypoint/services.py new file mode 100644 index 000000000..84061d6fd --- /dev/null +++ b/policyengine_api/services/v2/simulation_entrypoint/services.py @@ -0,0 +1,96 @@ +"""Ways of using the Simulation Entrypoint shared by the Public API's v2 services. + +``DeferredEconomyReportSimulation`` builds the adapter only when it is first +needed, and ``read_report_country`` finds a report's country without making +the entrypoint rebuild the report's stored result. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + from uuid import UUID + + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportRunRecord, + EconomyReportSimulation, + EconomyReportSubmission, + RunTrigger, + SimulationRecord, + ) + + +class DeferredEconomyReportSimulation: + """``EconomyReportSimulation`` that builds its adapter at its first call. + + The Public API decides some refusals itself, such as a household-only + country's economy report, and those must not depend on the simulation + client or its settings. A build that raises + ``SimulationServiceUnavailableError`` is tried again at the next call. + """ + + def __init__(self, build: Callable[[], EconomyReportSimulation]) -> None: + """Call ``build`` at the first call and reuse its adapter afterwards.""" + self._build = build + self._simulation: EconomyReportSimulation | None = None + + def _resolved(self) -> EconomyReportSimulation: + if self._simulation is None: + self._simulation = self._build() + return self._simulation + + def submit_report(self, submission: EconomyReportSubmission) -> EconomyReportRecord: + """Create a report and its first run in the entrypoint, and start it.""" + return self._resolved().submit_report(submission) + + def get_report(self, report_id: UUID) -> EconomyReportRecord: + """Return a report.""" + return self._resolved().get_report(report_id) + + def get_report_by_legacy_id(self, legacy_report_id: int) -> EconomyReportRecord: + """Return the report migrated from a v1 report output number.""" + return self._resolved().get_report_by_legacy_id(legacy_report_id) + + def run_report( + self, + report_id: UUID, + *, + policyengine_version: str, + trigger: RunTrigger, + force: bool, + ) -> EconomyReportRunRecord: + """Run a report again on the catalog version of ``policyengine_version``.""" + return self._resolved().run_report( + report_id, + policyengine_version=policyengine_version, + trigger=trigger, + force=force, + ) + + def list_report_runs(self, report_id: UUID) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, newest first.""" + return self._resolved().list_report_runs(report_id) + + def get_simulation(self, simulation_id: UUID) -> SimulationRecord: + """Return a simulation.""" + return self._resolved().get_simulation(simulation_id) + + +def read_report_country( + simulation: EconomyReportSimulation, + report_id: UUID, + runs: tuple[EconomyReportRunRecord, ...], +) -> str: + """Return a report's country, given its runs, without reading its result. + + A run names the report's simulations, whose country is the report's. A + report with no run has no stored result to rebuild, so reading it is + cheap, and reading a report whose stored result cannot be rebuilt is + avoided. Raise ``SimulationRecordNotFoundError`` for an unknown report. + """ + if runs: + return simulation.get_simulation(runs[0].baseline_run.simulation_id).country_id + return simulation.get_report(report_id).country_id diff --git a/policyengine_api/services/v2/simulation_entrypoint/types.py b/policyengine_api/services/v2/simulation_entrypoint/types.py new file mode 100644 index 000000000..58f1c6d6d --- /dev/null +++ b/policyengine_api/services/v2/simulation_entrypoint/types.py @@ -0,0 +1,250 @@ +"""The Simulation Entrypoint's economy report routes, as the Public API uses them. + +The Simulation Entrypoint owns economy reports and their simulations: their +definitions, runs and results live in its database, and only its routes +create or read them. ``EconomyReportSimulation`` is those routes, with the +records they answer and the four errors they raise. +``policyengine_api.libs.simulation_economy_reports`` implements it with the +generated simulation client, and tests replace it with a fake. The economy +report routes and the checks of saved v2 report and simulation IDs both use it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from pydantic import BaseModel, ConfigDict, JsonValue + +from policyengine_api.query_parameters import ResourceId + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + +JSONObject = dict[str, JsonValue] + + +class ReportStatus(StrEnum): + """Where a report or report run stands; a report's is its newest run's. + + A report with no run, such as one migrated from v1, is ``not_run``. + """ + + NOT_RUN = "not_run" + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class RunTrigger(StrEnum): + """Why a report runs again: a user's re-run, or one the Public API starts.""" + + MANUAL = "manual" + SYSTEM = "system" + + +class StrictSimulationInput(BaseModel): + """Reject fields outside the reviewed simulation contract.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class SimulationChoice(StrictSimulationInput): + """A simulation's v2 policy and dynamic; neither means current law.""" + + policy_id: ResourceId | None = None + dynamic_id: ResourceId | None = None + + +@dataclass(frozen=True) +class EconomyReportSubmission: + """A new economy report as the Simulation Entrypoint is asked to create it. + + ``policyengine_version`` selects the country's catalog version. + """ + + country_id: str + policyengine_version: str + year: int + region_code: str + dataset_name: str | None + baseline: SimulationChoice + reform: SimulationChoice + + +@dataclass(frozen=True) +class RunError: + """Why a run failed, without user data.""" + + code: str + message: str | None + + +@dataclass(frozen=True) +class SimulationRecord: + """A simulation: one policy and dynamic applied to one region.""" + + id: UUID + kind: str + country_id: str + policy_id: UUID | None + dynamic_id: UUID | None + region_code: str | None + dataset_name: str | None + legacy_simulation_id: int | None + created_at: datetime + + +@dataclass(frozen=True) +class SimulationRunRecord: + """One execution of a simulation for a catalog version and year.""" + + id: UUID + simulation_id: UUID + status: str + attempt: int + trigger: str + error: RunError | None + dispatched_at: datetime | None + completed_at: datetime | None + created_at: datetime + + +@dataclass(frozen=True) +class ComparisonRecord: + """A report run's comparison step, which computes the report's result.""" + + status: str + error: RunError | None + dispatched_at: datetime | None + completed_at: datetime | None + + +@dataclass(frozen=True) +class EconomyReportRunRecord: + """One calculation of a report: two simulation runs and their comparison. + + ``tax_benefit_model_version_id`` is the catalog version it ran on, and + ``policyengine_version`` that version's PolicyEngine.py version. ``stage`` + says what a running run is doing. + """ + + id: UUID + report_id: UUID + trigger: str + tax_benefit_model_version_id: UUID + policyengine_version: str | None + year: int + status: ReportStatus + stage: str | None + error: RunError | None + baseline_run: SimulationRunRecord + reform_run: SimulationRunRecord + comparison: ComparisonRecord + created_at: datetime + + +@dataclass(frozen=True) +class EconomyReportRecord: + """An economy report, its newest run and, once that run succeeded, its result. + + ``result`` is today's economy result format (``SingleYearMacroOutput``). + """ + + id: UUID + country_id: str + year: int + options: JSONObject + legacy_report_id: int | None + created_at: datetime + baseline_simulation: SimulationRecord + reform_simulation: SimulationRecord + status: ReportStatus + run: EconomyReportRunRecord | None + result: JSONObject | None + + +class SimulationRecordNotFoundError(LookupError): + """The Simulation Entrypoint has no such report or simulation.""" + + +class SimulationRecordUnreadableError(RuntimeError): + """The Simulation Entrypoint has the report but cannot read back its result. + + It answers this when a succeeded run's stored result cannot be rebuilt. + """ + + +class SimulationRequestRefusedError(ValueError): + """The Simulation Entrypoint cannot run the request; the message says why. + + For example: no catalog for the PolicyEngine.py version, no such region or + dataset, or a policy or dynamic that does not exist for the country. + """ + + +class SimulationServiceUnavailableError(RuntimeError): + """The Simulation Entrypoint cannot be used now. + + It is not configured or reachable here, it refused the Public API's + credentials, it answered that its own database or Modal environment is + unavailable, or it answered in a form the Public API does not understand. + """ + + +class EconomyReportSimulation(Protocol): + """The Simulation Entrypoint's economy report routes. + + Every method raises ``SimulationServiceUnavailableError`` when the + entrypoint cannot be used, and the methods that read a report raise + ``SimulationRecordUnreadableError`` when its stored result cannot be read + back. + """ + + def submit_report(self, submission: EconomyReportSubmission) -> EconomyReportRecord: + """Create the report's simulations, the report and its first run, and start it. + + Raise ``SimulationRequestRefusedError`` when it cannot run. + """ + ... + + def get_report(self, report_id: UUID) -> EconomyReportRecord: + """Return a report; raise ``SimulationRecordNotFoundError`` if unknown.""" + ... + + def get_report_by_legacy_id(self, legacy_report_id: int) -> EconomyReportRecord: + """Return the report migrated from a v1 report output number. + + Raise ``SimulationRecordNotFoundError`` when none was. + """ + ... + + def run_report( + self, + report_id: UUID, + *, + policyengine_version: str, + trigger: RunTrigger, + force: bool, + ) -> EconomyReportRunRecord: + """Run a report again on the catalog version of ``policyengine_version``. + + Existing simulation runs are reused unless ``force`` is set. Raise + ``SimulationRecordNotFoundError`` for an unknown report and + ``SimulationRequestRefusedError`` when it cannot run. + """ + ... + + def list_report_runs(self, report_id: UUID) -> tuple[EconomyReportRunRecord, ...]: + """Return a report's runs, newest first. + + Raise ``SimulationRecordNotFoundError`` for an unknown report. + """ + ... + + def get_simulation(self, simulation_id: UUID) -> SimulationRecord: + """Return a simulation; raise ``SimulationRecordNotFoundError`` if unknown.""" + ... diff --git a/policyengine_api/services/v2/simulation_refs/__init__.py b/policyengine_api/services/v2/simulation_refs/__init__.py new file mode 100644 index 000000000..b778e711b --- /dev/null +++ b/policyengine_api/services/v2/simulation_refs/__init__.py @@ -0,0 +1 @@ +"""Checks that v2 reports and simulations exist before a user saves them.""" diff --git a/policyengine_api/services/v2/simulation_refs/services.py b/policyengine_api/services/v2/simulation_refs/services.py new file mode 100644 index 000000000..44e0c24ca --- /dev/null +++ b/policyengine_api/services/v2/simulation_refs/services.py @@ -0,0 +1,124 @@ +"""Check v2 report and simulation references with the simulation service. + +Reports and simulations live in the Simulation Entrypoint's database, so the +v2 tables that save them keep no foreign key to them. Before a v2 ID is +saved, the entrypoint is asked whether that record exists in the country. +The check runs before the caller opens its v2 transaction. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.simulation_entrypoint.services import ( + read_report_country, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationRecordNotFoundError, +) +from policyengine_api.services.v2.simulation_refs.validators import ( + require_migrated_from, + require_reference_in_country, +) + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportRecord, + EconomyReportSimulation, + SimulationRecord, + ) + + +class SimulationReferenceService: + """Existence checks of v2 reports and simulations, one entrypoint read each.""" + + def __init__(self, simulation: EconomyReportSimulation) -> None: + """Read reports and simulations through ``simulation``.""" + self._simulation = simulation + + def require_report( + self, + *, + country_id: str, + report_id: UUID, + legacy_report_id: int | None = None, + ) -> None: + """Raise ``SimulationReferenceNotFoundError`` unless the country has it. + + With ``legacy_report_id``, also raise + ``AssociationReferenceConflictError`` unless the report was migrated + from that v1 report output; that check reads the whole report. Without + it, the country is read from the report's simulations, so a report + whose stored result cannot be read back can still be saved. + """ + if legacy_report_id is None: + self._require_report_country(country_id=country_id, report_id=report_id) + return + try: + report: EconomyReportRecord | None = self._simulation.get_report(report_id) + except SimulationRecordNotFoundError: + report = None + require_reference_in_country( + None if report is None else report.country_id, + kind="report", + reference_id=report_id, + country_id=country_id, + ) + if report is not None: + require_migrated_from( + report.legacy_report_id, + kind="report", + record_id=report_id, + legacy_id=legacy_report_id, + ) + + def _require_report_country(self, *, country_id: str, report_id: UUID) -> None: + try: + found_country_id: str | None = read_report_country( + self._simulation, + report_id, + self._simulation.list_report_runs(report_id), + ) + except SimulationRecordNotFoundError: + found_country_id = None + require_reference_in_country( + found_country_id, + kind="report", + reference_id=report_id, + country_id=country_id, + ) + + def require_simulation( + self, + *, + country_id: str, + simulation_id: UUID, + legacy_simulation_id: int | None = None, + ) -> None: + """Raise ``SimulationReferenceNotFoundError`` unless the country has it. + + With ``legacy_simulation_id``, also raise + ``AssociationReferenceConflictError`` unless the simulation was + migrated from that v1 simulation. + """ + try: + simulation: SimulationRecord | None = self._simulation.get_simulation( + simulation_id + ) + except SimulationRecordNotFoundError: + simulation = None + require_reference_in_country( + None if simulation is None else simulation.country_id, + kind="simulation", + reference_id=simulation_id, + country_id=country_id, + ) + if simulation is not None and legacy_simulation_id is not None: + require_migrated_from( + simulation.legacy_simulation_id, + kind="simulation", + record_id=simulation_id, + legacy_id=legacy_simulation_id, + ) diff --git a/policyengine_api/services/v2/simulation_refs/validators.py b/policyengine_api/services/v2/simulation_refs/validators.py new file mode 100644 index 000000000..ba3a745b4 --- /dev/null +++ b/policyengine_api/services/v2/simulation_refs/validators.py @@ -0,0 +1,52 @@ +"""Database-independent validation of v2 report and simulation references.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.associations.validators import ( + AssociationReferenceConflictError, +) + +if TYPE_CHECKING: + from uuid import UUID + + +class SimulationReferenceNotFoundError(LookupError): + """The simulation service has no such report or simulation in the country.""" + + +def require_reference_in_country( + found_country_id: str | None, + *, + kind: str, + reference_id: UUID, + country_id: str, +) -> None: + """Raise ``SimulationReferenceNotFoundError`` unless the record is the country's. + + ``found_country_id`` is the found record's country, or None when the + simulation service has no such record. + """ + if found_country_id != country_id: + msg = f"v2 {kind} {reference_id} was not found for country {country_id}" + raise SimulationReferenceNotFoundError(msg) + + +def require_migrated_from( + found_legacy_id: int | None, + *, + kind: str, + record_id: UUID, + legacy_id: int, +) -> None: + """Raise unless the v2 report or simulation was migrated from that v1 one. + + ``found_legacy_id`` is the v1 record the simulation service records the + v2 record was migrated from, or None for one created on v2. Raise + ``AssociationReferenceConflictError`` when they differ, because a saved + record holding both references records that migration. + """ + if found_legacy_id != legacy_id: + msg = f"v2 {kind} {record_id} was not migrated from v1 {kind} {legacy_id}" + raise AssociationReferenceConflictError(msg) diff --git a/policyengine_api/services/v2/user_reports/database_connectors/reads.py b/policyengine_api/services/v2/user_reports/database_connectors/reads.py index e39a1afe4..fd4828082 100644 --- a/policyengine_api/services/v2/user_reports/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_reports/database_connectors/reads.py @@ -20,15 +20,21 @@ def read_owned_user_report_row( user_id: UUID, country_id: str, association_id: UUID, + lock: bool = False, ) -> UserReportAssociation | None: - """Return the user's association with this id in this country, or None.""" - return session.exec( - select(UserReportAssociation).where( - UserReportAssociation.id == association_id, - UserReportAssociation.country_id == country_id, - UserReportAssociation.user_id == user_id, - ) - ).one_or_none() + """Return the user's association with this id in this country, or None. + + With ``lock``, the row is selected FOR UPDATE, so updates of one saved + report take turns and each sees the row as the one before it left it. + """ + statement = select(UserReportAssociation).where( + UserReportAssociation.id == association_id, + UserReportAssociation.country_id == country_id, + UserReportAssociation.user_id == user_id, + ) + if lock: + statement = statement.with_for_update() + return session.exec(statement).one_or_none() def read_user_report_row( diff --git a/policyengine_api/services/v2/user_reports/services.py b/policyengine_api/services/v2/user_reports/services.py index 8b1961eda..c8051d1be 100644 --- a/policyengine_api/services/v2/user_reports/services.py +++ b/policyengine_api/services/v2/user_reports/services.py @@ -1,9 +1,10 @@ """Application services for native v2 user-report associations. -Stage 12A anchors saved reports on the requested v1 report output ID. The -owner always comes from the verified session principal (a signed-in user or a -signed-out browser), and v1 references are checked in a short read-only v1 -transaction before the v2 transaction opens. Reads never provision a +A saved report names a v2 report, a v1 report output, or both. The owner +always comes from the verified session principal (a signed-in user or a +signed-out browser). Before the v2 transaction opens, a v1 reference is +checked in a short read-only v1 transaction, and a v2 report with the +simulation service, which owns v2 reports. Reads never provision a browser's unauthenticated user; one that does not exist yet owns no records. An optional ``user_policy_id`` (the saved reform the report was run from) and an optional ``baseline_user_policy_id`` (the saved reform it was compared with) @@ -20,6 +21,14 @@ from policyengine_api.services.v2.associations.services import ( require_legacy_reference, ) +from policyengine_api.services.v2.associations.transformations import ( + repointed_references, +) +from policyengine_api.services.v2.associations.types import RecordReferences +from policyengine_api.services.v2.associations.validators import ( + require_simulation_reference_checker, + require_unchanged_references, +) from policyengine_api.services.v2.identity.services import ( resolve_principal_user_for_read_in_session, resolve_principal_user_in_session, @@ -45,6 +54,7 @@ from policyengine_api.services.v2.user_reports.transformations import ( owned_user_report_input, shared_user_report_read, + stored_report_references, user_report_page, user_report_read, user_report_update_values, @@ -61,7 +71,10 @@ from sqlmodel import Session from policyengine_api.data.v2.models import UserReportAssociation - from policyengine_api.services.v2.associations.types import LegacyReferenceChecker + from policyengine_api.services.v2.associations.types import ( + LegacyReferenceChecker, + SimulationReferenceChecker, + ) from policyengine_api.services.v2.identity.types import AccountPrincipal from policyengine_api.services.v2.user_reports.database_session import ( UserReportDatabaseSession, @@ -81,8 +94,12 @@ def read_owned_user_report( user_id: UUID | None, country_id: str, association_id: UUID, + lock: bool = False, ) -> UserReportAssociation: - """``user_id`` is None for an unauthenticated user that owns nothing yet.""" + """``user_id`` is None for an unauthenticated user that owns nothing yet. + + With ``lock``, the row stays locked until the session's transaction ends. + """ return require_user_report( None if user_id is None @@ -91,6 +108,7 @@ def read_owned_user_report( user_id=user_id, country_id=country_id, association_id=association_id, + lock=lock, ), association_id=association_id, ) @@ -137,23 +155,65 @@ def __init__( self, database_session: UserReportDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, + simulation_references: SimulationReferenceChecker | None = None, ) -> None: """Open sessions through ``database_session``. - ``legacy_references`` checks v1 report outputs before a write. Without - it, creating a saved report or repointing one raises - LegacyReferenceCheckUnavailableError. + ``legacy_references`` checks v1 report outputs and + ``simulation_references`` v2 reports before a write. Saving or + repointing to a reference whose checker is missing raises + LegacyReferenceCheckUnavailableError or + SimulationReferenceCheckUnavailableError. """ self._database_session = database_session self._legacy_references = legacy_references + self._simulation_references = simulation_references - def _require_legacy_report(self, *, country_id: str, legacy_report_id: int) -> None: - require_legacy_reference( - self._legacy_references, - LegacyRecordKind.REPORT, - country_id=country_id, - legacy_id=legacy_report_id, - ) + def _require_references( + self, + *, + country_id: str, + references: RecordReferences, + check_legacy_report: bool, + ) -> None: + """Check a saved report's references before the v2 transaction opens. + + The v1 report output is checked when ``check_legacy_report`` is set. The + v2 report is checked whenever there is one, and when there are both, + it must be the report migrated from that v1 report output. + """ + if check_legacy_report and references.legacy_id is not None: + require_legacy_reference( + self._legacy_references, + LegacyRecordKind.REPORT, + country_id=country_id, + legacy_id=references.legacy_id, + ) + if references.v2_id is not None: + require_simulation_reference_checker( + self._simulation_references + ).require_report( + country_id=country_id, + report_id=references.v2_id, + legacy_report_id=references.legacy_id, + ) + + def _stored_references( + self, *, principal: AccountPrincipal, country_id: str, association_id: UUID + ) -> RecordReferences: + """Read the references the principal's saved report holds now.""" + with self._database_session.read() as session: + owner = resolve_principal_user_for_read_in_session( + session, principal, country_id=country_id + ) + return stored_report_references( + read_owned_user_report( + session, + user_id=None if owner is None else owner.user_id, + country_id=country_id, + association_id=association_id, + ) + ) def create_user_report( self, @@ -161,16 +221,22 @@ def create_user_report( principal: AccountPrincipal, association_input: UserReportCreateFields, ) -> UserReportRead: - """Save a v1 report for the principal, provisioning its v2 user if needed. + """Save a report for the principal, provisioning its v2 user if needed. - The v1 report is checked first, in a read-only v1 transaction, and - LegacyReferenceNotFoundError is raised if it does not exist in the - country. Each saved-reform link must name the owner's saved policy in - the same country, or AssociationLinkNotFoundError is raised. + Each supplied reference is checked first: a v1 report output in a + read-only v1 transaction (LegacyReferenceNotFoundError), and a v2 + report with the simulation service (SimulationReferenceNotFoundError), + each in the country. Each saved-reform link must name the owner's + saved policy in the same country, or AssociationLinkNotFoundError is + raised. """ - self._require_legacy_report( + self._require_references( country_id=association_input.country_id, - legacy_report_id=association_input.legacy_report_id, + references=RecordReferences( + v2_id=association_input.report_id, + legacy_id=association_input.legacy_report_id, + ), + check_legacy_report=True, ) with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( @@ -294,16 +360,35 @@ def patch_user_report( ) -> UserReportRead: """Update the supplied fields of one of the principal's saved reports. - A new ``legacy_report_id`` is checked in v1 first, as on create, and a + A new ``report_id`` or ``legacy_report_id`` is checked first, as on + create, together with the reference the saved report keeps: when it + will hold both, the v2 report must be the one migrated from that v1 + report output (AssociationReferenceConflictError). If the saved report + changes during that check, AssociationConflictError is raised. A supplied saved-reform link must name the owner's saved policy in this country. Raise UserReportNotFoundError if the principal has no such report in this country. """ - if association_input.legacy_report_id is not None: - self._require_legacy_report( + checked: RecordReferences | None = None + if ( + association_input.report_id is not None + or association_input.legacy_report_id is not None + ): + stored = self._stored_references( + principal=principal, + country_id=country_id, + association_id=association_id, + ) + self._require_references( country_id=country_id, - legacy_report_id=association_input.legacy_report_id, + references=repointed_references( + stored, + v2_id=association_input.report_id, + legacy_id=association_input.legacy_report_id, + ), + check_legacy_report=association_input.legacy_report_id is not None, ) + checked = stored with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id @@ -313,7 +398,14 @@ def patch_user_report( user_id=owner.user_id, country_id=country_id, association_id=association_id, + lock=checked is not None, ) + if checked is not None: + require_unchanged_references( + stored_report_references(association), + checked=checked, + association_id=association_id, + ) for link_field in ("user_policy_id", "baseline_user_policy_id"): if link_field in association_input.model_fields_set: validate_linked_user_policy( diff --git a/policyengine_api/services/v2/user_reports/transformations.py b/policyengine_api/services/v2/user_reports/transformations.py index b029e0727..77f3588a3 100644 --- a/policyengine_api/services/v2/user_reports/transformations.py +++ b/policyengine_api/services/v2/user_reports/transformations.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any +from policyengine_api.services.v2.associations.types import RecordReferences from policyengine_api.services.v2.user_reports.types import ( SharedBaselineReformRead, SharedUserReportRead, @@ -58,6 +59,7 @@ def shared_user_report_read( return SharedUserReportRead( id=association.id, country_id=association.country_id, + report_id=association.report_id, legacy_report_id=association.legacy_report_id, name=association.name, source_note=association.source_note, @@ -120,3 +122,10 @@ def user_report_update_values( if "validation" in association_input.model_fields_set: values["validation"] = stored_validation(association_input.validation) return values + + +def stored_report_references(association: UserReportAssociation) -> RecordReferences: + """Return the v2 report and v1 report output a stored saved report names.""" + return RecordReferences( + v2_id=association.report_id, legacy_id=association.legacy_report_id + ) diff --git a/policyengine_api/services/v2/user_reports/types.py b/policyengine_api/services/v2/user_reports/types.py index 5cc0574e9..2c6ba36db 100644 --- a/policyengine_api/services/v2/user_reports/types.py +++ b/policyengine_api/services/v2/user_reports/types.py @@ -187,21 +187,43 @@ def bound_provisions(self) -> _ReportPresentationFields: return self +def _require_report_reference( + report_id: object | None, legacy_report_id: object | None +) -> None: + if report_id is None and legacy_report_id is None: + msg = "At least one of report_id or legacy_report_id is required" + raise ValueError(msg) + + class UserReportCreateFields(_ReportPresentationFields): - """Caller-supplied fields; Stage 12A accepts only v1 report references.""" + """Caller-supplied fields; the owner comes from the session. + + The saved report names a v2 report (``report_id``), a v1 report output + (``legacy_report_id``), or both; each is checked before it is saved, and + with both, the v2 report must be the one migrated from that v1 report + output. + """ country_id: CountryId - legacy_report_id: LegacyRecordId + report_id: ResourceId | None = None + legacy_report_id: LegacyRecordId | None = None user_policy_id: ResourceId | None = None baseline_user_policy_id: ResourceId | None = None + @model_validator(mode="after") + def require_report_reference(self) -> UserReportCreateFields: + """Require a v2 report_id, a v1 legacy_report_id, or both.""" + _require_report_reference(self.report_id, self.legacy_report_id) + return self + class UserReportCreationInput(StrictUserReportModel): """Complete persisted association values, including the resolved owner.""" country_id: CountryId user_id: UserId - legacy_report_id: LegacyRecordId + report_id: ResourceId | None = None + legacy_report_id: LegacyRecordId | None = None name: AssociationName | None = None description: str | None = None user_policy_id: ResourceId | None = None @@ -212,23 +234,32 @@ class UserReportCreationInput(StrictUserReportModel): budget_window_years: int | None = None last_run_at: AwareDatetime | None = None + @model_validator(mode="after") + def require_report_reference(self) -> UserReportCreationInput: + """Require a v2 report_id, a v1 legacy_report_id, or both.""" + _require_report_reference(self.report_id, self.legacy_report_id) + return self + class UserReportUpdateInput(_ReportPresentationFields): """A partial update of a saved report; only supplied fields change. - At least one field must be supplied. ``legacy_report_id`` repoints the saved - report to another v1 report output and cannot be cleared; the saved-reform - links are cleared with null. + At least one field must be supplied. ``report_id`` repoints the saved + report to a v2 report and ``legacy_report_id`` to another v1 report + output; neither can be cleared. The saved-reform links are cleared with + null. """ + report_id: ResourceId | None = None legacy_report_id: LegacyRecordId | None = None user_policy_id: ResourceId | None = None baseline_user_policy_id: ResourceId | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserReportUpdateInput: - """Require a supplied field and reject a null ``legacy_report_id``.""" + """Require a supplied field and reject a null report reference.""" mutable_fields = { + "report_id", "legacy_report_id", "name", "description", @@ -244,9 +275,10 @@ def require_valid_supplied_field(self) -> UserReportUpdateInput: if not supplied: msg = "At least one mutable user-report field must be supplied" raise ValueError(msg) - if "legacy_report_id" in supplied and self.legacy_report_id is None: - msg = "legacy_report_id must not be null" - raise ValueError(msg) + for reference in ("report_id", "legacy_report_id"): + if reference in supplied and getattr(self, reference) is None: + msg = f"{reference} must not be null" + raise ValueError(msg) return self @@ -293,6 +325,7 @@ class SharedUserReportRead: id: UUID country_id: str + report_id: UUID | None legacy_report_id: int | None name: str | None source_note: str | None diff --git a/policyengine_api/services/v2/user_simulations/database_connectors/reads.py b/policyengine_api/services/v2/user_simulations/database_connectors/reads.py index 64609e364..68d400f0e 100644 --- a/policyengine_api/services/v2/user_simulations/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_simulations/database_connectors/reads.py @@ -22,15 +22,21 @@ def read_owned_user_simulation_row( user_id: UUID, country_id: str, association_id: UUID, + lock: bool = False, ) -> UserSimulationAssociation | None: - """Return the user's association with this id in this country, or None.""" - return session.exec( - select(UserSimulationAssociation).where( - UserSimulationAssociation.id == association_id, - UserSimulationAssociation.country_id == country_id, - UserSimulationAssociation.user_id == user_id, - ) - ).one_or_none() + """Return the user's association with this id in this country, or None. + + With ``lock``, the row is selected FOR UPDATE, so updates of one saved + simulation take turns and each sees the row as the one before it left it. + """ + statement = select(UserSimulationAssociation).where( + UserSimulationAssociation.id == association_id, + UserSimulationAssociation.country_id == country_id, + UserSimulationAssociation.user_id == user_id, + ) + if lock: + statement = statement.with_for_update() + return session.exec(statement).one_or_none() def read_user_simulation_rows( diff --git a/policyengine_api/services/v2/user_simulations/services.py b/policyengine_api/services/v2/user_simulations/services.py index 8e2682aa5..f967753f7 100644 --- a/policyengine_api/services/v2/user_simulations/services.py +++ b/policyengine_api/services/v2/user_simulations/services.py @@ -1,9 +1,10 @@ """Application services for native v2 user-simulation associations. -Stage 12A anchors saved simulations on v1 simulation IDs. The owner always -comes from the verified session principal (a signed-in user or a signed-out -browser), and v1 references are checked in a short read-only v1 transaction -before the v2 transaction opens. Reads never provision a browser's +A saved simulation names a v2 simulation, a v1 simulation, or both. The owner +always comes from the verified session principal (a signed-in user or a +signed-out browser). Before the v2 transaction opens, a v1 reference is +checked in a short read-only v1 transaction, and a v2 simulation with the +simulation service, which owns v2 simulations. Reads never provision a browser's unauthenticated user; one that does not exist yet owns no records. """ @@ -14,6 +15,14 @@ from policyengine_api.services.v2.associations.services import ( require_legacy_reference, ) +from policyengine_api.services.v2.associations.transformations import ( + repointed_references, +) +from policyengine_api.services.v2.associations.types import RecordReferences +from policyengine_api.services.v2.associations.validators import ( + require_simulation_reference_checker, + require_unchanged_references, +) from policyengine_api.services.v2.identity.services import ( resolve_principal_user_for_read_in_session, resolve_principal_user_in_session, @@ -35,6 +44,7 @@ ) from policyengine_api.services.v2.user_simulations.transformations import ( owned_user_simulation_input, + stored_simulation_references, user_simulation_page, user_simulation_read, ) @@ -49,7 +59,10 @@ from sqlmodel import Session from policyengine_api.data.v2.models import UserSimulationAssociation - from policyengine_api.services.v2.associations.types import LegacyReferenceChecker + from policyengine_api.services.v2.associations.types import ( + LegacyReferenceChecker, + SimulationReferenceChecker, + ) from policyengine_api.services.v2.identity.types import AccountPrincipal from policyengine_api.services.v2.user_simulations.database_session import ( UserSimulationDatabaseSession, @@ -68,8 +81,12 @@ def read_owned_user_simulation( user_id: UUID | None, country_id: str, association_id: UUID, + lock: bool = False, ) -> UserSimulationAssociation: - """``user_id`` is None for an unauthenticated user that owns nothing yet.""" + """``user_id`` is None for an unauthenticated user that owns nothing yet. + + With ``lock``, the row stays locked until the session's transaction ends. + """ return require_user_simulation( None if user_id is None @@ -78,6 +95,7 @@ def read_owned_user_simulation( user_id=user_id, country_id=country_id, association_id=association_id, + lock=lock, ), association_id=association_id, ) @@ -100,25 +118,65 @@ def __init__( self, database_session: UserSimulationDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, + simulation_references: SimulationReferenceChecker | None = None, ) -> None: """Open sessions through ``database_session``. - ``legacy_references`` checks v1 simulations before a write. Without it, - creating a saved simulation or repointing one raises - LegacyReferenceCheckUnavailableError. + ``legacy_references`` checks v1 simulations and + ``simulation_references`` v2 simulations before a write. Saving or + repointing to a reference whose checker is missing raises + LegacyReferenceCheckUnavailableError or + SimulationReferenceCheckUnavailableError. """ self._database_session = database_session self._legacy_references = legacy_references + self._simulation_references = simulation_references - def _require_legacy_simulation( - self, *, country_id: str, legacy_simulation_id: int + def _require_references( + self, + *, + country_id: str, + references: RecordReferences, + check_legacy_simulation: bool, ) -> None: - require_legacy_reference( - self._legacy_references, - LegacyRecordKind.SIMULATION, - country_id=country_id, - legacy_id=legacy_simulation_id, - ) + """Check a saved simulation's references before the v2 transaction opens. + + The v1 simulation is checked when ``check_legacy_simulation`` is set. + The v2 simulation is checked whenever there is one, and when there are + both, it must be the simulation migrated from that v1 simulation. + """ + if check_legacy_simulation and references.legacy_id is not None: + require_legacy_reference( + self._legacy_references, + LegacyRecordKind.SIMULATION, + country_id=country_id, + legacy_id=references.legacy_id, + ) + if references.v2_id is not None: + require_simulation_reference_checker( + self._simulation_references + ).require_simulation( + country_id=country_id, + simulation_id=references.v2_id, + legacy_simulation_id=references.legacy_id, + ) + + def _stored_references( + self, *, principal: AccountPrincipal, country_id: str, association_id: UUID + ) -> RecordReferences: + """Read the references the principal's saved simulation holds now.""" + with self._database_session.read() as session: + owner = resolve_principal_user_for_read_in_session( + session, principal, country_id=country_id + ) + return stored_simulation_references( + read_owned_user_simulation( + session, + user_id=None if owner is None else owner.user_id, + country_id=country_id, + association_id=association_id, + ) + ) def create_user_simulation( self, @@ -126,15 +184,20 @@ def create_user_simulation( principal: AccountPrincipal, association_input: UserSimulationCreateFields, ) -> UserSimulationRead: - """Save a v1 simulation for the principal, provisioning its v2 user if needed. + """Save a simulation for the principal, provisioning its v2 user if needed. - The v1 simulation is checked first, in a read-only v1 transaction, and - LegacyReferenceNotFoundError is raised if it does not exist in the - country. + Each supplied reference is checked first: a v1 simulation in a + read-only v1 transaction (LegacyReferenceNotFoundError), and a v2 + simulation with the simulation service + (SimulationReferenceNotFoundError), each in the country. """ - self._require_legacy_simulation( + self._require_references( country_id=association_input.country_id, - legacy_simulation_id=association_input.legacy_simulation_id, + references=RecordReferences( + v2_id=association_input.simulation_id, + legacy_id=association_input.legacy_simulation_id, + ), + check_legacy_simulation=True, ) with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( @@ -220,30 +283,55 @@ def patch_user_simulation( ) -> UserSimulationRead: """Update the supplied fields of one of the principal's associations. - A new ``legacy_simulation_id`` is checked in v1 first, as on create. Raise - UserSimulationNotFoundError if the principal has no such association in - this country. + A new ``simulation_id`` or ``legacy_simulation_id`` is checked first, + as on create, together with the reference the saved simulation keeps: + when it will hold both, the v2 simulation must be the one migrated from + that v1 simulation (AssociationReferenceConflictError). If the saved + simulation changes during that check, AssociationConflictError is + raised. Raise UserSimulationNotFoundError if the principal has no such + association in this country. """ - if association_input.legacy_simulation_id is not None: - self._require_legacy_simulation( + checked: RecordReferences | None = None + if ( + association_input.simulation_id is not None + or association_input.legacy_simulation_id is not None + ): + stored = self._stored_references( + principal=principal, + country_id=country_id, + association_id=association_id, + ) + self._require_references( country_id=country_id, - legacy_simulation_id=association_input.legacy_simulation_id, + references=repointed_references( + stored, + v2_id=association_input.simulation_id, + legacy_id=association_input.legacy_simulation_id, + ), + check_legacy_simulation=( + association_input.legacy_simulation_id is not None + ), ) + checked = stored with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id ) - return user_simulation_read( - update_user_simulation( - session, - read_owned_user_simulation( - session, - user_id=owner.user_id, - country_id=country_id, - association_id=association_id, - ), - association_input, + association = read_owned_user_simulation( + session, + user_id=owner.user_id, + country_id=country_id, + association_id=association_id, + lock=checked is not None, + ) + if checked is not None: + require_unchanged_references( + stored_simulation_references(association), + checked=checked, + association_id=association_id, ) + return user_simulation_read( + update_user_simulation(session, association, association_input) ) def delete_user_simulation( diff --git a/policyengine_api/services/v2/user_simulations/transformations.py b/policyengine_api/services/v2/user_simulations/transformations.py index 75216a5b6..7e7483f11 100644 --- a/policyengine_api/services/v2/user_simulations/transformations.py +++ b/policyengine_api/services/v2/user_simulations/transformations.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING +from policyengine_api.services.v2.associations.types import RecordReferences from policyengine_api.services.v2.user_simulations.types import ( UserSimulationCreateFields, UserSimulationCreationInput, @@ -59,3 +60,12 @@ def owned_user_simulation_input( ) -> UserSimulationCreationInput: """Add the owner resolved from the session principal to the caller's fields.""" return UserSimulationCreationInput(**fields.model_dump(), user_id=user_id) + + +def stored_simulation_references( + association: UserSimulationAssociation, +) -> RecordReferences: + """Return the v2 simulation and v1 simulation a stored saved simulation names.""" + return RecordReferences( + v2_id=association.simulation_id, legacy_id=association.legacy_simulation_id + ) diff --git a/policyengine_api/services/v2/user_simulations/types.py b/policyengine_api/services/v2/user_simulations/types.py index d456e0cee..9652b7d4d 100644 --- a/policyengine_api/services/v2/user_simulations/types.py +++ b/policyengine_api/services/v2/user_simulations/types.py @@ -7,7 +7,7 @@ from pydantic import model_validator -from policyengine_api.query_parameters import CountryId, UserId +from policyengine_api.query_parameters import CountryId, ResourceId, UserId from policyengine_api.services.v2.associations.types import ( AssociationDescription, AssociationName, @@ -24,51 +24,82 @@ class StrictUserSimulationInput(StrictAssociationModel): """Reject fields outside the reviewed user-simulation contract.""" +def _require_simulation_reference( + simulation_id: object | None, legacy_simulation_id: object | None +) -> None: + if simulation_id is None and legacy_simulation_id is None: + msg = "At least one of simulation_id or legacy_simulation_id is required" + raise ValueError(msg) + + class UserSimulationCreateFields(StrictUserSimulationInput): - """Caller-supplied fields; Stage 12A accepts only v1 simulation references.""" + """Caller-supplied fields; the owner comes from the session. + + The saved simulation names a v2 simulation (``simulation_id``), a v1 + simulation (``legacy_simulation_id``), or both; each is checked before it + is saved, and with both, the v2 simulation must be the one migrated from + that v1 simulation. + """ country_id: CountryId - legacy_simulation_id: LegacyRecordId + simulation_id: ResourceId | None = None + legacy_simulation_id: LegacyRecordId | None = None name: AssociationName | None = None description: AssociationDescription | None = None + @model_validator(mode="after") + def require_simulation_reference(self) -> UserSimulationCreateFields: + """Require a v2 simulation_id, a v1 legacy_simulation_id, or both.""" + _require_simulation_reference(self.simulation_id, self.legacy_simulation_id) + return self + class UserSimulationCreationInput(StrictUserSimulationInput): """Complete persisted association values, including the resolved owner.""" country_id: CountryId user_id: UserId - legacy_simulation_id: LegacyRecordId + simulation_id: ResourceId | None = None + legacy_simulation_id: LegacyRecordId | None = None name: AssociationName | None = None description: str | None = None + @model_validator(mode="after") + def require_simulation_reference(self) -> UserSimulationCreationInput: + """Require a v2 simulation_id, a v1 legacy_simulation_id, or both.""" + _require_simulation_reference(self.simulation_id, self.legacy_simulation_id) + return self + class UserSimulationUpdateInput(StrictUserSimulationInput): """A partial update of a saved simulation; only supplied fields change. - At least one field must be supplied. ``legacy_simulation_id`` repoints the - saved simulation to another v1 simulation and cannot be cleared. + At least one field must be supplied. ``simulation_id`` repoints the saved + simulation to a v2 simulation and ``legacy_simulation_id`` to another v1 + simulation; neither can be cleared. """ + simulation_id: ResourceId | None = None legacy_simulation_id: LegacyRecordId | None = None name: AssociationName | None = None description: AssociationDescription | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserSimulationUpdateInput: - """Require a supplied field and reject a null ``legacy_simulation_id``.""" + """Require a supplied field and reject a null simulation reference.""" supplied = self.model_fields_set.intersection( - {"legacy_simulation_id", "name", "description"} + {"simulation_id", "legacy_simulation_id", "name", "description"} ) if not supplied: msg = ( - "At least one of legacy_simulation_id, name, or description must " - "be supplied" + "At least one of simulation_id, legacy_simulation_id, name, or " + "description must be supplied" ) raise ValueError(msg) - if "legacy_simulation_id" in supplied and self.legacy_simulation_id is None: - msg = "legacy_simulation_id must not be null" - raise ValueError(msg) + for reference in ("simulation_id", "legacy_simulation_id"): + if reference in supplied and getattr(self, reference) is None: + msg = f"{reference} must not be null" + raise ValueError(msg) return self @@ -76,8 +107,8 @@ def require_valid_supplied_field(self) -> UserSimulationUpdateInput: class UserSimulationRead: """A stored user-simulation association as the service returns it to routes. - Stage 12A saves only v1 simulations, so ``legacy_simulation_id`` names the - simulation and ``simulation_id`` is a logical reference with no foreign key. + ``simulation_id`` names a v2 simulation, a logical reference with no + foreign key, and ``legacy_simulation_id`` a v1 simulation. """ id: UUID diff --git a/policyengine_api/startup_country_check.py b/policyengine_api/startup_country_check.py new file mode 100644 index 000000000..69f233988 --- /dev/null +++ b/policyengine_api/startup_country_check.py @@ -0,0 +1,106 @@ +"""Refuse to start when the v2 countries table and the country list differ. + +Which countries the Public API serves is decided in code +(``policyengine_api.supported_countries``). The v2 ``countries`` table records +what each supports, and must hold one row for each listed country and no +other. ``policyengine_api.asgi`` runs this check while the application loads, +so a disagreement stops the worker before it serves a request. + +The check runs only where a v2 runtime database is selected: on Cloud Run +(``K_SERVICE`` is set), or when a v2 runtime database URL or its secret +resource is configured. A local or test startup without v2 configuration +opens no connection. + +Startup stops when the database answers and its answer disagrees with the +list: different countries, a missing table, or no privilege to read it. +When the check cannot get an answer, it logs why and startup continues: the +v1 routes do not read the v2 database, and their copies into it already +continue without it. That covers a database that cannot be reached or +refuses the credentials, and v2 configuration that cannot be used, including +a secret that cannot be resolved; the readiness check already reports the +last as not ready, and recovers once the secret resolves. +""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING, Protocol + +from sqlalchemy.exc import InterfaceError, OperationalError + +from policyengine_api.data.v2.settings import ( + V2_RUNTIME_DATABASE_URL, + V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE, + V2ConfigurationError, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Mapping + +logger = logging.getLogger(__name__) + +CLOUD_RUN_SERVICE = "K_SERVICE" +V2_RUNTIME_SELECTION_VARIABLES = ( + CLOUD_RUN_SERVICE, + V2_RUNTIME_DATABASE_URL, + V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE, +) + + +class CountryListVerifier(Protocol): + """Compare the stored countries rows with a list of country IDs.""" + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + """Raise ``CountryListMismatchError`` unless the rows name exactly these.""" + ... + + +def _default_country_list_verifier() -> CountryListVerifier: + from policyengine_api.data.v2.database import get_v2_session_factory + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.services import V2CountryService + + return V2CountryService(CountryDatabaseSession(get_v2_session_factory())) + + +def v2_runtime_database_selected(environ: Mapping[str, str]) -> bool: + """Whether this process runs on Cloud Run or has a v2 runtime database.""" + return any(environ.get(name, "").strip() for name in V2_RUNTIME_SELECTION_VARIABLES) + + +def check_countries_table_at_startup( + environ: Mapping[str, str] | None = None, + verifier_factory: Callable[[], CountryListVerifier] | None = None, +) -> None: + """Raise unless the countries table names exactly the countries in code. + + Raises ``CountryListMismatchError`` on a disagreement and SQLAlchemy's + ``ProgrammingError`` when the table or the privilege to read it is + missing. Unusable v2 configuration is logged with its message, which + never contains a secret, and a database connection failure by its error + type only, as its detail can name the host; both skip the check. + """ + values = os.environ if environ is None else environ + if not v2_runtime_database_selected(values): + return + factory = verifier_factory or _default_country_list_verifier + try: + factory().verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + except V2ConfigurationError as error: + logger.warning( + "Skipped the countries table check at startup: the v2 runtime " + "configuration cannot be used (%s)", + error, + ) + except (OperationalError, InterfaceError) as error: + logger.warning( + "Skipped the countries table check at startup: the v2 database " + "could not be queried (%s)", + type(error).__name__, + ) diff --git a/policyengine_api/supported_countries.py b/policyengine_api/supported_countries.py new file mode 100644 index 000000000..0c5ad01a1 --- /dev/null +++ b/policyengine_api/supported_countries.py @@ -0,0 +1,33 @@ +"""The countries the Public API serves, listed once in code. + +``SupportedCountryId`` is the one country list. Every other list derives from +it: the v2 routes' ``CountryId`` type, the countries catalog extraction +publishes, and the countries the v1 routes copy into v2 (dual-write). Adding +a country starts here. + +The v2 ``countries`` table records what each country supports, and must hold +one row for each listed country and no other: a v2 Alembic revision inserts +the row when a country is added. The application refuses to start when they +differ (``policyengine_api.startup_country_check``). Because a deploy runs +that revision before it starts the new release, instances of the previous +release that start afterwards, or a rollback to it, refuse to start until +the revision is downgraded. + +A new country also needs the US-and-UK checks widened on the tables that +pair v1 records with their v2 copies (``legacy_policy_mappings``, +``legacy_user_policy_mappings``, ``legacy_household_mappings``) and on the +v1 ``household_mirror_events`` table, which the v1 household insert writes +in its own transaction. + +The module imports nothing from the package, so modules that the query +models import, such as the catalog package, can import it too. +""" + +from __future__ import annotations + +from typing import Final, Literal, get_args + +SupportedCountryId = Literal["us", "uk"] +SUPPORTED_COUNTRY_IDS: Final[tuple[SupportedCountryId, ...]] = get_args( + SupportedCountryId +) diff --git a/pyproject.toml b/pyproject.toml index ed101271e..66c6a592d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,8 @@ files = [ "policyengine_api/services/v2", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", + "policyengine_api/startup_country_check.py", + "policyengine_api/supported_countries.py", ] explicit_package_bases = true follow_imports = "silent" @@ -203,8 +205,10 @@ showcontent = true # never at module level (TID253 enforces it; TYPE_CHECKING imports are fine): # policyengine and spm_calculator load country models, the v2 database module # registers an exit hook at import, the v2 database sessions and services load -# persistence only when a default factory runs, and Auth0 Management and the -# Google Secret Manager client load only when a secret is read. Ruff's PLC0415 +# persistence only when a default factory runs, Auth0 Management and the +# Google Secret Manager client load only when a secret is read, and the +# economy report adapter needs the generated simulation client, which is not +# yet a dependency, only when an economy report route runs. Ruff's PLC0415 # accepts in-function imports of these without a suppression. A module belongs # here only if no application module imports it at module level. banned-module-level-imports = [ @@ -212,6 +216,11 @@ banned-module-level-imports = [ "policyengine", "policyengine_api.data.v2.database", "policyengine_api.libs.auth0_management", + "policyengine_api.libs.simulation_economy_reports", + "policyengine_api.services.v2.countries.database_session", + "policyengine_api.services.v2.countries.services", + "policyengine_api.services.v2.economy_reports.database_session", + "policyengine_api.services.v2.economy_reports.services", "policyengine_api.services.v2.households.database_session", "policyengine_api.services.v2.households.services", "policyengine_api.services.v2.legacy_refs.database_session", @@ -220,6 +229,7 @@ banned-module-level-imports = [ "policyengine_api.services.v2.metadata.services", "policyengine_api.services.v2.policies.database_session", "policyengine_api.services.v2.policies.services", + "policyengine_api.services.v2.simulation_refs.services", "policyengine_api.services.v2.user_accounts.database_session", "policyengine_api.services.v2.user_accounts.services", "policyengine_api.services.v2.user_households.database_session", @@ -251,6 +261,8 @@ runtime-evaluated-base-classes = [ "policyengine_api.fastapi_routes.v2.metadata.response_models.MetadataResourceSuccessResponse", "policyengine_api.fastapi_routes.v2.policies.response_models.PolicySuccessResponse", "policyengine_api.fastapi_routes.v2.policies.response_models.StrictPolicyAPIModel", + "policyengine_api.fastapi_routes.v2.reports.response_models.ReportSuccessResponse", + "policyengine_api.fastapi_routes.v2.reports.response_models.StrictReportAPIModel", "policyengine_api.fastapi_routes.v2.user_households.response_models.StrictUserHouseholdAPIModel", "policyengine_api.fastapi_routes.v2.user_households.response_models.UserHouseholdSuccessResponse", "policyengine_api.fastapi_routes.v2.user_policies.response_models.StrictUserPolicyAPIModel", @@ -266,6 +278,9 @@ runtime-evaluated-base-classes = [ "policyengine_api.query_parameters.PaginationQuery", "policyengine_api.query_parameters.StrictQueryParameters", "policyengine_api.services.v2.associations.types.StrictAssociationModel", + "policyengine_api.services.v2.economy_reports.types.EconomyReportFields", + "policyengine_api.services.v2.economy_reports.types.EconomyReportRunFields", + "policyengine_api.services.v2.economy_reports.types.StrictEconomyReportInput", "policyengine_api.services.v2.households.types.HouseholdCreationInput", "policyengine_api.services.v2.households.types.StrictHouseholdInput", "policyengine_api.services.v2.metadata.types.MetadataModelVersion", diff --git a/rules/registry.toml b/rules/registry.toml index e517ec50a..46718196b 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -396,7 +396,7 @@ invocation = "make test" [[rule]] id = "api.startup.no-import-side-effects" category = "startup" -statement = "Importing v2 modules, starting up and serving a request open no network connections and create no files or runtime SQLite databases." +statement = "Importing v2 modules opens no network connections, starting up and serving a request open none unless the process runs on Cloud Run or has a v2 runtime database configured, and none of them creates files or runtime SQLite databases." rationale = "docs/engineering/skills/testing.md#focused-commands" scope = ["policyengine_api/**/*.py"] tier = "mechanical" @@ -919,9 +919,9 @@ invocation = "ruff check ." [[rule]] id = "api.quality.mypy-v2" category = "quality" -statement = "Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed." +statement = "Shared query models, FastAPI dependencies, v2 routes and services, the user sign-in and account-deletion libraries, the country list and its startup check pass mypy with untyped definitions disallowed." rationale = "docs/engineering/skills/v2-code-organization.md" -scope = ["policyengine_api/fastapi_routes/v2/**/*.py", "policyengine_api/services/v2/**/*.py", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py"] +scope = ["policyengine_api/fastapi_routes/v2/**/*.py", "policyengine_api/services/v2/**/*.py", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", "policyengine_api/startup_country_check.py", "policyengine_api/supported_countries.py"] tier = "mechanical" status = "enforced" autofix = false @@ -935,6 +935,8 @@ paths = [ "policyengine_api/services/v2", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", + "policyengine_api/startup_country_check.py", + "policyengine_api/supported_countries.py", ] [[rule.ci]] workflow = ".github/workflows/pr.yml" diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index f7af31072..49717019b 100755 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -20,6 +20,7 @@ process_household_event_after_commit, ) from policyengine_api.services.household_service import HouseholdService +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS if TYPE_CHECKING: from collections.abc import Mapping, Sequence @@ -149,7 +150,7 @@ def _parser() -> argparse.ArgumentParser: required=True, choices=("staging", "production"), ) - parser.add_argument("--country-id", required=True, choices=("us", "uk")) + parser.add_argument("--country-id", required=True, choices=SUPPORTED_COUNTRY_IDS) parser.add_argument("--legacy-household-id", required=True, type=int) return parser diff --git a/tests/contract/registry.py b/tests/contract/registry.py index 778c3da26..250c3078a 100644 --- a/tests/contract/registry.py +++ b/tests/contract/registry.py @@ -389,7 +389,12 @@ class WorkflowContract: "name", "description", ), - patch_fields=("legacy_simulation_id", "name", "description"), + patch_fields=( + "simulation_id", + "legacy_simulation_id", + "name", + "description", + ), ), ), WorkflowContract( @@ -414,6 +419,7 @@ class WorkflowContract: "last_run_at", ), patch_fields=( + "report_id", "legacy_report_id", "name", "description", @@ -437,6 +443,7 @@ class WorkflowContract: "message", "result.item.id", "result.item.country_id", + "result.item.report_id", "result.item.legacy_report_id", "result.item.name", "result.item.source_note", @@ -587,6 +594,107 @@ class WorkflowContract: ), ), ), + WorkflowContract( + name="country_selection_v2", + current_contract="typed_v2_resources", + future_owner_pr="Countries table and one country list in code", + requests=( + ContractRequest( + method="GET", + path="/v2/countries", + expected_status=200, + stable_response_fields=( + "status", + "message", + "result.items", + ), + route_group="metadata", + ), + ), + ), + WorkflowContract( + name="economy_reports_v2", + current_contract="typed_v2_resources", + future_owner_pr="Economy reports on v2", + requests=( + ContractRequest( + method="POST", + path="/v2/reports?country_id=us", + expected_status=201, + stable_response_fields=( + "status", + "message", + "result.item.id", + "result.item.country_id", + "result.item.kind", + "result.item.year", + "result.item.baseline_simulation", + "result.item.reform_simulation", + "result.item.status", + "result.item.run", + "result.item.result", + ), + route_group="report", + ), + ContractRequest( + method="GET", + path="/v2/reports?country_id=us&report_id={report_id}", + expected_status=200, + stable_response_fields=("status", "message", "result.items"), + route_group="report", + ), + ContractRequest( + method="GET", + path="/v2/reports/{report_id}?country_id=us", + expected_status=200, + stable_response_fields=( + "status", + "message", + "result.item.id", + "result.item.legacy_report_id", + "result.item.status", + "result.item.run", + "result.item.result", + ), + route_group="report", + ), + ContractRequest( + method="GET", + path="/v2/reports/by-legacy-number/{legacy_report_id}?country_id=us", + expected_status=200, + stable_response_fields=( + "status", + "message", + "result.item.id", + "result.item.legacy_report_id", + "result.item.status", + ), + route_group="report", + ), + ContractRequest( + method="POST", + path="/v2/reports/{report_id}/runs?country_id=us", + expected_status=201, + stable_response_fields=( + "status", + "message", + "result.item.id", + "result.item.trigger", + "result.item.tax_benefit_model_version_id", + "result.item.policyengine_version", + "result.item.status", + ), + route_group="report", + ), + ContractRequest( + method="GET", + path="/v2/reports/{report_id}/runs?country_id=us", + expected_status=200, + stable_response_fields=("status", "message", "result.items"), + route_group="report", + ), + ), + ), WorkflowContract( name="metadata_resources_v2_preview", current_contract="typed_v2_resources", diff --git a/tests/contract/test_app_v2_workflow_contracts.py b/tests/contract/test_app_v2_workflow_contracts.py index ff861db14..abdcb3ba7 100644 --- a/tests/contract/test_app_v2_workflow_contracts.py +++ b/tests/contract/test_app_v2_workflow_contracts.py @@ -20,6 +20,8 @@ def test_app_v2_workflow_contract_registry_is_complete(): "user_account_v2", "household_calculate", "region_selection", + "country_selection_v2", + "economy_reports_v2", "metadata_resources_v2_preview", "simulation_submit_poll", "report_create_poll", @@ -31,6 +33,8 @@ def test_app_v2_workflow_contract_registry_is_complete(): "typed_v2_resources" if workflow.name in { + "country_selection_v2", + "economy_reports_v2", "metadata_resources_v2_preview", "policy_resources_v2", "user_policy_associations_v2", diff --git a/tests/fixtures/simulation_entrypoint/economy_result_us_congressional_district.json b/tests/fixtures/simulation_entrypoint/economy_result_us_congressional_district.json new file mode 100644 index 000000000..af251e848 --- /dev/null +++ b/tests/fixtures/simulation_entrypoint/economy_result_us_congressional_district.json @@ -0,0 +1,301 @@ +{ + "budget": { + "baseline_net_income": 11131057718479.627, + "benefit_spending_impact": 20784007719.238895, + "budgetary_impact": -52407074581.62173, + "households": 132514406.08216108, + "state_tax_revenue_impact": -260089666.9038415, + "tax_revenue_impact": 8845845059.19037 + }, + "cliff_impact": null, + "congressional_district_impact": { + "districts": [ + { + "average_household_income_change": -1950.8870104080981, + "district": "CA-01", + "loser_percentage": 0.04326821221926003, + "no_change_percentage": 0.9502674871120121, + "population": 1070424.2919654003, + "relative_household_income_change": -3.5e-09, + "winner_percentage": 0.14085523548621137 + } + ] + }, + "constituency_impact": null, + "data_version": "1.115.5", + "decile": { + "average": { + "1": -1192.394042905726, + "10": 1805.4526259113698, + "2": -2813.9294911815, + "3": 2193.1634218736735, + "4": -163.505468007199, + "5": 1312.9435443948184, + "6": 2272.87680153289, + "7": 1284.7769016672153, + "8": 2526.592005503247, + "9": -630.2195759955365 + }, + "relative": { + "1": -0.0, + "10": 0.0, + "2": 0.0, + "3": -3.5e-09, + "4": -3.5e-09, + "5": 1e-12, + "6": -3.5e-09, + "7": 0.5740227352895624, + "8": 0.5850741074053635, + "9": 1e-12 + } + }, + "detailed_budget": { + "ctc": { + "baseline": -698767151.9529521, + "difference": 736090614.2865936, + "reform": 269721316.5703769 + }, + "eitc": { + "baseline": 672922902.5487775, + "difference": 278136281.0883239, + "reform": -47293582.601330094 + }, + "income_tax": { + "baseline": -973664016.8902518, + "difference": -481291971.3439847, + "reform": 674938164.19292 + }, + "snap": { + "baseline": -531338077.9066073, + "difference": -59472984.95510411, + "reform": 991289671.0209255 + }, + "ssi": { + "baseline": 46362420.76660269, + "difference": 342822950.7391851, + "reform": 482503712.4029805 + }, + "state_income_tax": { + "baseline": -871937123.5460054, + "difference": 182199165.8626352, + "reform": 516460492.5736347 + } + }, + "inequality": { + "gini": { + "baseline": 0.8564005663967557, + "reform": 0.9909896448688151 + }, + "top_10_pct_share": { + "baseline": 0.6712735421625182, + "reform": 0.16309962197106975 + }, + "top_1_pct_share": { + "baseline": 0.8606375331162682, + "reform": 0.9646329473090614 + } + }, + "intra_decile": { + "all": { + "Gain less than 5%": 0.02692120500933004, + "Gain more than 5%": 0.772273412444121, + "Lose less than 5%": 0.7351731816785095, + "Lose more than 5%": 0.310016695410093, + "No change": 0.5692580020705822 + }, + "deciles": { + "Gain less than 5%": [ + 0.6438887183877269, + 0.5956502384019705, + 0.5592610620153905, + 0.620126135445262, + 0.9406212554239632, + 0.5070268159456598, + 0.43119155343093274, + 0.7203112521441384, + 0.23763561946478406, + 0.3010868611741494 + ], + "Gain more than 5%": [ + 0.6104671229673415, + 0.15619899101356471, + 0.04243582472120422, + 0.8677790339277224, + 0.3138305199160917, + 0.958659426408455, + 0.8966596414276016, + 0.3777892394121827, + 0.46040963284590475, + 0.5200729845925639 + ], + "Lose less than 5%": [ + 0.627341109010956, + 0.4662504296967078, + 0.6792813980381086, + 0.35257698304104734, + 0.706950249365899, + 0.7380342892520343, + 0.02218246899080989, + 0.06057680364579732, + 0.6760203094873768, + 0.9633055803862574 + ], + "Lose more than 5%": [ + 0.2511222781834702, + 0.45631212963637924, + 0.5926718756664892, + 0.32002538574800654, + 0.3639550893399569, + 0.3126706611969108, + 0.36915397527695404, + 0.5956215058575365, + 0.3004039733655893, + 0.37716034132891685 + ], + "No change": [ + 0.9777973164486353, + 0.521127293281206, + 0.5484304676868622, + 0.01145748636421906, + 0.415210343803882, + 0.5799652137970656, + 0.020052890304599336, + 0.6157979413062568, + 0.6321805352961154, + 0.060080510627723016 + ] + } + }, + "intra_wealth_decile": null, + "labor_supply_response": null, + "local_authority_impact": null, + "model_version": "1.715.2", + "poverty": { + "deep_poverty": { + "adult": { + "baseline": 0.08851809310972836, + "reform": 0.8005953212575019 + }, + "all": { + "baseline": 0.2938912468190622, + "reform": 0.7687918872773446 + }, + "child": { + "baseline": 0.8539424884226802, + "reform": 0.9898060149215813 + }, + "senior": { + "baseline": 0.41046182734590886, + "reform": 0.15076537445280958 + } + }, + "poverty": { + "adult": { + "baseline": 0.7138170201741992, + "reform": 0.2111249836755983 + }, + "all": { + "baseline": 0.2849574619862052, + "reform": 0.06346057714522935 + }, + "child": { + "baseline": 0.9046959845122367, + "reform": 0.5691075034743235 + }, + "senior": { + "baseline": 0.8316079302733542, + "reform": 0.5735323523512847 + } + } + }, + "poverty_by_gender": { + "deep_poverty": { + "female": { + "baseline": 0.880905307247358, + "reform": 0.9806357568433396 + }, + "male": { + "baseline": 0.7184404774485162, + "reform": 0.3309541460190075 + } + }, + "poverty": { + "female": { + "baseline": 0.6145325285318086, + "reform": 0.0449402434960362 + }, + "male": { + "baseline": 0.8727670246282013, + "reform": 0.04419006112954338 + } + } + }, + "poverty_by_race": { + "poverty": { + "black": { + "baseline": 0.309670053476339, + "reform": 0.0769707047054119 + }, + "hispanic": { + "baseline": 0.5997628087966007, + "reform": 0.031377762175317736 + }, + "other": { + "baseline": 0.1973848564284194, + "reform": 0.4079361356169908 + }, + "white": { + "baseline": 0.505420373648044, + "reform": 0.9985089453757765 + } + } + }, + "spm_provenance": { + "baseline": { + "execution_count": 2, + "receipt": { + "as_of": null, + "county_vintage": "2020", + "forecast_id": "spm-forecast-2022-2036", + "forecast_sha256": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "geography_id": null, + "geography_kind": "county", + "runtime_versions": { + "policyengine": "6.10.0", + "policyengine-core": "3.20.0", + "policyengine-us": "1.715.2", + "spm-calculator": "0.4.1" + }, + "scenario": "central", + "schema_version": "canonical-spm-provenance-v2", + "years": [ + "2026" + ] + } + }, + "reform": { + "execution_count": 2, + "receipt": { + "as_of": null, + "county_vintage": "2020", + "forecast_id": "spm-forecast-2022-2036", + "forecast_sha256": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "geography_id": null, + "geography_kind": "county", + "runtime_versions": { + "policyengine": "6.10.0", + "policyengine-core": "3.20.0", + "policyengine-us": "1.715.2", + "spm-calculator": "0.4.1" + }, + "scenario": "central", + "schema_version": "canonical-spm-provenance-v2", + "years": [ + "2026" + ] + } + }, + "schema_version": "canonical-spm-comparison-v2" + }, + "wealth_decile": null +} diff --git a/tests/fixtures/v2_countries.py b/tests/fixtures/v2_countries.py new file mode 100644 index 000000000..36d4b6604 --- /dev/null +++ b/tests/fixtures/v2_countries.py @@ -0,0 +1,28 @@ +"""Countries rows for SQLite tests that enforce v2 foreign keys. + +``V2_METADATA.create_all`` creates an empty countries table, while the v2 +Alembic revision that adds the table inserts the US and UK rows. Tests that +turn on SQLite foreign-key enforcement add the same rows before writing +records that name a country. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.data.v2.models import Country + +if TYPE_CHECKING: + from sqlmodel import Session + +# The rows the countries revision inserts, ordered by country ID. +LISTED_COUNTRY_ROWS = ( + {"id": "uk", "household_enabled": True, "society_wide_enabled": True}, + {"id": "us", "household_enabled": True, "society_wide_enabled": True}, +) + + +def add_listed_countries(session: Session) -> None: + """Insert and commit the US and UK countries rows.""" + session.add_all(Country(**row) for row in LISTED_COUNTRY_ROWS) + session.commit() diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 814971236..2df169511 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -3,6 +3,7 @@ import itertools import os from collections.abc import Iterator +from dataclasses import dataclass from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 @@ -30,6 +31,7 @@ STAGE_12_PREVIOUS_REVISION = "724b1b11a33e" STAGE_12_COMPARISON_PREVIOUS_REVISION = "439303be14fe" STAGE_12A_PREVIOUS_REVISION = "91af4fc8d8c1" +COUNTRIES_PREVIOUS_REVISION = "fbb7488ebe8a" V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) @@ -187,7 +189,6 @@ def _assert_stage_12a_association_shape(engine) -> None: for constraint in inspect(engine).get_check_constraints("users") } assert { - "ck_users_primary_country", "ck_users_user_type", "ck_users_standard_primary_country", } <= user_checks @@ -352,6 +353,75 @@ def _assert_base(engine: Engine) -> None: ) +# Every column that names a record's country, with the check it replaced. +COUNTRY_REFERENCE_COLUMNS = { + ("users", "primary_country"): "ck_users_primary_country", + ("households", "country_id"): "ck_households_country", + ("policies", "country_id"): "ck_policies_country", + ("user_household_associations", "country_id"): ( + "ck_user_household_associations_country" + ), + ("user_policies", "country_id"): "ck_user_policies_country", + ("user_simulation_associations", "country_id"): ( + "ck_user_simulation_associations_country" + ), + ("user_report_associations", "country_id"): "ck_user_report_associations_country", +} +# The rows the countries revision inserts: the countries the replaced checks +# allowed, both household and society-wide capable. +LISTED_COUNTRIES = [("uk", True, True), ("us", True, True)] + + +def _country_foreign_keys(engine, table_name: str) -> list[dict]: + return [ + foreign_key + for foreign_key in inspect(engine).get_foreign_keys(table_name) + if foreign_key["referred_table"] == "countries" + ] + + +def _check_names(engine, table_name: str) -> set[str]: + return { + constraint["name"] + for constraint in inspect(engine).get_check_constraints(table_name) + } + + +def _listed_countries(engine) -> list[tuple]: + with engine.connect() as connection: + return [ + tuple(row) + for row in connection.execute( + text( + "SELECT id, household_enabled, society_wide_enabled " + "FROM countries ORDER BY id" + ) + ) + ] + + +def _assert_countries_shape(engine) -> None: + assert _listed_countries(engine) == LISTED_COUNTRIES + assert "ck_countries_society_wide_requires_household" in _check_names( + engine, "countries" + ) + for (table_name, column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): + country_keys = _country_foreign_keys(engine, table_name) + assert len(country_keys) == 1, table_name + assert country_keys[0]["name"] == f"fk_{table_name}_{column_name}_countries" + assert country_keys[0]["constrained_columns"] == [column_name] + assert country_keys[0]["referred_columns"] == ["id"] + assert country_keys[0]["options"]["ondelete"] == "RESTRICT" + assert check_name not in _check_names(engine, table_name) + + +def _assert_pre_countries_shape(engine) -> None: + assert "countries" not in inspect(engine).get_table_names(schema="public") + for (table_name, _column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): + assert _country_foreign_keys(engine, table_name) == [] + assert check_name in _check_names(engine, table_name) + + def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: database_url = _disposable_url() config = _config() @@ -364,6 +434,13 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: } _upgrade_and_assert_head(config, engine) + _assert_countries_shape(engine) + + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + _assert_revision(engine, COUNTRIES_PREVIOUS_REVISION) + _assert_pre_countries_shape(engine) + _upgrade_and_assert_head(config, engine) + _assert_countries_shape(engine) command.downgrade(config, STAGE_12A_PREVIOUS_REVISION) _assert_revision(engine, STAGE_12A_PREVIOUS_REVISION) @@ -396,10 +473,10 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: engine.dispose() -# Dormant at the parent revision: its models map these tables, and deleting a -# user cascades through them, but no parent-revision code inserts into them. -# They are the only tables Stage 12A may give a new required column. -PARENT_REVISION_READ_ONLY_TABLES = frozenset( +# Dormant at the Stage 12A parent revision: its models map these tables, and +# deleting a user cascades through them, but no parent-revision code inserts +# into them. They are the only tables Stage 12A may give a new required column. +STAGE_12A_PARENT_READ_ONLY_TABLES = frozenset( {"user_simulation_associations", "user_report_associations"} ) USER_ASSOCIATION_TABLES = ( @@ -409,8 +486,18 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: "user_report_associations", ) COUNTRY_COLUMN_NAMES = frozenset({"country", "country_id", "primary_country"}) -# Nullable references that a parent-revision check constraint requires. -REFERENCES_REQUIRED_BY_CHECKS = frozenset({("simulations", "household_id")}) +# Nullable references that a parent-revision check constraint requires: a +# standard user's country, and one of each association's two references. +REFERENCES_REQUIRED_BY_CHECKS = frozenset( + { + ("simulations", "household_id"), + ("users", "primary_country"), + ("user_policies", "policy_id"), + ("user_household_associations", "household_id"), + ("user_simulation_associations", "simulation_id"), + ("user_report_associations", "report_id"), + } +) # Required values that a parent-revision check constraint restricts. CHECKED_VALUES = {("datasets", "year"): 2026} @@ -547,14 +634,33 @@ def _rewrite_mapped_columns(connection, table: sa.Table, row: dict) -> None: assert result.rowcount == 1, table.name +@dataclass(frozen=True) +class PreviousRevisionSchemas: + """A previous revision's reflected schema and head's, after the upgrade. + + ``read_only_tables`` are the previous revision's tables that only new code + writes, so their test rows are inserted through head's table. + """ + + parent: sa.MetaData + head: sa.MetaData + read_only_tables: frozenset[str] + + def insert_table(self, table_name: str) -> sa.Table | None: + """Head's table for a read-only table, otherwise ``None``.""" + if table_name in self.read_only_tables: + return self.head.tables[table_name] + return None + + def _assert_user_deletion_cascades( connection, - parent: sa.MetaData, - head: sa.MetaData, + schemas: PreviousRevisionSchemas, stored_rows: dict[str, dict], sequence: Iterator[int], ) -> None: - """Delete, through the parent's users table, a user who owns Stage 12A rows.""" + """Delete, through the parent's users table, a user who owns association rows.""" + parent, head = schemas.parent, schemas.head users = parent.tables["users"] owner = _insert_and_read_back( connection, users, _row_values(users, stored_rows, sequence) @@ -562,11 +668,7 @@ def _assert_user_deletion_cascades( owner_rows = {**stored_rows, "users": owner} for table_name in USER_ASSOCIATION_TABLES: table = parent.tables[table_name] - head_table = ( - head.tables[table_name] - if table_name in PARENT_REVISION_READ_ONLY_TABLES - else None - ) + head_table = schemas.insert_table(table_name) _insert_and_read_back( connection, table, @@ -596,22 +698,26 @@ def _assert_user_deletion_cascades( assert remaining == 0, table_name -def test_stage_12a_schema_serves_the_previous_revision() -> None: - """Run the previous Public API revision's table usage against head. +def _assert_head_serves_previous_revision( + previous_revision: str, + read_only_tables: frozenset[str], +) -> None: + """Run a previous Public API revision's table usage against head. The parent revision's schema is reflected and then used, after upgrading, to insert a row into every table it maps (supplying only what it requires), read the row back with every mapped column, rewrite those - columns, delete it, and delete a user who owns Stage 12A rows so the + columns, delete it, and delete a user who owns association rows so the cascades run. Head may not give any table the parent revision writes a - new required column without a server default. + new required column without a server default; ``read_only_tables`` are the + parent's tables that only new code writes. """ database_url = _disposable_url() config = _config() engine = create_engine(database_url) try: - command.downgrade(config, STAGE_12A_PREVIOUS_REVISION) + command.downgrade(config, previous_revision) parent = sa.MetaData() parent.reflect(engine) command.upgrade(config, "head") @@ -634,21 +740,18 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: for table in parent_tables if _new_required_columns(table, head.tables[table.name]) } - assert set(tables_with_new_required_columns) <= ( - PARENT_REVISION_READ_ONLY_TABLES - ), tables_with_new_required_columns + assert set(tables_with_new_required_columns) <= read_only_tables, ( + tables_with_new_required_columns + ) + schemas = PreviousRevisionSchemas(parent, head, read_only_tables) sequence = itertools.count() with engine.connect() as connection: transaction = connection.begin() try: stored_rows: dict[str, dict] = {} for table in parent_tables: - head_table = ( - head.tables[table.name] - if table.name in PARENT_REVISION_READ_ONLY_TABLES - else None - ) + head_table = schemas.insert_table(table.name) stored_rows[table.name] = _insert_and_read_back( connection, table, @@ -660,7 +763,7 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: _rewrite_mapped_columns(connection, table, stored_rows[table.name]) _assert_user_deletion_cascades( - connection, parent, head, stored_rows, sequence + connection, schemas, stored_rows, sequence ) for table in reversed(parent_tables): @@ -677,6 +780,21 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: engine.dispose() +def test_stage_12a_schema_serves_the_previous_revision() -> None: + _assert_head_serves_previous_revision( + STAGE_12A_PREVIOUS_REVISION, + STAGE_12A_PARENT_READ_ONLY_TABLES, + ) + + +def test_countries_schema_serves_the_previous_revision() -> None: + """The release before the countries table keeps working during a deploy.""" + _assert_head_serves_previous_revision( + COUNTRIES_PREVIOUS_REVISION, + frozenset(), + ) + + def test_upgrade_to_head_validates_the_resulting_schema_against_metadata() -> None: database_url = _disposable_url() config = _config() @@ -957,3 +1075,176 @@ def reflected(name: str) -> sa.Table: connection.execute(sa.delete(table).where(table.c.id == row_id)) command.upgrade(config, "head") engine.dispose() + + +def test_countries_revision_keeps_existing_rows_and_rejects_unlisted_countries() -> ( + None +): + """Rows saved before the revision survive it, and only listed countries pass.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + + def reflected(name: str) -> sa.Table: + return sa.Table(name, sa.MetaData(), autoload_with=engine) + + try: + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + with engine.begin() as connection: + connection.execute( + sa.insert(reflected("users")).values(id=user_id, primary_country="uk") + ) + connection.execute( + sa.insert(reflected("user_report_associations")).values( + id=uuid4(), + user_id=user_id, + country_id="uk", + legacy_report_id=7, + ) + ) + + command.upgrade(config, "head") + command.check(config) + _assert_countries_shape(engine) + + with engine.connect() as connection: + stored_countries = connection.execute( + text( + "SELECT users.primary_country, " + "user_report_associations.country_id FROM users " + "JOIN user_report_associations " + "ON user_report_associations.user_id = users.id " + "WHERE users.id = :user_id" + ), + {"user_id": user_id}, + ).one() + assert tuple(stored_countries) == ("uk", "uk") + + with engine.connect() as connection, pytest.raises(sa.exc.IntegrityError): + connection.execute( + sa.insert(reflected("user_report_associations")).values( + id=uuid4(), + user_id=user_id, + country_id="ca", + legacy_report_id=8, + ) + ) + + with engine.connect() as connection, pytest.raises(sa.exc.IntegrityError): + connection.execute(text("DELETE FROM countries WHERE id = 'uk'")) + finally: + with engine.begin() as connection: + users = reflected("users") + connection.execute(sa.delete(users).where(users.c.id == user_id)) + command.upgrade(config, "head") + engine.dispose() + + +def test_countries_downgrade_refuses_records_of_an_unlisted_country() -> None: + """The restored US-and-UK checks stop a downgrade that would orphan records.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + + try: + command.upgrade(config, "head") + with engine.begin() as connection: + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('be', true, false)" + ) + ) + connection.execute( + text("INSERT INTO users (id, primary_country) VALUES (:user_id, 'be')"), + {"user_id": user_id}, + ) + + with pytest.raises(sa.exc.IntegrityError): + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() == V2_ALEMBIC_HEAD_REVISION + finally: + with engine.begin() as connection: + connection.execute( + text("DELETE FROM users WHERE id = :user_id"), {"user_id": user_id} + ) + connection.execute(text("DELETE FROM countries WHERE id = 'be'")) + command.upgrade(config, "head") + engine.dispose() + + +def test_a_household_only_country_row_lets_records_name_that_country() -> None: + """Once a household-only country has a row, its records can be stored.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + household_id = uuid4() + + try: + command.upgrade(config, "head") + with engine.connect() as connection: + transaction = connection.begin() + try: + with pytest.raises(sa.exc.IntegrityError), connection.begin_nested(): + connection.execute( + text( + "INSERT INTO users (id, primary_country) " + "VALUES (:user_id, 'be')" + ), + {"user_id": user_id}, + ) + + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('be', true, false)" + ) + ) + connection.execute( + text( + "INSERT INTO users (id, primary_country) " + "VALUES (:user_id, 'be')" + ), + {"user_id": user_id}, + ) + connection.execute( + text( + "INSERT INTO households (id, country_id, household_data, " + "canonicalization_version, content_hash) VALUES " + "(:household_id, 'be', '{}'::jsonb, 1, :content_hash)" + ), + {"household_id": household_id, "content_hash": "c" * 64}, + ) + connection.execute( + text( + "INSERT INTO user_household_associations " + "(id, user_id, household_id, country_id) VALUES " + "(:association_id, :user_id, :household_id, 'be')" + ), + { + "association_id": uuid4(), + "user_id": user_id, + "household_id": household_id, + }, + ) + + with pytest.raises(sa.exc.IntegrityError), connection.begin_nested(): + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('xx', false, true)" + ) + ) + finally: + transaction.rollback() + finally: + engine.dispose() diff --git a/tests/integration/test_v2_country_routes.py b/tests/integration/test_v2_country_routes.py new file mode 100644 index 000000000..ee6408ddc --- /dev/null +++ b/tests/integration/test_v2_country_routes.py @@ -0,0 +1,111 @@ +"""Serve and check the v2 countries from a migrated disposable Postgres database. + +The route and the startup check run as the Public API's runtime role, so the +tests prove the role's declared privileges cover them (rule +``api.v2.runtime-privileges-tested-as-runtime-role``). +""" + +from __future__ import annotations + +import pytest +from alembic import command +from alembic.config import Config +from fastapi.testclient import TestClient +from flask import Flask +from sqlalchemy import create_engine, text +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool +from sqlmodel import Session + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.constants import REPO +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) +from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, +) +from policyengine_api.services.v2.countries.services import V2CountryService +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS +from tests.fixtures.v2_countries import LISTED_COUNTRY_ROWS + + +def test_postgres_countries_list_the_rows_the_revision_inserted( + managed_v2_roles, + v2_engine_as_role, +) -> None: + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + sessions = sessionmaker(runtime_engine, class_=Session) + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_country_service_factory=lambda: V2CountryService( + CountryDatabaseSession(sessions) + ), + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + client = TestClient( + create_asgi_app( + Flask(__name__), + dependencies=dependencies, + route_settings=settings, + ) + ) + + try: + response = client.get("/v2/countries") + finally: + runtime_engine.dispose() + + assert response.status_code == 200 + assert response.json() == { + "status": "ok", + "message": None, + "result": {"items": list(LISTED_COUNTRY_ROWS)}, + } + + +def test_postgres_countries_match_the_code_list_as_the_runtime_role( + managed_v2_roles, + v2_engine_as_role, + disposable_v2_database_url: str, +) -> None: + """The startup check passes on the revision's rows and fails on an extra row.""" + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(runtime_engine, class_=Session)) + ) + owner_engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + + try: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + with owner_engine.begin() as connection: + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('ca', true, false)" + ) + ) + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + assert raised.value.rows_not_listed == ("ca",) + finally: + with owner_engine.begin() as connection: + connection.execute(text("DELETE FROM countries WHERE id = 'ca'")) + owner_engine.dispose() + runtime_engine.dispose() diff --git a/tests/integration/test_v2_economy_reports.py b/tests/integration/test_v2_economy_reports.py new file mode 100644 index 000000000..5fe37379f --- /dev/null +++ b/tests/integration/test_v2_economy_reports.py @@ -0,0 +1,154 @@ +"""Economy reports against a migrated disposable Postgres database. + +The service and route read the countries row as the Public API's runtime +role, so the tests prove its declared privileges cover them (rule +``api.v2.runtime-privileges-tested-as-runtime-role``). The simulation service +is a fake: these tests check what the Public API decides before asking it. +""" + +from __future__ import annotations + +import pytest +from alembic import command +from alembic.config import Config +from fastapi.testclient import TestClient +from flask import Flask +from sqlalchemy import create_engine, text +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool +from sqlmodel import Session + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.constants import REPO +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) +from policyengine_api.services.v2.economy_reports.database_session import ( + EconomyReportDatabaseSession, +) +from policyengine_api.services.v2.economy_reports.services import ( + V2EconomyReportService, +) +from policyengine_api.services.v2.economy_reports.types import EconomyReportFields +from policyengine_api.services.v2.economy_reports.validators import ( + EconomyReportsUnsupportedError, +) +from policyengine_api.services.v2.simulation_entrypoint.types import SimulationChoice +from tests.unit.v2.economy_report_support import ( + CURRENT_POLICYENGINE_VERSION, + REPORT_ID, + FakeEconomyReportSimulation, + report_record, + report_run_record, +) + +# A country that has only household calculations; no real country does yet. +HOUSEHOLD_ONLY_COUNTRY_ID = "zz" +REPORT_FIELDS = EconomyReportFields( + kind="economy", + year=2026, + region_code="state/ca", + baseline=SimulationChoice(), + reform=SimulationChoice(), +) + + +@pytest.fixture +def runtime_sessions(managed_v2_roles, v2_engine_as_role, disposable_v2_database_url): + """Migrate, add the test-only household-only row, and yield runtime sessions. + + The row is removed afterwards, so later tests see only the listed + countries. + """ + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + owner_engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + with owner_engine.begin() as connection: + connection.execute( + text( + "INSERT INTO countries (id, household_enabled, society_wide_enabled) " + "VALUES (:country_id, true, false)" + ), + {"country_id": HOUSEHOLD_ONLY_COUNTRY_ID}, + ) + try: + yield sessionmaker(runtime_engine, class_=Session) + finally: + with owner_engine.begin() as connection: + connection.execute( + text("DELETE FROM countries WHERE id = :country_id"), + {"country_id": HOUSEHOLD_ONLY_COUNTRY_ID}, + ) + owner_engine.dispose() + runtime_engine.dispose() + + +def _service(sessions, simulation) -> V2EconomyReportService: + return V2EconomyReportService( + EconomyReportDatabaseSession(sessions), + simulation, + current_policyengine_version=CURRENT_POLICYENGINE_VERSION, + ) + + +def test_postgres_household_only_country_is_refused_before_the_simulation_service( + runtime_sessions, +) -> None: + simulation = FakeEconomyReportSimulation( + report_record(country_id=HOUSEHOLD_ONLY_COUNTRY_ID, run=report_run_record()) + ) + service = _service(runtime_sessions, simulation) + + with pytest.raises(EconomyReportsUnsupportedError) as submitted: + service.submit_report( + country_id=HOUSEHOLD_ONLY_COUNTRY_ID, report_input=REPORT_FIELDS + ) + with pytest.raises(EconomyReportsUnsupportedError) as rerun: + service.run_report( + country_id=HOUSEHOLD_ONLY_COUNTRY_ID, report_id=REPORT_ID, force=False + ) + + for refusal in (submitted, rerun): + assert str(refusal.value) == ( + f"Economy reports are not available for country " + f"{HOUSEHOLD_ONLY_COUNTRY_ID!r}" + ) + assert simulation.calls == [] + + +def test_postgres_society_wide_country_submits_through_the_route( + runtime_sessions, +) -> None: + simulation = FakeEconomyReportSimulation() + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_economy_report_service_factory=lambda: _service( + runtime_sessions, simulation + ), + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + client = TestClient( + create_asgi_app( + Flask(__name__), dependencies=dependencies, route_settings=settings + ) + ) + + response = client.post( + "/v2/reports?country_id=us", + json={ + "country_id": "us", + **REPORT_FIELDS.model_dump(mode="json"), + }, + ) + + assert response.status_code == 201 + assert [name for name, _ in simulation.calls] == ["submit_report"] diff --git a/tests/integration/test_v2_saved_reference_updates.py b/tests/integration/test_v2_saved_reference_updates.py new file mode 100644 index 000000000..40c00cf03 --- /dev/null +++ b/tests/integration/test_v2_saved_reference_updates.py @@ -0,0 +1,204 @@ +"""Repointing a saved report or simulation against migrated Postgres. + +The saved record's references are checked with the simulation service before +the update's transaction opens, so the update locks the saved row and refuses +it if another update changed the row meanwhile. The services run as the +Public API's runtime role. +""" + +from __future__ import annotations + +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from uuid import UUID + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, text +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool +from sqlmodel import Session + +from policyengine_api.constants import REPO +from policyengine_api.services.v2.associations.validators import ( + AssociationConflictError, +) +from policyengine_api.services.v2.simulation_refs.services import ( + SimulationReferenceService, +) +from policyengine_api.services.v2.user_reports.database_session import ( + UserReportDatabaseSession, +) +from policyengine_api.services.v2.user_reports.services import V2UserReportService +from policyengine_api.services.v2.user_reports.types import ( + UserReportCreateFields, + UserReportUpdateInput, +) +from policyengine_api.services.v2.user_simulations.database_session import ( + UserSimulationDatabaseSession, +) +from policyengine_api.services.v2.user_simulations.services import ( + V2UserSimulationService, +) +from policyengine_api.services.v2.user_simulations.types import ( + UserSimulationCreateFields, + UserSimulationUpdateInput, +) +from tests.unit.v2.economy_report_support import ( + REPORT_ID, + FakeEconomyReportSimulation, + report_record, + simulation_record, +) +from tests.unit.v2.user_service_support import ALICE, default_legacy_references + +MIGRATED_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000194") +LOCK_WAIT_SECONDS = 10 + + +@pytest.fixture +def engines(managed_v2_roles, v2_engine_as_role, disposable_v2_database_url): + """Migrate, then yield the runtime role's sessions and an owner engine.""" + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + owner_engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + try: + yield sessionmaker(runtime_engine, class_=Session), owner_engine + finally: + owner_engine.dispose() + runtime_engine.dispose() + + +def _entrypoint() -> FakeEconomyReportSimulation: + entrypoint = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + entrypoint.simulations[MIGRATED_SIMULATION_ID] = simulation_record( + MIGRATED_SIMULATION_ID, legacy_simulation_id=613 + ) + return entrypoint + + +def _wait_for_a_blocked_update(owner_engine) -> None: + """Wait until another session waits for a row lock.""" + deadline = time.monotonic() + LOCK_WAIT_SECONDS + while time.monotonic() < deadline: + with owner_engine.connect() as connection: + waiting = connection.execute( + text( + "SELECT count(*) FROM pg_stat_activity " + "WHERE wait_event_type = 'Lock'" + ) + ).scalar_one() + if waiting: + return + time.sleep(0.05) + pytest.fail("the update never waited for the locked row") + + +@dataclass(frozen=True) +class RowChange: + """Another update's change to one saved row: ``column`` set to ``value``.""" + + table: str + column: str + value: int + row_id: UUID + + +def _repoint_while_the_row_changes(owner_engine, change: RowChange, repoint) -> None: + """Lock the row, start ``repoint``, change the row once it waits, then commit.""" + with owner_engine.connect() as owner, ThreadPoolExecutor(max_workers=1) as pool: + owner.execute( + text(f"SELECT id FROM {change.table} WHERE id = :row_id FOR UPDATE"), + {"row_id": change.row_id}, + ) + repointed = pool.submit(repoint) + _wait_for_a_blocked_update(owner_engine) + owner.execute( + text( + f"UPDATE {change.table} SET {change.column} = :value WHERE id = :row_id" + ), + {"value": change.value, "row_id": change.row_id}, + ) + owner.commit() + with pytest.raises(AssociationConflictError): + repointed.result(timeout=LOCK_WAIT_SECONDS) + + +def test_postgres_saved_report_changed_during_its_check_is_not_repointed( + engines, +) -> None: + sessions, owner_engine = engines + service = V2UserReportService( + UserReportDatabaseSession(sessions), + default_legacy_references(), + SimulationReferenceService(_entrypoint()), + ) + saved = service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", legacy_report_id=615), + ) + + _repoint_while_the_row_changes( + owner_engine, + RowChange("user_report_associations", "legacy_report_id", 616, saved.id), + lambda: service.patch_user_report( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserReportUpdateInput(report_id=REPORT_ID), + ), + ) + + with owner_engine.connect() as connection: + stored = connection.execute( + text( + "SELECT report_id, legacy_report_id FROM user_report_associations " + "WHERE id = :row_id" + ), + {"row_id": saved.id}, + ).one() + assert tuple(stored) == (None, 616) + + +def test_postgres_saved_simulation_changed_during_its_check_is_not_repointed( + engines, +) -> None: + sessions, owner_engine = engines + service = V2UserSimulationService( + UserSimulationDatabaseSession(sessions), + default_legacy_references(), + SimulationReferenceService(_entrypoint()), + ) + saved = service.create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields( + country_id="us", legacy_simulation_id=613 + ), + ) + + _repoint_while_the_row_changes( + owner_engine, + RowChange( + "user_simulation_associations", "legacy_simulation_id", 614, saved.id + ), + lambda: service.patch_user_simulation( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserSimulationUpdateInput( + simulation_id=MIGRATED_SIMULATION_ID + ), + ), + ) + + with owner_engine.connect() as connection: + stored = connection.execute( + text( + "SELECT simulation_id, legacy_simulation_id " + "FROM user_simulation_associations WHERE id = :row_id" + ), + {"row_id": saved.id}, + ).one() + assert tuple(stored) == (None, 614) diff --git a/tests/unit/libs/test_gateway_auth.py b/tests/unit/libs/test_gateway_auth.py index 142361849..fba2f1e9f 100644 --- a/tests/unit/libs/test_gateway_auth.py +++ b/tests/unit/libs/test_gateway_auth.py @@ -21,6 +21,7 @@ GatewayBearerAuth, _require_all_or_none_gateway_auth_env, gateway_auth_required, + gateway_bearer_auth_from_environment, ) ISSUER = "https://policyengine.uk.auth0.com" @@ -485,3 +486,34 @@ def test__given_2xx_response__then_no_retry(self): flow.send(httpx.Response(200, request=request)) provider.invalidate.assert_not_called() + + +def _set_complete_gateway_auth(monkeypatch) -> None: + monkeypatch.setenv("GATEWAY_AUTH_ISSUER", ISSUER) + monkeypatch.setenv("GATEWAY_AUTH_AUDIENCE", AUDIENCE) + monkeypatch.setenv("GATEWAY_AUTH_CLIENT_ID", CLIENT_ID) + monkeypatch.setenv(GATEWAY_AUTH_CLIENT_SECRET_ENV, CLIENT_SECRET) + + +def test_the_environment_selects_bearer_auth_when_it_is_complete(monkeypatch) -> None: + _set_complete_gateway_auth(monkeypatch) + + assert isinstance(gateway_bearer_auth_from_environment(), GatewayBearerAuth) + + +def test_the_environment_selects_no_auth_where_auth_is_not_required() -> None: + assert gateway_bearer_auth_from_environment() is None + + +def test_required_auth_must_be_configured(monkeypatch) -> None: + monkeypatch.setenv(GATEWAY_AUTH_REQUIRED_ENV, "1") + + with pytest.raises(GatewayAuthError, match="GATEWAY_AUTH_ISSUER"): + gateway_bearer_auth_from_environment() + + +def test_partly_configured_auth_is_refused(monkeypatch) -> None: + monkeypatch.setenv("GATEWAY_AUTH_ISSUER", ISSUER) + + with pytest.raises(GatewayAuthError, match="partially configured"): + gateway_bearer_auth_from_environment() diff --git a/tests/unit/libs/test_simulation_economy_reports.py b/tests/unit/libs/test_simulation_economy_reports.py new file mode 100644 index 000000000..859c54ebc --- /dev/null +++ b/tests/unit/libs/test_simulation_economy_reports.py @@ -0,0 +1,628 @@ +"""The economy report adapter against the generated simulation client. + +The client is generated from the Simulation Entrypoint's OpenAPI document and +is not a dependency until a release contains the economy report routes, so +these tests skip when it is not installed. Requests go to an in-memory +transport that answers as the entrypoint does. +""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime +from pathlib import Path +from uuid import UUID + +import httpx +import pytest + +pytest.importorskip("policyengine_api_simulation_client") + +from policyengine_api.libs.gateway_auth import GatewayAuthError +from policyengine_api.libs.simulation_economy_reports import ( + SimulationEntrypointEconomyReports, + economy_reports_from_environment, + simulation_entrypoint_economy_reports, +) +from policyengine_api.request_context import _asgi_request_id +from policyengine_api.services.v2.simulation_entrypoint.types import ( + ComparisonRecord, + EconomyReportRecord, + EconomyReportRunRecord, + EconomyReportSubmission, + ReportStatus, + RunError, + RunTrigger, + SimulationChoice, + SimulationRecord, + SimulationRecordNotFoundError, + SimulationRecordUnreadableError, + SimulationRequestRefusedError, + SimulationRunRecord, + SimulationServiceUnavailableError, +) + +BASE_URL = "https://simulation.example.test" +REPORT_ID = UUID("00000000-0000-4000-8000-000000000001") +BASELINE_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000002") +REFORM_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000003") +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-000000000004") +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-000000000005") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-000000000006") +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-000000000007") +REFORM_POLICY_ID = UUID("00000000-0000-4000-8000-000000000008") +CREATED_AT = "2026-10-27T12:00:00Z" +CREATED_AT_VALUE = datetime(2026, 10, 27, 12, tzinfo=UTC) +COMPLETED_AT = "2026-10-27T12:05:00Z" +COMPLETED_AT_VALUE = datetime(2026, 10, 27, 12, 5, tzinfo=UTC) +RESULT = json.loads( + ( + Path(__file__).parents[2] + / "fixtures" + / "simulation_entrypoint" + / "economy_result_us_congressional_district.json" + ).read_text() +) + + +def _simulation_body(simulation_id: UUID, policy_id: UUID | None) -> dict: + return { + "id": str(simulation_id), + "kind": "economy", + "country": "us", + "policy_id": None if policy_id is None else str(policy_id), + "dynamic_id": None, + "region_code": "congressional_district/CA-12", + "dataset_name": None, + "legacy_simulation_id": None, + "created_at": CREATED_AT, + } + + +def _simulation_run_body(run_id: UUID, simulation_id: UUID) -> dict: + return { + "id": str(run_id), + "simulation_id": str(simulation_id), + "status": "succeeded", + "attempt": 1, + "trigger": "initial", + "error": None, + "dispatched_at": CREATED_AT, + "completed_at": COMPLETED_AT, + "created_at": CREATED_AT, + } + + +def _report_run_body(**changes) -> dict: + return { + "id": str(REPORT_RUN_ID), + "report_id": str(REPORT_ID), + "trigger": "initial", + "catalog_version_id": str(CATALOG_VERSION_ID), + "policyengine_version": "6.7.5", + "year": 2026, + "status": "succeeded", + "stage": None, + "error": None, + "baseline_run": _simulation_run_body(BASELINE_RUN_ID, BASELINE_SIMULATION_ID), + "reform_run": _simulation_run_body(REFORM_RUN_ID, REFORM_SIMULATION_ID), + "comparison": { + "status": "succeeded", + "error": None, + "dispatched_at": COMPLETED_AT, + "completed_at": COMPLETED_AT, + }, + "created_at": CREATED_AT, + **changes, + } + + +def _report_body(**changes) -> dict: + return { + "id": str(REPORT_ID), + "country": "us", + "year": 2026, + "options": {}, + "legacy_report_id": 615, + "created_at": CREATED_AT, + "baseline_simulation": _simulation_body(BASELINE_SIMULATION_ID, None), + "reform_simulation": _simulation_body(REFORM_SIMULATION_ID, REFORM_POLICY_ID), + "status": "succeeded", + "run": _report_run_body(), + "result": RESULT, + **changes, + } + + +def _simulation_record(simulation_id: UUID, policy_id: UUID | None): + return SimulationRecord( + id=simulation_id, + kind="economy", + country_id="us", + policy_id=policy_id, + dynamic_id=None, + region_code="congressional_district/CA-12", + dataset_name=None, + legacy_simulation_id=None, + created_at=CREATED_AT_VALUE, + ) + + +def _simulation_run_record(run_id: UUID, simulation_id: UUID): + return SimulationRunRecord( + id=run_id, + simulation_id=simulation_id, + status="succeeded", + attempt=1, + trigger="initial", + error=None, + dispatched_at=CREATED_AT_VALUE, + completed_at=COMPLETED_AT_VALUE, + created_at=CREATED_AT_VALUE, + ) + + +SUCCEEDED_RUN = EconomyReportRunRecord( + id=REPORT_RUN_ID, + report_id=REPORT_ID, + trigger="initial", + tax_benefit_model_version_id=CATALOG_VERSION_ID, + policyengine_version="6.7.5", + year=2026, + status=ReportStatus.SUCCEEDED, + stage=None, + error=None, + baseline_run=_simulation_run_record(BASELINE_RUN_ID, BASELINE_SIMULATION_ID), + reform_run=_simulation_run_record(REFORM_RUN_ID, REFORM_SIMULATION_ID), + comparison=ComparisonRecord( + status="succeeded", + error=None, + dispatched_at=COMPLETED_AT_VALUE, + completed_at=COMPLETED_AT_VALUE, + ), + created_at=CREATED_AT_VALUE, +) + + +class SimulationEntrypointStub: + """Answer each request with the next queued response and record the request.""" + + def __init__(self, *responses: httpx.Response) -> None: + self.responses = list(responses) + self.requests: list[httpx.Request] = [] + # The auth flow sends one request object again, so its header is + # recorded as it was when sent. + self.authorizations: list[str | None] = [] + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + self.authorizations.append(request.headers.get("Authorization")) + return self.responses.pop(0) + + def adapter(self, auth: httpx.Auth | None = None): + return simulation_entrypoint_economy_reports( + BASE_URL, + auth=auth, + transport=httpx.MockTransport(self), + ) + + +def _submission() -> EconomyReportSubmission: + return EconomyReportSubmission( + country_id="us", + policyengine_version="6.7.5", + year=2026, + region_code="congressional_district/CA-12", + dataset_name=None, + baseline=SimulationChoice(), + reform=SimulationChoice(policy_id=REFORM_POLICY_ID), + ) + + +def test_submit_sends_the_generated_request_and_reads_the_new_report() -> None: + pending_run = _report_run_body( + status="pending", + stage="queued", + comparison={ + "status": "pending", + "error": None, + "dispatched_at": None, + "completed_at": None, + }, + ) + entrypoint = SimulationEntrypointStub( + httpx.Response( + 201, + json=_report_body( + status="pending", + run=pending_run, + result=None, + legacy_report_id=None, + ), + ) + ) + + report = entrypoint.adapter().submit_report(_submission()) + + (request,) = entrypoint.requests + assert request.method == "POST" + assert str(request.url) == f"{BASE_URL}/v2/economy-reports" + assert json.loads(request.content) == { + "country": "us", + "policyengine_version": "6.7.5", + "year": 2026, + "region_code": "congressional_district/CA-12", + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": str(REFORM_POLICY_ID), "dynamic_id": None}, + } + assert report.status is ReportStatus.PENDING + assert report.result is None + assert report.legacy_report_id is None + assert report.run is not None + assert report.run.stage == "queued" + assert report.run.comparison.dispatched_at is None + + +def test_submit_sends_a_named_dataset() -> None: + entrypoint = SimulationEntrypointStub(httpx.Response(201, json=_report_body())) + submission = EconomyReportSubmission( + **{**_submission().__dict__, "dataset_name": "populace_us_2024"} + ) + + entrypoint.adapter().submit_report(submission) + + assert json.loads(entrypoint.requests[0].content)["dataset_name"] == ( + "populace_us_2024" + ) + + +def test_a_succeeded_report_reads_back_with_its_exact_result() -> None: + entrypoint = SimulationEntrypointStub(httpx.Response(200, json=_report_body())) + + report = entrypoint.adapter().get_report(REPORT_ID) + + assert str(entrypoint.requests[0].url) == ( + f"{BASE_URL}/v2/economy-reports/{REPORT_ID}" + ) + assert report == EconomyReportRecord( + id=REPORT_ID, + country_id="us", + year=2026, + options={}, + legacy_report_id=615, + created_at=CREATED_AT_VALUE, + baseline_simulation=_simulation_record(BASELINE_SIMULATION_ID, None), + reform_simulation=_simulation_record(REFORM_SIMULATION_ID, REFORM_POLICY_ID), + status=ReportStatus.SUCCEEDED, + run=SUCCEEDED_RUN, + result=RESULT, + ) + + +def test_a_migrated_report_reads_as_not_run_without_a_run() -> None: + entrypoint = SimulationEntrypointStub( + httpx.Response(200, json=_report_body(status="not_run", run=None, result=None)) + ) + + report = entrypoint.adapter().get_report_by_legacy_id(615) + + assert str(entrypoint.requests[0].url) == ( + f"{BASE_URL}/v2/economy-reports/by-legacy-id/615" + ) + assert report.status is ReportStatus.NOT_RUN + assert report.run is None + assert report.result is None + + +def test_a_run_sends_its_version_trigger_and_force() -> None: + failed_run = _report_run_body( + trigger="manual", + status="failed", + error={"code": "invalid_input", "message": "parameter gov.x is unknown"}, + ) + entrypoint = SimulationEntrypointStub(httpx.Response(201, json=failed_run)) + + run = entrypoint.adapter().run_report( + REPORT_ID, + policyengine_version="6.8.0", + trigger=RunTrigger.MANUAL, + force=True, + ) + + (request,) = entrypoint.requests + assert request.method == "POST" + assert str(request.url) == f"{BASE_URL}/v2/economy-reports/{REPORT_ID}/runs" + assert json.loads(request.content) == { + "policyengine_version": "6.8.0", + "trigger": "manual", + "force": True, + } + assert run.status is ReportStatus.FAILED + assert run.error == RunError( + code="invalid_input", message="parameter gov.x is unknown" + ) + + +def test_runs_are_listed_in_the_order_given() -> None: + newer = _report_run_body( + id="00000000-0000-4000-8000-000000000009", trigger="system" + ) + entrypoint = SimulationEntrypointStub( + httpx.Response(200, json={"runs": [newer, _report_run_body()]}) + ) + + runs = entrypoint.adapter().list_report_runs(REPORT_ID) + + assert str(entrypoint.requests[0].url) == ( + f"{BASE_URL}/v2/economy-reports/{REPORT_ID}/runs" + ) + assert [run.trigger for run in runs] == ["system", "initial"] + assert runs[1] == SUCCEEDED_RUN + + +def test_a_simulation_reads_back() -> None: + entrypoint = SimulationEntrypointStub( + httpx.Response( + 200, json=_simulation_body(REFORM_SIMULATION_ID, REFORM_POLICY_ID) + ) + ) + + simulation = entrypoint.adapter().get_simulation(REFORM_SIMULATION_ID) + + assert str(entrypoint.requests[0].url) == ( + f"{BASE_URL}/v2/simulations/{REFORM_SIMULATION_ID}" + ) + assert simulation == _simulation_record(REFORM_SIMULATION_ID, REFORM_POLICY_ID) + + +def test_an_unknown_report_or_simulation_is_not_found() -> None: + entrypoint = SimulationEntrypointStub( + httpx.Response(404, json={"detail": f"no economy report {REPORT_ID}"}), + httpx.Response(404, json={"detail": "no economy report 615"}), + httpx.Response(404, json={"detail": f"no report {REPORT_ID}"}), + httpx.Response(404, json={"detail": f"no economy report {REPORT_ID}"}), + httpx.Response(404, json={"detail": f"no simulation {REFORM_SIMULATION_ID}"}), + ) + adapter = entrypoint.adapter() + reads = ( + lambda: adapter.get_report(REPORT_ID), + lambda: adapter.get_report_by_legacy_id(615), + lambda: adapter.run_report( + REPORT_ID, + policyengine_version="6.7.5", + trigger=RunTrigger.MANUAL, + force=False, + ), + lambda: adapter.list_report_runs(REPORT_ID), + lambda: adapter.get_simulation(REFORM_SIMULATION_ID), + ) + + for read in reads: + with pytest.raises(SimulationRecordNotFoundError): + read() + + +def test_a_route_the_entrypoint_does_not_serve_means_it_is_unavailable() -> None: + """FastAPI answers 404 "Not Found" for a route sim does not serve.""" + entrypoint = SimulationEntrypointStub( + *(httpx.Response(404, json={"detail": "Not Found"}) for _ in range(3)) + ) + adapter = entrypoint.adapter() + + for read in ( + lambda: adapter.get_report(REPORT_ID), + lambda: adapter.get_simulation(REFORM_SIMULATION_ID), + lambda: adapter.submit_report(_submission()), + ): + with pytest.raises(SimulationServiceUnavailableError) as raised: + read() + assert "404" in str(raised.value) + + +def test_a_stored_result_that_cannot_be_read_back_is_reported() -> None: + detail = "The report's stored result cannot be read back." + entrypoint = SimulationEntrypointStub( + httpx.Response(500, json={"detail": detail}), + httpx.Response(500, json={"detail": detail}), + ) + adapter = entrypoint.adapter() + + for read in ( + lambda: adapter.get_report(REPORT_ID), + lambda: adapter.get_report_by_legacy_id(615), + ): + with pytest.raises(SimulationRecordUnreadableError) as raised: + read() + assert str(raised.value) == detail + + +def test_a_refusal_carries_the_entrypoint_reason() -> None: + reason = "us has no catalog for PolicyEngine.py 9.9.9" + entrypoint = SimulationEntrypointStub(httpx.Response(422, json={"detail": reason})) + + with pytest.raises(SimulationRequestRefusedError) as raised: + entrypoint.adapter().submit_report(_submission()) + + assert str(raised.value) == reason + + +def test_a_request_validation_refusal_names_each_problem() -> None: + entrypoint = SimulationEntrypointStub( + httpx.Response( + 422, + json={ + "detail": [ + { + "loc": ["body", "year"], + "msg": "Input should be less than or equal to 2200", + "type": "less_than_equal", + } + ] + }, + ) + ) + + with pytest.raises(SimulationRequestRefusedError) as raised: + entrypoint.adapter().submit_report(_submission()) + + assert str(raised.value) == "body.year: Input should be less than or equal to 2200" + + +@pytest.mark.parametrize( + "response", + [ + httpx.Response(503, json={"detail": "The simulation database is down."}), + httpx.Response(500, text="Internal Server Error"), + httpx.Response(403, json={"detail": "missing permission"}), + httpx.Response(401, json={"detail": "invalid token"}), + httpx.Response(200, json={"unexpected": True}), + httpx.Response(200, text="not json"), + ], + ids=["unavailable", "server-error", "forbidden", "unauthorized", "shape", "text"], +) +def test_an_unusable_answer_means_the_entrypoint_is_unavailable( + response: httpx.Response, +) -> None: + entrypoint = SimulationEntrypointStub(response) + + with pytest.raises(SimulationServiceUnavailableError): + entrypoint.adapter().get_report(REPORT_ID) + + +def test_a_network_failure_means_the_entrypoint_is_unavailable() -> None: + def refuse(request: httpx.Request) -> httpx.Response: + message = "connection refused" + raise httpx.ConnectError(message, request=request) + + adapter = simulation_entrypoint_economy_reports( + BASE_URL, auth=None, transport=httpx.MockTransport(refuse) + ) + + with pytest.raises(SimulationServiceUnavailableError): + adapter.list_report_runs(REPORT_ID) + + +class CountingBearerAuth(httpx.Auth): + """Send a numbered token, and a new one after a 401, as the gateway auth does.""" + + def __init__(self) -> None: + self.tokens_issued = 0 + + def auth_flow(self, request: httpx.Request): + self.tokens_issued += 1 + request.headers["Authorization"] = f"Bearer token-{self.tokens_issued}" + response = yield request + if response.status_code == 401: + self.tokens_issued += 1 + request.headers["Authorization"] = f"Bearer token-{self.tokens_issued}" + yield request + + +def test_an_expired_token_is_replaced_and_the_request_sent_again() -> None: + entrypoint = SimulationEntrypointStub( + httpx.Response(401, json={"detail": "expired"}), + httpx.Response(200, json=_report_body()), + ) + + report = entrypoint.adapter(auth=CountingBearerAuth()).get_report(REPORT_ID) + + assert report.id == REPORT_ID + assert entrypoint.authorizations == ["Bearer token-1", "Bearer token-2"] + + +class FailingTokenAuth(httpx.Auth): + """Fail to fetch a token, as the gateway auth does when Auth0 cannot answer.""" + + def auth_flow(self, request: httpx.Request): + message = "Auth0 token fetch failed: HTTP 503" + raise GatewayAuthError(message) + yield request # pragma: no cover - makes this a generator, as httpx expects + + +def test_a_failed_token_fetch_means_the_entrypoint_is_unavailable() -> None: + entrypoint = SimulationEntrypointStub(httpx.Response(200, json=_report_body())) + + with pytest.raises(SimulationServiceUnavailableError) as raised: + entrypoint.adapter(auth=FailingTokenAuth()).get_report(REPORT_ID) + + assert "token" in str(raised.value) + assert entrypoint.requests == [] + + +def test_requests_carry_the_current_request_id() -> None: + entrypoint = SimulationEntrypointStub(httpx.Response(200, json=_report_body())) + token = _asgi_request_id.set("api-request-id") + try: + entrypoint.adapter().get_report(REPORT_ID) + finally: + _asgi_request_id.reset(token) + + assert entrypoint.requests[0].headers["X-PolicyEngine-Request-Id"] == ( + "api-request-id" + ) + + +def _clear_gateway_auth(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "GATEWAY_AUTH_ISSUER", + "GATEWAY_AUTH_AUDIENCE", + "GATEWAY_AUTH_CLIENT_ID", + "GATEWAY_AUTH_CLIENT_SECRET", + "GATEWAY_AUTH_CLIENT_SECRET_RESOURCE", + "GATEWAY_AUTH_REQUIRED", + ): + monkeypatch.delenv(name, raising=False) + + +def test_the_environment_names_the_entrypoint_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_gateway_auth(monkeypatch) + monkeypatch.setenv("SIMULATION_ENTRYPOINT_URL", f"{BASE_URL}/") + monkeypatch.setenv("SIM_ENTRYPOINT", "old_gateway_direct") + + adapter = economy_reports_from_environment() + + assert isinstance(adapter, SimulationEntrypointEconomyReports) + assert adapter.base_url == BASE_URL + + +@pytest.mark.parametrize( + "url", + ["", "simulation.example.test", "https://user:secret@simulation.example.test"], + ids=["missing", "no-scheme", "credentials"], +) +def test_an_unusable_entrypoint_url_makes_reports_unavailable( + monkeypatch: pytest.MonkeyPatch, url: str +) -> None: + _clear_gateway_auth(monkeypatch) + monkeypatch.setenv("SIMULATION_ENTRYPOINT_URL", url) + + with pytest.raises(SimulationServiceUnavailableError) as raised: + economy_reports_from_environment() + + assert "SIMULATION_ENTRYPOINT_URL" in str(raised.value) + + +def test_required_gateway_auth_must_be_configured( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_gateway_auth(monkeypatch) + monkeypatch.setenv("SIMULATION_ENTRYPOINT_URL", BASE_URL) + monkeypatch.setenv("GATEWAY_AUTH_REQUIRED", "1") + + with pytest.raises(SimulationServiceUnavailableError) as raised: + economy_reports_from_environment() + + assert "GATEWAY_AUTH_ISSUER" in str(raised.value) + + +def test_partial_gateway_auth_is_refused(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_gateway_auth(monkeypatch) + monkeypatch.setenv("SIMULATION_ENTRYPOINT_URL", BASE_URL) + monkeypatch.setenv("GATEWAY_AUTH_ISSUER", "https://auth.example.test") + + with pytest.raises(SimulationServiceUnavailableError) as raised: + economy_reports_from_environment() + + assert "partially configured" in str(raised.value) diff --git a/tests/unit/libs/test_simulation_http.py b/tests/unit/libs/test_simulation_http.py new file mode 100644 index 000000000..43920c632 --- /dev/null +++ b/tests/unit/libs/test_simulation_http.py @@ -0,0 +1,77 @@ +"""HTTP conventions shared by the Public API's clients of the simulation service.""" + +from __future__ import annotations + +import httpx +import pytest + +from policyengine_api.libs.simulation_http import ( + attach_request_identifiers, + required_base_url, +) +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_api.request_context import ( + REQUEST_ID_HEADER, + _asgi_observability_id, + _asgi_request_id, +) + +ENVIRONMENT_NAME = "SIMULATION_ENTRYPOINT_URL" +OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000042" + + +def test_a_base_url_is_read_without_its_trailing_slash( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(ENVIRONMENT_NAME, " https://simulation.example.test/ ") + + assert required_base_url(ENVIRONMENT_NAME) == "https://simulation.example.test" + + +def test_a_missing_base_url_is_refused(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(ENVIRONMENT_NAME, raising=False) + + with pytest.raises(ValueError, match=f"{ENVIRONMENT_NAME} is required"): + required_base_url(ENVIRONMENT_NAME) + + +@pytest.mark.parametrize( + "url", + [ + "simulation.example.test", + "ftp://simulation.example.test", + "https://user:secret@simulation.example.test", + "https://simulation.example.test/?mode=test", + "https://simulation.example.test/#part", + ], +) +def test_an_unusable_base_url_is_refused( + monkeypatch: pytest.MonkeyPatch, url: str +) -> None: + monkeypatch.setenv(ENVIRONMENT_NAME, url) + + with pytest.raises(ValueError, match="absolute HTTP"): + required_base_url(ENVIRONMENT_NAME) + + +def test_a_request_carries_the_current_identifiers() -> None: + request = httpx.Request("GET", "https://simulation.example.test/health") + request_token = _asgi_request_id.set("api-request-id") + observability_token = _asgi_observability_id.set(OBSERVABILITY_ID) + try: + attach_request_identifiers(request) + finally: + _asgi_request_id.reset(request_token) + _asgi_observability_id.reset(observability_token) + + assert request.headers[REQUEST_ID_HEADER] == "api-request-id" + assert request.headers[OBSERVABILITY_ID_HEADER] == OBSERVABILITY_ID + + +def test_a_request_outside_an_api_request_carries_no_identifiers() -> None: + request = httpx.Request("GET", "https://simulation.example.test/health") + + attach_request_identifiers(request) + + assert REQUEST_ID_HEADER not in request.headers + assert OBSERVABILITY_ID_HEADER not in request.headers diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 797c82848..5f3812f00 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -221,6 +221,9 @@ def test_reusable_v2_integration_check_uses_disposable_databases_and_coverage(): assert "RUN_V2_CATALOG_COMPATIBILITY" in workflow assert "test_v2_catalog_publication.py" in workflow assert "test_v2_metadata_routes.py" in workflow + assert "test_v2_country_routes.py" in workflow + assert "test_v2_economy_reports.py" in workflow + assert "test_v2_saved_reference_updates.py" in workflow assert "test_v2_policy_persistence.py" in workflow assert "test_v1_policy_dual_write.py" in workflow assert "test_mysql_policy_dual_write.py" in workflow diff --git a/tests/unit/test_migration_contract_artifacts.py b/tests/unit/test_migration_contract_artifacts.py index d08090feb..353ae639f 100644 --- a/tests/unit/test_migration_contract_artifacts.py +++ b/tests/unit/test_migration_contract_artifacts.py @@ -13,8 +13,8 @@ def test_migration_contract_payload_summarizes_route_contracts(): assert payload["version"] == 1 assert payload["metadata"] == { "route_group_count": 9, - "workflow_count": 16, - "request_count": 89, + "workflow_count": 18, + "request_count": 96, "db_entity_count": 6, "sim_flow_count": 3, } @@ -31,6 +31,8 @@ def test_migration_contract_payload_summarizes_route_contracts(): "user_account_v2", "household_calculate", "region_selection", + "country_selection_v2", + "economy_reports_v2", "metadata_resources_v2_preview", "simulation_submit_poll", "report_create_poll", diff --git a/tests/unit/test_startup_country_check.py b/tests/unit/test_startup_country_check.py new file mode 100644 index 000000000..64bd0f219 --- /dev/null +++ b/tests/unit/test_startup_country_check.py @@ -0,0 +1,229 @@ +"""The startup check that the v2 countries table matches the country list in code.""" + +from __future__ import annotations + +import logging +import os +import subprocess +import sys +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError + +from policyengine_api import startup_country_check +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.startup_country_check import check_countries_table_at_startup +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS + +if TYPE_CHECKING: + from collections.abc import Iterable + from pathlib import Path + +V2_RUNTIME_SECRET_RESOURCE = "projects/example/secrets/v2-runtime/versions/latest" + + +@pytest.fixture(autouse=True) +def enabled_startup_check_logger(monkeypatch: pytest.MonkeyPatch) -> None: + """Re-enable the check's logger, which Alembic's logging setup disables. + + Tests that run Alembic commands load its logging configuration, which + disables every logger that already exists, including this module's. + """ + monkeypatch.setattr(startup_country_check.logger, "disabled", False) + + +class FakeCountryListVerifier: + """Record each verification, then raise ``error`` if one is set.""" + + def __init__(self, error: Exception | None = None) -> None: + self.error = error + self.verified_country_lists: list[tuple[str, ...]] = [] + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + self.verified_country_lists.append(tuple(listed_country_ids)) + if self.error is not None: + raise self.error + + +class CountingVerifierFactory: + """Return one verifier and count how often startup asks for it.""" + + def __init__(self, verifier: FakeCountryListVerifier) -> None: + self.verifier = verifier + self.calls = 0 + + def __call__(self) -> FakeCountryListVerifier: + self.calls += 1 + return self.verifier + + +def test_startup_without_a_v2_runtime_database_skips_the_check() -> None: + factory = CountingVerifierFactory(FakeCountryListVerifier()) + + check_countries_table_at_startup( + environ={"FLASK_DEBUG": "1"}, verifier_factory=factory + ) + + assert factory.calls == 0 + + +def test_blank_selection_variables_skip_the_check() -> None: + factory = CountingVerifierFactory(FakeCountryListVerifier()) + + check_countries_table_at_startup( + environ={ + "K_SERVICE": "", + "V2_RUNTIME_DATABASE_URL": " ", + "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE": "", + }, + verifier_factory=factory, + ) + + assert factory.calls == 0 + + +@pytest.mark.parametrize( + "environ", + [ + {"K_SERVICE": "policyengine-api"}, + {"V2_RUNTIME_DATABASE_URL": "postgresql+psycopg://configured"}, + {"V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE": V2_RUNTIME_SECRET_RESOURCE}, + ], + ids=["cloud-run", "runtime-url", "runtime-secret"], +) +def test_a_selected_v2_runtime_database_is_checked_against_the_code_list( + environ: dict[str, str], +) -> None: + verifier = FakeCountryListVerifier() + + check_countries_table_at_startup( + environ=environ, verifier_factory=CountingVerifierFactory(verifier) + ) + + assert verifier.verified_country_lists == [SUPPORTED_COUNTRY_IDS] + + +@pytest.mark.parametrize( + "error", + [ + CountryListMismatchError(listed_without_row=("us",), rows_not_listed=()), + ProgrammingError( + "SELECT countries.id FROM countries", + {}, + Exception('relation "countries" does not exist'), + ), + ProgrammingError( + "SELECT countries.id FROM countries", + {}, + Exception("permission denied for table countries"), + ), + ], + ids=["mismatch", "missing-table", "missing-privilege"], +) +def test_a_failed_check_stops_startup(error: Exception) -> None: + verifier = FakeCountryListVerifier(error=error) + + with pytest.raises(type(error)): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=CountingVerifierFactory(verifier), + ) + + +@pytest.mark.parametrize( + "message", + [ + "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host", + "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE could not be resolved", + ], + ids=["invalid-url", "secret-unavailable"], +) +def test_unusable_v2_configuration_is_logged_and_left_to_readiness( + message: str, + caplog: pytest.LogCaptureFixture, +) -> None: + """Readiness already reports not ready for it, and recovers when it does.""" + + def unusable_configuration() -> FakeCountryListVerifier: + raise V2ConfigurationError(message) + + with caplog.at_level(logging.WARNING): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=unusable_configuration, + ) + + assert [record.levelno for record in caplog.records] == [logging.WARNING] + assert message in caplog.records[0].getMessage() + + +@pytest.mark.parametrize( + "error_type", + [OperationalError, InterfaceError], + ids=["operational", "interface"], +) +def test_an_unreachable_database_is_logged_and_startup_continues( + error_type: type[OperationalError | InterfaceError], + caplog: pytest.LogCaptureFixture, +) -> None: + verifier = FakeCountryListVerifier( + error=error_type( + "SELECT countries.id FROM countries", + {}, + Exception("connection to db.secret-host failed: password=hunter2"), + ) + ) + + with caplog.at_level(logging.WARNING): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=CountingVerifierFactory(verifier), + ) + + assert verifier.verified_country_lists == [SUPPORTED_COUNTRY_IDS] + assert [record.levelno for record in caplog.records] == [logging.WARNING] + message = caplog.records[0].getMessage() + assert "countries table" in message + assert error_type.__name__ in message + assert "secret-host" not in message + assert "hunter2" not in message + + +def test_importing_the_asgi_application_runs_the_check(tmp_path: Path) -> None: + """The check runs at import; an unusable v2 runtime URL is only logged.""" + environment = os.environ.copy() + environment.update( + { + "POLICYENGINE_API_STARTUP_WARMUP": "0", + "V2_RUNTIME_DATABASE_URL": ( + "postgresql+psycopg://runtime:hunter2@localhost:5432/postgres" + "?sslmode=require" + ), + "V2_SUPABASE_PROJECT_REF": "abcdefghijklmnopqrst", + "V2_SUPABASE_ENVIRONMENT": "test-startup", + } + ) + for name in ("V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE", "K_SERVICE"): + environment.pop(name, None) + + result = subprocess.run( + [sys.executable, "-c", "import policyengine_api.asgi"], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + check=False, + ) + + # The application's log handler writes structured records to stdout. + output = result.stdout + result.stderr + assert result.returncode == 0, result.stderr + assert "Skipped the countries table check" in output + assert "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" in output + assert "hunter2" not in output diff --git a/tests/unit/test_supported_countries.py b/tests/unit/test_supported_countries.py new file mode 100644 index 000000000..429f27657 --- /dev/null +++ b/tests/unit/test_supported_countries.py @@ -0,0 +1,139 @@ +"""The Public API's one country list, and the lists that derive from it.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from typing import get_args + +import pytest +from pydantic import TypeAdapter, ValidationError + +from policyengine_api.constants import REPO +from policyengine_api.data.v2.catalog import extraction +from policyengine_api.query_parameters import CountryId +from policyengine_api.supported_countries import ( + SUPPORTED_COUNTRY_IDS, + SupportedCountryId, +) + +# Runs in a fresh interpreter: it adds a country to the one list before any +# module that derives from it is imported, then reports what each derived +# list contains. PolicyEngine.py is replaced by a stand-in, because importing +# the real package builds every country model. +EXTENDED_LIST_SCRIPT = """ +import json +import sys +from types import ModuleType, SimpleNamespace +from typing import Literal, get_args + +import policyengine_api.supported_countries as supported_countries + +supported_countries.SupportedCountryId = Literal["us", "uk", "zz"] +supported_countries.SUPPORTED_COUNTRY_IDS = get_args( + supported_countries.SupportedCountryId +) + +installed_policyengine = ModuleType("policyengine") +for country_id in ("us", "uk", "zz", "ca"): + setattr( + installed_policyengine, + country_id, + SimpleNamespace(model=f"{country_id} model"), + ) +installed_bundle = ModuleType("policyengine.bundle") +installed_bundle.get_current_bundle = dict +sys.modules["policyengine"] = installed_policyengine +sys.modules["policyengine.bundle"] = installed_bundle + +from pydantic import TypeAdapter + +from policyengine_api.data.v2.catalog import catalog_selection, extraction +from policyengine_api.query_parameters import CountryId +from scripts import process_v1_household_mirror_event + +captured_arguments = {} + + +def capture_extraction(**arguments): + captured_arguments.update(arguments) + + +extraction.extract_catalog = capture_extraction +extraction.extract_installed_catalog() +country_option = next( + action + for action in process_v1_household_mirror_event._parser()._actions + if action.dest == "country_id" +) +print( + json.dumps( + { + "country_type": TypeAdapter(CountryId).validate_python("ZZ"), + "extracted_countries": list(extraction.SUPPORTED_COUNTRIES), + "required_dependencies": list(extraction.REQUIRED_DEPENDENCIES), + "installed_models": captured_arguments["models"], + "dual_write_countries": sorted(catalog_selection.SUPPORTED_V2_COUNTRY_IDS), + "household_copy_countries": list(country_option.choices), + } + ) +) +""" + + +def test_the_country_list_is_the_us_and_the_uk_in_catalog_order() -> None: + assert SUPPORTED_COUNTRY_IDS == ("us", "uk") + assert get_args(SupportedCountryId) == SUPPORTED_COUNTRY_IDS + + +def test_the_v2_country_type_accepts_exactly_the_listed_countries() -> None: + country_id_adapter = TypeAdapter(CountryId) + + for country_id in SUPPORTED_COUNTRY_IDS: + assert country_id_adapter.validate_python(country_id) == country_id + assert country_id_adapter.validate_python(country_id.upper()) == country_id + for unlisted_country_id in ("ca", "ng", "il", ""): + with pytest.raises(ValidationError): + country_id_adapter.validate_python(unlisted_country_id) + + +def test_every_country_list_follows_a_country_added_to_the_one_list() -> None: + result = subprocess.run( + [sys.executable, "-c", EXTENDED_LIST_SCRIPT], + cwd=REPO, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + derived_lists = json.loads(result.stdout.strip().splitlines()[-1]) + assert derived_lists == { + "country_type": "zz", + "extracted_countries": ["us", "uk", "zz"], + "required_dependencies": [ + "policyengine-core", + "policyengine-us", + "policyengine-uk", + "policyengine-zz", + ], + # The installed package's other countries are not extracted. + "installed_models": { + "us": "us model", + "uk": "uk model", + "zz": "zz model", + }, + "dual_write_countries": ["uk", "us", "zz"], + "household_copy_countries": ["us", "uk", "zz"], + } + + +def test_catalog_settings_cover_exactly_the_listed_countries() -> None: + # Each listed country needs reviewed catalog settings, so adding a country + # without them fails here rather than during publication. + listed_countries = set(SUPPORTED_COUNTRY_IDS) + + assert set(extraction.REVIEWED_DEFAULT_DATASETS) == listed_countries + assert set(extraction.CURRENT_LAW_IDS) == listed_countries + assert set(extraction.METADATA_TIME_PERIODS) == listed_countries diff --git a/tests/unit/v2/economy_report_support.py b/tests/unit/v2/economy_report_support.py new file mode 100644 index 000000000..8a10c3eea --- /dev/null +++ b/tests/unit/v2/economy_report_support.py @@ -0,0 +1,303 @@ +"""A fake Simulation Entrypoint and economy report records for unit tests.""" + +from __future__ import annotations + +import threading +from dataclasses import replace +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import NAMESPACE_URL, UUID, uuid5 + +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import StaticPool +from sqlmodel import Session, create_engine + +from policyengine_api.data.v2.models import V2_METADATA, Country +from policyengine_api.services.v2.economy_reports.database_session import ( + EconomyReportDatabaseSession, +) +from policyengine_api.services.v2.economy_reports.services import ( + V2EconomyReportService, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + ComparisonRecord, + EconomyReportRecord, + EconomyReportRunRecord, + ReportStatus, + RunTrigger, + SimulationRecord, + SimulationRecordNotFoundError, + SimulationRecordUnreadableError, + SimulationRunRecord, +) +from tests.fixtures.v2_countries import add_listed_countries + +if TYPE_CHECKING: + from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportSubmission, + ) + +CURRENT_POLICYENGINE_VERSION = "6.7.5" +REPORT_ID = UUID("00000000-0000-4000-8000-000000000101") +BASELINE_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000102") +REFORM_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000103") +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-000000000104") +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-000000000105") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-000000000106") +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-000000000107") +REFORM_POLICY_ID = UUID("00000000-0000-4000-8000-000000000108") +CREATED_AT = datetime(2026, 10, 27, 12, tzinfo=UTC) +# A country that has only household calculations; no real country does yet. +HOUSEHOLD_ONLY_COUNTRY = Country( + id="zz", household_enabled=True, society_wide_enabled=False +) +RESULT = {"budget": {"budgetary_impact": -1250000.0}} + + +def simulation_record( + simulation_id: UUID, + *, + policy_id: UUID | None = None, + country_id: str = "us", + legacy_simulation_id: int | None = None, +) -> SimulationRecord: + return SimulationRecord( + id=simulation_id, + kind="economy", + country_id=country_id, + policy_id=policy_id, + dynamic_id=None, + region_code="state/ca", + dataset_name=None, + legacy_simulation_id=legacy_simulation_id, + created_at=CREATED_AT, + ) + + +def _simulation_run(run_id: UUID, simulation_id: UUID) -> SimulationRunRecord: + return SimulationRunRecord( + id=run_id, + simulation_id=simulation_id, + status="pending", + attempt=1, + trigger="initial", + error=None, + dispatched_at=None, + completed_at=None, + created_at=CREATED_AT, + ) + + +def report_run_record( + *, + run_id: UUID = REPORT_RUN_ID, + trigger: str = "initial", + status: ReportStatus = ReportStatus.PENDING, + policyengine_version: str = CURRENT_POLICYENGINE_VERSION, +) -> EconomyReportRunRecord: + return EconomyReportRunRecord( + id=run_id, + report_id=REPORT_ID, + trigger=trigger, + tax_benefit_model_version_id=CATALOG_VERSION_ID, + policyengine_version=policyengine_version, + year=2026, + status=status, + stage="queued" if status is ReportStatus.PENDING else None, + error=None, + baseline_run=_simulation_run(BASELINE_RUN_ID, BASELINE_SIMULATION_ID), + reform_run=_simulation_run(REFORM_RUN_ID, REFORM_SIMULATION_ID), + comparison=ComparisonRecord( + status="pending", error=None, dispatched_at=None, completed_at=None + ), + created_at=CREATED_AT, + ) + + +def report_record( + *, + report_id: UUID = REPORT_ID, + country_id: str = "us", + run: EconomyReportRunRecord | None = None, + legacy_report_id: int | None = None, +) -> EconomyReportRecord: + """Build a report: not run without a run, with a result once it succeeded. + + Each report has its own two simulations, which its run names. + """ + baseline_id, reform_id = _simulation_ids(report_id) + if run is not None: + run = replace( + run, + report_id=report_id, + baseline_run=replace(run.baseline_run, simulation_id=baseline_id), + reform_run=replace(run.reform_run, simulation_id=reform_id), + ) + status = ReportStatus.NOT_RUN if run is None else run.status + return EconomyReportRecord( + id=report_id, + country_id=country_id, + year=2026, + options={}, + legacy_report_id=legacy_report_id, + created_at=CREATED_AT, + baseline_simulation=simulation_record(baseline_id, country_id=country_id), + reform_simulation=simulation_record( + reform_id, policy_id=REFORM_POLICY_ID, country_id=country_id + ), + status=status, + run=run, + result=RESULT if status is ReportStatus.SUCCEEDED else None, + ) + + +def _simulation_ids(report_id: UUID) -> tuple[UUID, UUID]: + """Return a report's baseline and reform simulation IDs.""" + if report_id == REPORT_ID: + return BASELINE_SIMULATION_ID, REFORM_SIMULATION_ID + return ( + uuid5(NAMESPACE_URL, f"{report_id}/baseline"), + uuid5(NAMESPACE_URL, f"{report_id}/reform"), + ) + + +class FakeEconomyReportSimulation: + """An in-memory Simulation Entrypoint that records every call. + + ``reports`` and ``simulations`` hold what it knows; ``error`` is raised by + every call while set, and ``run_error`` by run starts only. Reading a + report in ``unreadable_results`` fails as the entrypoint does when it + cannot rebuild a stored result. As the entrypoint does, a ``system`` run + of a report that has a run returns its newest run instead of adding one, + under a lock, so simultaneous first runs record one run. + """ + + def __init__(self, *reports: EconomyReportRecord) -> None: + self.reports = {report.id: report for report in reports} + self.simulations: dict[UUID, SimulationRecord] = { + simulation.id: simulation + for report in reports + for simulation in (report.baseline_simulation, report.reform_simulation) + } + self.runs: dict[UUID, list[EconomyReportRunRecord]] = { + report.id: [report.run] for report in reports if report.run is not None + } + self.unreadable_results: set[UUID] = set() + self.calls: list[tuple[str, object]] = [] + self.error: Exception | None = None + self.run_error: Exception | None = None + self.next_run = report_run_record() + self._runs_lock = threading.Lock() + + def _record(self, name: str, argument: object) -> None: + self.calls.append((name, argument)) + if self.error is not None: + raise self.error + + def submit_report(self, submission: EconomyReportSubmission) -> EconomyReportRecord: + self._record("submit_report", submission) + report = report_record(country_id=submission.country_id, run=self.next_run) + self.reports[report.id] = report + for simulation in (report.baseline_simulation, report.reform_simulation): + self.simulations[simulation.id] = simulation + if report.run is not None: + self.runs[report.id] = [report.run] + return report + + def get_report(self, report_id: UUID) -> EconomyReportRecord: + self._record("get_report", report_id) + if report_id not in self.reports: + msg = f"no economy report {report_id}" + raise SimulationRecordNotFoundError(msg) + if report_id in self.unreadable_results: + msg = "The report's stored result cannot be read back." + raise SimulationRecordUnreadableError(msg) + return self.reports[report_id] + + def get_report_by_legacy_id(self, legacy_report_id: int) -> EconomyReportRecord: + self._record("get_report_by_legacy_id", legacy_report_id) + for report in self.reports.values(): + if report.legacy_report_id == legacy_report_id: + return report + msg = f"no economy report {legacy_report_id}" + raise SimulationRecordNotFoundError(msg) + + def run_report( + self, + report_id: UUID, + *, + policyengine_version: str, + trigger: RunTrigger, + force: bool, + ) -> EconomyReportRunRecord: + self._record( + "run_report", + { + "report_id": report_id, + "policyengine_version": policyengine_version, + "trigger": trigger, + "force": force, + }, + ) + if self.run_error is not None: + raise self.run_error + if report_id not in self.reports: + msg = f"no report {report_id}" + raise SimulationRecordNotFoundError(msg) + with self._runs_lock: + existing_runs = self.runs.get(report_id, []) + if trigger is RunTrigger.SYSTEM and existing_runs: + return existing_runs[0] + run = replace(self.next_run, trigger=trigger.value) + self.runs.setdefault(report_id, []).insert(0, run) + self.reports[report_id] = replace( + self.reports[report_id], run=run, status=run.status, result=None + ) + return run + + def list_report_runs(self, report_id: UUID) -> tuple[EconomyReportRunRecord, ...]: + self._record("list_report_runs", report_id) + if report_id not in self.reports: + msg = f"no economy report {report_id}" + raise SimulationRecordNotFoundError(msg) + return tuple(self.runs.get(report_id, ())) + + def get_simulation(self, simulation_id: UUID) -> SimulationRecord: + self._record("get_simulation", simulation_id) + if simulation_id not in self.simulations: + msg = f"no simulation {simulation_id}" + raise SimulationRecordNotFoundError(msg) + return self.simulations[simulation_id] + + +def countries_store(*extra_countries: Country) -> sessionmaker[Session]: + """Build an SQLite v2 schema with the US and UK rows and ``extra_countries``. + + Routes run in a worker thread, so every session shares one connection to + the in-memory database. + """ + engine = create_engine( + "sqlite://", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + V2_METADATA.create_all(engine) + sessions = sessionmaker(engine, class_=Session) + with sessions() as session: + add_listed_countries(session) + session.add_all(extra_countries) + session.commit() + return sessions + + +def economy_report_service( + simulation: FakeEconomyReportSimulation, + sessions: sessionmaker[Session] | None = None, +) -> V2EconomyReportService: + return V2EconomyReportService( + EconomyReportDatabaseSession( + countries_store(HOUSEHOLD_ONLY_COUNTRY) if sessions is None else sessions + ), + simulation, + current_policyengine_version=CURRENT_POLICYENGINE_VERSION, + ) diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 4c0e6ee09..56e31410b 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( { + "countries", "legacy_household_mappings", "legacy_policy_mappings", "legacy_user_mappings", @@ -235,6 +236,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, + "fbb7488ebe8a", "91af4fc8d8c1", "60d6518b6a98", "439303be14fe", @@ -495,7 +497,6 @@ def test_stage_12a_user_identity_and_association_revision_is_generated() -> None / "migrations/v2/versions/fbb7488ebe8a_link_signed_in_users_to_v1_anchored_.py" ).read_text(encoding="utf-8") - assert V2_ALEMBIC_HEAD_REVISION == "fbb7488ebe8a" assert ( "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" in revision @@ -565,6 +566,94 @@ def test_stage_12a_user_identity_and_association_revision_is_generated() -> None ) in upgrade +COUNTRY_REFERENCE_COLUMNS = ( + ("users", "primary_country"), + ("households", "country_id"), + ("policies", "country_id"), + ("user_household_associations", "country_id"), + ("user_policies", "country_id"), + ("user_simulation_associations", "country_id"), + ("user_report_associations", "country_id"), +) +# The generated downgrade restores each check in PostgreSQL's reflected form +# of `` IN ('us', 'uk')``. +US_AND_UK = "ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])" +REPLACED_COUNTRY_CHECKS = { + "ck_users_primary_country": ("users", f"primary_country::text = {US_AND_UK}"), + "ck_households_country": ("households", f"country_id::text = {US_AND_UK}"), + "ck_policies_country": ("policies", f"country_id::text = {US_AND_UK}"), + "ck_user_household_associations_country": ( + "user_household_associations", + f"country_id::text = {US_AND_UK}", + ), + "ck_user_policies_country": ("user_policies", f"country_id::text = {US_AND_UK}"), + "ck_user_simulation_associations_country": ( + "user_simulation_associations", + f"country_id::text = {US_AND_UK}", + ), + "ck_user_report_associations_country": ( + "user_report_associations", + f"country_id::text = {US_AND_UK}", + ), +} + + +def _countries_revision_source() -> str: + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + countries_revision = next( + revision + for revision in script.walk_revisions() + if revision.down_revision == "fbb7488ebe8a" + ) + assert countries_revision.revision == V2_ALEMBIC_HEAD_REVISION + assert countries_revision.path is not None + return Path(countries_revision.path).read_text(encoding="utf-8") + + +def test_countries_revision_inserts_us_and_uk_and_replaces_country_checks() -> None: + revision = _countries_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 revision.count("op.create_table(") == 1 + assert 'op.create_table(\n "countries"' in upgrade + assert "ck_countries_society_wide_requires_household" in upgrade + assert 'op.drop_table("countries")' in downgrade + # The one hand-written operation, an approved exception to generated-only + # revisions: the US and UK rows, inserted before the foreign keys that + # existing rows must satisfy. The generated downgrade drops them with the + # table. + assert revision.count("Approved hand-written data operation") == 1 + assert revision.count("Post-generation correction") == 0 + assert revision.count("op.bulk_insert(") == 1 + assert "op.execute(" not in revision + inserted_rows = upgrade[ + upgrade.index("op.bulk_insert(") : upgrade.index("op.create_foreign_key(") + ] + for country_id in ("uk", "us"): + assert f'"id": "{country_id}"' in inserted_rows + assert inserted_rows.count('"household_enabled": True') == 2 + assert inserted_rows.count('"society_wide_enabled": True') == 2 + for table_name, column_name in COUNTRY_REFERENCE_COLUMNS: + key_name = f"fk_{table_name}_{column_name}_countries" + assert key_name in upgrade + assert key_name in downgrade + assert upgrade.count('ondelete="RESTRICT"') == len(COUNTRY_REFERENCE_COLUMNS) + for check_name, (table_name, expression) in REPLACED_COUNTRY_CHECKS.items(): + assert check_name in upgrade + assert f'"{table_name}"' in upgrade + assert check_name in downgrade + assert expression in downgrade + assert "op.drop_column(" not in upgrade + assert "op.add_column(" not in upgrade + + def test_alembic_rejects_unknown_missing_and_divergent_history(tmp_path: Path) -> None: original = REPO / "migrations/v2" missing = tmp_path / "missing" diff --git a/tests/unit/v2/test_country_routes.py b/tests/unit/v2/test_country_routes.py new file mode 100644 index 000000000..7fd2e2b91 --- /dev/null +++ b/tests/unit/v2/test_country_routes.py @@ -0,0 +1,198 @@ +"""Route coverage for the native v2 countries list.""" + +from __future__ import annotations + +import pytest +from fastapi.testclient import TestClient +from flask import Flask, jsonify +from sqlalchemy.exc import OperationalError + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, + infer_route_group, +) +from policyengine_api.services.v2.countries.types import CountryRead + +UNITED_KINGDOM = CountryRead(id="uk", household_enabled=True, society_wide_enabled=True) +UNITED_STATES = CountryRead(id="us", household_enabled=True, society_wide_enabled=True) +HOUSEHOLD_ONLY_COUNTRY = CountryRead( + id="be", household_enabled=True, society_wide_enabled=False +) + + +class FakeCountryService: + """Return fixed countries, or raise ``error``, and count the calls.""" + + def __init__( + self, + countries: list[CountryRead], + *, + error: Exception | None = None, + ) -> None: + self.countries = countries + self.error = error + self.list_calls = 0 + + def list_countries(self) -> list[CountryRead]: + self.list_calls += 1 + if self.error is not None: + raise self.error + return self.countries + + +def _client(service: FakeCountryService) -> tuple[TestClient, dict[str, int]]: + flask_calls = {"count": 0} + flask_app = Flask(__name__) + + @flask_app.route("/", methods=["GET", "POST"]) + def fallback(resource: str): + flask_calls["count"] += 1 + return jsonify({"source": "flask", "resource": resource}) + + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_country_service_factory=lambda: service, + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + client = TestClient( + create_asgi_app(flask_app, dependencies=dependencies, route_settings=settings), + raise_server_exceptions=False, + ) + return client, flask_calls + + +def test_countries_list_every_country_with_its_capability() -> None: + service = FakeCountryService( + [HOUSEHOLD_ONLY_COUNTRY, UNITED_KINGDOM, UNITED_STATES] + ) + client, flask_calls = _client(service) + + response = client.get("/v2/countries") + + assert response.status_code == 200 + assert response.json() == { + "status": "ok", + "message": None, + "result": { + "items": [ + { + "id": "be", + "household_enabled": True, + "society_wide_enabled": False, + }, + { + "id": "uk", + "household_enabled": True, + "society_wide_enabled": True, + }, + { + "id": "us", + "household_enabled": True, + "society_wide_enabled": True, + }, + ] + }, + } + assert service.list_calls == 1 + assert flask_calls["count"] == 0 + + +def test_countries_list_needs_no_user_credential() -> None: + client, _flask_calls = _client(FakeCountryService([UNITED_STATES])) + + response = client.get("/v2/countries", headers={"Authorization": ""}) + + assert response.status_code == 200 + + +def test_countries_list_rejects_query_parameters() -> None: + service = FakeCountryService([UNITED_STATES]) + client, _flask_calls = _client(service) + + response = client.get("/v2/countries?country_id=us") + + assert response.status_code == 422 + assert response.json()["status"] == "error" + assert service.list_calls == 0 + + +@pytest.mark.parametrize( + "error", + [ + V2ConfigurationError("V2_RUNTIME_DATABASE_URL is required"), + OperationalError("SELECT", {}, Exception("connection refused")), + ], + ids=["configuration", "database"], +) +def test_countries_list_reports_an_unavailable_database_as_503( + error: Exception, +) -> None: + client, _flask_calls = _client(FakeCountryService([], error=error)) + + response = client.get("/v2/countries") + + assert response.status_code == 503 + assert response.json() == { + "status": "error", + "message": "V2 countries are unavailable", + } + + +def test_countries_list_hides_unexpected_failures_behind_a_500() -> None: + client, _flask_calls = _client( + FakeCountryService([], error=RuntimeError("secret detail")) + ) + + response = client.get("/v2/countries") + + assert response.status_code == 500 + assert response.json() == { + "status": "error", + "message": "V2 countries query failed", + } + + +def test_countries_list_supports_get_only() -> None: + service = FakeCountryService([UNITED_STATES]) + client, flask_calls = _client(service) + + response = client.post("/v2/countries", json={}) + + assert response.status_code == 405 + assert service.list_calls == 0 + assert flask_calls["count"] == 0 + + +def test_countries_list_is_documented_in_the_v2_openapi_document() -> None: + client, _flask_calls = _client(FakeCountryService([])) + + document = client.get("/v2/openapi.json").json() + operation = document["paths"]["/v2/countries"]["get"] + + assert operation.get("parameters", []) == [] + assert operation["responses"]["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/CountryListResponse" + } + assert {"422", "500", "503"} <= operation["responses"].keys() + item_schema = document["components"]["schemas"]["CountryItem"] + assert item_schema["required"] == [ + "id", + "household_enabled", + "society_wide_enabled", + ] + assert item_schema["additionalProperties"] is False + + +def test_countries_requests_are_logged_with_the_metadata_route_group() -> None: + assert infer_route_group("/v2/countries") == "metadata" diff --git a/tests/unit/v2/test_country_service.py b/tests/unit/v2/test_country_service.py new file mode 100644 index 000000000..253278599 --- /dev/null +++ b/tests/unit/v2/test_country_service.py @@ -0,0 +1,85 @@ +"""SQLite-backed coverage for the v2 countries service: the list and the check.""" + +from __future__ import annotations + +import pytest +from sqlalchemy.orm import sessionmaker +from sqlmodel import Session, create_engine + +from policyengine_api.data.v2.models import V2_METADATA, Country +from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, +) +from policyengine_api.services.v2.countries.services import V2CountryService +from policyengine_api.services.v2.countries.types import CountryRead +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS +from tests.fixtures.v2_countries import add_listed_countries + + +def _service_with_countries(*extra_countries: Country) -> V2CountryService: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + sessions = sessionmaker(engine, class_=Session) + with sessions() as session: + add_listed_countries(session) + session.add_all(extra_countries) + session.commit() + return V2CountryService(CountryDatabaseSession(sessions)) + + +def test_countries_are_listed_by_id_with_their_capability_flags() -> None: + service = _service_with_countries( + Country(id="be", household_enabled=True, society_wide_enabled=False) + ) + + countries = service.list_countries() + + assert countries == [ + CountryRead(id="be", household_enabled=True, society_wide_enabled=False), + CountryRead(id="uk", household_enabled=True, society_wide_enabled=True), + CountryRead(id="us", household_enabled=True, society_wide_enabled=True), + ] + + +def test_an_empty_countries_table_lists_no_countries() -> None: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(engine, class_=Session)) + ) + + assert service.list_countries() == [] + + +def test_rows_for_exactly_the_listed_countries_pass_verification() -> None: + service = _service_with_countries() + + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + +def test_a_row_for_an_unlisted_country_fails_verification() -> None: + service = _service_with_countries( + Country(id="ca", household_enabled=True, society_wide_enabled=False) + ) + + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + assert raised.value.rows_not_listed == ("ca",) + assert raised.value.listed_without_row == () + + +def test_an_empty_countries_table_fails_verification() -> None: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(engine, class_=Session)) + ) + + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + assert raised.value.listed_without_row == tuple(sorted(SUPPORTED_COUNTRY_IDS)) diff --git a/tests/unit/v2/test_country_validators.py b/tests/unit/v2/test_country_validators.py new file mode 100644 index 000000000..d56db5c00 --- /dev/null +++ b/tests/unit/v2/test_country_validators.py @@ -0,0 +1,68 @@ +"""The comparison of stored countries rows with the country list in code.""" + +from __future__ import annotations + +import pytest + +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, + validate_countries_match_code_list, +) + + +def test_the_same_countries_in_any_order_match() -> None: + validate_countries_match_code_list( + stored_country_ids=["us", "uk"], + listed_country_ids=("uk", "us"), + ) + + +def test_a_listed_country_without_a_row_is_a_mismatch() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["uk"], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == ("us",) + assert raised.value.rows_not_listed == () + assert str(raised.value) == ( + "the countries table does not match the country list in code: " + "listed without a row: us; rows not listed: none" + ) + + +def test_a_row_for_an_unlisted_country_is_a_mismatch() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["ca", "uk", "us"], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == () + assert raised.value.rows_not_listed == ("ca",) + + +def test_every_difference_is_reported_in_order() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["ng", "ca"], + listed_country_ids=("us", "uk"), + ) + + assert raised.value.listed_without_row == ("uk", "us") + assert raised.value.rows_not_listed == ("ca", "ng") + assert str(raised.value) == ( + "the countries table does not match the country list in code: " + "listed without a row: uk, us; rows not listed: ca, ng" + ) + + +def test_an_empty_table_misses_every_listed_country() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=[], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == ("uk", "us") diff --git a/tests/unit/v2/test_economy_report_routes.py b/tests/unit/v2/test_economy_report_routes.py new file mode 100644 index 000000000..1d5e7ed27 --- /dev/null +++ b/tests/unit/v2/test_economy_report_routes.py @@ -0,0 +1,788 @@ +"""Route coverage for the native v2 economy reports. + +The routes run on the real service, with SQLite countries rows and a fake +Simulation Entrypoint. +""" + +from __future__ import annotations + +import logging +import sys +from uuid import UUID + +import pytest +from fastapi.testclient import TestClient +from flask import Flask, jsonify +from sqlalchemy.exc import OperationalError +from sqlmodel import select + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.v2.models import Country +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + _default_economy_report_simulation, + _default_v2_economy_report_service_factory, +) +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, + infer_route_group, +) +from policyengine_api.services.v2.economy_reports.types import EconomyReportFields +from policyengine_api.services.v2.economy_reports.validators import ( + EconomyReportsUnsupportedError, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + ReportStatus, + SimulationChoice, + SimulationRecordUnreadableError, + SimulationRequestRefusedError, + SimulationServiceUnavailableError, +) +from tests.unit.v2.economy_report_support import ( + BASELINE_RUN_ID, + BASELINE_SIMULATION_ID, + CATALOG_VERSION_ID, + CURRENT_POLICYENGINE_VERSION, + HOUSEHOLD_ONLY_COUNTRY, + REFORM_POLICY_ID, + REFORM_RUN_ID, + REFORM_SIMULATION_ID, + REPORT_ID, + REPORT_RUN_ID, + RESULT, + FakeEconomyReportSimulation, + countries_store, + economy_report_service, + report_record, + report_run_record, +) + +SUBMISSION = { + "country_id": "us", + "kind": "economy", + "year": 2026, + "region_code": "state/ca", + "baseline": {}, + "reform": {"policy_id": str(REFORM_POLICY_ID)}, +} +CREATED_AT = "2026-10-27T12:00:00Z" + + +def _client( + simulation: FakeEconomyReportSimulation, + *, + service_factory=None, + sessions=None, +) -> TestClient: + flask_app = Flask(__name__) + + @flask_app.route("/", methods=["GET", "POST"]) + def fallback(resource: str): + return jsonify({"source": "flask", "resource": resource}) + + service = economy_report_service(simulation, sessions) + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_economy_report_service_factory=service_factory or (lambda: service), + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + return TestClient( + create_asgi_app(flask_app, dependencies=dependencies, route_settings=settings), + raise_server_exceptions=False, + ) + + +def _simulation_item(simulation_id, policy_id=None) -> dict: + return { + "id": str(simulation_id), + "kind": "economy", + "country_id": "us", + "policy_id": None if policy_id is None else str(policy_id), + "dynamic_id": None, + "region_code": "state/ca", + "dataset_name": None, + "legacy_simulation_id": None, + "created_at": CREATED_AT, + } + + +def _simulation_run_item(run_id, simulation_id) -> dict: + return { + "id": str(run_id), + "simulation_id": str(simulation_id), + "status": "pending", + "attempt": 1, + "trigger": "initial", + "error": None, + "dispatched_at": None, + "completed_at": None, + "created_at": CREATED_AT, + } + + +PENDING_RUN_ITEM = { + "id": str(REPORT_RUN_ID), + "report_id": str(REPORT_ID), + "trigger": "initial", + "tax_benefit_model_version_id": str(CATALOG_VERSION_ID), + "policyengine_version": CURRENT_POLICYENGINE_VERSION, + "year": 2026, + "status": "pending", + "stage": "queued", + "error": None, + "baseline_run": _simulation_run_item(BASELINE_RUN_ID, BASELINE_SIMULATION_ID), + "reform_run": _simulation_run_item(REFORM_RUN_ID, REFORM_SIMULATION_ID), + "comparison": { + "status": "pending", + "error": None, + "dispatched_at": None, + "completed_at": None, + }, + "created_at": CREATED_AT, +} + + +def test_submit_creates_the_report_and_answers_with_its_first_run() -> None: + simulation = FakeEconomyReportSimulation() + client = _client(simulation) + + response = client.post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 201 + assert response.json() == { + "status": "ok", + "message": None, + "result": { + "item": { + "id": str(REPORT_ID), + "country_id": "us", + "kind": "economy", + "year": 2026, + "options": {}, + "legacy_report_id": None, + "created_at": CREATED_AT, + "baseline_simulation": _simulation_item(BASELINE_SIMULATION_ID), + "reform_simulation": _simulation_item( + REFORM_SIMULATION_ID, REFORM_POLICY_ID + ), + "status": "pending", + "run": PENDING_RUN_ITEM, + "result": None, + } + }, + } + ((name, submission),) = simulation.calls + assert name == "submit_report" + assert submission.policyengine_version == CURRENT_POLICYENGINE_VERSION + + +def test_submit_needs_no_user_credential() -> None: + client = _client(FakeEconomyReportSimulation()) + + response = client.post( + "/v2/reports?country_id=us", json=SUBMISSION, headers={"Authorization": ""} + ) + + assert response.status_code == 201 + + +def test_submit_refuses_a_body_country_that_differs_from_the_query() -> None: + simulation = FakeEconomyReportSimulation() + client = _client(simulation) + + response = client.post("/v2/reports?country_id=uk", json=SUBMISSION) + + assert response.status_code == 400 + assert response.json() == { + "status": "error", + "message": "Body country_id must match query country_id", + } + assert simulation.calls == [] + + +@pytest.mark.parametrize( + "changes", + [ + {"kind": "household"}, + {"kind": None}, + {"year": 1899}, + {"year": 2201}, + {"region_code": ""}, + {"region_code": "state ca"}, + {"dataset_name": ""}, + {"options": {}}, + {"baseline": {"policy_id": "not-a-uuid"}}, + {"reform": {"policy_id": str(REFORM_POLICY_ID), "extra": 1}}, + {"country_id": "ca"}, + ], +) +def test_submit_refuses_a_body_outside_the_contract(changes: dict) -> None: + simulation = FakeEconomyReportSimulation() + client = _client(simulation) + + response = client.post("/v2/reports?country_id=us", json={**SUBMISSION, **changes}) + + assert response.status_code == 422 + assert response.json()["status"] == "error" + assert simulation.calls == [] + + +@pytest.mark.parametrize("missing", ["kind", "year", "region_code", "reform"]) +def test_submit_refuses_a_body_missing_a_required_field(missing: str) -> None: + simulation = FakeEconomyReportSimulation() + body = {name: value for name, value in SUBMISSION.items() if name != missing} + + response = _client(simulation).post("/v2/reports?country_id=us", json=body) + + assert response.status_code == 422 + assert simulation.calls == [] + + +def test_submit_for_a_household_only_country_is_refused_naming_it() -> None: + sessions = countries_store() + with sessions() as session: + united_states = session.exec(select(Country).where(Country.id == "us")).one() + united_states.society_wide_enabled = False + session.commit() + simulation = FakeEconomyReportSimulation() + client = _client(simulation, sessions=sessions) + + response = client.post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 422 + assert response.json() == { + "status": "error", + "message": "Economy reports are not available for country 'us'", + } + assert simulation.calls == [] + + +def test_submit_relays_the_entrypoint_reason_when_it_cannot_run_the_report() -> None: + simulation = FakeEconomyReportSimulation() + reason = ( + "region place/ca-los_angeles has no household filter in this catalog version" + ) + simulation.error = SimulationRequestRefusedError(reason) + + response = _client(simulation).post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 422 + assert response.json() == {"status": "error", "message": reason} + + +@pytest.mark.parametrize( + "error", + [ + SimulationServiceUnavailableError("the entrypoint answered HTTP 503"), + V2ConfigurationError("V2_RUNTIME_DATABASE_URL is not configured"), + OperationalError("SELECT 1", {}, Exception("connection refused")), + ], + ids=["entrypoint", "configuration", "database"], +) +def test_submit_is_unavailable_when_a_dependency_is(error: Exception) -> None: + def failing_service(): + raise error + + response = _client( + FakeEconomyReportSimulation(), service_factory=failing_service + ).post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 503 + assert response.json()["status"] == "error" + assert "connection refused" not in response.json()["message"] + + +def test_an_unexpected_failure_is_a_secret_safe_server_error() -> None: + simulation = FakeEconomyReportSimulation() + simulation.error = RuntimeError("token=secret") + + response = _client(simulation).post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 500 + assert response.json() == { + "status": "error", + "message": "V2 report operation failed", + } + + +def test_reports_are_unavailable_without_the_generated_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delitem( + sys.modules, "policyengine_api.libs.simulation_economy_reports", raising=False + ) + monkeypatch.setitem(sys.modules, "policyengine_api_simulation_client", None) + + with pytest.raises(SimulationServiceUnavailableError) as raised: + _default_economy_report_simulation() + + assert "policyengine-api-simulation-client" in str(raised.value) + + +def test_report_requests_are_logged_with_the_report_route_group() -> None: + assert infer_route_group("/v2/reports") == "report" + assert infer_route_group(f"/v2/reports/{REPORT_ID}/runs") == "report" + + +def test_openapi_publishes_the_submit_contract() -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + operation = schema["paths"]["/v2/reports"]["post"] + assert {parameter["name"] for parameter in operation["parameters"]} == { + "country_id" + } + assert set(operation["responses"]) >= {"201", "400", "422", "503"} + request_schema = operation["requestBody"]["content"]["application/json"]["schema"] + assert request_schema["$ref"].endswith("/EconomyReportCreateRequest") + body = schema["components"]["schemas"]["EconomyReportCreateRequest"] + assert set(body["required"]) == { + "country_id", + "kind", + "year", + "region_code", + "baseline", + "reform", + } + assert body["additionalProperties"] is False + + +def test_a_succeeded_report_is_read_with_its_result() -> None: + simulation = FakeEconomyReportSimulation( + report_record(run=report_run_record(status=ReportStatus.SUCCEEDED)) + ) + + response = _client(simulation).get(f"/v2/reports/{REPORT_ID}?country_id=us") + + assert response.status_code == 200 + item = response.json()["result"]["item"] + assert item["status"] == "succeeded" + assert item["run"]["status"] == "succeeded" + assert item["result"] == RESULT + + +def test_reading_a_migrated_report_starts_its_first_run() -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + response = _client(simulation).get(f"/v2/reports/{REPORT_ID}?country_id=us") + + assert response.status_code == 200 + item = response.json()["result"]["item"] + assert item["legacy_report_id"] == 615 + assert item["status"] == "pending" + assert item["run"]["trigger"] == "system" + assert item["result"] is None + + +@pytest.mark.parametrize( + "reports", + [(), (report_record(country_id="uk"),)], + ids=["unknown", "other-country"], +) +def test_a_report_outside_the_country_is_not_found(reports) -> None: + response = _client(FakeEconomyReportSimulation(*reports)).get( + f"/v2/reports/{REPORT_ID}?country_id=us" + ) + + assert response.status_code == 404 + assert response.json() == { + "status": "error", + "message": f"Economy report {REPORT_ID} was not found in country 'us'", + } + + +@pytest.mark.parametrize( + "path", + [ + "/v2/reports/not-a-uuid?country_id=us", + f"/v2/reports/{REPORT_ID}", + f"/v2/reports/{REPORT_ID}?country_id=us&country_id=uk", + f"/v2/reports/{REPORT_ID}?country_id=us&extra=1", + "/v2/reports/by-legacy-number/0?country_id=us", + "/v2/reports/by-legacy-number/first?country_id=us", + "/v2/reports/by-legacy-number/615", + ], +) +def test_a_read_outside_the_contract_is_refused(path: str) -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + response = _client(simulation).get(path) + + assert response.status_code == 422 + assert simulation.calls == [] + + +def test_a_migrated_report_is_found_by_its_v1_number() -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + response = _client(simulation).get("/v2/reports/by-legacy-number/615?country_id=us") + + assert response.status_code == 200 + assert response.json()["result"]["item"]["id"] == str(REPORT_ID) + assert simulation.calls[0] == ("get_report_by_legacy_id", 615) + + +def test_an_unknown_v1_number_is_not_found() -> None: + response = _client(FakeEconomyReportSimulation()).get( + "/v2/reports/by-legacy-number/615?country_id=us" + ) + + assert response.status_code == 404 + assert response.json()["message"] == ( + "Economy report migrated from v1 report 615 was not found in country 'us'" + ) + + +def test_openapi_publishes_the_read_contracts() -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + for path in ( + "/v2/reports/{report_id}", + "/v2/reports/by-legacy-number/{legacy_report_id}", + ): + operation = schema["paths"][path]["get"] + assert {parameter["name"] for parameter in operation["parameters"]} >= { + "country_id" + } + assert set(operation["responses"]) >= {"200", "404", "422", "503"} + legacy_number = next( + parameter + for parameter in schema["paths"][ + "/v2/reports/by-legacy-number/{legacy_report_id}" + ]["get"]["parameters"] + if parameter["name"] == "legacy_report_id" + ) + assert legacy_number["schema"]["minimum"] == 1 + + +def test_a_rerun_answers_with_the_new_run() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + response = _client(simulation).post( + f"/v2/reports/{REPORT_ID}/runs?country_id=us", json={"force": True} + ) + + assert response.status_code == 201 + assert response.json() == { + "status": "ok", + "message": None, + "result": {"item": {**PENDING_RUN_ITEM, "trigger": "manual"}}, + } + assert simulation.calls[-1][1]["force"] is True + + +def test_a_rerun_reuses_simulation_runs_unless_forced() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + response = _client(simulation).post( + f"/v2/reports/{REPORT_ID}/runs?country_id=us", json={} + ) + + assert response.status_code == 201 + assert simulation.calls[-1][1]["force"] is False + + +def test_a_rerun_in_a_household_only_country_is_refused_naming_it() -> None: + sessions = countries_store() + with sessions() as session: + united_states = session.exec(select(Country).where(Country.id == "us")).one() + united_states.society_wide_enabled = False + session.commit() + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + response = _client(simulation, sessions=sessions).post( + f"/v2/reports/{REPORT_ID}/runs?country_id=us", json={} + ) + + assert response.status_code == 422 + assert response.json()["message"] == ( + "Economy reports are not available for country 'us'" + ) + assert simulation.calls == [] + + +@pytest.mark.parametrize( + "body", + [{"force": "yes"}, {"force": None}, {"trigger": "system"}, None], + ids=["text", "null", "trigger", "no-body"], +) +def test_a_rerun_body_outside_the_contract_is_refused(body) -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + response = _client(simulation).post( + f"/v2/reports/{REPORT_ID}/runs?country_id=us", json=body + ) + + assert response.status_code == 422 + assert simulation.calls == [] + + +def test_a_reports_runs_are_listed() -> None: + simulation = FakeEconomyReportSimulation(report_record()) + client = _client(simulation) + client.get(f"/v2/reports/{REPORT_ID}?country_id=us") + client.post(f"/v2/reports/{REPORT_ID}/runs?country_id=us", json={}) + + response = client.get(f"/v2/reports/{REPORT_ID}/runs?country_id=us") + + assert response.status_code == 200 + assert response.json() == { + "status": "ok", + "message": None, + "result": { + "items": [ + {**PENDING_RUN_ITEM, "trigger": "manual"}, + {**PENDING_RUN_ITEM, "trigger": "system"}, + ] + }, + } + + +def test_the_runs_of_an_unknown_report_are_not_found() -> None: + response = _client(FakeEconomyReportSimulation()).get( + f"/v2/reports/{REPORT_ID}/runs?country_id=us" + ) + + assert response.status_code == 404 + + +def test_openapi_publishes_the_run_contracts() -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + runs = schema["paths"]["/v2/reports/{report_id}/runs"] + assert set(runs) == {"get", "post"} + assert "201" in runs["post"]["responses"] + assert set(runs["get"]["responses"]) >= {"200", "404", "422", "503"} + body = schema["components"]["schemas"]["EconomyReportRunRequest"] + assert body["additionalProperties"] is False + assert body["properties"]["force"]["default"] is False + + +SECOND_REPORT_ID = "00000000-0000-4000-8000-000000000201" + + +def test_reports_are_listed_by_id_without_their_results() -> None: + simulation = FakeEconomyReportSimulation( + report_record(run=report_run_record(status=ReportStatus.SUCCEEDED)), + report_record(report_id=UUID(SECOND_REPORT_ID), legacy_report_id=615), + ) + + response = _client(simulation).get( + f"/v2/reports?country_id=us&report_id={SECOND_REPORT_ID}&report_id={REPORT_ID}" + ) + + assert response.status_code == 200 + items = response.json()["result"]["items"] + assert [item["id"] for item in items] == [SECOND_REPORT_ID, str(REPORT_ID)] + assert [item["status"] for item in items] == ["not_run", "succeeded"] + assert all("result" not in item for item in items) + assert items[0]["run"] is None + assert [name for name, _ in simulation.calls] == ["get_report", "get_report"] + + +@pytest.mark.parametrize( + "query", + [ + "country_id=us", + "country_id=us&report_id=not-a-uuid", + "report_id=" + SECOND_REPORT_ID, + "country_id=us&country_id=uk&report_id=" + SECOND_REPORT_ID, + "country_id=us&" + + "&".join( + f"report_id=00000000-0000-4000-8000-{number:012d}" for number in range(26) + ), + ], + ids=["no-ids", "bad-id", "no-country", "two-countries", "too-many"], +) +def test_a_listing_outside_the_contract_is_refused(query: str) -> None: + simulation = FakeEconomyReportSimulation() + + response = _client(simulation).get(f"/v2/reports?{query}") + + assert response.status_code == 422 + assert simulation.calls == [] + + +def test_openapi_publishes_the_listing_contract() -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + operation = schema["paths"]["/v2/reports"]["get"] + report_ids = next( + parameter + for parameter in operation["parameters"] + if parameter["name"] == "report_id" + ) + assert report_ids["required"] is True + assert report_ids["schema"]["type"] == "array" + assert report_ids["schema"]["maxItems"] == 25 + item = schema["components"]["schemas"]["ReportSummaryItem"] + assert "result" not in item["properties"] + + +def _without_the_generated_client(monkeypatch: pytest.MonkeyPatch, sessions) -> None: + """Make the default factory's client import fail, and read countries from SQLite.""" + from policyengine_api.data.v2 import database + + monkeypatch.delitem( + sys.modules, "policyengine_api.libs.simulation_economy_reports", raising=False + ) + monkeypatch.setitem(sys.modules, "policyengine_api_simulation_client", None) + monkeypatch.setattr(database, "get_v2_session_factory", lambda: sessions) + + +def test_the_default_service_refuses_a_household_only_country_without_the_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _without_the_generated_client(monkeypatch, countries_store(HOUSEHOLD_ONLY_COUNTRY)) + service = _default_v2_economy_report_service_factory() + report_input = EconomyReportFields( + kind="economy", + year=2026, + region_code="state/ca", + baseline=SimulationChoice(), + reform=SimulationChoice(), + ) + + with pytest.raises(EconomyReportsUnsupportedError): + service.submit_report(country_id="zz", report_input=report_input) + with pytest.raises(EconomyReportsUnsupportedError): + service.run_report(country_id="zz", report_id=REPORT_ID, force=False) + with pytest.raises(SimulationServiceUnavailableError): + service.submit_report(country_id="us", report_input=report_input) + + +def test_a_household_only_submission_is_a_422_even_without_the_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sessions = countries_store() + with sessions() as session: + united_states = session.exec(select(Country).where(Country.id == "us")).one() + united_states.society_wide_enabled = False + session.commit() + _without_the_generated_client(monkeypatch, sessions) + + response = _client( + FakeEconomyReportSimulation(), + service_factory=_default_v2_economy_report_service_factory, + ).post("/v2/reports?country_id=us", json=SUBMISSION) + + assert response.status_code == 422 + assert response.json()["message"] == ( + "Economy reports are not available for country 'us'" + ) + + +def test_a_report_whose_result_cannot_be_read_back_is_a_bad_gateway() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + simulation.error = SimulationRecordUnreadableError( + "The report's stored result cannot be read back." + ) + + response = _client(simulation).get(f"/v2/reports/{REPORT_ID}?country_id=us") + + assert response.status_code == 502 + assert response.json() == { + "status": "error", + "message": "The simulation service cannot read back this report's result", + } + + +@pytest.mark.parametrize( + ("method", "path", "statuses"), + [ + ("post", "/v2/reports", {"201", "400", "413", "422", "500", "503"}), + ("get", "/v2/reports", {"200", "422", "500", "503"}), + ( + "get", + "/v2/reports/{report_id}", + {"200", "404", "422", "500", "502", "503"}, + ), + ( + "get", + "/v2/reports/by-legacy-number/{legacy_report_id}", + {"200", "404", "422", "500", "502", "503"}, + ), + ( + "post", + "/v2/reports/{report_id}/runs", + {"201", "404", "413", "422", "500", "503"}, + ), + ("get", "/v2/reports/{report_id}/runs", {"200", "404", "422", "500", "503"}), + ], +) +def test_openapi_declares_only_the_answers_each_route_can_give( + method: str, path: str, statuses: set[str] +) -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + assert set(schema["paths"][path][method]["responses"]) == statuses + + +class RecordingHandler(logging.Handler): + """Keep every record logged to the logger it is attached to.""" + + def __init__(self) -> None: + super().__init__(level=logging.DEBUG) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +def test_an_unexpected_failure_is_logged_with_where_it_happened() -> None: + simulation = FakeEconomyReportSimulation() + simulation.error = RuntimeError("token=secret") + routes_logger = logging.getLogger( + "policyengine_api.fastapi_routes.v2.reports.routes" + ) + handler = RecordingHandler() + # Alembic's logging configuration, run by earlier tests, disables loggers + # that already exist. + was_disabled = routes_logger.disabled + routes_logger.disabled = False + routes_logger.addHandler(handler) + try: + _client(simulation).post("/v2/reports?country_id=us", json=SUBMISSION) + finally: + routes_logger.removeHandler(handler) + routes_logger.disabled = was_disabled + + (record,) = [ + record + for record in handler.records + if record.getMessage() == "v2_report_operation_failed" + ] + attributes = record.policyengine_attributes + assert record.levelno == logging.ERROR + assert attributes["error_type"] == "RuntimeError" + assert "economy_report_support.py" in attributes["error_frames"] + assert "secret" not in str(attributes) + + +@pytest.mark.parametrize( + ("method", "path", "refuses_runs"), + [ + ("post", "/v2/reports", True), + ("post", "/v2/reports/{report_id}/runs", True), + ("get", "/v2/reports", False), + ("get", "/v2/reports/{report_id}", False), + ("get", "/v2/reports/by-legacy-number/{legacy_report_id}", False), + ("get", "/v2/reports/{report_id}/runs", False), + ], +) +def test_openapi_says_which_routes_refuse_a_report_that_cannot_run( + method: str, path: str, *, refuses_runs: bool +) -> None: + schema = _client(FakeEconomyReportSimulation()).get("/v2/openapi.json").json() + + description = schema["paths"][path][method]["responses"]["422"]["description"] + + assert ("economy reports" in description) is refuses_runs diff --git a/tests/unit/v2/test_economy_report_service.py b/tests/unit/v2/test_economy_report_service.py new file mode 100644 index 000000000..d82b6f6fa --- /dev/null +++ b/tests/unit/v2/test_economy_report_service.py @@ -0,0 +1,429 @@ +"""The v2 economy report service, with SQLite countries rows and a fake entrypoint.""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor +from uuid import UUID + +import pytest + +from policyengine_api.services.v2.economy_reports.types import EconomyReportFields +from policyengine_api.services.v2.economy_reports.validators import ( + EconomyReportNotFoundError, + EconomyReportsUnsupportedError, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + EconomyReportSubmission, + ReportStatus, + RunTrigger, + SimulationChoice, + SimulationRequestRefusedError, +) +from tests.unit.v2.economy_report_support import ( + BASELINE_SIMULATION_ID, + CURRENT_POLICYENGINE_VERSION, + REFORM_POLICY_ID, + REPORT_ID, + FakeEconomyReportSimulation, + economy_report_service, + report_record, + report_run_record, +) + + +def _report_fields(**changes) -> EconomyReportFields: + return EconomyReportFields( + **{ + "kind": "economy", + "year": 2026, + "region_code": "state/ca", + "baseline": SimulationChoice(), + "reform": SimulationChoice(policy_id=REFORM_POLICY_ID), + **changes, + } + ) + + +def test_submit_asks_the_entrypoint_for_the_report_on_the_current_catalog() -> None: + simulation = FakeEconomyReportSimulation() + service = economy_report_service(simulation) + + report = service.submit_report( + country_id="us", + report_input=_report_fields(dataset_name="populace_us_2024"), + ) + + assert simulation.calls == [ + ( + "submit_report", + EconomyReportSubmission( + country_id="us", + policyengine_version=CURRENT_POLICYENGINE_VERSION, + year=2026, + region_code="state/ca", + dataset_name="populace_us_2024", + baseline=SimulationChoice(), + reform=SimulationChoice(policy_id=REFORM_POLICY_ID), + ), + ) + ] + assert report == simulation.reports[report.id] + + +def test_a_household_only_country_is_refused_before_the_entrypoint_is_asked() -> None: + simulation = FakeEconomyReportSimulation() + service = economy_report_service(simulation) + + with pytest.raises(EconomyReportsUnsupportedError) as raised: + service.submit_report(country_id="zz", report_input=_report_fields()) + + assert "'zz'" in str(raised.value) + assert simulation.calls == [] + + +def test_a_country_without_a_countries_row_is_refused() -> None: + simulation = FakeEconomyReportSimulation() + service = economy_report_service(simulation) + + with pytest.raises(EconomyReportsUnsupportedError) as raised: + service.submit_report(country_id="xx", report_input=_report_fields()) + + assert "'xx'" in str(raised.value) + assert simulation.calls == [] + + +def test_the_entrypoint_refusal_reaches_the_caller_unchanged() -> None: + simulation = FakeEconomyReportSimulation() + reason = "region place/ca-los_angeles has no household filter" + simulation.error = SimulationRequestRefusedError(reason) + service = economy_report_service(simulation) + + with pytest.raises(SimulationRequestRefusedError, match=reason): + service.submit_report(country_id="us", report_input=_report_fields()) + + +def test_a_report_of_the_country_is_read_as_the_entrypoint_has_it() -> None: + succeeded = report_record(run=report_run_record(status=ReportStatus.SUCCEEDED)) + simulation = FakeEconomyReportSimulation(succeeded) + + report = economy_report_service(simulation).get_report( + country_id="us", report_id=REPORT_ID + ) + + assert report == succeeded + assert simulation.calls == [("get_report", REPORT_ID)] + + +@pytest.mark.parametrize( + "reports", + [(), (report_record(country_id="uk"),)], + ids=["unknown", "other-country"], +) +def test_a_report_not_in_the_country_is_not_found(reports) -> None: + simulation = FakeEconomyReportSimulation(*reports) + + with pytest.raises(EconomyReportNotFoundError) as raised: + economy_report_service(simulation).get_report( + country_id="us", report_id=REPORT_ID + ) + + assert str(raised.value) == ( + f"Economy report {REPORT_ID} was not found in country 'us'" + ) + assert [name for name, _ in simulation.calls] == ["get_report"] + + +def test_reading_a_report_with_no_run_starts_its_first_run() -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + report = economy_report_service(simulation).get_report( + country_id="us", report_id=REPORT_ID + ) + + assert simulation.calls[1] == ( + "run_report", + { + "report_id": REPORT_ID, + "policyengine_version": CURRENT_POLICYENGINE_VERSION, + "trigger": RunTrigger.SYSTEM, + "force": False, + }, + ) + assert [name for name, _ in simulation.calls] == [ + "get_report", + "run_report", + "get_report", + ] + assert report.status is ReportStatus.PENDING + assert report.run is not None + assert report.run.trigger == "system" + + +def test_a_report_with_a_run_is_not_run_again_when_read() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + economy_report_service(simulation).get_report(country_id="us", report_id=REPORT_ID) + + assert [name for name, _ in simulation.calls] == ["get_report"] + + +def test_a_household_only_country_report_is_not_run_when_read() -> None: + simulation = FakeEconomyReportSimulation(report_record(country_id="zz")) + + report = economy_report_service(simulation).get_report( + country_id="zz", report_id=REPORT_ID + ) + + assert report.status is ReportStatus.NOT_RUN + assert [name for name, _ in simulation.calls] == ["get_report"] + + +def test_a_first_run_the_entrypoint_refuses_leaves_the_report_not_run() -> None: + simulation = FakeEconomyReportSimulation(report_record()) + simulation.run_error = SimulationRequestRefusedError( + "region constituency/E14000530 is not in the catalog version" + ) + + report = economy_report_service(simulation).get_report( + country_id="us", report_id=REPORT_ID + ) + + assert report.status is ReportStatus.NOT_RUN + assert report.run is None + + +def test_a_migrated_report_is_found_by_its_v1_number_and_first_run() -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + report = economy_report_service(simulation).get_report_by_legacy_number( + country_id="us", legacy_report_id=615 + ) + + assert report.legacy_report_id == 615 + assert report.run is not None + assert report.run.trigger == "system" + assert simulation.calls[0] == ("get_report_by_legacy_id", 615) + + +def test_an_unknown_v1_number_is_not_found() -> None: + simulation = FakeEconomyReportSimulation(report_record(legacy_report_id=615)) + + with pytest.raises(EconomyReportNotFoundError) as raised: + economy_report_service(simulation).get_report_by_legacy_number( + country_id="us", legacy_report_id=616 + ) + + assert str(raised.value) == ( + "Economy report migrated from v1 report 616 was not found in country 'us'" + ) + + +def test_a_rerun_runs_on_the_current_catalog_with_the_manual_trigger() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + + run = economy_report_service(simulation).run_report( + country_id="us", report_id=REPORT_ID, force=True + ) + + assert simulation.calls == [ + ("list_report_runs", REPORT_ID), + ("get_simulation", BASELINE_SIMULATION_ID), + ( + "run_report", + { + "report_id": REPORT_ID, + "policyengine_version": CURRENT_POLICYENGINE_VERSION, + "trigger": RunTrigger.MANUAL, + "force": True, + }, + ), + ] + assert run.trigger == "manual" + + +def test_a_rerun_in_a_household_only_country_is_refused_before_the_entrypoint() -> None: + simulation = FakeEconomyReportSimulation(report_record(country_id="zz")) + + with pytest.raises(EconomyReportsUnsupportedError) as raised: + economy_report_service(simulation).run_report( + country_id="zz", report_id=REPORT_ID, force=False + ) + + assert "'zz'" in str(raised.value) + assert simulation.calls == [] + + +@pytest.mark.parametrize( + "reports", + [(), (report_record(country_id="uk"),)], + ids=["unknown", "other-country"], +) +def test_a_rerun_of_a_report_not_in_the_country_is_not_found(reports) -> None: + simulation = FakeEconomyReportSimulation(*reports) + + with pytest.raises(EconomyReportNotFoundError): + economy_report_service(simulation).run_report( + country_id="us", report_id=REPORT_ID, force=False + ) + + assert "run_report" not in [name for name, _ in simulation.calls] + + +def test_a_rerun_the_entrypoint_refuses_reaches_the_caller() -> None: + simulation = FakeEconomyReportSimulation(report_record(run=report_run_record())) + reason = "us has no catalog for PolicyEngine.py 6.7.5" + simulation.run_error = SimulationRequestRefusedError(reason) + + with pytest.raises(SimulationRequestRefusedError, match=reason): + economy_report_service(simulation).run_report( + country_id="us", report_id=REPORT_ID, force=False + ) + + +def test_a_reports_runs_are_listed_newest_first() -> None: + simulation = FakeEconomyReportSimulation(report_record()) + service = economy_report_service(simulation) + service.get_report(country_id="us", report_id=REPORT_ID) + service.run_report(country_id="us", report_id=REPORT_ID, force=True) + + runs = service.list_report_runs(country_id="us", report_id=REPORT_ID) + + assert [run.trigger for run in runs] == ["manual", "system"] + + +def test_the_runs_of_a_report_not_in_the_country_are_not_found() -> None: + simulation = FakeEconomyReportSimulation(report_record(country_id="uk")) + + with pytest.raises(EconomyReportNotFoundError): + economy_report_service(simulation).list_report_runs( + country_id="us", report_id=REPORT_ID + ) + + +def test_listed_reports_keep_the_asked_order_and_only_the_countrys() -> None: + second_id = UUID("00000000-0000-4000-8000-000000000201") + other_country_id = UUID("00000000-0000-4000-8000-000000000202") + unknown_id = UUID("00000000-0000-4000-8000-000000000203") + simulation = FakeEconomyReportSimulation( + report_record(), + report_record(report_id=second_id, run=report_run_record()), + report_record(report_id=other_country_id, country_id="uk"), + ) + + reports = economy_report_service(simulation).list_reports( + country_id="us", + report_ids=[second_id, unknown_id, REPORT_ID, other_country_id, second_id], + ) + + assert [report.id for report in reports] == [second_id, REPORT_ID] + assert reports[1].status is ReportStatus.NOT_RUN + assert [name for name, _ in simulation.calls] == ["get_report"] * 4 + + +class SimultaneousFirstReads(FakeEconomyReportSimulation): + """Hold each of the first two report reads until both have read the report.""" + + def __init__(self, *reports) -> None: + super().__init__(*reports) + self._first_reads = threading.Barrier(2, timeout=5) + self._reads_lock = threading.Lock() + self._reads = 0 + + def get_report(self, report_id): + report = super().get_report(report_id) + with self._reads_lock: + self._reads += 1 + first_read = self._reads <= 2 + if first_read: + self._first_reads.wait() + return report + + +def test_simultaneous_first_reads_share_one_system_run() -> None: + """Both reads see no run and start one; the entrypoint records one run. + + The entrypoint returns a report's newest run to a system run, under a + lock, so the Public API must start first runs only with the system + trigger and no force, and answer with the run the entrypoint returns. + """ + simulation = SimultaneousFirstReads(report_record(legacy_report_id=615)) + service = economy_report_service(simulation) + + with ThreadPoolExecutor(max_workers=2) as executor: + reads = [ + executor.submit(service.get_report, country_id="us", report_id=REPORT_ID) + for _ in range(2) + ] + reports = [read.result(timeout=10) for read in reads] + + run_starts = [call for name, call in simulation.calls if name == "run_report"] + assert len(run_starts) == 2 + assert all( + start["trigger"] is RunTrigger.SYSTEM and start["force"] is False + for start in run_starts + ) + assert len(simulation.runs[REPORT_ID]) == 1 + assert {report.run.id for report in reports if report.run is not None} == { + simulation.runs[REPORT_ID][0].id + } + assert all(report.status is ReportStatus.PENDING for report in reports) + + +def _unreadable_succeeded_report(country_id: str = "us"): + simulation = FakeEconomyReportSimulation( + report_record( + country_id=country_id, + run=report_run_record(status=ReportStatus.SUCCEEDED), + ) + ) + simulation.unreadable_results.add(REPORT_ID) + return simulation + + +def test_a_report_whose_result_cannot_be_read_back_can_run_again() -> None: + """Re-running is how such a report becomes readable again.""" + simulation = _unreadable_succeeded_report() + + run = economy_report_service(simulation).run_report( + country_id="us", report_id=REPORT_ID, force=False + ) + + assert run.trigger == "manual" + assert "get_report" not in [name for name, _ in simulation.calls] + + +def test_the_runs_of_a_report_whose_result_cannot_be_read_back_are_listed() -> None: + simulation = _unreadable_succeeded_report() + + runs = economy_report_service(simulation).list_report_runs( + country_id="us", report_id=REPORT_ID + ) + + assert [run.status for run in runs] == [ReportStatus.SUCCEEDED] + assert "get_report" not in [name for name, _ in simulation.calls] + + +def test_a_report_of_another_country_with_runs_is_not_found() -> None: + simulation = FakeEconomyReportSimulation( + report_record(country_id="uk", run=report_run_record()) + ) + + with pytest.raises(EconomyReportNotFoundError): + economy_report_service(simulation).run_report( + country_id="us", report_id=REPORT_ID, force=True + ) + + assert "run_report" not in [name for name, _ in simulation.calls] + + +def test_a_listing_leaves_out_a_report_whose_result_cannot_be_read_back() -> None: + readable_id = UUID("00000000-0000-4000-8000-000000000204") + simulation = _unreadable_succeeded_report() + simulation.reports[readable_id] = report_record(report_id=readable_id) + + reports = economy_report_service(simulation).list_reports( + country_id="us", report_ids=[REPORT_ID, readable_id] + ) + + assert [report.id for report in reports] == [readable_id] diff --git a/tests/unit/v2/test_household_models.py b/tests/unit/v2/test_household_models.py index f0798e35e..bef59c843 100644 --- a/tests/unit/v2/test_household_models.py +++ b/tests/unit/v2/test_household_models.py @@ -18,6 +18,7 @@ User, UserHouseholdAssociation, ) +from tests.fixtures.v2_countries import add_listed_countries def _engine(): @@ -30,6 +31,8 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: cursor.close() V2_METADATA.create_all(engine) + with Session(engine) as session: + add_listed_countries(session) return engine diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index 7b753004d..af6913897 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) == 39 +assert len(V2_METADATA.tables) == 40 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_routes.py b/tests/unit/v2/test_metadata_routes.py index 79ec7c1ed..b9294be3f 100644 --- a/tests/unit/v2/test_metadata_routes.py +++ b/tests/unit/v2/test_metadata_routes.py @@ -538,10 +538,15 @@ def test_openapi_references_explicit_resource_response_schemas() -> None: "/v2/variables/{variable_id}", } native_paths = { + "/v2/countries", "/v2/households", "/v2/households/{household_id}", "/v2/policies", "/v2/policies/{policy_id}", + "/v2/reports", + "/v2/reports/{report_id}", + "/v2/reports/{report_id}/runs", + "/v2/reports/by-legacy-number/{legacy_report_id}", "/v2/user-households", "/v2/user-households/{association_id}", "/v2/user-policies", diff --git a/tests/unit/v2/test_model_persistence.py b/tests/unit/v2/test_model_persistence.py index be7879424..732fa7d98 100644 --- a/tests/unit/v2/test_model_persistence.py +++ b/tests/unit/v2/test_model_persistence.py @@ -10,7 +10,9 @@ from policyengine_api.data.v2.models import ( V2_METADATA, + Country, Dynamic, + Household, LegacyPolicyMapping, LegacyUserMapping, LegacyUserPolicyMapping, @@ -20,9 +22,11 @@ TaxBenefitModel, TaxBenefitModelVersion, User, + UserHouseholdAssociation, UserPolicy, ) from policyengine_api.data.v2.models.base import DIRECT_SQLALCHEMY_EXCEPTIONS +from tests.fixtures.v2_countries import add_listed_countries def _relational_sqlite_engine(): @@ -35,6 +39,8 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: cursor.close() V2_METADATA.create_all(engine) + with Session(engine) as session: + add_listed_countries(session) return engine @@ -99,9 +105,8 @@ def test_direct_sqlalchemy_categories_are_complete_and_documented() -> None: assert all(DIRECT_SQLALCHEMY_EXCEPTIONS.values()) -def test_user_primary_country_can_change_between_us_and_uk_only() -> None: - engine = create_engine("sqlite://") - V2_METADATA.create_all(engine) +def test_user_primary_country_must_be_a_listed_country() -> None: + engine = _relational_sqlite_engine() with Session(engine) as session: user = User( @@ -126,6 +131,73 @@ def test_user_primary_country_can_change_between_us_and_uk_only() -> None: engine.dispose() +def test_a_country_row_lets_records_use_a_household_only_country() -> None: + engine = _relational_sqlite_engine() + household_only_country = Country( + id="be", + household_enabled=True, + society_wide_enabled=False, + ) + + with Session(engine) as session: + session.add(User(primary_country="be")) + with pytest.raises(IntegrityError): + session.commit() + session.rollback() + + session.add(household_only_country) + session.commit() + user = User(primary_country="be") + household = Household( + country_id="be", + household_data={"people": {}}, + canonicalization_version=1, + content_hash="b" * 64, + ) + session.add_all([user, household]) + session.flush() + session.add( + UserHouseholdAssociation( + user_id=user.id, + household_id=household.id, + country_id="be", + ) + ) + session.commit() + + stored_country = session.exec(select(UserHouseholdAssociation.country_id)).one() + assert stored_country == "be" + + engine.dispose() + + +def test_a_listed_country_cannot_be_deleted_while_records_use_it() -> None: + engine = _relational_sqlite_engine() + + with Session(engine) as session: + session.add(User(primary_country="uk")) + session.commit() + + session.delete(session.get(Country, "uk")) + with pytest.raises(IntegrityError): + session.commit() + + engine.dispose() + + +def test_a_society_wide_country_must_also_be_household_capable() -> None: + engine = _relational_sqlite_engine() + + with Session(engine) as session: + session.add( + Country(id="be", household_enabled=False, society_wide_enabled=True) + ) + with pytest.raises(IntegrityError): + session.commit() + + engine.dispose() + + def test_only_unauthenticated_users_may_omit_a_primary_country() -> None: engine = create_engine("sqlite://") V2_METADATA.create_all(engine) diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 20947ffcb..baa0e2590 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -8,6 +8,7 @@ from policyengine_api.data.v1_models import V1Base from policyengine_api.data.v2.models import ( V2_METADATA, + Country, DatasetVersion, Dynamic, Household, @@ -29,6 +30,18 @@ UserType, ) +# Every column that names the country a record belongs to. Each references the +# countries table instead of checking a fixed list of country IDs. +COUNTRY_REFERENCE_COLUMNS = ( + ("users", "primary_country"), + ("households", "country_id"), + ("policies", "country_id"), + ("user_household_associations", "country_id"), + ("user_policies", "country_id"), + ("user_simulation_associations", "country_id"), + ("user_report_associations", "country_id"), +) + RUN_OUTPUT_TABLES = frozenset( { "aggregates", @@ -48,6 +61,7 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: expected_modules = { + Country: "countries", User: "users", UserIdentity: "users", Policy: "policies", @@ -354,8 +368,8 @@ def test_standard_users_require_a_supported_primary_country() -> None: for constraint in users.constraints if isinstance(constraint, sa.CheckConstraint) } + # The countries table, not a check, lists the countries a user may choose. assert checks == { - "ck_users_primary_country": "primary_country IN ('us', 'uk')", "ck_users_user_type": "user_type IN ('unauthenticated', 'standard')", "ck_users_standard_primary_country": ( "user_type = 'unauthenticated' OR primary_country IS NOT NULL" @@ -474,7 +488,8 @@ def test_named_checks_and_required_indexes_cover_core_invariants() -> None: assert { "ck_regions_required_filter_values", "ck_simulations_type_input", - "ck_users_primary_country", + "ck_users_standard_primary_country", + "ck_countries_society_wide_requires_household", "ck_report_runs_terminal_completion", "ck_parameter_values_single_owner", }.issubset(constraint_names) @@ -607,10 +622,11 @@ def test_simulation_and_report_associations_have_no_cross_service_keys() -> None assert table.c.label.nullable assert { foreign_key.target_fullname for foreign_key in simulations.foreign_keys - } == {"users.id"} + } == {"users.id", "countries.id"} assert {foreign_key.target_fullname for foreign_key in reports.foreign_keys} == { "users.id", "user_policies.id", + "countries.id", } user_policy_key = next(iter(reports.c.user_policy_id.foreign_keys)) assert user_policy_key.ondelete == "SET NULL" @@ -635,3 +651,71 @@ def test_user_policy_records_v3_reform_fields() -> None: assert "current-policy" in str(baseline_check.sqltext) postgres_type = policies.c.provenance.type.dialect_impl(postgresql.dialect()) assert isinstance(postgres_type, postgresql.JSONB) + + +def test_countries_record_household_and_society_wide_capability() -> None: + countries = V2_METADATA.tables["countries"] + + assert countries.c.keys() == ["id", "household_enabled", "society_wide_enabled"] + assert countries.primary_key.columns.keys() == ["id"] + assert countries.primary_key.name == "pk_countries" + assert countries.c.id.type.length == 2 + for flag_name in ("household_enabled", "society_wide_enabled"): + assert isinstance(countries.c[flag_name].type, sa.Boolean) + assert not countries.c[flag_name].nullable + assert countries.c[flag_name].server_default is None + checks = { + constraint.name: str(constraint.sqltext) + for constraint in countries.constraints + if isinstance(constraint, sa.CheckConstraint) + } + # Society-wide reports need household calculations, so a society-wide + # country is always household capable too. + assert checks == { + "ck_countries_society_wide_requires_household": ( + "household_enabled OR NOT society_wide_enabled" + ), + } + + +def test_country_columns_reference_the_countries_table_instead_of_a_fixed_list() -> ( + None +): + for table_name, column_name in COUNTRY_REFERENCE_COLUMNS: + table = V2_METADATA.tables[table_name] + country_keys = [ + foreign_key + for foreign_key in table.c[column_name].foreign_keys + if foreign_key.column.table.name == "countries" + ] + assert len(country_keys) == 1, table_name + assert country_keys[0].target_fullname == "countries.id" + assert country_keys[0].ondelete == "RESTRICT" + assert country_keys[0].constraint.name == ( + f"fk_{table_name}_{column_name}_countries" + ) + fixed_country_checks = [ + constraint.name + for constraint in table.constraints + if isinstance(constraint, sa.CheckConstraint) + and "'us', 'uk'" in str(constraint.sqltext) + ] + assert fixed_country_checks == [], table_name + + +def test_v1_copy_mappings_keep_their_us_and_uk_checks() -> None: + # These tables pair v1 records with their v2 copies. Each already + # references a v2 record of the same country, and they are removed with + # dual-write, so they keep their fixed checks. + for table_name in ( + "legacy_household_mappings", + "legacy_policy_mappings", + "legacy_user_policy_mappings", + ): + table = V2_METADATA.tables[table_name] + assert f"ck_{table_name}_country" in { + constraint.name for constraint in table.constraints + } + assert "countries.id" not in { + foreign_key.target_fullname for foreign_key in table.foreign_keys + } diff --git a/tests/unit/v2/test_shared_report_routes.py b/tests/unit/v2/test_shared_report_routes.py index 2bf23f22a..8832fa3d4 100644 --- a/tests/unit/v2/test_shared_report_routes.py +++ b/tests/unit/v2/test_shared_report_routes.py @@ -23,6 +23,7 @@ REPORT_ASSOCIATION_ID = UUID("00000000-0000-0000-0000-000000000063") SHARED_PATH = f"/v2/shared-reports/{REPORT_ASSOCIATION_ID}" +V2_REPORT_ID = UUID("00000000-0000-4000-8000-000000000101") LAST_RUN_AT = datetime(2026, 2, 1, 12, tzinfo=UTC) PROVISIONS = [ { @@ -39,6 +40,7 @@ def _shared_report_read() -> SharedUserReportRead: return SharedUserReportRead( id=REPORT_ASSOCIATION_ID, country_id="us", + report_id=V2_REPORT_ID, legacy_report_id=615, name="CTC report", source_note="from chat", @@ -76,6 +78,7 @@ def test_anyone_reads_a_saved_report_by_its_link_without_the_owner() -> None: "item": { "id": str(REPORT_ASSOCIATION_ID), "country_id": "us", + "report_id": str(V2_REPORT_ID), "legacy_report_id": 615, "name": "CTC report", "source_note": "from chat", diff --git a/tests/unit/v2/test_simulation_entrypoint_contract.py b/tests/unit/v2/test_simulation_entrypoint_contract.py new file mode 100644 index 000000000..64a719ff7 --- /dev/null +++ b/tests/unit/v2/test_simulation_entrypoint_contract.py @@ -0,0 +1,50 @@ +"""The simulation service contract is shared without the economy report package. + +Saved reports and simulations check v2 references with the simulation +service, and the adapter implements the service's routes; neither depends on +how economy reports are served. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +from policyengine_api.constants import REPO + +ECONOMY_REPORT_PACKAGE = "policyengine_api.services.v2.economy_reports" +INDEPENDENT_MODULES = ( + "policyengine_api/libs/simulation_economy_reports.py", + "policyengine_api/fastapi_routes/v2/user_resources.py", + "policyengine_api/services/v2/associations/types.py", + "policyengine_api/services/v2/associations/validators.py", + "policyengine_api/services/v2/simulation_refs/services.py", + "policyengine_api/services/v2/simulation_refs/validators.py", + "policyengine_api/services/v2/user_reports/services.py", + "policyengine_api/services/v2/user_simulations/services.py", +) + + +def _imported_modules(source_path: Path) -> set[str]: + tree = ast.parse(source_path.read_text()) + imported: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module is not None: + imported.add(node.module) + elif isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + return imported + + +@pytest.mark.parametrize("module_path", INDEPENDENT_MODULES) +def test_module_uses_the_shared_simulation_contract(module_path: str) -> None: + imported = _imported_modules(Path(REPO) / module_path) + + assert not { + module + for module in imported + if module == ECONOMY_REPORT_PACKAGE + or module.startswith(f"{ECONOMY_REPORT_PACKAGE}.") + } diff --git a/tests/unit/v2/test_user_association_v2_references.py b/tests/unit/v2/test_user_association_v2_references.py new file mode 100644 index 000000000..f9bdb176a --- /dev/null +++ b/tests/unit/v2/test_user_association_v2_references.py @@ -0,0 +1,595 @@ +"""Saved reports and simulations that name v2 records, checked with the entrypoint. + +The v2 tables keep no foreign key to the Simulation Entrypoint's reports and +simulations, so each v2 ID is checked with the entrypoint before it is saved. +""" + +from __future__ import annotations + +from uuid import UUID + +import pytest +from pydantic import ValidationError +from sqlmodel import select + +from policyengine_api.data.v2.models import ( + UserReportAssociation, + UserSimulationAssociation, +) +from policyengine_api.services.v2.associations.validators import ( + AssociationConflictError, + AssociationReferenceConflictError, +) +from policyengine_api.services.v2.legacy_refs.validators import ( + LegacyReferenceNotFoundError, +) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + ReportStatus, + SimulationServiceUnavailableError, +) +from policyengine_api.services.v2.simulation_refs.services import ( + SimulationReferenceService, +) +from policyengine_api.services.v2.simulation_refs.validators import ( + SimulationReferenceNotFoundError, +) +from policyengine_api.services.v2.user_reports.database_session import ( + UserReportDatabaseSession, +) +from policyengine_api.services.v2.user_reports.services import V2UserReportService +from policyengine_api.services.v2.user_reports.types import ( + UserReportCreateFields, + UserReportUpdateInput, +) +from policyengine_api.services.v2.user_simulations.database_session import ( + UserSimulationDatabaseSession, +) +from policyengine_api.services.v2.user_simulations.services import ( + V2UserSimulationService, +) +from policyengine_api.services.v2.user_simulations.types import ( + UserSimulationCreateFields, + UserSimulationUpdateInput, +) +from tests.unit.v2.economy_report_support import ( + BASELINE_SIMULATION_ID, + REPORT_ID, + FakeEconomyReportSimulation, + report_record, + report_run_record, + simulation_record, +) +from tests.unit.v2.user_service_support import ( + ALICE, + default_legacy_references, + v2_store, +) + +OTHER_REPORT_ID = UUID("00000000-0000-4000-8000-000000000199") +NOT_MIGRATED_REPORT_ID = UUID("00000000-0000-4000-8000-000000000195") +UK_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000196") + + +@pytest.fixture +def entrypoint() -> FakeEconomyReportSimulation: + return FakeEconomyReportSimulation( + report_record(legacy_report_id=615), + report_record(report_id=NOT_MIGRATED_REPORT_ID), + report_record(report_id=OTHER_REPORT_ID, country_id="uk"), + ) + + +@pytest.fixture +def store(): + with v2_store() as v2: + yield v2 + + +def _service(store, entrypoint) -> V2UserReportService: + return V2UserReportService( + UserReportDatabaseSession(store.sessions), + default_legacy_references(), + SimulationReferenceService(entrypoint), + ) + + +def _saved_reports(store) -> list[UserReportAssociation]: + with store.sessions() as session: + return list(session.exec(select(UserReportAssociation)).all()) + + +def test_a_report_of_the_country_is_saved_by_its_v2_id(store, entrypoint) -> None: + saved = _service(store, entrypoint).create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", report_id=REPORT_ID), + ) + + assert saved.report_id == REPORT_ID + assert saved.legacy_report_id is None + assert entrypoint.calls == [ + ("list_report_runs", REPORT_ID), + ("get_report", REPORT_ID), + ] + + +@pytest.mark.parametrize( + "report_id", + [OTHER_REPORT_ID, UUID("00000000-0000-4000-8000-000000000198")], + ids=["other-country", "unknown"], +) +def test_a_report_the_entrypoint_lacks_in_the_country_is_not_saved( + store, entrypoint, report_id +) -> None: + with pytest.raises(SimulationReferenceNotFoundError) as raised: + _service(store, entrypoint).create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=report_id + ), + ) + + assert str(raised.value) == (f"v2 report {report_id} was not found for country us") + assert _saved_reports(store) == [] + + +def test_both_references_are_checked_before_saving(store, entrypoint) -> None: + service = _service(store, entrypoint) + + saved = service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=REPORT_ID, legacy_report_id=615 + ), + ) + with pytest.raises(LegacyReferenceNotFoundError): + service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=REPORT_ID, legacy_report_id=999 + ), + ) + + assert (saved.report_id, saved.legacy_report_id) == (REPORT_ID, 615) + assert len(_saved_reports(store)) == 1 + + +def test_a_saved_report_is_repointed_to_a_checked_v2_report(store, entrypoint) -> None: + service = _service(store, entrypoint) + saved = service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", legacy_report_id=615), + ) + + repointed = service.patch_user_report( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserReportUpdateInput(report_id=REPORT_ID), + ) + with pytest.raises(SimulationReferenceNotFoundError): + service.patch_user_report( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserReportUpdateInput(report_id=OTHER_REPORT_ID), + ) + + assert (repointed.report_id, repointed.legacy_report_id) == (REPORT_ID, 615) + assert _saved_reports(store)[0].report_id == REPORT_ID + + +def test_a_v2_report_cannot_be_saved_while_the_entrypoint_is_unavailable( + store, entrypoint +) -> None: + entrypoint.error = SimulationServiceUnavailableError("HTTP 503") + + with pytest.raises(SimulationServiceUnavailableError): + _service(store, entrypoint).create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=REPORT_ID + ), + ) + + assert _saved_reports(store) == [] + + +def test_a_v2_report_cannot_be_saved_without_the_entrypoint_check(store) -> None: + service = V2UserReportService( + UserReportDatabaseSession(store.sessions), default_legacy_references() + ) + + with pytest.raises(SimulationServiceUnavailableError): + service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=REPORT_ID + ), + ) + + +def test_a_saved_report_needs_a_report_reference() -> None: + with pytest.raises(ValidationError): + UserReportCreateFields(country_id="us") + with pytest.raises(ValidationError): + UserReportUpdateInput(report_id=None) + + +def test_a_shared_report_shows_its_v2_report(store, entrypoint) -> None: + service = _service(store, entrypoint) + saved = service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", report_id=REPORT_ID), + ) + + shared = service.get_shared_user_report(country_id="us", association_id=saved.id) + + assert shared.report_id == REPORT_ID + + +@pytest.fixture +def simulation_entrypoint() -> FakeEconomyReportSimulation: + entrypoint = FakeEconomyReportSimulation() + entrypoint.simulations = { + BASELINE_SIMULATION_ID: simulation_record( + BASELINE_SIMULATION_ID, legacy_simulation_id=613 + ), + UK_SIMULATION_ID: simulation_record(UK_SIMULATION_ID, country_id="uk"), + } + return entrypoint + + +def _simulation_service(store, entrypoint) -> V2UserSimulationService: + return V2UserSimulationService( + UserSimulationDatabaseSession(store.sessions), + default_legacy_references(), + SimulationReferenceService(entrypoint), + ) + + +def _saved_simulations(store) -> list[UserSimulationAssociation]: + with store.sessions() as session: + return list(session.exec(select(UserSimulationAssociation)).all()) + + +def test_a_simulation_of_the_country_is_saved_by_its_v2_id( + store, simulation_entrypoint +) -> None: + saved = _simulation_service(store, simulation_entrypoint).create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields( + country_id="us", simulation_id=BASELINE_SIMULATION_ID + ), + ) + + assert saved.simulation_id == BASELINE_SIMULATION_ID + assert saved.legacy_simulation_id is None + assert simulation_entrypoint.calls == [("get_simulation", BASELINE_SIMULATION_ID)] + + +@pytest.mark.parametrize( + "simulation_id", + [UK_SIMULATION_ID, UUID("00000000-0000-4000-8000-000000000197")], + ids=["other-country", "unknown"], +) +def test_a_simulation_the_entrypoint_lacks_in_the_country_is_not_saved( + store, simulation_entrypoint, simulation_id +) -> None: + with pytest.raises(SimulationReferenceNotFoundError) as raised: + _simulation_service(store, simulation_entrypoint).create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields( + country_id="us", simulation_id=simulation_id + ), + ) + + assert str(raised.value) == ( + f"v2 simulation {simulation_id} was not found for country us" + ) + assert _saved_simulations(store) == [] + + +def test_a_saved_simulation_is_repointed_to_a_checked_v2_simulation( + store, simulation_entrypoint +) -> None: + service = _simulation_service(store, simulation_entrypoint) + saved = service.create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields( + country_id="us", legacy_simulation_id=613 + ), + ) + + repointed = service.patch_user_simulation( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserSimulationUpdateInput( + simulation_id=BASELINE_SIMULATION_ID + ), + ) + with pytest.raises(SimulationReferenceNotFoundError): + service.patch_user_simulation( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserSimulationUpdateInput(simulation_id=UK_SIMULATION_ID), + ) + + assert (repointed.simulation_id, repointed.legacy_simulation_id) == ( + BASELINE_SIMULATION_ID, + 613, + ) + + +def test_a_saved_simulation_needs_a_simulation_reference() -> None: + with pytest.raises(ValidationError): + UserSimulationCreateFields(country_id="us") + with pytest.raises(ValidationError): + UserSimulationUpdateInput(simulation_id=None) + + +@pytest.mark.parametrize( + ("report_id", "legacy_report_id"), + [(REPORT_ID, 616), (NOT_MIGRATED_REPORT_ID, 615)], + ids=["migrated-from-another", "not-migrated"], +) +def test_both_references_must_name_the_same_report( + store, entrypoint, report_id, legacy_report_id +) -> None: + with pytest.raises(AssociationReferenceConflictError) as raised: + _service(store, entrypoint).create_user_report( + principal=ALICE, + association_input=UserReportCreateFields( + country_id="us", report_id=report_id, legacy_report_id=legacy_report_id + ), + ) + + assert str(raised.value) == ( + f"v2 report {report_id} was not migrated from v1 report {legacy_report_id}" + ) + assert _saved_reports(store) == [] + + +def _saved_report(service, **references): + return service.create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", **references), + ) + + +def _patch(service, saved, **changes): + return service.patch_user_report( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserReportUpdateInput(**changes), + ) + + +def test_a_repointed_reference_must_match_the_stored_one(store, entrypoint) -> None: + service = _service(store, entrypoint) + v1_only = _saved_report(service, legacy_report_id=616) + v2_only = _saved_report(service, report_id=REPORT_ID) + + with pytest.raises(AssociationReferenceConflictError): + _patch(service, v1_only, report_id=REPORT_ID) + with pytest.raises(AssociationReferenceConflictError): + _patch(service, v2_only, legacy_report_id=616) + paired = _patch(service, v2_only, legacy_report_id=615) + + stored = {report.id: report for report in _saved_reports(store)} + assert (stored[v1_only.id].report_id, stored[v1_only.id].legacy_report_id) == ( + None, + 616, + ) + assert (paired.report_id, paired.legacy_report_id) == (REPORT_ID, 615) + + +def test_a_patch_that_keeps_both_references_asks_nothing(store, entrypoint) -> None: + service = _service(store, entrypoint) + saved = _saved_report(service, report_id=REPORT_ID, legacy_report_id=615) + entrypoint.calls.clear() + + renamed = _patch(service, saved, name="Child tax credit") + + assert renamed.name == "Child tax credit" + assert entrypoint.calls == [] + + +class ChangingReferenceChecker: + """Check references, then change the saved report as a concurrent patch would.""" + + def __init__(self, checker, store, association_id) -> None: + self._checker = checker + self._store = store + self._association_id = association_id + + def require_report(self, *, country_id, report_id, legacy_report_id=None) -> None: + self._checker.require_report( + country_id=country_id, + report_id=report_id, + legacy_report_id=legacy_report_id, + ) + with self._store.sessions.begin() as session: + association = session.get(UserReportAssociation, self._association_id) + association.legacy_report_id = 616 + + def require_simulation(self, *, country_id, simulation_id) -> None: + self._checker.require_simulation( + country_id=country_id, simulation_id=simulation_id + ) + + +def test_a_report_changed_while_its_reference_is_checked_is_not_repointed( + store, entrypoint +) -> None: + saved = _saved_report(_service(store, entrypoint), legacy_report_id=615) + service = V2UserReportService( + UserReportDatabaseSession(store.sessions), + default_legacy_references(), + ChangingReferenceChecker( + SimulationReferenceService(entrypoint), store, saved.id + ), + ) + + with pytest.raises(AssociationConflictError): + _patch(service, saved, report_id=REPORT_ID) + + (stored,) = _saved_reports(store) + assert (stored.report_id, stored.legacy_report_id) == (None, 616) + + +def test_a_report_whose_result_cannot_be_read_back_is_saved_by_its_v2_id( + store, +) -> None: + entrypoint = FakeEconomyReportSimulation( + report_record(run=report_run_record(status=ReportStatus.SUCCEEDED)) + ) + entrypoint.unreadable_results.add(REPORT_ID) + + saved = _service(store, entrypoint).create_user_report( + principal=ALICE, + association_input=UserReportCreateFields(country_id="us", report_id=REPORT_ID), + ) + + assert saved.report_id == REPORT_ID + assert "get_report" not in [name for name, _ in entrypoint.calls] + + +MIGRATED_SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000194") + + +@pytest.fixture +def migrated_simulation_entrypoint() -> FakeEconomyReportSimulation: + entrypoint = FakeEconomyReportSimulation() + entrypoint.simulations = { + BASELINE_SIMULATION_ID: simulation_record(BASELINE_SIMULATION_ID), + MIGRATED_SIMULATION_ID: simulation_record( + MIGRATED_SIMULATION_ID, legacy_simulation_id=613 + ), + } + return entrypoint + + +@pytest.mark.parametrize( + ("simulation_id", "legacy_simulation_id"), + [(MIGRATED_SIMULATION_ID, 614), (BASELINE_SIMULATION_ID, 613)], + ids=["migrated-from-another", "not-migrated"], +) +def test_both_simulation_references_must_name_the_same_simulation( + store, migrated_simulation_entrypoint, simulation_id, legacy_simulation_id +) -> None: + with pytest.raises(AssociationReferenceConflictError) as raised: + _simulation_service( + store, migrated_simulation_entrypoint + ).create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields( + country_id="us", + simulation_id=simulation_id, + legacy_simulation_id=legacy_simulation_id, + ), + ) + + assert str(raised.value) == ( + f"v2 simulation {simulation_id} was not migrated from v1 simulation " + f"{legacy_simulation_id}" + ) + assert _saved_simulations(store) == [] + + +def _saved_simulation(service, **references): + return service.create_user_simulation( + principal=ALICE, + association_input=UserSimulationCreateFields(country_id="us", **references), + ) + + +def _patch_simulation(service, saved, **changes): + return service.patch_user_simulation( + principal=ALICE, + country_id="us", + association_id=saved.id, + association_input=UserSimulationUpdateInput(**changes), + ) + + +def test_a_repointed_simulation_reference_must_match_the_stored_one( + store, migrated_simulation_entrypoint +) -> None: + service = _simulation_service(store, migrated_simulation_entrypoint) + v1_only = _saved_simulation(service, legacy_simulation_id=614) + v2_only = _saved_simulation(service, simulation_id=MIGRATED_SIMULATION_ID) + + with pytest.raises(AssociationReferenceConflictError): + _patch_simulation(service, v1_only, simulation_id=MIGRATED_SIMULATION_ID) + with pytest.raises(AssociationReferenceConflictError): + _patch_simulation(service, v2_only, legacy_simulation_id=614) + paired = _patch_simulation(service, v2_only, legacy_simulation_id=613) + migrated_entrypoint_calls = len(migrated_simulation_entrypoint.calls) + renamed = _patch_simulation(service, paired, name="Baseline") + + stored = {row.id: row for row in _saved_simulations(store)} + assert ( + stored[v1_only.id].simulation_id, + stored[v1_only.id].legacy_simulation_id, + ) == (None, 614) + assert (paired.simulation_id, paired.legacy_simulation_id) == ( + MIGRATED_SIMULATION_ID, + 613, + ) + assert renamed.name == "Baseline" + assert len(migrated_simulation_entrypoint.calls) == migrated_entrypoint_calls + + +class ChangingSimulationReferenceChecker: + """Check references, then change the saved simulation as another patch would.""" + + def __init__(self, checker, store, association_id) -> None: + self._checker = checker + self._store = store + self._association_id = association_id + + def require_report(self, *, country_id, report_id, legacy_report_id=None) -> None: + self._checker.require_report( + country_id=country_id, + report_id=report_id, + legacy_report_id=legacy_report_id, + ) + + def require_simulation( + self, *, country_id, simulation_id, legacy_simulation_id=None + ) -> None: + self._checker.require_simulation( + country_id=country_id, + simulation_id=simulation_id, + legacy_simulation_id=legacy_simulation_id, + ) + with self._store.sessions.begin() as session: + association = session.get(UserSimulationAssociation, self._association_id) + association.legacy_simulation_id = 614 + + +def test_a_simulation_changed_while_its_reference_is_checked_is_not_repointed( + store, migrated_simulation_entrypoint +) -> None: + saved = _saved_simulation( + _simulation_service(store, migrated_simulation_entrypoint), + legacy_simulation_id=613, + ) + service = V2UserSimulationService( + UserSimulationDatabaseSession(store.sessions), + default_legacy_references(), + ChangingSimulationReferenceChecker( + SimulationReferenceService(migrated_simulation_entrypoint), + store, + saved.id, + ), + ) + + with pytest.raises(AssociationConflictError): + _patch_simulation(service, saved, simulation_id=MIGRATED_SIMULATION_ID) + + (stored,) = _saved_simulations(store) + assert (stored.simulation_id, stored.legacy_simulation_id) == (None, 614) diff --git a/tests/unit/v2/test_user_simulation_report_routes.py b/tests/unit/v2/test_user_simulation_report_routes.py index 9a5aca120..897f8ae69 100644 --- a/tests/unit/v2/test_user_simulation_report_routes.py +++ b/tests/unit/v2/test_user_simulation_report_routes.py @@ -9,11 +9,20 @@ import pytest from policyengine_api.services.v2.associations.validators import ( + AssociationConflictError, AssociationLinkNotFoundError, + AssociationReferenceConflictError, ) from policyengine_api.services.v2.legacy_refs.validators import ( LegacyReferenceNotFoundError, ) +from policyengine_api.services.v2.simulation_entrypoint.types import ( + SimulationRecordUnreadableError, + SimulationServiceUnavailableError, +) +from policyengine_api.services.v2.simulation_refs.validators import ( + SimulationReferenceNotFoundError, +) from policyengine_api.services.v2.user_reports.types import ( UserReportPage, UserReportRead, @@ -157,7 +166,7 @@ def _clients(): return client, simulations, reports -def test_simulation_routes_require_v1_references_and_scope_to_the_user() -> None: +def test_simulation_routes_take_v1_or_v2_references_and_scope_to_the_user() -> None: client, simulations, _reports = _clients() created = client.post( @@ -189,18 +198,22 @@ def test_simulation_routes_require_v1_references_and_scope_to_the_user() -> None "created_at": "2026-01-01T00:00:00Z", "updated_at": "2026-01-01T00:00:01Z", } - assert v2_reference.status_code == 422 + assert v2_reference.status_code == 201 assert listed.status_code == 200 assert patched.status_code == 200 assert deleted.status_code == 204 assert [call[0] for call in simulations.calls] == [ + "create_user_simulation", "create_user_simulation", "list_user_simulations", "patch_user_simulation", "delete_user_simulation", ] assert all(call[1]["principal"] == PRINCIPAL for call in simulations.calls) - assert simulations.calls[1][1]["legacy_simulation_id"] == 613 + assert simulations.calls[1][1]["association_input"].simulation_id == ( + SIMULATION_ASSOCIATION_ID + ) + assert simulations.calls[2][1]["legacy_simulation_id"] == 613 def test_report_routes_round_trip_v3_fields() -> None: @@ -243,7 +256,7 @@ def test_report_routes_round_trip_v3_fields() -> None: @pytest.mark.parametrize( "changes", [ - {"report_id": str(REPORT_ASSOCIATION_ID)}, + {"report_id": "not-a-uuid"}, {"legacy_report_id": None}, {"user_id": str(OWNER_ID)}, {"last_run_at": "2026-02-01T12:00:00"}, @@ -280,7 +293,17 @@ def test_report_create_rejects_invalid_v3_payloads(changes: dict) -> None: [ (AssociationLinkNotFoundError("user-policy association was not found"), 422), (LegacyReferenceNotFoundError("v1 report 615 was not found"), 422), + (SimulationReferenceNotFoundError("v2 report was not found"), 422), (UserReportNotFoundError("association was not found"), 404), + (SimulationServiceUnavailableError("HTTP 503"), 503), + (SimulationRecordUnreadableError("stored result cannot be read back"), 502), + ( + AssociationReferenceConflictError( + "v2 report was not migrated from v1 report 616" + ), + 400, + ), + (AssociationConflictError("association changed; try again"), 409), ], ) def test_report_failures_map_to_contract_statuses( @@ -457,3 +480,84 @@ def test_stored_snapshots_outside_input_bounds_still_serialize() -> None: assert first["validation"]["calibration"]["geography"] == long_label assert "retired_key" not in first["validation"] assert second["validation"] == VALIDATION + + +V2_REPORT_ID = UUID("00000000-0000-4000-8000-000000000101") + + +def test_a_report_is_saved_and_repointed_by_its_v2_id() -> None: + client, _simulations, reports = _clients() + + created = client.post( + "/v2/user-reports?country_id=us", + json={"country_id": "us", "report_id": str(V2_REPORT_ID)}, + ) + both = client.post( + "/v2/user-reports?country_id=us", + json={ + "country_id": "us", + "report_id": str(V2_REPORT_ID), + "legacy_report_id": 615, + }, + ) + neither = client.post( + "/v2/user-reports?country_id=us", json={"country_id": "us", "name": "x"} + ) + repointed = client.patch( + f"/v2/user-reports/{REPORT_ASSOCIATION_ID}?country_id=us", + json={"report_id": str(V2_REPORT_ID)}, + ) + cleared = client.patch( + f"/v2/user-reports/{REPORT_ASSOCIATION_ID}?country_id=us", + json={"report_id": None}, + ) + + assert [response.status_code for response in (created, both, repointed)] == [ + 201, + 201, + 200, + ] + assert neither.status_code == 422 + assert cleared.status_code == 422 + assert [name for name, _ in reports.calls] == [ + "create_user_report", + "create_user_report", + "patch_user_report", + ] + assert reports.calls[0][1]["association_input"].report_id == V2_REPORT_ID + assert reports.calls[0][1]["association_input"].legacy_report_id is None + assert reports.calls[2][1]["association_input"].model_fields_set == {"report_id"} + + +def test_a_simulation_is_repointed_by_its_v2_id_but_never_left_without_one() -> None: + client, simulations, _reports = _clients() + + repointed = client.patch( + f"/v2/user-simulations/{SIMULATION_ASSOCIATION_ID}?country_id=us", + json={"simulation_id": str(V2_REPORT_ID)}, + ) + cleared = client.patch( + f"/v2/user-simulations/{SIMULATION_ASSOCIATION_ID}?country_id=us", + json={"simulation_id": None}, + ) + neither = client.post( + "/v2/user-simulations?country_id=us", json={"country_id": "us"} + ) + + assert repointed.status_code == 200 + assert cleared.status_code == 422 + assert neither.status_code == 422 + assert [name for name, _ in simulations.calls] == ["patch_user_simulation"] + assert simulations.calls[0][1]["association_input"].model_fields_set == { + "simulation_id" + } + + +def test_openapi_declares_the_unreadable_report_answer_where_it_can_occur() -> None: + client, _simulations, _reports = _clients() + + paths = client.get("/v2/openapi.json").json()["paths"] + + assert "502" in paths["/v2/user-reports"]["post"]["responses"] + assert "502" in paths["/v2/user-reports/{association_id}"]["patch"]["responses"] + assert "502" not in paths["/v2/user-reports/{association_id}"]["get"]["responses"] diff --git a/tests/unit/v2/user_service_support.py b/tests/unit/v2/user_service_support.py index ce54237e3..209a63dc5 100644 --- a/tests/unit/v2/user_service_support.py +++ b/tests/unit/v2/user_service_support.py @@ -31,6 +31,7 @@ from policyengine_api.services.v2.legacy_refs.validators import ( require_legacy_references, ) +from tests.fixtures.v2_countries import add_listed_countries if TYPE_CHECKING: from collections.abc import Collection, Iterator, Mapping @@ -128,6 +129,8 @@ def begin_transaction(connection) -> None: V2_METADATA.create_all(engine) sessions = sessionmaker(engine, class_=Session, expire_on_commit=False) + with sessions() as session: + add_listed_countries(session) with sessions.begin() as session: model = TaxBenefitModel(name="policyengine-us") version = TaxBenefitModelVersion(