diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ec37ea2..4963bec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,6 +41,7 @@ jobs: cache: pip - run: python -m pip install -e '.[test]' - run: python scripts/smoke_container.py + - run: python scripts/security_container.py image-security: runs-on: ubuntu-latest diff --git a/README.md b/README.md index ccad0f4..1c16278 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,10 @@ the separate `/tmp` tmpfs can therefore remain small and constrained. ## Verification +See [the coordinated production foundations guide](docs/production-foundations.md) +for compatibility fixtures, leases and quota behavior, token rotation, mTLS, +administrative purge, joint testing and rollback. + ```bash pytest ruff check . diff --git a/compose.yaml b/compose.yaml index df420e7..57c1aee 100644 --- a/compose.yaml +++ b/compose.yaml @@ -1,5 +1,6 @@ services: data-query-worker: + init: true build: context: . image: omero-data-query-worker:smoke diff --git a/deploy/Dockerfile.tests b/deploy/Dockerfile.tests new file mode 100644 index 0000000..af1bf9a --- /dev/null +++ b/deploy/Dockerfile.tests @@ -0,0 +1,8 @@ +ARG WORKER_IMAGE=omero-data-query-worker:smoke +FROM ${WORKER_IMAGE} +USER root +RUN python -m pip install --no-cache-dir 'pytest>=8.4,<10' 'httpx>=0.28,<1' +COPY tests /app/tests +COPY scripts/security_inside.py /app/scripts/security_inside.py +USER 10001:10001 +CMD ["python", "-m", "pytest", "-q", "-p", "no:cacheprovider", "/app/tests"] diff --git a/deploy/Dockerfile.tests.dockerignore b/deploy/Dockerfile.tests.dockerignore new file mode 100644 index 0000000..3e26a18 --- /dev/null +++ b/deploy/Dockerfile.tests.dockerignore @@ -0,0 +1,7 @@ +** +!deploy +!deploy/Dockerfile.tests +!tests +!tests/** +!scripts +!scripts/security_inside.py diff --git a/deploy/compose.large-export.yaml b/deploy/compose.large-export.yaml new file mode 100644 index 0000000..4641364 --- /dev/null +++ b/deploy/compose.large-export.yaml @@ -0,0 +1,11 @@ +# Opt in after the target environment passes the large-export gate. +services: + data-query-worker: + init: true + environment: + DQW_MAX_RESULT_ROWS: "10000000" + DQW_DUCKDB_MEMORY_LIMIT: "256MB" + DQW_MAX_RESULT_BYTES: "2147483648" + DQW_QUERY_TIMEOUT_SECONDS: "900" + DQW_INGESTION_TIMEOUT_SECONDS: "1800" + DQW_RESULT_TTL_SECONDS: "86400" diff --git a/deploy/compose.mtls.yaml b/deploy/compose.mtls.yaml new file mode 100644 index 0000000..5ab5bc0 --- /dev/null +++ b/deploy/compose.mtls.yaml @@ -0,0 +1,24 @@ +# Use with compose.yaml. Neither service publishes a host port. +services: + data-query-worker: + environment: + DQW_TOKEN_KEYRING_FILE: /run/query-secrets/tokens.json + volumes: + - ${DQW_SECRETS_DIR:?Set DQW_SECRETS_DIR}:/run/query-secrets:ro + query-tls: + image: nginxinc/nginx-unprivileged:1.28-alpine + read_only: true + cap_drop: [ALL] + security_opt: [no-new-privileges:true] + pids_limit: 64 + mem_limit: 128m + tmpfs: + - /tmp:size=16m,mode=1777 + volumes: + - ./deploy/nginx-mtls.conf:/etc/nginx/conf.d/default.conf:ro + - ${DQW_TLS_DIR:?Set DQW_TLS_DIR}:/run/query-tls:ro + networks: [worker-private, query-broker] + depends_on: [data-query-worker] +networks: + query-broker: + internal: true diff --git a/deploy/nginx-mtls.conf b/deploy/nginx-mtls.conf new file mode 100644 index 0000000..4848d34 --- /dev/null +++ b/deploy/nginx-mtls.conf @@ -0,0 +1,22 @@ +server { + listen 8443 ssl; + server_name query-tls; + ssl_certificate /run/query-tls/server.crt; + ssl_certificate_key /run/query-tls/server.key; + ssl_client_certificate /run/query-tls/ca.crt; + ssl_verify_client on; + ssl_protocols TLSv1.2 TLSv1.3; + # Do not log result URLs, SQL bodies, Authorization, or client credentials. + access_log off; + client_max_body_size 21g; + client_body_timeout 1800s; + location / { + proxy_pass http://data-query-worker:8080; + proxy_request_buffering off; + proxy_buffering off; + proxy_read_timeout 1800s; + proxy_send_timeout 1800s; + proxy_set_header Host $host; + proxy_set_header X-Request-ID $http_x_request_id; + } +} diff --git a/docs/production-foundations.md b/docs/production-foundations.md new file mode 100644 index 0000000..caa1655 --- /dev/null +++ b/docs/production-foundations.md @@ -0,0 +1,86 @@ +# Coordinated release: worker 0.2.0 / Analysis 0.14.0 + +Analysis is the OMERO-aware gateway. The worker remains read-only and has no OMERO credentials or write API. This release retains `omero-data-query-worker-v1`, existing routes, typed `$name` parameters, preview rules and CSV bytes. Analysis retains `omero-data-query-v1`, result tokens, local/remote policy, notebook `ctx.query()` and saved Method bindings v1/v2. Saving is additive and depends on `features.result_provenance_v1`. + +## Compatibility contract + +The frozen worker source is commit `2115594`; the frozen Analysis broker is from `4db7f8b` on `analysis_integration`. `tests/fixtures/compatibility` contains the contract and original broker. The Analysis HTTP harness runs both brokers against both worker implementations, exercising DuckDB, SQLite, SQLite3, CSV, schema, typed parameters, Unicode/quoting, empty results, downloads, hit/miss/bypass and reuse after rollback. It starts real HTTP processes and separate caches, using the installed dependency set. It is not a substitute for testing every historical engine build. + +Tested engines: DuckDB 1.5.5, sqlglot 28.10.1, SQLite 3.45.3 (Windows) and 3.46.1 (Linux container). The retained local worker 0.1.0 image also contains DuckDB 1.5.5/sqlglot 28.10.1/SQLite 3.46.1. SQLite3 is a filename alias for SQLite. The existing browser runtime uses DuckDB 1.5.1. Supported files are standalone databases readable by the selected engine and UTF-8 CSV accepted by the existing ingestion rules; external files, extensions, attached databases and executable database features are outside the supported format. Copying an uncheckpointed database without its journal is unsupported. Original source bytes are never rewritten. + +Actual engine, parser, worker and policy versions, effective query limits, source/schema digests and typed parameters participate in result identity. Legacy manifests are readable; missing execution provenance causes query recomputation. Startup discards legacy entries lacking the digest needed for recovery validation, including converted CSV sources without a converted-database digest. Rollback can reuse readable sources and recompute results under the previous key; do not downgrade DuckDB separately without testing its file compatibility. + +## Resource and trust boundaries + +Use exactly one service process per cache volume. An OS lock rejects a second process. Source leases span execution; result leases span sending the file, including failures/disconnects. Cleanup, quota eviction and administrative purge skip leases. Duplicate deterministic queries and exhausted ingestion/query slots return 429 instead of accumulating lock waiters. + +Upload admission happens before multipart parsing. Spools, source copies, CSV conversion output, manifests and result data count toward storage admission. Source and result quotas remain separate. A source may require its spool, immutable copy and converted database simultaneously. If safe eviction cannot make space, the request fails with 507. Result rows are serialized with an exact UTF-8 byte check before each write, including headers and individual oversized values. Subprocess conversion files are observed periodically; use a dedicated filesystem quota in production to cap transient conversion growth and protect unrelated services from a native-engine failure. + +Ingestion runs off the async loop and defaults to one simultaneous ingestion; query concurrency remains four. Subprocess limits, cancellation monitoring and `finally` cleanup release pipes, processes, uploads, staging paths and slots. Engine children receive credential-free settings and an allowlisted environment before processing sources or SQL. They still run under the service UID: native-code compromise is contained by the container/host boundaries, not by a separate per-query security identity. Keep secrets and cache inaccessible to other users, worker egress blocked, root filesystem read-only, capabilities dropped, no-new-privileges enabled and cgroup memory/PID limits enforced. + +Cache storage contains plaintext scientific data. It is disposable, not an audit store or durable result archive. Persist promoted results and provenance in OMERO. No SQL or parameter values are included in metrics. The Analysis audit records identity, context, hashes, counts, duration and outcome; the full recipe is deliberately stored only in the protected provenance annotation. + +Source and result publication now fsyncs the staged files and directories, renames the entry, then fsyncs both parents before returning. On restart, a service/engine lock fences surviving work, incomplete staging is removed immediately and committed file checksums/lengths are verified before readiness. Corrupt entries are discarded for recomputation. Linux execution children install a parent-death signal. Each DuckDB query has a private spill directory inside its quota-accounted staging path, preventing concurrent queries from colliding in a source-adjacent spill directory. + +## Optional 10M export profile + +The [completed joint release evidence](https://github.com/NL-BioImaging/OMERO.Analysis/blob/297f5c44ed1b0501a6c19444a5eab7f9e8d0f686/docs/testing/query-large-export-2026-09-08.json) validates 540 10M requests: 468 successes, 72 expected admission rejections and zero unexpected failures. All 72 single-client 1M/4M regression requests also pass. Normal runs have no OOM kills and stay within cache quotas. Separate engine-OOM, exact row/byte boundary, live OMERO save and VM power-cut gates pass. See the [tested implementation pair](https://github.com/NL-BioImaging/OMERO.Analysis/blob/297f5c44ed1b0501a6c19444a5eab7f9e8d0f686/docs/testing/query-version-pair.json) for exact commits, engine versions and image/test scope. + +`deploy/compose.large-export.yaml` raises the row ceiling to 10,000,000, CSV byte limit to 2 GiB, query timeout to 900 seconds and ingestion timeout to 1,800 seconds. Normal defaults remain unchanged. Apply the matching Analysis profile, including its dedicated CSV promotion limit and web/proxy timeouts. The tested two-CPU, 2 GiB, four-query container uses 256 MB per DuckDB engine; four 512 MB engines exceeded the container memory in testing. Wide results must still fit the byte/time/storage limits or fail explicitly. + +The coordinated Analysis [large-export/recovery runbook](https://github.com/NL-BioImaging/OMERO.Analysis/blob/297f5c44ed1b0501a6c19444a5eab7f9e8d0f686/docs/data-query-large-export-recovery.md) provides reproducible 1M/4M/10M fixtures, real HTTP capacity tests, exact 10M/2-GiB boundary checks, real DuckDB/SQLite/CSV OOM, whole-worker OOM, active-query restart and isolated Hyper-V power-cut gates. Real OMERO saved the same verified 609,383,573-byte 10M CSV for all three formats. Analysis journals and an explicit administrator command reconcile interrupted promotions; the worker remains read-only and OMERO-agnostic. + +## Operator controls + +Additional environment settings: + +| Setting | Default / purpose | +| --- | --- | +| `DQW_MAX_CONCURRENT_INGESTIONS` | `1`, includes early upload admission | +| `DQW_CLEANUP_INTERVAL_SECONDS` | `60` | +| `DQW_TOKEN_KEYRING_FILE` | optional atomically replaced JSON keyring | + +`GET /health/live` is independent of storage. Readiness checks writable temporary storage, physical free space, shutdown state and headroom after safe eviction. `GET /v1/metrics` requires bearer authentication and exposes request durations, admission rejection, active operations, cache counters/bytes, failures and subprocess exits. Counters reset on restart; scrape externally. Periodic cleanup removes expired inactive entries and abandoned temporary directories older than twice the larger execution timeout. + +Run the local administrative command inside the running worker container: + +```sh +omero-data-query-worker purge --source-id src_ +omero-data-query-worker purge --scope-id --apply +omero-data-query-worker purge --result-id res_ --apply +``` + +The default is dry-run. Exactly one selector is required. Source/scope selection includes associated results. A busy selection is reported without partial deletion. This command contacts the live lease manager over loopback; do not manipulate active cache directories with shell deletion tools. + +## Transport and overlapping credentials + +Legacy private-network `DQW_API_TOKEN` remains supported. TLS is optional and must not be switched on automatically during a mixed-version upgrade. For rotation, mount a **directory**, set `DQW_TOKEN_KEYRING_FILE=/run/query-secrets/tokens.json`, and atomically rename a new JSON file within that directory. Format: `{"tokens":["old credential at least 16 characters","new credential at least 16 characters"]}`. Never commit real values. Mounting a single file may pin its old inode across host-side atomic replacement. Malformed keyrings fail closed. The environment token is accepted alongside the ring; remove it from deployment configuration when retiring that credential. + +`deploy/compose.mtls.yaml` and `deploy/nginx-mtls.conf` supply an optional private mTLS reverse proxy. Combine with `compose.yaml` as a starting profile and replace the smoke token/image with deployment secrets and a pinned candidate image. Mount server.crt, server.key and ca.crt; the server SAN must cover `query-tls`. Attach Analysis only to the broker-facing network. Neither service publishes a host port. Analysis supplies a CA bundle and client certificate/key. Retain bearer authentication behind mTLS. The base compose file is a smoke fixture, not production credential provisioning. + +## Reproducible gates and rollout + +```sh +python -m pip install -e '.[test]' +python -m pytest +ruff check . +ruff format --check . +mypy src +python scripts/smoke_container.py +python scripts/security_container.py +# In the adjacent Analysis checkout (its test dependencies installed): +DQW_RUN_INTEGRATION=1 python -m pytest tests/test_worker_contract.py +python scripts/test_local_omero_queries.py +python scripts/test_query_mtls.py +python scripts/benchmark_query_capacity.py --rows 100000 --repeats 3 +``` + +PowerShell: `$env:DQW_RUN_INTEGRATION='1'` before pytest. Set `DQW_REPO`/`DQW_PYTHON` for non-adjacent checkouts. The Linux gate uses disposable containers for regression, read-only filesystem, blocked egress, PID pressure and OOM probes; tests also cover real query timeout/crash/disconnect, corrupt sources, storage exhaustion and leased eviction. The allocator OOM probe validates the cgroup boundary; it does not prove recovery for every possible engine OOM. See the Analysis release document for live permission coverage and remaining acceptance boundaries. + +1. Retain both previous images/configuration and back up deployment configuration. Record exact candidate image digests and engine versions. Keep audit/state volumes separate from cache. +2. Upgrade only the worker. Exercise existing Analysis queries, schema and downloads. +3. Upgrade Analysis from `analysis_integration`. Check legacy notebooks and Methods, then verify exact CSV promotion and provenance. +4. Enable the chosen saving/transport deployment configuration after the permission, container and compatibility gates pass on the target environment. Test overlapping credentials before removing the old credential. +5. To roll back, restore the previous Analysis and/or worker image with its matching transport configuration. Keep OMERO annotations, audit/state and cache volumes. Do not delete caches to force rollback; allow safe reuse/recomputation. Old Analysis ignores additive metadata and offers its existing manual result saving. + +Production rollout remains gated on target-environment acceptance, including realistic large exports and restart/fault behavior. These changes do not introduce asynchronous jobs, cancellation APIs, Arrow/Parquet, multi-worker routing, OMERO.Tables materialization or grid discovery. diff --git a/pyproject.toml b/pyproject.toml index b9723be..913dbbd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "omero-data-query-worker" -version = "0.1.1" +version = "0.2.0" description = "A hardened, OMERO-agnostic service for bounded queries over DuckDB, SQLite, and CSV sources" readme = "README.md" requires-python = ">=3.11,<3.14" diff --git a/scripts/security_container.py b/scripts/security_container.py new file mode 100644 index 0000000..097923a --- /dev/null +++ b/scripts/security_container.py @@ -0,0 +1,91 @@ +"""Linux release gate. All destructive fault probes use new disposable containers.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import uuid +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def run(*args: str, check: bool = True) -> subprocess.CompletedProcess[str]: + return subprocess.run(["docker", *args], cwd=ROOT, check=check, text=True) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--image", default="omero-data-query-worker:smoke") + args = parser.parse_args() + image = "dqw-security-tests:" + uuid.uuid4().hex[:12] + run( + "build", + "-f", + "deploy/Dockerfile.tests", + "--build-arg", + f"WORKER_IMAGE={args.image}", + "-t", + image, + ".", + ) + security = [ + "--network", + "none", + "--read-only", + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--tmpfs", + "/tmp:size=256m,mode=1777", + ] + try: + run("run", "--rm", *security, "--memory", "2g", "--pids-limit", "128", image) + for case in ("isolation", "pids"): + run( + "run", + "--rm", + *security, + "--memory", + "256m", + "--pids-limit", + "32", + image, + "python", + "scripts/security_inside.py", + case, + ) + name = "dqw-oom-" + uuid.uuid4().hex[:12] + try: + result = run( + "run", + "--name", + name, + *security, + "--memory", + "96m", + "--memory-swap", + "96m", + "--pids-limit", + "32", + image, + "python", + "scripts/security_inside.py", + "oom", + check=False, + ) + state = json.loads(subprocess.check_output(["docker", "inspect", name], text=True))[0][ + "State" + ] + assert result.returncode == 137 and state["OOMKilled"], state + finally: + run("rm", "-f", "-v", name, check=False) + print("Linux security gates passed: tests, filesystem, egress, PID, OOM") + finally: + run("image", "rm", image, check=False) + + +if __name__ == "__main__": + main() diff --git a/scripts/security_inside.py b/scripts/security_inside.py new file mode 100644 index 0000000..a40464b --- /dev/null +++ b/scripts/security_inside.py @@ -0,0 +1,59 @@ +"""Bounded OS fault probes, executed only in disposable containers.""" + +from __future__ import annotations + +import argparse +import errno +import json +import os +import socket +import subprocess +import sys + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("case", choices=["isolation", "pids", "oom"]) + case = parser.parse_args().case + if case == "oom": + # Cgroup budget is 96 MiB; this container is expected to be OOM killed. + bytearray(512 * 1024 * 1024) + raise AssertionError("OOM limit was not enforced") + if case == "isolation": + assert os.getuid() != 0 + try: + with open("/app/forbidden", "w") as handle: + handle.write("unexpected") + except OSError as exc: + assert exc.errno in (errno.EROFS, errno.EACCES) + else: + raise AssertionError("Root filesystem is writable") + try: + socket.create_connection(("1.1.1.1", 443), timeout=1) + except OSError: + pass + else: + raise AssertionError("Unexpected network egress") + if case == "pids": + children = [] + try: + for _ in range(64): + try: + children.append( + subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + ) + except OSError as exc: + assert exc.errno == errno.EAGAIN + break + else: + raise AssertionError("PID limit was not enforced") + finally: + for child in children: + child.terminate() + for child in children: + child.wait(timeout=5) + print(json.dumps({"case": case, "status": "passed"})) + + +if __name__ == "__main__": + main() diff --git a/src/omero_data_query_worker/__init__.py b/src/omero_data_query_worker/__init__.py index 8bba50a..9da50c1 100644 --- a/src/omero_data_query_worker/__init__.py +++ b/src/omero_data_query_worker/__init__.py @@ -1,3 +1,3 @@ """Hardened tabular data query worker.""" -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/src/omero_data_query_worker/__main__.py b/src/omero_data_query_worker/__main__.py index 772909d..a856df5 100644 --- a/src/omero_data_query_worker/__main__.py +++ b/src/omero_data_query_worker/__main__.py @@ -1,9 +1,41 @@ from __future__ import annotations +import argparse +import json +import urllib.request + import uvicorn +from .config import Settings +from .operations import Credentials + def main() -> None: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command") + purge = commands.add_parser("purge", help="Purge using the running worker's lease manager") + selectors = purge.add_mutually_exclusive_group(required=True) + selectors.add_argument("--source-id") + selectors.add_argument("--scope-id") + selectors.add_argument("--result-id") + purge.add_argument("--apply", action="store_true", help="Default is dry-run") + args = parser.parse_args() + if args.command == "purge": + payload = { + "source_id": args.source_id, + "scope_id": args.scope_id, + "result_id": args.result_id, + "dry_run": not args.apply, + } + token = Credentials(Settings.from_env()).tokens()[0] + request = urllib.request.Request( + "http://127.0.0.1:8080/v1/cache/purge", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json", "Authorization": f"Bearer {token}"}, + ) + with urllib.request.urlopen(request, timeout=30) as response: + print(response.read().decode()) + return uvicorn.run( "omero_data_query_worker.app:app", host="0.0.0.0", # noqa: S104 - the container API must listen on its private network diff --git a/src/omero_data_query_worker/app.py b/src/omero_data_query_worker/app.py index 766ec36..f78edb8 100644 --- a/src/omero_data_query_worker/app.py +++ b/src/omero_data_query_worker/app.py @@ -1,11 +1,15 @@ from __future__ import annotations +import asyncio +import functools import logging -import secrets -import uuid -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager -from typing import Annotated +import shutil +import tempfile +import threading +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager, suppress +from pathlib import Path +from typing import Annotated, TypeVar from fastapi import ( APIRouter, @@ -18,13 +22,17 @@ UploadFile, ) from fastapi.responses import FileResponse, JSONResponse +from starlette.concurrency import run_in_threadpool +from starlette.responses import PlainTextResponse from . import __version__ from .config import Settings -from .errors import AuthenticationError, WorkerError +from .engines import engine_versions +from .errors import WorkerError from .models import ( CacheStatusResponse, ErrorResponse, + PurgeRequest, QueryRequest, QueryResponse, SchemaResponse, @@ -33,20 +41,42 @@ SourceResolveResponse, SourceSummary, ) +from .operations import Credentials, volume_lock +from .operations import request_id as safe_request_id from .policy import POLICY_VERSION from .service import QueryService +from .upload import UploadAdmission logger = logging.getLogger("omero_data_query_worker") +T = TypeVar("T") def create_app(settings: Settings | None = None) -> FastAPI: selected = settings or Settings.from_env() service = QueryService(selected) + credentials = Credentials(selected) + + def cleanup_loop() -> None: + while not service.stopping.wait(selected.cleanup_interval_seconds): + try: + service.cache.cleanup_sources() + service.cache.cleanup_results() + service.cache.cleanup_temporary() + except OSError: + service.metrics.add("cleanup_failures_total") @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: - service.prepare() - yield + with volume_lock(selected.cache_dir): + service.prepare() + credentials.tokens() + cleaner = threading.Thread(target=cleanup_loop, daemon=True) + cleaner.start() + try: + yield + finally: + service.stopping.set() + cleaner.join(5) application = FastAPI( title="OMERO.DataQueryWorker", @@ -58,7 +88,7 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: @application.middleware("http") async def request_context(request: Request, call_next): # type: ignore[no-untyped-def] - request_id = request.headers.get("X-Request-ID") or uuid.uuid4().hex + request_id = safe_request_id(request.headers.get("X-Request-ID")) request.state.request_id = request_id response = await call_next(request) response.headers["X-Request-ID"] = request_id @@ -68,7 +98,8 @@ async def request_context(request: Request, call_next): # type: ignore[no-untyp @application.exception_handler(WorkerError) async def worker_error(request: Request, exc: WorkerError) -> JSONResponse: - request_id = getattr(request.state, "request_id", uuid.uuid4().hex) + request_id = getattr(request.state, "request_id", safe_request_id(None)) + service.metrics.add(f"errors_{exc.code}_total") logger.warning( "request_failed request_id=%s code=%s path=%s", request_id, @@ -89,13 +120,41 @@ def live() -> dict[str, str]: return {"status": "live"} @application.get("/health/ready", include_in_schema=False) - def ready() -> dict[str, str]: - return {"status": "ready"} + def ready() -> JSONResponse: + try: + with tempfile.TemporaryFile(dir=selected.tmp_dir) as probe: + probe.write(b"ready") + probe.flush() + if ( + service.stopping.is_set() + or not service.cache.has_capacity() + or shutil.disk_usage(selected.cache_dir).free < 1024 * 1024 + ): + raise OSError + except OSError: + return JSONResponse({"status": "unavailable"}, status_code=503) + return JSONResponse({"status": "ready"}) def authenticate(authorization: Annotated[str | None, Header()] = None) -> None: - scheme, _, value = (authorization or "").partition(" ") - if scheme.lower() != "bearer" or not secrets.compare_digest(value, selected.api_token): - raise AuthenticationError("A valid worker bearer token is required") + credentials.authenticate(authorization) + + async def watch_disconnect(request: Request, cancelled: threading.Event) -> None: + while not cancelled.is_set(): + if await request.is_disconnected(): + cancelled.set() + return + await asyncio.sleep(0.1) + + async def operation(request: Request, function: Callable[..., T]) -> T: + cancelled = threading.Event() + watcher = asyncio.create_task(watch_disconnect(request, cancelled)) + try: + return await run_in_threadpool(functools.partial(function, cancelled=cancelled)) + finally: + cancelled.set() + watcher.cancel() + with suppress(asyncio.CancelledError): + await watcher router = APIRouter(prefix="/v1", dependencies=[Depends(authenticate)]) @@ -108,12 +167,16 @@ def capabilities() -> dict[str, object]: "formats": ["duckdb", "sqlite", "csv"], "csv_table": "data", "parameter_style": "$name", + "features": {"result_provenance_v1": True}, + "engine_versions": engine_versions(), "limits": { "query_timeout_seconds": selected.query_timeout_seconds, "max_result_rows": selected.max_result_rows, "max_result_bytes": selected.max_result_bytes, "max_source_bytes": selected.max_source_bytes, "max_concurrent_queries": selected.max_concurrent_queries, + "max_concurrent_ingestions": selected.max_concurrent_ingestions, + "ingestion_timeout_seconds": selected.ingestion_timeout_seconds, }, } @@ -123,6 +186,7 @@ def resolve_source(payload: SourceResolveRequest) -> SourceResolveResponse: @router.post("/sources", response_model=SourceSummary, status_code=201) async def ingest_source( + request: Request, scope_id: Annotated[str, Form()], source_ref: Annotated[str, Form()], format: Annotated[SourceFormat, Form()], # noqa: A002 - public wire name @@ -130,13 +194,17 @@ async def ingest_source( file: Annotated[UploadFile, File()], expected_sha256: Annotated[str | None, Form()] = None, ) -> SourceSummary: - return await service.ingest_source( - scope_id=scope_id, - source_ref=source_ref, - source_format=format, - declared_size=size, - expected_sha256=expected_sha256, - upload=file, + return await operation( + request, + functools.partial( + service.ingest_source, + scope_id=scope_id, + source_ref=source_ref, + source_format=format, + declared_size=size, + expected_sha256=expected_sha256, + upload=file, + ), ) @router.get("/sources/{source_id}/schema", response_model=SchemaResponse) @@ -144,24 +212,54 @@ def source_schema(source_id: str) -> SchemaResponse: return service.schema(source_id) @router.post("/sources/{source_id}/query", response_model=QueryResponse) - def source_query(source_id: str, payload: QueryRequest) -> QueryResponse: - return service.query(source_id, payload) + async def source_query( + request: Request, source_id: str, payload: QueryRequest + ) -> QueryResponse: + return await operation(request, functools.partial(service.query, source_id, payload)) @router.get("/results/{result_id}/download") def result_download(result_id: str) -> FileResponse: - path, record = service.result_file(result_id) - return FileResponse( - path, - media_type="text/csv; charset=utf-8", - filename=f"{record.result_id}.csv", - ) + record = service.cache.acquire_result(result_id) + try: + return LeasedFileResponse( + service.cache.result_path(result_id) / record.result_file, + media_type="text/csv; charset=utf-8", + filename=f"{record.result_id}.csv", + release=lambda: service.cache.release(result_id), + ) + except BaseException: + service.cache.release(result_id) + raise @router.get("/cache/status", response_model=CacheStatusResponse) def cache_status() -> dict[str, int]: return service.cache.status() + @router.get("/metrics", response_class=PlainTextResponse) + def metrics() -> str: + return service.metrics.render(service.cache.status()) + + @router.post("/cache/purge") + def purge(payload: PurgeRequest) -> dict[str, object]: + return service.cache.purge(**payload.model_dump()) + application.include_router(router) + application.add_middleware(UploadAdmission, service=service, credentials=credentials) return application +class LeasedFileResponse(FileResponse): + def __init__( + self, path: Path, *, release: Callable[[], None], media_type: str, filename: str + ) -> None: + super().__init__(path, media_type=media_type, filename=filename) + self.release = release + + async def __call__(self, scope, receive, send): # type: ignore[no-untyped-def] + try: + await super().__call__(scope, receive, send) + finally: + self.release() + + app = create_app() diff --git a/src/omero_data_query_worker/cache.py b/src/omero_data_query_worker/cache.py index 992de5a..acaa0a6 100644 --- a/src/omero_data_query_worker/cache.py +++ b/src/omero_data_query_worker/cache.py @@ -9,13 +9,20 @@ import time import uuid from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any from .config import Settings -from .errors import ResultNotFound, SourceNotFound +from .errors import ( + CacheCapacityExceeded, + InvalidQuery, + QueryCapacityExceeded, + ResultNotFound, + SourceNotFound, +) from .models import SourceFormat SOURCE_ID_PATTERN = re.compile(r"^src_[a-f0-9]{64}$") @@ -45,6 +52,29 @@ def _write_json_atomic(path: Path, value: dict[str, Any]) -> None: handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) + sync_directory(path.parent) + + +def sync_directory(path: Path) -> None: + if os.name == "posix": + fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def file_digest(path: Path) -> str: + with path.open("rb") as handle: + return hashlib.file_digest(handle, "sha256").hexdigest() + + +def sync_tree(path: Path) -> None: + for item in path.iterdir(): + if item.is_file(): + with item.open("r+b") as handle: + os.fsync(handle.fileno()) + sync_directory(path) @dataclass(slots=True) @@ -62,6 +92,7 @@ class SourceRecord: created_at: str accessed_at: str expires_at: str + data_sha256: str | None = None @classmethod def from_dict(cls, value: dict[str, Any]) -> SourceRecord: @@ -79,6 +110,7 @@ def from_dict(cls, value: dict[str, Any]) -> SourceRecord: created_at=str(value["created_at"]), accessed_at=str(value["accessed_at"]), expires_at=str(value["expires_at"]), + data_sha256=value.get("data_sha256"), ) def as_dict(self) -> dict[str, Any]: @@ -96,6 +128,7 @@ def as_dict(self) -> dict[str, Any]: "created_at": self.created_at, "accessed_at": self.accessed_at, "expires_at": self.expires_at, + "data_sha256": self.data_sha256, } @@ -115,6 +148,7 @@ class ResultRecord: created_at: str accessed_at: str expires_at: str + execution: dict[str, Any] | None = None @classmethod def from_dict(cls, value: dict[str, Any]) -> ResultRecord: @@ -133,6 +167,7 @@ def from_dict(cls, value: dict[str, Any]) -> ResultRecord: created_at=str(value["created_at"]), accessed_at=str(value["accessed_at"]), expires_at=str(value["expires_at"]), + execution=value.get("execution"), ) def as_dict(self) -> dict[str, Any]: @@ -151,6 +186,7 @@ def as_dict(self) -> dict[str, Any]: "created_at": self.created_at, "accessed_at": self.accessed_at, "expires_at": self.expires_at, + "execution": self.execution, } @@ -163,6 +199,60 @@ def __init__(self, settings: Settings): self.bypasses = 0 self.evictions = 0 self.active_queries = 0 + self._leases: dict[str, int] = {} + self._staging: dict[Path, tuple[str, int]] = {} + + @contextmanager + def lease_source(self, source_id: str) -> Iterator[SourceRecord]: + with self._lock: + record = self.get_source(source_id) + self._leases[source_id] = self._leases.get(source_id, 0) + 1 + try: + yield record + finally: + self.release(source_id) + + def acquire_result(self, result_id: str) -> ResultRecord: + with self._lock: + record = self.get_result(result_id) + self._leases[result_id] = self._leases.get(result_id, 0) + 1 + return record + + def release(self, identifier: str) -> None: + with self._lock: + count = self._leases.get(identifier, 0) + if count <= 1: + self._leases.pop(identifier, None) + else: + self._leases[identifier] = count - 1 + + def reserve(self, staging: Path, size: int) -> None: + """Account for concurrent staging bytes as well as committed cache entries.""" + with self._lock: + kind, old_size = self._staging[staging] + root = self.settings.sources_dir if kind == "source" else self.settings.results_dir + maximum = ( + self.settings.source_cache_max_bytes + if kind == "source" + else self.settings.result_cache_max_bytes + ) + others = sum(n for p, (k, n) in self._staging.items() if p != staging and k == kind) + if size + others > maximum: + raise CacheCapacityExceeded("Staged data exceeds the cache capacity") + if max(0, size - old_size) > shutil.disk_usage(self.settings.cache_dir).free: + raise CacheCapacityExceeded("Insufficient free space on the cache volume") + self._cleanup(root, maximum - size - others, set()) + if directory_bytes(root) + size + others > maximum: + raise CacheCapacityExceeded("Cache capacity is occupied by active work") + self._staging[staging] = kind, size + + def observe_staging(self, staging: Path) -> None: + self.reserve(staging, directory_bytes(staging)) + + def discard_staging(self, staging: Path) -> None: + with self._lock: + self._staging.pop(staging, None) + self._remove_directory(staging, self.settings.tmp_dir, count=False) @staticmethod def source_id(scope_id: str, source_ref: str) -> str: @@ -176,9 +266,11 @@ def result_path(self, result_id: str) -> Path: return self.settings.results_dir / result_id def staging_path(self, prefix: str) -> Path: - path = self.settings.tmp_dir / f"{prefix}-{uuid.uuid4().hex}" - path.mkdir(parents=True) - return path + with self._lock: + path = self.settings.tmp_dir / f"{prefix}-{uuid.uuid4().hex}" + path.mkdir(parents=True) + self._staging[path] = prefix, 0 + return path def _read_manifest(self, path: Path) -> dict[str, Any]: try: @@ -198,7 +290,11 @@ def get_source(self, source_id: str, *, touch: bool = True) -> SourceRecord: record = SourceRecord.from_dict(self._read_manifest(path / "manifest.json")) except (FileNotFoundError, KeyError, TypeError, ValueError) as exc: raise SourceNotFound("The source is not cached") from exc - if parse_iso(record.expires_at) <= utc_now() or not (path / record.data_file).is_file(): + if ( + (parse_iso(record.expires_at) <= utc_now() and not self._leases.get(source_id)) + or Path(record.data_file).name != record.data_file + or not (path / record.data_file).is_file() + ): self._remove_directory(path, self.settings.sources_dir) raise SourceNotFound("The source has expired or is incomplete") if touch: @@ -229,15 +325,23 @@ def resolve_source( def commit_source(self, staging: Path, record: SourceRecord) -> SourceRecord: target = self.source_path(record.source_id) with self._lock: + record.data_sha256 = file_digest(staging / record.data_file) _write_json_atomic(staging / "manifest.json", record.as_dict()) if target.exists(): - self._remove_directory(staging, self.settings.tmp_dir, count=False) + self.discard_staging(staging) return self.get_source(record.source_id) + self.observe_staging(staging) + sync_tree(staging) os.replace(staging, target) + sync_directory(target.parent) + sync_directory(staging.parent) + self._staging.pop(staging, None) self.cleanup_sources(protected={record.source_id}) return record - def get_result_by_query_key(self, query_key: str) -> ResultRecord | None: + def get_result_by_query_key( + self, query_key: str, *, require_provenance: bool = False + ) -> ResultRecord | None: result_id = f"res_{query_key}" try: record = self.get_result(result_id) @@ -245,6 +349,12 @@ def get_result_by_query_key(self, query_key: str) -> ResultRecord | None: return None if record.query_key != query_key: return None + if require_provenance and not record.execution: + with self._lock: + if self._leases.get(result_id): + raise QueryCapacityExceeded("Legacy result is being downloaded; retry later") + self._remove_directory(self.result_path(result_id), self.settings.results_dir) + return None return record def get_result(self, result_id: str, *, touch: bool = True) -> ResultRecord: @@ -257,7 +367,8 @@ def get_result(self, result_id: str, *, touch: bool = True) -> ResultRecord: except (FileNotFoundError, KeyError, TypeError, ValueError) as exc: raise ResultNotFound("The result is not cached") from exc if ( - parse_iso(record.expires_at) <= utc_now() + (parse_iso(record.expires_at) <= utc_now() and not self._leases.get(result_id)) + or Path(record.result_file).name != record.result_file or not (path / record.result_file).is_file() ): self._remove_directory(path, self.settings.results_dir) @@ -274,9 +385,14 @@ def commit_result(self, staging: Path, record: ResultRecord) -> ResultRecord: with self._lock: _write_json_atomic(staging / "manifest.json", record.as_dict()) if target.exists(): - self._remove_directory(staging, self.settings.tmp_dir, count=False) + self.discard_staging(staging) return self.get_result(record.result_id) + self.observe_staging(staging) + sync_tree(staging) os.replace(staging, target) + sync_directory(target.parent) + sync_directory(staging.parent) + self._staging.pop(staging, None) self.cleanup_results(protected={record.result_id}) return record @@ -289,7 +405,7 @@ def _manifest_directories(self, root: Path) -> Iterator[tuple[Path, dict[str, An try: yield path, self._read_manifest(path / "manifest.json") except FileNotFoundError: - self._remove_directory(path, root) + yield path, {} def _cleanup( self, @@ -298,6 +414,7 @@ def _cleanup( protected: set[str], ) -> None: with self._lock: + protected = protected | set(self._leases) now = utc_now() entries: list[tuple[datetime, Path, int]] = [] total = 0 @@ -336,7 +453,41 @@ def cleanup_results(self, protected: set[str] | None = None) -> None: protected or set(), ) - def cleanup_temporary(self) -> None: + def recover(self) -> None: + """Called under exclusive service/engine locks before accepting requests.""" + self.cleanup_temporary(startup=True) + for root, kind in ( + (self.settings.sources_dir, "source"), + (self.settings.results_dir, "result"), + ): + for path, value in self._manifest_directories(root): + try: + if kind == "source": + source = SourceRecord.from_dict(value) + filename = source.data_file + expected = source.data_sha256 or ( + source.sha256 if source.format is not SourceFormat.csv else None + ) + identifier = source.source_id + else: + result = ResultRecord.from_dict(value) + filename = result.result_file + expected = (result.execution or {}).get("result_sha256") + identifier = result.result_id + if (path / filename).stat().st_size != result.byte_count: + raise ValueError("Result length disagrees with its manifest") + if ( + identifier != path.name + or Path(filename).name != filename + or not expected + or (path / filename).is_symlink() + or file_digest(path / filename) != expected + ): + raise ValueError("Incomplete cache publication") + except (OSError, ValueError, KeyError, TypeError, AttributeError): + self._remove_directory(path, root) + + def cleanup_temporary(self, *, startup: bool = False) -> None: cutoff = ( time.time() - max( @@ -347,12 +498,18 @@ def cleanup_temporary(self) -> None: ) for path in self.settings.tmp_dir.iterdir(): try: - if path.is_dir() and path.stat().st_mtime < cutoff: + if ( + path not in self._staging + and path.is_dir() + and (startup or path.stat().st_mtime < cutoff) + ): self._remove_directory(path, self.settings.tmp_dir, count=False) except FileNotFoundError: continue def _remove_directory(self, path: Path, root: Path, *, count: bool = True) -> None: + if self._leases.get(path.name): + return root_resolved = root.resolve() path_resolved = path.resolve() if path_resolved.parent != root_resolved: @@ -362,6 +519,51 @@ def _remove_directory(self, path: Path, root: Path, *, count: bool = True) -> No if count: self.evictions += 1 + def purge( + self, + *, + source_id: str | None = None, + scope_id: str | None = None, + result_id: str | None = None, + dry_run: bool = True, + ) -> dict[str, Any]: + if sum(value is not None for value in (source_id, scope_id, result_id)) != 1: + raise InvalidQuery("Select exactly one source_id, scope_id, or result_id") + with self._lock: + sources = list(self._manifest_directories(self.settings.sources_dir)) + selected_sources = { + p.name + for p, m in sources + if p.name == source_id or (scope_id and m.get("scope_id") == scope_id) + } + entries = [ + (p, self.settings.sources_dir) for p, _ in sources if p.name in selected_sources + ] + entries.extend( + (p, self.settings.results_dir) + for p, m in self._manifest_directories(self.settings.results_dir) + if p.name == result_id or m.get("source_id") in selected_sources + ) + busy = {p.name for p, _ in entries if self._leases.get(p.name)} + # Do not partially purge a source while it is producing a new result. + if busy: + return { + "dry_run": dry_run, + "selected": [p.name for p, _ in entries], + "removed": [], + "busy": sorted(busy), + } + selected = [p.name for p, _ in entries] + if not dry_run: + for p, root in entries: + self._remove_directory(p, root) + return { + "dry_run": dry_run, + "selected": selected, + "removed": [] if dry_run else selected, + "busy": [], + } + def status(self) -> dict[str, int]: with self._lock: source_entries = list(self._manifest_directories(self.settings.sources_dir)) @@ -377,3 +579,18 @@ def status(self) -> dict[str, int]: "evictions": self.evictions, "active_queries": self.active_queries, } + + def has_capacity(self) -> bool: + """Readiness requires headroom after eviction of inactive entries.""" + with self._lock: + for kind, root, maximum in ( + ("source", self.settings.sources_dir, self.settings.source_cache_max_bytes), + ("result", self.settings.results_dir, self.settings.result_cache_max_bytes), + ): + pinned = sum( + directory_bytes(root / key) for key in self._leases if (root / key).is_dir() + ) + staged = sum(size for k, size in self._staging.values() if k == kind) + if pinned + staged >= maximum: + return False + return True diff --git a/src/omero_data_query_worker/config.py b/src/omero_data_query_worker/config.py index 3d53e24..7aaa736 100644 --- a/src/omero_data_query_worker/config.py +++ b/src/omero_data_query_worker/config.py @@ -2,7 +2,7 @@ import os import tempfile -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path GIB = 1024**3 @@ -38,6 +38,9 @@ class Settings: max_concurrent_queries: int = 4 duckdb_memory_limit: str = "1GB" duckdb_threads: int = 2 + token_keyring_file: str = "" + max_concurrent_ingestions: int = 1 + cleanup_interval_seconds: int = 60 @classmethod def from_env(cls) -> Settings: @@ -57,12 +60,15 @@ def from_env(cls) -> Settings: max_concurrent_queries=_positive_int("DQW_MAX_CONCURRENT_QUERIES", 4), duckdb_memory_limit=os.getenv("DQW_DUCKDB_MEMORY_LIMIT", "1GB"), duckdb_threads=_positive_int("DQW_DUCKDB_THREADS", 2), + token_keyring_file=os.getenv("DQW_TOKEN_KEYRING_FILE", ""), + max_concurrent_ingestions=_positive_int("DQW_MAX_CONCURRENT_INGESTIONS", 1), + cleanup_interval_seconds=_positive_int("DQW_CLEANUP_INTERVAL_SECONDS", 60), ) def prepare(self) -> None: - if not self.api_token: + if not self.api_token and not self.token_keyring_file: raise RuntimeError("DQW_API_TOKEN is required") - if len(self.api_token) < 16: + if self.api_token and len(self.api_token) < 16: raise RuntimeError("DQW_API_TOKEN must contain at least 16 characters") for path in (self.cache_dir, self.sources_dir, self.results_dir, self.tmp_dir): path.mkdir(parents=True, exist_ok=True) @@ -77,6 +83,10 @@ def prepare(self) -> None: probe.write_bytes(b"ready") probe.unlink() + def execution_settings(self) -> Settings: + """Never serialize service credentials into an engine process.""" + return replace(self, api_token="", token_keyring_file="") + @property def sources_dir(self) -> Path: return self.cache_dir / "sources" diff --git a/src/omero_data_query_worker/engines.py b/src/omero_data_query_worker/engines.py index 65a6545..9fc56e8 100644 --- a/src/omero_data_query_worker/engines.py +++ b/src/omero_data_query_worker/engines.py @@ -9,6 +9,7 @@ from collections.abc import Iterable from datetime import date, datetime, time from decimal import Decimal +from importlib.metadata import version from pathlib import Path from typing import Any @@ -23,6 +24,31 @@ CSV_EXTENSION = ".csv" +def engine_versions() -> dict[str, str]: + return { + "duckdb": duckdb.__version__, + "sqlite": sqlite3.sqlite_version, + "sqlglot": version("sqlglot"), + } + + +class BoundedCSVOutput: + def __init__(self, handle: Any, maximum: int): + self.handle = handle + self.maximum = maximum + self.size = 0 + self.digest = hashlib.sha256() + + def write(self, text: str) -> int: + content = text.encode("utf-8") + if self.size + len(content) > self.maximum: + raise QueryLimitExceeded(f"Query exceeds the {self.maximum} byte limit") + self.handle.write(content) + self.digest.update(content) + self.size += len(content) + return len(text) + + def safe_filename(filename: str) -> str: name = Path(filename.replace("\\", "/")).name.strip() if not name or name in {".", ".."} or "\x00" in name: @@ -275,8 +301,9 @@ def _write_result( ] preview: list[list[Any]] = [] row_count = 0 - with output_path.open("w", encoding="utf-8", newline="") as text_handle: - writer = csv.writer(text_handle, lineterminator="\n") + with output_path.open("wb") as text_handle: + bounded = BoundedCSVOutput(text_handle, max_bytes) + writer = csv.writer(bounded, lineterminator="\n") writer.writerow([column["name"] for column in columns]) while True: rows: Iterable[tuple[Any, ...]] = cursor.fetchmany(1000) @@ -295,8 +322,6 @@ def _write_result( preview.append([_json_value(value) for value in row]) writer.writerow([_csv_value(value) for value in row]) text_handle.flush() - if os.fstat(text_handle.fileno()).st_size > max_bytes: - raise QueryLimitExceeded(f"Query exceeds the {max_bytes} byte limit") text_handle.flush() os.fsync(text_handle.fileno()) byte_count = os.fstat(text_handle.fileno()).st_size @@ -305,6 +330,7 @@ def _write_result( "row_count": row_count, "byte_count": byte_count, "preview": preview, + "result_sha256": bounded.digest.hexdigest(), } @@ -316,7 +342,14 @@ def execute_duckdb_query( settings: Settings, ) -> dict[str, Any]: try: - connection = duckdb.connect(str(source_path), read_only=True) + # DuckDB's default spill path is derived from the shared source filename. + # Independent read-only processes must never share spill files. Keeping + # spill inside this query's staging directory also makes it quota-accounted. + connection = duckdb.connect( + str(source_path), + read_only=True, + config={"temp_directory": str(output_path.parent / "duckdb-spill")}, + ) configure_duckdb(connection, settings) cursor = connection.execute(sql, parameters) return _write_result( diff --git a/src/omero_data_query_worker/executor.py b/src/omero_data_query_worker/executor.py index 1524b2f..4a4166e 100644 --- a/src/omero_data_query_worker/executor.py +++ b/src/omero_data_query_worker/executor.py @@ -4,7 +4,9 @@ import multiprocessing import os import time +from collections.abc import Callable from multiprocessing.connection import Connection +from multiprocessing.process import BaseProcess from pathlib import Path from typing import Any @@ -12,6 +14,7 @@ from .engines import execute_duckdb_query, execute_sqlite_query from .errors import QueryExecutionError, QueryLimitExceeded, QueryTimedOut, WorkerError from .models import SourceFormat +from .operations import execution_guard, sanitize_execution_environment def apply_process_limits(settings: Settings, timeout_seconds: int) -> None: @@ -33,9 +36,12 @@ def _child_execute( parameters: dict[str, Any], output_path: str, settings: Settings, + parent_pid: int, ) -> None: + guard = None try: - os.environ.pop("DQW_API_TOKEN", None) + guard = execution_guard(parent_pid, settings.cache_dir) + sanitize_execution_environment() apply_process_limits(settings, settings.query_timeout_seconds) started = time.monotonic() if SourceFormat(source_format) is SourceFormat.sqlite: @@ -54,6 +60,8 @@ def _child_execute( child.send({"ok": False, "code": "query_execution_failed", "message": str(exc)}) finally: child.close() + if guard: + guard.close() def execute_in_subprocess( @@ -63,6 +71,8 @@ def execute_in_subprocess( parameters: dict[str, Any], output_path: Path, settings: Settings, + monitor: Callable[[], None] | None = None, + on_exit: Callable[[int | None], None] | None = None, ) -> dict[str, Any]: context = multiprocessing.get_context("spawn") parent, child = context.Pipe(duplex=False) @@ -75,31 +85,33 @@ def execute_in_subprocess( sql, parameters, str(output_path), - settings, + settings.execution_settings(), + os.getpid(), ), daemon=True, ) - process.start() - child.close() try: - if not parent.poll(settings.query_timeout_seconds): - process.terminate() - process.join(2) - if process.is_alive(): - process.kill() - process.join(2) - raise QueryTimedOut( - f"Query exceeded the {settings.query_timeout_seconds} second timeout" - ) + process.start() + child.close() + deadline = time.monotonic() + settings.query_timeout_seconds + while not parent.poll(0.1): + if monitor: + monitor() + if time.monotonic() >= deadline: + raise QueryTimedOut( + f"Query exceeded the {settings.query_timeout_seconds} second timeout" + ) payload = parent.recv() + if monitor: + monitor() except EOFError as exc: raise QueryExecutionError("The disposable query process exited unexpectedly") from exc finally: parent.close() - process.join(2) - if process.is_alive(): - process.terminate() - process.join(2) + child.close() + exit_code = stop_process(process) + if on_exit: + on_exit(exit_code) if payload.get("ok"): return dict(payload["result"]) if payload.get("code") == "query_limit_exceeded": @@ -107,3 +119,18 @@ def execute_in_subprocess( if payload.get("code") == "query_timeout": raise QueryTimedOut(str(payload.get("message") or "Query timed out")) raise QueryExecutionError(str(payload.get("message") or "Query execution failed")) + + +def stop_process(process: BaseProcess) -> int | None: + if process.pid is None: + return None + process.join(0.1) + if process.is_alive(): + process.terminate() + process.join(2) + if process.is_alive(): + process.kill() + process.join(2) + exit_code = process.exitcode + process.close() + return exit_code diff --git a/src/omero_data_query_worker/ingestion.py b/src/omero_data_query_worker/ingestion.py index c441fdb..94d577e 100644 --- a/src/omero_data_query_worker/ingestion.py +++ b/src/omero_data_query_worker/ingestion.py @@ -2,6 +2,8 @@ import multiprocessing import os +import time +from collections.abc import Callable from multiprocessing.connection import Connection from pathlib import Path from typing import Any @@ -9,8 +11,9 @@ from .config import Settings from .engines import convert_csv, inspect_source from .errors import IngestionTimedOut, InvalidSource, WorkerError -from .executor import apply_process_limits +from .executor import apply_process_limits, stop_process from .models import SourceFormat +from .operations import execution_guard, sanitize_execution_environment def _child_ingest( @@ -18,9 +21,12 @@ def _child_ingest( raw_path: str, source_format: str, settings: Settings, + parent_pid: int, ) -> None: + guard = None try: - os.environ.pop("DQW_API_TOKEN", None) + guard = execution_guard(parent_pid, settings.cache_dir) + sanitize_execution_environment() apply_process_limits(settings, settings.ingestion_timeout_seconds) raw = Path(raw_path) selected = SourceFormat(source_format) @@ -44,41 +50,53 @@ def _child_ingest( child.send({"ok": False, "message": str(exc)}) finally: child.close() + if guard: + guard.close() def ingest_in_subprocess( raw_path: Path, source_format: SourceFormat, settings: Settings, + monitor: Callable[[], None] | None = None, + on_exit: Callable[[int | None], None] | None = None, ) -> tuple[Path, dict[str, Any], str]: context = multiprocessing.get_context("spawn") parent, child = context.Pipe(duplex=False) process = context.Process( target=_child_ingest, - args=(child, str(raw_path), source_format.value, settings), + args=( + child, + str(raw_path), + source_format.value, + settings.execution_settings(), + os.getpid(), + ), daemon=True, ) - process.start() - child.close() try: - if not parent.poll(settings.ingestion_timeout_seconds): - process.terminate() - process.join(2) - if process.is_alive(): - process.kill() - process.join(2) - raise IngestionTimedOut( - f"Source ingestion exceeded the {settings.ingestion_timeout_seconds} second timeout" - ) + process.start() + child.close() + deadline = time.monotonic() + settings.ingestion_timeout_seconds + while not parent.poll(0.1): + if monitor: + monitor() + if time.monotonic() >= deadline: + raise IngestionTimedOut( + f"Source ingestion exceeded the {settings.ingestion_timeout_seconds} " + "second timeout" + ) payload = parent.recv() + if monitor: + monitor() except EOFError as exc: raise InvalidSource("The disposable ingestion process exited unexpectedly") from exc finally: parent.close() - process.join(2) - if process.is_alive(): - process.terminate() - process.join(2) + child.close() + exit_code = stop_process(process) + if on_exit: + on_exit(exit_code) if not payload.get("ok"): raise InvalidSource(str(payload.get("message") or "Source ingestion failed")) return ( diff --git a/src/omero_data_query_worker/models.py b/src/omero_data_query_worker/models.py index 4a65ec8..41ee32d 100644 --- a/src/omero_data_query_worker/models.py +++ b/src/omero_data_query_worker/models.py @@ -102,6 +102,15 @@ class QueryResponse(BaseModel): sql_sha256: str duration_ms: int cache_status: Literal["hit", "miss", "bypass"] + execution: dict[str, Any] | None = None + + +class PurgeRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + source_id: str | None = None + scope_id: str | None = Field(default=None, pattern=OPAQUE_PATTERN) + result_id: str | None = None + dry_run: bool = True class ErrorResponse(BaseModel): diff --git a/src/omero_data_query_worker/operations.py b/src/omero_data_query_worker/operations.py new file mode 100644 index 0000000..b7ac988 --- /dev/null +++ b/src/omero_data_query_worker/operations.py @@ -0,0 +1,167 @@ +"""Private operational controls; no SQL or data values in telemetry.""" + +from __future__ import annotations + +import importlib +import json +import os +import re +import secrets +import threading +import time +import uuid +from collections import Counter +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any, BinaryIO + +from .config import Settings +from .errors import AuthenticationError + + +def request_id(value: str | None) -> str: + return value if value and re.fullmatch(r"[A-Za-z0-9._-]{1,80}", value) else uuid.uuid4().hex + + +class Credentials: + def __init__(self, settings: Settings): + self.settings = settings + + def tokens(self) -> list[str]: + tokens = [self.settings.api_token] if self.settings.api_token else [] + if self.settings.token_keyring_file: + try: + # Atomic replacement by the operator makes each read a complete revision. + path = Path(self.settings.token_keyring_file) + if path.stat().st_size > 65536: + raise ValueError + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict) or set(value) != {"tokens"}: + raise ValueError + ring = value["tokens"] + if not isinstance(ring, list) or not 1 <= len(ring) <= 16: + raise ValueError + if any(not isinstance(t, str) or not 16 <= len(t) <= 512 for t in ring): + raise ValueError + tokens.extend(ring) + except (OSError, ValueError, TypeError) as exc: + raise AuthenticationError("Worker credential configuration is unavailable") from exc + return tokens + + def authenticate(self, authorization: str | None) -> None: + scheme, _, value = (authorization or "").partition(" ") + matches = [secrets.compare_digest(value.encode(), t.encode()) for t in self.tokens()] + if scheme.lower() != "bearer" or not any(matches): + raise AuthenticationError("A valid worker bearer token is required") + + +def sanitize_execution_environment() -> None: + # Preserve only OS/runtime essentials. In particular no OMERO/session/cloud credentials. + allowed = { + "PATH", + "SYSTEMROOT", + "WINDIR", + "TEMP", + "TMP", + "TMPDIR", + "LANG", + "LC_ALL", + "LD_LIBRARY_PATH", + "PYTHONPATH", + "PYTHONHOME", + "PYTHONDONTWRITEBYTECODE", + } + for name in list(os.environ): + if name.upper() not in allowed: + del os.environ[name] + + +def execution_guard(parent_pid: int, directory: Path) -> BinaryIO | None: + """Kill Linux engines with their service; fence startup against surviving children.""" + import sys + + if sys.platform != "linux": + return None + ctypes = importlib.import_module("ctypes") + signal = importlib.import_module("signal") + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(1, signal.SIGKILL, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "Cannot install engine parent-death guard") + if os.getppid() != parent_pid: + os._exit(125) + directory.mkdir(parents=True, exist_ok=True) + handle = (directory / ".engines.lock").open("a+b") + locking = importlib.import_module("fcntl") + locking.flock(handle.fileno(), locking.LOCK_SH) + return handle + + +@contextmanager +def recovery_lock(directory: Path) -> Iterator[None]: + if os.name != "posix": + yield + return + locking = importlib.import_module("fcntl") + with (directory / ".engines.lock").open("a+b") as handle: + deadline = time.monotonic() + 5 + while True: + try: + locking.flock(handle.fileno(), locking.LOCK_EX | locking.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() >= deadline: + raise RuntimeError( + "Previous engines still own this cache; retry startup" + ) from None + time.sleep(0.05) + try: + yield + finally: + locking.flock(handle.fileno(), locking.LOCK_UN) + + +@contextmanager +def volume_lock(directory: Path) -> Iterator[None]: + """Prevent two service processes from independently managing the same cache.""" + directory.mkdir(parents=True, exist_ok=True) + with (directory / ".worker.lock").open("a+b") as handle: + handle.seek(0) + if os.name == "nt": + locking: Any = importlib.import_module("msvcrt") + handle.write(b"0") + handle.flush() + handle.seek(0) + try: + locking.locking(handle.fileno(), locking.LK_NBLCK, 1) + except OSError as exc: + raise RuntimeError("Another worker owns this cache volume") from exc + else: + locking = importlib.import_module("fcntl") + try: + locking.flock(handle.fileno(), locking.LOCK_EX | locking.LOCK_NB) + except OSError as exc: + raise RuntimeError("Another worker owns this cache volume") from exc + try: + yield + finally: + if os.name == "nt": + handle.seek(0) + locking.locking(handle.fileno(), locking.LK_UNLCK, 1) + else: + locking.flock(handle.fileno(), locking.LOCK_UN) + + +class Metrics: + def __init__(self) -> None: + self._lock = threading.Lock() + self.values: Counter[str] = Counter() + + def add(self, name: str, value: float = 1) -> None: + with self._lock: + self.values[name] += value # type: ignore[assignment] + + def render(self, gauges: dict[str, int]) -> str: + with self._lock: + values = {**self.values, **gauges} + return "".join(f"dqw_{name} {value}\n" for name, value in sorted(values.items())) diff --git a/src/omero_data_query_worker/policy.py b/src/omero_data_query_worker/policy.py index 9b0ab17..13a3328 100644 --- a/src/omero_data_query_worker/policy.py +++ b/src/omero_data_query_worker/policy.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from sqlglot import exp, parse -from sqlglot.errors import ParseError +from sqlglot.errors import ParseError, TokenError from .errors import InvalidQuery from .models import SourceFormat @@ -105,7 +105,7 @@ def validate_query( raise InvalidQuery("SQL must not contain NUL bytes") try: statements = parse(sql, read=_dialect(source_format)) - except ParseError as exc: + except (ParseError, TokenError) as exc: raise InvalidQuery(f"SQL could not be parsed: {exc}") from exc if len(statements) != 1 or statements[0] is None: raise InvalidQuery("Exactly one SQL statement is required") diff --git a/src/omero_data_query_worker/service.py b/src/omero_data_query_worker/service.py index eba04cd..b511dc5 100644 --- a/src/omero_data_query_worker/service.py +++ b/src/omero_data_query_worker/service.py @@ -2,24 +2,29 @@ import hashlib import json -import shutil import threading +import time as clock import uuid +import weakref +from collections.abc import Iterator +from contextlib import contextmanager from datetime import date, datetime, time from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Any from fastapi import UploadFile +from pydantic import ValidationError from . import __version__ -from .cache import CacheManager, ResultRecord, SourceRecord, iso, utc_now +from .cache import CacheManager, ResultRecord, SourceRecord, iso, sync_directory, utc_now from .config import Settings -from .engines import validate_filename +from .engines import engine_versions, validate_filename from .errors import ( InvalidQuery, InvalidSource, QueryCapacityExceeded, + QueryExecutionError, SourceConflict, SourceNotFound, SourceTooLarge, @@ -37,6 +42,7 @@ SourceSummary, TypedParameter, ) +from .operations import Metrics from .policy import POLICY_VERSION, ValidatedQuery, validate_query @@ -45,17 +51,37 @@ def __init__(self, settings: Settings): self.settings = settings self.cache = CacheManager(settings) self._query_capacity = threading.BoundedSemaphore(settings.max_concurrent_queries) - self._source_locks: dict[str, threading.Lock] = {} - self._query_locks: dict[str, threading.Lock] = {} + self._source_locks: weakref.WeakValueDictionary[str, threading.Lock] = ( + weakref.WeakValueDictionary() + ) + self._query_locks: weakref.WeakValueDictionary[str, threading.Lock] = ( + weakref.WeakValueDictionary() + ) + self._ingest_capacity = threading.BoundedSemaphore(settings.max_concurrent_ingestions) + self.stopping = threading.Event() + self.metrics = Metrics() self._locks_guard = threading.Lock() def prepare(self) -> None: + missing = [] + parent = self.settings.cache_dir + while not parent.exists(): + missing.append(parent) + parent = parent.parent self.settings.prepare() - self.cache.cleanup_temporary() + for created in reversed(missing): + sync_directory(created.parent) + sync_directory(self.settings.cache_dir) + from .operations import recovery_lock + + with recovery_lock(self.settings.cache_dir): + self.cache.recover() self.cache.cleanup_sources() self.cache.cleanup_results() - def _named_lock(self, collection: dict[str, threading.Lock], key: str) -> threading.Lock: + def _named_lock( + self, collection: weakref.WeakValueDictionary[str, threading.Lock], key: str + ) -> threading.Lock: with self._locks_guard: return collection.setdefault(key, threading.Lock()) @@ -84,7 +110,7 @@ def resolve_source(self, request: SourceResolveRequest) -> SourceResolveResponse return SourceResolveResponse(cached=False) return SourceResolveResponse(cached=True, source=self._summary(record, "reused")) - async def ingest_source( + def ingest_source( self, *, scope_id: str, @@ -93,18 +119,50 @@ async def ingest_source( declared_size: int, expected_sha256: str | None, upload: UploadFile, + cancelled: threading.Event | None = None, ) -> SourceSummary: - source_request = SourceResolveRequest( - scope_id=scope_id, - source_ref=source_ref, - format=source_format, - size=declared_size, - expected_sha256=expected_sha256, - ) + started = clock.monotonic() + try: + with self._admit(self._ingest_capacity, "ingestion"): + return self._ingest_source( + scope_id=scope_id, + source_ref=source_ref, + source_format=source_format, + declared_size=declared_size, + expected_sha256=expected_sha256, + upload=upload, + cancelled=cancelled, + ) + finally: + upload.file.close() + self.metrics.add("ingestion_requests_total") + self.metrics.add("ingestion_duration_seconds_sum", clock.monotonic() - started) + + def _ingest_source( + self, + *, + scope_id: str, + source_ref: str, + source_format: SourceFormat, + declared_size: int, + expected_sha256: str | None, + upload: UploadFile, + cancelled: threading.Event | None = None, + ) -> SourceSummary: + try: + source_request = SourceResolveRequest( + scope_id=scope_id, + source_ref=source_ref, + format=source_format, + size=declared_size, + expected_sha256=expected_sha256, + ) + except ValidationError as exc: + raise InvalidSource("Invalid source reference, size, or checksum metadata") from exc filename = validate_filename(upload.filename or "", source_format) source_id = self.cache.source_id(scope_id, source_ref) lock = self._named_lock(self._source_locks, source_id) - with lock: + with self._try_lock(lock): try: existing = self.cache.get_source(source_id) except SourceNotFound: @@ -124,6 +182,8 @@ async def ingest_source( ) staging = self.cache.staging_path("source") try: + self._monitor(staging, cancelled) + self.cache.reserve(staging, declared_size) suffix = { SourceFormat.duckdb: ".duckdb", SourceFormat.sqlite: ".sqlite", @@ -133,12 +193,14 @@ async def ingest_source( digest = hashlib.sha256() actual_size = 0 with raw_path.open("wb") as output: - while chunk := await upload.read(1024 * 1024): + while chunk := upload.file.read(1024 * 1024): + self._check_cancelled(cancelled) actual_size += len(chunk) if actual_size > self.settings.max_source_bytes: raise SourceTooLarge( f"Source exceeds the {self.settings.max_source_bytes} byte limit" ) + self.cache.reserve(staging, actual_size) digest.update(chunk) output.write(chunk) output.flush() @@ -151,7 +213,11 @@ async def ingest_source( raise InvalidSource("Uploaded content does not match expected_sha256") data_path, schema, schema_digest = ingest_in_subprocess( - raw_path, source_format, self.settings + raw_path, + source_format, + self.settings, + monitor=lambda: self._monitor(staging, cancelled), + on_exit=lambda code: self._process_exit("ingestion", code), ) now = utc_now() from datetime import timedelta @@ -173,12 +239,8 @@ async def ingest_source( ) committed = self.cache.commit_source(staging, record) return self._summary(committed, "created") - except BaseException: - if staging.exists(): - shutil.rmtree(staging) - raise finally: - await upload.close() + self.cache.discard_staging(staging) def schema(self, source_id: str) -> SchemaResponse: record = self.cache.get_source(source_id) @@ -253,6 +315,10 @@ def _query_key( "parameters": parameters, "worker_version": __version__, "policy_version": POLICY_VERSION, + "engines": engine_versions(), + "timeout": self.settings.query_timeout_seconds, + "memory": self.settings.duckdb_memory_limit, + "threads": self.settings.duckdb_threads, "max_rows": self.settings.max_result_rows, "max_bytes": self.settings.max_result_bytes, } @@ -271,10 +337,27 @@ def _response(record: ResultRecord, cache_status: str) -> QueryResponse: sql_sha256=record.sql_sha256, duration_ms=record.duration_ms, cache_status=cache_status, # type: ignore[arg-type] + execution=record.execution, ) - def query(self, source_id: str, request: QueryRequest) -> QueryResponse: - source = self.cache.get_source(source_id) + def query( + self, source_id: str, request: QueryRequest, cancelled: threading.Event | None = None + ) -> QueryResponse: + started = clock.monotonic() + try: + with self.cache.lease_source(source_id) as source: + return self._query(source, request, cancelled) + except BaseException: + self.metrics.add("query_failures_total") + raise + finally: + self.metrics.add("query_requests_total") + self.metrics.add("query_duration_seconds_sum", clock.monotonic() - started) + + def _query( + self, source: SourceRecord, request: QueryRequest, cancelled: threading.Event | None + ) -> QueryResponse: + self._check_cancelled(cancelled) validated = validate_query(request.sql, source.format, set(request.parameters)) converted_parameters = { name: self._convert_parameter(parameter) @@ -293,7 +376,7 @@ def query(self, source_id: str, request: QueryRequest) -> QueryResponse: } query_key = self._query_key(source, validated, request) if validated.deterministic: - cached = self.cache.get_result_by_query_key(query_key) + cached = self.cache.get_result_by_query_key(query_key, require_provenance=True) if cached is not None: self.cache.hits += 1 return self._response(cached, "hit") @@ -302,9 +385,9 @@ def query(self, source_id: str, request: QueryRequest) -> QueryResponse: self.cache.bypasses += 1 query_lock = threading.Lock() - with query_lock: + with self._try_lock(query_lock): if validated.deterministic: - cached = self.cache.get_result_by_query_key(query_key) + cached = self.cache.get_result_by_query_key(query_key, require_provenance=True) if cached is not None: self.cache.hits += 1 return self._response(cached, "hit") @@ -317,50 +400,129 @@ def query(self, source_id: str, request: QueryRequest) -> QueryResponse: result_id = f"res_{uuid.uuid4().hex}" stored_query_key = None - if not self._query_capacity.acquire(blocking=False): - raise QueryCapacityExceeded("All disposable query slots are busy") - self.cache.active_queries += 1 - staging = self.cache.staging_path("result") - try: - source_path = self.cache.source_path(source.source_id) / source.data_file - output_path = staging / "result.csv" - executed = execute_in_subprocess( - source_path, - source.format, - validated.original_sql, + with self._admit(self._query_capacity, "query"): + return self._execute( + source, + validated, converted_parameters, - output_path, - self.settings, + result_id, + stored_query_key, + cache_status, + cancelled, ) - now = utc_now() - from datetime import timedelta - record = ResultRecord( - result_id=result_id, - query_key=stored_query_key, - source_id=source.source_id, - source_sha256=source.sha256, - sql_sha256=validated.sql_sha256, - columns=executed["columns"], - row_count=int(executed["row_count"]), - byte_count=int(executed["byte_count"]), - preview=executed["preview"], - duration_ms=int(executed["duration_ms"]), - result_file="result.csv", - created_at=iso(now), - accessed_at=iso(now), - expires_at=iso(now + timedelta(seconds=self.settings.result_ttl_seconds)), - ) - committed = self.cache.commit_result(staging, record) - return self._response(committed, cache_status) - except BaseException: - if staging.exists(): - shutil.rmtree(staging) - raise - finally: - self.cache.active_queries -= 1 - self._query_capacity.release() + def _execute( + self, + source: SourceRecord, + validated: ValidatedQuery, + converted_parameters: dict[str, Any], + result_id: str, + stored_query_key: str | None, + cache_status: str, + cancelled: threading.Event | None, + ) -> QueryResponse: + staging = self.cache.staging_path("result") + try: + source_path = self.cache.source_path(source.source_id) / source.data_file + output_path = staging / "result.csv" + executed = execute_in_subprocess( + source_path, + source.format, + validated.original_sql, + converted_parameters, + output_path, + self.settings, + monitor=lambda: self._monitor(staging, cancelled), + on_exit=lambda code: self._process_exit("query", code), + ) + now = utc_now() + from datetime import timedelta + + record = ResultRecord( + result_id=result_id, + query_key=stored_query_key, + source_id=source.source_id, + source_sha256=source.sha256, + sql_sha256=validated.sql_sha256, + columns=executed["columns"], + row_count=int(executed["row_count"]), + byte_count=int(executed["byte_count"]), + preview=executed["preview"], + duration_ms=int(executed["duration_ms"]), + result_file="result.csv", + execution={ + "schema": "result_provenance_v1", + "result_sha256": executed["result_sha256"], + "schema_digest": source.schema_digest, + "engine": "sqlite" if source.format is SourceFormat.sqlite else "duckdb", + "versions": { + **engine_versions(), + "worker": __version__, + "policy": POLICY_VERSION, + }, + "limits": { + "max_result_rows": self.settings.max_result_rows, + "max_result_bytes": self.settings.max_result_bytes, + "query_timeout_seconds": self.settings.query_timeout_seconds, + "duckdb_memory_limit": self.settings.duckdb_memory_limit, + "duckdb_threads": self.settings.duckdb_threads, + }, + "executed_at": iso(now - timedelta(milliseconds=int(executed["duration_ms"]))), + "completed_at": iso(now), + }, + created_at=iso(now), + accessed_at=iso(now), + expires_at=iso(now + timedelta(seconds=self.settings.result_ttl_seconds)), + ) + committed = self.cache.commit_result(staging, record) + return self._response(committed, cache_status) + finally: + self.cache.discard_staging(staging) def result_file(self, result_id: str) -> tuple[Path, ResultRecord]: record = self.cache.get_result(result_id) return self.cache.result_path(result_id) / record.result_file, record + + @contextmanager + def _try_lock(self, lock: threading.Lock) -> Iterator[None]: + if not lock.acquire(blocking=False): + self.metrics.add("admission_rejections_total") + raise QueryCapacityExceeded("Identical work is already running; retry later") + try: + yield + finally: + lock.release() + + @contextmanager + def _admit(self, capacity: threading.BoundedSemaphore, kind: str) -> Iterator[None]: + self._check_cancelled(None) + if not capacity.acquire(blocking=False): + self.metrics.add("admission_rejections_total") + raise QueryCapacityExceeded(f"All {kind} slots are busy; retry later") + self.metrics.add(f"active_{kind}s", 1) + if kind == "query": + with self.cache._lock: + self.cache.active_queries += 1 + try: + yield + finally: + if kind == "query": + with self.cache._lock: + self.cache.active_queries -= 1 + self.metrics.add(f"active_{kind}s", -1) + capacity.release() + + def _check_cancelled(self, cancelled: threading.Event | None) -> None: + if self.stopping.is_set() or (cancelled is not None and cancelled.is_set()): + raise QueryExecutionError( + "Operation interrupted; client disconnected or worker stopping" + ) + + def _monitor(self, staging: Path, cancelled: threading.Event | None) -> None: + self._check_cancelled(cancelled) + self.cache.observe_staging(staging) + + def _process_exit(self, kind: str, code: int | None) -> None: + self.metrics.add(f"subprocess_{kind}_exits_total") + if code != 0: + self.metrics.add(f"subprocess_{kind}_abnormal_exits_total") diff --git a/src/omero_data_query_worker/upload.py b/src/omero_data_query_worker/upload.py new file mode 100644 index 0000000..f0cc084 --- /dev/null +++ b/src/omero_data_query_worker/upload.py @@ -0,0 +1,83 @@ +"""Admission and spool accounting before Starlette parses multipart uploads.""" + +from __future__ import annotations + +import threading + +from starlette.concurrency import run_in_threadpool +from starlette.datastructures import Headers +from starlette.formparsers import MultiPartException +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from .errors import QueryCapacityExceeded, SourceTooLarge, WorkerError +from .operations import Credentials, request_id +from .service import QueryService + + +class UploadAdmission: + def __init__(self, app: ASGIApp, service: QueryService, credentials: Credentials): + self.app, self.service, self.credentials = app, service, credentials + self.capacity = threading.BoundedSemaphore(service.settings.max_concurrent_ingestions) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http" or scope["path"] != "/v1/sources" or scope["method"] != "POST": + await self.app(scope, receive, send) + return + headers = Headers(scope=scope) + correlation = request_id(headers.get("X-Request-ID")) + + async def reject(exc: WorkerError) -> None: + self.service.metrics.add(f"errors_{exc.code}_total") + response = JSONResponse( + {"code": exc.code, "message": str(exc), "request_id": correlation}, + status_code=exc.status_code, + headers={"X-Request-ID": correlation}, + ) + await response(scope, receive, send) + + try: + self.credentials.authenticate(headers.get("Authorization")) + if not self.capacity.acquire(blocking=False): + self.service.metrics.add("admission_rejections_total") + raise QueryCapacityExceeded("All upload slots are busy; retry later") + except WorkerError as exc: + await reject(exc) + return + staging = None + failure: WorkerError | None = None + sent_failure = False + received = 0 + try: + staging = await run_in_threadpool(self.service.cache.staging_path, "source") + + async def bounded_receive() -> Message: + nonlocal received, failure + message = await receive() + if message["type"] == "http.request": + received += len(message.get("body", b"")) + try: + # Multipart fields/boundaries have a separate bounded allowance. + if received > self.service.settings.max_source_bytes + 1024 * 1024: + raise SourceTooLarge("Upload exceeds the source size limit") + await run_in_threadpool(self.service.cache.reserve, staging, received) + except WorkerError as exc: + failure = exc + # Starlette closes every partially parsed upload on this exception. + raise MultiPartException("Upload capacity exceeded") from exc + return message + + async def bounded_send(message: Message) -> None: + nonlocal sent_failure + if failure: + if not sent_failure: + sent_failure = True + await reject(failure) + else: + await send(message) + + await self.app(scope, bounded_receive, bounded_send) + finally: + if staging is not None: + await run_in_threadpool(self.service.cache.discard_staging, staging) + self.capacity.release() diff --git a/tests/fixtures/compatibility/analysis_broker_v013.txt b/tests/fixtures/compatibility/analysis_broker_v013.txt new file mode 100644 index 0000000..79dd1a0 --- /dev/null +++ b/tests/fixtures/compatibility/analysis_broker_v013.txt @@ -0,0 +1,345 @@ +"""OMERO-authorized broker for the OMERO-agnostic DataQueryWorker.""" + +from __future__ import annotations + +import base64 +import hashlib +import hmac +import json +from pathlib import Path +from typing import Any + +import requests +from cryptography.fernet import Fernet, InvalidToken as InvalidFernetToken +from django.conf import settings as django_settings +from requests_toolbelt import MultipartEncoder + +from .errors import InvalidToken, RemoteQueryFailed, RemoteQueryUnavailable, UnsupportedMedia +from .services import object_group_id +from .settings import ( + data_query_request_timeout_seconds, + data_query_result_ttl_seconds, + data_query_source_upload_timeout_seconds, + data_query_worker_token, + data_query_worker_url, + remote_query_threshold_bytes, +) + +SUPPORTED_FORMATS = { + ".duckdb": "duckdb", + ".sqlite": "sqlite", + ".sqlite3": "sqlite", + ".csv": "csv", +} +CAPABILITY = "omero-data-query-v1" + + +class ChunkReader: + """Bounded file-like adapter over the OMERO chunk iterator.""" + + def __init__(self, chunks: Any, size: int) -> None: + self._chunks = iter(chunks) + self._buffer = b"" + self._remaining = size + + @property + def len(self) -> int: + return self._remaining + + def read(self, size: int = -1) -> bytes: + if self._remaining <= 0: + return b"" + target = self._remaining if size is None or size < 0 else min(size, self._remaining) + while len(self._buffer) < target: + try: + self._buffer += bytes(next(self._chunks)) + except StopIteration: + break + value, self._buffer = self._buffer[:target], self._buffer[target:] + self._remaining -= len(value) + return value + + +def source_format(name: str) -> str: + value = SUPPORTED_FORMATS.get(Path(name).suffix.lower()) + if value is None: + raise UnsupportedMedia( + "Remote queries support DuckDB, SQLite, and CSV attachments" + ) + return value + + +def query_policy(info: Any, worker_ready: bool) -> dict[str, Any]: + try: + format_name = source_format(info.name) + except UnsupportedMedia: + return {} + threshold = remote_query_threshold_bytes() + forced = threshold == 0 + default_mode = "remote" if forced or info.size >= threshold else "local" + allowed = ["remote"] if forced else ["local", "remote"] + if not worker_ready: + allowed = [mode for mode in allowed if mode != "remote"] + reason = ( + "remote-required-by-policy" + if forced + else "size-at-or-above-threshold" + if info.size >= threshold + else "below-threshold" + ) + return { + "query_format": format_name, + "default_mode": default_mode, + "allowed_modes": allowed, + "threshold_bytes": threshold, + "threshold_reason": reason, + "worker_ready": worker_ready, + } + + +class DataQueryBroker: + def __init__(self, session: requests.Session | None = None) -> None: + self.url = data_query_worker_url() + self.token = data_query_worker_token() + self.session = session or requests.Session() + + @property + def configured(self) -> bool: + return bool(self.url and self.token) + + def capabilities(self) -> dict[str, Any]: + threshold = remote_query_threshold_bytes() + unavailable = { + "available": bool(self.configured), + "ready": False, + "capability": CAPABILITY, + "formats": ["duckdb", "sqlite", "csv"], + "threshold_bytes": threshold, + "result_ttl_seconds": data_query_result_ttl_seconds(), + } + if not self.configured: + return unavailable + try: + health = self.session.get(f"{self.url}/health/ready", timeout=3) + health.raise_for_status() + payload = self._request("GET", "/v1/capabilities", timeout=5) + except (requests.RequestException, RemoteQueryFailed, RemoteQueryUnavailable): + return unavailable + return { + "available": True, + "ready": True, + "capability": CAPABILITY, + "formats": payload.get("formats", []), + "limits": payload.get("limits", {}), + "parameter_style": payload.get("parameter_style", "$name"), + "csv_table": payload.get("csv_table", "data"), + "threshold_bytes": threshold, + "result_ttl_seconds": data_query_result_ttl_seconds(), + } + + def schema( + self, annotation: Any, info: Any, scope: str, source_ref: str + ) -> dict[str, Any]: + source = self._source(annotation, info, scope, source_ref) + payload = self._request( + "GET", f"/v1/sources/{source['source_id']}/schema" + ) + return {key: value for key, value in payload.items() if key != "source_id"} + + def query( + self, + annotation: Any, + info: Any, + scope: str, + source_ref: str, + payload: dict[str, Any], + ) -> tuple[dict[str, Any], str]: + if set(payload) - {"sql", "parameters"}: + raise RemoteQueryFailed("Remote query request contains unsupported fields") + source = self._source(annotation, info, scope, source_ref) + result = self._request( + "POST", + f"/v1/sources/{source['source_id']}/query", + json=payload, + ) + result_id = str(result.pop("result_id")) + return result, result_id + + def download(self, result_id: str) -> requests.Response: + return self._stream("GET", f"/v1/results/{result_id}/download") + + def _source( + self, annotation: Any, info: Any, scope: str, source_ref: str + ) -> dict[str, Any]: + format_name = source_format(info.name) + resolve = { + "scope_id": scope, + "source_ref": source_ref, + "format": format_name, + "size": info.size, + } + cached = self._request("POST", "/v1/sources/resolve", json=resolve) + if cached.get("cached") and isinstance(cached.get("source"), dict): + return cached["source"] + reader = ChunkReader(annotation.getFileInChunks(), info.size) + encoder = MultipartEncoder( + fields={ + "scope_id": scope, + "source_ref": source_ref, + "format": format_name, + "size": str(info.size), + "file": (info.name, reader, info.mimetype), + } + ) + return self._request( + "POST", + "/v1/sources", + data=encoder, + headers={"Content-Type": encoder.content_type}, + timeout=data_query_source_upload_timeout_seconds(), + ) + + def _request(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: + response = self._stream(method, path, **kwargs) + try: + value = response.json() + except ValueError as exc: + raise RemoteQueryFailed( + "The data query worker returned an invalid response" + ) from exc + if not isinstance(value, dict): + raise RemoteQueryFailed("The data query worker returned an invalid response") + return value + + def _stream(self, method: str, path: str, **kwargs: Any) -> requests.Response: + if not self.configured: + raise RemoteQueryUnavailable( + "The remote data query worker is not configured" + ) + headers = { + "Authorization": f"Bearer {self.token}", + **kwargs.pop("headers", {}), + } + try: + response = self.session.request( + method, + f"{self.url}{path}", + headers=headers, + timeout=kwargs.pop("timeout", data_query_request_timeout_seconds()), + stream=True, + **kwargs, + ) + except requests.RequestException as exc: + raise RemoteQueryUnavailable( + "The remote data query worker is unavailable" + ) from exc + if not response.ok: + try: + error = response.json() + message = str(error.get("message") or error.get("detail") or "") + except (ValueError, AttributeError): + message = "" + raise RemoteQueryFailed(message or "Remote query failed") + return response + + +def opaque_references( + request: Any, conn: Any, claims: dict[str, Any], info: Any +) -> tuple[str, str]: + key = hashlib.sha256( + (str(django_settings.SECRET_KEY) + "\0data-query-refs-v1").encode() + ).digest() + session_key = str( + getattr(getattr(request, "session", None), "session_key", "") or "" + ) + scope_material = ( + f"{conn.getUserId()}:{claims['group_id']}:{session_key}:" + f"{claims['object_type']}:{claims['object_id']}" + ) + source_material = ( + f"{scope_material}:{info.annotation_id}:{info.file_id}:" + f"{info.size}:{info.name}" + ) + return _opaque(key, scope_material), _opaque(key, source_material) + + +def _opaque(key: bytes, value: str) -> str: + return ( + base64.urlsafe_b64encode(hmac.new(key, value.encode(), hashlib.sha256).digest()) + .decode() + .rstrip("=") + ) + + +def make_result_token( + request: Any, + conn: Any, + claims: dict[str, Any], + info: Any, + result_id: str, +) -> str: + payload = { + "v": 1, + "user": int(conn.getUserId()), + "group": claims["group_id"], + "session": str(getattr(request.session, "session_key", "") or ""), + "object_type": claims["object_type"], + "object_id": claims["object_id"], + "annotation_id": info.annotation_id, + "file_id": info.file_id, + "result_id": result_id, + } + return _fernet().encrypt( + json.dumps(payload, separators=(",", ":")).encode() + ).decode() + + +def validate_result_token( + request: Any, conn: Any, token: str, obj: Any, info: Any +) -> dict[str, Any]: + try: + payload = json.loads( + _fernet().decrypt( + token.encode(), ttl=data_query_result_ttl_seconds() + ) + ) + except (InvalidFernetToken, ValueError, json.JSONDecodeError) as exc: + raise InvalidToken( + "The data query result token is invalid or expired" + ) from exc + expected = { + "v": 1, + "user": int(conn.getUserId()), + "group": object_group_id(obj), + "session": str(getattr(request.session, "session_key", "") or ""), + "annotation_id": info.annotation_id, + "file_id": info.file_id, + } + if any(payload.get(key) != value for key, value in expected.items()): + raise InvalidToken( + "The data query result token is not valid for this context" + ) + return payload + + +def result_token_claims(token: str) -> dict[str, Any]: + try: + value = json.loads( + _fernet().decrypt( + token.encode(), ttl=data_query_result_ttl_seconds() + ) + ) + except (InvalidFernetToken, ValueError, json.JSONDecodeError) as exc: + raise InvalidToken( + "The data query result token is invalid or expired" + ) from exc + if not isinstance(value, dict): + raise InvalidToken("The data query result token is invalid") + return value + + +def _fernet() -> Fernet: + digest = hashlib.sha256( + (str(django_settings.SECRET_KEY) + "\0data-query-results-v1").encode() + ).digest() + return Fernet(base64.urlsafe_b64encode(digest)) diff --git a/tests/fixtures/compatibility/v1.json b/tests/fixtures/compatibility/v1.json new file mode 100644 index 0000000..ea093d9 --- /dev/null +++ b/tests/fixtures/compatibility/v1.json @@ -0,0 +1,27 @@ +{ + "worker_commit": "2115594b39108fb55605e4958e11bebf497ef282", + "analysis_commit": "4db7f8b25f5562ed3045fd116309d2b5db3d6440", + "worker_protocol": "omero-data-query-worker-v1", + "analysis_capability": "omero-data-query-v1", + "formats": [ + "duckdb", + "sqlite", + "csv" + ], + "query_fields": [ + "sql", + "parameters" + ], + "result_fields": [ + "result_id", + "columns", + "row_count", + "byte_count", + "preview", + "source_sha256", + "sql_sha256", + "duration_ms", + "cache_status" + ], + "csv": "id,area\n1,10.5\n2,20.5\n" +} diff --git a/tests/test_api.py b/tests/test_api.py index 0033524..f9da3f8 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -26,6 +26,20 @@ def test_capabilities(client: TestClient, auth: dict[str, str]) -> None: assert response.json()["csv_table"] == "data" +def test_invalid_upload_metadata_is_a_client_error(client, auth): + response = upload_source( + client, + auth, + source_ref="invalid:reference", + source_format="csv", + filename="source.csv", + content=b"id\n1\n", + ) + assert response.status_code == 422 + assert response.json()["code"] == "invalid_source" + assert client.get("/health/ready").status_code == 200 + + def test_duckdb_ingest_schema_query_cache_and_download( client: TestClient, auth: dict[str, str], diff --git a/tests/test_cache.py b/tests/test_cache.py index 55ad119..645f216 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -59,7 +59,8 @@ def test_incomplete_atomic_entry_is_cleaned(settings: Settings) -> None: def test_source_cache_evicts_least_recently_used(settings: Settings) -> None: - limited = replace(settings, source_cache_max_bytes=1000) + # Each entry includes ~500 bytes of manifest as well as 700 bytes of data. + limited = replace(settings, source_cache_max_bytes=1500) limited.prepare() cache = CacheManager(limited) diff --git a/tests/test_foundations.py b/tests/test_foundations.py new file mode 100644 index 0000000..aba940b --- /dev/null +++ b/tests/test_foundations.py @@ -0,0 +1,310 @@ +from __future__ import annotations + +import asyncio +import hashlib +import io +import json +import threading +from dataclasses import replace + +import pytest +from fastapi import UploadFile +from fastapi.testclient import TestClient + +from omero_data_query_worker.app import LeasedFileResponse, create_app +from omero_data_query_worker.cache import CacheManager, ResultRecord +from omero_data_query_worker.engines import _write_result, engine_versions +from omero_data_query_worker.errors import ( + AuthenticationError, + CacheCapacityExceeded, + QueryLimitExceeded, +) +from omero_data_query_worker.models import QueryRequest, SourceFormat +from omero_data_query_worker.operations import ( + Credentials, + sanitize_execution_environment, + volume_lock, +) +from omero_data_query_worker.service import QueryService + +from .conftest import upload_source + + +def test_provenance_survives_hit_and_legacy_manifest(settings, auth): + with TestClient(create_app(settings)) as client: + source = upload_source( + client, + auth, + source_ref="receipt", + source_format="csv", + filename="test.csv", + content=b"id,label\n1,alpha\n", + ).json() + endpoint = f"/v1/sources/{source['source_id']}/query" + payload = {"sql": "SELECT * FROM data", "parameters": {}} + first = client.post(endpoint, headers=auth, json=payload).json() + csv = client.get(f"/v1/results/{first['result_id']}/download", headers=auth).content + assert first["execution"]["result_sha256"] == hashlib.sha256(csv).hexdigest() + assert first["execution"]["versions"]["duckdb"] == engine_versions()["duckdb"] + second = client.post(endpoint, headers=auth, json=payload).json() + assert second["cache_status"] == "hit" + assert second["execution"] == first["execution"] + manifest = settings.results_dir / first["result_id"] / "manifest.json" + old = json.loads(manifest.read_text()) + old.pop("execution") + assert ResultRecord.from_dict(old).execution is None + manifest.write_text(json.dumps(old)) + recomputed = client.post(endpoint, headers=auth, json=payload).json() + assert recomputed["cache_status"] == "miss" + assert recomputed["execution"]["result_sha256"] == hashlib.sha256(csv).hexdigest() + + +@pytest.mark.parametrize("rows,limit", [([], 1), ([("a" * 5000,)], 100)]) +def test_byte_limit_includes_header_and_single_value(tmp_path, rows, limit): + class Cursor: + description = [("column", "VARCHAR")] + + def fetchmany(self, size): + nonlocal rows + result, rows = rows, [] + return result + + output = tmp_path / "bounded.csv" + with pytest.raises(QueryLimitExceeded): + _write_result(Cursor(), output, 100, limit, False) + assert output.stat().st_size <= limit + + +def test_exact_utf8_csv_limit_and_digest(tmp_path): + class Cursor: + description = [("label", "VARCHAR")] + calls = 0 + + def fetchmany(self, size): + self.calls += 1 + return [('é,"x"',)] if self.calls == 1 else [] + + expected = 'label\n"é,""x"""\n'.encode() + result = _write_result(Cursor(), tmp_path / "out", 1, len(expected), False) + assert (tmp_path / "out").read_bytes() == expected + assert result["byte_count"] == len(expected) + + +def test_leases_prevent_eviction_purge_and_expiry(client, settings, auth): + source = upload_source( + client, + auth, + source_ref="lease", + source_format="csv", + filename="test.csv", + content=b"id\n1\n", + ).json() + service = client.app.state.query_service + cache = service.cache + with cache.lease_source(source["source_id"]): + cache._cleanup(settings.sources_dir, 0, set()) + assert cache.source_path(source["source_id"]).exists() + assert cache.purge(source_id=source["source_id"], dry_run=False)["busy"] + result = service.query(source["source_id"], QueryRequest(sql="SELECT * FROM data")) + cache.acquire_result(result.result_id) + cache._cleanup(settings.results_dir, 0, set()) + assert cache.result_path(result.result_id).exists() + cache.release(result.result_id) + dry = cache.purge(scope_id="test-scope") + assert len(dry["selected"]) == 2 and not dry["removed"] + assert len(cache.purge(scope_id="test-scope", dry_run=False)["removed"]) == 2 + + +def test_concurrent_staging_capacity_is_accounted(settings): + settings = replace(settings, result_cache_max_bytes=100) + settings.prepare() + cache = CacheManager(settings) + a, b = cache.staging_path("result"), cache.staging_path("result") + cache.reserve(a, 60) + with pytest.raises(CacheCapacityExceeded): + cache.reserve(b, 41) + cache.discard_staging(a) + cache.reserve(b, 100) + cache.discard_staging(b) + + +def test_single_entry_including_manifest_must_fit(settings, auth): + with TestClient(create_app(replace(settings, source_cache_max_bytes=100))) as client: + response = upload_source( + client, + auth, + source_ref="too-large", + source_format="csv", + filename="test.csv", + content=b"id\n1\n", + ) + assert response.status_code == 507 + assert not list(settings.sources_dir.iterdir()) + assert not list(settings.tmp_dir.glob("source-*")) + + +def test_keyring_rotation_and_fail_closed(settings, tmp_path): + path = tmp_path / "keys.json" + old, new = "a" * 24, "b" * 24 + path.write_text(json.dumps({"tokens": [old]})) + credentials = Credentials(replace(settings, api_token="", token_keyring_file=str(path))) + credentials.authenticate(f"Bearer {old}") + next_path = path.with_suffix(".new") + next_path.write_text(json.dumps({"tokens": [old, new]})) + next_path.replace(path) + credentials.authenticate(f"Bearer {old}") + credentials.authenticate(f"Bearer {new}") + path.write_text(json.dumps({"tokens": [new]})) + with pytest.raises(AuthenticationError): + credentials.authenticate(f"Bearer {old}") + path.write_text("broken") + with pytest.raises(AuthenticationError): + credentials.authenticate(f"Bearer {new}") + + +def test_no_credentials_in_execution_configuration_or_environment(settings, monkeypatch): + safe = replace(settings, token_keyring_file="secret-path").execution_settings() + assert not safe.api_token and not safe.token_keyring_file + # Restore the whole environment after exercising the allowlist. + import os + + with monkeypatch.context() as scoped: + for name, value in list(os.environ.items()): + scoped.setenv(name, value) + scoped.setenv("OMERO_SESSION", "secret") + scoped.setenv("DQW_API_TOKEN", "secret") + sanitize_execution_environment() + assert "OMERO_SESSION" not in os.environ and "DQW_API_TOKEN" not in os.environ + + +def test_volume_owner_is_exclusive(tmp_path): + with ( + volume_lock(tmp_path), + pytest.raises(RuntimeError, match="Another worker"), + volume_lock(tmp_path), + ): + pass + with volume_lock(tmp_path): + pass + + +def test_busy_ingestion_closes_upload_and_lock_wait_is_bounded(settings): + settings.prepare() + service = QueryService(settings) + upload = UploadFile(io.BytesIO(b"id\n1\n"), filename="test.csv") + with ( + service._admit(service._ingest_capacity, "ingestion"), + pytest.raises(Exception, match="slots are busy"), + ): + service.ingest_source( + scope_id="s", + source_ref="r", + source_format=SourceFormat.csv, + declared_size=5, + expected_sha256=None, + upload=upload, + ) + assert upload.file.closed + lock = threading.Lock() + with lock, pytest.raises(Exception, match="already running"), service._try_lock(lock): + pass + + +def test_download_lease_released_if_send_fails(tmp_path): + path = tmp_path / "result.csv" + path.write_bytes(b"id\n1\n") + released = [] + response = LeasedFileResponse( + path, release=lambda: released.append(True), media_type="text/csv", filename="result.csv" + ) + + async def send(message): + raise OSError("disconnected") + + async def receive(): + return {"type": "http.disconnect"} + + with pytest.raises(OSError): + asyncio.run(response({"type": "http", "method": "GET", "headers": []}, receive, send)) + assert released == [True] + + +def test_metrics_private_and_correlation_validated(client, auth): + assert client.get("/v1/metrics").status_code == 401 + response = client.get("/v1/metrics", headers={**auth, "X-Request-ID": "bad/id"}) + assert response.status_code == 200 + assert response.headers["X-Request-ID"] != "bad/id" + assert "dqw_source_count" in response.text + + +def test_upload_spool_capacity_and_unauthorized_body_not_consumed(settings, auth): + limited = replace(settings, source_cache_max_bytes=100) + with TestClient(create_app(limited)) as client: + response = upload_source( + client, + auth, + source_ref="spool", + source_format="csv", + filename="test.csv", + content=b"id\n1\n", + ) + assert response.status_code == 507 + assert not client.app.state.query_service.cache._staging + assert not list(limited.tmp_dir.iterdir()) + + +def test_real_timeout_and_disconnect_leave_no_children_or_leases(settings, tmp_path): + import multiprocessing + import sqlite3 + + from omero_data_query_worker.errors import QueryExecutionError, QueryTimedOut + from omero_data_query_worker.executor import execute_in_subprocess + + path = tmp_path / "source.sqlite" + with sqlite3.connect(path) as connection: + connection.execute("CREATE TABLE data(id INTEGER)") + sql = ( + "WITH RECURSIVE t(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM t " + "WHERE x<1000000000) SELECT sum(x) FROM t" + ) + before = {p.pid for p in multiprocessing.active_children()} + with pytest.raises(QueryTimedOut): + execute_in_subprocess( + path, + SourceFormat.sqlite, + sql, + {}, + tmp_path / "out.csv", + replace(settings, query_timeout_seconds=1), + ) + service = QueryService(settings) + cancelled = threading.Event() + cancelled.set() + with pytest.raises(QueryExecutionError, match="interrupted"): + execute_in_subprocess( + path, + SourceFormat.sqlite, + sql, + {}, + tmp_path / "out.csv", + settings, + monitor=lambda: service._check_cancelled(cancelled), + ) + assert {p.pid for p in multiprocessing.active_children()} == before + + +def _crash_child(*args): + import os + + os._exit(17) + + +def test_subprocess_crash_is_reported_and_reaped(settings, monkeypatch, tmp_path): + from omero_data_query_worker import executor + from omero_data_query_worker.errors import QueryExecutionError + + monkeypatch.setattr(executor, "_child_execute", _crash_child) + with pytest.raises(QueryExecutionError, match="exited unexpectedly"): + executor.execute_in_subprocess( + tmp_path / "unused", SourceFormat.sqlite, "SELECT 1", {}, tmp_path / "out", settings + ) diff --git a/tests/test_policy.py b/tests/test_policy.py index e4f7e92..5cffa75 100644 --- a/tests/test_policy.py +++ b/tests/test_policy.py @@ -33,6 +33,7 @@ def test_accepts_select_with_named_parameters() -> None: "SELECT getenv('DQW_API_TOKEN')", "SELECT * FROM sqlite_scan('x.sqlite', 'data')", "SELECT 1; SELECT 2", + "SELECT 'unterminated", ], ) def test_rejects_unsafe_sql(sql: str) -> None: diff --git a/tests/test_recovery.py b/tests/test_recovery.py new file mode 100644 index 0000000..58abde2 --- /dev/null +++ b/tests/test_recovery.py @@ -0,0 +1,85 @@ +"""Recovery rejects corrupt publications without rewriting immutable inputs.""" + +import hashlib +import json + +import pytest + +from omero_data_query_worker.cache import CacheManager +from omero_data_query_worker.service import QueryService +from tests.conftest import upload_source + + +def test_restart_removes_fresh_staging_and_preserves_valid_cache(client, auth, settings): + uploaded = upload_source( + client, + auth, + source_ref="recovery", + source_format="csv", + filename="source.csv", + content=b"id\n1\n2\n", + ).json() + source_id = uploaded["source_id"] + result = client.post( + f"/v1/sources/{source_id}/query", headers=auth, json={"sql": "SELECT * FROM data"} + ).json() + cache = CacheManager(settings) + abandoned = settings.tmp_dir / "result-abandoned" + abandoned.mkdir() + (abandoned / "result.csv").write_bytes(b"id\n1\n") + source_file = cache.source_path(source_id) / "data.duckdb" + before = hashlib.sha256(source_file.read_bytes()).hexdigest() + QueryService(settings).prepare() + assert not abandoned.exists() + assert cache.get_result(result["result_id"]).row_count == 2 + assert hashlib.sha256(source_file.read_bytes()).hexdigest() == before + + +def test_restart_discards_same_size_corruption(client, auth, settings): + source = upload_source( + client, + auth, + source_ref="corrupt-recovery", + source_format="csv", + filename="source.csv", + content=b"id\n1\n", + ).json() + result = client.post( + f"/v1/sources/{source['source_id']}/query", headers=auth, json={"sql": "SELECT * FROM data"} + ).json() + cache = CacheManager(settings) + path = cache.result_path(result["result_id"]) + data = path / "result.csv" + data.write_bytes(data.read_bytes().replace(b"1", b"9")) + cache.recover() + assert not path.exists() + assert cache.source_path(source["source_id"]).exists() + + +@pytest.mark.parametrize("damage", ["execution", "byte_count", "source_bytes"]) +def test_restart_recovers_from_damaged_metadata_and_sources(client, auth, settings, damage): + source = upload_source( + client, + auth, + source_ref="damaged-manifest", + source_format="csv", + filename="source.csv", + content=b"id\n1\n", + ).json() + result = client.post( + f"/v1/sources/{source['source_id']}/query", headers=auth, json={"sql": "SELECT * FROM data"} + ).json() + cache = CacheManager(settings) + result_path = cache.result_path(result["result_id"]) + source_path = cache.source_path(source["source_id"]) + if damage == "source_bytes": + (source_path / "data.duckdb").write_bytes(b"corrupt") + else: + path = result_path / "manifest.json" + manifest = json.loads(path.read_text()) + manifest[damage] = [] if damage == "execution" else 999 + if damage == "execution": + manifest[damage] = ["invalid"] + path.write_text(json.dumps(manifest)) + cache.recover() + assert not (source_path if damage == "source_bytes" else result_path).exists()