diff --git a/.github/scripts/build_pr_summary_comment.py b/.github/scripts/build_pr_summary_comment.py index 852e7d49..c17cc34c 100644 --- a/.github/scripts/build_pr_summary_comment.py +++ b/.github/scripts/build_pr_summary_comment.py @@ -30,7 +30,9 @@ def build_comment_body(root: Path) -> str: for item in items: for warning in item.get("warnings", []): warnings.append(f"{item.get('name', 'unknown')}: {warning}") - sections.setdefault(item.get("section", "other"), []).append(item) + # Entries with no table data only contribute warnings. + if item.get("data"): + sections.setdefault(item.get("section", "other"), []).append(item) short_sha = "" for item in items: diff --git a/.github/workflows/build-and-deploy.yml b/.github/workflows/build-and-deploy.yml index 0dbc24b4..43130ca9 100644 --- a/.github/workflows/build-and-deploy.yml +++ b/.github/workflows/build-and-deploy.yml @@ -289,6 +289,27 @@ jobs: postgres-db: ${{ needs.deploy-db.outputs.postgres-db }} postgres-sslmode: ${{ needs.deploy-db.outputs.postgres-sslmode }} + # Rebuilds the API's resource-state read model for this commit (a no-op if + # it's already current). Runs after the API deploy, which waits until the + # previous revision has stopped serving (wait-for-cutover): that revision's + # writes refresh rows with the old code's rules, and could otherwise land + # after this rebuild marks the table ready. + # Nothing waits on this job: until it finishes, list reads are served live + # (correct, just slower), so a slow or failed rebuild never blocks a deploy. + rebuild-read-model: + name: "Rebuild Read Model" + needs: [resolve-context, build-api-docker-image, deploy-db, run-db-migrations, deploy-api] + uses: ./.github/workflows/rebuild-read-model.yml + secrets: + stitch-app-password: ${{ secrets.STITCH_APP_PASSWORD }} + with: + deployment-lane: ${{ needs.resolve-context.outputs.deployment-lane }} + full-image-name: ${{ needs.build-api-docker-image.outputs.digest-image-name }} + postgres-host: ${{ needs.deploy-db.outputs.postgres-host }} + postgres-port: ${{ needs.deploy-db.outputs.postgres-port }} + postgres-db: ${{ needs.deploy-db.outputs.postgres-db }} + postgres-sslmode: ${{ needs.deploy-db.outputs.postgres-sslmode }} + check-db-migrations-live: name: "Check DB Migrations (live)" needs: [resolve-context, build-api-docker-image, deploy-db, run-db-migrations] @@ -337,6 +358,9 @@ jobs: OTEL_SERVICE_NAME=stitch-api OTEL_RESOURCE_ATTRIBUTES=deployment.name=${{ needs.resolve-context.outputs.deployment-name }},deployment.lane=${{ needs.resolve-context.outputs.deployment-lane }},service.version=${{ github.sha }} deployment-label: api + # The read-model rebuild that runs after this deploy must not overlap + # the previous API revision's writes. + wait-for-cutover: true # Only the production release stays warm; everything else scales to zero. min-replicas: ${{ needs.resolve-context.outputs.always-on }} @@ -470,4 +494,4 @@ jobs: && !contains(github.event.pull_request.labels.*.name, 'no-deploy') }} uses: ./.github/workflows/comment-pr-summary.yml secrets: inherit - needs: [resolve-context, build-api-docker-image, build-stitch-llm-docker-image, build-seed-docker-image, deploy-db, run-db-migrations, deploy-api, deploy-stitch-llm, deploy-etl, run-seed, deploy-frontend] + needs: [resolve-context, build-api-docker-image, build-stitch-llm-docker-image, build-seed-docker-image, deploy-db, run-db-migrations, deploy-api, rebuild-read-model, deploy-stitch-llm, deploy-etl, run-seed, deploy-frontend] diff --git a/.github/workflows/deploy-container.yml b/.github/workflows/deploy-container.yml index c24a7b2f..e29d45c9 100644 --- a/.github/workflows/deploy-container.yml +++ b/.github/workflows/deploy-container.yml @@ -68,6 +68,11 @@ on: required: false default: "" type: string + wait-for-cutover: + description: "Optional. After deploying, wait until every active revision runs this image (the previous revision no longer receives requests), plus time for its in-flight requests to finish. Use when something run after this deploy must not overlap the previous version's writes (e.g. the API's read-model rebuild). Fails the deploy if the new revision fails to provision or doesn't take over within 10 minutes." + required: false + default: false + type: boolean secrets: azure-client-id: required: true @@ -345,6 +350,72 @@ jobs: --resource-group "${AZURE_RESOURCE_GROUP}" \ --yaml /tmp/app-cooldown.yaml + - name: Wait for the new revision to replace the previous one + # The deploy action can return while the new revision is still + # provisioning, with the previous revision still serving requests (see + # the cooldown step above). Callers that set wait-for-cutover need the + # previous version to have stopped handling requests -- the API's + # read-model rebuild must not overlap the old code's writes (see + # deployments/api/READ_MODEL.md). Done when no active revision runs + # another image and one running this image has finished provisioning; + # then wait for the old revision's in-flight requests (deactivated + # replicas get a 30 s termination grace period by default). + if: ${{ inputs.wait-for-cutover }} + shell: bash + env: + AZURE_RESOURCE_GROUP: ${{ vars.AZURE_RESOURCE_GROUP }} + APP_NAME: ${{ inputs.container-app-name }} + IMAGE: ${{ inputs.full-image-name }} + TIMEOUT_SECONDS: "600" + POLL_SECONDS: "10" + DRAIN_SECONDS: "45" + run: | + set -euo pipefail + deadline=$(( $(date +%s) + TIMEOUT_SECONDS )) + az_errors="$(mktemp)" + while true; do + # A failed poll (ARM throttling, a network blip, an unexpected + # payload) is logged and retried until the deadline rather than + # failing the deploy; only the outcomes below end the step. + if ! revisions="$( + az containerapp revision list \ + --name "$APP_NAME" \ + --resource-group "$AZURE_RESOURCE_GROUP" \ + --output json 2> "$az_errors" + )"; then + state="waiting: couldn't list revisions ($(tail -n 1 "$az_errors"))" + elif ! state="$( + jq -r --arg image "$IMAGE" ' + [.[] | select(.properties.active)] as $active + | [$active[] | select(any((.properties.template.containers // [])[]; .image == $image))] as $new + | [$active[] | select(all((.properties.template.containers // [])[]; .image != $image))] as $old + | if any($new[]; .properties.provisioningState == "Failed") then "failed" + elif ($old | length) == 0 and any($new[]; .properties.provisioningState == "Provisioned") then "done" + else "waiting: \($old | length) active revision(s) on another image; \($new | length) on this image (\([$new[].properties.provisioningState] | join(", ")))" + end + ' <<< "$revisions" 2>&1 + )"; then + state="waiting: couldn't read the revision list ($state)" + fi + case "$state" in + done) + echo "Only revisions running $IMAGE are active; waiting ${DRAIN_SECONDS}s for the previous revision's in-flight requests." + sleep "$DRAIN_SECONDS" + exit 0 + ;; + failed) + echo "::error::The new revision of $APP_NAME failed to provision; the previous revision is still serving." + exit 1 + ;; + esac + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "::error::The new revision of $APP_NAME did not replace the previous one within ${TIMEOUT_SECONDS}s ($state)." + exit 1 + fi + echo "$state" + sleep "$POLL_SECONDS" + done + - name: Export application URL id: export-app-url shell: bash diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml index 351bb21c..140e6c0b 100644 --- a/.github/workflows/python-test.yml +++ b/.github/workflows/python-test.yml @@ -5,6 +5,26 @@ jobs: test: name: "tests (pytest)" runs-on: ubuntu-latest + # Throwaway Postgres for the API tests SQLite can't exercise (row locks under + # concurrency, JSONB): see deployments/api/tests/db/read_model/postgres.py. + # Each test creates and drops its own schema. + services: + postgres: + image: postgres:17-alpine + env: + POSTGRES_PASSWORD: postgres + POSTGRES_DB: stitch_test + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U postgres -d stitch_test" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + STITCH_TEST_POSTGRES_URL: postgresql+psycopg://postgres:postgres@localhost:5432/stitch_test + # Fail rather than skip if the URL above ever goes missing. + STITCH_REQUIRE_POSTGRES_TESTS: "1" steps: - name: Checkout code uses: actions/checkout@v7 diff --git a/.github/workflows/rebuild-read-model.yml b/.github/workflows/rebuild-read-model.yml new file mode 100644 index 00000000..0aa3ad63 --- /dev/null +++ b/.github/workflows/rebuild-read-model.yml @@ -0,0 +1,121 @@ +--- +# Post-deploy step: rebuild the resource-state read model if it isn't already +# built by this commit. The API only serves list/filter-options from the read +# model when a full rebuild by its own commit (GIT_SHA, baked into the image) +# has completed; until then it answers live (correct, just slower). Running the +# deployed image here records the same commit the API runs. +# +# See deployments/api/READ_MODEL.md. +on: + workflow_call: + inputs: + deployment-lane: + description: "Deployment lane and GitHub Environment name" + required: true + type: string + full-image-name: + description: "Full pushed API image name including host/registry, name, and digest" + required: true + type: string + postgres-host: + required: true + type: string + postgres-port: + required: false + default: "5432" + type: string + postgres-db: + required: true + type: string + postgres-sslmode: + required: false + default: "require" + type: string + secrets: + stitch-app-password: + required: true + +name: Rebuild Read Model + +jobs: + rebuild-read-model: + runs-on: ubuntu-latest + environment: ${{ inputs.deployment-lane }} + permissions: + contents: read + packages: read + timeout-minutes: 30 + + steps: + - name: Validate lane secrets + shell: bash + env: + STITCH_APP_PASSWORD: ${{ secrets.stitch-app-password }} + run: | + set -euo pipefail + : "${STITCH_APP_PASSWORD:?Missing required secret STITCH_APP_PASSWORD}" + + - name: Log in to GHCR + run: echo "${{ github.token }}" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Pull API image + run: docker pull "${{ inputs.full-image-name }}" + + - name: Rebuild read model if stale + id: rebuild + env: + POSTGRES_HOST: ${{ inputs.postgres-host }} + POSTGRES_PORT: ${{ inputs.postgres-port }} + POSTGRES_DB: ${{ inputs.postgres-db }} + # Reads and rewrites rows only, so the app's own role is enough. + POSTGRES_USER: stitch_app + POSTGRES_PASSWORD: ${{ secrets.stitch-app-password }} + PGSSLMODE: ${{ inputs.postgres-sslmode }} + # A batch job: don't instrument every query for tracing. + OTEL_TRACES_EXPORTER: none + run: | + set -euo pipefail + docker run --rm \ + -e POSTGRES_HOST \ + -e POSTGRES_PORT \ + -e POSTGRES_DB \ + -e POSTGRES_USER \ + -e POSTGRES_PASSWORD \ + -e PGSSLMODE \ + -e OTEL_TRACES_EXPORTER \ + "${{ inputs.full-image-name }}" \ + python -m stitch.api.db.read_model.rebuild --if-stale + + # Only when the rebuild didn't succeed (failed, cancelled, or never ran): + # a warning-only entry for the PR summary comment, so reviewers know the + # preview's lists are on the slower live path. + - name: Write comment JSON + if: ${{ always() && github.event_name == 'pull_request' && steps.rebuild.outcome != 'success' }} + shell: bash + env: + OUTCOME: ${{ steps.rebuild.outcome }} + run: | + set -euo pipefail + mkdir -p /tmp/comment-json + json_file="/tmp/comment-json/comment-json-read-model.json" + + jq -n \ + --arg outcome "$OUTCOME" \ + '{ + "kind": "read_model_rebuild", + "name": "rebuild-read-model", + "warnings": [ + "the read model was not rebuilt (\($outcome)); list reads use the slower live path until it is" + ] + }' > "$json_file" + + jq empty "$json_file" + cat "$json_file" + + - name: Upload comment JSON artifact + if: ${{ always() && github.event_name == 'pull_request' && steps.rebuild.outcome != 'success' }} + uses: actions/upload-artifact@v7 + with: + name: comment-json-read-model + path: /tmp/comment-json/*.json + if-no-files-found: error diff --git a/HACKING.md b/HACKING.md index a8be3702..7be9a7ee 100644 --- a/HACKING.md +++ b/HACKING.md @@ -96,6 +96,29 @@ Depending on which entrypoint you use, these are the main local endpoints: * API docs: `http://localhost:8000/docs` * Adminer: `http://localhost:8081` +## Fast list reads locally + +The list and filter-options endpoints read from a precomputed table, but only +after a full rebuild *by the same code* has marked it ready; until then they use +the slower live path, with the same results. Run the rebuild where the API runs: + +- API in Docker (`make frontend-dev`, `make dev-docker`): the image records your + commit, so rebuild inside it, and again after rebuilding the image: + + ```bash + docker compose exec api python -m stitch.api.db.read_model.rebuild + ``` + +- API on your machine (`make api-dev`): the code version is `unknown`, so a + rebuild stays valid until you rebuild again (do so after changing how answers + are computed): + + ```bash + POSTGRES_HOST=127.0.0.1 POSTGRES_USER=stitch_app uv run --env-file .env -- python -m stitch.api.db.read_model.rebuild + ``` + +See [`deployments/api/READ_MODEL.md`](deployments/api/READ_MODEL.md) for details. + ## Other useful targets You do not need to memorize the whole Makefile, but these are worth knowing. diff --git a/deployments/CI_DEPLOYMENTS.md b/deployments/CI_DEPLOYMENTS.md index 02510dc2..d7b4448d 100644 --- a/deployments/CI_DEPLOYMENTS.md +++ b/deployments/CI_DEPLOYMENTS.md @@ -324,7 +324,15 @@ deployments, `pr-*`, `main` deploy to the same PostgreSQL instance, with separate logical DBs) It also handles running `db-migrations` (via Alembic in the `api` image) and -`seed`, both directly from GH Actions (rather than starting on Azure). +`seed`, both directly from GH Actions (rather than starting on Azure). The API +deploy waits until its new revision has fully replaced the previous one +(`wait-for-cutover` in `deploy-container.yml`; it fails the deploy if the new +revision fails to provision or doesn't take over within 10 minutes). After the +API deploy, `rebuild-read-model` runs the API's resource-state read-model +rebuild from the same `api` image (a no-op when it's already current for this +commit); nothing waits on it, so a failure leaves list reads on the slower live +path without blocking the deploy. See +[`api/READ_MODEL.md`](./api/READ_MODEL.md). ### Custom domains (RMI hostnames) diff --git a/deployments/api/README.md b/deployments/api/README.md index b93e20bf..3292880e 100644 --- a/deployments/api/README.md +++ b/deployments/api/README.md @@ -1,5 +1,13 @@ # api +## Resource-state read model + +The list and filter-options endpoints are served from a precomputed table that +must be **rebuilt after each deploy of new code** (CI does this automatically) +and after a data restore. Until then they use the slower live path, with the +same results. See [`READ_MODEL.md`](./READ_MODEL.md) for rollout, rebuild, and +verification steps. + ## Deployment: Azure Container App ### Prerequisites diff --git a/deployments/api/READ_MODEL.md b/deployments/api/READ_MODEL.md new file mode 100644 index 00000000..897788f5 --- /dev/null +++ b/deployments/api/READ_MODEL.md @@ -0,0 +1,278 @@ +# Resource-state read model + +The oil & gas field **list** and **filter-options** endpoints are served from a +precomputed table, `og_field_resource_state`, instead of working out every field's +winning value on each request. At ~50K resources this takes typical list requests +from ~2–3 s to under 10 ms (see [Performance](#performance-at-50k-resources)). + +The table is a cache, not a source of truth. Everything in it can be rebuilt from +the authoritative tables at any time, and when it can't be trusted the API simply +answers the slower way ("live"), with the same results. + +## How it works + +- **What's stored:** one row per resource *per permission profile*, holding each + field's winning value and which source it came from (`provenance`). There are + four profiles: all public sources, plus any combination of the two restricted + sources (`wm`, `ccr`). See `db/read_model/permissions.py`. +- **Who gets the fast path:** a request whose licensed sources *exactly* match one + of the four profiles, while the table is marked **ready**. Everyone else + (partial or unusual grants, internal unscoped calls) gets the live answer. +- **Readiness:** a single row in `og_field_resource_state_status`. Only a complete + full rebuild sets `ready = true`, and it records the commit of the code that + ran it (`code_version`, from the image's `GIT_SHA`). The API uses the table + only when it is ready **and** was built by the same commit the API is running; + otherwise it answers live. So every deploy of new code switches reads to live + until that code rebuilds the table — a change to how answers are computed can + never be served from rows built by older code. An operator can also switch + the fast path off (`disabled`, see [Verifying](#verifying)); that stays off + across rebuilds and deploys until someone switches it back on. Each API + process re-checks readiness at most every 5 seconds, so changes to it take + effect within that. +- **Staying current:** every write that can change a list answer (creating or + attaching sources, re-ordering a field's sources, merging) locks the affected + resources first, then recomputes their rows in the same transaction, so they + commit together. Rows written by an older API revision still serving during a + deploy carry that revision's commit, so the next rebuild finds and redoes + them. +- **Unchanged:** the detail endpoint (`GET /oil-gas-fields/{id}`) and per-field + source listings always compute live. + +## Rolling it out + +Every deploy of new code needs a full rebuild by that code before the fast path +comes back (each commit only trusts rows it built). CI does this automatically: + +1. **Deploy as usual.** CI runs the migrations and deploys the API, then the + **Rebuild Read Model** job (`.github/workflows/rebuild-read-model.yml`) runs + `python -m stitch.api.db.read_model.rebuild --if-stale` from the same API + image. The API deploy doesn't finish until the previous revision has stopped + serving (only revisions running the new image are active, plus 45 s for + in-flight requests), so the old code's writes, which refresh rows with the + old rules, can't land after the rebuild marks the table ready. Any that land + *during* the rebuild are caught by each row's `code_version` and redone. + Reads use live from the deploy until the rebuild finishes (~1.5–2 minutes at + 50K resources after the deploy) — correct, just slower. +2. **If that job fails,** the deploy still goes out and lists stay on live. On + a PR, the deploy summary comment shows a warning saying so. Re-run the job, + or run the rebuild by hand from the deployed image (below). +3. **Verify** (below). If anything looks wrong, switch the fast path off + (below) and reads go back to live within a few seconds, and stay there + through later deploys until you switch it back on. + +## Running a rebuild + +```bash +# Full rebuild (all resources). Safe while the API is serving traffic. +python -m stitch.api.db.read_model.rebuild + +# Full rebuild only if it isn't already built by this code (what CI runs). +python -m stitch.api.db.read_model.rebuild --if-stale + +# Recompute specific resources only (e.g. after fixing data by hand). +python -m stitch.api.db.read_model.rebuild --ids 123 456 +``` + +- **Where to run it:** with the *same code* the API is running, because the + rebuild records its commit and the API only trusts rows built by its own + commit. For a deployed lane, run the deployed `api` image (it carries the + commit as `GIT_SHA`) the same way CI runs migrations, with the lane's + `POSTGRES_*` variables, as `stitch_app` or `stitch_migrator`. A rebuild run + from a laptop records `unknown` and leaves the deployed API on live. + + ```bash + docker run --rm -e POSTGRES_HOST -e POSTGRES_PORT -e POSTGRES_DB -e POSTGRES_USER -e POSTGRES_PASSWORD -e PGSSLMODE python -m stitch.api.db.read_model.rebuild + ``` + + Locally, see `HACKING.md`: run it inside the `api` container if the API runs + in Docker, or on your machine if the API does: + + ```bash + POSTGRES_HOST=127.0.0.1 POSTGRES_USER=stitch_app uv run --env-file .env -- python -m stitch.api.db.read_model.rebuild + ``` + + Don't run a rebuild from a *different* image while the old one is still + serving: an old API process that checked readiness in the last 5 seconds + keeps trusting the table, and can read rows the new code has just rewritten. + Deploys wait for the previous revision to stop serving before rebuilding (see + [Rolling it out](#rolling-it-out)); do the same whenever you rebuild by hand. + +- **What happens:** readiness is cleared first, so reads switch to live for the + duration. Resources are rebuilt in committed batches (`--batch-size`, + default 500), each batch replacing its resources' rows. Then it checks for + resources whose rows were since rewritten by other code (an older API + revision still serving writes, or an `--ids` repair run from another image) + and redoes those, repeating until none are left. The last of these checks + happens in the same transaction that sets readiness, so no repair can slip in + between: rows for resources that are no longer listable are removed, the + result is checked (every listable resource has all four rows, all computed by + this code, and nothing else has any), and only then is readiness set. The + table is never emptied along the way, so a reader that checked readiness just + before the rebuild started still finds a row for every resource. At ~50K + resources it takes about 1.5 minutes. +- **Full rebuild vs `--if-stale`:** a plain full rebuild recomputes every + resource. `--if-stale` only recomputes resources not already computed by this + code, so re-running it after a failed or cancelled run picks up where that run + stopped. +- **If it fails:** readiness stays cleared and reads stay live; just run it again. + Whatever rows a failed run leaves are never served, and the next run replaces + or reuses them. If other code keeps rewriting rows, it stops after 5 passes + with an error asking whether an older API version is still serving writes. +- **Two at once:** a full rebuild started by hand while another is running + exits with "Another resource-state rebuild is already running." `--if-stale` + (the deploy step) instead waits up to 10 minutes for it to finish, then checks + again and does nothing if that rebuild already made the read model current. +- **A rebuild that dies mid-run** (e.g. a cancelled CI job) releases its lock + within about a minute: the lock's connection asks Postgres to drop vanished + clients quickly. If a lock still seems stuck, find the session holding it and, + once you're sure no rebuild is running, end it: + + ```sql + SELECT l.pid, a.state, a.backend_start, a.client_addr + FROM pg_locks l JOIN pg_stat_activity a USING (pid) + WHERE l.locktype = 'advisory' AND l.objid = 766001 AND l.granted; + + SELECT pg_terminate_backend(); + ``` +- `--ids` never changes readiness, and is safe while the API serves reads. It + reports which ids it recomputed and which aren't listed (merged away, or no + active sources — any stored rows for those are removed). It changes nothing + and exits with an error if any id doesn't exist (usually a typo), or if the + table is ready but was built by different code (even if switched off, since + switching it back on serves those rows): run it from the same image as the + API, so the table never mixes two versions. During a full rebuild the + table isn't ready, so a repair from any image runs; the rebuild waits for it + and redoes its rows before finishing. A repair that starts just as a rebuild + finishes waits for it, then refuses if that rebuild was by other code. +- Expected failures (another rebuild running, a failed integrity check, rows + still rewritten by other code after 5 passes, unknown ids, a table built by + other code) print as a single `ERROR` line; unexpected ones (e.g. the + database is unreachable) show a traceback. Either way it exits with status 1. + Invalid arguments print the usage and exit with status 2. + +## Verifying + +```sql +-- Ready, when it was last rebuilt, by which commit (compare with +-- build.git_sha from the API's health endpoint), and whether it's switched off +SELECT ready, rebuilt_at, code_version, disabled FROM og_field_resource_state_status; + +-- Exactly 4 rows per listable resource +SELECT count(*) AS state_rows, count(DISTINCT resource_id) AS resources +FROM og_field_resource_state; +SELECT count(DISTINCT r.id) AS listable_resources +FROM og_field_resources r +JOIN og_field_memberships m ON m.resource_id = r.id +WHERE r.repointed_id IS NULL AND m.status = 'ACTIVE'; + +-- No merged-away resources (expect 0) +SELECT count(*) FROM og_field_resource_state s +JOIN og_field_resources r ON r.id = s.resource_id +WHERE r.repointed_id IS NOT NULL; + +-- Every row computed by the serving commit (expect 0) +SELECT count(*) FROM og_field_resource_state s, og_field_resource_state_status st +WHERE s.code_version IS DISTINCT FROM st.code_version; +``` + +Then smoke-test the list and filter-options pages, and check the query timings +(see `deployments/PERFORMANCE.md`). Served from the read model, a list request +runs `resources.count` and `resources.list_page`; served live it runs +`resources.count`, `resources.list_ids` and `resources.list_hydrate`. +Filter-options is `resources.filter_options` either way. The cost the read +model adds to writes (and to rebuild batches) is labeled `read_model.refresh`. + +**To switch the fast path off** without deleting anything (API processes pick +this up within about 5 seconds): + +```sql +UPDATE og_field_resource_state_status SET disabled = true; +``` + +It stays off until you switch it back on. Rebuilds (including CI's after every +deploy) still run and keep the rows current, but never switch it back on; they +log a warning saying it's off. Once the problem is fixed and a rebuild by the +deployed code has finished: + +```sql +UPDATE og_field_resource_state_status SET disabled = false; +``` + +(Setting `ready = false` instead only lasts until the next rebuild, which CI +runs on every deploy.) + +## When a full rebuild is needed again + +Ordinary API writes keep the table current. Beyond that: + +- **Every deploy of new code** needs a rebuild before the fast path comes back + (see readiness above); CI's Rebuild Read Model job runs it. This covers changes + to how answers are computed and migrations alike — nothing else needs to clear + readiness. +- **A database restore, or manual SQL edits** to sources, memberships, or + priorities, change answers without the code or the writes noticing: run a + plain full rebuild (or `--ids` for the affected resources). Not `--if-stale`: + rows already computed by this code would be skipped, even though the data + under them changed. + +## Changing the code safely + +- **New write path** that can change a list answer (a source value, membership, + priority, or merge state): first call `lock_resources` from + `db/read_model/state.py` for the affected resources, before reading or + writing anything (so a writer that waited reads the other's committed + changes, e.g. a merge); then call + `refresh_resource_state(s)` in the same transaction; and add a case to + `tests/db/read_model/test_mutations.py`. Editing a source's values in place + affects *every* resource it is attached to — lock and refresh all of them. +- **New field** on `OilGasFieldBase`: add the column to `OGFieldResourceState` + (an import-time check fails until you do) plus a migration, then rebuild. +- **New source key:** decide whether it's public or restricted in + `db/read_model/permissions.py`. A test fails until you do; until then, anyone + holding the new key gets the live path. + +## Tests + +The normal suite runs on SQLite. A few tests need real Postgres (row locks under +concurrency, separate statements seeing different moments, JSONB output). CI +runs them on every PR against a throwaway Postgres; locally they're skipped +unless pointed at a disposable database. They create and drop their own schema: + +```bash +STITCH_TEST_POSTGRES_URL=postgresql+psycopg://postgres:pg@127.0.0.1:55432/stitch uv run --package stitch-api pytest deployments/api/tests/db/read_model -q +``` + +## Performance at 50K resources + +Measured October 2026 on a laptop (Postgres 17, local), against a restored copy of +the 50K-resource benchmark seed (52K listable resources, 1.4M source values). +Medians of 7 runs, full action call including the readiness check; both columns +return identical results. + +| Request | Live | Read model | +|---|---:|---:| +| List, default sort, page 1 | 2.0–2.7 s | 6–9 ms | +| List, page 500 | 2.1–2.9 s | 7–9 ms | +| List, filtered (country, basin) | 2.0–2.7 s | 4 ms | +| List, sorted by a field other than `name` | 2.0–2.6 s | 30–100 ms | +| List, `q=` text search | 2.2–2.9 s | 50–150 ms | +| Filter options | 0.4–0.5 s | 20–25 ms | + +Full rebuild: ~97 s, ~160 MB peak memory; the table is ~195 MB including indexes. + +## Known limitations and follow-ups + +- **Incremental `--since` repair** is not implemented: no timestamp can reliably + catch every kind of change (STIT-809). +- **Write overhead:** each write's refresh (query label `read_model.refresh`) + takes ~85 ms at 50K resources, almost all of it finding the resource's + memberships, because `og_field_memberships` has no index on `resource_id`. + The unique `(resource_id, source_pk)` constraint from the separate DB-cleanup + work (#322) provides one; with it the refresh drops to ~16 ms per resource + (~24 ms for a three-resource merge). +- **Sorting by fields other than `name`** has no index (30–100 ms). Add + per-direction indexes only if usage shows those sorts matter. +- **`q=` text search** scans the profile's rows (50–150 ms); a btree index can't + help a substring match. A trigram index (`pg_trgm`) is the likely fix if needed. +- **The live list's count query** sorts every matching row before counting (the + read model's count doesn't). Worth fixing on the live path separately. diff --git a/deployments/api/alembic/versions/2f2fb644b180_resource_state_disabled_flag.py b/deployments/api/alembic/versions/2f2fb644b180_resource_state_disabled_flag.py new file mode 100644 index 00000000..c2ebdc6b --- /dev/null +++ b/deployments/api/alembic/versions/2f2fb644b180_resource_state_disabled_flag.py @@ -0,0 +1,33 @@ +"""let operators switch the resource-state read model off + +Revision ID: 2f2fb644b180 +Revises: 48a49d197641 +Create Date: 2026-10-15 00:00:00.000000 + +Add ``disabled`` to ``og_field_resource_state_status``: an operator's switch to +keep the API on live reads. Rebuilds keep the rows current but never clear it, +so a deploy's automatic rebuild can't turn the read model back on. Existing rows +get ``false`` (unchanged behaviour). +""" + +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision = "2f2fb644b180" +down_revision = "48a49d197641" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column( + "og_field_resource_state_status", + sa.Column("disabled", sa.Boolean(), nullable=False, server_default=sa.false()), + ) + + +def downgrade() -> None: + raise RuntimeError(f"Irreversible migration: {revision}") diff --git a/deployments/api/alembic/versions/48a49d197641_add_resource_state_read_model.py b/deployments/api/alembic/versions/48a49d197641_add_resource_state_read_model.py new file mode 100644 index 00000000..760aec8c --- /dev/null +++ b/deployments/api/alembic/versions/48a49d197641_add_resource_state_read_model.py @@ -0,0 +1,122 @@ +"""add og_field_resource_state read model + +Revision ID: 48a49d197641 +Revises: 71aa7ef8150b +Create Date: 2026-10-15 00:00:00.000000 + +Create the precomputed current-state read model (STIT-766): + +* ``og_field_resource_state`` -- one row per ``(resource, permission_mask)`` with + the coalesced value of every ``OilGasFieldBase`` field plus a ``provenance`` map, + and the ``code_version`` (git SHA) of the code that computed the row. A full + rebuild recomputes rows computed by other code (for example an older API + revision still writing during a deploy) before marking the table ready. +* ``og_field_resource_state_status`` -- the single readiness row, created with + ``ready = false``. ``code_version`` records the commit whose full rebuild last + marked it ready; readers use the read model only when it matches their own + commit. + +Schema only: the tables start empty and not ready, and this revision changes no +answers. The table is derived data, filled by the application's full rebuild (run +outside Alembic because it reuses the app's async coalescer); until a rebuild +completes and sets ``ready = true``, reads use live coalescing. +""" + +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql, sqlite + +# revision identifiers, used by Alembic. +revision = "48a49d197641" +down_revision = "71aa7ef8150b" +branch_labels = None +depends_on = None + +_BIGINT = ( + sa.BigInteger() + .with_variant(postgresql.BIGINT(), "postgresql") + .with_variant(sqlite.INTEGER(), "sqlite") +) +_FLOAT = ( + sa.Float() + .with_variant(postgresql.DOUBLE_PRECISION(), "postgresql") + .with_variant(sqlite.REAL(), "sqlite") +) +_JSON = sa.JSON().with_variant(postgresql.JSONB(), "postgresql") + +_FILTER_OPTION_FIELDS = ( + "basin", + "country", + "field_status", + "primary_hydrocarbon_group", + "region", + "state_province", +) + + +def upgrade() -> None: + op.create_table( + "og_field_resource_state", + sa.Column("resource_id", _BIGINT, nullable=False), + sa.Column("permission_mask", sa.SmallInteger(), nullable=False), + sa.Column("name", sa.String(), nullable=True), + sa.Column("country", sa.String(), nullable=True), + sa.Column("name_local", sa.String(), nullable=True), + sa.Column("state_province", sa.String(), nullable=True), + sa.Column("region", sa.String(), nullable=True), + sa.Column("basin", sa.String(), nullable=True), + sa.Column("reservoir_formation", sa.String(), nullable=True), + sa.Column("location_type", sa.String(), nullable=True), + sa.Column("production_conventionality", sa.String(), nullable=True), + sa.Column("primary_hydrocarbon_group", sa.String(), nullable=True), + sa.Column("field_status", sa.String(), nullable=True), + sa.Column("latitude", _FLOAT, nullable=True), + sa.Column("longitude", _FLOAT, nullable=True), + sa.Column("discovery_year", sa.Integer(), nullable=True), + sa.Column("production_start_year", sa.Integer(), nullable=True), + sa.Column("fid_year", sa.Integer(), nullable=True), + sa.Column("owners", _JSON, nullable=True), + sa.Column("operators", _JSON, nullable=True), + sa.Column("provenance", _JSON, nullable=False), + sa.Column("code_version", sa.String(), nullable=True), + sa.CheckConstraint( + "permission_mask BETWEEN 0 AND 3", + name="ck_resource_state_permission_mask", + ), + sa.ForeignKeyConstraint( + ["resource_id"], ["og_field_resources.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("resource_id", "permission_mask"), + ) + op.create_index( + "ix_resource_state_mask_resource", + "og_field_resource_state", + ["permission_mask", "resource_id"], + ) + op.create_index( + "ix_resource_state_mask_name", + "og_field_resource_state", + ["permission_mask", "name", "resource_id"], + ) + for field in _FILTER_OPTION_FIELDS: + op.create_index( + f"ix_resource_state_mask_{field}", + "og_field_resource_state", + ["permission_mask", field], + ) + + status = op.create_table( + "og_field_resource_state_status", + sa.Column("key", sa.String(), nullable=False), + sa.Column("ready", sa.Boolean(), nullable=False), + sa.Column("rebuilt_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("code_version", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("key"), + ) + op.bulk_insert(status, [{"key": "resource_state", "ready": False}]) + + +def downgrade() -> None: + raise RuntimeError(f"Irreversible migration: {revision}") diff --git a/deployments/api/src/stitch/api/db/model/__init__.py b/deployments/api/src/stitch/api/db/model/__init__.py index 55cd212a..e4078bfc 100644 --- a/deployments/api/src/stitch/api/db/model/__init__.py +++ b/deployments/api/src/stitch/api/db/model/__init__.py @@ -1,6 +1,10 @@ from .common import Base as StitchBase from .og_field_source_priority import OGFieldSourcePriority from .og_field_resource_source_priority import OGFieldResourceSourcePriority +from .og_field_resource_state import ( + OGFieldResourceState, + OGFieldResourceStateStatus, +) from .oil_gas_field_source import OilGasFieldSourceModel from .oil_gas_field_source_value import OilGasFieldSourceValueModel from .membership import MembershipModel, MembershipStatus @@ -13,6 +17,8 @@ "MembershipStatus", "OGFieldSourcePriority", "OGFieldResourceSourcePriority", + "OGFieldResourceState", + "OGFieldResourceStateStatus", "OilGasFieldSourceModel", "OilGasFieldSourceValueModel", "MergeCandidateItemModel", diff --git a/deployments/api/src/stitch/api/db/model/og_field_resource_state.py b/deployments/api/src/stitch/api/db/model/og_field_resource_state.py new file mode 100644 index 00000000..2ff0b108 --- /dev/null +++ b/deployments/api/src/stitch/api/db/model/og_field_resource_state.py @@ -0,0 +1,146 @@ +"""Precomputed, permission-scoped current state of OG field resources (STIT-766). + +``og_field_resource_state`` holds one row per ``(resource, permission_mask)``: the +already-coalesced value of every ``OilGasFieldBase`` field as that permission +profile sees it, plus the field -> winning source key ``provenance`` map the list +view returns. ``list`` and ``filter-options`` read these rows instead of re-running +the coalescing window on every request. + +This is derived data. Every row is reconstructable from memberships, source values +and priorities using the live coalescer, so deleting or emptying the table loses +nothing. ``permission_mask`` identifies one of four exact +permission profiles (see ``db.read_model.permissions``). + +``og_field_resource_state_status`` is a single-row readiness record. Readers use +the state table only while ``ready`` is true, ``code_version`` names the commit +they are running, and ``disabled`` is false; a missing row means not ready. A +partial, failed, or in-progress rebuild leaves ``ready`` false, so incomplete +state is never served, and deploying new code makes the existing rows unusable +until that code rebuilds them. ``disabled`` is an operator's switch-off that +rebuilds never clear. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Final + +from sqlalchemy import ( + Boolean, + false, + CheckConstraint, + DateTime, + ForeignKey, + Index, + Integer, + SmallInteger, + String, +) +from sqlalchemy.orm import Mapped, mapped_column + +from stitch.api.entities import FILTER_OPTION_FIELDS + +from .common import Base +from .oil_gas_field_source_value import ATTRIBUTE_NAMES +from .types import PORTABLE_BIGINT, PORTABLE_FLOAT, PORTABLE_JSON, PORTABLE_JSON_NULL + +# Key of the single readiness row in ``og_field_resource_state_status``. +RESOURCE_STATE_STATUS_KEY: Final[str] = "resource_state" + + +class OGFieldResourceState(Base): + __tablename__ = "og_field_resource_state" + + resource_id: Mapped[int] = mapped_column( + PORTABLE_BIGINT, + ForeignKey("og_field_resources.id", ondelete="CASCADE"), + primary_key=True, + ) + # One of the exact permission profiles in read_model.permissions. + permission_mask: Mapped[int] = mapped_column(SmallInteger, primary_key=True) + + # --- Coalesced values, one column per OilGasFieldBase field --- + # Column types follow ATTRIBUTE_KINDS (text / int / float / json). + name: Mapped[str | None] = mapped_column(String, nullable=True) + country: Mapped[str | None] = mapped_column(String, nullable=True) + name_local: Mapped[str | None] = mapped_column(String, nullable=True) + state_province: Mapped[str | None] = mapped_column(String, nullable=True) + region: Mapped[str | None] = mapped_column(String, nullable=True) + basin: Mapped[str | None] = mapped_column(String, nullable=True) + reservoir_formation: Mapped[str | None] = mapped_column(String, nullable=True) + location_type: Mapped[str | None] = mapped_column(String, nullable=True) + production_conventionality: Mapped[str | None] = mapped_column( + String, nullable=True + ) + primary_hydrocarbon_group: Mapped[str | None] = mapped_column(String, nullable=True) + field_status: Mapped[str | None] = mapped_column(String, nullable=True) + latitude: Mapped[float | None] = mapped_column(PORTABLE_FLOAT, nullable=True) + longitude: Mapped[float | None] = mapped_column(PORTABLE_FLOAT, nullable=True) + discovery_year: Mapped[int | None] = mapped_column(Integer, nullable=True) + production_start_year: Mapped[int | None] = mapped_column(Integer, nullable=True) + fid_year: Mapped[int | None] = mapped_column(Integer, nullable=True) + owners: Mapped[Any | None] = mapped_column(PORTABLE_JSON_NULL, nullable=True) + operators: Mapped[Any | None] = mapped_column(PORTABLE_JSON_NULL, nullable=True) + + # field -> winning source key (or None). Returned with list items; never + # filtered or sorted on. + provenance: Mapped[dict[str, Any]] = mapped_column(PORTABLE_JSON, nullable=False) + + # The commit (GIT_SHA) of the code that computed this row. A full rebuild + # redoes rows computed by other code -- e.g. an older API revision still + # writing during a deploy -- before marking the table ready. + code_version: Mapped[str | None] = mapped_column(String, nullable=True) + + __table_args__ = ( + CheckConstraint( + "permission_mask BETWEEN 0 AND 3", + name="ck_resource_state_permission_mask", + ), + # Scans of a single profile. + Index("ix_resource_state_mask_resource", "permission_mask", "resource_id"), + # The list's default sort (name ASC NULLS LAST, then resource_id). + Index("ix_resource_state_mask_name", "permission_mask", "name", "resource_id"), + # filter-options does one DISTINCT per field within a profile. + *( + Index(f"ix_resource_state_mask_{field}", "permission_mask", field) + for field in FILTER_OPTION_FIELDS + ), + ) + + +# Fail fast if the value columns drift from the coalesced entity they mirror +# (like the ATTRIBUTE_KINDS guard against OilGasFieldBase). An explicit raise, +# not assert, so the check survives `python -O`. +_VALUE_COLUMNS = frozenset(OGFieldResourceState.__table__.columns.keys()) - { + "resource_id", + "permission_mask", + "provenance", + "code_version", +} +if _VALUE_COLUMNS != frozenset(ATTRIBUTE_NAMES): + raise RuntimeError( + "OGFieldResourceState value columns out of sync with ATTRIBUTE_NAMES: " + f"{_VALUE_COLUMNS ^ frozenset(ATTRIBUTE_NAMES)}" + ) + + +class OGFieldResourceStateStatus(Base): + """Readiness of ``og_field_resource_state`` (operational metadata, not domain data). + + The migration creates the single row (``key = RESOURCE_STATE_STATUS_KEY``) + with ``ready = false``. Only a complete full rebuild sets it true, recording + in ``code_version`` the commit of the code that built the rows. ``disabled`` + is set and cleared only by hand (see deployments/api/READ_MODEL.md). + """ + + __tablename__ = "og_field_resource_state_status" + + key: Mapped[str] = mapped_column(String, primary_key=True) + ready: Mapped[bool] = mapped_column(Boolean, nullable=False) + rebuilt_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + code_version: Mapped[str | None] = mapped_column(String, nullable=True) + disabled: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=False, server_default=false() + ) diff --git a/deployments/api/src/stitch/api/db/og_field_resource_actions.py b/deployments/api/src/stitch/api/db/og_field_resource_actions.py index 7b4dda3a..37b41eb5 100644 --- a/deployments/api/src/stitch/api/db/og_field_resource_actions.py +++ b/deployments/api/src/stitch/api/db/og_field_resource_actions.py @@ -1,5 +1,5 @@ from collections.abc import Collection, Sequence -from typing import get_args +from typing import Any, get_args from fastapi import HTTPException from sqlalchemy import delete, func, select @@ -27,6 +27,7 @@ OGFieldResource, OGFieldSourceValueView, ) +from stitch.ogsi.model.og_field import OilGasFieldBase from stitch.ogsi.model.types import OGSISrcKey from .model import ( @@ -43,6 +44,16 @@ base_resource_query, field_source_candidates, filter_option_rows, + resource_state_filter_option_query, + resource_state_page_query, + resource_state_query, +) +from .read_model.permissions import read_model_mask +from .read_model.state import ( + is_resource_state_ready_for_reads, + lock_resources, + refresh_resource_state, + refresh_resource_states, ) from .utils import ( coalesce_resources, @@ -61,10 +72,9 @@ async def query( ) -> tuple[list[OGFieldListItemView], int]: """Query coalesced resource list items, restricted to licensed sources. - A narrowed id-query (+ count) over the participating fields selects and - orders the page, then one SQL coalesce (``coalesce_resources``) hydrates - those ids -- values and provenance come straight from the query, with no - second coalesce pass. + Served from the precomputed read model when the caller's licensed sources + exactly match one of its permission profiles and the model is ready; + otherwise from live coalescing. Both return the same answer. """ if params.sort_by == "source": raise HTTPException( @@ -72,6 +82,73 @@ async def query( detail="sort_by=source is not supported for resource list queries.", ) + mask = await _ready_read_model_mask(session, licensed_sources) + if mask is not None: + return await _query_read_model(session, params, mask) + return await _query_live(session, params, licensed_sources) + + +async def _ready_read_model_mask( + session: AsyncSession, licensed_sources: Collection[OGSISrcKey] | None +) -> int | None: + """The permission mask to read, or ``None`` to use live coalescing. + + Readiness is re-checked at most every few seconds per process. When it is + asked, the lookup is unlabeled on purpose: it is not part of the list / + filter-options ``named_query`` contract. + """ + mask = read_model_mask(licensed_sources) + if mask is None or not await is_resource_state_ready_for_reads(session): + return None + return mask + + +async def _query_read_model( + session: AsyncSession, params: OGFieldQueryParams, mask: int +) -> tuple[list[OGFieldListItemView], int]: + """List page read from stored, already-coalesced rows for one profile. + + The page's rows are read by the same statement that picks them, so a write + that lands mid-request (e.g. a merge removing a resource's rows) can't + remove a row between choosing the page and loading it. + """ + # Count without the ORDER BY: Postgres would otherwise sort every matching + # row just to count them (~40ms at 50K resources for a non-indexed sort). + count_stmt = select(func.count()).select_from( + resource_state_query(params, mask).order_by(None).subquery() + ) + with named_query("resources.count"): + total = (await session.scalar(count_stmt)) or 0 + with named_query("resources.list_page"): + rows = (await session.execute(resource_state_page_query(params, mask))).all() + return [_state_row_to_list_item_view(row) for row in rows], total + + +def _state_row_to_list_item_view(row: Any) -> OGFieldListItemView: + """Build the list item from a stored row, as the live path builds it from + winners: the same ``OilGasFieldBase`` inputs, and provenance keyed in + attribute order (JSONB does not keep key order).""" + return OGFieldListItemView( + id=row.resource_id, + data=OilGasFieldBase( + **{field: getattr(row, field) for field in ATTRIBUTE_NAMES} + ), + provenance={field: row.provenance.get(field) for field in ATTRIBUTE_NAMES}, + ) + + +async def _query_live( + session: AsyncSession, + params: OGFieldQueryParams, + licensed_sources: Collection[OGSISrcKey] | None = None, +) -> tuple[list[OGFieldListItemView], int]: + """List page from live coalescing: any profile, any readiness. + + A narrowed id-query (+ count) over the participating fields selects and + orders the page, then one SQL coalesce (``coalesce_resources``) hydrates + those ids -- values and provenance come straight from the query, with no + second coalesce pass. + """ ids_stmt = base_resource_query(params, licensed_sources) count_stmt = select(func.count()).select_from(ids_stmt.subquery()) with named_query("resources.count"): @@ -95,7 +172,33 @@ async def filter_options( session: AsyncSession, licensed_sources: Collection[OGSISrcKey] | None = None, ) -> dict[str, list[str]]: - """Distinct coalesced values for every filterable field, in one query.""" + """Distinct coalesced values for every filterable field. + + Read from the precomputed read model under the same gate as ``query``; + otherwise from live coalescing. + """ + mask = await _ready_read_model_mask(session, licensed_sources) + if mask is None: + return await _filter_options_live(session, licensed_sources) + + # One DISTINCT per field, each answered from its (permission_mask, field) + # index. A single UNION ALL of all six was measured slower at 50K resources + # (the cost is the index scans, not round trips); that the six statements + # may see slightly different moments only affects dropdown suggestions, + # transiently. + options: dict[str, list[str]] = {} + with named_query("resources.filter_options"): + for field in FILTER_OPTION_FIELDS: + stmt = resource_state_filter_option_query(field, mask) + options[field] = list((await session.scalars(stmt)).all()) + return options + + +async def _filter_options_live( + session: AsyncSession, + licensed_sources: Collection[OGSISrcKey] | None = None, +) -> dict[str, list[str]]: + """Distinct coalesced values for every filterable field, in one live query.""" options: dict[str, list[str]] = {field: [] for field in FILTER_OPTION_FIELDS} with named_query("resources.filter_options"): for colname, value in await session.execute( @@ -255,8 +358,12 @@ async def set_field_source_priority( status_code=422, detail=f"field={field} is not a known resource field.", ) + # Lock before reading or writing anything, so concurrent writes to this + # resource queue here and each reads the other's committed changes; then + # re-read it in case it was merged while we waited. with named_query("resources.set_field_source_priority.load"): - resource = await session.get(ResourceModel, id) + await lock_resources(session, [id]) + resource = await session.get(ResourceModel, id, populate_existing=True) if resource is None: raise HTTPException( status_code=HTTP_404_NOT_FOUND, detail=f"No Resource with id `{id}` found." @@ -320,6 +427,8 @@ async def set_field_source_priority( ) ) await session.flush() + # The new order can change this resource's winners. + await refresh_resource_state(session, id) return await field_source_values(session, id, field, licensed_sources) @@ -373,9 +482,18 @@ async def apply_resource_merge( f"Merging only possible between multiple ids: received: {unique_ids}" ) - stmt = select(ResourceModel).where(ResourceModel.id.in_(unique_ids)) + # Lock the originals (in id order, so overlapping merges can't deadlock) + # before reading them, and re-read their current state rather than any copy + # the session already holds: a merge that waited on the lock must see that + # a concurrent merge already repointed them. + stmt = ( + select(ResourceModel) + .where(ResourceModel.id.in_(unique_ids)) + .execution_options(populate_existing=True) + ) with named_query("resources.merge.load"): + await lock_resources(session, unique_ids) results = (await session.scalars(stmt)).all() missing_ids = set(unique_ids).difference(set([r.id for r in results])) if len(missing_ids) > 0: @@ -410,6 +528,10 @@ async def apply_resource_merge( _ = await _repoint_memberships(session, user, new_resource.id, unique_ids) + # The originals stop being listable (their rows are removed) and the new + # resource becomes listable with their combined sources. + await refresh_resource_states(session, [*unique_ids, new_resource.id]) + # Return the canonical resource entity await session.refresh(new_resource, ["memberships"]) return await resource_model_to_entity(session, new_resource) diff --git a/deployments/api/src/stitch/api/db/og_field_source_actions.py b/deployments/api/src/stitch/api/db/og_field_source_actions.py index 2abd5eec..c501fba2 100644 --- a/deployments/api/src/stitch/api/db/og_field_source_actions.py +++ b/deployments/api/src/stitch/api/db/og_field_source_actions.py @@ -23,6 +23,7 @@ from .queries import ( base_source_query, ) +from .read_model.state import lock_resources, refresh_resource_state from .utils import resource_model_to_entity @@ -45,14 +46,21 @@ async def create_source( async def _get_attachable_resource( session: AsyncSession, resource_id: int ) -> ResourceModel: - """Fetch a resource that is a valid attachment target. + """Lock and fetch a resource that is a valid attachment target. + + The row lock is taken before anything is read or written for the attach, + and held to the end of the transaction: concurrent writes to the resource + then queue here, and the resource is re-read after locking, so a merge that + committed while we waited is seen (rather than attaching to a resource that + was just merged away). Raises: ResourceNotFoundError: if no resource exists for ``resource_id``. ResourceIntegrityError: if the resource has been merged (repointed); memberships on it would never surface in coalescing/queries. """ - resource = await session.get(ResourceModel, resource_id) + await lock_resources(session, [resource_id]) + resource = await session.get(ResourceModel, resource_id, populate_existing=True) if resource is None: raise ResourceNotFoundError(f"No resource found for id: {resource_id}") if resource.repointed_id is not None: @@ -182,7 +190,12 @@ async def _attach_source_models( src_models: Sequence[OilGasFieldSourceModel], user: User, ) -> None: - """Create ACTIVE memberships linking each source model to ``resource``.""" + """Create ACTIVE memberships linking each source model to ``resource``. + + The single place the attach paths (create, attach, create-and-attach) add + memberships, so it also refreshes the resource's read-model state in the + same transaction. + """ memberships = [ MembershipModel.create( created_by=user, @@ -194,6 +207,7 @@ async def _attach_source_models( ] session.add_all(memberships) await session.flush() + await refresh_resource_state(session, resource.id) async def attach_sources_to_resource( diff --git a/deployments/api/src/stitch/api/db/queries.py b/deployments/api/src/stitch/api/db/queries.py index 02ba5d5c..081bf49b 100644 --- a/deployments/api/src/stitch/api/db/queries.py +++ b/deployments/api/src/stitch/api/db/queries.py @@ -32,6 +32,7 @@ MembershipModel, MembershipStatus, OGFieldResourceSourcePriority, + OGFieldResourceState, OGFieldSourcePriority, OilGasFieldSourceModel, OilGasFieldSourceValueModel, @@ -388,11 +389,14 @@ def coalesced_candidate_rows( ) -def _resource_universe() -> Select[tuple[int]]: +def resource_universe() -> Select[tuple[int]]: """Resources eligible to appear in a list: any non-repointed resource with an active membership. Membership-derived and ungated by licensing/source, so a resource whose licensed values are all absent still appears as a null-shell on - an unfiltered list (and drops out once a field is filtered).""" + an unfiltered list (and drops out once a field is filtered). + + Also decides which resources get rows in the resource-state read model + (``read_model.state``), so cached and live listability can't diverge.""" m = MembershipModel r = ResourceModel return ( @@ -409,7 +413,7 @@ def base_resource_query( licensed_sources: Collection[OGSISrcKey] | None = None, ) -> Select[tuple[int]]: involved = _participating_columns(params) - universe = _resource_universe().cte("resource_universe") + universe = resource_universe().cte("resource_universe") if not involved: # No value field filtered or sorted -> the universe alone (every active @@ -448,6 +452,72 @@ def base_resource_query( return stmt.order_by(*_build_sort_clauses(base, params, "resource_id")) +def resource_state_query( + params: OGFieldQueryParams, permission_mask: int +) -> Select[tuple[int]]: + """``base_resource_query`` served from the precomputed read model. + + Same filters, sort, and id ordering, applied with the shared clause builders + to the stored typed columns of one permission profile -- no coalescing + window. The table only holds listable resources, so it is the universe. + """ + state = OGFieldResourceState + profile = _resource_state_profile(permission_mask) + stmt = select(state.resource_id).where(state.permission_mask == permission_mask) + for cond in _build_field_conditions(profile, params): + stmt = stmt.where(cond) + if params.id is not None: + stmt = stmt.where(state.resource_id == params.id) + return stmt.order_by(*_build_sort_clauses(profile, params, "resource_id")) + + +def resource_state_page_query( + params: OGFieldQueryParams, permission_mask: int +) -> Select[Any]: + """One page of stored rows, chosen and read in a single statement. + + The inner query picks the page's ids (``resource_state_query`` with the + page's limit/offset), which indexes can answer without touching skipped + rows; the outer query then reads just those rows, in the same order. Being + one statement, the page is chosen and read from one snapshot. + """ + state = OGFieldResourceState + page = ( + resource_state_query(params, permission_mask) + .limit(params.limit) + .offset(params.offset) + .subquery("page") + ) + profile = _resource_state_profile(permission_mask) + return profile.join(page, page.c.resource_id == state.resource_id).order_by( + *_build_sort_clauses(profile, params, "resource_id") + ) + + +def _resource_state_profile(permission_mask: int) -> Select[Any]: + """Every stored row for one profile. + + Also the column source for the shared clause builders: its columns are the + table's own, so conditions and sorts apply straight to the indexed table. + """ + state = OGFieldResourceState + return select(state.__table__).where(state.permission_mask == permission_mask) + + +def resource_state_filter_option_query( + field_name: str, permission_mask: int +) -> Select[tuple[str]]: + """Distinct non-null stored values of one filterable field for one profile.""" + state = OGFieldResourceState + col = getattr(state, field_name) + return ( + select(col) + .where(state.permission_mask == permission_mask, col.is_not(None)) + .distinct() + .order_by(col) + ) + + def _add_pivot_columns( stmt: Select, fields: Collection[str], diff --git a/deployments/api/src/stitch/api/db/read_model/__init__.py b/deployments/api/src/stitch/api/db/read_model/__init__.py new file mode 100644 index 00000000..c074a007 --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/__init__.py @@ -0,0 +1,6 @@ +"""Precomputed current-state read model for OG field resources (STIT-766). + +``permissions`` declares the exact permission profiles the read model stores. The +``og_field_resource_state`` table (``db.model.og_field_resource_state``) is derived +data, rebuildable at any time from authoritative tables via the live coalescer. +""" diff --git a/deployments/api/src/stitch/api/db/read_model/permissions.py b/deployments/api/src/stitch/api/db/read_model/permissions.py new file mode 100644 index 00000000..e324c6f3 --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/permissions.py @@ -0,0 +1,57 @@ +"""The exact permission profiles the resource-state read model stores. + +The read model keeps one coalesced row per resource for each of four profiles, +identified by ``permission_mask``: + + 0 = PUBLIC + 1 = PUBLIC + ccr + 2 = PUBLIC + wm + 3 = PUBLIC + wm + ccr + +A caller is served from the read model only when their licensed sources *exactly +equal* one of these sets. Anything else -- a partial public grant, an extra or +unknown key, or an unscoped (``None``) caller -- uses live coalescing. There is no +subset or closest-profile matching, so a cached answer can never show a caller more +(or less) than live coalescing would. + +The sets are spelled out literally rather than derived from ``OGSISrcKey``: a new +source key must be added here deliberately. Until then every caller holding it +falls back to live (fail closed), and a drift test flags the gap. +""" + +from __future__ import annotations + +from collections.abc import Collection +from typing import Final + +from stitch.ogsi.model.types import OGSISrcKey + +PUBLIC_SOURCES: Final[frozenset[OGSISrcKey]] = frozenset( + {"gem", "rmi", "llm", "alb", "bc", "nor"} +) +RESTRICTED_SOURCES: Final[frozenset[OGSISrcKey]] = frozenset({"wm", "ccr"}) + +PROFILES: Final[dict[int, frozenset[OGSISrcKey]]] = { + 0: PUBLIC_SOURCES, + 1: PUBLIC_SOURCES | {"ccr"}, + 2: PUBLIC_SOURCES | {"wm"}, + 3: PUBLIC_SOURCES | {"wm", "ccr"}, +} + +PERMISSION_MASKS: Final[tuple[int, ...]] = tuple(PROFILES) + +_MASK_BY_PROFILE: Final[dict[frozenset[OGSISrcKey], int]] = { + sources: mask for mask, sources in PROFILES.items() +} + + +def read_model_mask(sources: Collection[OGSISrcKey] | None) -> int | None: + """The mask whose profile exactly equals ``sources``, else ``None`` (use live).""" + if sources is None: + return None + return _MASK_BY_PROFILE.get(frozenset(sources)) + + +def visible_sources_for_mask(mask: int) -> frozenset[OGSISrcKey]: + """The exact source set a mask's rows are coalesced over.""" + return PROFILES[mask] diff --git a/deployments/api/src/stitch/api/db/read_model/rebuild.py b/deployments/api/src/stitch/api/db/read_model/rebuild.py new file mode 100644 index 00000000..4d228a7a --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/rebuild.py @@ -0,0 +1,196 @@ +"""Rebuild or repair the resource-state read model. + + python -m stitch.api.db.read_model.rebuild # full rebuild + python -m stitch.api.db.read_model.rebuild --if-stale # only if not current + python -m stitch.api.db.read_model.rebuild --ids 12 34 # repair these resources + +The API uses the read model only when a full rebuild by the same commit has +completed, so a full rebuild is needed after every deploy of new code (and after +a data restore). ``--if-stale`` is the deploy step: it rebuilds only when the +read model isn't already current for this code. ``--ids`` recomputes only the +given resources and is safe to run while the read model is serving reads. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +from collections.abc import Sequence + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from ..config import dispose_engine, get_engine, get_session_factory +from ..model import OGFieldResourceStateStatus, ResourceModel +from ..model.og_field_resource_state import RESOURCE_STATE_STATUS_KEY +from .state import ( + DEFAULT_BATCH_SIZE, + ResourceStateRebuildError, + current_code_version, + is_resource_state_disabled, + rebuild_all_resource_state, + rebuild_resource_state_if_stale, + refresh_resource_states, +) + +logger = logging.getLogger(__name__) + + +def _positive_int(raw: str) -> int: + value = int(raw) + if value < 1: + raise argparse.ArgumentTypeError(f"must be at least 1, got {value}") + return value + + +def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Rebuild the resource-state read model, or repair specific " + "resources." + ) + scope = parser.add_mutually_exclusive_group() + scope.add_argument( + "--ids", + nargs="+", + type=int, + metavar="RESOURCE_ID", + help="Only recompute these resource ids (safe while serving reads).", + ) + scope.add_argument( + "--if-stale", + action="store_true", + help="Do a full rebuild only if the read model isn't already built by " + "this code (the post-deploy step). Waits for a rebuild already running, " + "then checks again.", + ) + parser.add_argument( + "--batch-size", + type=_positive_int, + default=DEFAULT_BATCH_SIZE, + help=f"Resources per committed batch in a full rebuild " + f"(default: {DEFAULT_BATCH_SIZE}).", + ) + return parser.parse_args(argv) + + +async def _refuse_if_served_by_other_code(session: AsyncSession) -> None: + """Refuse to change a table that is ready as built by other code. + + The API trusts a ready table only as built entirely by its own commit, so + rows recomputed here by different code would be served under the wrong + version. A table that isn't ready isn't served, and the next full rebuild + replaces its rows. A switched-off table still counts: switching it back on + serves its rows again, and ``--if-stale`` from the commit that built it + would see nothing to redo. Holds a share lock on the status row + until this transaction ends, so readiness can't change meanwhile: a full + rebuild finishing now waits for the repair to commit, then sees its rows and + redoes them before marking the table ready. + """ + status = ( + await session.execute( + select( + OGFieldResourceStateStatus.ready, + OGFieldResourceStateStatus.code_version, + ) + .where(OGFieldResourceStateStatus.key == RESOURCE_STATE_STATUS_KEY) + .with_for_update(read=True) + ) + ).one_or_none() + if ( + status is not None + and status.ready + and status.code_version != current_code_version() + ): + raise ResourceStateRebuildError( + f"The read model is ready as built by code " + f"{status.code_version}, but this is code {current_code_version()}. " + "Run the repair from that code (or a full rebuild from this code); " + "nothing was changed." + ) + + +async def _repair( + session_factory: async_sessionmaker[AsyncSession], ids: list[int] +) -> None: + """Recompute the given resources, refusing unknown ids (likely typos) and + tables being served as built by other code.""" + requested = sorted(set(ids)) + async with session_factory.begin() as session: + await _refuse_if_served_by_other_code(session) + existing = set( + ( + await session.scalars( + select(ResourceModel.id).where(ResourceModel.id.in_(requested)) + ) + ).all() + ) + if missing := [rid for rid in requested if rid not in existing]: + raise ResourceStateRebuildError( + f"No resource with id(s) {missing}; nothing was changed." + ) + recomputed = await refresh_resource_states(session, requested) + if recomputed: + logger.info("Recomputed resource state for ids %s.", recomputed) + if unlisted := [rid for rid in requested if rid not in recomputed]: + logger.info( + "Ids %s are not listed (merged away, or no active sources), so they " + "have no stored state; any they had was removed.", + unlisted, + ) + + +async def _run(ids: list[int] | None, batch_size: int, if_stale: bool = False) -> None: + # Outside prod the app engine echoes every statement and its parameters; for + # a batch job over every resource that floods the logs. Progress is logged + # per batch instead. + get_engine().echo = False + session_factory = get_session_factory() + try: + if ids: + await _repair(session_factory, ids) + return + if if_stale: + count = await rebuild_resource_state_if_stale( + session_factory, batch_size=batch_size + ) + else: + count = await rebuild_all_resource_state( + session_factory, batch_size=batch_size + ) + if count is None: + logger.info( + "The read model is already built by this code (%s); nothing to do.", + current_code_version(), + ) + else: + logger.info( + "Rebuilt resource state for %d resources with code %s; the read " + "model is ready.", + count, + current_code_version(), + ) + async with session_factory() as session: + disabled = await is_resource_state_disabled(session) + if disabled: + logger.warning( + "The read model is switched off (disabled), so the API keeps " + "answering lists live. Rebuilds never switch it back on; to do so: " + "UPDATE og_field_resource_state_status SET disabled = false;" + ) + finally: + await dispose_engine() + + +def main(argv: Sequence[str] | None = None) -> None: + args = _parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + try: + asyncio.run(_run(args.ids, args.batch_size, args.if_stale)) + except ResourceStateRebuildError as exc: + logger.error("%s", exc) + raise SystemExit(1) from exc + + +if __name__ == "__main__": + main() diff --git a/deployments/api/src/stitch/api/db/read_model/state.py b/deployments/api/src/stitch/api/db/read_model/state.py new file mode 100644 index 00000000..69ec3f84 --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/state.py @@ -0,0 +1,612 @@ +"""Build and maintain ``og_field_resource_state`` from the live coalescer. + +Every stored row comes from ``queries.coalesced_winner_rows`` -- the same SQL +ranking the live list path hydrates with (through ``utils.coalesce_resources``) +-- run once per permission profile with that profile's exact source set. There +is no second ranking implementation, so cached state can't drift from live +coalescing. Rows hold the same materialized winning values the live path builds +its view from; like the live path, they are validated only when a list page is +built from them, so a stored value today's model rejects can't stop a rebuild +or a write. + +* ``refresh_resource_states`` recomputes the given resources inside the caller's + transaction: lock their rows, delete their state, and re-insert one row per + profile for those still listable. A resource that stopped being listable (e.g. + merged away) is left with no rows. +* ``rebuild_all_resource_state`` rebuilds everything in committed batches behind + the readiness flag, so new readers stay on the live path until it fully + succeeds. It never empties the table, so a reader already past the readiness + check sees each resource's old or new rows, never a missing one. +* ``is_resource_state_current`` checks a full rebuild by the same commit as the + running code completed. Any change to how rows are computed ships in a new + commit, so deploying it makes the table stale until that code rebuilds it. +* ``is_resource_state_ready`` checks readers may use it: current, and not + switched off by an operator (``disabled``, which rebuilds never clear). List + reads go through ``is_resource_state_ready_for_reads``, which asks at most + every few seconds. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from collections.abc import AsyncIterator, Collection +from contextlib import asynccontextmanager +from datetime import UTC, datetime, timedelta +from typing import Any, Final + +from sqlalchemy import Select, delete, distinct, func, insert, or_, select, text +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession, async_sessionmaker + +from stitch.api.observability.context import named_query +from stitch.api.settings import get_settings + +from ..model import OGFieldResourceState, OGFieldResourceStateStatus, ResourceModel +from ..model.og_field_resource_state import RESOURCE_STATE_STATUS_KEY +from ..model.oil_gas_field_source_value import ATTRIBUTE_NAMES, materialize_value +from ..queries import coalesced_winner_rows, resource_universe +from .permissions import PERMISSION_MASKS, visible_sources_for_mask + +logger = logging.getLogger(__name__) + +DEFAULT_BATCH_SIZE: Final[int] = 500 + +# Code version when the build sets no GIT_SHA (e.g. running the API directly on +# a developer machine). Deliberately not the local git HEAD: that would switch a +# local read model off on every commit. +UNKNOWN_CODE_VERSION: Final[str] = "unknown" + +# List reads re-check readiness at most this often per API process. The table +# is never emptied and each process only trusts rows built by its own commit, +# so a briefly out-of-date answer only delays switching between the read model +# and live reads; it never serves a broken page. The exception: a rebuild by +# other code while this process still serves can rewrite rows within that +# window, so deploys wait for the old revision to stop serving before +# rebuilding (see deployments/api/READ_MODEL.md). +READINESS_CACHE_SECONDS: Final[float] = 5.0 + +# Postgres advisory-lock key held for the duration of a full rebuild, so two +# rebuilds can't interleave. Arbitrary; only needs to be unique within the app. +_REBUILD_LOCK_KEY: Final[int] = 766_001 + +# A rebuild repeats passes over resources whose rows are missing or were computed +# by other code until none are left; this many passes without settling means +# something keeps writing with other code (e.g. an older API still serving). +_MAX_REBUILD_PASSES: Final[int] = 5 + +# How long a deploy's rebuild (``rebuild_resource_state_if_stale``) waits for a +# rebuild already holding the lock, and how often it checks. +DEFAULT_LOCK_WAIT: Final[timedelta] = timedelta(minutes=10) +_LOCK_POLL_SECONDS: Final[float] = 5.0 + +# Server-side TCP keepalives for the lock-holding connection: if the process +# running a rebuild vanishes without closing its connection (e.g. a cancelled CI +# job whose machine is torn down), Postgres ends the session -- releasing the +# lock -- after about a minute instead of the OS default of hours. +_KEEPALIVE_SETTINGS: Final[dict[str, str]] = { + "tcp_keepalives_idle": "30", + "tcp_keepalives_interval": "10", + "tcp_keepalives_count": "3", +} + + +class ResourceStateRebuildError(RuntimeError): + """A rebuild or repair could not run, or did not produce complete state.""" + + +def lock_resources_statement(resource_ids: Collection[int]) -> Select[tuple[int]]: + """``SELECT ... FOR NO KEY UPDATE`` over the resource rows, in ascending id + order. + + A stable order means two transactions locking overlapping sets (e.g. two + merges) always queue in the same order instead of deadlocking. + """ + return ( + select(ResourceModel.id) + .where(ResourceModel.id.in_(sorted(set(resource_ids)))) + .order_by(ResourceModel.id) + .with_for_update(key_share=True) + ) + + +async def lock_resources(session: AsyncSession, resource_ids: Collection[int]) -> None: + """Lock the resource rows until the caller's transaction ends. + + ``FOR NO KEY UPDATE`` is the lock Postgres already takes to update a + resource's non-key columns (e.g. a merge setting ``repointed_id``). It + conflicts with itself, so concurrent writers taking it queue, but not with + the key-share lock that inserting any row referencing a resource takes (a + membership, a merge candidate's items, ...). ``FOR UPDATE`` would conflict + with those: a writer holding one resource and waiting on a second, while an + insert holds a key-share lock on the second and waits on the first, + deadlock. + + Write actions take it at their start, before their own reads: the second of + two concurrent writers waits, then reads the first's committed changes -- + including a merge that repointed the resource -- so neither acts on stale + data or overwrites the other with a stale recompute. + ``refresh_resource_states`` locks again (a no-op when already held) to + protect callers like the full rebuild. SQLite has no row locks (it already + serializes writers), so this is a no-op there. + """ + if resource_ids: + await session.execute(lock_resources_statement(resource_ids)) + + +def _empty_state_row(resource_id: int, mask: int, code_version: str) -> dict[str, Any]: + """A row with every field (and its provenance) unset, stamped with the + code computing it.""" + return { + "resource_id": resource_id, + "permission_mask": mask, + **{field: None for field in ATTRIBUTE_NAMES}, + "provenance": {field: None for field in ATTRIBUTE_NAMES}, + "code_version": code_version, + } + + +async def _compute_state_rows( + session: AsyncSession, resource_ids: list[int] +) -> list[dict[str, Any]]: + """Rows for every ``(resource, profile)`` pair: one winner query per profile. + + Each winner's value is materialized exactly as the live path does before + building its view, and its source key recorded as provenance. Fields with no + winner stay unset. Every row records the commit computing it. + """ + code_version = current_code_version() + rows: list[dict[str, Any]] = [] + for mask in PERMISSION_MASKS: + by_id = {rid: _empty_state_row(rid, mask, code_version) for rid in resource_ids} + winners = await session.execute( + coalesced_winner_rows(resource_ids, visible_sources_for_mask(mask)) + ) + for rid, colname, value_text, value_num, value_json, source, _ in winners: + row = by_id[rid] + row[colname] = materialize_value( + colname, + value_text=value_text, + value_num=value_num, + value_json=value_json, + ) + row["provenance"][colname] = source + rows.extend(by_id.values()) + return rows + + +async def refresh_resource_states( + session: AsyncSession, resource_ids: Collection[int] +) -> list[int]: + """Make the given resources' stored state match authoritative data. + + Runs in the caller's transaction, so it commits (or rolls back) with the + write that triggered it. Resources that are no longer listable end up with + no rows; ids that don't exist are ignored. Returns the ids that were + recomputed (the listable ones), in ascending order. + + Its queries are labeled ``read_model.refresh`` (rather than the calling + write's label), so the cost the read model adds to writes and rebuilds can + be seen on its own. + """ + ids = sorted(set(resource_ids)) + if not ids: + return [] + with named_query("read_model.refresh"): + await lock_resources(session, ids) + await session.execute( + delete(OGFieldResourceState).where( + OGFieldResourceState.resource_id.in_(ids) + ) + ) + listable = sorted( + ( + await session.scalars( + resource_universe().where(ResourceModel.id.in_(ids)) + ) + ).all() + ) + if listable: + await session.execute( + insert(OGFieldResourceState), + await _compute_state_rows(session, listable), + ) + return listable + + +async def refresh_resource_state(session: AsyncSession, resource_id: int) -> None: + """``refresh_resource_states`` for a single resource.""" + await refresh_resource_states(session, [resource_id]) + + +def current_code_version() -> str: + """The commit the running code was built from (``GIT_SHA``).""" + return get_settings().git_sha or UNKNOWN_CODE_VERSION + + +async def _status(session: AsyncSession) -> Any | None: + return ( + await session.execute( + select( + OGFieldResourceStateStatus.ready, + OGFieldResourceStateStatus.code_version, + OGFieldResourceStateStatus.disabled, + ).where(OGFieldResourceStateStatus.key == RESOURCE_STATE_STATUS_KEY) + ) + ).one_or_none() + + +def _is_current(status: Any | None) -> bool: + return ( + status is not None + and status.ready is True + and status.code_version == current_code_version() + ) + + +async def is_resource_state_current(session: AsyncSession) -> bool: + """True when a full rebuild by this same commit completed (enabled or not). + + A missing status row, a cleared flag, or rows built by other code (an older + or newer deploy) all mean not current. + """ + return _is_current(await _status(session)) + + +async def is_resource_state_ready(session: AsyncSession) -> bool: + """True when readers may use the table: current, and not ``disabled``.""" + status = await _status(session) + return _is_current(status) and not status.disabled + + +async def is_resource_state_disabled(session: AsyncSession) -> bool: + """True when an operator has switched the read model off.""" + status = await _status(session) + return status is not None and status.disabled is True + + +class _ReadinessCache: + """This process's most recent readiness answer, and when it was checked.""" + + def __init__(self) -> None: + self.checked_at: float | None = None + self.ready = False + + def fresh(self) -> bool: + return ( + self.checked_at is not None + and time.monotonic() - self.checked_at < READINESS_CACHE_SECONDS + ) + + def remember(self, ready: bool) -> None: + self.checked_at = time.monotonic() + self.ready = ready + + def forget(self) -> None: + self.checked_at = None + + +_readiness_cache = _ReadinessCache() + + +async def is_resource_state_ready_for_reads(session: AsyncSession) -> bool: + """``is_resource_state_ready``, asked at most every few seconds per process. + + For the list/filter-options read path, so most requests don't pay a round + trip just to check readiness (see ``READINESS_CACHE_SECONDS``). + """ + if not _readiness_cache.fresh(): + _readiness_cache.remember(await is_resource_state_ready(session)) + return _readiness_cache.ready + + +def forget_cached_readiness() -> None: + """Make the next list read ask the database again.""" + _readiness_cache.forget() + + +async def _lock_status_row(session: AsyncSession) -> None: + """Hold the status row exclusively until the caller's transaction ends. + + ``--ids`` repairs hold it shared while they run, so this waits for any that + are running and makes new ones wait. No-op on SQLite (no row locks). + """ + await session.execute( + select(OGFieldResourceStateStatus.key) + .where(OGFieldResourceStateStatus.key == RESOURCE_STATE_STATUS_KEY) + .with_for_update() + ) + + +async def _set_ready(session: AsyncSession, ready: bool) -> None: + status = await session.get(OGFieldResourceStateStatus, RESOURCE_STATE_STATUS_KEY) + if status is None: + status = OGFieldResourceStateStatus(key=RESOURCE_STATE_STATUS_KEY, ready=ready) + session.add(status) + status.ready = ready + if ready: + status.rebuilt_at = datetime.now(UTC) + status.code_version = current_code_version() + + +async def _delete_unlistable_state(session: AsyncSession) -> None: + """Drop rows of resources that are no longer listable (e.g. merged away).""" + universe = resource_universe().subquery() + await session.execute( + delete(OGFieldResourceState).where( + OGFieldResourceState.resource_id.not_in(select(universe.c.resource_id)) + ) + ) + + +async def _integrity_problems(session: AsyncSession) -> list[str]: + """Ways the stored state fails to cover exactly the listable resources.""" + state = OGFieldResourceState + universe = resource_universe().subquery() + row_counts = ( + select(state.resource_id, func.count().label("n")) + .group_by(state.resource_id) + .subquery() + ) + incomplete = await session.scalar( + select(func.count()) + .select_from( + universe.outerjoin( + row_counts, row_counts.c.resource_id == universe.c.resource_id + ) + ) + .where(or_(row_counts.c.n.is_(None), row_counts.c.n != len(PERMISSION_MASKS))) + ) + unlistable = await session.scalar( + select(func.count(distinct(state.resource_id))).where( + state.resource_id.not_in(select(universe.c.resource_id)) + ) + ) + other_code = await session.scalar( + select(func.count(distinct(state.resource_id))).where( + state.resource_id.in_(select(universe.c.resource_id)), + or_( + state.code_version.is_(None), + state.code_version != current_code_version(), + ), + ) + ) + problems: list[str] = [] + if incomplete: + problems.append( + f"{incomplete} listable resource(s) lack a row for every permission profile" + ) + if other_code: + problems.append( + f"{other_code} listable resource(s) have rows computed by other code" + ) + if unlistable: + problems.append(f"{unlistable} non-listable resource(s) have stored state") + return problems + + +async def _resources_needing_rebuild(session: AsyncSession) -> list[int]: + """Listable resources lacking a row per profile computed by this code.""" + state = OGFieldResourceState + universe = resource_universe().subquery() + current_rows = ( + select(state.resource_id, func.count().label("n")) + .where(state.code_version == current_code_version()) + .group_by(state.resource_id) + .subquery() + ) + return list( + ( + await session.scalars( + select(universe.c.resource_id) + .select_from( + universe.outerjoin( + current_rows, + current_rows.c.resource_id == universe.c.resource_id, + ) + ) + .where( + or_( + current_rows.c.n.is_(None), + current_rows.c.n != len(PERMISSION_MASKS), + ) + ) + .order_by(universe.c.resource_id) + ) + ).all() + ) + + +async def _drop_vanished_clients_quickly(conn: AsyncConnection) -> None: + """Ask Postgres to notice quickly if this connection's client vanishes. + + Session settings, so they also apply if the pooled connection is reused -- + harmless (it only makes Postgres notice dead clients sooner). + """ + for name, value in _KEEPALIVE_SETTINGS.items(): + await conn.execute( + text("SELECT set_config(:name, :value, false)"), + {"name": name, "value": value}, + ) + + +@asynccontextmanager +async def _exclusive_rebuild( + session_factory: async_sessionmaker[AsyncSession], + *, + wait: timedelta | None = None, +) -> AsyncIterator[None]: + """Hold a Postgres advisory lock so only one full rebuild runs at a time. + + Without ``wait``, fails at once if another rebuild holds the lock; with it, + retries until the lock is free or ``wait`` has passed. Uses a session-level + lock on an autocommit connection, so it neither holds a transaction open for + the whole rebuild nor conflicts with the status row the rebuild updates. + No-op on other dialects. + """ + async with session_factory() as session: + if session.get_bind().dialect.name != "postgresql": + yield + return + conn = await session.connection( + execution_options={"isolation_level": "AUTOCOMMIT"} + ) + await _drop_vanished_clients_quickly(conn) + give_up_at = time.monotonic() + (wait.total_seconds() if wait else 0) + while not await conn.scalar( + text("SELECT pg_try_advisory_lock(:key)"), {"key": _REBUILD_LOCK_KEY} + ): + if time.monotonic() >= give_up_at: + raise ResourceStateRebuildError( + "Another resource-state rebuild is already running." + + (f" Waited {wait} for it to finish." if wait else "") + ) + logger.info("Waiting for another resource-state rebuild to finish.") + await asyncio.sleep(_LOCK_POLL_SECONDS) + try: + yield + finally: + await conn.execute( + text("SELECT pg_advisory_unlock(:key)"), {"key": _REBUILD_LOCK_KEY} + ) + + +async def rebuild_all_resource_state( + session_factory: async_sessionmaker[AsyncSession], + *, + batch_size: int = DEFAULT_BATCH_SIZE, +) -> int: + """Rebuild every listable resource's state; return how many were rebuilt. + + Fails at once if another rebuild is running. Each step commits on its own, + so the table is never locked for the whole run. Readiness is cleared first, + so new readers use the live path; it is set again only after every batch + commits, rows of no-longer-listable resources are removed, and the integrity + check passes. The table is never emptied: a reader that passed the readiness + check just before the rebuild started still finds a row for every resource. + A failure leaves readiness cleared; whatever rows remain are not served, and + the next run replaces them. Other refreshes (e.g. an ``--ids`` repair) may + run concurrently -- each locks a resource before recomputing it. + """ + _check_batch_size(batch_size) + async with _exclusive_rebuild(session_factory): + return await _rebuild_while_locked( + session_factory, batch_size, recompute_everything=True + ) + + +async def rebuild_resource_state_if_stale( + session_factory: async_sessionmaker[AsyncSession], + *, + batch_size: int = DEFAULT_BATCH_SIZE, + lock_wait: timedelta = DEFAULT_LOCK_WAIT, +) -> int | None: + """Rebuild unless the read model is already current for this code. + + The post-deploy step. Waits up to ``lock_wait`` for a rebuild already + running, then checks again: if that rebuild (by the same commit) made the + read model current, there is nothing to do. Returns how many resources were + rebuilt, or ``None`` if it was already current. + """ + _check_batch_size(batch_size) + async with _exclusive_rebuild(session_factory, wait=lock_wait): + async with session_factory() as session: + if await is_resource_state_current(session): + return None + return await _rebuild_while_locked( + session_factory, batch_size, recompute_everything=False + ) + + +def _check_batch_size(batch_size: int) -> None: + if batch_size < 1: + raise ValueError(f"batch_size must be at least 1, got {batch_size}") + + +async def _rebuild_while_locked( + session_factory: async_sessionmaker[AsyncSession], + batch_size: int, + *, + recompute_everything: bool, +) -> int: + """The rebuild itself; the caller holds the rebuild lock. Returns how many + resource recomputes it did. + + Clears readiness, then works in passes. The first pass recomputes every + listable resource (``recompute_everything``, e.g. after a restore) or only + those not already computed by this code (resuming a rebuild that stopped). + Each later pass recomputes resources whose rows were since rewritten by + other code -- an older API revision still serving writes, or an ``--ids`` + repair run from another image -- until none are left. + + The last check happens in the transaction that sets readiness, holding the + status row exclusively: a repair holds it shared while it runs, so this + waits for any repair already running to commit, and repairs that start + later wait for readiness and then see who built the table. Anything found + then gets another pass. Only when nothing is left are stray rows removed, + the result checked, and readiness set, all before the lock is released. + Forgets this process's cached readiness after each change commits, so reads + here switch immediately (other processes switch within the cache window). + """ + async with session_factory.begin() as session: + await _set_ready(session, False) + forget_cached_readiness() + + async with session_factory() as session: + if recompute_everything: + pending = list( + ( + await session.scalars( + resource_universe().order_by(ResourceModel.id) + ) + ).all() + ) + else: + pending = await _resources_needing_rebuild(session) + + recomputed = 0 + passes = 0 + while True: + if pending: + passes += 1 + if passes > _MAX_REBUILD_PASSES: + raise ResourceStateRebuildError( + f"After {_MAX_REBUILD_PASSES} passes, {len(pending)} " + "resource(s) still lack a row per profile computed by this " + "code: rows are missing, or keep being rewritten by other code " + "(is an older API version still serving writes?). The read " + "model was left not ready." + ) + if passes > 1: + logger.info( + "Pass %d: %d resource(s) were rewritten by other code during " + "the rebuild; recomputing them.", + passes, + len(pending), + ) + for start in range(0, len(pending), batch_size): + async with session_factory.begin() as session: + await refresh_resource_states( + session, pending[start : start + batch_size] + ) + logger.info( + "Rebuilt resource state for %d of %d resources.", + min(start + batch_size, len(pending)), + len(pending), + ) + recomputed += len(pending) + + async with session_factory.begin() as session: + await _lock_status_row(session) + pending = await _resources_needing_rebuild(session) + if pending: + continue + await _delete_unlistable_state(session) + if problems := await _integrity_problems(session): + raise ResourceStateRebuildError( + "Resource-state rebuild failed its integrity check: " + + "; ".join(problems) + ) + await _set_ready(session, True) + forget_cached_readiness() + return recomputed diff --git a/deployments/api/tests/conftest.py b/deployments/api/tests/conftest.py index a3045408..ee083b85 100644 --- a/deployments/api/tests/conftest.py +++ b/deployments/api/tests/conftest.py @@ -15,6 +15,7 @@ from stitch.ogsi.model import SOURCE_PRIORITY from stitch.api.db.config import UnitOfWork, get_uow +from stitch.api.db.read_model.state import forget_cached_readiness from stitch.api.db.model import ( OGFieldSourcePriority, StitchBase, @@ -104,6 +105,12 @@ def mock_uow(mock_session: MagicMock) -> MagicMock: return uow +@pytest.fixture(autouse=True) +def forget_cached_read_model_readiness(): + """Each test starts without another test's cached read-model readiness.""" + forget_cached_readiness() + + @pytest.fixture(autouse=True) def reset_dependency_overrides(): """Reset FastAPI dependency overrides and auth caches after each test.""" diff --git a/deployments/api/tests/db/read_model/__init__.py b/deployments/api/tests/db/read_model/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/deployments/api/tests/db/read_model/conftest.py b/deployments/api/tests/db/read_model/conftest.py new file mode 100644 index 00000000..80a4bf07 --- /dev/null +++ b/deployments/api/tests/db/read_model/conftest.py @@ -0,0 +1,69 @@ +"""Fixtures shared by the read-model tests.""" + +import uuid +from collections.abc import AsyncIterator + +import pytest +from sqlalchemy import insert, text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from stitch.api.db.model import OGFieldSourcePriority, StitchBase, UserModel +from stitch.api.db.read_model.state import rebuild_all_resource_state +from stitch.api.entities import User +from stitch.ogsi.model import SOURCE_PRIORITY + +from .dataset import Dataset, seed_dataset +from .postgres import POSTGRES_URL + + +@pytest.fixture +async def dataset(seeded_integration_session: AsyncSession, test_user: User) -> Dataset: + """The seeded dataset, with an empty (never rebuilt) read model.""" + return await seed_dataset(seeded_integration_session, test_user) + + +@pytest.fixture +async def rebuilt( + dataset: Dataset, + integration_session_factory: async_sessionmaker[AsyncSession], +) -> Dataset: + """The seeded dataset after a successful full rebuild.""" + await rebuild_all_resource_state(integration_session_factory) + return dataset + + +@pytest.fixture +async def pg_session_factory( + test_user_model: UserModel, +) -> AsyncIterator[async_sessionmaker[AsyncSession]]: + """Sessions on the opt-in Postgres, in a schema of their own (see ``postgres.py``). + + Only for tests marked ``requires_postgres``. + """ + assert POSTGRES_URL is not None, "mark the test with requires_postgres" + schema = f"stit766_test_{uuid.uuid4().hex[:8]}" + admin = create_async_engine(POSTGRES_URL) + async with admin.begin() as conn: + await conn.execute(text(f'CREATE SCHEMA "{schema}"')) + engine = create_async_engine( + POSTGRES_URL, connect_args={"options": f"-csearch_path={schema}"} + ) + try: + async with engine.begin() as conn: + await conn.run_sync(StitchBase.metadata.create_all) + await conn.execute( + insert(OGFieldSourcePriority), + [ + {"source": source, "priority": i + 1} + for i, source in enumerate(SOURCE_PRIORITY) + ], + ) + factory = async_sessionmaker(engine, expire_on_commit=False) + async with factory.begin() as session: + session.add(test_user_model) + yield factory + finally: + await engine.dispose() + async with admin.begin() as conn: + await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + await admin.dispose() diff --git a/deployments/api/tests/db/read_model/dataset.py b/deployments/api/tests/db/read_model/dataset.py new file mode 100644 index 00000000..32229d9b --- /dev/null +++ b/deployments/api/tests/db/read_model/dataset.py @@ -0,0 +1,220 @@ +"""A small deterministic dataset covering what can change a coalesced answer. + +Default source priority is rmi > wm > ccr > bc > alb > nor > gem > llm. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db.model import ( + MembershipModel, + MembershipStatus, + OGFieldResourceSourcePriority, + OGFieldResourceState, + ResourceModel, +) +from stitch.api.db.model.oil_gas_field_source_value import ATTRIBUTE_NAMES +from stitch.api.entities import User +from tests.utils import make_source_model + +ALPHA_OWNERS = [{"name": "Acme", "stake": 60.0}, {"name": "Globex", "stake": None}] +BETA_OPERATORS = [{"name": "Aramco", "stake": 100.0}] + + +@dataclass(frozen=True) +class Dataset: + alpha: int # public sources only; competing names; JSON + numeric values + beta: int # wm outranks gem for name + gamma: int # ccr outranks alb for basin; ccr alone supplies region + delta: int # wm only: an all-null shell for profiles without wm + epsilon: int # per-resource overrides invert the default order + empty: int # active membership whose source has no values + inactive: int # only an inactive membership: not listable + bare: int # no memberships: not listable + merged: int # current resource created by merging zeta_a + zeta_b + zeta_a: int # merged away + zeta_b: int # merged away + + @property + def listable(self) -> set[int]: + return { + self.alpha, + self.beta, + self.gamma, + self.delta, + self.epsilon, + self.empty, + self.merged, + } + + +async def new_resource(session: AsyncSession, user: User) -> int: + resource = ResourceModel.create(created_by=user) + session.add(resource) + await session.flush() + return resource.id + + +async def attach( + session: AsyncSession, + user: User, + resource_id: int, + source: str, + *, + status: MembershipStatus = MembershipStatus.ACTIVE, + **attrs: Any, +) -> int: + """Create a source with ``attrs`` and attach it directly (no action hooks).""" + model = make_source_model(source=source, created_by_id=user.id, **attrs) + session.add(model) + await session.flush() + session.add( + MembershipModel.create( + created_by=user, + resource_id=resource_id, + source=model.source, + source_pk=model.id, + status=status, + ) + ) + await session.flush() + return model.id + + +def override( + user: User, resource_id: int, source: str, source_pk: int, field: str, rank: int +) -> OGFieldResourceSourcePriority: + return OGFieldResourceSourcePriority.create( + created_by=user, + resource_id=resource_id, + source=source, + source_pk=source_pk, + colname=field, + priority=rank, + ) + + +async def seed_dataset(session: AsyncSession, user: User) -> Dataset: + """Seed the dataset and commit it, with no read-model state.""" + alpha = await new_resource(session, user) + await attach( + session, + user, + alpha, + "gem", + name="Alpha GEM", + country="USA", + basin="Permian", + latitude=31.5, + discovery_year=1921, + owners=ALPHA_OWNERS, + ) + await attach( + session, + user, + alpha, + "rmi", + name="Alpha RMI", + region="North America", + field_status="Producing", + ) + + beta = await new_resource(session, user) + await attach(session, user, beta, "gem", name="Beta GEM", country="SAU") + await attach( + session, + user, + beta, + "wm", + name="Beta WM", + operators=BETA_OPERATORS, + production_start_year=1951, + ) + + gamma = await new_resource(session, user) + await attach(session, user, gamma, "llm", name="Gamma", country="NOR") + await attach( + session, user, gamma, "ccr", region="North Sea", basin="Viking", fid_year=1970 + ) + await attach( + session, user, gamma, "alb", state_province="Rogaland", basin="Alb Basin" + ) + + delta = await new_resource(session, user) + await attach(session, user, delta, "wm", name="Delta", country="GBR") + + epsilon = await new_resource(session, user) + bc = await attach(session, user, epsilon, "bc", name="Epsilon BC", country="CAN") + nor = await attach(session, user, epsilon, "nor", name="Epsilon NOR") + wm = await attach(session, user, epsilon, "wm", basin="WM Basin") + ccr = await attach(session, user, epsilon, "ccr", basin="CCR Basin") + session.add_all( + [ + override(user, epsilon, "nor", nor, "name", 0), + override(user, epsilon, "bc", bc, "name", 1), + override(user, epsilon, "ccr", ccr, "basin", 0), + override(user, epsilon, "wm", wm, "basin", 1), + ] + ) + + empty = await new_resource(session, user) + await attach(session, user, empty, "gem") + + inactive = await new_resource(session, user) + await attach( + session, + user, + inactive, + "gem", + name="Inactive", + status=MembershipStatus.INACTIVE, + ) + + bare = await new_resource(session, user) + + zeta_a = await new_resource(session, user) + await attach(session, user, zeta_a, "gem", name="Zeta", country="IND") + zeta_b = await new_resource(session, user) + await attach(session, user, zeta_b, "rmi", name="Zeta RMI", longitude=78.9) + merged = await resource_actions.apply_resource_merge( + session, user, [zeta_a, zeta_b] + ) + assert merged.id is not None + + # The merge action maintains read-model state as it goes; start every test + # from an empty read model instead, as after the schema migration. + await session.execute(delete(OGFieldResourceState)) + await session.commit() + return Dataset( + alpha=alpha, + beta=beta, + gamma=gamma, + delta=delta, + epsilon=epsilon, + empty=empty, + inactive=inactive, + bare=bare, + merged=merged.id, + zeta_a=zeta_a, + zeta_b=zeta_b, + ) + + +async def dump_state(session: AsyncSession) -> dict[tuple[int, int], dict[str, Any]]: + """Every stored row as ``(resource_id, mask) -> {field..., provenance}``.""" + # Plain table rows, not ORM entities, so nothing stale comes from the + # session's identity map. + rows = (await session.execute(select(OGFieldResourceState.__table__))).all() + return { + (row.resource_id, row.permission_mask): { + **{field: getattr(row, field) for field in ATTRIBUTE_NAMES}, + "provenance": row.provenance, + } + for row in rows + } diff --git a/deployments/api/tests/db/read_model/postgres.py b/deployments/api/tests/db/read_model/postgres.py new file mode 100644 index 00000000..e3593ae0 --- /dev/null +++ b/deployments/api/tests/db/read_model/postgres.py @@ -0,0 +1,30 @@ +"""Postgres for the tests SQLite can't exercise. + +CI runs them against a throwaway Postgres (see ``.github/workflows/python-test.yml``). +Locally they're skipped unless ``STITCH_TEST_POSTGRES_URL`` points at a +disposable database:: + + STITCH_TEST_POSTGRES_URL=postgresql+psycopg://postgres:pg@127.0.0.1:55432/stitch + +Each test builds the tables in its own new schema and drops it afterwards +(``pg_session_factory`` in ``conftest.py``); existing tables are not touched. + +Setting ``STITCH_REQUIRE_POSTGRES_TESTS`` (as CI does) makes a missing URL an +error instead of a skip, so these tests can't silently stop running. +""" + +import os + +import pytest + +POSTGRES_URL = os.environ.get("STITCH_TEST_POSTGRES_URL") + +if os.environ.get("STITCH_REQUIRE_POSTGRES_TESTS") and not POSTGRES_URL: + raise RuntimeError( + "STITCH_REQUIRE_POSTGRES_TESTS is set but STITCH_TEST_POSTGRES_URL is not: " + "the Postgres tests would be skipped." + ) + +requires_postgres = pytest.mark.skipif( + not POSTGRES_URL, reason="set STITCH_TEST_POSTGRES_URL to run Postgres tests" +) diff --git a/deployments/api/tests/db/read_model/scenarios.py b/deployments/api/tests/db/read_model/scenarios.py new file mode 100644 index 00000000..94dbbfa5 --- /dev/null +++ b/deployments/api/tests/db/read_model/scenarios.py @@ -0,0 +1,138 @@ +"""Multi-step scenarios shared by the SQLite and Postgres read-model tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Collection + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db import og_field_source_actions as source_actions +from stitch.api.db.model import ResourceModel +from stitch.api.db.queries import resource_universe +from stitch.api.db.read_model import state +from stitch.api.db.read_model.permissions import PROFILES +from stitch.api.db.read_model.state import ( + _integrity_problems, + is_resource_state_ready, + rebuild_all_resource_state, + refresh_resource_states, +) +from stitch.api.entities import OGFieldQueryParams, User +from stitch.ogsi.model import GemSource, OGFieldResource, RMISource + +from tests.utils import make_source_record + +from .dataset import attach, dump_state, new_resource + +type SessionFactory = async_sessionmaker[AsyncSession] + + +async def rebuild_with_reads_and_writes_midway( + factory: SessionFactory, user: User, monkeypatch +) -> None: + """Pause a full rebuild after its first batch, read and write, then finish it. + + Six resources are rebuilt two per batch. While the rebuild waits before its + second batch: + + * a reader that passed the readiness check just before the rebuild started + lists resources straight from the table, and must get the full, correct + answer (not an empty or partial table); + * writes land on resources in the finished batch and in pending batches: + attaches, a merge across both, and a brand-new resource. + + The finished rebuild must be ready, complete, and match a fresh recompute. + """ + async with factory.begin() as session: + ids = [] + for i in range(6): + rid = await new_resource(session, user) + await attach(session, user, rid, "gem", name=f"Seed {i}", country="USA") + ids.append(rid) + await refresh_resource_states(session, ids) + + rebuild_paused = asyncio.Event() + resume_rebuild = asyncio.Event() + rebuild: asyncio.Task[int] | None = None + batches_started = 0 + real_refresh = state.refresh_resource_states + + async def pause_before_second_batch( + session: AsyncSession, resource_ids: Collection[int] + ) -> None: + nonlocal batches_started + # Only the rebuild's own batches pause; writes running meanwhile also + # refresh through this function and must go straight through. + if asyncio.current_task() is rebuild: + batches_started += 1 + if batches_started == 2: + rebuild_paused.set() + await resume_rebuild.wait() + await real_refresh(session, resource_ids) + + monkeypatch.setattr(state, "refresh_resource_states", pause_before_second_batch) + rebuild = asyncio.create_task(rebuild_all_resource_state(factory, batch_size=2)) + await asyncio.wait_for(rebuild_paused.wait(), timeout=30) + + # A reader that saw ready=true just before the rebuild began. + params = OGFieldQueryParams(page_size=50) + async with factory() as session: + cached = await resource_actions._query_read_model(session, params, mask=3) + live = await resource_actions._query_live(session, params, PROFILES[3]) + assert cached == live, "an in-flight reader saw a partially rebuilt table" + + done, pending = ids[:2], ids[2:] + async with factory.begin() as session: + await source_actions.create_and_attach_source( + session, + user, + RMISource( + name="Done RMI", country="CAN", source_record=make_source_record() + ), + done[0], + ) + async with factory.begin() as session: + await source_actions.create_and_attach_source( + session, + user, + RMISource( + name="Pending RMI", country="CAN", source_record=make_source_record() + ), + pending[0], + ) + async with factory.begin() as session: + await resource_actions.apply_resource_merge( + session, user, [done[1], pending[-1]] + ) + async with factory.begin() as session: + await resource_actions.create( + session, + user, + OGFieldResource( + source_data=[ + GemSource( + name="Created", + country="BRA", + source_record=make_source_record(), + ) + ], + constituents=frozenset(), + ), + ) + + resume_rebuild.set() + await asyncio.wait_for(rebuild, timeout=60) + + async with factory() as session: + assert await is_resource_state_ready(session) + assert await _integrity_problems(session) == [] + listable = set((await session.scalars(resource_universe())).all()) + assert {rid for rid, _ in await dump_state(session)} == listable + maintained = await dump_state(session) + all_ids = (await session.scalars(select(ResourceModel.id))).all() + await real_refresh(session, all_ids) + assert await dump_state(session) == maintained + await session.rollback() diff --git a/deployments/api/tests/db/read_model/test_concurrency_postgres.py b/deployments/api/tests/db/read_model/test_concurrency_postgres.py new file mode 100644 index 00000000..e21526ef --- /dev/null +++ b/deployments/api/tests/db/read_model/test_concurrency_postgres.py @@ -0,0 +1,421 @@ +"""Concurrent writes to the same resource, on a real Postgres (opt-in). + +Every write recomputes its resource's read-model rows in its own transaction. +Without the per-resource row lock, two concurrent writers would each recompute +from a snapshot missing the other's change, and the second insert would collide +with the first's rows. The lock is taken at the *start* of each write, so a +writer that waited reads the other's committed changes (e.g. a merge), and it +is ``FOR NO KEY UPDATE`` so it never conflicts with the key-share locks that +inserting rows referencing a resource takes (which would deadlock). SQLite +serializes all writers and has no row locks, so this can only be exercised on +Postgres (opt-in, see ``postgres.py``). +""" + +from __future__ import annotations + +import asyncio + +import pytest +from sqlalchemy import select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db import og_field_source_actions as source_actions +from stitch.api.db.errors import ResourceIntegrityError +from stitch.api.db.merge_candidate_actions import _fingerprint +from stitch.api.db.model import ( + MergeCandidateItemModel, + MergeCandidateModel, + ResourceModel, +) +from stitch.api.db.read_model import state +from stitch.api.db.read_model.state import refresh_resource_states +from stitch.api.entities import User +from stitch.ogsi.model import RMISource, WoodMacSource + +from tests.utils import make_source_record + +from .dataset import attach, dump_state, new_resource +from .postgres import requires_postgres +from .scenarios import rebuild_with_reads_and_writes_midway + +pytestmark = [pytest.mark.anyio, requires_postgres] + +# How long a writer must stay blocked to count as waiting on the other's lock. +_BLOCKED_FOR_SECONDS = 0.5 + +type SessionFactory = async_sessionmaker[AsyncSession] + + +async def _seed_resources(factory: SessionFactory, user: User, count: int) -> list[int]: + """``count`` listable resources, each with one gem source and current state.""" + async with factory.begin() as session: + ids = [] + for i in range(count): + rid = await new_resource(session, user) + await attach(session, user, rid, "gem", name=f"Seed {i}", country="USA") + ids.append(rid) + await refresh_resource_states(session, ids) + return ids + + +async def _assert_matches_recompute(factory: SessionFactory) -> None: + async with factory() as session: + maintained = await dump_state(session) + all_ids = (await session.scalars(select(ResourceModel.id))).all() + await refresh_resource_states(session, all_ids) + assert await dump_state(session) == maintained + await session.rollback() + + +async def _concurrent_attaches(factory: SessionFactory, user: User, rid: int) -> bool: + """Attach a source to ``rid`` in two overlapping transactions. + + The first finishes its write (and recompute) but holds its transaction open; + the second starts, then the first commits. Returns whether the second was + still waiting when the first committed. + """ + async with factory() as first, factory() as second: + await source_actions.create_and_attach_source( + first, + user, + RMISource( + name="From RMI", country="USA", source_record=make_source_record() + ), + rid, + ) + second_write = asyncio.create_task( + source_actions.create_and_attach_source( + second, + user, + WoodMacSource( + name="From WM", country="USA", source_record=make_source_record() + ), + rid, + ) + ) + await asyncio.sleep(_BLOCKED_FOR_SECONDS) + second_was_waiting = not second_write.done() + await first.commit() + await second_write + await second.commit() + return second_was_waiting + + +async def test_concurrent_writes_to_one_resource_serialize( + pg_session_factory: SessionFactory, test_user: User +): + [rid] = await _seed_resources(pg_session_factory, test_user, 1) + + assert await _concurrent_attaches(pg_session_factory, test_user, rid) + + # The second writer recomputed after seeing the first's committed source, so + # the stored state reflects both writes. + async with pg_session_factory() as session: + stored = await dump_state(session) + assert stored[(rid, 0)]["name"] == "From RMI" # wm hidden; rmi > gem + assert stored[(rid, 2)]["provenance"]["name"] == "rmi" # rmi > wm > gem + await _assert_matches_recompute(pg_session_factory) + + +async def test_without_the_lock_concurrent_writes_collide( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + """Shows concurrent writers really race: with the locks removed, two writers + that refresh together each miss the other's rows, and the second insert + collides with the first's.""" + [rid] = await _seed_resources(pg_session_factory, test_user, 1) + + async def no_lock(*_args: object, **_kwargs: object) -> None: + return None + + for module in (state, source_actions): + monkeypatch.setattr(module, "lock_resources", no_lock) + # Without locks neither writer waits on the other, so both always arrive: + # wait long enough that machine load can't make them miss each other. + _hold_writers_before_refresh(monkeypatch, _Rendezvous(timeout=30)) + with pytest.raises(IntegrityError): + await asyncio.gather( + _attach_in_own_transaction( + pg_session_factory, + test_user, + rid, + RMISource( + name="From RMI", country="USA", source_record=make_source_record() + ), + ), + _attach_in_own_transaction( + pg_session_factory, + test_user, + rid, + WoodMacSource( + name="From WM", country="USA", source_record=make_source_record() + ), + ), + ) + + +async def test_overlapping_merges_in_opposite_order_serialize( + pg_session_factory: SessionFactory, test_user: User +): + """Locking in id order means two merges sharing resources queue instead of + deadlocking, and the second sees the first already merged them.""" + a, b, c = await _seed_resources(pg_session_factory, test_user, 3) + + async with pg_session_factory() as first, pg_session_factory() as second: + # The second session already holds a (soon stale) copy of ``b``, as the + # merge-candidate approval path does before merging. + stale_b = await second.get(ResourceModel, b) + assert stale_b is not None and stale_b.repointed_id is None + + await resource_actions.apply_resource_merge(first, test_user, [a, b]) + second_merge = asyncio.create_task( + resource_actions.apply_resource_merge(second, test_user, [c, b]) + ) + await asyncio.sleep(_BLOCKED_FOR_SECONDS) + assert not second_merge.done() + await first.commit() + + with pytest.raises(ResourceIntegrityError, match="already been merged"): + await second_merge + assert stale_b.repointed_id is not None # re-read, not the stale copy + await second.rollback() + + async with pg_session_factory() as session: + stored_ids = {rid for rid, _ in await dump_state(session)} + assert a not in stored_ids and b not in stored_ids + assert c in stored_ids + await _assert_matches_recompute(pg_session_factory) + + +async def test_merge_candidate_creation_and_a_merge_in_opposite_order_do_not_deadlock( + pg_session_factory: SessionFactory, test_user: User +): + """Creating a merge candidate inserts its items in request order, each + taking a key-share lock on its resource; a merge locks its resources in id + order. A merge lock that conflicted with key-share (``FOR UPDATE``) would + deadlock here: the merge holds ``a`` and waits on ``b``, while the creation + holds ``b`` and waits on ``a``.""" + a, b = await _seed_resources(pg_session_factory, test_user, 2) + + async with pg_session_factory.begin() as creating: + candidate = MergeCandidateModel.create( + created_by=test_user, fingerprint=_fingerprint([b, a]) + ) + creating.add(candidate) + await creating.flush() + creating.add( + MergeCandidateItemModel( + merge_candidate_id=candidate.id, resource_id=b, position=0 + ) + ) + await creating.flush() + + async def merge() -> None: + async with pg_session_factory.begin() as merging: + await resource_actions.apply_resource_merge(merging, test_user, [a, b]) + + merging = asyncio.create_task(merge()) + await asyncio.sleep(_BLOCKED_FOR_SECONDS) + creating.add( + MergeCandidateItemModel( + merge_candidate_id=candidate.id, resource_id=a, position=1 + ) + ) + await creating.flush() + + await asyncio.wait_for(merging, timeout=30) + async with pg_session_factory() as session: + stored_ids = {rid for rid, _ in await dump_state(session)} + assert a not in stored_ids and b not in stored_ids + await _assert_matches_recompute(pg_session_factory) + + +class _Rendezvous: + """Holds each writer just before its read-model refresh until the other + writer gets there too (or ``timeout`` passes), so both have finished their + own inserts before either refreshes -- the interleaving most likely to + deadlock if a writer's lock conflicted with its inserts' key-share locks + (e.g. ``FOR UPDATE`` taken only at refresh time). + + The timeout lets a writer carry on when the other is (correctly) waiting on + its lock; under heavy machine load a short one can let the writers miss + each other, which makes a deadlock regression pass rather than fail. + """ + + def __init__(self, parties: int = 2, timeout: float = 1.0) -> None: + self._parties = parties + self._timeout = timeout + self._arrived = 0 + self._all_arrived = asyncio.Event() + + async def arrive(self) -> None: + self._arrived += 1 + if self._arrived >= self._parties: + self._all_arrived.set() + try: + await asyncio.wait_for(self._all_arrived.wait(), self._timeout) + except TimeoutError: + pass # The other writer is waiting on our lock; carry on. + + +def _hold_writers_before_refresh( + monkeypatch: pytest.MonkeyPatch, rendezvous: _Rendezvous +) -> None: + for module in (source_actions, resource_actions): + real_refresh = module.refresh_resource_state + + async def held(session: AsyncSession, resource_id: int, _real=real_refresh): + await rendezvous.arrive() + await _real(session, resource_id) + + monkeypatch.setattr(module, "refresh_resource_state", held) + + +async def _attach_in_own_transaction( + factory: SessionFactory, user: User, rid: int, source: RMISource | WoodMacSource +) -> None: + async with factory.begin() as session: + await source_actions.create_and_attach_source(session, user, source, rid) + + +async def _reverse_order_in_own_transaction( + factory: SessionFactory, user: User, rid: int, field: str +) -> None: + async with factory.begin() as session: + current = await resource_actions.field_source_values(session, rid, field) + await resource_actions.set_field_source_priority( + session, user, rid, field, [view.source_id for view in reversed(current)] + ) + + +async def _seed_resource_with_competing_sources( + factory: SessionFactory, user: User +) -> int: + """A resource whose gem and rmi sources both carry ``name`` and ``country``.""" + async with factory.begin() as session: + rid = await new_resource(session, user) + await attach(session, user, rid, "gem", name="GEM name", country="USA") + await attach(session, user, rid, "rmi", name="RMI name", country="CAN") + await refresh_resource_states(session, [rid]) + return rid + + +async def test_writers_that_both_insert_before_refreshing_do_not_deadlock( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + [rid] = await _seed_resources(pg_session_factory, test_user, 1) + _hold_writers_before_refresh(monkeypatch, _Rendezvous()) + + await asyncio.gather( + _attach_in_own_transaction( + pg_session_factory, + test_user, + rid, + RMISource( + name="From RMI", country="USA", source_record=make_source_record() + ), + ), + _attach_in_own_transaction( + pg_session_factory, + test_user, + rid, + WoodMacSource( + name="From WM", country="USA", source_record=make_source_record() + ), + ), + ) + + await _assert_matches_recompute(pg_session_factory) + + +async def test_attach_and_reprioritize_together_do_not_deadlock( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + rid = await _seed_resource_with_competing_sources(pg_session_factory, test_user) + _hold_writers_before_refresh(monkeypatch, _Rendezvous()) + + await asyncio.gather( + # No `name`, so the reprioritized field's eligible sources don't change. + _attach_in_own_transaction( + pg_session_factory, + test_user, + rid, + WoodMacSource(name=None, country="USA", source_record=make_source_record()), + ), + _reverse_order_in_own_transaction(pg_session_factory, test_user, rid, "name"), + ) + + async with pg_session_factory() as session: + stored = await dump_state(session) + assert stored[(rid, 0)]["name"] == "GEM name" # reordered: gem now first + await _assert_matches_recompute(pg_session_factory) + + +async def test_two_reprioritizations_together_do_not_deadlock( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + rid = await _seed_resource_with_competing_sources(pg_session_factory, test_user) + _hold_writers_before_refresh(monkeypatch, _Rendezvous()) + + # Different fields, so the two writes don't conflict for any other reason. + await asyncio.gather( + _reverse_order_in_own_transaction(pg_session_factory, test_user, rid, "name"), + _reverse_order_in_own_transaction( + pg_session_factory, test_user, rid, "country" + ), + ) + + async with pg_session_factory() as session: + stored = await dump_state(session) + assert stored[(rid, 0)]["name"] == "GEM name" + assert stored[(rid, 0)]["country"] == "USA" + await _assert_matches_recompute(pg_session_factory) + + +async def test_attach_waiting_on_a_merge_sees_the_resource_was_merged( + pg_session_factory: SessionFactory, test_user: User +): + """An attach queued behind a merge of its resource re-reads it after the + merge commits and refuses, instead of attaching to a merged-away resource + whose memberships would never be listed.""" + a, b = await _seed_resources(pg_session_factory, test_user, 2) + + async with pg_session_factory() as merging: + await resource_actions.apply_resource_merge(merging, test_user, [a, b]) + attach = asyncio.create_task( + _attach_in_own_transaction( + pg_session_factory, + test_user, + a, + RMISource( + name="Late", country="USA", source_record=make_source_record() + ), + ) + ) + await asyncio.sleep(_BLOCKED_FOR_SECONDS) + assert not attach.done() + await merging.commit() + + with pytest.raises(ResourceIntegrityError, match="has been merged"): + await attach + await _assert_matches_recompute(pg_session_factory) + + +async def test_reads_and_writes_midway_through_a_rebuild( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + await rebuild_with_reads_and_writes_midway( + pg_session_factory, test_user, monkeypatch + ) diff --git a/deployments/api/tests/db/read_model/test_mutations.py b/deployments/api/tests/db/read_model/test_mutations.py new file mode 100644 index 00000000..2f4f3ff0 --- /dev/null +++ b/deployments/api/tests/db/read_model/test_mutations.py @@ -0,0 +1,388 @@ +"""Every write that can change a cached answer keeps the read model in sync. + +Each test runs a real write action on a rebuilt dataset, commits, and checks the +maintained state equals a from-scratch recompute of every resource. The lock +tests check that recomputes take the resource lock before reading anything; the +lock itself only has an effect on Postgres (see ``test_concurrency_postgres``). +""" + +from __future__ import annotations + +from collections.abc import Collection +from typing import Any + +import pytest +from sqlalchemy import event, select, update +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + +from stitch.api.db import merge_candidate_actions +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db import og_field_source_actions as source_actions +from stitch.api.db.model import ( + MembershipModel, + OilGasFieldSourceValueModel, + ResourceModel, +) +from stitch.api.db.read_model import state +from stitch.api.db.read_model.permissions import PERMISSION_MASKS +from stitch.api.db.read_model.state import refresh_resource_states +from stitch.api.entities import MergeCandidateCreateRequest, User +from stitch.ogsi.model import GemSource, OGFieldResource, RMISource, WoodMacSource + +from tests.utils import make_source_record + +from .dataset import Dataset, dump_state + +pytestmark = pytest.mark.anyio + + +async def _assert_matches_recompute(session: AsyncSession) -> None: + """The committed, write-maintained state equals recomputing every resource.""" + maintained = await dump_state(session) + all_ids = (await session.scalars(select(ResourceModel.id))).all() + await refresh_resource_states(session, all_ids) + assert await dump_state(session) == maintained + await session.rollback() + + +def _resource_ids(stored: dict[tuple[int, int], Any]) -> set[int]: + return {rid for rid, _ in stored} + + +class TestWritesKeepStateInSync: + async def test_create_resource_with_sources( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + created = await resource_actions.create( + session, + test_user, + OGFieldResource( + source_data=[ + GemSource( + name="Created", + country="BRA", + source_record=make_source_record(), + ) + ], + constituents=frozenset(), + ), + ) + await session.commit() + + stored = await dump_state(session) + assert stored[(created.id, 0)]["name"] == "Created" + await _assert_matches_recompute(session) + + async def test_create_resource_without_sources_is_not_listed( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + created = await resource_actions.create( + session, + test_user, + OGFieldResource(source_data=[], constituents=frozenset()), + ) + await session.commit() + + assert created.id not in _resource_ids(await dump_state(session)) + await _assert_matches_recompute(session) + + async def test_create_and_attach_source( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + await source_actions.create_and_attach_source( + session, + test_user, + RMISource( + name="Beta RMI", country="SAU", source_record=make_source_record() + ), + rebuilt.beta, + ) + await session.commit() + + stored = await dump_state(session) + # rmi outranks both gem and wm. + for mask in PERMISSION_MASKS: + assert stored[(rebuilt.beta, mask)]["name"] == "Beta RMI" + await _assert_matches_recompute(session) + + async def test_attach_existing_source( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + alpha_gem = await session.scalar( + select(MembershipModel.source_pk).where( + MembershipModel.resource_id == rebuilt.alpha, + MembershipModel.source == "gem", + ) + ) + existing = (await source_actions.get_sources(session, [alpha_gem]))[0] + await source_actions.attach_sources_to_resource( + session, rebuilt.delta, [existing], test_user + ) + await session.commit() + + stored = await dump_state(session) + assert stored[(rebuilt.delta, 0)]["name"] == "Alpha GEM" + assert stored[(rebuilt.delta, 2)]["name"] == "Delta" + await _assert_matches_recompute(session) + + async def test_set_field_source_priority( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + current = await resource_actions.field_source_values( + session, rebuilt.alpha, "name" + ) + await resource_actions.set_field_source_priority( + session, + test_user, + rebuilt.alpha, + "name", + [view.source_id for view in reversed(current)], + ) + await session.commit() + + stored = await dump_state(session) + for mask in PERMISSION_MASKS: + assert stored[(rebuilt.alpha, mask)]["name"] == "Alpha GEM" + await _assert_matches_recompute(session) + + async def test_merge( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + merged = await resource_actions.apply_resource_merge( + session, test_user, [rebuilt.beta, rebuilt.alpha] + ) + await session.commit() + + stored_ids = _resource_ids(await dump_state(session)) + assert merged.id in stored_ids + assert rebuilt.alpha not in stored_ids + assert rebuilt.beta not in stored_ids + await _assert_matches_recompute(session) + + async def test_approve_merge_candidate( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + candidate = await merge_candidate_actions.create_merge_candidate( + session, + test_user, + MergeCandidateCreateRequest(resource_ids=[rebuilt.gamma, rebuilt.delta]), + ) + approved = await merge_candidate_actions.approve_merge_candidate( + session, test_user, candidate.id + ) + await session.commit() + + stored_ids = _resource_ids(await dump_state(session)) + assert approved.merged_resource_id in stored_ids + assert not {rebuilt.gamma, rebuilt.delta} & stored_ids + await _assert_matches_recompute(session) + + async def test_rolled_back_write_leaves_state_unchanged( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + ): + session = seeded_integration_session + before = await dump_state(session) + await source_actions.create_and_attach_source( + session, + test_user, + WoodMacSource( + name="Never committed", + country="USA", + source_record=make_source_record(), + ), + rebuilt.alpha, + ) + await session.rollback() + + assert await dump_state(session) == before + + +class TestInPlaceSourceValueEdit: + """No write path edits a stored source value today. If one is added, it changes + every resource the source is attached to, so it must refresh all of them -- + the per-resource hooks above don't cover it.""" + + async def test_must_refresh_every_resource_linked_to_the_source( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + # Zeta's gem source is attached to both the merged-away original and the + # merged resource. + zeta_gem = await session.scalar( + select(MembershipModel.source_pk).where( + MembershipModel.resource_id == rebuilt.zeta_a + ) + ) + await session.execute( + update(OilGasFieldSourceValueModel) + .where( + OilGasFieldSourceValueModel.source_pk == zeta_gem, + OilGasFieldSourceValueModel.colname == "country", + ) + .values(value_text="PAK") + ) + await session.commit() + + # Without a refresh the cache is stale. + assert (await dump_state(session))[(rebuilt.merged, 0)]["country"] == "IND" + + linked = ( + await session.scalars( + select(MembershipModel.resource_id).where( + MembershipModel.source_pk == zeta_gem + ) + ) + ).all() + assert set(linked) == {rebuilt.zeta_a, rebuilt.merged} + await refresh_resource_states(session, linked) + await session.commit() + + assert (await dump_state(session))[(rebuilt.merged, 0)]["country"] == "PAK" + await _assert_matches_recompute(session) + + +class TestLockBeforeRecompute: + """Option (a): runs everywhere, checks the lock is taken before any read. + + SQLite doesn't render row locks, so the lock call is recorded with a spy + alongside every SQL statement the engine executes. + """ + + @pytest.fixture + def events( + self, integration_engine: AsyncEngine, monkeypatch: pytest.MonkeyPatch + ) -> list[Any]: + recorded: list[Any] = [] + + def on_execute(conn, cursor, statement, parameters, context, executemany): + recorded.append(statement) + + event.listen( + integration_engine.sync_engine, "before_cursor_execute", on_execute + ) + + def spy(module: Any) -> None: + real = module.lock_resources + + async def recording_lock( + session: AsyncSession, resource_ids: Collection[int] + ) -> None: + recorded.append(("LOCK", sorted(set(resource_ids)))) + await real(session, resource_ids) + + monkeypatch.setattr(module, "lock_resources", recording_lock) + + spy(state) + spy(resource_actions) + spy(source_actions) + yield recorded + event.remove( + integration_engine.sync_engine, "before_cursor_execute", on_execute + ) + + async def test_refresh_locks_before_any_read( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset, events: list + ): + session = seeded_integration_session + events.clear() + await refresh_resource_states(session, [rebuilt.beta, rebuilt.alpha]) + + assert events[0] == ("LOCK", sorted([rebuilt.alpha, rebuilt.beta])) + assert "FROM og_field_resources" in events[1] + + async def test_merge_locks_originals_in_id_order_before_reading_them( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + events: list, + ): + session = seeded_integration_session + events.clear() + await resource_actions.apply_resource_merge( + session, test_user, [rebuilt.gamma, rebuilt.alpha, rebuilt.beta] + ) + + assert events[0] == ( + "LOCK", + sorted([rebuilt.alpha, rebuilt.beta, rebuilt.gamma]), + ) + locks = [e for e in events if isinstance(e, tuple)] + # Then one refresh covering the originals and the new resource. + assert len(locks) == 2 + assert set(locks[0][1]) < set(locks[1][1]) + + @pytest.mark.parametrize("action", ["attach", "reprioritize"]) + async def test_write_actions_lock_the_resource_before_anything_else( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + test_user: User, + events: list, + action: str, + ): + session = seeded_integration_session + if action == "attach": + events.clear() + await source_actions.create_and_attach_source( + session, + test_user, + RMISource( + name="New", country="SAU", source_record=make_source_record() + ), + rebuilt.beta, + ) + resource_id = rebuilt.beta + else: + current = await resource_actions.field_source_values( + session, rebuilt.alpha, "name" + ) + events.clear() + await resource_actions.set_field_source_priority( + session, + test_user, + rebuilt.alpha, + "name", + [view.source_id for view in reversed(current)], + ) + resource_id = rebuilt.alpha + + # Locking first -- before the action's own reads -- is what makes a + # writer that waited read the other's committed changes (e.g. a merge) + # instead of acting on what it read before. + assert events[0] == ("LOCK", [resource_id]) + state_writes = [ + e for e in events if isinstance(e, str) and "og_field_resource_state" in e + ] + assert state_writes, "the action did not touch the read model" diff --git a/deployments/api/tests/db/read_model/test_permissions.py b/deployments/api/tests/db/read_model/test_permissions.py new file mode 100644 index 00000000..e860287c --- /dev/null +++ b/deployments/api/tests/db/read_model/test_permissions.py @@ -0,0 +1,80 @@ +"""Exact permission-profile matching for the resource-state read model.""" + +from typing import get_args + +import pytest + +from stitch.api.db.read_model.permissions import ( + PERMISSION_MASKS, + PROFILES, + PUBLIC_SOURCES, + RESTRICTED_SOURCES, + read_model_mask, + visible_sources_for_mask, +) +from stitch.ogsi.model.types import OGSISrcKey + + +def test_source_universe_is_exactly_public_plus_restricted(): + """Adding a source key must be a deliberate change to the profiles. + + A new key fails closed on its own (no grant that holds it equals a profile), + but every caller holding it would silently lose the fast path. This test makes + that gap visible: decide whether the new key is public or restricted. + """ + assert frozenset(get_args(OGSISrcKey)) == PUBLIC_SOURCES | RESTRICTED_SOURCES + + +def test_public_and_restricted_are_disjoint(): + assert not PUBLIC_SOURCES & RESTRICTED_SOURCES + + +def test_profiles_are_the_four_documented_sets(): + assert PROFILES == { + 0: PUBLIC_SOURCES, + 1: PUBLIC_SOURCES | {"ccr"}, + 2: PUBLIC_SOURCES | {"wm"}, + 3: PUBLIC_SOURCES | {"wm", "ccr"}, + } + assert PERMISSION_MASKS == (0, 1, 2, 3) + + +@pytest.mark.parametrize("mask", PERMISSION_MASKS) +def test_exact_profile_maps_to_its_mask(mask: int): + assert read_model_mask(PROFILES[mask]) == mask + + +@pytest.mark.parametrize("mask", PERMISSION_MASKS) +def test_visible_sources_round_trip(mask: int): + assert read_model_mask(visible_sources_for_mask(mask)) == mask + + +def test_matching_ignores_order_and_duplicates(): + sources = [*sorted(PUBLIC_SOURCES), "wm", "gem", "wm"] + assert read_model_mask(sources) == 2 + + +def test_unscoped_caller_uses_live(): + assert read_model_mask(None) is None + + +def test_empty_grant_uses_live(): + assert read_model_mask(frozenset()) is None + + +@pytest.mark.parametrize("missing", sorted(PUBLIC_SOURCES)) +@pytest.mark.parametrize("restricted", [set(), {"wm"}, {"ccr"}, {"wm", "ccr"}]) +def test_partial_public_grant_uses_live(missing: str, restricted: set[str]): + """Never round a partial public grant to a nearby profile.""" + sources = (PUBLIC_SOURCES - {missing}) | restricted + assert read_model_mask(sources) is None + + +@pytest.mark.parametrize("restricted", [{"wm"}, {"ccr"}, {"wm", "ccr"}]) +def test_restricted_only_grant_uses_live(restricted: set[str]): + assert read_model_mask(restricted) is None + + +@pytest.mark.parametrize("mask", PERMISSION_MASKS) +def test_unknown_extra_key_uses_live(mask: int): + assert read_model_mask(PROFILES[mask] | {"not-a-source"}) is None diff --git a/deployments/api/tests/db/read_model/test_read_path.py b/deployments/api/tests/db/read_model/test_read_path.py new file mode 100644 index 00000000..7d1a5751 --- /dev/null +++ b/deployments/api/tests/db/read_model/test_read_path.py @@ -0,0 +1,419 @@ +"""list / filter-options served from the read model: parity with live, and fallback.""" + +from __future__ import annotations + +from collections.abc import Callable, Collection, Iterator +from typing import Any, get_args + +import pytest +from httpx import AsyncClient +from sqlalchemy import event, update +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db.model import OGFieldResourceState, OGFieldResourceStateStatus +from stitch.api.db.read_model import state +from stitch.api.db.read_model.permissions import PROFILES, PUBLIC_SOURCES +from stitch.api.db.read_model.state import ( + rebuild_all_resource_state, + refresh_resource_states, +) +from stitch.api.entities import OGFieldQueryParams, SortableField, User +from stitch.api.settings import get_settings + +from .dataset import Dataset, seed_dataset + +pytestmark = pytest.mark.anyio + +_PROFILE_IDS = {0: "public", 1: "public+ccr", 2: "public+wm", 3: "all"} + + +def _fail(name: str): + async def fail(*_args: Any, **_kwargs: Any): + raise AssertionError(f"{name} should not have been called") + + return fail + + +async def _assert_query_served_from_read_model( + session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, + params: OGFieldQueryParams, + licensed: frozenset[str], +) -> None: + """``query`` answers from the read model, identically to live coalescing.""" + live = await resource_actions._query_live(session, params, licensed) + with monkeypatch.context() as m: + m.setattr(resource_actions, "_query_live", _fail("_query_live")) + cached = await resource_actions.query(session, params, licensed) + assert cached == live + + +_P = OGFieldQueryParams + +# Each scenario builds its params from the dataset (some filter on its ids). +_SCENARIOS: dict[str, Callable[[Dataset], OGFieldQueryParams]] = { + "default": lambda _: _P(), + "second_page": lambda _: _P(page=2, page_size=2), + "past_last_page": lambda _: _P(page=10, page_size=5), + "q_name_case_insensitive": lambda _: _P(q="alpha"), + "q_region": lambda _: _P(q="Sea"), + "q_no_match": lambda _: _P(q="zzz"), + "filter_name": lambda _: _P(name="Gamma"), + "filter_name_local": lambda _: _P(name_local="Nowhere"), + "filter_basin_multi": lambda _: _P(basin=["Viking", "Alb Basin", "Permian"]), + "filter_state_province": lambda _: _P(state_province=["Rogaland"]), + "filter_region": lambda _: _P(region=["North Sea", "North America"]), + "filter_country_multi": lambda _: _P(country=["USA", "SAU", "GBR"]), + "filter_field_status": lambda _: _P(field_status=["Producing"]), + "filter_location_type": lambda _: _P(location_type="Onshore"), + "filter_production_conventionality": lambda _: _P( + production_conventionality="Mixed" + ), + "filter_primary_hydrocarbon_group": lambda _: _P( + primary_hydrocarbon_group=["Dry Gas"] + ), + "filter_id": lambda ds: _P(id=ds.gamma), + "filter_id_merged_away": lambda ds: _P(id=ds.zeta_a), + "combined": lambda _: _P( + country=["NOR", "USA", "CAN"], + q="a", + sort_by="basin", + sort_order="desc", + page=1, + page_size=2, + ), +} + +_SORT_FIELDS = [f for f in get_args(SortableField)] + + +class TestListParity: + @pytest.mark.parametrize("mask", PROFILES, ids=_PROFILE_IDS.get) + @pytest.mark.parametrize("scenario", list(_SCENARIOS)) + async def test_scenario_matches_live( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + monkeypatch: pytest.MonkeyPatch, + mask: int, + scenario: str, + ): + await _assert_query_served_from_read_model( + seeded_integration_session, + monkeypatch, + _SCENARIOS[scenario](rebuilt), + PROFILES[mask], + ) + + # Sorting is profile-independent SQL; 0 (most nulls) and 3 (most values) suffice. + @pytest.mark.parametrize("mask", [0, 3], ids=_PROFILE_IDS.get) + @pytest.mark.parametrize("sort_order", ["asc", "desc"]) + @pytest.mark.parametrize("sort_by", _SORT_FIELDS) + async def test_sort_matches_live( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + monkeypatch: pytest.MonkeyPatch, + mask: int, + sort_order: str, + sort_by: str, + ): + await _assert_query_served_from_read_model( + seeded_integration_session, + monkeypatch, + OGFieldQueryParams(sort_by=sort_by, sort_order=sort_order), + PROFILES[mask], + ) + + +class TestFilterOptionsParity: + @pytest.mark.parametrize("mask", PROFILES, ids=_PROFILE_IDS.get) + async def test_matches_live( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + monkeypatch: pytest.MonkeyPatch, + mask: int, + ): + session = seeded_integration_session + licensed = PROFILES[mask] + live = await resource_actions._filter_options_live(session, licensed) + monkeypatch.setattr( + resource_actions, "_filter_options_live", _fail("_filter_options_live") + ) + cached = await resource_actions.filter_options(session, licensed) + + assert cached == live + assert list(cached) == list(live) # same field order in the response + + +class TestFallback: + """When the read model can't answer exactly, ``query`` and ``filter_options`` + use live coalescing.""" + + @pytest.fixture + def read_model_forbidden(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + resource_actions, "_query_read_model", _fail("_query_read_model") + ) + monkeypatch.setattr( + resource_actions, + "resource_state_filter_option_query", + _fail("resource_state_filter_option_query"), + ) + + async def _assert_live_answer( + self, session: AsyncSession, licensed: frozenset[str] | None + ) -> None: + params = OGFieldQueryParams() + assert await resource_actions.query( + session, params, licensed + ) == await resource_actions._query_live(session, params, licensed) + assert await resource_actions.filter_options( + session, licensed + ) == await resource_actions._filter_options_live(session, licensed) + + async def test_never_rebuilt( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + read_model_forbidden: None, + ): + await self._assert_live_answer(seeded_integration_session, PROFILES[3]) + + async def test_not_ready_with_some_write_populated_rows( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + read_model_forbidden: None, + ): + """Rows written by write-path refreshes before the first full rebuild + must not be served -- they cover only some resources.""" + session = seeded_integration_session + await refresh_resource_states(session, [dataset.alpha]) + await session.commit() + + items, total = await resource_actions.query( + session, OGFieldQueryParams(), PROFILES[3] + ) + assert total == len(dataset.listable) + await self._assert_live_answer(session, PROFILES[3]) + + async def test_readiness_cleared_after_a_rebuild( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + read_model_forbidden: None, + ): + session = seeded_integration_session + await session.execute(update(OGFieldResourceStateStatus).values(ready=False)) + await session.commit() + await self._assert_live_answer(session, PROFILES[3]) + + async def test_read_model_switched_off( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + read_model_forbidden: None, + ): + session = seeded_integration_session + await session.execute(update(OGFieldResourceStateStatus).values(disabled=True)) + await session.commit() + await self._assert_live_answer(session, PROFILES[3]) + + async def test_read_model_built_by_other_code( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + read_model_forbidden: None, + monkeypatch: pytest.MonkeyPatch, + ): + """After a deploy, rows built by the previous commit are not used.""" + monkeypatch.setenv("GIT_SHA", "a-newer-commit") + get_settings.cache_clear() + await self._assert_live_answer(seeded_integration_session, PROFILES[3]) + + @pytest.mark.parametrize( + "licensed", + [ + None, + PUBLIC_SOURCES - {"gem"}, + PUBLIC_SOURCES - {"nor"} | {"wm", "ccr"}, + frozenset({"wm"}), + frozenset(), + ], + ids=[ + "unscoped", + "partial_public", + "partial_public_plus_both", + "wm_only", + "none", + ], + ) + async def test_unsupported_profile_while_ready( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + read_model_forbidden: None, + licensed: frozenset[str] | None, + ): + await self._assert_live_answer(seeded_integration_session, licensed) + + +class TestReadinessCache: + """List reads ask the database about readiness at most every few seconds.""" + + @pytest.fixture + def readiness_lookups(self, integration_engine: AsyncEngine) -> Iterator[list[str]]: + seen: list[str] = [] + + def record(conn, cursor, statement, params, context, executemany): + if "og_field_resource_state_status" in statement: + seen.append(statement) + + event.listen(integration_engine.sync_engine, "before_cursor_execute", record) + yield seen + event.remove(integration_engine.sync_engine, "before_cursor_execute", record) + + @pytest.fixture + def paths_taken(self, monkeypatch: pytest.MonkeyPatch) -> list[str]: + taken: list[str] = [] + for name, label in (("_query_read_model", "cached"), ("_query_live", "live")): + real = getattr(resource_actions, name) + + async def tracked(*args: Any, _real=real, _label=label, **kwargs: Any): + taken.append(_label) + return await _real(*args, **kwargs) + + monkeypatch.setattr(resource_actions, name, tracked) + return taken + + async def _list(self, session: AsyncSession) -> None: + await resource_actions.query(session, OGFieldQueryParams(), PROFILES[3]) + + async def test_back_to_back_requests_ask_once( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + readiness_lookups: list[str], + ): + session = seeded_integration_session + await self._list(session) + await self._list(session) + await resource_actions.filter_options(session, PROFILES[3]) + assert len(readiness_lookups) == 1 + + async def test_readiness_change_is_seen_once_the_window_passes( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + paths_taken: list[str], + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + await self._list(session) + await session.execute(update(OGFieldResourceStateStatus).values(ready=False)) + await session.commit() + + await self._list(session) # within the window: still the earlier answer + # The window has passed (shrinking it, rather than moving the shared + # clock the event loop also uses). + monkeypatch.setattr(state, "READINESS_CACHE_SECONDS", 0) + await self._list(session) + + assert paths_taken == ["cached", "cached", "live"] + + async def test_a_rebuild_in_this_process_switches_reads_immediately( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: async_sessionmaker[AsyncSession], + paths_taken: list[str], + ): + session = seeded_integration_session + await self._list(session) # not built yet: live, and remembered + await rebuild_all_resource_state(integration_session_factory) + await self._list(session) + + assert paths_taken == ["live", "cached"] + + async def test_a_read_during_the_rebuild_is_forgotten_when_it_finishes( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: async_sessionmaker[AsyncSession], + paths_taken: list[str], + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + real_refresh = state.refresh_resource_states + + async def list_during_first_batch( + batch_session: AsyncSession, resource_ids: Collection[int] + ) -> list[int]: + if not paths_taken: + await self._list(session) # mid-rebuild: live, and remembered + return await real_refresh(batch_session, resource_ids) + + monkeypatch.setattr(state, "refresh_resource_states", list_during_first_batch) + await rebuild_all_resource_state(integration_session_factory) + await self._list(session) + + assert paths_taken == ["live", "cached"] + + +class TestServedFromStoredState: + async def test_list_reads_stored_rows_and_detail_stays_live( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + """Overwrite a stored value: the list reflects it (proving it reads the + table), the detail endpoint does not (it always coalesces live).""" + session = seeded_integration_session + await session.execute( + update(OGFieldResourceState) + .where(OGFieldResourceState.resource_id == rebuilt.gamma) + .values(name="Stored") + ) + await session.commit() + + items, _ = await resource_actions.query( + session, OGFieldQueryParams(id=rebuilt.gamma), PROFILES[3] + ) + assert [item.data.name for item in items] == ["Stored"] + + detail = await resource_actions.get(session, rebuilt.gamma, PROFILES[3]) + assert detail.view is not None and detail.view.name == "Gamma" + + +class TestHttpResponsesUnchanged: + """The serialized responses are byte-for-byte the same once the read model + serves them (e.g. provenance key order survives JSONB storage).""" + + async def test_list_and_filter_options( + self, + integration_client: AsyncClient, + integration_session_factory: async_sessionmaker[AsyncSession], + test_user: User, + monkeypatch: pytest.MonkeyPatch, + ): + async with integration_session_factory() as session: + await seed_dataset(session, test_user) + requests = [ + "/oil-gas-fields/?page_size=50", + "/oil-gas-fields/?sort_by=latitude&sort_order=desc&page_size=3", + "/oil-gas-fields/?country=USA&country=NOR&q=a", + "/oil-gas-fields/filter-options", + ] + + live = [(await integration_client.get(url)) for url in requests] + await rebuild_all_resource_state(integration_session_factory) + monkeypatch.setattr(resource_actions, "_query_live", _fail("_query_live")) + monkeypatch.setattr( + resource_actions, "_filter_options_live", _fail("_filter_options_live") + ) + cached = [(await integration_client.get(url)) for url in requests] + + for url, before, after in zip(requests, live, cached, strict=True): + assert before.status_code == after.status_code == 200, url + assert after.text == before.text, url diff --git a/deployments/api/tests/db/read_model/test_read_path_postgres.py b/deployments/api/tests/db/read_model/test_read_path_postgres.py new file mode 100644 index 00000000..13ce2d7c --- /dev/null +++ b/deployments/api/tests/db/read_model/test_read_path_postgres.py @@ -0,0 +1,108 @@ +"""Read-model list / filter-options output on a real Postgres (opt-in). + +JSONB stores object keys in its own order, unlike SQLite's JSON, so only +Postgres shows whether serialized responses stay identical to live coalescing +(e.g. the provenance map's key order). +""" + +from __future__ import annotations + +import json +from typing import get_args + +import pytest +from sqlalchemy import delete, event +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db.model import OGFieldResourceState +from stitch.api.db.read_model.permissions import PROFILES +from stitch.api.db.read_model.state import rebuild_all_resource_state +from stitch.api.entities import OGFieldQueryParams, SortableField, User +from stitch.api.observability.context import query_name_var + +from .dataset import seed_dataset +from .postgres import requires_postgres + +pytestmark = [pytest.mark.anyio, requires_postgres] + +_SCENARIOS = [ + OGFieldQueryParams(page_size=50), + OGFieldQueryParams(q="a", page=2, page_size=2), + OGFieldQueryParams(country=["USA", "NOR"], basin=["Viking", "Permian"]), + *( + OGFieldQueryParams(sort_by=sort_by, sort_order=sort_order) + for sort_by in get_args(SortableField) + for sort_order in ("asc", "desc") + ), +] + + +@pytest.mark.parametrize("mask", PROFILES) +async def test_serialized_output_matches_live( + pg_session_factory: async_sessionmaker[AsyncSession], test_user: User, mask: int +): + async with pg_session_factory() as session: + await seed_dataset(session, test_user) + await rebuild_all_resource_state(pg_session_factory) + + licensed = PROFILES[mask] + async with pg_session_factory() as session: + for params in _SCENARIOS: + cached, cached_total = await resource_actions._query_read_model( + session, params, mask + ) + live, live_total = await resource_actions._query_live( + session, params, licensed + ) + assert cached_total == live_total + assert json.dumps( + [item.model_dump(mode="json") for item in cached] + ) == json.dumps([item.model_dump(mode="json") for item in live]), params + + assert json.dumps( + await resource_actions.filter_options(session, licensed) + ) == json.dumps(await resource_actions._filter_options_live(session, licensed)) + + +async def test_rows_removed_after_the_page_is_picked_do_not_break_the_list( + pg_session_factory: async_sessionmaker[AsyncSession], test_user: User +): + """A concurrent merge removes a resource's rows (via its refresh) while a list + request is running. Landing right after the request has picked its page, it + must not make the request fail: the page's rows are read together with the + choice of page, from one snapshot.""" + async with pg_session_factory() as session: + dataset = await seed_dataset(session, test_user) + await rebuild_all_resource_state(pg_session_factory) + engine = pg_session_factory.kw["bind"].sync_engine + labeled_statements: list[str] = [] + removed: list[int] = [] + + def remove_rows_after_page_is_picked(conn, cursor, statement, params, ctx, many): + label = query_name_var.get() + if label is None or removed: + return + labeled_statements.append(label) + if len(labeled_statements) == 2: # the count, then the page selection + with engine.connect() as other: + other.execute( + delete(OGFieldResourceState).where( + OGFieldResourceState.resource_id == dataset.gamma + ) + ) + other.commit() + removed.append(dataset.gamma) + + event.listen(engine, "after_cursor_execute", remove_rows_after_page_is_picked) + try: + async with pg_session_factory() as session: + items, total = await resource_actions.query( + session, OGFieldQueryParams(page_size=50), PROFILES[3] + ) + finally: + event.remove(engine, "after_cursor_execute", remove_rows_after_page_is_picked) + + assert removed == [dataset.gamma] + assert total == len(dataset.listable) + assert dataset.gamma in [item.id for item in items] diff --git a/deployments/api/tests/db/read_model/test_rebuild_and_repair_postgres.py b/deployments/api/tests/db/read_model/test_rebuild_and_repair_postgres.py new file mode 100644 index 00000000..a216a25e --- /dev/null +++ b/deployments/api/tests/db/read_model/test_rebuild_and_repair_postgres.py @@ -0,0 +1,159 @@ +"""A full rebuild finishing while an ``--ids`` repair runs, on a real Postgres (opt-in). + +The repair holds the status row shared for its whole transaction; the rebuild's +last check, and setting readiness, hold it exclusively. Whichever gets there +first, the table is never marked ready while holding a repair's rows from other +code: a repair already running is waited for and its rows redone, and a repair +starting later waits for readiness, then refuses because other code built the +table. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable + +import pytest +from sqlalchemy import select, update +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from stitch.api.db.model import OGFieldResourceState, OGFieldResourceStateStatus +from stitch.api.db.read_model import rebuild as rebuild_cli +from stitch.api.db.read_model import state +from stitch.api.db.read_model.state import ( + ResourceStateRebuildError, + rebuild_all_resource_state, +) +from stitch.api.entities import User +from stitch.api.settings import get_settings + +from .dataset import attach, new_resource +from .postgres import requires_postgres + +pytestmark = [pytest.mark.anyio, requires_postgres] + +type SessionFactory = async_sessionmaker[AsyncSession] + + +@pytest.fixture(autouse=True) +def _rebuild_runs_commit_b(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GIT_SHA", "commit-b") + get_settings.cache_clear() + + +async def _seed(factory: SessionFactory, user: User) -> list[int]: + ids = [] + async with factory.begin() as session: + for i in range(3): + rid = await new_resource(session, user) + await attach(session, user, rid, "gem", name=f"Seed {i}", country="USA") + ids.append(rid) + return ids + + +def _pause_at_final_check( + monkeypatch: pytest.MonkeyPatch, *, after_locking: bool +) -> tuple[asyncio.Event, asyncio.Event]: + """Pause the rebuild the first time it reaches its final, locked check, + just before (or just after) it takes the status row lock.""" + paused, resume = asyncio.Event(), asyncio.Event() + real_lock: Callable[[AsyncSession], Awaitable[None]] = state._lock_status_row + calls = 0 + + async def pausing_lock(session: AsyncSession) -> None: + nonlocal calls + calls += 1 + if calls > 1: + await real_lock(session) + return + if after_locking: + await real_lock(session) + paused.set() + await resume.wait() + if not after_locking: + await real_lock(session) + + monkeypatch.setattr(state, "_lock_status_row", pausing_lock) + return paused, resume + + +async def _row_versions(factory: SessionFactory) -> set[str | None]: + async with factory() as session: + return set( + (await session.scalars(select(OGFieldResourceState.code_version))).all() + ) + + +async def _status(factory: SessionFactory): + async with factory() as session: + return ( + await session.execute(select(OGFieldResourceStateStatus.__table__)) + ).one() + + +async def test_repair_already_running_is_waited_for_and_redone( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + ids = await _seed(pg_session_factory, test_user) + paused, resume = _pause_at_final_check(monkeypatch, after_locking=False) + rebuild = asyncio.create_task(rebuild_all_resource_state(pg_session_factory)) + await asyncio.wait_for(paused.wait(), timeout=30) + + # An --ids repair from the old image (commit-a): its guard allows it, since + # the rebuild cleared readiness, and it rewrites a resource's rows with the + # old code but hasn't committed yet. + async with pg_session_factory.begin() as repair: + await rebuild_cli._refuse_if_served_by_other_code(repair) + await repair.execute( + update(OGFieldResourceState) + .where(OGFieldResourceState.resource_id == ids[0]) + .values(name="Old code's answer", code_version="commit-a") + ) + resume.set() + await asyncio.sleep(0.5) + assert not rebuild.done(), "the rebuild didn't wait for the running repair" + + # The repair committed; the rebuild sees its rows and redoes them. + assert await asyncio.wait_for(rebuild, timeout=30) == len(ids) + 1 + status = await _status(pg_session_factory) + assert (status.ready, status.code_version) == (True, "commit-b") + assert await _row_versions(pg_session_factory) == {"commit-b"} + async with pg_session_factory() as session: + assert ( + await session.scalar( + select(OGFieldResourceState.name).where( + OGFieldResourceState.resource_id == ids[0], + OGFieldResourceState.permission_mask == 0, + ) + ) + == "Seed 0" + ) + + +async def test_repair_starting_later_waits_then_refuses( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + ids = await _seed(pg_session_factory, test_user) + paused, resume = _pause_at_final_check(monkeypatch, after_locking=True) + rebuild = asyncio.create_task(rebuild_all_resource_state(pg_session_factory)) + await asyncio.wait_for(paused.wait(), timeout=30) + + # The rebuild holds the status row; an --ids repair from the old image + # starts now. + monkeypatch.setattr(rebuild_cli, "current_code_version", lambda: "commit-a") + repair = asyncio.create_task(rebuild_cli._repair(pg_session_factory, [ids[0]])) + await asyncio.sleep(0.5) + assert not repair.done(), "the repair didn't wait for the rebuild to finish" + + resume.set() + await asyncio.wait_for(rebuild, timeout=30) + with pytest.raises(ResourceStateRebuildError, match="built by code commit-b"): + await asyncio.wait_for(repair, timeout=30) + + status = await _status(pg_session_factory) + assert (status.ready, status.code_version) == (True, "commit-b") + assert await _row_versions(pg_session_factory) == {"commit-b"} diff --git a/deployments/api/tests/db/read_model/test_rebuild_lock_postgres.py b/deployments/api/tests/db/read_model/test_rebuild_lock_postgres.py new file mode 100644 index 00000000..ab9ab49b --- /dev/null +++ b/deployments/api/tests/db/read_model/test_rebuild_lock_postgres.py @@ -0,0 +1,193 @@ +"""The lock that keeps full rebuilds from overlapping, on a real Postgres (opt-in). + +A deploy's ``--if-stale`` rebuild waits for a rebuild already holding the lock, +then rebuilds only if the read model is still stale. The lock-holding connection +asks Postgres to notice a vanished client quickly, so a rebuild whose machine +disappears (e.g. a cancelled CI job) doesn't hold the lock for hours. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import timedelta +from typing import Any + +import pytest +from sqlalchemy import select, text +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession, async_sessionmaker + +from stitch.api.db.model import OGFieldResourceStateStatus +from stitch.api.db.read_model import state +from stitch.api.db.read_model.state import ( + ResourceStateRebuildError, + is_resource_state_current, + rebuild_all_resource_state, + rebuild_resource_state_if_stale, +) +from stitch.api.entities import User + +from .dataset import attach, new_resource +from .postgres import requires_postgres + +pytestmark = [pytest.mark.anyio, requires_postgres] + +type SessionFactory = async_sessionmaker[AsyncSession] + + +@pytest.fixture(autouse=True) +def _poll_quickly(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(state, "_LOCK_POLL_SECONDS", 0.05) + + +async def _seed(factory: SessionFactory, user: User, count: int = 3) -> None: + async with factory.begin() as session: + for i in range(count): + rid = await new_resource(session, user) + await attach(session, user, rid, "gem", name=f"Seed {i}", country="USA") + + +@asynccontextmanager +async def _another_rebuild_holds_the_lock( + factory: SessionFactory, +) -> AsyncIterator[AsyncConnection]: + engine = factory.kw["bind"] + async with engine.connect() as conn: + conn = await conn.execution_options(isolation_level="AUTOCOMMIT") + await conn.execute( + text("SELECT pg_advisory_lock(:key)"), {"key": state._REBUILD_LOCK_KEY} + ) + try: + yield conn + finally: + await conn.execute( + text("SELECT pg_advisory_unlock(:key)"), + {"key": state._REBUILD_LOCK_KEY}, + ) + + +async def _rebuilt_at(factory: SessionFactory): + async with factory() as session: + return await session.scalar(select(OGFieldResourceStateStatus.rebuilt_at)) + + +async def test_if_stale_waits_for_a_running_rebuild_then_rebuilds( + pg_session_factory: SessionFactory, test_user: User +): + await _seed(pg_session_factory, test_user) + + async with _another_rebuild_holds_the_lock(pg_session_factory): + rebuild = asyncio.create_task( + rebuild_resource_state_if_stale( + pg_session_factory, lock_wait=timedelta(seconds=30) + ) + ) + await asyncio.sleep(0.5) + assert not rebuild.done(), "gave up instead of waiting for the lock" + + assert await asyncio.wait_for(rebuild, timeout=30) == 3 + async with pg_session_factory() as session: + assert await is_resource_state_current(session) + + +async def test_if_stale_skips_when_the_rebuild_it_waited_for_made_it_current( + pg_session_factory: SessionFactory, test_user: User +): + await _seed(pg_session_factory, test_user) + + async with _another_rebuild_holds_the_lock(pg_session_factory): + rebuild = asyncio.create_task( + rebuild_resource_state_if_stale( + pg_session_factory, lock_wait=timedelta(seconds=30) + ) + ) + await asyncio.sleep(0.2) + # The other rebuild (same commit) finishes while we wait. + async with pg_session_factory.begin() as session: + await state._set_ready(session, True) + finished_at = await _rebuilt_at(pg_session_factory) + + assert await asyncio.wait_for(rebuild, timeout=30) is None + assert await _rebuilt_at(pg_session_factory) == finished_at + + +async def test_if_stale_gives_up_after_waiting_its_limit( + pg_session_factory: SessionFactory, test_user: User +): + await _seed(pg_session_factory, test_user) + + async with _another_rebuild_holds_the_lock(pg_session_factory): + with pytest.raises(ResourceStateRebuildError, match="already running"): + await rebuild_resource_state_if_stale( + pg_session_factory, lock_wait=timedelta(seconds=0.5) + ) + + +async def test_full_rebuild_still_fails_straight_away( + pg_session_factory: SessionFactory, test_user: User +): + await _seed(pg_session_factory, test_user) + + async with _another_rebuild_holds_the_lock(pg_session_factory): + started = time.monotonic() + with pytest.raises(ResourceStateRebuildError, match="already running"): + await rebuild_all_resource_state(pg_session_factory) + assert time.monotonic() - started < 1 + + +async def test_lock_connection_asks_postgres_to_drop_vanished_clients_quickly( + pg_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +): + """The connection holding the lock has short server-side TCP keepalives, so + Postgres ends its session (releasing the lock) soon after its client + vanishes, rather than after the OS default of hours.""" + await _seed(pg_session_factory, test_user) + applied: list[dict[str, object]] = [] + real_configure = state._drop_vanished_clients_quickly + + async def spy_configure(conn: AsyncConnection) -> None: + await real_configure(conn) + rows = await conn.execute( + text( + "SELECT name, setting FROM pg_settings WHERE name IN " + "('tcp_keepalives_idle', 'tcp_keepalives_interval', " + "'tcp_keepalives_count')" + ) + ) + settings: dict[str, object] = {name: value for name, value in rows} + settings["pid"] = await conn.scalar(text("SELECT pg_backend_pid()")) + applied.append(settings) + + lock_holders: list[int] = [] + real_rebuild = state._rebuild_while_locked + + async def record_lock_holder( + factory: SessionFactory, batch_size: int, **kwargs: Any + ) -> int: + async with factory() as session: + lock_holders.append( + await session.scalar( + text( + "SELECT pid FROM pg_locks WHERE locktype = 'advisory' " + "AND objid = :key AND granted" + ), + {"key": state._REBUILD_LOCK_KEY}, + ) + ) + return await real_rebuild(factory, batch_size, **kwargs) + + monkeypatch.setattr(state, "_drop_vanished_clients_quickly", spy_configure) + monkeypatch.setattr(state, "_rebuild_while_locked", record_lock_holder) + await rebuild_all_resource_state(pg_session_factory) + + [settings] = applied + assert {k: v for k, v in settings.items() if k != "pid"} == { + "tcp_keepalives_idle": "30", + "tcp_keepalives_interval": "10", + "tcp_keepalives_count": "3", + } + assert lock_holders == [settings["pid"]] diff --git a/deployments/api/tests/db/read_model/test_resource_state_schema.py b/deployments/api/tests/db/read_model/test_resource_state_schema.py new file mode 100644 index 00000000..a4c852c3 --- /dev/null +++ b/deployments/api/tests/db/read_model/test_resource_state_schema.py @@ -0,0 +1,147 @@ +"""Schema and drift guards for ``og_field_resource_state``.""" + +import pytest +from sqlalchemy import JSON, Float, Integer, String, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from stitch.api.db.model import OGFieldResourceState, ResourceModel +from stitch.api.db.model.oil_gas_field_source_value import ( + ATTRIBUTE_KINDS, + ATTRIBUTE_NAMES, + ValueKind, +) +from stitch.api.entities import FILTER_OPTION_FIELDS, User + +_STATE_TABLE = OGFieldResourceState.__table__ +_NON_VALUE_COLUMNS = {"resource_id", "permission_mask", "provenance", "code_version"} + +_EXPECTED_TYPE: dict[ValueKind, type] = { + ValueKind.TEXT: String, + ValueKind.INT: Integer, + ValueKind.FLOAT: Float, + ValueKind.JSON: JSON, +} + + +def test_value_columns_match_flattened_attributes(): + value_columns = set(_STATE_TABLE.columns.keys()) - _NON_VALUE_COLUMNS + assert value_columns == set(ATTRIBUTE_NAMES) + + +@pytest.mark.parametrize("field", ATTRIBUTE_NAMES) +def test_value_column_type_follows_attribute_kind(field: str): + column = _STATE_TABLE.columns[field] + assert isinstance(column.type, _EXPECTED_TYPE[ATTRIBUTE_KINDS[field]]) + assert column.nullable + + +def test_primary_key_is_resource_and_mask(): + assert [c.name for c in _STATE_TABLE.primary_key.columns] == [ + "resource_id", + "permission_mask", + ] + + +def test_resource_fk_cascades(): + [fk] = _STATE_TABLE.c.resource_id.foreign_keys + assert fk.target_fullname == "og_field_resources.id" + assert fk.ondelete == "CASCADE" + + +def _index_columns() -> set[tuple[str, ...]]: + return {tuple(c.name for c in index.columns) for index in _STATE_TABLE.indexes} + + +@pytest.mark.parametrize("field", FILTER_OPTION_FIELDS) +def test_each_filter_option_field_has_a_mask_leading_index(field: str): + assert ("permission_mask", field) in _index_columns() + + +def test_single_profile_and_default_sort_indexes(): + indexes = _index_columns() + assert ("permission_mask", "resource_id") in indexes + assert ("permission_mask", "name", "resource_id") in indexes + + +async def _new_resource(session: AsyncSession, user: User) -> int: + resource = ResourceModel.create(created_by=user) + session.add(resource) + await session.flush() + return resource.id + + +@pytest.mark.anyio +async def test_typed_values_round_trip( + seeded_integration_session: AsyncSession, test_user: User +): + session = seeded_integration_session + rid = await _new_resource(session, test_user) + owners = [{"name": "Acme", "share": 0.6}, {"name": "Globex", "share": 0.4}] + provenance = {name: None for name in ATTRIBUTE_NAMES} | { + "name": "gem", + "owners": "wm", + "latitude": "rmi", + "discovery_year": "ccr", + } + session.add( + OGFieldResourceState( + resource_id=rid, + permission_mask=3, + name="Alpha", + owners=owners, + operators=["Acme Ops"], + latitude=12.5, + discovery_year=1987, + provenance=provenance, + ) + ) + await session.commit() + session.expunge_all() + + row = await session.get(OGFieldResourceState, (rid, 3)) + assert row is not None + assert row.name == "Alpha" + assert row.owners == owners + assert row.operators == ["Acme Ops"] + assert row.latitude == 12.5 + assert row.discovery_year == 1987 + assert row.provenance == provenance + assert row.country is None + + +@pytest.mark.anyio +async def test_unset_json_value_is_sql_null( + seeded_integration_session: AsyncSession, test_user: User +): + """An unset JSON field must be SQL NULL, not JSON ``null``, so it matches + ``IS NULL`` / ``NULLS LAST`` the same way the live path's absent row does.""" + session = seeded_integration_session + rid = await _new_resource(session, test_user) + session.add( + OGFieldResourceState( + resource_id=rid, permission_mask=0, owners=None, provenance={} + ) + ) + await session.commit() + + is_null = await session.scalar( + text( + "SELECT owners IS NULL FROM og_field_resource_state WHERE resource_id = :id" + ).bindparams(id=rid) + ) + assert is_null + + +@pytest.mark.anyio +@pytest.mark.parametrize("mask", [-1, 4]) +async def test_permission_mask_out_of_range_is_rejected( + seeded_integration_session: AsyncSession, test_user: User, mask: int +): + session = seeded_integration_session + rid = await _new_resource(session, test_user) + session.add( + OGFieldResourceState(resource_id=rid, permission_mask=mask, provenance={}) + ) + with pytest.raises(IntegrityError): + await session.flush() diff --git a/deployments/api/tests/db/read_model/test_state.py b/deployments/api/tests/db/read_model/test_state.py new file mode 100644 index 00000000..1a9cbd3c --- /dev/null +++ b/deployments/api/tests/db/read_model/test_state.py @@ -0,0 +1,1100 @@ +"""Building, refreshing, and rebuilding the resource-state read model.""" + +from __future__ import annotations + +import logging +from collections.abc import Collection +from typing import Any + +import pytest +from pydantic import ValidationError +from sqlalchemy import delete, insert, select, update +from sqlalchemy.dialects import postgresql +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from stitch.api.db.model import ( + MembershipModel, + MembershipStatus, + OGFieldResourceState, + OGFieldResourceStateStatus, +) +from stitch.api.db import og_field_resource_actions as resource_actions +from stitch.api.db.model.oil_gas_field_source_value import ATTRIBUTE_NAMES +from stitch.api.db.queries import resource_universe +from stitch.api.db.read_model import rebuild as rebuild_cli +from stitch.api.db.read_model import state +from stitch.api.db.read_model.permissions import ( + PERMISSION_MASKS, + PROFILES, + visible_sources_for_mask, +) +from stitch.api.db.read_model.state import ( + ResourceStateRebuildError, + current_code_version, + rebuild_resource_state_if_stale, + _integrity_problems, + is_resource_state_current, + is_resource_state_ready, + lock_resources_statement, + rebuild_all_resource_state, + refresh_resource_state, + refresh_resource_states, +) +from stitch.api.db.utils import coalesce_resources, resource_to_list_item_view +from stitch.api.entities import OGFieldQueryParams, User +from stitch.api.settings import get_settings +from stitch.ogsi.model import OGFieldListItemView +from stitch.ogsi.model.og_field import OilGasFieldBase + +from .dataset import ( + ALPHA_OWNERS, + BETA_OPERATORS, + Dataset, + attach, + dump_state, + new_resource, +) +from .scenarios import rebuild_with_reads_and_writes_midway + +pytestmark = pytest.mark.anyio + +type SessionFactory = async_sessionmaker[AsyncSession] + + +def _cached_item(resource_id: int, stored: dict[str, Any]) -> OGFieldListItemView: + return OGFieldListItemView( + id=resource_id, + data=OilGasFieldBase(**{field: stored[field] for field in ATTRIBUTE_NAMES}), + provenance=stored["provenance"], + ) + + +async def _corrupt(session: AsyncSession, resource_id: int) -> None: + await session.execute( + update(OGFieldResourceState) + .where(OGFieldResourceState.resource_id == resource_id) + .values(name="corrupted") + ) + await session.commit() + + +class TestBuilderParity: + async def test_rows_exist_exactly_for_listable_resources( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + universe = set((await session.scalars(resource_universe())).all()) + stored = await dump_state(session) + + assert universe == rebuilt.listable + assert set(stored) == { + (rid, mask) for rid in rebuilt.listable for mask in PERMISSION_MASKS + } + + @pytest.mark.parametrize("mask", PERMISSION_MASKS) + async def test_stored_state_matches_live_coalescing( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset, mask: int + ): + session = seeded_integration_session + stored = await dump_state(session) + ids = sorted(rebuilt.listable) + live = await coalesce_resources(session, ids, visible_sources_for_mask(mask)) + + for rid in ids: + row = stored[(rid, mask)] + # Raw values, including their Python types (int vs float, JSON lists). + for field in ATTRIBUTE_NAMES: + winner = live[rid].provenance[field] + expected = None if winner is None else winner[0] + assert row[field] == expected, (rid, mask, field) + assert type(row[field]) is type(expected), (rid, mask, field) + # The list item the read path will build from the row. + assert _cached_item(rid, row) == resource_to_list_item_view(live[rid]) + + async def test_restricted_values_only_appear_in_profiles_that_see_them( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + stored = await dump_state(seeded_integration_session) + + beta = {mask: stored[(rebuilt.beta, mask)] for mask in PERMISSION_MASKS} + assert [beta[m]["name"] for m in (0, 1, 2, 3)] == [ + "Beta GEM", + "Beta GEM", + "Beta WM", + "Beta WM", + ] + assert [beta[m]["provenance"]["name"] for m in (0, 1, 2, 3)] == [ + "gem", + "gem", + "wm", + "wm", + ] + assert beta[0]["operators"] is None + assert beta[2]["operators"] == BETA_OPERATORS + + gamma = {mask: stored[(rebuilt.gamma, mask)] for mask in PERMISSION_MASKS} + assert [gamma[m]["basin"] for m in (0, 1, 2, 3)] == [ + "Alb Basin", + "Viking", + "Alb Basin", + "Viking", + ] + assert gamma[0]["region"] is None + assert gamma[0]["provenance"]["region"] is None + assert gamma[1]["region"] == "North Sea" + + async def test_wm_only_resource_is_a_null_shell_without_wm( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + stored = await dump_state(seeded_integration_session) + for mask in (0, 1): + row = stored[(rebuilt.delta, mask)] + assert all(row[field] is None for field in ATTRIBUTE_NAMES) + assert set(row["provenance"].values()) == {None} + assert stored[(rebuilt.delta, 2)]["name"] == "Delta" + + async def test_per_resource_overrides_win( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + stored = await dump_state(seeded_integration_session) + epsilon = {mask: stored[(rebuilt.epsilon, mask)] for mask in PERMISSION_MASKS} + assert {epsilon[m]["name"] for m in PERMISSION_MASKS} == {"Epsilon NOR"} + # By default wm outranks ccr; the override puts ccr first. + assert [epsilon[m]["basin"] for m in (0, 1, 2, 3)] == [ + None, + "CCR Basin", + "WM Basin", + "CCR Basin", + ] + + async def test_json_and_numeric_values_round_trip( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + stored = await dump_state(seeded_integration_session) + alpha = stored[(rebuilt.alpha, 0)] + assert alpha["owners"] == ALPHA_OWNERS + assert alpha["latitude"] == 31.5 + assert alpha["discovery_year"] == 1921 + assert isinstance(alpha["discovery_year"], int) + assert alpha["provenance"]["owners"] == "gem" + assert alpha["name"] == "Alpha RMI" + + async def test_merged_resource_combines_its_originals( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + stored = await dump_state(seeded_integration_session) + merged = stored[(rebuilt.merged, 0)] + assert merged["name"] == "Zeta RMI" + assert merged["country"] == "IND" + assert merged["longitude"] == 78.9 + + +class TestTargetedRefresh: + async def test_recomputes_only_requested_resources( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + pristine = await dump_state(session) + await _corrupt(session, rebuilt.alpha) + await _corrupt(session, rebuilt.beta) + + await refresh_resource_state(session, rebuilt.alpha) + await session.commit() + + after = await dump_state(session) + for mask in PERMISSION_MASKS: + assert after[(rebuilt.alpha, mask)] == pristine[(rebuilt.alpha, mask)] + assert after[(rebuilt.beta, mask)]["name"] == "corrupted" + + async def test_resource_that_stops_being_listable_loses_its_rows( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + await session.execute( + update(MembershipModel) + .where(MembershipModel.resource_id == rebuilt.delta) + .values(status=MembershipStatus.INACTIVE) + ) + await refresh_resource_state(session, rebuilt.delta) + await session.commit() + + assert not any(rid == rebuilt.delta for rid, _ in await dump_state(session)) + + async def test_unlisted_and_unknown_ids_get_no_rows( + self, seeded_integration_session: AsyncSession, dataset: Dataset + ): + session = seeded_integration_session + recomputed = await refresh_resource_states( + session, + [dataset.inactive, dataset.bare, dataset.zeta_a, 999_999, dataset.alpha], + ) + await session.commit() + assert recomputed == [dataset.alpha] + assert {rid for rid, _ in await dump_state(session)} == {dataset.alpha} + + async def test_does_not_mark_the_read_model_ready( + self, seeded_integration_session: AsyncSession, dataset: Dataset + ): + """Refreshes (e.g. an ``--ids`` repair) may populate rows before the + first full rebuild; that partial state must stay invisible.""" + session = seeded_integration_session + await refresh_resource_states(session, [dataset.alpha, dataset.beta]) + await session.commit() + + assert len(await dump_state(session)) == 2 * len(PERMISSION_MASKS) + assert not await is_resource_state_current(session) + + +class TestLocking: + def test_locks_resource_rows_in_ascending_id_order(self): + sql = str( + lock_resources_statement([30, 10, 20, 10]).compile( + dialect=postgresql.dialect(), + compile_kwargs={"literal_binds": True}, + ) + ) + assert "FROM og_field_resources" in sql + assert "IN (10, 20, 30)" in sql + # NO KEY UPDATE, not UPDATE: it mustn't conflict with the key-share + # locks that inserting rows referencing a resource takes. + assert sql.rstrip().endswith( + "ORDER BY og_field_resources.id FOR NO KEY UPDATE" + ), sql + + +def _deploy(monkeypatch: pytest.MonkeyPatch, git_sha: str | None) -> None: + """Run the rest of the test as code built from ``git_sha`` (None: no SHA).""" + if git_sha is None: + monkeypatch.delenv("GIT_SHA", raising=False) + else: + monkeypatch.setenv("GIT_SHA", git_sha) + get_settings.cache_clear() + + +async def _set_disabled(session: AsyncSession, disabled: bool) -> None: + """What the runbook's switch-off / switch-on SQL does.""" + await session.execute(update(OGFieldResourceStateStatus).values(disabled=disabled)) + await session.commit() + + +class TestDisabled: + """An operator's switch-off sticks: rebuilds keep the rows current but never + turn the read model back on; only clearing ``disabled`` does.""" + + async def test_disabled_read_model_is_not_ready( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + await _set_disabled(session, True) + assert not await is_resource_state_ready(session) + + await _set_disabled(session, False) + assert await is_resource_state_ready(session) + + async def test_rebuilds_do_not_re_enable_it( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + await _set_disabled(session, True) + + await rebuild_all_resource_state(integration_session_factory) + _deploy(monkeypatch, "a-later-commit") + await rebuild_resource_state_if_stale(integration_session_factory) + + status = ( + await session.execute(select(OGFieldResourceStateStatus.__table__)) + ).one() + assert (status.ready, status.code_version, status.disabled) == ( + True, + "a-later-commit", + True, + ) + assert not await is_resource_state_ready(session) + + async def test_if_stale_does_not_rebuild_a_current_disabled_table( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + ): + """Disabled isn't stale: deploys of the same commit don't rebuild.""" + session = seeded_integration_session + await _set_disabled(session, True) + built_at = await session.scalar(select(OGFieldResourceStateStatus.rebuilt_at)) + + assert ( + await rebuild_resource_state_if_stale(integration_session_factory) is None + ) + assert ( + await session.scalar(select(OGFieldResourceStateStatus.rebuilt_at)) + == built_at + ) + + async def test_cli_says_the_rebuilt_read_model_is_still_disabled( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + monkeypatch: pytest.MonkeyPatch, + integration_engine: AsyncEngine, + integration_session_factory: SessionFactory, + caplog: pytest.LogCaptureFixture, + ): + async def _no_dispose() -> None: + return None + + monkeypatch.setattr( + rebuild_cli, "get_session_factory", lambda: integration_session_factory + ) + monkeypatch.setattr(rebuild_cli, "dispose_engine", _no_dispose) + monkeypatch.setattr(rebuild_cli, "get_engine", lambda: integration_engine) + caplog.set_level(logging.INFO) + await _set_disabled(seeded_integration_session, True) + + await rebuild_cli._run(ids=None, batch_size=500) + + assert "disabled" in caplog.text + + async def test_ids_repair_from_other_code_is_refused_on_a_disabled_table( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + integration_engine: AsyncEngine, + ): + """Switched off isn't the same as not ready: switching it back on would + serve the repair's rows under the commit that built the rest, and + ``--if-stale`` from that commit would see nothing to redo.""" + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + await _set_disabled(session, True) + before = await dump_state(session) + + async def _no_dispose() -> None: + return None + + monkeypatch.setattr( + rebuild_cli, "get_session_factory", lambda: integration_session_factory + ) + monkeypatch.setattr(rebuild_cli, "dispose_engine", _no_dispose) + monkeypatch.setattr(rebuild_cli, "get_engine", lambda: integration_engine) + _deploy(monkeypatch, "commit-b") + with pytest.raises(ResourceStateRebuildError, match="built by code commit-a"): + await rebuild_cli._run(ids=[dataset.alpha], batch_size=1) + assert await dump_state(session) == before + + # Back under commit A and switched back on: every served row is A's. + _deploy(monkeypatch, "commit-a") + assert ( + await rebuild_resource_state_if_stale(integration_session_factory) is None + ) + await _set_disabled(session, False) + assert await is_resource_state_ready(session) + assert set().union(*(await _row_versions(session)).values()) == {"commit-a"} + + +class TestCodeVersion: + """The read model is used only by the commit whose rebuild produced it.""" + + def test_code_version_is_the_builds_git_sha(self, monkeypatch: pytest.MonkeyPatch): + _deploy(monkeypatch, "abc123") + assert current_code_version() == "abc123" + + def test_code_version_without_a_git_sha_is_unknown( + self, monkeypatch: pytest.MonkeyPatch + ): + _deploy(monkeypatch, None) + assert current_code_version() == "unknown" + + async def test_rebuild_records_the_commit_that_built_it( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + + session = seeded_integration_session + status = ( + await session.execute(select(OGFieldResourceStateStatus.__table__)) + ).one() + assert status.code_version == "commit-a" + assert await is_resource_state_current(session) + + async def test_deploying_other_code_turns_the_read_model_off_until_it_rebuilds( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + + _deploy(monkeypatch, "commit-b") + assert not await is_resource_state_current(session) + + await rebuild_all_resource_state(integration_session_factory) + assert await is_resource_state_current(session) + + # Code from the previous deploy, still running during cutover, doesn't + # trust rows built by the new commit. + _deploy(monkeypatch, "commit-a") + assert not await is_resource_state_current(session) + + async def test_ready_row_without_a_code_version_is_not_current( + self, seeded_integration_session: AsyncSession + ): + """E.g. a status row marked ready by hand rather than by a rebuild.""" + session = seeded_integration_session + session.add( + OGFieldResourceStateStatus( + key="resource_state", ready=True, code_version=None + ) + ) + await session.commit() + assert not await is_resource_state_current(session) + + async def test_targeted_refresh_does_not_change_the_code_version( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + + _deploy(monkeypatch, "commit-b") + await refresh_resource_states(session, [dataset.alpha]) + await session.commit() + assert not await is_resource_state_current(session) + + +async def _row_versions(session: AsyncSession) -> dict[int, set[str | None]]: + """Each resource's set of row code versions.""" + rows = await session.execute( + select(OGFieldResourceState.resource_id, OGFieldResourceState.code_version) + ) + versions: dict[int, set[str | None]] = {} + for rid, version in rows: + versions.setdefault(rid, set()).add(version) + return versions + + +class TestRowCodeVersions: + """Each row records the commit that computed it; a rebuild redoes any row + computed by other code, so an older API still writing during a cutover + can't leave its rows in a table marked ready for the new commit.""" + + async def test_rows_carry_the_code_that_computed_them( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + assert set().union(*(await _row_versions(session)).values()) == {"commit-a"} + + _deploy(monkeypatch, "commit-b") + await refresh_resource_states(session, [dataset.alpha]) + await session.commit() + versions = await _row_versions(session) + assert versions[dataset.alpha] == {"commit-b"} + assert versions[dataset.beta] == {"commit-a"} + + async def test_rebuild_redoes_rows_other_code_rewrote_during_it( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + """GPT's cutover scenario: the new commit's rebuild has done a resource, + then the old revision (still serving) writes it with the old rules.""" + session = seeded_integration_session + _deploy(monkeypatch, "commit-b") + first = min(dataset.listable) + real_refresh = state.refresh_resource_states + batches = 0 + + async def old_code_writes_after_the_first_batch( + batch_session: AsyncSession, resource_ids: Collection[int] + ) -> list[int]: + nonlocal batches + batches += 1 + if batches == 2: + with monkeypatch.context() as old_code: + old_code.setattr(state, "current_code_version", lambda: "commit-a") + async with integration_session_factory.begin() as other: + await real_refresh(other, [first]) + return await real_refresh(batch_session, resource_ids) + + monkeypatch.setattr( + state, "refresh_resource_states", old_code_writes_after_the_first_batch + ) + await rebuild_all_resource_state(integration_session_factory, batch_size=1) + + assert await is_resource_state_current(session) + versions = await _row_versions(session) + assert versions[first] == {"commit-b"} + assert set().union(*versions.values()) == {"commit-b"} + + async def test_rebuild_gives_up_if_other_code_keeps_writing( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + _deploy(monkeypatch, "commit-b") + real_refresh = state.refresh_resource_states + + async def other_code_rewrites_a_row_every_time( + batch_session: AsyncSession, resource_ids: Collection[int] + ) -> list[int]: + recomputed = await real_refresh(batch_session, resource_ids) + await batch_session.execute( + update(OGFieldResourceState) + .where(OGFieldResourceState.resource_id == min(dataset.listable)) + .values(code_version="commit-a") + ) + return recomputed + + monkeypatch.setattr( + state, "refresh_resource_states", other_code_rewrites_a_row_every_time + ) + with pytest.raises(ResourceStateRebuildError, match="other code"): + await rebuild_all_resource_state(integration_session_factory) + assert not await is_resource_state_current(session) + + async def test_if_stale_picks_up_where_a_failed_run_stopped( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + _deploy(monkeypatch, "commit-b") + real_compute = state._compute_state_rows + computed: list[int] = [] + fail_on_batch: int | None = 2 + + async def recording_compute(session: AsyncSession, ids: list[int]): + if fail_on_batch is not None and len(computed) // 2 + 1 == fail_on_batch: + raise RuntimeError("boom") + computed.extend(ids) + return await real_compute(session, ids) + + monkeypatch.setattr(state, "_compute_state_rows", recording_compute) + with pytest.raises(RuntimeError, match="boom"): + await rebuild_resource_state_if_stale( + integration_session_factory, batch_size=2 + ) + done_first_time = list(computed) + assert len(done_first_time) == 2 + + computed.clear() + fail_on_batch = None + rebuilt_count = await rebuild_resource_state_if_stale( + integration_session_factory, batch_size=2 + ) + + assert sorted(computed) == sorted(dataset.listable - set(done_first_time)) + assert rebuilt_count == len(dataset.listable) - 2 + assert await is_resource_state_current(seeded_integration_session) + + async def test_full_rebuild_recomputes_rows_already_from_this_code( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + ): + """A plain full rebuild is what you run after a restore or manual SQL: + rows stamped with this commit can still be wrong then.""" + session = seeded_integration_session + await _corrupt(session, rebuilt.alpha) + await rebuild_all_resource_state(integration_session_factory) + assert (await dump_state(session))[(rebuilt.alpha, 0)]["name"] == "Alpha RMI" + + +class TestFullRebuild: + async def test_missing_status_row_is_not_ready( + self, seeded_integration_session: AsyncSession + ): + assert not await is_resource_state_current(seeded_integration_session) + + async def test_status_row_present_but_false_is_not_ready( + self, seeded_integration_session: AsyncSession + ): + session = seeded_integration_session + session.add(OGFieldResourceStateStatus(key="resource_state", ready=False)) + await session.commit() + assert not await is_resource_state_current(session) + + async def test_marks_ready_and_records_time( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + ): + count = await rebuild_all_resource_state(integration_session_factory) + + session = seeded_integration_session + status = ( + await session.execute(select(OGFieldResourceStateStatus.__table__)) + ).one() + assert count == len(dataset.listable) + assert status.ready + assert status.rebuilt_at is not None + assert await is_resource_state_current(session) + + async def test_is_idempotent( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + ): + first = await dump_state(seeded_integration_session) + await rebuild_all_resource_state(integration_session_factory) + assert await dump_state(seeded_integration_session) == first + + async def test_batch_size_does_not_change_the_result( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + ): + expected = await dump_state(seeded_integration_session) + await rebuild_all_resource_state(integration_session_factory, batch_size=1) + assert await dump_state(seeded_integration_session) == expected + + async def test_rejects_non_positive_batch_size( + self, integration_session_factory: SessionFactory + ): + with pytest.raises(ValueError): + await rebuild_all_resource_state(integration_session_factory, batch_size=0) + + async def test_repairs_corrupted_missing_and_stray_rows( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + ): + session = seeded_integration_session + pristine = await dump_state(session) + await _corrupt(session, rebuilt.alpha) + await session.execute( + delete(OGFieldResourceState).where( + OGFieldResourceState.resource_id == rebuilt.beta, + OGFieldResourceState.permission_mask == 2, + ) + ) + await session.execute( + insert(OGFieldResourceState).values( + resource_id=rebuilt.zeta_a, permission_mask=0, provenance={} + ) + ) + await session.commit() + + await rebuild_all_resource_state(integration_session_factory) + assert await dump_state(session) == pristine + + async def test_failed_rebuild_leaves_read_model_not_ready( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + session = seeded_integration_session + assert await is_resource_state_current(session) + + real_compute = state._compute_state_rows + batches = 0 + + async def fail_on_second_batch(*args: Any, **kwargs: Any): + nonlocal batches + batches += 1 + if batches == 2: + raise RuntimeError("boom") + return await real_compute(*args, **kwargs) + + monkeypatch.setattr(state, "_compute_state_rows", fail_on_second_batch) + with pytest.raises(RuntimeError, match="boom"): + await rebuild_all_resource_state(integration_session_factory, batch_size=1) + + assert not await is_resource_state_current(session) + # The table was never emptied: every resource still has rows (rebuilt + # or from before), but none are served until a rebuild succeeds. + assert {rid for rid, _ in await dump_state(session)} == rebuilt.listable + + async def test_incomplete_result_fails_integrity_check( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + real_compute = state._compute_state_rows + + async def drop_one_profile(session: AsyncSession, ids: list[int]): + rows = await real_compute(session, ids) + return [row for row in rows if row["permission_mask"] != 3] + + monkeypatch.setattr(state, "_compute_state_rows", drop_one_profile) + with pytest.raises(ResourceStateRebuildError, match="lack a row"): + await rebuild_all_resource_state(integration_session_factory) + + assert not await is_resource_state_current(seeded_integration_session) + + +class TestValuesTodaysModelRejects: + """Stored values can predate today's validation (e.g. a country code that + is no longer valid). Building state must not validate them: one such value + would otherwise stop every full rebuild from finishing and fail every write + to its resource, while the live list only fails on a page that shows it.""" + + @pytest.fixture + async def good_and_legacy( + self, seeded_integration_session: AsyncSession, test_user: User + ) -> tuple[int, int]: + """A valid resource, and one whose winning country is no longer valid.""" + session = seeded_integration_session + good = await new_resource(session, test_user) + await attach(session, test_user, good, "gem", name="Good", country="USA") + legacy = await new_resource(session, test_user) + await attach( + session, test_user, legacy, "gem", name="Old GEM", country="Gemland" + ) + await attach(session, test_user, legacy, "rmi", name="Old RMI") + await session.commit() + return good, legacy + + async def test_full_rebuild_completes( + self, + seeded_integration_session: AsyncSession, + integration_session_factory: SessionFactory, + good_and_legacy: tuple[int, int], + ): + _, legacy = good_and_legacy + await rebuild_all_resource_state(integration_session_factory) + + session = seeded_integration_session + assert await is_resource_state_current(session) + assert (await dump_state(session))[(legacy, 0)]["country"] == "Gemland" + + async def test_writes_to_the_resource_still_succeed( + self, + seeded_integration_session: AsyncSession, + test_user: User, + good_and_legacy: tuple[int, int], + ): + _, legacy = good_and_legacy + session = seeded_integration_session + current = await resource_actions.field_source_values(session, legacy, "name") + await resource_actions.set_field_source_priority( + session, + test_user, + legacy, + "name", + [view.source_id for view in reversed(current)], + ) + await session.commit() + + stored = await dump_state(session) + assert stored[(legacy, 0)]["name"] == "Old GEM" + assert stored[(legacy, 0)]["country"] == "Gemland" + + async def test_list_fails_only_where_live_fails( + self, + seeded_integration_session: AsyncSession, + integration_session_factory: SessionFactory, + good_and_legacy: tuple[int, int], + ): + good, legacy = good_and_legacy + await rebuild_all_resource_state(integration_session_factory) + session = seeded_integration_session + licensed = PROFILES[3] + + good_page = OGFieldQueryParams(id=good) + assert await resource_actions.query( + session, good_page, licensed + ) == await resource_actions._query_live(session, good_page, licensed) + + legacy_page = OGFieldQueryParams(id=legacy) + with pytest.raises(ValidationError): + await resource_actions._query_live(session, legacy_page, licensed) + with pytest.raises(ValidationError): + await resource_actions.query(session, legacy_page, licensed) + + +class TestRebuildAlongsideTraffic: + async def test_reads_and_writes_midway_through_a_rebuild( + self, + seeded_integration_session: AsyncSession, + integration_session_factory: SessionFactory, + test_user: User, + monkeypatch: pytest.MonkeyPatch, + ): + await rebuild_with_reads_and_writes_midway( + integration_session_factory, test_user, monkeypatch + ) + + +class TestIntegrityCheck: + async def test_clean_state_has_no_problems( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + assert await _integrity_problems(seeded_integration_session) == [] + + async def test_reports_missing_profile_rows( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + await session.execute( + delete(OGFieldResourceState).where( + OGFieldResourceState.resource_id == rebuilt.alpha, + OGFieldResourceState.permission_mask == 1, + ) + ) + assert await _integrity_problems(session) == [ + "1 listable resource(s) lack a row for every permission profile" + ] + + async def test_reports_rows_computed_by_other_code( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + await session.execute( + update(OGFieldResourceState) + .where(OGFieldResourceState.resource_id == rebuilt.alpha) + .values(code_version="some-other-commit") + ) + assert await _integrity_problems(session) == [ + "1 listable resource(s) have rows computed by other code" + ] + + async def test_reports_rows_for_merged_away_resources( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + await session.execute( + insert(OGFieldResourceState).values( + resource_id=rebuilt.zeta_b, permission_mask=0, provenance={} + ) + ) + assert await _integrity_problems(session) == [ + "1 non-listable resource(s) have stored state" + ] + + +class TestRebuildCliArgs: + def test_parses_ids_and_batch_size(self): + args = rebuild_cli._parse_args(["--ids", "3", "1", "--batch-size", "50"]) + assert args.ids == [3, 1] + assert args.batch_size == 50 + + def test_rejects_non_positive_batch_size(self): + with pytest.raises(SystemExit): + rebuild_cli._parse_args(["--batch-size", "0"]) + + def test_if_stale_is_off_by_default(self): + assert rebuild_cli._parse_args([]).if_stale is False + assert rebuild_cli._parse_args(["--if-stale"]).if_stale is True + + def test_if_stale_cannot_be_combined_with_ids(self): + with pytest.raises(SystemExit): + rebuild_cli._parse_args(["--if-stale", "--ids", "1"]) + + def test_rebuild_errors_exit_non_zero_with_one_clear_line( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ): + async def refuse(*_args: Any) -> None: + raise ResourceStateRebuildError("Another rebuild is already running.") + + monkeypatch.setattr(rebuild_cli, "_run", refuse) + with pytest.raises(SystemExit) as exited: + rebuild_cli.main([]) + + assert exited.value.code == 1 + assert "Another rebuild is already running." in caplog.text + + +class TestRebuildCli: + @pytest.fixture(autouse=True) + def _use_test_database( + self, + monkeypatch: pytest.MonkeyPatch, + integration_engine: AsyncEngine, + integration_session_factory: SessionFactory, + ): + async def _no_dispose() -> None: + return None + + monkeypatch.setattr( + rebuild_cli, "get_session_factory", lambda: integration_session_factory + ) + monkeypatch.setattr(rebuild_cli, "dispose_engine", _no_dispose) + monkeypatch.setattr(rebuild_cli, "get_engine", lambda: integration_engine) + + async def test_full_rebuild( + self, seeded_integration_session: AsyncSession, dataset: Dataset + ): + await rebuild_cli._run(ids=None, batch_size=2) + + session = seeded_integration_session + assert await is_resource_state_current(session) + assert {rid for rid, _ in await dump_state(session)} == dataset.listable + + async def test_if_stale_does_nothing_when_built_by_this_code( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + caplog: pytest.LogCaptureFixture, + ): + caplog.set_level(logging.INFO) + session = seeded_integration_session + built_at = await session.scalar(select(OGFieldResourceStateStatus.rebuilt_at)) + + await rebuild_cli._run(ids=None, batch_size=2, if_stale=True) + + assert "already built by this code" in caplog.text + assert ( + await session.scalar(select(OGFieldResourceStateStatus.rebuilt_at)) + == built_at + ) + assert await is_resource_state_current(session) + + @pytest.mark.parametrize("previously", ["never built", "built by other code"]) + async def test_if_stale_rebuilds_when_not_current( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + previously: str, + ): + caplog.set_level(logging.INFO) + if previously == "built by other code": + _deploy(monkeypatch, "previous-commit") + await rebuild_all_resource_state(integration_session_factory) + _deploy(monkeypatch, "this-commit") + + await rebuild_cli._run(ids=None, batch_size=2, if_stale=True) + + session = seeded_integration_session + assert await is_resource_state_current(session) + assert {rid for rid, _ in await dump_state(session)} == dataset.listable + assert "this-commit" in caplog.text + + async def test_ids_refresh_only_those_resources( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + caplog: pytest.LogCaptureFixture, + ): + caplog.set_level(logging.INFO) + await rebuild_cli._run(ids=[dataset.gamma], batch_size=1) + + session = seeded_integration_session + assert {rid for rid, _ in await dump_state(session)} == {dataset.gamma} + assert not await is_resource_state_current(session) + assert f"Recomputed resource state for ids [{dataset.gamma}]" in caplog.text + + async def test_ids_that_are_not_listed_are_reported_not_recomputed( + self, + seeded_integration_session: AsyncSession, + rebuilt: Dataset, + caplog: pytest.LogCaptureFixture, + ): + caplog.set_level(logging.INFO) + # A stale row for a merged-away resource gets removed. + await seeded_integration_session.execute( + insert(OGFieldResourceState).values( + resource_id=rebuilt.zeta_a, permission_mask=0, provenance={} + ) + ) + await seeded_integration_session.commit() + + await rebuild_cli._run(ids=[rebuilt.alpha, rebuilt.zeta_a], batch_size=1) + + stored = {rid for rid, _ in await dump_state(seeded_integration_session)} + assert rebuilt.zeta_a not in stored + assert f"Recomputed resource state for ids [{rebuilt.alpha}]" in caplog.text + assert f"[{rebuilt.zeta_a}]" in caplog.text + assert "not listed" in caplog.text + + async def test_ids_refused_when_the_served_table_is_from_other_code( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + """Repairing from commit B would put B-computed rows into a table that + commit A is serving as built entirely by A.""" + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + await _corrupt(session, dataset.alpha) + before = await dump_state(session) + + _deploy(monkeypatch, "commit-b") + with pytest.raises(ResourceStateRebuildError, match="commit-a"): + await rebuild_cli._run(ids=[dataset.alpha], batch_size=1) + + assert await dump_state(session) == before + status = ( + await session.execute(select(OGFieldResourceStateStatus.__table__)) + ).one() + assert (status.ready, status.code_version) == (True, "commit-a") + + async def test_ids_allowed_from_any_code_while_not_ready( + self, + seeded_integration_session: AsyncSession, + dataset: Dataset, + integration_session_factory: SessionFactory, + monkeypatch: pytest.MonkeyPatch, + ): + """A table that isn't ready isn't being served, and the next rebuild + replaces its rows, so a repair from any code is harmless.""" + session = seeded_integration_session + _deploy(monkeypatch, "commit-a") + await rebuild_all_resource_state(integration_session_factory) + await session.execute(update(OGFieldResourceStateStatus).values(ready=False)) + await session.commit() + + _deploy(monkeypatch, "commit-b") + await rebuild_cli._run(ids=[dataset.alpha], batch_size=1) + + assert {rid for rid, _ in await dump_state(session)} >= {dataset.alpha} + + async def test_unknown_ids_fail_without_changing_anything( + self, seeded_integration_session: AsyncSession, rebuilt: Dataset + ): + session = seeded_integration_session + before = await dump_state(session) + await _corrupt(session, rebuilt.alpha) + corrupted = await dump_state(session) + assert corrupted != before + + with pytest.raises(ResourceStateRebuildError, match="999999"): + await rebuild_cli._run(ids=[rebuilt.alpha, 999_999], batch_size=1) + + # The valid id was not repaired either: nothing changed. + assert await dump_state(session) == corrupted diff --git a/deployments/api/tests/observability/test_query_name_actions.py b/deployments/api/tests/observability/test_query_name_actions.py index 619d2af3..d8bd8a6e 100644 --- a/deployments/api/tests/observability/test_query_name_actions.py +++ b/deployments/api/tests/observability/test_query_name_actions.py @@ -25,6 +25,7 @@ from stitch.api.db import og_field_source_actions as source_actions from stitch.api.db.model import MembershipModel, MembershipStatus, ResourceModel from stitch.api.db.model.oil_gas_field_source_value import ATTRIBUTE_NAMES +from stitch.api.db.read_model.state import rebuild_all_resource_state from stitch.api.entities import ( MergeCandidateCreateRequest, MergeCandidateReviewRequest, @@ -59,6 +60,15 @@ def _assert_labels(events: list[dict], expected: set[str]) -> None: assert actual == expected, f"unexpected labels: {actual ^ expected}" +def _labeled_statements_use_read_model(events: list[dict]) -> bool: + """Every labeled query read the precomputed table (not live coalescing).""" + return all( + "og_field_resource_state" in event["statement"] + for event in events + if "query_name" in event + ) + + class TestEndpointQueryLabels: """Drive the real HTTP read endpoints and assert the exact labels each emits.""" @@ -97,6 +107,39 @@ async def test_filter_options_endpoint_labels( _assert_labels(captured_query_events, {"resources.filter_options"}) + @pytest.mark.anyio + async def test_read_model_list_and_filter_options_labels( + self, + integration_client: AsyncClient, + integration_session_factory, + og_create_res_fact, + captured_query_events: list[dict], + ): + """Served from the precomputed read model, the list reads its page's rows + in the same statement that picks them (``resources.list_page``, in place + of the live path's ``list_ids`` + ``list_hydrate``); filter-options keeps + its label. The readiness lookup is unlabeled.""" + create = await integration_client.post( + "/oil-gas-fields/", + json=og_create_res_fact(name="Labeled Resource").model_dump(mode="json"), + ) + assert create.status_code == 200, create.text + await rebuild_all_resource_state(integration_session_factory) + + captured_query_events.clear() + response = await integration_client.get("/oil-gas-fields/") + assert response.status_code == 200, response.text + _assert_labels( + captured_query_events, {"resources.count", "resources.list_page"} + ) + assert _labeled_statements_use_read_model(captured_query_events) + + captured_query_events.clear() + response = await integration_client.get("/oil-gas-fields/filter-options") + assert response.status_code == 200, response.text + _assert_labels(captured_query_events, {"resources.filter_options"}) + assert _labeled_statements_use_read_model(captured_query_events) + @pytest.mark.anyio async def test_detail_endpoint_labels( self, @@ -207,7 +250,12 @@ async def test_create_label( # create() with source_data fans out to the source helpers. _assert_labels( captured_query_events, - {"resources.create", "sources.get_or_create", "sources.attach"}, + { + "resources.create", + "sources.get_or_create", + "sources.attach", + "read_model.refresh", + }, ) @pytest.mark.anyio @@ -232,7 +280,8 @@ async def test_merge_labels( session, test_user, [first.id, second.id] ) _assert_labels( - captured_query_events, {"resources.merge.load", "resources.merge.apply"} + captured_query_events, + {"resources.merge.load", "resources.merge.apply", "read_model.refresh"}, ) @pytest.mark.anyio @@ -275,11 +324,51 @@ async def test_set_field_source_priority_labels( "resources.set_field_source_priority.load", "resources.set_field_source_priority.candidates", "resources.set_field_source_priority.persist", + "read_model.refresh", "resources.resolve_root", "resources.field_source_values", }, ) + @pytest.mark.anyio + async def test_read_model_refresh_has_its_own_label( + self, + seeded_integration_session: AsyncSession, + test_user: User, + og_create_res_fact, + source_maker, + captured_query_events: list[dict], + ): + """The read-model refresh inside a write is labeled on its own, and the + write's label is restored afterwards (the attach then loads the entity + it returns under ``sources.attach``).""" + session = seeded_integration_session + parent = await resource_actions.create( + session, test_user, og_create_res_fact(name="Parent") + ) + await session.commit() + existing = ( + await source_actions.create_and_attach_sources( + session, test_user, [source_maker(managed=False)], parent.id + ) + )[0] + + captured_query_events.clear() + await source_actions.attach_sources_to_resource( + session, parent.id, [existing], test_user + ) + + labels = [e["query_name"] for e in captured_query_events if "query_name" in e] + state_labels = { + e["query_name"] + for e in captured_query_events + if "og_field_resource_state" in e["statement"] + } + assert state_labels == {"read_model.refresh"} + refreshed_at = labels.index("read_model.refresh") + assert set(labels[:refreshed_at]) == {"sources.attach"} + assert labels[-1] == "sources.attach" + @pytest.mark.anyio async def test_source_action_labels( self, @@ -305,7 +394,9 @@ async def test_source_action_labels( await source_actions.create_and_attach_sources( session, test_user, [source_maker(managed=False, source="rmi")], parent.id ) - _assert_labels(captured_query_events, {"sources.create_and_attach"}) + _assert_labels( + captured_query_events, {"sources.create_and_attach", "read_model.refresh"} + ) captured_query_events.clear() await source_actions.get_or_create_sources( @@ -317,7 +408,7 @@ async def test_source_action_labels( await source_actions.attach_sources_to_resource( session, parent.id, [source_maker(managed=False, source="bc")], test_user ) - _assert_labels(captured_query_events, {"sources.attach"}) + _assert_labels(captured_query_events, {"sources.attach", "read_model.refresh"}) captured_query_events.clear() await source_actions.get_source(session, created.id) @@ -399,6 +490,7 @@ async def test_merge_candidate_action_labels( "merge_candidates.approve.persist", "resources.merge.load", "resources.merge.apply", + "read_model.refresh", }, )