diff --git a/.env.example b/.env.example index e580076..f0b8ed0 100644 --- a/.env.example +++ b/.env.example @@ -12,6 +12,12 @@ ALLOWED_ORIGINS= # Password for the runtime web role. Generate with: openssl rand -hex 32 APP_DB_PASSWORD= +# Passwords for the least-privilege queue roles the API, worker, and retention +# pass connect as. Generate each with: openssl rand -hex 32 +COMPUTE_PRODUCER_PASSWORD= +COMPUTE_WORKER_PASSWORD= +COMPUTE_PURGER_PASSWORD= + # Browser-reachable MinIO endpoint for output download URLs. MINIO_PUBLIC_ENDPOINT=localhost:9000 diff --git a/.github/workflows/crash-recovery.yml b/.github/workflows/crash-recovery.yml index e15984a..5a33e57 100644 --- a/.github/workflows/crash-recovery.yml +++ b/.github/workflows/crash-recovery.yml @@ -20,6 +20,9 @@ jobs: env: COMPUTE_API_TOKEN: compose-e2e-token APP_DB_PASSWORD: compose-e2e-app-db-password + COMPUTE_PRODUCER_PASSWORD: compose-e2e-producer-password + COMPUTE_WORKER_PASSWORD: compose-e2e-worker-password + COMPUTE_PURGER_PASSWORD: compose-e2e-purger-password BETTER_AUTH_SECRET: compose-e2e-better-auth-secret ORIGIN: http://localhost:3000 MINIO_ACCESS_KEY: minioadmin diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2fdca08..cc0bc66 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -235,6 +235,9 @@ jobs: SIMULATION_ID: 4cfe522f-7e7d-46e0-96ca-7b98743fb9f5 COMPUTE_API_TOKEN: compose-e2e-token APP_DB_PASSWORD: compose-e2e-app-db-password + COMPUTE_PRODUCER_PASSWORD: compose-e2e-producer-password + COMPUTE_WORKER_PASSWORD: compose-e2e-worker-password + COMPUTE_PURGER_PASSWORD: compose-e2e-purger-password BETTER_AUTH_SECRET: compose-e2e-better-auth-secret ORIGIN: http://localhost:3000 MINIO_ACCESS_KEY: minioadmin diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index f195b79..fa25c7f 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -103,6 +103,31 @@ The web runtime role can read and write the web tables and read changes. The table definition in `apps/web/src/lib/server/db/compute.ts` is used only for reads and is excluded from web migrations. +The compute service has three restricted runtime roles, provisioned after the +compute and queue migrations by `tsdhn-queue-grants`. None can run schema +changes, and row-level security scopes each to the deployment's queue: + +| Role | Process | Queue capability | `compute.jobs` | +| --- | --- | --- | --- | +| `COMPUTE_PRODUCER_ROLE` | API | enqueue and read (`PRODUCE`) | `SELECT`, `INSERT` | +| `COMPUTE_WORKER_ROLE` | worker | claim and transition (`CONSUME`) | `SELECT`, `UPDATE` | +| `COMPUTE_PURGER_ROLE` | worker retention | inspect plus delete queue jobs | none | + +The producer cannot claim a job, the worker cannot delete one, and the purger +cannot reach `compute.jobs`. Deleting a queue job cascades to its attempt +history, so retention uses its own credential rather than the consumer role. + +## Queue retention + +Terminal `task_queue.jobs` rows whose matching `compute.jobs` row is also +terminal are deleted after seven days by an hourly pass in the worker process. +Rows whose compute counterpart is still running, missing, or malformed remain +available for reconciliation. The queue row has no application value once the +handler records the outcome; the week preserves attempt history for operational +inspection over a working week and weekend. Attempt and occurrence rows follow +through `ON DELETE CASCADE`. `compute.jobs` is not purged because it holds the +simulation result record. + ## Identifiers | Name | Owner | Purpose | diff --git a/docker-compose.yml b/docker-compose.yml index 283a04c..c5589a9 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -44,14 +44,21 @@ services: condition: service_healthy environment: COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn + COMPUTE_QUEUE: ${COMPUTE_QUEUE:-simulations} COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue} APP_DB_ROLE: ${APP_DB_ROLE:-tsdhn_app} APP_DB_PASSWORD: ${APP_DB_PASSWORD:?set APP_DB_PASSWORD in .env} + COMPUTE_PRODUCER_ROLE: ${COMPUTE_PRODUCER_ROLE:-tsdhn_producer} + COMPUTE_PRODUCER_PASSWORD: ${COMPUTE_PRODUCER_PASSWORD:?set COMPUTE_PRODUCER_PASSWORD in .env} + COMPUTE_WORKER_ROLE: ${COMPUTE_WORKER_ROLE:-tsdhn_worker} + COMPUTE_WORKER_PASSWORD: ${COMPUTE_WORKER_PASSWORD:?set COMPUTE_WORKER_PASSWORD in .env} + COMPUTE_PURGER_ROLE: ${COMPUTE_PURGER_ROLE:-tsdhn_purger} + COMPUTE_PURGER_PASSWORD: ${COMPUTE_PURGER_PASSWORD:?set COMPUTE_PURGER_PASSWORD in .env} command: [ "sh", "-lc", - "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev rqueue --database-url \"$$COMPUTE_DATABASE_URL\" --schema \"$$COMPUTE_QUEUE_SCHEMA\" migrate", + "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev rqueue --database-url \"$$COMPUTE_DATABASE_URL\" --schema \"$$COMPUTE_QUEUE_SCHEMA\" migrate && uv run --no-dev tsdhn-queue-grants", ] restart: "no" @@ -79,6 +86,8 @@ services: MINIO_BUCKET: ${MINIO_BUCKET:-tsdhn-results} COMPUTE_API_TOKEN: ${COMPUTE_API_TOKEN:?set COMPUTE_API_TOKEN in .env} ALLOWED_ORIGINS: ${ALLOWED_ORIGINS:-} + COMPUTE_PRODUCER_ROLE: ${COMPUTE_PRODUCER_ROLE:-tsdhn_producer} + COMPUTE_PRODUCER_PASSWORD: ${COMPUTE_PRODUCER_PASSWORD:?set COMPUTE_PRODUCER_PASSWORD in .env} ports: - "8000:8000" @@ -105,6 +114,10 @@ services: MINIO_BUCKET: ${MINIO_BUCKET:-tsdhn-results} TSDHN_MODEL_DIR: /app/model TSDHN_JOBS_DIR: /var/tmp/jobs + COMPUTE_WORKER_ROLE: ${COMPUTE_WORKER_ROLE:-tsdhn_worker} + COMPUTE_WORKER_PASSWORD: ${COMPUTE_WORKER_PASSWORD:?set COMPUTE_WORKER_PASSWORD in .env} + COMPUTE_PURGER_ROLE: ${COMPUTE_PURGER_ROLE:-tsdhn_purger} + COMPUTE_PURGER_PASSWORD: ${COMPUTE_PURGER_PASSWORD:?set COMPUTE_PURGER_PASSWORD in .env} command: ["uv", "run", "--no-dev", "tsdhn-worker"] volumes: - jobs-data:/var/tmp/jobs diff --git a/mise.toml b/mise.toml index ab047db..36c3969 100644 --- a/mise.toml +++ b/mise.toml @@ -156,6 +156,7 @@ depends = ["db:start"] run = [ "uv run tsdhn-compute-migrate", "uv run rqueue --database-url \"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" --schema \"${COMPUTE_QUEUE_SCHEMA:-task_queue}\" migrate", + "uv run tsdhn-queue-grants", "DATABASE_URL=\"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" bun --filter web db:migrate", "uv run tsdhn-web-grants", ] diff --git a/packages/api/api/core/db.py b/packages/api/api/core/db.py index 43807a6..0db26d3 100644 --- a/packages/api/api/core/db.py +++ b/packages/api/api/core/db.py @@ -7,6 +7,7 @@ `settings.worker_pool_size()` sizes the worker's pool the way it does. """ +import logging import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -15,9 +16,10 @@ import asyncpg from api.core.errors import TransientInfraError -from api.core.settings import COMPUTE_DATABASE_URL +from api.core.settings import COMPUTE_DATABASE_URL, role_database_url __all__ = [ + "COMPUTE_DATABASE_URL", "CONNECT_TIMEOUT", "JobRow", "acquire", @@ -27,14 +29,21 @@ "is_transient", "notify_channel", "open_pool", + "runtime_dsn", "transient_connection_errors", ] +logger = logging.getLogger(__name__) + JobRow = dict[str, Any] CONNECT_TIMEOUT = 2 _pool: asyncpg.Pool | None = None +# The DSN the open pool used, so `connect()` reaches the same role rather than +# quietly falling back to the owner behind the pool's back. +_dsn: str | None = None + # Client-side failures: the server was never reached, or the socket died. _CLIENT_ERRORS = (ConnectionError, OSError, TimeoutError) @@ -71,12 +80,27 @@ def is_transient(exc: BaseException) -> bool: return isinstance(sqlstate, str) and sqlstate[:2] in _TRANSIENT_SQLSTATE_CLASSES -async def open_pool(*, min_size: int, max_size: int) -> asyncpg.Pool: +def runtime_dsn(role: str, password: str) -> str | None: + """Resolve a runtime role's DSN, warning when it is not provisioned.""" + dsn = role_database_url(role, password) + if dsn is None: + logger.warning( + "no password configured for role %s; connecting as the database " + "owner instead of the least-privilege role", + role or "", + ) + return dsn + + +async def open_pool( + *, min_size: int, max_size: int, dsn: str | None = None +) -> asyncpg.Pool: """Create the process-wide pool. Idempotent within one process.""" - global _pool + global _pool, _dsn if _pool is None: + _dsn = dsn or COMPUTE_DATABASE_URL _pool = await asyncpg.create_pool( - COMPUTE_DATABASE_URL, + _dsn, min_size=min_size, max_size=max_size, timeout=CONNECT_TIMEOUT, @@ -85,10 +109,11 @@ async def open_pool(*, min_size: int, max_size: int) -> asyncpg.Pool: async def close_pool() -> None: - global _pool + global _pool, _dsn if _pool is not None: await _pool.close() _pool = None + _dsn = None def get_pool() -> asyncpg.Pool: @@ -178,7 +203,9 @@ async def connect() -> asyncpg.Connection: watching browsers starve every other route. """ try: - return await asyncpg.connect(COMPUTE_DATABASE_URL, timeout=CONNECT_TIMEOUT) + return await asyncpg.connect( + _dsn or COMPUTE_DATABASE_URL, timeout=CONNECT_TIMEOUT + ) except Exception as e: if not is_transient(e): raise diff --git a/packages/api/api/core/settings.py b/packages/api/api/core/settings.py index 37ab65e..d03c069 100644 --- a/packages/api/api/core/settings.py +++ b/packages/api/api/core/settings.py @@ -1,12 +1,19 @@ import os from pathlib import Path +from urllib.parse import quote, urlsplit, urlunsplit __all__ = [ "APP_DB_PASSWORD", "APP_DB_ROLE", "COMPUTE_DATABASE_URL", + "COMPUTE_PRODUCER_PASSWORD", + "COMPUTE_PRODUCER_ROLE", + "COMPUTE_PURGER_PASSWORD", + "COMPUTE_PURGER_ROLE", "COMPUTE_QUEUE", "COMPUTE_QUEUE_SCHEMA", + "COMPUTE_WORKER_PASSWORD", + "COMPUTE_WORKER_ROLE", "DB_POOL_MAX_SIZE", "DB_POOL_MIN_SIZE", "JOBS_DIR", @@ -81,6 +88,43 @@ def worker_pool_size() -> tuple[int, int]: APP_DB_ROLE = os.environ.get("APP_DB_ROLE", "tsdhn_app") APP_DB_PASSWORD = os.environ.get("APP_DB_PASSWORD", "") +# COMPUTE_DATABASE_URL names the schema owner: the role that runs migrations +# and owns both `compute` and the queue schema. Runtime processes use these +# roles instead, each provisioned with only the grants its process needs. +COMPUTE_PRODUCER_ROLE = os.environ.get("COMPUTE_PRODUCER_ROLE", "tsdhn_producer") +COMPUTE_PRODUCER_PASSWORD = os.environ.get("COMPUTE_PRODUCER_PASSWORD", "") + +COMPUTE_WORKER_ROLE = os.environ.get("COMPUTE_WORKER_ROLE", "tsdhn_worker") +COMPUTE_WORKER_PASSWORD = os.environ.get("COMPUTE_WORKER_PASSWORD", "") + +# Retention deletes queue history, which no consuming role may do. Keeping it +# on its own credential means a deployment can withhold it or move retention +# to a maintenance container without changing the worker role. +COMPUTE_PURGER_ROLE = os.environ.get("COMPUTE_PURGER_ROLE", "tsdhn_purger") +COMPUTE_PURGER_PASSWORD = os.environ.get("COMPUTE_PURGER_PASSWORD", "") + + +def role_database_url(role: str, password: str) -> str | None: + """Return the compute URL rewritten to connect as `role`. + + A missing role or password means the role has not been provisioned. The + runtime caller can then retain the owner URL while warning the operator. + Credentials are percent-encoded because passwords commonly contain URL + punctuation. + """ + if not role or not password: + return None + parts = urlsplit(COMPUTE_DATABASE_URL) + credentials = f"{quote(role, safe='')}:{quote(password, safe='')}" + # Preserve the authority verbatim. asyncpg accepts socket URLs with no + # hostname and multi-host authorities; accessing ``parts.hostname`` or + # ``parts.port`` would reject those valid DSN forms before asyncpg sees + # them. Existing userinfo is replaced by taking everything after the last + # @, while the raw host list (if any) remains untouched. + authority = parts.netloc.rsplit("@", 1)[-1] + return urlunsplit(parts._replace(netloc=f"{credentials}@{authority}")) + + MINIO_ENDPOINT = os.environ.get("MINIO_ENDPOINT", "localhost:9000") # Public endpoint differs from API endpoint for browser downloads. MINIO_PUBLIC_ENDPOINT = os.environ.get("MINIO_PUBLIC_ENDPOINT", MINIO_ENDPOINT) diff --git a/packages/api/api/core/tasks.py b/packages/api/api/core/tasks.py index d44a2ef..e57cd21 100644 --- a/packages/api/api/core/tasks.py +++ b/packages/api/api/core/tasks.py @@ -89,29 +89,41 @@ import threading import uuid from collections.abc import Awaitable, Callable -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any import asyncpg -from rqueue import CancelJob, Job, PermanentFailure, Queue, RetryPolicy, TaskContext +from rqueue import ( + Admin, + CancelJob, + Job, + PermanentFailure, + Queue, + RetryPolicy, + TaskContext, +) +from rqueue.models import TERMINAL_STATES from api.core import db, repository from api.core.errors import TransientInfraError from api.core.queue import get_queue -from api.core.settings import JOBS_DIR +from api.core.settings import COMPUTE_QUEUE, JOBS_DIR from tsdhn.domain import EarthquakeInput, JobStatus from tsdhn.engine import run_simulation __all__ = [ + "JOB_RETENTION", "MAX_ATTEMPTS", "RUN_SIMULATION", "TRANSIENT_RETRY", "AbandonedAttempt", "decode_payload", "enqueue_simulation", + "purge_finished_jobs", "reconcile_terminal_jobs", "register_tasks", + "run_periodic_purge", "run_periodic_reconcile", "run_periodic_sweep", "run_simulation_task", @@ -179,6 +191,28 @@ class AbandonedAttempt(Exception): # plus the grace, not by the interval. RECONCILE_INTERVAL_SECONDS = 60.0 +# A terminal queue row has no application value after the handler records its +# outcome. Keep the queue-side attempt history for a working week plus the +# weekend, while leaving compute.jobs (the result record) untouched. +JOB_RETENTION = timedelta(days=7) +PURGE_INTERVAL_SECONDS = 3600.0 + +# Bound each purge transaction so a first run against a long-unpurged table +# does not hold one enormous delete open; the next hourly pass takes the rest. +PURGE_LIMIT = 10000 + +# `compute.jobs` uses application states rather than rqueue's states. A queue +# row is safe to remove only after reconciliation (or the simulation itself) +# has recorded one of these terminal outcomes on its compute counterpart. +COMPUTE_TERMINAL_STATUSES = ( + JobStatus.COMPLETED.value, + JobStatus.FAILED.value, +) + +# The cast below is protected by a materialized CTE and this canonical UUID +# check, just as repository reconciliation protects its own payload cast. +COMPUTE_JOB_ID_RE = "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$" + def decode_payload(payload: Any) -> uuid.UUID: """Turn a queued payload into the compute job id it names. @@ -546,6 +580,92 @@ async def sweep_abandoned_work_dirs() -> None: await asyncio.to_thread(remove_workspace, JOBS_DIR / simulation_id) +async def _eligible_queue_job_ids( + admin: Admin, compute_pool: asyncpg.Pool, cutoff: datetime +) -> list[uuid.UUID]: + """Find old queue rows whose compute records are already terminal. + + The check runs through the worker pool, which can read both schemas. The + purger role deliberately cannot read ``compute.jobs``; it only performs + the final DELETE after this join has selected safe queue rows. Invalid + payloads and missing/nonterminal compute rows are conservatively skipped. + """ + async with compute_pool.acquire() as connection: + rows = await connection.fetch( + f""" + WITH candidates AS MATERIALIZED ( + SELECT q.id, + (q.payload->>'compute_job_id')::uuid AS compute_job_id, + q.finished_at + FROM {admin.schema}.jobs AS q + WHERE q.queue = $1 + AND q.state = ANY($2::text[]) + AND q.finished_at < $3 + AND q.payload->>'compute_job_id' ~* $4 + ) + SELECT candidates.id + FROM candidates + JOIN compute.jobs AS c ON c.id = candidates.compute_job_id + WHERE c.status = ANY($5::text[]) + ORDER BY candidates.finished_at + LIMIT $6 + """, # noqa: S608 - admin.schema is validated by rqueue + COMPUTE_QUEUE, + list(TERMINAL_STATES), + cutoff, + COMPUTE_JOB_ID_RE, + list(COMPUTE_TERMINAL_STATUSES), + PURGE_LIMIT, + ) + return [row["id"] for row in rows] + + +async def _delete_queue_job_ids( + admin: Admin, job_ids: list[uuid.UUID], cutoff: datetime +) -> int: + """Delete the prechecked queue rows using the separately scoped role.""" + if not job_ids: + return 0 + async with admin.pool.acquire() as connection: + # The purger has DELETE on jobs and the cascade removes queue history. + # A single statement keeps the selected IDs' deletion atomic. + removed: int = await connection.fetchval( + f""" + WITH removed AS ( + DELETE FROM {admin.schema}.jobs + WHERE id = ANY($1::uuid[]) + AND queue = $2 + AND state = ANY($3::text[]) + AND finished_at < $4 + RETURNING id + ) + SELECT count(*)::int FROM removed + """, # noqa: S608 - admin.schema is validated by rqueue + job_ids, + COMPUTE_QUEUE, + list(TERMINAL_STATES), + cutoff, + ) + return removed + + +async def purge_finished_jobs( + admin: Admin, *, compute_pool: asyncpg.Pool | None = None +) -> int: + """Delete old queue rows only after their compute jobs are terminal. + + ``Admin.purge`` cannot express the cross-schema safety join, so retention + first checks candidates using the worker's compute-readable pool and then + deletes the selected queue IDs using the purger pool. A compute row that is + still pending or running remains paired with its queue evidence for the + next reconciliation pass. + """ + compute_pool = compute_pool or db.get_pool() + cutoff = datetime.now(UTC) - JOB_RETENTION + job_ids = await _eligible_queue_job_ids(admin, compute_pool, cutoff) + return await _delete_queue_job_ids(admin, job_ids, cutoff) + + async def reconcile_terminal_jobs(*, grace: timedelta = RECONCILE_GRACE) -> None: """Sync compute.jobs to jobs the queue finished without the run reporting. @@ -621,6 +741,32 @@ async def run_periodic_reconcile( ) +async def run_periodic_purge( + admin: Admin, + stop: asyncio.Event, + *, + compute_pool: asyncpg.Pool | None = None, + interval: float = PURGE_INTERVAL_SECONDS, +) -> None: + """Purge finished queue rows until ``stop`` is set.""" + + async def purge() -> None: + removed = await purge_finished_jobs(admin, compute_pool=compute_pool) + if removed: + logger.info( + "purged %d finished queue job(s) older than %s", + removed, + JOB_RETENTION, + ) + + await _run_periodically( + "Purge of finished queue jobs", + purge, + stop, + interval, + ) + + def register_tasks(queue: Queue) -> Queue: """Register every task this deployment runs, and return the queue. diff --git a/packages/api/api/main.py b/packages/api/api/main.py index d9c837a..4927cb3 100644 --- a/packages/api/api/main.py +++ b/packages/api/api/main.py @@ -10,7 +10,12 @@ from api import __version__ from api.core import db from api.core.queue import build_queue -from api.core.settings import LOG_LEVEL, api_pool_size +from api.core.settings import ( + COMPUTE_PRODUCER_PASSWORD, + COMPUTE_PRODUCER_ROLE, + LOG_LEVEL, + api_pool_size, +) from api.core.tasks import register_tasks from api.routes import get_calculator, ops_router, router @@ -27,7 +32,11 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: get_calculator() min_size, max_size = api_pool_size() # A zero floor keeps this from waiting on Postgres to start. - pool = await db.open_pool(min_size=min_size, max_size=max_size) + pool = await db.open_pool( + min_size=min_size, + max_size=max_size, + dsn=db.runtime_dsn(COMPUTE_PRODUCER_ROLE, COMPUTE_PRODUCER_PASSWORD), + ) register_tasks(build_queue(pool)) logger.info("TSDHN API ready") try: diff --git a/packages/api/api/queue_grants.py b/packages/api/api/queue_grants.py new file mode 100644 index 0000000..246e29a --- /dev/null +++ b/packages/api/api/queue_grants.py @@ -0,0 +1,244 @@ +"""Provision the least-privilege roles the API and worker connect as. + +`tsdhn-compute-migrate` owns the compute schema and the web role; +`rqueue ... migrate` owns the queue schema. This runs after both, for the same +reason `tsdhn-web-grants` runs after the web migrations: a GRANT needs the +tables to exist. Everything here is idempotent, and re-running it narrows a +role that has been over-granted back to the table below. + +| role | queue schema | compute schema | +| ------------------- | ------------------------------------- | -------------- | +| COMPUTE_PRODUCER | `Capability.PRODUCE` | SELECT, INSERT | +| COMPUTE_WORKER | `Capability.CONSUME` | SELECT, UPDATE | +| COMPUTE_PURGER | `Capability.INSPECT` + DELETE on jobs | none | + +`provision_role` grants no DDL to any of them, and row-level security scopes +`jobs` and `job_attempts` to the one queue this deployment runs; +`concurrency_slots` and `runtime_heartbeats` have no per-queue RLS in rqueue +today, which is fine for this single-queue deployment. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from urllib.parse import urlsplit + +import asyncpg +from rqueue.models import TERMINAL_STATES +from rqueue.roles import Capability, provision_role + +from api.core.settings import ( + APP_DB_ROLE, + COMPUTE_DATABASE_URL, + COMPUTE_PRODUCER_PASSWORD, + COMPUTE_PRODUCER_ROLE, + COMPUTE_PURGER_PASSWORD, + COMPUTE_PURGER_ROLE, + COMPUTE_QUEUE, + COMPUTE_QUEUE_SCHEMA, + COMPUTE_WORKER_PASSWORD, + COMPUTE_WORKER_ROLE, +) + +__all__ = ["QueueRole", "provision_queue_roles", "queue_roles"] + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class QueueRole: + """One runtime role: what it may do to the queue, and to `compute.jobs`.""" + + name: str + password: str + #: The environment variable the password comes from, for error messages. + env_var: str + capabilities: tuple[Capability, ...] + #: Privileges on `compute.jobs`; empty means the role never reaches it. + compute_privileges: str = "" + #: Queue-table privileges no capability's grant set covers. + extra_queue_privileges: tuple[tuple[str, str], ...] = field(default_factory=tuple) + + +def queue_roles() -> tuple[QueueRole, ...]: + """The roles this deployment runs, read from settings at call time.""" + return ( + QueueRole( + name=COMPUTE_PRODUCER_ROLE, + password=COMPUTE_PRODUCER_PASSWORD, + env_var="COMPUTE_PRODUCER_PASSWORD", + # The API only enqueues and reads back. It never touches schedules + # or queue pauses, which is all INSPECT would add over PRODUCE. + capabilities=(Capability.PRODUCE,), + # `create_or_get_job` inserts and reads; no route updates a row. + compute_privileges="SELECT, INSERT", + ), + QueueRole( + name=COMPUTE_WORKER_ROLE, + password=COMPUTE_WORKER_PASSWORD, + env_var="COMPUTE_WORKER_PASSWORD", + capabilities=(Capability.CONSUME,), + # The worker only ever advances rows the API created. + compute_privileges="SELECT, UPDATE", + ), + QueueRole( + name=COMPUTE_PURGER_ROLE, + password=COMPUTE_PURGER_PASSWORD, + env_var="COMPUTE_PURGER_PASSWORD", + # `Admin.purge` reads the terminal rows it is about to remove. + capabilities=(Capability.INSPECT,), + # Retention is a queue-only concern: compute.jobs keeps its own + # history, and this task is not the one to start expiring it. + compute_privileges="", + # No capability carries DELETE, on purpose -- a consumer that can + # delete a job can erase its own attempt history with it. Purge + # needs exactly this one grant and nothing else. + extra_queue_privileges=(("jobs", "DELETE"),), + ), + ) + + +async def _ddl(connection: asyncpg.Connection, template: str, *args: str) -> None: + """Run one DDL statement with PostgreSQL doing the identifier quoting. + + The same trick `rqueue.roles` uses: role, schema and table names have no + bind-parameter form, so `format(..., %I)` is evaluated server-side and only + its already-quoted result is executed. + """ + placeholders = ", ".join(f"${index + 2}::text" for index in range(len(args))) + statement = await connection.fetchval( + f"SELECT format($1::text, {placeholders})", template, *args + ) + await connection.execute(statement) + + +async def _grant_compute_access( + connection: asyncpg.Connection, role: QueueRole +) -> None: + """Give one role its `compute.jobs` privileges, and only those. + + The revoke comes first so re-running repairs a role that was widened by + hand, matching what `migrate.provision_web_role` does for the web role. + """ + await _ddl(connection, "REVOKE ALL PRIVILEGES ON compute.jobs FROM %I", role.name) + await _ddl(connection, "REVOKE CREATE ON SCHEMA compute FROM %I", role.name) + if not role.compute_privileges: + await _ddl(connection, "REVOKE USAGE ON SCHEMA compute FROM %I", role.name) + return + await _ddl(connection, "GRANT USAGE ON SCHEMA compute TO %I", role.name) + await _ddl( + connection, + f"GRANT {role.compute_privileges} ON compute.jobs TO %I", + role.name, + ) + + +async def _grant_purger_delete_policy( + connection: asyncpg.Connection, role: QueueRole +) -> None: + """Restrict the purger's direct queue deletes to terminal rows.""" + policy_name = "compute_purger_terminal_delete" + await _ddl( + connection, + "DROP POLICY IF EXISTS %I ON %I.%I", + policy_name, + COMPUTE_QUEUE_SCHEMA, + "jobs", + ) + state_placeholders = ", ".join("%L" for _ in TERMINAL_STATES) + await _ddl( + connection, + "CREATE POLICY %I ON %I.%I AS RESTRICTIVE FOR DELETE TO %I " + f"USING (state = ANY(ARRAY[{state_placeholders}]::text[]))", + policy_name, + COMPUTE_QUEUE_SCHEMA, + "jobs", + role.name, + *TERMINAL_STATES, + ) + + +def _validate_role_names(roles: tuple[QueueRole, ...]) -> None: + """Reject runtime roles that could overwrite another database role.""" + role_env_vars = ( + "COMPUTE_PRODUCER_ROLE", + "COMPUTE_WORKER_ROLE", + "COMPUTE_PURGER_ROLE", + ) + names = [role.name for role in roles] + if len(names) != len(set(names)): + raise ValueError( + "COMPUTE_PRODUCER_ROLE, COMPUTE_WORKER_ROLE, and " + "COMPUTE_PURGER_ROLE must be pairwise distinct" + ) + protected = ( + ("APP_DB_ROLE", APP_DB_ROLE), + ("COMPUTE_DATABASE_URL username", urlsplit(COMPUTE_DATABASE_URL).username), + ) + for env_var, role in zip(role_env_vars, roles, strict=True): + for protected_name, protected_value in protected: + if protected_value and role.name == protected_value: + raise ValueError( + f"{env_var}={role.name!r} collides with " + f"{protected_name}={protected_value!r}" + ) + + +async def provision_queue_roles(connection: asyncpg.Connection) -> None: + """Create or repair every runtime role, scoped to COMPUTE_QUEUE.""" + roles = queue_roles() + _validate_role_names(roles) + + for role in roles: + if not role.password: + raise RuntimeError( + f"{role.env_var} must be set: it is the password for " + f"the database role {role.name}." + ) + await provision_role( + connection, + role=role.name, + capabilities=role.capabilities, + schema=COMPUTE_QUEUE_SCHEMA, + # One queue, so the row-level-security scope is that queue rather + # than '*'. A second queue would need a row here, not a code change. + queues=(COMPUTE_QUEUE,), + password=role.password, + ) + for table, privileges in role.extra_queue_privileges: + await _ddl( + connection, + f"GRANT {privileges} ON %I.%I TO %I", + COMPUTE_QUEUE_SCHEMA, + table, + role.name, + ) + if role.name == COMPUTE_PURGER_ROLE: + await _grant_purger_delete_policy(connection, role) + await _grant_compute_access(connection, role) + logger.info( + "provisioned %s with %s on queue %s", + role.name, + ", ".join(capability.value for capability in role.capabilities), + COMPUTE_QUEUE, + ) + + +async def _run() -> None: + connection = await asyncpg.connect(COMPUTE_DATABASE_URL) + try: + async with connection.transaction(): + await provision_queue_roles(connection) + finally: + await connection.close() + + +def main() -> None: # pragma: no cover + logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") + asyncio.run(_run()) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/packages/api/api/worker.py b/packages/api/api/worker.py index 258dccc..c0435f2 100644 --- a/packages/api/api/worker.py +++ b/packages/api/api/worker.py @@ -3,14 +3,22 @@ import os import signal import socket +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +import asyncpg import numba -from rqueue import Worker +from rqueue import Admin, Worker from api.core import db from api.core.queue import build_queue from api.core.settings import ( + COMPUTE_PURGER_PASSWORD, + COMPUTE_PURGER_ROLE, COMPUTE_QUEUE, + COMPUTE_QUEUE_SCHEMA, + COMPUTE_WORKER_PASSWORD, + COMPUTE_WORKER_ROLE, LOG_LEVEL, NUMBA_THREADS, WORKER_CONCURRENCY, @@ -20,6 +28,7 @@ ) from api.core.tasks import ( register_tasks, + run_periodic_purge, run_periodic_reconcile, run_periodic_sweep, ) @@ -42,9 +51,35 @@ def worker_id() -> str: return f"tsdhn-worker-{host or 'unknown'}-{os.getpid()}"[:128] +@asynccontextmanager +async def purge_pool() -> AsyncIterator[asyncpg.Pool]: + """Yield the pool used by retention, separately credentialed from consume. + + When the purger role is not configured, retain the same development + fallback as ``db.open_pool``: a separate owner-DSN pool. Reusing the worker + pool here would use the CONSUME role, which intentionally cannot delete + queue jobs. + """ + dsn = db.runtime_dsn(COMPUTE_PURGER_ROLE, COMPUTE_PURGER_PASSWORD) + fallback_dsn = dsn or db.COMPUTE_DATABASE_URL + # Retention is hourly, so one lazy connection is enough. It is always a + # separate pool, including the owner fallback above. + pool = await asyncpg.create_pool( + fallback_dsn, min_size=0, max_size=1, timeout=db.CONNECT_TIMEOUT + ) + try: + yield pool + finally: + await pool.close() + + async def run() -> None: min_size, max_size = worker_pool_size() - pool = await db.open_pool(min_size=min_size, max_size=max_size) + pool = await db.open_pool( + min_size=min_size, + max_size=max_size, + dsn=db.runtime_dsn(COMPUTE_WORKER_ROLE, COMPUTE_WORKER_PASSWORD), + ) try: worker = Worker( register_tasks(build_queue(pool)), @@ -69,22 +104,27 @@ async def run() -> None: # here; in practice the orchestrator's own SIGKILL bounds it. loop.add_signal_handler(received, worker.stop) - # The worker process owns maintenance because it is the process that - # owns recovery: reconciliation exists to repair the compute.jobs rows - # rqueue's own lease recovery leaves behind. stop_maintenance = asyncio.Event() - maintenance = [ - asyncio.create_task(run_periodic_sweep(stop_maintenance)), - asyncio.create_task(run_periodic_reconcile(stop_maintenance)), - ] logger.info("simulation worker serving queue %s", COMPUTE_QUEUE) - try: - await worker.run() - finally: - stop_maintenance.set() - for task in maintenance: - task.cancel() - await asyncio.gather(*maintenance, return_exceptions=True) + async with purge_pool() as retention_pool: + maintenance = [ + asyncio.create_task(run_periodic_sweep(stop_maintenance)), + asyncio.create_task(run_periodic_reconcile(stop_maintenance)), + asyncio.create_task( + run_periodic_purge( + Admin(retention_pool, schema=COMPUTE_QUEUE_SCHEMA), + stop_maintenance, + compute_pool=pool, + ) + ), + ] + try: + await worker.run() + finally: + stop_maintenance.set() + for task in maintenance: + task.cancel() + await asyncio.gather(*maintenance, return_exceptions=True) finally: await db.close_pool() diff --git a/packages/api/pyproject.toml b/packages/api/pyproject.toml index ff95ef6..2574912 100644 --- a/packages/api/pyproject.toml +++ b/packages/api/pyproject.toml @@ -25,6 +25,7 @@ tsdhn-api = "api.main:start_app" tsdhn-worker = "api.worker:main" tsdhn-compute-migrate = "api.migrate:main" tsdhn-web-grants = "api.web_grants:main" +tsdhn-queue-grants = "api.queue_grants:main" [build-system] requires = ["uv_build==0.11.28"] diff --git a/packages/api/readme.md b/packages/api/readme.md index 68deb00..9ce616f 100644 --- a/packages/api/readme.md +++ b/packages/api/readme.md @@ -18,25 +18,32 @@ uv run tsdhn-api uv run tsdhn-worker ``` -Create the compute and queue tables before starting the API or worker. -`tsdhn-compute-migrate` also provisions the web application's database role, so -it needs a password for that role; there is no default, because a default -password is worse than an error: +Create the compute and queue tables before starting the API or worker. The +migration connection is the schema owner; runtime processes use three roles +with separate passwords from `.env`: ```sh -export APP_DB_PASSWORD="$(openssl rand -hex 32)" # or set it in .env - uv run tsdhn-compute-migrate uv run rqueue \ --database-url "${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@localhost:5432/tsdhn}" \ --schema "${COMPUTE_QUEUE_SCHEMA:-task_queue}" \ migrate +uv run tsdhn-queue-grants ``` -Everything else defaults to the values `api/core/settings.py` uses, so with -`APP_DB_PASSWORD` set the block works on a fresh clone with nothing else -exported. `rqueue`'s CLI reads `RQUEUE_DATABASE_URL`, not -`COMPUTE_DATABASE_URL`, which is why the URL is passed explicitly. +`tsdhn-compute-migrate` also provisions the web application's database role, so +set `APP_DB_PASSWORD` alongside `COMPUTE_PRODUCER_PASSWORD`, +`COMPUTE_WORKER_PASSWORD`, and `COMPUTE_PURGER_PASSWORD` before running these +commands. `tsdhn-queue-grants` provisions the API producer, worker consumer, +and retention purger roles. +`COMPUTE_DATABASE_URL` remains the owner URL used only for migrations and +grants. If a runtime role password is absent, that process warns and falls +back to the owner URL so an unprovisioned development database can still start; +the worker's retention path opens its own owner-URL pool rather than reusing +the consumer pool. + +`rqueue`'s CLI reads `RQUEUE_DATABASE_URL`, not `COMPUTE_DATABASE_URL`, which is +why the URL is passed explicitly. `rqueue` owns the queue tables and their migrations; `--schema` is a global flag, before the subcommand. They live in `COMPUTE_QUEUE_SCHEMA` (default @@ -72,8 +79,10 @@ download flows. - `api/core/db.py` owns the process-wide asyncpg pool. - `api/migrate.py` and `api/web_grants.py` apply the database changes described in `DEPLOY.md`; `rqueue migrate` applies the queue's own. -- `api/worker.py` starts an `rqueue.Worker` for the configured queue and the - periodic workspace sweep. +- `api/queue_grants.py` provisions the least-privilege queue roles and their + `compute.jobs` grants. +- `api/worker.py` starts an `rqueue.Worker` for the configured queue, the + periodic workspace sweep, and the hourly queue-retention purge. ## Tests diff --git a/packages/api/tests/test_db.py b/packages/api/tests/test_db.py index f37617a..dcf3794 100644 --- a/packages/api/tests/test_db.py +++ b/packages/api/tests/test_db.py @@ -117,6 +117,79 @@ def test_notify_channel_is_a_bare_identifier_per_job() -> None: assert db.notify_channel(uuid.uuid4()) != channel +def test_role_database_url_swaps_only_the_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + settings, + "COMPUTE_DATABASE_URL", + "postgresql://owner:secret@db.internal:5432/tsdhn", + ) + + url = settings.role_database_url("tsdhn_producer", "p@ss word") + + assert url == "postgresql://tsdhn_producer:p%40ss%20word@db.internal:5432/tsdhn" + + +@pytest.mark.parametrize( + ("database_url", "expected"), + [ + ( + "postgresql:///tsdhn?host=/var/run/postgresql", + "postgresql://tsdhn_worker:secret@/tsdhn?host=/var/run/postgresql", + ), + ( + "postgresql://owner:old@db1:5432,db2:5432/tsdhn", + "postgresql://tsdhn_worker:secret@db1:5432,db2:5432/tsdhn", + ), + ], +) +def test_role_database_url_preserves_asyncpg_dsn_authorities( + monkeypatch: pytest.MonkeyPatch, database_url: str, expected: str +) -> None: + monkeypatch.setattr(settings, "COMPUTE_DATABASE_URL", database_url) + + assert settings.role_database_url("tsdhn_worker", "secret") == expected + + +def test_an_unprovisioned_role_leaves_the_caller_on_the_owner_connection() -> None: + assert settings.role_database_url("tsdhn_producer", "") is None + assert settings.role_database_url("", "secret") is None + + +def test_runtime_dsn_warns_when_a_role_is_unprovisioned( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level("WARNING"): + assert db.runtime_dsn("tsdhn_producer", "") is None + + assert "tsdhn_producer" in caplog.text + + +@pytest.mark.asyncio +async def test_connect_reuses_the_dsn_the_pool_was_opened_with( + monkeypatch: pytest.MonkeyPatch, +) -> None: + seen: list[str] = [] + + async def create_pool(dsn: str, **_kwargs: Any) -> object: + return object() + + async def connect(dsn: str, **_kwargs: Any) -> object: + seen.append(dsn) + return object() + + monkeypatch.setattr(db_module.asyncpg, "create_pool", create_pool) + monkeypatch.setattr(db_module.asyncpg, "connect", connect) + monkeypatch.setattr(db_module, "_pool", None) + monkeypatch.setattr(db_module, "_dsn", None) + + await db.open_pool(min_size=0, max_size=1, dsn="postgresql://role:pw@host/db") + await db.connect() + + assert seen == ["postgresql://role:pw@host/db"] + + class _Connection: """A connection that reports whether it is closed, like asyncpg's.""" diff --git a/packages/api/tests/test_queue_grants.py b/packages/api/tests/test_queue_grants.py new file mode 100644 index 0000000..187209b --- /dev/null +++ b/packages/api/tests/test_queue_grants.py @@ -0,0 +1,50 @@ +"""Database-independent checks for queue-role configuration.""" + +from typing import Any, cast + +import pytest + +from api import queue_grants + + +@pytest.mark.asyncio +async def test_queue_role_names_must_be_pairwise_distinct( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(queue_grants, "COMPUTE_PRODUCER_ROLE", "same-role") + monkeypatch.setattr(queue_grants, "COMPUTE_WORKER_ROLE", "same-role") + monkeypatch.setattr(queue_grants, "COMPUTE_PURGER_ROLE", "purger-role") + + with pytest.raises(ValueError, match="pairwise distinct"): + await queue_grants.provision_queue_roles(cast(Any, object())) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("field", "value", "conflict"), + [ + ("COMPUTE_PRODUCER_ROLE", "web-role", "APP_DB_ROLE"), + ("COMPUTE_WORKER_ROLE", "db-owner", "COMPUTE_DATABASE_URL username"), + ], +) +async def test_queue_role_names_cannot_overwrite_protected_roles( + monkeypatch: pytest.MonkeyPatch, + field: str, + value: str, + conflict: str, +) -> None: + monkeypatch.setattr(queue_grants, "COMPUTE_PRODUCER_ROLE", "producer-role") + monkeypatch.setattr(queue_grants, "COMPUTE_WORKER_ROLE", "worker-role") + monkeypatch.setattr(queue_grants, "COMPUTE_PURGER_ROLE", "purger-role") + monkeypatch.setattr(queue_grants, field, value) + if conflict == "APP_DB_ROLE": + monkeypatch.setattr(queue_grants, "APP_DB_ROLE", value) + else: + monkeypatch.setattr( + queue_grants, + "COMPUTE_DATABASE_URL", + "postgresql://db-owner:secret@localhost:5432/tsdhn", + ) + + with pytest.raises(ValueError, match=conflict): + await queue_grants.provision_queue_roles(cast(Any, object())) diff --git a/packages/api/tests/test_queue_roles_integration.py b/packages/api/tests/test_queue_roles_integration.py new file mode 100644 index 0000000..079c3b1 --- /dev/null +++ b/packages/api/tests/test_queue_roles_integration.py @@ -0,0 +1,559 @@ +"""What the least-privilege queue roles can and cannot do, against PostgreSQL. + +Every test here connects as a real provisioned role and asserts the boundary +from the grant table in `api.queue_grants`, in both directions: the statement +each process actually runs must succeed, and the neighbouring one it must +never run has to be refused by PostgreSQL, not by application code. +""" + +from __future__ import annotations + +import uuid +from collections.abc import AsyncIterator +from datetime import UTC, datetime, timedelta +from typing import Any +from urllib.parse import urlsplit + +import asyncpg +import pytest +import pytest_asyncio +from rqueue import Admin, Queue, RetryPolicy, TaskContext, Worker +from rqueue.roles import revoke_role + +from api import queue_grants +from api import worker as worker_module +from api.core import db +from api.core.settings import COMPUTE_QUEUE, COMPUTE_QUEUE_SCHEMA +from api.core.tasks import JOB_RETENTION, RUN_SIMULATION, purge_finished_jobs +from scripts.database import database_target + +pytestmark = pytest.mark.integration + +SCHEMA = COMPUTE_QUEUE_SCHEMA +# A disposable role in a disposable database; the cluster trusts localhost. +ROLE_PASSWORD = "queue-role-test-password" + + +class Roles: + """The three provisioned role names, and how to connect as one.""" + + def __init__(self, database_url: str, suffix: str) -> None: + self.database_url = database_url + self.producer = f"tsdhn_test_producer_{suffix}" + self.worker = f"tsdhn_test_worker_{suffix}" + self.purger = f"tsdhn_test_purger_{suffix}" + + @property + def all(self) -> tuple[str, ...]: + return (self.producer, self.worker, self.purger) + + def url(self, role: str) -> str: + name = urlsplit(self.database_url).path.lstrip("/") + return database_target( + self.database_url, name, user=role, password=ROLE_PASSWORD + ).database_url + + async def connect(self, role: str) -> asyncpg.Connection: + return await asyncpg.connect(self.url(role)) + + +@pytest_asyncio.fixture +async def roles( + queue: Queue, isolated_database: str, monkeypatch: pytest.MonkeyPatch +) -> AsyncIterator[Roles]: + """Provision the three roles the way `tsdhn-queue-grants` does.""" + provisioned = Roles(isolated_database, uuid.uuid4().hex[:12]) + monkeypatch.setattr(queue_grants, "COMPUTE_PRODUCER_ROLE", provisioned.producer) + monkeypatch.setattr(queue_grants, "COMPUTE_WORKER_ROLE", provisioned.worker) + monkeypatch.setattr(queue_grants, "COMPUTE_PURGER_ROLE", provisioned.purger) + for setting in ( + "COMPUTE_PRODUCER_PASSWORD", + "COMPUTE_WORKER_PASSWORD", + "COMPUTE_PURGER_PASSWORD", + ): + monkeypatch.setattr(queue_grants, setting, ROLE_PASSWORD) + + async with db.acquire() as connection: + await queue_grants.provision_queue_roles(connection) + try: + yield provisioned + finally: + async with db.acquire() as connection: + for role in provisioned.all: + await revoke_role(connection, role=role, schema=SCHEMA, drop=True) + + +@pytest.mark.asyncio +async def test_unprovisioned_purger_uses_a_separate_owner_pool( + queue: Queue, roles: Roles, monkeypatch: pytest.MonkeyPatch +) -> None: + """The development fallback still has purge privileges without a role.""" + job_id = await _enqueue(queue) + await _finish(job_id, age=JOB_RETENTION + timedelta(days=1)) + + worker_pool = await asyncpg.create_pool( + roles.url(roles.worker), min_size=1, max_size=1 + ) + assert worker_pool is not None + monkeypatch.setattr(worker_module, "COMPUTE_PURGER_PASSWORD", "") + try: + async with worker_module.purge_pool() as fallback_pool: + assert fallback_pool is not worker_pool + async with fallback_pool.acquire() as connection: + assert await connection.fetchval("SELECT current_user") == "tsdhn" + assert ( + await purge_finished_jobs( + Admin(fallback_pool, schema=SCHEMA), compute_pool=worker_pool + ) + == 1 + ) + finally: + await worker_pool.close() + + +async def _enqueue(queue: Queue, **kwargs: Any) -> uuid.UUID: + compute_job_id = uuid.uuid4() + async with db.acquire() as connection: + await connection.execute( + """ + INSERT INTO compute.jobs (id, simulation_id, status, input_params) + VALUES ($1, gen_random_uuid(), 'queued', '{}'::jsonb) + """, + compute_job_id, + ) + job = await queue.enqueue( + connection, + task=RUN_SIMULATION, + payload={"compute_job_id": str(compute_job_id)}, + **kwargs, + ) + job_id: uuid.UUID = job.id + return job_id + + +async def _finish( + job_id: uuid.UUID, + *, + age: timedelta, + state: str = "succeeded", + reconcile: bool = True, +) -> None: + """Put one job in a terminal state, finished `age` ago.""" + async with db.acquire() as connection: + await connection.execute( + f""" + UPDATE {SCHEMA}.jobs + SET state = $2, attempt = 1, started_at = $3, finished_at = $3, + lease_token = NULL, leased_until = NULL + WHERE id = $1 + """, # noqa: S608 - the schema is a validated identifier, not input + job_id, + state, + datetime.now(UTC) - age, + ) + if reconcile: + compute_job_id = await connection.fetchval( + f"SELECT payload->>'compute_job_id' FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + await connection.execute( + "UPDATE compute.jobs SET status = $2, finished_at = $3 WHERE id = $1", + uuid.UUID(compute_job_id), + "completed" if state == "succeeded" else "failed", + datetime.now(UTC) - age, + ) + await connection.execute( + f""" + INSERT INTO {SCHEMA}.job_attempts + (job_id, queue, task, attempt, worker_id, lease_token, + finished_at, outcome) + VALUES ($1, $2, $3, 1, 'test-worker', gen_random_uuid(), now(), $4) + """, # noqa: S608 - the schema is a validated identifier, not input + job_id, + COMPUTE_QUEUE, + RUN_SIMULATION, + state, + ) + + +# --------------------------------------------------------------- the producer + + +@pytest.mark.asyncio +async def test_the_producer_can_enqueue_and_read_but_not_transition( + queue: Queue, roles: Roles +) -> None: + connection = await roles.connect(roles.producer) + try: + job_id = uuid.uuid4() + await connection.execute( + f""" + INSERT INTO {SCHEMA}.jobs + (id, queue, task, payload, state, max_attempts) + VALUES ($1, $2, $3, '{{}}'::jsonb, 'pending', 3) + """, # noqa: S608 - the schema is a validated identifier, not input + job_id, + COMPUTE_QUEUE, + RUN_SIMULATION, + ) + assert ( + await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + == 1 + ) + + # The dedupe-conflict path is DO UPDATE SET updated_at, so this exact + # column-scoped write is the one UPDATE a producer is allowed. + await connection.execute( + f"UPDATE {SCHEMA}.jobs SET updated_at = now() WHERE id = $1", # noqa: S608 + job_id, + ) + + # Claiming is a whole-row UPDATE, which it must not be able to make. + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute( + f"UPDATE {SCHEMA}.jobs SET state = 'leased' WHERE id = $1", # noqa: S608 + job_id, + ) + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute( + f"DELETE FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + finally: + await connection.close() + + +@pytest.mark.asyncio +async def test_the_producer_reaches_compute_jobs_for_reads_and_inserts_only( + queue: Queue, roles: Roles +) -> None: + connection = await roles.connect(roles.producer) + try: + await connection.fetchval("SELECT count(*) FROM compute.jobs") + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute("UPDATE compute.jobs SET status = 'running'") + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute("DELETE FROM compute.jobs") + finally: + await connection.close() + + +# ----------------------------------------------------------------- the worker + + +@pytest.mark.asyncio +async def test_the_worker_can_transition_a_job_but_never_delete_one( + queue: Queue, roles: Roles +) -> None: + job_id = await _enqueue(queue) + connection = await roles.connect(roles.worker) + try: + await connection.execute( + f""" + UPDATE {SCHEMA}.jobs + SET state = 'leased', lease_token = gen_random_uuid(), + leased_until = now() + interval '1 minute', attempt = 1 + WHERE id = $1 + """, # noqa: S608 - the schema is a validated identifier, not input + job_id, + ) + # Purge is not the worker's to run: deleting a job takes its attempt + # history with it, which a consuming role must not be able to erase. + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute(f"DELETE FROM {SCHEMA}.jobs") # noqa: S608 + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute(f"DELETE FROM {SCHEMA}.job_attempts") # noqa: S608 + finally: + await connection.close() + + +@pytest.mark.asyncio +async def test_the_worker_updates_compute_jobs_but_cannot_create_one( + queue: Queue, roles: Roles +) -> None: + connection = await roles.connect(roles.worker) + try: + await connection.execute("UPDATE compute.jobs SET status = status") + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute( + "INSERT INTO compute.jobs (id, simulation_id, status, input_params) " + "VALUES (gen_random_uuid(), gen_random_uuid(), 'queued', '{}'::jsonb)" + ) + finally: + await connection.close() + + +@pytest.mark.asyncio +async def test_reprovisioning_removes_compute_schema_create_drift( + queue: Queue, roles: Roles +) -> None: + """Role repair must remove schema-level CREATE as well as table grants.""" + async with db.acquire() as connection: + for role in (roles.worker, roles.purger): + await connection.execute(f'GRANT CREATE ON SCHEMA compute TO "{role}"') + assert await connection.fetchval( + "SELECT has_schema_privilege($1, 'compute', 'CREATE')", role + ) + + await queue_grants.provision_queue_roles(connection) + + for role in (roles.worker, roles.purger): + assert not await connection.fetchval( + "SELECT has_schema_privilege($1, 'compute', 'CREATE')", role + ) + + +# ----------------------------------------------------------------- the purger + + +@pytest.mark.asyncio +async def test_the_purger_may_only_read_and_delete_queue_rows( + queue: Queue, roles: Roles +) -> None: + job_id = await _enqueue(queue) + connection = await roles.connect(roles.purger) + try: + assert ( + await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.jobs" # noqa: S608 + ) + == 1 + ) + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute( + f"UPDATE {SCHEMA}.jobs SET state = 'cancelled'" # noqa: S608 + ) + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute( + f""" + INSERT INTO {SCHEMA}.jobs + (id, queue, task, payload, state, max_attempts) + VALUES (gen_random_uuid(), $1, $2, '{{}}'::jsonb, 'pending', 1) + """, # noqa: S608 - the schema is a validated identifier + COMPUTE_QUEUE, + RUN_SIMULATION, + ) + # The RLS defense in depth rejects direct deletion of active work. + await connection.execute( + f"DELETE FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + assert ( + await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + == 1 + ) + await _finish(job_id, age=timedelta(days=1)) + # Retention itself, which is the one thing this role exists for. + await connection.execute( + f"DELETE FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + job_id, + ) + finally: + await connection.close() + + +@pytest.mark.asyncio +async def test_the_purger_cannot_see_the_compute_schema_at_all( + queue: Queue, roles: Roles +) -> None: + connection = await roles.connect(roles.purger) + try: + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.fetchval("SELECT count(*) FROM compute.jobs") + finally: + await connection.close() + + +# ------------------------------------------------------------- shared floors + + +@pytest.mark.asyncio +@pytest.mark.parametrize("which", ["producer", "worker", "purger"]) +async def test_no_runtime_role_can_run_ddl( + queue: Queue, roles: Roles, which: str +) -> None: + connection = await roles.connect(getattr(roles, which)) + try: + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute(f"CREATE TABLE {SCHEMA}.forbidden (id integer)") + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute(f"ALTER TABLE {SCHEMA}.jobs ADD COLUMN x integer") + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await connection.execute("CREATE SCHEMA forbidden") + finally: + await connection.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("which", ["producer", "worker", "purger"]) +async def test_a_role_is_scoped_to_this_deployments_queue( + queue: Queue, roles: Roles, which: str +) -> None: + # Row-level security, not a grant: the roles are granted COMPUTE_QUEUE + # rather than '*', so a job on another queue is invisible to all of them. + await _enqueue(queue) + async with db.acquire() as connection: + await connection.execute( + f""" + INSERT INTO {SCHEMA}.jobs (id, queue, task, payload, state, max_attempts) + VALUES (gen_random_uuid(), 'other', $1, '{{}}'::jsonb, 'pending', 3) + """, # noqa: S608 - the schema is a validated identifier, not input + RUN_SIMULATION, + ) + + connection = await roles.connect(getattr(roles, which)) + try: + visible = await connection.fetch( + f"SELECT queue FROM {SCHEMA}.jobs" # noqa: S608 + ) + finally: + await connection.close() + + assert [row["queue"] for row in visible] == [COMPUTE_QUEUE] + + +# ------------------------------------------------- the real end-to-end paths + + +@pytest.mark.asyncio +async def test_the_provisioned_roles_run_the_real_enqueue_and_claim_paths( + queue: Queue, roles: Roles +) -> None: + """A grant table that passes SQL probes but not rqueue's own SQL is no use. + + So this drives the code the two processes actually run: `Queue.enqueue` + with the dedupe conflict path the producer needs `UPDATE (updated_at)` + for, then a real `Worker` claiming, heartbeating and finalizing the job. + """ + probe = "tests.role_probe" + ran: list[str] = [] + + async def handler(_payload: Any, _context: TaskContext) -> None: + ran.append("yes") + + producer_pool = await asyncpg.create_pool( + roles.url(roles.producer), min_size=1, max_size=2 + ) + assert producer_pool is not None + try: + producer = Queue(producer_pool, name=COMPUTE_QUEUE, schema=SCHEMA) + producer.register( + name=probe, handler=handler, retry=RetryPolicy(max_attempts=1) + ) + async with producer_pool.acquire() as connection: + first = await producer.enqueue( + connection, + task=probe, + payload={}, + dedupe_key="role-probe", + on_conflict="return_existing", + ) + # The second enqueue takes ON CONFLICT ... DO UPDATE SET + # updated_at, which is why PRODUCE is granted that one column. + second = await producer.enqueue( + connection, + task=probe, + payload={}, + dedupe_key="role-probe", + on_conflict="return_existing", + ) + finally: + await producer_pool.close() + + worker_pool = await asyncpg.create_pool( + roles.url(roles.worker), min_size=1, max_size=6 + ) + assert worker_pool is not None + try: + consumer = Queue(worker_pool, name=COMPUTE_QUEUE, schema=SCHEMA) + consumer.register( + name=probe, handler=handler, retry=RetryPolicy(max_attempts=1) + ) + await Worker( + consumer, worker_id="role-probe-worker", concurrency=1, lease_duration=30.0 + ).drain(timeout=30) + finally: + await worker_pool.close() + + async with db.acquire() as connection: + row = await connection.fetchrow( + f"SELECT state, attempt FROM {SCHEMA}.jobs WHERE id = $1", # noqa: S608 + first.id, + ) + attempts = await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.job_attempts WHERE job_id = $1", # noqa: S608 + first.id, + ) + + assert first.id == second.id + assert ran == ["yes"] + assert (row["state"], row["attempt"]) == ("succeeded", 1) + assert attempts == 1 + + +# ------------------------------------------------------------------ retention + + +@pytest.mark.asyncio +async def test_purge_is_refused_on_the_worker_role(queue: Queue, roles: Roles) -> None: + """Why retention has a third role rather than riding on the worker's pool.""" + job_id = await _enqueue(queue) + await _finish(job_id, age=JOB_RETENTION + timedelta(days=1)) + + pool = await asyncpg.create_pool(roles.url(roles.worker), min_size=1, max_size=1) + assert pool is not None + try: + with pytest.raises(asyncpg.InsufficientPrivilegeError): + await purge_finished_jobs(Admin(pool, schema=SCHEMA)) + finally: + await pool.close() + + +@pytest.mark.asyncio +async def test_purge_removes_only_jobs_past_the_retention_window( + queue: Queue, roles: Roles +) -> None: + old = await _enqueue(queue) + recent = await _enqueue(queue) + pending = await _enqueue(queue) + unreconciled = await _enqueue(queue) + await _finish(old, age=JOB_RETENTION + timedelta(days=1)) + await _finish(recent, age=JOB_RETENTION - timedelta(days=1), state="failed") + await _finish( + unreconciled, + age=JOB_RETENTION + timedelta(days=1), + reconcile=False, + ) + + pool = await asyncpg.create_pool(roles.url(roles.purger), min_size=1, max_size=1) + assert pool is not None + try: + removed = await purge_finished_jobs(Admin(pool, schema=SCHEMA)) + finally: + await pool.close() + + async with db.acquire() as connection: + remaining = { + row["id"] + for row in await connection.fetch(f"SELECT id FROM {SCHEMA}.jobs") # noqa: S608 + } + orphaned = await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.job_attempts WHERE job_id = $1", # noqa: S608 + old, + ) + kept_attempts = await connection.fetchval( + f"SELECT count(*) FROM {SCHEMA}.job_attempts WHERE job_id = $1", # noqa: S608 + recent, + ) + + assert removed == 1 + # The unfinished job is untouched however old it is; only terminal rows go. + assert remaining == {recent, pending, unreconciled} + assert orphaned == 0 + assert kept_attempts == 1 diff --git a/packages/api/tests/test_tasks.py b/packages/api/tests/test_tasks.py index ea58edd..a3a3cb7 100644 --- a/packages/api/tests/test_tasks.py +++ b/packages/api/tests/test_tasks.py @@ -4,7 +4,7 @@ import threading import uuid from pathlib import Path -from typing import Any +from typing import Any, cast import pytest from rqueue import PermanentFailure, TaskContext @@ -372,6 +372,81 @@ async def sweep() -> None: assert len(passes) == 2 +@pytest.mark.asyncio +async def test_the_periodic_purge_keeps_running_after_a_failed_pass( + monkeypatch: pytest.MonkeyPatch, +) -> None: + passes: list[int] = [] + stop = asyncio.Event() + + async def purge(_admin: Any, *, compute_pool: Any = None) -> int: + passes.append(len(passes)) + if len(passes) == 1: + raise RuntimeError("purge role is not provisioned") + stop.set() + return 3 + + monkeypatch.setattr(tasks, "purge_finished_jobs", purge) + + await tasks.run_periodic_purge(cast(Any, object()), stop, interval=0.01) + + assert len(passes) == 2 + + +@pytest.mark.asyncio +async def test_purge_checks_compute_state_before_deleting_queue_rows() -> None: + job_id = uuid.uuid4() + queries: list[tuple[str, tuple[Any, ...]]] = [] + + class _Connection: + async def fetch(self, query: str, *args: Any) -> list[dict[str, Any]]: + queries.append((query, args)) + return [{"id": job_id}] + + async def fetchval(self, query: str, *args: Any) -> int: + queries.append((query, args)) + return 1 + + class _Acquire: + def __init__(self, connection: _Connection) -> None: + self.connection = connection + + async def __aenter__(self) -> _Connection: + return self.connection + + async def __aexit__(self, *_exc: object) -> None: + return None + + class _Pool: + def __init__(self, connection: _Connection) -> None: + self.connection = connection + + def acquire(self) -> _Acquire: + return _Acquire(self.connection) + + class _Admin: + schema = "task_queue" + + def __init__(self, pool: _Pool) -> None: + self.pool = pool + + pool = _Pool(_Connection()) + removed = await tasks.purge_finished_jobs( + cast(Any, _Admin(pool)), compute_pool=cast(Any, pool) + ) + + assert removed == 1 + assert len(queries) == 2 + assert "JOIN compute.jobs" in queries[0][0] + assert queries[0][1][0] == COMPUTE_QUEUE + assert queries[0][1][-1] == tasks.PURGE_LIMIT + assert "DELETE FROM task_queue.jobs" in queries[1][0] + assert "state = ANY($3::text[])" in queries[1][0] + assert "finished_at < $4" in queries[1][0] + assert queries[1][1][0] == [job_id] + assert queries[1][1][1] == COMPUTE_QUEUE + + @pytest.mark.asyncio async def test_the_periodic_reconcile_keeps_running_after_a_failed_pass( monkeypatch: pytest.MonkeyPatch, diff --git a/readme.md b/readme.md index bc45748..018348d 100644 --- a/readme.md +++ b/readme.md @@ -25,8 +25,9 @@ uv run tsdhn doctor uv run tsdhn calc --mw 8.0 --lat -20.5 --lon -70.5 ``` -Create `.env`, set `COMPUTE_API_TOKEN`, `BETTER_AUTH_SECRET`, and -`APP_DB_PASSWORD`, then run the self-hosted stack: +Create `.env` and fill in `COMPUTE_API_TOKEN`, `BETTER_AUTH_SECRET`, +`APP_DB_PASSWORD`, and the three queue-role passwords it lists, then run the +self-hosted stack: ```sh cp .env.example .env diff --git a/scripts/database.py b/scripts/database.py index e705df5..1795a56 100644 --- a/scripts/database.py +++ b/scripts/database.py @@ -78,9 +78,18 @@ def drop_database(base_url: str, name: str) -> None: def drop_role(base_url: str, role: str) -> None: - """Remove a role created for an isolated integration run.""" + """Remove a role created for an isolated integration run. + + Cleanup also runs when provisioning failed before the role was created; + ``DROP OWNED BY`` has no ``IF EXISTS`` form. + """ target = database_target(base_url, role) with psycopg.connect(target.maintenance_url, autocommit=True) as connection: + exists = connection.execute( + "SELECT 1 FROM pg_roles WHERE rolname = %s", [role] + ).fetchone() + if exists is None: + return identifier = sql.Identifier(role) connection.execute(sql.SQL("DROP OWNED BY {}").format(identifier)) connection.execute(sql.SQL("DROP ROLE IF EXISTS {}").format(identifier)) @@ -103,7 +112,7 @@ def main() -> None: drop = subparsers.add_parser("drop") drop.add_argument("--base-url", required=True) drop.add_argument("--name", required=True) - drop.add_argument("--role") + drop.add_argument("--role", action="append", default=[]) args = parser.parse_args() if args.command == "create": @@ -119,8 +128,8 @@ def main() -> None: ) else: drop_database(args.base_url, args.name) - if args.role: - drop_role(args.base_url, args.role) + for role in args.role: + drop_role(args.base_url, role) if __name__ == "__main__": diff --git a/scripts/integration.sh b/scripts/integration.sh index 621a930..2d29a21 100755 --- a/scripts/integration.sh +++ b/scripts/integration.sh @@ -7,6 +7,13 @@ database_name="tsdhn_integration_$(date +%s)_$$" app_role="${database_name}_role" app_password="tsdhn-web-test-password" queue_schema="${COMPUTE_QUEUE_SCHEMA:-task_queue}" +queue_name="${COMPUTE_QUEUE:-simulations}" + +# Roles are cluster-wide, so each disposable database gets unique role names. +producer_role="${database_name}_producer" +worker_role="${database_name}_worker" +purger_role="${database_name}_purger" +queue_password="tsdhn-queue-test-password" case "${1:-}" in "") coverage=0 ;; @@ -36,7 +43,10 @@ cleanup() { uv run python -m scripts.database drop \ --base-url "$base_url" \ --name "$database_name" \ - --role "$app_role" >/dev/null + --role "$app_role" \ + --role "$producer_role" \ + --role "$worker_role" \ + --role "$purger_role" >/dev/null } trap cleanup EXIT @@ -47,6 +57,17 @@ uv run tsdhn-compute-migrate uv run rqueue --database-url "$admin_url" --schema "$queue_schema" migrate +COMPUTE_DATABASE_URL="$admin_url" \ +COMPUTE_QUEUE="$queue_name" \ +COMPUTE_QUEUE_SCHEMA="$queue_schema" \ +COMPUTE_PRODUCER_ROLE="$producer_role" \ +COMPUTE_PRODUCER_PASSWORD="$queue_password" \ +COMPUTE_WORKER_ROLE="$worker_role" \ +COMPUTE_WORKER_PASSWORD="$queue_password" \ +COMPUTE_PURGER_ROLE="$purger_role" \ +COMPUTE_PURGER_PASSWORD="$queue_password" \ +uv run tsdhn-queue-grants + # Schema changes run with the database-owner connection. The app role below # is intentionally limited to runtime DML and compute-state reads. DATABASE_URL="$admin_url" bun --filter web db:migrate