diff --git a/.env.example b/.env.example index 562132a89..1c6ef9ce3 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,9 @@ # Password for connecting to the PolicyEngine database POLICYENGINE_DB_PASSWORD=policyengine_db_password +# Cloud SQL instance targeted by remote database connections +POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=policyengine-api:us-central1:policyengine-api-data + # Github Microdata Token POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN=policyengine_github_token diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 8342cdc3f..eb1e5e7ab 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -10,5 +10,8 @@ or migration guard changes, read For tests, read `docs/engineering/skills/testing.md` before adding, moving, or reviewing test files. +For SQLAlchemy model or Alembic migration work, read +`docs/engineering/skills/alembic-migrations.md`. + For pull requests, read `docs/engineering/skills/github-prs.md` before opening, replacing, or sharing a PR. diff --git a/.github/scripts/cloud_run_env.sh b/.github/scripts/cloud_run_env.sh index 893934f83..2263b3f23 100755 --- a/.github/scripts/cloud_run_env.sh +++ b/.github/scripts/cloud_run_env.sh @@ -9,7 +9,6 @@ cloud_run_set_defaults() { # image built by the staging track, so it must not embed the service name. CLOUD_RUN_IMAGE_NAME="${CLOUD_RUN_IMAGE_NAME:-policyengine-api}" CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT="${CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT:-policyengine-api-cr-runtime@policyengine-api.iam.gserviceaccount.com}" - CLOUD_RUN_CLOUD_SQL_INSTANCE="${CLOUD_RUN_CLOUD_SQL_INSTANCE:-policyengine-api:us-central1:policyengine-api-data}" CLOUD_RUN_CPU="${CLOUD_RUN_CPU:-4}" CLOUD_RUN_MEMORY="${CLOUD_RUN_MEMORY:-16Gi}" CLOUD_RUN_TIMEOUT="${CLOUD_RUN_TIMEOUT:-300}" @@ -57,7 +56,6 @@ cloud_run_set_defaults() { export CLOUD_RUN_ARTIFACT_REPOSITORY export CLOUD_RUN_IMAGE_NAME export CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT - export CLOUD_RUN_CLOUD_SQL_INSTANCE export CLOUD_RUN_CPU export CLOUD_RUN_MEMORY export CLOUD_RUN_TIMEOUT diff --git a/.github/scripts/create_cloud_sql_backup.sh b/.github/scripts/create_cloud_sql_backup.sh new file mode 100644 index 000000000..44407c234 --- /dev/null +++ b/.github/scripts/create_cloud_sql_backup.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" + +instance_id="${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME##*:}" +description="policyengine-api-v1-alembic-${GITHUB_SHA:-manual}-${GITHUB_RUN_ID:-local}-${GITHUB_RUN_ATTEMPT:-1}" +gcloud sql backups create \ + --project policyengine-api \ + --instance "${instance_id}" \ + --description "${description}" \ + --quiet >&2 + +# `gcloud sql backups create` waits for completion but does not consistently +# emit the created resource with value-format output. Recover the ID from the +# unique workflow description and require the service-reported successful state. +backup_id="$( + gcloud sql backups list \ + --project policyengine-api \ + --instance "${instance_id}" \ + --filter="description=${description} AND status=SUCCESSFUL" \ + --sort-by='~startTime' \ + --limit=1 \ + --format='value(id)' +)" + +if [[ -z "${backup_id}" ]]; then + echo "Cloud SQL did not return a completed backup ID." >&2 + exit 1 +fi + +printf '%s\n' "${backup_id}" diff --git a/.github/scripts/deploy_cloud_run_candidate.sh b/.github/scripts/deploy_cloud_run_candidate.sh index 04d288397..ffbe6dd8d 100755 --- a/.github/scripts/deploy_cloud_run_candidate.sh +++ b/.github/scripts/deploy_cloud_run_candidate.sh @@ -8,7 +8,7 @@ cloud_run_set_defaults bash .github/scripts/validate_cloud_run_deploy_env.sh env_vars=( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=${CLOUD_RUN_CLOUD_SQL_INSTANCE}" + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" "POLICYENGINE_DB_USER=${POLICYENGINE_DB_USER:-policyengine}" "POLICYENGINE_DB_NAME=${POLICYENGINE_DB_NAME:-policyengine}" "GATEWAY_AUTH_REQUIRED=1" @@ -54,7 +54,7 @@ cloud_run_run gcloud run deploy "${CLOUD_RUN_SERVICE}" \ --allow-unauthenticated \ --execution-environment gen2 \ --service-account "${CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT}" \ - --add-cloudsql-instances "${CLOUD_RUN_CLOUD_SQL_INSTANCE}" \ + --add-cloudsql-instances "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" \ --port "${CLOUD_RUN_PORT}" \ --cpu "${CLOUD_RUN_CPU}" \ --cpu-boost \ diff --git a/.github/scripts/migrate_v1_cloud_sql.sh b/.github/scripts/migrate_v1_cloud_sql.sh new file mode 100644 index 000000000..77d2bb23d --- /dev/null +++ b/.github/scripts/migrate_v1_cloud_sql.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" + +readonly_password="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-readonly-password \ + --project policyengine-api +)" +migration_password="$( + gcloud secrets versions access latest \ + --secret policyengine-api-prod-db-migration-password \ + --project policyengine-api +)" + +if [[ -z "${readonly_password}" || -z "${migration_password}" ]]; then + echo "Cloud SQL database credentials must not be empty." >&2 + exit 1 +fi + +printf '::add-mask::%s\n' "${readonly_password}" +printf '::add-mask::%s\n' "${migration_password}" +export POLICYENGINE_DB_READONLY_PASSWORD="${readonly_password}" +export POLICYENGINE_DB_MIGRATION_PASSWORD="${migration_password}" + +# Never allow a job-level URL to bypass the credentials fetched for this release. +unset STAGE7_EXISTING_DATABASE_URL ALEMBIC_DATABASE_URL + +database_state="$(python scripts/v1_database_migration.py --mode state)" +echo "Detected v1 database state: ${database_state}" + +case "${database_state}" in + head) + ;; + pending) + backup_id="$(bash .github/scripts/create_cloud_sql_backup.sh)" + python scripts/v1_database_migration.py \ + --mode upgrade \ + --backup-id "${backup_id}" + ;; + unversioned | invalid) + echo "database is unversioned or has invalid Alembic state; automatic baseline stamping is disabled and manual recovery is required" >&2 + exit 1 + ;; + *) + echo "Unrecognized v1 database state: ${database_state}" >&2 + exit 1 + ;; +esac + +python scripts/v1_database_migration.py --mode verify-head diff --git a/.github/scripts/start_cloud_sql_proxy.sh b/.github/scripts/start_cloud_sql_proxy.sh new file mode 100644 index 000000000..38bf8d64c --- /dev/null +++ b/.github/scripts/start_cloud_sql_proxy.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME:?POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is required}" + +proxy_version="2.25.0" +proxy_sha256="091a9a12eddab6c028b6c563a4f2dacd067e8f7689c25a3fb4afce397e1f0c60" +proxy_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy" +pid_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.pid" +log_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.log" + +curl -fsSL \ + "https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v${proxy_version}/cloud-sql-proxy.linux.amd64" \ + --output "${proxy_path}" +printf '%s %s\n' "${proxy_sha256}" "${proxy_path}" | sha256sum --check --status +chmod +x "${proxy_path}" + +"${proxy_path}" \ + --quota-project policyengine-api \ + --address 127.0.0.1 \ + --port 3307 \ + "${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" \ + >"${log_path}" 2>&1 & +proxy_pid="$!" +printf '%s\n' "${proxy_pid}" >"${pid_path}" + +for _ in $(seq 1 30); do + if ! kill -0 "${proxy_pid}" 2>/dev/null; then + echo "Cloud SQL Auth Proxy exited before becoming ready." >&2 + sed -n '1,120p' "${log_path}" >&2 + exit 1 + fi + if python -c 'import socket; socket.create_connection(("127.0.0.1", 3307), 1).close()' 2>/dev/null; then + exit 0 + fi + sleep 1 +done + +echo "Cloud SQL Auth Proxy did not become ready." >&2 +exit 1 diff --git a/.github/scripts/stop_cloud_sql_proxy.sh b/.github/scripts/stop_cloud_sql_proxy.sh new file mode 100644 index 000000000..89f24b26c --- /dev/null +++ b/.github/scripts/stop_cloud_sql_proxy.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash + +set -euo pipefail + +pid_path="${RUNNER_TEMP:-/tmp}/cloud-sql-proxy.pid" +if [[ ! -f "${pid_path}" ]]; then + exit 0 +fi + +proxy_pid="$(cat "${pid_path}")" +if kill -0 "${proxy_pid}" 2>/dev/null; then + kill "${proxy_pid}" + wait "${proxy_pid}" 2>/dev/null || true +fi diff --git a/.github/scripts/validate_app_engine_deploy_env.sh b/.github/scripts/validate_app_engine_deploy_env.sh index 4e081a4b0..95fdd9a2b 100644 --- a/.github/scripts/validate_app_engine_deploy_env.sh +++ b/.github/scripts/validate_app_engine_deploy_env.sh @@ -9,6 +9,7 @@ selected_url_env="$( )" required=( + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME SIM_ENTRYPOINT "${selected_url_env}" GATEWAY_AUTH_ISSUER diff --git a/.github/scripts/validate_cloud_run_deploy_env.sh b/.github/scripts/validate_cloud_run_deploy_env.sh index 389d18f7e..aed2784ed 100755 --- a/.github/scripts/validate_cloud_run_deploy_env.sh +++ b/.github/scripts/validate_cloud_run_deploy_env.sh @@ -23,7 +23,7 @@ cloud_run_require_env \ CLOUD_RUN_IMAGE_URI \ CLOUD_RUN_TAG \ CLOUD_RUN_RUNTIME_SERVICE_ACCOUNT \ - CLOUD_RUN_CLOUD_SQL_INSTANCE \ + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME \ CLOUD_RUN_POLICYENGINE_DB_PASSWORD_SECRET \ CLOUD_RUN_GITHUB_MICRODATA_TOKEN_SECRET \ CLOUD_RUN_ANTHROPIC_API_KEY_SECRET \ diff --git a/.github/workflows/alembic-v1-check.yml b/.github/workflows/alembic-v1-check.yml new file mode 100644 index 000000000..6e9094b16 --- /dev/null +++ b/.github/workflows/alembic-v1-check.yml @@ -0,0 +1,49 @@ +name: Alembic v1 checks + +on: + workflow_call: + workflow_dispatch: + +jobs: + mysql-lifecycle: + name: Alembic MySQL lifecycle + runs-on: ubuntu-latest + services: + mysql: + image: mysql:8.4 + # This database exists only for this job, so these test credentials are + # intentionally non-secret and safe to expose in the workflow. + env: + MYSQL_ROOT_PASSWORD: policyengine_test + MYSQL_DATABASE: policyengine_alembic_test + ports: + - 3306:3306 + options: >- + --health-cmd="mysqladmin ping --silent" + --health-interval=5s + --health-timeout=5s + --health-retries=20 + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install dependencies + run: make install + - name: Test Alembic configuration and MySQL lifecycle + run: >- + python -m pytest tests/unit/data/test_alembic_baseline.py + tests/unit/data/test_v1_database_migration.py + tests/integration/test_alembic_mysql_lifecycle.py -q + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test + - name: Require database at all v1 heads + run: python -m alembic -c alembic-v1.ini current --check-heads + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test + - name: Require no ungenerated v1 operations + run: python -m alembic -c alembic-v1.ini check + env: + ALEMBIC_DATABASE_URL: mysql+pymysql://root:policyengine_test@127.0.0.1:3306/policyengine_alembic_test diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 2dd9b5235..83d23f875 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -50,6 +50,10 @@ jobs: python-version: "3.12" - name: Run quality guards run: python scripts/run_quality_guards.py + + alembic-v1-check: + name: Alembic v1 qualification + uses: ./.github/workflows/alembic-v1-check.yml check-changelog: name: Check changelog fragment runs-on: ubuntu-latest diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index f0696abd5..cab5ff0ff 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -14,11 +14,10 @@ concurrency: cancel-in-progress: false jobs: - Lint: + lint: + name: Lint runs-on: ubuntu-latest - if: | - (github.repository == 'PolicyEngine/policyengine-uk') - && (github.event.head_commit.message == 'Update PolicyEngine API') + if: github.repository == 'PolicyEngine/policyengine-api' steps: - name: Checkout repo uses: actions/checkout@v4 @@ -29,6 +28,11 @@ jobs: - name: Format check with ruff run: ruff format --check . + alembic-v1-check: + name: Alembic v1 qualification + if: github.repository == 'PolicyEngine/policyengine-api' + uses: ./.github/workflows/alembic-v1-check.yml + ensure-staging-model-version-aligns-with-sim-api: name: Ensure staging model version aligns with simulation API runs-on: ubuntu-latest @@ -49,6 +53,7 @@ jobs: versioning: name: Update versioning + needs: [lint, alembic-v1-check] if: | (github.repository == 'PolicyEngine/policyengine-api') && !(github.event.head_commit.message == 'Update PolicyEngine API') @@ -86,7 +91,10 @@ jobs: publish-git-tag: name: Publish Git Tag runs-on: ubuntu-latest - needs: ensure-staging-model-version-aligns-with-sim-api + needs: + - ensure-staging-model-version-aligns-with-sim-api + - lint + - alembic-v1-check if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') @@ -100,12 +108,52 @@ jobs: - name: Publish Git Tag run: ".github/publish-git-tag.sh" + migrate-v1-cloud-sql: + name: Upgrade v1 Cloud SQL schema + runs-on: ubuntu-latest + needs: + - ensure-staging-model-version-aligns-with-sim-api + - publish-git-tag + if: | + (github.repository == 'PolicyEngine/policyengine-api') + && (github.event.head_commit.message == 'Update PolicyEngine API') + environment: production-database + permissions: + contents: read + id-token: write + env: + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} + steps: + - name: Checkout repo + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Authenticate to GCP + uses: google-github-actions/auth@v2 + with: + workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} + service_account: ${{ vars.GCP_DB_MIGRATION_SERVICE_ACCOUNT }} + - name: Set up GCloud + uses: google-github-actions/setup-gcloud@v2 + - name: Install dependencies + run: make install + - name: Start Cloud SQL Auth Proxy + run: bash .github/scripts/start_cloud_sql_proxy.sh + - name: Upgrade and verify v1 database + run: bash .github/scripts/migrate_v1_cloud_sql.sh + - name: Stop Cloud SQL Auth Proxy + if: always() + run: bash .github/scripts/stop_cloud_sql_proxy.sh + deploy-staging: name: Deploy staging App Engine version runs-on: ubuntu-latest needs: - ensure-staging-model-version-aligns-with-sim-api - publish-git-tag + - migrate-v1-cloud-sql if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') @@ -113,6 +161,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} permissions: contents: read id-token: write @@ -203,6 +252,8 @@ jobs: APP_ENGINE_VERSION: ${{ steps.version.outputs.version }} - name: Wait for staging version health run: bash .github/scripts/health_check.sh "${{ steps.version_url.outputs.url }}/readiness-check" + env: + HEALTH_CHECK_TIMEOUT_SECONDS: "1800" deploy-cloud-run-staging: name: Deploy staging Cloud Run candidate @@ -210,6 +261,7 @@ jobs: needs: - ensure-staging-model-version-aligns-with-sim-api - publish-git-tag + - migrate-v1-cloud-sql if: | (github.repository == 'PolicyEngine/policyengine-api') && (github.event.head_commit.message == 'Update PolicyEngine API') @@ -217,6 +269,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} ROUTE_IMPL_HEALTH: ${{ vars.ROUTE_IMPL_HEALTH }} ROUTE_IMPL_SPECIFICATION: ${{ vars.ROUTE_IMPL_SPECIFICATION }} ROUTE_IMPL_METADATA: ${{ vars.ROUTE_IMPL_METADATA }} @@ -451,6 +504,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} permissions: contents: read id-token: write @@ -511,6 +565,8 @@ jobs: APP_ENGINE_VERSION: ${{ steps.version.outputs.version }} - name: Wait for production version health run: bash .github/scripts/health_check.sh "${{ steps.version_url.outputs.url }}/readiness-check" + env: + HEALTH_CHECK_TIMEOUT_SECONDS: "1800" promote-production: name: Promote production App Engine candidate @@ -580,6 +636,7 @@ jobs: env: SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }} ROUTE_IMPL_HEALTH: ${{ vars.ROUTE_IMPL_HEALTH }} ROUTE_IMPL_SPECIFICATION: ${{ vars.ROUTE_IMPL_SPECIFICATION }} ROUTE_IMPL_METADATA: ${{ vars.ROUTE_IMPL_METADATA }} diff --git a/AGENTS.md b/AGENTS.md index 3113e7459..3837d2f22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,6 +15,10 @@ cutover plans, generated migration docs, or migration guard scripts, read When adding, moving, or reviewing tests, read `docs/engineering/skills/testing.md`. +When changing SQLAlchemy models, Alembic configuration, database schemas, or +migration revisions, read +`docs/engineering/skills/alembic-migrations.md`. + ## GitHub PRs Read `docs/engineering/skills/github-prs.md` before opening, replacing, or diff --git a/CLAUDE.md b/CLAUDE.md index bb654f145..6627f4b40 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -23,6 +23,10 @@ generated migration docs, or migration guard scripts, read When adding, moving, or reviewing tests, read `docs/engineering/skills/testing.md`. +When changing SQLAlchemy models, Alembic configuration, database schemas, or +migration revisions, read +`docs/engineering/skills/alembic-migrations.md`. + ## Safety Boundaries Do not claim a route, database table, compute path, or deployment surface has diff --git a/README.md b/README.md index 724a0e408..925405cd9 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,7 @@ make setup-env `make setup-env` creates a local `.env` from `.env.example`. At minimum, local development expects values for: - `POLICYENGINE_DB_PASSWORD` +- `POLICYENGINE_DB_INSTANCE_CONNECTION_NAME` - `POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN` - `ANTHROPIC_API_KEY` - `OPENAI_API_KEY` diff --git a/alembic-v1.ini b/alembic-v1.ini new file mode 100644 index 000000000..3b9f297d1 --- /dev/null +++ b/alembic-v1.ini @@ -0,0 +1,37 @@ +[alembic] +script_location = migrations/v1 +prepend_sys_path = . +path_separator = os + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/changelog.d/3788.changed.md b/changelog.d/3788.changed.md new file mode 100644 index 000000000..e86a42214 --- /dev/null +++ b/changelog.d/3788.changed.md @@ -0,0 +1 @@ +Migrate API v1 persistence to service-owned SQLAlchemy 2 sessions, direct mapped models, and a dedicated Cloud SQL/MySQL Alembic chain while preserving public API contracts. Add conditional pull-request migration qualification, mandatory release qualification, backup-before-DDL, and a fail-closed pre-deployment migration gate. Generated migrations remove the orphaned `question` prototype table, keep tracer storage local-only, require reform-impact execution IDs, and require callers to provide the reform-impact dataset explicitly. diff --git a/dashboard/app.py b/dashboard/app.py deleted file mode 100644 index cdfcf33e2..000000000 --- a/dashboard/app.py +++ /dev/null @@ -1,58 +0,0 @@ -from policyengine_api.data.data import database -import streamlit as st - -st.title("PolicyEngine API dashboard") - -# Add a text box that the user can enter a SQL query into, with a submit button and a table showing the results. - -st.subheader("Run a SQL query") - -query = st.text_area("Enter a SQL query", "SELECT * FROM policy LIMIT 10;") -if st.button("Submit"): - try: - results = database.query(query) - st.table(results.fetchall()) - except Exception as e: - st.error(e) - -# Enable the user to look up a policy by ID. - -st.subheader("Look up a policy") - -policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_lookup_text")) -country_id = st.text_input("Enter a country ID", "uk", key="policy_lookup_country") -if st.button("Look up policy", key="policy_lookup"): - try: - results = database.query( - f"SELECT * FROM policy WHERE id IS '{policy_id}' AND country_id IS '{country_id}' LIMIT 10;" - ) - st.table(results.fetchall()) - except Exception as e: - st.error(e) - -# Enable the user to set the label of a policy. - -st.subheader("Set a policy's label") - -policy_id = int(st.text_input("Enter a policy ID", "1")) -country_id = st.text_input("Enter a country ID", "uk") -new_label = st.text_input("Enter a new label", "New label", key="policy_label_text") -if st.button("Set policy label", key="policy_label"): - try: - database.set_policy_label(policy_id, country_id, new_label) - st.success("Success!") - except Exception as e: - st.error(e) - -# Enable the user to delete a policy. - -st.subheader("Delete a policy") - -policy_id = int(st.text_input("Enter a policy ID", "1", key="policy_delete_text")) -country_id = st.text_input("Enter a country ID", "uk", key="policy_delete_country") -if st.button("Delete policy", key="policy_delete"): - try: - database.delete_policy(policy_id, country_id) - st.success("Success!") - except Exception as e: - st.error(e) diff --git a/docs/engineering/skills/README.md b/docs/engineering/skills/README.md index 2f3fea74b..a3248f5b2 100644 --- a/docs/engineering/skills/README.md +++ b/docs/engineering/skills/README.md @@ -9,6 +9,8 @@ first, then keep adapters thin. Current skills: +- `alembic-migrations.md`: mandatory autogenerated Alembic revision workflow, + deployment safeguards, and migration validation. - `github-prs.md`: PR workflow and migration PR handoff expectations. - `migration_contracts.md`: API v2 migration route contracts, route-group metadata, generated migration artifacts, and quality guards. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md new file mode 100644 index 000000000..ea6891559 --- /dev/null +++ b/docs/engineering/skills/alembic-migrations.md @@ -0,0 +1,116 @@ +# Alembic Migrations + +Use this skill whenever adding or changing SQLAlchemy models, Alembic +configuration, database schema, or migration revisions. + +## Mandatory generation rule + +AI **MUST NOT manually author Alembic revision scripts**. Generate every schema +revision from reviewed SQLAlchemy metadata: + +```bash +uv run alembic -c alembic-v1.ini revision --autogenerate -m "" +``` + +The mandatory generation operation is `alembic revision --autogenerate`; the +explicit v1 configuration keeps these revisions in the Cloud SQL/MySQL chain. + +If generated operations are wrong, first correct the model metadata and +regenerate. AI may make minimal post-generation corrections only for dialect compatibility, +reversibility, or an objectively incorrect autogenerated +operation. Document each correction and cover it with upgrade and downgrade +tests. Do not add arbitrary schema or data operations to a generated revision. + +If the required migration cannot be expressed safely by autogeneration plus +those narrow review corrections, stop and request a human migration decision. + +## Required checks + +Before committing a migration: + +1. Run `uv run alembic -c alembic-v1.ini check` and review the generated + operations. +2. Upgrade a fresh database to `head`. +3. Compare the deployed database to the ORM metadata before release. +4. Downgrade one revision and upgrade to `head` again in an isolated database. +5. Confirm application startup performs no implicit DDL. + +Never print database credentials or embed them in Alembic configuration. The +deployed v1 database is already versioned; never recreate or restamp its +baseline. Missing or invalid revision history must fail closed for manual +recovery. + +## API v1 database targets + +Alembic manages the API v1 MySQL schema only. The local SQLite database is a +temporary application cache and is deliberately outside Alembic's lifecycle. +Every Alembic command must therefore receive an explicit MySQL URL: + +```bash +export ALEMBIC_DATABASE_URL="mysql+pymysql://..." +uv run alembic -c alembic-v1.ini upgrade head +``` + +The v1 configuration is `alembic-v1.ini`; its revision tree is `migrations/v1`. +Do not put a database URL in the configuration. Keeping it without a default +prevents an omitted environment variable from silently migrating the wrong +database. + +## Separate v1 and v2 migration domains + +The current chain manages API v1 Cloud SQL/MySQL only. Stage 8 must introduce a +different configuration and a separate revision chain for the v2 +Supabase/Postgres schema. Never point `alembic-v1.ini` at Supabase, append v2 +Postgres revisions under `migrations/v1`, or use both Alembic and Supabase CLI +schema migrations as authorities for the same tables. + +### Fresh database qualification + +Use a disposable schema named `policyengine_alembic_test` on local MySQL, then +run the lifecycle integration test: + +```bash +ALEMBIC_DATABASE_URL="mysql+pymysql://.../policyengine_alembic_test" \ + uv run pytest tests/integration/test_alembic_mysql_lifecycle.py -q +``` + +The test must upgrade from an empty schema, run `alembic check`, compare the +live schema with the ORM metadata, downgrade one revision, and upgrade to +`head` again. Its host and schema-name checks prevent it from running against a +shared or deployed database. + +### CI/CD behavior + +- Pull requests first detect whether the v1 Alembic surface changed. Relevant + PRs invoke `.github/workflows/alembic-v1-check.yml`; unrelated PRs skip it. +- Every push to `master` runs the reusable disposable-MySQL qualification next + to lint, regardless of which paths changed. +- The release migration job runs before either staging deployment. It refuses + unversioned or invalid history, takes a backup only when revisions are + pending, upgrades under the advisory lock, and requires zero post-migration + drift. +- Database credentials and derived URLs must never be printed. The application + runtime account must not be granted schema-migration privileges. +- Production downgrades are never automatic. MySQL DDL is non-transactional, so + rollback means application compatibility plus a reviewed forward fix or a + database restore. + +### Production-only table cleanup + +Revision `01e49b3a056e` was autogenerated after reflecting the orphaned +production `question` table. Its only post-generation correction is +`if_exists=True` on the generated drop: deployed databases contain the table, +while fresh databases built from the baseline do not. Integration tests cover +both paths and verify the generated downgrade recreates the table structure. +The nine removed prototype rows are not recreated by downgrade, so retain a +database backup if their contents may be needed later. + +### Local-only tracer cleanup + +Revision `17bb32415f97` was autogenerated after moving the local SQLite +`tracers` cache table out of production `V1Base` metadata and aligning +`reform_impact.execution_id` with its required application contract. Its only +post-generation correction is `if_exists=True` on the generated tracer drop: +fresh databases built from the original baseline contain the table, while +deployed MySQL databases do not. Integration tests cover both fresh and +production-shaped schemas and verify that `execution_id` becomes non-nullable. diff --git a/gcp/cloud_run/start.sh b/gcp/cloud_run/start.sh index 31e2b188b..45a704d8a 100755 --- a/gcp/cloud_run/start.sh +++ b/gcp/cloud_run/start.sh @@ -61,7 +61,7 @@ done # gunicorn's master binds the listen socket before forking workers, so the # Cloud Run TCP startup probe passes immediately instead of racing the # multi-minute app import (which happens in the worker, post-fork, because -# --preload is NOT set). --timeout 0 is required: a worker mid-import does +# application preloading is disabled). --timeout 0 is required: a worker mid-import does # not heartbeat, and the default 30s watchdog would kill it before boot. gunicorn policyengine_api.asgi:app \ --worker-class uvicorn.workers.UvicornWorker \ diff --git a/gcp/export.py b/gcp/export.py index a56bcbf39..5a726e7d5 100644 --- a/gcp/export.py +++ b/gcp/export.py @@ -1,6 +1,9 @@ import os DB_PD = os.environ["POLICYENGINE_DB_PASSWORD"] +POLICYENGINE_DB_INSTANCE_CONNECTION_NAME = os.environ[ + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" +] GITHUB_MICRODATA_TOKEN = os.environ["POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN"] ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"] OPENAI_API_KEY = os.environ["OPENAI_API_KEY"] @@ -35,6 +38,15 @@ with open(".dbpw", "w") as f: f.write(DB_PD) +app_config_location = "gcp/policyengine_api/app.yaml" +with open(app_config_location) as f: + app_config = f.read().replace( + ".policyengine_db_instance_connection_name", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME, + ) +with open(app_config_location, "w") as f: + f.write(app_config) + # in gcp/compute_api/Dockerfile, replace .github_microdata_token with the contents of the file for dockerfile_location in [ "gcp/policyengine_api/Dockerfile", diff --git a/gcp/policyengine_api/app.yaml b/gcp/policyengine_api/app.yaml index 67e1af80a..2db2d191a 100644 --- a/gcp/policyengine_api/app.yaml +++ b/gcp/policyengine_api/app.yaml @@ -16,14 +16,16 @@ liveness_check: timeout_sec: 30 # Allow 30 seconds for a response failure_threshold: 5 success_threshold: 2 - initial_delay_sec: 60 # Don't check for first 60 seconds to allow full boot + initial_delay_sec: 1800 # Allow non-preloaded workers to finish importing runtime_config: operating_system: "ubuntu22" runtime_version: "22" +env_variables: + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: ".policyengine_db_instance_connection_name" readiness_check: path: "/readiness-check" check_interval_sec: 30 timeout_sec: 30 failure_threshold: 5 success_threshold: 2 - app_start_timeout_sec: 900 + app_start_timeout_sec: 1800 diff --git a/gcp/policyengine_api/start.sh b/gcp/policyengine_api/start.sh index 92818ba81..3fee8e4fe 100644 --- a/gcp/policyengine_api/start.sh +++ b/gcp/policyengine_api/start.sh @@ -19,7 +19,7 @@ until redis-cli -h "$CACHE_REDIS_HOST" -p "$CACHE_REDIS_PORT" ping >/dev/null 2> done # Start the API -gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 --preload & +gunicorn -b :"$PORT" policyengine_api.api --timeout 900 --workers 5 & # Keep the script running and handle shutdown gracefully trap "pkill -P $$; exit 1" INT TERM diff --git a/migrations/v1/env.py b/migrations/v1/env.py new file mode 100644 index 000000000..8fce659c9 --- /dev/null +++ b/migrations/v1/env.py @@ -0,0 +1,77 @@ +"""Alembic environment for the API v1 schema.""" + +from logging.config import fileConfig +import os + +from alembic import context +from sqlalchemy import engine_from_config, pool +from sqlalchemy.engine import make_url + +from policyengine_api.data.v1_models import V1Base + + +config = context.config +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +provided_connection = config.attributes.get("connection") +if provided_connection is not None: + if provided_connection.dialect.name != "mysql": + raise RuntimeError("Alembic migrations must target a MySQL database") +else: + database_url = os.environ.get("ALEMBIC_DATABASE_URL") or config.get_main_option( + "sqlalchemy.url" + ) + if not database_url: + raise RuntimeError( + "ALEMBIC_DATABASE_URL is required; the local SQLite cache is not an " + "Alembic migration target" + ) + if make_url(database_url).get_backend_name() != "mysql": + raise RuntimeError("Alembic migrations must target a MySQL database") + config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) + +target_metadata = V1Base.metadata + + +def run_migrations_offline() -> None: + context.configure( + url=config.get_main_option("sqlalchemy.url"), + target_metadata=target_metadata, + literal_binds=True, + compare_type=True, + compare_server_default=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + if provided_connection is not None: + _run_migrations_with_connection(provided_connection) + return + + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + with connectable.connect() as connection: + _run_migrations_with_connection(connection) + + +def _run_migrations_with_connection(connection) -> None: + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + compare_server_default=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/migrations/v1/script.py.mako b/migrations/v1/script.py.mako new file mode 100644 index 000000000..1ba49a84b --- /dev/null +++ b/migrations/v1/script.py.mako @@ -0,0 +1,25 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/tests/unit/endpoints/__init__.py b/migrations/v1/versions/.gitkeep similarity index 100% rename from tests/unit/endpoints/__init__.py rename to migrations/v1/versions/.gitkeep diff --git a/migrations/v1/versions/01e49b3a056e_remove_orphaned_question_table.py b/migrations/v1/versions/01e49b3a056e_remove_orphaned_question_table.py new file mode 100644 index 000000000..e0869cff7 --- /dev/null +++ b/migrations/v1/versions/01e49b3a056e_remove_orphaned_question_table.py @@ -0,0 +1,43 @@ +"""remove orphaned question table + +Revision ID: 01e49b3a056e +Revises: eafc2a547a4e +Create Date: 2026-08-11 16:46:20.938200 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "01e49b3a056e" +down_revision: Union[str, None] = "eafc2a547a4e" +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! ### + # Production has this orphaned prototype table; fresh baseline schemas do not. + op.drop_table("question", if_exists=True) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "question", + sa.Column("question_id", mysql.INTEGER(), autoincrement=True, nullable=False), + sa.Column("question", mysql.LONGTEXT(), nullable=False), + sa.Column("answer", mysql.LONGTEXT(), nullable=True), + sa.Column("policy_id", mysql.INTEGER(), autoincrement=False, nullable=True), + sa.Column("country_id", mysql.VARCHAR(length=3), nullable=False), + sa.Column("subtask", mysql.VARCHAR(length=32), nullable=False), + sa.Column("status", mysql.VARCHAR(length=32), nullable=False), + sa.PrimaryKeyConstraint("question_id"), + mysql_collate="utf8mb4_0900_ai_ci", + mysql_default_charset="utf8mb4", + mysql_engine="InnoDB", + ) + # ### end Alembic commands ### diff --git a/migrations/v1/versions/17bb32415f97_separate_local_tracer_and_require_.py b/migrations/v1/versions/17bb32415f97_separate_local_tracer_and_require_.py new file mode 100644 index 000000000..f75a69049 --- /dev/null +++ b/migrations/v1/versions/17bb32415f97_separate_local_tracer_and_require_.py @@ -0,0 +1,54 @@ +"""separate local tracer and require execution id + +Revision ID: 17bb32415f97 +Revises: 01e49b3a056e +Create Date: 2026-08-11 18:17:19.903942 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "17bb32415f97" +down_revision: Union[str, None] = "01e49b3a056e" +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! ### + # Production never had this local-only table; fresh baseline schemas do. + op.drop_table("tracers", if_exists=True) + op.alter_column( + "reform_impact", + "execution_id", + existing_type=mysql.VARCHAR(length=255), + nullable=False, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "execution_id", + existing_type=mysql.VARCHAR(length=255), + nullable=True, + ) + op.create_table( + "tracers", + sa.Column("id", mysql.INTEGER(), autoincrement=True, nullable=False), + sa.Column("household_id", mysql.INTEGER(), autoincrement=False, nullable=False), + sa.Column("policy_id", mysql.INTEGER(), autoincrement=False, nullable=False), + sa.Column("country_id", mysql.VARCHAR(length=3), nullable=False), + sa.Column("api_version", mysql.VARCHAR(length=10), nullable=False), + sa.Column("tracer_output", mysql.JSON(), nullable=False), + sa.PrimaryKeyConstraint("id"), + mysql_collate="utf8mb4_0900_ai_ci", + mysql_default_charset="utf8mb4", + mysql_engine="InnoDB", + ) + # ### end Alembic commands ### diff --git a/migrations/v1/versions/1914c0422236_require_explicit_reform_impact_dataset.py b/migrations/v1/versions/1914c0422236_require_explicit_reform_impact_dataset.py new file mode 100644 index 000000000..49f1e3f98 --- /dev/null +++ b/migrations/v1/versions/1914c0422236_require_explicit_reform_impact_dataset.py @@ -0,0 +1,41 @@ +"""require explicit reform impact dataset + +Revision ID: 1914c0422236 +Revises: 17bb32415f97 +Create Date: 2026-08-11 19:01:40.778647 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "1914c0422236" +down_revision: Union[str, None] = "17bb32415f97" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "dataset", + existing_type=mysql.VARCHAR(length=255), + server_default=None, + existing_nullable=False, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "reform_impact", + "dataset", + existing_type=mysql.VARCHAR(length=255), + server_default=sa.text("'default'"), + existing_nullable=False, + ) + # ### end Alembic commands ### diff --git a/migrations/v1/versions/eafc2a547a4e_baseline_existing_v1_schema.py b/migrations/v1/versions/eafc2a547a4e_baseline_existing_v1_schema.py new file mode 100644 index 000000000..3a744c562 --- /dev/null +++ b/migrations/v1/versions/eafc2a547a4e_baseline_existing_v1_schema.py @@ -0,0 +1,280 @@ +"""baseline existing v1 schema + +Revision ID: eafc2a547a4e +Revises: +Create Date: 2026-08-06 22:40:07.810466 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +revision: str = "eafc2a547a4e" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "analysis", + sa.Column("prompt_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column( + "prompt", sa.Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=False + ), + sa.Column( + "analysis", sa.Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=True + ), + sa.Column("status", sa.String(length=32), nullable=False), + sa.PrimaryKeyConstraint("prompt_id"), + ) + op.create_table( + "computed_household", + sa.Column("household_id", sa.Integer(), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("computed_household_json", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=True), + sa.PrimaryKeyConstraint("household_id", "policy_id", "country_id"), + ) + op.create_table( + "economy", + sa.Column("economy_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("region", sa.String(length=32), nullable=True), + sa.Column("time_period", sa.String(length=32), nullable=True), + sa.Column("options_json", sa.JSON(), nullable=False), + sa.Column("options_hash", sa.String(length=255), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("economy_json", sa.JSON(none_as_null=True), nullable=True), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("message", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("economy_id"), + ) + op.create_table( + "household", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("label", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=255), nullable=False), + sa.Column("household_json", sa.JSON(), nullable=False), + sa.Column("household_hash", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "legacy_report_output_aliases", + sa.Column( + "legacy_report_output_id", sa.Integer(), autoincrement=False, nullable=False + ), + sa.Column("canonical_report_output_id", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("legacy_report_output_id"), + ) + op.create_table( + "policy", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("label", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("policy_json", sa.JSON(), nullable=False), + sa.Column("policy_hash", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("id", "country_id", "policy_hash"), + ) + op.create_table( + "reform_impact", + sa.Column("reform_impact_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("baseline_policy_id", sa.Integer(), nullable=False), + sa.Column("reform_policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("region", sa.String(length=32), nullable=False), + sa.Column("dataset", sa.String(length=255), nullable=False), + sa.Column("time_period", sa.String(length=32), nullable=False), + sa.Column("options_json", sa.JSON(none_as_null=True), nullable=True), + sa.Column("options_hash", sa.String(length=255), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("reform_impact_json", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("message", sa.String(length=255), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column("execution_id", sa.String(length=255), nullable=False), + sa.PrimaryKeyConstraint("reform_impact_id"), + ) + op.create_table( + "report_output_runs", + sa.Column("id", sa.CHAR(length=36), nullable=False), + sa.Column("report_output_id", sa.Integer(), nullable=False), + sa.Column("run_sequence", sa.Integer(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("trigger_type", sa.String(length=32), nullable=False), + sa.Column("requested_at", sa.DateTime(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("source_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "report_spec_snapshot_json", sa.JSON(none_as_null=True), nullable=True + ), + sa.Column("country_package_version", sa.String(length=255), nullable=True), + sa.Column("policyengine_version", sa.String(length=255), nullable=True), + sa.Column("data_version", sa.String(length=255), nullable=True), + sa.Column("runtime_app_name", sa.String(length=255), nullable=True), + sa.Column("report_cache_version", sa.String(length=255), nullable=True), + sa.Column("simulation_cache_version", sa.String(length=255), nullable=True), + sa.Column("requested_version_override", sa.String(length=255), nullable=True), + sa.Column("resolved_dataset", sa.String(length=255), nullable=True), + sa.Column("resolved_options_hash", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "report_output_id", "run_sequence", name="report_output_run_sequence_idx" + ), + ) + op.create_table( + "report_outputs", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("simulation_1_id", sa.Integer(), nullable=False), + sa.Column("simulation_2_id", sa.Integer(), nullable=True), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column( + "status", + sa.String(length=32), + server_default=sa.text("'pending'"), + nullable=False, + ), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "year", + sa.String(length=255), + server_default=sa.text("'2025'"), + nullable=True, + ), + sa.Column("report_kind", sa.String(length=64), nullable=True), + sa.Column("report_spec_json", sa.JSON(none_as_null=True), nullable=True), + sa.Column("report_spec_schema_version", sa.Integer(), nullable=True), + sa.Column("report_spec_status", sa.String(length=32), nullable=True), + sa.Column("active_run_id", sa.CHAR(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.CHAR(length=36), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "simulation_runs", + sa.Column("id", sa.CHAR(length=36), nullable=False), + sa.Column("simulation_id", sa.Integer(), nullable=False), + sa.Column("report_output_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "input_position", + sa.Integer().with_variant(mysql.TINYINT(), "mysql"), + nullable=True, + ), + sa.Column("run_sequence", sa.Integer(), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("trigger_type", sa.String(length=32), nullable=False), + sa.Column("requested_at", sa.DateTime(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("source_run_id", sa.CHAR(length=36), nullable=True), + sa.Column( + "simulation_spec_snapshot_json", sa.JSON(none_as_null=True), nullable=True + ), + sa.Column("country_package_version", sa.String(length=255), nullable=True), + sa.Column("policyengine_version", sa.String(length=255), nullable=True), + sa.Column("data_version", sa.String(length=255), nullable=True), + sa.Column("runtime_app_name", sa.String(length=255), nullable=True), + sa.Column("simulation_cache_version", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "simulation_id", "run_sequence", name="simulation_run_sequence_idx" + ), + ) + op.create_table( + "simulations", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("population_id", sa.String(length=255), nullable=False), + sa.Column("population_type", sa.String(length=50), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column( + "status", + sa.String(length=32), + server_default=sa.text("'pending'"), + nullable=False, + ), + sa.Column("output", sa.JSON(none_as_null=True), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("simulation_spec_json", sa.JSON(none_as_null=True), nullable=True), + sa.Column("simulation_spec_schema_version", sa.Integer(), nullable=True), + sa.Column("active_run_id", sa.CHAR(length=36), nullable=True), + sa.Column("latest_successful_run_id", sa.CHAR(length=36), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "tracers", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("household_id", sa.Integer(), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("api_version", sa.String(length=10), nullable=False), + sa.Column("tracer_output", sa.JSON(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "user_policies", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("country_id", sa.String(length=3), nullable=False), + sa.Column("reform_id", sa.Integer(), nullable=False), + sa.Column("reform_label", sa.String(length=255), nullable=True), + sa.Column("baseline_id", sa.Integer(), nullable=False), + sa.Column("baseline_label", sa.String(length=255), nullable=True), + sa.Column("user_id", sa.String(length=255), nullable=False), + sa.Column("year", sa.String(length=32), nullable=False), + sa.Column("geography", sa.String(length=255), nullable=False), + sa.Column("dataset", sa.String(length=255), nullable=True), + sa.Column("number_of_provisions", sa.Integer(), nullable=False), + sa.Column("api_version", sa.String(length=32), nullable=False), + sa.Column("added_date", sa.BigInteger(), nullable=False), + sa.Column("updated_date", sa.BigInteger(), nullable=False), + sa.Column("budgetary_impact", sa.String(length=255), nullable=True), + sa.Column("type", sa.String(length=255), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "user_profiles", + sa.Column("user_id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("auth0_id", sa.String(length=255), nullable=False), + sa.Column("username", sa.String(length=255), nullable=True), + sa.Column("primary_country", sa.String(length=3), nullable=False), + sa.Column("user_since", sa.BigInteger(), nullable=False), + sa.PrimaryKeyConstraint("user_id"), + sa.UniqueConstraint("auth0_id"), + sa.UniqueConstraint("username"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("user_profiles") + op.drop_table("user_policies") + op.drop_table("tracers") + op.drop_table("simulations") + op.drop_table("simulation_runs") + op.drop_table("report_outputs") + op.drop_table("report_output_runs") + op.drop_table("reform_impact") + op.drop_table("policy") + op.drop_table("legacy_report_output_aliases") + op.drop_table("household") + op.drop_table("economy") + op.drop_table("computed_household") + op.drop_table("analysis") + # ### end Alembic commands ### diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 9f052a93f..b0730e0ae 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -1,3 +1,4 @@ +# ruff: noqa: E402 """ This is the main Flask app for the PolicyEngine API. """ @@ -21,8 +22,7 @@ def log_timing(message): log_timing("Flask imports completed") -from flask_caching import Cache -from policyengine_api.utils import make_cache_key +from policyengine_api.extensions import cache from policyengine_api.migration_logging import register_migration_request_logging log_timing("Caching utilities import completed") @@ -33,6 +33,9 @@ def log_timing(message): from policyengine_api.routes.error_routes import error_bp log_timing("Error routes import completed") +from policyengine_api.routes.home_routes import home_bp + +log_timing("Home routes import completed") from policyengine_api.routes.economy_routes import economy_bp log_timing("Economy routes import completed") @@ -59,24 +62,11 @@ def log_timing(message): from policyengine_api.routes.ai_prompt_routes import ai_prompt_bp from policyengine_api.routes.simulation_routes import simulation_bp from policyengine_api.routes.report_output_routes import report_output_bp +from policyengine_api.routes.reform_impact_routes import reform_impact_bp +from policyengine_api.routes.system_routes import system_bp log_timing("Base AI routes import completed") -from .endpoints import ( - get_home, - get_policy_search, - get_household_under_policy, - get_calculate, - set_user_policy, - get_user_policy, - update_user_policy, - get_simulations, -) - -log_timing("Legacy endpoints import completed") - -from policyengine_api.readiness import is_ready - log_timing("Initialising API...") app = application = flask.Flask(__name__) @@ -92,7 +82,7 @@ def log_timing(message): "CACHE_DEFAULT_TIMEOUT": 300, } ) -cache = Cache(app) +cache.init_app(app) log_timing("Caching initialised") CORS(app) @@ -104,7 +94,7 @@ def log_timing(message): app.register_blueprint(error_bp) log_timing("Error routes registered") -app.route("/", methods=["GET"])(get_home) +app.register_blueprint(home_bp) log_timing("Home routes registered") app.register_blueprint(metadata_bp) @@ -117,27 +107,6 @@ def log_timing(message): app.register_blueprint(policy_bp) log_timing("Policy routes registered") -app.route("//policies", methods=["GET"])(get_policy_search) -log_timing("Policy search endpoint registered") - -app.route( - "//household//policy/", - methods=["GET"], -)(get_household_under_policy) -log_timing("Household under policy endpoint registered") - -app.route("//calculate", methods=["POST"])( - cache.cached(make_cache_key=make_cache_key)(get_calculate) -) -log_timing("Calculate endpoint registered") - -app.route("//calculate-full", methods=["POST"])( - cache.cached(make_cache_key=make_cache_key)( - lambda *args, **kwargs: get_calculate(*args, **kwargs, add_missing=True) - ) -) -log_timing("Calculate-full endpoint registered") - # Routes for economy microsimulation app.register_blueprint(economy_bp) log_timing("Economy routes registered") @@ -146,19 +115,10 @@ def log_timing(message): app.register_blueprint(simulation_analysis_bp) log_timing("Simulation analysis routes registered") -app.route("//user-policy", methods=["POST"])(set_user_policy) -log_timing("User policy set endpoint registered") - -app.route("//user-policy", methods=["PUT"])(update_user_policy) -log_timing("User policy update endpoint registered") - -app.route("//user-policy/", methods=["GET"])(get_user_policy) -log_timing("User policy get endpoint registered") - app.register_blueprint(user_profile_bp) log_timing("User profile routes registered") -app.route("/simulations", methods=["GET"])(get_simulations) +app.register_blueprint(reform_impact_bp) log_timing("Simulations endpoint registered") app.register_blueprint(tracer_analysis_bp) @@ -170,42 +130,8 @@ def log_timing(message): app.register_blueprint(simulation_bp) app.register_blueprint(report_output_bp) - - -@app.route("/liveness-check", methods=["GET"]) -def liveness_check(): - return flask.Response("OK", status=200, headers={"Content-Type": "text/plain"}) - - -log_timing("Liveness check endpoint registered") - - -@app.route("/readiness-check", methods=["GET"]) -def readiness_check(): - # 503 until the startup warmup has compiled the simulation machinery - # (policyengine_api.readiness); /liveness-check stays unconditional. - if not is_ready(): - return flask.Response( - "NOT READY", status=503, headers={"Content-Type": "text/plain"} - ) - return flask.Response("OK", status=200, headers={"Content-Type": "text/plain"}) - - -log_timing("Readiness check endpoint registered") - - -from policyengine_api.specification import OPENAPI_SPECIFICATION - -openapi_spec = OPENAPI_SPECIFICATION -log_timing("OpenAPI spec loaded") - - -@app.route("/specification", methods=["GET"]) -def get_specification(): - return flask.jsonify(openapi_spec) - - -log_timing("Specification endpoint registered") +app.register_blueprint(system_bp) +log_timing("System routes registered") log_timing("API initialised.") diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index 901f60b59..8843ad65d 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -6,10 +6,14 @@ from policyengine_api.api import app as flask_app from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.orm import close_v1_engines from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.warmup import run_startup_warmup -app = application = create_asgi_app(flask_app) +app = application = create_asgi_app( + flask_app, + shutdown_callback=close_v1_engines, +) # Warm the simulation machinery before serving (see policyengine_api.warmup). # POLICYENGINE_API_STARTUP_WARMUP=0 skips it. diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index f65893864..c2929d4a6 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -2,6 +2,8 @@ from __future__ import annotations +from collections.abc import Callable +from contextlib import asynccontextmanager import time from a2wsgi import WSGIMiddleware @@ -60,6 +62,7 @@ def create_asgi_app( *, route_settings: RouteImplementationSettings | None = None, dependencies: NativeRouteDependencies | None = None, + shutdown_callback: Callable[[], None] | None = None, ) -> FastAPI: """Create the Stage 2 FastAPI shell around the existing Flask app.""" @@ -68,12 +71,21 @@ def create_asgi_app( if dependencies is None: dependencies = NativeRouteDependencies.defaults() + @asynccontextmanager + async def lifespan(_app: FastAPI): + try: + yield + finally: + if shutdown_callback is not None: + shutdown_callback() + app = FastAPI( title="PolicyEngine API", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None, + lifespan=lifespan, ) # compresslevel 4 instead of starlette's default 9: /us/metadata is ~70MB # raw, and level-9 compression inside the request costs seconds of CPU on diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 6c3acc6b9..42980dc1b 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -2,7 +2,7 @@ import inspect import json from policyengine_core.taxbenefitsystems import TaxBenefitSystem -from typing import Union, Optional +from typing import Union from policyengine_api.utils import get_safe_json from policyengine_core.parameters import ( ParameterNode, @@ -23,11 +23,10 @@ build_congressional_district_metadata, ) -from policyengine_api.data import local_database from policyengine_api.constants import ( - COUNTRY_PACKAGE_VERSIONS, get_bundle_default_dataset_option, ) +from policyengine_api.services.household_calculation_service import CalculationResult class PolicyEngineCountry: @@ -361,9 +360,7 @@ def calculate( self, household: dict, reform: Union[dict, None], - household_id: Optional[int] = None, - policy_id: Optional[int] = None, - ): + ) -> CalculationResult: simulation, system = self._create_simulation(household, reform) household = json.loads(json.dumps(household)) @@ -430,25 +427,11 @@ def calculate( tracer_output = simulation.tracer.computation_log log_lines = tracer_output.lines(aggregate=False, max_depth=10) - log_json = json.dumps(log_lines) - - if household_id is not None and policy_id is not None: - # write to local database - local_database.query( - """ - INSERT INTO tracers (household_id, policy_id, country_id, api_version, tracer_output) - VALUES (?, ?, ?, ?, ?) - """, - ( - household_id, - policy_id, - self.country_id, - COUNTRY_PACKAGE_VERSIONS[self.country_id], - log_json, - ), - ) - return household + return CalculationResult( + household=household, + tracer_output=log_lines, + ) def _create_simulation( self, diff --git a/policyengine_api/data/__init__.py b/policyengine_api/data/__init__.py index 15673afdb..b412d8246 100644 --- a/policyengine_api/data/__init__.py +++ b/policyengine_api/data/__init__.py @@ -1 +1 @@ -from .data import PolicyEngineDatabase, database, local_database +"""SQLAlchemy persistence models, engine configuration, and migrations.""" diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index f5218fd36..6dbac4be1 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -74,7 +74,7 @@ class CongressionalDistrictMetadataItem(BaseModel): Uses Pydantic BaseModel for: - Runtime validation of data integrity - Automatic serialization/deserialization - - Consistency with existing codebase patterns (see policyengine_api/endpoints/economy/compare.py) + - Consistency with existing codebase patterns (see services/economy_comparison.py) - Self-documenting schema with type hints """ diff --git a/policyengine_api/data/data.py b/policyengine_api/data/data.py deleted file mode 100644 index 83f7eb820..000000000 --- a/policyengine_api/data/data.py +++ /dev/null @@ -1,289 +0,0 @@ -import fcntl -import sqlite3 -from policyengine_api.constants import REPO, COUNTRY_PACKAGE_VERSIONS -from policyengine_api.utils import hash_object -from pathlib import Path -from dotenv import load_dotenv -import json -from google.cloud.sql.connector import Connector -import sqlalchemy -import sqlalchemy.exc -import os -import sys - -load_dotenv() - -DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME = ( - "policyengine-api:us-central1:policyengine-api-data" -) -DEFAULT_REMOTE_DB_USER = "policyengine" -DEFAULT_REMOTE_DB_NAME = "policyengine" - - -def get_remote_database_config() -> dict[str, str]: - return { - "instance_connection_name": os.environ.get( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", - DEFAULT_REMOTE_DB_INSTANCE_CONNECTION_NAME, - ), - "db_user": os.environ.get("POLICYENGINE_DB_USER", DEFAULT_REMOTE_DB_USER), - "db_name": os.environ.get("POLICYENGINE_DB_NAME", DEFAULT_REMOTE_DB_NAME), - } - - -class _ResultProxy: - """Lightweight wrapper that eagerly fetches results from a - SQLAlchemy CursorResult so they survive connection closure. - Provides fetchone()/fetchall() with dict-like row access.""" - - def __init__(self, cursor_result): - try: - # Use .mappings() so rows behave like dicts - self._rows = list(cursor_result.mappings()) - except Exception: - # For non-SELECT statements (INSERT/UPDATE/DELETE) - # there are no rows to fetch - self._rows = [] - self._index = 0 - - def fetchone(self): - if self._index < len(self._rows): - row = self._rows[self._index] - self._index += 1 - return row - return None - - def fetchall(self): - remaining = self._rows[self._index :] - self._index = len(self._rows) - return remaining - - -class _TransactionProxy: - """Execute queries against an existing connection inside a transaction.""" - - def __init__(self, connection, local: bool): - self._connection = connection - self._local = local - - def query(self, *query): - if self._local: - cursor = self._connection.cursor() - return cursor.execute(*query) - - query = list(query) - main_query = query[0].replace("?", "%s") - query[0] = main_query - params = query[1] if len(query) > 1 else None - if params is not None: - result = self._connection.exec_driver_sql(main_query, params) - else: - result = self._connection.exec_driver_sql(main_query) - return _ResultProxy(result) - - -class PolicyEngineDatabase: - """ - A wrapper around the database connection. - - It uses the Python package sqlite3. - """ - - household_cache: dict = {} - - @staticmethod - def _dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - - def __init__( - self, - local: bool = False, - initialize: bool = False, - ): - self.local = local - if local: - # Local development uses a sqlite database. - self.db_url = REPO / "policyengine_api" / "data" / "policyengine.db" - # Serialize the exists-check + initialize under an exclusive file - # lock: with multiple gunicorn workers importing concurrently on a - # fresh instance, both can otherwise pass the exists() check and - # race initialize() (seed INSERTs collide -> worker dies at boot), - # or one can observe a created-but-unseeded file and skip - # initialization entirely. - lock_path = str(self.db_url) + ".init.lock" - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file, fcntl.LOCK_EX) - try: - if initialize or not Path(self.db_url).exists(): - self.initialize() - finally: - fcntl.flock(lock_file, fcntl.LOCK_UN) - else: - self._create_pool() - if initialize: - self.initialize() - - def _create_pool(self): - db_config = get_remote_database_config() - self.connector = Connector() - db_pass = os.environ["POLICYENGINE_DB_PASSWORD"] - if db_pass == ".dbpw": - with open(".dbpw") as f: - db_pass = f.read().strip() - conn = self.connector.connect( - instance_connection_string=db_config["instance_connection_name"], - driver="pymysql", - db=db_config["db_name"], - user=db_config["db_user"], - password=db_pass, - ) - self.pool = sqlalchemy.create_engine( - "mysql+pymysql://", - creator=lambda: conn, - ) - - def _close_pool(self): - try: - self.pool.dispose() - self.connector.close() - except Exception: - pass - - def _execute_remote(self, query_args): - """Execute a query against the remote database using - SQLAlchemy v2 connection-based execution.""" - main_query = query_args[0] - params = query_args[1] if len(query_args) > 1 else None - with self.pool.connect() as conn: - if params is not None: - result = conn.exec_driver_sql(main_query, params) - else: - result = conn.exec_driver_sql(main_query) - conn.commit() - # Return a lightweight wrapper that holds - # the fetched results so they survive the - # connection context closing - return _ResultProxy(result) - - def _execute_remote_transaction(self, callback): - with self.pool.connect() as conn: - transaction = conn.begin() - proxy = _TransactionProxy(conn, local=False) - try: - result = callback(proxy) - transaction.commit() - return result - except Exception: - transaction.rollback() - raise - - def query(self, *query): - if self.local: - with sqlite3.connect(self.db_url) as conn: - conn.row_factory = self._dict_factory - cursor = conn.cursor() - return cursor.execute(*query) - else: - query = list(query) - main_query = query[0] - main_query = main_query.replace("?", "%s") - query[0] = main_query - try: - return self._execute_remote(query) - # Except InterfaceError and OperationalError, which are thrown when the connection is lost. - except ( - sqlalchemy.exc.InterfaceError, - sqlalchemy.exc.OperationalError, - ): - try: - self._close_pool() - self._create_pool() - return self._execute_remote(query) - except Exception as e: - raise e - - def transaction(self, callback): - if self.local: - connection = getattr(self, "_connection", None) - owns_connection = connection is None - if owns_connection: - connection = sqlite3.connect(self.db_url) - connection.row_factory = self._dict_factory - try: - connection.execute("BEGIN IMMEDIATE") - proxy = _TransactionProxy(connection, local=True) - result = callback(proxy) - connection.commit() - return result - except Exception: - connection.rollback() - raise - finally: - if owns_connection: - connection.close() - - try: - return self._execute_remote_transaction(callback) - except ( - sqlalchemy.exc.InterfaceError, - sqlalchemy.exc.OperationalError, - ): - self._close_pool() - self._create_pool() - return self._execute_remote_transaction(callback) - - def initialize(self): - """ - Create the database tables. - """ - if self.local: - # If the db_url exists, delete it. - if Path(self.db_url).exists(): - Path(self.db_url).unlink() - # If the db_url doesn't exist, create it. - if not Path(self.db_url).exists(): - Path(self.db_url).touch() - - with open( - REPO - / "policyengine_api" - / "data" - / f"initialise{'_local' if self.local else ''}.sql", - "r", - ) as f: - full_query = f.read() - # Split the query into individual queries. - queries = full_query.split(";") - for query in queries: - print(query, sys.stdout) - # Execute each query. - self.query(query) - - # Insert the UK, US and Canadian 'current law' policies. e.g. the UK policy table must have a row with id=1, country_id="uk", label="Current law", api_version=COUNTRY_PACKAGE_VERSIONS["uk"], policy_json="{}", policy_hash=hash_object({}) - for country_id, policy_id in zip( - COUNTRY_PACKAGE_VERSIONS.keys(), - range(1, 1 + len(COUNTRY_PACKAGE_VERSIONS)), - ): - self.query( - "INSERT INTO policy (id, country_id, label, api_version, policy_json, policy_hash) VALUES (?, ?, ?, ?, ?, ?)", - ( - policy_id, - country_id, - "Current law", - COUNTRY_PACKAGE_VERSIONS[country_id], - json.dumps({}), - hash_object({}), - ), - ) - - -# Determine if app is in debug mode, and if so, do not attempt connection with remote db -if os.environ.get("FLASK_DEBUG") == "1": - database = PolicyEngineDatabase(local=True, initialize=False) -else: - database = PolicyEngineDatabase(local=False, initialize=False) - -local_database = PolicyEngineDatabase(local=True, initialize=False) diff --git a/policyengine_api/data/initialise.sql b/policyengine_api/data/initialise.sql deleted file mode 100644 index 085f31c0b..000000000 --- a/policyengine_api/data/initialise.sql +++ /dev/null @@ -1,193 +0,0 @@ -CREATE TABLE IF NOT EXISTS household ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(255) NOT NULL, - household_json JSON NOT NULL, - household_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS computed_household ( - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - computed_household_json JSON NOT NULL, - status VARCHAR(32), - PRIMARY KEY (household_id, policy_id, country_id) -); - -CREATE TABLE IF NOT EXISTS policy ( - id INTEGER AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - policy_json JSON NOT NULL, - policy_hash VARCHAR(255) NOT NULL, - PRIMARY KEY (id, country_id, policy_hash) -); - -CREATE TABLE IF NOT EXISTS economy ( - economy_id INTEGER PRIMARY KEY AUTO_INCREMENT, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32), - time_period VARCHAR(32), - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - economy_json JSON, - status VARCHAR(32) NOT NULL, - message VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - dataset VARCHAR(255) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON, - options_hash VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME, - execution_id VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS analysis ( - prompt_id INTEGER PRIMARY KEY AUTO_INCREMENT, - prompt LONGTEXT NOT NULL, - analysis LONGTEXT, - status VARCHAR(32) NOT NULL -) - --- The dataset row below was added while the table is in prod; --- we must allow NULL values for this column -CREATE TABLE IF NOT EXISTS user_policies ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - reform_id INTEGER NOT NULL, - reform_label VARCHAR(255), - baseline_id INTEGER NOT NULL, - baseline_label VARCHAR(255), - user_id VARCHAR(255) NOT NULL, - year VARCHAR(32) NOT NULL, - geography VARCHAR(255) NOT NULL, - dataset VARCHAR(255), - number_of_provisions INTEGER NOT NULL, - api_version VARCHAR(32) NOT NULL, - added_date BIGINT NOT NULL, - updated_date BIGINT NOT NULL, - budgetary_impact VARCHAR(255), - type VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS user_profiles ( - user_id INTEGER PRIMARY KEY AUTO_INCREMENT, - auth0_id VARCHAR(255) NOT NULL UNIQUE, - username VARCHAR(255) UNIQUE, - primary_country VARCHAR(3) NOT NULL, - user_since BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS tracers ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - tracer_output JSON NOT NULL -); - -CREATE TABLE IF NOT EXISTS simulations ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - -- VARCHAR(255) to accommodate both household IDs and geography codes - population_id VARCHAR(255) NOT NULL, - population_type VARCHAR(50) NOT NULL, - policy_id INT NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - simulation_spec_json JSON DEFAULT NULL, - simulation_spec_schema_version INT DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_outputs ( - id INTEGER PRIMARY KEY AUTO_INCREMENT, - country_id VARCHAR(3) NOT NULL, - simulation_1_id INT NOT NULL, - simulation_2_id INT DEFAULT NULL, - api_version VARCHAR(10) NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - year VARCHAR(255) DEFAULT '2025', - report_kind VARCHAR(64) DEFAULT NULL, - report_spec_json JSON DEFAULT NULL, - report_spec_schema_version INT DEFAULT NULL, - report_spec_status VARCHAR(32) DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_output_runs ( - id CHAR(36) PRIMARY KEY, - report_output_id INT NOT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - report_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - report_cache_version VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - requested_version_override VARCHAR(255) DEFAULT NULL, - resolved_dataset VARCHAR(255) DEFAULT NULL, - resolved_options_hash VARCHAR(255) DEFAULT NULL, - UNIQUE KEY report_output_run_sequence_idx (report_output_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS simulation_runs ( - id CHAR(36) PRIMARY KEY, - simulation_id INT NOT NULL, - report_output_run_id CHAR(36) DEFAULT NULL, - input_position TINYINT DEFAULT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - simulation_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - UNIQUE KEY simulation_run_sequence_idx (simulation_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS legacy_report_output_aliases ( - legacy_report_output_id INT PRIMARY KEY, - canonical_report_output_id INT NOT NULL -); diff --git a/policyengine_api/data/initialise_local.sql b/policyengine_api/data/initialise_local.sql deleted file mode 100644 index 53a37b4c8..000000000 --- a/policyengine_api/data/initialise_local.sql +++ /dev/null @@ -1,205 +0,0 @@ -DROP TABLE IF EXISTS household; -DROP TABLE IF EXISTS computed_household; -DROP TABLE IF EXISTS policy; -DROP TABLE IF EXISTS economy; -DROP TABLE IF EXISTS reform_impact; -DROP TABLE IF EXISTS analysis; -DROP TABLE IF EXISTS user_policies; -DROP TABLE IF EXISTS tracers; -DROP TABLE IF EXISTS report_output_runs; -DROP TABLE IF EXISTS simulation_runs; -DROP TABLE IF EXISTS legacy_report_output_aliases; - -CREATE TABLE IF NOT EXISTS household ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(255) NOT NULL, - household_json JSONB NOT NULL, - household_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS computed_household ( - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - computed_household_json JSONB NOT NULL, - status VARCHAR(32), - PRIMARY KEY (household_id, policy_id, country_id) -); - -CREATE TABLE IF NOT EXISTS policy ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - label VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - policy_json JSONB NOT NULL, - policy_hash VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS economy ( - economy_id INTEGER PRIMARY KEY, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32), - time_period VARCHAR(32), - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - economy_json JSON, - status VARCHAR(32) NOT NULL, - message VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - dataset VARCHAR(255) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON NOT NULL, - options_hash VARCHAR(255) NOT NULL, - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME NOT NULL, - end_time DATETIME, - execution_id VARCHAR(255) NOT NULL -); - -CREATE TABLE IF NOT EXISTS analysis ( - prompt_id INTEGER PRIMARY KEY, - prompt LONGTEXT NOT NULL, - analysis LONGTEXT, - status VARCHAR(32) NOT NULL -); - --- The dataset row below was added while the table is in prod; --- we must allow NULL values for this column -CREATE TABLE IF NOT EXISTS user_policies ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - reform_id INTEGER NOT NULL, - reform_label VARCHAR(255), - baseline_id INTEGER NOT NULL, - baseline_label VARCHAR(255), - user_id VARCHAR(255) NOT NULL, - year VARCHAR(32) NOT NULL, - geography VARCHAR(255) NOT NULL, - dataset VARCHAR(255), - number_of_provisions INTEGER NOT NULL, - api_version VARCHAR(32) NOT NULL, - added_date BIGINT NOT NULL, - updated_date BIGINT NOT NULL, - budgetary_impact VARCHAR(255), - type VARCHAR(255) -); - -CREATE TABLE IF NOT EXISTS user_profiles ( - user_id INTEGER PRIMARY KEY, - auth0_id VARCHAR(255) NOT NULL UNIQUE, - username VARCHAR(255) UNIQUE, - primary_country VARCHAR(3) NOT NULL, - user_since BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS tracers ( - id INTEGER PRIMARY KEY, - household_id INT NOT NULL, - policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - tracer_output JSON NOT NULL -); - -CREATE TABLE IF NOT EXISTS simulations ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - api_version VARCHAR(10) NOT NULL, - -- VARCHAR(255) to accommodate both household IDs and geography codes - population_id VARCHAR(255) NOT NULL, - population_type VARCHAR(50) NOT NULL, - policy_id INT NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - simulation_spec_json JSON DEFAULT NULL, - simulation_spec_schema_version INT DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_outputs ( - id INTEGER PRIMARY KEY, - country_id VARCHAR(3) NOT NULL, - simulation_1_id INT NOT NULL, - simulation_2_id INT DEFAULT NULL, - api_version VARCHAR(10) NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - year VARCHAR(255) DEFAULT '2025', - report_kind VARCHAR(64) DEFAULT NULL, - report_spec_json JSON DEFAULT NULL, - report_spec_schema_version INT DEFAULT NULL, - report_spec_status VARCHAR(32) DEFAULT NULL, - active_run_id CHAR(36) DEFAULT NULL, - latest_successful_run_id CHAR(36) DEFAULT NULL -); - -CREATE TABLE IF NOT EXISTS report_output_runs ( - id CHAR(36) PRIMARY KEY, - report_output_id INT NOT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - report_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - report_cache_version VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - requested_version_override VARCHAR(255) DEFAULT NULL, - resolved_dataset VARCHAR(255) DEFAULT NULL, - resolved_options_hash VARCHAR(255) DEFAULT NULL, - UNIQUE (report_output_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS simulation_runs ( - id CHAR(36) PRIMARY KEY, - simulation_id INT NOT NULL, - report_output_run_id CHAR(36) DEFAULT NULL, - input_position TINYINT DEFAULT NULL, - run_sequence INT NOT NULL, - status VARCHAR(32) NOT NULL, - output JSON DEFAULT NULL, - error_message TEXT DEFAULT NULL, - trigger_type VARCHAR(32) NOT NULL, - requested_at DATETIME DEFAULT NULL, - started_at DATETIME DEFAULT NULL, - finished_at DATETIME DEFAULT NULL, - source_run_id CHAR(36) DEFAULT NULL, - simulation_spec_snapshot_json JSON DEFAULT NULL, - country_package_version VARCHAR(255) DEFAULT NULL, - policyengine_version VARCHAR(255) DEFAULT NULL, - data_version VARCHAR(255) DEFAULT NULL, - runtime_app_name VARCHAR(255) DEFAULT NULL, - simulation_cache_version VARCHAR(255) DEFAULT NULL, - UNIQUE (simulation_id, run_sequence) -); - -CREATE TABLE IF NOT EXISTS legacy_report_output_aliases ( - legacy_report_output_id INT PRIMARY KEY, - canonical_report_output_id INT NOT NULL -); diff --git a/policyengine_api/data/local_database.py b/policyengine_api/data/local_database.py new file mode 100644 index 000000000..2d29c0210 --- /dev/null +++ b/policyengine_api/data/local_database.py @@ -0,0 +1,45 @@ +"""Temporary SQLAlchemy schema bootstrap for the local SQLite cache. + +The production ``policy`` table has an autoincrementing column inside a +composite primary key. SQLite only autoincrements an ``INTEGER PRIMARY KEY`` +column when it is the sole primary-key column, so the local cache has always +used ``policy.id`` as its database primary key. Keep that one physical-schema +exception explicit while deriving every other local table from the production +ORM metadata. +""" + +from __future__ import annotations + +from sqlalchemy import Column, Engine, Integer, JSON, MetaData, String, Table + +from policyengine_api.data.local_models import LocalV1Base +from policyengine_api.data.v1_models import Policy, V1Base + + +_sqlite_policy_metadata = MetaData() +Table( + Policy.__tablename__, + _sqlite_policy_metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("country_id", String(3), nullable=False), + Column("label", String(255)), + Column("api_version", String(10), nullable=False), + Column("policy_json", JSON, nullable=False), + Column("policy_hash", String(255), nullable=False), +) + + +def create_local_v1_schema(engine: Engine) -> None: + """Create the temporary local schema through SQLAlchemy DDL constructs.""" + + if engine.dialect.name != "sqlite": + raise ValueError("The local v1 schema is only defined for SQLite") + + production_tables = [ + table + for table in V1Base.metadata.sorted_tables + if table is not Policy.__table__ + ] + V1Base.metadata.create_all(engine, tables=production_tables) + LocalV1Base.metadata.create_all(engine) + _sqlite_policy_metadata.create_all(engine) diff --git a/policyengine_api/data/local_models.py b/policyengine_api/data/local_models.py new file mode 100644 index 000000000..35ca26074 --- /dev/null +++ b/policyengine_api/data/local_models.py @@ -0,0 +1,26 @@ +"""Declarative mappings for the temporary local SQLite cache only. + +These tables are deliberately outside the production API v1 metadata and +Alembic lifecycle. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import Integer, JSON, String +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + +class LocalV1Base(DeclarativeBase): + pass + + +class Tracer(LocalV1Base): + __tablename__ = "tracers" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + household_id: Mapped[int] + policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + api_version: Mapped[str] = mapped_column(String(10)) + tracer_output: Mapped[Any] = mapped_column(JSON) diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py new file mode 100644 index 000000000..88f4c9426 --- /dev/null +++ b/policyengine_api/data/orm.py @@ -0,0 +1,179 @@ +"""Canonical SQLAlchemy engine and Session configuration for API v1.""" + +from __future__ import annotations + +import atexit +import fcntl +import os +from pathlib import Path + +from dotenv import load_dotenv +from google.cloud.sql.connector import Connector, IPTypes +from sqlalchemy import Engine, create_engine, select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, REPO +from policyengine_api.data.local_database import create_local_v1_schema +from policyengine_api.data.v1_models import Policy +from policyengine_api.utils import hash_object + + +load_dotenv() + +DEFAULT_REMOTE_DB_USER = "policyengine" +DEFAULT_REMOTE_DB_NAME = "policyengine" +CLOUD_SQL_IP_TYPE = IPTypes.PUBLIC +DATABASE_POOL_RECYCLE_SECONDS = 1800 +DATABASE_POOL_SIZE = 5 +DATABASE_POOL_MAX_OVERFLOW = 2 +DATABASE_POOL_TIMEOUT_SECONDS = 30 +LOCAL_DATABASE_PATH = REPO / "policyengine_api" / "data" / "policyengine.db" + +_v1_engines: dict[bool, Engine] = {} +_v1_session_factories: dict[bool, sessionmaker[Session]] = {} +_cloud_sql_connectors: dict[bool, Connector] = {} + + +def get_remote_database_config() -> dict[str, str]: + return { + "instance_connection_name": os.environ[ + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" + ], + "db_user": os.environ.get("POLICYENGINE_DB_USER", DEFAULT_REMOTE_DB_USER), + "db_name": os.environ.get("POLICYENGINE_DB_NAME", DEFAULT_REMOTE_DB_NAME), + } + + +def build_session_factory(engine: Engine) -> sessionmaker[Session]: + """Return SQLAlchemy's standard configurable Session factory.""" + + return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) + + +def _initialize_local_database(engine: Engine) -> None: + create_local_v1_schema(engine) + + current_law_policies = [ + Policy( + id=policy_id, + country_id=country_id, + label="Current law", + api_version=COUNTRY_PACKAGE_VERSIONS[country_id], + policy_json={}, + policy_hash=hash_object({}), + ) + for policy_id, country_id in enumerate(COUNTRY_PACKAGE_VERSIONS, start=1) + ] + policy_ids = [policy.id for policy in current_law_policies] + with build_session_factory(engine).begin() as session: + existing_policy_ids = set( + session.scalars(select(Policy.id).where(Policy.id.in_(policy_ids))) + ) + session.add_all( + policy + for policy in current_law_policies + if policy.id not in existing_policy_ids + ) + + +# TODO: Remove this local-database initialization pattern and replace the local +# persistence path with a traditional cache. Application imports should +# eventually neither create a database file nor bootstrap a schema. +def _ensure_local_database( + engine: Engine, + database_path: Path = LOCAL_DATABASE_PATH, +) -> None: + lock_path = Path(f"{database_path}.init.lock") + with lock_path.open("w") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + _initialize_local_database(engine) + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) + + +def _build_local_engine() -> Engine: + engine = create_engine(f"sqlite+pysqlite:///{LOCAL_DATABASE_PATH}") + try: + _ensure_local_database(engine) + except Exception: + engine.dispose() + raise + return engine + + +def _database_password() -> str: + password = os.environ["POLICYENGINE_DB_PASSWORD"] + if password == ".dbpw": + return Path(".dbpw").read_text(encoding="utf-8").strip() + return password + + +def _build_remote_engine() -> Engine: + config = get_remote_database_config() + connector = Connector( + ip_type=CLOUD_SQL_IP_TYPE, + refresh_strategy="LAZY", + ) + password = _database_password() + + def get_connection(): + return connector.connect( + instance_connection_string=config["instance_connection_name"], + driver="pymysql", + db=config["db_name"], + user=config["db_user"], + password=password, + ) + + engine = create_engine( + "mysql+pymysql://", + creator=get_connection, + pool_pre_ping=True, + pool_recycle=DATABASE_POOL_RECYCLE_SECONDS, + pool_size=DATABASE_POOL_SIZE, + max_overflow=DATABASE_POOL_MAX_OVERFLOW, + pool_timeout=DATABASE_POOL_TIMEOUT_SECONDS, + ) + _cloud_sql_connectors[False] = connector + return engine + + +def get_v1_engine(*, local: bool = False) -> Engine: + """Return one process-owned SQLAlchemy Engine for the selected runtime.""" + + use_local = local or os.environ.get("FLASK_DEBUG") == "1" + if use_local not in _v1_engines: + _v1_engines[use_local] = ( + _build_local_engine() if use_local else _build_remote_engine() + ) + return _v1_engines[use_local] + + +def get_v1_session_factory(*, local: bool = False) -> sessionmaker[Session]: + """Return one configured Session factory per process-owned v1 Engine.""" + + if local not in _v1_session_factories: + _v1_session_factories[local] = build_session_factory(get_v1_engine(local=local)) + return _v1_session_factories[local] + + +def clear_v1_session_factories() -> None: + """Forget cached factories, primarily after replacing an Engine in tests.""" + + _v1_session_factories.clear() + + +def close_v1_engines() -> None: + """Release process-owned SQLAlchemy pools and Cloud SQL connectors.""" + + clear_v1_session_factories() + for engine in _v1_engines.values(): + engine.dispose() + _v1_engines.clear() + for connector in _cloud_sql_connectors.values(): + connector.close() + _cloud_sql_connectors.clear() + + +atexit.register(close_v1_engines) diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py new file mode 100644 index 000000000..6e848d106 --- /dev/null +++ b/policyengine_api/data/v1_models.py @@ -0,0 +1,252 @@ +"""Declarative mappings for the existing API v1 schema. + +These mappings describe the legacy tables; importing this module never emits +DDL. Alembic is the only schema-management entrypoint. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from sqlalchemy import ( + BigInteger, + CHAR, + DateTime, + Integer, + JSON, + String, + Text, + UniqueConstraint, + text, +) +from sqlalchemy.dialects.mysql import LONGTEXT, TINYINT +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + +class V1Base(DeclarativeBase): + pass + + +class Household(V1Base): + __tablename__ = "household" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + label: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(255)) + household_json: Mapped[Any] = mapped_column(JSON) + household_hash: Mapped[str] = mapped_column(String(255)) + + +class ComputedHousehold(V1Base): + __tablename__ = "computed_household" + household_id: Mapped[int] = mapped_column(Integer, primary_key=True) + policy_id: Mapped[int] = mapped_column(Integer, primary_key=True) + country_id: Mapped[str] = mapped_column(String(3), primary_key=True) + api_version: Mapped[str] = mapped_column(String(10)) + computed_household_json: Mapped[Any] = mapped_column(JSON) + status: Mapped[str | None] = mapped_column(String(32)) + + +class Policy(V1Base): + __tablename__ = "policy" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3), primary_key=True) + label: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + policy_json: Mapped[Any] = mapped_column(JSON) + policy_hash: Mapped[str] = mapped_column(String(255), primary_key=True) + + +class Economy(V1Base): + __tablename__ = "economy" + economy_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) + policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + region: Mapped[str | None] = mapped_column(String(32)) + time_period: Mapped[str | None] = mapped_column(String(32)) + options_json: Mapped[Any] = mapped_column(JSON) + options_hash: Mapped[str] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + economy_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + status: Mapped[str] = mapped_column(String(32)) + message: Mapped[str | None] = mapped_column(String(255)) + + +class ReformImpact(V1Base): + __tablename__ = "reform_impact" + reform_impact_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) + baseline_policy_id: Mapped[int] + reform_policy_id: Mapped[int] + country_id: Mapped[str] = mapped_column(String(3)) + region: Mapped[str] = mapped_column(String(32)) + dataset: Mapped[str] = mapped_column(String(255)) + time_period: Mapped[str] = mapped_column(String(32)) + options_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + options_hash: Mapped[str | None] = mapped_column(String(255)) + api_version: Mapped[str] = mapped_column(String(10)) + reform_impact_json: Mapped[Any] = mapped_column(JSON) + status: Mapped[str] = mapped_column(String(32)) + message: Mapped[str | None] = mapped_column(String(255)) + start_time: Mapped[datetime | None] = mapped_column(DateTime) + end_time: Mapped[datetime | None] = mapped_column(DateTime) + execution_id: Mapped[str] = mapped_column(String(255)) + + +class Analysis(V1Base): + __tablename__ = "analysis" + prompt_id: Mapped[int] = mapped_column( + Integer, primary_key=True, autoincrement=True + ) + prompt: Mapped[str] = mapped_column(Text().with_variant(LONGTEXT(), "mysql")) + analysis: Mapped[str | None] = mapped_column( + Text().with_variant(LONGTEXT(), "mysql") + ) + status: Mapped[str] = mapped_column(String(32)) + + +class UserPolicy(V1Base): + __tablename__ = "user_policies" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + reform_id: Mapped[int] + reform_label: Mapped[str | None] = mapped_column(String(255)) + baseline_id: Mapped[int] + baseline_label: Mapped[str | None] = mapped_column(String(255)) + user_id: Mapped[str] = mapped_column(String(255)) + year: Mapped[str] = mapped_column(String(32)) + geography: Mapped[str] = mapped_column(String(255)) + dataset: Mapped[str | None] = mapped_column(String(255)) + number_of_provisions: Mapped[int] + api_version: Mapped[str] = mapped_column(String(32)) + added_date: Mapped[int] = mapped_column(BigInteger) + updated_date: Mapped[int] = mapped_column(BigInteger) + budgetary_impact: Mapped[str | None] = mapped_column(String(255)) + type: Mapped[str | None] = mapped_column(String(255)) + + +class UserProfile(V1Base): + __tablename__ = "user_profiles" + user_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + auth0_id: Mapped[str] = mapped_column(String(255), unique=True) + username: Mapped[str | None] = mapped_column(String(255), unique=True) + primary_country: Mapped[str] = mapped_column(String(3)) + user_since: Mapped[int] = mapped_column(BigInteger) + + +class Simulation(V1Base): + __tablename__ = "simulations" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + api_version: Mapped[str] = mapped_column(String(10)) + population_id: Mapped[str] = mapped_column(String(255)) + population_type: Mapped[str] = mapped_column(String(50)) + policy_id: Mapped[int] + status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + error_message: Mapped[str | None] = mapped_column(Text) + simulation_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + simulation_spec_schema_version: Mapped[int | None] + active_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + + +class ReportOutput(V1Base): + __tablename__ = "report_outputs" + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + country_id: Mapped[str] = mapped_column(String(3)) + simulation_1_id: Mapped[int] + simulation_2_id: Mapped[int | None] + api_version: Mapped[str] = mapped_column(String(10)) + status: Mapped[str] = mapped_column(String(32), server_default=text("'pending'")) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + error_message: Mapped[str | None] = mapped_column(Text) + year: Mapped[str | None] = mapped_column(String(255), server_default=text("'2025'")) + report_kind: Mapped[str | None] = mapped_column(String(64)) + report_spec_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + report_spec_schema_version: Mapped[int | None] + report_spec_status: Mapped[str | None] = mapped_column(String(32)) + active_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + latest_successful_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + + +class ReportOutputRun(V1Base): + __tablename__ = "report_output_runs" + __table_args__ = ( + UniqueConstraint( + "report_output_id", + "run_sequence", + name="report_output_run_sequence_idx", + ), + ) + id: Mapped[str] = mapped_column(CHAR(36), primary_key=True) + report_output_id: Mapped[int] + run_sequence: Mapped[int] + status: Mapped[str] = mapped_column(String(32)) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + error_message: Mapped[str | None] = mapped_column(Text) + trigger_type: Mapped[str] = mapped_column(String(32)) + requested_at: Mapped[datetime | None] = mapped_column(DateTime) + started_at: Mapped[datetime | None] = mapped_column(DateTime) + finished_at: Mapped[datetime | None] = mapped_column(DateTime) + source_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + report_spec_snapshot_json: Mapped[Any | None] = mapped_column( + JSON(none_as_null=True) + ) + country_package_version: Mapped[str | None] = mapped_column(String(255)) + policyengine_version: Mapped[str | None] = mapped_column(String(255)) + data_version: Mapped[str | None] = mapped_column(String(255)) + runtime_app_name: Mapped[str | None] = mapped_column(String(255)) + report_cache_version: Mapped[str | None] = mapped_column(String(255)) + simulation_cache_version: Mapped[str | None] = mapped_column(String(255)) + requested_version_override: Mapped[str | None] = mapped_column(String(255)) + resolved_dataset: Mapped[str | None] = mapped_column(String(255)) + resolved_options_hash: Mapped[str | None] = mapped_column(String(255)) + + +class SimulationRun(V1Base): + __tablename__ = "simulation_runs" + __table_args__ = ( + UniqueConstraint( + "simulation_id", + "run_sequence", + name="simulation_run_sequence_idx", + ), + ) + id: Mapped[str] = mapped_column(CHAR(36), primary_key=True) + simulation_id: Mapped[int] + report_output_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + input_position: Mapped[int | None] = mapped_column( + Integer().with_variant(TINYINT(), "mysql") + ) + run_sequence: Mapped[int] + status: Mapped[str] = mapped_column(String(32)) + output: Mapped[Any | None] = mapped_column(JSON(none_as_null=True)) + error_message: Mapped[str | None] = mapped_column(Text) + trigger_type: Mapped[str] = mapped_column(String(32)) + requested_at: Mapped[datetime | None] = mapped_column(DateTime) + started_at: Mapped[datetime | None] = mapped_column(DateTime) + finished_at: Mapped[datetime | None] = mapped_column(DateTime) + source_run_id: Mapped[str | None] = mapped_column(CHAR(36)) + simulation_spec_snapshot_json: Mapped[Any | None] = mapped_column( + JSON(none_as_null=True) + ) + country_package_version: Mapped[str | None] = mapped_column(String(255)) + policyengine_version: Mapped[str | None] = mapped_column(String(255)) + data_version: Mapped[str | None] = mapped_column(String(255)) + runtime_app_name: Mapped[str | None] = mapped_column(String(255)) + simulation_cache_version: Mapped[str | None] = mapped_column(String(255)) + + +class LegacyReportOutputAlias(V1Base): + __tablename__ = "legacy_report_output_aliases" + legacy_report_output_id: Mapped[int] = mapped_column( + Integer, + primary_key=True, + autoincrement=False, + ) + canonical_report_output_id: Mapped[int] diff --git a/policyengine_api/endpoints/__init__.py b/policyengine_api/endpoints/__init__.py deleted file mode 100644 index 9639c2317..000000000 --- a/policyengine_api/endpoints/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .home import get_home -from .household import ( - get_household_under_policy, - get_calculate, -) -from .policy import ( - get_policy_search, - set_user_policy, - get_user_policy, - update_user_policy, -) - -from .simulation import get_simulations diff --git a/policyengine_api/endpoints/economy/reform_impact.py b/policyengine_api/endpoints/economy/reform_impact.py deleted file mode 100644 index 42795243d..000000000 --- a/policyengine_api/endpoints/economy/reform_impact.py +++ /dev/null @@ -1,32 +0,0 @@ -from policyengine_api.data import local_database - - -def set_comment_on_job( - comment: str, - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options_hash, -): - query = ( - "UPDATE reform_impact SET message = ? WHERE country_id = ? AND " - "reform_policy_id = ? AND baseline_policy_id = ? AND region = ? AND " - "time_period = ? AND options_hash = ? AND dataset = ?" - ) - - local_database.query( - query, - ( - comment, - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - ), - ) diff --git a/policyengine_api/endpoints/home.py b/policyengine_api/endpoints/home.py deleted file mode 100644 index 33017204c..000000000 --- a/policyengine_api/endpoints/home.py +++ /dev/null @@ -1,7 +0,0 @@ -def get_home() -> str: - """Get the home page of the PolicyEngine API. - - Returns: - str: The home page. - """ - return f"

PolicyEngine households API

Use this API to compute the impact of public policy on individual households.

" diff --git a/policyengine_api/endpoints/household.py b/policyengine_api/endpoints/household.py deleted file mode 100644 index 616da54e7..000000000 --- a/policyengine_api/endpoints/household.py +++ /dev/null @@ -1,315 +0,0 @@ -from policyengine_api.data import database, local_database -import json -from flask import Response, request -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -import logging -from datetime import date -from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs -from policyengine_api.utils.input_validation import ( - find_unrecognized_inputs, - format_unrecognized_inputs_message, -) -from policyengine_api.utils.payload_validators import validate_country -from policyengine_core.errors import SituationParsingError - - -def get_countries(): - from policyengine_api.country import COUNTRIES - - return COUNTRIES - - -def add_yearly_variables(household, country_id, countries=None): - """ - Add yearly variables to a household dict before enqueueing calculation - """ - metadata = (countries or get_countries()).get(country_id).metadata - - variables = metadata["variables"] - entities = metadata["entities"] - household_year = get_household_year(household) - - for variable in variables: - if variables[variable]["definitionPeriod"] in ( - "year", - "month", - "eternity", - ): - entity_plural = entities[variables[variable]["entity"]]["plural"] - if entity_plural in household: - possible_entities = household[entity_plural].keys() - for entity in possible_entities: - if ( - variables[variable]["name"] - not in household[entity_plural][entity] - ): - if variables[variable]["isInputVariable"]: - household[entity_plural][entity][ - variables[variable]["name"] - ] = {household_year: variables[variable]["defaultValue"]} - else: - household[entity_plural][entity][ - variables[variable]["name"] - ] = {household_year: None} - return household - - -def get_invalid_inputs_response(household_json, policy_json, country): - invalid_inputs = find_unrecognized_inputs( - household_json, - policy_json, - country.metadata, - ) - if not invalid_inputs: - return None - - response_body = dict( - status="error", - message=format_unrecognized_inputs_message(invalid_inputs), - result=None, - errors=[invalid_input.to_dict() for invalid_input in invalid_inputs], - ) - return Response( - json.dumps(response_body), - status=400, - mimetype="application/json", - ) - - -def get_household_year(household): - """Given a household dict, get the household's year - - Args: - household (dict): The household itself - """ - - # Set household_year based on current year - household_year = date.today().year - - # Determine if "age" variable present within household and return list of values at it - household_age_list = list( - household.get("people", {}).get("you", {}).get("age", {}).keys() - ) - # If it is, overwrite household_year with the value present - if len(household_age_list) > 0: - household_year = household_age_list[0] - - return household_year - - -@validate_country -def get_household_under_policy(country_id: str, household_id: str, policy_id: str): - """Get a household's output data under a given policy. - - Args: - country_id (str): The country ID. - household_id (str): The household ID. - policy_id (str): The policy ID. - """ - - api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - # Look in computed_households to see if already computed - - row = local_database.query( - "SELECT * FROM computed_household WHERE household_id = ? AND policy_id = ? AND api_version = ?", - (household_id, policy_id, api_version), - ).fetchone() - - if row is not None: - result = dict( - policy_id=row["policy_id"], - household_id=row["household_id"], - country_id=row["country_id"], - api_version=row["api_version"], - computed_household_json=row["computed_household_json"], - status=row["status"], - ) - result["result"] = json.loads(result["computed_household_json"]) - del result["computed_household_json"] - return dict( - status="ok", - message=None, - result=result["result"], - ) - - # Retrieve from the household table - - row = database.query( - "SELECT * FROM household WHERE id = ? AND country_id = ?", - (household_id, country_id), - ).fetchone() - - if row is not None: - household = dict(row) - household["household_json"] = json.loads(household["household_json"]) - else: - response_body = dict( - status="error", - message=f"Household #{household_id} not found.", - ) - return Response( - json.dumps(response_body), - status=404, - mimetype="application/json", - ) - - # Add in any missing yearly variables - household["household_json"] = add_yearly_variables( - household["household_json"], country_id - ) - deprecated_inputs = drop_deprecated_inputs(household["household_json"]) - household["household_json"] = deprecated_inputs.household - - # Retrieve from the policy table - - row = database.query( - "SELECT * FROM policy WHERE id = ? AND country_id = ?", - (policy_id, country_id), - ).fetchone() - - if row is not None: - policy = dict(row) - policy["policy_json"] = json.loads(policy["policy_json"]) - else: - response_body = dict( - status="error", - message=f"Policy #{policy_id} not found.", - ) - return Response( - json.dumps(response_body), - status=404, - mimetype="application/json", - ) - - country = get_countries().get(country_id) - invalid_inputs_response = get_invalid_inputs_response( - household["household_json"], - policy["policy_json"], - country, - ) - if invalid_inputs_response is not None: - return invalid_inputs_response - - try: - result = country.calculate( - household["household_json"], - policy["policy_json"], - household_id, - policy_id, - ) - except Exception as e: - logging.exception(e) - response_body = dict( - status="error", - message=f"Error calculating household #{household_id} under policy #{policy_id}: {e}", - ) - return Response( - json.dumps(response_body), - status=500, - mimetype="application/json", - ) - - # Store the result in the computed_household table - - try: - local_database.query( - "INSERT INTO computed_household (country_id, household_id, policy_id, computed_household_json, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - household_id, - policy_id, - json.dumps(result), - api_version, - ), - ) - except Exception: - # Update the result if it already exists - local_database.query( - "UPDATE computed_household SET computed_household_json = ? WHERE country_id = ? AND household_id = ? AND policy_id = ?", - (json.dumps(result), country_id, household_id, policy_id), - ) - - response_body = dict( - status="ok", - message=None, - result=result, - ) - warning_messages = [w.message for w in deprecated_inputs.warnings] - if warning_messages: - response_body["warnings"] = warning_messages - return response_body - - -@validate_country -def get_calculate(country_id: str, add_missing: bool = False) -> dict: - """Lightweight endpoint for passing in household and policy JSON objects and calculating without storing data. - - Args: - country_id (str): The country ID. - """ - - payload = request.json - household_json = payload.get("household", {}) - policy_json = payload.get("policy", {}) - - if add_missing: - # Add in any missing yearly variables to household_json - household_json = add_yearly_variables(household_json, country_id) - - # Strip deprecated inputs from a copy before the engine runs so - # partners who still pass removed/renamed variables get a warning + - # working response instead of a `VariableNotFoundError` HTTP 500. - deprecated_inputs = drop_deprecated_inputs(household_json) - household_json = deprecated_inputs.household - deprecation_warnings = deprecated_inputs.warnings - - country = get_countries().get(country_id) - invalid_inputs_response = get_invalid_inputs_response( - household_json, - policy_json, - country, - ) - if invalid_inputs_response is not None: - return invalid_inputs_response - - try: - result = country.calculate(household_json, policy_json) - except SituationParsingError as e: - # Malformed household payloads (e.g. a dict where a number belongs) - # are client errors, not server errors — mostly bot traffic. - response_body = dict( - status="error", - message=f"Invalid household payload: {e}", - result=None, - ) - return Response( - json.dumps(response_body), - status=400, - mimetype="application/json", - ) - except Exception as e: - logging.exception(e) - response_body = dict( - status="error", - message=f"Error calculating household under policy: {e}", - ) - return Response( - json.dumps(response_body), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message=None, - result=result, - ) - - warning_messages = [w.message for w in deprecation_warnings] - if warning_messages: - # Serialize to strings on the wire; the structured dataclasses - # stay available for any future caller that wants the fields. - response_body["warnings"] = warning_messages - - return response_body diff --git a/policyengine_api/endpoints/policy.py b/policyengine_api/endpoints/policy.py deleted file mode 100644 index f5d33e938..000000000 --- a/policyengine_api/endpoints/policy.py +++ /dev/null @@ -1,419 +0,0 @@ -from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data import database -from policyengine_api.utils import hash_object -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -import json -from flask import Response, request - - -@validate_country -def get_policy_search(country_id: str) -> dict: - """ - Search for policies for a specified country - - Args: - country_id (str): The country ID. - - Query Parameters: - query (str): Optional search term to filter policies - unique_only (bool): If true, return only unique policy-label combinations - - Returns: - Response: Json response with: - - On success: list of policies with id and label - - On failure: error message and appropriate status code - - Example: - GET /api/policies/us?query=tax&unique_only=true - """ - - query = request.args.get("query", "") - # The "json.loads" default type is added to convert lowercase - # "true" and "false" to Python-friendly bool values - unique_only = request.args.get("unique_only", default=False, type=json.loads) - - try: - results = database.query( - "SELECT id, label, policy_hash FROM policy WHERE country_id = ? AND label LIKE ?", - (country_id, f"%{query}%"), - ) - - results = results.fetchall() - - if not results: - body = dict( - status="error", - message=f"No policies found for country {country_id} for query '{query}", - ) - return Response(json.dumps(body), status=404, mimetype="application/json") - - # If unique_only is true, filter results to only include - # items where everything except ID is unique - if unique_only: - processed_vals = set() - new_results = [] - - # Compare every label and hash to what's contained in processed_vals - # If a label-hash set aren't already in processed_vals, - # add them to new_results - for policy in results: - comparison_vals = policy["label"], policy["policy_hash"] - if comparison_vals not in processed_vals: - new_results.append(policy) - processed_vals.add(comparison_vals) - - # Overwrite results with new_results - results = new_results - - # Format into: [{ id: 1, label: "My policy" }, ...] - policies = [dict(id=result["id"], label=result["label"]) for result in results] - body = dict( - status="ok", - message="Policies found", - result=policies, - ) - return Response(json.dumps(body), status=200, mimetype="application/json") - except Exception as e: - body = dict(status="error", message=f"Internal server error: {e}") - return Response(json.dumps(body), status=500, mimetype="application/json") - - -@validate_country -def set_user_policy(country_id: str) -> dict: - """ - Adds a record (if unique, barring type) to the user_policy table - that defines a particular policy as saved by a user to "their - policies"; this table also contains an optional "type" column that - is currently unused - """ - - payload = request.json - reform_label = payload.pop("reform_label", None) - reform_id = payload.pop("reform_id") - baseline_label = payload.pop("baseline_label", None) - baseline_id = payload.pop("baseline_id") - user_id = payload.pop("user_id") - year = payload.pop("year") - geography = payload.pop("geography") - dataset = payload.pop("dataset", None) - number_of_provisions = payload.pop("number_of_provisions") - api_version = payload.pop("api_version") - added_date = payload.pop("added_date") - updated_date = payload.pop("updated_date") - budgetary_impact = payload.pop("budgetary_impact", None) - type = payload.pop("type", None) - - # The following code is a workaround to the fact that - # SQLite's cursor method does not properly convert - # 'WHERE x = None' to 'WHERE x IS NULL'; though SQLite - # supports searching and setting with 'WHERE x IS y', - # the production MySQL does not, requiring this - - # This workaround should be removed if and when a proper - # ORM package is added to the API, and this package's - # sanitization methods should be utilized instead - nullable_keys = [] - not_null_values = [] - possible_nulls = { - "reform_label": reform_label, - "baseline_label": baseline_label, - "dataset": dataset, - } - - for key, value in possible_nulls.items(): - if not value: - nullable_keys.append(f"{key} IS NULL") - else: - nullable_keys.append(f"{key} = ?") - not_null_values.append(value) - - nullable_key_string = " AND ".join(nullable_keys) - - # When setting a user policy, "unique" records contain - # a unique set of the following pieces of data: - # country_id, reform_id, baseline_id, user_id, year, - # geography, reform_label, baseline_label, dataset; - # added_date, budgetary_impact, updated_date, - # number_of_provisions, and api_version are - # all changeable, and thus do not need - # to be tested; type is not yet implemented - - try: - row = database.query( - f"SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? AND {nullable_key_string}", - ( - country_id, - reform_id, - baseline_id, - user_id, - year, - geography, - *not_null_values, - ), - ).fetchone() - if row is not None: - readable_row = dict(row) - - response = dict( - status="ok", - message=f"The reform #{reform_id} / baseline #{baseline_id} pair already exists for user {user_id}", - result=dict(id=readable_row["id"]), - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - try: - # Unfortunately, it's not possible to use RETURNING - # with SQLite3 without rewriting the PolicyEngineDatabase - # object or implementing a true ORM, thus the double query - - query = ( - "INSERT INTO user_policies (country_id, reform_label, " - "reform_id, baseline_label, baseline_id, user_id, year, " - "geography, number_of_provisions, api_version, added_date, " - "updated_date, budgetary_impact, type, dataset) VALUES " - f"(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) - - database.query( - query, - ( - country_id, - reform_label, - reform_id, - baseline_label, - baseline_id, - user_id, - year, - geography, - number_of_provisions, - api_version, - added_date, - updated_date, - budgetary_impact, - type, - dataset, - ), - ) - - # "IS NULL" is not treated as the same as - # "= None" in SQL - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) - - row = database.query( - query, - tuple(params), - ).fetchone() - - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message="Record created successfully", - result=dict( - id=row["id"], - country_id=row["country_id"], - reform_id=row["reform_id"], - reform_label=row["reform_label"], - baseline_id=row["baseline_id"], - baseline_label=row["baseline_label"], - user_id=row["user_id"], - year=row["year"], - geography=row["geography"], - dataset=row["dataset"], - number_of_provisions=row["number_of_provisions"], - api_version=row["api_version"], - added_date=row["added_date"], - updated_date=row["updated_date"], - budgetary_impact=row["budgetary_impact"], - type=row["type"], - ), - ) - - return Response( - json.dumps(response_body), - status=201, - mimetype="application/json", - ) - - -@validate_country -def get_user_policy(country_id: str, user_id: str) -> dict: - """ - Fetch all saved user policies by user id - """ - - # Get the policy record for a given policy ID. - rows = database.query( - f"SELECT * FROM user_policies WHERE country_id = ? AND user_id = ?", - (country_id, user_id), - ).fetchall() - - rows_parsed = [ - dict( - id=row["id"], - country_id=row["country_id"], - reform_id=row["reform_id"], - reform_label=row["reform_label"], - baseline_id=row["baseline_id"], - baseline_label=row["baseline_label"], - user_id=row["user_id"], - year=row["year"], - geography=row["geography"], - dataset=row["dataset"], - number_of_provisions=row["number_of_provisions"], - api_version=row["api_version"], - added_date=row["added_date"], - updated_date=row["updated_date"], - budgetary_impact=row["budgetary_impact"], - type=row["type"], - ) - for row in rows - ] - - if rows_parsed is None: - response = dict( - status="ok", - message=f"No saved policies found for user {user_id}", - ) - return Response( - json.dumps(response), - status=200, - mimetype="application/json", - ) - return dict( - status="ok", - message=None, - result=rows_parsed, - ) - - -# Whitelist of columns that callers are allowed to modify via -# update_user_policy. Identity columns (id, country_id, user_id, -# reform_id, baseline_id) are intentionally excluded because they -# define the record; allowing clients to rewrite them would both -# break referential assumptions and let the column name be used -# as a SQL injection vector (keys are interpolated into the -# UPDATE statement below). -UPDATE_USER_POLICY_ALLOWED_FIELDS = frozenset( - { - "reform_label", - "baseline_label", - "year", - "geography", - "dataset", - "number_of_provisions", - "api_version", - "added_date", - "updated_date", - "budgetary_impact", - "type", - } -) - - -@validate_country -def update_user_policy(country_id: str) -> dict: - """ - Update any parts of a user_policy, given a user_policy ID - """ - - payload = request.json - if not isinstance(payload, dict) or "id" not in payload: - return Response( - json.dumps({"message": "Request body must include an 'id' field."}), - status=400, - mimetype="application/json", - ) - - user_policy_id = payload.pop("id") - - # Reject any unknown/unsafe keys. The keys end up interpolated - # into a SQL UPDATE statement, so we must validate them against - # a static whitelist instead of trusting the JSON payload. - unknown_keys = [ - key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS - ] - if unknown_keys: - return Response( - json.dumps( - { - "message": ( - "Request body contains unsupported fields: " - f"{sorted(unknown_keys)}" - ) - } - ), - status=400, - mimetype="application/json", - ) - - if not payload: - return Response( - json.dumps( - {"message": "Request body must include at least one field to update."} - ), - status=400, - mimetype="application/json", - ) - - # Construct the relevant UPDATE request from whitelisted keys. - setter_array = [] - args = [] - for key in payload: - setter_array.append(f"{key} = ?") - args.append(payload[key]) - setter_phrase = ", ".join(setter_array) - - args.append(user_policy_id) - sql_request = f"UPDATE user_policies SET {setter_phrase} WHERE id = ?" - - try: - database.query(sql_request, (tuple(args))) - except Exception as e: - return Response( - json.dumps( - {"message": f"Internal database error: {e}; please try again later."} - ), - status=500, - mimetype="application/json", - ) - - response_body = dict( - status="ok", - message="Record updated successfully", - result=dict(id=user_policy_id), - ) - - return Response( - json.dumps(response_body), - status=200, - mimetype="application/json", - ) diff --git a/policyengine_api/endpoints/simulation.py b/policyengine_api/endpoints/simulation.py deleted file mode 100644 index a0d9bd70d..000000000 --- a/policyengine_api/endpoints/simulation.py +++ /dev/null @@ -1,52 +0,0 @@ -from policyengine_api.data import local_database - -""" - -CREATE TABLE IF NOT EXISTS reform_impact ( - reform_impact_id INTEGER PRIMARY KEY AUTO_INCREMENT, - baseline_policy_id INT NOT NULL, - reform_policy_id INT NOT NULL, - country_id VARCHAR(3) NOT NULL, - region VARCHAR(32) NOT NULL, - time_period VARCHAR(32) NOT NULL, - options_json JSON, - options_hash VARCHAR(255), - api_version VARCHAR(10) NOT NULL, - reform_impact_json JSON NOT NULL, - status VARCHAR(32) NOT NULL, - message VARCHAR(255), - start_time DATETIME -); - -""" - - -_MAX_SIMULATION_RESULTS = 1000 -_DEFAULT_SIMULATION_RESULTS = 100 - - -def get_simulations( - max_results: int | None = 100, -): - # Get the last N simulations ordered by start time. - # - # LIMIT is always applied (unbounded scans against reform_impact - # are expensive) and max_results is clamped to [1, - # _MAX_SIMULATION_RESULTS] before being bound as a parameter, so - # the value can never be interpolated into the SQL string. - if max_results is None: - max_results = _DEFAULT_SIMULATION_RESULTS - try: - max_results = int(max_results) - except (TypeError, ValueError): - max_results = _DEFAULT_SIMULATION_RESULTS - max_results = max(1, min(max_results, _MAX_SIMULATION_RESULTS)) - - result = local_database.query( - "SELECT * FROM reform_impact ORDER BY start_time DESC LIMIT ?", - (max_results,), - ).fetchall() - - # Format into [{}] - - return {"result": [dict(r) for r in result]} diff --git a/policyengine_api/extensions.py b/policyengine_api/extensions.py new file mode 100644 index 000000000..064b55d22 --- /dev/null +++ b/policyengine_api/extensions.py @@ -0,0 +1,6 @@ +"""Shared Flask extensions initialized by the application factory module.""" + +from flask_caching import Cache + + +cache = Cache() diff --git a/policyengine_api/response_factory.py b/policyengine_api/response_factory.py new file mode 100644 index 000000000..84cdd4240 --- /dev/null +++ b/policyengine_api/response_factory.py @@ -0,0 +1,26 @@ +"""Factories for consistent Flask response construction.""" + +import json +from typing import Any + +from flask import Response + + +def _make_error_response( + message: object, + status_code: int, + *, + include_status: bool = True, + mimetype: str | None = "application/json", + **payload_fields: Any, +) -> Response: + """Build a JSON error response while preserving legacy payload variants.""" + payload = {"message": str(message), **payload_fields} + if include_status: + payload = {"status": "error", **payload} + + return Response( + json.dumps(payload), + status=status_code, + mimetype=mimetype, + ) diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 8507cfe3e..56bb35e12 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -4,6 +4,7 @@ EconomicImpactResult, BudgetWindowEconomicImpactResult, ) +from policyengine_api.response_factory import _make_error_response from policyengine_api.utils import get_current_law_policy_id from policyengine_api.utils.payload_validators import validate_country from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS @@ -24,14 +25,7 @@ def _json_response(payload: dict, status: int = 200) -> Response: def _bad_request_response(message: str) -> Response: - return _json_response( - { - "status": "error", - "message": message, - "result": None, - }, - status=400, - ) + return _make_error_response(message, 400, result=None) @economy_bp.route( diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py index e9fced1c0..a4801bb09 100644 --- a/policyengine_api/routes/error_routes.py +++ b/policyengine_api/routes/error_routes.py @@ -1,67 +1,50 @@ -import json -from flask import Response, Blueprint +from flask import Blueprint, Response from werkzeug.exceptions import ( HTTPException, ) +from policyengine_api.response_factory import _make_error_response + error_bp = Blueprint("error", __name__) @error_bp.app_errorhandler(404) def response_404(error) -> Response: """Specific handler for 404 Not Found errors""" - return make_error_response(error, 404) + return _make_error_response(error, 404, result=None) @error_bp.app_errorhandler(400) def response_400(error) -> Response: """Specific handler for 400 Bad Request errors""" - return make_error_response(error, 400) + return _make_error_response(error, 400, result=None) @error_bp.app_errorhandler(401) def response_401(error) -> Response: """Specific handler for 401 Unauthorized errors""" - return make_error_response(error, 401) + return _make_error_response(error, 401, result=None) @error_bp.app_errorhandler(403) def response_403(error) -> Response: """Specific handler for 403 Forbidden errors""" - return make_error_response(error, 403) + return _make_error_response(error, 403, result=None) @error_bp.app_errorhandler(500) def response_500(error) -> Response: """Specific handler for 500 Internal Server errors""" - return make_error_response(error, 500) + return _make_error_response(error, 500, result=None) @error_bp.app_errorhandler(HTTPException) def response_http_exception(error: HTTPException) -> Response: """Generic handler for HTTPException; should be raised if no specific handler is found""" - return make_error_response(str(error), error.code) + return _make_error_response(error, error.code, result=None) @error_bp.app_errorhandler(Exception) def response_generic_error(error: Exception) -> Response: """Handler for any unhandled exceptions""" - return make_error_response(str(error), 500) - - -def make_error_response( - error, - status_code: int, -) -> Response: - """Create a generic error response""" - return Response( - json.dumps( - { - "status": "error", - "message": str(error), - "result": None, - } - ), - status_code, - mimetype="application/json", - ) + return _make_error_response(error, 500, result=None) diff --git a/policyengine_api/routes/home_routes.py b/policyengine_api/routes/home_routes.py new file mode 100644 index 000000000..fdc230705 --- /dev/null +++ b/policyengine_api/routes/home_routes.py @@ -0,0 +1,14 @@ +from flask import Blueprint + + +home_bp = Blueprint("home", __name__) + + +@home_bp.route("/", methods=["GET"]) +def get_home() -> str: + """Get the home page of the PolicyEngine API.""" + return ( + "

PolicyEngine households API

" + "

Use this API to compute the impact of public policy on individual " + "households.

" + ) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index d7420b39e..731895889 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -1,15 +1,41 @@ -from flask import Blueprint, Response, request -from werkzeug.exceptions import NotFound, BadRequest import json +import logging + +from flask import Blueprint, Response, request +from policyengine_core.errors import SituationParsingError +from werkzeug.exceptions import BadRequest, NotFound +from policyengine_api.data.v1_models import Household +from policyengine_api.extensions import cache +from policyengine_api.response_factory import _make_error_response +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, + HouseholdNotFoundError, + InvalidHouseholdInputsError, + PolicyNotFoundError, +) from policyengine_api.services.household_service import HouseholdService +from policyengine_api.utils import make_cache_key +from policyengine_api.utils.input_validation import format_unrecognized_inputs_message from policyengine_api.utils.payload_validators import ( - validate_household_payload, validate_country, + validate_household_payload, ) household_bp = Blueprint("household", __name__) household_service = HouseholdService() +household_calculation_service = HouseholdCalculationService() + + +def _serialize_household(household: Household) -> dict: + return { + "id": household.id, + "country_id": household.country_id, + "label": household.label, + "api_version": household.api_version, + "household_json": household.household_json, + "household_hash": household.household_hash, + } @household_bp.route("//household/", methods=["GET"]) @@ -24,8 +50,9 @@ def get_household(country_id: str, household_id: int) -> Response: """ print(f"Got request for household {household_id} in country {country_id}") - household: dict | None = household_service.get_household(country_id, household_id) - if household is None: + household = household_service.get_household(country_id, household_id) + result = None if household is None else _serialize_household(household) + if result is None: raise NotFound(f"Household #{household_id} not found.") else: return Response( @@ -33,7 +60,7 @@ def get_household(country_id: str, household_id: int) -> Response: { "status": "ok", "message": None, - "result": household, + "result": result, } ), status=200, @@ -62,7 +89,12 @@ def post_household(country_id: str) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - household_id = household_service.create_household(country_id, household_json, label) + household = household_service.create_household( + country_id, + household_json, + label, + ) + household_id = household.id return Response( json.dumps( @@ -102,14 +134,16 @@ def update_household(country_id: str, household_id: int) -> Response: label: str | None = payload.get("label") household_json: dict = payload.get("data") - household: dict | None = household_service.get_household(country_id, household_id) - if household is None: - raise NotFound(f"Household #{household_id} not found.") - - # Next, update the household - updated_household: dict = household_service.update_household( - country_id, household_id, household_json, label - ) + try: + updated_household = household_service.update_household( + country_id, + household_id, + household_json, + label, + ) + except LookupError: + raise NotFound(f"Household #{household_id} not found.") from None + updated_household_json = updated_household.household_json return Response( json.dumps( { @@ -117,10 +151,108 @@ def update_household(country_id: str, household_id: int) -> Response: "message": None, "result": { "household_id": household_id, - "household_json": updated_household["household_json"], + "household_json": updated_household_json, }, } ), status=200, mimetype="application/json", ) + + +@household_bp.route( + "//household//policy/", + methods=["GET"], +) +@validate_country +def get_household_under_policy(country_id: str, household_id: str, policy_id: str): + """Get a stored household's output under a stored policy.""" + try: + calculation = household_calculation_service.calculate_stored_household( + country_id, + int(household_id), + int(policy_id), + ) + except HouseholdNotFoundError: + return _make_error_response( + f"Household #{household_id} not found.", + 404, + ) + except PolicyNotFoundError: + return _make_error_response( + f"Policy #{policy_id} not found.", + 404, + ) + except InvalidHouseholdInputsError as error: + return _make_error_response( + format_unrecognized_inputs_message(error.invalid_inputs), + 400, + result=None, + errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], + ) + except Exception as error: + logging.exception(error) + return _make_error_response( + f"Error calculating household #{household_id} under policy " + f"#{policy_id}: {error}", + 500, + ) + + response_body = dict(status="ok", message=None, result=calculation.household) + if calculation.warnings: + response_body["warnings"] = list(calculation.warnings) + return response_body + + +def _calculate(country_id: str, *, add_missing: bool) -> dict | Response: + payload = request.json + household_json = payload.get("household", {}) + policy_json = payload.get("policy", {}) + + try: + calculation = household_calculation_service.calculate_household( + country_id, + household_json, + policy_json, + add_missing=add_missing, + ) + except InvalidHouseholdInputsError as error: + return _make_error_response( + format_unrecognized_inputs_message(error.invalid_inputs), + 400, + result=None, + errors=[invalid_input.to_dict() for invalid_input in error.invalid_inputs], + ) + except SituationParsingError as error: + return _make_error_response( + f"Invalid household payload: {error}", + 400, + result=None, + ) + except Exception as error: + logging.exception(error) + return _make_error_response( + f"Error calculating household under policy: {error}", + 500, + ) + + response_body = dict(status="ok", message=None, result=calculation.household) + if calculation.warnings: + response_body["warnings"] = list(calculation.warnings) + return response_body + + +@household_bp.route("//calculate", methods=["POST"]) +@cache.cached(make_cache_key=make_cache_key) +@validate_country +def get_calculate(country_id: str) -> dict | Response: + """Calculate a household without adding omitted yearly variables.""" + return _calculate(country_id, add_missing=False) + + +@household_bp.route("//calculate-full", methods=["POST"]) +@cache.cached(make_cache_key=make_cache_key) +@validate_country +def get_calculate_full(country_id: str) -> dict | Response: + """Calculate a household after adding omitted yearly variables.""" + return _calculate(country_id, add_missing=True) diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 3fc88fbf4..f857954a4 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,8 +1,12 @@ -from flask import Blueprint, Response, request import json +from flask import Blueprint, Response, request +from werkzeug.exceptions import BadRequest, NotFound + +from policyengine_api.data.v1_models import Policy, UserPolicy +from policyengine_api.response_factory import _make_error_response from policyengine_api.services.policy_service import PolicyService -from werkzeug.exceptions import NotFound, BadRequest +from policyengine_api.services.user_policy_service import UserPolicyService from policyengine_api.utils.payload_validators import ( validate_country, validate_set_policy_payload, @@ -10,6 +14,18 @@ policy_bp = Blueprint("policy", __name__) policy_service = PolicyService() +user_policy_service = UserPolicyService() + + +def _serialize_policy(policy: Policy) -> dict: + return { + "id": policy.id, + "country_id": policy.country_id, + "label": policy.label, + "api_version": policy.api_version, + "policy_json": policy.policy_json, + "policy_hash": policy.policy_hash, + } @policy_bp.route("//policy/", methods=["GET"]) @@ -30,13 +46,14 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: # Specifically cast policy_id to an integer policy_id = int(policy_id) - policy: dict | None = policy_service.get_policy(country_id, policy_id) + policy = policy_service.get_policy(country_id, policy_id) + result = None if policy is None else _serialize_policy(policy) - if policy is None: + if result is None: raise NotFound(f"Policy #{policy_id} not found.") return Response( - json.dumps({"status": "ok", "message": None, "result": policy}), + json.dumps({"status": "ok", "message": None, "result": result}), status=200, ) @@ -77,3 +94,216 @@ def set_policy(country_id: str) -> Response: code = 200 if is_existing_policy else 201 return Response(json.dumps(response_body), status=code, mimetype="application/json") + + +def _serialize_user_policy(user_policy: UserPolicy) -> dict: + return { + column.name: getattr(user_policy, column.name) + for column in UserPolicy.__table__.columns + } + + +@policy_bp.route("//policies", methods=["GET"]) +@validate_country +def get_policy_search(country_id: str) -> Response: + """Search policies for a country.""" + query = request.args.get("query", "") + unique_only = request.args.get("unique_only", default=False, type=json.loads) + + try: + results = policy_service.search_policies( + country_id, + query, + unique_only=unique_only, + ) + if not results: + return _make_error_response( + f"No policies found for country {country_id} for query '{query}", + 404, + ) + + policies = [dict(id=result.id, label=result.label) for result in results] + return Response( + json.dumps( + dict( + status="ok", + message="Policies found", + result=policies, + ) + ), + status=200, + mimetype="application/json", + ) + except Exception as error: + return _make_error_response( + f"Internal server error: {error}", + 500, + ) + + +@policy_bp.route("//user-policy", methods=["POST"]) +@validate_country +def set_user_policy(country_id: str) -> Response: + """Create a saved policy for a user, or return its existing ID.""" + payload = request.json + reform_label = payload.pop("reform_label", None) + reform_id = payload.pop("reform_id") + baseline_label = payload.pop("baseline_label", None) + baseline_id = payload.pop("baseline_id") + user_id = payload.pop("user_id") + year = payload.pop("year") + geography = payload.pop("geography") + dataset = payload.pop("dataset", None) + number_of_provisions = payload.pop("number_of_provisions") + api_version = payload.pop("api_version") + added_date = payload.pop("added_date") + updated_date = payload.pop("updated_date") + budgetary_impact = payload.pop("budgetary_impact", None) + policy_type = payload.pop("type", None) + + values = { + "country_id": country_id, + "reform_id": reform_id, + "reform_label": reform_label, + "baseline_id": baseline_id, + "baseline_label": baseline_label, + "user_id": user_id, + "year": year, + "geography": geography, + "dataset": dataset, + "number_of_provisions": number_of_provisions, + "api_version": api_version, + "added_date": added_date, + "updated_date": updated_date, + "budgetary_impact": budgetary_impact, + "type": policy_type, + } + + try: + creation = user_policy_service.create_or_get_user_policy(values) + user_policy = creation.user_policy + if not creation.created: + return Response( + json.dumps( + dict( + status="ok", + message=( + f"The reform #{reform_id} / baseline #{baseline_id} pair " + f"already exists for user {user_id}" + ), + result=dict(id=user_policy.id), + ) + ), + status=200, + mimetype="application/json", + ) + except Exception as error: + return _make_error_response( + f"Internal database error: {error}; please try again later.", + 500, + include_status=False, + ) + + return Response( + json.dumps( + dict( + status="ok", + message="Record created successfully", + result=_serialize_user_policy(user_policy), + ) + ), + status=201, + mimetype="application/json", + ) + + +@policy_bp.route("//user-policy/", methods=["GET"]) +@validate_country +def get_user_policy(country_id: str, user_id: str) -> dict: + """Fetch all saved policies for a user.""" + user_policies = user_policy_service.list_user_policies(country_id, user_id) + return dict( + status="ok", + message=None, + result=[_serialize_user_policy(row) for row in user_policies], + ) + + +UPDATE_USER_POLICY_ALLOWED_FIELDS = frozenset( + { + "reform_label", + "baseline_label", + "year", + "geography", + "dataset", + "number_of_provisions", + "api_version", + "added_date", + "updated_date", + "budgetary_impact", + "type", + } +) + + +@policy_bp.route("//user-policy", methods=["PUT"]) +@validate_country +def update_user_policy(country_id: str) -> Response: + """Update mutable fields on a saved user policy.""" + payload = request.json + if not isinstance(payload, dict) or "id" not in payload: + return _make_error_response( + "Request body must include an 'id' field.", + 400, + include_status=False, + ) + + user_policy_id = payload.pop("id") + unknown_keys = [ + key for key in payload if key not in UPDATE_USER_POLICY_ALLOWED_FIELDS + ] + if unknown_keys: + return _make_error_response( + f"Request body contains unsupported fields: {sorted(unknown_keys)}", + 400, + include_status=False, + ) + + if not payload: + return _make_error_response( + "Request body must include at least one field to update.", + 400, + include_status=False, + ) + + try: + user_policy = user_policy_service.update_user_policy( + country_id, + user_policy_id, + payload, + ) + except Exception as error: + return _make_error_response( + f"Internal database error: {error}; please try again later.", + 500, + include_status=False, + ) + + if user_policy is None: + return _make_error_response( + f"User policy #{user_policy_id} not found.", + 404, + include_status=False, + ) + + return Response( + json.dumps( + dict( + status="ok", + message="Record updated successfully", + result=dict(id=user_policy_id), + ) + ), + status=200, + mimetype="application/json", + ) diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py new file mode 100644 index 000000000..e49076304 --- /dev/null +++ b/policyengine_api/routes/reform_impact_routes.py @@ -0,0 +1,57 @@ +from datetime import datetime +import json + +from flask import Blueprint, Response, request + +from policyengine_api.data.v1_models import ReformImpact +from policyengine_api.services.reform_impacts_service import ReformImpactsService + + +reform_impact_bp = Blueprint("reform_impact", __name__) +reform_impacts_service = ReformImpactsService() + +_MAX_SIMULATION_RESULTS = 1000 +_DEFAULT_SIMULATION_RESULTS = 100 + + +def _serialize_v1_reform_impact(impact: ReformImpact) -> dict: + """Project canonical ORM values onto the historical v1 response shape.""" + + result = { + column.name: getattr(impact, column.name) + for column in ReformImpact.__table__.columns + } + for field in ("options_json", "reform_impact_json"): + value = result[field] + if value is not None and not isinstance(value, str): + result[field] = json.dumps(value) + for field in ("start_time", "end_time"): + value = result[field] + if isinstance(value, datetime): + result[field] = str(value) + return result + + +def _parse_result_limit(value: str | None) -> int: + if value is None: + return _DEFAULT_SIMULATION_RESULTS + try: + result_limit = int(value) + except (TypeError, ValueError): + return _DEFAULT_SIMULATION_RESULTS + return max(1, min(result_limit, _MAX_SIMULATION_RESULTS)) + + +@reform_impact_bp.route("/simulations", methods=["GET"]) +def get_simulations() -> Response: + """Return recent reform impacts, bounded to protect the database query.""" + result_limit = _parse_result_limit(request.args.get("max_results")) + impacts = reform_impacts_service.get_recent_reform_impacts(result_limit) + + return Response( + json.dumps( + {"result": [_serialize_v1_reform_impact(impact) for impact in impacts]} + ), + status=200, + mimetype="application/json", + ) diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index 48a2ac43a..2b2bc21e7 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -5,14 +5,38 @@ import jsonschema import pydantic -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.services.report_output_service import ( + ReportOutputService, + ReportOutputView, +) from policyengine_api.constants import CURRENT_YEAR +from policyengine_api.data.v1_models import ReportOutput from policyengine_api.utils.payload_validators import validate_country report_output_bp = Blueprint("report_output", __name__) report_output_service = ReportOutputService() +def _serialize_v1_report_output(view: ReportOutputView) -> dict: + """Project mapped report state onto the historical v1 response shape.""" + + report_output = view.report_output + result = { + column.name: getattr(report_output, column.name) + for column in ReportOutput.__table__.columns + } + if view.response_id is not None: + result["id"] = view.response_id + if result.get("output") is not None and not isinstance(result["output"], str): + result["output"] = json.dumps(result["output"]) + if view.display_run is not None: + for field in ("requested_at", "started_at", "finished_at"): + result[field] = report_output_service.format_run_timestamp( + getattr(view.display_run, field) + ) + return result + + @report_output_bp.route("//report", methods=["POST"]) @validate_country def create_report_output(country_id: str) -> Response: @@ -48,51 +72,29 @@ def create_report_output(country_id: str) -> Response: raise BadRequest("year must be a string") try: - # Check if report already exists with these simulation IDs and year - existing_report = report_output_service.find_existing_report_output( + creation = report_output_service.create_or_reuse_report_output( country_id=country_id, simulation_1_id=simulation_1_id, simulation_2_id=simulation_2_id, year=year, ) - - if existing_report: - existing_report = ( - report_output_service.ensure_report_output_dual_write_state( - existing_report["id"], - country_id=country_id, - ) - ) - # Report already exists, return it with 200 status - response_body = dict( - status="ok", - message="Report output already exists", - result=existing_report, - ) - - return Response( - json.dumps(response_body), - status=200, - mimetype="application/json", - ) - - # Create new report output - created_report = report_output_service.create_report_output( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, + result = _serialize_v1_report_output(creation.view) + message = ( + "Report output created successfully" + if creation.created + else "Report output already exists" ) + status_code = 201 if creation.created else 200 response_body = dict( status="ok", - message="Report output created successfully", - result=created_report, + message=message, + result=result, ) return Response( json.dumps(response_body), - status=201, + status=status_code, mimetype="application/json", ) @@ -129,17 +131,15 @@ def get_report_output(country_id: str, report_id: int) -> Response: """ print(f"Getting report output {report_id} for country {country_id}") - report_output: dict | None = report_output_service.get_report_output( - country_id, report_id - ) - - if report_output is None: + view = report_output_service.resolve_report_output(country_id, report_id) + if view is None: raise NotFound(f"Report #{report_id} not found.") + result = _serialize_v1_report_output(view) response_body = dict( status="ok", message=None, - result=report_output, + result=result, ) return Response( @@ -191,34 +191,21 @@ def update_report_output(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - # First check if the report output exists without running pointer sync: - # syncing a completed parent before this mutation can clear an active - # pending rerun that this PATCH is about to mark as running. - if not report_output_service.report_output_exists(country_id, report_id): - raise NotFound(f"Report #{report_id} not found.") - - # Update the report output - success = report_output_service.update_report_output( + view = report_output_service.update_report_output( country_id=country_id, report_id=report_id, status=status, output=output, error_message=error_message, ) - - if not success: + if view is None: raise BadRequest("No fields to update") - - # Get the updated stored record so stale-runtime jobs do not appear to - # complete the current runtime lineage in the PATCH response. - updated_report = report_output_service.get_stored_report_output( - country_id, report_id - ) + result = _serialize_v1_report_output(view) response_body = dict( status="ok", message="Report output updated successfully", - result=updated_report, + result=result, ) return Response( @@ -227,6 +214,8 @@ def update_report_output(country_id: str) -> Response: mimetype="application/json", ) + except LookupError: + raise NotFound(f"Report #{report_id} not found.") from None except HTTPException: # Let explicit client-error responses (BadRequest/NotFound/etc.) pass # through without being logged as "Unexpected error". diff --git a/policyengine_api/routes/simulation_analysis_routes.py b/policyengine_api/routes/simulation_analysis_routes.py index 5157b807d..d9ddc0b97 100644 --- a/policyengine_api/routes/simulation_analysis_routes.py +++ b/policyengine_api/routes/simulation_analysis_routes.py @@ -1,6 +1,5 @@ from flask import Blueprint, request, Response, stream_with_context from werkzeug.exceptions import BadRequest -from policyengine_api.utils.payload_validators import validate_country from policyengine_api.services.simulation_analysis_service import ( SimulationAnalysisService, ) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index f2bacd6cb..9476d6d67 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -5,6 +5,7 @@ import jsonschema import pydantic +from policyengine_api.data.v1_models import Simulation from policyengine_api.services.simulation_service import SimulationService from policyengine_api.utils.payload_validators import validate_country @@ -12,6 +13,24 @@ simulation_service = SimulationService() +def _serialize_v1_simulation(simulation: Simulation) -> dict: + """Project canonical ORM JSON objects onto the legacy v1 response shape.""" + + response = { + column.name: getattr(simulation, column.name) + for column in Simulation.__table__.columns + } + for field in ("output", "simulation_spec_json"): + value = response.get(field) + if value is not None and not isinstance(value, str): + response[field] = json.dumps(value) + + # TODO: Remove this compatibility projection when the v1 contract is + # retired. New v2 response models should expose these fields as native JSON + # objects instead of carrying the historical JSON-in-a-string shape forward. + return response + + @simulation_bp.route("//simulation", methods=["POST"]) @validate_country def create_simulation(country_id: str) -> Response: @@ -50,49 +69,29 @@ def create_simulation(country_id: str) -> Response: raise BadRequest("policy_id must be an integer") try: - # Check if simulation already exists with these parameters - existing_simulation = simulation_service.find_existing_simulation( + creation = simulation_service.get_or_create_simulation( country_id=country_id, population_id=population_id, population_type=population_type, policy_id=policy_id, ) - - if existing_simulation: - existing_simulation = simulation_service.ensure_simulation_dual_write_state( - existing_simulation["id"], - country_id=country_id, - ) - # Simulation already exists, return it with 200 status - response_body = dict( - status="ok", - message="Simulation already exists", - result=existing_simulation, - ) - - return Response( - json.dumps(response_body), - status=200, - mimetype="application/json", - ) - - # Create new simulation - created_simulation = simulation_service.create_simulation( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, + result = _serialize_v1_simulation(creation.simulation) + message = ( + "Simulation created successfully" + if creation.created + else "Simulation already exists" ) + status_code = 201 if creation.created else 200 response_body = dict( status="ok", - message="Simulation created successfully", - result=created_simulation, + message=message, + result=result, ) return Response( json.dumps(response_body), - status=201, + status=status_code, mimetype="application/json", ) @@ -129,17 +128,16 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: if simulation_id <= 0: raise BadRequest("simulation_id must be a positive integer") - simulation: dict | None = simulation_service.get_simulation( - country_id, simulation_id - ) + simulation = simulation_service.get_simulation(country_id, simulation_id) + result = None if simulation is None else _serialize_v1_simulation(simulation) - if simulation is None: + if result is None: raise NotFound(f"Simulation #{simulation_id} not found.") response_body = dict( status="ok", message=None, - result=simulation, + result=result, ) return Response( @@ -186,15 +184,7 @@ def update_simulation(country_id: str) -> Response: raise BadRequest("output is required when status is 'complete'") try: - # First check if the simulation exists - existing_simulation = simulation_service.get_simulation( - country_id, simulation_id - ) - if existing_simulation is None: - raise NotFound(f"Simulation #{simulation_id} not found.") - - # Update the simulation - success = simulation_service.update_simulation( + simulation = simulation_service.update_simulation( country_id=country_id, simulation_id=simulation_id, status=status, @@ -202,18 +192,15 @@ def update_simulation(country_id: str) -> Response: error_message=error_message, ) - if not success: + if simulation is None: raise BadRequest("No fields to update") - # Get the updated record - updated_simulation = simulation_service.get_simulation( - country_id, simulation_id - ) + result = _serialize_v1_simulation(simulation) response_body = dict( status="ok", message="Simulation updated successfully", - result=updated_simulation, + result=result, ) return Response( @@ -222,6 +209,8 @@ def update_simulation(country_id: str) -> Response: mimetype="application/json", ) + except LookupError: + raise NotFound(f"Simulation #{simulation_id} not found.") from None except HTTPException: # Let explicit client-error responses (BadRequest/NotFound/etc.) pass # through without being logged as "Unexpected error". diff --git a/policyengine_api/routes/system_routes.py b/policyengine_api/routes/system_routes.py new file mode 100644 index 000000000..f0e2748c0 --- /dev/null +++ b/policyengine_api/routes/system_routes.py @@ -0,0 +1,30 @@ +from flask import Blueprint, Response, jsonify + +from policyengine_api.readiness import is_ready +from policyengine_api.specification import OPENAPI_SPECIFICATION + + +system_bp = Blueprint("system", __name__) + + +@system_bp.route("/liveness-check", methods=["GET"]) +def liveness_check() -> Response: + return Response("OK", status=200, headers={"Content-Type": "text/plain"}) + + +@system_bp.route("/readiness-check", methods=["GET"]) +def readiness_check() -> Response: + # The service is not ready until startup warmup compiles the simulation + # machinery. Liveness remains unconditional so the worker is not restarted. + if not is_ready(): + return Response( + "NOT READY", + status=503, + headers={"Content-Type": "text/plain"}, + ) + return Response("OK", status=200, headers={"Content-Type": "text/plain"}) + + +@system_bp.route("/specification", methods=["GET"]) +def get_specification() -> Response: + return jsonify(OPENAPI_SPECIFICATION) diff --git a/policyengine_api/routes/tracer_analysis_routes.py b/policyengine_api/routes/tracer_analysis_routes.py index 6638282a4..43695d45c 100644 --- a/policyengine_api/routes/tracer_analysis_routes.py +++ b/policyengine_api/routes/tracer_analysis_routes.py @@ -8,8 +8,6 @@ TracerAnalysisService, ) import json -from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS -import re tracer_analysis_bp = Blueprint("tracer_analysis", __name__) tracer_analysis_service = TracerAnalysisService() @@ -27,8 +25,6 @@ def execute_tracer_analysis(country_id): household_id = payload.get("household_id") policy_id = payload.get("policy_id") variable = payload.get("variable") - api_version = COUNTRY_PACKAGE_VERSIONS[country_id] - if not isinstance(variable, str): raise BadRequest("variable must be a string") diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index d859629c6..77175feff 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,6 +1,6 @@ from flask import Blueprint, Response, request +from policyengine_api.data.v1_models import UserProfile from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.data import database import json from policyengine_api.services.user_service import UserService from werkzeug.exceptions import BadRequest, NotFound @@ -9,6 +9,22 @@ user_service = UserService() +def _serialize_user_profile( + profile: UserProfile, + *, + include_auth0_id: bool, +) -> dict: + result = { + "user_id": profile.user_id, + "primary_country": profile.primary_country, + "username": profile.username, + "user_since": profile.user_since, + } + if include_auth0_id: + result["auth0_id"] = profile.auth0_id + return result + + @user_profile_bp.route("//user-profile", methods=["POST"]) @validate_country def set_user_profile(country_id: str) -> Response: @@ -24,22 +40,18 @@ def set_user_profile(country_id: str) -> Response: username = payload.pop("username", None) user_since = payload.pop("user_since") - created, row = user_service.create_profile( + created, profile = user_service.create_profile( primary_country=country_id, auth0_id=auth0_id, username=username, user_since=user_since, ) + result = _serialize_user_profile(profile, include_auth0_id=False) response = dict( status="ok", message="Record created successfully" if created else "Record exists", - result=dict( - user_id=row["user_id"], - primary_country=row["primary_country"], - username=row["username"], - user_since=row["user_since"], - ), + result=result, ) return Response( json.dumps(response), @@ -57,21 +69,23 @@ def get_user_profile(country_id: str) -> Response: if (auth0_id is None) and (user_id is None): raise BadRequest("auth0_id or user_id must be provided") - row = ( + profile = ( user_service.get_profile(user_id=user_id) if auth0_id is None else user_service.get_profile(auth0_id=auth0_id) ) + readable_row = ( + None + if profile is None + else _serialize_user_profile( + profile, + include_auth0_id=auth0_id is not None, + ) + ) - if row is None: + if readable_row is None: raise NotFound("No such user") - readable_row = dict(row) - # Delete auth0_id value if querying from user_id, as that value - # is a more private attribute than all others - if auth0_id is None: - del readable_row["auth0_id"] - response_body = dict( status="ok", message=f"User #{readable_row['user_id']} found successfully", @@ -94,9 +108,6 @@ def update_user_profile(country_id: str) -> Response: will assume malicious intent and 403 """ - # Construct the relevant UPDATE request - setter_array = [] - args = [] payload = request.json if payload is None: diff --git a/policyengine_api/services/ai_analysis_service.py b/policyengine_api/services/ai_analysis_service.py index ef0f6ce6a..6dd0bee93 100644 --- a/policyengine_api/services/ai_analysis_service.py +++ b/policyengine_api/services/ai_analysis_service.py @@ -1,9 +1,15 @@ -import anthropic -import os import json -from typing import Generator, Optional -from policyengine_api.data import local_database +import os +from collections.abc import Generator +from typing import Callable + +import anthropic from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Analysis class StreamEvent(BaseModel): @@ -21,34 +27,53 @@ class ErrorEvent(StreamEvent): class AIAnalysisService: - """ - Base class for various AI analysis-based services, - including SimulationAnalysisService, that connects with the analysis - local database table - """ - - def get_existing_analysis(self, prompt: str) -> Optional[str]: - """ - Get existing analysis from the local database - """ - - analysis = local_database.query( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ).fetchone() - - if analysis is None: - return None - - return json.dumps(analysis["analysis"]) - - def trigger_ai_analysis(self, prompt: str) -> Generator[str, None, None]: - # Configure a Claude client - claude_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + """AI analysis operations with short, service-owned ORM scopes.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + claude_client_factory: Callable[[], anthropic.Anthropic] | None = None, + ) -> None: + self._injected_session_factory = session_factory + self._claude_client_factory = claude_client_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory(local=True) + + def get_existing_analysis( + self, + prompt: str, + ) -> Analysis | None: + with self._sessions() as session: + return self._get_existing_analysis(session, prompt) + + @staticmethod + def _get_existing_analysis( + session: Session, + prompt: str, + ) -> Analysis | None: + return session.scalar( + select(Analysis) + .where( + Analysis.prompt == prompt, + Analysis.status.in_(("complete", "ok")), + ) + .order_by(Analysis.prompt_id.desc()) + ) + + def trigger_ai_analysis( + self, + prompt: str, + ) -> Generator[str, None, None]: + claude_client = ( + self._claude_client_factory() + if self._claude_client_factory is not None + else anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + ) def generate(): response_text = "" - with claude_client.messages.stream( model="claude-sonnet-4-20250514", max_tokens=1500, @@ -57,22 +82,26 @@ def generate(): messages=[{"role": "user", "content": prompt}], ) as stream: for event in stream: - # Docs on structure of Anthropic error events at https://docs.anthropic.com/en/api/messages-streaming#error-events if event.type == "error": - error: dict[str, str] = event.error - error_type: str = error["type"] - return_event = ErrorEvent(error=error_type) - yield json.dumps(return_event.model_dump()) + "\n" + yield ( + json.dumps( + ErrorEvent(error=event.error["type"]).model_dump() + ) + + "\n" + ) return if event.type == "text": response_text += event.text - return_event = TextEvent(stream=event.text) - yield json.dumps(return_event.model_dump()) + "\n" - - # Update the analysis record and return if no error occurred - local_database.query( - f"INSERT INTO analysis (prompt, analysis, status) VALUES (?, ?, ?)", - (prompt, response_text, "ok"), - ) + yield ( + json.dumps(TextEvent(stream=event.text).model_dump()) + "\n" + ) + with self._sessions.begin() as session: + session.add( + Analysis( + prompt=prompt, + analysis=response_text, + status="ok", + ) + ) return generate() diff --git a/policyengine_api/endpoints/economy/compare.py b/policyengine_api/services/economy_comparison.py similarity index 99% rename from policyengine_api/endpoints/economy/compare.py rename to policyengine_api/services/economy_comparison.py index bb10ed7f6..a50f909f8 100644 --- a/policyengine_api/endpoints/economy/compare.py +++ b/policyengine_api/services/economy_comparison.py @@ -1,12 +1,13 @@ +"""Reusable economy comparison calculations and response models.""" + import logging -from microdf import MicroDataFrame, MicroSeries + +import h5py import numpy as np -import sys -from policyengine_core.tools.hugging_face import download_huggingface_dataset import pandas as pd -import h5py +from microdf import MicroSeries +from policyengine_core.tools.hugging_face import download_huggingface_dataset from pydantic import BaseModel -from typing import Any logger = logging.getLogger(__name__) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index d896d6d01..24b88b9a9 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -3,7 +3,7 @@ import json import uuid from enum import Enum -from typing import Annotated, Any, Literal, Optional +from typing import Any, Literal, Optional import httpx import numpy as np @@ -22,6 +22,7 @@ get_valid_state_codes, normalize_us_region, ) +from policyengine_api.data.v1_models import ReformImpact from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint @@ -35,8 +36,6 @@ load_dotenv() -policy_service = PolicyService() -reform_impacts_service = ReformImpactsService() budget_window_cache = BudgetWindowCache() @@ -235,6 +234,73 @@ class EconomyService: with other services to access their respective tables """ + def __init__( + self, + *, + primary_session_factory=None, + local_session_factory=None, + policy_service_: PolicyService | None = None, + reform_impacts_service_: ReformImpactsService | None = None, + budget_window_cache_: BudgetWindowCache | None = None, + simulation_entrypoint_=None, + ) -> None: + self._primary_session_factory = primary_session_factory + self._local_session_factory = local_session_factory + self._injected_policy_service = policy_service_ + self._injected_reform_impacts_service = reform_impacts_service_ + self._injected_budget_window_cache = budget_window_cache_ + self._injected_simulation_entrypoint = simulation_entrypoint_ + + @property + def _policies(self) -> PolicyService: + if self._injected_policy_service is None: + self._injected_policy_service = PolicyService(self._primary_session_factory) + return self._injected_policy_service + + @property + def _reform_impacts(self) -> ReformImpactsService: + if self._injected_reform_impacts_service is None: + self._injected_reform_impacts_service = ReformImpactsService( + self._local_session_factory + ) + return self._injected_reform_impacts_service + + @property + def _budget_window_cache(self) -> BudgetWindowCache: + return self._injected_budget_window_cache or budget_window_cache + + @property + def _simulation_gateway(self): + return self._injected_simulation_entrypoint or simulation_entrypoint + + def _get_policy_jsons( + self, + country_id: str, + baseline_policy_id: int, + reform_policy_id: int, + ) -> tuple[dict[str, Any], dict[str, Any]]: + baseline = self._policies.get_policy_json( + country_id, + baseline_policy_id, + ) + reform = self._policies.get_policy_json( + country_id, + reform_policy_id, + ) + return ( + self._parse_json_object(baseline), + self._parse_json_object(reform), + ) + + @staticmethod + def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: + """Accept ORM-decoded objects and legacy JSON text at a read boundary.""" + + parsed = json.loads(value) if isinstance(value, str) else value + if not isinstance(parsed, dict): + raise TypeError("Expected a JSON object") + return parsed + def get_economic_impact( self, country_id: str, @@ -320,14 +386,14 @@ def get_budget_window_economic_impact( ) cache_key = self._build_budget_window_cache_key(setup_options) - cached_result = budget_window_cache.get_completed_result(cache_key) + cached_result = self._budget_window_cache.get_completed_result(cache_key) if cached_result is not None: return BudgetWindowEconomicImpactResult.completed( cached_result, cache_status="result-hit", ) - batch_job_id = budget_window_cache.get_batch_job_id(cache_key) + batch_job_id = self._budget_window_cache.get_batch_job_id(cache_key) if batch_job_id: return self._get_budget_window_result_from_batch_job_id( batch_job_id=batch_job_id, @@ -339,7 +405,7 @@ def get_budget_window_economic_impact( claim_token = setup_options.process_id cache_status = "starting-claim-hit" - if budget_window_cache.claim_batch_start(cache_key, claim_token): + if self._budget_window_cache.claim_batch_start(cache_key, claim_token): cache_status = "miss" try: batch_execution = self._start_budget_window_batch( @@ -348,11 +414,13 @@ def get_budget_window_economic_impact( window_size=window_size, max_parallel=max_active_years, ) - budget_window_cache.store_batch_job_id( + self._budget_window_cache.store_batch_job_id( cache_key, batch_execution.batch_job_id ) except httpx.HTTPStatusError as error: - budget_window_cache.clear_starting_claim(cache_key, claim_token) + self._budget_window_cache.clear_starting_claim( + cache_key, claim_token + ) if ( error.response.status_code in BUDGET_WINDOW_SUBMISSION_VALIDATION_ERROR_STATUS_CODES @@ -364,7 +432,9 @@ def get_budget_window_economic_impact( ) raise except Exception: - budget_window_cache.clear_starting_claim(cache_key, claim_token) + self._budget_window_cache.clear_starting_claim( + cache_key, claim_token + ) raise return self._build_budget_window_computing_result( @@ -383,7 +453,7 @@ def _build_budget_window_cache_key( self, setup_options: EconomicImpactSetupOptions, ) -> str: - return budget_window_cache.build_key( + return self._budget_window_cache.build_key( country_id=setup_options.country_id, reform_policy_id=setup_options.reform_policy_id, baseline_policy_id=setup_options.baseline_policy_id, @@ -402,12 +472,9 @@ def _build_budget_window_batch_payload( window_size: int, max_parallel: int, ) -> dict[str, Any]: - baseline_policy = policy_service.get_policy_json( + baseline_policy, reform_policy = self._get_policy_jsons( setup_options.country_id, setup_options.baseline_policy_id, - ) - reform_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.reform_policy_id, ) sim_config: SimulationOptions = self._setup_sim_options( @@ -457,7 +524,7 @@ def _start_budget_window_batch( severity="INFO", ) - return simulation_entrypoint.run_budget_window_batch(sim_params) + return self._simulation_gateway.run_budget_window_batch(sim_params) def _build_budget_window_submission_error_message( self, error: httpx.HTTPStatusError @@ -488,14 +555,14 @@ def _get_budget_window_result_from_batch_job_id( queued_years_on_submit: list[str], cache_status: Optional[str] = None, ) -> BudgetWindowEconomicImpactResult: - batch_execution = simulation_entrypoint.get_budget_window_batch_by_id( + batch_execution = self._simulation_gateway.get_budget_window_batch_by_id( batch_job_id ) if batch_execution.status in EXECUTION_STATUSES_SUCCESS: result = batch_execution.result if not isinstance(result, dict) or not result: - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.failed( "Budget-window batch completed without a result", completed_years=batch_execution.completed_years, @@ -503,8 +570,8 @@ def _get_budget_window_result_from_batch_job_id( queued_years=batch_execution.queued_years or queued_years_on_submit, cache_status=cache_status, ) - budget_window_cache.set_completed_result(cache_key, result) - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.set_completed_result(cache_key, result) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.completed( result, cache_status=cache_status, @@ -512,7 +579,7 @@ def _get_budget_window_result_from_batch_job_id( if batch_execution.status in EXECUTION_STATUSES_FAILURE: error_message = batch_execution.error or "Budget-window batch failed" - budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.failed( error_message, completed_years=batch_execution.completed_years, @@ -634,7 +701,7 @@ def _get_or_create_economic_impact( most_recent_impact = self._get_most_recent_impact(setup_options) if ( not most_recent_impact - or most_recent_impact.get("options_hash") != setup_options.options_hash + or most_recent_impact.options_hash != setup_options.options_hash ): most_recent_impact = None @@ -691,7 +758,7 @@ def _resolve_runtime_bundle_for_setup_options( ( setup_options.runtime_app_name, setup_options.model_version, - ) = simulation_entrypoint.resolve_app_name( + ) = self._simulation_gateway.resolve_app_name( setup_options.country_id, setup_options.model_version, policyengine_version=setup_options.policyengine_version, @@ -741,8 +808,9 @@ def _get_previous_impacts( Fetch any previous simulation runs for the given policy reform. """ - previous_impacts: list[Any] = ( - reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix( + previous_impacts: list[Any] = [] + previous_impacts = ( + self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( country_id, policy_id, baseline_policy_id, @@ -779,19 +847,19 @@ def _get_most_recent_impact( return None for impact in previous_impacts: - if impact.get("options_hash") == setup_options.options_hash: + if impact.options_hash == setup_options.options_hash: return impact return previous_impacts[0] def _determine_impact_action( self, - most_recent_impact: dict | None, + most_recent_impact: ReformImpact | None, ) -> ImpactAction: if not most_recent_impact: return ImpactAction.CREATE - status = most_recent_impact.get("status") + status = most_recent_impact.status if status in [ImpactStatus.OK.value, ImpactStatus.ERROR.value]: return ImpactAction.COMPLETED elif status == ImpactStatus.COMPUTING.value: @@ -803,7 +871,7 @@ def _handle_execution_state( self, setup_options: EconomicImpactSetupOptions, execution_state: str, - reform_impact: dict, + reform_impact: ReformImpact, execution: Optional[Any] = None, ) -> EconomicImpactResult: """ @@ -814,14 +882,14 @@ def _handle_execution_state( """ if execution_state in EXECUTION_STATUSES_SUCCESS: result = self._with_policyengine_bundle( - result=simulation_entrypoint.get_execution_result(execution), + result=self._simulation_gateway.get_execution_result(execution), setup_options=setup_options, execution=execution, ) self._set_reform_impact_complete( setup_options=setup_options, - reform_impact_json=json.dumps(result), - execution_id=reform_impact["execution_id"], + reform_impact_json=result, + execution_id=reform_impact.execution_id, ) logger.log_struct( {"message": "Sim API execution completed"}, @@ -844,7 +912,7 @@ def _handle_execution_state( self._set_reform_impact_error( setup_options=setup_options, message=error_message, - execution_id=reform_impact["execution_id"], + execution_id=reform_impact.execution_id, ) logger.log_struct( {"message": error_message}, @@ -865,9 +933,9 @@ def _handle_execution_state( def _handle_completed_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> EconomicImpactResult: - result = json.loads(most_recent_impact["reform_impact_json"]) + result = self._parse_json_object(most_recent_impact.reform_impact_json) return EconomicImpactResult.completed( data=self._with_policyengine_bundle( result=result, @@ -878,12 +946,12 @@ def _handle_completed_impact( def _handle_computing_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> EconomicImpactResult: - execution = simulation_entrypoint.get_execution_by_id( - most_recent_impact["execution_id"] + execution = self._simulation_gateway.get_execution_by_id( + most_recent_impact.execution_id ) - execution_state = simulation_entrypoint.get_execution_status(execution) + execution_state = self._simulation_gateway.get_execution_status(execution) return self._handle_execution_state( execution_state=execution_state, setup_options=setup_options, @@ -895,11 +963,10 @@ def _handle_create_impact( self, setup_options: EconomicImpactSetupOptions, ) -> EconomicImpactResult: - baseline_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.baseline_policy_id - ) - reform_policy = policy_service.get_policy_json( - setup_options.country_id, setup_options.reform_policy_id + baseline_policy, reform_policy = self._get_policy_jsons( + setup_options.country_id, + setup_options.baseline_policy_id, + setup_options.reform_policy_id, ) sim_config: SimulationOptions = self._setup_sim_options( @@ -952,8 +1019,8 @@ def _handle_create_impact( if sim_params.get("time_period") is not None: sim_params["time_period"] = str(sim_params["time_period"]) - entrypoint_execution = simulation_entrypoint.run(sim_params) - execution_id = simulation_entrypoint.get_execution_id(entrypoint_execution) + entrypoint_execution = self._simulation_gateway.run(sim_params) + execution_id = self._simulation_gateway.get_execution_id(entrypoint_execution) run_id = getattr(entrypoint_execution, "run_id", None) or telemetry["run_id"] @@ -975,8 +1042,8 @@ def _handle_create_impact( def _setup_sim_options( self, country_id: str, - reform_policy: Annotated[str, "String-formatted JSON"], - baseline_policy: Annotated[str, "String-formatted JSON"], + reform_policy: dict[str, Any] | str, + baseline_policy: dict[str, Any] | str, region: str, time_period: str, scope: Literal["macro", "household"] = "macro", @@ -994,8 +1061,8 @@ def _setup_sim_options( { "country": country_id, "scope": scope, - "reform": json.loads(reform_policy), - "baseline": json.loads(baseline_policy), + "reform": self._parse_json_object(reform_policy), + "baseline": self._parse_json_object(baseline_policy), "time_period": time_period, "include_cliffs": include_cliffs, "region": self._setup_region(country_id=country_id, region=region), @@ -1045,25 +1112,25 @@ def _extract_dataset_version(self, dataset: str | None) -> str | None: return None return dataset.rsplit("@", 1)[1] - def _extract_cached_result(self, most_recent_impact: dict) -> dict: + def _extract_cached_result(self, most_recent_impact: ReformImpact) -> dict: try: - return json.loads(most_recent_impact["reform_impact_json"]) + return self._parse_json_object(most_recent_impact.reform_impact_json) except (TypeError, ValueError): return {} def _should_refresh_cached_impact( self, setup_options: EconomicImpactSetupOptions, - most_recent_impact: dict, + most_recent_impact: ReformImpact, ) -> bool: - if most_recent_impact.get("status") == ImpactStatus.COMPUTING.value: + if most_recent_impact.status == ImpactStatus.COMPUTING.value: return False cached_result = self._extract_cached_result(most_recent_impact) cached_resolved_app_name = cached_result.get("resolved_app_name") try: runtime_app_name, resolved_model_version = ( - simulation_entrypoint.resolve_app_name( + self._simulation_gateway.resolve_app_name( setup_options.country_id, setup_options.model_version, policyengine_version=setup_options.policyengine_version, @@ -1307,18 +1374,18 @@ def _set_reform_impact_computing( In the reform_impact table, set the status of the impact to "computing". """ try: - reform_impacts_service.set_reform_impact( + self._reform_impacts.set_reform_impact( country_id=setup_options.country_id, policy_id=setup_options.reform_policy_id, baseline_policy_id=setup_options.baseline_policy_id, region=setup_options.region, dataset=setup_options.dataset, time_period=setup_options.time_period, - options=json.dumps(setup_options.options), + options=setup_options.options, options_hash=setup_options.options_hash, status=ImpactStatus.COMPUTING.value, api_version=setup_options.api_version, - reform_impact_json=json.dumps({}), + reform_impact_json={}, start_time=datetime.datetime.now(), execution_id=execution_id, ) @@ -1334,14 +1401,14 @@ def _set_reform_impact_computing( def _set_reform_impact_complete( self, setup_options: EconomicImpactSetupOptions, - reform_impact_json: str, + reform_impact_json: dict[str, Any], execution_id: str, ): """ In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON. """ try: - reform_impacts_service.set_complete_reform_impact( + self._reform_impacts.set_complete_reform_impact( country_id=setup_options.country_id, reform_policy_id=setup_options.reform_policy_id, baseline_policy_id=setup_options.baseline_policy_id, @@ -1371,7 +1438,7 @@ def _set_reform_impact_error( In the reform_impact table, set the status of the impact to "error" and store the error message. """ try: - reform_impacts_service.set_error_reform_impact( + self._reform_impacts.set_error_reform_impact( country_id=setup_options.country_id, policy_id=setup_options.reform_policy_id, baseline_policy_id=setup_options.baseline_policy_id, diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py new file mode 100644 index 000000000..0179dc2f9 --- /dev/null +++ b/policyengine_api/services/household_calculation_service.py @@ -0,0 +1,301 @@ +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from datetime import date +from typing import Any, Callable + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.local_models import Tracer +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, +) +from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs +from policyengine_api.utils.input_validation import find_unrecognized_inputs + + +@dataclass(frozen=True) +class CalculationResult: + household: dict + tracer_output: list[str] + + +@dataclass(frozen=True) +class HouseholdCalculationResult: + household: dict + warnings: tuple[str, ...] = () + cached: bool = False + + +class HouseholdNotFoundError(LookupError): + pass + + +class PolicyNotFoundError(LookupError): + pass + + +class InvalidHouseholdInputsError(ValueError): + def __init__(self, invalid_inputs: list[Any]) -> None: + self.invalid_inputs = invalid_inputs + super().__init__("Household or policy contains unrecognized inputs") + + +def get_household_year(household: dict) -> int | str: + household_year: int | str = date.today().year + household_age_list = list( + household.get("people", {}).get("you", {}).get("age", {}).keys() + ) + if household_age_list: + household_year = household_age_list[0] + return household_year + + +def add_yearly_variables( + household: dict, + country_id: str, + countries: dict | None = None, +) -> dict: + if countries is None: + from policyengine_api.country import COUNTRIES + + countries = COUNTRIES + metadata = countries.get(country_id).metadata + variables = metadata["variables"] + entities = metadata["entities"] + household_year = get_household_year(household) + + for variable in variables.values(): + if variable["definitionPeriod"] not in ("year", "month", "eternity"): + continue + entity_plural = entities[variable["entity"]]["plural"] + for entity in household.get(entity_plural, {}).values(): + if variable["name"] not in entity: + entity[variable["name"]] = { + household_year: ( + variable["defaultValue"] + if variable["isInputVariable"] + else None + ) + } + return household + + +class HouseholdCalculationService: + """Orchestrate stored-household calculations with short DB scopes.""" + + def __init__( + self, + primary_session_factory: sessionmaker[Session] | None = None, + local_session_factory: sessionmaker[Session] | None = None, + country_provider: Callable[[], dict] | None = None, + ) -> None: + self._injected_primary_session_factory = primary_session_factory + self._injected_local_session_factory = local_session_factory + self._country_provider = country_provider + + @property + def _primary_sessions(self) -> sessionmaker[Session]: + return self._injected_primary_session_factory or get_v1_session_factory() + + @property + def _local_sessions(self) -> sessionmaker[Session]: + return self._injected_local_session_factory or get_v1_session_factory( + local=True + ) + + def _countries(self) -> dict: + if self._country_provider is not None: + return self._country_provider() + from policyengine_api.country import COUNTRIES + + return COUNTRIES + + def _get_cached_household( + self, + country_id: str, + household_id: int, + policy_id: int, + api_version: str, + ) -> ComputedHousehold | None: + with self._local_sessions() as session: + return session.scalar( + select(ComputedHousehold).where( + ComputedHousehold.household_id == household_id, + ComputedHousehold.policy_id == policy_id, + ComputedHousehold.country_id == country_id, + ComputedHousehold.api_version == api_version, + ) + ) + + def _get_inputs( + self, + country_id: str, + household_id: int, + policy_id: int, + ) -> tuple[Household | None, Policy | None]: + with self._primary_sessions() as session: + household = session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) + policy = session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, + ) + ) + return household, policy + + def _store_result( + self, + country_id: str, + household_id: int, + policy_id: int, + api_version: str, + calculation: CalculationResult, + ) -> None: + with self._local_sessions.begin() as session: + identity = (household_id, policy_id, country_id) + computed_household = session.get(ComputedHousehold, identity) + if computed_household is None: + computed_household = ComputedHousehold( + country_id=country_id, + household_id=household_id, + policy_id=policy_id, + computed_household_json=calculation.household, + api_version=api_version, + status="complete", + ) + session.add(computed_household) + else: + computed_household.computed_household_json = calculation.household + computed_household.api_version = api_version + computed_household.status = "complete" + if calculation.tracer_output: + session.add( + Tracer( + household_id=household_id, + policy_id=policy_id, + country_id=country_id, + api_version=api_version, + tracer_output=calculation.tracer_output, + ) + ) + + def calculate_stored_household( + self, + country_id: str, + household_id: int, + policy_id: int, + ) -> HouseholdCalculationResult: + api_version = COUNTRY_PACKAGE_VERSIONS[country_id] + cached = self._get_cached_household( + country_id, + household_id, + policy_id, + api_version, + ) + if cached is not None: + return HouseholdCalculationResult( + household=cached.computed_household_json, + cached=True, + ) + + household, policy = self._get_inputs(country_id, household_id, policy_id) + if household is None: + raise HouseholdNotFoundError(household_id) + if policy is None: + raise PolicyNotFoundError(policy_id) + + countries = self._countries() + country = countries.get(country_id) + household_json = add_yearly_variables( + deepcopy(household.household_json), + country_id, + countries, + ) + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( + household_json, + policy.policy_json, + country.metadata, + ) + if invalid_inputs: + raise InvalidHouseholdInputsError(invalid_inputs) + + raw_calculation = country.calculate(household_json, policy.policy_json) + if isinstance(raw_calculation, CalculationResult): + calculation = raw_calculation + elif hasattr(raw_calculation, "household"): + calculation = CalculationResult( + household=raw_calculation.household, + tracer_output=raw_calculation.tracer_output, + ) + else: + # Temporary compatibility for test doubles and country packages + # that have not yet adopted CalculationResult. + calculation = CalculationResult( + household=raw_calculation, + tracer_output=[], + ) + self._store_result( + country_id, + household_id, + policy_id, + api_version, + calculation, + ) + return HouseholdCalculationResult( + household=calculation.household, + warnings=tuple(warning.message for warning in deprecated_inputs.warnings), + ) + + def calculate_household( + self, + country_id: str, + household_json: dict, + policy_json: dict, + *, + add_missing: bool = False, + ) -> HouseholdCalculationResult: + """Validate and calculate request-provided household and policy data.""" + countries = self._countries() + country = countries.get(country_id) + household_json = deepcopy(household_json) + if add_missing: + household_json = add_yearly_variables( + household_json, + country_id, + countries, + ) + + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( + household_json, + policy_json, + country.metadata, + ) + if invalid_inputs: + raise InvalidHouseholdInputsError(invalid_inputs) + + raw_calculation = country.calculate(household_json, policy_json) + household = ( + raw_calculation + if isinstance(raw_calculation, dict) + else raw_calculation.household + ) + return HouseholdCalculationResult( + household=household, + warnings=tuple(warning.message for warning in deprecated_inputs.warnings), + ) diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 2d3601737..441b9b078 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,141 +1,116 @@ -import json -from sqlalchemy.engine.row import Row +from __future__ import annotations + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker -from policyengine_api.data import database -from policyengine_api.utils import hash_object from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Household +from policyengine_api.utils import hash_object class HouseholdService: - def get_household(self, country_id: str, household_id: int) -> dict | None: - """ - Get a household's input data with a given ID. - - Args: - country_id (str): The country ID. - household_id (int): The household ID. - """ - print("Getting household data") + """Household operations with service-owned ORM transaction boundaries.""" - try: - if type(household_id) is not int or household_id < 0: - raise Exception( - f"Invalid household ID: {household_id}. Must be a positive integer." - ) + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory - row: Row | None = database.query( - f"SELECT * FROM household WHERE id = ? AND country_id = ?", - (household_id, country_id), - ).fetchone() + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() - # If row is present, we must JSON.loads the household_json - household = None - if row is not None: - household = dict(row) - if household["household_json"]: - household["household_json"] = json.loads( - household["household_json"] - ) - return household + def get_household( + self, + country_id: str, + household_id: int, + ) -> Household | None: + if type(household_id) is not int or household_id < 0: + raise Exception( + f"Invalid household ID: {household_id}. Must be a positive integer." + ) + with self._sessions() as session: + return self._get_household(session, country_id, household_id) - except Exception as e: - print(f"Error fetching household #{household_id}. Details: {str(e)}") - raise e + @staticmethod + def _get_household( + session: Session, + country_id: str, + household_id: int, + ) -> Household | None: + return session.scalar( + select(Household).where( + Household.country_id == country_id, + Household.id == household_id, + ) + ) def create_household( self, country_id: str, household_json: dict, label: str | None, - ) -> int: - """ - Create a new household with the given data. - - Args: - country_id (str): The country ID. - household_json (dict): The household data. - household_hash (int): The hash of the household data. - label (str): The label for the household. - api_version (str): The API version. - """ - - print("Creating new household") - - try: - household_hash: str = hash_object(household_json) - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - database.query( - f"INSERT INTO household (country_id, household_json, household_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - json.dumps(household_json), - household_hash, - label, - api_version, - ), + ) -> Household: + with self._sessions.begin() as session: + return self._create_household( + session, + country_id, + household_json, + label, ) - household_id = database.query( - f"SELECT id FROM household WHERE country_id = ? AND household_hash = ?", - (country_id, household_hash), - ).fetchone()["id"] - - return household_id - except Exception as e: - print(f"Error creating household. Details: {str(e)}") - raise e + @staticmethod + def _create_household( + session: Session, + country_id: str, + household_json: dict, + label: str | None, + ) -> Household: + household = Household( + country_id=country_id, + label=label, + household_json=household_json, + household_hash=hash_object(household_json), + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + ) + session.add(household) + session.flush() + return household def update_household( self, country_id: str, household_id: int, household_json: dict, - label: str, - ) -> dict: - """ - Update a household with the given data. - - Args: - country_id (str): The country ID. - household_id (int): The household ID. - payload (dict): The data to update the household with. - """ - print("Updating household") - - try: - household_hash: str = hash_object(household_json) - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - # WHERE must include country_id so an update scoped to - # one country cannot silently overwrite a household that - # happens to share the same numeric id under another - # country. - database.query( - "UPDATE household " - "SET household_json = ?, household_hash = ?, label = ?, api_version = ? " - "WHERE id = ? AND country_id = ?", - ( - json.dumps(household_json), - household_hash, - label, - api_version, - household_id, - country_id, - ), + label: str | None, + ) -> Household: + with self._sessions.begin() as session: + return self._update_household( + session, + country_id, + household_id, + household_json, + label, ) - # Fetch the updated JSON back from the table. If the - # household did not exist for this country, get_household - # returns None. - updated_household: dict | None = self.get_household( - country_id, household_id + @classmethod + def _update_household( + cls, + session: Session, + country_id: str, + household_id: int, + household_json: dict, + label: str | None, + ) -> Household: + household = cls._get_household(session, country_id, household_id) + if household is None: + raise LookupError( + f"Household #{household_id} not found for country {country_id}." ) - if updated_household is None: - raise LookupError( - f"Household #{household_id} not found for country {country_id}." - ) - return updated_household - except Exception as e: - print(f"Error updating household #{household_id}. Details: {str(e)}") - raise e + household.label = label + household.household_json = household_json + household.household_hash = hash_object(household_json) + household.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) + return household diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index bc63bc34c..728565882 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,211 +1,175 @@ -import json -from sqlalchemy.engine.row import Row +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker -from policyengine_api.data import database -from policyengine_api.utils import hash_object from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy +from policyengine_api.utils import hash_object class PolicyService: - """ - Service for storing and retrieving policies; - this is connected to the /policy route and - to the policy database table - """ - - def get_policy(self, country_id: str, policy_id: int) -> dict | None: - """ - Fetch policy based only on policy ID and country ID - - Arguments - country_id: str - policy_id: int - - Returns - dict | None -- the policy data, or None if not found - """ - print(f"Getting policy {policy_id} for {country_id}") - - try: - if type(policy_id) is not int or policy_id < 0: - raise Exception( - f"Invalid policy ID: {policy_id}. Must be a positive integer." - ) + """Policy operations with service-owned ORM transaction boundaries.""" - if not country_id: - # This will check for None and empty string - raise ValueError("country_id cannot be empty or None") - - # If no policy found, this will return None - row: Row | None = database.query( - "SELECT * FROM policy WHERE country_id = ? AND id = ?", - (country_id, policy_id), - ).fetchone() - - # policy_json is JSON and must be loaded, if present; to enable, - # we must convert the row to a dictionary - policy = None - if row: - policy = dict(row) - if policy["policy_json"]: - policy["policy_json"] = json.loads(policy["policy_json"]) - return policy - except Exception as e: - print(f"Error getting policy: {str(e)}") - raise e - - def get_policy_json(self, country_id: str, policy_id: int) -> str: - """ - Fetch policy JSON based only on policy ID and country ID - """ - print("Getting policy json") - try: - if type(policy_id) is not int or policy_id < 0: - raise Exception( - f"Invalid policy ID: {policy_id}. Must be a positive integer." - ) - policy = database.query( - f"SELECT policy_json FROM policy WHERE country_id = ? AND id = ?", - (country_id, policy_id), - ).fetchone() + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory - # Handle nonexisiting record case - if policy is None: - return None + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() - return policy["policy_json"] - except Exception as e: - print(f"Error getting policy json: {str(e)}") - raise e + @staticmethod + def _validate_policy_id(policy_id: int) -> None: + if type(policy_id) is not int or policy_id < 0: + raise Exception( + f"Invalid policy ID: {policy_id}. Must be a positive integer." + ) - def set_policy( - self, country_id: str, label: str, policy_json: dict - ) -> tuple[int, str, bool]: - """ - Insert a new policy into the database - - Arguments - country_id: str - label: str - policy_json: dict - - Returns - tuple[int, str, bool] -- the new policy ID, a message, and whether or not - the policy already existed - """ - print("Setting new policy") - - try: - # Convert country_id to lowercase - if not country_id.islower(): - country_id = country_id.lower() - - # Validate country_id - if country_id not in COUNTRY_PACKAGE_VERSIONS: - raise ValueError(f"Invalid country_id: {country_id}") - - policy_hash = hash_object(policy_json) - api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) - # Check if policy already exists - print("Checking if policy exists") - existing_policy = self._get_unique_policy_with_label( - country_id, policy_hash, label + def get_policy( + self, + country_id: str, + policy_id: int, + ) -> Policy | None: + self._validate_policy_id(policy_id) + if not country_id: + raise ValueError("country_id cannot be empty or None") + with self._sessions() as session: + return self._get_policy(session, country_id, policy_id) + + @staticmethod + def _get_policy( + session: Session, + country_id: str, + policy_id: int, + ) -> Policy | None: + return session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.id == policy_id, ) + ) - # If so, pass appropriate values back - if existing_policy: - print("Policy already exists") - return existing_policy["id"], "Policy already exists", True + def get_policy_json( + self, + country_id: str, + policy_id: int, + ) -> Any | None: + policy = self.get_policy(country_id, policy_id) + return None if policy is None else policy.policy_json - # Otherwise, insert the new policy... - print("Policy does not exist; creating new policy") - self._create_new_policy( - country_id, policy_json, policy_hash, label, api_version + def search_policies( + self, + country_id: str, + query: str = "", + *, + unique_only: bool = False, + ) -> list[Policy]: + with self._sessions() as session: + results = list( + session.scalars( + select(Policy).where( + Policy.country_id == country_id, + Policy.label.contains(query, autoescape=True), + ) + ) ) + if not unique_only: + return results + + unique_results = [] + processed_values = set() + for policy in results: + identity = policy.label, policy.policy_hash + if identity not in processed_values: + unique_results.append(policy) + processed_values.add(identity) + return unique_results - # And then fetch it back out; once an ORM is added, we can - # just return policy ID directly from the insert - new_policy = self._get_unique_policy_with_label( - country_id, policy_hash, label + def set_policy( + self, + country_id: str, + label: str | None, + policy_json: dict, + ) -> tuple[int, str, bool]: + country_id = country_id.lower() + if country_id not in COUNTRY_PACKAGE_VERSIONS: + raise ValueError(f"Invalid country_id: {country_id}") + + policy_hash = hash_object(policy_json) + with self._sessions.begin() as session: + return self._set_policy( + session, + country_id, + label, + policy_json, + policy_hash, ) - return int(new_policy["id"]), "Policy created", False - - except Exception as e: - print(f"Error setting policy: {str(e)}") - raise e + def _set_policy( + self, + session: Session, + country_id: str, + label: str | None, + policy_json: dict, + policy_hash: str, + ) -> tuple[int, str, bool]: + existing = self._get_unique_policy_with_label( + session, + country_id, + policy_hash, + label or None, + ) + if existing is not None: + return existing.id, "Policy already exists", True + + policy = Policy( + country_id=country_id, + label=label, + policy_json=policy_json, + policy_hash=policy_hash, + api_version=COUNTRY_PACKAGE_VERSIONS[country_id], + ) + session.add(policy) + session.flush() + return policy.id, "Policy created", False def _create_new_policy( self, + session: Session, country_id: str, policy_json: dict, policy_hash: str, label: str | None, api_version: str, - ) -> None: - """ - Create new policy and insert into database - - Arguments - country_id: str - policy_json: dict - policy_hash: str - label: str | None - api_version: str - - """ - try: - database.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - country_id, - json.dumps(policy_json), - policy_hash, - label, - api_version, - ), - ) - except Exception as e: - print(f"Error creating new policy: {str(e)}") - raise e + ) -> Policy: + policy = Policy( + country_id=country_id, + label=label, + policy_json=policy_json, + policy_hash=policy_hash, + api_version=api_version, + ) + session.add(policy) + session.flush() + return policy def _get_unique_policy_with_label( - self, country_id: str, policy_hash: str, label: str - ) -> dict | None: - """ - Given policy content (represented as a hash) and a label, fetch the policy; - this method ensures that both policy content and label are unique, a workaround - to an old issue whereby multiple copies of the same policy/label pair could be created - - Arguments - country_id: str - policy_hash: str - policy_label: str - - Returns - dict | None -- the policy data, or None if not found - """ - # The code in get_policy_with_label is a workaround - # to the fact that SQLite's cursor method does not properly - # convert 'WHERE x = None' to 'WHERE x IS NULL'; - # though SQLite supports searching and setting with 'WHERE - # x IS y', the production MySQL does not, requiring this - - # This workaround should be removed if and when a proper - # ORM package is added to the API, and this package's - # sanitization methods should be utilized instead - - try: - label_value = "IS NULL" if not label else "= ?" - args = [country_id, policy_hash] - if label: - args.append(label) - - policy = database.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label {label_value}", - tuple(args), - ).fetchone() - return policy - except Exception as e: - print(f"Error getting unique policy with label: {str(e)}") - raise e + self, + session: Session, + country_id: str, + policy_hash: str, + label: str | None, + ) -> Policy | None: + return session.scalar( + select(Policy).where( + Policy.country_id == country_id, + Policy.policy_hash == policy_hash, + Policy.label == label, + ) + ) diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 0f41352f3..d338e8687 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -1,13 +1,35 @@ -from policyengine_api.data import local_database import datetime +from typing import Any + +from sqlalchemy import delete, or_, select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReformImpact class ReformImpactsService: - """ - Service for storing and retrieving economy-wide reform impacts; - this is connected to the locally-stored reform_impact table - and no existing route - """ + """Reform-impact operations with service-owned local transactions.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory(local=True) + + def get_recent_reform_impacts(self, max_results: int) -> list[ReformImpact]: + with self._sessions() as session: + return list( + session.scalars( + select(ReformImpact) + .order_by(ReformImpact.start_time.desc()) + .limit(max_results) + ) + ) def get_all_reform_impacts( self, @@ -19,30 +41,19 @@ def get_all_reform_impacts( time_period, options_hash, api_version, - ): - try: - query = ( - "SELECT reform_impact_json, status, message, start_time, execution_id FROM " - "reform_impact WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "options_hash = ? AND api_version = ? AND dataset = ?" + ) -> list[ReformImpact]: + with self._sessions() as session: + return self._get_all_reform_impacts( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + api_version, ) - return local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - api_version, - dataset, - ), - ).fetchall() - except Exception as e: - print(f"Error getting all reform impacts: {str(e)}") - raise e def get_all_reform_impacts_by_options_hash_prefix( self, @@ -55,33 +66,20 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash, options_hash_prefix, api_version, - ): - try: - query = ( - "SELECT reform_impact_json, status, message, start_time, execution_id, options_hash FROM " - "reform_impact WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "(options_hash = ? OR options_hash LIKE ? ESCAPE '\\') AND api_version = ? AND dataset = ? " - "ORDER BY CASE WHEN options_hash = ? THEN 0 ELSE 1 END, start_time DESC" + ) -> list[ReformImpact]: + with self._sessions() as session: + return self._get_all_reform_impacts_by_options_hash_prefix( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + options_hash_prefix, + api_version, ) - return local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - options_hash_prefix, - api_version, - dataset, - options_hash, - ), - ).fetchall() - except Exception as e: - print(f"Error getting reform impacts by prefix: {str(e)}") - raise e def set_reform_impact( self, @@ -91,41 +89,31 @@ def set_reform_impact( region, dataset, time_period, - options, + options: dict[str, Any], options_hash, status, api_version, - reform_impact_json, + reform_impact_json: dict[str, Any], start_time, execution_id: str, - ): - try: - query = ( - "INSERT INTO reform_impact (country_id, reform_policy_id, baseline_policy_id, " - "region, dataset, time_period, options_json, options_hash, status, api_version, " - "reform_impact_json, start_time, execution_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) - local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options, - options_hash, - status, - api_version, - reform_impact_json, - start_time, - execution_id, - ), + ) -> ReformImpact: + with self._sessions.begin() as session: + return self._set_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options, + options_hash, + status, + api_version, + reform_impact_json, + start_time, + execution_id, ) - except Exception as e: - print(f"Error setting reform impact: {str(e)}") - raise e def delete_reform_impact( self, @@ -136,31 +124,19 @@ def delete_reform_impact( dataset, time_period, options_hash, - ): - try: - query = ( - "DELETE FROM reform_impact WHERE country_id = ? AND " - "reform_policy_id = ? AND baseline_policy_id = ? AND " - "region = ? AND time_period = ? AND options_hash = ? AND " - "dataset = ? AND status = 'computing'" + ) -> None: + with self._sessions.begin() as session: + self._delete_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, ) - local_database.query( - query, - ( - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - ), - ) - except Exception as e: - print(f"Error deleting reform impact: {str(e)}") - raise e - def set_error_reform_impact( self, country_id, @@ -172,38 +148,20 @@ def set_error_reform_impact( options_hash, message, execution_id: str, - ): - try: - query = ( - "UPDATE reform_impact SET status = ?, message = ?, end_time = ? WHERE " - "country_id = ? AND reform_policy_id = ? AND baseline_policy_id = ? AND " - "region = ? AND time_period = ? AND options_hash = ? AND dataset = ? AND " - "execution_id = ?" + ) -> ReformImpact | None: + with self._sessions.begin() as session: + return self._set_error_reform_impact( + session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + message, + execution_id, ) - local_database.query( - query, - ( - "error", - message, - datetime.datetime.strftime( - datetime.datetime.now(datetime.timezone.utc), - "%Y-%m-%d %H:%M:%S.%f", - ), - country_id, - policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - execution_id, - ), - ) - except Exception as e: - print( - f"Error setting error reform impact (something must be REALLY wrong): {str(e)}" - ) - raise e def set_complete_reform_impact( self, @@ -214,36 +172,245 @@ def set_complete_reform_impact( dataset, time_period, options_hash, - reform_impact_json, + reform_impact_json: dict[str, Any], execution_id, + ) -> ReformImpact | None: + with self._sessions.begin() as session: + return self._set_complete_reform_impact( + session, + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + reform_impact_json, + execution_id, + ) + + @staticmethod + def _filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version=None, ): - try: - query = ( - "UPDATE reform_impact SET status = ?, message = ?, end_time = ?, " - "reform_impact_json = ? WHERE country_id = ? AND reform_policy_id = ? AND " - "baseline_policy_id = ? AND region = ? AND time_period = ? AND " - "options_hash = ? AND dataset = ? AND execution_id = ?" + filters = { + "country_id": country_id, + "reform_policy_id": policy_id, + "baseline_policy_id": baseline_policy_id, + "region": region, + "dataset": dataset, + "time_period": time_period, + } + if api_version is not None: + filters["api_version"] = api_version + return filters + + @staticmethod + def _scope(statement, **filters): + return statement.where( + *(getattr(ReformImpact, key) == value for key, value in filters.items()) + ) + + def _get_all_reform_impacts( + self, + session: Session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + api_version, + ) -> list[ReformImpact]: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ) + statement = self._scope(select(ReformImpact), **filters).where( + ReformImpact.options_hash == options_hash + ) + return list(session.scalars(statement.order_by(ReformImpact.start_time.desc()))) + + def _get_all_reform_impacts_by_options_hash_prefix( + self, + session: Session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + options_hash_prefix, + api_version, + ) -> list[ReformImpact]: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + api_version, + ) + statement = self._scope(select(ReformImpact), **filters).where( + or_( + ReformImpact.options_hash == options_hash, + ReformImpact.options_hash.like(options_hash_prefix, escape="\\"), + ) + ) + return list( + session.scalars( + statement.order_by( + (ReformImpact.options_hash == options_hash).desc(), + ReformImpact.start_time.desc(), + ) ) - local_database.query( - query, - ( - "ok", - "Completed", - datetime.datetime.strftime( - datetime.datetime.now(datetime.timezone.utc), - "%Y-%m-%d %H:%M:%S.%f", - ), - reform_impact_json, - country_id, - reform_policy_id, - baseline_policy_id, - region, - time_period, - options_hash, - dataset, - execution_id, - ), + ) + + def _set_reform_impact( + self, + session: Session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options: dict[str, Any], + options_hash, + status, + api_version, + reform_impact_json: dict[str, Any], + start_time, + execution_id: str, + ) -> ReformImpact: + impact = ReformImpact( + country_id=country_id, + reform_policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + region=region, + dataset=dataset, + time_period=time_period, + options_json=options, + options_hash=options_hash, + status=status, + api_version=api_version, + reform_impact_json=reform_impact_json, + start_time=start_time, + execution_id=execution_id, + ) + session.add(impact) + session.flush() + return impact + + def _delete_reform_impact( + self, + session: Session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) -> None: + filters = self._filters( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + ) + session.execute( + self._scope(delete(ReformImpact), **filters).where( + ReformImpact.options_hash == options_hash, + ReformImpact.status == "computing", ) - except Exception as e: - print(f"Error setting completed reform impact: {str(e)}") - raise e + ) + + def _set_error_reform_impact( + self, + session: Session, + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + message, + execution_id: str, + ) -> ReformImpact | None: + del ( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) + impact = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if impact is None: + return None + impact.status = "error" + impact.message = message + impact.end_time = self._now() + return impact + + def _set_complete_reform_impact( + self, + session: Session, + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + reform_impact_json: dict[str, Any], + execution_id, + ) -> ReformImpact | None: + del ( + country_id, + reform_policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + ) + impact = session.scalar( + select(ReformImpact) + .where(ReformImpact.execution_id == execution_id) + .order_by(ReformImpact.reform_impact_id.desc()) + ) + if impact is None: + return None + impact.status = "ok" + impact.message = "Completed" + impact.reform_impact_json = reform_impact_json + impact.end_time = self._now() + return impact + + @staticmethod + def _now() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 9440cfdfd..90d23c6bc 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,97 +1,75 @@ -from sqlalchemy.engine.row import Row +from sqlalchemy.orm import Session -from policyengine_api.data import database +from policyengine_api.data.v1_models import LegacyReportOutputAlias, ReportOutput class ReportOutputAliasService: - def _get_report_output_row(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT id, country_id, simulation_1_id, simulation_2_id, year - FROM report_outputs - WHERE id = ? - """, - (report_output_id,), - ).fetchone() - return dict(row) if row is not None else None + """Legacy report-ID aliases persisted through mapped ORM models.""" - def get_alias(self, legacy_report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM legacy_report_output_aliases - WHERE legacy_report_output_id = ? - """, - (legacy_report_output_id,), - ).fetchone() - return dict(row) if row is not None else None + @staticmethod + def get_alias( + session: Session, legacy_report_output_id: int + ) -> LegacyReportOutputAlias | None: + return session.get(LegacyReportOutputAlias, legacy_report_output_id) def resolve_canonical_report_output_id( - self, requested_report_output_id: int + self, session: Session, requested_report_output_id: int ) -> int | None: - alias = self.get_alias(requested_report_output_id) + alias = self.get_alias(session, requested_report_output_id) if alias is not None: - canonical_report_output_id = alias["canonical_report_output_id"] - if self._get_report_output_row(canonical_report_output_id) is None: + canonical_id = alias.canonical_report_output_id + if session.get(ReportOutput, canonical_id) is None: raise ValueError( - "Alias points to missing canonical report output " - f"#{canonical_report_output_id}" + f"Alias points to missing canonical report output #{canonical_id}" ) - return canonical_report_output_id - - row: Row | None = database.query( - "SELECT id FROM report_outputs WHERE id = ?", - (requested_report_output_id,), - ).fetchone() - return row["id"] if row is not None else None + return canonical_id + report = session.get(ReportOutput, requested_report_output_id) + return report.id if report is not None else None def set_alias( self, + session: Session, legacy_report_output_id: int, canonical_report_output_id: int, ) -> bool: - legacy_report_output = self._get_report_output_row(legacy_report_output_id) - if legacy_report_output is None: + legacy = session.get(ReportOutput, legacy_report_output_id) + if legacy is None: raise ValueError( f"Legacy report output #{legacy_report_output_id} not found" ) - - canonical_report_output = self._get_report_output_row( - canonical_report_output_id - ) - if canonical_report_output is None: + canonical = session.get(ReportOutput, canonical_report_output_id) + if canonical is None: raise ValueError( f"Canonical report output #{canonical_report_output_id} not found" ) if legacy_report_output_id == canonical_report_output_id: raise ValueError("Legacy and canonical report outputs must be different") - - existing_alias = self.get_alias(legacy_report_output_id) - if existing_alias is not None: - if ( - existing_alias["canonical_report_output_id"] - == canonical_report_output_id - ): + existing = self.get_alias(session, legacy_report_output_id) + if existing is not None: + if existing.canonical_report_output_id == canonical_report_output_id: return True - raise ValueError( "Legacy report output alias already points to canonical report output " - f"#{existing_alias['canonical_report_output_id']}" + f"#{existing.canonical_report_output_id}" ) - - logical_key = ("country_id", "simulation_1_id", "simulation_2_id", "year") + logical_fields = ( + "country_id", + "simulation_1_id", + "simulation_2_id", + "year", + ) if any( - legacy_report_output[field] != canonical_report_output[field] - for field in logical_key + getattr(legacy, field) != getattr(canonical, field) + for field in logical_fields ): raise ValueError( "Legacy and canonical report outputs must describe the same report" ) - database.query( - """ - INSERT INTO legacy_report_output_aliases - (legacy_report_output_id, canonical_report_output_id) - VALUES (?, ?) - """, - (legacy_report_output_id, canonical_report_output_id), + session.add( + LegacyReportOutputAlias( + legacy_report_output_id=legacy_report_output_id, + canonical_report_output_id=canonical_report_output_id, + ) ) + session.flush() return True diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 38b5704fa..28bdd3e9e 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,40 +1,59 @@ -import uuid +import json +from dataclasses import dataclass from datetime import datetime, timezone -from sqlalchemy.engine.row import Row +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.data import database +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun, Simulation +from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.report_spec_service import ( ECONOMY_REPORT_KINDS, ReportSpec, ReportSpecService, ) -from policyengine_api.services.run_sync_utils import ( - determine_parent_pointers, - parse_json_field, - run_matches_report_result, - select_display_report_run, - serialize_json_field, -) from policyengine_api.services.simulation_service import SimulationService +@dataclass(frozen=True) +class ReportOutputView: + report_output: ReportOutput + display_run: ReportOutputRun | None + response_id: int | None = None + + +@dataclass(frozen=True) +class ReportCreateResult: + view: ReportOutputView + created: bool + + class ReportOutputService: - def __init__(self): + """Report-output orchestration with service-owned transactions.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory self.report_spec_service = ReportSpecService() - self.simulation_service = SimulationService() + self.report_run_service = ReportRunService() + self.simulation_service = SimulationService(session_factory) - def _lock_clause(self) -> str: - return "" if database.local else " FOR UPDATE" + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() - def _utc_timestamp(self) -> str: - return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + @staticmethod + def _utc_timestamp() -> datetime: + return datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) - def _format_run_timestamp(self, value) -> str | None: + @staticmethod + def format_run_timestamp(value: datetime | str | None) -> str | None: if value is None: return None - if isinstance(value, datetime): timestamp = value if timestamp.tzinfo is None: @@ -45,22 +64,17 @@ def _format_run_timestamp(self, value) -> str | None: .isoformat() .replace("+00:00", "Z") ) - timestamp = str(value).strip() if not timestamp: return None - normalized = timestamp.replace(" ", "T", 1) - parseable_timestamp = ( + parseable = ( f"{normalized[:-1]}+00:00" if normalized.endswith("Z") else normalized ) try: - parsed = datetime.fromisoformat(parseable_timestamp) + parsed = datetime.fromisoformat(parseable) except ValueError: - if "T" in normalized: - return normalized if normalized.endswith("Z") else f"{normalized}Z" - return f"{normalized}Z" - + return normalized if normalized.endswith("Z") else f"{normalized}Z" if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) return ( @@ -70,894 +84,574 @@ def _format_run_timestamp(self, value) -> str | None: .replace("+00:00", "Z") ) - def _get_report_output_row( - self, + @staticmethod + def _select_report_output( + session: Session, report_output_id: int, - *, - queryer=None, country_id: str | None = None, + *, for_update: bool = False, - ) -> dict | None: - queryer = queryer or database - query = "SELECT * FROM report_outputs WHERE id = ?" - params: list[int | str] = [report_output_id] + ) -> ReportOutput | None: + statement = select(ReportOutput).where(ReportOutput.id == report_output_id) if country_id is not None: - query += " AND country_id = ?" - params.append(country_id) + statement = statement.where(ReportOutput.country_id == country_id) if for_update: - query += self._lock_clause() - - row: Row | None = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None + statement = statement.with_for_update() + return session.scalar(statement) + + @staticmethod + def _list_runs_descending( + session: Session, report_output_id: int + ) -> list[ReportOutputRun]: + return list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + ) def _get_linked_simulations( self, - report_output: dict, + session: Session, + report_output: ReportOutput, *, - queryer=None, - bootstrap_dual_write_state: bool = False, - ) -> tuple[dict, dict | None]: - queryer = queryer or database - if bootstrap_dual_write_state: - simulation_1 = self.simulation_service._ensure_simulation_dual_write_state_in_transaction( - queryer, - report_output["simulation_1_id"], - country_id=report_output["country_id"], - ) - else: - simulation_1 = self.simulation_service._get_simulation_row( - report_output["simulation_1_id"], - queryer=queryer, - country_id=report_output["country_id"], + bootstrap_dual_write_state: bool, + ) -> tuple[Simulation, Simulation | None]: + def get_simulation(simulation_id: int) -> Simulation | None: + if bootstrap_dual_write_state: + try: + return self.simulation_service._ensure_simulation_dual_write_state( + session, + simulation_id, + report_output.country_id, + ) + except ValueError: + return None + return session.scalar( + select(Simulation).where( + Simulation.id == simulation_id, + Simulation.country_id == report_output.country_id, + ) ) + + simulation_1 = get_simulation(report_output.simulation_1_id) if simulation_1 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_1_id']}" + f"#{report_output.simulation_1_id}" ) - simulation_2 = None - if report_output["simulation_2_id"] is not None: - if bootstrap_dual_write_state: - simulation_2 = self.simulation_service._ensure_simulation_dual_write_state_in_transaction( - queryer, - report_output["simulation_2_id"], - country_id=report_output["country_id"], - ) - else: - simulation_2 = self.simulation_service._get_simulation_row( - report_output["simulation_2_id"], - queryer=queryer, - country_id=report_output["country_id"], - ) + if report_output.simulation_2_id is not None: + simulation_2 = get_simulation(report_output.simulation_2_id) if simulation_2 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_2_id']}" + f"#{report_output.simulation_2_id}" ) - return simulation_1, simulation_2 - def _require_simulation_exists( - self, - tx, - *, - country_id: str, - simulation_id: int, - ) -> dict: - simulation = self.simulation_service._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - ) - if simulation is None: - raise ValueError( - f"Report output references missing simulation #{simulation_id}" + @staticmethod + def _select_mutable_run( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] + ) -> ReportOutputRun | None: + if report_output.status == "running": + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + and run.status in {"pending", "running"} + ), + None, + ) + if active is not None: + return active + return next( + ( + run + for run in runs_descending + if run.status in {"pending", "running"} + ), + None, ) - return simulation - - def _list_report_runs_descending( - self, report_output_id: int, *, queryer=None - ) -> list[dict]: - queryer = queryer or database - rows = queryer.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence DESC - """, - (report_output_id,), - ).fetchall() - - runs = [] - for row in rows: - run = dict(row) - run["report_spec_snapshot_json"] = parse_json_field( - run.get("report_spec_snapshot_json") + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + ), + None, ) - runs.append(run) - return runs - - def _select_mutable_run( - self, report_output: dict, runs_descending: list[dict] - ) -> dict | None: - active_run_id = report_output.get("active_run_id") - if report_output["status"] == "running": - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id and run["status"] in ( - "pending", - "running", - ): - return run - for run in runs_descending: - if run["status"] in ("pending", "running"): - return run - return None - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run + if active is not None: + return active return runs_descending[0] if runs_descending else None - def _has_mutable_running_run(self, report_output: dict, *, queryer=None) -> bool: - runs_descending = self._list_report_runs_descending( - report_output["id"], queryer=queryer - ) - if not runs_descending: + @staticmethod + def _run_needs_timestamp_sync(run: ReportOutputRun, status: str) -> bool: + if run.requested_at is None: return True + if status in {"complete", "error"}: + return run.started_at is None or run.finished_at is None + if status == "running": + return run.started_at is None or run.finished_at is not None + return run.started_at is not None or run.finished_at is not None - active_run_id = report_output.get("active_run_id") - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run["status"] in ("pending", "running") - return False - - return any(run["status"] in ("pending", "running") for run in runs_descending) - - def _run_needs_timestamp_sync(self, run: dict, status: str) -> bool: - if run.get("requested_at") is None: + @staticmethod + def _has_mutable_running_run( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] + ) -> bool: + if not runs_descending: return True - if status in ("complete", "error"): - return run.get("started_at") is None or run.get("finished_at") is None - if status == "running": - return run.get("started_at") is None or run.get("finished_at") is not None - return run.get("started_at") is not None or run.get("finished_at") is not None - - def _with_display_run_timestamps( - self, report_output: dict, *, queryer=None - ) -> dict: - """ - Overlay selected run timestamps onto the legacy report response shape. - - This is a response-compatibility bridge for app-v2 while report output - reads still return a report_outputs row. The authoritative timestamp - values live on report_output_runs; this helper chooses the display run, - formats its requested/started/finished timestamps, and returns an - enriched copy of the report output dict. It intentionally does not - mutate database state. - - These timestamps describe the selected base report execution. They are - not user-report association metadata and should not be treated as a - user-specific "last run" value. - - TODO: When report output reads are cut over to canonical run-backed - resolution, move this projection into the final response serializer - instead of keeping it as an ad hoc enrichment helper. - """ - runs_descending = self._list_report_runs_descending( - report_output["id"], queryer=queryer - ) - display_run = select_display_report_run(report_output, runs_descending) - if display_run is None: - return report_output - - enriched_report_output = dict(report_output) - for field in ("requested_at", "started_at", "finished_at"): - enriched_report_output[field] = self._format_run_timestamp( - display_run.get(field) + if report_output.active_run_id is not None: + active = next( + ( + run + for run in runs_descending + if run.id == report_output.active_run_id + ), + None, ) - return enriched_report_output + return active is not None and active.status in {"pending", "running"} + return any(run.status in {"pending", "running"} for run in runs_descending) - def _derive_report_country_package_version( - self, - simulation_1: dict | None, - simulation_2: dict | None = None, + @staticmethod + def _derive_country_package_version( + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> str | None: versions = [ - simulation["api_version"] + simulation.api_version for simulation in (simulation_1, simulation_2) - if simulation is not None and simulation.get("api_version") is not None + if simulation is not None and simulation.api_version is not None ] - if not versions: - return None - if len(set(versions)) == 1: - return versions[0] - return None + return versions[0] if versions and len(set(versions)) == 1 else None def _build_version_manifest( self, - report_output: dict, + report_output: ReportOutput, report_spec: ReportSpec | None, - simulation_1: dict | None = None, - simulation_2: dict | None = None, + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> dict[str, str | None]: - resolved_dataset = None - if report_spec is not None and report_spec.report_kind in ECONOMY_REPORT_KINDS: - resolved_dataset = report_spec.dataset - return { - "country_package_version": self._derive_report_country_package_version( + "country_package_version": self._derive_country_package_version( simulation_1, simulation_2 ), "policyengine_version": None, "data_version": None, "runtime_app_name": None, - "report_cache_version": report_output.get("api_version"), + "report_cache_version": report_output.api_version, "simulation_cache_version": None, "requested_version_override": None, - "resolved_dataset": resolved_dataset, + "resolved_dataset": ( + report_spec.dataset + if report_spec is not None + and report_spec.report_kind in ECONOMY_REPORT_KINDS + else None + ), "resolved_options_hash": None, } - def _get_report_spec_status(self, report_spec: ReportSpec) -> str: - if report_spec.report_kind in ECONOMY_REPORT_KINDS: - return "backfilled_assumed" - return "explicit" + @staticmethod + def _report_spec_status(report_spec: ReportSpec) -> str: + return ( + "backfilled_assumed" + if report_spec.report_kind in ECONOMY_REPORT_KINDS + else "explicit" + ) - def _upsert_report_spec_in_transaction( + def _upsert_report_spec( self, - tx, - report_output: dict, - simulation_1: dict | None, - simulation_2: dict | None, + report_output: ReportOutput, + simulation_1: Simulation | None, + simulation_2: Simulation | None, ) -> ReportSpec | None: if simulation_1 is None: return None - try: report_spec = self.report_spec_service.build_report_spec( - report_output=report_output, - simulation_1=simulation_1, - simulation_2=simulation_2, - ) - except ValueError as exc: - print( - "Skipping report spec sync for report output " - f"#{report_output['id']}. Details: {str(exc)}" + report_output, simulation_1, simulation_2 ) + except ValueError: return None - - report_spec_status = self._get_report_spec_status(report_spec) - existing_spec = parse_json_field(report_output.get("report_spec_json")) - if ( - existing_spec != report_spec.model_dump() - or report_output.get("report_kind") != report_spec.report_kind - or report_output.get("report_spec_schema_version") != 1 - or report_output.get("report_spec_status") != report_spec_status - ): - tx.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, - report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - report_spec.report_kind, - report_spec.model_dump_json(), - 1, - report_spec_status, - report_output["id"], - ), - ) - report_output["report_kind"] = report_spec.report_kind - report_output["report_spec_json"] = report_spec.model_dump() - report_output["report_spec_schema_version"] = 1 - report_output["report_spec_status"] = report_spec_status - + report_output.report_kind = report_spec.report_kind + report_output.report_spec_json = report_spec.model_dump() + report_output.report_spec_schema_version = 1 + report_output.report_spec_status = self._report_spec_status(report_spec) return report_spec + @staticmethod def _run_matches_parent( - self, - run: dict, - report_output: dict, + run: ReportOutputRun, + report_output: ReportOutput, report_spec: ReportSpec | None, version_manifest: dict[str, str | None], ) -> bool: - expected_snapshot = ( - report_spec.model_dump() if report_spec is not None else None - ) return ( - run["status"] == report_output["status"] - and run.get("output") == report_output.get("output") - and run.get("error_message") == report_output.get("error_message") - and run.get("report_spec_snapshot_json") == expected_snapshot - and run.get("country_package_version") - == version_manifest["country_package_version"] - and run.get("policyengine_version") - == version_manifest["policyengine_version"] - and run.get("data_version") == version_manifest["data_version"] - and run.get("runtime_app_name") == version_manifest["runtime_app_name"] - and run.get("report_cache_version") - == version_manifest["report_cache_version"] - and run.get("simulation_cache_version") - == version_manifest["simulation_cache_version"] - and run.get("requested_version_override") - == version_manifest["requested_version_override"] - and run.get("resolved_dataset") == version_manifest["resolved_dataset"] - and run.get("resolved_options_hash") - == version_manifest["resolved_options_hash"] - ) - - def _insert_bootstrap_report_run( - self, - tx, - report_output: dict, - report_spec: ReportSpec | None, - version_manifest: dict[str, str | None], - ) -> None: - requested_at = self._utc_timestamp() - is_terminal = report_output["status"] in ("complete", "error") - has_started = report_output["status"] in ("running", "complete", "error") - started_at = requested_at if has_started else None - finished_at = requested_at if is_terminal else None - - tx.query( - """ - INSERT INTO report_output_runs ( - id, report_output_id, run_sequence, status, output, error_message, - trigger_type, requested_at, started_at, finished_at, source_run_id, - report_spec_snapshot_json, country_package_version, policyengine_version, - data_version, runtime_app_name, report_cache_version, - simulation_cache_version, requested_version_override, resolved_dataset, - resolved_options_hash - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - str(uuid.uuid4()), - report_output["id"], - 1, - report_output["status"], - serialize_json_field(report_output.get("output")), - report_output.get("error_message"), - "initial", - requested_at, - started_at, - finished_at, - None, - (report_spec.model_dump_json() if report_spec is not None else None), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["report_cache_version"], - version_manifest["simulation_cache_version"], - version_manifest["requested_version_override"], - version_manifest["resolved_dataset"], - version_manifest["resolved_options_hash"], - ), + run.status == report_output.status + and run.output == report_output.output + and run.error_message == report_output.error_message + and run.report_spec_snapshot_json + == (report_spec.model_dump() if report_spec else None) + and all( + getattr(run, field) == value + for field, value in version_manifest.items() + ) ) - def _update_report_run_in_transaction( + def _update_run_from_parent( self, - tx, - run_id: str, - report_output: dict, + run: ReportOutputRun, + report_output: ReportOutput, report_spec: ReportSpec | None, version_manifest: dict[str, str | None], - preserve_terminal_finished_at: bool = False, + *, + preserve_terminal_finished_at: bool, ) -> None: - fallback_timestamp = self._utc_timestamp() - timestamp_updates = [ - "requested_at = COALESCE(requested_at, started_at, finished_at, ?)" - ] - timestamp_values = [fallback_timestamp] - if report_output["status"] in ("complete", "error"): - finished_at = self._utc_timestamp() - timestamp_updates.append( - "started_at = COALESCE(started_at, finished_at, requested_at, ?)" + now = self._utc_timestamp() + run.requested_at = run.requested_at or run.started_at or run.finished_at or now + if report_output.status in {"complete", "error"}: + run.started_at = ( + run.started_at or run.finished_at or run.requested_at or now ) - timestamp_values.append(finished_at) - if preserve_terminal_finished_at: - timestamp_updates.append("finished_at = COALESCE(finished_at, ?)") - else: - timestamp_updates.append("finished_at = ?") - timestamp_values.append(finished_at) - elif report_output["status"] == "running": - started_at = self._utc_timestamp() - timestamp_updates.extend( - [ - "started_at = COALESCE(started_at, requested_at, ?)", - "finished_at = NULL", - ] - ) - timestamp_values.append(started_at) + if not preserve_terminal_finished_at or run.finished_at is None: + run.finished_at = now + elif report_output.status == "running": + run.started_at = run.started_at or run.requested_at or now + run.finished_at = None else: - timestamp_updates.extend(["started_at = NULL", "finished_at = NULL"]) - - tx.query( - f""" - UPDATE report_output_runs - SET status = ?, output = ?, error_message = ?, - {", ".join(timestamp_updates)}, - report_spec_snapshot_json = ?, country_package_version = ?, - policyengine_version = ?, data_version = ?, runtime_app_name = ?, - report_cache_version = ?, simulation_cache_version = ?, - requested_version_override = ?, resolved_dataset = ?, - resolved_options_hash = ? - WHERE id = ? - """, - ( - report_output["status"], - serialize_json_field(report_output.get("output")), - report_output.get("error_message"), - *timestamp_values, - (report_spec.model_dump_json() if report_spec is not None else None), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["report_cache_version"], - version_manifest["simulation_cache_version"], - version_manifest["requested_version_override"], - version_manifest["resolved_dataset"], - version_manifest["resolved_options_hash"], - run_id, - ), + run.started_at = None + run.finished_at = None + run.status = report_output.status + run.output = report_output.output + run.error_message = report_output.error_message + run.report_spec_snapshot_json = ( + report_spec.model_dump() if report_spec else None ) + for field, value in version_manifest.items(): + setattr(run, field, value) - def _sync_parent_pointers_in_transaction( - self, - tx, - report_output: dict, - runs_descending: list[dict], + @staticmethod + def _sync_parent_pointers( + report_output: ReportOutput, runs_descending: list[ReportOutputRun] ) -> None: - desired_active_run_id, desired_latest_successful_run_id = ( - determine_parent_pointers(report_output["status"], runs_descending) - ) - if ( - report_output.get("active_run_id") == desired_active_run_id - and report_output.get("latest_successful_run_id") - == desired_latest_successful_run_id - ): - return - - tx.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ( - desired_active_run_id, - desired_latest_successful_run_id, - report_output["id"], - ), + latest_successful = next( + (run.id for run in runs_descending if run.status == "complete"), None ) - report_output["active_run_id"] = desired_active_run_id - report_output["latest_successful_run_id"] = desired_latest_successful_run_id + if report_output.status in {"pending", "running"} and runs_descending: + report_output.active_run_id = runs_descending[0].id + else: + report_output.active_run_id = None + if report_output.status == "complete" and latest_successful is None: + latest_successful = runs_descending[0].id if runs_descending else None + report_output.latest_successful_run_id = latest_successful - def _ensure_report_output_dual_write_state_in_transaction( + def _ensure_report_output_dual_write_state( self, - tx, + session: Session, report_output_id: int, - *, country_id: str | None = None, - ) -> dict: - report_output = self._get_report_output_row( - report_output_id, - queryer=tx, - country_id=country_id, - for_update=True, + ) -> ReportOutput: + report_output = self._select_report_output( + session, report_output_id, country_id, for_update=True ) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") - try: simulation_1, simulation_2 = self._get_linked_simulations( + session, report_output, - queryer=tx, bootstrap_dual_write_state=True, ) - except ValueError as exc: - print( - "Skipping linked simulation sync for report output " - f"#{report_output_id}. Details: {str(exc)}" - ) + except ValueError: simulation_1, simulation_2 = None, None - - report_spec = self._upsert_report_spec_in_transaction( - tx, - report_output, - simulation_1, - simulation_2, + report_spec = self._upsert_report_spec( + report_output, simulation_1, simulation_2 ) - version_manifest = self._build_version_manifest( - report_output, - report_spec=report_spec, - simulation_1=simulation_1, - simulation_2=simulation_2, + manifest = self._build_version_manifest( + report_output, report_spec, simulation_1, simulation_2 ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx - ) - if not runs_descending: - self._insert_bootstrap_report_run( - tx, - report_output, - report_spec, - version_manifest, - ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx + runs = self._list_runs_descending(session, report_output_id) + if not runs: + self.report_run_service.create_report_output_run( + session, + report_output_id, + status=report_output.status, + output=report_output.output, + error_message=report_output.error_message, + trigger_type="initial", + report_spec_snapshot=( + report_spec.model_dump() if report_spec else None + ), + version_manifest=manifest, ) + runs = self._list_runs_descending(session, report_output_id) else: - mutable_run = self._select_mutable_run(report_output, runs_descending) - if mutable_run is not None: - run_matches_parent = self._run_matches_parent( - mutable_run, - report_output, - report_spec, - version_manifest, - ) - needs_timestamp_sync = self._run_needs_timestamp_sync( - mutable_run, report_output["status"] + mutable = self._select_mutable_run(report_output, runs) + if mutable is not None: + matches_result = ( + mutable.status == report_output.status + and mutable.output == report_output.output + and mutable.error_message == report_output.error_message ) - if not run_matches_parent or needs_timestamp_sync: - run_matches_result = run_matches_report_result( - mutable_run, report_output - ) - self._update_report_run_in_transaction( - tx, - run_id=mutable_run["id"], - report_output=report_output, - report_spec=report_spec, - version_manifest=version_manifest, - preserve_terminal_finished_at=run_matches_result, - ) - runs_descending = self._list_report_runs_descending( - report_output_id, queryer=tx + if not self._run_matches_parent( + mutable, report_output, report_spec, manifest + ) or self._run_needs_timestamp_sync(mutable, report_output.status): + self._update_run_from_parent( + mutable, + report_output, + report_spec, + manifest, + preserve_terminal_finished_at=matches_result, ) + session.flush() + runs = self._list_runs_descending(session, report_output_id) + self._sync_parent_pointers(report_output, runs) + session.flush() + return report_output - self._sync_parent_pointers_in_transaction(tx, report_output, runs_descending) - refreshed_report_output = self._get_report_output_row( - report_output_id, - queryer=tx, - country_id=country_id, + def _find_existing_report_output( + self, + session: Session, + country_id: str, + simulation_1_id: int, + simulation_2_id: int | None = None, + year: str = "2025", + ) -> ReportOutput | None: + return session.scalar( + select(ReportOutput) + .where( + ReportOutput.country_id == country_id, + ReportOutput.simulation_1_id == simulation_1_id, + ReportOutput.simulation_2_id == simulation_2_id, + ReportOutput.year == year, + ReportOutput.api_version == get_report_output_cache_version(country_id), + ) + .order_by(ReportOutput.id.desc()) ) - if refreshed_report_output is None: - raise ValueError(f"Report output #{report_output_id} not found after sync") - return self._with_display_run_timestamps(refreshed_report_output, queryer=tx) - def ensure_report_output_dual_write_state( - self, - report_output_id: int, - country_id: str | None = None, - ) -> dict: - return database.transaction( - lambda tx: self._ensure_report_output_dual_write_state_in_transaction( - tx, - report_output_id, - country_id=country_id, + @staticmethod + def _require_simulation( + session: Session, country_id: str, simulation_id: int + ) -> Simulation: + simulation = session.scalar( + select(Simulation).where( + Simulation.id == simulation_id, + Simulation.country_id == country_id, ) ) + if simulation is None: + raise ValueError( + f"Report output references missing simulation #{simulation_id}" + ) + return simulation - def get_stored_report_output( - self, country_id: str, report_output_id: int - ) -> dict | None: - """ - Get a stored report output row without aliasing to current runtime lineage. - - This is used by mutation paths that must address the originally - requested row. It still runs dual-write synchronization, so it may - bootstrap or repair run/spec metadata and returns the display-run - timestamp projection. It is therefore not a raw database read. - - TODO: Split raw storage lookup from synchronized response projection in - a later run-backed read migration PR. - """ - report_output = self._get_report_output_row( - report_output_id, country_id=country_id + def _create_report_output( + self, + session: Session, + country_id: str, + simulation_1_id: int, + simulation_2_id: int | None = None, + year: str = "2025", + ) -> ReportOutput: + existing = self._find_existing_report_output( + session, country_id, simulation_1_id, simulation_2_id, year ) - if report_output is None: - return None - return self.ensure_report_output_dual_write_state( - report_output_id, + if existing is not None: + return self._ensure_report_output_dual_write_state( + session, existing.id, country_id + ) + self._require_simulation(session, country_id, simulation_1_id) + if simulation_2_id is not None: + self._require_simulation(session, country_id, simulation_2_id) + report_output = ReportOutput( country_id=country_id, + simulation_1_id=simulation_1_id, + simulation_2_id=simulation_2_id, + api_version=get_report_output_cache_version(country_id), + status="pending", + year=year, + ) + session.add(report_output) + session.flush() + return self._ensure_report_output_dual_write_state( + session, report_output.id, country_id ) - def report_output_exists(self, country_id: str, report_output_id: int) -> bool: - return ( - self._get_report_output_row(report_output_id, country_id=country_id) - is not None + def _get_report_output( + self, session: Session, country_id: str, report_output_id: int + ) -> ReportOutput | None: + if type(report_output_id) is not int or report_output_id < 0: + raise Exception( + f"Invalid report output ID: {report_output_id}. " + "Must be a positive integer." + ) + return self._select_report_output(session, report_output_id, country_id) + + @staticmethod + def _is_current_report_output(report_output: ReportOutput) -> bool: + return report_output.api_version == get_report_output_cache_version( + report_output.country_id ) - def _is_current_report_output(self, report_output: dict) -> bool: - return report_output.get("api_version") == get_report_output_cache_version( - report_output["country_id"] + def _get_or_create_current_report_output( + self, session: Session, report_output: ReportOutput + ) -> ReportOutput: + existing = self._find_existing_report_output( + session, + report_output.country_id, + report_output.simulation_1_id, + report_output.simulation_2_id, + report_output.year, + ) + if existing is not None: + return self._ensure_report_output_dual_write_state( + session, existing.id, report_output.country_id + ) + return self._create_report_output( + session, + report_output.country_id, + report_output.simulation_1_id, + report_output.simulation_2_id, + report_output.year, ) - def _find_existing_report_output_row( + def _build_view( self, + session: Session, + report_output: ReportOutput, *, - country_id: str, - simulation_1_id: int, - simulation_2_id: int | None, - year: str, - queryer=None, - ) -> dict | None: - queryer = queryer or database - api_version = get_report_output_cache_version(country_id) - query = """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND year = ? AND api_version = ? - """ - params: list[int | str] = [country_id, simulation_1_id, year, api_version] - if simulation_2_id is not None: - query += " AND simulation_2_id = ?" - params.append(simulation_2_id) - else: - query += " AND simulation_2_id IS NULL" - query += " ORDER BY id DESC" - - row = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None - - def _get_or_create_current_report_output(self, report_output: dict) -> dict: - current_report = self.find_existing_report_output( - country_id=report_output["country_id"], - simulation_1_id=report_output["simulation_1_id"], - simulation_2_id=report_output["simulation_2_id"], - year=report_output["year"], - ) - if current_report is not None: - return self._with_display_run_timestamps(current_report) - - return self.create_report_output( - country_id=report_output["country_id"], - simulation_1_id=report_output["simulation_1_id"], - simulation_2_id=report_output["simulation_2_id"], - year=report_output["year"], + response_id: int | None = None, + ) -> ReportOutputView: + return ReportOutputView( + report_output=report_output, + display_run=self.report_run_service.select_display_run( + session, report_output + ), + response_id=response_id, ) - def _alias_report_output(self, report_output_id: int, report_output: dict) -> dict: - aliased_report = dict(report_output) - aliased_report["id"] = report_output_id - return aliased_report - - def find_existing_report_output( + def create_or_reuse_report_output( self, country_id: str, simulation_1_id: int, simulation_2_id: int | None = None, year: str = "2025", - ) -> dict | None: - """ - Find an existing report output with the same simulation IDs and year. - """ - print("Checking for existing report output") - - try: - existing_report = self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, + ) -> ReportCreateResult: + with self._sessions.begin() as session: + existing = self._find_existing_report_output( + session, + country_id, + simulation_1_id, + simulation_2_id, + year, ) - if existing_report is not None: - print(f"Found existing report output with ID: {existing_report['id']}") - return self.ensure_report_output_dual_write_state( - existing_report["id"], - country_id=country_id, + created = existing is None + report_output = ( + self._create_report_output( + session, + country_id, + simulation_1_id, + simulation_2_id, + year, ) - return None - - except Exception as e: - print(f"Error checking for existing report output. Details: {str(e)}") - raise e + if existing is None + else self._ensure_report_output_dual_write_state( + session, existing.id, country_id + ) + ) + return ReportCreateResult( + view=self._build_view(session, report_output), + created=created, + ) - def create_report_output( + def resolve_report_output( self, country_id: str, - simulation_1_id: int, - simulation_2_id: int | None = None, - year: str = "2025", - ) -> dict: - """ - Create a new report output record with pending status. - """ - print("Creating new report output") - api_version = get_report_output_cache_version(country_id) - - try: - - def tx_callback(tx): - existing_report = self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - queryer=tx, - ) - if existing_report is not None: - print( - f"Reusing existing report output with ID: {existing_report['id']}" - ) - return self._ensure_report_output_dual_write_state_in_transaction( - tx, - existing_report["id"], - country_id=country_id, - ) - - self._require_simulation_exists( - tx, - country_id=country_id, - simulation_id=simulation_1_id, - ) - if simulation_2_id is not None: - self._require_simulation_exists( - tx, - country_id=country_id, - simulation_id=simulation_2_id, - ) - - if simulation_2_id is not None: - tx.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - country_id, - simulation_1_id, - simulation_2_id, - api_version, - "pending", - year, - ), - ) - else: - tx.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?) - """, - ( - country_id, - simulation_1_id, - api_version, - "pending", - year, - ), - ) - - created_report = self._find_existing_report_output_row( - country_id=country_id, - simulation_1_id=simulation_1_id, - simulation_2_id=simulation_2_id, - year=year, - queryer=tx, - ) - if created_report is None: - raise Exception("Failed to retrieve created report output") - - print(f"Created report output with ID: {created_report['id']}") - return self._ensure_report_output_dual_write_state_in_transaction( - tx, - created_report["id"], - country_id=country_id, - ) - - return database.transaction(tx_callback) - - except Exception as e: - print(f"Error creating report output. Details: {str(e)}") - raise e - - def get_report_output(self, country_id: str, report_output_id: int) -> dict | None: - """ - Get a report output record by ID. - """ - print(f"Getting report output {report_output_id}") - - try: - if type(report_output_id) is not int or report_output_id < 0: - raise Exception( - f"Invalid report output ID: {report_output_id}. Must be a positive integer." - ) - - report_output = self._get_report_output_row( - report_output_id, - country_id=country_id, + report_output_id: int, + ) -> ReportOutputView | None: + if type(report_output_id) is not int or report_output_id < 0: + raise Exception( + f"Invalid report output ID: {report_output_id}. " + "Must be a positive integer." ) - if report_output is None: + with self._sessions.begin() as session: + requested = self._get_report_output(session, country_id, report_output_id) + if requested is None: return None - - if self._is_current_report_output(report_output): - return self.ensure_report_output_dual_write_state( - report_output_id, - country_id=country_id, + if self._is_current_report_output(requested): + report_output = self._ensure_report_output_dual_write_state( + session, report_output_id, country_id ) - - current_report = self._get_or_create_current_report_output(report_output) - return self._alias_report_output(report_output_id, current_report) - - except Exception as e: - print( - f"Error fetching report output #{report_output_id}. Details: {str(e)}" + response_id = None + else: + report_output = self._get_or_create_current_report_output( + session, requested + ) + response_id = report_output_id + return self._build_view( + session, + report_output, + response_id=response_id, ) - raise e def update_report_output( self, country_id: str, report_id: int, status: str | None = None, - output: str | None = None, + output: dict | list | str | None = None, error_message: str | None = None, - ) -> bool: - """ - Update a report output record with results or error. - """ - print(f"Updating report output {report_id}") - - try: - update_fields = [] - update_values = [] - - if status is not None: - update_fields.append("status = ?") - update_values.append(status) - - if output is not None: - update_fields.append("output = ?") - update_values.append(output) - - if error_message is not None: - update_fields.append("error_message = ?") - update_values.append(error_message) - - if not update_fields: - print("No fields to update") - return False - - def tx_callback(tx): - requested_report = self._get_report_output_row( - report_id, - queryer=tx, - country_id=country_id, - for_update=True, - ) - if requested_report is None: - raise ValueError(f"Report output #{report_id} not found") - - if status == "running" and not self._has_mutable_running_run( - requested_report, queryer=tx - ): - raise ValueError( - "Cannot mark report output running without an active " - "pending or running report run" - ) + ) -> ReportOutputView | None: + values = { + key: value + for key, value in { + "status": status, + "output": output, + "error_message": error_message, + }.items() + if value is not None + } + if not values: + return None + if isinstance(values.get("output"), str): + values["output"] = json.loads(values["output"]) + with self._sessions.begin() as session: + report_output = self._select_report_output( + session, report_id, country_id, for_update=True + ) + if report_output is None: + raise LookupError(f"Report output #{report_id} not found") + self._update_report_output( + session, + report_output, + values, + requested_status=status, + ) + return self._build_view(session, report_output) - tx.query( - f"UPDATE report_outputs SET {', '.join(update_fields)} WHERE id = ? AND country_id = ?", - (*update_values, report_id, country_id), - ) - self._ensure_report_output_dual_write_state_in_transaction( - tx, - report_id, - country_id=country_id, + def _update_report_output( + self, + session: Session, + report_output: ReportOutput, + values: dict, + *, + requested_status: str | None, + ) -> None: + if requested_status == "running": + runs = self._list_runs_descending(session, report_output.id) + if not self._has_mutable_running_run(report_output, runs): + raise ValueError( + "Cannot mark report output running without an active pending " + "or running report run" ) - - database.transaction(tx_callback) - - print(f"Successfully updated report output #{report_id}") - return True - - except Exception as e: - print(f"Error updating report output #{report_id}. Details: {str(e)}") - raise e + for field, value in values.items(): + setattr(report_output, field, value) + self._ensure_report_output_dual_write_state( + session, + report_output.id, + report_output.country_id, + ) diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index 9899f6cc9..e0c8127ff 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -1,12 +1,11 @@ -import json import uuid from datetime import datetime, timezone from typing import Any -from sqlalchemy.engine.row import Row +from sqlalchemy import func, select +from sqlalchemy.orm import Session -from policyengine_api.data import database -from policyengine_api.services.run_sync_utils import select_display_report_run +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun REPORT_RUN_VERSION_FIELDS = ( @@ -23,134 +22,137 @@ class ReportRunService: - def _utc_timestamp(self) -> str: - return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") - - def _serialize_json( - self, value: dict[str, Any] | list[Any] | str | None - ) -> str | None: - if value is None or isinstance(value, str): - return value - return json.dumps(value) - - def _parse_run_row(self, row: Row | dict | None) -> dict | None: - if row is None: - return None - - run = dict(row) - if isinstance(run.get("report_spec_snapshot_json"), str): - run["report_spec_snapshot_json"] = json.loads( - run["report_spec_snapshot_json"] - ) - return run + """Report-run operations performed through a caller-owned ORM Session.""" + + @staticmethod + def _matches_report_result(run: ReportOutputRun, report: ReportOutput) -> bool: + return ( + run.status == report.status + and run.output == report.output + and run.error_message == report.error_message + ) def create_report_output_run( self, + session: Session, report_output_id: int, status: str = "pending", trigger_type: str = "initial", - output: dict[str, Any] | list[Any] | str | None = None, + output: dict[str, Any] | list[Any] | None = None, error_message: str | None = None, source_run_id: str | None = None, - report_spec_snapshot: dict[str, Any] | str | None = None, + report_spec_snapshot: dict[str, Any] | None = None, version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, - ) -> dict: - run_id = run_id or str(uuid.uuid4()) - version_manifest = version_manifest or {} - lock_clause = "" if database.local else " FOR UPDATE" - - def create_run_transaction(tx) -> None: - parent_row: Row | None = tx.query( - f"SELECT id FROM report_outputs WHERE id = ?{lock_clause}", - (report_output_id,), - ).fetchone() - if parent_row is None: - raise ValueError(f"Report output #{report_output_id} not found") + ) -> ReportOutputRun: + parent = session.scalar( + select(ReportOutput) + .where(ReportOutput.id == report_output_id) + .with_for_update() + ) + if parent is None: + raise ValueError(f"Report output #{report_output_id} not found") + sequence = ( + session.scalar( + select(func.max(ReportOutputRun.run_sequence)).where( + ReportOutputRun.report_output_id == report_output_id + ) + ) + or 0 + ) + 1 + now = datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) + values = { + "status": status, + "output": output, + "error_message": error_message, + "trigger_type": trigger_type, + "requested_at": now, + "started_at": now if status in {"running", "complete", "error"} else None, + "finished_at": now if status in {"complete", "error"} else None, + "source_run_id": source_run_id, + "report_spec_snapshot_json": report_spec_snapshot, + } + values.update( + { + field: (version_manifest or {}).get(field) + for field in REPORT_RUN_VERSION_FIELDS + } + ) + run = ReportOutputRun( + id=run_id or str(uuid.uuid4()), + report_output_id=report_output_id, + run_sequence=sequence, + **values, + ) + session.add(run) + session.flush() + return run - run_sequence_row: Row | None = tx.query( - """ - SELECT COALESCE(MAX(run_sequence), 0) AS max_run_sequence - FROM report_output_runs - WHERE report_output_id = ? - """, - (report_output_id,), - ).fetchone() - run_sequence = ( - int(run_sequence_row["max_run_sequence"]) + 1 - if run_sequence_row is not None - else 1 + def get_report_output_run( + self, session: Session, run_id: str + ) -> ReportOutputRun | None: + return session.get(ReportOutputRun, run_id) + + def list_report_output_runs( + self, session: Session, report_output_id: int + ) -> list[ReportOutputRun]: + return list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.asc()) ) + ) - requested_at = self._utc_timestamp() - is_terminal = status in ("complete", "error") - has_started = status in ("running", "complete", "error") - started_at = requested_at if has_started else None - finished_at = requested_at if is_terminal else None + def get_newest_report_output_run( + self, session: Session, report_output_id: int + ) -> ReportOutputRun | None: + return session.scalar( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output_id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) - tx.query( - f""" - INSERT INTO report_output_runs ( - id, report_output_id, run_sequence, status, output, error_message, - trigger_type, requested_at, started_at, finished_at, source_run_id, - report_spec_snapshot_json, {", ".join(REPORT_RUN_VERSION_FIELDS)} - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, + def select_display_run( + self, session: Session, report_output: ReportOutput + ) -> ReportOutputRun | None: + runs = list( + session.scalars( + select(ReportOutputRun) + .where(ReportOutputRun.report_output_id == report_output.id) + .order_by(ReportOutputRun.run_sequence.desc()) + ) + ) + if report_output.active_run_id is not None: + active = next( + (run for run in runs if run.id == report_output.active_run_id), None + ) + if active is not None: + return active + if report_output.status == "error": + matching_error = next( ( - run_id, - report_output_id, - run_sequence, - status, - self._serialize_json(output), - error_message, - trigger_type, - requested_at, - started_at, - finished_at, - source_run_id, - self._serialize_json(report_spec_snapshot), - *[ - version_manifest.get(field) - for field in REPORT_RUN_VERSION_FIELDS - ], + run + for run in runs + if self._matches_report_result(run, report_output) ), + None, ) - - database.transaction(create_run_transaction) - return self.get_report_output_run(run_id) - - def get_report_output_run(self, run_id: str) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (run_id,), - ).fetchone() - return self._parse_run_row(row) - - def list_report_output_runs(self, report_output_id: int) -> list[dict]: - rows = database.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (report_output_id,), - ).fetchall() - return [self._parse_run_row(row) for row in rows] - - def get_newest_report_output_run(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence DESC - LIMIT 1 - """, - (report_output_id,), - ).fetchone() - return self._parse_run_row(row) - - def select_display_run(self, report_output: dict) -> dict | None: - runs_descending = list( - reversed(self.list_report_output_runs(report_output["id"])) + if matching_error is not None: + return matching_error + if report_output.latest_successful_run_id is not None: + successful = next( + ( + run + for run in runs + if run.id == report_output.latest_successful_run_id + ), + None, + ) + if successful is not None: + return successful + matching = next( + (run for run in runs if self._matches_report_result(run, report_output)), + None, ) - return select_display_report_run(report_output, runs_descending) + return matching or (runs[0] if runs else None) diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index b81cc566f..38b014257 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -2,9 +2,10 @@ from typing import Any, Literal from pydantic import BaseModel, Field -from sqlalchemy.engine.row import Row +from sqlalchemy.orm import Session + +from policyengine_api.data.v1_models import ReportOutput, Simulation -from policyengine_api.data import database REPORT_SPEC_SCHEMA_VERSION = 1 REPORT_SPEC_STATUSES = {"explicit", "backfilled_assumed"} @@ -42,265 +43,99 @@ class EconomyReportSpec(BaseModel): class ReportSpecService: - def _validate_schema_version(self, schema_version: int | None) -> None: + @staticmethod + def _validate_schema_version(schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: raise ValueError( f"Unsupported report spec schema version: {schema_version}" ) - def _get_report_output_row(self, report_output_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output_id,), - ).fetchone() - return dict(row) if row is not None else None - - def _get_simulation_row(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_id,), - ).fetchone() - return dict(row) if row is not None else None - - def _get_linked_simulations(self, report_output: dict) -> tuple[dict, dict | None]: - simulation_1 = self._get_simulation_row(report_output["simulation_1_id"]) + @staticmethod + def _get_linked_simulations( + session: Session, report_output: ReportOutput + ) -> tuple[Simulation, Simulation | None]: + simulation_1 = session.get(Simulation, report_output.simulation_1_id) if simulation_1 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_1_id']}" + f"#{report_output.simulation_1_id}" ) - simulation_2 = None - if report_output["simulation_2_id"] is not None: - simulation_2 = self._get_simulation_row(report_output["simulation_2_id"]) + if report_output.simulation_2_id is not None: + simulation_2 = session.get(Simulation, report_output.simulation_2_id) if simulation_2 is None: raise ValueError( "Report output references missing simulation " - f"#{report_output['simulation_2_id']}" + f"#{report_output.simulation_2_id}" ) - return simulation_1, simulation_2 + @staticmethod def _validate_report_simulation_linkage( - self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None, ) -> None: - if simulation_1.get("id") != report_output["simulation_1_id"]: + if simulation_1.id != report_output.simulation_1_id: raise ValueError( "Simulation 1 must match report_output.simulation_1_id to build a " "report spec" ) - - report_simulation_2_id = report_output["simulation_2_id"] - if report_simulation_2_id is None: + if report_output.simulation_2_id is None: if simulation_2 is not None: raise ValueError("Report output does not reference a second simulation") return - if simulation_2 is None: raise ValueError( "Report output requires a second simulation to build a comparison " "report spec" ) - if simulation_2.get("id") != report_simulation_2_id: + if simulation_2.id != report_output.simulation_2_id: raise ValueError( "Simulation 2 must match report_output.simulation_2_id to build a " "report spec" ) + @staticmethod def _validate_report_country( - self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None, ) -> None: - report_country_id = report_output["country_id"] - if simulation_1["country_id"] != report_country_id: + if simulation_1.country_id != report_output.country_id: raise ValueError( "Simulation 1 country must match report output country to build a " "report spec" ) - if simulation_2 is not None and simulation_2["country_id"] != report_country_id: - raise ValueError( - "Simulation 2 country must match report output country to build a " - "report spec" - ) - - def _build_household_report_spec( - self, - report_output: dict, - report_kind: str, - simulation_1: dict, - simulation_2: dict | None, - time_period: str, - ) -> HouseholdReportSpec: - if simulation_1["population_type"] != "household": - raise ValueError("Household report specs require household simulations") - if ( - simulation_2 is not None - and simulation_2["population_id"] != simulation_1["population_id"] - ): - raise ValueError( - "Household comparison report specs require matching household IDs" - ) - - return HouseholdReportSpec.model_validate( - { - "country_id": report_output["country_id"], - "report_kind": report_kind, - "time_period": time_period, - "simulation_1": { - "population_type": simulation_1["population_type"], - "population_id": simulation_1["population_id"], - "policy_id": simulation_1["policy_id"], - }, - "simulation_2": ( - { - "population_type": simulation_2["population_type"], - "population_id": simulation_2["population_id"], - "policy_id": simulation_2["policy_id"], - } - if simulation_2 is not None - else None - ), - } - ) - - def _build_economy_report_spec( - self, - report_output: dict, - report_kind: str, - simulation_1: dict, - simulation_2: dict | None, - time_period: str, - dataset: str, - target: Literal["general", "cliff"], - options: dict[str, Any] | None, - ) -> EconomyReportSpec: - if simulation_1["population_type"] != "geography": - raise ValueError("Economy report specs require geography simulations") if ( simulation_2 is not None - and simulation_2["population_id"] != simulation_1["population_id"] + and simulation_2.country_id != report_output.country_id ): raise ValueError( - "Economy comparison report specs require matching geography IDs" - ) - - return EconomyReportSpec.model_validate( - { - "country_id": report_output["country_id"], - "report_kind": report_kind, - "time_period": time_period, - "region": simulation_1["population_id"], - "baseline_policy_id": simulation_1["policy_id"], - "reform_policy_id": ( - simulation_2["policy_id"] - if simulation_2 is not None - else simulation_1["policy_id"] - ), - "dataset": dataset, - "target": target, - "options": options or {}, - } - ) - - def _validate_report_spec_matches_row( - self, report_output: dict, report_spec: ReportSpec - ) -> None: - simulation_1, simulation_2 = self._get_linked_simulations(report_output) - inferred_report_kind = self.infer_report_kind(simulation_1, simulation_2) - if report_spec.country_id != report_output["country_id"]: - raise ValueError("Report spec country must match report output country") - if report_spec.time_period != report_output["year"]: - raise ValueError("Report spec time_period must match report output year") - if report_spec.report_kind != inferred_report_kind: - raise ValueError("Report spec kind must match linked simulations") - - if isinstance(report_spec, HouseholdReportSpec): - if report_spec.simulation_1.model_dump() != { - "population_type": simulation_1["population_type"], - "population_id": simulation_1["population_id"], - "policy_id": simulation_1["policy_id"], - }: - raise ValueError( - "Report spec simulation_1 must match linked simulation 1" - ) - - expected_simulation_2 = ( - { - "population_type": simulation_2["population_type"], - "population_id": simulation_2["population_id"], - "policy_id": simulation_2["policy_id"], - } - if simulation_2 is not None - else None - ) - actual_simulation_2 = ( - report_spec.simulation_2.model_dump() - if report_spec.simulation_2 is not None - else None - ) - if actual_simulation_2 != expected_simulation_2: - raise ValueError( - "Report spec simulation_2 must match linked simulation 2" - ) - return - - expected_region = simulation_1["population_id"] - expected_baseline_policy_id = simulation_1["policy_id"] - expected_reform_policy_id = ( - simulation_2["policy_id"] - if simulation_2 is not None - else simulation_1["policy_id"] - ) - - if report_spec.region != expected_region: - raise ValueError("Report spec region must match linked simulations") - if report_spec.baseline_policy_id != expected_baseline_policy_id: - raise ValueError( - "Report spec baseline_policy_id must match linked simulations" - ) - if report_spec.reform_policy_id != expected_reform_policy_id: - raise ValueError( - "Report spec reform_policy_id must match linked simulations" + "Simulation 2 country must match report output country to build a " + "report spec" ) + @staticmethod def infer_report_kind( - self, - simulation_1: dict, - simulation_2: dict | None = None, + simulation_1: Simulation, simulation_2: Simulation | None = None ) -> str: - population_type = simulation_1["population_type"] - if ( - simulation_2 is not None - and simulation_2["population_type"] != population_type - ): + population_type = simulation_1.population_type + if simulation_2 is not None and simulation_2.population_type != population_type: raise ValueError( "Simulation population types must match to build a report spec" ) - if population_type == "household": - return ( - "household_comparison" - if simulation_2 is not None - else "household_single" - ) - + return "household_comparison" if simulation_2 else "household_single" if population_type == "geography": - return ( - "economy_comparison" if simulation_2 is not None else "economy_single" - ) - + return "economy_comparison" if simulation_2 else "economy_single" raise ValueError(f"Unsupported simulation population type: {population_type}") def build_report_spec( self, - report_output: dict, - simulation_1: dict, - simulation_2: dict | None = None, + report_output: ReportOutput, + simulation_1: Simulation, + simulation_2: Simulation | None = None, dataset: str = "default", target: Literal["general", "cliff"] = "general", options: dict[str, Any] | None = None, @@ -308,57 +143,115 @@ def build_report_spec( self._validate_report_simulation_linkage( report_output, simulation_1, simulation_2 ) - report_kind = self.infer_report_kind(simulation_1, simulation_2) - time_period = report_output["year"] self._validate_report_country(report_output, simulation_1, simulation_2) - + report_kind = self.infer_report_kind(simulation_1, simulation_2) if report_kind in HOUSEHOLD_REPORT_KINDS: - return self._build_household_report_spec( - report_output=report_output, + if ( + simulation_2 is not None + and simulation_2.population_id != simulation_1.population_id + ): + raise ValueError( + "Household comparison report specs require matching household IDs" + ) + return HouseholdReportSpec( + country_id=report_output.country_id, report_kind=report_kind, - simulation_1=simulation_1, - simulation_2=simulation_2, - time_period=time_period, + time_period=report_output.year, + simulation_1=ReportSimulationInput( + population_type=simulation_1.population_type, + population_id=simulation_1.population_id, + policy_id=simulation_1.policy_id, + ), + simulation_2=( + ReportSimulationInput( + population_type=simulation_2.population_type, + population_id=simulation_2.population_id, + policy_id=simulation_2.policy_id, + ) + if simulation_2 + else None + ), ) - - return self._build_economy_report_spec( - report_output=report_output, + if ( + simulation_2 is not None + and simulation_2.population_id != simulation_1.population_id + ): + raise ValueError( + "Economy comparison report specs require matching geography IDs" + ) + return EconomyReportSpec( + country_id=report_output.country_id, report_kind=report_kind, - simulation_1=simulation_1, - simulation_2=simulation_2, - time_period=time_period, + time_period=report_output.year, + region=simulation_1.population_id, + baseline_policy_id=simulation_1.policy_id, + reform_policy_id=( + simulation_2.policy_id if simulation_2 else simulation_1.policy_id + ), dataset=dataset, target=target, - options=options, + options=options or {}, ) - def _parse_json_field(self, value: str | dict | None) -> dict | None: - if value is None: - return None - if isinstance(value, str): - return json.loads(value) - return value + def _validate_report_spec_matches_model( + self, + session: Session, + report_output: ReportOutput, + report_spec: ReportSpec, + ) -> None: + simulation_1, simulation_2 = self._get_linked_simulations( + session, report_output + ) + expected = self.build_report_spec( + report_output, + simulation_1, + simulation_2, + dataset=( + report_spec.dataset + if isinstance(report_spec, EconomyReportSpec) + else "default" + ), + target=( + report_spec.target + if isinstance(report_spec, EconomyReportSpec) + else "general" + ), + options=( + report_spec.options + if isinstance(report_spec, EconomyReportSpec) + else None + ), + ) + if report_spec != expected: + raise ValueError("Report spec must match the linked report and simulations") - def _parse_report_spec(self, report_kind: str, raw_spec: dict) -> ReportSpec: + @staticmethod + def _parse_report_spec(report_kind: str, raw_spec: dict) -> ReportSpec: if report_kind in HOUSEHOLD_REPORT_KINDS: return HouseholdReportSpec.model_validate(raw_spec) if report_kind in ECONOMY_REPORT_KINDS: return EconomyReportSpec.model_validate(raw_spec) raise ValueError(f"Unsupported report kind: {report_kind}") - def get_report_spec(self, report_output_id: int) -> ReportSpec | None: - report_output = self._get_report_output_row(report_output_id) - if report_output is None or report_output["report_spec_json"] is None: + def get_report_spec( + self, session: Session, report_output_id: int + ) -> ReportSpec | None: + report_output = session.get(ReportOutput, report_output_id) + if report_output is None or report_output.report_spec_json is None: return None - - self._validate_schema_version(report_output["report_spec_schema_version"]) - raw_spec = self._parse_json_field(report_output["report_spec_json"]) - report_spec = self._parse_report_spec(report_output["report_kind"], raw_spec) - self._validate_report_spec_matches_row(report_output, report_spec) + self._validate_schema_version(report_output.report_spec_schema_version) + raw_spec = report_output.report_spec_json + if isinstance(raw_spec, str): + # Existing databases may contain pre-ORM JSON text. New writes below + # always assign Python objects and leave conversion to SQLAlchemy. + raw_spec = json.loads(raw_spec) + report_spec = self._parse_report_spec(report_output.report_kind, raw_spec) + self._validate_report_spec_matches_model(session, report_output, report_spec) return report_spec def set_report_spec( self, + session: Session, report_output_id: int, report_spec: ReportSpec, report_spec_status: Literal["explicit", "backfilled_assumed"], @@ -367,24 +260,12 @@ def set_report_spec( if report_spec_status not in REPORT_SPEC_STATUSES: raise ValueError(f"Unsupported report spec status: {report_spec_status}") self._validate_schema_version(schema_version) - - report_output = self._get_report_output_row(report_output_id) + report_output = session.get(ReportOutput, report_output_id) if report_output is None: raise ValueError(f"Report output #{report_output_id} not found") - self._validate_report_spec_matches_row(report_output, report_spec) - - database.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - report_spec.report_kind, - report_spec.model_dump_json(), - schema_version, - report_spec_status, - report_output_id, - ), - ) + self._validate_report_spec_matches_model(session, report_output, report_spec) + report_output.report_kind = report_spec.report_kind + report_output.report_spec_json = report_spec.model_dump() + report_output.report_spec_schema_version = schema_version + report_output.report_spec_status = report_spec_status return True diff --git a/policyengine_api/services/simulation_analysis_service.py b/policyengine_api/services/simulation_analysis_service.py index 8738dd625..bb6382c87 100644 --- a/policyengine_api/services/simulation_analysis_service.py +++ b/policyengine_api/services/simulation_analysis_service.py @@ -12,9 +12,6 @@ class SimulationAnalysisService(AIAnalysisService): analysis database table """ - def __init__(self): - super().__init__() - def execute_analysis( self, country_id: str, @@ -63,7 +60,7 @@ def execute_analysis( # streaming response existing_analysis = self.get_existing_analysis(prompt) if existing_analysis is not None: - return existing_analysis, "static" + return existing_analysis.analysis, "static" print("Found no existing AI analysis; triggering new analysis with Claude") # Otherwise, pass prompt to Claude, then return streaming function diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 544aca9c2..61b3a970b 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -1,10 +1,10 @@ -import json import uuid from typing import Any -from sqlalchemy.engine.row import Row +from sqlalchemy import func, select +from sqlalchemy.orm import Session -from policyengine_api.data import database +from policyengine_api.data.v1_models import Simulation, SimulationRun SIMULATION_RUN_VERSION_FIELDS = ( @@ -17,26 +17,9 @@ class SimulationRunService: - def _serialize_json( - self, value: dict[str, Any] | list[Any] | str | None - ) -> str | None: - if value is None or isinstance(value, str): - return value - return json.dumps(value) - - def _parse_run_row(self, row: Row | dict | None) -> dict | None: - if row is None: - return None - - run = dict(row) - if isinstance(run.get("simulation_spec_snapshot_json"), str): - run["simulation_spec_snapshot_json"] = json.loads( - run["simulation_spec_snapshot_json"] - ) - return run - def create_simulation_run( self, + session: Session, simulation_id: int, report_output_run_id: str | None = None, input_position: int | None = None, @@ -48,106 +31,94 @@ def create_simulation_run( simulation_spec_snapshot: dict[str, Any] | str | None = None, version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, - ) -> dict: - run_id = run_id or str(uuid.uuid4()) - version_manifest = version_manifest or {} - lock_clause = "" if database.local else " FOR UPDATE" - - def create_run_transaction(tx) -> None: - parent_row: Row | None = tx.query( - f"SELECT id FROM simulations WHERE id = ?{lock_clause}", - (simulation_id,), - ).fetchone() - if parent_row is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - run_sequence_row: Row | None = tx.query( - """ - SELECT COALESCE(MAX(run_sequence), 0) AS max_run_sequence - FROM simulation_runs - WHERE simulation_id = ? - """, - (simulation_id,), - ).fetchone() - run_sequence = ( - int(run_sequence_row["max_run_sequence"]) + 1 - if run_sequence_row is not None - else 1 + ) -> SimulationRun: + parent = session.scalar( + select(Simulation).where(Simulation.id == simulation_id).with_for_update() + ) + if parent is None: + raise ValueError(f"Simulation #{simulation_id} not found") + sequence = ( + session.scalar( + select(func.max(SimulationRun.run_sequence)).where( + SimulationRun.simulation_id == simulation_id + ) ) + or 0 + ) + 1 + values = { + "report_output_run_id": report_output_run_id, + "input_position": input_position, + "status": status, + "output": output, + "error_message": error_message, + "trigger_type": trigger_type, + "requested_at": None, + "started_at": None, + "finished_at": None, + "source_run_id": source_run_id, + "simulation_spec_snapshot_json": simulation_spec_snapshot, + } + values.update( + { + field: (version_manifest or {}).get(field) + for field in SIMULATION_RUN_VERSION_FIELDS + } + ) + run = SimulationRun( + id=run_id or str(uuid.uuid4()), + simulation_id=simulation_id, + run_sequence=sequence, + **values, + ) + session.add(run) + session.flush() + return run - tx.query( - f""" - INSERT INTO simulation_runs ( - id, simulation_id, report_output_run_id, input_position, run_sequence, - status, output, error_message, trigger_type, requested_at, started_at, - finished_at, source_run_id, simulation_spec_snapshot_json, - {", ".join(SIMULATION_RUN_VERSION_FIELDS)} - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - run_id, - simulation_id, - report_output_run_id, - input_position, - run_sequence, - status, - self._serialize_json(output), - error_message, - trigger_type, - None, - None, - None, - source_run_id, - self._serialize_json(simulation_spec_snapshot), - *[ - version_manifest.get(field) - for field in SIMULATION_RUN_VERSION_FIELDS - ], - ), - ) - - database.transaction(create_run_transaction) - return self.get_simulation_run(run_id) - - def get_simulation_run(self, run_id: str) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulation_runs WHERE id = ?", - (run_id,), - ).fetchone() - return self._parse_run_row(row) - - def list_simulation_runs(self, simulation_id: int) -> list[dict]: - rows = database.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence ASC - """, - (simulation_id,), - ).fetchall() - return [self._parse_run_row(row) for row in rows] + def get_simulation_run( + self, + session: Session, + run_id: str, + ) -> SimulationRun | None: + return session.get(SimulationRun, run_id) - def get_newest_simulation_run(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence DESC - LIMIT 1 - """, - (simulation_id,), - ).fetchone() - return self._parse_run_row(row) + def list_simulation_runs( + self, + session: Session, + simulation_id: int, + ) -> list[SimulationRun]: + return list( + session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.asc()) + ) + ) - def select_display_run(self, simulation: dict) -> dict | None: - if simulation.get("active_run_id"): - active_run = self.get_simulation_run(simulation["active_run_id"]) + def get_newest_simulation_run( + self, + session: Session, + simulation_id: int, + ) -> SimulationRun | None: + return session.scalar( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) + + def select_display_run( + self, + session: Session, + simulation: Simulation, + ) -> SimulationRun | None: + if simulation.active_run_id: + active_run = self.get_simulation_run(session, simulation.active_run_id) if active_run is not None: return active_run - if simulation.get("latest_successful_run_id"): - latest_successful_run = self.get_simulation_run( - simulation["latest_successful_run_id"] + if simulation.latest_successful_run_id: + successful = self.get_simulation_run( + session, + simulation.latest_successful_run_id, ) - if latest_successful_run is not None: - return latest_successful_run - return self.get_newest_simulation_run(simulation["id"]) + if successful is not None: + return successful + return self.get_newest_simulation_run(session, simulation.id) diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index e5582ee17..5e073440f 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,534 +1,247 @@ +import json import uuid +from dataclasses import dataclass -from sqlalchemy.engine.row import Row +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS -from policyengine_api.data import database -from policyengine_api.services.run_sync_utils import ( - determine_parent_pointers, - parse_json_field, - serialize_json_field, -) -from policyengine_api.services.simulation_spec_service import ( - SimulationSpec, - SimulationSpecService, -) +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Simulation, SimulationRun -class SimulationService: - def __init__(self): - self.simulation_spec_service = SimulationSpecService() +@dataclass(frozen=True) +class SimulationCreateResult: + simulation: Simulation + created: bool + - def _lock_clause(self) -> str: - return "" if database.local else " FOR UPDATE" +class SimulationService: + """Simulation operations with service-owned ORM transaction boundaries.""" - def _get_simulation_row( + def __init__( self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() + + @staticmethod + def _select_simulation( + session: Session, simulation_id: int, - *, - queryer=None, country_id: str | None = None, + *, for_update: bool = False, - ) -> dict | None: - queryer = queryer or database - query = "SELECT * FROM simulations WHERE id = ?" - params: list[int | str] = [simulation_id] + ) -> Simulation | None: + statement = select(Simulation).where(Simulation.id == simulation_id) if country_id is not None: - query += " AND country_id = ?" - params.append(country_id) + statement = statement.where(Simulation.country_id == country_id) if for_update: - query += self._lock_clause() - - row: Row | None = queryer.query(query, tuple(params)).fetchone() - return dict(row) if row is not None else None - - def _find_existing_simulation_row( - self, - *, - country_id: str, - population_id: str, - population_type: str, - policy_id: int, - queryer=None, - ) -> dict | None: - queryer = queryer or database - row: Row | None = queryer.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - ORDER BY id DESC - """, - (country_id, population_id, population_type, policy_id), - ).fetchone() - return dict(row) if row is not None else None - - def _build_version_manifest(self, simulation: dict) -> dict[str, str | None]: - return { - "country_package_version": simulation.get("api_version"), - "policyengine_version": None, - "data_version": None, - "runtime_app_name": None, - "simulation_cache_version": None, - } - - def _list_simulation_runs_descending( - self, simulation_id: int, *, queryer=None - ) -> list[dict]: - queryer = queryer or database - rows = queryer.query( - """ - SELECT * FROM simulation_runs - WHERE simulation_id = ? - ORDER BY run_sequence DESC - """, - (simulation_id,), - ).fetchall() - - runs = [] - for row in rows: - run = dict(row) - run["simulation_spec_snapshot_json"] = parse_json_field( - run.get("simulation_spec_snapshot_json") - ) - runs.append(run) - return runs + statement = statement.with_for_update() + return session.scalar(statement) - def _select_mutable_run( - self, simulation: dict, runs_descending: list[dict] - ) -> dict | None: - active_run_id = simulation.get("active_run_id") - if active_run_id is not None: - for run in runs_descending: - if run["id"] == active_run_id: - return run - return runs_descending[0] if runs_descending else None + @staticmethod + def _latest_successful_run_id(runs: list[SimulationRun]) -> str | None: + return next((run.id for run in runs if run.status == "complete"), None) - def _upsert_simulation_spec_in_transaction( - self, tx, simulation: dict - ) -> SimulationSpec: - expected_spec = self.simulation_spec_service.build_simulation_spec(simulation) - existing_spec = parse_json_field(simulation.get("simulation_spec_json")) - if ( - existing_spec != expected_spec.model_dump() - or simulation.get("simulation_spec_schema_version") != 1 - ): - tx.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - expected_spec.model_dump_json(), - 1, - simulation["id"], - ), - ) - simulation["simulation_spec_json"] = expected_spec.model_dump() - simulation["simulation_spec_schema_version"] = 1 - - return expected_spec - - def _run_matches_parent( + def _ensure_simulation_dual_write_state( self, - run: dict, - simulation: dict, - simulation_spec: SimulationSpec, - ) -> bool: - version_manifest = self._build_version_manifest(simulation) - return ( - run["status"] == simulation["status"] - and run.get("output") == simulation.get("output") - and run.get("error_message") == simulation.get("error_message") - and run.get("simulation_spec_snapshot_json") == simulation_spec.model_dump() - and run.get("country_package_version") - == version_manifest["country_package_version"] - and run.get("policyengine_version") - == version_manifest["policyengine_version"] - and run.get("data_version") == version_manifest["data_version"] - and run.get("runtime_app_name") == version_manifest["runtime_app_name"] - and run.get("simulation_cache_version") - == version_manifest["simulation_cache_version"] - ) - - def _insert_bootstrap_run( - self, tx, simulation: dict, simulation_spec: SimulationSpec - ) -> None: - version_manifest = self._build_version_manifest(simulation) - tx.query( - """ - INSERT INTO simulation_runs ( - id, simulation_id, report_output_run_id, input_position, run_sequence, - status, output, error_message, trigger_type, requested_at, started_at, - finished_at, source_run_id, simulation_spec_snapshot_json, - country_package_version, policyengine_version, data_version, - runtime_app_name, simulation_cache_version - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - str(uuid.uuid4()), - simulation["id"], - None, - None, - 1, - simulation["status"], - serialize_json_field(simulation.get("output")), - simulation.get("error_message"), - "initial", - None, - None, - None, - None, - simulation_spec.model_dump_json(), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["simulation_cache_version"], - ), - ) - - def _update_simulation_run_in_transaction( - self, - tx, - run_id: str, - simulation: dict, - simulation_spec: SimulationSpec, - ) -> None: - version_manifest = self._build_version_manifest(simulation) - tx.query( - """ - UPDATE simulation_runs - SET status = ?, output = ?, error_message = ?, - simulation_spec_snapshot_json = ?, country_package_version = ?, - policyengine_version = ?, data_version = ?, runtime_app_name = ?, - simulation_cache_version = ? - WHERE id = ? - """, - ( - simulation["status"], - serialize_json_field(simulation.get("output")), - simulation.get("error_message"), - simulation_spec.model_dump_json(), - version_manifest["country_package_version"], - version_manifest["policyengine_version"], - version_manifest["data_version"], - version_manifest["runtime_app_name"], - version_manifest["simulation_cache_version"], - run_id, - ), - ) - - def _sync_parent_pointers_in_transaction( - self, tx, simulation: dict, runs_descending: list[dict] - ) -> None: - desired_active_run_id, desired_latest_successful_run_id = ( - determine_parent_pointers(simulation["status"], runs_descending) - ) - if ( - simulation.get("active_run_id") == desired_active_run_id - and simulation.get("latest_successful_run_id") - == desired_latest_successful_run_id - ): - return - - tx.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ( - desired_active_run_id, - desired_latest_successful_run_id, - simulation["id"], - ), - ) - simulation["active_run_id"] = desired_active_run_id - simulation["latest_successful_run_id"] = desired_latest_successful_run_id - - def _ensure_simulation_dual_write_state_in_transaction( - self, - tx, + session: Session, simulation_id: int, - *, country_id: str | None = None, - ) -> dict: - simulation = self._get_simulation_row( + ) -> Simulation: + simulation = self._select_simulation( + session, simulation_id, - queryer=tx, - country_id=country_id, + country_id, for_update=True, ) if simulation is None: raise ValueError(f"Simulation #{simulation_id} not found") - simulation_spec = self._upsert_simulation_spec_in_transaction(tx, simulation) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx + spec = { + "country_id": simulation.country_id, + "population_id": simulation.population_id, + "population_type": simulation.population_type, + "policy_id": simulation.policy_id, + } + simulation.simulation_spec_json = spec + simulation.simulation_spec_schema_version = 1 + runs = list( + session.scalars( + select(SimulationRun) + .where(SimulationRun.simulation_id == simulation_id) + .order_by(SimulationRun.run_sequence.desc()) + ) ) - if not runs_descending: - self._insert_bootstrap_run(tx, simulation, simulation_spec) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx + if not runs: + run = SimulationRun( + id=str(uuid.uuid4()), + simulation_id=simulation_id, + run_sequence=1, + status=simulation.status, + output=simulation.output, + error_message=simulation.error_message, + trigger_type="initial", + simulation_spec_snapshot_json=spec, + country_package_version=simulation.api_version, ) + session.add(run) + session.flush() + runs = [run] else: - mutable_run = self._select_mutable_run(simulation, runs_descending) - if mutable_run is not None and not self._run_matches_parent( - mutable_run, - simulation, - simulation_spec, - ): - self._update_simulation_run_in_transaction( - tx, - run_id=mutable_run["id"], - simulation=simulation, - simulation_spec=simulation_spec, - ) - runs_descending = self._list_simulation_runs_descending( - simulation_id, queryer=tx - ) - - self._sync_parent_pointers_in_transaction(tx, simulation, runs_descending) - refreshed_simulation = self._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, + mutable = next( + (run for run in runs if run.id == simulation.active_run_id), + runs[0], + ) + mutable.status = simulation.status + mutable.output = simulation.output + mutable.error_message = simulation.error_message + mutable.simulation_spec_snapshot_json = spec + mutable.country_package_version = simulation.api_version + + latest_successful = self._latest_successful_run_id(runs) + simulation.active_run_id = ( + runs[0].id if simulation.status in {"pending", "running"} else None ) - if refreshed_simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found after sync") - return refreshed_simulation - - def ensure_simulation_dual_write_state( - self, simulation_id: int, country_id: str | None = None - ) -> dict: - return database.transaction( - lambda tx: self._ensure_simulation_dual_write_state_in_transaction( - tx, - simulation_id, - country_id=country_id, + if simulation.status == "complete" and latest_successful is None: + latest_successful = runs[0].id + simulation.latest_successful_run_id = latest_successful + session.flush() + return simulation + + @staticmethod + def _find_existing_simulation( + session: Session, + country_id: str, + population_id: str, + population_type: str, + policy_id: int, + *, + for_update: bool = False, + ) -> Simulation | None: + statement = ( + select(Simulation) + .where( + Simulation.country_id == country_id, + Simulation.population_id == population_id, + Simulation.population_type == population_type, + Simulation.policy_id == policy_id, ) + .order_by(Simulation.id.desc()) ) + if for_update: + statement = statement.with_for_update() + return session.scalar(statement) - def find_existing_simulation( + def _create_simulation( self, + session: Session, country_id: str, population_id: str, population_type: str, policy_id: int, - ) -> dict | None: - """ - Find an existing simulation with the same parameters. - - Args: - country_id (str): The country ID. - population_id (str): The population identifier (household or geography ID). - population_type (str): Type of population ('household' or 'geography'). - policy_id (int): The policy ID. - - Returns: - dict | None: The existing simulation data or None if not found. - """ - print("Checking for existing simulation") - - try: - existing_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - ) - if existing_simulation is not None: - print(f"Found existing simulation with ID: {existing_simulation['id']}") - return existing_simulation - - except Exception as e: - print(f"Error checking for existing simulation. Details: {str(e)}") - raise e + ) -> Simulation: + simulation = Simulation( + country_id=country_id, + api_version=COUNTRY_PACKAGE_VERSIONS.get(country_id), + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", + ) + session.add(simulation) + session.flush() + return self._ensure_simulation_dual_write_state( + session, + simulation.id, + country_id, + ) - def create_simulation( + def get_or_create_simulation( self, country_id: str, population_id: str, population_type: str, policy_id: int, - ) -> dict: - """ - Create a new simulation record with pending status. - - Args: - country_id (str): The country ID. - population_id (str): The population identifier (household or geography ID). - population_type (str): Type of population ('household' or 'geography'). - policy_id (int): The policy ID. - - Returns: - dict: The created simulation record. - """ - print("Creating new simulation") - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - try: - - def tx_callback(tx): - existing_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - queryer=tx, - ) - if existing_simulation is not None: - print( - f"Reusing existing simulation with ID: {existing_simulation['id']}" - ) - return self._ensure_simulation_dual_write_state_in_transaction( - tx, - existing_simulation["id"], - country_id=country_id, - ) - - tx.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - country_id, - api_version, - population_id, - population_type, - policy_id, - "pending", - ), - ) - - created_simulation = self._find_existing_simulation_row( - country_id=country_id, - population_id=population_id, - population_type=population_type, - policy_id=policy_id, - queryer=tx, - ) - if created_simulation is None: - raise Exception("Failed to retrieve created simulation") - - print(f"Created simulation with ID: {created_simulation['id']}") - return self._ensure_simulation_dual_write_state_in_transaction( - tx, - created_simulation["id"], - country_id=country_id, + ) -> SimulationCreateResult: + with self._sessions.begin() as session: + simulation = self._find_existing_simulation( + session, + country_id, + population_id, + population_type, + policy_id, + for_update=True, + ) + created = simulation is None + if simulation is None: + simulation = self._create_simulation( + session, + country_id, + population_id, + population_type, + policy_id, ) - - return database.transaction(tx_callback) - - except Exception as e: - print(f"Error creating simulation. Details: {str(e)}") - raise e - - def get_simulation(self, country_id: str, simulation_id: int) -> dict | None: - """ - Get a simulation record by ID. - - Args: - country_id (str): The country ID. - simulation_id (int): The simulation ID. - - Returns: - dict | None: The simulation data or None if not found. - """ - print(f"Getting simulation {simulation_id}") - - try: - if type(simulation_id) is not int or simulation_id < 0: - raise Exception( - f"Invalid simulation ID: {simulation_id}. Must be a positive integer." + else: + simulation = self._ensure_simulation_dual_write_state( + session, + simulation.id, + country_id, ) + return SimulationCreateResult(simulation=simulation, created=created) - return self._get_simulation_row(simulation_id, country_id=country_id) - - except Exception as e: - print(f"Error fetching simulation #{simulation_id}. Details: {str(e)}") - raise e + def get_simulation( + self, + country_id: str, + simulation_id: int, + ) -> Simulation | None: + if type(simulation_id) is not int or simulation_id < 0: + raise Exception( + f"Invalid simulation ID: {simulation_id}. Must be a positive integer." + ) + with self._sessions() as session: + return self._select_simulation(session, simulation_id, country_id) def update_simulation( self, country_id: str, simulation_id: int, status: str | None = None, - output: str | None = None, + output: dict | list | str | None = None, error_message: str | None = None, - ) -> bool: - """ - Update a simulation record with results or error. - - Args: - country_id (str): The country ID. - simulation_id (int): The simulation ID. - status (str | None): The new status ('complete' or 'error'). - output (str | None): The result output as JSON string (for complete status). - error_message (str | None): The error message (for error status). - - Returns: - bool: True if update was successful. - """ - print(f"Updating simulation {simulation_id}") - api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id) - - try: - update_fields = [] - update_values = [] - - if status is not None: - update_fields.append("status = ?") - update_values.append(status) - - if output is not None: - update_fields.append("output = ?") - update_values.append(output) - - if error_message is not None: - update_fields.append("error_message = ?") - update_values.append(error_message) - - # Only refresh api_version when the caller is actually - # changing one of the user-supplied fields above. The - # previous code appended api_version unconditionally, so - # the "no fields to update" guard below never fired and a - # PATCH with an empty body still touched the row. - if not update_fields: - print("No fields to update") - return False - - update_fields.append("api_version = ?") - update_values.append(api_version) - - def tx_callback(tx): - simulation = self._get_simulation_row( - simulation_id, - queryer=tx, - country_id=country_id, - for_update=True, - ) - if simulation is None: - raise ValueError(f"Simulation #{simulation_id} not found") - - tx.query( - f"UPDATE simulations SET {', '.join(update_fields)} WHERE id = ? AND country_id = ?", - (*update_values, simulation_id, country_id), - ) - self._ensure_simulation_dual_write_state_in_transaction( - tx, - simulation_id, - country_id=country_id, - ) - - database.transaction(tx_callback) - - print(f"Successfully updated simulation #{simulation_id}") - return True - - except Exception as e: - print(f"Error updating simulation #{simulation_id}. Details: {str(e)}") - raise e + ) -> Simulation | None: + values = { + key: value + for key, value in { + "status": status, + "output": output, + "error_message": error_message, + }.items() + if value is not None + } + if not values: + return None + if isinstance(values.get("output"), str): + values["output"] = json.loads(values["output"]) + with self._sessions.begin() as session: + simulation = self._select_simulation( + session, + simulation_id, + country_id, + for_update=True, + ) + if simulation is None: + raise LookupError(f"Simulation #{simulation_id} not found") + for key, value in values.items(): + setattr(simulation, key, value) + simulation.api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) + return self._ensure_simulation_dual_write_state( + session, + simulation_id, + country_id, + ) diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index d5f6e86f7..f6c5a7373 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -2,9 +2,9 @@ from typing import Literal from pydantic import BaseModel -from sqlalchemy.engine.row import Row +from sqlalchemy.orm import Session -from policyengine_api.data import database +from policyengine_api.data.v1_models import Simulation SIMULATION_SPEC_SCHEMA_VERSION = 1 @@ -23,70 +23,59 @@ def _validate_schema_version(self, schema_version: int | None) -> None: f"Unsupported simulation spec schema version: {schema_version}" ) - def _get_simulation_row(self, simulation_id: int) -> dict | None: - row: Row | None = database.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_id,), - ).fetchone() - return dict(row) if row is not None else None - - def _validate_simulation_spec_matches_row( - self, simulation: dict, simulation_spec: SimulationSpec + def _validate_simulation_spec_matches_model( + self, + simulation: Simulation, + simulation_spec: SimulationSpec, ) -> None: expected_spec = { - "country_id": simulation["country_id"], - "population_id": simulation["population_id"], - "population_type": simulation["population_type"], - "policy_id": simulation["policy_id"], + "country_id": simulation.country_id, + "population_id": simulation.population_id, + "population_type": simulation.population_type, + "policy_id": simulation.policy_id, } if simulation_spec.model_dump() != expected_spec: raise ValueError("Simulation spec must match the linked simulation row") - def build_simulation_spec(self, simulation: dict) -> SimulationSpec: - return SimulationSpec.model_validate( - { - "country_id": simulation["country_id"], - "population_id": simulation["population_id"], - "population_type": simulation["population_type"], - "policy_id": simulation["policy_id"], - } + def build_simulation_spec(self, simulation: Simulation) -> SimulationSpec: + return SimulationSpec( + country_id=simulation.country_id, + population_id=simulation.population_id, + population_type=simulation.population_type, + policy_id=simulation.policy_id, ) - def get_simulation_spec(self, simulation_id: int) -> SimulationSpec | None: - simulation = self._get_simulation_row(simulation_id) - if simulation is None or simulation["simulation_spec_json"] is None: + def get_simulation_spec( + self, + session: Session, + simulation_id: int, + ) -> SimulationSpec | None: + simulation = session.get(Simulation, simulation_id) + if simulation is None or simulation.simulation_spec_json is None: return None - self._validate_schema_version(simulation["simulation_spec_schema_version"]) - raw_spec = simulation["simulation_spec_json"] + self._validate_schema_version(simulation.simulation_spec_schema_version) + raw_spec = simulation.simulation_spec_json if isinstance(raw_spec, str): + # Existing databases may contain pre-ORM JSON text. New writes below + # always assign Python objects and leave conversion to SQLAlchemy. raw_spec = json.loads(raw_spec) simulation_spec = SimulationSpec.model_validate(raw_spec) - self._validate_simulation_spec_matches_row(simulation, simulation_spec) + self._validate_simulation_spec_matches_model(simulation, simulation_spec) return simulation_spec def set_simulation_spec( self, + session: Session, simulation_id: int, simulation_spec: SimulationSpec, schema_version: int = SIMULATION_SPEC_SCHEMA_VERSION, ) -> bool: self._validate_schema_version(schema_version) - simulation = self._get_simulation_row(simulation_id) + simulation = session.get(Simulation, simulation_id) if simulation is None: raise ValueError(f"Simulation #{simulation_id} not found") - self._validate_simulation_spec_matches_row(simulation, simulation_spec) - - database.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - simulation_spec.model_dump_json(), - schema_version, - simulation_id, - ), - ) + self._validate_simulation_spec_matches_model(simulation, simulation_spec) + simulation.simulation_spec_json = simulation_spec.model_dump() + simulation.simulation_spec_schema_version = schema_version return True diff --git a/policyengine_api/services/tracer_analysis_service.py b/policyengine_api/services/tracer_analysis_service.py index 2fd072f83..575cb6a4f 100644 --- a/policyengine_api/services/tracer_analysis_service.py +++ b/policyengine_api/services/tracer_analysis_service.py @@ -1,17 +1,15 @@ -from policyengine_api.data import local_database -import json -from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from typing import Generator, Literal import re import anthropic from policyengine_api.services.ai_analysis_service import AIAnalysisService from werkzeug.exceptions import NotFound +from sqlalchemy import select +from policyengine_api.data.local_models import Tracer -class TracerAnalysisService(AIAnalysisService): - def __init__(self): - super().__init__() +class TracerAnalysisService(AIAnalysisService): def execute_analysis( self, country_id: str, @@ -56,9 +54,9 @@ def execute_analysis( ) # If a calculated record exists for this prompt, return it as a string - existing_analysis: str = self.get_existing_analysis(prompt) + existing_analysis = self.get_existing_analysis(prompt) if existing_analysis is not None: - return existing_analysis, "static" + return existing_analysis.analysis, "static" # Otherwise, pass prompt to Claude, then return streaming function try: @@ -78,20 +76,22 @@ def get_tracer( api_version: str, ) -> list: try: - # Retrieve from the tracers table in the local database - row = local_database.query( - """ - SELECT * FROM tracers - WHERE household_id = ? AND policy_id = ? AND country_id = ? AND api_version = ? - """, - (household_id, policy_id, country_id, api_version), - ).fetchone() - - if row is None: + with self._sessions() as session: + tracer = session.scalar( + select(Tracer) + .where( + Tracer.household_id == int(household_id), + Tracer.policy_id == int(policy_id), + Tracer.country_id == country_id, + Tracer.api_version == api_version, + ) + .order_by(Tracer.id.desc()) + ) + + if tracer is None: raise NotFound("No household simulation tracer found") - tracer_output_list = json.loads(row["tracer_output"]) - return tracer_output_list + return tracer.tracer_output except Exception as e: print(f"Error getting existing tracer analysis: {str(e)}") diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py new file mode 100644 index 000000000..07339310f --- /dev/null +++ b/policyengine_api/services/user_policy_service.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserPolicy + + +USER_POLICY_IDENTITY_FIELDS = ( + "country_id", + "reform_id", + "baseline_id", + "user_id", + "year", + "geography", + "reform_label", + "baseline_label", + "dataset", +) + + +@dataclass(frozen=True) +class UserPolicyCreateResult: + user_policy: UserPolicy + created: bool + + +class UserPolicyService: + """Saved-policy operations with service-owned ORM transactions.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() + + @staticmethod + def _find_matching_user_policy( + session: Session, + values: Mapping[str, Any], + ) -> UserPolicy | None: + return session.scalar( + select(UserPolicy).where( + *( + getattr(UserPolicy, field) == values[field] + for field in USER_POLICY_IDENTITY_FIELDS + ) + ) + ) + + def create_or_get_user_policy( + self, + values: Mapping[str, Any], + ) -> UserPolicyCreateResult: + with self._sessions.begin() as session: + user_policy = self._find_matching_user_policy(session, values) + created = user_policy is None + if user_policy is None: + user_policy = UserPolicy(**values) + session.add(user_policy) + session.flush() + return UserPolicyCreateResult( + user_policy=user_policy, + created=created, + ) + + def list_user_policies( + self, + country_id: str, + user_id: str, + ) -> list[UserPolicy]: + with self._sessions() as session: + return list( + session.scalars( + select(UserPolicy).where( + UserPolicy.country_id == country_id, + UserPolicy.user_id == user_id, + ) + ) + ) + + def update_user_policy( + self, + country_id: str, + user_policy_id: int, + values: Mapping[str, Any], + ) -> UserPolicy | None: + with self._sessions.begin() as session: + user_policy = session.scalar( + select(UserPolicy).where( + UserPolicy.id == user_policy_id, + UserPolicy.country_id == country_id, + ) + ) + if user_policy is None: + return None + for field, value in values.items(): + setattr(user_policy, field, value) + return user_policy diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index 0bdfb0cdd..d7c11193d 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,76 +1,119 @@ -import json -from typing import Any -from policyengine_api.data import database +from __future__ import annotations + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserProfile class UserService: + """User-profile operations with service-owned ORM transaction boundaries.""" + + def __init__( + self, + session_factory: sessionmaker[Session] | None = None, + ) -> None: + self._injected_session_factory = session_factory + + @property + def _sessions(self) -> sessionmaker[Session]: + return self._injected_session_factory or get_v1_session_factory() + def create_profile( self, primary_country: str, auth0_id: str, username: str | None, - user_since: str, - ) -> tuple[bool, Any]: - """ - returns true if a new record was created and false otherwise. - """ - # TODO: this is not written as an atomic operation. This will cause intermittent errors - # in some cases - # https://github.com/PolicyEngine/policyengine-api/issues/2058 to resolve after - # this refactor. - row = self.get_profile(auth0_id=auth0_id) - if row is not None: - return False, row - # Unfortunately, it's not possible to use RETURNING - # with SQLite3 without rewriting the PolicyEngineDatabase - # object or implementing a true ORM, thus the double query - database.query( - f"INSERT INTO user_profiles (primary_country, auth0_id, username, user_since) VALUES (?, ?, ?, ?)", - (primary_country, auth0_id, username, user_since), - ) + user_since: int, + ) -> tuple[bool, UserProfile]: + with self._sessions.begin() as session: + return self._create_profile( + session, + primary_country, + auth0_id, + username, + user_since, + ) - row = self.get_profile(auth0_id=auth0_id) - - return (True, row) + @classmethod + def _create_profile( + cls, + session: Session, + primary_country: str, + auth0_id: str, + username: str | None, + user_since: int, + ) -> tuple[bool, UserProfile]: + existing = cls._get_profile(session, auth0_id=auth0_id) + if existing is not None: + return False, existing + profile = UserProfile( + auth0_id=auth0_id, + username=username, + primary_country=primary_country, + user_since=user_since, + ) + session.add(profile) + session.flush() + return True, profile def get_profile( - self, auth0_id: str | None = None, user_id: str | None = None - ) -> Any | None: - key = "user_id" if auth0_id is None else "auth0_id" - value = user_id if auth0_id is None else auth0_id - if value is None: + self, + auth0_id: str | None = None, + user_id: int | str | None = None, + ) -> UserProfile | None: + if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") - row = database.query( - f"SELECT * FROM user_profiles WHERE {key} = ?", - (value,), - ).fetchone() + with self._sessions() as session: + return self._get_profile(session, auth0_id, user_id) - return row + @staticmethod + def _get_profile( + session: Session, + auth0_id: str | None = None, + user_id: int | str | None = None, + ) -> UserProfile | None: + condition = ( + UserProfile.user_id == user_id + if user_id is not None + else UserProfile.auth0_id == auth0_id + ) + return session.scalar(select(UserProfile).where(condition)) def update_profile( self, - user_id: str, + user_id: int, primary_country: str | None, username: str | None, - user_since: str, - ) -> bool: - fields = dict( - primary_country=primary_country, - username=username, - user_since=user_since, - ) - if self.get_profile(user_id=user_id) is None: - return False - - with_values = [key for key in fields if fields[key] is not None] - fields_update = ",".join([f"{key} = ?" for key in with_values]) - query = f"UPDATE user_profiles SET {fields_update} WHERE user_id = ?" - values = [fields[key] for key in with_values] + [user_id] + user_since: int | None, + ) -> UserProfile | None: + if user_id is None: + raise ValueError("you must specify either auth0_id or user_id") + with self._sessions.begin() as session: + return self._update_profile( + session, + user_id, + primary_country, + username, + user_since, + ) - print(f"Updating record {user_id}") - try: - database.query(query, (tuple(values))) - except Exception as ex: - print(f"ERROR: unable to update user record: {ex}") - raise - return True + @staticmethod + def _update_profile( + session: Session, + user_id: int, + primary_country: str | None, + username: str | None, + user_since: int | None, + ) -> UserProfile | None: + profile = session.get(UserProfile, user_id) + if profile is None: + return None + if primary_country is not None: + profile.primary_country = primary_country + if username is not None: + profile.username = username + if user_since is not None: + profile.user_since = user_since + return profile diff --git a/policyengine_api/setup_data.py b/policyengine_api/setup_data.py deleted file mode 100644 index 47be04faf..000000000 --- a/policyengine_api/setup_data.py +++ /dev/null @@ -1,2 +0,0 @@ -def setup_data(): - import policyengine_api.endpoints.search diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index cfd6c61f0..346840d34 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -1,11 +1,13 @@ from functools import wraps from typing import Union + from flask import Response -import json + from policyengine_api.country_validation import ( InvalidCountryError, ensure_supported_country, ) +from policyengine_api.response_factory import _make_error_response def validate_country(func): @@ -26,7 +28,9 @@ def validate_country_wrapper( try: ensure_supported_country(country_id) except InvalidCountryError as error: - return Response(json.dumps(error.to_payload()), status=400) + # Preserve the legacy v1 content type while native FastAPI routes + # are contractually required to match the Flask response. + return _make_error_response(error, 400, mimetype=None) return func(country_id, *args, **kwargs) return validate_country_wrapper diff --git a/pyproject.toml b/pyproject.toml index 8cac617b7..0c11a1518 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ classifiers = [ ] dependencies = [ "a2wsgi>=1.10,<2", + "alembic>=1.14,<2", "anthropic", "assertpy", "click>=8,<9", @@ -64,9 +65,6 @@ dev = [ "towncrier>=24.8.0", ] -[project.scripts] -policyengine-api-setup = "policyengine_api.setup_data:setup_data" - [tool.hatch.build.targets.wheel] packages = ["policyengine_api"] diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py new file mode 100644 index 000000000..80611b9e8 --- /dev/null +++ b/scripts/v1_database_migration.py @@ -0,0 +1,235 @@ +"""Safely inspect and upgrade the API v1 Cloud SQL schema.""" + +from __future__ import annotations + +import argparse +from enum import StrEnum +import os +from typing import Any + +from alembic import command +from alembic.autogenerate import compare_metadata +from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.script import ScriptDirectory +from sqlalchemy import Connection, create_engine, inspect, text +from sqlalchemy.engine import URL +from sqlalchemy.pool import NullPool + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base + + +ALEMBIC_CONFIG = REPO / "alembic-v1.ini" +MIGRATION_LOCK_NAME = "policyengine-api-v1-alembic" +DATABASE_NAME = "policyengine" +READONLY_USER = "policyengine_schema_reader" +MIGRATION_USER = "policyengine_schema_migrator" +PROXY_HOST = "127.0.0.1" +PROXY_PORT = 3307 + + +class DatabaseState(StrEnum): + UNVERSIONED = "unversioned" + INVALID = "invalid" + PENDING = "pending" + HEAD = "head" + + +def build_database_url( + *, + username: str, + password: str, + host: str, + port: int, + database: str, +) -> URL: + """Build a SQLAlchemy URL without stringifying its credentials.""" + + return URL.create( + drivername="mysql+pymysql", + username=username, + password=password, + host=host, + port=port, + database=database, + ) + + +def classify_database_state( + *, + version_table_exists: bool, + current_heads: set[str], + script_heads: set[str] | None = None, +) -> DatabaseState: + if not version_table_exists: + return DatabaseState.UNVERSIONED + if not current_heads: + return DatabaseState.INVALID + if script_heads is not None and current_heads == script_heads: + return DatabaseState.HEAD + return DatabaseState.PENDING + + +def describe_metadata_difference(difference: Any) -> str: + """Return a stable, data-free description of an Alembic metadata diff.""" + + item = difference + if isinstance(item, list) and len(item) == 1: + item = item[0] + if not isinstance(item, tuple) or not item: + raise ValueError("Unsupported metadata difference shape") + + operation = item[0] + if operation == "remove_table" and len(item) >= 2: + return f"extra_table:{item[1].name}" + if operation == "modify_nullable" and len(item) >= 7: + return ( + f"nullable:{item[2]}.{item[3]}:" + f"{str(item[5]).lower()}->{str(item[6]).lower()}" + ) + if operation == "modify_default" and len(item) >= 7: + existing = "none" if item[5] is None else "present" + target = "none" if item[6] is None else "present" + return f"default:{item[2]}.{item[3]}:{existing}->{target}" + raise ValueError(f"Unsupported metadata difference operation: {operation}") + + +def _config(connection: Connection) -> Config: + config = Config(str(ALEMBIC_CONFIG)) + config.attributes["connection"] = connection + return config + + +def _script_heads(config: Config) -> set[str]: + return set(ScriptDirectory.from_config(config).get_heads()) + + +def database_state(connection: Connection) -> DatabaseState: + inspector = inspect(connection) + version_table_exists = "alembic_version" in inspector.get_table_names() + context = MigrationContext.configure(connection) + config = _config(connection) + return classify_database_state( + version_table_exists=version_table_exists, + current_heads=set(context.get_current_heads()), + script_heads=_script_heads(config), + ) + + +def metadata_differences(connection: Connection) -> set[str]: + context = MigrationContext.configure( + connection, + opts={"compare_type": True, "compare_server_default": True}, + ) + return { + describe_metadata_difference(difference) + for difference in compare_metadata(context, V1Base.metadata) + } + + +def verify_head_schema(connection: Connection) -> None: + state = database_state(connection) + if state is not DatabaseState.HEAD: + raise RuntimeError(f"Database is not at the Alembic head; found {state}") + differences = metadata_differences(connection) + if differences: + raise RuntimeError( + f"Database metadata drift remains after migration: {sorted(differences)}" + ) + + +def _acquire_lock(connection: Connection) -> None: + acquired = connection.scalar( + text("SELECT GET_LOCK(:name, 60)"), {"name": MIGRATION_LOCK_NAME} + ) + if acquired != 1: + raise RuntimeError("Could not acquire the v1 Alembic migration lock") + + +def _release_lock(connection: Connection) -> None: + connection.execute( + text("SELECT RELEASE_LOCK(:name)"), {"name": MIGRATION_LOCK_NAME} + ) + + +def upgrade_database(connection: Connection, *, backup_id: str) -> None: + _acquire_lock(connection) + try: + state = database_state(connection) + if state in {DatabaseState.UNVERSIONED, DatabaseState.INVALID}: + raise RuntimeError( + "database has no valid Alembic revision; automatic baseline " + "stamping is disabled" + ) + if state is DatabaseState.HEAD: + verify_head_schema(connection) + return + if not backup_id.strip(): + raise ValueError("A completed Cloud SQL backup ID is required") + + command.upgrade(_config(connection), "head") + verify_head_schema(connection) + finally: + _release_lock(connection) + + +def _database_target(mode: str) -> str | URL: + url_env_name = ( + "STAGE7_EXISTING_DATABASE_URL" + if mode in {"verify-head", "state"} + else "ALEMBIC_DATABASE_URL" + ) + if explicit_url := os.environ.get(url_env_name): + return explicit_url + + readonly = mode in {"verify-head", "state"} + password_env_name = ( + "POLICYENGINE_DB_READONLY_PASSWORD" + if readonly + else "POLICYENGINE_DB_MIGRATION_PASSWORD" + ) + password = os.environ.get(password_env_name) + if not password: + raise RuntimeError(f"{url_env_name} or {password_env_name} is required") + return build_database_url( + username=READONLY_USER if readonly else MIGRATION_USER, + password=password, + host=PROXY_HOST, + port=PROXY_PORT, + database=DATABASE_NAME, + ) + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--mode", + required=True, + choices=("state", "verify-head", "upgrade"), + ) + parser.add_argument("--backup-id", default="") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + engine = create_engine(_database_target(args.mode), poolclass=NullPool) + try: + connection_context = ( + engine.begin() if args.mode == "upgrade" else engine.connect() + ) + with connection_context as connection: + if args.mode == "state": + print(database_state(connection).value) + elif args.mode == "verify-head": + verify_head_schema(connection) + else: + upgrade_database(connection, backup_id=args.backup_id) + finally: + engine.dispose() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index 6b3407bbe..eae014cf1 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -7,12 +7,27 @@ import pytest from flask import Flask, Response -from policyengine_api.endpoints.household import get_calculate -from policyengine_api.endpoints.policy import get_policy_search +from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.v1_models import ( + Household, + Policy, + ReportOutput, + ReportOutputRun, + Simulation, +) +from policyengine_api.extensions import cache from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.simulation_routes import simulation_bp +from policyengine_api.services.report_output_service import ( + ReportCreateResult, + ReportOutputView, +) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationResult, +) +from policyengine_api.services.simulation_service import SimulationCreateResult from tests.contract.clients import ( ASGIContractClient, ContractClient, @@ -91,13 +106,17 @@ def _load_blueprint_with_fake_service( blueprint_name: str, ): sentinel = object() + original_route_module = sys.modules.get(route_module_name, sentinel) original_service_module = sys.modules.get(service_module_name, sentinel) sys.modules.pop(route_module_name, None) sys.modules[service_module_name] = fake_service_module try: return getattr(importlib.import_module(route_module_name), blueprint_name) finally: - sys.modules.pop(route_module_name, None) + if original_route_module is sentinel: + sys.modules.pop(route_module_name, None) + else: + sys.modules[route_module_name] = original_route_module if original_service_module is sentinel: sys.modules.pop(service_module_name, None) else: @@ -128,15 +147,14 @@ def _load_contract_economy_blueprint(): def create_contract_flask_app() -> Flask: app = Flask(__name__) - app.config["TESTING"] = True + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) app.register_blueprint(_load_contract_metadata_blueprint()) app.register_blueprint(policy_bp) app.register_blueprint(household_bp) app.register_blueprint(_load_contract_economy_blueprint()) app.register_blueprint(simulation_bp) app.register_blueprint(report_output_bp) - app.route("//policies", methods=["GET"])(get_policy_search) - app.route("//calculate", methods=["POST"])(get_calculate) @app.route("/liveness-check") def liveness_check(): @@ -189,15 +207,6 @@ def _json_payload(contract: ContractRequest) -> dict | None: return None -def _policy_search_rows(): - return SimpleNamespace( - fetchall=lambda: [ - {"id": 123, "label": "Tax reform", "policy_hash": "hash-1"}, - {"id": 124, "label": "Tax reform", "policy_hash": "hash-1"}, - ] - ) - - def _fake_country(): return SimpleNamespace( metadata={}, @@ -213,7 +222,14 @@ def _patched_route_dependencies(): stack.enter_context( patch( "policyengine_api.routes.policy_routes.policy_service.get_policy", - return_value={"id": 22, "label": "Current law", "policy_json": {}}, + return_value=Policy( + id=22, + country_id="us", + label="Current law", + api_version="1", + policy_json={}, + policy_hash="hash-22", + ), ) ) stack.enter_context( @@ -224,88 +240,143 @@ def _patched_route_dependencies(): ) stack.enter_context( patch( - "policyengine_api.endpoints.policy.database.query", - return_value=_policy_search_rows(), + "policyengine_api.routes.policy_routes.policy_service.search_policies", + return_value=[ + Policy( + id=123, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + Policy( + id=124, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="hash-1", + ), + ], ) ) stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.create_household", - return_value=456, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={}, + household_hash="hash-456", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.get_household", - return_value={"id": 456, "label": "Empty household", "household_json": {}}, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={}, + household_hash="hash-456", + ), ) ) stack.enter_context( patch( "policyengine_api.routes.household_routes.household_service.update_household", - return_value={"household_json": {"people": {"you": {}}}}, - ) - ) - stack.enter_context( - patch( - "policyengine_api.endpoints.household.get_countries", - return_value={"us": _fake_country()}, - ) - ) - stack.enter_context( - patch( - "policyengine_api.endpoints.household.get_invalid_inputs_response", - return_value=None, + return_value=Household( + id=456, + country_id="us", + label="Empty household", + api_version="1", + household_json={"people": {"you": {}}}, + household_hash="hash-456", + ), ) ) stack.enter_context( patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", - return_value=None, + "policyengine_api.routes.household_routes.household_calculation_service.calculate_household", + return_value=HouseholdCalculationResult( + household={ + "people": {"you": {"age": {"2026": 40}}}, + } + ), ) ) stack.enter_context( patch( - "policyengine_api.routes.simulation_routes.simulation_service.create_simulation", - return_value={ - "id": 11, - "country_id": "us", - "population_id": "household-1", - "population_type": "household", - "policy_id": 22, - "status": "pending", - }, + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", + return_value=SimulationCreateResult( + simulation=Simulation( + id=11, + country_id="us", + population_id="household-1", + population_type="household", + policy_id=22, + status="pending", + ), + created=True, + ), ) ) stack.enter_context( patch( "policyengine_api.routes.simulation_routes.simulation_service.get_simulation", - return_value={"id": 11, "status": "pending", "country_id": "us"}, + return_value=Simulation(id=11, status="pending", country_id="us"), ) ) stack.enter_context( patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", - return_value=None, - ) - ) - stack.enter_context( - patch( - "policyengine_api.routes.report_output_routes.report_output_service.create_report_output", - return_value={ - "id": 33, - "country_id": "us", - "simulation_1_id": 11, - "simulation_2_id": None, - "status": "pending", - "year": "2026", - }, + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", + return_value=ReportCreateResult( + view=ReportOutputView( + report_output=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), + display_run=ReportOutputRun( + id="run-33", + report_output_id=33, + run_sequence=1, + status="pending", + ), + ), + created=True, + ), ) ) stack.enter_context( patch( - "policyengine_api.routes.report_output_routes.report_output_service.get_report_output", - return_value={"id": 33, "status": "pending", "country_id": "us"}, + "policyengine_api.routes.report_output_routes.report_output_service.resolve_report_output", + return_value=ReportOutputView( + report_output=ReportOutput( + id=33, + country_id="us", + simulation_1_id=11, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2026", + ), + display_run=ReportOutputRun( + id="run-33", + report_output_id=33, + run_sequence=1, + status="pending", + ), + ), ) ) return stack diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index 741676047..b0fe44636 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -2,9 +2,9 @@ Test fixtures and constants for axes calculation tests """ -import pytest -from unittest.mock import Mock, MagicMock, patch -from policyengine_api.endpoints.household import add_yearly_variables +from policyengine_api.services.household_calculation_service import ( + add_yearly_variables, +) STANDARD_AXES_COUNT = ( 401 # Not formally defined anywhere, but this value is used throughout the API diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index e4061efe3..8e1bfe0e9 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -7,6 +7,7 @@ MODAL_EXECUTION_STATUS_RUNNING, MODAL_EXECUTION_STATUS_SUBMITTED, ) +from policyengine_api.data.v1_models import ReformImpact # Mock data constants MOCK_COUNTRY_ID = "us" @@ -95,19 +96,19 @@ def mock_policyengine_version(): @pytest.fixture def mock_policy_service(): - """Mock PolicyService with get_policy_json method.""" + """Mock the ORM-facing PolicyService with decoded JSON objects.""" mock_service = MagicMock() mock_service.get_policy_json.side_effect = lambda country_id, policy_id: ( - MOCK_REFORM_POLICY_JSON + json.loads(MOCK_REFORM_POLICY_JSON) if policy_id == MOCK_POLICY_ID - else MOCK_BASELINE_POLICY_JSON + else json.loads(MOCK_BASELINE_POLICY_JSON) ) with patch( - "policyengine_api.services.economy_service.policy_service", - mock_service, - ) as mock: - yield mock + "policyengine_api.services.economy_service.PolicyService", + return_value=mock_service, + ): + yield mock_service @pytest.fixture @@ -121,10 +122,10 @@ def mock_reform_impacts_service(): mock_service.set_error_reform_impact.return_value = None with patch( - "policyengine_api.services.economy_service.reform_impacts_service", - mock_service, - ) as mock: - yield mock + "policyengine_api.services.economy_service.ReformImpactsService", + return_value=mock_service, + ): + yield mock_service @pytest.fixture @@ -223,25 +224,24 @@ def create_mock_reform_impact( }, } ) - return { - "id": 1, - "country_id": MOCK_COUNTRY_ID, - "policy_id": MOCK_POLICY_ID, - "baseline_policy_id": MOCK_BASELINE_POLICY_ID, - "region": MOCK_REGION, - "dataset": MOCK_RESOLVED_DATASET, - "time_period": time_period, - "options_hash": options_hash, - "status": status, - "api_version": MOCK_API_VERSION, - "reform_impact_json": reform_impact_json or default_reform_impact_json, - "execution_id": execution_id, - "message": message, - "start_time": start_time or datetime.datetime(2025, 6, 26, 12, 0, 0), - "end_time": ( - datetime.datetime(2025, 6, 26, 12, 5, 0) if status == "ok" else None - ), - } + return ReformImpact( + reform_impact_id=1, + country_id=MOCK_COUNTRY_ID, + reform_policy_id=MOCK_POLICY_ID, + baseline_policy_id=MOCK_BASELINE_POLICY_ID, + region=MOCK_REGION, + dataset=MOCK_RESOLVED_DATASET, + time_period=time_period, + options_json=MOCK_OPTIONS, + options_hash=options_hash, + status=status, + api_version=MOCK_API_VERSION, + reform_impact_json=reform_impact_json or default_reform_impact_json, + execution_id=execution_id, + message=message, + start_time=start_time or datetime.datetime(2025, 6, 26, 12, 0, 0), + end_time=(datetime.datetime(2025, 6, 26, 12, 5, 0) if status == "ok" else None), + ) def create_mock_modal_execution( diff --git a/tests/fixtures/services/household_fixtures.py b/tests/fixtures/services/household_fixtures.py index 063778c7a..23602dbdf 100644 --- a/tests/fixtures/services/household_fixtures.py +++ b/tests/fixtures/services/household_fixtures.py @@ -2,6 +2,8 @@ import json from unittest.mock import patch +from policyengine_api.data.v1_models import Household + valid_request_body = { "data": {"people": {"person1": {"age": 30, "income": 50000}}}, "label": "Test Household", @@ -28,22 +30,16 @@ def mock_hash_object(): @pytest.fixture -def existing_household_record(test_db): +def existing_household_record(orm_session): """Insert an existing household record into the database.""" - test_db.query( - "INSERT INTO household (id, country_id, household_json, household_hash, label, api_version) VALUES (?, ?, ?, ?, ?, ?)", - ( - valid_db_row["id"], - valid_db_row["country_id"], - valid_db_row["household_json"], - valid_db_row["household_hash"], - valid_db_row["label"], - valid_db_row["api_version"], - ), + household = Household( + id=valid_db_row["id"], + country_id=valid_db_row["country_id"], + household_json=json.loads(valid_db_row["household_json"]), + household_hash=valid_db_row["household_hash"], + label=valid_db_row["label"], + api_version=valid_db_row["api_version"], ) - - inserted_row = test_db.query( - "SELECT * FROM household WHERE id = ?", (valid_db_row["id"],) - ).fetchone() - - return inserted_row + orm_session.add(household) + orm_session.commit() + return household diff --git a/tests/fixtures/services/policy_service.py b/tests/fixtures/services/policy_service.py index 6c4a27f66..9b655232c 100644 --- a/tests/fixtures/services/policy_service.py +++ b/tests/fixtures/services/policy_service.py @@ -2,6 +2,8 @@ import json from unittest.mock import patch +from policyengine_api.data.v1_models import Policy + valid_policy_json = { "data": {"gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433}}, } @@ -19,13 +21,6 @@ } -@pytest.fixture -def mock_database(): - """Mock the database module.""" - with patch("policyengine_api.services.policy_service.database") as mock_db: - yield mock_db - - @pytest.fixture def mock_hash_object(): """Mock the hash_object function.""" @@ -35,22 +30,16 @@ def mock_hash_object(): @pytest.fixture -def existing_policy_record(test_db): +def existing_policy_record(orm_session): """Insert an existing policy record into the database.""" - test_db.query( - "INSERT INTO policy (id, country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?, ?)", - ( - valid_policy_data["id"], - valid_policy_data["country_id"], - valid_policy_data["policy_json"], - valid_policy_data["policy_hash"], - valid_policy_data["label"], - valid_policy_data["api_version"], - ), + policy = Policy( + id=valid_policy_data["id"], + country_id=valid_policy_data["country_id"], + policy_json=json.loads(valid_policy_data["policy_json"]), + policy_hash=valid_policy_data["policy_hash"], + label=valid_policy_data["label"], + api_version=valid_policy_data["api_version"], ) - - inserted_row = test_db.query( - "SELECT * FROM policy WHERE id = ?", (valid_policy_data["id"],) - ).fetchone() - - return inserted_row + orm_session.add(policy) + orm_session.commit() + return policy diff --git a/tests/fixtures/services/report_output_fixtures.py b/tests/fixtures/services/report_output_fixtures.py deleted file mode 100644 index 3d4ca5a14..000000000 --- a/tests/fixtures/services/report_output_fixtures.py +++ /dev/null @@ -1,55 +0,0 @@ -import pytest -import json - -from policyengine_api.constants import get_report_output_cache_version - -valid_report_data = { - "country_id": "us", - "simulation_1_id": 1, - "simulation_2_id": None, - "api_version": get_report_output_cache_version("us"), - "status": "pending", - "output": None, - "error_message": None, - "year": "2025", -} - -sample_report_output = { - "population_impact": { - "baseline": {"decile_1": 1000}, - "reform": {"decile_1": 1100}, - }, - "budget_impact": 5000000, -} - - -@pytest.fixture -def existing_report_record(test_db): - """Insert an existing report output record into the database.""" - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, api_version, status, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ( - valid_report_data["country_id"], - valid_report_data["simulation_1_id"], - valid_report_data["simulation_2_id"], - valid_report_data["api_version"], - valid_report_data["status"], - valid_report_data["year"], - ), - ) - - # Get the inserted record - inserted_row = test_db.query( - """SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND status = ? - ORDER BY id DESC LIMIT 1""", - ( - valid_report_data["country_id"], - valid_report_data["simulation_1_id"], - valid_report_data["status"], - ), - ).fetchone() - - return inserted_row diff --git a/tests/fixtures/services/simulation_fixtures.py b/tests/fixtures/services/simulation_fixtures.py deleted file mode 100644 index 672553636..000000000 --- a/tests/fixtures/services/simulation_fixtures.py +++ /dev/null @@ -1,52 +0,0 @@ -import pytest - -valid_simulation_data = { - "country_id": "us", - "api_version": "1.0.0", - "population_id": "household_test_123", - "population_type": "household", - "policy_id": 1, -} - -duplicate_simulation_data = { - "country_id": "us", - "api_version": "1.0.0", - "population_id": "household_test_123", - "population_type": "household", - "policy_id": 1, -} - - -@pytest.fixture -def existing_simulation_record(test_db): - """Insert an existing simulation record into the database.""" - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ( - valid_simulation_data["country_id"], - valid_simulation_data["api_version"], - valid_simulation_data["population_id"], - valid_simulation_data["population_type"], - valid_simulation_data["policy_id"], - "pending", - ), - ) - - # Get the inserted record - inserted_row = test_db.query( - """SELECT * FROM simulations - WHERE country_id = ? AND api_version = ? - AND population_id = ? AND population_type = ? - AND policy_id = ?""", - ( - valid_simulation_data["country_id"], - valid_simulation_data["api_version"], - valid_simulation_data["population_id"], - valid_simulation_data["population_type"], - valid_simulation_data["policy_id"], - ), - ).fetchone() - - return inserted_row diff --git a/tests/fixtures/services/tracer_analysis_service.py b/tests/fixtures/services/tracer_analysis_service.py index 1a1262dc0..4118f233d 100644 --- a/tests/fixtures/services/tracer_analysis_service.py +++ b/tests/fixtures/services/tracer_analysis_service.py @@ -1,9 +1,9 @@ import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) from unittest.mock import patch +from policyengine_api.data.v1_models import Analysis valid_tracer_output = [ " snap<2027, (default)> = [6769.799]", @@ -17,10 +17,10 @@ " snap_fpg<2027-01, (default)> = [1806.4779]", ] -invalid_tracer_output = { - "variable": "only_government_benefit <1500>", - "variable": " market_income <1000>", -} +invalid_tracer_output = [ + "only_government_benefit <1500>", + " market_income <1000>", +] spliced_valid_tracer_output_root_variable = valid_tracer_output[0:] @@ -68,7 +68,11 @@ def mock_get_existing_analysis(): with patch.object( TracerAnalysisService, "get_existing_analysis", - return_value="Existing static analysis", + return_value=Analysis( + prompt="prompt", + analysis="Existing static analysis", + status="ok", + ), ) as mock: yield mock diff --git a/tests/fixtures/services/tracer_fixture_service.py b/tests/fixtures/services/tracer_fixture_service.py index bc4742a25..2d3d015d0 100644 --- a/tests/fixtures/services/tracer_fixture_service.py +++ b/tests/fixtures/services/tracer_fixture_service.py @@ -1,8 +1,6 @@ import pytest import json -from policyengine_api.services.tracer_analysis_service import ( - TracerAnalysisService, -) +from policyengine_api.data.local_models import Tracer valid_tracer = { "tracer_output": [ @@ -25,31 +23,14 @@ @pytest.fixture -def test_tracer_data(test_db): - # Insert data using query() - test_db.query( - """ - INSERT INTO tracers (household_id, policy_id, country_id, api_version, tracer_output) - VALUES (?, ?, ?, ?, ?) - """, - ( - valid_tracer_row["household_id"], - valid_tracer_row["policy_id"], - valid_tracer_row["country_id"], - valid_tracer_row["api_version"], - valid_tracer_row["tracer_output"], - ), +def test_tracer_data(orm_session): + tracer = Tracer( + household_id=int(valid_tracer_row["household_id"]), + policy_id=int(valid_tracer_row["policy_id"]), + country_id=valid_tracer_row["country_id"], + api_version=valid_tracer_row["api_version"], + tracer_output=json.loads(valid_tracer_row["tracer_output"]), ) - - # Verify that the data has been inserted - inserted_row = test_db.query( - "SELECT * FROM tracers WHERE household_id = ? AND policy_id = ? AND country_id = ? AND api_version = ?", - ( - valid_tracer_row["household_id"], - valid_tracer_row["policy_id"], - valid_tracer_row["country_id"], - valid_tracer_row["api_version"], - ), - ).fetchone() - - return inserted_row + orm_session.add(tracer) + orm_session.commit() + return tracer diff --git a/tests/fixtures/services/user_service.py b/tests/fixtures/services/user_service.py index 1ab78bc80..1520cad85 100644 --- a/tests/fixtures/services/user_service.py +++ b/tests/fixtures/services/user_service.py @@ -1,5 +1,7 @@ import pytest +from policyengine_api.data.v1_models import UserProfile + valid_user_record = { "user_id": 1, "auth0_id": "123", @@ -10,21 +12,15 @@ @pytest.fixture -def existing_user_profile(test_db): +def existing_user_profile(orm_session): """Insert an existing user record into the database.""" - test_db.query( - "INSERT INTO user_profiles (user_id, auth0_id, username, primary_country, user_since) VALUES (?, ?, ?, ?, ?)", - ( - valid_user_record["user_id"], - valid_user_record["auth0_id"], - valid_user_record["username"], - valid_user_record["primary_country"], - valid_user_record["user_since"], - ), + profile = UserProfile( + user_id=valid_user_record["user_id"], + auth0_id=valid_user_record["auth0_id"], + username=valid_user_record["username"], + primary_country=valid_user_record["primary_country"], + user_since=valid_user_record["user_since"], ) - inserted_row = test_db.query( - "SELECT * FROM user_profiles WHERE auth0_id = ?", - (valid_user_record["auth0_id"],), - ).fetchone() - - return inserted_row + orm_session.add(profile) + orm_session.commit() + return profile diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py new file mode 100644 index 000000000..9e4f0344f --- /dev/null +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -0,0 +1,147 @@ +"""Exercise the complete Alembic lifecycle against an ephemeral MySQL schema.""" + +import os + +from alembic import command +from alembic.autogenerate import compare_metadata +from alembic.config import Config +from alembic.migration import MigrationContext +import pytest +from sqlalchemy import ( + Column, + Integer, + MetaData, + String, + Table, + Text, + create_engine, + inspect, +) +from sqlalchemy.dialects.mysql import LONGTEXT +from sqlalchemy.engine import make_url + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base +from scripts.v1_database_migration import ( + DatabaseState, + database_state, + upgrade_database, +) + + +BASELINE_REVISION = "eafc2a547a4e" +PREVIOUS_REVISION = "17bb32415f97" + + +def _deployed_question_table() -> Table: + """Describe the orphaned production table without adding it to ORM metadata.""" + + metadata = MetaData() + return Table( + "question", + metadata, + Column("question_id", Integer, primary_key=True, autoincrement=True), + Column("question", Text().with_variant(LONGTEXT(), "mysql"), nullable=False), + Column("answer", Text().with_variant(LONGTEXT(), "mysql")), + Column("policy_id", Integer), + Column("country_id", String(3), nullable=False), + Column("subtask", String(32), nullable=False), + Column("status", String(32), nullable=False), + ) + + +def _ephemeral_mysql_url() -> str: + database_url = os.environ.get("ALEMBIC_DATABASE_URL", "") + if not database_url: + pytest.skip("ALEMBIC_DATABASE_URL is not set") + + url = make_url(database_url) + if url.get_backend_name() != "mysql": + pytest.fail("ALEMBIC_DATABASE_URL must use MySQL for this test") + if url.host not in {"127.0.0.1", "localhost"}: + pytest.fail("Alembic lifecycle tests may only target local MySQL") + if url.database != "policyengine_alembic_test": + pytest.fail( + "Alembic lifecycle tests require the policyengine_alembic_test schema" + ) + return database_url + + +def _alembic_config(database_url: str) -> Config: + config = Config(str(REPO / "alembic-v1.ini")) + config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) + return config + + +def test_fresh_upgrade_check_downgrade_and_reupgrade(): + database_url = _ephemeral_mysql_url() + config = _alembic_config(database_url) + engine = create_engine(database_url) + question = _deployed_question_table() + + try: + command.downgrade(config, "base") + question.drop(engine, checkfirst=True) + command.upgrade(config, "head") + command.check(config) + + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() is not None + assert compare_metadata(context, V1Base.metadata) == [] + + inspector = inspect(engine) + assert "tracers" not in inspector.get_table_names() + reform_impact_columns = { + column["name"]: column for column in inspector.get_columns("reform_impact") + } + assert reform_impact_columns["dataset"]["default"] is None + assert reform_impact_columns["execution_id"]["nullable"] is False + + command.downgrade(config, BASELINE_REVISION) + assert "question" in inspect(engine).get_table_names() + + command.upgrade(config, "head") + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() is not None + assert compare_metadata(context, V1Base.metadata) == [] + finally: + engine.dispose() + + +def test_pending_upgrade_commits_head_revision(monkeypatch): + database_url = _ephemeral_mysql_url() + config = _alembic_config(database_url) + engine = create_engine(database_url) + question = _deployed_question_table() + + try: + command.downgrade(config, "base") + question.drop(engine, checkfirst=True) + command.upgrade(config, "head") + command.downgrade(config, PREVIOUS_REVISION) + + with engine.connect() as connection: + assert database_state(connection) is DatabaseState.PENDING + + with engine.begin() as connection: + monkeypatch.delenv("ALEMBIC_DATABASE_URL") + upgrade_database(connection, backup_id="test-backup") + + with engine.connect() as connection: + assert database_state(connection) is DatabaseState.HEAD + context = MigrationContext.configure(connection) + assert compare_metadata(context, V1Base.metadata) == [] + + inspector = inspect(engine) + assert "question" not in inspector.get_table_names() + assert "tracers" not in inspector.get_table_names() + reform_impact_columns = { + column["name"]: column for column in inspector.get_columns("reform_impact") + } + assert reform_impact_columns["dataset"]["default"] is None + assert reform_impact_columns["execution_id"]["nullable"] is False + finally: + command.upgrade(config, "head") + engine.dispose() diff --git a/tests/integration/test_budget_window_in_flight_dedupe.py b/tests/integration/test_budget_window_in_flight_dedupe.py index 14a273214..5d6cb82ea 100644 --- a/tests/integration/test_budget_window_in_flight_dedupe.py +++ b/tests/integration/test_budget_window_in_flight_dedupe.py @@ -35,9 +35,11 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( from policyengine_api.libs.simulation_entrypoint import ( ModalBudgetWindowBatchExecution, ) - from policyengine_api.routes.economy_routes import economy_bp - from policyengine_api.services import economy_service as economy_service_module + from policyengine_api.routes import economy_routes from policyengine_api.services.budget_window_cache import BudgetWindowCache + from policyengine_api.services.economy_service import EconomyService + + economy_bp = economy_routes.economy_bp fake_cache = BudgetWindowCache(client=FakeRedis()) simulation_entrypoint = MagicMock() @@ -60,17 +62,17 @@ def test_budget_window_in_flight_dedupe_uses_existing_batch_without_live_db( ) ) - monkeypatch.setattr(economy_service_module, "budget_window_cache", fake_cache) - monkeypatch.setattr( - economy_service_module, "simulation_entrypoint", simulation_entrypoint - ) monkeypatch.setattr( - economy_service_module, - "reform_impacts_service", - reform_impacts_service, + economy_routes, + "economy_service", + EconomyService( + reform_impacts_service_=reform_impacts_service, + budget_window_cache_=fake_cache, + simulation_entrypoint_=simulation_entrypoint, + ), ) monkeypatch.setattr( - economy_service_module.EconomyService, + EconomyService, "_build_budget_window_batch_payload", lambda self, **kwargs: { "country_id": "us", diff --git a/tests/integration/test_simulations.py b/tests/integration/test_simulations.py index 9be9ba7e2..77937a260 100644 --- a/tests/integration/test_simulations.py +++ b/tests/integration/test_simulations.py @@ -1,8 +1,4 @@ -import pytest -from unittest.mock import Mock, patch, MagicMock -import numpy as np from policyengine_api.country import COUNTRIES -from policyengine_api.endpoints.household import add_yearly_variables from tests.fixtures.integration.simulations import ( TEST_COUNTRY_ID, SMALL_AXES_COUNT, @@ -23,7 +19,7 @@ def test__given_any_number_of_axes__sim_returns_valid_arrays( base_household, small_axes_config ) country = COUNTRIES.get(TEST_COUNTRY_ID) - result = country.calculate(household_with_axes, {}) + result = country.calculate(household_with_axes, {}).household # This variable does not function like others; it is a list of member names and is not calculated FORBIDDEN_VARIABLES = ["members"] diff --git a/tests/integration/test_v1_schema_metadata_compatibility.py b/tests/integration/test_v1_schema_metadata_compatibility.py new file mode 100644 index 000000000..ed16a0247 --- /dev/null +++ b/tests/integration/test_v1_schema_metadata_compatibility.py @@ -0,0 +1,28 @@ +"""Compare ORM metadata with an existing v1 schema without mutating it.""" + +import os + +import pytest +from sqlalchemy import create_engine + +from scripts.v1_database_migration import metadata_differences + + +def compare_existing_schema(database_url: str) -> list[str]: + """Return metadata drift without mutating the target database.""" + + engine = create_engine(database_url) + try: + with engine.connect() as connection: + return sorted(metadata_differences(connection)) + finally: + engine.dispose() + + +@pytest.mark.skipif( + "STAGE7_EXISTING_DATABASE_URL" not in os.environ, + reason="A read-only existing Cloud SQL URL was not supplied", +) +def test_live_existing_schema_matches_metadata_without_mutation(): + database_url = os.environ["STAGE7_EXISTING_DATABASE_URL"] + assert compare_existing_schema(database_url) == [] diff --git a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py index 5fa6af91c..8dbd4bb23 100644 --- a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py +++ b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py @@ -16,19 +16,11 @@ "api_version": "3.0.0", } -valid_hash_value = "some-hash" - - -@pytest.fixture -def mock_hash_object(): - """Mock the hash_object function.""" - with patch("policyengine_api.services.household_service.hash_object") as mock: - mock.return_value = valid_hash_value - yield mock - @pytest.fixture def mock_database(): - """Mock the database module.""" - with patch("policyengine_api.services.household_service.database") as mock_db: - yield mock_db + """Replace the route's service with its typed persistence boundary.""" + with patch( + "policyengine_api.routes.household_routes.household_service" + ) as household_service: + yield household_service diff --git a/tests/to_refactor/python/test_ai_analysis_service_old.py b/tests/to_refactor/python/test_ai_analysis_service_old.py deleted file mode 100644 index 0df3928ca..000000000 --- a/tests/to_refactor/python/test_ai_analysis_service_old.py +++ /dev/null @@ -1,58 +0,0 @@ -import pytest -from unittest.mock import patch, MagicMock -import json -import os -from policyengine_api.services.ai_analysis_service import AIAnalysisService - -test_ai_service = AIAnalysisService() - - -@patch("policyengine_api.services.ai_analysis_service.local_database") -def test_get_existing_analysis_found(mock_db): - mock_db.query.return_value.fetchone.return_value = {"analysis": "Existing analysis"} - - prompt = "Test prompt" - output = test_ai_service.get_existing_analysis(prompt) - - assert output == json.dumps("Existing analysis") - - # Check database query - mock_db.query.assert_called_once_with( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ) - - -@patch("policyengine_api.services.ai_analysis_service.local_database") -def test_get_existing_analysis_not_found(mock_db): - mock_db.query.return_value.fetchone.return_value = None - - prompt = "Test prompt" - result = test_ai_service.get_existing_analysis(prompt) - - assert result is None - mock_db.query.assert_called_once_with( - f"SELECT analysis FROM analysis WHERE prompt = ?", - (prompt,), - ) - - -# Additional test to check environment variable -def test_anthropic_api_key(): - with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test_key"}): - assert os.getenv("ANTHROPIC_API_KEY") == "test_key" - - -# Test error handling in trigger_ai_analysis -@patch("policyengine_api.services.ai_analysis_service.anthropic.Anthropic") -def test_trigger_ai_analysis_error(mock_anthropic): - mock_client = MagicMock() - mock_anthropic.return_value = mock_client - mock_client.messages.stream.side_effect = Exception("API Error") - - prompt = "Test prompt" - generator = test_ai_service.trigger_ai_analysis(prompt) - - # The generator should stop after the initial yield due to the error - with pytest.raises(Exception, match="API Error"): - list(generator) diff --git a/tests/to_refactor/python/test_data.py b/tests/to_refactor/python/test_data.py deleted file mode 100644 index a3f8e7c32..000000000 --- a/tests/to_refactor/python/test_data.py +++ /dev/null @@ -1,187 +0,0 @@ -import pytest -import json - -from policyengine_api.data import PolicyEngineDatabase - - -# Test the query method using the db's policy table -class TestQuery: - # Set shared variables - country_id = "us" - placeholder = "placeholder" - first_updated_placeholder = "maxwell" - second_updated_placeholder = "dworkin" - - # Initialize db connection - db = PolicyEngineDatabase(local=True, initialize=True) - - # Test INSERT and SELECT statements - def test_insert(self): - self.db.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label = ? AND api_version = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - - # Test INSERT and SELECT statements with None - def test_insert_null(self): - self.db.query( - f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - self.country_id, - self.placeholder, - self.placeholder, - None, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label IS NULL AND api_version = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - - # Test UPDATE - def test_update(self): - self.db.query( - f"UPDATE policy SET policy_json = ? WHERE policy_json = ? AND label = ? AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label = ? AND api_version = ?", - ( - self.country_id, - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["policy_json"] == self.first_updated_placeholder - - # Test UPDATE with None as search param - def test_update_none_param(self): - self.db.query( - f"UPDATE policy SET policy_json = ? WHERE policy_json = ? AND label IS NULL AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.second_updated_placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND policy_json = ? AND policy_hash = ? AND label IS NULL AND api_version = ?", - ( - self.country_id, - self.second_updated_placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - assert str(row["policy_json"]) == self.second_updated_placeholder - - # Test UPDATE with None as set value - def test_update_set_none(self): - self.db.query( - f"UPDATE policy SET label = ? WHERE policy_json = ? AND api_version = ? AND label = ? AND country_id = ? AND policy_hash = ? ", - ( - None, - self.first_updated_placeholder, - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - row = self.db.query( - f"SELECT * FROM policy WHERE country_id = ? AND label IS NULL AND policy_hash = ? AND api_version = ? AND policy_json = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.first_updated_placeholder, - ), - ).fetchone() - - assert row is not None - assert row["label"] is None - - # Test DELETE - def test_delete(self): - # Clean up the first record that was added - self.db.query( - f"DELETE FROM policy WHERE policy_json = ? AND label IS NULL AND api_version = ? AND country_id = ? AND policy_hash = ? ", - ( - self.second_updated_placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - # Delete the second for testing purposes - self.db.query( - f"DELETE FROM policy WHERE label IS NULL AND api_version = ? AND policy_json = ? AND country_id = ? AND policy_hash = ? ", - ( - self.placeholder, - self.placeholder, - self.country_id, - self.placeholder, - ), - ) - - # Confirm that it no longer exists - row = self.db.query( - f"SELECT * FROM policy WHERE label IS NULL AND country_id = ? AND api_version = ? AND policy_hash = ? AND policy_json = ?", - ( - self.country_id, - self.placeholder, - self.placeholder, - self.placeholder, - ), - ).fetchone() - - assert row is None diff --git a/tests/to_refactor/python/test_household_routes.py b/tests/to_refactor/python/test_household_routes.py index 3456429dc..fae4a1cad 100644 --- a/tests/to_refactor/python/test_household_routes.py +++ b/tests/to_refactor/python/test_household_routes.py @@ -1,26 +1,21 @@ -import pytest import json -from unittest.mock import MagicMock, patch - -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from unittest.mock import patch +from policyengine_api.data.v1_models import Household from tests.to_refactor.fixtures.to_refactor_household_fixtures import ( valid_request_body, valid_db_row, - mock_database, - mock_hash_object, ) +pytest_plugins = ["tests.to_refactor.fixtures.to_refactor_household_fixtures"] + class TestGetHousehold: def test_get_existing_household(self, rest_client, mock_database): """Test getting an existing household.""" - # Mock database response as a dict-like object - # (SQLAlchemy v2 Row objects support dict() via ._mapping) - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: valid_db_row[x] - mock_row.keys.return_value = valid_db_row.keys() - mock_database.query().fetchone.return_value = mock_row + mock_database.get_household.return_value = Household( + **{**valid_db_row, "household_json": valid_request_body["data"]} + ) # Make request response = rest_client.get("/us/household/1") @@ -32,7 +27,7 @@ def test_get_existing_household(self, rest_client, mock_database): def test_get_nonexistent_household(self, rest_client, mock_database): """Test getting a non-existent household.""" - mock_database.query().fetchone.return_value = None + mock_database.get_household.return_value = None response = rest_client.get("/us/household/999") data = json.loads(response.data) @@ -50,14 +45,11 @@ def test_get_household_invalid_id(self, rest_client): class TestCreateHousehold: - def test_create_household_success( - self, rest_client, mock_database, mock_hash_object - ): + def test_create_household_success(self, rest_client, mock_database): """Test successfully creating a new household.""" - # Mock database responses - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: {"id": 1}[x] - mock_database.query().fetchone.return_value = mock_row + mock_database.create_household.return_value = Household( + **{**valid_db_row, "id": 1, "household_json": valid_request_body["data"]} + ) response = rest_client.post( "/us/household", @@ -104,15 +96,11 @@ def test_create_household_invalid_label(self, rest_client): class TestUpdateHousehold: - def test_update_household_success( - self, rest_client, mock_database, mock_hash_object - ): + def test_update_household_success(self, rest_client, mock_database): """Test successfully updating an existing household.""" - # Mock getting existing household - mock_row = MagicMock() - mock_row.__getitem__.side_effect = lambda x: valid_db_row[x] - mock_row.keys.return_value = valid_db_row.keys() - mock_database.query().fetchone.return_value = mock_row + mock_database.get_household.return_value = Household( + **{**valid_db_row, "household_json": valid_request_body["data"]} + ) updated_household = {"people": {"person1": {"age": 31, "income": 55000}}} @@ -120,6 +108,9 @@ def test_update_household_success( "data": updated_household, "label": valid_request_body["label"], } + mock_database.update_household.return_value = Household( + **{**valid_db_row, "household_json": updated_household} + ) response = rest_client.put( "/us/household/1", @@ -131,25 +122,17 @@ def test_update_household_success( assert response.status_code == 200 assert data["status"] == "ok" assert data["result"]["household_id"] == 1 - # assert data["result"]["household_json"] == updated_data["data"] - # WHERE now includes country_id (issue #3447). - mock_database.query.assert_any_call( - "UPDATE household " - "SET household_json = ?, household_hash = ?, label = ?, api_version = ? " - "WHERE id = ? AND country_id = ?", - ( - json.dumps(updated_household), - "some-hash", - valid_request_body["label"], - COUNTRY_PACKAGE_VERSIONS.get("us"), - 1, - "us", - ), + assert data["result"]["household_json"] == updated_data["data"] + mock_database.update_household.assert_called_once_with( + "us", + 1, + updated_household, + valid_request_body["label"], ) def test_update_nonexistent_household(self, rest_client, mock_database): """Test updating a non-existent household.""" - mock_database.query().fetchone.return_value = None + mock_database.update_household.side_effect = LookupError("No household") response = rest_client.put( "/us/household/999", @@ -219,22 +202,16 @@ def test_put_household_service_error(self, mock_update, rest_client): """Test PUT endpoint when service raises an error.""" mock_update.side_effect = Exception("Failed to update household") - # First mock the get_household call that checks existence - with patch( - "policyengine_api.services.household_service.HouseholdService.get_household" - ) as mock_get: - mock_get.return_value = {"id": 1} # Simulate existing household - - response = rest_client.put( - "/us/household/1", - json={"data": {"valid": "payload"}}, - content_type="application/json", - ) - data = json.loads(response.data) - - assert response.status_code == 500 - assert data["status"] == "error" - assert "Failed to update household" in data["message"] + response = rest_client.put( + "/us/household/1", + json={"data": {"valid": "payload"}}, + content_type="application/json", + ) + data = json.loads(response.data) + + assert response.status_code == 500 + assert data["status"] == "error" + assert "Failed to update household" in data["message"] def test_missing_json_body(self, rest_client): """Test endpoints when JSON body is missing.""" diff --git a/tests/to_refactor/python/test_policy.py b/tests/to_refactor/python/test_policy.py index 0c66c09df..d5b57eb80 100644 --- a/tests/to_refactor/python/test_policy.py +++ b/tests/to_refactor/python/test_policy.py @@ -1,10 +1,9 @@ -import pytest import json -import time -import sqlite3 -from policyengine_api.data import database +from sqlalchemy import delete + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import Policy from policyengine_api.utils import hash_object -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS class TestPolicyCreation: @@ -23,10 +22,14 @@ class TestPolicyCreation: """ def test_create_unique_policy(self, rest_client): - database.query( - f"DELETE FROM policy WHERE policy_hash = ? AND label = ? AND country_id = ?", - (self.policy_hash, self.label, self.country_id), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(Policy).where( + Policy.policy_hash == self.policy_hash, + Policy.label == self.label, + Policy.country_id == self.country_id, + ) + ) res = rest_client.post("/us/policy", json=self.test_policy) return_object = json.loads(res.text) @@ -41,10 +44,14 @@ def test_create_nonunique_policy(self, rest_client): assert return_object["status"] == "ok" assert res.status_code == 200 - database.query( - f"DELETE FROM policy WHERE policy_hash = ? AND label = ? AND country_id = ?", - (self.policy_hash, self.label, self.country_id), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(Policy).where( + Policy.policy_hash == self.policy_hash, + Policy.label == self.label, + Policy.country_id == self.country_id, + ) + ) def test_create_policy_invalid_country(self, rest_client): res = rest_client.post("/au/policy", json=self.test_policy) diff --git a/tests/to_refactor/python/test_policy_service_old.py b/tests/to_refactor/python/test_policy_service_old.py deleted file mode 100644 index a84e1b1b0..000000000 --- a/tests/to_refactor/python/test_policy_service_old.py +++ /dev/null @@ -1,197 +0,0 @@ -import pytest -from assertpy import assert_that -from unittest.mock import patch, MagicMock, ANY, call -import json -from policyengine_api.services.policy_service import PolicyService - - -@pytest.fixture -def mock_database(): - with patch("policyengine_api.services.policy_service.database") as mock_db: - yield mock_db - - -@pytest.fixture -def sample_policy_data(): - return { - "id": 1, - "country_id": "US", - "policy_json": json.dumps({"param": "value"}), - "policy_hash": "hash123", - "label": "test_policy", - "api_version": "1.0.0", - } - - -@pytest.fixture -def policy_service(): - return PolicyService() - - -class TestPolicyService: - a_test_policy_id = 8 # Pre-seeded current law policies occupy IDs 1 through 5 - - def test_get_policy_success( - self, policy_service, mock_database, sample_policy_data - ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data - - # Test - result = policy_service.get_policy("us", self.a_test_policy_id) - - # Verify - assert_that(result).contains_entry({"policy_json": {"param": "value"}}) - mock_database.query.assert_called_once_with( - "SELECT * FROM policy WHERE country_id = ? AND id = ?", - ("us", self.a_test_policy_id), - ) - - def test_get_policy_not_found(self, policy_service, mock_database): - # Setup mock - mock_database.query.return_value.fetchone.return_value = None - - # Test - garbage_id = 999 - result = policy_service.get_policy("us", garbage_id) - - # Verify - assert result is None - mock_database.query.assert_called_once() - - def test_get_policy_json(self, policy_service, mock_database, sample_policy_data): - # Setup mock - mock_database.query.return_value.fetchone.return_value = { - "policy_json": sample_policy_data["policy_json"] - } - - # Test - result = policy_service.get_policy_json("us", self.a_test_policy_id) - - # Verify - assert result == sample_policy_data["policy_json"] - mock_database.query.assert_called_once() - - def test_set_policy_new(self, policy_service, mock_database): - new_policy_id = 10 - - # Setup mocks - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call for existing policy check - {"id": new_policy_id}, # Second call to get new policy - ] - - test_policy = {"param": "value"} - test_label = "new_policy" - test_country_id = "us" - - expected_calls = [ - # First call - check if policy exists - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, ANY, test_label), - ), - # Second call - insert new policy - call( - "INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - test_country_id, - json.dumps(test_policy), - ANY, - test_label, - ANY, - ), - ), - # Third call - get the newly created policy - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, ANY, test_label), - ), - ] - - # Test - policy_id, message, exists = policy_service.set_policy( - test_country_id, test_label, test_policy - ) - - # Verify - assert policy_id == new_policy_id - assert message == "Policy created" - assert exists is False - assert mock_database.query.call_args_list == expected_calls - - def test_set_policy_existing( - self, policy_service, mock_database, sample_policy_data - ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data - - # Test - policy_id, message, exists = policy_service.set_policy( - "us", - sample_policy_data["label"], - json.loads(sample_policy_data["policy_json"]), - ) - - # Verify - assert policy_id == sample_policy_data["id"] - assert message == "Policy already exists" - assert exists is True - mock_database.query.assert_called_once() - - def test_get_unique_policy_with_label( - self, policy_service, mock_database, sample_policy_data - ): - # Setup mock - mock_database.query.return_value.fetchone.return_value = sample_policy_data - - # Test - result = policy_service._get_unique_policy_with_label( - "us", - sample_policy_data["policy_hash"], - sample_policy_data["label"], - ) - - # Verify - assert result == sample_policy_data - mock_database.query.assert_called_once() - - def test_get_unique_policy_with_null_label(self, policy_service, mock_database): - # Setup mock - mock_database.query.return_value.fetchone.return_value = None - - # Test - result = policy_service._get_unique_policy_with_label("us", "hash123", None) - - # Verify - assert result is None - mock_database.query.assert_called_once_with( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label IS NULL", - ("us", "hash123"), - ) - - @pytest.mark.parametrize( - "error_method", - [ - "get_policy", - "get_policy_json", - "set_policy", - "_get_unique_policy_with_label", - ], - ) - def test_error_handling(self, policy_service, mock_database, error_method): - # Setup mock to raise exception - mock_database.query.side_effect = Exception("Database error") - - # Test - with pytest.raises(Exception) as exc_info: - if error_method == "get_policy": - policy_service.get_policy("us", 1) - elif error_method == "get_policy_json": - policy_service.get_policy_json("us", 1) - elif error_method == "set_policy": - policy_service.set_policy("us", "label", {}) - else: - policy_service._get_unique_policy_with_label("us", "hash", "label") - - assert str(exc_info.value) == "Database error" diff --git a/tests/to_refactor/python/test_simulation_analysis_routes.py b/tests/to_refactor/python/test_simulation_analysis_routes.py index 7e7d5ea34..7d0086a7b 100644 --- a/tests/to_refactor/python/test_simulation_analysis_routes.py +++ b/tests/to_refactor/python/test_simulation_analysis_routes.py @@ -1,14 +1,9 @@ -import pytest from unittest.mock import patch -from flask import Flask +from policyengine_api.data.v1_models import Analysis from policyengine_api.services.simulation_analysis_service import ( SimulationAnalysisService, ) -from policyengine_api.routes.simulation_analysis_routes import ( - execute_simulation_analysis, -) - from tests.to_refactor.fixtures.simulation_analysis_fixtures import ( test_json, test_impact, @@ -21,7 +16,11 @@ def test_execute_simulation_analysis_existing_analysis(rest_client): with patch( "policyengine_api.services.ai_analysis_service.AIAnalysisService.get_existing_analysis" ) as mock_get_existing: - mock_get_existing.return_value = "Existing analysis" + mock_get_existing.return_value = Analysis( + prompt="prompt", + analysis="Existing analysis", + status="ok", + ) response = rest_client.post("/us/simulation-analysis", json=test_json) @@ -82,7 +81,7 @@ def test_execute_simulation_analysis_custom_dataset(rest_client): } with patch( "policyengine_api.services.simulation_analysis_service.SimulationAnalysisService._generate_simulation_analysis_prompt" - ) as mock_generate_prompt: + ): with patch( "policyengine_api.services.ai_analysis_service.AIAnalysisService.get_existing_analysis" ) as mock_get_existing: diff --git a/tests/to_refactor/python/test_tracer_analysis_routes.py b/tests/to_refactor/python/test_tracer_analysis_routes.py index 83f7bde23..851c67bef 100644 --- a/tests/to_refactor/python/test_tracer_analysis_routes.py +++ b/tests/to_refactor/python/test_tracer_analysis_routes.py @@ -1,7 +1,6 @@ -import pytest from flask import json from unittest.mock import patch -from werkzeug.exceptions import BadRequest +from werkzeug.exceptions import NotFound # constants VALID_HOUSEHOLD_ID = 123 @@ -12,19 +11,12 @@ INVALID_VARIABLE = 123 -@patch("policyengine_api.services.tracer_analysis_service.local_database") -@patch( - "policyengine_api.services.tracer_analysis_service.TracerAnalysisService.trigger_ai_analysis" -) -def test_execute_tracer_analysis_success( - mock_trigger_ai_analysis, mock_db, rest_client -): - mock_db.query.return_value.fetchone.return_value = { - "tracer_output": json.dumps( - ["disposable_income <1000>", " market_income <1000>"] - ) - } - mock_trigger_ai_analysis.return_value = "AI analysis result" +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_success(mock_service, rest_client): + mock_service.execute_analysis.return_value = ( + iter(["AI analysis result"]), + "streaming", + ) test_household_id = 1500 # Set this to US current law @@ -43,9 +35,11 @@ def test_execute_tracer_analysis_success( assert b"AI analysis result" in response.data -@patch("policyengine_api.services.tracer_analysis_service.local_database") -def test_execute_tracer_analysis_no_tracer(mock_db, rest_client): - mock_db.query.return_value.fetchone.return_value = None +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_no_tracer(mock_service, rest_client): + mock_service.execute_analysis.side_effect = NotFound( + "No household simulation tracer found" + ) response = rest_client.post( "/us/tracer-analysis", @@ -62,19 +56,9 @@ def test_execute_tracer_analysis_no_tracer(mock_db, rest_client): ) -@patch("policyengine_api.services.tracer_analysis_service.local_database") -@patch( - "policyengine_api.services.tracer_analysis_service.TracerAnalysisService.trigger_ai_analysis" -) -def test_execute_tracer_analysis_ai_error( - mock_trigger_ai_analysis, mock_db, rest_client -): - mock_db.query.return_value.fetchone.return_value = { - "tracer_output": json.dumps( - ["disposable_income <1000>", " market_income <1000>"] - ) - } - mock_trigger_ai_analysis.side_effect = Exception(KeyError) +@patch("policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service") +def test_execute_tracer_analysis_ai_error(mock_service, rest_client): + mock_service.execute_analysis.side_effect = Exception(KeyError) test_household_id = 1500 test_policy_id = 2 @@ -93,8 +77,7 @@ def test_execute_tracer_analysis_ai_error( assert json.loads(response.data)["status"] == "error" -@patch("policyengine_api.services.tracer_analysis_service.local_database") -def test_invalid_variable_types(mock_db, rest_client): +def test_invalid_variable_types(rest_client): """Test that different non-string variable types are rejected""" invalid_variables = [ 123, diff --git a/tests/to_refactor/python/test_user_profile_routes.py b/tests/to_refactor/python/test_user_profile_routes.py index ec30f9eef..3b2296b4f 100644 --- a/tests/to_refactor/python/test_user_profile_routes.py +++ b/tests/to_refactor/python/test_user_profile_routes.py @@ -1,8 +1,11 @@ import json -from datetime import datetime -from policyengine_api.data import database import time +from sqlalchemy import delete, select + +from policyengine_api.data.orm import get_v1_session_factory +from policyengine_api.data.v1_models import UserProfile + class TestUserProfiles: # Define the profile to test against @@ -22,13 +25,13 @@ class TestUserProfiles: """ def test_set_and_get_record(self, rest_client): - database.query( - f"DELETE FROM user_profiles WHERE auth0_id = ? AND primary_country = ?", - ( - self.auth0_id, - self.primary_country, - ), - ) + with get_v1_session_factory().begin() as session: + session.execute( + delete(UserProfile).where( + UserProfile.auth0_id == self.auth0_id, + UserProfile.primary_country == self.primary_country, + ) + ) res = rest_client.post("/us/user-profile", json=self.test_profile) return_object = json.loads(res.text) @@ -43,7 +46,7 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert return_object["result"]["auth0_id"] == self.auth0_id assert return_object["result"]["primary_country"] == self.primary_country - assert return_object["result"]["username"] == None + assert return_object["result"]["username"] is None user_id = return_object["result"]["user_id"] @@ -54,7 +57,7 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert return_object["result"]["primary_country"] == self.primary_country assert return_object["result"].get("auth0_id") is None - assert return_object["result"]["username"] == None + assert return_object["result"]["username"] is None test_username = "maxwell" updated_profile = {"user_id": user_id, "username": test_username} @@ -65,11 +68,14 @@ def test_set_and_get_record(self, rest_client): assert return_object["status"] == "ok" assert res.status_code == 200 - row = database.query( - f"SELECT * FROM user_profiles WHERE user_id = ? AND username = ?", - (user_id, test_username), - ).fetchone() - assert row is not None + with get_v1_session_factory()() as session: + row = session.scalar( + select(UserProfile).where( + UserProfile.user_id == user_id, + UserProfile.username == test_username, + ) + ) + assert row is not None malicious_updated_profile = {**updated_profile, "auth0_id": "BOGUS"} @@ -78,22 +84,15 @@ def test_set_and_get_record(self, rest_client): assert res.status_code == 200 - row = database.query( - f"SELECT * FROM user_profiles WHERE username = ?", - (test_username,), - ).fetchone() - - assert row["auth0_id"] == self.auth0_id - - database.query( - f"DELETE FROM user_profiles WHERE user_id = ? AND auth0_id = ? AND primary_country = ?", - (user_id, self.auth0_id, self.primary_country), - ) + with get_v1_session_factory().begin() as session: + row = session.scalar( + select(UserProfile).where(UserProfile.username == test_username) + ) + assert row.auth0_id == self.auth0_id + session.delete(row) def test_non_existent_record(self, rest_client): non_existent_auth0_id = "non-existent-auth0-id" res = rest_client.get(f"/us/user-profile?auth0_id={non_existent_auth0_id}") - return_object = json.loads(res.text) - assert res.status_code == 404 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 336fe07f0..b703a0b76 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,135 +1,54 @@ import os -import sqlite3 import pytest +from sqlalchemy import create_engine +from sqlalchemy.pool import StaticPool -# The legacy SQL module creates its default database object at import time. -# Keep unit tests on SQLite until the SQL layer is broadly refactored in a -# later migration stage. -os.environ.setdefault("FLASK_DEBUG", "1") - -from policyengine_api.constants import REPO -from policyengine_api.data import PolicyEngineDatabase - - -class TestPolicyEngineDatabase(PolicyEngineDatabase): - """Test version of PolicyEngineDatabase that uses in-memory SQLite""" - - def __init__(self, initialize: bool = True): - self.local = True # Always use SQLite for tests - if initialize: - self._setup_connection() - self.initialize() - - def _setup_connection(self): - """Setup the in-memory connection""" - if not hasattr(self, "_connection"): - self._connection = sqlite3.connect(":memory:") - - def dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - - self._connection.row_factory = dict_factory - - def initialize(self): - """ - Override initialize to avoid file operations from parent class - """ - self._setup_connection() - - # Read the SQL initialization file - init_file = ( - REPO - / "policyengine_api" - / "data" - / f"initialise{'_local' if self.local else ''}.sql" - ) - with open(init_file) as f: - full_query = f.read() - - # Split and execute the queries - queries = full_query.split(";") - for query in queries: - if query.strip(): # Skip empty queries - self.query(query) - - def query(self, *query): - """Override query method to use in-memory connection""" - if not hasattr(self, "_connection"): - # Create a persistent connection for the in-memory database - self._connection = sqlite3.connect(self.db_url) - - def dict_factory(cursor, row): - d = {} - for idx, col in enumerate(cursor.description): - d[col[0]] = row[idx] - return d - self._connection.row_factory = dict_factory - - cursor = self._connection.cursor() - result = cursor.execute(*query) - self._connection.commit() - return result - - def clean(self): - """Clear all data from tables while preserving the schema""" - if hasattr(self, "_connection"): - cursor = self._connection.cursor() - - # Get all table names - tables = cursor.execute( - "SELECT name FROM sqlite_master WHERE type='table'" - ).fetchall() - - # Disable foreign key checks temporarily - cursor.execute("PRAGMA foreign_keys=OFF") - - # Delete all data from each table - for table in tables: - table_name = table["name"] - cursor.execute(f"DELETE FROM {table_name}") - - # Re-enable foreign key checks - cursor.execute("PRAGMA foreign_keys=ON") +os.environ.setdefault("FLASK_DEBUG", "1") - self._connection.commit() +from policyengine_api.data import orm +from policyengine_api.data.local_database import create_local_v1_schema +from policyengine_api.data.local_models import LocalV1Base +from policyengine_api.data.v1_models import V1Base @pytest.fixture(scope="session") -def test_db(): - """Create a test database instance that persists for the whole test session""" - db = TestPolicyEngineDatabase(initialize=True) - yield db - # Clean up the connection when done - if hasattr(db, "_connection"): - db._connection.close() +def test_engine(): + engine = create_engine( + "sqlite+pysqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + create_local_v1_schema(engine) + yield engine + engine.dispose() @pytest.fixture(autouse=True) -def override_database(test_db, monkeypatch): - """ - Global database override that affects all imports of the database. - This fixture automatically applies to all tests. - """ - test_db.clean() - - # Patch at the root module level where database is defined - import policyengine_api.data - - monkeypatch.setattr(policyengine_api.data, "database", test_db) - - # Also patch the module-level variable for any existing imports - import sys - - for module_name, module in list(sys.modules.items()): - if module_name.startswith("policyengine_api."): - if hasattr(module, "database"): - monkeypatch.setattr(module, "database", test_db) - if hasattr(module, "local_database"): - monkeypatch.setattr(module, "local_database", test_db) - - yield test_db +def isolated_orm_database(test_engine, monkeypatch): + """Bind runtime factories to a clean in-memory ORM database per test.""" + + monkeypatch.setattr(orm, "get_v1_engine", lambda *, local=False: test_engine) + orm.clear_v1_session_factories() + factory = orm.get_v1_session_factory() + with factory.begin() as session: + local_tables = LocalV1Base.metadata.sorted_tables + production_tables = V1Base.metadata.sorted_tables + for table in reversed([*production_tables, *local_tables]): + session.execute(table.delete()) + try: + yield + finally: + orm.clear_v1_session_factories() + + +@pytest.fixture +def orm_session_factory(isolated_orm_database): + return orm.get_v1_session_factory() + + +@pytest.fixture +def orm_session(orm_session_factory): + with orm_session_factory() as session: + yield session diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py new file mode 100644 index 000000000..a6bf8aeb5 --- /dev/null +++ b/tests/unit/data/test_alembic_baseline.py @@ -0,0 +1,66 @@ +from io import StringIO + +from alembic import command +from alembic.config import Config +from alembic.script import ScriptDirectory +import pytest + +from policyengine_api.constants import REPO +from policyengine_api.data.v1_models import V1Base + + +def _mysql_offline_config() -> tuple[Config, StringIO]: + output = StringIO() + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=output) + config.set_main_option( + "sqlalchemy.url", + "mysql+pymysql://offline:offline@localhost/offline", + ) + return config, output + + +def test_default_configuration_requires_an_explicit_database_url(monkeypatch): + monkeypatch.delenv("ALEMBIC_DATABASE_URL", raising=False) + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=StringIO()) + + with pytest.raises(RuntimeError, match="ALEMBIC_DATABASE_URL"): + command.upgrade(config, "head", sql=True) + + +def test_sqlite_is_rejected_as_a_migration_target(monkeypatch): + monkeypatch.setenv("ALEMBIC_DATABASE_URL", "sqlite+pysqlite:///:memory:") + config = Config(str(REPO / "alembic-v1.ini"), output_buffer=StringIO()) + + with pytest.raises(RuntimeError, match="MySQL"): + command.upgrade(config, "head", sql=True) + + +def test_baseline_renders_the_v1_schema_for_mysql_without_connecting(): + config, output = _mysql_offline_config() + + command.upgrade(config, "head", sql=True) + + rendered_sql = output.getvalue() + for table_name in V1Base.metadata.tables: + assert f"CREATE TABLE {table_name}" in rendered_sql + assert "CREATE TABLE alembic_version" in rendered_sql + + +def test_baseline_is_the_single_root_revision(): + config, _ = _mysql_offline_config() + scripts = ScriptDirectory.from_config(config) + baseline = scripts.get_revision("eafc2a547a4e") + + assert scripts.get_bases() == ["eafc2a547a4e"] + assert baseline is not None + assert baseline.down_revision is None + + +def test_baseline_renders_a_complete_mysql_downgrade_without_connecting(): + config, output = _mysql_offline_config() + + command.downgrade(config, "head:base", sql=True) + + rendered_sql = output.getvalue() + for table_name in V1Base.metadata.tables: + assert f"DROP TABLE {table_name}" in rendered_sql diff --git a/tests/unit/data/test_orm_fixtures.py b/tests/unit/data/test_orm_fixtures.py new file mode 100644 index 000000000..771c1ecc7 --- /dev/null +++ b/tests/unit/data/test_orm_fixtures.py @@ -0,0 +1,33 @@ +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from policyengine_api.data.v1_models import Household + + +def test_orm_session_factory_uses_mapped_models_and_python_json( + orm_session_factory, +): + assert isinstance(orm_session_factory, sessionmaker) + + with orm_session_factory.begin() as session: + household = Household( + country_id="uk", + label="Fixture household", + api_version="1.0.0", + household_json={"people": {"you": {"age": {"2025": 40}}}}, + household_hash="fixture-hash", + ) + session.add(household) + + with orm_session_factory() as session: + stored = session.scalar( + select(Household).where(Household.household_hash == "fixture-hash") + ) + + assert stored is not None + assert stored.household_json == {"people": {"you": {"age": {"2025": 40}}}} + + +def test_orm_session_fixture_is_a_live_session(orm_session): + assert isinstance(orm_session, Session) + assert orm_session.is_active diff --git a/tests/unit/data/test_orm_sessions.py b/tests/unit/data/test_orm_sessions.py new file mode 100644 index 000000000..681593c7b --- /dev/null +++ b/tests/unit/data/test_orm_sessions.py @@ -0,0 +1,69 @@ +import pytest +from sqlalchemy import create_engine, text +from sqlalchemy.orm import Session + +import policyengine_api.data.orm as orm_module +from policyengine_api.data.orm import ( + build_session_factory, + get_v1_session_factory, +) + + +def test_session_factory_creates_distinct_sessions_bound_to_one_engine(): + engine = create_engine("sqlite+pysqlite:///:memory:") + factory = build_session_factory(engine) + + first = factory() + second = factory() + try: + assert isinstance(first, Session) + assert isinstance(second, Session) + assert first is not second + assert first.get_bind() is engine + assert second.get_bind() is engine + assert first.expire_on_commit is False + finally: + first.close() + second.close() + + +def test_session_factory_begin_commits_and_rolls_back(): + engine = create_engine("sqlite+pysqlite:///:memory:") + factory = build_session_factory(engine) + with engine.begin() as connection: + connection.execute(text("CREATE TABLE item (id INTEGER PRIMARY KEY)")) + + with factory.begin() as session: + session.execute(text("INSERT INTO item (id) VALUES (1)")) + + with pytest.raises(RuntimeError, match="stop"): + with factory.begin() as session: + session.execute(text("INSERT INTO item (id) VALUES (2)")) + raise RuntimeError("stop") + + with factory() as session: + assert session.scalars(text("SELECT id FROM item ORDER BY id")).all() == [1] + + +def test_runtime_factories_are_cached_and_separate(monkeypatch): + remote_engine = create_engine("sqlite+pysqlite:///:memory:") + local_engine = create_engine("sqlite+pysqlite:///:memory:") + + monkeypatch.setattr( + orm_module, + "get_v1_engine", + lambda *, local=False: local_engine if local else remote_engine, + ) + orm_module.clear_v1_session_factories() + + try: + remote = get_v1_session_factory() + local = get_v1_session_factory(local=True) + + assert remote is get_v1_session_factory() + assert local is get_v1_session_factory(local=True) + assert remote is not local + assert remote.kw["bind"] is remote_engine + assert local.kw["bind"] is local_engine + finally: + orm_module.clear_v1_session_factories() diff --git a/tests/unit/data/test_remote_database_config.py b/tests/unit/data/test_remote_database_config.py index 33ce245a4..118594f27 100644 --- a/tests/unit/data/test_remote_database_config.py +++ b/tests/unit/data/test_remote_database_config.py @@ -2,11 +2,23 @@ os.environ.setdefault("FLASK_DEBUG", "1") -from policyengine_api.data.data import get_remote_database_config +import pytest +from policyengine_api.data.orm import get_remote_database_config -def test_remote_database_config_defaults_to_current_production_values(monkeypatch): + +def test_remote_database_config_requires_instance_connection_name(monkeypatch): monkeypatch.delenv("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", raising=False) + + with pytest.raises(KeyError, match="POLICYENGINE_DB_INSTANCE_CONNECTION_NAME"): + get_remote_database_config() + + +def test_remote_database_config_defaults_non_target_settings(monkeypatch): + monkeypatch.setenv( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", + "policyengine-api:us-central1:policyengine-api-data", + ) monkeypatch.delenv("POLICYENGINE_DB_USER", raising=False) monkeypatch.delenv("POLICYENGINE_DB_NAME", raising=False) diff --git a/tests/unit/data/test_run_schema.py b/tests/unit/data/test_run_schema.py index 2bcba1eff..09d7740c5 100644 --- a/tests/unit/data/test_run_schema.py +++ b/tests/unit/data/test_run_schema.py @@ -1,15 +1,12 @@ -from pathlib import Path +from sqlalchemy import inspect -from policyengine_api.constants import REPO +def _column_names(test_engine, table_name: str) -> set[str]: + return {column["name"] for column in inspect(test_engine).get_columns(table_name)} -def _column_names(test_db, table_name: str) -> set[str]: - rows = test_db.query(f"PRAGMA table_info({table_name})").fetchall() - return {row["name"] for row in rows} - -def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): - report_output_columns = _column_names(test_db, "report_outputs") +def test_stage_one_run_schema_is_initialized_in_local_test_db(test_engine): + report_output_columns = _column_names(test_engine, "report_outputs") assert { "report_kind", "report_spec_json", @@ -19,7 +16,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "latest_successful_run_id", }.issubset(report_output_columns) - simulation_columns = _column_names(test_db, "simulations") + simulation_columns = _column_names(test_engine, "simulations") assert { "simulation_spec_json", "simulation_spec_schema_version", @@ -27,7 +24,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "latest_successful_run_id", }.issubset(simulation_columns) - report_run_columns = _column_names(test_db, "report_output_runs") + report_run_columns = _column_names(test_engine, "report_output_runs") assert { "id", "report_output_id", @@ -45,7 +42,7 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "resolved_options_hash", }.issubset(report_run_columns) - simulation_run_columns = _column_names(test_db, "simulation_runs") + simulation_run_columns = _column_names(test_engine, "simulation_runs") assert { "id", "simulation_id", @@ -61,28 +58,5 @@ def test_stage_one_run_schema_is_initialized_in_local_test_db(test_db): "simulation_cache_version", }.issubset(simulation_run_columns) - alias_columns = _column_names(test_db, "legacy_report_output_aliases") + alias_columns = _column_names(test_engine, "legacy_report_output_aliases") assert {"legacy_report_output_id", "canonical_report_output_id"} == alias_columns - - -def test_stage_one_schema_is_defined_in_both_sql_initializers(): - sql_paths = [ - REPO / "policyengine_api" / "data" / "initialise.sql", - REPO / "policyengine_api" / "data" / "initialise_local.sql", - ] - - required_snippets = [ - "CREATE TABLE IF NOT EXISTS report_output_runs", - "CREATE TABLE IF NOT EXISTS simulation_runs", - "CREATE TABLE IF NOT EXISTS legacy_report_output_aliases", - "report_spec_json", - "report_spec_status", - "simulation_spec_json", - "active_run_id", - "latest_successful_run_id", - ] - - for sql_path in sql_paths: - sql_text = Path(sql_path).read_text() - for snippet in required_snippets: - assert snippet in sql_text, f"{snippet} missing from {sql_path.name}" diff --git a/tests/unit/data/test_runtime_sql_boundaries.py b/tests/unit/data/test_runtime_sql_boundaries.py new file mode 100644 index 000000000..6f42a5708 --- /dev/null +++ b/tests/unit/data/test_runtime_sql_boundaries.py @@ -0,0 +1,57 @@ +"""Guards that confine runtime SQL access to approved persistence layers.""" + +from pathlib import Path + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +def test_runtime_sql_is_confined_to_the_data_access_layer(): + offenders = [] + for path in PACKAGE_ROOT.rglob("*.py"): + if "data" in path.relative_to(PACKAGE_ROOT).parts: + continue + source = path.read_text(encoding="utf-8") + if any( + token in source + for token in ( + "database.query(", + "local_database.query(", + "database.transaction(", + "local_database.transaction(", + ) + ): + offenders.append(str(path.relative_to(PACKAGE_ROOT))) + assert offenders == [] + + +def test_ordinary_runtime_modules_no_longer_use_raw_sql_facade(): + relative_paths = ( + "routes/household_routes.py", + "routes/policy_routes.py", + "routes/reform_impact_routes.py", + "country.py", + "services/ai_analysis_service.py", + "services/reform_impacts_service.py", + "services/tracer_analysis_service.py", + ) + for relative_path in relative_paths: + source = (PACKAGE_ROOT / relative_path).read_text(encoding="utf-8") + assert "runtime_sqlalchemy_dao" not in source + + +def test_report_orchestration_uses_sessions_and_models_not_raw_sql(): + source = (PACKAGE_ROOT / "services/report_output_service.py").read_text( + encoding="utf-8" + ) + assert "SQLAlchemyDAO" not in source + assert ".query(" not in source + assert "exec_driver_sql" not in source + + +def test_legacy_persistence_compatibility_layers_are_absent(): + assert not (PACKAGE_ROOT / "data/v1_daos.py").exists() + assert not (PACKAGE_ROOT / "data/data.py").exists() + orm_source = (PACKAGE_ROOT / "data/orm.py").read_text(encoding="utf-8") + assert "SessionManager" not in orm_source + assert "PolicyEngineDatabase" not in orm_source diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 1f380b9df..89d9689e2 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -1,259 +1,138 @@ -"""Tests for SQLAlchemy v2 compatibility. +from unittest.mock import Mock -These tests verify that the database layer works correctly with -SQLAlchemy v2, specifically: -- The _ResultProxy wrapper provides fetchone()/fetchall() on eagerly - fetched results. -- The remote (non-local) query path uses connection-based execution - instead of the removed engine.execute(). -- Row objects returned from the remote path support dict-like access - (dict(row) and row["key"]). -""" - -import policyengine_api.data.data as data_module import sqlalchemy -from policyengine_api.data.data import PolicyEngineDatabase, _ResultProxy - - -class TestSQLAlchemyVersion: - """Verify that SQLAlchemy v2 is installed.""" - - def test_sqlalchemy_version_is_v2(self): - major = int(sqlalchemy.__version__.split(".")[0]) - assert major >= 2, f"Expected SQLAlchemy v2+, got {sqlalchemy.__version__}" - - -class TestResultProxy: - """Test the _ResultProxy wrapper that bridges SQLAlchemy v2 - connection-scoped results with the existing query() API.""" - - def test_fetchone_returns_dict_like_rows(self): - """Rows returned by fetchone() should support dict() and - key-based access.""" - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql( - "CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)" - ) - conn.exec_driver_sql("INSERT INTO test VALUES (1, 'hello')") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - row = proxy.fetchone() - assert row is not None - assert dict(row) == {"id": 1, "name": "hello"} - assert row["id"] == 1 - assert row["name"] == "hello" - - def test_fetchone_returns_none_when_exhausted(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - assert proxy.fetchone() is None - - def test_fetchall_returns_all_rows(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY, val TEXT)") - conn.exec_driver_sql("INSERT INTO test VALUES (1, 'a')") - conn.exec_driver_sql("INSERT INTO test VALUES (2, 'b')") - conn.exec_driver_sql("INSERT INTO test VALUES (3, 'c')") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - rows = proxy.fetchall() - assert len(rows) == 3 - assert dict(rows[0]) == {"id": 1, "val": "a"} - assert dict(rows[2]) == {"id": 3, "val": "c"} - - def test_fetchone_then_fetchall_respects_cursor_position(self): - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - conn.exec_driver_sql("INSERT INTO test VALUES (1)") - conn.exec_driver_sql("INSERT INTO test VALUES (2)") - conn.exec_driver_sql("INSERT INTO test VALUES (3)") - result = conn.exec_driver_sql("SELECT * FROM test") - proxy = _ResultProxy(result) - - first = proxy.fetchone() - assert dict(first) == {"id": 1} - remaining = proxy.fetchall() - assert len(remaining) == 2 - assert dict(remaining[0]) == {"id": 2} - - def test_result_proxy_for_insert_statement(self): - """INSERT statements produce no rows; _ResultProxy should - handle this gracefully.""" - engine = sqlalchemy.create_engine("sqlite://") - with engine.connect() as conn: - conn.exec_driver_sql("CREATE TABLE test (id INTEGER PRIMARY KEY)") - result = conn.exec_driver_sql("INSERT INTO test VALUES (1)") - proxy = _ResultProxy(result) - - assert proxy.fetchone() is None - assert proxy.fetchall() == [] - - -class TestRemoteQueryPath: - """Test the non-local query path that uses SQLAlchemy engine - with connection-based execution (v2 pattern).""" - - def _make_remote_db(self): - """Create a PolicyEngineDatabase-like object that uses - a SQLAlchemy engine (the 'remote' path) but backed by - in-memory SQLite for testing.""" - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db.local = False - db.pool = sqlalchemy.create_engine("sqlite://") - # Initialize schema using the remote path - with db.pool.connect() as conn: - conn.exec_driver_sql( - "CREATE TABLE test_table " - "(id INTEGER PRIMARY KEY, name TEXT, value REAL)" +from sqlalchemy import create_engine, func, inspect, select +from sqlalchemy.orm import Session + +import policyengine_api.data.orm as orm +from policyengine_api.data.v1_models import Policy + + +def test_sqlalchemy_v2_or_newer_is_installed(): + assert int(sqlalchemy.__version__.split(".")[0]) >= 2 + + +def test_remote_engine_delegates_connections_and_pooling_to_sqlalchemy(monkeypatch): + connector_calls = [] + engine_calls = [] + + class FakeConnector: + def __init__(self, **options): + self.options = options + + def connect(self, **options): + connector_calls.append(options) + return object() + + def close(self): + pass + + connector = FakeConnector() + + def connector_factory(**options): + connector.options = options + return connector + + monkeypatch.setattr(orm, "Connector", connector_factory) + monkeypatch.setattr( + orm, + "create_engine", + lambda url, **options: engine_calls.append((url, options)) or "engine", + ) + monkeypatch.setenv( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", "project:region:instance" + ) + monkeypatch.setenv("POLICYENGINE_DB_USER", "user") + monkeypatch.setenv("POLICYENGINE_DB_NAME", "database") + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "password") + + engine = orm._build_remote_engine() + + assert engine == "engine" + assert connector.options == { + "ip_type": orm.IPTypes.PUBLIC, + "refresh_strategy": "LAZY", + } + url, options = engine_calls[0] + assert url == "mysql+pymysql://" + assert options["pool_pre_ping"] is True + assert options["pool_recycle"] == 1800 + assert options["pool_size"] == 5 + assert options["max_overflow"] == 2 + assert options["pool_timeout"] == 30 + assert connector_calls == [] + options["creator"]() + assert connector_calls == [ + { + "instance_connection_string": "project:region:instance", + "driver": "pymysql", + "db": "database", + "user": "user", + "password": "password", + } + ] + + +def test_database_password_can_be_loaded_from_file(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", ".dbpw") + (tmp_path / ".dbpw").write_text("file-password\n", encoding="utf-8") + + assert orm._database_password() == "file-password" + + +def test_local_schema_preserves_the_documented_sqlite_policy_key_exception(): + from policyengine_api.data.local_database import create_local_v1_schema + + engine = create_engine("sqlite+pysqlite:///:memory:") + try: + create_local_v1_schema(engine) + + policy_key = inspect(engine).get_pk_constraint("policy") + assert policy_key["constrained_columns"] == ["id"] + assert "tracers" in inspect(engine).get_table_names() + finally: + engine.dispose() + + +def test_local_initializer_bootstraps_schema_and_current_law_rows(tmp_path): + database_path = tmp_path / "local.db" + engine = create_engine(f"sqlite+pysqlite:///{database_path}") + + try: + orm._initialize_local_database(engine) + with Session(engine) as session: + assert session.scalar(select(func.count()).select_from(Policy)) > 0 + assert session.scalars(select(Policy)).first().policy_json == {} + finally: + engine.dispose() + + +def test_local_initializer_is_idempotent(tmp_path): + engine = create_engine(f"sqlite+pysqlite:///{tmp_path / 'local.db'}") + try: + orm._initialize_local_database(engine) + orm._initialize_local_database(engine) + + with Session(engine) as session: + assert session.scalar(select(func.count()).select_from(Policy)) == len( + orm.COUNTRY_PACKAGE_VERSIONS ) - conn.commit() - return db - - def test_remote_insert_and_select(self): - """Test INSERT then SELECT through the remote query path.""" - db = self._make_remote_db() - - # Note: remote path converts ? to %s for MySQL, but SQLite - # uses ? natively. Since exec_driver_sql passes to the DBAPI - # driver directly and SQLite's driver uses ?, we need to - # test with the actual query() method which does the conversion. - # For SQLite DBAPI, ? is the native marker. - - # Use exec_driver_sql directly to bypass ?->%s conversion - # (which would break SQLite) - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "test", 3.14), - ] - ) - - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - row = result.fetchone() - assert row is not None - assert row["id"] == 1 - assert row["name"] == "test" - assert row["value"] == 3.14 - assert dict(row) == {"id": 1, "name": "test", "value": 3.14} - - def test_remote_select_no_results(self): - db = self._make_remote_db() - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (999,)]) - assert result.fetchone() is None - - def test_remote_update(self): - db = self._make_remote_db() - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "original", 1.0), - ] - ) - db._execute_remote( - [ - "UPDATE test_table SET name = ? WHERE id = ?", - ("updated", 1), - ] - ) - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - row = result.fetchone() - assert row["name"] == "updated" - - def test_remote_delete(self): - db = self._make_remote_db() - db._execute_remote( - [ - "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)", - (1, "to_delete", 0.0), - ] - ) - db._execute_remote(["DELETE FROM test_table WHERE id = ?", (1,)]) - result = db._execute_remote(["SELECT * FROM test_table WHERE id = ?", (1,)]) - assert result.fetchone() is None - - -class TestRemotePoolSetup: - """Test remote pool setup without opening a real Cloud SQL connection.""" - - def _stub_remote_pool(self, monkeypatch): - fake_connection = object() - connector_calls = [] - engine_calls = [] - - class FakeConnector: - def connect(self, **kwargs): - connector_calls.append(kwargs) - return fake_connection - - def fake_create_engine(url, creator): - engine_calls.append((url, creator)) - assert creator() is fake_connection - return "fake-engine" - - fake_connector = FakeConnector() - monkeypatch.setattr(data_module, "Connector", lambda: fake_connector) - monkeypatch.setattr(data_module.sqlalchemy, "create_engine", fake_create_engine) - return fake_connector, connector_calls, engine_calls - - def test_create_pool_uses_remote_database_config(self, monkeypatch): - fake_connector, connector_calls, engine_calls = self._stub_remote_pool( - monkeypatch - ) - monkeypatch.setenv( - "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", - "test-project:us-central1:test-db", - ) - monkeypatch.setenv("POLICYENGINE_DB_USER", "test-user") - monkeypatch.setenv("POLICYENGINE_DB_NAME", "test-db") - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", "test-password") - - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db._create_pool() - - assert db.connector is fake_connector - assert db.pool == "fake-engine" - assert connector_calls == [ - { - "instance_connection_string": "test-project:us-central1:test-db", - "driver": "pymysql", - "db": "test-db", - "user": "test-user", - "password": "test-password", - } - ] - assert engine_calls[0][0] == "mysql+pymysql://" - - def test_create_pool_reads_dot_dbpw_file(self, monkeypatch, tmp_path): - _, connector_calls, _ = self._stub_remote_pool(monkeypatch) - monkeypatch.chdir(tmp_path) - monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", ".dbpw") - (tmp_path / ".dbpw").write_text("file-password\n") - - db = PolicyEngineDatabase.__new__(PolicyEngineDatabase) - db._create_pool() - - assert connector_calls[0]["password"] == "file-password" - - def test_remote_constructor_initializes_pool_without_local_database( - self, monkeypatch - ): - calls = [] + finally: + engine.dispose() - def fake_create_pool(self): - calls.append(("pool", self.local)) - monkeypatch.setattr(PolicyEngineDatabase, "_create_pool", fake_create_pool) +def test_close_v1_engines_disposes_pools_and_connectors(monkeypatch): + engine = Mock() + connector = Mock() + monkeypatch.setattr(orm, "_v1_engines", {False: engine}) + monkeypatch.setattr(orm, "_cloud_sql_connectors", {False: connector}) + monkeypatch.setattr(orm, "_v1_session_factories", {False: Mock()}) - db = PolicyEngineDatabase(local=False, initialize=False) + orm.close_v1_engines() - assert db.local is False - assert calls == [("pool", False)] + engine.dispose.assert_called_once_with() + connector.close.assert_called_once_with() + assert orm._v1_engines == {} + assert orm._cloud_sql_connectors == {} + assert orm._v1_session_factories == {} diff --git a/tests/unit/data/test_v1_database_migration.py b/tests/unit/data/test_v1_database_migration.py new file mode 100644 index 000000000..b4271cb16 --- /dev/null +++ b/tests/unit/data/test_v1_database_migration.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from contextlib import nullcontext + +import pytest +from sqlalchemy import Column, Integer, MetaData, String, Table, text +from sqlalchemy.engine import URL + +import scripts.v1_database_migration as migration +from scripts.v1_database_migration import ( + DatabaseState, + build_database_url, + classify_database_state, + describe_metadata_difference, + upgrade_database, +) + + +def test_database_url_uses_sqlalchemy_url_without_stringifying_credentials(): + url = build_database_url( + username="schema reader", + password="p@ss:/word", + host="127.0.0.1", + port=3307, + database="policyengine", + ) + + assert isinstance(url, URL) + assert url.username == "schema reader" + assert url.password == "p@ss:/word" + assert url.host == "127.0.0.1" + assert url.port == 3307 + assert url.database == "policyengine" + + +@pytest.mark.parametrize( + ("mode", "password_name", "expected_user"), + [ + ("state", "POLICYENGINE_DB_READONLY_PASSWORD", "policyengine_schema_reader"), + ( + "verify-head", + "POLICYENGINE_DB_READONLY_PASSWORD", + "policyengine_schema_reader", + ), + ( + "upgrade", + "POLICYENGINE_DB_MIGRATION_PASSWORD", + "policyengine_schema_migrator", + ), + ], +) +def test_database_target_builds_in_memory_url_for_each_database_role( + monkeypatch, + mode, + password_name, + expected_user, +): + monkeypatch.delenv("STAGE7_EXISTING_DATABASE_URL", raising=False) + monkeypatch.delenv("ALEMBIC_DATABASE_URL", raising=False) + monkeypatch.setenv(password_name, "p@ssword") + + target = migration._database_target(mode) + + assert isinstance(target, URL) + assert target.username == expected_user + assert target.password == "p@ssword" + assert target.host == "127.0.0.1" + assert target.port == 3307 + assert target.database == "policyengine" + + +def test_database_target_preserves_explicit_url_override(monkeypatch): + explicit_url = "mysql+pymysql://root:test@127.0.0.1/test" + monkeypatch.setenv("ALEMBIC_DATABASE_URL", explicit_url) + monkeypatch.delenv("POLICYENGINE_DB_MIGRATION_PASSWORD", raising=False) + + assert migration._database_target("upgrade") == explicit_url + + +def test_database_state_is_unversioned_without_a_version_table(): + assert ( + classify_database_state(version_table_exists=False, current_heads=set()) + is DatabaseState.UNVERSIONED + ) + + +def test_database_state_is_invalid_when_version_table_has_no_revision(): + assert ( + classify_database_state(version_table_exists=True, current_heads=set()) + is DatabaseState.INVALID + ) + + +def test_database_state_is_head_only_when_all_script_heads_are_applied(): + assert ( + classify_database_state( + version_table_exists=True, + current_heads={"head-a"}, + script_heads={"head-a"}, + ) + is DatabaseState.HEAD + ) + assert ( + classify_database_state( + version_table_exists=True, + current_heads={"previous-revision"}, + script_heads={"head-a"}, + ) + is DatabaseState.PENDING + ) + + +def test_metadata_difference_descriptions_are_stable_and_do_not_include_data(): + metadata = MetaData() + question = Table( + "question", + metadata, + Column("question_id", Integer, primary_key=True), + Column("question", String(255)), + ) + + assert describe_metadata_difference(("remove_table", question)) == ( + "extra_table:question" + ) + assert ( + describe_metadata_difference( + [ + ( + "modify_nullable", + None, + "reform_impact", + "execution_id", + {}, + True, + False, + ) + ] + ) + == "nullable:reform_impact.execution_id:true->false" + ) + assert ( + describe_metadata_difference( + [ + ( + "modify_default", + None, + "reform_impact", + "dataset", + {}, + text("'default'"), + None, + ) + ] + ) + == "default:reform_impact.dataset:present->none" + ) + + +def test_metadata_difference_rejects_unknown_shapes_without_repr_leakage(): + secret = "do-not-print-this-value" + difference = ("unknown", secret) + + with pytest.raises(ValueError, match="Unsupported metadata difference") as error: + describe_metadata_difference(difference) + + assert secret not in str(error.value) + + +def test_upgrade_cli_commits_the_externally_supplied_connection(monkeypatch): + connection = object() + + class FakeEngine: + def begin(self): + return nullcontext(connection) + + def connect(self): + raise AssertionError("upgrade must use a committing transaction") + + def dispose(self): + pass + + calls = [] + monkeypatch.setattr( + migration, "create_engine", lambda *args, **kwargs: FakeEngine() + ) + monkeypatch.setattr( + migration, + "upgrade_database", + lambda supplied_connection, **kwargs: calls.append( + (supplied_connection, kwargs) + ), + ) + monkeypatch.setenv("ALEMBIC_DATABASE_URL", "mysql+pymysql://unused") + + assert ( + migration.main( + [ + "--mode", + "upgrade", + "--backup-id", + "verified-backup", + ] + ) + == 0 + ) + assert calls == [ + ( + connection, + { + "backup_id": "verified-backup", + }, + ) + ] + + +@pytest.mark.parametrize("state", [DatabaseState.UNVERSIONED, DatabaseState.INVALID]) +def test_upgrade_refuses_databases_without_valid_revision_history(monkeypatch, state): + class FakeConnection: + def scalar(self, *args, **kwargs): + return 1 + + def execute(self, *args, **kwargs): + pass + + monkeypatch.setattr(migration, "database_state", lambda connection: state) + + with pytest.raises(RuntimeError, match="no valid Alembic revision"): + upgrade_database(FakeConnection(), backup_id="verified-backup") diff --git a/tests/unit/data/test_v1_models.py b/tests/unit/data/test_v1_models.py new file mode 100644 index 000000000..24155d04d --- /dev/null +++ b/tests/unit/data/test_v1_models.py @@ -0,0 +1,59 @@ +from sqlalchemy.dialects import mysql +from sqlalchemy.schema import CreateTable + +from policyengine_api.data.local_models import LocalV1Base +from policyengine_api.data.v1_models import V1Base + + +EXPECTED_TABLES = { + "analysis", + "computed_household", + "economy", + "household", + "legacy_report_output_aliases", + "policy", + "reform_impact", + "report_output_runs", + "report_outputs", + "simulation_runs", + "simulations", + "user_policies", + "user_profiles", +} + + +def test_v1_metadata_contains_every_legacy_table(): + assert set(V1Base.metadata.tables) == EXPECTED_TABLES + + +def test_tracer_metadata_is_local_only(): + assert "tracers" not in V1Base.metadata.tables + assert set(LocalV1Base.metadata.tables) == {"tracers"} + + +def test_v1_metadata_compiles_for_the_production_mysql_dialect(): + for table in V1Base.metadata.sorted_tables: + assert str(CreateTable(table).compile(dialect=mysql.dialect())) + + +def test_v1_composite_and_unique_keys_match_legacy_contract(): + policy = V1Base.metadata.tables["policy"] + assert policy.c.id.autoincrement is True + assert [column.name for column in policy.primary_key.columns] == [ + "id", + "country_id", + "policy_hash", + ] + computed = V1Base.metadata.tables["computed_household"] + assert [column.name for column in computed.primary_key.columns] == [ + "household_id", + "policy_id", + "country_id", + ] + assert V1Base.metadata.tables["user_profiles"].c.auth0_id.unique + assert ( + V1Base.metadata.tables[ + "legacy_report_output_aliases" + ].c.legacy_report_output_id.autoincrement + is False + ) diff --git a/tests/unit/endpoints/economy/__init__.py b/tests/unit/endpoints/economy/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/tests/unit/endpoints/economy/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/unit/endpoints/test_get_simulations.py b/tests/unit/endpoints/test_get_simulations.py deleted file mode 100644 index 2fa061bff..000000000 --- a/tests/unit/endpoints/test_get_simulations.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Regression tests for issue #3451. - -get_simulations built its LIMIT via an f-string -(`f"DESC LIMIT {max_results}"`), which is a SQL injection vector -(max_results flows in from a caller) and had no cap, so a tall -integer could drop unbounded rows on a production MySQL. The fix: -always LIMIT, clamp to [1, 1000], and bind as a parameter. -""" - -from policyengine_api.endpoints.simulation import get_simulations - - -def _seed_reform_impacts(test_db, n: int) -> None: - for i in range(n): - test_db.query( - """INSERT INTO reform_impact - (baseline_policy_id, reform_policy_id, country_id, region, dataset, - time_period, options_json, options_hash, api_version, - reform_impact_json, status, start_time, execution_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - i + 1, - i + 2, - "us", - "us", - "custom_dataset", - "2025", - "{}", - f"hash-{i}", - "1.0.0", - "{}", - "complete", - f"2026-01-01 00:00:{i:02d}", - f"exec-{i}", - ), - ) - - -def test_get_simulations_default_limit_caps_at_100(test_db): - _seed_reform_impacts(test_db, 150) - result = get_simulations() - assert len(result["result"]) == 100 - - -def test_get_simulations_clamps_huge_max_results(test_db): - _seed_reform_impacts(test_db, 50) - # A caller passing an absurdly large value must not crash and - # must not cause a full scan; the value is clamped at 1000. - result = get_simulations(max_results=10**9) - assert len(result["result"]) == 50 # only 50 seeded - - -def test_get_simulations_clamps_negative_max_results(test_db): - _seed_reform_impacts(test_db, 5) - # max_results of 0 or negative must still return something sane. - result = get_simulations(max_results=0) - assert 1 <= len(result["result"]) <= 5 - - -def test_get_simulations_defaults_when_none(test_db): - _seed_reform_impacts(test_db, 10) - result = get_simulations(max_results=None) - assert len(result["result"]) == 10 # fewer than the default 100 - - -def test_get_simulations_rejects_non_integer_gracefully(test_db): - _seed_reform_impacts(test_db, 5) - # A string like "100; DROP TABLE reform_impact" must not reach - # the SQL statement; it falls back to the default. - result = get_simulations(max_results="100; DROP TABLE reform_impact") - assert len(result["result"]) == 5 - - # And the table must still exist. - rows = test_db.query("SELECT COUNT(*) AS c FROM reform_impact").fetchone() - assert rows["c"] == 5 diff --git a/tests/unit/endpoints/test_set_user_policy_dataset.py b/tests/unit/endpoints/test_set_user_policy_dataset.py deleted file mode 100644 index dcc6082e2..000000000 --- a/tests/unit/endpoints/test_set_user_policy_dataset.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Regression test for issue #3310. - -When dataset is truthy in set_user_policy (policy.py:224-237), the query -builds "AND dataset = ?" (7 placeholders) but the params must include -the dataset value. Previously, the dataset was never appended, causing -a parameter binding crash. -""" - -import time - - -def test_set_user_policy_dataset_param_included(test_db): - """Verify that the SELECT after INSERT in set_user_policy correctly - includes the dataset value in params when dataset is truthy. - - Reproduces the exact code path from policy.py:224-237. - """ - now = int(time.time()) - - country_id = "us" - reform_id = 2 - baseline_id = 1 - user_id = "user1" - year = "2025" - geography = "us" - dataset = "custom_dataset" - - # Insert a user_policy with a non-null dataset - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - country_id, - None, - reform_id, - None, - baseline_id, - user_id, - year, - geography, - dataset, - 3, - "1.0.0", - now, - now, - ), - ) - - # Reproduce the exact code from policy.py:224-237 - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) - - # This must not crash — 7 placeholders, 7 params - row = test_db.query(query, tuple(params)).fetchone() - - assert row is not None - assert row["dataset"] == "custom_dataset" - assert row["reform_id"] == reform_id - assert row["user_id"] == user_id - - -def test_set_user_policy_dataset_null_still_works(test_db): - """When dataset is None/falsy, the query uses IS NULL with 6 params.""" - now = int(time.time()) - - country_id = "us" - reform_id = 2 - baseline_id = 1 - user_id = "user1" - year = "2025" - geography = "us" - dataset = None - - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - country_id, - None, - reform_id, - None, - baseline_id, - user_id, - year, - geography, - dataset, - 3, - "1.0.0", - now, - now, - ), - ) - - dataset_select_str = "IS NULL" if not dataset else "= ?" - query = ( - "SELECT * FROM user_policies WHERE country_id = ? AND reform_id = ? " - "AND baseline_id = ? AND user_id = ? AND year = ? AND geography = ? " - f"AND dataset {dataset_select_str}" - ) - - params = [country_id, reform_id, baseline_id, user_id, year, geography] - if dataset: - params.append(dataset) - - row = test_db.query(query, tuple(params)).fetchone() - - assert row is not None - assert row["dataset"] is None diff --git a/tests/unit/endpoints/test_update_user_policy.py b/tests/unit/endpoints/test_update_user_policy.py deleted file mode 100644 index 556f93a23..000000000 --- a/tests/unit/endpoints/test_update_user_policy.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Regression tests for issue #3445. - -update_user_policy (policy.py) previously interpolated untrusted -payload keys directly into an UPDATE statement, allowing arbitrary -SQL fragments (and identity-column tampering) via the JSON body. - -The fix rejects unknown keys with a 400 response and restricts -writable columns to a static whitelist. -""" - -import time - -from flask import Flask - -from policyengine_api.endpoints import update_user_policy - - -def _create_test_client() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.route("//user-policy", methods=["PUT"])(update_user_policy) - return app.test_client() - - -def _insert_user_policy(test_db) -> int: - now = int(time.time()) - test_db.query( - "INSERT INTO user_policies (country_id, reform_label, reform_id, " - "baseline_label, baseline_id, user_id, year, geography, dataset, " - "number_of_provisions, api_version, added_date, updated_date) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - "us", - "old label", - 2, - None, - 1, - "user1", - "2025", - "us", - "custom_dataset", - 3, - "1.0.0", - now, - now, - ), - ) - row = test_db.query( - "SELECT id FROM user_policies ORDER BY id DESC LIMIT 1" - ).fetchone() - return row["id"] - - -def test_update_user_policy_rejects_sql_injection_key(test_db): - """Unknown keys (including SQL injection attempts) must be rejected.""" - policy_id = _insert_user_policy(test_db) - - client = _create_test_client() - response = client.put( - "/us/user-policy", - json={ - "id": policy_id, - "username; DROP TABLE x --": "x", - }, - ) - - assert response.status_code == 400 - body = response.get_json() - assert "unsupported fields" in body["message"] - - # The row must be untouched. - row = test_db.query( - "SELECT reform_label FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["reform_label"] == "old label" - - -def test_update_user_policy_rejects_identity_column(test_db): - """Identity columns (user_id, country_id, ...) must not be writable.""" - policy_id = _insert_user_policy(test_db) - - client = _create_test_client() - response = client.put( - "/us/user-policy", - json={"id": policy_id, "user_id": "attacker"}, - ) - - assert response.status_code == 400 - row = test_db.query( - "SELECT user_id FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["user_id"] == "user1" - - -def test_update_user_policy_allows_whitelisted_field(test_db): - """Whitelisted fields (e.g. reform_label) can still be updated.""" - policy_id = _insert_user_policy(test_db) - - client = _create_test_client() - response = client.put( - "/us/user-policy", - json={"id": policy_id, "reform_label": "new label"}, - ) - - assert response.status_code == 200 - row = test_db.query( - "SELECT reform_label FROM user_policies WHERE id = ?", - (policy_id,), - ).fetchone() - assert row["reform_label"] == "new label" - - -def test_update_user_policy_requires_id(test_db): - client = _create_test_client() - response = client.put("/us/user-policy", json={"reform_label": "x"}) - assert response.status_code == 400 - - -def test_update_user_policy_requires_at_least_one_field(test_db): - policy_id = _insert_user_policy(test_db) - client = _create_test_client() - response = client.put("/us/user-policy", json={"id": policy_id}) - assert response.status_code == 400 diff --git a/tests/unit/endpoints/test_calculate_deprecated_inputs.py b/tests/unit/routes/test_calculate_deprecated_inputs.py similarity index 90% rename from tests/unit/endpoints/test_calculate_deprecated_inputs.py rename to tests/unit/routes/test_calculate_deprecated_inputs.py index 0d7799abd..aeb0af3e8 100644 --- a/tests/unit/endpoints/test_calculate_deprecated_inputs.py +++ b/tests/unit/routes/test_calculate_deprecated_inputs.py @@ -1,7 +1,11 @@ from flask import Flask import pytest -from policyengine_api.endpoints import household as household_endpoint +from policyengine_api.routes import household_routes +from policyengine_api.extensions import cache +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) class DummyCountry: @@ -43,18 +47,15 @@ def calculate(self, household, policy): def calculate_client(monkeypatch): country = DummyCountry() monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) app = Flask(__name__) - app.add_url_rule( - "//calculate", - "calculate", - household_endpoint.get_calculate, - methods=["POST"], - ) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) + app.register_blueprint(household_routes.household_bp) return app.test_client(), country @@ -261,14 +262,13 @@ def test__calculate__preserves_relationship_fields(calculate_client): def test__calculate_full__drops_deprecated_input_after_add_missing(monkeypatch): country = DummyCountry() monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) monkeypatch.setattr( - household_endpoint, - "add_yearly_variables", - lambda household, country_id: { + "policyengine_api.services.household_calculation_service.add_yearly_variables", + lambda household, country_id, countries=None: { **household, "people": { **household["people"], @@ -281,16 +281,9 @@ def test__calculate_full__drops_deprecated_input_after_add_missing(monkeypatch): ) app = Flask(__name__) - - def calculate_full(country_id): - return household_endpoint.get_calculate(country_id, add_missing=True) - - app.add_url_rule( - "//calculate-full", - "calculate_full", - calculate_full, - methods=["POST"], - ) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) + app.register_blueprint(household_routes.household_bp) client = app.test_client() household = { "people": { diff --git a/tests/unit/endpoints/test_calculate_error_statuses.py b/tests/unit/routes/test_calculate_error_statuses.py similarity index 74% rename from tests/unit/endpoints/test_calculate_error_statuses.py rename to tests/unit/routes/test_calculate_error_statuses.py index 286bab6a3..6f8f6cd19 100644 --- a/tests/unit/endpoints/test_calculate_error_statuses.py +++ b/tests/unit/routes/test_calculate_error_statuses.py @@ -1,8 +1,11 @@ from flask import Flask -import pytest from policyengine_core.errors import SituationParsingError -from policyengine_api.endpoints import household as household_endpoint +from policyengine_api.extensions import cache +from policyengine_api.routes import household_routes +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) from .test_calculate_deprecated_inputs import DummyCountry @@ -31,34 +34,19 @@ def calculate(self, household, policy): def make_client(monkeypatch, country, add_missing=False): monkeypatch.setattr( - household_endpoint, - "get_countries", - lambda: {"us": country}, + household_routes, + "household_calculation_service", + HouseholdCalculationService(country_provider=lambda: {"us": country}), ) app = Flask(__name__) + app.config.update(TESTING=True, CACHE_TYPE="NullCache") + cache.init_app(app) if add_missing: monkeypatch.setattr( - household_endpoint, - "add_yearly_variables", - lambda household, country_id: household, - ) - - def handler(country_id): - return household_endpoint.get_calculate(country_id, add_missing=True) - - app.add_url_rule( - "//calculate-full", - "calculate_full", - handler, - methods=["POST"], - ) - else: - app.add_url_rule( - "//calculate", - "calculate", - household_endpoint.get_calculate, - methods=["POST"], + "policyengine_api.services.household_calculation_service.add_yearly_variables", + lambda household, country_id, countries=None: household, ) + app.register_blueprint(household_routes.household_bp) return app.test_client() diff --git a/tests/unit/routes/test_flask_route_boundaries.py b/tests/unit/routes/test_flask_route_boundaries.py new file mode 100644 index 000000000..ae2f01b60 --- /dev/null +++ b/tests/unit/routes/test_flask_route_boundaries.py @@ -0,0 +1,91 @@ +from pathlib import Path + +from flask import Flask + +from policyengine_api.routes.home_routes import home_bp +from policyengine_api.routes.household_routes import household_bp +from policyengine_api.routes.policy_routes import policy_bp +from policyengine_api.routes.reform_impact_routes import reform_impact_bp +from policyengine_api.routes.system_routes import system_bp + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +def test_legacy_flask_urls_are_owned_by_blueprints(): + app = Flask(__name__) + app.register_blueprint(home_bp) + app.register_blueprint(household_bp) + app.register_blueprint(policy_bp) + app.register_blueprint(reform_impact_bp) + app.register_blueprint(system_bp) + + rules = { + (rule.rule, method) + for rule in app.url_map.iter_rules() + for method in rule.methods + if method not in {"HEAD", "OPTIONS"} + } + + assert { + ("/", "GET"), + ("//policies", "GET"), + ("//household//policy/", "GET"), + ("//calculate", "POST"), + ("//calculate-full", "POST"), + ("//user-policy", "POST"), + ("//user-policy", "PUT"), + ("//user-policy/", "GET"), + ("/simulations", "GET"), + ("/liveness-check", "GET"), + ("/readiness-check", "GET"), + ("/specification", "GET"), + } <= rules + + +def test_flask_app_assembles_blueprints_without_legacy_endpoint_wiring(): + source = (PACKAGE_ROOT / "api.py").read_text(encoding="utf-8") + + assert "policyengine_api.endpoints" not in source + assert "from .endpoints" not in source + assert "Legacy endpoints" not in source + assert "get_policy_search" not in source + assert "get_household_under_policy" not in source + assert "get_calculate" not in source + assert "set_user_policy" not in source + assert "get_user_policy" not in source + assert "update_user_policy" not in source + assert "get_simulations" not in source + assert "@app.route" not in source + assert "app.route(" not in source + + +def test_legacy_endpoints_package_is_removed(): + assert not any((PACKAGE_ROOT / "endpoints").rglob("*.py")) + + +def test_economy_comparison_logic_is_not_in_the_http_layer(): + comparison_module = PACKAGE_ROOT / "services/economy_comparison.py" + + assert comparison_module.exists() + source = comparison_module.read_text(encoding="utf-8") + assert "from flask" not in source + assert "import flask" not in source + + +def test_calculation_route_delegates_domain_processing_to_its_service(): + source = (PACKAGE_ROOT / "routes/household_routes.py").read_text(encoding="utf-8") + + assert "country.calculate(" not in source + assert "find_unrecognized_inputs(" not in source + assert "drop_deprecated_inputs(" not in source + + +def test_routes_do_not_construct_json_error_responses_inline(): + offenders = [] + for path in sorted((PACKAGE_ROOT / "routes").glob("*.py")): + source = path.read_text(encoding="utf-8") + if '"status": "error"' in source or 'status="error"' in source: + offenders.append(path.name) + + assert offenders == [] diff --git a/tests/unit/routes/test_household_and_user_policy_orm_routes.py b/tests/unit/routes/test_household_and_user_policy_orm_routes.py new file mode 100644 index 000000000..5850aac74 --- /dev/null +++ b/tests/unit/routes/test_household_and_user_policy_orm_routes.py @@ -0,0 +1,133 @@ +"""ORM integration behavior for household and saved-policy routes.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from flask import Flask +from sqlalchemy import select + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, + UserPolicy, +) +from policyengine_api.routes.household_routes import get_household_under_policy +from policyengine_api.routes.policy_routes import ( + get_user_policy, + set_user_policy, + update_user_policy, +) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) + + +def test_household_under_policy_returns_cached_json_object(orm_session_factory): + stored_result = {"people": {"you": {"net_income": {"2026": 42}}}} + with orm_session_factory.begin() as session: + session.add( + ComputedHousehold( + household_id=1, + policy_id=2, + country_id="us", + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + computed_household_json=stored_result, + status="complete", + ) + ) + + response = get_household_under_policy("us", "1", "2") + + assert response["result"] == {"people": {"you": {"net_income": {"2026": 42}}}} + + +def test_household_under_policy_calculates_and_caches_json_as_an_object( + orm_session_factory, +): + with orm_session_factory.begin() as session: + session.add_all( + [ + Household( + id=1, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + household_json={"people": {"you": {}}}, + household_hash="household-hash", + ), + Policy( + id=2, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + policy_json={"gov.example.parameter": 1}, + policy_hash="policy-hash", + ), + ] + ) + calculated = {"people": {"you": {"net_income": {"2026": 42}}}} + country = SimpleNamespace( + calculate=Mock(return_value=calculated), + metadata={ + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + "parameters": {"gov.example.parameter": {}}, + }, + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) + + with patch( + "policyengine_api.routes.household_routes.household_calculation_service", + service, + ): + response = get_household_under_policy("us", "1", "2") + + assert response["result"] == calculated + country.calculate.assert_called_once_with( + {"people": {"you": {}}}, + {"gov.example.parameter": 1}, + ) + with orm_session_factory() as session: + cached = session.scalar(select(ComputedHousehold)) + assert cached.computed_household_json == calculated + + +def test_user_policy_endpoints_round_trip_through_orm_session_factory( + orm_session_factory, +): + app = Flask(__name__) + payload = { + "reform_label": "Reform", + "reform_id": 2, + "baseline_label": "Current law", + "baseline_id": 1, + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": "default", + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 1, + "budgetary_impact": None, + "type": None, + } + + with app.test_request_context(json=payload): + created = set_user_policy("us") + listed = get_user_policy("us", "auth0|one") + with app.test_request_context(json={"id": 1, "reform_label": "Updated"}): + updated = update_user_policy("us") + + assert created.status_code == 201 + assert created.get_json()["result"]["dataset"] == "default" + assert listed["result"][0]["reform_label"] == "Reform" + assert updated.status_code == 200 + with orm_session_factory() as session: + assert session.get(UserPolicy, 1).reform_label == "Updated" diff --git a/tests/unit/routes/test_reform_impact_routes.py b/tests/unit/routes/test_reform_impact_routes.py new file mode 100644 index 000000000..4fe072ee4 --- /dev/null +++ b/tests/unit/routes/test_reform_impact_routes.py @@ -0,0 +1,115 @@ +"""Regression tests for issue #3451. + +get_simulations built its LIMIT via an f-string +(`f"DESC LIMIT {max_results}"`), which is a SQL injection vector +(max_results flows in from a caller) and had no cap, so a tall +integer could drop unbounded rows on a production MySQL. The fix: +always LIMIT, clamp to [1, 1000], and bind as a parameter. +""" + +from datetime import datetime +import json + +from fastapi.testclient import TestClient +from flask import Flask +from sqlalchemy import func, select + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.v1_models import ReformImpact +from policyengine_api.routes.reform_impact_routes import reform_impact_bp + + +def _get_simulations(max_results=100): + app = _create_app() + query = "" if max_results is None else f"?max_results={max_results}" + return app.test_client().get(f"/simulations{query}").get_json() + + +def _create_app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(reform_impact_bp) + return app + + +def _seed_reform_impacts(orm_session, n: int) -> None: + for i in range(n): + orm_session.add( + ReformImpact( + baseline_policy_id=i + 1, + reform_policy_id=i + 2, + country_id="us", + region="us", + dataset="custom_dataset", + time_period="2025", + options_json={}, + options_hash=f"hash-{i}", + api_version="1.0.0", + reform_impact_json={}, + status="complete", + start_time=datetime(2026, 1, 1, 0, i // 60, i % 60), + end_time=datetime(2026, 1, 1, 1, i // 60, i % 60), + execution_id=f"exec-{i}", + ) + ) + orm_session.commit() + + +def test_get_simulations_default_limit_caps_at_100(orm_session): + _seed_reform_impacts(orm_session, 150) + result = _get_simulations() + assert len(result["result"]) == 100 + + +def test_get_simulations_clamps_huge_max_results(orm_session): + _seed_reform_impacts(orm_session, 50) + # A caller passing an absurdly large value must not crash and + # must not cause a full scan; the value is clamped at 1000. + result = _get_simulations(max_results=10**9) + assert len(result["result"]) == 50 # only 50 seeded + + +def test_get_simulations_clamps_negative_max_results(orm_session): + _seed_reform_impacts(orm_session, 5) + # max_results of 0 or negative must still return something sane. + result = _get_simulations(max_results=0) + assert 1 <= len(result["result"]) <= 5 + + +def test_get_simulations_defaults_when_none(orm_session): + _seed_reform_impacts(orm_session, 10) + result = _get_simulations(max_results=None) + assert len(result["result"]) == 10 # fewer than the default 100 + + +def test_get_simulations_rejects_non_integer_gracefully(orm_session): + _seed_reform_impacts(orm_session, 5) + # A string like "100; DROP TABLE reform_impact" must not reach + # the SQL statement; it falls back to the default. + result = _get_simulations(max_results="100; DROP TABLE reform_impact") + assert len(result["result"]) == 5 + + # And the table must still exist. + assert orm_session.scalar(select(func.count()).select_from(ReformImpact)) == 5 + + +def test_get_simulations_preserves_v1_json_and_timestamp_fields(orm_session): + _seed_reform_impacts(orm_session, 1) + + impact = _get_simulations(max_results=1)["result"][0] + + assert json.loads(impact["options_json"]) == {} + assert json.loads(impact["reform_impact_json"]) == {} + assert impact["start_time"] == "2026-01-01 00:00:00" + assert impact["end_time"] == "2026-01-01 01:00:00" + + +def test_get_simulations_matches_through_fastapi_fallback(orm_session): + _seed_reform_impacts(orm_session, 1) + app = _create_app() + + flask_response = app.test_client().get("/simulations?max_results=1") + asgi_response = TestClient(create_asgi_app(app)).get("/simulations?max_results=1") + + assert asgi_response.status_code == flask_response.status_code == 200 + assert asgi_response.json() == flask_response.get_json() diff --git a/tests/unit/routes/test_route_exception_handling.py b/tests/unit/routes/test_route_exception_handling.py index b6fb38a28..b67e97274 100644 --- a/tests/unit/routes/test_route_exception_handling.py +++ b/tests/unit/routes/test_route_exception_handling.py @@ -29,7 +29,7 @@ def _client_with(*blueprints): def test_simulation_create_runtime_error_becomes_500(): client = _client_with(simulation_bp) with patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", side_effect=RuntimeError("db went away"), ): response = client.post( @@ -46,7 +46,7 @@ def test_simulation_create_runtime_error_becomes_500(): def test_simulation_create_value_error_still_400(): client = _client_with(simulation_bp) with patch( - "policyengine_api.routes.simulation_routes.simulation_service.find_existing_simulation", + "policyengine_api.routes.simulation_routes.simulation_service.get_or_create_simulation", side_effect=ValueError("bad input"), ): response = client.post( @@ -63,7 +63,7 @@ def test_simulation_create_value_error_still_400(): def test_report_create_runtime_error_becomes_500(): client = _client_with(report_output_bp) with patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", side_effect=RuntimeError("db went away"), ): response = client.post( @@ -76,7 +76,7 @@ def test_report_create_runtime_error_becomes_500(): def test_report_create_value_error_still_400(): client = _client_with(report_output_bp) with patch( - "policyengine_api.routes.report_output_routes.report_output_service.find_existing_report_output", + "policyengine_api.routes.report_output_routes.report_output_service.create_or_reuse_report_output", side_effect=ValueError("bad input"), ): response = client.post( @@ -86,7 +86,7 @@ def test_report_create_value_error_still_400(): assert response.status_code == 400 -def test_simulation_patch_empty_body_returns_400(test_db): +def test_simulation_patch_empty_body_returns_400(orm_session_factory): """Regression for issue #3449. PATCH /{country}/simulation with a body that only contains the @@ -95,14 +95,14 @@ def test_simulation_patch_empty_body_returns_400(test_db): """ from policyengine_api.services.simulation_service import SimulationService - simulation_service = SimulationService() - created = simulation_service.create_simulation( + simulation_service = SimulationService(orm_session_factory) + created = simulation_service.get_or_create_simulation( country_id="us", population_id="household_patch_empty", population_type="household", policy_id=50, - ) + ).simulation client = _client_with(simulation_bp) - response = client.patch("/us/simulation", json={"id": created["id"]}) + response = client.patch("/us/simulation", json={"id": created.id}) assert response.status_code == 400 diff --git a/tests/unit/routes/test_set_user_policy_dataset.py b/tests/unit/routes/test_set_user_policy_dataset.py new file mode 100644 index 000000000..887dd653d --- /dev/null +++ b/tests/unit/routes/test_set_user_policy_dataset.py @@ -0,0 +1,40 @@ +import time + +import pytest +from flask import Flask +from sqlalchemy import select + +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.routes.policy_routes import policy_bp + + +def create_client(): + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(policy_bp) + return app.test_client() + + +@pytest.mark.parametrize("dataset", ["custom_dataset", None]) +def test_set_user_policy_persists_dataset_with_orm(orm_session, dataset): + now = int(time.time()) + response = create_client().post( + "/us/user-policy", + json={ + "reform_id": 2, + "baseline_id": 1, + "user_id": "user1", + "year": "2025", + "geography": "us", + "dataset": dataset, + "number_of_provisions": 3, + "api_version": "1.0.0", + "added_date": now, + "updated_date": now, + }, + ) + + assert response.status_code == 201 + orm_session.expire_all() + policy = orm_session.scalar(select(UserPolicy).where(UserPolicy.user_id == "user1")) + assert policy.dataset == dataset diff --git a/tests/unit/routes/test_update_user_policy.py b/tests/unit/routes/test_update_user_policy.py new file mode 100644 index 000000000..262f6137e --- /dev/null +++ b/tests/unit/routes/test_update_user_policy.py @@ -0,0 +1,133 @@ +"""Regression tests for issue #3445. + +update_user_policy (policy.py) previously interpolated untrusted +payload keys directly into an UPDATE statement, allowing arbitrary +SQL fragments (and identity-column tampering) via the JSON body. + +The fix rejects unknown keys with a 400 response and restricts +writable columns to a static whitelist. +""" + +import time + +from flask import Flask + +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.routes.policy_routes import policy_bp + + +def _create_test_client() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(policy_bp) + return app.test_client() + + +def _insert_user_policy(orm_session, *, country_id: str = "us") -> int: + now = int(time.time()) + policy = UserPolicy( + country_id=country_id, + reform_label="old label", + reform_id=2, + baseline_label=None, + baseline_id=1, + user_id="user1", + year="2025", + geography="us", + dataset="custom_dataset", + number_of_provisions=3, + api_version="1.0.0", + added_date=now, + updated_date=now, + ) + orm_session.add(policy) + orm_session.commit() + return policy.id + + +def test_update_user_policy_rejects_sql_injection_key(orm_session): + """Unknown keys (including SQL injection attempts) must be rejected.""" + policy_id = _insert_user_policy(orm_session) + + client = _create_test_client() + response = client.put( + "/us/user-policy", + json={ + "id": policy_id, + "username; DROP TABLE x --": "x", + }, + ) + + assert response.status_code == 400 + body = response.get_json() + assert "unsupported fields" in body["message"] + + # The row must be untouched. + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "old label" + + +def test_update_user_policy_rejects_identity_column(orm_session): + """Identity columns (user_id, country_id, ...) must not be writable.""" + policy_id = _insert_user_policy(orm_session) + + client = _create_test_client() + response = client.put( + "/us/user-policy", + json={"id": policy_id, "user_id": "attacker"}, + ) + + assert response.status_code == 400 + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).user_id == "user1" + + +def test_update_user_policy_allows_whitelisted_field(orm_session): + """Whitelisted fields (e.g. reform_label) can still be updated.""" + policy_id = _insert_user_policy(orm_session) + + client = _create_test_client() + response = client.put( + "/us/user-policy", + json={"id": policy_id, "reform_label": "new label"}, + ) + + assert response.status_code == 200 + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "new label" + + +def test_update_user_policy_requires_id(): + client = _create_test_client() + response = client.put("/us/user-policy", json={"reform_label": "x"}) + assert response.status_code == 400 + + +def test_update_user_policy_requires_at_least_one_field(orm_session): + policy_id = _insert_user_policy(orm_session) + client = _create_test_client() + response = client.put("/us/user-policy", json={"id": policy_id}) + assert response.status_code == 400 + + +def test_update_user_policy_returns_not_found_for_missing_id(): + response = _create_test_client().put( + "/us/user-policy", + json={"id": 999_999, "reform_label": "missing"}, + ) + + assert response.status_code == 404 + assert response.get_json()["message"] == "User policy #999999 not found." + + +def test_update_user_policy_does_not_update_another_country(orm_session): + policy_id = _insert_user_policy(orm_session, country_id="uk") + + response = _create_test_client().put( + "/us/user-policy", + json={"id": policy_id, "reform_label": "wrong country"}, + ) + + assert response.status_code == 404 + orm_session.expire_all() + assert orm_session.get(UserPolicy, policy_id).reform_label == "old label" diff --git a/tests/unit/services/test_ai_analysis_service.py b/tests/unit/services/test_ai_analysis_service.py index e853a1f68..13867f6e3 100644 --- a/tests/unit/services/test_ai_analysis_service.py +++ b/tests/unit/services/test_ai_analysis_service.py @@ -1,29 +1,88 @@ import json -from policyengine_api.services.ai_analysis_service import AIAnalysisService -from tests.fixtures.services.ai_analysis_service import ( - mock_stream_text_events, - mock_stream_error_event, - patch_anthropic, - parse_to_chunks, -) +from contextlib import contextmanager +from types import SimpleNamespace + import pytest +from sqlalchemy import select -# Initialize the service -service = AIAnalysisService() +from policyengine_api.data.v1_models import Analysis +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from tests.fixtures.services.ai_analysis_service import parse_to_chunks + +pytest_plugins = ["tests.fixtures.services.ai_analysis_service"] class TestTriggerAIAnalysis: + def test_claude_stream_runs_without_an_open_database_session( + self, + orm_session_factory, + ): + class TrackingSessions: + def __init__(self, delegate): + self.delegate = delegate + self.active = 0 + + @contextmanager + def __call__(self): + self.active += 1 + try: + with self.delegate() as session: + yield session + finally: + self.active -= 1 + + @contextmanager + def begin(self): + self.active += 1 + try: + with self.delegate.begin() as session: + yield session + finally: + self.active -= 1 + + sessions = TrackingSessions(orm_session_factory) + + class ClaudeStream: + def __enter__(self): + assert sessions.active == 0 + return self + + def __exit__(self, *args): + return None + + def __iter__(self): + assert sessions.active == 0 + yield SimpleNamespace(type="text", text="analysis") + + claude_client = SimpleNamespace( + messages=SimpleNamespace(stream=lambda **kwargs: ClaudeStream()) + ) + service = AIAnalysisService( + sessions, + claude_client_factory=lambda: claude_client, + ) + + assert list(service.trigger_ai_analysis("prompt")) == [ + json.dumps({"type": "text", "stream": "analysis"}) + "\n" + ] + assert sessions.active == 0 + + with orm_session_factory() as session: + stored = session.scalar(select(Analysis).where(Analysis.prompt == "prompt")) + assert stored is not None + assert stored.analysis == "analysis" + def test_trigger_ai_analysis_given_successful_streaming( - self, mock_stream_text_events, test_db + self, mock_stream_text_events, orm_session_factory ): # GIVEN a series of successful text messages from the Claude API expected_response = "This is a historical quote." text_chunks = parse_to_chunks(expected_response) - mock_client = mock_stream_text_events(text_chunks=text_chunks) + mock_stream_text_events(text_chunks=text_chunks) # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote" - generator = service.trigger_ai_analysis(prompt) + generator = AIAnalysisService(orm_session_factory).trigger_ai_analysis(prompt) # THEN it should yield the expected chunks results = list(generator) @@ -37,13 +96,14 @@ def test_trigger_ai_analysis_given_successful_streaming( assert chunk == expected_chunk # Verify the database was updated with the complete response - analysis_record = test_db.query( - "SELECT * FROM analysis WHERE prompt = ?", (prompt,) - ).fetchone() + with orm_session_factory() as session: + analysis_record = session.scalar( + select(Analysis).where(Analysis.prompt == prompt) + ) assert analysis_record is not None - assert analysis_record["analysis"] == expected_response - assert analysis_record["status"] == "ok" + assert analysis_record.analysis == expected_response + assert analysis_record.status == "ok" @pytest.mark.parametrize( "error_type", @@ -54,14 +114,14 @@ def test_trigger_ai_analysis_given_successful_streaming( ], ) def test_trigger_ai_analysis_given_error( - self, mock_stream_error_event, test_db, error_type + self, mock_stream_error_event, orm_session_factory, error_type ): # GIVEN an overloaded_error event from the Claude API - mock_client = mock_stream_error_event(error_type) + mock_stream_error_event(error_type) # WHEN we call trigger_ai_analysis prompt = "Tell me a historical quote about erroneous systems" - generator = service.trigger_ai_analysis(prompt) + generator = AIAnalysisService(orm_session_factory).trigger_ai_analysis(prompt) # THEN it should yield the expected error message results = list(generator) @@ -79,8 +139,9 @@ def test_trigger_ai_analysis_given_error( assert results[0] == expected_error # Verify the database was not updated - analysis_record = test_db.query( - "SELECT * FROM analysis WHERE prompt = ?", (prompt,) - ).fetchone() + with orm_session_factory() as session: + analysis_record = session.scalar( + select(Analysis).where(Analysis.prompt == prompt) + ) assert analysis_record is None diff --git a/tests/unit/services/test_create_profile.py b/tests/unit/services/test_create_profile.py deleted file mode 100644 index 6217d7633..000000000 --- a/tests/unit/services/test_create_profile.py +++ /dev/null @@ -1,63 +0,0 @@ -import pytest -import unittest.mock as mock -import time -from policyengine_api.services.user_service import UserService - -userService = UserService() - - -class TestCreateProfile: - def test_create_profile_valid(self): - auth0_id = "test-auth-id" - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - - result = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - - assert result[0] is True - user_record = userService.get_profile(auth0_id) - assert user_record is not None - assert user_record["auth0_id"] == auth0_id - assert user_record["primary_country"] == primary_country - assert user_record["username"] == username - assert user_record["user_since"] == user_since - - def test_create_profile_invalid(self): - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - with pytest.raises( - Exception, - match=r"UserService.create_profile\(\) missing 1 required positional argument: 'auth0_id'", - ): - userService.create_profile( - primary_country=primary_country, - username=username, - user_since=user_since, - ) - - def test_create_profile_duplicate(self): - auth0_id = "test-auth-id" - primary_country = "United States" - username = "test_username" - user_since = int(time.time() * 1000) - result1 = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - - result2 = userService.create_profile( - primary_country=primary_country, - auth0_id=auth0_id, - username=username, - user_since=user_since, - ) - assert result2[0] == False diff --git a/tests/unit/services/test_direct_orm_local_analysis.py b/tests/unit/services/test_direct_orm_local_analysis.py new file mode 100644 index 000000000..c6a4b0ef6 --- /dev/null +++ b/tests/unit/services/test_direct_orm_local_analysis.py @@ -0,0 +1,72 @@ +from datetime import datetime + +from policyengine_api.data.local_models import Tracer +from policyengine_api.data.v1_models import Analysis, ReformImpact +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService + + +def test_ai_analysis_service_returns_the_latest_mapped_analysis( + orm_session, + orm_session_factory, +): + orm_session.add_all( + [ + Analysis(prompt="prompt", analysis="old", status="ok"), + Analysis(prompt="prompt", analysis="new", status="complete"), + ] + ) + orm_session.flush() + + orm_session.commit() + analysis = AIAnalysisService(orm_session_factory).get_existing_analysis("prompt") + + assert isinstance(analysis, Analysis) + assert analysis.analysis == "new" + + +def test_reform_impact_service_writes_mapped_entity(orm_session_factory): + impact = ReformImpactsService(orm_session_factory).set_reform_impact( + country_id="us", + policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options={"dataset": "default"}, + options_hash="hash", + status="computing", + api_version="1", + reform_impact_json={}, + start_time=datetime(2026, 1, 1), + execution_id="job", + ) + + assert isinstance(impact, ReformImpact) + assert impact.options_json == {"dataset": "default"} + + +def test_tracer_service_reads_python_json_from_mapped_entity( + orm_session, + orm_session_factory, +): + orm_session.add( + Tracer( + household_id=1, + policy_id=2, + country_id="us", + api_version="1", + tracer_output=["net_income <2026>", " dependency"], + ) + ) + orm_session.commit() + + tracer = TracerAnalysisService(orm_session_factory).get_tracer( + "us", + "1", + "2", + "1", + ) + + assert tracer == ["net_income <2026>", " dependency"] diff --git a/tests/unit/services/test_direct_orm_policy_household.py b/tests/unit/services/test_direct_orm_policy_household.py new file mode 100644 index 000000000..bcf480e10 --- /dev/null +++ b/tests/unit/services/test_direct_orm_policy_household.py @@ -0,0 +1,62 @@ +from sqlalchemy import select + +from policyengine_api.data.v1_models import Household, Policy +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService + + +def test_policy_service_reads_and_writes_mapped_models( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "policy-hash", + ) + service = PolicyService(orm_session_factory) + + policy_id, message, existed = service.set_policy( + "us", + "Direct ORM policy", + {"gov.example.rate": {"2026": 0.2}}, + ) + policy = service.get_policy("us", policy_id) + + assert isinstance(policy, Policy) + assert policy.policy_json == {"gov.example.rate": {"2026": 0.2}} + assert message == "Policy created" + assert existed is False + + +def test_household_service_reads_updates_and_writes_mapped_models( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + service = HouseholdService(orm_session_factory) + payload = {"people": {"you": {"age": {"2026": 40}}}} + + household = service.create_household( + "us", + payload, + "Direct ORM household", + ) + with orm_session_factory() as session: + stored = session.scalar(select(Household).where(Household.id == household.id)) + + assert household.id == stored.id + assert stored.household_json == payload + + updated = service.update_household( + "us", + stored.id, + {"people": {"you": {"age": {"2026": 41}}}}, + "Updated", + ) + + assert isinstance(updated, Household) + assert updated.label == "Updated" + assert updated.household_json["people"]["you"]["age"]["2026"] == 41 diff --git a/tests/unit/services/test_direct_orm_users.py b/tests/unit/services/test_direct_orm_users.py new file mode 100644 index 000000000..d06282d81 --- /dev/null +++ b/tests/unit/services/test_direct_orm_users.py @@ -0,0 +1,46 @@ +from policyengine_api.data.v1_models import UserProfile +from policyengine_api.services.user_service import UserService + + +def test_user_service_reads_and_writes_mapped_profiles(orm_session_factory): + service = UserService(orm_session_factory) + + created, profile = service.create_profile( + primary_country="us", + auth0_id="auth0|direct", + username="direct-user", + user_since=123, + ) + duplicate_created, duplicate = service.create_profile( + primary_country="us", + auth0_id="auth0|direct", + username="ignored", + user_since=456, + ) + + assert created is True + assert duplicate_created is False + assert isinstance(profile, UserProfile) + assert duplicate.user_id == profile.user_id + assert service.get_profile(auth0_id="auth0|direct").user_id == profile.user_id + + +def test_user_service_updates_the_mapped_profile(orm_session_factory): + service = UserService(orm_session_factory) + _, profile = service.create_profile( + primary_country="us", + auth0_id="auth0|update", + username=None, + user_since=123, + ) + + updated = service.update_profile( + user_id=profile.user_id, + primary_country="uk", + username="updated-user", + user_since=456, + ) + + assert updated.user_id == profile.user_id + assert updated.primary_country == "uk" + assert updated.username == "updated-user" diff --git a/tests/unit/endpoints/economy/test_compare.py b/tests/unit/services/test_economy_comparison.py similarity index 92% rename from tests/unit/endpoints/economy/test_compare.py rename to tests/unit/services/test_economy_comparison.py index 86d2187dd..f0e451d45 100644 --- a/tests/unit/endpoints/economy/test_compare.py +++ b/tests/unit/services/test_economy_comparison.py @@ -4,7 +4,7 @@ import pandas as pd from pydantic import ValidationError -from policyengine_api.endpoints.economy.compare import ( +from policyengine_api.services.economy_comparison import ( UKConstituencyBreakdownByConstituency, UKConstituencyBreakdown, UKLocalAuthorityBreakdownByLA, @@ -121,9 +121,9 @@ def test__given_non_uk_country_canada__returns_none(self): result = uk_local_authority_breakdown({}, {}, "ca") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_country__returns_breakdown( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -228,9 +228,9 @@ def test__outcome_bucket_categorization_logic(self): f"Failed for {percent_change}: expected {expected_bucket}, got {bucket}" ) - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__outcome_buckets_are_correct( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -264,9 +264,9 @@ def test__outcome_buckets_are_correct( assert result.outcomes_by_region["uk"]["Gain more than 5%"] == 1 assert result.outcomes_by_region["uk"]["Gain less than 5%"] == 0 - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__downloads_from_correct_repos( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -313,9 +313,9 @@ def test__given_constituency_region_with_code__returns_none(self): result = uk_local_authority_breakdown({}, {}, "uk", "constituency/E12345678") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_specific_la_region__returns_only_that_la( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -354,9 +354,9 @@ def test__given_specific_la_region__returns_only_that_la( assert "Aberdeen City" not in result.by_local_authority assert "Isle of Anglesey" not in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_country_scotland_region__returns_only_scottish_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -395,9 +395,9 @@ def test__given_country_scotland_region__returns_only_scottish_las( assert "Hartlepool" not in result.by_local_authority assert "Isle of Anglesey" not in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_region__returns_all_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -434,9 +434,9 @@ def test__given_uk_region__returns_all_las( assert "Aberdeen City" in result.by_local_authority assert "Isle of Anglesey" in result.by_local_authority - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_no_region__returns_all_las( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -520,9 +520,9 @@ def test__given_local_authority_region_with_code__returns_none(self): result = uk_constituency_breakdown({}, {}, "uk", "local_authority/E06000016") assert result is None - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_specific_constituency_region__returns_only_that_constituency( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -563,9 +563,9 @@ def test__given_specific_constituency_region__returns_only_that_constituency( assert "Edinburgh East" not in result.by_constituency assert "Cardiff South" not in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_country_scotland_region__returns_only_scottish_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -602,9 +602,9 @@ def test__given_country_scotland_region__returns_only_scottish_constituencies( assert "Aldershot" not in result.by_constituency assert "Cardiff South" not in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_uk_region__returns_all_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -641,9 +641,9 @@ def test__given_uk_region__returns_all_constituencies( assert "Edinburgh East" in result.by_constituency assert "Cardiff South" in result.by_constituency - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__given_no_region__returns_all_constituencies( self, mock_read_csv, mock_h5py_file, mock_download ): @@ -677,9 +677,9 @@ def test__given_no_region__returns_all_constituencies( assert result is not None assert len(result.by_constituency) == 3 - @patch("policyengine_api.endpoints.economy.compare.download_huggingface_dataset") - @patch("policyengine_api.endpoints.economy.compare.h5py.File") - @patch("policyengine_api.endpoints.economy.compare.pd.read_csv") + @patch("policyengine_api.services.economy_comparison.download_huggingface_dataset") + @patch("policyengine_api.services.economy_comparison.h5py.File") + @patch("policyengine_api.services.economy_comparison.pd.read_csv") def test__country_filter_uses_prefix_not_substring( self, mock_read_csv, mock_h5py_file, mock_download ): diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 9600f5124..af9f4b14b 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -13,6 +13,7 @@ ImpactAction, ImpactStatus, ) +from policyengine_api.services.policy_service import PolicyService from tests.fixtures.services.economy_service import ( MOCK_API_VERSION, MOCK_BASELINE_POLICY_ID, @@ -134,6 +135,36 @@ def test__given_completed_impact__returns_completed_result( ) mock_simulation_entrypoint.run.assert_not_called() + def test__given_orm_decoded_completed_impact__returns_completed_result( + self, + economy_service, + base_params, + mock_country_package_versions, + mock_policyengine_version, + mock_policy_service, + mock_reform_impacts_service, + mock_simulation_entrypoint, + mock_logger, + mock_datetime, + mock_numpy_random, + ): + completed_impact = create_mock_reform_impact(status="ok") + completed_impact.reform_impact_json = json.loads( + completed_impact.reform_impact_json + ) + mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ + completed_impact + ] + + result = economy_service.get_economic_impact(**base_params) + + assert result.status == ImpactStatus.OK + assert ( + result.data["poverty_impact"] + == MOCK_REFORM_IMPACT_DATA["poverty_impact"] + ) + mock_simulation_entrypoint.run.assert_not_called() + def test__given_legacy_completed_impact__refreshes_cache( self, economy_service, @@ -275,6 +306,67 @@ def test__given_no_previous_impact__creates_new_simulation( assert result.data is None mock_simulation_entrypoint.run.assert_called_once() mock_reform_impacts_service.set_reform_impact.assert_called_once() + write_values = ( + mock_reform_impacts_service.set_reform_impact.call_args.kwargs + ) + assert write_values["options"] == MOCK_OPTIONS + assert write_values["reform_impact_json"] == {} + + def test__given_policies_created_through_orm__submits_decoded_json( + self, + orm_session_factory, + monkeypatch, + ): + policy_service = PolicyService(orm_session_factory) + baseline_policy_id, _, _ = policy_service.set_policy( + "us", + "ORM baseline", + {}, + ) + reform = {"gov.example.parameter": {"2026": 1}} + reform_policy_id, _, _ = policy_service.set_policy( + "us", + "ORM reform", + reform, + ) + + reform_impacts = MagicMock() + reform_impacts.get_all_reform_impacts_by_options_hash_prefix.return_value = [] + simulation_gateway = MagicMock() + simulation_gateway.resolve_app_name.return_value = ( + "policyengine-simulation-test", + MOCK_MODEL_VERSION, + ) + simulation_gateway.get_execution_id.return_value = "execution-1" + simulation_gateway.run.return_value.run_id = "run-1" + monkeypatch.setattr( + "policyengine_api.services.economy_service.logger", + MagicMock(), + ) + + service = EconomyService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + policy_service_=policy_service, + reform_impacts_service_=reform_impacts, + simulation_entrypoint_=simulation_gateway, + ) + + result = service.get_economic_impact( + country_id="us", + policy_id=reform_policy_id, + baseline_policy_id=baseline_policy_id, + region="us", + dataset="default", + time_period="2026", + options={}, + api_version="test", + ) + + assert result.status is ImpactStatus.COMPUTING + submitted = simulation_gateway.run.call_args.args[0] + assert submitted["baseline"] == {} + assert submitted["reform"] == reform def test__given_no_previous_impact__includes_metadata_in_simulation_params( self, @@ -1413,6 +1505,14 @@ def test__given_succeeded_state__returns_completed_result( "dataset": MOCK_RESOLVED_DATASET, } mock_reform_impacts_service.set_complete_reform_impact.assert_called_once() + write_values = ( + mock_reform_impacts_service.set_complete_reform_impact.call_args.kwargs + ) + assert isinstance(write_values["reform_impact_json"], dict) + assert ( + write_values["reform_impact_json"]["poverty_impact"] + == (MOCK_REFORM_IMPACT_DATA["poverty_impact"]) + ) def test__given_failed_state__returns_error_result( self, diff --git a/tests/unit/services/test_execute_analysis.py b/tests/unit/services/test_execute_analysis.py index f0c2ff622..e4e2e7714 100644 --- a/tests/unit/services/test_execute_analysis.py +++ b/tests/unit/services/test_execute_analysis.py @@ -1,20 +1,9 @@ -import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) -from werkzeug.exceptions import NotFound -from tests.fixtures.services.tracer_analysis_service import ( - sample_tracer_data, - sample_expected_segment, - mock_get_tracer, - mock_get_existing_analysis, - mock_parse_tracer_output, - mock_trigger_ai_analysis, -) +pytest_plugins = ["tests.fixtures.services.tracer_analysis_service"] -service = TracerAnalysisService() country_id = "us" household_id = "71424" policy_id = "2" @@ -24,6 +13,7 @@ class TestExecuteAnalysis: def test_execute_analysis_static( self, + orm_session_factory, mock_get_tracer, mock_parse_tracer_output, mock_get_existing_analysis, @@ -35,15 +25,16 @@ def test_execute_analysis_static( THEN then a static analysis with the "static" flag should be returned. """ - analysis, analysis_type = service.execute_analysis( - country_id, household_id, policy_id, target_variable - ) + analysis, analysis_type = TracerAnalysisService( + orm_session_factory + ).execute_analysis(country_id, household_id, policy_id, target_variable) assert analysis == "Existing static analysis" assert analysis_type == "static" def test_execute_analysis_streaming( self, + orm_session_factory, mock_get_tracer, mock_parse_tracer_output, mock_get_existing_analysis, @@ -59,9 +50,9 @@ def test_execute_analysis_streaming( # When existing analysis value is None mock_get_existing_analysis.return_value = None - analysis, analysis_type = service.execute_analysis( - country_id, household_id, policy_id, target_variable - ) + analysis, analysis_type = TracerAnalysisService( + orm_session_factory + ).execute_analysis(country_id, household_id, policy_id, target_variable) expected_streaming_output = ["stream chunk 1", "stream chunk 2"] streaming_output = list(analysis) diff --git a/tests/unit/services/test_household_calculation_service.py b/tests/unit/services/test_household_calculation_service.py new file mode 100644 index 000000000..e849861b8 --- /dev/null +++ b/tests/unit/services/test_household_calculation_service.py @@ -0,0 +1,189 @@ +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +from sqlalchemy import select + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.local_models import Tracer +from policyengine_api.data.v1_models import ( + ComputedHousehold, + Household, + Policy, +) +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) + + +PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" + + +class TrackingSessionFactory: + def __init__(self, factory): + self.factory = factory + self.active_scopes = 0 + + @contextmanager + def __call__(self): + self.active_scopes += 1 + try: + with self.factory() as session: + yield session + finally: + self.active_scopes -= 1 + + @contextmanager + def begin(self): + self.active_scopes += 1 + try: + with self.factory.begin() as session: + yield session + finally: + self.active_scopes -= 1 + + +def _seed_inputs(factory): + with factory.begin() as session: + session.add_all( + [ + Household( + id=1, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + household_json={"people": {"you": {}}}, + household_hash="household-hash", + ), + Policy( + id=2, + country_id="us", + label=None, + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + policy_json={}, + policy_hash="policy-hash", + ), + ] + ) + + +def test_household_route_and_country_do_not_manage_persistence(): + route_source = (PACKAGE_ROOT / "routes" / "household_routes.py").read_text( + encoding="utf-8" + ) + country_source = (PACKAGE_ROOT / "country.py").read_text(encoding="utf-8") + assert "get_v1_session_factory" not in route_source + assert "from sqlalchemy" not in route_source + assert "select(" not in route_source + assert "get_v1_session_factory" not in country_source + assert "Tracer(" not in country_source + + +def test_calculation_closes_reads_before_compute_and_persists_local_results( + orm_session_factory, +): + _seed_inputs(orm_session_factory) + primary = TrackingSessionFactory(orm_session_factory) + local = TrackingSessionFactory(orm_session_factory) + + class Country: + metadata = { + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + } + + def calculate(self, household, policy): + assert primary.active_scopes == 0 + assert local.active_scopes == 0 + return SimpleNamespace( + household={"people": {"you": {"net_income": {"2026": 42}}}}, + tracer_output=["net_income <2026>"], + ) + + service = HouseholdCalculationService( + primary_session_factory=primary, + local_session_factory=local, + country_provider=lambda: {"us": Country()}, + ) + + result = service.calculate_stored_household("us", 1, 2) + + assert result.household["people"]["you"]["net_income"]["2026"] == 42 + with orm_session_factory() as session: + cached = session.scalar(select(ComputedHousehold)) + tracer = session.scalar(select(Tracer)) + assert cached.computed_household_json == result.household + assert tracer.tracer_output == ["net_income <2026>"] + + +def test_calculation_uses_local_cache_without_recomputing(orm_session_factory): + _seed_inputs(orm_session_factory) + calculated = {"people": {"you": {"net_income": {"2026": 42}}}} + with orm_session_factory.begin() as session: + session.add( + ComputedHousehold( + household_id=1, + policy_id=2, + country_id="us", + api_version=COUNTRY_PACKAGE_VERSIONS["us"], + computed_household_json=calculated, + status="complete", + ) + ) + country = SimpleNamespace( + metadata={"variables": {}, "entities": {}}, + calculate=lambda *_: (_ for _ in ()).throw( + AssertionError("cache hit should not calculate") + ), + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) + + result = service.calculate_stored_household("us", 1, 2) + + assert result.household == calculated + assert result.cached is True + + +def test_local_computed_household_and_tracer_write_roll_back_together( + orm_session_factory, + monkeypatch, +): + _seed_inputs(orm_session_factory) + country = SimpleNamespace( + metadata={ + "variables": {}, + "entities": {"person": {"plural": "people", "roles": {}}}, + }, + calculate=lambda *_: SimpleNamespace( + household={"people": {"you": {}}}, + tracer_output=["trace"], + ), + ) + service = HouseholdCalculationService( + primary_session_factory=orm_session_factory, + local_session_factory=orm_session_factory, + country_provider=lambda: {"us": country}, + ) + session_type = orm_session_factory.class_ + original_flush = session_type.flush + + def fail_tracer_flush(session, *args, **kwargs): + has_tracer = any(isinstance(value, Tracer) for value in session.new) + original_flush(session, *args, **kwargs) + if has_tracer: + raise RuntimeError("tracer insert failed") + + monkeypatch.setattr(session_type, "flush", fail_tracer_flush) + + with pytest.raises(RuntimeError, match="tracer insert failed"): + service.calculate_stored_household("us", 1, 2) + + monkeypatch.setattr(session_type, "flush", original_flush) + with orm_session_factory() as session: + assert session.scalar(select(ComputedHousehold)) is None + assert session.scalar(select(Tracer)) is None diff --git a/tests/unit/services/test_household_service.py b/tests/unit/services/test_household_service.py index 97d25b58e..8f676480d 100644 --- a/tests/unit/services/test_household_service.py +++ b/tests/unit/services/test_household_service.py @@ -1,193 +1,88 @@ import pytest -import json -from unittest.mock import MagicMock -import re +from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdService -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS - from tests.fixtures.services.household_fixtures import ( - valid_request_body, valid_db_row, - valid_hash_value, - existing_household_record, - mock_hash_object, + valid_request_body, ) -service = HouseholdService() - - -class TestGetHousehold: - def test_get_household_given_existing_record( - self, test_db, existing_household_record - ): - # GIVEN an existing record... (included as fixture) - - # WHEN we call get_household for this record... - result = service.get_household(valid_db_row["country_id"], valid_db_row["id"]) - - valid_household_json = valid_request_body["data"] - - # THEN the result should be the expected household data - assert result["household_json"] == valid_household_json - - def test_get_household_given_nonexistent_record(self, test_db): - # GIVEN an empty database (this is created by default)... - - # WHEN we call get_household for a nonexistent record... - NO_SUCH_RECORD_ID = 999 - result = service.get_household("us", NO_SUCH_RECORD_ID) +pytest_plugins = ["tests.fixtures.services.household_fixtures"] - # THEN the result should be None - assert result is None - def test_get_household_given_str_id(self, test_db): - # GIVEN an invalid ID... +@pytest.fixture +def service(orm_session_factory): + return HouseholdService(orm_session_factory) - INVALID_RECORD_ID = "invalid" - with pytest.raises( - Exception, - match=f"Invalid household ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_household with the invalid ID... - # THEN an exception should be raised - service.get_household("us", INVALID_RECORD_ID) +def test_get_household_returns_mapped_entity(service, existing_household_record): + household = service.get_household( + valid_db_row["country_id"], + valid_db_row["id"], + ) - def test_get_household_given_negative_int_id(self, test_db): - # GIVEN an invalid ID... - INVALID_RECORD_ID = -1 + assert isinstance(household, Household) + assert household.household_json == valid_request_body["data"] - with pytest.raises( - Exception, - match=f"Invalid household ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_household with the invalid ID... - # THEN an exception should be raised - service.get_household("us", INVALID_RECORD_ID) +def test_get_household_returns_none_for_missing_entity(service): + assert service.get_household("us", 999) is None -class TestCreateHousehold: - service = HouseholdService() - def test_create_household_given_valid_data(self, test_db): - def fetch_created_record(): - row = test_db.query( - "SELECT * FROM household", - ).fetchone() - return row +@pytest.mark.parametrize("household_id", ["invalid", -1]) +def test_get_household_rejects_invalid_id(service, household_id): + with pytest.raises(Exception, match="Invalid household ID"): + service.get_household("us", household_id) - # GIVEN valid household data and an empty database... - # WHEN we call create_household with this data... - country_id = "us" - valid_json = valid_request_body["data"] - valid_label = valid_request_body["label"] - test_id = service.create_household(country_id, valid_json, valid_label) +def test_create_household_adds_mapped_entity(service, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "some-hash", + ) - # THEN there should only be one record, and if we re-fetch it, - # it should match the data we provided - test_row = fetch_created_record() + household = service.create_household( + "us", + valid_request_body["data"], + valid_request_body["label"], + ) - valid_json_in_db = json.dumps(valid_request_body["data"]) - valid_label_in_db = valid_request_body["label"] + assert isinstance(household, Household) + assert household.id is not None + assert household.household_json == valid_request_body["data"] - assert test_id == test_row["id"] - assert test_row["household_json"] == valid_json_in_db - assert test_row["label"] == valid_label_in_db - - def test_create_household_given_missing_data(self, test_db): - # GIVEN an empty database... - - # WHEN we call create_household with missing required data... - country_id = "us" - valid_label = valid_request_body["label"] - - with pytest.raises( - Exception, - match=re.escape( - "HouseholdService.create_household() missing 1 required positional argument: 'household_json'" - ), - ): - # THEN an exception should be raised - service.create_household(country_id, label=valid_label) - - -class TestUpdateHousehold: - def test_update_household_given_existing_record( - self, test_db, mock_hash_object, existing_household_record - ): - def fetch_updated_record(): - row = test_db.query( - "SELECT * FROM household WHERE id = ?", (valid_db_row["id"],) - ).fetchone() - return row - - # GIVEN an existing record...(included as fixture) - - # WHEN we call update_household for this record's label and fill other necessary info... - test_update_label = "Updated Household" - - existing_country_id = valid_db_row["country_id"] - existing_record_id = valid_db_row["id"] - existing_data = valid_db_row["household_json"] +def test_update_household_mutates_mapped_entity( + service, + existing_household_record, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "updated-hash", + ) + + household = service.update_household( + "us", + valid_db_row["id"], + {"people": {"person1": {"age": 31}}}, + "Updated Household", + ) + + assert household.label == "Updated Household" + assert household.household_hash == "updated-hash" + assert household.household_json == {"people": {"person1": {"age": 31}}} + + +def test_update_household_rejects_missing_or_cross_country_entity( + service, + existing_household_record, +): + with pytest.raises(LookupError): service.update_household( - existing_country_id, - existing_record_id, - existing_data, - test_update_label, + "uk", + valid_db_row["id"], + {}, + "Wrong country", ) - - # THEN the database should be updated with the new data - test_row = fetch_updated_record() - assert test_row["label"] == test_update_label - - def test_update_household_given_nonexistent_record(self, test_db): - # GIVEN an empty database... - - # WHEN we call update_household for a nonexistent record... - NO_SUCH_RECORD_ID = 999 - - existing_country_id = valid_db_row["country_id"] - existing_data = valid_db_row["household_json"] - existing_label = valid_db_row["label"] - - # THEN update_household raises LookupError because the id - # does not exist for this country (issue #3447). - with pytest.raises(LookupError): - service.update_household( - existing_country_id, - NO_SUCH_RECORD_ID, - existing_data, - existing_label, - ) - - def test_update_household_rejects_cross_country_id( - self, test_db, mock_hash_object, existing_household_record - ): - """Regression for issue #3447. - - An existing US household must not be overwritten by a request - that targets the same numeric id under a different country. - """ - - existing_record_id = valid_db_row["id"] - existing_data = valid_db_row["household_json"] - - with pytest.raises(LookupError): - service.update_household( - "uk", # wrong country - existing_record_id, - existing_data, - "Attacker label", - ) - - # The original US row must be untouched. - row = test_db.query( - "SELECT label, country_id FROM household WHERE id = ?", - (existing_record_id,), - ).fetchone() - assert row["country_id"] == "us" - assert row["label"] == valid_db_row["label"] diff --git a/tests/unit/services/test_local_data_service_boundaries.py b/tests/unit/services/test_local_data_service_boundaries.py new file mode 100644 index 000000000..cdf0ab4f7 --- /dev/null +++ b/tests/unit/services/test_local_data_service_boundaries.py @@ -0,0 +1,86 @@ +"""Persistence boundaries for services backed by local data.""" + +import inspect +from pathlib import Path + +import pytest + +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.simulation_analysis_service import ( + SimulationAnalysisService, +) +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + [ + "ai_analysis_service.py", + "reform_impacts_service.py", + "tracer_analysis_service.py", + "report_output_alias_service.py", + ], +) +def test_local_data_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "local_database" not in source + assert "from policyengine_api.data import database" not in source + + +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + (AIAnalysisService, ("get_existing_analysis", "trigger_ai_analysis")), + (SimulationAnalysisService, ("execute_analysis",)), + (TracerAnalysisService, ("execute_analysis", "get_tracer")), + ( + ReformImpactsService, + ( + "get_recent_reform_impacts", + "get_all_reform_impacts", + "get_all_reform_impacts_by_options_hash_prefix", + "set_reform_impact", + "delete_reform_impact", + "set_error_reform_impact", + "set_complete_reform_impact", + ), + ), + ], +) +def test_local_service_public_methods_do_not_accept_persistence( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_analysis_routes_do_not_manage_sessions(): + route_root = SERVICE_ROOT.parent / "routes" + for module_name in ( + "simulation_analysis_routes.py", + "tracer_analysis_routes.py", + ): + source = (route_root / module_name).read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_reform_impact_route_does_not_manage_sessions(): + source = (SERVICE_ROOT.parent / "routes" / "reform_impact_routes.py").read_text( + encoding="utf-8" + ) + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_economy_service_delegates_persistence_without_opening_sessions(): + source = (SERVICE_ROOT / "economy_service.py").read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source diff --git a/tests/unit/services/test_orchestration_service_database_boundaries.py b/tests/unit/services/test_orchestration_service_database_boundaries.py new file mode 100644 index 000000000..c16e42049 --- /dev/null +++ b/tests/unit/services/test_orchestration_service_database_boundaries.py @@ -0,0 +1,17 @@ +"""Database-access boundaries for orchestration services.""" + +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", ["simulation_service.py", "report_output_service.py"] +) +def test_orchestration_services_do_not_access_database_connections(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert "from policyengine_api.data import database" not in source + assert "database.transaction(" not in source diff --git a/tests/unit/services/test_orm_service_boundaries.py b/tests/unit/services/test_orm_service_boundaries.py new file mode 100644 index 000000000..1c48c8494 --- /dev/null +++ b/tests/unit/services/test_orm_service_boundaries.py @@ -0,0 +1,33 @@ +"""ORM and removed-DAO boundaries for entity services.""" + +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + ["household_service.py", "policy_service.py", "user_service.py"], +) +def test_migrated_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "from policyengine_api.data import database" not in source + + +@pytest.mark.parametrize( + "module_name", + ["household_service.py", "policy_service.py", "user_service.py"], +) +def test_migrated_services_use_sessions_and_mapped_models(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert "from sqlalchemy.orm import Session" in source + assert "from policyengine_api.data.v1_daos" not in source + assert "build_v1_session_manager" not in source + + +def test_temporary_dao_module_has_been_removed(): + assert not (SERVICE_ROOT.parent / "data" / "v1_daos.py").exists() diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index 86d509ee8..a2b1c8062 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -1,304 +1,157 @@ import pytest -import json -from unittest.mock import call +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS +from policyengine_api.data.v1_models import Policy from policyengine_api.services.policy_service import PolicyService +from tests.fixtures.services.policy_service import valid_policy_data -from tests.fixtures.services.policy_service import ( - valid_policy_data, - valid_hash_value, - mock_hash_object, - mock_database, - existing_policy_record, -) -service = PolicyService() +pytest_plugins = ["tests.fixtures.services.policy_service"] -class TestGetPolicy: - def test_get_policy_given_existing_record(self, test_db, existing_policy_record): - # GIVEN an existing record... (included as fixture) +@pytest.fixture +def service(orm_session_factory): + return PolicyService(orm_session_factory) - # WHEN we call get_policy for this record... - result = service.get_policy( - valid_policy_data["country_id"], valid_policy_data["id"] - ) - expected_result = { - "id": valid_policy_data["id"], - "country_id": valid_policy_data["country_id"], - "label": valid_policy_data["label"], - "api_version": valid_policy_data["api_version"], - "policy_json": json.loads(valid_policy_data["policy_json"]), - "policy_hash": valid_policy_data["policy_hash"], - } - - # THEN the result should contain the expected policy data - assert result == expected_result - - def test_get_policy_given_nonexistent_record(self, test_db): - # GIVEN an empty database (this is created by default) - - # WHEN we call get_policy for a nonexistent record - NO_SUCH_RECORD_ID = 999 - result = service.get_policy(valid_policy_data["country_id"], NO_SUCH_RECORD_ID) - - # THEN the result should be None - assert result is None - - def test_get_policy_given_str_id(self): - # GIVEN an invalid ID - INVALID_RECORD_ID = "invalid" - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy with the invalid ID - # THEN an exception should be raised - service.get_policy(valid_policy_data["country_id"], INVALID_RECORD_ID) - - def test_get_policy_given_negative_int_id(self): - # GIVEN an invalid ID - INVALID_RECORD_ID = -1 - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy with the invalid ID - # THEN an exception should be raised - service.get_policy(valid_policy_data["country_id"], INVALID_RECORD_ID) - - def test_get_policy_given_invalid_country_id(self): - # GIVEN an invalid country_id - INVALID_COUNTRY_ID = "xx" # Unsupported country code - - # WHEN we call get_policy with the invalid country_id - result = service.get_policy(INVALID_COUNTRY_ID, valid_policy_data["id"]) - - # THEN the result should be None or raise an exception - assert result is None - - def test_get_policy_given_empty_string_country_id(self): - # GIVEN an empty string as country_id - EMPTY_COUNTRY_ID = "" - - # WHEN we call get_policy with country_id = "" - with pytest.raises( - Exception, - match="country_id cannot be empty or None", - ): - # THEN an exception should be raised - service.get_policy(EMPTY_COUNTRY_ID, valid_policy_data["id"]) - - def test_get_policy_given_none_country_id(self): - # GIVEN a country_id of None - NONE_COUNTRY_ID = None - - # WHEN we call get_policy with country_id = None - with pytest.raises( - Exception, - match="country_id cannot be empty or None", - ): - # THEN an exception should be raised - service.get_policy(NONE_COUNTRY_ID, valid_policy_data["id"]) - - -class TestGetPolicyJson: - def test_get_policy_json_given_existing_record( - self, test_db, existing_policy_record - ): - # GIVEN an existing record... (included as fixture) - - # WHEN we call get_policy_json for this record... - result = service.get_policy_json( - valid_policy_data["country_id"], valid_policy_data["id"] - ) +def test_get_policy_returns_mapped_entity(service, existing_policy_record): + policy = service.get_policy( + valid_policy_data["country_id"], + valid_policy_data["id"], + ) + + assert isinstance(policy, Policy) + assert policy.id == valid_policy_data["id"] + assert policy.policy_json == { + "gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433} + } + + +def test_get_policy_returns_none_for_missing_entity(service): + assert service.get_policy("us", 999) is None + + +@pytest.mark.parametrize("policy_id", ["invalid", -1]) +def test_get_policy_rejects_invalid_id(service, policy_id): + with pytest.raises(Exception, match="Invalid policy ID"): + service.get_policy("us", policy_id) + + +@pytest.mark.parametrize("country_id", ["", None]) +def test_get_policy_rejects_empty_country(service, country_id): + with pytest.raises(ValueError, match="country_id cannot be empty or None"): + service.get_policy(country_id, 1) + + +def test_get_policy_json_returns_python_object(service, existing_policy_record): + result = service.get_policy_json("us", valid_policy_data["id"]) - valid_policy_json = valid_policy_data["policy_json"] - - # THEN result should be the expected policy json - assert result == valid_policy_json - - def test_get_policy_json_given_nonexisting_record(self, test_db): - # GIVEN an empty database... (created by default) - - # WHEN we call get_policy_json for nonexistent record... - NO_SUCH_RECORD_ID = 999 - result = service.get_policy_json("us", NO_SUCH_RECORD_ID) - - # THEN result should be None - assert result is None - - def test_get_policy_json_given_str_id(self, test_db): - # GIVEN an invalid ID... - - INVALID_RECORD_ID = "invalid" - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy_json with the invalid ID... - # THEN an exception should be raised - service.get_policy_json("us", INVALID_RECORD_ID) - - def test_get_policy_json_given_negative_int_id(self, test_db): - # GIVEN an invalid ID... - - INVALID_RECORD_ID = -1 - - with pytest.raises( - Exception, - match=f"Invalid policy ID: {INVALID_RECORD_ID}. Must be a positive integer.", - ): - # WHEN we call get_policy_json with the invalid ID... - # THEN an exception should be raised - service.get_policy_json("us", INVALID_RECORD_ID) - - -class TestSetPolicy: - def test_set_policy_new(self, mock_database, mock_hash_object): - # GIVEN a new policy to insert - new_policy_id = 12 # Different from existing fixture ID - test_policy = {"param": "value"} - test_label = "new_policy" - test_country_id = "us" - - # Get current API version dynamically - from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS - - current_api_version = COUNTRY_PACKAGE_VERSIONS.get(test_country_id) - - # Setup mocks - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call: policy does not exist - {"id": new_policy_id}, # Second call: fetch newly inserted policy - ] - - # Define expected database calls - expected_calls = [ - # First call - check if policy exists - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, valid_hash_value, test_label), - ), - # Second call - insert new policy - call( - "INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)", - ( - test_country_id, - json.dumps(test_policy), - valid_hash_value, - test_label, - current_api_version, # From COUNTRY_PACKAGE_VERSIONS for 'us' + assert result == { + "gov.irs.income.bracket.rates.2": {"2024-01-01.2024-12-31": 0.2433} + } + + +def test_search_policies_filters_and_deduplicates(service, orm_session_factory): + with orm_session_factory.begin() as session: + session.add_all( + [ + Policy( + id=31, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={}, + policy_hash="same-hash", + ), + Policy( + id=32, + country_id="us", + label="Tax reform", + api_version="1", + policy_json={"different": True}, + policy_hash="same-hash", ), - ), - # Third call - fetch the newly created policy - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label = ?", - (test_country_id, valid_hash_value, test_label), - ), - ] - - # WHEN we call set_policy - policy_id, message, exists = service.set_policy( - test_country_id, test_label, test_policy + Policy( + id=33, + country_id="us", + label="Benefit reform", + api_version="1", + policy_json={}, + policy_hash="other-hash", + ), + ] ) - # THEN the result should indicate a new policy was created - assert policy_id == new_policy_id - assert message == "Policy created" - assert exists is False - - # Verify the database queries were called as expected - assert mock_database.query.call_args_list == expected_calls - - def test_set_policy_existing( - self, mock_database, mock_hash_object, existing_policy_record - ): - # GIVEN an existing policy record - existing_policy = existing_policy_record - - # Setup mock - mock_database.query.return_value.fetchone.return_value = existing_policy - - # Define expected database calls - matches actual implementation - expected_calls = [ - call( - "SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label IS NULL", - ( - existing_policy["country_id"], - valid_hash_value, - ), # No None parameter for IS NULL - ), - ] - - # WHEN we call set_policy with existing policy data - policy_id, message, exists = service.set_policy( - existing_policy["country_id"], - existing_policy["label"], - json.loads(existing_policy["policy_json"]), - ) + all_results = service.search_policies("us", "Tax", unique_only=False) + unique_results = service.search_policies("us", "Tax", unique_only=True) + escaped_wildcard_results = service.search_policies("us", "Tax%", unique_only=False) - # THEN the result should indicate the policy already exists - assert policy_id == existing_policy["id"] - assert message == "Policy already exists" - assert exists is True - - # Verify the database query was called as expected - assert mock_database.query.call_args_list == expected_calls - - def test_set_policy_given_database_insert_failure( - self, mock_database, mock_hash_object - ): - # GIVEN a database insertion failure - test_policy = {"param": "value"} - test_label = "test_policy" - test_country_id = "us" - - # Setup mock to raise exception on insert - mock_database.query.return_value.fetchone.side_effect = [ - None, # First call: policy does not exist - Exception("Database insertion failed"), # Second call: insertion fails - ] - - # WHEN we call set_policy - with pytest.raises(Exception, match="Database insertion failed"): - # THEN an exception should be raised - service.set_policy(test_country_id, test_label, test_policy) - - def test_set_policy_given_invalid_country_id(self, mock_hash_object): - # GIVEN an invalid country_id - INVALID_COUNTRY_ID = "xx" # Unsupported country code - test_policy = {"param": "value"} - test_label = "test_policy" - - # WHEN we call set_policy with an invalid country_id - with pytest.raises( - ValueError, match=f"Invalid country_id: {INVALID_COUNTRY_ID}" - ): - # THEN an exception should be raised - service.set_policy(INVALID_COUNTRY_ID, test_label, test_policy) - - def test_set_policy_given_empty_label(self, mock_database, mock_hash_object): - # GIVEN an empty label - EMPTY_LABEL = "" - test_policy = {"param": "value"} - test_country_id = "us" - - # Setup mock - mock_database.query.return_value.fetchone.side_effect = [ - None, # Policy does not exist - {"id": 13}, # Return mock policy after creation - ] - - # WHEN we call set_policy with an empty label - policy_id, message, exists = service.set_policy( - test_country_id, EMPTY_LABEL, test_policy - ) + assert len(all_results) == 2 + assert len(unique_results) == 1 + assert unique_results[0].label == "Tax reform" + assert escaped_wildcard_results == [] - # THEN the result should indicate a new policy was created - assert policy_id == 13 - assert message == "Policy created" - assert exists is False + +def test_set_policy_adds_mapped_entity(service, monkeypatch): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "new-hash", + ) + + policy_id, message, exists = service.set_policy( + "US", + "New policy", + {"parameter": 1}, + ) + + policy = service.get_policy("us", policy_id) + assert policy.policy_json == {"parameter": 1} + assert policy.api_version == COUNTRY_PACKAGE_VERSIONS["us"] + assert message == "Policy created" + assert exists is False + + +def test_set_policy_returns_existing_mapped_entity( + service, + existing_policy_record, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: valid_policy_data["policy_hash"], + ) + + policy_id, message, exists = service.set_policy( + "us", + None, + {}, + ) + + assert policy_id == valid_policy_data["id"] + assert message == "Policy already exists" + assert exists is True + + +def test_set_policy_rejects_invalid_country(service): + with pytest.raises(ValueError, match="Invalid country_id: xx"): + service.set_policy("xx", "Policy", {}) + + +def test_set_policy_propagates_flush_failure( + service, + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.policy_service.hash_object", + lambda value: "new-hash", + ) + monkeypatch.setattr( + orm_session_factory.class_, + "flush", + lambda self: (_ for _ in ()).throw(SQLAlchemyError("insert failed")), + ) + + with pytest.raises(SQLAlchemyError, match="insert failed"): + service.set_policy("us", "Policy", {}) diff --git a/tests/unit/services/test_reform_impacts_service.py b/tests/unit/services/test_reform_impacts_service.py new file mode 100644 index 000000000..15b5029d5 --- /dev/null +++ b/tests/unit/services/test_reform_impacts_service.py @@ -0,0 +1,182 @@ +from datetime import datetime + +import pytest +from sqlalchemy import select + +from policyengine_api.data.v1_models import ReformImpact +from policyengine_api.services.reform_impacts_service import ReformImpactsService + + +@pytest.fixture +def service(orm_session_factory): + return ReformImpactsService(orm_session_factory) + + +def _create_impact( + service, + *, + execution_id: str, + options_hash: str, + day: int, +): + return service.set_reform_impact( + country_id="us", + policy_id=2, + baseline_policy_id=1, + region="us", + dataset="default", + time_period="2026", + options={"hash": options_hash}, + options_hash=options_hash, + status="computing", + api_version="1", + reform_impact_json={}, + start_time=datetime(2026, 1, day), + execution_id=execution_id, + ) + + +def test_get_recent_reform_impacts_orders_and_limits_results(service): + older = _create_impact( + service, + execution_id="older-job", + options_hash="older", + day=1, + ) + newer = _create_impact( + service, + execution_id="newer-job", + options_hash="newer", + day=2, + ) + + assert [ + impact.reform_impact_id for impact in service.get_recent_reform_impacts(1) + ] == [newer.reform_impact_id] + assert older.reform_impact_id != newer.reform_impact_id + + +def test_reform_impact_service_round_trips_models_and_transitions(service): + exact = _create_impact( + service, + execution_id="exact-job", + options_hash="hash-exact", + day=1, + ) + compatible = _create_impact( + service, + execution_id="compatible-job", + options_hash="hash-compatible", + day=2, + ) + + exact_results = service.get_all_reform_impacts( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + "1", + ) + compatible_results = service.get_all_reform_impacts_by_options_hash_prefix( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + "hash-%", + "1", + ) + + assert [impact.reform_impact_id for impact in exact_results] == [ + exact.reform_impact_id + ] + assert [impact.reform_impact_id for impact in compatible_results] == [ + exact.reform_impact_id, + compatible.reform_impact_id, + ] + completed = service.set_complete_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-exact", + {"result": 1}, + "exact-job", + ) + failed = service.set_error_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash-compatible", + "failed", + "compatible-job", + ) + + assert completed.status == "ok" + assert completed.reform_impact_json == {"result": 1} + assert failed.status == "error" + assert failed.message == "failed" + + +def test_reform_impact_service_deletes_only_matching_computing_rows( + service, + orm_session_factory, +): + _create_impact( + service, + execution_id="delete-job", + options_hash="delete-hash", + day=1, + ) + retained = _create_impact( + service, + execution_id="retain-job", + options_hash="retain-hash", + day=2, + ) + + service.delete_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "delete-hash", + ) + + with orm_session_factory() as session: + assert ( + session.scalar( + select(ReformImpact).where(ReformImpact.execution_id == "delete-job") + ) + is None + ) + assert session.get(ReformImpact, retained.reform_impact_id) is not None + + +def test_reform_impact_transitions_return_none_for_missing_execution(service): + assert ( + service.set_error_reform_impact( + "us", + 2, + 1, + "us", + "default", + "2026", + "hash", + "missing", + "missing-job", + ) + is None + ) diff --git a/tests/unit/services/test_report_output_alias_service.py b/tests/unit/services/test_report_output_alias_service.py index e4e28c916..9377fe06b 100644 --- a/tests/unit/services/test_report_output_alias_service.py +++ b/tests/unit/services/test_report_output_alias_service.py @@ -1,279 +1,91 @@ import pytest +from policyengine_api.data.v1_models import ( + LegacyReportOutputAlias, + ReportOutput, +) from policyengine_api.services.report_output_alias_service import ( ReportOutputAliasService, ) -from policyengine_api.services.report_output_service import ReportOutputService -from policyengine_api.services.simulation_service import SimulationService -alias_service = ReportOutputAliasService() -report_output_service = ReportOutputService() -simulation_service = SimulationService() +service = ReportOutputAliasService() -class TestReportOutputAliasService: - def _insert_legacy_report_output( - self, - test_db, - legacy_report_output_id: int, - canonical_report: dict, - api_version: str = "legacy-version", - ) -> None: - test_db.query( - """ - INSERT INTO report_outputs ( - id, country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - legacy_report_output_id, - canonical_report["country_id"], - canonical_report["simulation_1_id"], - canonical_report["simulation_2_id"], - api_version, - canonical_report["status"], - canonical_report["year"], - ), - ) - def test_resolves_to_canonical_report_output_id_when_alias_exists(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 999, canonical_report) +def add_report(orm_session, report_id): + report = ReportOutput( + id=report_id, + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + status="pending", + year="2025", + ) + orm_session.add(report) + orm_session.flush() + return report - alias_service.set_alias( - legacy_report_output_id=999, - canonical_report_output_id=canonical_report["id"], - ) - resolved_id = alias_service.resolve_canonical_report_output_id(999) +def test_sets_and_resolves_mapped_alias(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) - assert resolved_id == canonical_report["id"] + assert service.set_alias(orm_session, 100, 200) is True - def test_returns_requested_id_when_alias_is_not_needed(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) + alias = service.get_alias(orm_session, 100) + assert isinstance(alias, LegacyReportOutputAlias) + assert alias.canonical_report_output_id == 200 + assert service.resolve_canonical_report_output_id(orm_session, 100) == 200 + assert service.resolve_canonical_report_output_id(orm_session, 200) == 200 - resolved_id = alias_service.resolve_canonical_report_output_id( - report_output["id"] - ) - assert resolved_id == report_output["id"] +def test_setting_same_alias_is_idempotent(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) - def test_returns_none_for_unknown_report_output(self, test_db): - assert alias_service.resolve_canonical_report_output_id(123456) is None + service.set_alias(orm_session, 100, 200) - def test_set_alias_is_idempotent_for_same_canonical_report_output(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1001, canonical_report) + assert service.set_alias(orm_session, 100, 200) is True - assert ( - alias_service.set_alias( - legacy_report_output_id=1001, - canonical_report_output_id=canonical_report["id"], - ) - is True - ) - assert ( - alias_service.set_alias( - legacy_report_output_id=1001, - canonical_report_output_id=canonical_report["id"], - ) - is True - ) - def test_rejects_alias_to_missing_canonical_report_output(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3a", - population_type="household", - policy_id=3, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1002, canonical_report) +def test_rejects_conflicting_alias(orm_session): + add_report(orm_session, 100) + add_report(orm_session, 200) + add_report(orm_session, 300) + service.set_alias(orm_session, 100, 200) - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=1002, - canonical_report_output_id=999999, - ) + with pytest.raises(ValueError, match="already points"): + service.set_alias(orm_session, 100, 300) - assert "Canonical report output #999999 not found" in str(exc_info.value) - def test_rejects_conflicting_alias_remap(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - other_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2026", - ) - self._insert_legacy_report_output(test_db, 1003, canonical_report) - alias_service.set_alias( - legacy_report_output_id=1003, - canonical_report_output_id=canonical_report["id"], - ) +def test_rejects_missing_and_self_aliases(orm_session): + add_report(orm_session, 100) - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=1003, - canonical_report_output_id=other_report["id"], - ) + with pytest.raises(ValueError, match="Canonical report output #999 not found"): + service.set_alias(orm_session, 100, 999) + with pytest.raises(ValueError, match="must be different"): + service.set_alias(orm_session, 100, 100) - assert ( - "Legacy report output alias already points to canonical report output " - f"#{canonical_report['id']}" - ) in str(exc_info.value) - def test_rejects_alias_when_legacy_report_output_is_missing(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4a", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) +def test_rejects_reports_with_different_logical_keys(orm_session): + add_report(orm_session, 100) + different = add_report(orm_session, 200) + different.year = "2026" - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=10030, - canonical_report_output_id=canonical_report["id"], - ) + with pytest.raises(ValueError, match="must describe the same report"): + service.set_alias(orm_session, 100, 200) - assert "Legacy report output #10030 not found" in str(exc_info.value) - def test_rejects_alias_when_legacy_and_canonical_reports_do_not_match( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4b", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - mismatched_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2026", - ) - - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=mismatched_report["id"], - canonical_report_output_id=canonical_report["id"], - ) - - assert "must describe the same report" in str(exc_info.value) - - def test_rejects_alias_when_legacy_and_canonical_ids_match(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4c", - population_type="household", - policy_id=4, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", +def test_rejects_alias_pointing_to_missing_canonical_report(orm_session): + add_report(orm_session, 100) + orm_session.add( + LegacyReportOutputAlias( + legacy_report_output_id=100, + canonical_report_output_id=999, ) + ) + orm_session.flush() - with pytest.raises(ValueError) as exc_info: - alias_service.set_alias( - legacy_report_output_id=canonical_report["id"], - canonical_report_output_id=canonical_report["id"], - ) - - assert "must be different" in str(exc_info.value) - - def test_rejects_alias_resolution_when_canonical_report_output_is_missing( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5", - population_type="household", - policy_id=5, - ) - canonical_report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - self._insert_legacy_report_output(test_db, 1004, canonical_report) - alias_service.set_alias( - legacy_report_output_id=1004, - canonical_report_output_id=canonical_report["id"], - ) - test_db.query( - "DELETE FROM report_outputs WHERE id = ?", - (canonical_report["id"],), - ) - - with pytest.raises(ValueError) as exc_info: - alias_service.resolve_canonical_report_output_id(1004) - - assert ( - f"Alias points to missing canonical report output #{canonical_report['id']}" - ) in str(exc_info.value) + with pytest.raises(ValueError, match="missing canonical report output #999"): + service.resolve_canonical_report_output_id(orm_session, 100) diff --git a/tests/unit/services/test_report_output_service.py b/tests/unit/services/test_report_output_service.py index 55ee2ff62..f4e6d67ea 100644 --- a/tests/unit/services/test_report_output_service.py +++ b/tests/unit/services/test_report_output_service.py @@ -1,1897 +1,252 @@ import pytest -import json -from datetime import datetime, timezone +from sqlalchemy import func, select from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun from policyengine_api.services.report_output_service import ReportOutputService from policyengine_api.services.report_run_service import ReportRunService -from policyengine_api.services.run_sync_utils import select_display_report_run from policyengine_api.services.simulation_service import SimulationService -from tests.fixtures.services import report_output_fixtures -pytest_plugins = ("tests.fixtures.services.report_output_fixtures",) +run_service = ReportRunService() -service = ReportOutputService() -report_run_service = ReportRunService() -simulation_service = SimulationService() +@pytest.fixture +def service(orm_session_factory): + return ReportOutputService(orm_session_factory) -class TestReportOutputRunTimestamps: - def test_format_run_timestamp_handles_supported_values(self): - assert ( - service._format_run_timestamp(datetime(2026, 5, 4, 12, 0, 0)) - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp( - datetime(2026, 5, 4, 12, 0, 0, tzinfo=timezone.utc) - ) - == "2026-05-04T12:00:00Z" - ) - assert service._format_run_timestamp("") is None - assert ( - service._format_run_timestamp("2026-05-04T12:00:00") - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04T12:00:00Z") - == "2026-05-04T12:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04T12:00:00+01:00") - == "2026-05-04T11:00:00Z" - ) - assert ( - service._format_run_timestamp("2026-05-04 12:00:00.123456") - == "2026-05-04T12:00:00Z" - ) - - def test_select_display_run_uses_matching_result_before_newest_fallback(self): - report_output = { - "id": 1, - "status": "complete", - "output": '{"ok": true}', - "error_message": None, - "active_run_id": None, - "latest_successful_run_id": None, - } - matching_run = { - "id": "matching", - "status": "complete", - "output": '{"ok": true}', - "error_message": None, - } - newest_non_matching_run = { - "id": "newest", - "status": "pending", - "output": None, - "error_message": None, - } - - selected_run = select_display_report_run( - report_output, [newest_non_matching_run, matching_run] - ) - - assert selected_run["id"] == "matching" - - -class TestFindExistingReportOutput: - """Test finding existing report outputs in the database.""" - - def test_find_existing_report_output_found(self, test_db, existing_report_record): - """Test finding an existing report output.""" - # GIVEN an existing report record (from fixture) - - # WHEN we search for a report with matching simulation IDs - result = service.find_existing_report_output( - country_id=existing_report_record["country_id"], - simulation_1_id=existing_report_record["simulation_1_id"], - simulation_2_id=existing_report_record["simulation_2_id"], - ) - - # THEN the result should contain the existing report - assert result is not None - assert result["id"] == existing_report_record["id"] - assert ( - result["country_id"] - == report_output_fixtures.valid_report_data["country_id"] - ) - assert result["simulation_1_id"] == existing_report_record["simulation_1_id"] - assert result["status"] == existing_report_record["status"] - - def test_find_existing_report_output_not_found(self, test_db): - """Test that None is returned when no report exists.""" - # GIVEN an empty database - - # WHEN we search for a non-existent report - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=999, - simulation_2_id=888, - year="2025", - ) - - # THEN None should be returned - assert result is None - - def test_find_existing_report_output_with_null_simulation2(self, test_db): - """Test finding reports where simulation_2_id is NULL.""" - api_version = get_report_output_cache_version("us") - # GIVEN a report with NULL simulation_2_id - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 100, None, "complete", api_version, "2025"), - ) - - # WHEN we search for it - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=100, - simulation_2_id=None, - year="2025", - ) - - # THEN we should find it - assert result is not None - assert result["simulation_1_id"] == 100 - assert result["simulation_2_id"] is None - assert result["year"] == "2025" - - def test_find_existing_report_output_with_year(self, test_db): - """Test finding reports with different years.""" - api_version = get_report_output_cache_version("us") - # GIVEN reports with different years for the same simulation - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 101, None, "complete", api_version, "2025"), - ) - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 101, None, "complete", api_version, "2024"), - ) - - # WHEN we search for the 2025 report - result_2025 = service.find_existing_report_output( - country_id="us", - simulation_1_id=101, - simulation_2_id=None, - year="2025", - ) - - # THEN we should find the 2025 report - assert result_2025 is not None - assert result_2025["simulation_1_id"] == 101 - assert result_2025["year"] == "2025" - - # WHEN we search for the 2024 report - result_2024 = service.find_existing_report_output( - country_id="us", - simulation_1_id=101, - simulation_2_id=None, - year="2024", - ) - - # THEN we should find the 2024 report - assert result_2024 is not None - assert result_2024["simulation_1_id"] == 101 - assert result_2024["year"] == "2024" - - # AND the two reports should have different IDs - assert result_2025["id"] != result_2024["id"] - - def test_find_existing_report_output_ignores_stale_runtime_version(self, test_db): - current_version = get_report_output_cache_version("us") - stale_version = "r0stale1" - assert stale_version != current_version - - test_db.query( - "INSERT INTO report_outputs (country_id, simulation_1_id, simulation_2_id, status, api_version, year) VALUES (?, ?, ?, ?, ?, ?)", - ("us", 102, None, "complete", stale_version, "2025"), - ) - - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=102, - simulation_2_id=None, - year="2025", - ) - - assert result is None - - -class TestCreateReportOutput: - """Test creating new report outputs in the database.""" - - def test_create_report_output_single_simulation(self, test_db): - """Test creating a report output with a single simulation.""" - # GIVEN an empty database - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_single_create", - population_type="household", - policy_id=1, - ) - - # WHEN we create a report output with one simulation - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert isinstance(created_report, dict) - assert created_report["id"] > 0 - assert created_report["simulation_1_id"] == simulation["id"] - assert created_report["simulation_2_id"] is None - assert created_report["status"] == "pending" - assert created_report["year"] == "2025" - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result is not None - assert result["simulation_1_id"] == simulation["id"] - assert result["simulation_2_id"] is None - assert result["status"] == "pending" - assert result["year"] == "2025" - - def test_create_report_output_comparison(self, test_db): - """Test creating a report output comparing two simulations.""" - # GIVEN an empty database - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_report_comparison", - population_type="household", - policy_id=2, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_report_comparison", - population_type="household", - policy_id=3, - ) - - # WHEN we create a report output with two simulations - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert created_report["simulation_1_id"] == simulation_1["id"] - assert created_report["simulation_2_id"] == simulation_2["id"] - assert created_report["status"] == "pending" - assert created_report["year"] == "2025" - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result["simulation_1_id"] == simulation_1["id"] - assert result["simulation_2_id"] == simulation_2["id"] - assert result["status"] == "pending" - assert result["year"] == "2025" - - def test_create_report_output_retrieves_correct_id(self, test_db): - """Test that create_report_output retrieves the correct ID without race conditions.""" - # GIVEN we create multiple reports rapidly - - # WHEN we create reports with different parameters - created_reports = [] - simulation_ids = [] - for i in range(3): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id=f"household_report_id_{i}", - population_type="household", - policy_id=100 + i, - ) - simulation_2 = None - if i % 2 != 0: - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id=f"household_report_id_{i}", - population_type="household", - policy_id=200 + i, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None if simulation_2 is None else simulation_2["id"], - year="2025", - ) - created_reports.append(report) - simulation_ids.append( - ( - simulation_1["id"], - None if simulation_2 is None else simulation_2["id"], - ) - ) - - # THEN all IDs should be unique - ids = [report["id"] for report in created_reports] - assert len(set(ids)) == 3 - - # AND each report should have the correct data - for i, report in enumerate(created_reports): - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report["id"],) - ).fetchone() - expected_sim1, expected_sim2 = simulation_ids[i] - assert result["simulation_1_id"] == expected_sim1 - assert result["simulation_2_id"] == expected_sim2 - assert result["year"] == "2025" - - def test_create_report_output_with_different_year(self, test_db): - """Test creating a report output with a different year.""" - # GIVEN an empty database - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_other_year", - population_type="household", - policy_id=4, - ) - - # WHEN we create a report output with year 2024 - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2024", - ) - - # THEN a valid report record should be returned - assert created_report is not None - assert created_report["year"] == "2024" - assert created_report["simulation_1_id"] == simulation["id"] - - # AND the report should be in the database - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert result["year"] == "2024" - assert result["simulation_1_id"] == simulation["id"] - - def test_create_report_output_populates_dual_write_state(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_dual_write", - population_type="household", - policy_id=21, - ) - - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["report_kind"] == "household_single" - assert stored_report["report_spec_json"] is not None - assert stored_report["report_spec_schema_version"] == 1 - assert stored_report["report_spec_status"] == "explicit" - assert stored_report["active_run_id"] is not None - assert stored_report["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["trigger_type"] == "initial" - assert run["requested_at"] is not None - assert created_report["requested_at"] is not None - assert created_report["started_at"] is None - assert created_report["finished_at"] is None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "household_single" - - def test_create_report_output_reuses_existing_row_and_bootstraps_dual_write( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_existing_report", - population_type="household", - policy_id=22, - ) - - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND simulation_2_id IS NULL AND year = ? - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(rows) == 1 - assert created_report["id"] == rows[0]["id"] - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - - def test_create_report_output_populates_economy_comparison_report_spec( - self, test_db - ): - baseline_simulation = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=30, - ) - reform_simulation = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=31, - ) - created_report = service.create_report_output( +def create_simulation( + orm_session_factory, + *, + policy_id=1, + population_id="household-1", +): + return ( + SimulationService(orm_session_factory) + .get_or_create_simulation( country_id="us", - simulation_1_id=baseline_simulation["id"], - simulation_2_id=reform_simulation["id"], - year="2025", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["report_kind"] == "economy_comparison" - assert stored_report["report_spec_status"] == "backfilled_assumed" - - report_spec = stored_report["report_spec_json"] - if isinstance(report_spec, str): - report_spec = json.loads(report_spec) - assert report_spec["region"] == "state/ca" - assert report_spec["baseline_policy_id"] == 30 - assert report_spec["reform_policy_id"] == 31 - assert report_spec["dataset"] == "default" - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "economy_comparison" - assert snapshot["region"] == "state/ca" - - -class TestGetReportOutput: - """Test retrieving report outputs from the database.""" - - def test_get_report_output_existing(self, test_db, existing_report_record): - """Test retrieving an existing report output.""" - # GIVEN an existing report record - - # WHEN we retrieve the report - result = service.get_report_output( - country_id=existing_report_record["country_id"], - report_output_id=existing_report_record["id"], - ) - - # THEN the correct report should be returned - assert result is not None - assert result["id"] == existing_report_record["id"] - assert result["simulation_1_id"] == existing_report_record["simulation_1_id"] - assert result["status"] == existing_report_record["status"] - - def test_get_report_output_nonexistent(self, test_db): - """Test retrieving a non-existent report returns None.""" - # GIVEN an empty database - - # WHEN we try to retrieve a non-existent report - result = service.get_report_output(country_id="us", report_output_id=999) - - # THEN None should be returned - assert result is None - - def test_get_report_output_with_json_output(self, test_db): - """Test that JSON output is properly parsed when retrieved.""" - # GIVEN a report with JSON output - test_output = {"key": "value", "nested": {"data": 123}} - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 1, - None, - "complete", - json.dumps(test_output), - get_report_output_cache_version("us"), - "2025", - ), - ) - - # Get the ID of the inserted record - record = test_db.query( - "SELECT id FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - # WHEN we retrieve the report - result = service.get_report_output( - country_id="us", report_output_id=record["id"] - ) - - # THEN the output should be returned as JSON string (not parsed) - assert result["output"] == json.dumps(test_output) - assert result["year"] == "2025" - # Frontend will parse this string - - def test_get_report_output_includes_display_run_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_display_timestamps", + population_id=population_id, population_type="household", - policy_id=40, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 12:00:00", - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - run["id"], - ), + policy_id=policy_id, ) + .simulation + ) - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - assert result["requested_at"] == "2026-05-04T12:00:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" +def create_report(service, simulation_id, *, year="2025"): + return service.create_or_reuse_report_output( + "us", simulation_id, year=year + ).view.report_output - def test_update_report_output_sets_finished_at_on_display_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_finished_timestamp", - population_type="household", - policy_id=41, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - success = service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - assert success is True - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - assert result["status"] == "complete" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is not None - - def test_error_rerun_uses_error_run_timestamp_over_previous_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_error_rerun_timestamp", - population_type="household", - policy_id=42, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - previous_success_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 10:00:00", - "2026-05-04 10:01:00", - "2026-05-04 10:02:00", - previous_success_id, - ), - ) - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], previous_success_id, report["id"]), - ) +def test_creates_mapped_report_with_spec_and_initial_run( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="error", - error_message="rerun failed", - ) + creation = service.create_or_reuse_report_output( + country_id="us", + simulation_1_id=simulation.id, + year="2025", + ) + report = creation.view.report_output - updated_rerun = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) + assert creation.created is True + assert isinstance(report, ReportOutput) + assert report.report_kind == "household_single" + assert isinstance(report.report_spec_json, dict) + assert isinstance(creation.view.display_run, ReportOutputRun) + assert creation.view.display_run.status == "pending" + assert ( + creation.view.display_run.report_spec_snapshot_json == report.report_spec_json + ) - assert result["status"] == "error" - assert result["finished_at"] == service._format_run_timestamp( - updated_rerun["finished_at"] - ) - assert result["finished_at"] != "2026-05-04T10:02:00Z" - def test_pending_update_clears_terminal_display_timestamps(self, test_db): - simulation = simulation_service.create_simulation( +def test_create_reuses_current_report_and_repairs_dual_state( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + with orm_session_factory.begin() as session: + existing = ReportOutput( country_id="us", - population_id="household_report_pending_timestamp_reset", - population_type="household", - policy_id=43, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], + simulation_1_id=simulation.id, simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], + api_version=get_report_output_cache_version("us"), status="pending", - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "pending" - assert result["requested_at"] is not None - assert result["started_at"] is None - assert result["finished_at"] is None - - def test_running_update_sets_started_at_without_finished_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_timestamp", - population_type="household", - policy_id=44, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "running" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is None - - def test_running_update_requires_non_terminal_run_after_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_without_active_run", - population_type="household", - policy_id=45, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - - with pytest.raises(ValueError, match="active pending or running"): - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - assert stored_report["status"] == "complete" - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] == successful_run_id - assert successful_run["status"] == "complete" - assert successful_run["output"] == json.dumps({"ok": True}) - assert successful_run["finished_at"] is not None - - def test_running_update_uses_active_rerun_without_rewriting_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_active_running_rerun", - population_type="household", - policy_id=46, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - service.update_report_output( - country_id="us", - report_id=report["id"], - status="running", - ) - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - active_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - assert successful_run["status"] == "complete" - assert successful_run["finished_at"] is not None - assert active_run["status"] == "running" - assert active_run["started_at"] is not None - assert active_run["finished_at"] is None - assert stored_report["active_run_id"] == rerun["id"] - assert stored_report["latest_successful_run_id"] == successful_run_id - - def test_get_report_output_does_not_rewrite_terminal_active_run_for_running_parent( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_running_bad_active_run", - population_type="household", - policy_id=47, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - output_json = json.dumps({"ok": True}) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=output_json, - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_outputs - SET status = ?, active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("running", successful_run_id, successful_run_id, report["id"]), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - assert result["status"] == "running" - assert successful_run["status"] == "complete" - assert successful_run["output"] == output_json - assert successful_run["finished_at"] is not None - - def test_get_stored_report_output_includes_display_run_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_stored_timestamp", - population_type="household", - policy_id=48, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - result = service.get_stored_report_output("us", report["id"]) - - assert result is not None - assert result["requested_at"] is not None - assert result["started_at"] is None - assert result["finished_at"] is None - - def test_get_stored_report_output_returns_none_when_missing(self, test_db): - assert service.get_stored_report_output("us", 999999) is None - - def test_get_report_output_backfills_missing_timestamps_on_matching_legacy_run( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_timestamp_get", - population_type="household", - policy_id=46, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, started_at = NULL, finished_at = NULL - WHERE report_output_id = ? - """, - (report["id"],), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is not None - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert stored_run["requested_at"] is not None - assert stored_run["started_at"] is not None - assert stored_run["finished_at"] is not None - - def test_get_report_output_preserves_existing_finished_at_during_backfill( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_finished_at", - population_type="household", - policy_id=47, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, started_at = ?, finished_at = ? - WHERE report_output_id = ? - """, - ( - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - report["id"], - ), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] == "2026-05-04T12:01:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert service._format_run_timestamp(stored_run["requested_at"]) == ( - "2026-05-04T12:01:00Z" - ) - assert service._format_run_timestamp(stored_run["finished_at"]) == ( - "2026-05-04T12:02:00Z" - ) - - def test_get_report_output_preserves_finished_at_while_backfilling_metadata( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_finished_metadata", - population_type="household", - policy_id=48, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL, - started_at = ?, - finished_at = ?, - report_spec_snapshot_json = NULL, - country_package_version = NULL - WHERE report_output_id = ? - """, - ( - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - report["id"], - ), - ) - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["requested_at"] == "2026-05-04T12:01:00Z" - assert result["started_at"] == "2026-05-04T12:01:00Z" - assert result["finished_at"] == "2026-05-04T12:02:00Z" - stored_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert stored_run["report_spec_snapshot_json"] is not None - assert stored_run["country_package_version"] is not None - assert service._format_run_timestamp(stored_run["requested_at"]) == ( - "2026-05-04T12:01:00Z" - ) - assert service._format_run_timestamp(stored_run["finished_at"]) == ( - "2026-05-04T12:02:00Z" - ) - - def test_get_report_output_bootstraps_running_legacy_run_started_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_running", - population_type="household", - policy_id=49, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, status, api_version, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - "running", - get_report_output_cache_version("us"), - "2025", - ), - ) - report = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=report["id"] - ) - - assert result["status"] == "running" - assert result["requested_at"] is not None - assert result["started_at"] is not None - assert result["finished_at"] is None - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - assert run["status"] == "running" - assert run["started_at"] is not None - assert run["finished_at"] is None - - def test_find_existing_report_output_backfills_missing_timestamps(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_legacy_timestamp_find", - population_type="household", - policy_id=50, - ) - report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = NULL - WHERE report_output_id = ? - """, - (report["id"],), - ) - - result = service.find_existing_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - assert result is not None - assert result["requested_at"] is not None - - def test_get_report_output_resolves_stale_id_to_current_runtime_row(self, test_db): - stale_output = { - "budget": {"budgetary_impact": 1}, - "congressional_district_impact": { - "districts": [ - { - "district": "AL-01", - "average_household_income_change": 120, - "relative_household_income_change": 0.01, - } - ] - }, - } - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 2, - None, - "complete", - json.dumps(stale_output), - "r0stale1", - "2025", - ), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - current_version = get_report_output_cache_version("us") - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, output, api_version, year) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - "us", - 2, - None, - "complete", - json.dumps({"budget": {"budgetary_impact": 2}}), - current_version, - "2025", - ), - ) - - current_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=stale_record["id"] - ) - assert result is not None - assert result["id"] == stale_record["id"] - assert result["api_version"] == current_record["api_version"] - assert result["output"] == current_record["output"] - - def test_get_report_output_creates_current_runtime_row_for_stale_id(self, test_db): - stale_version = "r0stale1" - current_version = get_report_output_cache_version("us") - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_stale_runtime_create", - population_type="household", - policy_id=5, - ) - - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", simulation["id"], None, "complete", stale_version, "2025"), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - result = service.get_report_output( - country_id="us", report_output_id=stale_record["id"] - ) - - assert result is not None - assert result["id"] == stale_record["id"] - assert result["api_version"] == current_version - assert result["status"] == "pending" - assert result["output"] is None - - current_rows = test_db.query( - "SELECT * FROM report_outputs WHERE country_id = ? AND simulation_1_id = ? AND year = ? ORDER BY id ASC", - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(current_rows) == 2 - assert current_rows[0]["api_version"] == stale_version - assert current_rows[1]["api_version"] == current_version - - def test_get_report_output_invalid_id(self, test_db): - """Test that invalid report IDs are handled properly.""" - # GIVEN any database state - - # WHEN we call get_report_output with invalid ID types - # THEN an exception should be raised - with pytest.raises(Exception) as exc_info: - service.get_report_output(country_id="us", report_output_id=-1) - assert "Invalid report output ID" in str(exc_info.value) - - with pytest.raises(Exception) as exc_info: - service.get_report_output(country_id="us", report_output_id="not_an_int") - assert "Invalid report output ID" in str(exc_info.value) - - def test_get_report_output_wrong_country_returns_none( - self, test_db, existing_report_record - ): - result = service.get_report_output( - country_id="uk", - report_output_id=existing_report_record["id"], - ) - - assert result is None - - -class TestUniqueConstraint: - """Test that the unique constraint on report outputs works correctly.""" - - def test_duplicate_report_returns_existing(self, test_db): - """Test that creating duplicate reports returns the existing record.""" - # GIVEN we create a report - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_duplicate_report", - population_type="household", - policy_id=50, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_duplicate_report", - population_type="household", - policy_id=60, - ) - first_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - # WHEN we try to create an identical report - second_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], year="2025", ) + session.add(existing) + session.flush() + report_id = existing.id - # THEN the same report should be returned (no duplicate created) - assert first_report["id"] == second_report["id"] - assert first_report["country_id"] == second_report["country_id"] - assert first_report["simulation_1_id"] == second_report["simulation_1_id"] - assert first_report["simulation_2_id"] == second_report["simulation_2_id"] - assert first_report["year"] == second_report["year"] - + result = service.create_or_reuse_report_output("us", simulation.id, year="2025") -class TestUpdateReportOutput: - """Test updating report outputs in the database.""" + assert result.created is False + assert result.view.report_output.id == report_id + assert result.view.report_output.active_run_id is not None + assert result.view.report_output.report_spec_json is not None - def test_update_report_output_to_complete(self, test_db, existing_report_record): - """Test updating a report to complete status with output.""" - # GIVEN an existing pending report - report_id = existing_report_record["id"] - test_output = {"result": "success", "data": [1, 2, 3]} - test_output_json = json.dumps(test_output) - # WHEN we update it to complete with output (as JSON string) - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="complete", - output=test_output_json, - ) - - # THEN the update should succeed - assert success is True - - # AND the database should reflect the changes - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "complete" - assert result["output"] == test_output_json - - def test_update_report_output_to_error(self, test_db, existing_report_record): - """Test updating a report to error status with message.""" - # GIVEN an existing pending report - report_id = existing_report_record["id"] - error_msg = "Calculation failed due to invalid input" +@pytest.mark.parametrize("missing_secondary", [False, True]) +def test_create_rejects_missing_linked_simulation( + service, + orm_session_factory, + missing_secondary, +): + simulation = create_simulation(orm_session_factory) if missing_secondary else None - # WHEN we update it to error status - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="error", - error_message=error_msg, + with pytest.raises(ValueError, match="references missing simulation"): + service.create_or_reuse_report_output( + "us", + simulation.id if simulation else 999, + 999 if missing_secondary else None, ) - # THEN the update should succeed - assert success is True - - # AND the error should be recorded - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "error" - assert result["error_message"] == error_msg - - def test_update_report_output_partial_update(self, test_db, existing_report_record): - """Test that partial updates work correctly.""" - # GIVEN an existing report - report_id = existing_report_record["id"] - - # WHEN we update only the status - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=report_id, - status="complete", - ) + with orm_session_factory() as session: + assert session.scalar(select(func.count()).select_from(ReportOutput)) == 0 - # THEN the update should succeed - assert success is True - # AND only the status should change - result = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", (report_id,) - ).fetchone() - assert result["status"] == "complete" - assert result["output"] is None # Should remain unchanged - - def test_update_report_output_updates_dual_write_state(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_update", - population_type="household", - policy_id=23, - ) - created_report = service.create_report_output( +def test_create_ignores_stale_cache_version(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + with orm_session_factory.begin() as session: + stale = ReportOutput( country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - output_json = json.dumps({"distribution": [1, 2, 3]}) - - success = service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=output_json, - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] is not None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run["status"] == "complete" - assert run["output"] == output_json - assert run["id"] == stored_report["latest_successful_run_id"] - - def test_update_report_output_bootstraps_missing_run_state(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_legacy_report", - population_type="household", - policy_id=24, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_legacy_report", - population_type="household", - policy_id=25, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=report_output["id"], - status="error", - error_message="legacy report failure", - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is not None - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "error" - assert run["error_message"] == "legacy report failure" - - def test_update_report_output_keeps_invalid_legacy_linkage_working(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="us", - population_type="geography", - policy_id=26, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="state/ca", - population_type="geography", - policy_id=27, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=report_output["id"], - status="complete", - output=json.dumps({"budget": {"budgetary_impact": 1}}), - ) - - assert success is True - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - assert run["report_spec_snapshot_json"] is None - - def test_update_report_output_stale_id_updates_only_stale_lineage_run( - self, test_db - ): - stale_version = "r0stale1" - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_stale_lineage", - population_type="household", - policy_id=32, - ) - test_db.query( - """ - INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", simulation["id"], None, "pending", stale_version, "2025"), - ) - stale_report = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=stale_report["id"], - status="complete", - output=json.dumps({"result": "stale"}), - ) - - assert success is True - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND year = ? - ORDER BY id ASC - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert len(rows) == 1 - assert rows[0]["id"] == stale_report["id"] - assert rows[0]["status"] == "complete" - - runs = test_db.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (stale_report["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["status"] == "complete" - assert runs[0]["output"] == json.dumps({"result": "stale"}) - - def test_update_report_output_does_not_append_extra_run_for_legacy_patch_traffic( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_single_run", - population_type="household", - policy_id=33, - ) - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert first_run is not None - - success = service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=json.dumps({"distribution": [4, 5, 6]}), - ) - - assert success is True - - runs = test_db.query( - """ - SELECT * FROM report_output_runs - WHERE report_output_id = ? - ORDER BY run_sequence ASC - """, - (created_report["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["id"] == first_run["id"] - assert runs[0]["status"] == "complete" - - def test_update_report_output_no_fields_returns_false( - self, test_db, existing_report_record - ): - success = service.update_report_output( - country_id=existing_report_record["country_id"], - report_id=existing_report_record["id"], - ) - - assert success is False - - def test_update_report_output_stale_id_keeps_stale_output_quarantined( - self, test_db - ): - stale_version = "r0stale1" - output_json = json.dumps({"result": "fresh"}) - - test_db.query( - """INSERT INTO report_outputs - (country_id, simulation_1_id, simulation_2_id, status, api_version, year) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", 4, None, "pending", stale_version, "2025"), - ) - - stale_record = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - success = service.update_report_output( - country_id="us", - report_id=stale_record["id"], - status="complete", - output=output_json, - ) - - assert success is True - - rows = test_db.query( - "SELECT * FROM report_outputs WHERE country_id = ? AND simulation_1_id = ? AND year = ? ORDER BY id ASC", - ("us", 4, "2025"), - ).fetchall() - - assert len(rows) == 1 - assert rows[0]["id"] == stale_record["id"] - assert rows[0]["api_version"] == stale_version - assert rows[0]["status"] == "complete" - assert rows[0]["output"] == output_json - - def test_create_report_output_rolls_back_parent_insert_on_dual_write_failure( - self, test_db, monkeypatch - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_create_rollback", - population_type="household", - policy_id=34, - ) - - def fail_dual_write(tx, report_output_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_report_output_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - rows = test_db.query( - """ - SELECT * FROM report_outputs - WHERE country_id = ? AND simulation_1_id = ? AND simulation_2_id IS NULL AND year = ? - """, - ("us", simulation["id"], "2025"), - ).fetchall() - assert rows == [] - - def test_update_report_output_rolls_back_parent_update_on_dual_write_failure( - self, test_db, monkeypatch - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_report_update_rollback", - population_type="household", - policy_id=35, - ) - created_report = service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], + simulation_1_id=simulation.id, simulation_2_id=None, + api_version="stale", + status="pending", year="2025", ) + session.add(stale) + session.flush() + stale_id = stale.id - def fail_dual_write(tx, report_output_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_report_output_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.update_report_output( - country_id="us", - report_id=created_report["id"], - status="complete", - output=json.dumps({"rolled_back": True}), - ) - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (created_report["id"],), - ).fetchone() - assert stored_report["status"] == "pending" - assert stored_report["output"] is None + current = service.create_or_reuse_report_output("us", simulation.id, year="2025") - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (created_report["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["output"] is None - - def test_ensure_report_output_dual_write_state_bootstraps_linked_simulations( - self, test_db - ): - test_db.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - "us-system-1.0.0", - "household_stage5_linked", - "household", - 36, - "pending", - ), - ) - simulation_1 = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - test_db.query( - """ - INSERT INTO simulations ( - country_id, api_version, population_id, population_type, policy_id, status - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - "us-system-1.0.0", - "household_stage5_linked", - "household", - 37, - "pending", - ), - ) - simulation_2 = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - simulation_2["id"], - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - synced_report = service.ensure_report_output_dual_write_state( - report_output["id"], - country_id="us", - ) + assert current.created is True + assert current.view.report_output.id != stale_id + assert current.view.report_output.api_version == get_report_output_cache_version( + "us" + ) - assert synced_report["active_run_id"] is not None - stored_simulation_1 = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_1["id"],), - ).fetchone() - stored_simulation_2 = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation_2["id"],), - ).fetchone() - assert stored_simulation_1["simulation_spec_json"] is not None - assert stored_simulation_1["active_run_id"] is not None - assert stored_simulation_2["simulation_spec_json"] is not None - assert stored_simulation_2["active_run_id"] is not None +def test_resolve_is_scoped_to_country_and_validates_id( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + + assert service.resolve_report_output("us", report.id).report_output.id == report.id + assert service.resolve_report_output("uk", report.id) is None + with pytest.raises(Exception, match="Invalid report output ID"): + service.resolve_report_output("us", -1) - simulation_1_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation_1["id"],), - ).fetchone() - simulation_2_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation_2["id"],), - ).fetchone() - assert simulation_1_run is not None - assert simulation_2_run is not None + +def test_update_complete_stores_python_json_and_promotes_run( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + active_run_id = report.active_run_id + + view = service.update_report_output( + "us", + report.id, + status="complete", + output={"ok": True}, + ) + + assert view.report_output.output == {"ok": True} + assert view.report_output.active_run_id is None + assert view.report_output.latest_successful_run_id == active_run_id + assert view.display_run.status == "complete" + assert view.display_run.output == {"ok": True} + assert view.display_run.started_at is not None + assert view.display_run.finished_at is not None + + +def test_update_accepts_existing_v1_json_string_boundary( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + + view = service.update_report_output("us", report.id, output='{"ok": true}') + + assert view.report_output.output == {"ok": True} + + +def test_update_running_requires_mutable_run(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + service.update_report_output( + "us", report.id, status="complete", output={"ok": True} + ) + + with pytest.raises(ValueError, match="without an active pending or running"): + service.update_report_output("us", report.id, status="running") + + +def _add_rerun(orm_session_factory, report_id): + with orm_session_factory.begin() as session: + report = session.get(ReportOutput, report_id) + successful_run_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + session, report.id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + return rerun.id, successful_run_id + + +def test_update_running_targets_active_rerun(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + service.update_report_output( + "us", report.id, status="complete", output={"old": True} + ) + rerun_id, successful_run_id = _add_rerun(orm_session_factory, report.id) + + view = service.update_report_output("us", report.id, status="running") + + assert view.display_run.id == rerun_id + assert view.display_run.status == "running" + assert view.display_run.started_at is not None + assert view.report_output.active_run_id == rerun_id + assert view.report_output.latest_successful_run_id == successful_run_id + + +def test_failed_rerun_preserves_latest_successful_pointer( + service, + orm_session_factory, +): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + service.update_report_output( + "us", report.id, status="complete", output={"old": True} + ) + rerun_id, successful_run_id = _add_rerun(orm_session_factory, report.id) + + view = service.update_report_output( + "us", report.id, status="error", error_message="failed" + ) + + assert view.report_output.active_run_id is None + assert view.report_output.latest_successful_run_id == successful_run_id + assert view.display_run.id == rerun_id + assert view.display_run.status == "error" + assert view.display_run.finished_at is not None + + +def test_noop_and_missing_updates(service, orm_session_factory): + simulation = create_simulation(orm_session_factory) + report = create_report(service, simulation.id) + + assert service.update_report_output("us", report.id) is None + with pytest.raises(LookupError, match="Report output #999 not found"): + service.update_report_output("us", 999, status="pending") diff --git a/tests/unit/services/test_report_run_service.py b/tests/unit/services/test_report_run_service.py index 68aedb48f..890ef5e5d 100644 --- a/tests/unit/services/test_report_run_service.py +++ b/tests/unit/services/test_report_run_service.py @@ -1,382 +1,98 @@ import pytest -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun from policyengine_api.services.report_run_service import ReportRunService -from policyengine_api.services.simulation_service import SimulationService -report_output_service = ReportOutputService() -report_run_service = ReportRunService() -simulation_service = SimulationService() - -class TestCreateReportOutputRun: - def test_creates_report_runs_with_incrementing_sequence(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = report_run_service.create_report_output_run( - report_output["id"], - trigger_type="initial", - report_spec_snapshot={"country_id": "us"}, - version_manifest={ - "country_package_version": "us-1.0.0", - "report_cache_version": "r123", - }, - ) - second_run = report_run_service.create_report_output_run( - report_output["id"], - trigger_type="rerun", - ) - - assert first_run["run_sequence"] == 2 - assert first_run["trigger_type"] == "initial" - assert first_run["requested_at"] is not None - assert first_run["started_at"] is None - assert first_run["finished_at"] is None - assert first_run["report_spec_snapshot_json"] == {"country_id": "us"} - assert first_run["country_package_version"] == "us-1.0.0" - assert first_run["report_cache_version"] == "r123" - assert second_run["run_sequence"] == 3 - assert second_run["trigger_type"] == "rerun" - - def test_lists_report_runs_in_sequence_order(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - - runs = report_run_service.list_report_output_runs(report_output["id"]) - - assert [run["run_sequence"] for run in runs] == [1, 2, 3] - - def test_allocates_run_sequence_transactionally(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_7", - population_type="household", - policy_id=7, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - first_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - second_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - - assert first_run["run_sequence"] == 2 - assert second_run["run_sequence"] == 3 - - def test_raises_when_parent_report_output_is_missing(self, test_db): - with pytest.raises(ValueError) as exc_info: - report_run_service.create_report_output_run(999999, trigger_type="initial") - - assert "Report output #999999 not found" in str(exc_info.value) - - def test_running_report_run_sets_started_at(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_running_run_timestamp", - population_type="household", - policy_id=8, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - run = report_run_service.create_report_output_run( - report_output["id"], - status="running", - trigger_type="rerun", - ) - - assert run["requested_at"] is not None - assert run["started_at"] is not None - assert run["finished_at"] is None - - -class TestSelectDisplayReportRun: - def test_prefers_active_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - latest_successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - active_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (active_run["id"], latest_successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == active_run["id"] - - def test_falls_back_to_latest_successful_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - (successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == successful_run["id"] - - def test_prefers_matching_error_run_over_previous_success(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5b", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], - status="complete", - trigger_type="initial", - output={"ok": True}, - ) - error_run = report_run_service.create_report_output_run( - report_output["id"], - status="error", - trigger_type="rerun", - error_message="rerun failed", - ) - test_db.query( - """ - UPDATE report_outputs - SET status = ?, error_message = ?, active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("error", "rerun failed", successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == error_run["id"] - - def test_falls_back_when_active_run_pointer_is_stale(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_5a", - population_type="household", - policy_id=5, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - successful_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", successful_run["id"], report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_to_newest_run_when_no_pointers_exist(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - first_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="initial" - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - (report_output["id"],), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert first_run["run_sequence"] == 2 - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_latest_successful_pointer_is_stale( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6a", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_no_pointer_or_result_match(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_6b", - population_type="household", - policy_id=6, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - newest_run = report_run_service.create_report_output_run( - report_output["id"], status="pending", trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET status = ?, active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - ("complete", report_output["id"]), - ) - updated_report_output = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - - selected_run = report_run_service.select_display_run(updated_report_output) - - assert selected_run["id"] == newest_run["id"] +service = ReportRunService() + + +def create_report(orm_session, *, status="pending"): + report = ReportOutput( + country_id="us", + simulation_1_id=1, + simulation_2_id=None, + api_version="1", + status=status, + year="2025", + ) + orm_session.add(report) + orm_session.flush() + return report + + +def test_creates_mapped_runs_with_incrementing_sequence_and_python_json(orm_session): + report = create_report(orm_session) + + first = service.create_report_output_run( + orm_session, + report.id, + trigger_type="initial", + report_spec_snapshot={"country_id": "us"}, + version_manifest={"report_cache_version": "r123"}, + ) + second = service.create_report_output_run( + orm_session, report.id, trigger_type="rerun" + ) + + assert isinstance(first, ReportOutputRun) + assert first.run_sequence == 1 + assert first.requested_at is not None + assert first.started_at is None + assert first.finished_at is None + assert first.report_spec_snapshot_json == {"country_id": "us"} + assert first.report_cache_version == "r123" + assert second.run_sequence == 2 + + +@pytest.mark.parametrize( + ("status", "has_started", "has_finished"), + [ + ("pending", False, False), + ("running", True, False), + ("complete", True, True), + ("error", True, True), + ], +) +def test_sets_run_timestamps_from_status( + orm_session, status, has_started, has_finished +): + report = create_report(orm_session) + + run = service.create_report_output_run(orm_session, report.id, status=status) + + assert (run.started_at is not None) is has_started + assert (run.finished_at is not None) is has_finished + + +def test_raises_when_parent_report_is_missing(orm_session): + with pytest.raises(ValueError, match="Report output #999 not found"): + service.create_report_output_run(orm_session, 999) + + +def test_gets_lists_and_selects_mapped_runs(orm_session): + report = create_report(orm_session) + first = service.create_report_output_run(orm_session, report.id, status="complete") + second = service.create_report_output_run(orm_session, report.id, status="running") + report.latest_successful_run_id = first.id + report.active_run_id = second.id + + assert service.get_report_output_run(orm_session, first.id) is first + assert service.list_report_output_runs(orm_session, report.id) == [first, second] + assert service.get_newest_report_output_run(orm_session, report.id) is second + assert service.select_display_run(orm_session, report) is second + + +def test_select_display_run_falls_back_to_matching_error(orm_session): + report = create_report(orm_session, status="error") + matching = service.create_report_output_run( + orm_session, + report.id, + status="error", + error_message="failed", + ) + service.create_report_output_run(orm_session, report.id) + report.error_message = "failed" + report.active_run_id = None + + assert service.select_display_run(orm_session, report) is matching diff --git a/tests/unit/services/test_report_service_owned_sessions.py b/tests/unit/services/test_report_service_owned_sessions.py new file mode 100644 index 000000000..e9103a07b --- /dev/null +++ b/tests/unit/services/test_report_service_owned_sessions.py @@ -0,0 +1,110 @@ +import inspect +from pathlib import Path + +import pytest +from sqlalchemy import func, select + +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun +from policyengine_api.services.report_output_service import ( + ReportCreateResult, + ReportOutputService, + ReportOutputView, +) +from policyengine_api.services.simulation_service import SimulationService + + +ROUTE_PATH = ( + Path(__file__).parents[3] + / "policyengine_api" + / "routes" + / "report_output_routes.py" +) + + +@pytest.fixture +def simulation_id(orm_session_factory): + return ( + SimulationService(orm_session_factory) + .get_or_create_simulation("us", "household-1", "household", 1) + .simulation.id + ) + + +@pytest.fixture +def service(orm_session_factory): + return ReportOutputService(orm_session_factory) + + +def test_report_public_methods_do_not_accept_sessions(): + for method_name in ( + "create_or_reuse_report_output", + "resolve_report_output", + "update_report_output", + ): + parameters = inspect.signature( + getattr(ReportOutputService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_report_routes_do_not_manage_sessions_or_query_run_services(): + source = ROUTE_PATH.read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "report_run_service" not in source + assert "sqlalchemy" not in source + + +def test_create_or_reuse_returns_report_and_display_run(service, simulation_id): + created = service.create_or_reuse_report_output("us", simulation_id, year="2025") + reused = service.create_or_reuse_report_output("us", simulation_id, year="2025") + + assert isinstance(created, ReportCreateResult) + assert created.created is True + assert isinstance(created.view, ReportOutputView) + assert isinstance(created.view.report_output, ReportOutput) + assert isinstance(created.view.display_run, ReportOutputRun) + assert created.view.response_id is None + assert reused.created is False + assert reused.view.report_output.id == created.view.report_output.id + + +def test_create_rolls_back_all_rows_when_dual_write_fails( + service, + simulation_id, + orm_session_factory, + monkeypatch, +): + def fail_dual_write(*args, **kwargs): + raise RuntimeError("report dual write failed") + + monkeypatch.setattr( + service, "_ensure_report_output_dual_write_state", fail_dual_write + ) + + with pytest.raises(RuntimeError, match="report dual write failed"): + service.create_or_reuse_report_output("us", simulation_id, year="2026") + + with orm_session_factory() as session: + report_count = session.scalar(select(func.count()).select_from(ReportOutput)) + run_count = session.scalar(select(func.count()).select_from(ReportOutputRun)) + assert report_count == 0 + assert run_count == 0 + + +def test_update_returns_view_from_the_same_atomic_operation( + service, + simulation_id, +): + created = service.create_or_reuse_report_output("us", simulation_id) + + view = service.update_report_output( + "us", + created.view.report_output.id, + status="complete", + output={"result": 42}, + ) + + assert isinstance(view, ReportOutputView) + assert view.report_output.output == {"result": 42} + assert view.display_run.status == "complete" diff --git a/tests/unit/services/test_report_spec_service.py b/tests/unit/services/test_report_spec_service.py index f924df8db..3db91765e 100644 --- a/tests/unit/services/test_report_spec_service.py +++ b/tests/unit/services/test_report_spec_service.py @@ -1,469 +1,159 @@ import pytest -from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.data.v1_models import ReportOutput, Simulation from policyengine_api.services.report_spec_service import ( EconomyReportSpec, HouseholdReportSpec, ReportSpecService, ) -from policyengine_api.services.simulation_service import SimulationService -report_output_service = ReportOutputService() -report_spec_service = ReportSpecService() -simulation_service = SimulationService() - -class TestBuildReportSpec: - def test_builds_household_comparison_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2026", - ) - - report_spec = report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert isinstance(report_spec, HouseholdReportSpec) - assert report_spec.report_kind == "household_comparison" - assert report_spec.time_period == "2026" - assert report_spec.simulation_1.policy_id == 1 - assert report_spec.simulation_2.policy_id == 2 - - def test_builds_default_economy_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=11, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2027", - ) - - report_spec = report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert isinstance(report_spec, EconomyReportSpec) - assert report_spec.report_kind == "economy_comparison" - assert report_spec.region == "ca" - assert report_spec.dataset == "default" - assert report_spec.target == "general" - assert report_spec.options == {} - - def test_raises_for_mixed_population_types(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "population types must match" in str(exc_info.value) - - def test_raises_for_mismatched_household_ids(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", +service = ReportSpecService() + + +def add_simulation( + orm_session, + *, + population_type="household", + population_id="household-1", + policy_id=1, + country_id="us", +): + simulation = Simulation( + country_id=country_id, + api_version="1", + population_id=population_id, + population_type=population_type, + policy_id=policy_id, + status="pending", + ) + orm_session.add(simulation) + orm_session.flush() + return simulation + + +def add_report(orm_session, simulation_1, simulation_2=None, *, country_id="us"): + report = ReportOutput( + country_id=country_id, + simulation_1_id=simulation_1.id, + simulation_2_id=simulation_2.id if simulation_2 else None, + api_version="1", + status="pending", + year="2026", + ) + orm_session.add(report) + orm_session.flush() + return report + + +def test_builds_household_comparison_spec_from_models(orm_session): + first = add_simulation(orm_session, policy_id=1) + second = add_simulation(orm_session, policy_id=2) + report = add_report(orm_session, first, second) + + spec = service.build_report_spec(report, first, second) + + assert isinstance(spec, HouseholdReportSpec) + assert spec.report_kind == "household_comparison" + assert spec.simulation_1.policy_id == 1 + assert spec.simulation_2.policy_id == 2 + + +def test_builds_economy_spec_from_models(orm_session): + first = add_simulation( + orm_session, + population_type="geography", + population_id="ca", + policy_id=10, + ) + second = add_simulation( + orm_session, + population_type="geography", + population_id="ca", + policy_id=11, + ) + report = add_report(orm_session, first, second) + + spec = service.build_report_spec( + report, first, second, dataset="cps", options={"foo": "bar"} + ) + + assert isinstance(spec, EconomyReportSpec) + assert spec.report_kind == "economy_comparison" + assert spec.region == "ca" + assert spec.dataset == "cps" + assert spec.options == {"foo": "bar"} + + +def test_sets_and_gets_report_spec_as_python_json(orm_session): + simulation = add_simulation(orm_session) + report = add_report(orm_session, simulation) + spec = service.build_report_spec(report, simulation) + + assert service.set_report_spec( + orm_session, report.id, spec, report_spec_status="explicit" + ) + loaded = service.get_report_spec(orm_session, report.id) + + assert report.report_spec_json == spec.model_dump() + assert report.report_spec_schema_version == 1 + assert report.report_spec_status == "explicit" + assert loaded == spec + + +def test_rejects_missing_linked_simulation(orm_session): + report = ReportOutput( + country_id="us", + simulation_1_id=999, + simulation_2_id=None, + api_version="1", + status="pending", + year="2026", + ) + orm_session.add(report) + orm_session.flush() + spec = HouseholdReportSpec.model_validate( + { + "country_id": "us", + "report_kind": "household_single", + "time_period": "2026", + "simulation_1": { + "population_type": "household", + "population_id": "household-1", + "policy_id": 1, + }, + } + ) + + with pytest.raises(ValueError, match="references missing simulation #999"): + service.set_report_spec(orm_session, report.id, spec, "explicit") + + +def test_rejects_mismatched_country(orm_session): + simulation = add_simulation(orm_session, country_id="uk") + report = add_report(orm_session, simulation, country_id="us") + + with pytest.raises(ValueError, match="country must match"): + service.build_report_spec(report, simulation) + + +def test_rejects_mismatched_comparison_population(orm_session): + first = add_simulation(orm_session, population_id="household-1") + second = add_simulation(orm_session, population_id="household-2", policy_id=2) + report = add_report(orm_session, first, second) + + with pytest.raises(ValueError, match="matching household IDs"): + service.build_report_spec(report, first, second) + + +def test_rejects_unsupported_schema_version_and_status(orm_session): + simulation = add_simulation(orm_session) + report = add_report(orm_session, simulation) + spec = service.build_report_spec(report, simulation) + + with pytest.raises(ValueError, match="Unsupported report spec status"): + service.set_report_spec(orm_session, report.id, spec, "unknown") + with pytest.raises(ValueError, match="Unsupported report spec schema version"): + service.set_report_spec( + orm_session, report.id, spec, "explicit", schema_version=2 ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "matching household IDs" in str(exc_info.value) - - def test_raises_for_mismatched_geography_ids(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="ny", - population_type="geography", - policy_id=11, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2027", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec( - report_output, simulation_1, simulation_2 - ) - - assert "matching geography IDs" in str(exc_info.value) - - def test_raises_for_country_mismatch_between_report_and_simulation(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="uk", - population_id="household_1", - population_type="household", - policy_id=1, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation_1["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, simulation_1) - - assert "Simulation 1 country must match report output country" in str( - exc_info.value - ) - - def test_raises_when_simulation_1_does_not_match_report_output_linkage( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - other_simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, other_simulation) - - assert "Simulation 1 must match report_output.simulation_1_id" in str( - exc_info.value - ) - - def test_raises_when_report_requires_second_simulation_but_none_is_supplied( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - simulation_2 = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=2, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=simulation_2["id"], - year="2025", - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.build_report_spec(report_output, simulation_1) - - assert "requires a second simulation" in str(exc_info.value) - - -class TestPersistReportSpec: - def test_sets_and_gets_report_spec(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - result = report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert result is True - stored_report = test_db.query( - """ - SELECT report_kind, report_spec_json, report_spec_schema_version, report_spec_status - FROM report_outputs WHERE id = ? - """, - (report_output["id"],), - ).fetchone() - assert stored_report["report_kind"] == "economy_single" - assert stored_report["report_spec_schema_version"] == 1 - assert stored_report["report_spec_status"] == "explicit" - - loaded_report_spec = report_spec_service.get_report_spec(report_output["id"]) - assert loaded_report_spec is not None - assert loaded_report_spec.model_dump() == report_spec.model_dump() - - def test_rejects_report_spec_write_when_region_does_not_match_report(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ny", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert "Report spec region must match linked simulations" in str(exc_info.value) - - def test_rejects_report_spec_write_when_time_period_does_not_match_report( - self, test_db - ): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2028", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], report_spec, report_spec_status="explicit" - ) - - assert "time_period must match report output year" in str(exc_info.value) - - def test_rejects_inconsistent_stored_report_spec_on_read(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - test_db.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - "economy_single", - '{"country_id":"us","report_kind":"economy_single","time_period":"2027","region":"ny","baseline_policy_id":10,"reform_policy_id":10,"dataset":"default","target":"general","options":{}}', - 1, - "explicit", - report_output["id"], - ), - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.get_report_spec(report_output["id"]) - - assert "Report spec region must match linked simulations" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_write(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - report_spec = EconomyReportSpec.model_validate( - { - "country_id": "us", - "report_kind": "economy_single", - "time_period": "2027", - "region": "ca", - "baseline_policy_id": 10, - "reform_policy_id": 10, - "dataset": "default", - "target": "general", - "options": {}, - } - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.set_report_spec( - report_output["id"], - report_spec, - report_spec_status="explicit", - schema_version=2, - ) - - assert "Unsupported report spec schema version" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_read(self, test_db): - simulation_1 = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=10, - ) - report_output = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation_1["id"], - simulation_2_id=None, - year="2027", - ) - test_db.query( - """ - UPDATE report_outputs - SET report_kind = ?, report_spec_json = ?, report_spec_schema_version = ?, report_spec_status = ? - WHERE id = ? - """, - ( - "economy_single", - '{"country_id":"us","report_kind":"economy_single","time_period":"2027","region":"ca","baseline_policy_id":10,"reform_policy_id":10,"dataset":"default","target":"general","options":{}}', - 2, - "explicit", - report_output["id"], - ), - ) - - with pytest.raises(ValueError) as exc_info: - report_spec_service.get_report_spec(report_output["id"]) - - assert "Unsupported report spec schema version" in str(exc_info.value) diff --git a/tests/unit/services/test_run_and_spec_service_database_boundaries.py b/tests/unit/services/test_run_and_spec_service_database_boundaries.py new file mode 100644 index 000000000..fcf6cf398 --- /dev/null +++ b/tests/unit/services/test_run_and_spec_service_database_boundaries.py @@ -0,0 +1,23 @@ +"""Database-access boundaries for run and specification services.""" + +from pathlib import Path + +import pytest + + +SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" + + +@pytest.mark.parametrize( + "module_name", + [ + "simulation_run_service.py", + "simulation_spec_service.py", + "report_run_service.py", + "report_spec_service.py", + ], +) +def test_run_and_spec_services_do_not_issue_queries_directly(module_name): + source = (SERVICE_ROOT / module_name).read_text(encoding="utf-8") + assert ".query(" not in source + assert "from policyengine_api.data import database" not in source diff --git a/tests/unit/services/test_service_owned_session_architecture.py b/tests/unit/services/test_service_owned_session_architecture.py new file mode 100644 index 000000000..4f41dd600 --- /dev/null +++ b/tests/unit/services/test_service_owned_session_architecture.py @@ -0,0 +1,130 @@ +"""Architecture guards for service-owned SQLAlchemy sessions.""" + +import inspect +from pathlib import Path + +import pytest + +from policyengine_api.services.ai_analysis_service import AIAnalysisService +from policyengine_api.services.household_calculation_service import ( + HouseholdCalculationService, +) +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.reform_impacts_service import ReformImpactsService +from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.services.simulation_analysis_service import ( + SimulationAnalysisService, +) +from policyengine_api.services.simulation_service import SimulationService +from policyengine_api.services.tracer_analysis_service import TracerAnalysisService +from policyengine_api.services.user_policy_service import UserPolicyService +from policyengine_api.services.user_service import UserService + + +PROJECT_ROOT = Path(__file__).parents[3] +PACKAGE_ROOT = PROJECT_ROOT / "policyengine_api" + + +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + ( + HouseholdService, + ("get_household", "create_household", "update_household"), + ), + ( + HouseholdCalculationService, + ("calculate_stored_household", "calculate_household"), + ), + ( + PolicyService, + ("get_policy", "get_policy_json", "search_policies", "set_policy"), + ), + ( + UserPolicyService, + ( + "create_or_get_user_policy", + "list_user_policies", + "update_user_policy", + ), + ), + (UserService, ("get_profile", "create_profile", "update_profile")), + ( + SimulationService, + ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", + ), + ), + ( + ReportOutputService, + ( + "create_or_reuse_report_output", + "resolve_report_output", + "update_report_output", + ), + ), + (AIAnalysisService, ("get_existing_analysis", "trigger_ai_analysis")), + (SimulationAnalysisService, ("execute_analysis",)), + (TracerAnalysisService, ("execute_analysis", "get_tracer")), + ( + ReformImpactsService, + ( + "get_recent_reform_impacts", + "get_all_reform_impacts", + "get_all_reform_impacts_by_options_hash_prefix", + "set_reform_impact", + "delete_reform_impact", + "set_error_reform_impact", + "set_complete_reform_impact", + ), + ), + ], +) +def test_route_facing_service_methods_hide_persistence_dependencies( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_presentation_layer_does_not_import_or_create_sqlalchemy_sessions(): + offenders = [] + presentation_paths = [ + *sorted((PACKAGE_ROOT / "routes").glob("*.py")), + PACKAGE_ROOT / "country.py", + ] + banned_tokens = ( + "get_v1_session_factory", + "from sqlalchemy", + "import sqlalchemy", + "sessionmaker", + ) + for path in presentation_paths: + source = path.read_text(encoding="utf-8") + if any(token in source for token in banned_tokens): + offenders.append(str(path.relative_to(PACKAGE_ROOT))) + + assert offenders == [] + + +def test_removed_persistence_abstractions_stay_removed(): + removed_files = ( + "data/v1_daos.py", + "data/data.py", + ) + assert [path for path in removed_files if (PACKAGE_ROOT / path).exists()] == [] + assert not any((PACKAGE_ROOT / "endpoints").rglob("*.py")) + + +def test_changelog_describes_service_owned_sessions(): + changelog = (PROJECT_ROOT / "changelog.d/3788.changed.md").read_text( + encoding="utf-8" + ) + assert "service-owned" in changelog + assert "caller-owned" not in changelog diff --git a/tests/unit/services/test_service_owned_sessions.py b/tests/unit/services/test_service_owned_sessions.py new file mode 100644 index 000000000..1976013ad --- /dev/null +++ b/tests/unit/services/test_service_owned_sessions.py @@ -0,0 +1,111 @@ +import inspect +from pathlib import Path + +import pytest +from sqlalchemy import func, select + +from policyengine_api.data.v1_models import Household +from policyengine_api.services.household_service import HouseholdService +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.simulation_service import SimulationService +from policyengine_api.services.user_service import UserService + + +ROUTE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "routes" + + +@pytest.mark.parametrize( + ("service_type", "method_names"), + [ + ( + HouseholdService, + ("get_household", "create_household", "update_household"), + ), + ( + PolicyService, + ("get_policy", "get_policy_json", "search_policies", "set_policy"), + ), + ( + SimulationService, + ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", + ), + ), + (UserService, ("get_profile", "create_profile", "update_profile")), + ], +) +def test_core_service_public_methods_do_not_accept_sessions( + service_type, + method_names, +): + for method_name in method_names: + parameters = inspect.signature(getattr(service_type, method_name)).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +@pytest.mark.parametrize( + "module_name", + [ + "household_routes.py", + "policy_routes.py", + "simulation_routes.py", + "user_profile_routes.py", + ], +) +def test_core_routes_do_not_manage_sessions(module_name): + source = (ROUTE_ROOT / module_name).read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "sqlalchemy" not in source + + +def test_core_services_commit_writes_and_return_generated_ids( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + service = HouseholdService(orm_session_factory) + + household = service.create_household( + "us", + {"people": {"you": {"age": {"2026": 40}}}}, + "Service-owned transaction", + ) + + assert household.id is not None + with orm_session_factory() as session: + stored = session.get(Household, household.id) + assert stored is not None + assert stored.household_json == {"people": {"you": {"age": {"2026": 40}}}} + + +def test_core_services_roll_back_failed_writes( + orm_session_factory, + monkeypatch, +): + monkeypatch.setattr( + "policyengine_api.services.household_service.hash_object", + lambda value: "household-hash", + ) + session_type = orm_session_factory.class_ + original_flush = session_type.flush + + def fail_after_flush(session, *args, **kwargs): + original_flush(session, *args, **kwargs) + raise RuntimeError("forced failure") + + monkeypatch.setattr(session_type, "flush", fail_after_flush) + service = HouseholdService(orm_session_factory) + + with pytest.raises(RuntimeError, match="forced failure"): + service.create_household("us", {}, "Rolled back") + + monkeypatch.setattr(session_type, "flush", original_flush) + with orm_session_factory() as session: + count = session.scalar(select(func.count()).select_from(Household)) + assert count == 0 diff --git a/tests/unit/services/test_simulation_run_service.py b/tests/unit/services/test_simulation_run_service.py index f47e13a0b..43ba3b80b 100644 --- a/tests/unit/services/test_simulation_run_service.py +++ b/tests/unit/services/test_simulation_run_service.py @@ -1,214 +1,90 @@ import pytest +from policyengine_api.data.v1_models import SimulationRun from policyengine_api.services.simulation_run_service import SimulationRunService from policyengine_api.services.simulation_service import SimulationService -simulation_run_service = SimulationRunService() + +run_service = SimulationRunService() simulation_service = SimulationService() -class TestCreateSimulationRun: - def test_creates_simulation_runs_with_incrementing_sequence(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1", - population_type="household", - policy_id=1, - ) - - first_run = simulation_run_service.create_simulation_run( - simulation["id"], - input_position=1, - trigger_type="initial", - simulation_spec_snapshot={"population_id": "household_1"}, - version_manifest={"simulation_cache_version": "s123"}, - ) - second_run = simulation_run_service.create_simulation_run( - simulation["id"], - input_position=1, - trigger_type="rerun", - ) - - assert first_run["run_sequence"] == 2 - assert first_run["trigger_type"] == "initial" - assert first_run["simulation_spec_snapshot_json"] == { - "population_id": "household_1" - } - assert first_run["simulation_cache_version"] == "s123" - assert second_run["run_sequence"] == 3 - assert second_run["trigger_type"] == "rerun" - - def test_allocates_run_sequence_transactionally(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_1a", - population_type="household", - policy_id=1, - ) - - first_run = simulation_run_service.create_simulation_run( - simulation["id"], input_position=1, trigger_type="initial" - ) - second_run = simulation_run_service.create_simulation_run( - simulation["id"], input_position=1, trigger_type="rerun" - ) - - assert first_run["run_sequence"] == 2 - assert second_run["run_sequence"] == 3 - - def test_raises_when_parent_simulation_is_missing(self, test_db): - with pytest.raises(ValueError) as exc_info: - simulation_run_service.create_simulation_run( - 999999, input_position=1, trigger_type="initial" - ) - - assert "Simulation #999999 not found" in str(exc_info.value) - - -class TestSelectDisplaySimulationRun: - def test_prefers_active_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_2", - population_type="household", - policy_id=2, - ) - latest_successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - active_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (active_run["id"], latest_successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == active_run["id"] - - def test_falls_back_to_latest_successful_run(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3", - population_type="household", - policy_id=3, - ) - successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - (successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_when_active_run_pointer_is_stale(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_3a", - population_type="household", - policy_id=3, - ) - successful_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", successful_run["id"], simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == successful_run["id"] - - def test_falls_back_to_newest_run_when_no_pointers_exist(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4", - population_type="household", - policy_id=4, - ) - first_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="initial" - ) - newest_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = NULL - WHERE id = ? - """, - (simulation["id"],), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert first_run["run_sequence"] == 2 - assert selected_run["id"] == newest_run["id"] - - def test_falls_back_to_newest_run_when_latest_successful_pointer_is_stale( - self, test_db - ): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_4a", - population_type="household", - policy_id=4, - ) - newest_run = simulation_run_service.create_simulation_run( - simulation["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE simulations - SET active_run_id = NULL, latest_successful_run_id = ? - WHERE id = ? - """, - ("missing-run", simulation["id"]), - ) - updated_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - - selected_run = simulation_run_service.select_display_run(updated_simulation) - - assert selected_run["id"] == newest_run["id"] +def create_simulation(orm_session, population_id="household-1"): + return simulation_service._create_simulation( + orm_session, "us", population_id, "household", 1 + ) + + +def test_creates_mapped_runs_with_incrementing_sequence(orm_session): + simulation = create_simulation(orm_session) + + first = run_service.create_simulation_run( + orm_session, + simulation.id, + input_position=1, + trigger_type="rerun", + simulation_spec_snapshot={"population_id": "household-1"}, + version_manifest={"simulation_cache_version": "s123"}, + ) + second = run_service.create_simulation_run( + orm_session, simulation.id, input_position=1, trigger_type="rerun" + ) + + assert isinstance(first, SimulationRun) + assert first.run_sequence == 2 + assert first.simulation_spec_snapshot_json == {"population_id": "household-1"} + assert first.simulation_cache_version == "s123" + assert second.run_sequence == 3 + + +def test_raises_when_parent_simulation_is_missing(orm_session): + with pytest.raises(ValueError, match="Simulation #999999 not found"): + run_service.create_simulation_run(orm_session, 999999) + + +def test_gets_and_lists_runs_as_models(orm_session): + simulation = create_simulation(orm_session) + first = run_service.get_simulation_run(orm_session, simulation.active_run_id) + second = run_service.create_simulation_run(orm_session, simulation.id) + + assert run_service.get_simulation_run(orm_session, second.id) is second + assert run_service.list_simulation_runs(orm_session, simulation.id) == [ + first, + second, + ] + assert run_service.get_newest_simulation_run(orm_session, simulation.id) is second + + +def test_display_run_prefers_active_run(orm_session): + simulation = create_simulation(orm_session) + successful = run_service.create_simulation_run( + orm_session, simulation.id, status="complete" + ) + active = run_service.create_simulation_run( + orm_session, simulation.id, status="running" + ) + simulation.latest_successful_run_id = successful.id + simulation.active_run_id = active.id + + assert run_service.select_display_run(orm_session, simulation) is active + + +def test_display_run_falls_back_to_latest_successful_run(orm_session): + simulation = create_simulation(orm_session) + successful = run_service.create_simulation_run( + orm_session, simulation.id, status="complete" + ) + run_service.create_simulation_run(orm_session, simulation.id) + simulation.active_run_id = None + simulation.latest_successful_run_id = successful.id + + assert run_service.select_display_run(orm_session, simulation) is successful + + +def test_display_run_falls_back_to_newest_for_stale_pointers(orm_session): + simulation = create_simulation(orm_session) + newest = run_service.create_simulation_run(orm_session, simulation.id) + simulation.active_run_id = "missing-active-run" + simulation.latest_successful_run_id = "missing-successful-run" + + assert run_service.select_display_run(orm_session, simulation) is newest diff --git a/tests/unit/services/test_simulation_service.py b/tests/unit/services/test_simulation_service.py index 34116287f..0f8a3e81a 100644 --- a/tests/unit/services/test_simulation_service.py +++ b/tests/unit/services/test_simulation_service.py @@ -1,522 +1,175 @@ -import pytest -import json - -from policyengine_api.services.simulation_service import SimulationService - -from tests.fixtures.services import simulation_fixtures - -pytest_plugins = ("tests.fixtures.services.simulation_fixtures",) - -service = SimulationService() - - -class TestFindExistingSimulation: - """Test finding existing simulations in the database.""" - - def test_find_existing_simulation_given_existing_record( - self, test_db, existing_simulation_record - ): - """Test that find_existing_simulation returns the existing simulation.""" - # GIVEN an existing simulation record (from fixture) - - # WHEN we search for a simulation with matching parameters - result = service.find_existing_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - population_id=simulation_fixtures.valid_simulation_data["population_id"], - population_type=simulation_fixtures.valid_simulation_data[ - "population_type" - ], - policy_id=simulation_fixtures.valid_simulation_data["policy_id"], - ) - - # THEN the result should contain the existing simulation - assert result is not None - assert result["id"] == existing_simulation_record["id"] - assert ( - result["country_id"] - == simulation_fixtures.valid_simulation_data["country_id"] - ) - assert ( - result["population_id"] - == simulation_fixtures.valid_simulation_data["population_id"] - ) - assert ( - result["policy_id"] - == simulation_fixtures.valid_simulation_data["policy_id"] - ) - - def test_find_existing_simulation_given_no_match(self, test_db): - """Test that find_existing_simulation returns None when no match exists.""" - # GIVEN an empty database (default test state) - - # WHEN we search for a non-existent simulation - result = service.find_existing_simulation( - country_id="uk", - population_id="nonexistent_123", - population_type="household", - policy_id=999, - ) - - # THEN the result should be None - assert result is None - - def test_find_existing_simulation_ignores_api_version( - self, test_db, existing_simulation_record - ): - """Test that simulations are found regardless of API version.""" - # GIVEN an existing simulation record - - # WHEN we search for the same simulation (API version is ignored) - result = service.find_existing_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - population_id=simulation_fixtures.valid_simulation_data["population_id"], - population_type=simulation_fixtures.valid_simulation_data[ - "population_type" - ], - policy_id=simulation_fixtures.valid_simulation_data["policy_id"], - ) - - # THEN the existing record should be found (API version ignored) - assert result is not None - assert result["id"] == existing_simulation_record["id"] +import inspect +import pytest +from sqlalchemy import func, select -class TestCreateSimulation: - """Test creating new simulations in the database.""" +from policyengine_api.data.v1_models import Simulation, SimulationRun +from policyengine_api.services.simulation_service import ( + SimulationCreateResult, + SimulationService, +) - def test_create_simulation_success(self, test_db): - """Test successful creation of a new simulation.""" - # GIVEN an empty database - # WHEN we create a new simulation - created_simulation = service.create_simulation( - country_id="us", - population_id="household_123", - population_type="household", - policy_id=1, - ) +@pytest.fixture +def service(orm_session_factory): + return SimulationService(orm_session_factory) - # THEN a valid simulation record should be returned - assert created_simulation is not None - assert isinstance(created_simulation, dict) - assert created_simulation["id"] > 0 - assert created_simulation["country_id"] == "us" - assert created_simulation["population_id"] == "household_123" - - # AND the simulation should be retrievable from database - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert result is not None - assert result["country_id"] == "us" - assert result["population_id"] == "household_123" - - def test_create_simulation_with_geography_type(self, test_db): - """Test creating a simulation with geography population type.""" - # GIVEN an empty database - - # WHEN we create a simulation with geography type - created_simulation = service.create_simulation( - country_id="uk", - population_id="geo_code_456", - population_type="geography", - policy_id=2, - ) - - # THEN the simulation should be created successfully - assert created_simulation is not None - assert created_simulation["population_type"] == "geography" - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert result["population_type"] == "geography" - - def test_create_simulation_retrieves_correct_id(self, test_db): - """Test that create_simulation retrieves the correct ID without race conditions.""" - # GIVEN we create multiple simulations rapidly - - # WHEN we create simulations with different parameters - created_sims = [] - for i in range(3): - sim = service.create_simulation( - country_id="us", - population_id=f"household_{i}", - population_type="household", - policy_id=i, - ) - created_sims.append(sim) - - # THEN all IDs should be unique and sequential - ids = [sim["id"] for sim in created_sims] - assert len(set(ids)) == 3 # All IDs are unique - assert ids == sorted(ids) # IDs are in order - - # AND each simulation should have the correct data - for i, sim in enumerate(created_sims): - result = test_db.query( - "SELECT * FROM simulations WHERE id = ?", (sim["id"],) - ).fetchone() - assert result["population_id"] == f"household_{i}" - assert result["policy_id"] == i - - def test_create_simulation_populates_dual_write_state(self, test_db): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_dual_write", - population_type="household", - policy_id=3, - ) - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["simulation_spec_schema_version"] == 1 - assert stored_simulation["active_run_id"] is not None - assert stored_simulation["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["trigger_type"] == "initial" - snapshot = run["simulation_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["population_id"] == "household_dual_write" - assert snapshot["policy_id"] == 3 - - def test_create_simulation_reuses_existing_row_and_bootstraps_dual_write( - self, test_db +def test_public_simulation_methods_do_not_accept_sessions(): + for method_name in ( + "get_or_create_simulation", + "get_simulation", + "update_simulation", ): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_bootstrap", "household", 7, "pending"), - ) - - created_simulation = service.create_simulation( + parameters = inspect.signature( + getattr(SimulationService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_get_or_create_builds_simulation_spec_and_initial_run(service): + result = service.get_or_create_simulation( + country_id="us", + population_id="household-1", + population_type="household", + policy_id=1, + ) + + assert isinstance(result, SimulationCreateResult) + assert result.created is True + assert isinstance(result.simulation, Simulation) + assert result.simulation.simulation_spec_json == { + "country_id": "us", + "population_id": "household-1", + "population_type": "household", + "policy_id": 1, + } + assert result.simulation.simulation_spec_schema_version == 1 + assert result.simulation.active_run_id is not None + + +def test_get_or_create_reuses_existing_row_and_repairs_dual_write_state( + service, + orm_session_factory, +): + with orm_session_factory.begin() as session: + existing = Simulation( country_id="us", - population_id="household_bootstrap", + api_version="old-version", + population_id="household-1", population_type="household", policy_id=7, - ) - - rows = test_db.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - """, - ("us", "household_bootstrap", "household", 7), - ).fetchall() - assert len(rows) == 1 - assert created_simulation["id"] == rows[0]["id"] - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - - def test_create_simulation_rolls_back_parent_insert_on_dual_write_failure( - self, test_db, monkeypatch - ): - def fail_dual_write(tx, simulation_id, *, country_id=None): - raise RuntimeError("dual write sync failed") - - monkeypatch.setattr( - service, - "_ensure_simulation_dual_write_state_in_transaction", - fail_dual_write, - ) - - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.create_simulation( - country_id="us", - population_id="household_create_rollback", - population_type="household", - policy_id=8, - ) - - rows = test_db.query( - """ - SELECT * FROM simulations - WHERE country_id = ? AND population_id = ? AND population_type = ? AND policy_id = ? - """, - ("us", "household_create_rollback", "household", 8), - ).fetchall() - assert rows == [] - - -class TestGetSimulation: - """Test retrieving simulations from the database.""" - - def test_get_simulation_existing(self, test_db, existing_simulation_record): - """Test retrieving an existing simulation.""" - # GIVEN an existing simulation record - - # WHEN we retrieve the simulation - result = service.get_simulation( - country_id=simulation_fixtures.valid_simulation_data["country_id"], - simulation_id=existing_simulation_record["id"], - ) - - # THEN the correct simulation should be returned - assert result is not None - assert result["id"] == existing_simulation_record["id"] - assert ( - result["country_id"] - == simulation_fixtures.valid_simulation_data["country_id"] - ) - - def test_get_simulation_nonexistent(self, test_db): - """Test retrieving a non-existent simulation returns None.""" - # GIVEN an empty database - - # WHEN we try to retrieve a non-existent simulation - result = service.get_simulation(country_id="us", simulation_id=999) - - # THEN None should be returned - assert result is None - - def test_get_simulation_wrong_country(self, test_db, existing_simulation_record): - """Test that simulations are country-specific.""" - # GIVEN an existing simulation for 'us' - - # WHEN we try to retrieve it with a different country - result = service.get_simulation( - country_id="uk", # Wrong country - simulation_id=existing_simulation_record["id"], - ) - - # THEN None should be returned - assert result is None - - def test_get_simulation_invalid_id(self, test_db): - """Test that invalid simulation IDs are handled properly.""" - # GIVEN any database state - - # WHEN we call get_simulation with invalid ID types - # THEN an exception should be raised - with pytest.raises(Exception) as exc_info: - service.get_simulation(country_id="us", simulation_id=-1) - assert "Invalid simulation ID" in str(exc_info.value) - - with pytest.raises(Exception) as exc_info: - service.get_simulation(country_id="us", simulation_id="not_an_int") - assert "Invalid simulation ID" in str(exc_info.value) - - -class TestUniqueConstraint: - """Test that the unique constraint on simulations works correctly.""" - - def test_duplicate_simulation_returns_existing(self, test_db): - """Test that creating duplicate simulations returns the existing record.""" - # GIVEN we create a simulation - first_simulation = service.create_simulation( + status="pending", + ) + session.add(existing) + session.flush() + simulation_id = existing.id + + result = service.get_or_create_simulation( + country_id="us", + population_id="household-1", + population_type="household", + policy_id=7, + ) + + assert result.created is False + assert result.simulation.id == simulation_id + with orm_session_factory() as session: + stored = session.get(Simulation, simulation_id) + run = session.get(SimulationRun, stored.active_run_id) + assert isinstance(run, SimulationRun) + assert run.simulation_id == simulation_id + + +def test_get_or_create_rolls_back_simulation_when_dual_write_fails( + service, + orm_session_factory, + monkeypatch, +): + def fail_dual_write(*args, **kwargs): + raise RuntimeError("dual write sync failed") + + monkeypatch.setattr(service, "_ensure_simulation_dual_write_state", fail_dual_write) + + with pytest.raises(RuntimeError, match="dual write sync failed"): + service.get_or_create_simulation( country_id="us", - population_id="household_123", + population_id="rollback", population_type="household", - policy_id=1, + policy_id=8, ) - # WHEN we try to create an identical simulation - second_simulation = service.create_simulation( - country_id="us", - population_id="household_123", - population_type="household", - policy_id=1, + with orm_session_factory() as session: + count = session.scalar( + select(func.count()) + .select_from(Simulation) + .where(Simulation.population_id == "rollback") ) + assert count == 0 - # THEN the same simulation should be returned (no duplicate created) - assert first_simulation["id"] == second_simulation["id"] - assert first_simulation["country_id"] == second_simulation["country_id"] - assert first_simulation["population_id"] == second_simulation["population_id"] - assert first_simulation["policy_id"] == second_simulation["policy_id"] +def test_get_simulation_returns_model_scoped_to_country(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation -class TestUpdateSimulation: - def test_update_simulation_updates_dual_write_state(self, test_db): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_update", - population_type="household", - policy_id=11, - ) - output_json = json.dumps({"result": "ok"}) + assert service.get_simulation("us", simulation.id).id == simulation.id + assert service.get_simulation("uk", simulation.id) is None - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=output_json, - ) - assert success is True - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["active_run_id"] is None - assert stored_simulation["latest_successful_run_id"] is not None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run["status"] == "complete" - assert run["output"] == output_json - assert run["id"] == stored_simulation["latest_successful_run_id"] - - def test_update_simulation_bootstraps_missing_run_state(self, test_db): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_legacy", "household", 13, "pending"), - ) - simulation = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() +@pytest.mark.parametrize("simulation_id", [-1, "1", None]) +def test_get_simulation_rejects_invalid_ids(service, simulation_id): + with pytest.raises(Exception, match="Invalid simulation ID"): + service.get_simulation("us", simulation_id) - success = service.update_simulation( - country_id="us", - simulation_id=simulation["id"], - status="error", - error_message="legacy failure", - ) - - assert success is True - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["active_run_id"] is None - assert stored_simulation["latest_successful_run_id"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "error" - assert run["error_message"] == "legacy failure" - - def test_update_simulation_does_not_append_extra_run_for_legacy_patch_traffic( - self, test_db - ): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_single_run", - population_type="household", - policy_id=14, - ) - first_run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert first_run is not None +def test_update_simulation_updates_model_and_run_with_python_json( + service, + orm_session_factory, +): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=json.dumps({"value": 1}), - ) + updated = service.update_simulation( + "us", + simulation.id, + status="complete", + output={"result": 42}, + ) - assert success is True + assert isinstance(updated, Simulation) + assert updated.output == {"result": 42} + assert updated.active_run_id is None + with orm_session_factory() as session: + run = session.get(SimulationRun, updated.latest_successful_run_id) + assert run.output == {"result": 42} + assert run.status == "complete" - runs = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ? ORDER BY run_sequence", - (created_simulation["id"],), - ).fetchall() - assert len(runs) == 1 - assert runs[0]["id"] == first_run["id"] - assert runs[0]["status"] == "complete" - def test_update_simulation_rolls_back_parent_update_on_dual_write_failure( - self, test_db, monkeypatch - ): - created_simulation = service.create_simulation( - country_id="us", - population_id="household_update_rollback", - population_type="household", - policy_id=15, - ) +def test_update_simulation_accepts_legacy_json_text_at_wire_boundary(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - def fail_dual_write(tx, simulation_id, *, country_id=None): - raise RuntimeError("dual write sync failed") + updated = service.update_simulation( + "us", + simulation.id, + output='{"result": 42}', + ) - monkeypatch.setattr( - service, - "_ensure_simulation_dual_write_state_in_transaction", - fail_dual_write, - ) + assert updated.output == {"result": 42} - with pytest.raises(RuntimeError, match="dual write sync failed"): - service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - status="complete", - output=json.dumps({"rolled_back": True}), - ) - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert stored_simulation["status"] == "pending" - assert stored_simulation["output"] is None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (created_simulation["id"],), - ).fetchone() - assert run is not None - assert run["status"] == "pending" - assert run["output"] is None - - def test_update_simulation_with_no_user_fields_returns_false(self, test_db): - """Regression for issue #3449. - - update_fields used to always append api_version, so a PATCH - with no status/output/error_message still passed the - "no fields to update" guard and rewrote the row. The guard - must fire before api_version is appended so an empty PATCH - returns False (and the route converts that to a 400). - """ - created_simulation = service.create_simulation( - country_id="us", - population_id="household_empty_patch", - population_type="household", - policy_id=16, - ) - pre_row = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() +def test_update_simulation_without_values_is_a_noop(service): + simulation = service.get_or_create_simulation( + "us", "household-1", "household", 1 + ).simulation - success = service.update_simulation( - country_id="us", - simulation_id=created_simulation["id"], - ) + assert service.update_simulation("us", simulation.id) is None - assert success is False - post_row = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (created_simulation["id"],), - ).fetchone() - assert post_row["api_version"] == pre_row["api_version"] - assert post_row["status"] == pre_row["status"] +def test_update_missing_simulation_raises(service): + with pytest.raises(LookupError, match="Simulation #999 not found"): + service.update_simulation("us", 999, status="complete") diff --git a/tests/unit/services/test_simulation_spec_service.py b/tests/unit/services/test_simulation_spec_service.py index 89dce070c..26f06d02f 100644 --- a/tests/unit/services/test_simulation_spec_service.py +++ b/tests/unit/services/test_simulation_spec_service.py @@ -6,160 +6,97 @@ SimulationSpecService, ) + simulation_service = SimulationService() -simulation_spec_service = SimulationSpecService() +spec_service = SimulationSpecService() -class TestSimulationSpecService: - def test_builds_simulation_spec_from_row(self, test_db): - simulation = simulation_service.create_simulation( - country_id="uk", - population_id="household_42", - population_type="household", - policy_id=7, - ) +def create_simulation(orm_session): + return simulation_service._create_simulation( + orm_session, "us", "ca", "geography", 3 + ) - simulation_spec = simulation_spec_service.build_simulation_spec(simulation) - assert isinstance(simulation_spec, SimulationSpec) - assert simulation_spec.country_id == "uk" - assert simulation_spec.population_id == "household_42" - assert simulation_spec.policy_id == 7 +def test_builds_spec_from_mapped_simulation(orm_session): + simulation = create_simulation(orm_session) - def test_sets_and_gets_simulation_spec(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ca", - "population_type": "geography", - "policy_id": 3, - } - ) + spec = spec_service.build_simulation_spec(simulation) - result = simulation_spec_service.set_simulation_spec( - simulation["id"], simulation_spec - ) + assert isinstance(spec, SimulationSpec) + assert spec.model_dump() == { + "country_id": "us", + "population_id": "ca", + "population_type": "geography", + "policy_id": 3, + } - assert result is True - stored_simulation = test_db.query( - """ - SELECT simulation_spec_json, simulation_spec_schema_version - FROM simulations WHERE id = ? - """, - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_schema_version"] == 1 - - loaded_simulation_spec = simulation_spec_service.get_simulation_spec( - simulation["id"] - ) - assert loaded_simulation_spec is not None - assert loaded_simulation_spec.model_dump() == simulation_spec.model_dump() - - def test_rejects_unsupported_schema_version_on_write(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ca", - "population_type": "geography", - "policy_id": 3, - } - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.set_simulation_spec( - simulation["id"], - simulation_spec, - schema_version=2, - ) - - assert "Unsupported simulation spec schema version" in str(exc_info.value) - - def test_rejects_unsupported_schema_version_on_read(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - test_db.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - '{"country_id":"us","population_id":"ca","population_type":"geography","policy_id":3}', - 2, - simulation["id"], - ), - ) +def test_sets_and_gets_python_json_spec(orm_session): + simulation = create_simulation(orm_session) + spec = spec_service.build_simulation_spec(simulation) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.get_simulation_spec(simulation["id"]) + assert spec_service.set_simulation_spec(orm_session, simulation.id, spec) is True + loaded = spec_service.get_simulation_spec(orm_session, simulation.id) - assert "Unsupported simulation spec schema version" in str(exc_info.value) + assert simulation.simulation_spec_json == spec.model_dump() + assert loaded == spec - def test_rejects_simulation_spec_write_when_fields_do_not_match_row(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - simulation_spec = SimulationSpec.model_validate( - { - "country_id": "us", - "population_id": "ny", - "population_type": "geography", - "policy_id": 3, - } - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.set_simulation_spec( - simulation["id"], simulation_spec - ) +def test_rejects_unsupported_schema_version_on_write(orm_session): + simulation = create_simulation(orm_session) + spec = spec_service.build_simulation_spec(simulation) - assert "Simulation spec must match the linked simulation row" in str( - exc_info.value + with pytest.raises(ValueError, match="Unsupported simulation spec schema version"): + spec_service.set_simulation_spec( + orm_session, simulation.id, spec, schema_version=2 ) - def test_rejects_inconsistent_stored_simulation_spec_on_read(self, test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="ca", - population_type="geography", - policy_id=3, - ) - test_db.query( - """ - UPDATE simulations - SET simulation_spec_json = ?, simulation_spec_schema_version = ? - WHERE id = ? - """, - ( - '{"country_id":"us","population_id":"ny","population_type":"geography","policy_id":3}', - 1, - simulation["id"], - ), - ) - with pytest.raises(ValueError) as exc_info: - simulation_spec_service.get_simulation_spec(simulation["id"]) +def test_rejects_unsupported_schema_version_on_read(orm_session): + simulation = create_simulation(orm_session) + simulation.simulation_spec_schema_version = 2 - assert "Simulation spec must match the linked simulation row" in str( - exc_info.value - ) + with pytest.raises(ValueError, match="Unsupported simulation spec schema version"): + spec_service.get_simulation_spec(orm_session, simulation.id) + + +def test_rejects_spec_that_does_not_match_simulation(orm_session): + simulation = create_simulation(orm_session) + mismatched = SimulationSpec( + country_id="us", + population_id="ny", + population_type="geography", + policy_id=3, + ) + + with pytest.raises(ValueError, match="must match the linked simulation"): + spec_service.set_simulation_spec(orm_session, simulation.id, mismatched) + + +def test_rejects_inconsistent_stored_spec(orm_session): + simulation = create_simulation(orm_session) + simulation.simulation_spec_json = { + "country_id": "us", + "population_id": "ny", + "population_type": "geography", + "policy_id": 3, + } + + with pytest.raises(ValueError, match="must match the linked simulation"): + spec_service.get_simulation_spec(orm_session, simulation.id) + + +def test_missing_simulation_has_no_spec(orm_session): + assert spec_service.get_simulation_spec(orm_session, 999) is None + + +def test_setting_spec_for_missing_simulation_raises(orm_session): + spec = SimulationSpec( + country_id="us", + population_id="ca", + population_type="geography", + policy_id=3, + ) + + with pytest.raises(ValueError, match="Simulation #999 not found"): + spec_service.set_simulation_spec(orm_session, 999, spec) diff --git a/tests/unit/services/test_tracer_service.py b/tests/unit/services/test_tracer_service.py index 84ece8df3..8f5202fa7 100644 --- a/tests/unit/services/test_tracer_service.py +++ b/tests/unit/services/test_tracer_service.py @@ -1,27 +1,25 @@ import pytest -import json from policyengine_api.services.tracer_analysis_service import ( TracerAnalysisService, ) from werkzeug.exceptions import NotFound -from tests.fixtures.services.tracer_fixture_service import ( - test_tracer_data, - valid_tracer_row, - valid_tracer, -) +from tests.fixtures.services.tracer_fixture_service import valid_tracer -tracer_service = TracerAnalysisService() +pytest_plugins = ["tests.fixtures.services.tracer_fixture_service"] -def test_get_tracer_valid(test_tracer_data): +def test_get_tracer_valid( + test_tracer_data, + orm_session_factory, +): # Test get_tracer successfully retrieves valid data from the database. - result = tracer_service.get_tracer( - test_tracer_data["country_id"], - test_tracer_data["household_id"], - test_tracer_data["policy_id"], - test_tracer_data["api_version"], + result = TracerAnalysisService(orm_session_factory).get_tracer( + test_tracer_data.country_id, + test_tracer_data.household_id, + test_tracer_data.policy_id, + test_tracer_data.api_version, ) # match the valid output as collected from fixture @@ -29,7 +27,7 @@ def test_get_tracer_valid(test_tracer_data): assert result == valid_output -def test_get_tracer_not_found(): +def test_get_tracer_not_found(orm_session_factory): # Test get_tracer raises NotFound when no matching record exists. valid_country_val_in_db = "us" invalid_household_not_in_db = "9999999" @@ -42,10 +40,10 @@ def test_get_tracer_not_found(): invalid_api_version, ] with pytest.raises(NotFound): - tracer_service.get_tracer(*data_not_in_db) + TracerAnalysisService(orm_session_factory).get_tracer(*data_not_in_db) -def test_get_tracer_database_error(test_db): +def test_get_tracer_database_error(orm_session_factory): # Test get_tracer handles database errors properly. missing_country_id = "" valid_householdID = "71424" @@ -58,4 +56,6 @@ def test_get_tracer_database_error(test_db): valid_api_version, ] with pytest.raises(Exception): - tracer_service.get_tracer(*missing_parameter_causing_database_exception) + TracerAnalysisService(orm_session_factory).get_tracer( + *missing_parameter_causing_database_exception, + ) diff --git a/tests/unit/services/test_update_profile_service.py b/tests/unit/services/test_update_profile_service.py deleted file mode 100644 index 5c6016899..000000000 --- a/tests/unit/services/test_update_profile_service.py +++ /dev/null @@ -1,116 +0,0 @@ -import pytest -from policyengine_api.services.user_service import UserService - -from tests.fixtures.services.user_service import ( - valid_user_record, - existing_user_profile, -) - -service = UserService() - - -class TestUpdateProfile: - def test_update_profile_given_existing_record(self, test_db, existing_user_profile): - # GIVEN an existing profile record (from fixture) - - # WHEN we call update_profile with new data - updated_username = "updated_username" - updated_country = "uk" - - result = service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country=updated_country, - username=updated_username, - user_since=existing_user_profile["user_since"], - ) - - # THEN the method should return True for successful update - assert result is True - - # AND the database should be updated with new values - updated_record = test_db.query( - "SELECT * FROM user_profiles WHERE user_id = ?", - (existing_user_profile["user_id"],), - ).fetchone() - - assert updated_record["username"] == updated_username - assert updated_record["primary_country"] == updated_country - - def test_update_profile_given_nonexistent_record(self, test_db): - # GIVEN a nonexistent profile record id - NONEXISTENT_ID = 999 - - # WHEN we call update_profile for this nonexistent record - result = service.update_profile( - user_id=NONEXISTENT_ID, - primary_country="uk", - username="newuser", - user_since="2024-01-01", - ) - - # THEN the result should be False - assert result is False - - def test_update_profile_with_partial_fields(self, test_db, existing_user_profile): - # GIVEN an existing profile record (from fixture) - - # WHEN we call update_profile with only some fields provided - updated_country = "CA" - original_username = existing_user_profile["username"] - - result = service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country=updated_country, - username=None, - user_since=existing_user_profile["user_since"], - ) - - # THEN the method should return True for successful update - assert result is True - - # AND only the provided fields should be updated - updated_record = test_db.query( - "SELECT * FROM user_profiles WHERE user_id = ?", - (existing_user_profile["user_id"],), - ).fetchone() - - assert updated_record["primary_country"] == updated_country - assert ( - updated_record["username"] == original_username - ) # Username should remain unchanged - - def test_update_profile_with_database_error( - self, monkeypatch, existing_user_profile - ): - # GIVEN an existing profile record (from fixture) - - # AND a database that raises an exception - def mock_db_query_error(*args, **kwargs): - raise Exception("Database error") - - monkeypatch.setattr("policyengine_api.data.database.query", mock_db_query_error) - - # WHEN we call update_profile - # THEN an exception should be raised - with pytest.raises(Exception, match="Database error"): - service.update_profile( - user_id=existing_user_profile["user_id"], - primary_country="US", - username="testuser", - user_since="2023-01-01", - ) - - def test_update_profile_id_not_specified(self): - # GIVEN no user_id specified - - # WHEN we call update_profile with None as user_id - # THEN a ValueError should be raised - with pytest.raises( - ValueError, match="you must specify either auth0_id or user_id" - ): - service.update_profile( - user_id=None, - primary_country="US", - username="testuser", - user_since="2023-01-01", - ) diff --git a/tests/unit/services/test_user_policy_service.py b/tests/unit/services/test_user_policy_service.py new file mode 100644 index 000000000..44a96f363 --- /dev/null +++ b/tests/unit/services/test_user_policy_service.py @@ -0,0 +1,95 @@ +import inspect +from pathlib import Path + +from policyengine_api.data.v1_models import UserPolicy +from policyengine_api.services.user_policy_service import ( + UserPolicyCreateResult, + UserPolicyService, +) + + +ROUTE_PATH = ( + Path(__file__).parents[3] / "policyengine_api" / "routes" / "policy_routes.py" +) + + +def _values(**overrides): + values = { + "country_id": "us", + "reform_id": 2, + "reform_label": "Reform", + "baseline_id": 1, + "baseline_label": "Current law", + "user_id": "auth0|one", + "year": "2026", + "geography": "us", + "dataset": "enhanced_cps_2024", + "number_of_provisions": 3, + "api_version": "1", + "added_date": 1, + "updated_date": 2, + "budgetary_impact": None, + "type": None, + } + values.update(overrides) + return values + + +def test_user_policy_public_methods_do_not_accept_sessions(): + for method_name in ( + "create_or_get_user_policy", + "list_user_policies", + "update_user_policy", + ): + parameters = inspect.signature( + getattr(UserPolicyService, method_name) + ).parameters + assert "session" not in parameters + assert "session_factory" not in parameters + + +def test_policy_routes_do_not_manage_sessions_or_queries(): + source = ROUTE_PATH.read_text(encoding="utf-8") + assert "get_v1_session_factory" not in source + assert "from sqlalchemy" not in source + assert "select(" not in source + + +def test_create_reuse_list_and_update_user_policy(orm_session_factory): + service = UserPolicyService(orm_session_factory) + + created = service.create_or_get_user_policy(_values()) + reused = service.create_or_get_user_policy( + _values(number_of_provisions=99, updated_date=99) + ) + listed = service.list_user_policies("us", "auth0|one") + updated = service.update_user_policy( + "us", + created.user_policy.id, + {"reform_label": "Updated", "updated_date": 3}, + ) + + assert isinstance(created, UserPolicyCreateResult) + assert created.created is True + assert reused.created is False + assert reused.user_policy.id == created.user_policy.id + assert len(listed) == 1 + assert isinstance(listed[0], UserPolicy) + assert updated.reform_label == "Updated" + assert updated.updated_date == 3 + + +def test_update_user_policy_requires_matching_country(orm_session_factory): + service = UserPolicyService(orm_session_factory) + created = service.create_or_get_user_policy(_values(country_id="uk")) + + assert ( + service.update_user_policy( + "us", + created.user_policy.id, + {"reform_label": "Wrong country"}, + ) + is None + ) + stored = service.list_user_policies("uk", "auth0|one")[0] + assert stored.reform_label == "Reform" diff --git a/tests/unit/services/test_user_service.py b/tests/unit/services/test_user_service.py index 26875c8ca..f7804fba5 100644 --- a/tests/unit/services/test_user_service.py +++ b/tests/unit/services/test_user_service.py @@ -1,60 +1,89 @@ import pytest + +from policyengine_api.data.v1_models import UserProfile from policyengine_api.services.user_service import UserService +from tests.fixtures.services.user_service import valid_user_record -from tests.fixtures.services.user_service import ( - valid_user_record, - existing_user_profile, -) -service = UserService() +pytest_plugins = ["tests.fixtures.services.user_service"] + + +@pytest.fixture +def service(orm_session_factory): + return UserService(orm_session_factory) + +def test_get_profile_requires_an_identifier(service): + with pytest.raises( + ValueError, + match="you must specify either auth0_id or user_id", + ): + service.get_profile() -class TestGetProfile: - def test_get_profile_id_not_specified(self): - # GIVEN no ID - # WHEN we call get_profile with no auth0_id or user_id - # Then a ValueError should be raised - with pytest.raises( - ValueError, match="you must specify either auth0_id or user_id" - ): - service.get_profile() +def test_get_profile_returns_none_for_unknown_auth0_id(service): + assert service.get_profile(auth0_id="missing") is None - def test_get_profile_nonexistent_record(self): - # GIVEN nonexistent record - INVALID_RECORD_ID = "invalid" - # WHEN we call get_profile with nonexistent user - result = service.get_profile(auth0_id=INVALID_RECORD_ID) +def test_get_profile_returns_mapped_entity_by_either_identifier( + service, + existing_user_profile, +): + by_auth0 = service.get_profile( + auth0_id=valid_user_record["auth0_id"], + ) + by_id = service.get_profile( + user_id=valid_user_record["user_id"], + ) + + assert isinstance(by_auth0, UserProfile) + assert by_auth0.user_id == by_id.user_id + assert by_auth0.username == valid_user_record["username"] + - # THEN result is None - assert result is None +def test_create_profile_returns_existing_entity_for_duplicate_auth0_id(service): + created, profile = service.create_profile( + "us", + "auth0|duplicate", + "first", + 1, + ) + duplicate_created, duplicate = service.create_profile( + "uk", + "auth0|duplicate", + "second", + 2, + ) - def test_get_profile_auth0_id(self, existing_user_profile): - # WHEN we call get_profile with auth0_id - result = service.get_profile(auth0_id=existing_user_profile["auth0_id"]) + assert created is True + assert duplicate_created is False + assert duplicate.user_id == profile.user_id + assert duplicate.username == "first" - # THEN returns record - assert result == existing_user_profile - def test_get_profile_user_id(self, existing_user_profile): - # WHEN we call get_profile with user_id - result = service.get_profile(user_id=existing_user_profile["user_id"]) +def test_update_profile_returns_none_for_missing_entity(service): + assert service.update_profile(999, "uk", "missing", 2) is None - # THEN returns record - assert result == existing_user_profile - def test_get_profile_id_priority(self, test_db, existing_user_profile): - # WHEN we call get_profile with auth0_id and user_id - result = service.get_profile( - auth0_id=existing_user_profile["auth0_id"], - user_id=existing_user_profile["user_id"], - ) +def test_update_profile_only_changes_non_null_fields( + service, + existing_user_profile, +): + profile = service.update_profile( + valid_user_record["user_id"], + "uk", + None, + valid_user_record["user_since"] + 1, + ) + + assert profile.primary_country == "uk" + assert profile.username == valid_user_record["username"] + assert profile.user_since == valid_user_record["user_since"] + 1 - # THEN returns record using auth0_id - record = test_db.query( - "SELECT * FROM user_profiles WHERE auth0_id = ?", - (valid_user_record["auth0_id"],), - ).fetchone() - assert result == record +def test_update_profile_requires_user_id(service): + with pytest.raises( + ValueError, + match="you must specify either auth0_id or user_id", + ): + service.update_profile(None, "us", "name", 1) diff --git a/tests/unit/test_alembic_skill.py b/tests/unit/test_alembic_skill.py new file mode 100644 index 000000000..af743778c --- /dev/null +++ b/tests/unit/test_alembic_skill.py @@ -0,0 +1,26 @@ +from policyengine_api.constants import REPO + + +SKILL = REPO / "docs" / "engineering" / "skills" / "alembic-migrations.md" + + +def test_model_agnostic_alembic_skill_is_discoverable(): + assert SKILL.exists() + assert ( + "alembic-migrations.md" + in (REPO / "docs" / "engineering" / "skills" / "README.md").read_text() + ) + for adapter in ("AGENTS.md", "CLAUDE.md", ".github/copilot-instructions.md"): + assert ( + "docs/engineering/skills/alembic-migrations.md" + in (REPO / adapter).read_text() + ) + + +def test_alembic_skill_forbids_handwritten_ai_revisions(): + guidance = SKILL.read_text() + assert "MUST NOT manually author Alembic revision scripts" in guidance + assert "alembic revision --autogenerate" in guidance + assert "dialect compatibility" in guidance + assert "reversibility" in guidance + assert "request a human migration decision" in guidance diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py new file mode 100644 index 000000000..deab1d7be --- /dev/null +++ b/tests/unit/test_alembic_workflows.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import os +from pathlib import Path +import subprocess + +import pytest + + +REPO = Path(__file__).resolve().parents[2] + + +def _workflow(name: str) -> str: + return (REPO / ".github" / "workflows" / name).read_text(encoding="utf-8") + + +def _long_inline_run_blocks() -> list[str]: + offenders = [] + block_markers = {"|", "|-", "|+", ">", ">-", ">+"} + for path in sorted((REPO / ".github" / "workflows").glob("*.y*ml")): + lines = path.read_text(encoding="utf-8").splitlines() + line_index = 0 + while line_index < len(lines): + line = lines[line_index] + stripped = line.lstrip() + indentation = len(line) - len(stripped) + if not ( + stripped.startswith("run:") + and stripped.removeprefix("run:").strip() in block_markers + ): + line_index += 1 + continue + + body_start = line_index + 1 + line_index = body_start + while line_index < len(lines): + candidate = lines[line_index] + candidate_indentation = len(candidate) - len(candidate.lstrip()) + if candidate.strip() and candidate_indentation <= indentation: + break + line_index += 1 + + substantive_lines = [ + candidate + for candidate in lines[body_start:line_index] + if candidate.strip() and not candidate.lstrip().startswith("#") + ] + if len(substantive_lines) > 3: + offenders.append(f"{path.name}:{body_start}") + return offenders + + +def test_workflows_do_not_inline_long_shell_programs(): + assert _long_inline_run_blocks() == [] + + +def test_pr_always_runs_reusable_alembic_check(): + workflow = _workflow("pr.yml") + + assert "alembic-v1-check:" in workflow + assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow + assert "detect-v1-alembic-changes:" not in workflow + assert "needs.detect-v1-alembic-changes" not in workflow + + +def test_push_always_runs_lint_and_alembic_qualification_before_versioning(): + workflow = _workflow("push.yml") + + assert "lint:" in workflow + assert "alembic-v1-check:" in workflow + assert "uses: ./.github/workflows/alembic-v1-check.yml" in workflow + assert "needs: [lint, alembic-v1-check]" in workflow + assert "github.repository == 'PolicyEngine/policyengine-uk'" not in workflow + + +def test_release_migration_uses_the_installed_python_environment(): + workflow = _workflow("push.yml") + migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] + migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") + + assert "bash .github/scripts/migrate_v1_cloud_sql.sh" in migration_job + assert "python scripts/v1_database_migration.py" in orchestration_script + assert "uv run" not in orchestration_script + + +def test_reusable_alembic_check_uses_only_disposable_mysql(): + workflow = _workflow("alembic-v1-check.yml") + + assert "workflow_call:" in workflow + assert "workflow_dispatch:" in workflow + assert "mysql:8.4" in workflow + assert "policyengine_alembic_test" in workflow + assert "alembic-v1.ini" in workflow + assert "STAGE7_EXISTING_DATABASE_URL" not in workflow + assert "POLICYENGINE_DB_MIGRATION_PASSWORD" not in workflow + + +def test_reusable_alembic_check_uses_the_installed_python_environment(): + workflow = _workflow("alembic-v1-check.yml") + + assert "python -m pytest" in workflow + assert "python -m alembic" in workflow + assert "uv run" not in workflow + + +def test_release_migration_fails_closed_and_gates_both_staging_deploys(): + workflow = _workflow("push.yml") + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") + + assert "migrate-v1-cloud-sql:" in workflow + assert "environment: production-database" in workflow + assert "--mode state" in orchestration_script + assert "--mode upgrade" in orchestration_script + assert "--mode verify-head" in orchestration_script + assert "database is unversioned" in orchestration_script + assert "create_cloud_sql_backup.sh" in orchestration_script + + app_engine_job = workflow[workflow.index(" deploy-staging:") :] + app_engine_job = app_engine_job[ + : app_engine_job.index("\n deploy-cloud-run-staging:") + ] + assert "migrate-v1-cloud-sql" in app_engine_job + + cloud_run_job = workflow[workflow.index(" deploy-cloud-run-staging:") :] + cloud_run_job = cloud_run_job[ + : cloud_run_job.index("\n integration-tests-staging:") + ] + assert "migrate-v1-cloud-sql" in cloud_run_job + + +def test_cloud_sql_workflow_uses_oidc_and_separate_database_credentials(): + workflow = _workflow("push.yml") + migration_job = workflow[workflow.index(" migrate-v1-cloud-sql:") :] + migration_job = migration_job[: migration_job.index("\n deploy-staging:")] + orchestration_script = ( + REPO / ".github" / "scripts" / "migrate_v1_cloud_sql.sh" + ).read_text(encoding="utf-8") + + assert "google-github-actions/auth@v2" in migration_job + assert "GCP_DB_MIGRATION_SERVICE_ACCOUNT" in migration_job + assert "migrate_v1_cloud_sql.sh" in migration_job + assert "prepare_v1_database_urls.sh" not in migration_job + assert "GITHUB_ENV" not in orchestration_script + assert "secrets.POLICYENGINE_DB_READONLY_PASSWORD" not in migration_job + assert "secrets.POLICYENGINE_DB_MIGRATION_PASSWORD" not in migration_job + assert ( + "POLICYENGINE_DB_PASSWORD: ${{ secrets.POLICYENGINE_DB_PASSWORD }}" + not in migration_job + ) + + +def _write_fake_migration_commands(tmp_path: Path) -> tuple[Path, Path]: + bin_path = tmp_path / "bin" + bin_path.mkdir() + gcloud_path = bin_path / "gcloud" + gcloud_path.write_text( + "#!/usr/bin/env bash\n" + 'printf "%s\\n" "$*" >> "${GCLOUD_CALLS}"\n' + 'if [[ "$*" == *"readonly-password"* ]]; then\n' + ' printf "reader-p@ss\\n"\n' + 'elif [[ "$*" == *"migration-password"* ]]; then\n' + ' printf "migrator-p@ss\\n"\n' + 'elif [[ "$*" == *"sql backups list"* ]]; then\n' + ' printf "backup-123\\n"\n' + "fi\n", + encoding="utf-8", + ) + gcloud_path.chmod(0o755) + + python_path = bin_path / "python" + python_path.write_text( + "#!/usr/bin/env bash\n" + ': "${POLICYENGINE_DB_READONLY_PASSWORD:?}"\n' + ': "${POLICYENGINE_DB_MIGRATION_PASSWORD:?}"\n' + 'printf "%s\\n" "$*" >> "${PYTHON_CALLS}"\n' + 'if [[ "$*" == *"--mode state"* ]]; then\n' + ' printf "%s\\n" "${DATABASE_STATE}"\n' + "fi\n", + encoding="utf-8", + ) + python_path.chmod(0o755) + return bin_path, gcloud_path + + +def _run_migration_orchestrator(tmp_path: Path, database_state: str): + bin_path, _ = _write_fake_migration_commands(tmp_path) + python_calls = tmp_path / "python-calls" + gcloud_calls = tmp_path / "gcloud-calls" + result = subprocess.run( + ["bash", ".github/scripts/migrate_v1_cloud_sql.sh"], + cwd=REPO, + env={ + **os.environ, + "DATABASE_STATE": database_state, + "GCLOUD_CALLS": str(gcloud_calls), + "PYTHON_CALLS": str(python_calls), + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": "project:region:instance", + "PATH": f"{bin_path}:{os.environ['PATH']}", + }, + capture_output=True, + text=True, + check=False, + ) + return result, python_calls, gcloud_calls + + +def test_migration_orchestrator_keeps_credentials_local_and_upgrades_pending_schema( + tmp_path, +): + result, python_calls, gcloud_calls = _run_migration_orchestrator( + tmp_path, "pending" + ) + + assert result.returncode == 0, result.stderr + assert "::add-mask::reader-p@ss" in result.stdout + assert "::add-mask::migrator-p@ss" in result.stdout + assert "STAGE7_EXISTING_DATABASE_URL" not in result.stdout + assert "ALEMBIC_DATABASE_URL" not in result.stdout + calls = python_calls.read_text(encoding="utf-8").splitlines() + assert calls == [ + "scripts/v1_database_migration.py --mode state", + "scripts/v1_database_migration.py --mode upgrade --backup-id backup-123", + "scripts/v1_database_migration.py --mode verify-head", + ] + assert "sql backups create" in gcloud_calls.read_text(encoding="utf-8") + + +def test_migration_orchestrator_skips_backup_and_upgrade_at_head(tmp_path): + result, python_calls, gcloud_calls = _run_migration_orchestrator(tmp_path, "head") + + assert result.returncode == 0, result.stderr + assert python_calls.read_text(encoding="utf-8").splitlines() == [ + "scripts/v1_database_migration.py --mode state", + "scripts/v1_database_migration.py --mode verify-head", + ] + assert "sql backups create" not in gcloud_calls.read_text(encoding="utf-8") + + +@pytest.mark.parametrize("database_state", ["unversioned", "invalid"]) +def test_migration_orchestrator_refuses_unsafe_database_states( + tmp_path, + database_state, +): + result, python_calls, gcloud_calls = _run_migration_orchestrator( + tmp_path, database_state + ) + + assert result.returncode != 0 + assert "automatic baseline stamping is disabled" in result.stderr + assert python_calls.read_text(encoding="utf-8").splitlines() == [ + "scripts/v1_database_migration.py --mode state" + ] + assert "sql backups create" not in gcloud_calls.read_text(encoding="utf-8") + + +def test_backup_helper_recovers_and_verifies_the_created_backup_id(): + script = (REPO / ".github" / "scripts" / "create_cloud_sql_backup.sh").read_text( + encoding="utf-8" + ) + + assert "GITHUB_RUN_ID" in script + assert "gcloud sql backups create" in script + assert "gcloud sql backups list" in script + assert "status=SUCCESSFUL" in script + assert "GITHUB_OUTPUT" not in script + assert script.index("gcloud sql backups create") < script.index( + "gcloud sql backups list" + ) + + +def test_v1_and_future_v2_alembic_domains_are_explicitly_separate(): + assert (REPO / "alembic-v1.ini").exists() + assert (REPO / "migrations" / "v1" / "env.py").exists() + + guidance = REPO / "docs" / "engineering" / "skills" / "alembic-migrations.md" + text = guidance.read_text(encoding="utf-8") + assert "alembic-v1.ini" in text + assert "migrations/v1" in text + assert "Supabase/Postgres" in text + assert "separate revision chain" in text diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index 762c2925a..92d2649ce 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -4,7 +4,7 @@ import threading from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import policyengine_api.asgi_factory as asgi_factory @@ -409,6 +409,19 @@ def test_fastapi_documentation_routes_fall_through_to_flask_404(): assert "swagger" not in response.text.lower() +def test_lifespan_closes_runtime_resources_on_shutdown(): + close_runtime_resources = Mock() + with TestClient( + create_asgi_app( + create_test_wsgi_app(), + shutdown_callback=close_runtime_resources, + ) + ) as client: + assert client.get("/liveness-check").status_code == 200 + close_runtime_resources.assert_not_called() + close_runtime_resources.assert_called_once_with() + + def test_flask_fallback_preserves_status_body_headers_and_cookies(): client = TestClient(create_asgi_app(create_test_wsgi_app())) diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index 68dfcb4c3..d891c15bd 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -65,6 +65,7 @@ def _gateway_auth_env() -> dict[str, str]: def _required_runtime_env() -> dict[str, str]: return { + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": PRODUCTION_CLOUD_SQL_INSTANCE, "POLICYENGINE_DB_PASSWORD": "raw-db-secret-value", "POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN": ("raw-github-secret-value"), "ANTHROPIC_API_KEY": "raw-anthropic-secret-value", @@ -483,6 +484,35 @@ def test_cloud_run_startup_supervises_redis_and_server_children(): assert re.search(r"(?m)^ *wait 2>/dev/null", start_script) is None +def test_production_gunicorn_workers_do_not_inherit_database_pools(): + for relative_path in ("gcp/cloud_run/start.sh", "gcp/policyengine_api/start.sh"): + start_script = (REPO / relative_path).read_text(encoding="utf-8") + commands = "\n".join( + line + for line in start_script.splitlines() + if not line.lstrip().startswith("#") + ) + assert "--preload" not in commands + + +def test_app_engine_startup_allows_all_workers_to_finish_booting(): + start_script = (REPO / "gcp/policyengine_api/start.sh").read_text(encoding="utf-8") + app_config = (REPO / "gcp/policyengine_api/app.yaml").read_text(encoding="utf-8") + + assert "--timeout 900" in start_script + assert "--workers 5" in start_script + assert "initial_delay_sec: 1800" in app_config + assert "app_start_timeout_sec: 1800" in app_config + + +def test_app_engine_deploy_health_checks_allow_full_startup_window(): + workflow = _push_workflow() + + for job_name in ("deploy-staging", "deploy-production-candidate"): + job = _workflow_job_block(workflow, job_name) + assert 'HEALTH_CHECK_TIMEOUT_SECONDS: "1800"' in job + + def test_validate_cloud_run_deploy_env_requires_selector_environment_variable(): result = _run_script( ".github/scripts/validate_cloud_run_deploy_env.sh", @@ -505,6 +535,7 @@ def test_validate_cloud_run_deploy_env_accepts_direct_mode_from_environment(): ROUTE_IMPL_HEALTH="fastapi_native", ROUTE_IMPL_SPECIFICATION="fastapi_native", ROUTE_IMPL_METADATA="fastapi_native", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ), ) @@ -584,6 +615,7 @@ def test_validate_cloud_run_deploy_env_requires_only_selected_url( ROUTE_IMPL_HEALTH="fastapi_native", ROUTE_IMPL_SPECIFICATION="fastapi_native", ROUTE_IMPL_METADATA="fastapi_native", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ) missing_result = _run_script( @@ -619,6 +651,7 @@ def test_validate_app_engine_deploy_env_accepts_direct_mode_from_environment(): _script_env( SIM_ENTRYPOINT="old_gateway_direct", OLD_SIMULATION_GATEWAY_URL="https://old-gateway.example.test", + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ), ) @@ -648,6 +681,7 @@ def test_validate_app_engine_deploy_env_requires_only_selected_url( ): env = _script_env( SIM_ENTRYPOINT=entrypoint, + POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=PRODUCTION_CLOUD_SQL_INSTANCE, **_gateway_auth_env(), ) missing_result = _run_script( @@ -664,15 +698,41 @@ def test_validate_app_engine_deploy_env_requires_only_selected_url( assert valid_result.returncode == 0, valid_result.stderr -def test_app_engine_image_contains_simulation_routing_environment_placeholders(): +def test_app_engine_bundle_contains_runtime_environment_placeholders(): dockerfile = (REPO / "gcp/policyengine_api/Dockerfile").read_text(encoding="utf-8") + app_config = (REPO / "gcp/policyengine_api/app.yaml").read_text(encoding="utf-8") export_script = (REPO / "gcp/export.py").read_text(encoding="utf-8") assert 'ENV SIMULATION_ENTRYPOINT_URL=".simulation_entrypoint_url"' in dockerfile assert 'ENV OLD_SIMULATION_GATEWAY_URL=".old_simulation_gateway_url"' in dockerfile assert 'ENV SIM_ENTRYPOINT=".sim_entrypoint"' in dockerfile + assert ( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + '".policyengine_db_instance_connection_name"' in app_config + ) assert '".old_simulation_gateway_url", OLD_SIMULATION_GATEWAY_URL' in export_script assert '".sim_entrypoint", SIM_ENTRYPOINT' in export_script + assert '".policyengine_db_instance_connection_name",' in export_script + assert "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME," in export_script + + +@pytest.mark.parametrize( + "validation_script", + [ + ".github/scripts/validate_app_engine_deploy_env.sh", + ".github/scripts/validate_cloud_run_deploy_env.sh", + ], +) +def test_deployment_validation_requires_database_instance_connection_name( + validation_script, +): + env = _script_env(**_required_runtime_env()) + env.pop("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME") + + result = _run_script(validation_script, env) + + assert result.returncode == 1 + assert "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" in result.stderr @pytest.mark.parametrize( @@ -701,6 +761,10 @@ def test_app_engine_export_requires_only_selected_url( ): (tmp_path / "gcp/policyengine_api").mkdir(parents=True) shutil.copy2(REPO / "gcp/export.py", tmp_path / "gcp/export.py") + shutil.copy2( + REPO / "gcp/policyengine_api/app.yaml", + tmp_path / "gcp/policyengine_api/app.yaml", + ) shutil.copy2( REPO / "gcp/policyengine_api/Dockerfile", tmp_path / "gcp/policyengine_api/Dockerfile", @@ -726,8 +790,15 @@ def test_app_engine_export_requires_only_selected_url( rendered = (tmp_path / "gcp/policyengine_api/Dockerfile").read_text( encoding="utf-8" ) + rendered_app_config = (tmp_path / "gcp/policyengine_api/app.yaml").read_text( + encoding="utf-8" + ) assert f'ENV {selected_url_env}="{selected_url}"' in rendered assert f'ENV {unselected_url_env}=""' in rendered + assert ( + f"POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + f'"{PRODUCTION_CLOUD_SQL_INSTANCE}"' in rendered_app_config + ) def test_build_cloud_run_image_dry_run_uses_cloud_run_dockerfile(): @@ -796,6 +867,30 @@ def test_deploy_cloud_run_candidate_dry_run_never_shifts_traffic(): assert result.stdout.count(f"{selector}=fastapi_native") == 1 +def test_deploy_cloud_run_candidate_uses_configured_database_instance(): + configured_instance = "project:region:configured-instance" + env = { + **_required_runtime_env(), + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME": configured_instance, + } + result = _run_script( + ".github/scripts/deploy_cloud_run_candidate.sh", + _script_env( + **env, + CLOUD_RUN_IMAGE_URI="us-central1-docker.pkg.dev/project/repo/api:sha", + CLOUD_RUN_TAG="stage3-test", + ), + ) + + assert result.returncode == 0, result.stderr + assert f"--add-cloudsql-instances {configured_instance}" in result.stdout + assert ( + f"POLICYENGINE_DB_INSTANCE_CONNECTION_NAME={configured_instance}" + in result.stdout + ) + assert PRODUCTION_CLOUD_SQL_INSTANCE not in result.stdout + + @pytest.mark.parametrize( ("entrypoint", "selected_url_env", "selected_url", "unselected_url_env"), [ @@ -1472,6 +1567,22 @@ def test_cloud_run_deploy_jobs_use_environment_scoped_stage6_route_selectors(): assert f"{selector}: ${{{{ vars.{selector} }}}}" in job +def test_all_deploy_jobs_use_github_database_instance_variable(): + workflow = _push_workflow() + instance_env = ( + "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME: " + "${{ vars.POLICYENGINE_DB_INSTANCE_CONNECTION_NAME }}" + ) + + for job_name in ( + "deploy-staging", + "deploy-cloud-run-staging", + "deploy-production-candidate", + "deploy-cloud-run-candidate", + ): + assert instance_env in _workflow_job_block(workflow, job_name) + + def test_deployment_consumers_require_selector_from_environment(): consumers = ( ".github/request-simulation-model-versions.sh", diff --git a/tests/unit/test_response_factory.py b/tests/unit/test_response_factory.py new file mode 100644 index 000000000..b8a27c9e5 --- /dev/null +++ b/tests/unit/test_response_factory.py @@ -0,0 +1,44 @@ +from policyengine_api.response_factory import _make_error_response + + +def test_error_response_factory_builds_standard_json_response(): + response = _make_error_response("Policy #2 not found.", 404) + + assert response.status_code == 404 + assert response.mimetype == "application/json" + assert response.get_json() == { + "status": "error", + "message": "Policy #2 not found.", + } + + +def test_error_response_factory_includes_structured_payload_fields(): + response = _make_error_response( + "Invalid inputs", + 400, + result=None, + errors=[{"name": "unknown_variable"}], + ) + + assert response.get_json() == { + "status": "error", + "message": "Invalid inputs", + "result": None, + "errors": [{"name": "unknown_variable"}], + } + + +def test_error_response_factory_can_preserve_message_only_v1_payloads(): + response = _make_error_response( + "Database error", + 500, + include_status=False, + ) + + assert response.get_json() == {"message": "Database error"} + + +def test_error_response_factory_can_preserve_a_legacy_default_mimetype(): + response = _make_error_response("Invalid country", 400, mimetype=None) + + assert response.content_type == "text/html; charset=utf-8" diff --git a/tests/unit/test_simulation_and_report_routes.py b/tests/unit/test_simulation_and_report_routes.py new file mode 100644 index 000000000..2d2502f7c --- /dev/null +++ b/tests/unit/test_simulation_and_report_routes.py @@ -0,0 +1,282 @@ +"""Route behavior for simulation and report lifecycle operations.""" + +import json +from datetime import datetime + +from flask import Flask + +from policyengine_api.constants import get_report_output_cache_version +from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun, Simulation +from policyengine_api.routes.report_output_routes import report_output_bp +from policyengine_api.routes.simulation_routes import simulation_bp +from policyengine_api.services.report_output_service import ReportOutputService +from policyengine_api.services.report_run_service import ReportRunService +from policyengine_api.services.simulation_service import SimulationService + + +simulation_service = SimulationService() +report_service = ReportOutputService() +run_service = ReportRunService() + + +def create_test_client() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(simulation_bp) + app.register_blueprint(report_output_bp) + return app.test_client() + + +def create_simulation(factory, *, population_id="household-1", policy_id=1): + simulation = ( + SimulationService(factory) + .get_or_create_simulation("us", population_id, "household", policy_id) + .simulation + ) + return simulation.id + + +def create_report(factory, simulation_id): + report = ( + ReportOutputService(factory) + .create_or_reuse_report_output("us", simulation_id, year="2025") + .view.report_output + ) + return report.id + + +def test_create_simulation_existing_row_repairs_dual_write_state( + orm_session_factory, +): + with orm_session_factory.begin() as session: + simulation = Simulation( + country_id="us", + api_version="old", + population_id="household-route-repair", + population_type="household", + policy_id=40, + status="pending", + ) + session.add(simulation) + session.flush() + simulation_id = simulation.id + + response = create_test_client().post( + "/us/simulation", + json={ + "population_id": "household-route-repair", + "population_type": "household", + "policy_id": 40, + }, + ) + + assert response.status_code == 200 + assert response.get_json()["result"]["id"] == simulation_id + with orm_session_factory() as session: + simulation = session.get(Simulation, simulation_id) + assert simulation.simulation_spec_json is not None + assert simulation.active_run_id is not None + + +def test_create_report_existing_row_repairs_dual_write_state(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=41) + with orm_session_factory.begin() as session: + report = ReportOutput( + country_id="us", + simulation_1_id=simulation_id, + simulation_2_id=None, + api_version=get_report_output_cache_version("us"), + status="pending", + year="2025", + ) + session.add(report) + session.flush() + report_id = report.id + + response = create_test_client().post( + "/us/report", + json={"simulation_1_id": simulation_id, "year": "2025"}, + ) + + assert response.status_code == 200 + assert response.get_json()["result"]["id"] == report_id + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + assert report.report_spec_json is not None + assert report.active_run_id is not None + + +def test_report_post_returns_run_timestamps_for_new_and_existing( + orm_session_factory, +): + simulation_id = create_simulation(orm_session_factory, policy_id=42) + client = create_test_client() + + created = client.post( + "/us/report", json={"simulation_1_id": simulation_id, "year": "2025"} + ) + existing = client.post( + "/us/report", json={"simulation_1_id": simulation_id, "year": "2025"} + ) + + assert created.status_code == 201 + assert existing.status_code == 200 + for response in (created, existing): + result = response.get_json()["result"] + assert result["requested_at"] is not None + assert result["started_at"] is None + assert result["finished_at"] is None + + +def test_report_post_rejects_missing_linked_simulations(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=43) + client = create_test_client() + + missing_primary = client.post( + "/us/report", json={"simulation_1_id": 999999, "year": "2025"} + ) + missing_secondary = client.post( + "/us/report", + json={ + "simulation_1_id": simulation_id, + "simulation_2_id": 999999, + "year": "2025", + }, + ) + + assert missing_primary.status_code == 400 + assert missing_secondary.status_code == 400 + + +def test_simulation_routes_scope_reads_and_writes_to_country(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=44) + client = create_test_client() + + assert client.get(f"/uk/simulation/{simulation_id}").status_code == 404 + response = client.patch( + "/uk/simulation", + json={"id": simulation_id, "status": "complete", "output": {"bad": True}}, + ) + + assert response.status_code == 404 + with orm_session_factory() as session: + assert session.get(Simulation, simulation_id).status == "pending" + + +def test_report_routes_scope_reads_and_writes_to_country(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=45) + report_id = create_report(orm_session_factory, simulation_id) + client = create_test_client() + + assert client.get(f"/uk/report/{report_id}").status_code == 404 + response = client.patch( + "/uk/report", + json={"id": report_id, "status": "complete", "output": {"bad": True}}, + ) + + assert response.status_code == 404 + with orm_session_factory() as session: + assert session.get(ReportOutput, report_id).status == "pending" + + +def test_report_get_serializes_display_run_timestamps(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=46) + report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"ok": True} + ) + with orm_session_factory.begin() as session: + report = session.get(ReportOutput, report_id) + run = session.get(ReportOutputRun, report.latest_successful_run_id) + run.requested_at = datetime(2026, 5, 4, 12, 0) + run.started_at = datetime(2026, 5, 4, 12, 1) + run.finished_at = datetime(2026, 5, 4, 12, 2) + + result = create_test_client().get(f"/us/report/{report_id}").get_json()["result"] + + assert result["requested_at"] == "2026-05-04T12:00:00Z" + assert result["started_at"] == "2026-05-04T12:01:00Z" + assert result["finished_at"] == "2026-05-04T12:02:00Z" + + +def test_report_patch_updates_active_rerun_and_preserves_success( + orm_session_factory, +): + simulation_id = create_simulation(orm_session_factory, policy_id=47) + report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"old": True} + ) + with orm_session_factory.begin() as session: + report = session.get(ReportOutput, report_id) + successful_id = report.latest_successful_run_id + rerun = run_service.create_report_output_run( + session, report_id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + rerun_id = rerun.id + + running = create_test_client().patch( + "/us/report", json={"id": report_id, "status": "running"} + ) + failed = create_test_client().patch( + "/us/report", + json={"id": report_id, "status": "error", "error_message": "failed"}, + ) + + assert running.status_code == 200 + assert running.get_json()["result"]["started_at"] is not None + assert failed.status_code == 200 + assert failed.get_json()["result"]["finished_at"] is not None + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + rerun = session.get(ReportOutputRun, rerun_id) + assert report.latest_successful_run_id == successful_id + assert rerun.status == "error" + + +def test_report_patch_complete_promotes_active_rerun(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=48) + report_id = create_report(orm_session_factory, simulation_id) + ReportOutputService(orm_session_factory).update_report_output( + "us", report_id, status="complete", output={"old": True} + ) + with orm_session_factory.begin() as session: + report = session.get(ReportOutput, report_id) + rerun = run_service.create_report_output_run( + session, report_id, trigger_type="rerun" + ) + report.active_run_id = rerun.id + rerun_id = rerun.id + + response = create_test_client().patch( + "/us/report", + json={"id": report_id, "status": "complete", "output": {"new": True}}, + ) + + assert response.status_code == 200 + with orm_session_factory() as session: + report = session.get(ReportOutput, report_id) + assert report.active_run_id is None + assert report.latest_successful_run_id == rerun_id + + +def test_simulation_v1_routes_keep_json_fields_as_strings(orm_session_factory): + simulation_id = create_simulation(orm_session_factory, policy_id=49) + output = {"result": "ok", "values": [1, 2, 3]} + + patched = create_test_client().patch( + "/us/simulation", + json={"id": simulation_id, "status": "complete", "output": output}, + ) + fetched = create_test_client().get(f"/us/simulation/{simulation_id}") + + for response in (patched, fetched): + result = response.get_json()["result"] + assert isinstance(result["output"], str) + assert json.loads(result["output"]) == output + assert isinstance(result["simulation_spec_json"], str) + with orm_session_factory() as session: + simulation = session.get(Simulation, simulation_id) + assert simulation.output == output + assert isinstance(simulation.simulation_spec_json, dict) diff --git a/tests/unit/test_stage5_routes.py b/tests/unit/test_stage5_routes.py deleted file mode 100644 index ec9f34e1b..000000000 --- a/tests/unit/test_stage5_routes.py +++ /dev/null @@ -1,582 +0,0 @@ -import json - -from flask import Flask - -from policyengine_api.constants import get_report_output_cache_version -from policyengine_api.routes.report_output_routes import report_output_bp -from policyengine_api.routes.simulation_routes import simulation_bp -from policyengine_api.services.report_output_service import ReportOutputService -from policyengine_api.services.report_run_service import ReportRunService -from policyengine_api.services.simulation_service import SimulationService - - -simulation_service = SimulationService() -report_output_service = ReportOutputService() -report_run_service = ReportRunService() - - -def create_test_client() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(simulation_bp) - app.register_blueprint(report_output_bp) - return app.test_client() - - -def test_create_simulation_existing_row_repairs_dual_write_state(test_db): - test_db.query( - """INSERT INTO simulations - (country_id, api_version, population_id, population_type, policy_id, status) - VALUES (?, ?, ?, ?, ?, ?)""", - ("us", "us-system-1.0.0", "household_route_repair", "household", 40, "pending"), - ) - simulation = test_db.query( - "SELECT * FROM simulations ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() - response = client.post( - "/us/simulation", - json={ - "population_id": "household_route_repair", - "population_type": "household", - "policy_id": 40, - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["message"] == "Simulation already exists" - assert payload["result"]["id"] == simulation["id"] - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["simulation_spec_json"] is not None - assert stored_simulation["active_run_id"] is not None - - run = test_db.query( - "SELECT * FROM simulation_runs WHERE simulation_id = ?", - (simulation["id"],), - ).fetchone() - assert run is not None - - -def test_create_report_output_existing_row_repairs_dual_write_state(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_report", - population_type="household", - policy_id=41, - ) - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ( - "us", - simulation["id"], - None, - get_report_output_cache_version("us"), - "pending", - "2025", - ), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["message"] == "Report output already exists" - assert payload["result"]["id"] == report_output["id"] - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["report_spec_json"] is not None - assert stored_report["active_run_id"] is not None - - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report_output["id"],), - ).fetchone() - assert run is not None - snapshot = run["report_spec_snapshot_json"] - if isinstance(snapshot, str): - snapshot = json.loads(snapshot) - assert snapshot["report_kind"] == "household_single" - - -def test_post_report_output_returns_timestamp_fields_for_new_and_existing_report( - test_db, -): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_report_timestamps", - population_type="household", - policy_id=46, - ) - - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, - ) - - assert response.status_code == 201 - payload = response.get_json() - created_report = payload["result"] - assert created_report["requested_at"] is not None - assert created_report["started_at"] is None - assert created_report["finished_at"] is None - - existing_response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": None, - "year": "2025", - }, - ) - - assert existing_response.status_code == 200 - existing_payload = existing_response.get_json() - existing_report = existing_payload["result"] - assert existing_report["id"] == created_report["id"] - assert existing_report["requested_at"] is not None - assert existing_report["started_at"] is None - assert existing_report["finished_at"] is None - - -def test_create_report_output_missing_primary_simulation_returns_bad_request(test_db): - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": 999999, - "simulation_2_id": None, - "year": "2025", - }, - ) - - assert response.status_code == 400 - - report_rows = test_db.query("SELECT * FROM report_outputs").fetchall() - report_run_rows = test_db.query("SELECT * FROM report_output_runs").fetchall() - assert report_rows == [] - assert report_run_rows == [] - - -def test_create_report_output_missing_secondary_simulation_returns_bad_request(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_missing_secondary", - population_type="household", - policy_id=42, - ) - - client = create_test_client() - response = client.post( - "/us/report", - json={ - "simulation_1_id": simulation["id"], - "simulation_2_id": simulation["id"] + 999999, - "year": "2025", - }, - ) - - assert response.status_code == 400 - - report_rows = test_db.query( - "SELECT * FROM report_outputs WHERE simulation_1_id = ?", - (simulation["id"],), - ).fetchall() - report_run_rows = test_db.query("SELECT * FROM report_output_runs").fetchall() - assert report_rows == [] - assert report_run_rows == [] - - -def test_get_simulation_wrong_country_returns_not_found(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_wrong_country_get", - population_type="household", - policy_id=43, - ) - - client = create_test_client() - response = client.get(f"/uk/simulation/{simulation['id']}") - - assert response.status_code == 404 - - -def test_patch_simulation_wrong_country_returns_not_found_and_does_not_mutate(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_wrong_country_patch", - population_type="household", - policy_id=44, - ) - - client = create_test_client() - response = client.patch( - "/uk/simulation", - json={ - "id": simulation["id"], - "status": "complete", - "output": json.dumps({"should_not": "persist"}), - }, - ) - - assert response.status_code == 404 - - stored_simulation = test_db.query( - "SELECT * FROM simulations WHERE id = ?", - (simulation["id"],), - ).fetchone() - assert stored_simulation["country_id"] == "us" - assert stored_simulation["status"] == "pending" - assert stored_simulation["output"] is None - - -def test_get_report_output_wrong_country_returns_not_found(test_db): - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", 999, None, get_report_output_cache_version("us"), "pending", "2025"), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() - response = client.get(f"/uk/report/{report_output['id']}") - - assert response.status_code == 404 - - -def test_patch_report_output_wrong_country_returns_not_found_and_does_not_mutate( - test_db, -): - test_db.query( - """ - INSERT INTO report_outputs ( - country_id, simulation_1_id, simulation_2_id, api_version, status, year - ) VALUES (?, ?, ?, ?, ?, ?) - """, - ("us", 1000, None, get_report_output_cache_version("us"), "pending", "2025"), - ) - report_output = test_db.query( - "SELECT * FROM report_outputs ORDER BY id DESC LIMIT 1" - ).fetchone() - - client = create_test_client() - response = client.patch( - "/uk/report", - json={ - "id": report_output["id"], - "status": "complete", - "output": json.dumps({"should_not": "persist"}), - }, - ) - - assert response.status_code == 404 - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report_output["id"],), - ).fetchone() - assert stored_report["country_id"] == "us" - assert stored_report["status"] == "pending" - assert stored_report["output"] is None - - -def test_patch_report_output_accepts_running_status(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_running_report", - population_type="household", - policy_id=45, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "running", - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "running" - assert payload["result"]["requested_at"] is not None - assert payload["result"]["started_at"] is not None - assert payload["result"]["finished_at"] is None - - -def test_get_report_output_serializes_display_run_timestamps(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_get_timestamp", - population_type="household", - policy_id=47, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - run = test_db.query( - "SELECT * FROM report_output_runs WHERE report_output_id = ?", - (report["id"],), - ).fetchone() - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 12:00:00", - "2026-05-04 12:01:00", - "2026-05-04 12:02:00", - run["id"], - ), - ) - - client = create_test_client() - response = client.get(f"/us/report/{report['id']}") - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["requested_at"] == "2026-05-04T12:00:00Z" - assert payload["result"]["started_at"] == "2026-05-04T12:01:00Z" - assert payload["result"]["finished_at"] == "2026-05-04T12:02:00Z" - - -def test_patch_report_output_running_uses_active_rerun_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_running_rerun", - population_type="household", - policy_id=48, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "running", - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "running" - assert payload["result"]["started_at"] is not None - assert payload["result"]["finished_at"] is None - - successful_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (successful_run_id,), - ).fetchone() - active_run = test_db.query( - "SELECT * FROM report_output_runs WHERE id = ?", - (rerun["id"],), - ).fetchone() - assert successful_run["status"] == "complete" - assert successful_run["finished_at"] is not None - assert active_run["status"] == "running" - assert active_run["started_at"] is not None - assert active_run["finished_at"] is None - - -def test_patch_report_output_error_uses_active_rerun_timestamp_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_error_rerun", - population_type="household", - policy_id=49, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - test_db.query( - """ - UPDATE report_output_runs - SET requested_at = ?, started_at = ?, finished_at = ? - WHERE id = ? - """, - ( - "2026-05-04 10:00:00", - "2026-05-04 10:01:00", - "2026-05-04 10:02:00", - successful_run_id, - ), - ) - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "error", - "error_message": "rerun failed", - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "error" - assert payload["result"]["finished_at"] is not None - assert payload["result"]["finished_at"] != "2026-05-04T10:02:00Z" - - -def test_patch_report_output_complete_promotes_active_rerun_route_path(test_db): - simulation = simulation_service.create_simulation( - country_id="us", - population_id="household_route_active_complete_rerun", - population_type="household", - policy_id=50, - ) - report = report_output_service.create_report_output( - country_id="us", - simulation_1_id=simulation["id"], - simulation_2_id=None, - year="2025", - ) - report_output_service.update_report_output( - country_id="us", - report_id=report["id"], - status="complete", - output=json.dumps({"ok": True}), - ) - completed_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - successful_run_id = completed_report["latest_successful_run_id"] - rerun = report_run_service.create_report_output_run( - report["id"], trigger_type="rerun" - ) - test_db.query( - """ - UPDATE report_outputs - SET active_run_id = ?, latest_successful_run_id = ? - WHERE id = ? - """, - (rerun["id"], successful_run_id, report["id"]), - ) - - client = create_test_client() - response = client.patch( - "/us/report", - json={ - "id": report["id"], - "status": "complete", - "output": json.dumps({"ok": "rerun"}), - }, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["result"]["status"] == "complete" - assert payload["result"]["finished_at"] is not None - - stored_report = test_db.query( - "SELECT * FROM report_outputs WHERE id = ?", - (report["id"],), - ).fetchone() - assert stored_report["active_run_id"] is None - assert stored_report["latest_successful_run_id"] == rerun["id"] diff --git a/tests/unit/test_warmup.py b/tests/unit/test_warmup.py index 32aea3579..8bf1f9c2f 100644 --- a/tests/unit/test_warmup.py +++ b/tests/unit/test_warmup.py @@ -101,6 +101,8 @@ def test_asgi_runs_warmup_and_marks_ready(): def test_readiness_check_gates_on_readiness(): - src = (REPO / "policyengine_api/api.py").read_text(encoding="utf-8") + src = (REPO / "policyengine_api/routes/system_routes.py").read_text( + encoding="utf-8" + ) assert "is_ready" in src assert "status=503" in src diff --git a/tests/unit/test_yearly_var_removal.py b/tests/unit/test_yearly_var_removal.py index c5d81009c..9c13fd378 100644 --- a/tests/unit/test_yearly_var_removal.py +++ b/tests/unit/test_yearly_var_removal.py @@ -1,7 +1,9 @@ import copy from types import SimpleNamespace -from policyengine_api.endpoints.household import add_yearly_variables +from policyengine_api.services.household_calculation_service import ( + add_yearly_variables, +) TEST_YEAR = "2023" diff --git a/uv.lock b/uv.lock index c345d2531..0a548c102 100644 --- a/uv.lock +++ b/uv.lock @@ -155,6 +155,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "alembic" +version = "1.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mako" }, + { name = "sqlalchemy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fb/01/a48dab7827ac4421272399f7ed9a2ec17edd12c8bcde4417bd7b6821b71a/alembic-1.19.0.tar.gz", hash = "sha256:6487c612fc719dcfa22b17d2dd5b2b458929641e6aa2f0b65b135727f5e6d501", size = 2069906, upload-time = "2026-08-04T18:57:04.599Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/79/59ab6f1fe72eee229de91aa393225389a85df12a0fbbbfa264efcf6d7872/alembic-1.19.0-py3-none-any.whl", hash = "sha256:cf839d3849116aab3cc047e09c6968b9bd6b2fde61b6bb7c1e97352fe5503580", size = 265738, upload-time = "2026-08-04T18:57:06.219Z" }, +] + [[package]] name = "altair" version = "6.1.0" @@ -1793,6 +1807,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/91/6c074015990f4f656f7b69a5c2d15924906ce0bc19c7014ac953493c0cf0/linecheck-0.1.0-py3-none-any.whl", hash = "sha256:73c6b29790521fa711b00df7cd60af4caf7004337d8710606881fbecb0d1bc83", size = 2767, upload-time = "2022-07-16T13:06:17.01Z" }, ] +[[package]] +name = "mako" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/12/b5fa2353e2754cd67fb9f83793fa48ff42c213a5da7e719869d2301f6ab8/mako-1.4.1.tar.gz", hash = "sha256:d7904710b662996425a21627710c4777c45053146942cf8a7aebf757c92b8c27", size = 410165, upload-time = "2026-08-05T06:10:56.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/54/12ed58d458474aaab5c3d180173e745a4fe131bb330370596876d19ff60f/mako-1.4.1-py3-none-any.whl", hash = "sha256:a359d9a94a541213958742b2698d0a7757bb83551767bc468a74b9905aba9617", size = 80010, upload-time = "2026-08-05T06:10:58.248Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -2461,7 +2487,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "ptyprocess" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -2616,10 +2642,11 @@ models = [ [[package]] name = "policyengine-api" -version = "3.46.2" +version = "3.48.0" source = { editable = "." } dependencies = [ { name = "a2wsgi" }, + { name = "alembic" }, { name = "anthropic" }, { name = "assertpy" }, { name = "click" }, @@ -2665,6 +2692,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "a2wsgi", specifier = ">=1.10,<2" }, + { name = "alembic", specifier = ">=1.14,<2" }, { name = "anthropic" }, { name = "assertpy" }, { name = "build", marker = "extra == 'dev'" },