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/deployments/api/READ_MODEL.md b/deployments/api/READ_MODEL.md index 5db23ecf..7c80df7c 100644 --- a/deployments/api/READ_MODEL.md +++ b/deployments/api/READ_MODEL.md @@ -9,8 +9,8 @@ The table is a cache, not a source of truth. Everything in it can be rebuilt fro 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. -So far only the tables exist: they start empty and not ready, and nothing writes -or reads them yet, so answers are unchanged. +So far the tables and the rebuild command exist. Nothing reads the table yet, and +API writes don't keep it current yet, so answers are unchanged. ## How it works @@ -26,10 +26,104 @@ or reads them yet, so answers are unchanged. may be used only when it is ready **and** was built by the same commit the API is running. +## 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 (the post-deploy step). +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 read model only counts as ready for the + commit that built it. 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`, which matches no deployed commit. + + ```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, 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 + ``` + +- **What happens:** readiness is cleared first, so the table isn't used while + it's rebuilt. 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; just run it again. Whatever rows a + failed run leaves are never used, 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: 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. + ## Changing the code safely - **New field** on `OilGasFieldBase`: add the column to `OGFieldResourceState` - (an import-time check fails until you do) plus a migration. + (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 would get the live path. + +## Tests + +The normal suite runs on SQLite. A few tests need real Postgres (e.g. the rebuild +lock). 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 +``` diff --git a/deployments/api/src/stitch/api/db/queries.py b/deployments/api/src/stitch/api/db/queries.py index 02ba5d5c..744958f2 100644 --- a/deployments/api/src/stitch/api/db/queries.py +++ b/deployments/api/src/stitch/api/db/queries.py @@ -388,11 +388,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 +412,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 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..bde57168 --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/rebuild.py @@ -0,0 +1,185 @@ +"""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 read model counts as ready only once 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 never changes readiness. +""" + +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, + 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 being served 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. 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 being served 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(), + ) + 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..968fa29b --- /dev/null +++ b/deployments/api/src/stitch/api/db/read_model/state.py @@ -0,0 +1,530 @@ +"""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. +""" + +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" + +# 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 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() + ) + + +async def lock_resources(session: AsyncSession, resource_ids: Collection[int]) -> None: + """Lock the resource rows until the caller's transaction ends. + + ``refresh_resource_states`` takes it before recomputing, so two + transactions recomputing the same resource (e.g. a full rebuild and an + ``--ids`` repair) queue instead of one overwriting the other with a stale + recompute. 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, + ).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. + + 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 _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. + """ + async with session_factory.begin() as session: + await _set_ready(session, False) + + 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) + return recomputed 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..6a11fc6b --- /dev/null +++ b/deployments/api/tests/db/read_model/dataset.py @@ -0,0 +1,216 @@ +"""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 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 + 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/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_state.py b/deployments/api/tests/db/read_model/test_state.py new file mode 100644 index 00000000..4d72b62b --- /dev/null +++ b/deployments/api/tests/db/read_model/test_state.py @@ -0,0 +1,906 @@ +"""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 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.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, + 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, + 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 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, +) + +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 + assert sql.rstrip().endswith("ORDER BY og_field_resources.id FOR 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() + + +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.""" + + @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" + + +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