+ Approved cohort (frozen at approval) +
+.`), which
`trust/Makefile` `-include`s so every trust-internal container inherits it.
diff --git a/trust/data-access-api/data_access_api/config.py b/trust/data-access-api/data_access_api/config.py
index 8c3a4e687..4101ced3e 100644
--- a/trust/data-access-api/data_access_api/config.py
+++ b/trust/data-access-api/data_access_api/config.py
@@ -78,6 +78,25 @@ def coerce_empty_cohort_query_threshold(cls, v: object) -> object:
CACHE_MAX_RESULT_ROWS: PositiveInt = 50_000 # Max rows per cached result; larger results skip caching
CACHE_MAX_ENTRIES: PositiveInt = 64 # Max number of cached query results
+ # Approved-cohort snapshot store (FLIP#857). Directory where the cohort frozen at project
+ # approval is persisted, one sub-directory per hub project id — a dedicated bind mount in
+ # the compose files. The row-level routes serve ONLY these frozen artefacts (a project
+ # with no snapshot is refused), so an unset/unwritable directory means no row-level data
+ # can be released until the store is fixed — deliberately fail-closed.
+ COHORT_SNAPSHOT_DIR: str = ""
+
+ # Hard cap on one serialized snapshot. Over the cap the snapshot is REFUSED, never
+ # truncated — a partial cohort would silently poison training data.
+ SNAPSHOT_MAX_BYTES: PositiveInt = 536_870_912 # 512 MiB
+
+ @field_validator("SNAPSHOT_MAX_BYTES", mode="before")
+ @classmethod
+ def coerce_empty_snapshot_max_bytes(cls, v: object) -> object:
+ """Treat an empty-string cap as the default (same kit-file gotcha as the threshold)."""
+ if v is None or v == "":
+ return 536_870_912
+ return v
+
#
OMOP_DB_SERVICE_NAME: str = "omop-db" # The name of the OMOP database service in Docker Compose or Kubernetes
OMOP_DB_PORT: int = 5432 # Calls from another container use port 5432 (e.g. http://omop-db:5432)
@@ -96,6 +115,17 @@ def coerce_empty_cohort_query_threshold(cls, v: object) -> object:
TRUST_INTERNAL_SERVICE_KEY_HEADER: str = "X-Trust-Internal-Service-Key"
TRUST_INTERNAL_SERVICE_KEY: str = ""
+ # Cohort-write authorisation (FLIP#857). The snapshot create/delete routes carry a second
+ # gate on top of the trust-internal key: the caller must prove possession of AES_KEY_BASE64
+ # by sending the SHA-256 of the key in this header. This separates the services trusted to
+ # DEFINE a project's approved cohort (trust-api, data-access-api — both hold the AES key)
+ # from fl-client, which runs researcher training code, holds the trust-internal key for its
+ # imaging reads, and deliberately has no AES key. The read routes are unaffected, so
+ # fl-client's ``get_dataframe`` keeps working; only the cohort-defining writes are locked
+ # down. The digest — not the raw key — is transmitted, so a captured header or a log line
+ # never reveals the encryption key.
+ COHORT_ADMIN_KEY_HEADER: str = "X-Cohort-Admin-Key"
+
# Define the database URL for the OMOP database
@property
def OMOP_DATABASE_URL(self) -> SecretStr:
diff --git a/trust/data-access-api/data_access_api/main.py b/trust/data-access-api/data_access_api/main.py
index 102412382..6df094c2b 100644
--- a/trust/data-access-api/data_access_api/main.py
+++ b/trust/data-access-api/data_access_api/main.py
@@ -10,14 +10,30 @@
# limitations under the License.
#
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
from fastapi import FastAPI
from log_config import LoggingMiddleware
# Ensure structured logging is configured on import
import data_access_api.utils.logger # noqa: F401
from data_access_api.config import get_settings
-from data_access_api.routers.cohort import router as cohort_router
+from data_access_api.routers.cohort import read_router as cohort_read_router
+from data_access_api.routers.cohort import write_router as cohort_write_router
from data_access_api.routers.health import router as health_router
+from data_access_api.services.cohort_snapshot import ensure_store
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI) -> AsyncIterator[None]:
+ # Boot-time check of the approved-cohort snapshot store (FLIP#857): creates the
+ # directory, sweeps stale write debris, probes writability. Never raises — a broken
+ # store must not take the service down; the row-level routes then refuse projects
+ # whose artefact cannot be read (fail-closed) while statistics keep serving.
+ ensure_store()
+ yield
+
# Disable Swagger / OpenAPI / ReDoc in production. Data-access-api executes SQL
# against OMOP under a service account; leaking its route + schema map to anyone
@@ -32,9 +48,11 @@
docs_url="/docs" if _docs_enabled else None,
openapi_url="/openapi.json" if _docs_enabled else None,
redoc_url="/redoc" if _docs_enabled else None,
+ lifespan=lifespan,
)
app.add_middleware(LoggingMiddleware)
-app.include_router(cohort_router)
+app.include_router(cohort_read_router)
+app.include_router(cohort_write_router)
app.include_router(health_router)
diff --git a/trust/data-access-api/data_access_api/routers/cohort.py b/trust/data-access-api/data_access_api/routers/cohort.py
index 8df5cb4db..c76ed354a 100644
--- a/trust/data-access-api/data_access_api/routers/cohort.py
+++ b/trust/data-access-api/data_access_api/routers/cohort.py
@@ -20,11 +20,22 @@
AccessionIdsResponse,
CohortQueryInput,
DataframeQuery,
+ SnapshotDeleteRequest,
+ SnapshotResponse,
StatisticsResponse,
)
from data_access_api.services.cohort import get_records, get_statistics, validate_query
+from data_access_api.services.cohort_snapshot import (
+ Snapshot,
+ SnapshotTooLarge,
+ delete_snapshot,
+ get_snapshot,
+ normalised_query_hash,
+ save_snapshot,
+ snapshot_enabled,
+)
from data_access_api.utils.encryption import decrypt
-from data_access_api.utils.internal_auth import authenticate_internal_service
+from data_access_api.utils.internal_auth import authenticate_cohort_admin, authenticate_internal_service
from data_access_api.utils.logger import logger
# Returned instead of row-level data when a cohort is smaller than
@@ -34,16 +45,92 @@
# becomes a one-row oracle for probing the database.
_BELOW_THRESHOLD_DETAIL = "Cohort is too small for row-level data to be released."
+# Returned by the row-level routes for a project with no frozen cohort artefact.
+# Row-level data is released ONLY from the snapshot persisted at approval
+# (FLIP#857); there is no live-SQL serving path. Deliberately generic: it must
+# not reveal whether the project exists.
+_NO_SNAPSHOT_DETAIL = "No approved cohort snapshot exists for this project."
+
+
+# Two routers under the same /cohort prefix, split by privilege (FLIP#857).
+#
+# read_router — statistics + the two row-level serve routes. Gated on the trust-internal key
+# alone: every trust-internal service (trust-api, imaging-api, fl-client) legitimately reads
+# here, and the routes serve only the frozen, project-scoped, threshold-gated snapshot.
+#
+# write_router — the routes that DEFINE/destroy the frozen artefact the row-level routes then
+# serve. Gated on the trust-internal key AND cohort-admin (proof of possessing AES_KEY_BASE64),
+# so a caller must be both a trust-internal service and one of the two trusted to define cohorts.
+# fl-client holds the trust-internal key but no AES key, so it can read its cohort but never
+# rewrite or delete it. The dependency order matters: trust-internal runs first, so a caller
+# without it gets 401 before the cohort-admin check ever runs (which returns 403).
+read_router = APIRouter(prefix="/cohort", tags=["Cohort"], dependencies=[Depends(authenticate_internal_service)])
+write_router = APIRouter(
+ prefix="/cohort",
+ tags=["Cohort"],
+ dependencies=[Depends(authenticate_internal_service), Depends(authenticate_cohort_admin)],
+)
+
+
+def _require_snapshot(project_id: str) -> Snapshot:
+ """The project's frozen cohort, or the fixed 403 when there is none.
+
+ The snapshot persisted at approval is the ONLY source of row-level data: a project
+ that was never approved (or whose snapshot was purged, or whose trust has the store
+ unconfigured/unwritable) is refused. Fail-closed by construction — there is no
+ live-SQL fallback for these routes.
+ """
+ snapshot = get_snapshot(project_id)
+ if snapshot is None:
+ logger.warning(f"Refusing row-level data for project {project_id}: no approved cohort snapshot")
+ raise HTTPException(status_code=403, detail=_NO_SNAPSHOT_DETAIL)
+ return snapshot
+
+
+def _check_frozen_threshold(project_id: str, snapshot: Snapshot) -> None:
+ """Apply ``COHORT_QUERY_THRESHOLD`` to the FROZEN row count.
+
+ The gate reads the snapshot, not live OMOP: the frozen cohort is immutable, so a project
+ that cleared the threshold at approval keeps serving even while the live database drifts
+ underneath (pre-FLIP#857, the SQL was re-run on every call and the answer could change).
+ The refusal reuses the fixed below-threshold text so a zero-row snapshot and a
+ below-threshold one stay indistinguishable. The threshold itself is still read live, so
+ an operator RAISING their disclosure floor takes effect on already-approved projects.
+ """
+ if snapshot.meta.row_count < get_settings().COHORT_QUERY_THRESHOLD:
+ logger.warning(
+ f"Withholding frozen cohort for project {project_id}: snapshot of "
+ f"{snapshot.meta.row_count} rows is below the minimum size of "
+ f"{get_settings().COHORT_QUERY_THRESHOLD}"
+ )
+ raise HTTPException(status_code=403, detail=_BELOW_THRESHOLD_DETAIL)
+
+
+def _log_ignored_client_query(project_id: str, snapshot: Snapshot, client_query: str) -> None:
+ """Record that the caller-supplied SQL was ignored in favour of the frozen cohort.
-# Create Router
-router = APIRouter(prefix="/cohort", tags=["Cohort"], dependencies=[Depends(authenticate_internal_service)])
+ Serving the snapshot regardless of the submitted SQL is what closes the arbitrary-SQL
+ exposure on these routes (see FLIP#857's audit note): the only row-level data obtainable
+ under a project's id is the cohort that was approved. The hash comparison exists purely
+ so a mismatch is visible in the trust's logs; the ``query`` field stays in the request
+ schema because the FL client library sends it.
+ """
+ if normalised_query_hash(client_query) != snapshot.meta.query_hash:
+ logger.warning(
+ f"Client-supplied query for project {project_id} differs from the frozen cohort "
+ "query — ignored; serving the approved snapshot."
+ )
-@router.post("", response_model=StatisticsResponse)
+@read_router.post("", response_model=StatisticsResponse)
def receive_cohort_query(query_input: CohortQueryInput) -> StatisticsResponse:
"""
Receives a cohort query and returns the aggregated statistics.
+ This is the one route that always evaluates LIVE OMOP: it runs pre-approval by
+ definition (it is how a proposed cohort is sized in the first place), and it
+ releases aggregates only.
+
Below-threshold results are privacy-suppressed: any count below the threshold —
including a genuine zero — comes back as a normal ``StatisticsResponse`` with
``record_count=0``, empty ``data`` and ``suppressed=True``, *not* an HTTP error.
@@ -95,154 +182,206 @@ def receive_cohort_query(query_input: CohortQueryInput) -> StatisticsResponse:
return results
-@router.post("/dataframe")
+@read_router.post("/dataframe")
def get_dataframe(query_input: DataframeQuery) -> dict[str, list[Any]]:
"""
- Retrieves query results in a DataFrame-like structure (column-oriented dictionary).
-
- This is the training-data path: user-supplied FL code inside the trust's
- fl-client reaches it through ``flip.get_dataframe(...)``, so it necessarily
- returns row-level records — a model trains on rows. The data stays inside
- the trust; only model updates leave it.
-
- Because it is row-level, the cohort must clear ``COHORT_QUERY_THRESHOLD``
- before anything is released, mirroring the suppression the ``/cohort``
- statistics route already applies. A cohort below the threshold is refused
- outright rather than returned truncated: there is no training value in it,
- and releasing a handful of identifiable rows is exactly the disclosure the
- threshold exists to prevent.
-
- Note there is deliberately no column allowlist here. ``accession_id`` is
- load-bearing — it is how the returned rows join to the imaging studies
- pulled into XNAT — and shipped tutorials legitimately select ``*``, so a
- column filter would break every FL app on the platform while a caller could
- trivially alias around it. Column-level minimisation belongs in the cohort
- query the project submits and in project approval, not here.
+ Serves the project's FROZEN cohort in a DataFrame-like structure (column-oriented dict).
+
+ This is the training-data path: user-supplied FL code inside the trust's fl-client
+ reaches it through ``flip.get_dataframe(...)``, so it necessarily returns row-level
+ records — a model trains on rows. The data stays inside the trust; only model updates
+ leave it.
+
+ What it serves is the cohort snapshot persisted at project approval (FLIP#857), keyed
+ on the project id — **the caller-supplied SQL is ignored** (logged when it differs from
+ the frozen query). Consequences, all deliberate: the cohort cannot grow after approval;
+ every training round fetches an identical frame; and researcher code cannot execute
+ arbitrary SQL under an approved project's id. A project with no snapshot is refused —
+ there is no live-SQL serving path on this route.
+
+ Because it is row-level, the frozen cohort must clear ``COHORT_QUERY_THRESHOLD``
+ before anything is released, mirroring the suppression the ``/cohort`` statistics
+ route applies. The threshold is read live, so an operator raising their disclosure
+ floor takes effect on already-approved projects.
Args:
- query_input (DataframeQuery): The input data for the DataFrame query.
+ query_input (DataframeQuery): The encrypted project id (and the advisory query).
Returns:
- dict[str, list[Any]]: The query results in a DataFrame-like structure.
+ dict[str, list[Any]]: The frozen cohort in a DataFrame-like structure.
Raises:
- HTTPException: 400 if the query is invalid, 403 if the cohort is below
- the disclosure threshold, 500 if the query fails to execute.
+ HTTPException: 403 if the project has no cohort snapshot or the frozen cohort is
+ below the disclosure threshold.
"""
project_id = decrypt(query_input.encrypted_project_id)
logger.info(f"Received DataFrame query for project {project_id}")
- safe_query = validate_query(query_input.query)
-
- try:
- df = get_records(safe_query)
- except HTTPException:
- # get_records already converts driver errors into category-only
- # HTTPExceptions; re-wrapping them below would discard that work and
- # turn every categorised 400 into an opaque 500.
- raise
- except SQLAlchemyError:
- logger.exception("DataFrame query failed with a database error")
- raise HTTPException(status_code=500, detail="Query execution failed.")
- except Exception:
- # Detail is category-only: the trust forwards it to the hub, which shows
- # it to every project member, and raw exception text can carry row
- # values and connection internals.
- logger.exception("DataFrame query failed unexpectedly")
- raise HTTPException(status_code=500, detail="Query execution failed.")
-
- minimum_cohort_size = get_settings().COHORT_QUERY_THRESHOLD
- if len(df) < minimum_cohort_size:
- logger.warning(
- f"Withholding row-level data for project {project_id}: "
- f"cohort below the minimum size of {minimum_cohort_size}"
- )
- raise HTTPException(status_code=403, detail=_BELOW_THRESHOLD_DETAIL)
-
- return df.to_dict(orient="list")
+ snapshot = _require_snapshot(project_id)
+ _log_ignored_client_query(project_id, snapshot, query_input.query)
+ _check_frozen_threshold(project_id, snapshot)
+ logger.info(f"Serving frozen cohort snapshot for project {project_id}: {snapshot.meta.row_count} rows")
+ return snapshot.df.to_dict(orient="list")
-@router.post("/accession-ids", response_model=AccessionIdsResponse)
+@read_router.post("/accession-ids", response_model=AccessionIdsResponse)
def get_accession_ids(query_input: DataframeQuery) -> AccessionIdsResponse:
"""
- Returns only the ``accession_id`` column of the cohort, projected server-side.
+ Returns only the ``accession_id`` column of the project's FROZEN cohort.
- The caller's query is wrapped as ``SELECT accession_id FROM () sub`` so
- no other columns ever cross the trust boundary. This is the minimal-disclosure
- endpoint used by imaging-api to fetch the accession numbers it needs to import
- studies from PACS — it does not expose row-level patient attributes.
+ This is the minimal-disclosure endpoint used by imaging-api to fetch the accession
+ numbers it needs to import studies from PACS — it does not expose row-level patient
+ attributes. Like ``/cohort/dataframe`` it serves the snapshot persisted at approval and
+ **ignores the caller-supplied SQL** (FLIP#857): the imaging status poll (roughly every
+ 10 s while a project page is open) and reimport read a stable pointer set instead of
+ re-running the cohort SQL against a live OMOP that changes underneath.
Accession IDs are still row-level identifiers, and they are the pointer set into
the imaging data: they decide whose studies get pulled into XNAT, where project
- members view them. So the cohort must clear ``COHORT_QUERY_THRESHOLD`` here just
- as it must on ``/cohort/dataframe``, and the refusal reuses that route's fixed
- text so a zero-row cohort and a below-threshold one are indistinguishable.
-
- The trust applies this itself rather than relying on the hub's staging guard
- (``flip_api.project_services.stage_project``, which refuses to stage a trust whose
- cohort came back empty or suppressed). The hub is a separate administrative domain;
- a trust must stay safe regardless of what the hub checked, and the hub's own
- ``start_project_imaging_creation`` endpoint does not re-check staging.
-
- **This is evaluated against the live cohort on every call, not once at approval.**
- The endpoint is re-invoked by the imaging status poll roughly every 10 s while a
- user has the project page open, and again on reimport — each time re-running the
- cohort SQL against OMOP, which changes underneath. A project can therefore pull
- cleanly at approval and later start refusing if its cohort shrinks below the
- threshold. FLIP has no frozen approved-cohort artefact; see FLIP#857.
+ members view them. So the frozen cohort must clear ``COHORT_QUERY_THRESHOLD`` here
+ just as it must on ``/cohort/dataframe``, and the refusal reuses that route's fixed
+ text so a zero-row cohort and a below-threshold one are indistinguishable. The trust
+ applies this itself rather than relying on the hub's staging guard — the hub is a
+ separate administrative domain, and a trust must stay safe regardless of what the hub
+ checked.
+
+ A frozen cohort with no ``accession_id`` column returns an EMPTY list rather than an
+ error: a tabular/OMOP-only project legitimately has no imaging to pull.
Args:
- query_input (DataframeQuery): The cohort query.
+ query_input (DataframeQuery): The encrypted project id (and the advisory query).
Returns:
- AccessionIdsResponse: The accession IDs returned by the cohort query.
+ AccessionIdsResponse: The frozen cohort's accession IDs.
Raises:
- HTTPException: 400 if the query is invalid or does not select an
- ``accession_id`` column, 403 if the cohort is below the disclosure
- threshold, 500 if the query fails to execute.
+ HTTPException: 403 if the project has no cohort snapshot or the frozen cohort is
+ below the disclosure threshold.
"""
project_id = decrypt(query_input.encrypted_project_id)
logger.info(f"Received accession-ids query for project {project_id}")
- # validate_query returns the caller's SQL re-emitted from its parsed AST,
- # which breaks any injection taint chain and strips trailing semicolons so
- # the inner query composes cleanly inside the outer SELECT subquery.
- safe_inner = validate_query(query_input.query)
- wrapped_query = f"SELECT accession_id FROM ({safe_inner}) AS cohort_subquery"
+ snapshot = _require_snapshot(project_id)
+ _log_ignored_client_query(project_id, snapshot, query_input.query)
+ # Threshold before the column check: nothing about the cohort's shape is revealed for a
+ # below-threshold snapshot.
+ _check_frozen_threshold(project_id, snapshot)
+ if not snapshot.meta.has_accessions:
+ logger.info(f"Frozen cohort for project {project_id} has no accession_id column (tabular project)")
+ return AccessionIdsResponse(accession_ids=[])
+ frozen_ids = [str(value) for value in snapshot.df["accession_id"].tolist()]
+ logger.info(f"Serving {len(frozen_ids)} frozen accession ids for project {project_id}")
+ return AccessionIdsResponse(accession_ids=frozen_ids)
+
+
+@write_router.post("/snapshot", response_model=SnapshotResponse)
+def create_snapshot(query_input: DataframeQuery) -> SnapshotResponse:
+ """
+ Materialises the cohort ONCE and persists it as this project's frozen artefact (FLIP#857).
+
+ Called by trust-api when the hub approves a project. It defines the artefact everyone
+ then trains on, so it sits on the ``write_router`` and requires cohort-admin authorisation
+ (proof of possessing ``AES_KEY_BASE64``) in addition to the trust-internal key — fl-client
+ holds the latter but not the former, so researcher code cannot reach here (FLIP#857). From
+ that point the two row-level routes serve the persisted dataframe keyed on the project id
+ and ignore caller SQL, so the cohort a project trains on is exactly the cohort that was
+ approved — it cannot grow with the live database, and it is identical on every fetch.
+ Re-approval calls this again and atomically replaces the artefact (which is also how
+ OMOP-side removals and opt-outs propagate into an approved project: at explicit re-snapshot
+ events, never silently mid-training).
+
+ This route (with the statistics route) is where live OMOP is evaluated, exactly once
+ per approval; ``validate_query`` remains the authority on the SQL executed here. The
+ disclosure threshold is enforced BEFORE anything is persisted: a below-threshold cohort
+ leaves no artefact and returns the same fixed refusal as the row-level routes. The
+ response carries aggregates only (count, column names, timestamps) — the row-level data
+ stays on this trust's disk.
+
+ Args:
+ query_input (DataframeQuery): The approved cohort query and encrypted project id.
+
+ Returns:
+ SnapshotResponse: What was frozen.
+
+ Raises:
+ HTTPException: 400 if the query is invalid or the project id is not a UUID, 403 if
+ the cohort is below the disclosure threshold, 413 if the serialized snapshot
+ exceeds ``SNAPSHOT_MAX_BYTES``, 500 if the query fails to execute, 503 if the
+ snapshot store is not configured.
+ """
+ if not snapshot_enabled():
+ raise HTTPException(status_code=503, detail="Cohort snapshot store is not configured on this trust.")
+
+ project_id = decrypt(query_input.encrypted_project_id)
+ logger.info(f"Received cohort snapshot request for project {project_id}")
+ safe_query = validate_query(query_input.query)
try:
- df = get_records(wrapped_query)
+ # use_cache=False: the artefact must freeze LIVE OMOP at this moment. The statistics
+ # run at submission caches this same SQL for CACHE_TTL_DAYS, and on re-approval —
+ # the event that propagates OMOP-side removals/opt-outs into the snapshot — a cached
+ # read would silently re-freeze the pre-removal cohort.
+ df = get_records(safe_query, use_cache=False)
except HTTPException:
# get_records already converts driver errors into category-only
# HTTPExceptions; re-wrapping them below would discard that work and
# turn every categorised 400 into an opaque 500.
raise
except SQLAlchemyError:
- logger.exception("Accession-ids query failed with a database error")
+ logger.exception("Snapshot cohort query failed with a database error")
raise HTTPException(status_code=500, detail="Query execution failed.")
except Exception:
- # Detail is category-only: the trust forwards it to the hub, which shows
- # it to every project member, and raw exception text can carry row
- # values and connection internals.
- logger.exception("Accession-ids query failed unexpectedly")
+ # Detail is category-only, as on the other routes: it travels back to the hub.
+ logger.exception("Snapshot cohort query failed unexpectedly")
raise HTTPException(status_code=500, detail="Query execution failed.")
- minimum_cohort_size = get_settings().COHORT_QUERY_THRESHOLD
- if len(df) < minimum_cohort_size:
+ if len(df) < get_settings().COHORT_QUERY_THRESHOLD:
logger.warning(
- f"Withholding accession IDs for project {project_id}: "
- f"cohort below the minimum size of {minimum_cohort_size}"
+ f"Refusing to snapshot project {project_id}: cohort below the minimum size of "
+ f"{get_settings().COHORT_QUERY_THRESHOLD}"
)
raise HTTPException(status_code=403, detail=_BELOW_THRESHOLD_DETAIL)
- # No "did the DataFrame come back with accession_id?" guard here: it could never fire.
- # The wrapper above selects the column explicitly, so a cohort that does not project it
- # fails inside get_records with UndefinedColumn — surfacing as a category 400 through the
- # `except HTTPException: raise` branch, before any DataFrame exists. Pinned by
- # tests/integration/test_cohort_endpoint.py::test_accession_ids_missing_column_surfaces_get_records_400.
- accession_ids = [str(value) for value in df["accession_id"].tolist()]
- logger.info(f"accession-ids query for project {project_id} returned {len(accession_ids)} ids")
- return AccessionIdsResponse(accession_ids=accession_ids)
+ try:
+ # The hash is of the RAW submitted SQL (not the validator's re-emission) so serving
+ # can compare it against the raw query the hub injects into FL job configs.
+ meta = save_snapshot(project_id, df, query_hash=normalised_query_hash(query_input.query))
+ except ValueError:
+ raise HTTPException(status_code=400, detail="Project id is not a valid UUID.")
+ except SnapshotTooLarge as err:
+ logger.error(str(err))
+ raise HTTPException(status_code=413, detail="Cohort snapshot exceeds the configured size limit.")
+ except OSError:
+ logger.exception("Cohort snapshot store write failed")
+ raise HTTPException(status_code=500, detail="Snapshot persistence failed.")
+
+ return SnapshotResponse(
+ row_count=meta.row_count,
+ columns=meta.columns,
+ has_accessions=meta.has_accessions,
+ snapshot_at=meta.created_at,
+ query_hash=meta.query_hash,
+ )
+
+
+@write_router.post("/snapshot/delete")
+def remove_snapshot(query_input: SnapshotDeleteRequest) -> dict[str, bool]:
+ """
+ Removes a project's frozen cohort artefact. Idempotent.
+
+ The teardown hook for the project purge path (FLIP#997 — which has no hub-side caller
+ yet, so nothing invokes this in the current lifecycle). After deletion the project's
+ row-level routes refuse until a re-approval creates a fresh snapshot.
+
+ Args:
+ query_input (SnapshotDeleteRequest): The encrypted project id.
+
+ Returns:
+ dict[str, bool]: ``{"deleted": bool}`` — False when no snapshot existed.
+ """
+ project_id = decrypt(query_input.encrypted_project_id)
+ deleted = delete_snapshot(project_id)
+ logger.info(f"Snapshot delete for project {project_id}: {'removed' if deleted else 'nothing to remove'}")
+ return {"deleted": deleted}
diff --git a/trust/data-access-api/data_access_api/routers/schema.py b/trust/data-access-api/data_access_api/routers/schema.py
index 42c256363..add5a7771 100644
--- a/trust/data-access-api/data_access_api/routers/schema.py
+++ b/trust/data-access-api/data_access_api/routers/schema.py
@@ -82,3 +82,29 @@ class AccessionIdsResponse(BaseModel):
...,
description="The accession IDs of the cohort, in query order.",
)
+
+
+class SnapshotResponse(BaseModel):
+ """What ``POST /cohort/snapshot`` persisted — aggregates only, no row-level data."""
+
+ row_count: int = Field(..., description="Number of cohort rows frozen in the snapshot")
+ columns: list[str] = Field(..., description="Ordered column names of the frozen dataframe")
+ has_accessions: bool = Field(
+ ...,
+ description="Whether the frozen cohort carries an accession_id column (i.e. pulls imaging)",
+ )
+ snapshot_at: str = Field(..., description="ISO-8601 UTC timestamp of snapshot creation")
+ query_hash: str = Field(
+ ...,
+ description="SHA-256 of the normalised SQL the snapshot froze (drift detection, not a control)",
+ )
+
+
+class SnapshotDeleteRequest(BaseModel):
+ """Input for ``POST /cohort/snapshot/delete``."""
+
+ encrypted_project_id: str = Field(
+ ...,
+ description="The encrypted identifier for the central hub project",
+ json_schema_extra={"example": "encrypted_12345"},
+ )
diff --git a/trust/data-access-api/data_access_api/services/cohort.py b/trust/data-access-api/data_access_api/services/cohort.py
index 13855d211..bf25efdec 100644
--- a/trust/data-access-api/data_access_api/services/cohort.py
+++ b/trust/data-access-api/data_access_api/services/cohort.py
@@ -430,6 +430,7 @@ def _validate_query_ast(query: str) -> str:
def get_records(
query: str | TextClause,
params: Mapping[str, Any] | None = None,
+ use_cache: bool = True,
) -> pd.DataFrame:
"""
Executes a SQL query and returns results.
@@ -439,6 +440,12 @@ def get_records(
parameters when the query is parameterized.
params (Mapping[str, Any] | None): Optional mapping of bind parameter names to values
for parameterized queries.
+ use_cache (bool): Whether a cached result may be served. Cohort SNAPSHOT creation
+ passes False: the statistics run at submission caches the same SQL for
+ ``CACHE_TTL_DAYS``, so a cached read would freeze a stale frame — and on
+ RE-approval, which is precisely the event that must propagate OMOP-side
+ removals/opt-outs into the artefact (FLIP#857), it would silently re-freeze the
+ pre-removal cohort. A fresh read still updates the cache afterwards.
Returns:
pd.DataFrame: The results of the query as a DataFrame.
@@ -448,9 +455,10 @@ def get_records(
"""
logger.info("Executing SQL query")
- cached = get_cached_result(query, params)
- if cached is not None:
- return cached
+ if use_cache:
+ cached = get_cached_result(query, params)
+ if cached is not None:
+ return cached
try:
# TODO: Trace the query filtering to understand what the final user can see.
diff --git a/trust/data-access-api/data_access_api/services/cohort_snapshot.py b/trust/data-access-api/data_access_api/services/cohort_snapshot.py
new file mode 100644
index 000000000..b5b8a70cc
--- /dev/null
+++ b/trust/data-access-api/data_access_api/services/cohort_snapshot.py
@@ -0,0 +1,296 @@
+# Copyright (c) Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Durable, project-keyed store for the approved-cohort snapshot (FLIP#857).
+
+At project approval the trust materialises its cohort dataframe once; this module persists
+that artefact on the local filesystem (``COHORT_SNAPSHOT_DIR``, a dedicated bind mount) so
+the row-level routes can serve a frozen, immutable cohort instead of re-running the SQL
+against live OMOP on every call. Deliberately a file store: data-access-api keeps zero
+write access to any database, and researcher SQL (pinned to the ``omop`` schema by
+``validate_query``) cannot reach a filesystem at all.
+
+Layout, one directory per hub project id::
+
+ //
+ dataframe.parquet # the frozen cohort, dtype-faithful (parquet, no index)
+ meta.json # row_count / columns / query_hash / created_at / format_version
+
+Writes are atomic at directory granularity: everything lands in a ``.tmp-*`` sibling first
+and is activated with ``os.replace`` renames, so a reader never observes a half-written
+snapshot and a crash mid-write leaves (at worst) a stale temp directory that the boot-time
+sweep removes. There is deliberately NO TTL and no in-place mutation — a snapshot is
+immutable until it is overwritten by a re-approval or deleted; an approved cohort must not
+silently vanish or drift mid-training (contrast ``services/query_cache.py``, the volatile
+per-process cache this module's API is modelled on).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import os
+import shutil
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+
+import pandas as pd
+
+from data_access_api.config import get_settings
+from data_access_api.utils.logger import logger
+
+# Bumped when the on-disk layout changes; a snapshot with an unknown version is treated as
+# absent (legacy live-SQL fall-through) rather than mis-read.
+_FORMAT_VERSION = 1
+_DATA_FILENAME = "dataframe.parquet"
+_META_FILENAME = "meta.json"
+# Work-in-progress / superseded directories. Never valid snapshots; swept at startup.
+_TMP_PREFIX = ".tmp-"
+_OLD_PREFIX = ".old-"
+
+
+class SnapshotStoreDisabled(Exception):
+ """Raised on writes when ``COHORT_SNAPSHOT_DIR`` is not configured."""
+
+
+class SnapshotTooLarge(Exception):
+ """Raised when the serialized snapshot exceeds ``SNAPSHOT_MAX_BYTES`` (never truncated)."""
+
+
+@dataclass(frozen=True)
+class SnapshotMeta:
+ """The snapshot's serving-relevant facts, readable without deserialising the frame."""
+
+ row_count: int
+ columns: list[str]
+ query_hash: str
+ created_at: str # ISO-8601 UTC
+ format_version: int = _FORMAT_VERSION
+
+ @property
+ def has_accessions(self) -> bool:
+ return "accession_id" in self.columns
+
+
+@dataclass(frozen=True)
+class Snapshot:
+ df: pd.DataFrame
+ meta: SnapshotMeta
+
+
+def normalised_query_hash(query: str) -> str:
+ """SHA-256 of the whitespace-normalised, lowercased SQL text.
+
+ Used only to *detect and log* when a caller-supplied query differs from the one the
+ snapshot froze — never as a security control (the frozen artefact is served either
+ way; that is the point). Hashes the raw submitted text, not the validator's re-emitted
+ form, so the hub-side string of record compares equal across submission and serving.
+ """
+ normalised = " ".join(query.strip().lower().split())
+ return hashlib.sha256(normalised.encode()).hexdigest()
+
+
+def snapshot_enabled() -> bool:
+ """Whether a snapshot directory is configured (empty ``COHORT_SNAPSHOT_DIR`` = disabled)."""
+ return bool(get_settings().COHORT_SNAPSHOT_DIR)
+
+
+def _store_dir() -> Path:
+ configured = get_settings().COHORT_SNAPSHOT_DIR
+ if not configured:
+ raise SnapshotStoreDisabled("COHORT_SNAPSHOT_DIR is not configured")
+ return Path(configured)
+
+
+def _canonical_project_id(project_id: str) -> str | None:
+ """The project id as a canonical UUID string, or None when it is not a UUID.
+
+ The id becomes a directory name, so only a parsed-and-re-emitted UUID is ever used as
+ a path component — nothing else reaches the filesystem (no traversal surface, even
+ though every caller is already authenticated and the id is hub-encrypted).
+ """
+ try:
+ return str(uuid.UUID(str(project_id)))
+ except (ValueError, AttributeError, TypeError):
+ return None
+
+
+def ensure_store() -> None:
+ """Boot-time store check: create the directory, sweep stale temp dirs, probe writability.
+
+ Never raises — a missing or unwritable store must not take the service (and the
+ statistics route) down. Failures log at ERROR with the remediation; every subsequent
+ write fails loudly per-call and every read returns None, which the row-level routes
+ refuse (fail-closed).
+ """
+ if not snapshot_enabled():
+ logger.error(
+ "Cohort snapshot store DISABLED (COHORT_SNAPSHOT_DIR unset): snapshots cannot be "
+ "created and the row-level routes will refuse every project (fail-closed). Set "
+ "COHORT_SNAPSHOT_DIR / mount the snapshot volume."
+ )
+ return
+
+ base = _store_dir()
+ try:
+ base.mkdir(parents=True, exist_ok=True)
+ # Sweep leftovers from crashed writes: only this service writes here, and no write
+ # can be in flight during startup.
+ for stale in base.iterdir():
+ if stale.name.startswith((_TMP_PREFIX, _OLD_PREFIX)):
+ shutil.rmtree(stale, ignore_errors=True)
+ logger.warning(f"Removed stale snapshot work directory {stale.name}")
+ probe = base / f"{_TMP_PREFIX}write-probe"
+ probe.mkdir(exist_ok=True)
+ probe.rmdir()
+ except OSError:
+ logger.exception(
+ f"Cohort snapshot store at {base} is not writable — snapshots cannot be created and "
+ "the row-level routes will refuse projects whose artefact cannot be read (fail-closed). "
+ f"Remediation: create the directory on the host and chown it to this service's uid "
+ f"(uid {os.getuid()})."
+ )
+ return
+
+ logger.info(f"Cohort snapshot store ready at {base}")
+
+
+def save_snapshot(project_id: str, df: pd.DataFrame, query_hash: str) -> SnapshotMeta:
+ """Persist the cohort dataframe for ``project_id``, atomically replacing any predecessor.
+
+ Args:
+ project_id (str): The decrypted hub project id (must be a UUID).
+ df (pd.DataFrame): The cohort exactly as ``get_records`` returned it.
+ query_hash (str): ``normalised_query_hash`` of the raw SQL that produced ``df``.
+
+ Returns:
+ SnapshotMeta: What was written.
+
+ Raises:
+ SnapshotStoreDisabled: When no store directory is configured.
+ SnapshotTooLarge: When the serialized frame exceeds ``SNAPSHOT_MAX_BYTES``.
+ ValueError: When ``project_id`` is not a UUID.
+ OSError: When the store directory is not writable.
+ """
+ base = _store_dir()
+ canonical = _canonical_project_id(project_id)
+ if canonical is None:
+ raise ValueError("project_id must be a UUID to key a cohort snapshot")
+
+ # Parquet keeps pandas dtypes (datetimes, nullable ints) so every training-time fetch of
+ # the frozen cohort deserialises to an identical frame. The index is dropped — the serve
+ # routes emit to_dict(orient="list"), which never includes it.
+ buffer = io.BytesIO()
+ df.to_parquet(buffer, engine="pyarrow", index=False)
+ payload = buffer.getvalue()
+
+ max_bytes = get_settings().SNAPSHOT_MAX_BYTES
+ if len(payload) > max_bytes:
+ # Refuse rather than truncate: a partial cohort silently poisons training data.
+ raise SnapshotTooLarge(
+ f"Serialized cohort snapshot is {len(payload)} bytes, over the {max_bytes}-byte limit "
+ "(SNAPSHOT_MAX_BYTES). Narrow the cohort query's columns, or raise the limit."
+ )
+
+ meta = SnapshotMeta(
+ row_count=len(df),
+ columns=[str(column) for column in df.columns],
+ query_hash=query_hash,
+ created_at=datetime.now(UTC).isoformat(),
+ )
+
+ base.mkdir(parents=True, exist_ok=True)
+ nonce = uuid.uuid4().hex[:8]
+ workdir = base / f"{_TMP_PREFIX}{canonical}-{nonce}"
+ final = base / canonical
+ superseded = base / f"{_OLD_PREFIX}{canonical}-{nonce}"
+ try:
+ workdir.mkdir()
+ (workdir / _DATA_FILENAME).write_bytes(payload)
+ (workdir / _META_FILENAME).write_text(json.dumps(meta.__dict__))
+
+ # Two atomic renames. A crash between them leaves no active snapshot (readers fall
+ # back to live SQL and the boot sweep clears the debris) — never a half-written one.
+ if final.exists():
+ os.replace(final, superseded)
+ os.replace(workdir, final)
+ finally:
+ shutil.rmtree(workdir, ignore_errors=True)
+ shutil.rmtree(superseded, ignore_errors=True)
+
+ logger.info(
+ f"Cohort snapshot saved for project {canonical}: {meta.row_count} rows, "
+ f"{len(meta.columns)} columns, {len(payload)} bytes"
+ )
+ return meta
+
+
+def get_snapshot(project_id: str) -> Snapshot | None:
+ """The frozen cohort for ``project_id``, or None when there is none to serve.
+
+ None covers every no-artefact case — store disabled, non-UUID project id, no snapshot
+ yet, unreadable/corrupt artefact (logged at ERROR). The row-level routes refuse on
+ None (fail-closed); a corrupt artefact is never partially served.
+ """
+ if not snapshot_enabled():
+ return None
+ canonical = _canonical_project_id(project_id)
+ if canonical is None:
+ logger.debug(f"Project id {project_id!r} is not a UUID; no snapshot lookup")
+ return None
+
+ snapshot_dir = _store_dir() / canonical
+ meta_path = snapshot_dir / _META_FILENAME
+ if not meta_path.exists():
+ return None
+
+ try:
+ raw_meta = json.loads(meta_path.read_text())
+ meta = SnapshotMeta(**raw_meta)
+ if meta.format_version != _FORMAT_VERSION:
+ logger.error(
+ f"Cohort snapshot for project {canonical} has format_version "
+ f"{meta.format_version} (expected {_FORMAT_VERSION}) — treating as absent"
+ )
+ return None
+ df = pd.read_parquet(snapshot_dir / _DATA_FILENAME, engine="pyarrow")
+ except Exception:
+ logger.exception(
+ f"Cohort snapshot for project {canonical} is unreadable — treating as absent "
+ "(row-level routes refuse the project rather than serve a partial artefact)"
+ )
+ return None
+
+ return Snapshot(df=df, meta=meta)
+
+
+def delete_snapshot(project_id: str) -> bool:
+ """Remove the snapshot for ``project_id``. Idempotent; True if one existed."""
+ if not snapshot_enabled():
+ return False
+ canonical = _canonical_project_id(project_id)
+ if canonical is None:
+ return False
+
+ snapshot_dir = _store_dir() / canonical
+ if not snapshot_dir.exists():
+ return False
+ # Move aside first so a concurrent reader sees either the intact snapshot or none —
+ # never a directory whose files are vanishing under it mid-read.
+ tomb = _store_dir() / f"{_OLD_PREFIX}{canonical}-{uuid.uuid4().hex[:8]}"
+ os.replace(snapshot_dir, tomb)
+ shutil.rmtree(tomb, ignore_errors=True)
+ logger.info(f"Cohort snapshot deleted for project {canonical}")
+ return True
diff --git a/trust/data-access-api/data_access_api/utils/internal_auth.py b/trust/data-access-api/data_access_api/utils/internal_auth.py
index 13302d60c..d95691b6b 100644
--- a/trust/data-access-api/data_access_api/utils/internal_auth.py
+++ b/trust/data-access-api/data_access_api/utils/internal_auth.py
@@ -23,8 +23,20 @@
The key is held in plaintext by every trust-internal service (sender or
receiver). See ``imaging_api/utils/internal_auth.py`` for the rationale —
the same module-level docstring applies here.
+
+The trust-internal key gates every ``/cohort`` route, but it does not
+distinguish callers: fl-client holds it (legitimately — it reads the frozen
+cohort via ``get_dataframe`` and pulls imaging), so it alone cannot separate
+"may read the approved cohort" from "may DEFINE the approved cohort". The
+snapshot create/delete routes — which materialise and destroy the artefact
+every training round then trains on — therefore carry a second gate,
+``authenticate_cohort_admin``: proof of possessing ``AES_KEY_BASE64``. trust-api
+and data-access-api hold that key (they encrypt/decrypt hub payloads with it);
+fl-client deliberately does not. The proof is the SHA-256 of the key, never the
+key itself (FLIP#857).
"""
+import hashlib
import hmac
from fastapi import HTTPException, Security, status
@@ -40,6 +52,11 @@
auto_error=False,
)
+cohort_admin_header_scheme = APIKeyHeader(
+ name=_settings.COHORT_ADMIN_KEY_HEADER,
+ auto_error=False,
+)
+
def authenticate_internal_service(api_key: str | None = Security(internal_key_header_scheme)) -> None:
"""Authenticate a trust-internal caller (trust-api, imaging-api, fl-client).
@@ -73,3 +90,54 @@ def authenticate_internal_service(api_key: str | None = Security(internal_key_he
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid trust-internal service key.",
)
+
+
+# Returned when a caller is trust-internal-authenticated but is not a cohort admin. Fixed and
+# generic — it must not reveal whether the proof was absent or merely wrong, nor confirm the
+# project exists.
+_NOT_COHORT_ADMIN_DETAIL = "Not authorised to define the cohort snapshot."
+
+
+def cohort_admin_proof() -> str:
+ """The credential the cohort-write routes require: SHA-256 of the trust's AES key.
+
+ Possession of ``AES_KEY_BASE64`` is what separates the services trusted to define a
+ project's approved cohort (trust-api, data-access-api) from those that only consume it
+ (fl-client, which runs researcher code and holds no AES key). Sending the digest rather
+ than the key keeps the encryption key off the wire and out of logs — a captured header is
+ a replayable token scoped to two routes, never the key itself.
+
+ Returns:
+ str: The hex SHA-256 digest of ``AES_KEY_BASE64``.
+ """
+ return hashlib.sha256(get_settings().AES_KEY_BASE64.encode()).hexdigest()
+
+
+def authenticate_cohort_admin(proof: str | None = Security(cohort_admin_header_scheme)) -> None:
+ """Authorise a cohort-defining write (snapshot create/delete) on top of trust-internal auth.
+
+ Layered under ``authenticate_internal_service`` on the write router, so it only runs for
+ callers that already passed the trust-internal key check — hence 403 (authenticated, not
+ authorised) rather than 401. fl-client passes the first gate and fails this one; trust-api
+ passes both.
+
+ Args:
+ proof (str | None): The SHA-256-of-AES-key digest from the request header.
+
+ Raises:
+ HTTPException: 403 if the proof is missing, invalid, or the AES key is unconfigured.
+ """
+ aes_key = get_settings().AES_KEY_BASE64
+ if not aes_key:
+ # Fail closed. AES_KEY_BASE64 is a required setting, so this is defensive: an empty
+ # value would make sha256("") a universally-guessable proof. Refuse rather than admit it.
+ logger.warning("Cohort-admin authorization failed: AES_KEY_BASE64 not configured.")
+ raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=_NOT_COHORT_ADMIN_DETAIL)
+
+ if not proof:
+ logger.warning("Cohort-admin authorization failed: proof missing from request.")
+ raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=_NOT_COHORT_ADMIN_DETAIL)
+
+ if not hmac.compare_digest(proof, cohort_admin_proof()):
+ logger.warning("Cohort-admin authorization failed: invalid proof.")
+ raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=_NOT_COHORT_ADMIN_DETAIL)
diff --git a/trust/data-access-api/pyproject.toml b/trust/data-access-api/pyproject.toml
index 6cc509479..185995235 100644
--- a/trust/data-access-api/pyproject.toml
+++ b/trust/data-access-api/pyproject.toml
@@ -25,6 +25,8 @@ dependencies = [
"idna>=3.15",
"pandas>=2.2.3",
"psycopg2-binary>=2.9.10",
+ # Parquet (de)serialisation of the approved-cohort snapshots (FLIP#857).
+ "pyarrow>=19.0.0",
"pydantic-settings>=2.9.1",
"python-multipart>=0.0.26",
"sqlalchemy>=2.0.39",
diff --git a/trust/data-access-api/tests/conftest.py b/trust/data-access-api/tests/conftest.py
index e961ef56e..dc4f69aa3 100644
--- a/trust/data-access-api/tests/conftest.py
+++ b/trust/data-access-api/tests/conftest.py
@@ -10,6 +10,7 @@
# limitations under the License.
#
+import hashlib
import os
# Plaintext test key used by tests to authenticate against the /cohort router.
@@ -18,6 +19,17 @@
TEST_TRUST_INTERNAL_SERVICE_KEY = "test-trust-internal-key"
AUTH_HEADERS = {"X-Trust-Internal-Service-Key": TEST_TRUST_INTERNAL_SERVICE_KEY}
+# 32-byte AES key (base64), single source for both the env default below and the
+# cohort-admin proof the snapshot WRITE routes require (FLIP#857).
+TEST_AES_KEY_BASE64 = "QgZ+TBA0lUxcuCiRPLneFe/JjMaUEUJWHACHHGz2gGA=" # pragma: allowlist secret
+
+# Cohort-admin authorisation for the snapshot create/delete routes: proof of possessing
+# AES_KEY_BASE64, sent as the SHA-256 of the key. The read routes never need this header.
+COHORT_ADMIN_PROOF = hashlib.sha256(TEST_AES_KEY_BASE64.encode()).hexdigest()
+COHORT_ADMIN_HEADERS = {"X-Cohort-Admin-Key": COHORT_ADMIN_PROOF}
+# The write routes require BOTH the trust-internal key AND the cohort-admin proof.
+WRITE_AUTH_HEADERS = {**AUTH_HEADERS, **COHORT_ADMIN_HEADERS}
+
# Set dummy environment variables required by Settings() before any app code
# is imported. These are only used in tests; real values come from Docker
# Compose environment or .env files in deployed environments.
@@ -25,7 +37,7 @@
"DATA_ACCESS_POSTGRES_USER": "test_user",
"DATA_ACCESS_POSTGRES_PASSWORD": "test_password",
"OMOP_POSTGRES_DB": "test_omop_db",
- "AES_KEY_BASE64": "QgZ+TBA0lUxcuCiRPLneFe/JjMaUEUJWHACHHGz2gGA=", # 32-byte key, base64
+ "AES_KEY_BASE64": TEST_AES_KEY_BASE64,
"TRUST_INTERNAL_SERVICE_KEY": TEST_TRUST_INTERNAL_SERVICE_KEY,
}
diff --git a/trust/data-access-api/tests/integration/conftest.py b/trust/data-access-api/tests/integration/conftest.py
index c7b59c364..dc4787d78 100644
--- a/trust/data-access-api/tests/integration/conftest.py
+++ b/trust/data-access-api/tests/integration/conftest.py
@@ -24,6 +24,7 @@
process and the container always agree on the key the compose stack uses.
"""
+import hashlib
import os
# Pin the AES key BEFORE any data_access_api module loads its Settings singleton.
@@ -50,6 +51,11 @@
TRUST_INTERNAL_KEY = "test-trust-internal-service-key" # pragma: allowlist secret
AUTH_HEADERS = {"X-Trust-Internal-Service-Key": TRUST_INTERNAL_KEY}
+# Cohort-admin proof for the snapshot WRITE routes (FLIP#857): SHA-256 of the AES key the
+# container runs with (pinned above / baked into compose.test.yml). The read routes never
+# need it. Merged per-request onto the write calls, on top of the client's AUTH_HEADERS.
+COHORT_ADMIN_HEADERS = {"X-Cohort-Admin-Key": hashlib.sha256(os.environ["AES_KEY_BASE64"].encode()).hexdigest()}
+
@pytest.fixture(scope="session")
def compose_stack() -> Generator[DockerCompose, None, None]:
@@ -58,6 +64,10 @@ def compose_stack() -> Generator[DockerCompose, None, None]:
context=str(_COMPOSE_DIR),
compose_file_name=_COMPOSE_FILE,
wait=True,
+ # Rebuild the from-source service on every session: without it, `docker compose up`
+ # reuses a previously built image, so a dependency change (pyproject/uv.lock) never
+ # reaches the stack and the suite green-lights code that cannot run in a fresh build.
+ build=True,
) as compose:
yield compose
diff --git a/trust/data-access-api/tests/integration/test_cohort_endpoint.py b/trust/data-access-api/tests/integration/test_cohort_endpoint.py
index 69d938dfd..40b5c786e 100644
--- a/trust/data-access-api/tests/integration/test_cohort_endpoint.py
+++ b/trust/data-access-api/tests/integration/test_cohort_endpoint.py
@@ -21,7 +21,7 @@
import httpx
import pytest
-from tests.integration.conftest import AUTH_HEADERS
+from tests.integration.conftest import AUTH_HEADERS, COHORT_ADMIN_HEADERS
@pytest.fixture
@@ -41,17 +41,23 @@ def _cohort_payload(query: str) -> dict:
}
-def _dataframe_payload(query: str) -> dict:
- """Payload for the two row-level routes, which take ``DataframeQuery`` (two fields only).
+# Row-level routes key everything on the hub project id, which must be a UUID (it becomes a
+# directory name in the snapshot store).
+_PROJECT_A = "0b91a3f2-30c3-4bd5-9a1e-2f24c7f5a111"
+_PROJECT_B = "4d5e6f70-8192-4a3b-bc4d-5e6f70819222"
+
+
+def _dataframe_payload(query: str, project_id: str = _PROJECT_A) -> dict:
+ """Payload for the row-level + snapshot routes, which take ``DataframeQuery``.
``encrypted_project_id`` is encrypted with the same AES key the container is configured
- with, so the decrypt path stays real rather than mocked. The import is function-local for
- the same reason it is in ``test_dataframe_endpoint_returns_seeded_columns``: the conftest
- pins ``AES_KEY_BASE64`` before any ``data_access_api`` module builds its Settings singleton.
+ with, so the decrypt path stays real rather than mocked. The import is function-local:
+ the conftest pins ``AES_KEY_BASE64`` before any ``data_access_api`` module builds its
+ Settings singleton.
"""
from data_access_api.utils.encryption import encrypt
- return {"encrypted_project_id": encrypt("integration-project-1"), "query": query}
+ return {"encrypted_project_id": encrypt(project_id), "query": query}
def test_cohort_endpoint_returns_aggregates_for_image_occurrences(http_client):
@@ -132,112 +138,168 @@ def test_cohort_endpoint_requires_auth_header(data_access_api_url):
assert response.status_code == 401
-def test_dataframe_endpoint_returns_seeded_columns(http_client):
- """``/cohort/dataframe`` returns column-oriented data straight from Postgres.
+# ---------------------------------------------------------------------------
+# Snapshot lifecycle: the row-level routes serve ONLY the frozen approved cohort
+# (FLIP#857). Each test uses its own project UUID so the per-session store never
+# couples tests; snapshots are cleaned up where a later test would collide.
+# ---------------------------------------------------------------------------
- The endpoint decrypts ``encrypted_project_id`` with the shared AES key, so the test
- encrypts a real project id with the same key the container is configured with. This
- keeps the encryption path real instead of mocking ``decrypt``.
- """
+_SEED_COHORT_QUERY = (
+ "SELECT c.concept_code AS modality, io.accession_id "
+ "FROM omop.image_occurrence io "
+ "LEFT JOIN omop.concept c ON c.concept_id = io.modality_concept_id"
+)
+
+
+def _create_snapshot(http_client, query: str, project_id: str) -> httpx.Response:
+ # The write route needs the cohort-admin proof on top of the client's trust-internal key.
+ return http_client.post(
+ "/cohort/snapshot", json=_dataframe_payload(query, project_id), headers=COHORT_ADMIN_HEADERS
+ )
+
+
+def test_row_level_routes_refuse_without_a_snapshot(http_client):
+ """Fail-closed: no approved snapshot ⇒ no row-level data, on both routes."""
+ payload = _dataframe_payload("SELECT * FROM omop.image_occurrence", "97fca5ab-0000-4000-8000-000000000001")
+ for path in ("/cohort/dataframe", "/cohort/accession-ids"):
+ response = http_client.post(path, json=payload)
+ assert response.status_code == 403, response.text
+ assert response.json()["detail"] == "No approved cohort snapshot exists for this project."
+
+
+def test_snapshot_write_routes_reject_trust_internal_key_without_cohort_admin_proof(http_client):
+ """End-to-end: the container refuses a cohort-defining write from a caller that holds only the
+ shared trust-internal key (what fl-client has) — the AES-possession gate closes that path
+ (FLIP#857). ``http_client`` carries AUTH_HEADERS but not COHORT_ADMIN_HEADERS."""
from data_access_api.utils.encryption import encrypt
- payload = {
- "encrypted_project_id": encrypt("integration-project-1"),
- "query": (
- "SELECT c.concept_code AS modality, io.accession_id "
- "FROM omop.image_occurrence io "
- "LEFT JOIN omop.concept c ON c.concept_id = io.modality_concept_id"
- ),
- }
- response = http_client.post("/cohort/dataframe", json=payload)
- assert response.status_code == 200, response.text
+ create = http_client.post("/cohort/snapshot", json=_dataframe_payload(_SEED_COHORT_QUERY, _PROJECT_B))
+ assert create.status_code == 403, create.text
+ delete = http_client.post("/cohort/snapshot/delete", json={"encrypted_project_id": encrypt(_PROJECT_B)})
+ assert delete.status_code == 403, delete.text
- body = response.json()
- assert set(body.keys()) == {"modality", "accession_id"}
- assert len(body["modality"]) == 24
- # Sanity: counts in the dataframe should match the seed totals.
- assert body["modality"].count("CT") == 12
- assert body["modality"].count("MR") == 8
- assert body["modality"].count("XR") == 4
- # accession_id is unique per row in the seed.
- assert len(set(body["accession_id"])) == 24
-
-
-def test_accession_ids_returns_seeded_ids(http_client):
- """``/cohort/accession-ids`` projects the id column out of the caller's cohort.
-
- The route wraps the caller's (re-emitted) SQL as ``SELECT accession_id FROM (...) AS
- cohort_subquery``, so only that one column ever crosses the trust boundary regardless of
- what the inner query selected.
- """
+
+def test_snapshot_then_dataframe_serves_the_frozen_cohort(http_client):
+ """The full approval-time flow: freeze once, then serve — ignoring the client's SQL."""
+ created = _create_snapshot(http_client, _SEED_COHORT_QUERY, _PROJECT_A)
+ assert created.status_code == 200, created.text
+ body = created.json()
+ assert body["row_count"] == 24
+ assert body["has_accessions"] is True
+ assert set(body["columns"]) == {"modality", "accession_id"}
+
+ # The served query is hostile-looking and never executed: the frozen frame comes back.
response = http_client.post(
- "/cohort/accession-ids",
- json=_dataframe_payload("SELECT person_id, accession_id FROM omop.image_occurrence"),
+ "/cohort/dataframe",
+ json=_dataframe_payload("SELECT * FROM omop.person", _PROJECT_A),
)
assert response.status_code == 200, response.text
-
- accession_ids = response.json()["accession_ids"]
+ frame = response.json()
+ assert set(frame.keys()) == {"modality", "accession_id"}
+ assert len(frame["modality"]) == 24
+ assert frame["modality"].count("CT") == 12
+ assert frame["modality"].count("MR") == 8
+ assert frame["modality"].count("XR") == 4
+ assert len(set(frame["accession_id"])) == 24
+
+ # And the frozen accession pointer set serves from the same artefact.
+ ids_response = http_client.post(
+ "/cohort/accession-ids",
+ json=_dataframe_payload("SELECT 1 AS nothing FROM omop.person", _PROJECT_A),
+ )
+ assert ids_response.status_code == 200, ids_response.text
+ accession_ids = ids_response.json()["accession_ids"]
assert len(accession_ids) == 24
- # accession_id is unique per row in the seed (ACC-1001 … ACC-1024).
- assert len(set(accession_ids)) == 24
assert "ACC-1001" in accession_ids
-def test_accession_ids_missing_column_surfaces_get_records_400(http_client):
- """A cohort that does not project ``accession_id`` fails inside ``get_records``, not after it.
+def test_tabular_snapshot_serves_empty_accession_list(http_client):
+ """A frozen cohort without accession_id is a tabular project: imaging no-ops, no error.
- Pins which 400 actually reaches the caller. Because the route wraps the inner query in
- ``SELECT accession_id FROM (...)``, Postgres raises ``UndefinedColumn`` while the query is
- executing — so ``get_records`` converts it to a category 400 and the router's
- ``except HTTPException: raise`` branch re-raises it before any DataFrame is returned. A
- router-level "did the DataFrame come back with the column?" guard could never fire, which
- is why there isn't one; this test is the standing proof of that. The cohort here is 16
- rows — comfortably over the stack's threshold — so the refusal is about the column, not
- the cohort size.
+ (Pre-snapshot, a cohort not projecting accession_id produced an UndefinedColumn 400 out
+ of the live wrapped query; with frozen serving the column's absence is a legitimate
+ project shape, recorded in the snapshot's metadata.)
"""
+ created = _create_snapshot(http_client, "SELECT person_id FROM omop.person", _PROJECT_B)
+ assert created.status_code == 200, created.text
+ assert created.json()["has_accessions"] is False
+
response = http_client.post(
"/cohort/accession-ids",
- json=_dataframe_payload("SELECT person_id FROM omop.person"),
+ json=_dataframe_payload("SELECT person_id FROM omop.person", _PROJECT_B),
)
- assert response.status_code == 400, response.text
- assert response.json()["detail"] == "The column 'accession_id' does not exist."
+ assert response.status_code == 200, response.text
+ assert response.json()["accession_ids"] == []
-def test_accession_ids_below_threshold_is_refused(http_client):
- """A below-threshold cohort is refused outright — no partial list, no count.
+def test_below_threshold_snapshot_is_refused_and_persists_nothing(http_client):
+ """The disclosure floor is enforced at freeze time; a refused freeze leaves no artefact.
modality_concept_id 4013632 = 'XR'; the seed has 4 such rows, under the stack's
- COHORT_QUERY_THRESHOLD of 5.
+ COHORT_QUERY_THRESHOLD of 5. The refusal text matches the row-level routes' fixed
+ below-threshold string, and a zero-row cohort refuses byte-identically so the snapshot
+ route cannot act as a row-count oracle either.
"""
+ project_id = "97fca5ab-0000-4000-8000-000000000002"
+ below = _create_snapshot(
+ http_client,
+ "SELECT accession_id FROM omop.image_occurrence WHERE modality_concept_id = 4013632",
+ project_id,
+ )
+ zero = _create_snapshot(
+ http_client,
+ "SELECT accession_id FROM omop.image_occurrence WHERE accession_id = 'NONEXISTENT'",
+ project_id,
+ )
+ assert below.status_code == zero.status_code == 403
+ assert below.text == zero.text
+ assert below.json()["detail"] == "Cohort is too small for row-level data to be released."
+
+ # Nothing was persisted: the project still refuses row-level serving outright.
response = http_client.post(
- "/cohort/accession-ids",
- json=_dataframe_payload(
- "SELECT accession_id FROM omop.image_occurrence WHERE modality_concept_id = 4013632"
- ),
+ "/cohort/dataframe", json=_dataframe_payload("SELECT 1 AS one FROM omop.person", project_id)
)
- assert response.status_code == 403, response.text
- assert response.json()["detail"] == "Cohort is too small for row-level data to be released."
+ assert response.status_code == 403
+ assert response.json()["detail"] == "No approved cohort snapshot exists for this project."
-def test_accession_ids_zero_rows_indistinguishable_from_below_threshold(http_client):
- """Privacy regression, against real Postgres rather than a mocked DataFrame.
+def test_snapshot_route_rejects_unsafe_sql(http_client):
+ """validate_query remains the authority on the one query the snapshot route executes."""
+ response = _create_snapshot(
+ http_client,
+ "INSERT INTO omop.person (person_id) VALUES (999)",
+ "97fca5ab-0000-4000-8000-000000000003",
+ )
+ assert response.status_code == 400, response.text
- A cohort matching nothing and a cohort of 1-4 rows must produce byte-identical refusals,
- or the response itself becomes a row-count oracle: a caller could binary-search a
- predicate and learn whether *any* patient matches it, one bit at a time.
- """
- below_threshold = http_client.post(
- "/cohort/accession-ids",
- json=_dataframe_payload(
- "SELECT accession_id FROM omop.image_occurrence WHERE modality_concept_id = 4013632"
- ),
+
+def test_reapproval_replaces_the_snapshot_and_delete_removes_it(http_client):
+ """Overwrite-on-reapproval and the FLIP#997 teardown hook, end to end."""
+ from data_access_api.utils.encryption import encrypt
+
+ project_id = "97fca5ab-0000-4000-8000-000000000004"
+ first = _create_snapshot(http_client, _SEED_COHORT_QUERY, project_id)
+ assert first.status_code == 200, first.text
+
+ second = _create_snapshot(http_client, "SELECT person_id, accession_id FROM omop.image_occurrence", project_id)
+ assert second.status_code == 200, second.text
+
+ served = http_client.post(
+ "/cohort/dataframe", json=_dataframe_payload("SELECT 1 AS one FROM omop.person", project_id)
)
- genuine_zero = http_client.post(
- "/cohort/accession-ids",
- json=_dataframe_payload(
- "SELECT accession_id FROM omop.image_occurrence WHERE accession_id = 'NONEXISTENT'"
- ),
+ assert set(served.json().keys()) == {"person_id", "accession_id"}
+
+ deleted = http_client.post(
+ "/cohort/snapshot/delete", json={"encrypted_project_id": encrypt(project_id)}, headers=COHORT_ADMIN_HEADERS
)
+ assert deleted.status_code == 200
+ assert deleted.json() == {"deleted": True}
+ again = http_client.post(
+ "/cohort/snapshot/delete", json={"encrypted_project_id": encrypt(project_id)}, headers=COHORT_ADMIN_HEADERS
+ )
+ assert again.json() == {"deleted": False}
- assert below_threshold.status_code == genuine_zero.status_code == 403
- assert below_threshold.text == genuine_zero.text
+ refused = http_client.post(
+ "/cohort/dataframe", json=_dataframe_payload("SELECT 1 AS one FROM omop.person", project_id)
+ )
+ assert refused.status_code == 403
diff --git a/trust/data-access-api/tests/routers/test_cohort.py b/trust/data-access-api/tests/routers/test_cohort.py
index e6bf2280b..805ce6254 100644
--- a/trust/data-access-api/tests/routers/test_cohort.py
+++ b/trust/data-access-api/tests/routers/test_cohort.py
@@ -16,7 +16,6 @@
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
-from sqlalchemy.exc import SQLAlchemyError
from data_access_api.main import app
from data_access_api.routers.schema import StatisticsResponse
@@ -187,318 +186,11 @@ def test_receive_cohort_query_execution_error(mock_validate_query, mock_get_sett
"query": "SELECT age, gender FROM dummy_table",
}
-# Expected output
-sample_df_dict = {
- "age": [25, 30, 40],
- "gender": ["M", "F", "M"],
-}
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-@patch("data_access_api.routers.cohort.validate_query")
-def test_get_dataframe_success(mock_validate_query, mock_get_records, mock_decrypt, mock_get_settings):
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame(sample_df_dict)
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 200
- assert response.json() == sample_df_dict
- mock_decrypt.assert_called_once_with("encrypted-id")
- mock_validate_query.assert_called_once_with(sample_dataframe_query["query"])
- # The engine receives what validate_query emitted from the checked AST,
- # never the caller's raw string.
- mock_get_records.assert_called_once_with(mock_validate_query.return_value)
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.validate_query")
-def test_get_dataframe_invalid_query(mock_validate_query, mock_decrypt):
- mock_decrypt.return_value = "decrypted-id"
- mock_validate_query.side_effect = HTTPException(status_code=400, detail="Invalid query syntax")
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 400
- assert response.json()["detail"] == "Invalid query syntax"
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_sqlalchemy_error(mock_get_records, mock_decrypt):
- """Driver text must not reach the caller — the hub relays this detail to the UI."""
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = SQLAlchemyError("relation omop.person has 42 rows for patient Bob")
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 500
- assert "Bob" not in response.json()["detail"]
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_generic_error(mock_get_records, mock_decrypt):
- """Unexpected exception text must not reach the caller either."""
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = RuntimeError("connection string postgres://user:hunter2@omop-db")
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 500
- assert "hunter2" not in response.json()["detail"]
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_preserves_http_exception_from_get_records(mock_get_records, mock_decrypt):
- """A categorised 400 from get_records must not be re-wrapped as an opaque 500.
-
- ``get_records`` deliberately converts driver errors into category-only
- HTTPExceptions. ``except Exception`` also catches HTTPException, so the
- route used to swallow that work and re-raise every one of them as a 500
- whose detail was the *repr of the HTTPException*.
- """
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = HTTPException(status_code=400, detail="Column 'nope' does not exist")
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 400
- assert response.json()["detail"] == "Column 'nope' does not exist"
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_rejects_cohort_below_threshold(mock_get_records, mock_decrypt, mock_get_settings):
- """Row-level data is withheld for cohorts smaller than COHORT_QUERY_THRESHOLD."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"age": range(9)})
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 403
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_below_threshold_does_not_disclose_row_count(
- mock_get_records, mock_decrypt, mock_get_settings
-):
- """The refusal must not reveal how many rows matched — 0 and 9 look identical."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
-
- details = []
- for row_count in (0, 9):
- mock_get_records.return_value = pd.DataFrame({"age": range(row_count)})
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
- details.append(response.json()["detail"])
-
- assert details[0] == details[1]
- assert "9" not in details[1]
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_dataframe_allows_cohort_at_threshold(mock_get_records, mock_decrypt, mock_get_settings):
- """A cohort exactly at the threshold is released — the gate is not off-by-one."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"age": range(10)})
-
- response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 200
- assert len(response.json()["age"]) == 10
-
-
-# ---------------------------------------------------------------------------
-# /cohort/accession-ids
-# ---------------------------------------------------------------------------
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_success(mock_get_records, mock_decrypt, mock_get_settings):
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"accession_id": ["ACC1", "ACC2", "ACC3"]})
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 200
- assert response.json() == {"accession_ids": ["ACC1", "ACC2", "ACC3"]}
- mock_decrypt.assert_called_once_with("encrypted-id")
- # The caller's query must be wrapped server-side so only accession_id is projected.
- called_query = mock_get_records.call_args[0][0]
- assert called_query.startswith("SELECT accession_id FROM (")
- # What gets wrapped is validate_query's *emitted* SQL, not the caller's raw string — so the
- # unqualified `dummy_table` arrives pinned to the omop schema.
- assert "SELECT age, gender FROM omop.dummy_table" in called_query
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_strips_trailing_semicolon(mock_get_records, mock_decrypt, mock_get_settings):
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 1
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"accession_id": ["ACC1"]})
-
- response = client.post(
- "/cohort/accession-ids",
- json={"encrypted_project_id": "encrypted-id", "query": "SELECT * FROM cohort; "},
- headers=AUTH_HEADERS,
- )
-
- assert response.status_code == 200
- called_query = mock_get_records.call_args[0][0]
- # No bare semicolon should leak into the wrapped subquery. The table arrives schema-pinned
- # because the wrapper wraps the emitted SQL, not the caller's raw string.
- assert "SELECT * FROM omop.cohort)" in called_query
- assert ";" not in called_query
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.validate_query")
-def test_get_accession_ids_invalid_query(mock_validate_query, mock_decrypt):
- mock_decrypt.return_value = "decrypted-id"
- mock_validate_query.side_effect = HTTPException(status_code=400, detail="Invalid query syntax")
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 400
- assert response.json()["detail"] == "Invalid query syntax"
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_missing_column_propagates_400(mock_get_records, mock_decrypt):
- """A cohort that does not project ``accession_id`` is refused by ``get_records``, not the router.
-
- The route wraps the inner query as ``SELECT accession_id FROM (...)``, so Postgres raises
- ``UndefinedColumn`` during execution and ``get_records`` turns it into a category 400 —
- there is no DataFrame to inspect afterwards, which is why the router carries no
- column-presence guard. This mocks the conversion ``get_records`` performs; the real
- Postgres behaviour it stands in for is pinned by
- ``tests/integration/test_cohort_endpoint.py::test_accession_ids_missing_column_surfaces_get_records_400``.
- """
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = HTTPException(
- status_code=400, detail="The column 'accession_id' does not exist."
- )
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 400
- assert response.json()["detail"] == "The column 'accession_id' does not exist."
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_propagates_http_exception(mock_get_records, mock_decrypt):
- """``get_records`` raises HTTPException for things like undefined tables/columns;
- the wrapping ``except`` clauses must not swallow that into a 500."""
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = HTTPException(
- status_code=400, detail="The table 'omop.bogus' does not exist."
- )
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 400
- assert response.json()["detail"] == "The table 'omop.bogus' does not exist."
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_sqlalchemy_error_does_not_leak(mock_get_records, mock_decrypt):
- """Driver text must not reach the caller: the trust forwards this detail to the hub,
- which shows it to every project member (FLIP-PT-016)."""
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = SQLAlchemyError("relation omop.secret_table row 42")
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 500
- assert response.json()["detail"] == "Query execution failed."
- assert "secret_table" not in response.text
-
-
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_generic_error_does_not_leak(mock_get_records, mock_decrypt):
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.side_effect = RuntimeError("connection to 10.0.0.5 failed for user svc_omop")
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 500
- assert response.json()["detail"] == "Query execution failed."
- assert "svc_omop" not in response.text
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_rejects_cohort_below_threshold(mock_get_records, mock_decrypt, mock_get_settings):
- """Accession IDs are row-level identifiers and decide whose imaging is pulled into XNAT,
- so a below-threshold cohort is refused just as it is on /cohort/dataframe."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"accession_id": [f"ACC{i}" for i in range(9)]})
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 403
- assert response.json()["detail"] == "Cohort is too small for row-level data to be released."
- # No identifier may appear in the refusal.
- assert "ACC" not in response.text
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_below_threshold_is_indistinguishable_from_zero(
- mock_get_records, mock_decrypt, mock_get_settings
-):
- """A zero-row cohort and a below-threshold one must return byte-identical responses,
- or the refusal itself becomes a row-count oracle."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
-
- mock_get_records.return_value = pd.DataFrame({"accession_id": []})
- zero_response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- mock_get_records.return_value = pd.DataFrame({"accession_id": [f"ACC{i}" for i in range(9)]})
- below_response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert zero_response.status_code == below_response.status_code == 403
- assert zero_response.text == below_response.text
-
-
-@patch("data_access_api.routers.cohort.get_settings")
-@patch("data_access_api.routers.cohort.decrypt")
-@patch("data_access_api.routers.cohort.get_records")
-def test_get_accession_ids_allows_cohort_at_threshold(mock_get_records, mock_decrypt, mock_get_settings):
- """Exactly at the threshold is allowed — the gate is `<`, not `<=`."""
- mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
- mock_decrypt.return_value = "decrypted-id"
- mock_get_records.return_value = pd.DataFrame({"accession_id": [f"ACC{i}" for i in range(10)]})
-
- response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
-
- assert response.status_code == 200
- assert len(response.json()["accession_ids"]) == 10
+# The /cohort/dataframe and /cohort/accession-ids behavioural tests live in
+# tests/routers/test_cohort_snapshot.py: both routes serve ONLY the frozen
+# approved-cohort snapshot (FLIP#857) — the live-SQL serving path they used to
+# have (and the tests that pinned it) is gone. Live SQL now runs only on
+# /cohort (statistics) above and /cohort/snapshot (creation).
# ---------------------------------------------------------------------------
@@ -514,6 +206,8 @@ def test_get_accession_ids_allows_cohort_at_threshold(mock_get_records, mock_dec
("/cohort", sample_query_input),
("/cohort/dataframe", sample_dataframe_query),
("/cohort/accession-ids", sample_dataframe_query),
+ ("/cohort/snapshot", sample_dataframe_query),
+ ("/cohort/snapshot/delete", {"encrypted_project_id": "encrypted-id"}),
],
)
def test_cohort_route_rejects_missing_key(path, payload):
@@ -528,6 +222,8 @@ def test_cohort_route_rejects_missing_key(path, payload):
("/cohort", sample_query_input),
("/cohort/dataframe", sample_dataframe_query),
("/cohort/accession-ids", sample_dataframe_query),
+ ("/cohort/snapshot", sample_dataframe_query),
+ ("/cohort/snapshot/delete", {"encrypted_project_id": "encrypted-id"}),
],
)
def test_cohort_route_rejects_wrong_key(path, payload):
@@ -544,6 +240,53 @@ def test_health_does_not_require_auth():
assert response.status_code == 200
+# ---------------------------------------------------------------------------
+# Cohort-admin auth — the snapshot WRITE routes require proof of possessing
+# AES_KEY_BASE64 on top of the trust-internal key (FLIP#857), so a caller that
+# holds only the shared trust-internal key (fl-client) cannot DEFINE or destroy a
+# project's frozen cohort. The read routes must stay reachable with the
+# trust-internal key alone.
+# ---------------------------------------------------------------------------
+
+_WRITE_ROUTES = [
+ ("/cohort/snapshot", sample_dataframe_query),
+ ("/cohort/snapshot/delete", {"encrypted_project_id": "encrypted-id"}),
+]
+
+
+@pytest.mark.parametrize(("path", "payload"), _WRITE_ROUTES)
+def test_write_routes_reject_trust_internal_key_without_cohort_admin_proof(path, payload):
+ """A valid trust-internal key alone (what fl-client holds) is refused with 403 —
+ authenticated but not authorised to define the cohort."""
+ response = client.post(path, json=payload, headers=AUTH_HEADERS)
+ assert response.status_code == 403
+ assert "authorised" in response.json()["detail"].lower()
+
+
+@pytest.mark.parametrize(("path", "payload"), _WRITE_ROUTES)
+def test_write_routes_reject_wrong_cohort_admin_proof(path, payload):
+ """A wrong AES-possession proof is refused with the same fixed 403 as a missing one,
+ so the refusal never reveals whether the proof was absent or merely invalid."""
+ headers = {**AUTH_HEADERS, "X-Cohort-Admin-Key": "not-the-real-proof"}
+ response = client.post(path, json=payload, headers=headers)
+ assert response.status_code == 403
+ assert "authorised" in response.json()["detail"].lower()
+
+
+@pytest.mark.parametrize("path", ["/cohort/dataframe", "/cohort/accession-ids"])
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+def test_read_routes_do_not_require_cohort_admin_proof(mock_get_snapshot, mock_decrypt, path):
+ """The read routes must NOT gain the cohort-admin gate: the trust-internal key alone must
+ get past auth into the handler (a cohort-admin 403 here would mean fl-client's get_dataframe
+ broke). With no snapshot the handler reaches its own fail-closed 403 — distinct text — which
+ proves auth let the caller through rather than blocking on cohort-admin."""
+ mock_decrypt.return_value = "my_project"
+ mock_get_snapshot.return_value = None
+ response = client.post(path, json=sample_dataframe_query, headers=AUTH_HEADERS)
+ assert "authorised" not in response.json().get("detail", "").lower()
+
+
# ---------------------------------------------------------------------------
# Parse-then-emit tests
#
diff --git a/trust/data-access-api/tests/routers/test_cohort_snapshot.py b/trust/data-access-api/tests/routers/test_cohort_snapshot.py
new file mode 100644
index 000000000..187d8ceb1
--- /dev/null
+++ b/trust/data-access-api/tests/routers/test_cohort_snapshot.py
@@ -0,0 +1,304 @@
+# Copyright (c) Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Route-level tests for approved-cohort snapshot serving and creation (FLIP#857)."""
+
+from datetime import UTC, datetime
+from unittest.mock import patch
+
+import pandas as pd
+import pytest
+from fastapi.testclient import TestClient
+
+from data_access_api.main import app
+from data_access_api.routers.cohort import _BELOW_THRESHOLD_DETAIL, _NO_SNAPSHOT_DETAIL
+from data_access_api.services.cohort_snapshot import (
+ Snapshot,
+ SnapshotMeta,
+ SnapshotTooLarge,
+ normalised_query_hash,
+)
+from tests.conftest import AUTH_HEADERS, WRITE_AUTH_HEADERS
+
+client = TestClient(app)
+
+FROZEN_QUERY = "SELECT * FROM omop.person"
+
+sample_dataframe_query = {
+ "encrypted_project_id": "encrypted_my_project",
+ "query": FROZEN_QUERY,
+}
+
+
+def _snapshot(df: pd.DataFrame, query: str = FROZEN_QUERY) -> Snapshot:
+ return Snapshot(
+ df=df,
+ meta=SnapshotMeta(
+ row_count=len(df),
+ columns=[str(column) for column in df.columns],
+ query_hash=normalised_query_hash(query),
+ created_at=datetime.now(UTC).isoformat(),
+ ),
+ )
+
+
+# ---------------------------------------------------------------------------
+# Frozen serving on /cohort/dataframe
+# ---------------------------------------------------------------------------
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.get_records")
+def test_dataframe_serves_frozen_snapshot_and_ignores_client_sql(
+ mock_get_records, mock_validate_query, mock_get_snapshot, mock_decrypt, mock_get_settings
+):
+ """With a snapshot present, even hostile SQL is never validated or executed."""
+ mock_decrypt.return_value = "my_project"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ frozen = pd.DataFrame({"accession_id": ["A1", "A2", "A3"], "label": [0, 1, 0]})
+ mock_get_snapshot.return_value = _snapshot(frozen)
+
+ body = {**sample_dataframe_query, "query": "SELECT * FROM omop.person; DROP TABLE omop.person"}
+ response = client.post("/cohort/dataframe", json=body, headers=AUTH_HEADERS)
+
+ assert response.status_code == 200
+ assert response.json() == frozen.to_dict(orient="list")
+ mock_validate_query.assert_not_called()
+ mock_get_records.assert_not_called()
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+def test_dataframe_frozen_below_threshold_uses_the_fixed_refusal(
+ mock_get_snapshot, mock_decrypt, mock_get_settings
+):
+ """The frozen count is gated with the same fixed text as the live path — the threshold
+ is read live, so an operator raising their floor bites already-approved projects."""
+ mock_decrypt.return_value = "my_project"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
+ mock_get_snapshot.return_value = _snapshot(pd.DataFrame({"accession_id": ["A1"]}))
+
+ response = client.post("/cohort/dataframe", json=sample_dataframe_query, headers=AUTH_HEADERS)
+
+ assert response.status_code == 403
+ assert response.json()["detail"] == _BELOW_THRESHOLD_DETAIL
+
+
+@pytest.mark.parametrize("path", ["/cohort/dataframe", "/cohort/accession-ids"])
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.get_records")
+def test_row_level_routes_refuse_projects_without_a_snapshot(
+ mock_get_records, mock_validate_query, mock_get_snapshot, mock_decrypt, path
+):
+ """No snapshot ⇒ no row-level data, fail-closed: there is no live-SQL serving path."""
+ mock_decrypt.return_value = "my_project"
+ mock_get_snapshot.return_value = None
+
+ response = client.post(path, json=sample_dataframe_query, headers=AUTH_HEADERS)
+
+ assert response.status_code == 403
+ assert response.json()["detail"] == _NO_SNAPSHOT_DETAIL
+ mock_validate_query.assert_not_called()
+ mock_get_records.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# Frozen serving on /cohort/accession-ids
+# ---------------------------------------------------------------------------
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+@patch("data_access_api.routers.cohort.get_records")
+def test_accession_ids_serves_frozen_pointer_set(mock_get_records, mock_get_snapshot, mock_decrypt, mock_get_settings):
+ mock_decrypt.return_value = "my_project"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ mock_get_snapshot.return_value = _snapshot(pd.DataFrame({"accession_id": [101, 102], "label": [0, 1]}))
+
+ response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
+
+ assert response.status_code == 200
+ assert response.json() == {"accession_ids": ["101", "102"]}
+ mock_get_records.assert_not_called()
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+def test_accession_ids_tabular_snapshot_returns_empty_list_not_an_error(
+ mock_get_snapshot, mock_decrypt, mock_get_settings
+):
+ """A frozen cohort with no accession_id column is a tabular project: imaging no-ops."""
+ mock_decrypt.return_value = "my_project"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ mock_get_snapshot.return_value = _snapshot(pd.DataFrame({"person_id": [1, 2, 3], "label": [0, 1, 0]}))
+
+ response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
+
+ assert response.status_code == 200
+ assert response.json() == {"accession_ids": []}
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.get_snapshot")
+def test_accession_ids_frozen_below_threshold_is_indistinguishable_from_zero(
+ mock_get_snapshot, mock_decrypt, mock_get_settings
+):
+ mock_decrypt.return_value = "my_project"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
+
+ details = []
+ for rows in (0, 9):
+ frozen = pd.DataFrame({"accession_id": [f"A{i}" for i in range(rows)]})
+ mock_get_snapshot.return_value = _snapshot(frozen)
+ response = client.post("/cohort/accession-ids", json=sample_dataframe_query, headers=AUTH_HEADERS)
+ assert response.status_code == 403
+ details.append(response.json()["detail"])
+
+ assert details[0] == details[1] == _BELOW_THRESHOLD_DETAIL
+
+
+# ---------------------------------------------------------------------------
+# POST /cohort/snapshot (creation) and /cohort/snapshot/delete
+# ---------------------------------------------------------------------------
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.save_snapshot")
+@patch("data_access_api.routers.cohort.get_records")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.snapshot_enabled")
+def test_create_snapshot_freezes_the_validated_query_result(
+ mock_snapshot_enabled, mock_decrypt, mock_validate_query, mock_get_records, mock_save_snapshot, mock_get_settings
+):
+ mock_snapshot_enabled.return_value = True
+ mock_decrypt.return_value = "8b2e9d6e-5a53-4f2e-9c37-2c8f4f0f2d11"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ frozen = pd.DataFrame({"accession_id": ["A1", "A2", "A3"]})
+ mock_get_records.return_value = frozen
+ mock_save_snapshot.return_value = SnapshotMeta(
+ row_count=3,
+ columns=["accession_id"],
+ query_hash=normalised_query_hash(FROZEN_QUERY),
+ created_at="2026-08-26T00:00:00+00:00",
+ )
+
+ response = client.post("/cohort/snapshot", json=sample_dataframe_query, headers=WRITE_AUTH_HEADERS)
+
+ assert response.status_code == 200
+ payload = response.json()
+ assert payload["row_count"] == 3
+ assert payload["has_accessions"] is True
+ assert payload["query_hash"] == normalised_query_hash(FROZEN_QUERY)
+ # The frame is what validate_query's emission produced, read FRESH (use_cache=False —
+ # a re-approval must not re-freeze the stale frame the statistics run cached);
+ # the hash is of the RAW query.
+ mock_get_records.assert_called_once_with(mock_validate_query.return_value, use_cache=False)
+ mock_save_snapshot.assert_called_once()
+ assert mock_save_snapshot.call_args.kwargs["query_hash"] == normalised_query_hash(FROZEN_QUERY)
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.save_snapshot")
+@patch("data_access_api.routers.cohort.get_records")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.snapshot_enabled")
+def test_create_snapshot_below_threshold_persists_nothing(
+ mock_snapshot_enabled, mock_decrypt, mock_validate_query, mock_get_records, mock_save_snapshot, mock_get_settings
+):
+ mock_snapshot_enabled.return_value = True
+ mock_decrypt.return_value = "8b2e9d6e-5a53-4f2e-9c37-2c8f4f0f2d11"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 10
+ mock_get_records.return_value = pd.DataFrame({"accession_id": ["A1"]})
+
+ response = client.post("/cohort/snapshot", json=sample_dataframe_query, headers=WRITE_AUTH_HEADERS)
+
+ assert response.status_code == 403
+ assert response.json()["detail"] == _BELOW_THRESHOLD_DETAIL
+ mock_save_snapshot.assert_not_called()
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.save_snapshot")
+@patch("data_access_api.routers.cohort.get_records")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.snapshot_enabled")
+def test_create_snapshot_oversize_returns_413_without_detail_leakage(
+ mock_snapshot_enabled, mock_decrypt, mock_validate_query, mock_get_records, mock_save_snapshot, mock_get_settings
+):
+ mock_snapshot_enabled.return_value = True
+ mock_decrypt.return_value = "8b2e9d6e-5a53-4f2e-9c37-2c8f4f0f2d11"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ mock_get_records.return_value = pd.DataFrame({"accession_id": ["A1", "A2", "A3"]})
+ mock_save_snapshot.side_effect = SnapshotTooLarge("snapshot is 999 bytes, over the 10-byte limit")
+
+ response = client.post("/cohort/snapshot", json=sample_dataframe_query, headers=WRITE_AUTH_HEADERS)
+
+ assert response.status_code == 413
+ assert "byte" not in response.json()["detail"] # category-only, no internals
+
+
+@patch("data_access_api.routers.cohort.snapshot_enabled")
+def test_create_snapshot_store_disabled_returns_503(mock_snapshot_enabled):
+ mock_snapshot_enabled.return_value = False
+ response = client.post("/cohort/snapshot", json=sample_dataframe_query, headers=WRITE_AUTH_HEADERS)
+ assert response.status_code == 503
+
+
+@patch("data_access_api.routers.cohort.get_settings")
+@patch("data_access_api.routers.cohort.save_snapshot")
+@patch("data_access_api.routers.cohort.get_records")
+@patch("data_access_api.routers.cohort.validate_query")
+@patch("data_access_api.routers.cohort.decrypt")
+@patch("data_access_api.routers.cohort.snapshot_enabled")
+def test_create_snapshot_non_uuid_project_id_returns_400(
+ mock_snapshot_enabled, mock_decrypt, mock_validate_query, mock_get_records, mock_save_snapshot, mock_get_settings
+):
+ mock_snapshot_enabled.return_value = True
+ mock_decrypt.return_value = "not-a-uuid"
+ mock_get_settings.return_value.COHORT_QUERY_THRESHOLD = 2
+ mock_get_records.return_value = pd.DataFrame({"accession_id": ["A1", "A2", "A3"]})
+ mock_save_snapshot.side_effect = ValueError("project_id must be a UUID")
+
+ response = client.post("/cohort/snapshot", json=sample_dataframe_query, headers=WRITE_AUTH_HEADERS)
+
+ assert response.status_code == 400
+
+
+@patch("data_access_api.routers.cohort.delete_snapshot")
+@patch("data_access_api.routers.cohort.decrypt")
+def test_delete_snapshot_route_is_idempotent(mock_decrypt, mock_delete_snapshot):
+ mock_decrypt.return_value = "8b2e9d6e-5a53-4f2e-9c37-2c8f4f0f2d11"
+ mock_delete_snapshot.side_effect = [True, False]
+
+ first = client.post("/cohort/snapshot/delete", json={"encrypted_project_id": "enc"}, headers=WRITE_AUTH_HEADERS)
+ second = client.post("/cohort/snapshot/delete", json={"encrypted_project_id": "enc"}, headers=WRITE_AUTH_HEADERS)
+
+ assert first.json() == {"deleted": True}
+ assert second.json() == {"deleted": False}
+
+
+# (Auth coverage for the snapshot routes lives in test_cohort.py: the parametrised
+# missing-key / wrong-key tests cover the trust-internal gate alongside every other
+# /cohort route, and the cohort-admin tests cover the extra AES-possession gate the
+# WRITE routes carry — a valid trust-internal key without the proof is refused 403.)
diff --git a/trust/data-access-api/tests/services/test_cohort.py b/trust/data-access-api/tests/services/test_cohort.py
index fe83147de..66cb78940 100644
--- a/trust/data-access-api/tests/services/test_cohort.py
+++ b/trust/data-access-api/tests/services/test_cohort.py
@@ -171,6 +171,47 @@ def test_get_statistics_genuine_zero_is_suppressed(mock_read_sql):
assert stats.suppressed is True
+@patch("pandas.read_sql")
+def test_get_records_use_cache_false_bypasses_a_poisoned_cache(mock_read_sql):
+ """Snapshot creation must freeze LIVE OMOP: a cached frame (e.g. from the statistics
+ run at submission) is skipped when use_cache=False, and the fresh read replaces it."""
+ import pandas as pd
+
+ from data_access_api.services.cohort import get_records
+ from data_access_api.services.query_cache import get_cached_result, set_cached_result
+
+ query = "SELECT person_id FROM omop.person"
+ stale = pd.DataFrame({"person_id": [1, 2]})
+ fresh = pd.DataFrame({"person_id": [1]}) # a person was removed from OMOP since
+ set_cached_result(query, stale)
+ mock_read_sql.return_value = fresh
+
+ result = get_records(query, use_cache=False)
+
+ pd.testing.assert_frame_equal(result, fresh)
+ mock_read_sql.assert_called_once()
+ # The fresh read updates the cache, so subsequent cached reads see the new state too.
+ pd.testing.assert_frame_equal(get_cached_result(query), fresh)
+
+
+@patch("pandas.read_sql")
+def test_get_records_serves_cache_by_default(mock_read_sql):
+ """Default behaviour is unchanged: a cached frame short-circuits the database read."""
+ import pandas as pd
+
+ from data_access_api.services.cohort import get_records
+ from data_access_api.services.query_cache import set_cached_result
+
+ query = "SELECT person_id FROM omop.person"
+ cached = pd.DataFrame({"person_id": [1, 2]})
+ set_cached_result(query, cached)
+
+ result = get_records(query)
+
+ pd.testing.assert_frame_equal(result, cached)
+ mock_read_sql.assert_not_called()
+
+
@patch("pandas.read_sql")
def test_get_records_undefined_table_error(mock_read_sql):
"""
diff --git a/trust/data-access-api/tests/services/test_cohort_snapshot.py b/trust/data-access-api/tests/services/test_cohort_snapshot.py
new file mode 100644
index 000000000..fc3899a96
--- /dev/null
+++ b/trust/data-access-api/tests/services/test_cohort_snapshot.py
@@ -0,0 +1,174 @@
+# Copyright (c) Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Unit tests for the approved-cohort snapshot file store (FLIP#857)."""
+
+import json
+import uuid
+from unittest.mock import patch
+
+import pandas as pd
+import pytest
+
+from data_access_api.services import cohort_snapshot
+from data_access_api.services.cohort_snapshot import (
+ SnapshotStoreDisabled,
+ SnapshotTooLarge,
+ delete_snapshot,
+ ensure_store,
+ get_snapshot,
+ normalised_query_hash,
+ save_snapshot,
+ snapshot_enabled,
+)
+
+PROJECT_ID = "8b2e9d6e-5a53-4f2e-9c37-2c8f4f0f2d11"
+QUERY_HASH = normalised_query_hash("SELECT * FROM omop.person")
+
+
+@pytest.fixture
+def store(tmp_path):
+ """A configured, writable snapshot store rooted in a per-test temp directory."""
+ with patch("data_access_api.services.cohort_snapshot.get_settings") as mock_settings:
+ mock_settings.return_value.COHORT_SNAPSHOT_DIR = str(tmp_path)
+ mock_settings.return_value.SNAPSHOT_MAX_BYTES = 536_870_912
+ yield tmp_path
+
+
+def _sample_df() -> pd.DataFrame:
+ return pd.DataFrame(
+ {
+ "accession_id": ["A1", "A2", "A3"],
+ "age": pd.array([40, None, 61], dtype="Int64"),
+ "measured_at": pd.to_datetime(["2021-01-01", "2022-06-30", "2023-12-31"]),
+ }
+ )
+
+
+def test_save_then_get_round_trips_the_frame_dtype_faithfully(store):
+ df = _sample_df()
+ meta = save_snapshot(PROJECT_ID, df, query_hash=QUERY_HASH)
+
+ snapshot = get_snapshot(PROJECT_ID)
+ assert snapshot is not None
+ pd.testing.assert_frame_equal(snapshot.df, df)
+ assert snapshot.meta.row_count == 3
+ assert snapshot.meta.columns == ["accession_id", "age", "measured_at"]
+ assert snapshot.meta.query_hash == QUERY_HASH
+ assert snapshot.meta.has_accessions is True
+ assert meta.row_count == 3
+
+
+def test_save_overwrites_atomically_on_reapproval(store):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ replacement = pd.DataFrame({"person_id": [1, 2]})
+ save_snapshot(PROJECT_ID, replacement, query_hash=normalised_query_hash("SELECT person_id FROM omop.person"))
+
+ snapshot = get_snapshot(PROJECT_ID)
+ assert snapshot is not None
+ pd.testing.assert_frame_equal(snapshot.df, replacement)
+ assert snapshot.meta.has_accessions is False
+ # No write debris left behind after the swap.
+ leftovers = [p.name for p in store.iterdir() if p.name.startswith((".tmp-", ".old-"))]
+ assert leftovers == []
+
+
+def test_get_returns_none_for_unknown_project_and_non_uuid_ids(store):
+ assert get_snapshot(str(uuid.uuid4())) is None
+ # A non-UUID id must never touch the filesystem — it is a path component.
+ assert get_snapshot("../../etc/passwd") is None
+ assert get_snapshot("my_project") is None
+
+
+def test_save_rejects_non_uuid_project_id(store):
+ with pytest.raises(ValueError, match="UUID"):
+ save_snapshot("../escape", _sample_df(), query_hash=QUERY_HASH)
+ assert list(store.iterdir()) == []
+
+
+def test_save_refuses_oversized_snapshot_without_truncating(store):
+ with patch("data_access_api.services.cohort_snapshot.get_settings") as mock_settings:
+ mock_settings.return_value.COHORT_SNAPSHOT_DIR = str(store)
+ mock_settings.return_value.SNAPSHOT_MAX_BYTES = 10
+ with pytest.raises(SnapshotTooLarge, match="SNAPSHOT_MAX_BYTES"):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ # Refused means nothing persisted — no partial artefact to poison training data.
+ assert get_snapshot(PROJECT_ID) is None
+
+
+def test_disabled_store_reads_none_and_refuses_writes():
+ with patch("data_access_api.services.cohort_snapshot.get_settings") as mock_settings:
+ mock_settings.return_value.COHORT_SNAPSHOT_DIR = ""
+ assert snapshot_enabled() is False
+ assert get_snapshot(PROJECT_ID) is None
+ assert delete_snapshot(PROJECT_ID) is False
+ with pytest.raises(SnapshotStoreDisabled):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+
+
+def test_corrupt_meta_is_treated_as_absent(store):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ (store / PROJECT_ID / "meta.json").write_text("{not json")
+ assert get_snapshot(PROJECT_ID) is None
+
+
+def test_unknown_format_version_is_treated_as_absent(store):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ meta_path = store / PROJECT_ID / "meta.json"
+ meta = json.loads(meta_path.read_text())
+ meta["format_version"] = 999
+ meta_path.write_text(json.dumps(meta))
+ assert get_snapshot(PROJECT_ID) is None
+
+
+def test_delete_is_idempotent(store):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ assert delete_snapshot(PROJECT_ID) is True
+ assert get_snapshot(PROJECT_ID) is None
+ assert delete_snapshot(PROJECT_ID) is False
+
+
+def test_ensure_store_sweeps_stale_write_debris(store):
+ save_snapshot(PROJECT_ID, _sample_df(), query_hash=QUERY_HASH)
+ (store / ".tmp-crashed-write").mkdir()
+ (store / ".old-crashed-swap").mkdir()
+
+ ensure_store()
+
+ survivors = sorted(p.name for p in store.iterdir())
+ assert survivors == [PROJECT_ID]
+ assert get_snapshot(PROJECT_ID) is not None
+
+
+def test_ensure_store_survives_unwritable_directory(tmp_path):
+ target = tmp_path / "readonly"
+ target.mkdir()
+ target.chmod(0o500)
+ try:
+ with patch("data_access_api.services.cohort_snapshot.get_settings") as mock_settings:
+ mock_settings.return_value.COHORT_SNAPSHOT_DIR = str(target)
+ # Must log and return, never raise: a broken store cannot take OMOP serving down.
+ ensure_store()
+ finally:
+ target.chmod(0o700)
+
+
+def test_normalised_query_hash_ignores_case_and_whitespace_only():
+ base = normalised_query_hash("SELECT * FROM omop.person")
+ assert normalised_query_hash(" select *\n FROM omop.person ") == base
+ assert normalised_query_hash("SELECT person_id FROM omop.person") != base
+
+
+def test_hash_key_matches_module_constant_shape():
+ # cohort_snapshot deliberately does not import query_cache: pin that its normalisation
+ # stays self-contained and deterministic.
+ assert len(cohort_snapshot.normalised_query_hash("x")) == 64
diff --git a/trust/data-access-api/uv.lock b/trust/data-access-api/uv.lock
index 106973c06..3067cc170 100644
--- a/trust/data-access-api/uv.lock
+++ b/trust/data-access-api/uv.lock
@@ -282,6 +282,7 @@ dependencies = [
{ name = "idna" },
{ name = "pandas" },
{ name = "psycopg2-binary" },
+ { name = "pyarrow" },
{ name = "pydantic-settings" },
{ name = "python-multipart" },
{ name = "sqlalchemy" },
@@ -307,6 +308,7 @@ requires-dist = [
{ name = "idna", specifier = ">=3.15" },
{ name = "pandas", specifier = ">=2.2.3" },
{ name = "psycopg2-binary", specifier = ">=2.9.10" },
+ { name = "pyarrow", specifier = ">=19.0.0" },
{ name = "pydantic-settings", specifier = ">=2.9.1" },
{ name = "python-multipart", specifier = ">=0.0.26" },
{ name = "sqlalchemy", specifier = ">=2.0.39" },
@@ -877,6 +879,28 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/42/ee/dee8dcaad07f735824de3d6563bc67119fa6c28257b17977a8d624f02fab/psycopg2_binary-2.9.12-cp313-cp313-win_amd64.whl", hash = "sha256:b6937f5fe4e180aeee87de907a2fa982ded6f7f15d7218f78a083e4e1d68f2a0", size = 2757347, upload-time = "2026-04-20T23:35:21.283Z" },
]
+[[package]]
+name = "pyarrow"
+version = "25.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3d/e3/27f57f80141379d60defe6703eb50a707325706f07fedfd1312c7a751995/pyarrow-25.0.1.tar.gz", hash = "sha256:9150a83248bfed9813ea3c3af74c3856c1984d444aa28e58bf7733b9750ddf6a", size = 1201653, upload-time = "2026-08-10T12:40:53.904Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/a6/e2/9ab15b88cbfac28e16419ce5439ec29234c5172cb8259301b4ba639bdec0/pyarrow-25.0.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:df961f2e7ae9cf496459259d798652c70625f6c080650d6952f8c04053c58ee9", size = 35861559, upload-time = "2026-08-10T12:38:02.567Z" },
+ { url = "https://files.pythonhosted.org/packages/58/79/a0036dbe1eabe1f73127427342f1d99982584c4a2cde2651d6c93499c6f6/pyarrow-25.0.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:cc4aa407fde9fc660be3939e49ea31f50f3e9fec17c0ec63159f7711edd3efc9", size = 37628383, upload-time = "2026-08-10T12:38:09.083Z" },
+ { url = "https://files.pythonhosted.org/packages/13/49/d93a57d375f4bf0cf82913dd6bb54acafde83dd993be2282c81ac5616cad/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:4340f0ba6c1d2e13f21658de1d7c662ca2545018568d0030a1e9afca159d87e3", size = 46820190, upload-time = "2026-08-10T12:38:15.458Z" },
+ { url = "https://files.pythonhosted.org/packages/60/c9/711ca85d79f1ec98f29a5eae2b051e25b4ecec5de3e3c0e2d5c5dcb15664/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5389cdf79447ed1515c9e31620e6e1e2302249564d603f2ad727d4f6d313e4c3", size = 50102437, upload-time = "2026-08-10T12:38:22.487Z" },
+ { url = "https://files.pythonhosted.org/packages/80/53/8fb8359ff17cfb6263a1cf3ebf7caec9fe197de118719e84fcb1d0618026/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d51592cb7561e87877c506113e7adbf1342ab579e6c21f0ef44b8ba41cb74c80", size = 49942424, upload-time = "2026-08-10T12:38:28.755Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/83/4e5ae02a9341571b18a6fca380ac7a58ce6ddae7ab3c060208c0a1e79f02/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6109c94d8b9f3b17a041daca16cacb2f651ad8f1ef70a4232c2c0f37a23da2a8", size = 53144206, upload-time = "2026-08-10T12:38:34.862Z" },
+ { url = "https://files.pythonhosted.org/packages/65/ee/197cbf47e49f83e6ebeb946a5259a48a638dea27ac774db42fe78022179d/pyarrow-25.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:8858d7bfc22e3f51529aeaa4077225029724623e4595dc9eff8c793935c34140", size = 27953934, upload-time = "2026-08-10T12:38:39.808Z" },
+ { url = "https://files.pythonhosted.org/packages/cc/8d/8f271a7a034c834910ec925d56fa4b29733b1380f5289419f5aaa3b02777/pyarrow-25.0.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:c7c534ec03c358a76ea3e505e74c1b6aef290af90c444dfd092dbfe23e755b85", size = 35855328, upload-time = "2026-08-10T12:38:45.489Z" },
+ { url = "https://files.pythonhosted.org/packages/d2/cd/5bac242f4e841b9971d5eb94fdfe2577e2b70be983e27401e72055786037/pyarrow-25.0.1-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:dda9470024204d7bbf2042b47c6e8a0e47a3eeb8e34405882dfaea6577e0c153", size = 37622415, upload-time = "2026-08-10T12:38:51.107Z" },
+ { url = "https://files.pythonhosted.org/packages/63/1f/96d03b4e1506524f7087adb0fd6b2f69f0c9c7aaff1ec36d8030082e15a5/pyarrow-25.0.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:44a9120ce5bd81936b8ab9a88076e3fd47c2c6838e0e43630fed83626aca81d9", size = 46813813, upload-time = "2026-08-10T12:38:57.773Z" },
+ { url = "https://files.pythonhosted.org/packages/98/d6/33a411115b61dbfc16ad6ad73e71730f6fea654ee3667673bc53ab0e2fe7/pyarrow-25.0.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:0befcf816e45a1af33ac775a9970b749e4868a230c7372f0ae5e932bee27039f", size = 50104452, upload-time = "2026-08-10T12:39:04.579Z" },
+ { url = "https://files.pythonhosted.org/packages/33/ae/b1b97c9ca87f9f9ddbb5230c798df94eccce61bd79b9b45458c69a478588/pyarrow-25.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3f89685964f46e4216103c75483aac0c0692a5f72212d7ca835adba5ede56ce3", size = 49951343, upload-time = "2026-08-10T12:39:11.8Z" },
+ { url = "https://files.pythonhosted.org/packages/98/9e/a112df5cfd5a68cb1d9fc31cfe38c28d5aec9f10865ce37ecef2e4450873/pyarrow-25.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6943e2fe7954d29d84de45d29d34c8dc36ce96570e67d89aa9976e650a4a9138", size = 53144784, upload-time = "2026-08-10T12:39:20.503Z" },
+ { url = "https://files.pythonhosted.org/packages/31/24/97e8bd98f1e3b07e2ba08bcdff690674fbe16d69a7d2712cc3884665e615/pyarrow-25.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:31e49a7888fcdf3a835da33ae777f6bb9a866334e5a789282fc26dcf426f7f15", size = 27870159, upload-time = "2026-08-10T12:39:26.161Z" },
+]
+
[[package]]
name = "pycparser"
version = "3.0"
diff --git a/trust/deploy/compose.test.yml b/trust/deploy/compose.test.yml
index a5dec7f1e..18ad78f31 100644
--- a/trust/deploy/compose.test.yml
+++ b/trust/deploy/compose.test.yml
@@ -82,6 +82,9 @@ services:
# sex/age distribution assertions test the real per-bucket counts rather
# than always rolling into the "Other" bucket.
COHORT_QUERY_THRESHOLD: "5"
+ # Approved-cohort snapshot store (FLIP#857). /tmp is writable by the container
+ # user and ephemeral per run — each pytest session gets a clean store.
+ COHORT_SNAPSHOT_DIR: /tmp/snapshots
ports:
- "8000"
volumes:
diff --git a/trust/deploy/compose_trust.development.yml b/trust/deploy/compose_trust.development.yml
index dac7ecaa9..887fe3392 100644
--- a/trust/deploy/compose_trust.development.yml
+++ b/trust/deploy/compose_trust.development.yml
@@ -175,6 +175,9 @@ services:
# kit file to release less. The `:-10` keeps an unset variable from arriving as an
# empty string, which pydantic would reject against `int`.
COHORT_QUERY_THRESHOLD: ${COHORT_QUERY_THRESHOLD:-10}
+ # Approved-cohort snapshot store (FLIP#857): the container-side path of the bind
+ # mount below. Fixed here; the HOST side is the kit's COHORT_SNAPSHOT_STORAGE_DIR.
+ COHORT_SNAPSHOT_DIR: /snapshots
# Dev-on-pull (see `user:` above): import via PYTHONPATH, skip the editable
# install (UV_NO_SYNC), keep uv's cache + HOME writable.
HOME: /tmp
@@ -185,6 +188,11 @@ services:
- ./data-access-api/data_access_api:/app/data_access_api
- ./data-access-api/tests:/app/tests
- ./observability/log_config:/app/log_config:ro
+ # Approved-cohort snapshots: per-trust host dir from the kit's Host-local profile
+ # (pre-created by `make up-trust` so it is host-user-owned, not docker-root-owned —
+ # the service runs non-root and disables the store if it cannot write). Holds
+ # row-level patient data; gitignored, never committed.
+ - ${COHORT_SNAPSHOT_STORAGE_DIR:-./data-access-api/.snapshots/Trust_${TRUST_NUMBER}}:/snapshots
trust-api:
# Pull from GHCR by default; local trust_api/ bind-mount overlays it for live
diff --git a/trust/deploy/compose_trust.production.yml b/trust/deploy/compose_trust.production.yml
index 1a62ac6a8..2c32577a6 100644
--- a/trust/deploy/compose_trust.production.yml
+++ b/trust/deploy/compose_trust.production.yml
@@ -139,6 +139,15 @@ services:
# kit file to release less. The `:-10` keeps an unset variable from arriving as an
# empty string, which pydantic would reject against `int`.
COHORT_QUERY_THRESHOLD: ${COHORT_QUERY_THRESHOLD:-10}
+ # Approved-cohort snapshot store (FLIP#857): the container-side path of the bind
+ # mount below. Fixed here; the HOST side is the kit's COHORT_SNAPSHOT_STORAGE_DIR.
+ COHORT_SNAPSHOT_DIR: /snapshots
+ volumes:
+ # Approved-cohort snapshots: host dir from the kit's Host-local profile, pre-created
+ # host-user-owned by `make up-trust` (the service runs non-root and disables the
+ # store if it cannot write). Holds row-level patient data — include it in the
+ # trust's backup/retention practice alongside the OMOP data dir.
+ - ${COHORT_SNAPSHOT_STORAGE_DIR:-./data-access-api/.snapshots/trust}:/snapshots
trust-api:
image: ghcr.io/londonaicentre/trust-api:${DOCKER_TAG}
diff --git a/trust/imaging-api/imaging_api/services/retrieval.py b/trust/imaging-api/imaging_api/services/retrieval.py
index 586178d76..96ad483d2 100644
--- a/trust/imaging-api/imaging_api/services/retrieval.py
+++ b/trust/imaging-api/imaging_api/services/retrieval.py
@@ -79,16 +79,17 @@ async def retrieve_images_for_project(project_id: str, query: str, headers: XNAT
"""
# Check if project exists
try:
- get_project(project_id, headers)
+ project = get_project(project_id, headers)
except NotFoundError as e:
raise HTTPException(status_code=404, detail=str(e))
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
- # Get accession IDs from data access API. The endpoint projects the cohort
- # query to the accession_id column server-side, so no other columns are
- # transmitted across the trust boundary.
- encrypted_project_id = encrypt(project_id)
+ # Get accession IDs from data access API — the frozen approved-cohort pointer set
+ # (FLIP#857), keyed on the HUB project id. XNAT projects mint their own uuid and stash
+ # the hub id in secondary_ID, so that is what data-access-api's snapshot store is
+ # keyed on; sending the XNAT id would never match a snapshot.
+ encrypted_project_id = encrypt(project.secondary_ID)
try:
accession_ids: list[str] = await get_accession_ids(encrypted_project_id, query)
except CohortBelowThresholdError:
@@ -185,17 +186,21 @@ async def get_import_status(project_id: str, query: str, headers: XNATAuthHeader
ImportStatus: An object containing the status of study imports.
Raises:
- HTTPException: 403 if the cohort has fallen below the trust's ``COHORT_QUERY_THRESHOLD``
- so its accession IDs cannot be released — note the cohort query is re-run against
- live OMOP on every status check, so a project that imported cleanly can start
- refusing later if its cohort shrinks (see FLIP#857). Otherwise, if the request
- cannot be processed.
+ HTTPException: 403 if the frozen cohort is below the trust's ``COHORT_QUERY_THRESHOLD``
+ (or the project has no approved-cohort snapshot) so its accession IDs cannot be
+ released. Otherwise, if the request cannot be processed.
"""
- # Encrypt project ID to send to the data access API
- encrypted_project_id = encrypt(project_id)
+ # Resolve the HUB project id: data-access-api serves the frozen approved-cohort
+ # snapshot (FLIP#857) keyed on it, and the XNAT project stores it as secondary_ID.
+ try:
+ project = get_project(project_id, headers)
+ except NotFoundError as e:
+ raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+ encrypted_project_id = encrypt(project.secondary_ID)
- # Get accession IDs from data access API (server-side projection — no other
- # cohort columns leave the trust).
+ # Get accession IDs from data access API — the frozen pointer set; the query travels
+ # only as an advisory field (the snapshot serve ignores it), so the status poll no
+ # longer re-runs cohort SQL against live OMOP.
try:
accession_ids: list[str] = await get_accession_ids(encrypted_project_id, query)
except CohortBelowThresholdError as e:
diff --git a/trust/imaging-api/imaging_api/services_external/data_access.py b/trust/imaging-api/imaging_api/services_external/data_access.py
index 1607070dd..f0941433c 100644
--- a/trust/imaging-api/imaging_api/services_external/data_access.py
+++ b/trust/imaging-api/imaging_api/services_external/data_access.py
@@ -72,13 +72,18 @@ async def get_accession_ids(encrypted_project_id: str, query: str) -> list[str]:
except httpx.HTTPStatusError as exc:
if exc.response.status_code == httpx.codes.FORBIDDEN:
- # A deliberate refusal, not a failure: the cohort is below the trust's disclosure
- # threshold. Typed separately so callers can report it as a settled outcome rather
- # than as a transport error they might retry.
- message = (
- "get_accession_ids: the Data Access API refused to release accession IDs — "
- "the cohort is below the trust's minimum size."
- )
+ # A deliberate refusal, not a failure — typed separately so callers can report it
+ # as a settled outcome rather than as a transport error they might retry. Two
+ # policy refusals arrive as 403: the frozen cohort is below the trust's disclosure
+ # threshold, or the project has no approved-cohort snapshot yet (FLIP#857 —
+ # e.g. imaging creation raced ahead of the snapshot task). Both mean "no
+ # identifiers releasable right now"; relay data-access-api's own detail so the
+ # two stay distinguishable in status reporting.
+ try:
+ detail = exc.response.json().get("detail", "")
+ except ValueError:
+ detail = ""
+ message = f"get_accession_ids: the Data Access API refused to release accession IDs — {detail}"
logger.warning(message)
raise CohortBelowThresholdError(message) from exc
error_message = f"get_accession_ids: HTTP error occurred while calling the Data Access API: {exc}"
diff --git a/trust/imaging-api/tests/services/test_retrieval.py b/trust/imaging-api/tests/services/test_retrieval.py
index 75c8dc9f4..cbf220637 100644
--- a/trust/imaging-api/tests/services/test_retrieval.py
+++ b/trust/imaging-api/tests/services/test_retrieval.py
@@ -133,10 +133,13 @@ async def test_retrieve_images_below_threshold_queues_nothing(
mock_queue.assert_not_called()
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
-async def test_get_import_status_below_threshold_raises_403(mock_encrypt, mock_get_accession_ids, headers):
+async def test_get_import_status_below_threshold_raises_403(
+ mock_encrypt, mock_get_accession_ids, mock_get_project, headers
+):
"""The status path has a caller waiting, so the refusal must reach it as a 403 with a
readable reason — trust-api relays this detail to the hub for the per-trust status."""
mock_encrypt.return_value = "encrypted_id"
@@ -256,12 +259,13 @@ async def fake_session():
)
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_all_successful(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
mock_encrypt.return_value = "encrypted_id"
mock_get_accession_ids.return_value = ["ACC1", "ACC2"]
@@ -282,12 +286,13 @@ async def test_get_import_status_all_successful(
assert status.queue_failed == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_mixed(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
mock_encrypt.return_value = "encrypted_id"
mock_get_accession_ids.return_value = ["ACC_OK", "ACC_EXEC", "ACC_QUEUED", "ACC_UNKNOWN"]
@@ -313,12 +318,13 @@ async def test_get_import_status_mixed(
assert status.queue_failed == ["ACC_UNKNOWN"]
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_no_experiments(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
mock_encrypt.return_value = "encrypted_id"
mock_get_accession_ids.return_value = ["ACC1"]
@@ -345,12 +351,13 @@ def _direct_archive(accession_number: str, status: str) -> DirectArchiveSession:
)
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_executed_failed_is_failed(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A PACS retrieval that errored (executed status FAILED) is `failed`, not `processing`."""
mock_encrypt.return_value = "encrypted_id"
@@ -366,12 +373,13 @@ async def test_get_import_status_executed_failed_is_failed(
assert status.queue_failed == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_direct_archive_error_is_failed(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A directArchive build that errored (status ERROR) is `failed`, not `queue_failed`/`processing`."""
mock_encrypt.return_value = "encrypted_id"
@@ -389,12 +397,13 @@ async def test_get_import_status_direct_archive_error_is_failed(
assert status.queue_failed == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_direct_archive_error_overrides_received(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A study whose transfer RECEIVED but whose directArchive build then errored is `failed`.
@@ -415,12 +424,13 @@ async def test_get_import_status_direct_archive_error_overrides_received(
assert status.processing == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_direct_archive_receiving_is_processing(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A directArchive build still in progress (RECEIVING) is `processing`, not `failed`/`queue_failed`."""
mock_encrypt.return_value = "encrypted_id"
@@ -438,12 +448,13 @@ async def test_get_import_status_direct_archive_receiving_is_processing(
assert status.queue_failed == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_direct_archive_error_without_name_falls_back_to_folder_name(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A directArchive ERROR whose `name` is NULL is still attributed via its folder_name.
@@ -467,12 +478,13 @@ async def test_get_import_status_direct_archive_error_without_name_falls_back_to
assert status.queue_failed == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_failed_precedes_queued(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""A terminal-failed accession that also has a (re-)queued row is reported `failed`, not `queued`."""
mock_encrypt.return_value = "encrypted_id"
@@ -490,12 +502,13 @@ async def test_get_import_status_failed_precedes_queued(
assert status.queued == []
+@patch("imaging_api.services.retrieval.get_project")
@pytest.mark.asyncio
@patch("imaging_api.services.retrieval.get_experiments")
@patch("imaging_api.services.retrieval.get_accession_ids", new_callable=AsyncMock)
@patch("imaging_api.services.retrieval.encrypt")
async def test_get_import_status_experiment_present_beats_stale_direct_archive_error(
- mock_encrypt, mock_get_accession_ids, mock_get_experiments, headers,
+ mock_encrypt, mock_get_accession_ids, mock_get_experiments, mock_get_project, headers,
):
"""`successful` (archived as an experiment) wins over a stale directArchive ERROR for the same accession."""
mock_encrypt.return_value = "encrypted_id"
diff --git a/trust/imaging-api/tests/services_external/test_data_access.py b/trust/imaging-api/tests/services_external/test_data_access.py
index 26737899c..765c1cb7e 100644
--- a/trust/imaging-api/tests/services_external/test_data_access.py
+++ b/trust/imaging-api/tests/services_external/test_data_access.py
@@ -80,10 +80,10 @@ async def test_http_error_raises_runtime_error(self, mock_client_cls):
await get_accession_ids("encrypted-proj-id", "SELECT * FROM cohort")
@staticmethod
- def _client_returning_status(mock_client_cls, status_code: int):
+ def _client_returning_status(mock_client_cls, status_code: int, json_body: dict | None = None):
"""Wires the mocked client so ``raise_for_status`` raises for ``status_code``."""
request = httpx.Request("POST", "http://data-access-api/cohort/accession-ids")
- response = httpx.Response(status_code, request=request)
+ response = httpx.Response(status_code, request=request, json=json_body)
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock(
@@ -98,15 +98,28 @@ def _client_returning_status(mock_client_cls, status_code: int):
@pytest.mark.asyncio
@patch("imaging_api.services_external.data_access.httpx.AsyncClient")
- async def test_403_raises_cohort_below_threshold(self, mock_client_cls):
- """A 403 is the trust refusing to release identifiers for a too-small cohort.
+ async def test_403_raises_cohort_below_threshold_with_relayed_detail(self, mock_client_cls):
+ """A 403 is the trust refusing to release identifiers — a policy decision, not a failure.
It must be typed distinctly from a transport failure: callers report it as a settled
- outcome rather than an error to retry, and retrying cannot change the answer.
+ outcome rather than an error to retry, and retrying cannot change the answer. Two
+ refusals arrive this way (below-threshold cohort, or no approved-cohort snapshot —
+ FLIP#857), so data-access-api's own detail is relayed to keep them distinguishable.
"""
+ self._client_returning_status(
+ mock_client_cls, 403, json_body={"detail": "Cohort is too small for row-level data to be released."}
+ )
+
+ with pytest.raises(CohortBelowThresholdError, match="Cohort is too small"):
+ await get_accession_ids("encrypted-proj-id", "SELECT * FROM cohort")
+
+ @pytest.mark.asyncio
+ @patch("imaging_api.services_external.data_access.httpx.AsyncClient")
+ async def test_403_without_json_body_still_raises_typed_refusal(self, mock_client_cls):
+ """A body-less 403 must not crash the detail relay — the typed refusal still surfaces."""
self._client_returning_status(mock_client_cls, 403)
- with pytest.raises(CohortBelowThresholdError, match="below the trust's minimum size"):
+ with pytest.raises(CohortBelowThresholdError, match="refused to release accession IDs"):
await get_accession_ids("encrypted-proj-id", "SELECT * FROM cohort")
@pytest.mark.asyncio
diff --git a/trust/trust-api/tests/integration/conftest.py b/trust/trust-api/tests/integration/conftest.py
index c10d7465a..235121189 100644
--- a/trust/trust-api/tests/integration/conftest.py
+++ b/trust/trust-api/tests/integration/conftest.py
@@ -71,6 +71,10 @@ def compose_stack() -> Generator[DockerCompose, None, None]:
# Wait for the compose-level healthchecks rather than `wait_for_logs` —
# logs are noisy and fragile across image versions.
wait=True,
+ # Rebuild the from-source service on every session: without it, `docker compose up`
+ # reuses a previously built image, so a dependency change (pyproject/uv.lock) never
+ # reaches the stack and the suite green-lights code that cannot run in a fresh build.
+ build=True,
) as compose:
yield compose
diff --git a/trust/trust-api/tests/services/test_task_handlers.py b/trust/trust-api/tests/services/test_task_handlers.py
index 1a00d785d..da2129235 100644
--- a/trust/trust-api/tests/services/test_task_handlers.py
+++ b/trust/trust-api/tests/services/test_task_handlers.py
@@ -10,12 +10,16 @@
# limitations under the License.
#
+import hashlib
+import json
from unittest.mock import AsyncMock, patch
from uuid import uuid4
import pytest
from trust_api.services.task_handlers import (
+ AES_KEY_BASE64,
+ COHORT_ADMIN_KEY_HEADER,
TASK_HANDLERS,
TRUST_INTERNAL_SERVICE_KEY,
TRUST_INTERNAL_SERVICE_KEY_HEADER,
@@ -23,6 +27,7 @@
handle_create_imaging,
handle_delete_imaging,
handle_get_imaging_status,
+ handle_persist_cohort,
handle_reimport_studies,
handle_update_user_profile,
)
@@ -35,6 +40,16 @@ def _assert_trust_internal_auth_header(call_args) -> None:
assert headers.get(TRUST_INTERNAL_SERVICE_KEY_HEADER) == TRUST_INTERNAL_SERVICE_KEY
+def _assert_cohort_admin_header(call_args) -> None:
+ """A cohort-DEFINING write (snapshot create/delete) must additionally carry the
+ cohort-admin proof: the SHA-256 of AES_KEY_BASE64, never the key itself (FLIP#857)."""
+ headers = call_args.kwargs.get("headers") or {}
+ expected = hashlib.sha256(AES_KEY_BASE64.encode()).hexdigest()
+ assert headers.get(COHORT_ADMIN_KEY_HEADER) == expected
+ # Defence in depth: the raw key must never travel as a header value.
+ assert AES_KEY_BASE64 not in headers.values()
+
+
@pytest.fixture
def mock_make_request():
with patch("trust_api.services.task_handlers.make_request", new_callable=AsyncMock) as mock:
@@ -53,6 +68,7 @@ def test_task_handlers_registry():
"get_imaging_status",
"reimport_studies",
"update_user_profile",
+ "persist_cohort",
}
assert set(TASK_HANDLERS.keys()) == expected_types
@@ -327,3 +343,67 @@ async def test_handle_update_user_profile_error(mock_make_request):
assert result["success"] is False
assert "Service unavailable" in result["error"]
+
+
+# ---- Persist cohort handler (FLIP#857) ----
+
+
+@pytest.mark.asyncio
+async def test_handle_persist_cohort_freezes_via_data_access_and_returns_facts(mock_make_request):
+ """Forwards the snapshot request to data-access-api and returns its response as the result."""
+ snapshot_facts = {
+ "row_count": 24,
+ "columns": ["modality", "accession_id"],
+ "has_accessions": True,
+ "snapshot_at": "2026-08-26T00:00:00+00:00",
+ "query_hash": "abc123",
+ }
+ mock_make_request.return_value = snapshot_facts
+
+ payload = {
+ "project_id": str(uuid4()),
+ "trust_id": str(uuid4()),
+ "encrypted_project_id": "enc123",
+ "query": "SELECT * FROM omop.image_occurrence",
+ "query_id": str(uuid4()),
+ }
+ result = await handle_persist_cohort(payload)
+
+ assert result["success"] is True
+ assert json.loads(result["result"]) == snapshot_facts
+
+ call = mock_make_request.call_args_list[0]
+ assert call.kwargs["method"] == "POST"
+ assert call.kwargs["url"].endswith("/cohort/snapshot")
+ # Only the encrypted id + query go to data-access-api (its DataframeQuery schema).
+ assert call.kwargs["json_body"] == {"encrypted_project_id": "enc123", "query": payload["query"]}
+ _assert_trust_internal_auth_header(call)
+ # Snapshot creation is a cohort-DEFINING write: it must also carry the cohort-admin proof
+ # (FLIP#857), which fl-client cannot produce.
+ _assert_cohort_admin_header(call)
+
+
+@pytest.mark.asyncio
+async def test_handle_persist_cohort_invalid_payload():
+ """A malformed task payload fails validation before any request is made."""
+ result = await handle_persist_cohort({"query": "SELECT 1"})
+
+ assert result["success"] is False
+ assert "validation error" in result["error"].lower()
+
+
+@pytest.mark.asyncio
+async def test_handle_persist_cohort_reports_data_access_failure(mock_make_request):
+ """A refused snapshot (e.g. below-threshold 403) marks the task FAILED at the hub."""
+ mock_make_request.side_effect = Exception("403: Cohort is too small for row-level data to be released.")
+
+ payload = {
+ "project_id": str(uuid4()),
+ "trust_id": str(uuid4()),
+ "encrypted_project_id": "enc123",
+ "query": "SELECT * FROM omop.image_occurrence",
+ }
+ result = await handle_persist_cohort(payload)
+
+ assert result["success"] is False
+ assert "too small" in result["error"]
diff --git a/trust/trust-api/trust_api/config.py b/trust/trust-api/trust_api/config.py
index d457e0995..5804db158 100644
--- a/trust/trust-api/trust_api/config.py
+++ b/trust/trust-api/trust_api/config.py
@@ -58,6 +58,14 @@ def coerce_empty_env(cls, v: str) -> str:
TRUST_INTERNAL_SERVICE_KEY: str = ""
TRUST_INTERNAL_SERVICE_KEY_HEADER: str = "X-Trust-Internal-Service-Key"
+ # Cohort-write authorisation (FLIP#857). data-access-api's snapshot create/delete routes
+ # require proof of possessing AES_KEY_BASE64 in addition to the trust-internal key. trust-api
+ # sends the SHA-256 of the key in this header when it forwards an approval-time snapshot
+ # request; fl-client cannot (it holds no AES key), which is what keeps researcher code from
+ # rewriting or deleting a project's frozen cohort. Header name must match data-access-api's
+ # COHORT_ADMIN_KEY_HEADER.
+ COHORT_ADMIN_KEY_HEADER: str = "X-Cohort-Admin-Key"
+
# Polling configuration
POLL_INTERVAL_SECONDS: int = 5 # How often to poll the hub for tasks (seconds)
diff --git a/trust/trust-api/trust_api/routers/schemas.py b/trust/trust-api/trust_api/routers/schemas.py
index fd3b76f09..a8ef1f609 100644
--- a/trust/trust-api/trust_api/routers/schemas.py
+++ b/trust/trust-api/trust_api/routers/schemas.py
@@ -29,6 +29,20 @@ class CohortQueryInput(BaseModel):
trust_id: str = Field(..., description="The unique identifier for the trust")
+class PersistCohortInput(BaseModel):
+ """Payload of the approval-time PERSIST_COHORT task (FLIP#857).
+
+ Mirrors flip-api's ``IPersistCohort``: the approved cohort query of record plus the hub
+ project id, encrypted for forwarding to data-access-api and in the clear for logging.
+ """
+
+ project_id: str = Field(..., description="The central hub project id (clear, for logging)")
+ trust_id: str = Field(..., description="The unique identifier for the trust")
+ encrypted_project_id: str = Field(..., description="The hub project id, encrypted for data-access-api")
+ query: str = Field(..., description="The approved cohort SQL to freeze")
+ query_id: str | None = Field(default=None, description="The hub Queries row being frozen")
+
+
# #########################
# Imaging
# #########################
diff --git a/trust/trust-api/trust_api/services/task_handlers.py b/trust/trust-api/trust_api/services/task_handlers.py
index b7a558348..09dfafd83 100644
--- a/trust/trust-api/trust_api/services/task_handlers.py
+++ b/trust/trust-api/trust_api/services/task_handlers.py
@@ -17,6 +17,7 @@
"""
import datetime
+import hashlib
import json
from typing import Any
@@ -26,6 +27,7 @@
CohortQueryInput,
DeleteImagingInput,
GetImagingStatusInput,
+ PersistCohortInput,
ReimportStudiesInput,
UpdateProfileRequest,
)
@@ -39,6 +41,8 @@
TRUST_API_KEY_HEADER = get_settings().TRUST_API_KEY_HEADER
TRUST_INTERNAL_SERVICE_KEY = get_settings().TRUST_INTERNAL_SERVICE_KEY
TRUST_INTERNAL_SERVICE_KEY_HEADER = get_settings().TRUST_INTERNAL_SERVICE_KEY_HEADER
+AES_KEY_BASE64 = get_settings().AES_KEY_BASE64
+COHORT_ADMIN_KEY_HEADER = get_settings().COHORT_ADMIN_KEY_HEADER
def trust_internal_headers() -> dict[str, str]:
@@ -54,6 +58,22 @@ def trust_internal_headers() -> dict[str, str]:
return {TRUST_INTERNAL_SERVICE_KEY_HEADER: TRUST_INTERNAL_SERVICE_KEY}
+def cohort_admin_headers() -> dict[str, str]:
+ """Headers for data-access-api's cohort-DEFINING write routes (snapshot create/delete).
+
+ Those routes require the trust-internal key AND proof of possessing ``AES_KEY_BASE64``
+ (FLIP#857) — the second gate is what stops fl-client's researcher code from rewriting or
+ deleting a project's frozen cohort, since fl-client holds no AES key. The proof is the
+ SHA-256 of the key, never the key itself, so it stays off the wire and out of logs. Layers
+ the cohort-admin header on top of the trust-internal one.
+
+ Returns:
+ dict[str, str]: The trust-internal header plus the cohort-admin proof header.
+ """
+ proof = hashlib.sha256(AES_KEY_BASE64.encode()).hexdigest()
+ return {**trust_internal_headers(), COHORT_ADMIN_KEY_HEADER: proof}
+
+
# Task type constants — must match TaskType enum in flip-api/src/flip_api/domain/schemas/status.py
TASK_COHORT_QUERY = "cohort_query"
TASK_CREATE_IMAGING = "create_imaging"
@@ -61,6 +81,55 @@ def trust_internal_headers() -> dict[str, str]:
TASK_GET_IMAGING_STATUS = "get_imaging_status"
TASK_REIMPORT_STUDIES = "reimport_studies"
TASK_UPDATE_USER_PROFILE = "update_user_profile"
+TASK_PERSIST_COHORT = "persist_cohort"
+
+
+async def handle_persist_cohort(payload: dict[str, Any]) -> dict[str, Any]:
+ """
+ Freeze the approved cohort on this trust (FLIP#857).
+
+ Forwards the approval-time snapshot request to the local data-access-api, which runs
+ the approved query ONCE and persists the resulting dataframe as the project's frozen
+ artefact — the only thing the row-level routes serve from then on. The snapshot
+ response (aggregates only: row count, column names, timestamps) is returned as the
+ task result verbatim, so the hub can record its frozen-cohort audit row from it.
+
+ Failure (including a below-threshold cohort, which data-access-api refuses with 403)
+ marks the task FAILED at the hub with the category-only detail — nothing is persisted
+ trust-side in that case, and the project's row-level routes keep refusing.
+
+ Args:
+ payload: Task payload matching ``PersistCohortInput``.
+
+ Returns:
+ dict with success status, the snapshot facts as ``result``, or error details.
+ """
+ logger.info(f"Processing cohort snapshot task: project_id={payload.get('project_id')}")
+
+ try:
+ request = PersistCohortInput(**payload)
+ response = await make_request(
+ method="POST",
+ url=f"{DATA_ACCESS_API_URL}/cohort/snapshot",
+ json_body={
+ "encrypted_project_id": request.encrypted_project_id,
+ "query": request.query,
+ },
+ # Snapshot creation is a cohort-DEFINING write: it needs the cohort-admin proof on
+ # top of the trust-internal key (FLIP#857).
+ headers=cohort_admin_headers(),
+ # Snapshot creation runs the full cohort query, so it inherits the cohort
+ # query's timeout rather than the default request timeout.
+ timeout_seconds=get_settings().COHORT_QUERY_TIMEOUT_SECONDS,
+ )
+ logger.info(
+ f"Cohort snapshot persisted for project {request.project_id}: "
+ f"{response.get('row_count')} rows" # type: ignore[union-attr]
+ )
+ return {"success": True, "result": json.dumps(response)}
+ except Exception as e:
+ logger.error(f"Error persisting cohort snapshot: {e}")
+ return {"success": False, "error": str(e)}
async def handle_cohort_query(payload: dict[str, Any]) -> dict[str, Any]:
@@ -311,4 +380,5 @@ async def handle_update_user_profile(payload: dict[str, Any]) -> dict[str, Any]:
TASK_GET_IMAGING_STATUS: handle_get_imaging_status,
TASK_REIMPORT_STUDIES: handle_reimport_studies,
TASK_UPDATE_USER_PROFILE: handle_update_user_profile,
+ TASK_PERSIST_COHORT: handle_persist_cohort,
}