diff --git a/api/ROW_LEVEL_SECURITY.md b/api/ROW_LEVEL_SECURITY.md new file mode 100644 index 0000000..e69d39e --- /dev/null +++ b/api/ROW_LEVEL_SECURITY.md @@ -0,0 +1,159 @@ +# Row-level security + +Postgres row-level security (RLS) enforces who can read and change +experiment data and project membership, in the database, in addition to the +API's own checks. This page covers the access model, the database roles, and +how to roll it out or roll it back. + +## Database roles + +| Role | Used by | Privileges | +|---|---|---| +| `api_client` | Every user request, via `get_pg_connection(user)` | Table grants only; subject to RLS. The connection carries the caller in session settings. | +| `api_worker` | Background and system jobs, via `get_pg_worker_connection()`: ingestion, optimizer jobs, model registry, audits | Data access on all tables, explicit `TO api_worker` policies on RLS tables. Not superuser, not `BYPASSRLS`, owns nothing, cannot run DDL. | +| owner (`postgres`) | Schema migrations only, via `get_pg_migration_connection()` | Table owner. Never used to serve requests. | + +A user-scoped connection sets three session settings that the policies read: +`app.current_user` (user id), `app.current_user_role` (`admin`, `pi`, +`equipment_owner`, `researcher`) and `app.current_user_org`. They are set at +session scope and committed, so a commit or rollback part-way through a +request cannot drop the caller's identity. + +Until `PG_WORKER_PASS` is configured, `get_pg_worker_connection()` falls back +to the owner credentials and logs a warning once, so existing deployments keep +working. + +## Access rules + +The policies mirror the API's existing checks +(`api/scripts/add_row_level_security.sql`): + +- **Projects:** members and admins see private projects; anyone signed in sees + open projects. PIs and admins create projects; a project's PI or an admin + deletes it. +- **Project members:** visible to anyone who can read the project (as before) + and to admins. A PI adds plain members; the creator of a project enrols + themselves as its PI in the same transaction that created it; only admins + change PI assignments. A PI removes non-PI members. Always enforced. +- **Experiment definitions:** readable by the owner, anyone who can read the + project, and facility managers (equipment owners) for execution requests on + their equipment. Created by the owner into projects they can write, with a + template they can read. A trigger fixes an experiment's owner, project, + equipment, template and request id after creation (admins and background + jobs excepted), and limits facility managers' updates to the execution + status columns. +- **Experiment templates:** published templates are public; otherwise the + owner, shared templates, the owner's organization, admins, and anyone who can + read a project or experiment that uses the template. Only the owner or an + admin edits or deletes. Version history follows template visibility. +- **Samples and sample links:** follow project access; facility managers can + read samples linked to execution requests they manage. +- **Recipe batches and proposals:** readable with the experiment; written only + by the proposal worker; status updates by the experiment owner, project + members or admins. +- **Run tables** (`etcher_runs`, `equipment_runs`, trace tables, run files): + unchanged member/open read policies, plus admin read and worker access. + +Two operations deliberately cross tenants and are implemented as narrow +`SECURITY DEFINER` functions that re-check permission in the database: +`app_detach_experiment_type` (deleting a template clears references in other +projects) and `app_experiment_project_name` (a facility manager sees the name +of the project behind an execution request, and nothing else about it). + +Operations that intentionally act across all tenants run on the worker +connection after the API's own authorization: ingestion and run sync, the +optimizer and follow-up loops, equipment deletion (which detaches projects), +equipment fleet metrics, process-definition reads, identity merges and model +registry work. + +## Configuration + +| Variable | Where | Meaning | +|---|---|---| +| `PG_WORKER_USER`, `PG_WORKER_PASS` | API | `api_worker` credentials for background jobs. | +| `PG_SUPER_USER`, `PG_SUPER_PASS` | migrations (and the API until the worker role exists) | Owner credentials. | +| `DB_MIGRATE_ON_STARTUP` | API | `false` skips migrations at startup; default runs them when owner credentials are present. | +| `DB_RLS_ENFORCEMENT` | wherever migrations run | `enforce` turns RLS on for the nine experiment tables; `disable` turns it off for them (rollback); unset leaves them unchanged. `project_members` and the six original tables are always enforced. | +| `DB_MIGRATION_LOCK_TIMEOUT` | wherever migrations run | Maximum wait for a migration lock (default `30s`). | +| `PROPOSAL_JOB_SWEEP_SECONDS` | API | How often each API worker runs queued/expired proposal jobs (default 60, `0` disables). | +| `PROPOSAL_JOB_LEASE_SECONDS`, `PROPOSAL_JOB_MAX_ATTEMPTS`, `PROPOSAL_JOB_MAX_AGE_HOURS` | API | Proposal job lease (900 s), retry limit (3) and maximum age (72 h). | + +## Rollout + +Each step is safe on its own and can be verified before the next. + +1. **Deploy the new API image.** Startup migrations add the new columns, + helper functions, grants, policies and guards, and enforce RLS on + `project_members` (together with the membership write grants). Enforcement + on the nine experiment tables stays off (`DB_RLS_ENFORCEMENT` unset); the + API's own checks, including the new template-visibility checks, still + apply, as do the atomic experiment creation, durable proposal jobs and + idempotent retries. Check `GET /api/health` → `postgres.row_level_security`: + `projects` and `project_members` are `true`, the experiment tables `false`. +2. **Provision the worker role.** As `postgres`: + `ALTER ROLE api_worker WITH LOGIN PASSWORD '';` + Add `PG_WORKER_USER`/`PG_WORKER_PASS` to `dt-api-secret` and restart the API. + Health shows `runtime.postgres_worker_role_configured: true`, and the + "falling back to migration credentials" warning disappears from the logs. +3. **Enforce RLS.** Set `DB_RLS_ENFORCEMENT=enforce` in `dt-api-secret` (or in + the migration Job's secret) and restart / re-run migrations. Health shows + `row_level_security` all `true`. Smoke-test as a PI, a researcher, an + equipment owner and an admin: projects, experiments, samples, recipe + proposals and the facility execution queue. +4. **Take owner credentials out of the API.** Create `dt-db-migration-secret` + (see `geddes/k8s/01-secrets.yaml.example`), run migrations with + `geddes/k8s/09-db-migrate-job.yaml` before each rollout, set + `DB_MIGRATE_ON_STARTUP=false` and remove `PG_SUPER_USER`/`PG_SUPER_PASS` + from `dt-api-secret`. Health stays `ok` without them. + +## Rollback + +- **RLS causing problems:** set `DB_RLS_ENFORCEMENT=disable` and re-run + migrations (restart the API while it still runs them, or run the Job). This + turns RLS off on the nine experiment tables; `project_members` and the six + original tables stay enforced, and the API's own checks remain. +- **Worker role problems:** remove `PG_WORKER_PASS`; background jobs fall back + to the owner credentials (requires `PG_SUPER_PASS` in the API secret). +- **Rolling back to an API image from before this change:** turn experiment + enforcement off *first*, while the new image is still running + (`DB_RLS_ENFORCEMENT=disable`, restart), then roll back the image. The old + code creates experiments and templates without a user identity, which the + enforced policies reject. If the new image is already gone, run as the owner: + + ```sql + ALTER TABLE experiment_definitions DISABLE ROW LEVEL SECURITY; + ALTER TABLE experiment_types DISABLE ROW LEVEL SECURITY; + ALTER TABLE experiment_type_versions DISABLE ROW LEVEL SECURITY; + ALTER TABLE project_experiments DISABLE ROW LEVEL SECURITY; + ALTER TABLE samples DISABLE ROW LEVEL SECURITY; + ALTER TABLE experiment_samples DISABLE ROW LEVEL SECURITY; + ALTER TABLE run_samples DISABLE ROW LEVEL SECURITY; + ALTER TABLE experiment_recipe_batches DISABLE ROW LEVEL SECURITY; + ALTER TABLE experiment_recipe_proposals DISABLE ROW LEVEL SECURITY; + ``` + + After rollout step 4 (owner credentials removed from the API and + `DB_MIGRATE_ON_STARTUP=false`), restarting the API no longer changes + enforcement: set `DB_RLS_ENFORCEMENT=disable` in `dt-db-migration-secret` + and run the migration Job (or the SQL above as the owner), confirm health + shows the nine tables as `false`, and put `PG_SUPER_USER`/`PG_SUPER_PASS` + back into `dt-api-secret`, because older images require them to start. + Roll the `dt-api-retrain` CronJob image back together with the Deployment. + + The old code works with everything else this migration leaves behind, + including enforced `project_members` (verified by running the previous API + code against a migrated database). + +## Applying the SQL by hand + +Prefer `python -m db_migrations`, which applies the script in one +transaction. If you apply it with psql, keep it atomic: +`psql --single-transaction -v ON_ERROR_STOP=1 -f api/scripts/add_row_level_security.sql` + +## Verifying locally + +`api/tests/test_row_level_security_pg.py` runs the migrations with enforcement +on and exercises the API functions as different users against a disposable +Postgres 15. It is skipped unless `DT_RLS_TEST_PG_PORT` is set and refuses to +run against a database without an `rls_test_marker` table. See the module +docstring for setup. diff --git a/api/ai_readiness.py b/api/ai_readiness.py index 92590ef..9f42969 100644 --- a/api/ai_readiness.py +++ b/api/ai_readiness.py @@ -93,14 +93,11 @@ def export_ml_matrix( "column_types": {"run_id": "int", ...}, } """ - from data_loader_pg import get_pg_connection, get_pg_superuser_connection + from data_loader_pg import _scoped_connection from psycopg2.extras import RealDictCursor try: - if is_admin: - conn = get_pg_superuser_connection() - else: - conn = get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) # Build the flat query — no nested JSON, just raw columns select_cols = METADATA_COLUMNS + FEATURE_COLUMNS + TARGET_COLUMNS if include_metadata else FEATURE_COLUMNS + TARGET_COLUMNS @@ -220,7 +217,7 @@ def record_dataset_snapshot( SHA-256 hash of the full dataset. This lets you trace any model back to the exact data that trained it. """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection snapshot_record = { "snapshot_hash": snapshot_hash, @@ -231,7 +228,7 @@ def record_dataset_snapshot( } try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor() as cur: cur.execute( """ @@ -271,11 +268,11 @@ def record_dataset_snapshot( def get_latest_snapshots(limit: int = 10) -> list[dict[str, Any]]: """Fetch the most recent dataset snapshots.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( """ @@ -312,11 +309,11 @@ def generate_data_card( - Temporal coverage (runs per month) - Known limitations / bias warnings """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: # ── Overview statistics ────────────────────────────────────── diff --git a/api/data_loader_pg.py b/api/data_loader_pg.py index 5c51454..191a87d 100644 --- a/api/data_loader_pg.py +++ b/api/data_loader_pg.py @@ -7,6 +7,7 @@ from typing import List, Dict, Any, Optional import re import uuid +from types import SimpleNamespace import psycopg2 from psycopg2.extras import Json, RealDictCursor import pandas as pd @@ -14,6 +15,7 @@ MissingRuntimeConfiguration, get_pg_api_password, get_pg_superuser_password, + get_pg_worker_password, ) from db_mock import get_mock_connection @@ -56,52 +58,164 @@ def _default_pg_host() -> str: return "localhost" -def _get_pg_connection_kwargs() -> dict[str, str]: +def _pg_endpoint_kwargs() -> dict[str, str]: return { "host": os.getenv("PG_HOST", _default_pg_host()), "port": os.getenv("PG_PORT", "5432"), "dbname": os.getenv("PG_DB", "digital_twin"), + } + + +def _get_pg_connection_kwargs() -> dict[str, str]: + return { + **_pg_endpoint_kwargs(), "user": os.getenv("PG_USER", "api_client"), "password": get_pg_api_password(), + "application_name": "dt-api", } def _get_pg_superuser_connection_kwargs() -> dict[str, str]: return { - "host": os.getenv("PG_HOST", _default_pg_host()), - "port": os.getenv("PG_PORT", "5432"), - "dbname": os.getenv("PG_DB", "digital_twin"), + **_pg_endpoint_kwargs(), "user": os.getenv("PG_SUPER_USER", "postgres"), "password": get_pg_superuser_password(), + "application_name": "dt-api-migrations", + } + + +def _get_pg_migration_connection_kwargs() -> dict[str, str]: + # Schema changes take exclusive locks; give up instead of queueing behind + # long-running queries (and blocking everything queued behind the DDL). + lock_timeout = os.getenv("DB_MIGRATION_LOCK_TIMEOUT", "30s").strip() or "30s" + return { + **_get_pg_superuser_connection_kwargs(), + "options": f"-c lock_timeout={lock_timeout}", } -def get_pg_superuser_connection(): + +_worker_fallback_warned = False + + +def _get_pg_worker_connection_kwargs() -> dict[str, str]: + global _worker_fallback_warned + worker_password = get_pg_worker_password() + if worker_password: + return { + **_pg_endpoint_kwargs(), + "user": os.getenv("PG_WORKER_USER", "api_worker"), + "password": worker_password, + "application_name": "dt-api-worker", + } + if not _worker_fallback_warned: + logger.warning( + "PG_WORKER_PASS is not configured; background jobs are using the " + "migration (table-owner) credentials. Provision the api_worker role " + "and remove PG_SUPER_PASS from the API environment." + ) + _worker_fallback_warned = True + kwargs = _get_pg_superuser_connection_kwargs() + kwargs["application_name"] = "dt-api-worker" + return kwargs + + +# Session settings read by the row-level security policies. They are set at +# session scope (not transaction scope) so a commit or rollback part-way +# through a request cannot silently drop the caller's identity. +SESSION_USER_SETTING = "app.current_user" +SESSION_ROLE_SETTING = "app.current_user_role" +SESSION_ORG_SETTING = "app.current_user_org" + + +def _session_identity(identity: Any, is_admin: bool) -> tuple[str, str, str]: + if identity is None: + return "", "admin" if is_admin else "", "" + if isinstance(identity, str): + return identity.strip(), "admin" if is_admin else "", "" + user_id = str(getattr(identity, "id", "") or "").strip() + role = "admin" if is_admin else str(getattr(identity, "role", "") or "").strip() + organization = str(getattr(identity, "organization", "") or "").strip() + return user_id, role, organization + + +def get_pg_connection(identity: Any = None, *, is_admin: bool = False): """ - Establish a connection as the postgres table owner. - This bypasses RLS and is used exclusively by the protected ingestion endpoint. + Open an ``api_client`` connection scoped to the requesting user. + + ``identity`` is a ``PlatformUser`` (preferred: carries id, role and + organization) or a bare user id. Row-level security policies read the + resulting ``app.current_user*`` session settings; a connection opened + without an identity only sees rows the policies expose publicly. """ if is_mock_db_enabled(): return get_mock_connection() - - conn = psycopg2.connect(**_get_pg_superuser_connection_kwargs()) + + conn = psycopg2.connect(**_get_pg_connection_kwargs()) + user_id, role, organization = _session_identity(identity, is_admin) + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT + set_config(%s, %s, false), + set_config(%s, %s, false), + set_config(%s, %s, false) + """, + ( + SESSION_USER_SETTING, + user_id, + SESSION_ROLE_SETTING, + role, + SESSION_ORG_SETTING, + organization, + ), + ) + # Commit so a later rollback cannot revert the session settings. + conn.commit() + except Exception: + conn.close() + raise return conn -def get_pg_connection(nanohub_user_id: Optional[str] = None): + +def _scoped_connection(nanohub_user_id: Optional[str], is_admin: bool): + """ + Connection for a read that accepts ``(nanohub_user_id, is_admin)``. + + Requests always carry a user id, so they get a user-scoped ``api_client`` + connection (admins pass the row-level security admin policies). An admin + call without a user id comes from a system job (retraining, snapshots) and + uses the worker connection. """ - Establish a connection to PostgreSQL and set the session variable to the user ID. - This automatically triggers Row-Level Security! + if is_admin and not str(nanohub_user_id or "").strip(): + return get_pg_worker_connection() + return get_pg_connection(nanohub_user_id, is_admin=is_admin) + + +def get_pg_worker_connection(): + """ + Open a connection for background and system jobs (ingestion, optimizer + jobs, model registry, audits). + + Connects as the ``api_worker`` role, which is neither superuser nor + BYPASSRLS; the row-level security migration grants it explicit per-table + policies. Falls back to the migration credentials until ``PG_WORKER_PASS`` + is provisioned. """ if is_mock_db_enabled(): return get_mock_connection() + return psycopg2.connect(**_get_pg_worker_connection_kwargs()) - conn = psycopg2.connect(**_get_pg_connection_kwargs()) - - if nanohub_user_id: - with conn.cursor() as cur: - # Set local variable for RLS using set_config safely - cur.execute("SELECT set_config('app.current_user', %s, true);", (nanohub_user_id,)) - - return conn + +def get_pg_migration_connection(): + """ + Open a table-owner connection for schema migrations (DDL) only. + + Request handling and background jobs must never use this connection. + """ + if is_mock_db_enabled(): + return get_mock_connection() + return psycopg2.connect(**_get_pg_migration_connection_kwargs()) def _json_safe_payload(row: Dict[str, Any]) -> Dict[str, Any]: return json.loads(json.dumps(row, default=str)) @@ -298,7 +412,7 @@ def _attach_run_file_refs(conn, records: List[Dict[str, Any]]) -> None: def ensure_etcher_run_file_refs_pg() -> None: - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -373,7 +487,7 @@ def ensure_etcher_recipe_name_column_pg() -> None: Idempotent — safe to call on every startup. """ - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -394,7 +508,7 @@ def ensure_equipment_runs_pg() -> None: visibility model used by ``etcher_runs`` so generic runs respect the same open/shared/private project access rules. """ - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -578,7 +692,7 @@ def ensure_equipment_run_trace_samples_pg() -> None: provide. ``sample_time_raw`` preserves the exact source text (including any future fractional seconds or explicit offset) for provenance and display. """ - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -665,7 +779,7 @@ def ensure_equipment_run_trace_samples_pg() -> None: def ensure_equipment_run_trace_events_pg() -> None: """Create the normalized provenance table for events attached to a run.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -742,7 +856,7 @@ def ensure_equipment_run_trace_events_pg() -> None: def ensure_glance_ingestion_audit_pg() -> None: """Create connector audit storage; it is intentionally not API-readable.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -785,7 +899,7 @@ def ensure_glance_ingestion_audit_pg() -> None: def record_glance_ingestion_audit_pg(**record: Any) -> None: """Persist one redacted connector handoff result independently of trace data.""" - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -825,7 +939,7 @@ def record_glance_ingestion_audit_pg(**record: Any) -> None: def get_glance_ingestion_project_pg(project_id: str) -> Optional[Dict[str, Any]]: """Return the authoritative project/equipment assignment for ingestion.""" - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -1007,7 +1121,7 @@ def sync_equipment_runs_pg( if not rows: return 0 - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: insert_rows = [] for index, row in enumerate(rows): @@ -1510,7 +1624,7 @@ def sync_equipment_run_traces_pg( # therefore the same exposed catalog id on a retry. normalized_runs.sort(key=lambda run: run["source_run_id"]) - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: identity_lock = ( @@ -2177,7 +2291,7 @@ def merge_equipment_runs_pg( side_key = "inputs" if kind == "input" else "outputs" raw_key = "input_raw" if kind == "input" else "output_raw" - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: affected = 0 with conn.cursor() as cur: @@ -2415,8 +2529,8 @@ def get_runs_list_pg( Normal callers connect as ``api_client`` and are scoped by RLS using ``app.current_user`` (the project_members membership of ``nanohub_user_id``). - Admins (``is_admin=True``) connect via the superuser connection, which - bypasses RLS so they can see every run for moderation/oversight. + Admins (``is_admin=True``) use the same connection with the admin role + set, which the admin RLS policy lets see every run for moderation. ``project_id`` (when provided) filters to a single project. RLS still applies to non-admins, so a PI asking for a project they don't belong to @@ -2424,10 +2538,7 @@ def get_runs_list_pg( list for any private/shared project they can't see. """ try: - if is_admin: - conn = get_pg_superuser_connection() - else: - conn = get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) # Build the shared WHERE clause for both physical run tables. The # column names line up (is_outlier, is_calibration_recipe, project_id, @@ -2566,7 +2677,7 @@ def get_run_detail_pg( Fetch a single etcher_runs row by ``idruns``, joined with project context. Visibility rules match the list endpoint: - - admins use the superuser connection (RLS bypass) + - admins pass the admin RLS policy - everyone else uses the api_client connection with app.current_user set, so RLS filters out runs the user isn't allowed to see @@ -2576,10 +2687,7 @@ def get_run_detail_pg( can't be probed by guessing run IDs). """ try: - if is_admin: - conn = get_pg_superuser_connection() - else: - conn = get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) # run_id values at/above the offset belong to the generic # equipment_runs table (see EQUIPMENT_RUN_ID_OFFSET); everything below @@ -2676,7 +2784,7 @@ def get_equipment_run_trace_pg( requested_keys.append(key) seen_requested.add(key) - conn = get_pg_superuser_connection() if is_admin else get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) try: internal_run_id = run_id - EQUIPMENT_RUN_ID_OFFSET with conn.cursor(cursor_factory=RealDictCursor) as cur: @@ -2985,9 +3093,7 @@ def list_recipe_equipment_pg( ) -> List[Dict[str, Any]]: """List equipment with recipe files visible under catalog RLS.""" - conn = get_pg_superuser_connection() if is_admin else get_pg_connection( - nanohub_user_id - ) + conn = _scoped_connection(nanohub_user_id, is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3045,9 +3151,7 @@ def list_equipment_recipes_pg( raise ValueError("equipment_id is required") safe_limit = max(1, min(int(limit), 1_000)) safe_offset = max(0, int(offset)) - conn = get_pg_superuser_connection() if is_admin else get_pg_connection( - nanohub_user_id - ) + conn = _scoped_connection(nanohub_user_id, is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3201,9 +3305,7 @@ def get_equipment_recipe_pg( if int(source_recipe_id) < 1: raise ValueError("source_recipe_id must be positive") - conn = get_pg_superuser_connection() if is_admin else get_pg_connection( - nanohub_user_id - ) + conn = _scoped_connection(nanohub_user_id, is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3281,13 +3383,10 @@ def get_projects_list_pg(nanohub_user_id: Optional[str] = None, is_admin: bool = Fetch projects from Postgres. Normal callers are RLS-scoped by their project_members membership. - Admins bypass RLS via the superuser connection so they see every project. + Admins pass the admin RLS policy so they see every project. """ try: - if is_admin: - conn = get_pg_superuser_connection() - else: - conn = get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) query = """ SELECT @@ -3326,6 +3425,10 @@ def get_projects_list_pg(nanohub_user_id: Optional[str] = None, is_admin: bool = return [] +def _creator_identity(nanohub_user_id: str, role: str) -> SimpleNamespace: + return SimpleNamespace(id=nanohub_user_id, role=role or "", organization="") + + def create_project_pg( *, name: str, @@ -3335,23 +3438,28 @@ def create_project_pg( access_mode: str, nanohub_user_id: str, pi_name: str, + creator_role: str = "pi", ) -> Dict[str, Any]: """ Create a project and grant the creator membership as the PI. - Uses the superuser connection because the read-only app user is subject to RLS. + + Runs as the creator. Row-level security only lets PIs and admins insert + projects, and only lets the creator enrol themselves as PI of a project + that has no members yet. The row is read back after the PI membership + exists because a private project is invisible to its creator before that. """ project_slug = (re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_") or "project")[:40].strip("_") or "project" project_id = f"{project_slug}_{uuid.uuid4().hex[:8]}" - conn = get_pg_superuser_connection() + conn = get_pg_connection( + _creator_identity(nanohub_user_id, creator_role), + ) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( """ INSERT INTO projects (id, name, description, equipment_id, equipment_name, pi_name, access_mode) VALUES (%s, %s, %s, %s, %s, %s, %s) - RETURNING id, name, description, equipment_id, equipment_name, pi_name as pi, - access_mode as access, created_at as created """, ( project_id, @@ -3363,16 +3471,24 @@ def create_project_pg( access_mode, ), ) - project_row = cur.fetchone() cur.execute( """ INSERT INTO project_members (project_id, nanohub_user_id, role) VALUES (%s, %s, %s) - ON CONFLICT (project_id, nanohub_user_id) DO UPDATE SET role = EXCLUDED.role """, (project_id, nanohub_user_id, "pi"), ) + cur.execute( + """ + SELECT id, name, description, equipment_id, equipment_name, pi_name as pi, + access_mode as access, created_at as created + FROM projects + WHERE id = %s + """, + (project_id,), + ) + project_row = cur.fetchone() conn.commit() finally: @@ -3413,10 +3529,12 @@ def delete_project_pg( equipment did and must not disappear just because someone removes a project. - Uses the superuser connection so an admin can delete a project they - aren't a member of, and so RLS cannot silently turn this into a no-op. + Runs as the caller: row-level security lets admins delete any project + and PIs delete their own, and a project the caller cannot see reports + not_found. Foreign-key actions (the cascade and SET NULL above) are + enforced by Postgres regardless of row-level security. """ - conn = get_pg_superuser_connection() + conn = get_pg_connection(nanohub_user_id, is_admin=is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3471,8 +3589,8 @@ def delete_project_pg( # - The recorded PI of the project (``project_members.role = 'pi'``) # can manage members on their own project only. # The helpers raise PermissionError so the router can translate to HTTP 403. -# We use the superuser connection for writes because RLS will hide projects -# the caller doesn't own otherwise, which would mask a real "not found". +# They run as the caller, so row-level security enforces the same rules in the +# database; a private project the caller cannot see reports "not found". def _resolve_project_pi(cur, project_id: str) -> Dict[str, Any] | None: @@ -3517,7 +3635,7 @@ def list_project_members_pg( List members of a project. Visible to any member of the project; admins always see it. Researchers who are not members get PermissionError. """ - conn = get_pg_superuser_connection() + conn = get_pg_connection(nanohub_user_id, is_admin=is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: project = _resolve_project_pi(cur, project_id) @@ -3585,7 +3703,7 @@ def add_project_member_pg( f"Cannot assign role '{member_role}' through project invite" ) - conn = get_pg_superuser_connection() + conn = get_pg_connection(nanohub_user_id, is_admin=is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: project = _resolve_project_pi(cur, project_id) @@ -3654,7 +3772,7 @@ def remove_project_member_pg( removed via this path — use ``delete_project_pg`` instead, which has the correct cascade semantics for the historical-runs invariant. """ - conn = get_pg_superuser_connection() + conn = get_pg_connection(nanohub_user_id, is_admin=is_admin) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: project = _resolve_project_pi(cur, project_id) @@ -3700,7 +3818,7 @@ def remove_project_member_pg( def sync_runs_pg(runs_data: List[Dict[str, Any]]): """ - Bulk insert runs from Azure pipeline using the superuser connection. + Bulk insert runs from Azure pipeline using the worker connection. Also automatically seeds missing projects. """ from psycopg2.extras import execute_values @@ -3832,7 +3950,7 @@ def sync_runs_pg(runs_data: List[Dict[str, Any]]): if file_ref is not None: file_ref_rows.append(file_ref) - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor() as cur: # The legacy scalar table has only one row per run number. Guard # that limitation under the global run lock shared with full-trace @@ -4004,7 +4122,7 @@ def get_training_df_pg( Fetch etcher runs from Postgres and return a pandas DataFrame with the SAME column names that ml_engine.py / data_loader.py expect. - Uses the superuser connection only for admin/maintenance callers. User + Uses the worker connection only for system callers (admin without a user id). User requests pass nanohub_user_id with is_admin=False so Postgres RLS limits training rows to projects the caller can actually read. With ``project_id`` it trains from that project, optionally plus unassigned legacy history. @@ -4018,7 +4136,7 @@ def get_training_df_pg( if not is_admin and not nanohub_user_id: logger.warning("Non-admin training data request missing user context.") return None - conn = get_pg_superuser_connection() if is_admin else get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) numeric_pattern = r"^[+-]?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$" @@ -4312,7 +4430,7 @@ def get_etcher_run_stats_pg() -> Dict[str, Any]: """ Total run count and most recent run timestamp for the etcher. - Always uses the superuser connection so the equipment detail card shows + Always uses the worker connection so the equipment detail card shows the *true* fleet-wide count, independent of the caller's RLS scope. The etcher_runs table has no equipment_id column today — every row in it belongs to the etcher — so a plain COUNT(*) is the correct answer. @@ -4322,7 +4440,7 @@ def get_etcher_run_stats_pg() -> Dict[str, Any]: can render "—" without crashing. """ try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( "SELECT count(*) AS total_runs, max(run_date) AS last_data_at FROM etcher_runs" @@ -4348,7 +4466,7 @@ def get_summary_stats_pg( Uses user-scoped connection so counts reflect RLS visibility. """ try: - conn = get_pg_superuser_connection() if is_admin else get_pg_connection(nanohub_user_id) + conn = _scoped_connection(nanohub_user_id, is_admin) with conn.cursor(cursor_factory=RealDictCursor) as cur: # Aggregate canonical etcher_runs together with generic equipment_runs diff --git a/api/db_migrations.py b/api/db_migrations.py new file mode 100644 index 0000000..9f190ab --- /dev/null +++ b/api/db_migrations.py @@ -0,0 +1,241 @@ +""" +Schema migrations for the platform database. + +Everything here runs with the migration (table-owner) credentials, +``PG_SUPER_USER`` / ``PG_SUPER_PASS``. Request handling and background jobs +never use them: requests connect as ``api_client`` scoped to the caller, and +background jobs as ``api_worker``. + +Run from the API process at startup (when the migration credentials are +present and ``DB_MIGRATE_ON_STARTUP`` is not ``false``) or on its own, for +example from a Kubernetes Job, so the API can run without owner credentials: + + python -m db_migrations +""" +import hashlib +import logging +import os +from pathlib import Path +from typing import Any + +from runtime_config import pg_migration_credentials_configured + +logger = logging.getLogger("dt.api.migrations") + +STARTUP_MIGRATION_LOCK_ID = 6420260506 +RLS_SCRIPT_PATH = Path(__file__).resolve().parent / "scripts" / "add_row_level_security.sql" +RLS_MIGRATION_NAME = "row_level_security" + +# Tables whose row-level security is switched by DB_RLS_ENFORCEMENT. The six +# tables that already enforced RLS before this migration (projects and the run +# tables) and project_members (enforced by the script itself, together with +# its write grants) stay enforced regardless of the switch. +RLS_ENFORCED_TABLES = ( + "experiment_definitions", + "experiment_types", + "experiment_type_versions", + "project_experiments", + "samples", + "experiment_samples", + "run_samples", + "experiment_recipe_batches", + "experiment_recipe_proposals", +) + + +# Reported by /api/health. +RLS_REPORTED_TABLES = ("projects", "project_members") + RLS_ENFORCED_TABLES + + +def rls_enforcement_mode() -> str: + """ + ``enforce`` enables RLS on RLS_ENFORCED_TABLES, ``disable`` turns it off + (the rollback switch), anything else leaves the current state unchanged. + """ + mode = os.getenv("DB_RLS_ENFORCEMENT", "").strip().lower() + return mode if mode in {"enforce", "disable"} else "" + + +def should_run_startup_migrations() -> bool: + setting = os.getenv("DB_MIGRATE_ON_STARTUP", "").strip().lower() + if setting in {"false", "0", "no"}: + return False + if setting in {"true", "1", "yes"}: + return True + return pg_migration_credentials_configured() + + +def _apply_rls_enforcement(cur, mode: str) -> dict[str, bool]: + cur.execute( + """ + SELECT c.relname, c.relrowsecurity + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relkind = 'r' + AND c.relname = ANY(%s) + """, + (list(RLS_ENFORCED_TABLES),), + ) + state = {name: bool(enabled) for name, enabled in cur.fetchall()} + for table_name in RLS_ENFORCED_TABLES: + if table_name not in state: + logger.warning("RLS enforcement skipped: table %s does not exist", table_name) + continue + enabled = state[table_name] + # ALTER TABLE takes an exclusive lock, so only issue it on a change. + if mode == "enforce" and not enabled: + cur.execute(f'ALTER TABLE public."{table_name}" ENABLE ROW LEVEL SECURITY') + state[table_name] = True + elif mode == "disable" and enabled: + cur.execute(f'ALTER TABLE public."{table_name}" DISABLE ROW LEVEL SECURITY') + state[table_name] = False + return state + + +def ensure_row_level_security_pg() -> dict[str, Any]: + """ + Apply scripts/add_row_level_security.sql and the DB_RLS_ENFORCEMENT mode. + + The script is re-applied only when its checksum changes, so routine + restarts do not take exclusive locks to recreate identical policies. + """ + from data_loader_pg import get_pg_migration_connection + + script = RLS_SCRIPT_PATH.read_text() + checksum = hashlib.sha256(script.encode("utf-8")).hexdigest() + mode = rls_enforcement_mode() + + conn = get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + CREATE TABLE IF NOT EXISTS app_schema_migrations ( + name TEXT PRIMARY KEY, + checksum TEXT NOT NULL, + applied_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ) + """ + ) + cur.execute( + "SELECT checksum FROM app_schema_migrations WHERE name = %s", + (RLS_MIGRATION_NAME,), + ) + row = cur.fetchone() + script_applied = False + if not row or row[0] != checksum: + cur.execute(script) + cur.execute( + """ + INSERT INTO app_schema_migrations (name, checksum) + VALUES (%s, %s) + ON CONFLICT (name) DO UPDATE + SET checksum = EXCLUDED.checksum, applied_at = CURRENT_TIMESTAMP + """, + (RLS_MIGRATION_NAME, checksum), + ) + script_applied = True + enforcement = _apply_rls_enforcement(cur, mode) + conn.commit() + finally: + conn.close() + + logger.info( + "Row-level security migration: script_applied=%s mode=%s enforced=%s", + script_applied, + mode or "unchanged", + sorted(name for name, enabled in enforcement.items() if enabled), + ) + return { + "script_applied": script_applied, + "mode": mode or "unchanged", + "enforced_tables": enforcement, + } + + +def run_schema_migrations() -> None: + """ + Run idempotent schema migrations under a Postgres advisory lock. + + The API image starts Uvicorn with multiple worker processes, and each + worker enters FastAPI lifespan. Without a cross-process lock, concurrent + ALTER TABLE / ENABLE RLS statements can deadlock during rollout. + """ + from data_loader_pg import ( + ensure_equipment_run_trace_events_pg, + ensure_equipment_run_trace_samples_pg, + ensure_equipment_runs_pg, + ensure_etcher_recipe_name_column_pg, + ensure_etcher_run_file_refs_pg, + ensure_glance_ingestion_audit_pg, + get_pg_migration_connection, + ) + from metadata_pg import ( + ensure_computational_execution_columns_pg, + ensure_data_uploads_kind_column_pg, + ensure_data_uploads_processing_started_at_pg, + ensure_equipment_access_table_pg, + ensure_execution_queue_columns_pg, + ensure_execution_request_id_columns_pg, + ensure_experiment_definition_optimization_context_pg, + ensure_experiment_definition_snapshot_columns_pg, + ensure_experiment_proposal_generation_columns_pg, + ensure_experiment_type_reuse_columns_pg, + ensure_experiment_type_versioning_columns_pg, + ensure_process_definition_tables_pg, + ensure_project_experiments_table_pg, + ensure_run_version_tracking_columns_pg, + ensure_sample_and_publication_tables_pg, + ensure_user_role_audit_table_pg, + ) + from ml_response_cache import ensure_ml_endpoint_cache_table_pg + + lock_conn = get_pg_migration_connection() + try: + with lock_conn.cursor() as cur: + cur.execute("SELECT pg_advisory_lock(%s)", (STARTUP_MIGRATION_LOCK_ID,)) + lock_conn.commit() + + ensure_etcher_run_file_refs_pg() + ensure_etcher_recipe_name_column_pg() + ensure_equipment_runs_pg() + ensure_equipment_run_trace_samples_pg() + ensure_equipment_run_trace_events_pg() + ensure_glance_ingestion_audit_pg() + ensure_experiment_proposal_generation_columns_pg() + ensure_experiment_type_reuse_columns_pg() + ensure_experiment_type_versioning_columns_pg() + ensure_experiment_definition_snapshot_columns_pg() + ensure_experiment_definition_optimization_context_pg() + ensure_execution_request_id_columns_pg() + ensure_execution_queue_columns_pg() + ensure_computational_execution_columns_pg() + ensure_project_experiments_table_pg() + ensure_process_definition_tables_pg() + ensure_run_version_tracking_columns_pg() + ensure_user_role_audit_table_pg() + ensure_sample_and_publication_tables_pg() + ensure_equipment_access_table_pg() + ensure_data_uploads_kind_column_pg() + ensure_data_uploads_processing_started_at_pg() + ensure_ml_endpoint_cache_table_pg() + # Last: the policies reference tables and columns created above. + ensure_row_level_security_pg() + finally: + try: + with lock_conn.cursor() as cur: + cur.execute("SELECT pg_advisory_unlock(%s)", (STARTUP_MIGRATION_LOCK_ID,)) + lock_conn.commit() + finally: + lock_conn.close() + + +def main() -> int: + logging.basicConfig(level=logging.INFO) + run_schema_migrations() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/api/fair_archiver.py b/api/fair_archiver.py index 4c4b4ec..72aa6d5 100644 --- a/api/fair_archiver.py +++ b/api/fair_archiver.py @@ -119,10 +119,10 @@ def _get_project_metadata(project_id: str) -> dict[str, Any]: Returns sensible defaults if lookup fails. """ try: - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( """ diff --git a/api/glance_closed_loop.py b/api/glance_closed_loop.py index 81b8b80..632dc13 100644 --- a/api/glance_closed_loop.py +++ b/api/glance_closed_loop.py @@ -22,7 +22,7 @@ from data_loader import FEATURES, PRIMARY_TARGET, SECONDARY_TARGET from data_loader_pg import ( equipment_training_revision_sha256, - get_pg_superuser_connection, + get_pg_worker_connection, glance_source_run_lock_key, ) @@ -193,7 +193,7 @@ def enrich_glance_trace_outcomes_pg( return {} records_by_project: dict[str, list[dict[str, Any]]] = {} - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: for ( @@ -725,7 +725,7 @@ def reconcile_glance_trace_runs_pg( training_projects: set[str] = set() training_experiments: set[str] = set() - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: for run_record in run_records: @@ -1246,7 +1246,7 @@ def finish_glance_followup_claims_pg( """Complete or release persisted follow-up claims after synchronous work.""" if not claims: return - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: for claim in claims: diff --git a/api/legacy_etcher_processor.py b/api/legacy_etcher_processor.py index e9090ce..0b56a3e 100644 --- a/api/legacy_etcher_processor.py +++ b/api/legacy_etcher_processor.py @@ -140,11 +140,11 @@ def run_legacy_etcher_processor() -> dict[str, Any]: **status, } - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection lock_conn = None try: - lock_conn = get_pg_superuser_connection() + lock_conn = get_pg_worker_connection() with lock_conn.cursor() as cur: cur.execute("SELECT pg_try_advisory_lock(%s)", (PROCESSOR_ADVISORY_LOCK_ID,)) got_lock = bool(cur.fetchone()[0]) diff --git a/api/main.py b/api/main.py index 9a8bedc..5e29165 100644 --- a/api/main.py +++ b/api/main.py @@ -24,7 +24,9 @@ from fastapi.responses import JSONResponse from middleware.logging import RequestTracingMiddleware -from middleware.error_handler import global_exception_handler +import psycopg2.errors + +from middleware.error_handler import global_exception_handler, insufficient_privilege_handler from runtime_config import MissingRuntimeConfiguration, validate_runtime_configuration from routers import ml, dataset, dataset_v2, equipment, upload, admin, processing @@ -62,29 +64,40 @@ "equipment_access", ) -STARTUP_MIGRATION_LOCK_ID = 6420260506 - - def _env_present(name: str) -> bool: return bool(os.getenv(name, "").strip()) def _runtime_dependency_checks() -> dict[str, bool]: + """Credentials the API needs to serve traffic; all must be present.""" return { "system_token_configured": _env_present("DT_SYSTEM_TOKEN"), "ingestion_token_configured": _env_present("INGESTION_TOKEN"), "postgres_api_password_configured": _env_present("PG_PASS"), + # Background jobs use PG_WORKER_PASS, or the migration credentials + # until the api_worker role is provisioned. + "postgres_worker_credentials_configured": ( + _env_present("PG_WORKER_PASS") or _env_present("PG_SUPER_PASS") + ), + } + + +def _runtime_informational_checks() -> dict[str, bool]: + """Reported by /api/health but not required for readiness.""" + return { "postgres_superuser_password_configured": _env_present("PG_SUPER_PASS"), + "postgres_worker_role_configured": _env_present("PG_WORKER_PASS"), } def _postgres_health() -> dict[str, object]: from data_loader_pg import get_pg_connection + from db_migrations import RLS_REPORTED_TABLES runtime_checks = _runtime_dependency_checks() credentials_configured = ( runtime_checks["postgres_api_password_configured"] - and runtime_checks["postgres_superuser_password_configured"] + and runtime_checks["postgres_worker_credentials_configured"] ) if not credentials_configured: return { @@ -113,6 +126,16 @@ def _postgres_health() -> dict[str, object]: """ ) available_tables = {row[0] for row in cur.fetchall()} + cur.execute( + """ + SELECT c.relname, c.relrowsecurity + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' AND c.relname = ANY(%s) + """, + (list(RLS_REPORTED_TABLES),), + ) + row_level_security = {row[0]: bool(row[1]) for row in cur.fetchall()} conn.commit() missing_tables = [ @@ -126,6 +149,7 @@ def _postgres_health() -> dict[str, object]: "database": database_name, "user": user_name, "missing_tables": missing_tables, + "row_level_security": row_level_security, } except Exception as exc: logger.warning("Health check Postgres probe failed: %s", exc) @@ -141,76 +165,26 @@ def _postgres_health() -> dict[str, object]: def _run_startup_migrations() -> None: - """ - Run idempotent schema checks under a Postgres advisory lock. + from db_migrations import run_schema_migrations, should_run_startup_migrations + + if not should_run_startup_migrations(): + logger.info( + "Skipping schema migrations at startup (no migration credentials or " + "DB_MIGRATE_ON_STARTUP=false); run `python -m db_migrations` separately" + ) + return + run_schema_migrations() - The API image starts Uvicorn with multiple worker processes, and each - worker enters FastAPI lifespan. Without a cross-process lock, concurrent - ALTER TABLE / ENABLE RLS statements can deadlock during rollout. - """ - from data_loader_pg import get_pg_superuser_connection - from data_loader_pg import ensure_etcher_run_file_refs_pg - from data_loader_pg import ensure_equipment_run_trace_samples_pg - from data_loader_pg import ensure_equipment_run_trace_events_pg - from data_loader_pg import ensure_equipment_runs_pg - from data_loader_pg import ensure_glance_ingestion_audit_pg - from data_loader_pg import ensure_etcher_recipe_name_column_pg - from metadata_pg import ( - ensure_experiment_proposal_generation_columns_pg, - ensure_experiment_type_reuse_columns_pg, - ensure_experiment_type_versioning_columns_pg, - ensure_experiment_definition_snapshot_columns_pg, - ensure_experiment_definition_optimization_context_pg, - ensure_execution_request_id_columns_pg, - ensure_execution_queue_columns_pg, - ensure_computational_execution_columns_pg, - ensure_project_experiments_table_pg, - ensure_process_definition_tables_pg, - ensure_run_version_tracking_columns_pg, - ensure_user_role_audit_table_pg, - ensure_sample_and_publication_tables_pg, - ensure_equipment_access_table_pg, - ensure_data_uploads_kind_column_pg, - ensure_data_uploads_processing_started_at_pg, - recover_stale_db3_uploads_pg, - ) - lock_conn = get_pg_superuser_connection() +def _start_background_jobs() -> None: + """Recover interrupted background work and start the proposal job sweeper.""" + from metadata_pg import recover_stale_db3_uploads_pg, start_proposal_job_sweeper + try: - with lock_conn.cursor() as cur: - cur.execute("SELECT pg_advisory_lock(%s)", (STARTUP_MIGRATION_LOCK_ID,)) - lock_conn.commit() - - ensure_etcher_run_file_refs_pg() - ensure_etcher_recipe_name_column_pg() - ensure_equipment_runs_pg() - ensure_equipment_run_trace_samples_pg() - ensure_equipment_run_trace_events_pg() - ensure_glance_ingestion_audit_pg() - ensure_experiment_proposal_generation_columns_pg() - ensure_experiment_type_reuse_columns_pg() - ensure_experiment_type_versioning_columns_pg() - ensure_experiment_definition_snapshot_columns_pg() - ensure_experiment_definition_optimization_context_pg() - ensure_execution_request_id_columns_pg() - ensure_execution_queue_columns_pg() - ensure_computational_execution_columns_pg() - ensure_project_experiments_table_pg() - ensure_process_definition_tables_pg() - ensure_run_version_tracking_columns_pg() - ensure_user_role_audit_table_pg() - ensure_sample_and_publication_tables_pg() - ensure_equipment_access_table_pg() - ensure_data_uploads_kind_column_pg() - ensure_data_uploads_processing_started_at_pg() recover_stale_db3_uploads_pg() - finally: - try: - with lock_conn.cursor() as cur: - cur.execute("SELECT pg_advisory_unlock(%s)", (STARTUP_MIGRATION_LOCK_ID,)) - lock_conn.commit() - finally: - lock_conn.close() + except Exception: + logger.warning("Could not recover stale DB3 uploads at startup", exc_info=True) + start_proposal_job_sweeper() # ── Startup / shutdown ─────────────────────────────────────────────────────── @@ -226,6 +200,7 @@ async def lifespan(app: FastAPI): logger.error("Missing runtime configuration: %s", exc) raise _run_startup_migrations() + _start_background_jobs() logger.info("=" * 60) logger.info("Birck Digital Twin API starting") @@ -259,6 +234,7 @@ async def lifespan(app: FastAPI): # Global unhandled-exception handler app.add_exception_handler(Exception, global_exception_handler) +app.add_exception_handler(psycopg2.errors.InsufficientPrivilege, insufficient_privilege_handler) # Mount all routers under /api app.include_router(ml.router, prefix="/api") @@ -290,12 +266,13 @@ def health(request: Request): runtime_checks = _runtime_dependency_checks() postgres = _postgres_health() healthy = all(runtime_checks.values()) and postgres["reachable"] and not postgres["missing_tables"] + runtime_report = {**runtime_checks, **_runtime_informational_checks()} payload = { "status": "ok" if healthy else "degraded", "production_db": has_production_db(), "csv_available": DATASET_PATH.exists(), - "runtime": runtime_checks, + "runtime": runtime_report, "postgres": postgres, "trace_id": getattr(request.state, "trace_id", "n/a"), } diff --git a/api/metadata_pg.py b/api/metadata_pg.py index 07e1913..82ae9a6 100644 --- a/api/metadata_pg.py +++ b/api/metadata_pg.py @@ -1,9 +1,12 @@ +import hashlib import json import logging import os import re +import time import uuid from datetime import date, datetime +from threading import Thread from pathlib import Path from typing import Any @@ -12,7 +15,8 @@ from data_loader_pg import ( EQUIPMENT_RUN_ID_OFFSET, get_pg_connection, - get_pg_superuser_connection, + get_pg_migration_connection, + get_pg_worker_connection, ) from security import PlatformUser @@ -22,6 +26,52 @@ RECIPE_TERMINAL_STATUSES = {"completed", "rejected"} RECIPE_ACTIVE_STATUSES = {"pending", "accepted", "attempted"} PROPOSAL_GENERATION_STATUSES = {"not_started", "queued", "generating", "completed", "failed", "skipped"} +SKIPPED_PROPOSAL_GENERATION_NOTE = "Recipe generation was not requested for this project run." + + +class ExperimentRequestConflictError(Exception): + """A client_request_id was reused for a different experiment payload.""" + + +def _experiment_request_fingerprint(payload: dict[str, Any]) -> str: + """Stable hash of the fields that define an experiment create request.""" + fields = { + key: payload.get(key) + for key in ( + "experiment_name", + "selected_equipment", + "project_id", + "type_id", + "planned_date", + "planned_parameters", + "optimization_context", + "sample_ids", + "generate_proposals", + ) + } + encoded = json.dumps(fields, sort_keys=True, default=str, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _env_int(name: str, default: int, minimum: int) -> int: + try: + return max(minimum, int(os.getenv(name, str(default)))) + except (TypeError, ValueError): + return default + + +def proposal_job_lease_seconds() -> int: + """How long a claimed proposal job may run before another worker retries it.""" + return _env_int("PROPOSAL_JOB_LEASE_SECONDS", 900, 30) + + +def proposal_job_max_attempts() -> int: + return _env_int("PROPOSAL_JOB_MAX_ATTEMPTS", 3, 1) + + +def proposal_job_max_age_hours() -> int: + """Unfinished jobs older than this are failed instead of run.""" + return _env_int("PROPOSAL_JOB_MAX_AGE_HOURS", 72, 1) EXECUTION_STATUSES = { "requested", "under_review", @@ -41,8 +91,11 @@ def _make_execution_request_id() -> str: def ensure_experiment_proposal_generation_columns_pg() -> None: - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: + # Each statement commits on its own so the ALTER TABLE's exclusive + # lock is released before the indexes are built; the migration + # connection's lock_timeout bounds every lock acquisition. with conn.cursor() as cur: cur.execute( """ @@ -51,7 +104,33 @@ def ensure_experiment_proposal_generation_columns_pg() -> None: ADD COLUMN IF NOT EXISTS proposal_generation_error TEXT DEFAULT '', ADD COLUMN IF NOT EXISTS proposal_generation_started_at TIMESTAMPTZ, ADD COLUMN IF NOT EXISTS proposal_generation_completed_at TIMESTAMPTZ, - ADD COLUMN IF NOT EXISTS proposal_generation_updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP + ADD COLUMN IF NOT EXISTS proposal_generation_updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + ADD COLUMN IF NOT EXISTS proposal_generation_attempts INTEGER NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS proposal_generation_claim_id UUID, + ADD COLUMN IF NOT EXISTS proposal_generation_lease_expires_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS client_request_id VARCHAR(200) NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS client_request_fingerprint VARCHAR(64) NOT NULL DEFAULT '' + """ + ) + conn.commit() + with conn.cursor() as cur: + # Retried create requests carry the same client_request_id and + # must resolve to the experiment the first attempt created. + cur.execute( + """ + CREATE UNIQUE INDEX IF NOT EXISTS uq_experiment_definitions_client_request + ON experiment_definitions (owner_id, client_request_id) + WHERE client_request_id <> '' + """ + ) + conn.commit() + with conn.cursor() as cur: + # The proposal job sweeper only scans unfinished jobs. + cur.execute( + """ + CREATE INDEX IF NOT EXISTS idx_experiment_definitions_open_proposal_jobs + ON experiment_definitions (proposal_generation_updated_at) + WHERE proposal_generation_status IN ('queued', 'generating') """ ) conn.commit() @@ -60,7 +139,7 @@ def ensure_experiment_proposal_generation_columns_pg() -> None: def ensure_experiment_type_reuse_columns_pg() -> None: - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -79,7 +158,7 @@ def ensure_experiment_type_reuse_columns_pg() -> None: def ensure_experiment_type_versioning_columns_pg() -> None: """Add versioning, PID, authorship, and optimization target columns.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -129,7 +208,7 @@ def ensure_experiment_type_versioning_columns_pg() -> None: def ensure_run_version_tracking_columns_pg() -> None: """Track which equipment/experiment definition version produced each run.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -146,7 +225,7 @@ def ensure_run_version_tracking_columns_pg() -> None: def ensure_user_role_audit_table_pg() -> None: """Track admin role changes for platform users.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -180,7 +259,7 @@ def ensure_user_role_audit_table_pg() -> None: def ensure_experiment_definition_snapshot_columns_pg() -> None: """Persist the reusable experiment definition snapshot used by each run.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -198,7 +277,7 @@ def ensure_experiment_definition_snapshot_columns_pg() -> None: def ensure_experiment_definition_optimization_context_pg() -> None: """Persist optional run-level optimization/campaign intent.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -214,7 +293,7 @@ def ensure_experiment_definition_optimization_context_pg() -> None: def ensure_execution_request_id_columns_pg() -> None: """Attach a visible execution/request ID to project runs and outputs.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -328,7 +407,7 @@ def ensure_execution_request_id_columns_pg() -> None: def ensure_execution_queue_columns_pg() -> None: """Track facility-manager execution status for project runs.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -363,7 +442,7 @@ def ensure_execution_queue_columns_pg() -> None: def ensure_computational_execution_columns_pg() -> None: """Track automated computational submissions alongside physical runs.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -395,7 +474,7 @@ def ensure_computational_execution_columns_pg() -> None: def ensure_project_experiments_table_pg() -> None: """Create the project-to-experiment-template join table.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -443,7 +522,7 @@ def ensure_project_experiments_table_pg() -> None: def ensure_process_definition_tables_pg() -> None: """Create reusable process definitions separate from projects.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -506,7 +585,7 @@ def ensure_process_definition_tables_pg() -> None: def ensure_sample_and_publication_tables_pg() -> None: """Create sample registry and PI-controlled publication request tables.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -613,7 +692,7 @@ def ensure_sample_and_publication_tables_pg() -> None: def ensure_equipment_access_table_pg() -> None: """Create the equipment_access table for trusted maintainers.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -813,12 +892,18 @@ def update_experiment_proposal_generation_status_pg( status: str, error: str = "", ) -> None: + """ + Record proposal-generation progress driven outside the proposal job queue + (the ingestion follow-up loop). ``generating`` takes a lease so the queue + sweeper leaves the row alone while that work runs; any other status + releases the lease and any queue claim. + """ if not experiment_id: return if status not in PROPOSAL_GENERATION_STATUSES: raise ValueError("Invalid proposal generation status") - conn = get_pg_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -835,10 +920,24 @@ def update_experiment_proposal_generation_status_pg( WHEN %s IN ('completed', 'failed', 'skipped') THEN CURRENT_TIMESTAMP ELSE proposal_generation_completed_at END, + proposal_generation_lease_expires_at = CASE + WHEN %s = 'generating' + THEN CURRENT_TIMESTAMP + make_interval(secs => %s) + ELSE NULL + END, + proposal_generation_claim_id = NULL, proposal_generation_updated_at = CURRENT_TIMESTAMP WHERE id = %s """, - (status, error, status, status, experiment_id), + ( + status, + error, + status, + status, + status, + proposal_job_lease_seconds(), + experiment_id, + ), ) conn.commit() finally: @@ -1155,7 +1254,7 @@ def _fetch_equipment_run_metrics(_conn=None) -> dict[str, dict[str, Any]]: RLS-scoped connection passed by callers would otherwise hide runs in private/shared projects, even from members and admins, and undercount. """ - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -1924,7 +2023,7 @@ def delete_equipment_pg( ``projects`` (api_client has SELECT-only on that table) and RLS would otherwise make the update a silent no-op. """ - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -2124,6 +2223,46 @@ def _experiment_type_row_to_dict(row: dict) -> dict[str, Any]: """ +def _visible_experiment_type_filter( + user: PlatformUser, + alias: str = "experiment_types", +) -> tuple[str, list[Any]]: + """ + ``AND ...`` clause limiting ``experiment_types`` (referenced as ``alias``) + to templates ``user`` may see or use: their own, shared or published ones, + their organization's, and templates attached to a project they can read. + Row-level security enforces the same rule in the database; this keeps it + when enforcement is off. + """ + if user.role == "admin": + return "", [] + return ( + f""" + AND ( + {alias}.owner_id = %s + OR {alias}.visibility IN ('shared', 'published') + OR ({alias}.owner_org <> '' AND {alias}.owner_org = %s) + OR EXISTS ( + SELECT 1 + FROM project_experiments vis_pe + JOIN projects vis_p ON vis_p.id = vis_pe.project_id + WHERE vis_pe.type_id = {alias}.id + AND ( + vis_p.access_mode = 'open' + OR EXISTS ( + SELECT 1 + FROM project_members vis_pm + WHERE vis_pm.project_id = vis_p.id + AND vis_pm.nanohub_user_id = %s + ) + ) + ) + ) + """, + [user.id, user.organization, user.id], + ) + + def create_experiment_type_pg( *, payload: dict[str, Any], @@ -2132,7 +2271,7 @@ def create_experiment_type_pg( type_slug = _slugify(payload.get("type_name", ""), "experiment_type") type_id = f"{type_slug}_{uuid.uuid4().hex[:8]}" pid = _generate_pid() - conn = get_pg_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -2194,7 +2333,7 @@ def list_experiment_types_pg( user: PlatformUser, equipment_id: str = "", ) -> list[dict[str, Any]]: - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) params: list[Any] = [] conditions: list[str] = [] if equipment_id: @@ -2261,7 +2400,7 @@ def update_experiment_type_pg( user: PlatformUser, ) -> dict[str, Any]: """Edit an experiment type, incrementing version and saving a snapshot.""" - conn = get_pg_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: # Fetch current row @@ -2357,7 +2496,7 @@ def delete_experiment_type_pg( if not type_id: return {"status": "not_found", "type_id": type_id} - conn = get_pg_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -2370,40 +2509,16 @@ def delete_experiment_type_pg( if existing.get("owner_id") != user.id and user.role != "admin": raise PermissionError("Only the owner or an admin can delete this template") + # References to the template live in other users' projects and + # processes, which the caller's row-level security scope cannot + # see or update. A SECURITY DEFINER function re-checks ownership + # in the database and detaches them in this same transaction + # (see scripts/add_row_level_security.sql). cur.execute( - "SELECT COUNT(*) AS count FROM project_experiments WHERE type_id = %s", - (type_id,), - ) - project_attachment_count = int((cur.fetchone() or {}).get("count") or 0) - cur.execute( - "SELECT COUNT(*) AS count FROM process_steps WHERE type_id = %s", - (type_id,), - ) - process_step_count = int((cur.fetchone() or {}).get("count") or 0) - cur.execute( - "SELECT COUNT(*) AS count FROM experiment_definitions WHERE type_id = %s", - (type_id,), - ) - historical_run_count = int((cur.fetchone() or {}).get("count") or 0) - - cur.execute( - """ - UPDATE project_experiments - SET upstream_type_id = '', updated_at = CURRENT_TIMESTAMP - WHERE upstream_type_id = %s - """, - (type_id,), - ) - cleared_project_upstreams = cur.rowcount - cur.execute( - """ - UPDATE process_steps - SET upstream_type_id = '', updated_at = CURRENT_TIMESTAMP - WHERE upstream_type_id = %s - """, + "SELECT * FROM app_detach_experiment_type(%s)", (type_id,), ) - cleared_process_upstreams = cur.rowcount + detached = cur.fetchone() or {} cur.execute( """ @@ -2414,27 +2529,40 @@ def delete_experiment_type_pg( (type_id,), ) deleted = cur.fetchone() + if not deleted: + raise PermissionError("Only the owner or an admin can delete this template") conn.commit() return { "status": "deleted", "type_id": str(deleted["id"]), "name": deleted.get("name"), - "detached_project_experiments": project_attachment_count, - "detached_process_steps": process_step_count, - "detached_experiment_runs": historical_run_count, - "cleared_project_upstreams": cleared_project_upstreams, - "cleared_process_upstreams": cleared_process_upstreams, + "detached_project_experiments": int(detached.get("project_attachment_count") or 0), + "detached_process_steps": int(detached.get("process_step_count") or 0), + "detached_experiment_runs": int(detached.get("historical_run_count") or 0), + "cleared_project_upstreams": int(detached.get("cleared_project_upstreams") or 0), + "cleared_process_upstreams": int(detached.get("cleared_process_upstreams") or 0), "message": "Unit experiment deleted. Historical runs keep their captured parameter snapshots.", } finally: conn.close() -def get_experiment_type_versions_pg(type_id: str) -> list[dict[str, Any]]: - """Return version history for an experiment type, newest first.""" - conn = get_pg_connection() +def get_experiment_type_versions_pg( + type_id: str, + *, + user: PlatformUser, +) -> list[dict[str, Any]]: + """Return version history for an experiment type the caller can see, newest first.""" + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: + visibility_sql, visibility_params = _visible_experiment_type_filter(user) + cur.execute( + f"SELECT 1 FROM experiment_types WHERE id = %s {visibility_sql}", + [type_id, *visibility_params], + ) + if not cur.fetchone(): + return [] cur.execute( """ SELECT id, type_id, version, change_summary, @@ -2468,12 +2596,18 @@ def fork_experiment_type_pg( user: PlatformUser, ) -> dict[str, Any]: """Clone an experiment type for the calling user (fork).""" - conn = get_pg_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: + visibility_sql, visibility_params = _visible_experiment_type_filter(user) cur.execute( - f"SELECT {_EXPERIMENT_TYPE_SELECT_COLS} FROM experiment_types WHERE id = %s", - (type_id,), + f""" + SELECT {_EXPERIMENT_TYPE_SELECT_COLS} + FROM experiment_types + WHERE id = %s + {visibility_sql} + """, + [type_id, *visibility_params], ) source = cur.fetchone() if not source: @@ -2622,7 +2756,7 @@ def list_samples_pg( project_id: str = "", include_archived: bool = False, ) -> list[dict[str, Any]]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: params: list[Any] = [] @@ -2680,7 +2814,7 @@ def create_sample_pg( sample_slug = _slugify(name, "sample") sample_id = f"smp_{sample_slug}_{uuid.uuid4().hex[:8]}" - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) @@ -2732,7 +2866,7 @@ def update_sample_pg( payload: dict[str, Any], user: PlatformUser, ) -> dict[str, Any]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT * FROM samples WHERE id = %s", (sample_id,)) @@ -2788,7 +2922,7 @@ def archive_sample_pg( sample_id: str, user: PlatformUser, ) -> dict[str, Any]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT project_id FROM samples WHERE id = %s", (sample_id,)) @@ -2843,7 +2977,7 @@ def list_publication_requests_pg( project_id: str, user: PlatformUser, ) -> list[dict[str, Any]]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_access(cur, project_id=project_id, user=user) @@ -2870,7 +3004,7 @@ def create_publication_request_pg( title = str(payload.get("title") or "").strip() if not title: raise ValueError("Publication title is required") - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_publish_access(cur, project_id=project_id, user=user) @@ -3007,7 +3141,7 @@ def export_publication_request_pg( request_id: str, user: PlatformUser, ) -> dict[str, Any]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3063,7 +3197,7 @@ def update_publication_request_status_pg( "mark_submitted": "submitted", "cancel": "cancelled", } - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3182,11 +3316,15 @@ def add_project_experiment_pg( if not type_id: raise ValueError("experiment type_id is required") - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) - cur.execute("SELECT id FROM experiment_types WHERE id = %s", (type_id,)) + visibility_sql, visibility_params = _visible_experiment_type_filter(user) + cur.execute( + f"SELECT id FROM experiment_types WHERE id = %s {visibility_sql}", + [type_id, *visibility_params], + ) if not cur.fetchone(): return {"status": "not_found", "type_id": type_id} cur.execute( @@ -3234,7 +3372,7 @@ def list_project_experiments_pg( if not project_id: return [] - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_access(cur, project_id=project_id, user=user) @@ -3288,7 +3426,7 @@ def update_project_experiment_workflow_pg( if not isinstance(output_mapping, dict): output_mapping = {} - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) @@ -3409,7 +3547,7 @@ def _ensure_process_write_access(cur, *, process_id: str, user: PlatformUser) -> def list_process_definitions_pg(*, user: PlatformUser) -> list[dict[str, Any]]: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: params: list[Any] = [] @@ -3422,6 +3560,10 @@ def list_process_definitions_pg(*, user: PlatformUser) -> list[dict[str, Any]]: OR (p.visibility = 'shared' AND (p.owner_org = '' OR p.owner_org = %s)) """ params.extend([user.id, user.organization]) + type_visibility_sql, type_visibility_params = _visible_experiment_type_filter( + user, + alias="et", + ) cur.execute( f""" SELECT @@ -3443,14 +3585,15 @@ def list_process_definitions_pg(*, user: PlatformUser) -> list[dict[str, Any]]: ORDER BY ps.sequence_index ASC, ps.created_at ASC ) FROM process_steps ps - JOIN experiment_types et ON et.id = ps.type_id + LEFT JOIN experiment_types et + ON et.id = ps.type_id {type_visibility_sql} WHERE ps.process_id = p.id ), '[]'::json) AS step_summary FROM process_definitions p {where_sql} ORDER BY p.updated_at DESC, p.created_at DESC """, - params, + [*type_visibility_params, *params], ) rows = cur.fetchall() return [ @@ -3465,13 +3608,17 @@ def list_process_definitions_pg(*, user: PlatformUser) -> list[dict[str, Any]]: def get_process_definition_pg(*, process_id: str, user: PlatformUser) -> dict[str, Any] | None: - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: try: process = _ensure_process_access(cur, process_id=process_id, user=user) except LookupError: return None + type_visibility_sql, type_visibility_params = _visible_experiment_type_filter( + user, + alias="t", + ) cur.execute( f""" SELECT @@ -3487,11 +3634,14 @@ def get_process_definition_pg(*, process_id: str, user: PlatformUser) -> dict[st ps.updated_at AS process_step_updated_at, {_EXPERIMENT_TYPE_SELECT_COLS_T} FROM process_steps ps - JOIN experiment_types t ON t.id = ps.type_id + -- LEFT JOIN: a step whose template the caller cannot read + -- still appears, without the template's details. + LEFT JOIN experiment_types t + ON t.id = ps.type_id {type_visibility_sql} WHERE ps.process_id = %s ORDER BY ps.sequence_index ASC, ps.created_at ASC """, - (process_id,), + [*type_visibility_params, process_id], ) steps = [_process_step_row_to_dict(dict(row)) for row in cur.fetchall()] return _process_row_to_dict(process, steps) @@ -3509,7 +3659,7 @@ def create_process_definition_pg(*, payload: dict[str, Any], user: PlatformUser) process_slug = _slugify(name, "process") process_id = f"{process_slug}_{uuid.uuid4().hex[:8]}" - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -3536,7 +3686,11 @@ def create_process_definition_pg(*, payload: dict[str, Any], user: PlatformUser) type_id = str(step.get("type_id") or "").strip() if not type_id: raise ValueError("Every process step requires a unit experiment") - cur.execute("SELECT 1 FROM experiment_types WHERE id = %s", (type_id,)) + visibility_sql, visibility_params = _visible_experiment_type_filter(user) + cur.execute( + f"SELECT 1 FROM experiment_types WHERE id = %s {visibility_sql}", + [type_id, *visibility_params], + ) if not cur.fetchone(): raise ValueError(f"Unit experiment '{type_id}' was not found") output_mapping = step.get("output_mapping") or {} @@ -3582,7 +3736,7 @@ def instantiate_process_in_project_pg( if not project_id or not process_id: raise ValueError("project_id and process_id are required") - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) @@ -3600,6 +3754,19 @@ def instantiate_process_in_project_pg( if not steps: raise ValueError("Process has no unit experiment steps") + # Every step's template must be visible to the caller; a shared + # process cannot be used to attach someone's private template. + visibility_sql, visibility_params = _visible_experiment_type_filter(user) + for step in steps: + cur.execute( + f"SELECT 1 FROM experiment_types WHERE id = %s {visibility_sql}", + [str(step.get("type_id") or ""), *visibility_params], + ) + if not cur.fetchone(): + raise PermissionError( + "This process uses a unit experiment you do not have access to" + ) + attached = 0 for step in steps: type_id = str(step.get("type_id") or "") @@ -3803,32 +3970,38 @@ def _ensure_workflow_example_type( additional_inputs: list[dict[str, Any]], user: PlatformUser, ) -> str: - cur.execute("SELECT id FROM experiment_types WHERE id = %s", (type_id,)) + # Refresh the shared example template only when the caller owns it or is + # an admin (checked here, not only by row-level security). Otherwise reuse + # whichever template holds this id, if the caller may see it, or create it. + cur.execute( + """ + UPDATE experiment_types + SET + name = %s, + description = %s, + equipment_id = %s, + type_parameters_json = %s, + outputs_json = %s, + additional_inputs_json = %s, + visibility = CASE WHEN visibility = 'private' THEN 'shared' ELSE visibility END, + updated_at = CURRENT_TIMESTAMP + WHERE id = %s + AND (owner_id = %s OR %s) + RETURNING id + """, + ( + name, + description, + equipment_id, + Json(parameters), + Json(outputs), + Json(additional_inputs), + type_id, + user.id, + user.role == "admin", + ), + ) if cur.fetchone(): - cur.execute( - """ - UPDATE experiment_types - SET - name = %s, - description = %s, - equipment_id = %s, - type_parameters_json = %s, - outputs_json = %s, - additional_inputs_json = %s, - visibility = CASE WHEN visibility = 'private' THEN 'shared' ELSE visibility END, - updated_at = CURRENT_TIMESTAMP - WHERE id = %s - """, - ( - name, - description, - equipment_id, - Json(parameters), - Json(outputs), - Json(additional_inputs), - type_id, - ), - ) return type_id cur.execute( @@ -3853,6 +4026,7 @@ def _ensure_workflow_example_type( '', '', %s, %s ) + ON CONFLICT (id) DO NOTHING """, ( type_id, @@ -3869,6 +4043,16 @@ def _ensure_workflow_example_type( user.organization, ), ) + # Someone else may hold this id as a private template; never attach it. + visibility_sql, visibility_params = _visible_experiment_type_filter(user) + cur.execute( + f"SELECT 1 FROM experiment_types WHERE id = %s {visibility_sql}", + [type_id, *visibility_params], + ) + if not cur.fetchone(): + raise PermissionError( + f"Example unit experiment '{type_id}' exists but is not shared with you" + ) return type_id @@ -3982,7 +4166,7 @@ def create_etcher_profilometer_workflow_example_pg( } ] - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) @@ -4149,7 +4333,7 @@ def create_rcac_closed_loop_demo_pg( ), } - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: _ensure_project_write_access(cur, project_id=project_id, user=user) @@ -4307,16 +4491,79 @@ def create_rcac_closed_loop_demo_pg( conn.close() +def _replay_experiment_request( + cur, + *, + user: PlatformUser, + client_request_id: str, + request_fingerprint: str, +) -> dict[str, Any] | None: + """Return the experiment an earlier attempt with this request id created.""" + cur.execute( + """ + SELECT id, execution_request_id, proposal_generation_status, + client_request_fingerprint + FROM experiment_definitions + WHERE owner_id = %s AND client_request_id = %s + """, + (user.id, client_request_id), + ) + existing = cur.fetchone() + if not existing: + return None + stored_fingerprint = existing.get("client_request_fingerprint") or "" + if stored_fingerprint and stored_fingerprint != request_fingerprint: + raise ExperimentRequestConflictError( + "This request id was already used for a different experiment" + ) + return { + "status": "created", + "idempotent_replay": True, + "experiment_id": str(existing["id"]), + "execution_request_id": str(existing.get("execution_request_id") or ""), + "proposal_generation_status": existing.get("proposal_generation_status") or "", + } + + def create_experiment_definition_pg( *, payload: dict[str, Any], user: PlatformUser, ) -> dict[str, Any]: - conn = get_pg_connection() + """ + Create an experiment definition as ``user`` in a single transaction. + + The initial proposal-generation status is written with the row: ``queued`` + rows are the durable job the proposal worker picks up, so no follow-up + write can fail after the experiment exists. A repeated + ``client_request_id`` returns the experiment the first attempt created; + reusing it for a different payload raises ExperimentRequestConflictError. + """ + type_id = str(payload.get("type_id") or "").strip() + project_id = str(payload.get("project_id") or "").strip() + client_request_id = str(payload.get("client_request_id") or "").strip() + request_fingerprint = ( + _experiment_request_fingerprint(payload) if client_request_id else "" + ) + generate_proposals = bool(payload.get("generate_proposals")) + proposal_status = "queued" if generate_proposals else "skipped" + proposal_note = "" if generate_proposals else SKIPPED_PROPOSAL_GENERATION_NOTE + + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: - type_id = str(payload.get("type_id") or "").strip() - project_id = str(payload.get("project_id") or "").strip() + if client_request_id: + # A retry must resolve to the original experiment even if the + # template changed since, so check before validating anew. + replay = _replay_experiment_request( + cur, + user=user, + client_request_id=client_request_id, + request_fingerprint=request_fingerprint, + ) + if replay: + conn.commit() + return replay execution_request_id = ( str(payload.get("execution_request_id") or "").strip() or (_make_execution_request_id() if project_id else "") @@ -4325,8 +4572,9 @@ def create_experiment_definition_pg( output_snapshot: list[Any] = [] additional_inputs_snapshot: list[Any] = [] if type_id: + visibility_sql, visibility_params = _visible_experiment_type_filter(user) cur.execute( - """ + f""" SELECT equipment_id, type_parameters_json, @@ -4334,8 +4582,9 @@ def create_experiment_definition_pg( additional_inputs_json FROM experiment_types WHERE id = %s + {visibility_sql} """, - (type_id,), + [type_id, *visibility_params], ) type_row = cur.fetchone() if not type_row: @@ -4370,10 +4619,24 @@ def create_experiment_definition_pg( output_snapshot_json, additional_inputs_snapshot_json, owner_id, - owner_org + owner_org, + client_request_id, + client_request_fingerprint, + proposal_generation_status, + proposal_generation_error, + proposal_generation_completed_at, + proposal_generation_updated_at + ) + VALUES ( + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, + %s, %s, + CASE WHEN %s = 'skipped' THEN CURRENT_TIMESTAMP END, + CURRENT_TIMESTAMP ) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - RETURNING id, execution_request_id + ON CONFLICT (owner_id, client_request_id) + WHERE client_request_id <> '' + DO NOTHING + RETURNING id, execution_request_id, proposal_generation_status """, ( project_id or None, @@ -4389,9 +4652,27 @@ def create_experiment_definition_pg( Json(additional_inputs_snapshot), user.id, user.organization, + client_request_id, + request_fingerprint, + proposal_status, + proposal_note, + proposal_status, ), ) row = cur.fetchone() + if not row: + # A concurrent attempt with the same client_request_id won + # the insert; return its experiment instead of a duplicate. + replay = _replay_experiment_request( + cur, + user=user, + client_request_id=client_request_id, + request_fingerprint=request_fingerprint, + ) + conn.commit() + if not replay: + raise ValueError("Duplicate experiment request could not be resolved") + return replay experiment_id = str(row["id"]) execution_request_id = str(row.get("execution_request_id") or execution_request_id) if project_id and type_id: @@ -4443,8 +4724,10 @@ def create_experiment_definition_pg( conn.commit() return { "status": "created", + "idempotent_replay": False, "experiment_id": experiment_id, "execution_request_id": execution_request_id, + "proposal_generation_status": row.get("proposal_generation_status") or proposal_status, } finally: conn.close() @@ -4454,7 +4737,7 @@ def get_experiment_type_constraints_pg(type_id: str) -> dict[str, list[float | N if not type_id: return {} - conn = get_pg_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -4557,133 +4840,315 @@ def save_experiment_proposal_batches_pg( if not experiment_id or not isinstance(batches, list): return 0 - conn = get_pg_connection() - inserted = 0 + conn = get_pg_worker_connection() try: with conn.cursor() as cur: - if not execution_request_id: - cur.execute( - "SELECT execution_request_id FROM experiment_definitions WHERE id = %s", - (experiment_id,), - ) - request_row = cur.fetchone() - execution_request_id = str(request_row[0] or "") if request_row else "" - for batch in batches: - if not isinstance(batch, dict): - continue - proposals = batch.get("proposals") or [] + inserted = _insert_experiment_proposal_batches( + cur, + experiment_id=experiment_id, + project_id=project_id, + optimizer_result=optimizer_result, + iteration_offset=iteration_offset, + execution_request_id=execution_request_id, + ) + conn.commit() + return inserted + finally: + conn.close() + + +def _insert_experiment_proposal_batches( + cur, + *, + experiment_id: str, + project_id: str, + optimizer_result: dict[str, Any], + iteration_offset: int = 0, + execution_request_id: str = "", +) -> int: + """Insert optimizer batches on ``cur``; duplicate iterations are skipped.""" + batches = optimizer_result.get("batches") or [] + if not experiment_id or not isinstance(batches, list): + return 0 + + def first_value(row: Any) -> Any: + if not row: + return None + return next(iter(row.values())) if isinstance(row, dict) else row[0] + + inserted = 0 + if not execution_request_id: + cur.execute( + "SELECT execution_request_id FROM experiment_definitions WHERE id = %s", + (experiment_id,), + ) + execution_request_id = str(first_value(cur.fetchone()) or "") + for batch in batches: + if not isinstance(batch, dict): + continue + proposals = batch.get("proposals") or [] + cur.execute( + """ + INSERT INTO experiment_recipe_batches ( + experiment_id, + project_id, + execution_request_id, + iteration, + optimizer_source, + model, + n_train, + proposals_json, + optimizer_result_json + ) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (experiment_id, iteration) DO NOTHING + RETURNING id + """, + ( + experiment_id, + project_id or None, + execution_request_id, + iteration_offset + int(batch.get("iteration") or 1), + optimizer_result.get("source", ""), + batch.get("model", ""), + int(batch.get("n_train") or 0), + Json(proposals), + Json( + { + "source": optimizer_result.get("source"), + "constraints_applied": optimizer_result.get("constraints_applied", {}), + "constraints_ignored": optimizer_result.get("constraints_ignored", {}), + "training_context": optimizer_result.get("training_context", {}), + "training_snapshot_hash": optimizer_result.get("training_snapshot_hash"), + "training_row_count": optimizer_result.get("training_row_count"), + "training_run_ids": optimizer_result.get("training_run_ids", []), + "model_version": optimizer_result.get("model_version"), + "model_trained_at": optimizer_result.get("model_trained_at"), + "model_snapshot_hash": optimizer_result.get("model_snapshot_hash"), + "auto_proposal_reason": optimizer_result.get("auto_proposal_reason"), + } + ), + ), + ) + batch_id = first_value(cur.fetchone()) + if not batch_id: + continue + if isinstance(proposals, list): + for index, proposal in enumerate(proposals): cur.execute( """ - INSERT INTO experiment_recipe_batches ( + INSERT INTO experiment_recipe_proposals ( + batch_id, experiment_id, project_id, execution_request_id, - iteration, - optimizer_source, - model, - n_train, - proposals_json, - optimizer_result_json + proposal_index, + proposal_json ) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (experiment_id, iteration) DO NOTHING - RETURNING id + VALUES (%s, %s, %s, %s, %s, %s) + ON CONFLICT (batch_id, proposal_index) DO NOTHING """, ( + batch_id, experiment_id, project_id or None, execution_request_id, - iteration_offset + int(batch.get("iteration") or 1), - optimizer_result.get("source", ""), - batch.get("model", ""), - int(batch.get("n_train") or 0), - Json(proposals), - Json( - { - "source": optimizer_result.get("source"), - "constraints_applied": optimizer_result.get("constraints_applied", {}), - "constraints_ignored": optimizer_result.get("constraints_ignored", {}), - "training_context": optimizer_result.get("training_context", {}), - "training_snapshot_hash": optimizer_result.get("training_snapshot_hash"), - "training_row_count": optimizer_result.get("training_row_count"), - "training_run_ids": optimizer_result.get("training_run_ids", []), - "model_version": optimizer_result.get("model_version"), - "model_trained_at": optimizer_result.get("model_trained_at"), - "model_snapshot_hash": optimizer_result.get("model_snapshot_hash"), - "auto_proposal_reason": optimizer_result.get("auto_proposal_reason"), - } - ), + index, + Json(proposal if isinstance(proposal, dict) else {}), ), ) - batch_row = cur.fetchone() - batch_id = batch_row[0] if batch_row else None - if not batch_id: - continue - if isinstance(proposals, list): - for index, proposal in enumerate(proposals): - cur.execute( - """ - INSERT INTO experiment_recipe_proposals ( - batch_id, - experiment_id, - project_id, - execution_request_id, - proposal_index, - proposal_json + inserted += 1 + return inserted + + +# ── Proposal job queue ────────────────────────────────────────────────────── +# +# Rows in experiment_definitions with proposal_generation_status = 'queued' +# are durable proposal jobs. Workers claim a job by moving it to 'generating' +# with a claim id and a lease; results are committed together with the +# 'completed' status only while the claim is still held, so a job whose +# worker died is retried after its lease expires and a stale worker can never +# overwrite a newer attempt. + + +def _fail_exhausted_proposal_jobs(cur, *, experiment_id: str = "") -> int: + """Fail unfinished jobs that ran out of attempts or are too old to run.""" + max_attempts = proposal_job_max_attempts() + id_filter = "AND id = %s" if experiment_id else "" + params: list[Any] = [max_attempts, max_attempts, proposal_job_max_age_hours()] + if experiment_id: + params.append(experiment_id) + cur.execute( + f""" + UPDATE experiment_definitions + SET + proposal_generation_status = 'failed', + proposal_generation_error = CASE + WHEN proposal_generation_attempts >= %s + THEN 'Proposal generation did not finish after repeated attempts.' + ELSE 'Proposal generation expired before it could run.' + END, + proposal_generation_completed_at = CURRENT_TIMESTAMP, + proposal_generation_updated_at = CURRENT_TIMESTAMP, + proposal_generation_claim_id = NULL, + proposal_generation_lease_expires_at = NULL + WHERE ( + proposal_generation_status = 'queued' + OR ( + proposal_generation_status = 'generating' + AND ( + proposal_generation_lease_expires_at IS NULL + OR proposal_generation_lease_expires_at < CURRENT_TIMESTAMP + ) + ) + ) + AND ( + proposal_generation_attempts >= %s + OR created_at < CURRENT_TIMESTAMP - make_interval(hours => %s) + ) + {id_filter} + """, + params, + ) + failed = cur.rowcount or 0 + # An expired 'generating' lease on an experiment that already has recipe + # batches is interrupted follow-up generation from the ingestion loop, not + # an initial job: the initial runner cannot redo it. Record it as failed; + # the follow-up loop retries on its next trigger regardless of status. + followup_params: list[Any] = [experiment_id] if experiment_id else [] + cur.execute( + f""" + UPDATE experiment_definitions e + SET + proposal_generation_status = 'failed', + proposal_generation_error = + 'Follow-up recipe generation was interrupted; it is retried when new results arrive.', + proposal_generation_completed_at = CURRENT_TIMESTAMP, + proposal_generation_updated_at = CURRENT_TIMESTAMP, + proposal_generation_claim_id = NULL, + proposal_generation_lease_expires_at = NULL + WHERE e.proposal_generation_status = 'generating' + AND ( + e.proposal_generation_lease_expires_at IS NULL + OR e.proposal_generation_lease_expires_at < CURRENT_TIMESTAMP + ) + AND EXISTS ( + SELECT 1 FROM experiment_recipe_batches b WHERE b.experiment_id = e.id + ) + {id_filter.replace("id =", "e.id =")} + """, + followup_params, + ) + return failed + (cur.rowcount or 0) + + +def claim_experiment_proposal_job_pg(*, experiment_id: str = "") -> dict[str, Any] | None: + """ + Claim the next runnable proposal job (or the given experiment's job). + + Runnable means ``queued``, or ``generating`` with an expired (or missing) + lease, within the attempt and age limits. ``SKIP LOCKED`` lets several + API workers sweep concurrently without claiming the same job. + """ + conn = get_pg_worker_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + _fail_exhausted_proposal_jobs(cur, experiment_id=experiment_id) + id_filter = "AND id = %s" if experiment_id else "" + params: list[Any] = [proposal_job_max_attempts()] + if experiment_id: + params.append(experiment_id) + params.extend([str(uuid.uuid4()), proposal_job_lease_seconds()]) + cur.execute( + f""" + WITH candidate AS ( + SELECT id + FROM experiment_definitions + WHERE ( + proposal_generation_status = 'queued' + OR ( + proposal_generation_status = 'generating' + AND ( + proposal_generation_lease_expires_at IS NULL + OR proposal_generation_lease_expires_at < CURRENT_TIMESTAMP + ) ) - VALUES (%s, %s, %s, %s, %s, %s) - ON CONFLICT (batch_id, proposal_index) DO NOTHING - """, - ( - batch_id, - experiment_id, - project_id or None, - execution_request_id, - index, - Json(proposal if isinstance(proposal, dict) else {}), - ), ) - inserted += 1 + AND proposal_generation_attempts < %s + {id_filter} + ORDER BY proposal_generation_updated_at NULLS FIRST, created_at + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + UPDATE experiment_definitions e + SET + proposal_generation_status = 'generating', + proposal_generation_attempts = e.proposal_generation_attempts + 1, + proposal_generation_claim_id = %s::uuid, + proposal_generation_lease_expires_at = + CURRENT_TIMESTAMP + make_interval(secs => %s), + proposal_generation_started_at = CURRENT_TIMESTAMP, + proposal_generation_updated_at = CURRENT_TIMESTAMP, + proposal_generation_error = '' + FROM candidate + WHERE e.id = candidate.id + RETURNING + e.id::text AS experiment_id, + e.project_id, + e.type_id, + e.execution_request_id, + e.planned_parameters_json, + e.input_snapshot_json, + e.proposal_generation_attempts, + e.proposal_generation_claim_id::text AS claim_id + """, + params, + ) + row = cur.fetchone() conn.commit() - return inserted + return dict(row) if row else None finally: conn.close() -def generate_initial_experiment_proposals_pg( - *, - experiment_id: str, - project_id: str, - type_id: str, - planned_parameters: list[dict[str, Any]], -) -> dict[str, Any]: - if not experiment_id: - return {"status": "skipped", "saved_recipe_batches": 0, "proposal_count": 0} +def _finish_proposal_job_failed(*, experiment_id: str, claim_id: str, error: str) -> bool: + conn = get_pg_worker_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE experiment_definitions + SET + proposal_generation_status = 'failed', + proposal_generation_error = %s, + proposal_generation_completed_at = CURRENT_TIMESTAMP, + proposal_generation_updated_at = CURRENT_TIMESTAMP, + proposal_generation_claim_id = NULL, + proposal_generation_lease_expires_at = NULL + WHERE id = %s + AND proposal_generation_status = 'generating' + AND proposal_generation_claim_id = %s::uuid + """, + (error[:2000], experiment_id, claim_id), + ) + updated = bool(cur.rowcount) + conn.commit() + return updated + finally: + conn.close() + + +def run_experiment_proposal_job_pg(job: dict[str, Any]) -> dict[str, Any]: + """Compute proposals for a claimed job and commit them with its completion.""" + experiment_id = str(job.get("experiment_id") or "") + claim_id = str(job.get("claim_id") or "") + project_id = str(job.get("project_id") or "") + type_id = str(job.get("type_id") or "") + execution_request_id = str(job.get("execution_request_id") or "") + planned_parameters = _coerce_json(job.get("planned_parameters_json"), []) + input_snapshot = _coerce_json(job.get("input_snapshot_json"), []) - update_experiment_proposal_generation_status_pg( - experiment_id=experiment_id, - status="generating", - ) - optimizer_result: dict[str, Any] = {} try: - input_snapshot: list[Any] = [] - execution_request_id = "" - conn = get_pg_connection() - try: - with conn.cursor(cursor_factory=RealDictCursor) as cur: - cur.execute( - """ - SELECT input_snapshot_json, execution_request_id - FROM experiment_definitions - WHERE id = %s - """, - (experiment_id,), - ) - row = cur.fetchone() - input_snapshot = _coerce_json(row.get("input_snapshot_json"), []) if row else [] - execution_request_id = str(row.get("execution_request_id") or "") if row else "" - finally: - conn.close() constraints = _merge_planned_parameters_into_constraints( _constraints_from_type_parameters(input_snapshot) if input_snapshot @@ -4704,27 +5169,6 @@ def generate_initial_experiment_proposals_pg( "training_scope": "project_plus_global_history", }, ) - saved_batches = save_experiment_proposal_batches_pg( - experiment_id=experiment_id, - project_id=project_id, - optimizer_result=optimizer_result, - execution_request_id=execution_request_id, - ) - proposal_count = sum( - len(batch.get("proposals") or []) - for batch in optimizer_result.get("batches", []) - if isinstance(batch, dict) - ) - update_experiment_proposal_generation_status_pg( - experiment_id=experiment_id, - status="completed", - ) - return { - "status": "completed", - "saved_recipe_batches": saved_batches, - "proposal_count": proposal_count, - "optimizer_result": optimizer_result, - } except Exception as exc: logger.warning( "Failed to generate proposals for experiment %s: %s", @@ -4732,18 +5176,161 @@ def generate_initial_experiment_proposals_pg( exc, exc_info=True, ) - update_experiment_proposal_generation_status_pg( + _finish_proposal_job_failed( experiment_id=experiment_id, - status="failed", - error=str(exc)[:2000], + claim_id=claim_id, + error=str(exc), ) - return { - "status": "failed", - "saved_recipe_batches": 0, - "proposal_count": 0, - "optimizer_result": optimizer_result, - "error": str(exc), - } + return {"experiment_id": experiment_id, "status": "failed", "error": str(exc)} + + conn = get_pg_worker_connection() + try: + with conn.cursor(cursor_factory=RealDictCursor) as cur: + cur.execute( + """ + SELECT id + FROM experiment_definitions + WHERE id = %s + AND proposal_generation_status = 'generating' + AND proposal_generation_claim_id = %s::uuid + FOR UPDATE + """, + (experiment_id, claim_id), + ) + if not cur.fetchone(): + # The lease expired and another worker took over (or the + # ingestion loop finished this experiment); drop our result. + conn.rollback() + logger.info( + "Proposal job for experiment %s lost its claim; discarding result", + experiment_id, + ) + return {"experiment_id": experiment_id, "status": "claim_lost"} + saved_batches = _insert_experiment_proposal_batches( + cur, + experiment_id=experiment_id, + project_id=project_id, + optimizer_result=optimizer_result, + execution_request_id=execution_request_id, + ) + cur.execute( + """ + UPDATE experiment_definitions + SET + proposal_generation_status = 'completed', + proposal_generation_error = '', + proposal_generation_completed_at = CURRENT_TIMESTAMP, + proposal_generation_updated_at = CURRENT_TIMESTAMP, + proposal_generation_claim_id = NULL, + proposal_generation_lease_expires_at = NULL + WHERE id = %s + """, + (experiment_id,), + ) + conn.commit() + except Exception: + # Leave the claim in place: the job is retried once its lease expires. + conn.rollback() + logger.warning( + "Could not persist proposals for experiment %s; job will be retried", + experiment_id, + exc_info=True, + ) + raise + finally: + conn.close() + + proposal_count = sum( + len(batch.get("proposals") or []) + for batch in optimizer_result.get("batches", []) + if isinstance(batch, dict) + ) + return { + "experiment_id": experiment_id, + "status": "completed", + "saved_recipe_batches": saved_batches, + "proposal_count": proposal_count, + } + + +def process_experiment_proposal_jobs_pg( + *, + experiment_id: str = "", + limit: int = 5, +) -> dict[str, Any]: + """Claim and run up to ``limit`` proposal jobs (or one experiment's job).""" + results: list[dict[str, Any]] = [] + for _ in range(max(1, limit)): + job = claim_experiment_proposal_job_pg(experiment_id=experiment_id) + if not job: + break + try: + results.append(run_experiment_proposal_job_pg(job)) + except Exception as exc: + results.append( + { + "experiment_id": job.get("experiment_id"), + "status": "retry_pending", + "error": str(exc), + } + ) + if experiment_id: + break + return {"processed": len(results), "results": results} + + +def start_experiment_proposal_job(experiment_id: str) -> bool: + """ + Run one experiment's proposal job on a background thread. + + This is only a fast path: the job is already durable in the database, so + if the thread cannot start or the process exits, the sweeper runs it. + """ + def _run() -> None: + try: + process_experiment_proposal_jobs_pg(experiment_id=experiment_id, limit=1) + except Exception: + logger.warning( + "Proposal job fast path failed for experiment %s; the sweeper will retry", + experiment_id, + exc_info=True, + ) + + try: + Thread(target=_run, name=f"proposal-job-{experiment_id[:8]}", daemon=True).start() + return True + except Exception: + logger.warning( + "Could not start proposal job thread for experiment %s; the sweeper will run it", + experiment_id, + exc_info=True, + ) + return False + + +def start_proposal_job_sweeper() -> Thread | None: + """ + Periodically run queued and expired proposal jobs in this process. + + Every API worker may run a sweeper; claims use SKIP LOCKED and leases, so + concurrent sweepers never run the same job at once. Set + PROPOSAL_JOB_SWEEP_SECONDS=0 to disable. + """ + interval = _env_int("PROPOSAL_JOB_SWEEP_SECONDS", 60, 0) + if interval <= 0 or os.getenv("USE_MOCK_DATABASE", "false").lower() == "true": + return None + + def _loop() -> None: + while True: + try: + process_experiment_proposal_jobs_pg(limit=5) + except Exception: + logger.warning("Proposal job sweep failed", exc_info=True) + time.sleep(interval) + + thread = Thread(target=_loop, name="proposal-job-sweeper", daemon=True) + thread.start() + return thread def reconcile_ingested_runs_with_recipe_proposals_pg( @@ -4755,7 +5342,7 @@ def reconcile_ingested_runs_with_recipe_proposals_pg( if not runs_data: return {"matched": 0, "checked": 0} - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() matched = 0 checked = 0 try: @@ -4908,7 +5495,7 @@ def auto_generate_recipe_proposals_for_projects_pg( "generated": 0, } - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() generated = 0 skipped: list[dict[str, Any]] = [] processed_projects: set[str] = set() @@ -5266,7 +5853,7 @@ def list_experiment_definitions_pg( visible_project_ids: list[str], project_id: str = "", ) -> list[dict[str, Any]]: - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) params: list[Any] = [] conditions: list[str] = [] @@ -5478,7 +6065,7 @@ def list_execution_queue_pg(user: PlatformUser) -> list[dict[str, Any]]: if user.role not in {"admin", "equipment_owner"}: raise PermissionError("Facility execution queue access required") - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) params: list[Any] = [] conditions = ["e.project_id IS NOT NULL", "e.execution_request_id <> ''"] if user.role != "admin": @@ -5506,7 +6093,7 @@ def list_execution_queue_pg(user: PlatformUser) -> list[dict[str, Any]]: e.data_ingested_at, e.execution_updated_at, e.project_id, - p.name AS project_name, + COALESCE(p.name, app_experiment_project_name(e.id)) AS project_name, e.equipment_id, COALESCE(em.equipment_name, p.equipment_name, e.equipment_id) AS equipment_name, e.type_id, @@ -5618,7 +6205,7 @@ def update_execution_queue_status_pg( if status not in EXECUTION_STATUSES: raise ValueError("Invalid execution status") - conn = get_pg_superuser_connection() + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -5702,7 +6289,7 @@ def update_recipe_proposal_status_pg( if status not in {"pending", "accepted", "attempted", "completed", "rejected"}: raise ValueError("Invalid recipe proposal status") - conn = get_pg_connection(user.id) + conn = get_pg_connection(user) try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -5808,7 +6395,7 @@ def drain_recipe_optimizer_wakeups_pg( if str(proposal_id or "").strip() } ) - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: id_filter = ( @@ -5918,7 +6505,7 @@ def drain_recipe_optimizer_wakeups_pg( def ensure_data_uploads_kind_column_pg() -> None: """Add the data_uploads.kind column on databases created before split input/output uploads existed. Idempotent — safe to call on every startup.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( @@ -5932,7 +6519,7 @@ def ensure_data_uploads_kind_column_pg() -> None: def ensure_data_uploads_processing_started_at_pg() -> None: """Add a processing lease timestamp for restart-safe upload retries.""" - conn = get_pg_superuser_connection() + conn = get_pg_migration_connection() try: with conn.cursor() as cur: cur.execute( diff --git a/api/middleware/error_handler.py b/api/middleware/error_handler.py index f4f37af..4135ba2 100644 --- a/api/middleware/error_handler.py +++ b/api/middleware/error_handler.py @@ -40,3 +40,26 @@ async def global_exception_handler(request: Request, exc: Exception) -> JSONResp "path": str(request.url.path), }, ) + + +async def insufficient_privilege_handler(request: Request, exc: Exception) -> JSONResponse: + """ + Map a database permission denial (a row-level security policy rejecting a + write, SQLSTATE 42501) to 403. The API's own checks should reject these + first; reaching the database means a check is missing, so log it. + """ + trace_id = getattr(request.state, "trace_id", uuid.uuid4().hex[:8]) + logger.warning( + "[%s] Database denied %s %s: %s", + trace_id, + request.method, + request.url.path, + str(exc).strip().splitlines()[0] if str(exc).strip() else type(exc).__name__, + ) + return JSONResponse( + status_code=403, + content={ + "detail": "Not authorized to perform this operation", + "trace_id": trace_id, + }, + ) diff --git a/api/ml_response_cache.py b/api/ml_response_cache.py index ec84e44..ada9772 100644 --- a/api/ml_response_cache.py +++ b/api/ml_response_cache.py @@ -21,7 +21,7 @@ from psycopg2.extras import Json, RealDictCursor -from data_loader_pg import get_pg_superuser_connection +from data_loader_pg import get_pg_migration_connection, get_pg_worker_connection from security import PlatformUser @@ -60,7 +60,41 @@ def _lock_for(cache_key: str) -> threading.Lock: return lock +def ensure_ml_endpoint_cache_table_pg() -> None: + """Create the ML response cache table (schema migration; owner connection).""" + conn = get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + CREATE TABLE IF NOT EXISTS ml_endpoint_cache ( + cache_key TEXT PRIMARY KEY, + endpoint TEXT NOT NULL, + user_scope TEXT NOT NULL, + project_id TEXT NOT NULL DEFAULT '', + constraints_hash TEXT NOT NULL DEFAULT '', + data_signature TEXT NOT NULL, + payload_json JSONB NOT NULL, + computed_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + expires_at TIMESTAMPTZ NOT NULL, + compute_seconds DOUBLE PRECISION NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ) + """ + ) + cur.execute( + """ + CREATE INDEX IF NOT EXISTS idx_ml_endpoint_cache_lookup + ON ml_endpoint_cache(endpoint, user_scope, project_id, expires_at DESC) + """ + ) + conn.commit() + finally: + conn.close() + + def _ensure_table() -> None: + """Confirm the cache table exists; the schema migration creates it.""" global _table_ready if _table_ready: return @@ -68,41 +102,24 @@ def _ensure_table() -> None: with _table_guard: if _table_ready: return - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: - cur.execute( - """ - CREATE TABLE IF NOT EXISTS ml_endpoint_cache ( - cache_key TEXT PRIMARY KEY, - endpoint TEXT NOT NULL, - user_scope TEXT NOT NULL, - project_id TEXT NOT NULL DEFAULT '', - constraints_hash TEXT NOT NULL DEFAULT '', - data_signature TEXT NOT NULL, - payload_json JSONB NOT NULL, - computed_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - expires_at TIMESTAMPTZ NOT NULL, - compute_seconds DOUBLE PRECISION NOT NULL DEFAULT 0, - updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) - """ - ) - cur.execute( - """ - CREATE INDEX IF NOT EXISTS idx_ml_endpoint_cache_lookup - ON ml_endpoint_cache(endpoint, user_scope, project_id, expires_at DESC) - """ - ) + cur.execute("SELECT to_regclass('public.ml_endpoint_cache')") + exists = (cur.fetchone() or [None])[0] is not None conn.commit() - _table_ready = True finally: conn.close() + if not exists: + raise RuntimeError( + "ml_endpoint_cache table is missing; run the schema migrations" + ) + _table_ready = True def _data_signature(project_id: str) -> str: """Cheap signature for the training data behind the ML pages.""" - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: if project_id: @@ -143,7 +160,7 @@ def _data_signature(project_id: str) -> str: def _read_cache(cache_key: str, *, include_expired: bool = False) -> dict[str, Any] | None: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: if include_expired: @@ -198,7 +215,7 @@ def _write_cache( ttl_seconds: int, compute_seconds: float, ) -> None: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( diff --git a/api/model_registry.py b/api/model_registry.py index b0abade..cfd5de8 100644 --- a/api/model_registry.py +++ b/api/model_registry.py @@ -81,10 +81,10 @@ def save_model( Does NOT promote the row; call `promote_model` after metrics comparison. """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection artifact = _pickle_bundle(bundle) - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -154,9 +154,9 @@ def promote_model( - "new_r2": float | None - "prev_r2": float | None """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -259,9 +259,9 @@ def update_retrain_state( reason: str, ) -> None: """Reset `runs_since_last_train` to 0 and stamp a successful retrain.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -311,9 +311,9 @@ def deactivate_active_models( ) if not normalized_targets: return [] - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -351,9 +351,9 @@ def deactivate_active_models( def bump_runs_counter(*, domain_id: str, n_new: int) -> int: """Increment the debounce counter. Returns the new total.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -395,9 +395,9 @@ def credit_runs_counter_once( ), } - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: credited = 0 @@ -483,9 +483,9 @@ def load_active_bundle(*, domain_id: str, target: str) -> Optional[dict[str, Any "scaler": fitted StandardScaler | None, } """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -539,10 +539,10 @@ def load_active_bundle(*, domain_id: str, target: str) -> Optional[dict[str, Any def list_active_model_metadata(*, domain_id: str) -> list[dict[str, Any]]: """Return active model identity and metrics without loading pickle artifacts.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -580,11 +580,11 @@ def list_models( limit: int = 50, ) -> list[dict[str, Any]]: """List recent models (metadata only, no artifacts).""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: if domain_id: @@ -625,11 +625,11 @@ def list_models( def get_retrain_state(domain_id: str) -> Optional[dict[str, Any]]: """Fetch the debounce row for one domain.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: # Defensive SELECT *: the alert columns are added by a later @@ -667,11 +667,11 @@ def get_active_model_metadata( Used by the drift-alert logic (needs R² and version, not the pickle). """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection from psycopg2.extras import RealDictCursor - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( @@ -706,9 +706,9 @@ def increment_non_promotion_streak(domain_id: str) -> int: applied yet, this is a silent no-op (returns 0) so the retrain path keeps working. """ - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: try: @@ -742,9 +742,9 @@ def increment_non_promotion_streak(domain_id: str) -> int: def reset_non_promotion_streak(domain_id: str) -> None: """Reset ``consecutive_non_promotions`` to 0. Called after a clean retrain.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: try: @@ -770,9 +770,9 @@ def reset_non_promotion_streak(domain_id: str) -> None: def record_alert_sent(*, domain_id: str, reason: str) -> None: """Stamp ``last_alert_sent_at`` for cooldown throttling.""" - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: try: diff --git a/api/routers/admin.py b/api/routers/admin.py index fb3f034..8813b3d 100644 --- a/api/routers/admin.py +++ b/api/routers/admin.py @@ -4,7 +4,7 @@ from fastapi import APIRouter, Depends, HTTPException from psycopg2.extras import RealDictCursor from pydantic import BaseModel, EmailStr -from data_loader_pg import get_pg_connection, get_pg_superuser_connection +from data_loader_pg import get_pg_connection, get_pg_worker_connection from domain_configs import CONFIGS_DIR from metadata_pg import ( list_execution_queue_pg, @@ -40,7 +40,7 @@ def get_or_create_user_role( If the user does not exist, registers them automatically as 'researcher'. """ try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: user = sync_canonical_user( cur, @@ -113,7 +113,7 @@ def invite_user( ) -> Dict[str, Any]: """Create or update a platform user record before first login.""" try: - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() with conn.cursor(cursor_factory=RealDictCursor) as cur: inferred_user_id = payload.email.split("@")[0] user = sync_canonical_user( diff --git a/api/routers/dataset_v2.py b/api/routers/dataset_v2.py index 41a4241..1a7e42a 100644 --- a/api/routers/dataset_v2.py +++ b/api/routers/dataset_v2.py @@ -270,6 +270,7 @@ def create_project( access_mode=payload.access, nanohub_user_id=user.id, pi_name=user.name, + creator_role=user.role, ) diff --git a/api/routers/equipment.py b/api/routers/equipment.py index fb0f2b1..8067d8b 100644 --- a/api/routers/equipment.py +++ b/api/routers/equipment.py @@ -5,7 +5,6 @@ from their own organization, while owners/admins can also see pending items. """ import logging -from threading import Thread from typing import List, Optional from fastapi import APIRouter, Depends, HTTPException @@ -19,6 +18,7 @@ summarize_domain_config, ) from metadata_pg import ( + ExperimentRequestConflictError, approved_equipment_available_to_user, create_experiment_definition_pg, create_experiment_type_pg, @@ -29,7 +29,6 @@ fork_experiment_type_pg, get_equipment_pg, get_experiment_type_versions_pg, - generate_initial_experiment_proposals_pg, grant_equipment_access_pg, list_equipment_access_pg, list_experiment_definitions_pg, @@ -39,7 +38,7 @@ register_equipment_pg, revoke_equipment_access_pg, search_users_pg, - update_experiment_proposal_generation_status_pg, + start_experiment_proposal_job, update_experiment_type_pg, update_equipment_pg, update_recipe_proposal_status_pg, @@ -201,6 +200,9 @@ class ExperimentPayload(BaseModel): default_factory=OptimizationContextPayload, ) sample_ids: List[str] = [] + # Client-generated key reused on retries so a resubmitted request + # returns the experiment the first attempt created. + client_request_id: str = Field(default="", max_length=200) owner_id: str = "" owner_org: str = "" @@ -508,10 +510,10 @@ def delete_experiment_type( @router.get("/experiment-type/{type_id}/versions") def get_experiment_type_versions( type_id: str, - _user: PlatformUser = Depends(get_platform_user), + user: PlatformUser = Depends(get_platform_user), ): - """Get version history for an experiment type.""" - return get_experiment_type_versions_pg(type_id) + """Get version history for an experiment type the caller can see.""" + return get_experiment_type_versions_pg(type_id, user=user) @router.post("/experiment-type/{type_id}/fork") @@ -616,42 +618,30 @@ def create_experiment_definition( "planned_parameters": [param.model_dump() for param in payload.planned_parameters], "optimization_context": optimization_context, "sample_ids": payload.sample_ids, + "client_request_id": payload.client_request_id, + "generate_proposals": should_generate_proposals, }, user=user, ) + except ExperimentRequestConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc experiment_id = experiment_result.get("experiment_id", "") - if experiment_id and should_generate_proposals: - update_experiment_proposal_generation_status_pg( - experiment_id=experiment_id, - status="queued", - ) - proposal_thread = Thread( - target=generate_initial_experiment_proposals_pg, - kwargs={ - "experiment_id": experiment_id, - "project_id": project_id, - "type_id": payload.type_id, - "planned_parameters": [ - param.model_dump() for param in payload.planned_parameters - ], - }, - daemon=True, - ) - proposal_thread.start() - elif experiment_id: - update_experiment_proposal_generation_status_pg( - experiment_id=experiment_id, - status="skipped", - error="Recipe generation was not requested for this project run.", - ) + generation_status = str( + experiment_result.get("proposal_generation_status") + or ("queued" if should_generate_proposals else "skipped") + ) + if experiment_id and generation_status == "queued": + # The queued job is already committed with the experiment; this only + # starts it sooner. If the thread never runs, the sweeper does. + start_experiment_proposal_job(experiment_id) return { **experiment_result, "project_id": project_id, - "optimization_status": "queued" if experiment_id and should_generate_proposals else "skipped", - "proposal_generation_status": "queued" if experiment_id and should_generate_proposals else "skipped", + "optimization_status": generation_status, + "proposal_generation_status": generation_status, "saved_recipe_batches": 0, "proposal_count": 0, "optimizer_result": {}, diff --git a/api/runtime_config.py b/api/runtime_config.py index 66919e5..7e6b0b2 100644 --- a/api/runtime_config.py +++ b/api/runtime_config.py @@ -9,9 +9,13 @@ class MissingRuntimeConfiguration(RuntimeError): "DT_SYSTEM_TOKEN", "INGESTION_TOKEN", "PG_PASS", - "PG_SUPER_PASS", ) +# Background jobs connect as the dedicated ``api_worker`` role. Until that role +# is provisioned, they fall back to the migration (table-owner) credentials so +# existing deployments keep working. One of the two must be configured. +WORKER_CREDENTIAL_ENV_VARS = ("PG_WORKER_PASS", "PG_SUPER_PASS") + def require_env(name: str) -> str: value = os.getenv(name, "").strip() @@ -38,9 +42,27 @@ def get_pg_superuser_password() -> str: return require_env("PG_SUPER_PASS") +def get_pg_worker_password() -> str: + """Return the api_worker password, or empty string when not provisioned.""" + return os.getenv("PG_WORKER_PASS", "").strip() + + +def pg_worker_role_configured() -> bool: + return bool(get_pg_worker_password()) + + +def pg_migration_credentials_configured() -> bool: + return bool(os.getenv("PG_SUPER_PASS", "").strip()) + + def validate_runtime_configuration() -> None: for name in REQUIRED_RUNTIME_ENV_VARS: require_env(name) + if not any(os.getenv(name, "").strip() for name in WORKER_CREDENTIAL_ENV_VARS): + raise MissingRuntimeConfiguration( + "Missing background worker database credentials: set PG_WORKER_PASS " + "(preferred) or PG_SUPER_PASS" + ) # ── Optional: FAIR Data Archival (PURR) ────────────────────────────────────── diff --git a/api/scripts/add_experiment_proposal_jobs.sql b/api/scripts/add_experiment_proposal_jobs.sql new file mode 100644 index 0000000..3cf231b --- /dev/null +++ b/api/scripts/add_experiment_proposal_jobs.sql @@ -0,0 +1,20 @@ +-- Durable proposal jobs and idempotent experiment creation. +-- Also applied at startup by metadata_pg.ensure_experiment_proposal_generation_columns_pg(). + +ALTER TABLE experiment_definitions + ADD COLUMN IF NOT EXISTS proposal_generation_attempts INTEGER NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS proposal_generation_claim_id UUID, + ADD COLUMN IF NOT EXISTS proposal_generation_lease_expires_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS client_request_id VARCHAR(200) NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS client_request_fingerprint VARCHAR(64) NOT NULL DEFAULT ''; + +-- A retried create request (same owner and client_request_id) resolves to the +-- experiment the first attempt created. +CREATE UNIQUE INDEX IF NOT EXISTS uq_experiment_definitions_client_request +ON experiment_definitions (owner_id, client_request_id) +WHERE client_request_id <> ''; + +-- The proposal job sweeper only scans unfinished jobs. +CREATE INDEX IF NOT EXISTS idx_experiment_definitions_open_proposal_jobs +ON experiment_definitions (proposal_generation_updated_at) +WHERE proposal_generation_status IN ('queued', 'generating'); diff --git a/api/scripts/add_row_level_security.sql b/api/scripts/add_row_level_security.sql new file mode 100644 index 0000000..01036e8 --- /dev/null +++ b/api/scripts/add_row_level_security.sql @@ -0,0 +1,748 @@ +-- Row-level security for experiment data and project membership. +-- +-- Applied by db_migrations.ensure_row_level_security_pg() with the migration +-- (table-owner) connection, in one transaction; it is idempotent and safe to +-- re-run. Prefer `python -m db_migrations`. If you must apply it by hand, run +-- it atomically so membership grants never exist without their RLS: +-- psql --single-transaction -v ON_ERROR_STOP=1 -f api/scripts/add_row_level_security.sql +-- +-- This script creates the helper functions, grants, policies and guards. It +-- enforces RLS on project_members immediately, together with the write grants +-- it adds there, because membership is what every other policy trusts. It does +-- NOT switch enforcement on for the experiment tables: the migration runner +-- does that when DB_RLS_ENFORCEMENT=enforce, so those policies can be deployed +-- first and enforced once verified (see api/ROW_LEVEL_SECURITY.md). The admin, +-- worker and write policies on the six tables that already enforce RLS +-- (projects and the run tables) also take effect immediately. +-- +-- Identity comes from session settings the API sets on every user-scoped +-- api_client connection (data_loader_pg.get_pg_connection): +-- app.current_user platform user id +-- app.current_user_role admin | pi | equipment_owner | researcher +-- app.current_user_org organization +-- Background jobs connect as api_worker (not superuser, not BYPASSRLS) and get +-- explicit per-table policies below. + +-- Ledger used by db_migrations to skip re-applying an unchanged script. +CREATE TABLE IF NOT EXISTS app_schema_migrations ( + name TEXT PRIMARY KEY, + checksum TEXT NOT NULL, + applied_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +-- ── Roles ─────────────────────────────────────────────────────────────────── + +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'api_worker') THEN + CREATE ROLE api_worker NOLOGIN NOSUPERUSER NOBYPASSRLS NOCREATEDB NOCREATEROLE; + END IF; +END +$$; + +-- ── Identity helpers ─────────────────────────────────────────────────────── + +CREATE OR REPLACE FUNCTION app_current_user_id() RETURNS text +LANGUAGE sql STABLE +AS $$ SELECT NULLIF(current_setting('app.current_user', true), '') $$; + +CREATE OR REPLACE FUNCTION app_current_user_role() RETURNS text +LANGUAGE sql STABLE +AS $$ SELECT NULLIF(current_setting('app.current_user_role', true), '') $$; + +CREATE OR REPLACE FUNCTION app_current_user_org() RETURNS text +LANGUAGE sql STABLE +AS $$ SELECT NULLIF(current_setting('app.current_user_org', true), '') $$; + +CREATE OR REPLACE FUNCTION app_is_admin() RETURNS boolean +LANGUAGE sql STABLE +AS $$ + SELECT app_current_user_id() IS NOT NULL + AND COALESCE(app_current_user_role() = 'admin', false) +$$; + +-- ── Access helpers ───────────────────────────────────────────────────────── +-- SECURITY DEFINER so a policy can consult other protected tables without +-- recursing into their policies. Each one only answers a question about the +-- current session user, and pins search_path. + +-- Mirrors metadata_pg._project_access_role: 'admin', the member's role, or +-- 'viewer' for a signed-in user on an open project; NULL means no access. +CREATE OR REPLACE FUNCTION app_project_role(p_project_id text) RETURNS text +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT CASE + WHEN p_project_id IS NULL OR app_current_user_id() IS NULL THEN NULL + WHEN app_is_admin() THEN 'admin' + ELSE COALESCE( + ( + SELECT COALESCE(NULLIF(pm.role, ''), 'member') + FROM project_members pm + WHERE pm.project_id = p_project_id + AND pm.nanohub_user_id = app_current_user_id() + ), + ( + SELECT 'viewer' + FROM projects p + WHERE p.id = p_project_id + AND p.access_mode = 'open' + ) + ) + END +$$; + +-- Mirrors metadata_pg._ensure_project_write_access. +CREATE OR REPLACE FUNCTION app_can_write_project(p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ SELECT COALESCE(app_project_role(p_project_id) IN ('admin', 'pi', 'member'), false) $$; + +CREATE OR REPLACE FUNCTION app_is_project_member(p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_is_admin() OR EXISTS ( + SELECT 1 + FROM project_members pm + WHERE pm.project_id = p_project_id + AND pm.nanohub_user_id = app_current_user_id() + ) +$$; + +CREATE OR REPLACE FUNCTION app_is_project_pi(p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_is_admin() OR EXISTS ( + SELECT 1 + FROM project_members pm + WHERE pm.project_id = p_project_id + AND pm.nanohub_user_id = app_current_user_id() + AND pm.role = 'pi' + ) +$$; + +CREATE OR REPLACE FUNCTION app_project_has_members(p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ SELECT EXISTS (SELECT 1 FROM project_members WHERE project_id = p_project_id) $$; + +-- True only for a project row inserted by the current transaction, so the +-- "creator enrols as PI" rule cannot be used to claim an existing project +-- that happens to have no members (for example one seeded by ingestion). +CREATE OR REPLACE FUNCTION app_project_created_in_current_transaction(p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT EXISTS ( + SELECT 1 + FROM projects p + WHERE p.id = p_project_id + AND p.xmin = pg_current_xact_id()::xid + ) +$$; + +-- Mirrors metadata_pg.list_execution_queue_pg: facility managers act on +-- execution requests for equipment they (or their organization) own. +CREATE OR REPLACE FUNCTION app_manages_equipment(p_equipment_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_is_admin() OR ( + COALESCE(app_current_user_role() = 'equipment_owner', false) + AND EXISTS ( + SELECT 1 + FROM equipment_metadata em + WHERE em.domain_id = p_equipment_id + AND ( + em.owner_id = app_current_user_id() + OR (COALESCE(em.owner_org, '') <> '' AND em.owner_org = app_current_user_org()) + ) + ) + ) +$$; + +-- Experiment definition read rule (mirrors list_experiment_definitions_pg +-- and the facility execution queue). +CREATE OR REPLACE FUNCTION app_experiment_readable( + p_owner_id text, + p_project_id text, + p_equipment_id text, + p_execution_request_id text +) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_current_user_id() IS NOT NULL AND ( + app_is_admin() + OR p_owner_id = app_current_user_id() + OR (p_project_id IS NOT NULL AND app_project_role(p_project_id) IS NOT NULL) + OR ( + p_project_id IS NOT NULL + AND COALESCE(p_execution_request_id, '') <> '' + AND app_manages_equipment(p_equipment_id) + ) + ) +$$; + +-- Experiment definition write rule: owner, project writers, and facility +-- managers of the equipment for execution requests (the API limits which +-- columns each of them changes). +CREATE OR REPLACE FUNCTION app_experiment_writable( + p_owner_id text, + p_project_id text, + p_equipment_id text, + p_execution_request_id text +) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_current_user_id() IS NOT NULL AND ( + app_is_admin() + OR p_owner_id = app_current_user_id() + OR (p_project_id IS NOT NULL AND app_can_write_project(p_project_id)) + OR ( + p_project_id IS NOT NULL + AND COALESCE(p_execution_request_id, '') <> '' + AND app_manages_equipment(p_equipment_id) + ) + ) +$$; + +CREATE OR REPLACE FUNCTION app_experiment_readable_by_id(p_experiment_id uuid) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT COALESCE(( + SELECT app_experiment_readable(e.owner_id, e.project_id, e.equipment_id, e.execution_request_id) + FROM experiment_definitions e + WHERE e.id = p_experiment_id + ), false) +$$; + +CREATE OR REPLACE FUNCTION app_experiment_writable_by_id(p_experiment_id uuid) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT COALESCE(( + SELECT app_experiment_writable(e.owner_id, e.project_id, e.equipment_id, e.execution_request_id) + FROM experiment_definitions e + WHERE e.id = p_experiment_id + ), false) +$$; + +-- Name of the project an experiment belongs to, for callers who can read the +-- experiment but not the project itself (a facility manager reviewing an +-- execution request). Exposes only the name, never project visibility. +CREATE OR REPLACE FUNCTION app_experiment_project_name(p_experiment_id uuid) RETURNS text +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT p.name + FROM experiment_definitions e + JOIN projects p ON p.id = e.project_id + WHERE e.id = p_experiment_id + AND app_experiment_readable(e.owner_id, e.project_id, e.equipment_id, e.execution_request_id) +$$; + +-- Mirrors update_recipe_proposal_status_pg: admin, experiment owner, or any +-- member of the experiment's project. +CREATE OR REPLACE FUNCTION app_can_update_recipe_proposals(p_experiment_id uuid) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_current_user_id() IS NOT NULL AND COALESCE(( + SELECT app_is_admin() + OR e.owner_id = app_current_user_id() + OR (e.project_id IS NOT NULL AND app_is_project_member(e.project_id)) + FROM experiment_definitions e + WHERE e.id = p_experiment_id + ), false) +$$; + +-- Template read rule: list_experiment_types_pg (owner, shared, published, +-- same organization), plus templates attached to a project or used by an +-- experiment the caller can read, so joined names do not disappear. +-- Published templates are readable without a signed-in user. +CREATE OR REPLACE FUNCTION app_experiment_type_readable( + p_type_id text, + p_owner_id text, + p_owner_org text, + p_visibility text +) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT p_visibility = 'published' OR ( + app_current_user_id() IS NOT NULL AND ( + app_is_admin() + OR p_owner_id = app_current_user_id() + OR p_visibility = 'shared' + OR (COALESCE(p_owner_org, '') <> '' AND p_owner_org = app_current_user_org()) + OR EXISTS ( + SELECT 1 + FROM project_experiments pe + WHERE pe.type_id = p_type_id + AND app_project_role(pe.project_id) IS NOT NULL + ) + OR EXISTS ( + SELECT 1 + FROM experiment_definitions e + WHERE e.type_id = p_type_id + AND app_experiment_readable(e.owner_id, e.project_id, e.equipment_id, e.execution_request_id) + ) + ) + ) +$$; + +CREATE OR REPLACE FUNCTION app_experiment_type_readable_by_id(p_type_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT COALESCE(( + SELECT app_experiment_type_readable(t.id, t.owner_id, t.owner_org, t.visibility) + FROM experiment_types t + WHERE t.id = p_type_id + ), false) +$$; + +-- Mirrors update/delete_experiment_type_pg: only the owner or an admin. +CREATE OR REPLACE FUNCTION app_experiment_type_editable_by_id(p_type_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_current_user_id() IS NOT NULL AND COALESCE(( + SELECT app_is_admin() OR t.owner_id = app_current_user_id() + FROM experiment_types t + WHERE t.id = p_type_id + ), false) +$$; + +-- Sample read rule, evaluated on the row's own columns so that +-- INSERT ... RETURNING can see the row it just created: members and viewers +-- of the sample's project, plus anyone who can read an experiment the sample +-- is linked to (for example a facility manager's execution queue). +CREATE OR REPLACE FUNCTION app_sample_readable(p_sample_id text, p_project_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT app_current_user_id() IS NOT NULL AND ( + app_project_role(p_project_id) IS NOT NULL + OR EXISTS ( + SELECT 1 + FROM experiment_samples es + JOIN experiment_definitions e ON e.id = es.experiment_id + WHERE es.sample_id = p_sample_id + AND app_experiment_readable(e.owner_id, e.project_id, e.equipment_id, e.execution_request_id) + ) + ) +$$; + +CREATE OR REPLACE FUNCTION app_sample_readable_by_id(p_sample_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT COALESCE(( + SELECT app_sample_readable(s.id, s.project_id) + FROM samples s + WHERE s.id = p_sample_id + ), false) +$$; + +CREATE OR REPLACE FUNCTION app_sample_writable_by_id(p_sample_id text) RETURNS boolean +LANGUAGE sql STABLE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ + SELECT COALESCE(( + SELECT app_can_write_project(s.project_id) + FROM samples s + WHERE s.id = p_sample_id + ), false) +$$; + +-- Deleting a template detaches references in other users' projects and +-- processes, which the caller cannot see. This function re-checks that the +-- caller owns the template (or is an admin) before touching them. +CREATE OR REPLACE FUNCTION app_detach_experiment_type(p_type_id text) +RETURNS TABLE ( + project_attachment_count integer, + process_step_count integer, + historical_run_count integer, + cleared_project_upstreams integer, + cleared_process_upstreams integer +) +LANGUAGE plpgsql VOLATILE SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ +DECLARE + v_owner text; +BEGIN + SELECT t.owner_id INTO v_owner FROM experiment_types t WHERE t.id = p_type_id; + IF NOT FOUND THEN + RAISE EXCEPTION 'experiment type % not found', p_type_id USING ERRCODE = 'no_data_found'; + END IF; + -- IS TRUE: a NULL owner or identity must deny, not fall through. + IF NOT COALESCE( + app_is_admin() + OR (app_current_user_id() IS NOT NULL AND v_owner = app_current_user_id()), + false + ) THEN + RAISE EXCEPTION 'only the owner or an admin can delete this template' + USING ERRCODE = 'insufficient_privilege'; + END IF; + + SELECT count(*) INTO project_attachment_count FROM project_experiments pe WHERE pe.type_id = p_type_id; + SELECT count(*) INTO process_step_count FROM process_steps ps WHERE ps.type_id = p_type_id; + SELECT count(*) INTO historical_run_count FROM experiment_definitions e WHERE e.type_id = p_type_id; + + UPDATE project_experiments pe + SET upstream_type_id = '', updated_at = CURRENT_TIMESTAMP + WHERE pe.upstream_type_id = p_type_id; + GET DIAGNOSTICS cleared_project_upstreams = ROW_COUNT; + + UPDATE process_steps ps + SET upstream_type_id = '', updated_at = CURRENT_TIMESTAMP + WHERE ps.upstream_type_id = p_type_id; + GET DIAGNOSTICS cleared_process_upstreams = ROW_COUNT; + + RETURN NEXT; +END +$$; + +REVOKE ALL ON FUNCTION app_detach_experiment_type(text) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION app_detach_experiment_type(text) TO api_client; + +-- ── Guards ───────────────────────────────────────────────────────────────── +-- An experiment's scope (owner, project, equipment, template, execution +-- request) is fixed once created for signed-in non-admin sessions. RLS +-- WITH CHECK cannot compare old and new rows, so without this an owner could +-- move an experiment into another tenant's project. Background jobs (no user +-- identity) and foreign-key actions (ON DELETE SET NULL runs at trigger depth +-- above 1) are unaffected. Always active, independent of DB_RLS_ENFORCEMENT. +CREATE OR REPLACE FUNCTION app_guard_experiment_scope() RETURNS trigger +LANGUAGE plpgsql +SET search_path = public, pg_temp +AS $$ +BEGIN + IF pg_trigger_depth() = 1 + AND app_current_user_id() IS NOT NULL + AND NOT app_is_admin() + AND ( + NEW.owner_id IS DISTINCT FROM OLD.owner_id + OR NEW.project_id IS DISTINCT FROM OLD.project_id + OR NEW.equipment_id IS DISTINCT FROM OLD.equipment_id + OR NEW.type_id IS DISTINCT FROM OLD.type_id + OR NEW.execution_request_id IS DISTINCT FROM OLD.execution_request_id + ) + THEN + RAISE EXCEPTION 'experiment owner, project, equipment, template and request id cannot be changed' + USING ERRCODE = 'insufficient_privilege'; + END IF; + -- A facility manager's access comes only from managing the equipment: + -- they may move an execution request through the queue, nothing else. + IF pg_trigger_depth() = 1 + AND app_current_user_id() IS NOT NULL + AND NOT app_is_admin() + AND OLD.owner_id IS DISTINCT FROM app_current_user_id() + AND NOT (OLD.project_id IS NOT NULL AND app_can_write_project(OLD.project_id)) + AND ( + to_jsonb(NEW) - ARRAY[ + 'execution_status', 'execution_status_note', 'facility_reviewer_id', + 'facility_reviewed_at', 'scheduled_for', 'executed_at', + 'data_ingested_at', 'execution_updated_at' + ] + ) IS DISTINCT FROM ( + to_jsonb(OLD) - ARRAY[ + 'execution_status', 'execution_status_note', 'facility_reviewer_id', + 'facility_reviewed_at', 'scheduled_for', 'executed_at', + 'data_ingested_at', 'execution_updated_at' + ] + ) + THEN + RAISE EXCEPTION 'facility managers can only change execution status fields' + USING ERRCODE = 'insufficient_privilege'; + END IF; + RETURN NEW; +END +$$; + +DROP TRIGGER IF EXISTS dt_experiment_definitions_scope_guard ON experiment_definitions; +CREATE TRIGGER dt_experiment_definitions_scope_guard + BEFORE UPDATE ON experiment_definitions + FOR EACH ROW EXECUTE FUNCTION app_guard_experiment_scope(); + +-- ── Grants ───────────────────────────────────────────────────────────────── + +GRANT USAGE ON SCHEMA public TO api_client, api_worker; + +-- api_worker runs every background job; it gets data access (never DDL) on +-- all current tables, and on tables the migration role creates later. +GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO api_worker; +GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA public TO api_worker; +ALTER DEFAULT PRIVILEGES IN SCHEMA public + GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO api_worker; +ALTER DEFAULT PRIVILEGES IN SCHEMA public + GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO api_worker; +-- The migration ledger decides whether this script re-runs; owner only. +REVOKE ALL ON app_schema_migrations FROM PUBLIC, api_client, api_worker; + +-- api_client now creates projects and manages membership as the caller. +GRANT SELECT, INSERT, DELETE ON projects TO api_client; +GRANT SELECT, INSERT, DELETE ON project_members TO api_client; +GRANT SELECT, INSERT, UPDATE, DELETE ON + experiment_definitions, + experiment_types, + experiment_type_versions, + project_experiments, + samples, + experiment_samples, + run_samples +TO api_client; +GRANT SELECT ON experiment_recipe_batches TO api_client; +GRANT SELECT, UPDATE ON experiment_recipe_proposals TO api_client; +GRANT SELECT ON process_definitions, process_steps TO api_client; +GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO api_client; + +-- ── Policies: tables that already enforce RLS ────────────────────────────── +-- Their existing member/open SELECT policies are unchanged. + +DROP POLICY IF EXISTS dt_projects_admin ON projects; +CREATE POLICY dt_projects_admin ON projects + FOR ALL USING (app_is_admin()) WITH CHECK (app_is_admin()); +DROP POLICY IF EXISTS dt_projects_insert ON projects; +CREATE POLICY dt_projects_insert ON projects + FOR INSERT WITH CHECK ( + app_current_user_id() IS NOT NULL + AND COALESCE(app_current_user_role() IN ('admin', 'pi'), false) + ); +DROP POLICY IF EXISTS dt_projects_pi_delete ON projects; +CREATE POLICY dt_projects_pi_delete ON projects + FOR DELETE USING (app_is_project_pi(id)); +DROP POLICY IF EXISTS dt_projects_worker ON projects; +CREATE POLICY dt_projects_worker ON projects + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_etcher_runs_admin ON etcher_runs; +CREATE POLICY dt_etcher_runs_admin ON etcher_runs FOR SELECT USING (app_is_admin()); +DROP POLICY IF EXISTS dt_etcher_runs_worker ON etcher_runs; +CREATE POLICY dt_etcher_runs_worker ON etcher_runs + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_etcher_run_files_admin ON etcher_run_files; +CREATE POLICY dt_etcher_run_files_admin ON etcher_run_files FOR SELECT USING (app_is_admin()); +DROP POLICY IF EXISTS dt_etcher_run_files_worker ON etcher_run_files; +CREATE POLICY dt_etcher_run_files_worker ON etcher_run_files + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_equipment_runs_admin ON equipment_runs; +CREATE POLICY dt_equipment_runs_admin ON equipment_runs FOR SELECT USING (app_is_admin()); +DROP POLICY IF EXISTS dt_equipment_runs_worker ON equipment_runs; +CREATE POLICY dt_equipment_runs_worker ON equipment_runs + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_trace_samples_admin ON equipment_run_trace_samples; +CREATE POLICY dt_trace_samples_admin ON equipment_run_trace_samples FOR SELECT USING (app_is_admin()); +DROP POLICY IF EXISTS dt_trace_samples_worker ON equipment_run_trace_samples; +CREATE POLICY dt_trace_samples_worker ON equipment_run_trace_samples + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_trace_events_admin ON equipment_run_trace_events; +CREATE POLICY dt_trace_events_admin ON equipment_run_trace_events FOR SELECT USING (app_is_admin()); +DROP POLICY IF EXISTS dt_trace_events_worker ON equipment_run_trace_events; +CREATE POLICY dt_trace_events_worker ON equipment_run_trace_events + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +-- ── Policies: project membership ─────────────────────────────────────────── +-- Enforced immediately (see the header). Anyone who can read a project can +-- read its membership, matching what api_client could read before, so +-- member counts on open projects keep working. + +DROP POLICY IF EXISTS dt_project_members_select ON project_members; +CREATE POLICY dt_project_members_select ON project_members + FOR SELECT USING ( + app_is_admin() + OR nanohub_user_id = app_current_user_id() + OR app_project_role(project_id) IS NOT NULL + ); +-- A PI adds plain members; the creator of a project, in the same +-- transaction that created it, enrols themselves as its PI. PI changes stay +-- admin-only. +DROP POLICY IF EXISTS dt_project_members_insert ON project_members; +CREATE POLICY dt_project_members_insert ON project_members + FOR INSERT WITH CHECK ( + app_is_admin() + OR (app_is_project_pi(project_id) AND role = 'member') + OR ( + nanohub_user_id = app_current_user_id() + AND role = 'pi' + AND COALESCE(app_current_user_role() IN ('admin', 'pi'), false) + AND NOT app_project_has_members(project_id) + AND app_project_created_in_current_transaction(project_id) + ) + ); +DROP POLICY IF EXISTS dt_project_members_delete ON project_members; +CREATE POLICY dt_project_members_delete ON project_members + FOR DELETE USING ( + app_is_admin() + OR (app_is_project_pi(project_id) AND role <> 'pi') + ); +DROP POLICY IF EXISTS dt_project_members_admin_update ON project_members; +CREATE POLICY dt_project_members_admin_update ON project_members + FOR UPDATE USING (app_is_admin()) WITH CHECK (app_is_admin()); +DROP POLICY IF EXISTS dt_project_members_worker ON project_members; +CREATE POLICY dt_project_members_worker ON project_members + FOR ALL TO api_worker USING (true) WITH CHECK (true); +ALTER TABLE project_members ENABLE ROW LEVEL SECURITY; + +-- ── Policies: experiment definitions ─────────────────────────────────────── + +DROP POLICY IF EXISTS dt_experiment_definitions_select ON experiment_definitions; +CREATE POLICY dt_experiment_definitions_select ON experiment_definitions + FOR SELECT USING ( + app_experiment_readable(owner_id, project_id, equipment_id, execution_request_id) + ); +DROP POLICY IF EXISTS dt_experiment_definitions_insert ON experiment_definitions; +CREATE POLICY dt_experiment_definitions_insert ON experiment_definitions + FOR INSERT WITH CHECK ( + app_is_admin() + OR ( + app_current_user_id() IS NOT NULL + AND owner_id = app_current_user_id() + AND (project_id IS NULL OR app_can_write_project(project_id)) + AND (COALESCE(type_id, '') = '' OR app_experiment_type_readable_by_id(type_id)) + ) + ); +DROP POLICY IF EXISTS dt_experiment_definitions_update ON experiment_definitions; +CREATE POLICY dt_experiment_definitions_update ON experiment_definitions + FOR UPDATE + USING (app_experiment_writable(owner_id, project_id, equipment_id, execution_request_id)) + WITH CHECK (app_experiment_writable(owner_id, project_id, equipment_id, execution_request_id)); +DROP POLICY IF EXISTS dt_experiment_definitions_delete ON experiment_definitions; +CREATE POLICY dt_experiment_definitions_delete ON experiment_definitions + FOR DELETE USING ( + app_is_admin() + OR owner_id = app_current_user_id() + OR (project_id IS NOT NULL AND app_can_write_project(project_id)) + ); +DROP POLICY IF EXISTS dt_experiment_definitions_worker ON experiment_definitions; +CREATE POLICY dt_experiment_definitions_worker ON experiment_definitions + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +-- ── Policies: experiment templates and versions ──────────────────────────── + +DROP POLICY IF EXISTS dt_experiment_types_select ON experiment_types; +CREATE POLICY dt_experiment_types_select ON experiment_types + FOR SELECT USING (app_experiment_type_readable(id, owner_id, owner_org, visibility)); +DROP POLICY IF EXISTS dt_experiment_types_insert ON experiment_types; +CREATE POLICY dt_experiment_types_insert ON experiment_types + FOR INSERT WITH CHECK ( + app_is_admin() + OR (app_current_user_id() IS NOT NULL AND owner_id = app_current_user_id()) + ); +DROP POLICY IF EXISTS dt_experiment_types_update ON experiment_types; +CREATE POLICY dt_experiment_types_update ON experiment_types + FOR UPDATE + USING (app_is_admin() OR owner_id = app_current_user_id()) + WITH CHECK (app_is_admin() OR owner_id = app_current_user_id()); +DROP POLICY IF EXISTS dt_experiment_types_delete ON experiment_types; +CREATE POLICY dt_experiment_types_delete ON experiment_types + FOR DELETE USING (app_is_admin() OR owner_id = app_current_user_id()); +DROP POLICY IF EXISTS dt_experiment_types_worker ON experiment_types; +CREATE POLICY dt_experiment_types_worker ON experiment_types + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_experiment_type_versions_select ON experiment_type_versions; +CREATE POLICY dt_experiment_type_versions_select ON experiment_type_versions + FOR SELECT USING (app_experiment_type_readable_by_id(type_id)); +DROP POLICY IF EXISTS dt_experiment_type_versions_write ON experiment_type_versions; +CREATE POLICY dt_experiment_type_versions_write ON experiment_type_versions + FOR ALL + USING (app_experiment_type_editable_by_id(type_id)) + WITH CHECK (app_experiment_type_editable_by_id(type_id)); +DROP POLICY IF EXISTS dt_experiment_type_versions_worker ON experiment_type_versions; +CREATE POLICY dt_experiment_type_versions_worker ON experiment_type_versions + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +-- ── Policies: project workflow attachments ───────────────────────────────── + +DROP POLICY IF EXISTS dt_project_experiments_select ON project_experiments; +CREATE POLICY dt_project_experiments_select ON project_experiments + FOR SELECT USING (app_project_role(project_id) IS NOT NULL); +-- Attaching a template requires being able to read it, so a private +-- template cannot be pulled into a project by id. +DROP POLICY IF EXISTS dt_project_experiments_write ON project_experiments; +CREATE POLICY dt_project_experiments_write ON project_experiments + FOR ALL + USING (app_can_write_project(project_id)) + WITH CHECK ( + app_can_write_project(project_id) + AND app_experiment_type_readable_by_id(type_id) + ); +DROP POLICY IF EXISTS dt_project_experiments_worker ON project_experiments; +CREATE POLICY dt_project_experiments_worker ON project_experiments + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +-- ── Policies: samples and sample links ───────────────────────────────────── + +DROP POLICY IF EXISTS dt_samples_select ON samples; +CREATE POLICY dt_samples_select ON samples + FOR SELECT USING (app_sample_readable(id, project_id)); +DROP POLICY IF EXISTS dt_samples_write ON samples; +CREATE POLICY dt_samples_write ON samples + FOR ALL + USING (app_can_write_project(project_id)) + WITH CHECK (app_can_write_project(project_id)); +DROP POLICY IF EXISTS dt_samples_worker ON samples; +CREATE POLICY dt_samples_worker ON samples + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_experiment_samples_select ON experiment_samples; +CREATE POLICY dt_experiment_samples_select ON experiment_samples + FOR SELECT USING (app_experiment_readable_by_id(experiment_id)); +DROP POLICY IF EXISTS dt_experiment_samples_write ON experiment_samples; +CREATE POLICY dt_experiment_samples_write ON experiment_samples + FOR ALL + USING (app_experiment_writable_by_id(experiment_id)) + WITH CHECK ( + app_experiment_writable_by_id(experiment_id) + AND app_sample_readable_by_id(sample_id) + ); +DROP POLICY IF EXISTS dt_experiment_samples_worker ON experiment_samples; +CREATE POLICY dt_experiment_samples_worker ON experiment_samples + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_run_samples_select ON run_samples; +CREATE POLICY dt_run_samples_select ON run_samples + FOR SELECT USING (app_sample_readable_by_id(sample_id)); +DROP POLICY IF EXISTS dt_run_samples_write ON run_samples; +CREATE POLICY dt_run_samples_write ON run_samples + FOR ALL + USING (app_sample_writable_by_id(sample_id)) + WITH CHECK (app_sample_writable_by_id(sample_id)); +DROP POLICY IF EXISTS dt_run_samples_worker ON run_samples; +CREATE POLICY dt_run_samples_worker ON run_samples + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +-- ── Policies: optimizer recipe batches and proposals ─────────────────────── +-- Batches and proposals are written by the proposal worker; users read them +-- and update proposal status. + +DROP POLICY IF EXISTS dt_recipe_batches_select ON experiment_recipe_batches; +CREATE POLICY dt_recipe_batches_select ON experiment_recipe_batches + FOR SELECT USING (app_experiment_readable_by_id(experiment_id)); +DROP POLICY IF EXISTS dt_recipe_batches_worker ON experiment_recipe_batches; +CREATE POLICY dt_recipe_batches_worker ON experiment_recipe_batches + FOR ALL TO api_worker USING (true) WITH CHECK (true); + +DROP POLICY IF EXISTS dt_recipe_proposals_select ON experiment_recipe_proposals; +CREATE POLICY dt_recipe_proposals_select ON experiment_recipe_proposals + FOR SELECT USING (app_experiment_readable_by_id(experiment_id)); +DROP POLICY IF EXISTS dt_recipe_proposals_update ON experiment_recipe_proposals; +CREATE POLICY dt_recipe_proposals_update ON experiment_recipe_proposals + FOR UPDATE + USING (app_can_update_recipe_proposals(experiment_id)) + WITH CHECK (app_can_update_recipe_proposals(experiment_id)); +DROP POLICY IF EXISTS dt_recipe_proposals_worker ON experiment_recipe_proposals; +CREATE POLICY dt_recipe_proposals_worker ON experiment_recipe_proposals + FOR ALL TO api_worker USING (true) WITH CHECK (true); diff --git a/api/scripts/init_db.sql b/api/scripts/init_db.sql index 5151a2a..3ab0d28 100644 --- a/api/scripts/init_db.sql +++ b/api/scripts/init_db.sql @@ -12,16 +12,15 @@ CREATE EXTENSION IF NOT EXISTS "uuid-ossp"; CREATE EXTENSION IF NOT EXISTS pgcrypto; --- Create a purely unprivileged application token who is subject to RLS -DO $$ -BEGIN - IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'api_client') THEN - EXECUTE format('ALTER ROLE api_client WITH PASSWORD %L', :'api_client_password'); - ELSE - EXECUTE format('CREATE USER api_client WITH PASSWORD %L', :'api_client_password'); - END IF; -END -$$; +-- Create a purely unprivileged application token who is subject to RLS. +-- psql does not substitute variables inside dollar-quoted DO blocks, so the +-- statements are generated with format() and run with \gexec. +SELECT format('CREATE USER api_client WITH PASSWORD %L', :'api_client_password') +WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'api_client') +\gexec +SELECT format('ALTER ROLE api_client WITH PASSWORD %L', :'api_client_password') +WHERE EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'api_client') +\gexec -- Table: projects CREATE TABLE IF NOT EXISTS projects ( @@ -115,6 +114,7 @@ GRANT SELECT ON project_members TO api_client; -- before executing SELECT queries: `SELECT set_config('app.current_user', 'ngholiza', true);` -- Projects Policy: User can view if project is open, OR user is a member +DROP POLICY IF EXISTS project_visibility_policy ON projects; CREATE POLICY project_visibility_policy ON projects FOR SELECT USING ( @@ -128,6 +128,7 @@ CREATE POLICY project_visibility_policy ON projects ); -- Etcher Runs Policy: User can view a run if they have access to its parent project +DROP POLICY IF EXISTS run_visibility_policy ON etcher_runs; CREATE POLICY run_visibility_policy ON etcher_runs FOR SELECT USING ( @@ -146,6 +147,7 @@ CREATE POLICY run_visibility_policy ON etcher_runs ) ); +DROP POLICY IF EXISTS run_file_visibility_policy ON etcher_run_files; CREATE POLICY run_file_visibility_policy ON etcher_run_files FOR SELECT USING ( @@ -321,6 +323,11 @@ CREATE TABLE IF NOT EXISTS experiment_definitions ( proposal_generation_started_at TIMESTAMP WITH TIME ZONE, proposal_generation_completed_at TIMESTAMP WITH TIME ZONE, proposal_generation_updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + proposal_generation_attempts INTEGER NOT NULL DEFAULT 0, + proposal_generation_claim_id UUID, + proposal_generation_lease_expires_at TIMESTAMP WITH TIME ZONE, + client_request_id VARCHAR(200) NOT NULL DEFAULT '', + client_request_fingerprint VARCHAR(64) NOT NULL DEFAULT '', execution_status VARCHAR(50) NOT NULL DEFAULT 'requested', execution_status_note TEXT DEFAULT '', execution_backend VARCHAR(100) NOT NULL DEFAULT 'physical', @@ -336,6 +343,14 @@ CREATE TABLE IF NOT EXISTS experiment_definitions ( created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP ); +CREATE UNIQUE INDEX IF NOT EXISTS uq_experiment_definitions_client_request +ON experiment_definitions (owner_id, client_request_id) +WHERE client_request_id <> ''; + +CREATE INDEX IF NOT EXISTS idx_experiment_definitions_open_proposal_jobs +ON experiment_definitions (proposal_generation_updated_at) +WHERE proposal_generation_status IN ('queued', 'generating'); + CREATE TABLE IF NOT EXISTS samples ( id VARCHAR(80) PRIMARY KEY, project_id VARCHAR(50) NOT NULL REFERENCES projects(id) ON DELETE CASCADE, diff --git a/api/tests/E2E_README.md b/api/tests/E2E_README.md index c65c9ca..3d62d07 100644 --- a/api/tests/E2E_README.md +++ b/api/tests/E2E_README.md @@ -109,6 +109,11 @@ PG_SUPER_PASS=... \ python tests/test_e2e_admin_flow.py ``` +Run it directly as above. Under pytest, `api/tests/conftest.py` points +unconfigured Postgres connections at a dead port so unit tests can never reach +a port-forwarded production database; set `DT_ALLOW_LIVE_DB_TESTS=1` if you +deliberately run a live suite through pytest. + Output is a printed phase-by-phase log ending with a pass/fail summary plus a block of "fixture breadcrumbs" (domain ids, project ids, run ids) so you can look them up manually in the UI. diff --git a/api/tests/conftest.py b/api/tests/conftest.py new file mode 100644 index 0000000..6ac8a15 --- /dev/null +++ b/api/tests/conftest.py @@ -0,0 +1,20 @@ +""" +Keep unit tests away from real databases. + +The API defaults to ``localhost:5432``, which on a developer machine is often +a ``kubectl port-forward`` to the production database. Some unit tests reach +code paths that open a connection without mocking it (they tolerate the +failure), so without this guard they would read from and write to whatever +database is listening there. + +Unless ``DT_ALLOW_LIVE_DB_TESTS=1`` is set (the live end-to-end suites need +it; see E2E_README.md), point Postgres connections at a port nothing listens +on, even if the shell already sets PG_HOST/PG_PORT, so they fail immediately. +The disposable-database integration test (test_row_level_security_pg.py) +applies its own separately validated settings afterwards. +""" +import os + +if os.getenv("DT_ALLOW_LIVE_DB_TESTS", "").strip() != "1": + os.environ["PG_HOST"] = "127.0.0.1" + os.environ["PG_PORT"] = "1" diff --git a/api/tests/test_db_connection_model.py b/api/tests/test_db_connection_model.py new file mode 100644 index 0000000..f673e6e --- /dev/null +++ b/api/tests/test_db_connection_model.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import os +import sys +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +API_ROOT = Path(__file__).resolve().parents[1] +if str(API_ROOT) not in sys.path: + sys.path.insert(0, str(API_ROOT)) + +os.environ.setdefault("DT_SYSTEM_TOKEN", "test-system-token") +os.environ.setdefault("INGESTION_TOKEN", "test-ingestion-token") +os.environ.setdefault("PG_PASS", "test-api-password") +os.environ.setdefault("PG_SUPER_PASS", "test-super-password") + +from fastapi import FastAPI # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +import data_loader_pg # noqa: E402 +import db_migrations # noqa: E402 +import runtime_config # noqa: E402 +from routers import equipment # noqa: E402 +from security import PlatformUser # noqa: E402 + +PI = PlatformUser(id="pi-1", email="pi@example.test", name="PI", role="pi", organization="LabA") + + +class SessionIdentityTests(unittest.TestCase): + def test_platform_user_carries_id_role_and_org(self): + self.assertEqual(data_loader_pg._session_identity(PI, False), ("pi-1", "pi", "LabA")) + + def test_bare_user_id_is_admin_only_when_flagged(self): + self.assertEqual(data_loader_pg._session_identity("u-1", False), ("u-1", "", "")) + self.assertEqual(data_loader_pg._session_identity("u-1", True), ("u-1", "admin", "")) + + def test_no_identity_sets_empty_settings(self): + self.assertEqual(data_loader_pg._session_identity(None, False), ("", "", "")) + + def test_connection_sets_session_scoped_settings_and_commits(self): + connection = MagicMock() + cursor = connection.cursor.return_value.__enter__.return_value + with patch.object(data_loader_pg, "is_mock_db_enabled", return_value=False), patch.object( + data_loader_pg.psycopg2, "connect", return_value=connection + ): + result = data_loader_pg.get_pg_connection(PI) + + self.assertIs(result, connection) + sql, params = cursor.execute.call_args.args + self.assertIn("set_config(%s, %s, false)", sql) + self.assertEqual( + params, + ( + "app.current_user", "pi-1", + "app.current_user_role", "pi", + "app.current_user_org", "LabA", + ), + ) + connection.commit.assert_called_once() + + def test_connection_is_closed_if_identity_cannot_be_set(self): + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value.execute.side_effect = RuntimeError("boom") + with patch.object(data_loader_pg, "is_mock_db_enabled", return_value=False), patch.object( + data_loader_pg.psycopg2, "connect", return_value=connection + ): + with self.assertRaises(RuntimeError): + data_loader_pg.get_pg_connection(PI) + connection.close.assert_called_once() + + +class ConnectionRoutingTests(unittest.TestCase): + def test_admin_without_user_is_a_system_job(self): + with patch.object(data_loader_pg, "get_pg_worker_connection", return_value="worker") as worker, patch.object( + data_loader_pg, "get_pg_connection", return_value="user" + ) as user_conn: + self.assertEqual(data_loader_pg._scoped_connection(None, True), "worker") + self.assertEqual(data_loader_pg._scoped_connection("admin-1", True), "user") + self.assertEqual(data_loader_pg._scoped_connection("pi-1", False), "user") + worker.assert_called_once() + user_conn.assert_any_call("admin-1", is_admin=True) + user_conn.assert_any_call("pi-1", is_admin=False) + + def test_worker_uses_worker_role_when_configured(self): + with patch.dict(os.environ, {"PG_WORKER_PASS": "worker-pw", "PG_WORKER_USER": "api_worker"}): + kwargs = data_loader_pg._get_pg_worker_connection_kwargs() + self.assertEqual(kwargs["user"], "api_worker") + self.assertEqual(kwargs["password"], "worker-pw") + + def test_worker_falls_back_to_migration_credentials(self): + env = {k: v for k, v in os.environ.items() if k != "PG_WORKER_PASS"} + env.update({"PG_SUPER_USER": "postgres", "PG_SUPER_PASS": "owner-pw"}) + with patch.dict(os.environ, env, clear=True): + kwargs = data_loader_pg._get_pg_worker_connection_kwargs() + self.assertEqual(kwargs["user"], "postgres") + self.assertEqual(kwargs["password"], "owner-pw") + self.assertEqual(kwargs["application_name"], "dt-api-worker") + + +class RuntimeConfigurationTests(unittest.TestCase): + BASE = {"DT_SYSTEM_TOKEN": "t", "INGESTION_TOKEN": "i", "PG_PASS": "p"} + + def test_worker_password_alone_is_enough(self): + with patch.dict(os.environ, {**self.BASE, "PG_WORKER_PASS": "w"}, clear=True): + runtime_config.validate_runtime_configuration() + + def test_legacy_superuser_password_alone_is_enough(self): + with patch.dict(os.environ, {**self.BASE, "PG_SUPER_PASS": "s"}, clear=True): + runtime_config.validate_runtime_configuration() + + def test_some_worker_credential_is_required(self): + with patch.dict(os.environ, dict(self.BASE), clear=True): + with self.assertRaises(runtime_config.MissingRuntimeConfiguration): + runtime_config.validate_runtime_configuration() + + +class MigrationSwitchTests(unittest.TestCase): + def test_enforcement_mode_values(self): + for value, expected in (("enforce", "enforce"), ("DISABLE", "disable"), ("", ""), ("on", "")): + with patch.dict(os.environ, {"DB_RLS_ENFORCEMENT": value}): + self.assertEqual(db_migrations.rls_enforcement_mode(), expected) + + def test_startup_migrations_follow_credentials_unless_overridden(self): + with patch.dict(os.environ, {"PG_SUPER_PASS": "s", "DB_MIGRATE_ON_STARTUP": ""}): + self.assertTrue(db_migrations.should_run_startup_migrations()) + with patch.dict(os.environ, {"PG_SUPER_PASS": "s", "DB_MIGRATE_ON_STARTUP": "false"}): + self.assertFalse(db_migrations.should_run_startup_migrations()) + env = {k: v for k, v in os.environ.items() if k not in {"PG_SUPER_PASS", "DB_MIGRATE_ON_STARTUP"}} + with patch.dict(os.environ, env, clear=True): + self.assertFalse(db_migrations.should_run_startup_migrations()) + + def test_enforcement_only_alters_tables_whose_state_changes(self): + cursor = MagicMock() + cursor.fetchall.return_value = [ + (name, name == "samples") for name in db_migrations.RLS_ENFORCED_TABLES + ] + state = db_migrations._apply_rls_enforcement(cursor, "enforce") + altered = [c.args[0] for c in cursor.execute.call_args_list if c.args[0].startswith("ALTER TABLE")] + self.assertEqual(len(altered), len(db_migrations.RLS_ENFORCED_TABLES) - 1) + self.assertFalse(any('"samples"' in sql for sql in altered)) + self.assertTrue(all(state.values())) + + cursor.reset_mock() + cursor.fetchall.return_value = [(name, True) for name in db_migrations.RLS_ENFORCED_TABLES] + db_migrations._apply_rls_enforcement(cursor, "") + self.assertFalse(any(c.args[0].startswith("ALTER TABLE") for c in cursor.execute.call_args_list)) + + +class ExperimentEndpointJobTests(unittest.TestCase): + HEADERS = { + "X-System-Token": os.environ["DT_SYSTEM_TOKEN"], + "X-User-Id": "pi-1", + "X-User-Email": "pi@example.test", + "X-User-Role": "pi", + } + + def setUp(self): + app = FastAPI() + app.include_router(equipment.router, prefix="/api") + self.client = TestClient(app) + + def _post(self, body, create_result): + with patch.object( + equipment, + "get_equipment_pg", + return_value={"domain_id": "etcher_01", "status": "approved", "owner_id": "pi-1", "owner_org": ""}, + ), patch.object( + equipment, "create_experiment_definition_pg", return_value=create_result + ) as create, patch.object(equipment, "start_experiment_proposal_job") as start_job: + response = self.client.post("/api/equipment/experiment", headers=self.HEADERS, json=body) + return response, create, start_job + + def test_queued_job_is_started_and_payload_carries_job_intent(self): + response, create, start_job = self._post( + { + "experiment_name": "run", + "selected_equipment": "etcher_01", + "type_id": "type-1", + "optimization_context": {"mode": "optimization", "primary_target": "AvgEtchRate"}, + "client_request_id": "req-1", + }, + {"status": "created", "experiment_id": "exp-1", "proposal_generation_status": "queued"}, + ) + self.assertEqual(response.status_code, 200, response.text) + payload = create.call_args.kwargs["payload"] + self.assertTrue(payload["generate_proposals"]) + self.assertEqual(payload["client_request_id"], "req-1") + start_job.assert_called_once_with("exp-1") + self.assertEqual(response.json()["proposal_generation_status"], "queued") + + def test_replayed_request_reports_existing_status_without_starting_a_job(self): + response, _create, start_job = self._post( + { + "experiment_name": "run", + "selected_equipment": "etcher_01", + "type_id": "type-1", + "optimization_context": {"mode": "optimization", "primary_target": "AvgEtchRate"}, + "client_request_id": "req-1", + }, + { + "status": "created", + "experiment_id": "exp-1", + "idempotent_replay": True, + "proposal_generation_status": "completed", + }, + ) + self.assertEqual(response.status_code, 200, response.text) + start_job.assert_not_called() + self.assertEqual(response.json()["proposal_generation_status"], "completed") + self.assertTrue(response.json()["idempotent_replay"]) + + +class HealthReadinessTests(unittest.TestCase): + def test_superuser_password_is_not_required_for_readiness(self): + import main + + env = { + "DT_SYSTEM_TOKEN": "t", + "INGESTION_TOKEN": "i", + "PG_PASS": "p", + "PG_WORKER_PASS": "w", + } + with patch.dict(os.environ, env, clear=True): + checks = main._runtime_dependency_checks() + informational = main._runtime_informational_checks() + self.assertTrue(all(checks.values())) + self.assertFalse(informational["postgres_superuser_password_configured"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/test_glance_closed_loop.py b/api/tests/test_glance_closed_loop.py index e8b72f1..d76b223 100644 --- a/api/tests/test_glance_closed_loop.py +++ b/api/tests/test_glance_closed_loop.py @@ -392,7 +392,7 @@ def close(self): def test_database_reconciliation_and_replay_are_idempotent(self): first_connection, first_statements = self._connection(linked=False) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=first_connection, ): first = reconcile_glance_trace_runs_pg( @@ -412,7 +412,7 @@ def test_database_reconciliation_and_replay_are_idempotent(self): replay_connection, _ = self._connection(linked=True) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=replay_connection, ): replay = reconcile_glance_trace_runs_pg( @@ -430,7 +430,7 @@ def test_replay_preserves_human_rejection(self): status="rejected", ) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = reconcile_glance_trace_runs_pg( @@ -563,7 +563,7 @@ def close(self): pass with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=Connection(), ): result = reconcile_glance_trace_runs_pg( @@ -590,7 +590,7 @@ def test_summary_first_completed_proposal_moves_to_exact_trace_followup( legacy_run_id=161, ) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = reconcile_glance_trace_runs_pg( @@ -654,7 +654,7 @@ def close(self): connection = Connection() with ( patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ), self.assertRaisesRegex( @@ -690,7 +690,7 @@ def test_eligible_run_corrected_to_outlier_forces_dataset_refresh(self): "is_outlier": True, } with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = reconcile_glance_trace_runs_pg( @@ -723,7 +723,7 @@ def test_failed_attempt_then_source_correction_forces_refresh(self): **correction, } with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = reconcile_glance_trace_runs_pg( @@ -769,7 +769,7 @@ def test_failed_dataset_correction_retries_same_revision(self): attempted_followup_kind="dataset_change", ) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = reconcile_glance_trace_runs_pg( @@ -791,7 +791,7 @@ def test_first_outlier_result_wakes_optimizer_without_retrain_credit(self): "is_outlier": True, } with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): reconciliation = reconcile_glance_trace_runs_pg( @@ -863,7 +863,7 @@ def test_excluded_run_restored_to_eligible_forces_snapshot_refresh(self): completed_followup_revision=excluded_revision, ) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): restored = reconcile_glance_trace_runs_pg( @@ -923,7 +923,7 @@ def read_sql(query, connection, params=None): with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch( @@ -957,7 +957,7 @@ def close(self): with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch( @@ -1003,7 +1003,7 @@ def read_sql(query, connection, params=None): experiment_id = "22222222-2222-4222-8222-222222222222" with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch( @@ -1151,7 +1151,7 @@ def close(self): pass with patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ): records = get_runs_list_pg( @@ -1229,7 +1229,7 @@ def close(self): pass with patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ): record = get_run_detail_pg( @@ -1304,7 +1304,7 @@ def close(self): connection = Connection() with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=connection, ): result = enrich_glance_trace_outcomes_pg( @@ -1365,7 +1365,7 @@ def close(self): completed_followup_revision=previous_revision, ) with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=reconcile_connection, ): correction = reconcile_glance_trace_runs_pg( @@ -1451,7 +1451,7 @@ def close(self): pass with patch( - "glance_closed_loop.get_pg_superuser_connection", + "glance_closed_loop.get_pg_worker_connection", return_value=Connection(), ): result = enrich_glance_trace_outcomes_pg( @@ -1526,7 +1526,7 @@ def close(self): return_value=frame, ), patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch("ai_readiness.record_dataset_snapshot"), @@ -1574,7 +1574,7 @@ def close(self): pass with patch( - "metadata_pg.get_pg_connection", + "metadata_pg.get_pg_worker_connection", return_value=Connection(), ): inserted = save_experiment_proposal_batches_pg( @@ -1641,7 +1641,7 @@ def close(self): pass with patch( - "metadata_pg.get_pg_superuser_connection", + "metadata_pg.get_pg_worker_connection", return_value=Connection(), ): result = auto_generate_recipe_proposals_for_projects_pg( @@ -1723,7 +1723,7 @@ def close(self): } with ( patch( - "metadata_pg.get_pg_superuser_connection", + "metadata_pg.get_pg_worker_connection", return_value=Connection(), ), patch( @@ -1811,7 +1811,7 @@ def close(self): } with ( patch( - "metadata_pg.get_pg_superuser_connection", + "metadata_pg.get_pg_worker_connection", side_effect=lambda: Connection(), ), patch( @@ -2124,7 +2124,7 @@ def close(self): pass with patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", side_effect=lambda: Connection(), ): first = credit_runs_counter_once( diff --git a/api/tests/test_glance_db3_trace_import.py b/api/tests/test_glance_db3_trace_import.py index a7ea722..25d2ff0 100644 --- a/api/tests/test_glance_db3_trace_import.py +++ b/api/tests/test_glance_db3_trace_import.py @@ -673,7 +673,7 @@ def capture_execute_values(cursor, query, values, **kwargs): with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=connection, ), patch( diff --git a/api/tests/test_phase3_security.py b/api/tests/test_phase3_security.py index 6b8c855..d6e30b5 100644 --- a/api/tests/test_phase3_security.py +++ b/api/tests/test_phase3_security.py @@ -123,7 +123,7 @@ def test_role_sync_creates_new_user_with_system_token(self): with patch.object( admin, - "get_pg_superuser_connection", + "get_pg_worker_connection", return_value=connection, ), patch.object( admin, @@ -283,19 +283,15 @@ def test_project_creation_uses_authenticated_identity(self): self.assertEqual(create_project.call_args.kwargs["pi_name"], "Principal Investigator") self.assertEqual(create_project.call_args.kwargs["access_mode"], "shared") - def test_project_creation_requires_approved_equipment(self): + def test_project_creation_passes_creator_role_to_database_layer(self): + # Projects no longer require a registered, approved equipment (it is an + # optional label since "Align project experiment workflow"). The + # database's row-level security only lets PIs and admins insert + # projects, so the creator's role must reach create_project_pg. with patch.object( - dataset_v2, - "get_equipment_pg", - return_value={ - "domain_id": "etcher_01", - "owner_id": "owner-1", - "owner_org": "Birck", - "status": "pending", - }, - ), patch.object( dataset_v2, "create_project_pg", + return_value={"id": "proj-1", "name": "PI Project"}, ) as create_project: response = self.client.post( "/api/dataset/v2/projects", @@ -303,18 +299,13 @@ def test_project_creation_requires_approved_equipment(self): json={ "name": "PI Project", "description": "Created from the dashboard", - "equipment_id": "etcher_01", - "equipment_name": "Etcher", "access": "shared", }, ) - self.assertEqual(response.status_code, 403) - self.assertEqual( - response.json()["detail"], - "Equipment must be approved before it can be used in projects", - ) - self.assertEqual(create_project.call_count, 0) + self.assertEqual(response.status_code, 200) + self.assertEqual(create_project.call_args.kwargs["creator_role"], "pi") + self.assertEqual(create_project.call_args.kwargs["equipment_id"], "") def test_equipment_registration_uses_authenticated_owner_context(self): with patch.object( @@ -631,7 +622,9 @@ def test_experiment_creation_links_visible_matching_project(self): "project_id": "project-1", "type_id": "recipe_sweep", "planned_date": "2026-04-15", - "planned_parameters": [], + "planned_parameters": [ + {"name": "Etch_AvgPres", "planned_value": "25"}, + ], }, ) @@ -641,7 +634,9 @@ def test_experiment_creation_links_visible_matching_project(self): "project-1", ) - def test_experiment_creation_rejects_project_for_different_equipment(self): + def test_project_run_requires_planned_parameter_values(self): + # Projects may span several tools, so a project on other equipment is + # accepted; a project run still needs at least one planned value. with patch.object( equipment, "get_equipment_pg", @@ -679,10 +674,7 @@ def test_experiment_creation_rejects_project_for_different_equipment(self): ) self.assertEqual(response.status_code, 400) - self.assertEqual( - response.json()["detail"], - "Selected project is not associated with the selected equipment", - ) + self.assertIn("At least one planned parameter value is required", response.json()["detail"]) self.assertEqual(create_experiment_definition_pg.call_count, 0) def test_experiment_creation_requires_approved_equipment(self): diff --git a/api/tests/test_production_glance_ingestion.py b/api/tests/test_production_glance_ingestion.py index 8d7c1e1..5018728 100644 --- a/api/tests/test_production_glance_ingestion.py +++ b/api/tests/test_production_glance_ingestion.py @@ -235,7 +235,7 @@ def summary(tool_id: int) -> dict: with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", side_effect=lambda: Connection(), ), patch( @@ -333,7 +333,7 @@ def execute_values(cursor, query, values, **kwargs): connection = Connection() with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=connection, ), patch("psycopg2.extras.execute_values", side_effect=execute_values), @@ -552,7 +552,7 @@ def run_payload(*, revision: datetime, pressure: float, sample: float): older = datetime(2026, 7, 1, tzinfo=timezone.utc) with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch( diff --git a/api/tests/test_retrain.py b/api/tests/test_retrain.py index ee12f3c..f789a65 100644 --- a/api/tests/test_retrain.py +++ b/api/tests/test_retrain.py @@ -101,7 +101,7 @@ def close(self): with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=Connection(), ), patch.object( @@ -205,7 +205,7 @@ def close(self): with ( patch( - "data_loader_pg.get_pg_superuser_connection", + "data_loader_pg.get_pg_worker_connection", return_value=LockConnection(), ), patch.object( diff --git a/api/tests/test_row_level_security_pg.py b/api/tests/test_row_level_security_pg.py new file mode 100644 index 0000000..e014286 --- /dev/null +++ b/api/tests/test_row_level_security_pg.py @@ -0,0 +1,1065 @@ +""" +Row-level security integration tests against a real, disposable Postgres. + +These run the schema migrations with DB_RLS_ENFORCEMENT=enforce and then +exercise the API's database functions as different users. They are skipped +unless DT_RLS_TEST_PG_PORT is set, and they refuse to run against a database +that does not contain the ``rls_test_marker`` table, so they cannot touch a +real deployment. + +Setup (Postgres 15, matching production): + + docker run -d --name dt-rls-test -e POSTGRES_PASSWORD=localtest \\ + -e POSTGRES_DB=digital_twin -p 127.0.0.1:56543:5432 postgres:15 + # create api_client, load api/scripts/init_db.sql, then: + psql ... -c "CREATE TABLE rls_test_marker ()" \\ + -c "ALTER ROLE api_worker ..." # created by the migration + DT_RLS_TEST_PG_HOST=127.0.0.1 DT_RLS_TEST_PG_PORT=56543 \\ + DT_RLS_TEST_API_PASS=... DT_RLS_TEST_SUPER_PASS=... DT_RLS_TEST_WORKER_PASS=... \\ + python -m pytest api/tests/test_row_level_security_pg.py +""" +from __future__ import annotations + +import os +import sys +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + +API_ROOT = Path(__file__).resolve().parents[1] +if str(API_ROOT) not in sys.path: + sys.path.insert(0, str(API_ROOT)) + +TEST_PORT = os.getenv("DT_RLS_TEST_PG_PORT", "").strip() + +if TEST_PORT: + os.environ["PG_HOST"] = os.getenv("DT_RLS_TEST_PG_HOST", "127.0.0.1") + os.environ["PG_PORT"] = TEST_PORT + os.environ["PG_DB"] = os.getenv("DT_RLS_TEST_PG_DB", "digital_twin") + os.environ["PG_USER"] = "api_client" + os.environ["PG_PASS"] = os.environ["DT_RLS_TEST_API_PASS"] + os.environ["PG_SUPER_USER"] = os.getenv("DT_RLS_TEST_SUPER_USER", "postgres") + os.environ["PG_SUPER_PASS"] = os.environ["DT_RLS_TEST_SUPER_PASS"] + os.environ["PG_WORKER_USER"] = "api_worker" + os.environ["PG_WORKER_PASS"] = os.environ["DT_RLS_TEST_WORKER_PASS"] + os.environ["USE_MOCK_DATABASE"] = "false" + os.environ.setdefault("DT_SYSTEM_TOKEN", "rls-test-system-token") + os.environ.setdefault("INGESTION_TOKEN", "rls-test-ingestion-token") + +import psycopg2 # noqa: E402 +import psycopg2.errors # noqa: E402 + +from security import PlatformUser # noqa: E402 + + +def _user(user_id: str, role: str, org: str) -> PlatformUser: + return PlatformUser( + id=user_id, + email=f"{user_id}@example.test", + name=user_id.split("-")[0].title(), + role=role, + organization=org, + ) + + +@unittest.skipUnless(TEST_PORT, "set DT_RLS_TEST_PG_PORT to run RLS integration tests") +class RowLevelSecurityIntegrationTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + import data_loader_pg + import db_migrations + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute("SELECT to_regclass('public.rls_test_marker')") + if cur.fetchone()[0] is None: + raise RuntimeError( + "Refusing to run: the target database has no rls_test_marker table" + ) + finally: + conn.close() + + with patch.dict(os.environ, {"DB_RLS_ENFORCEMENT": "enforce"}): + db_migrations.run_schema_migrations() + + suffix = uuid.uuid4().hex[:6] + cls.suffix = suffix + cls.alice = _user(f"alice-{suffix}", "pi", "LabA") + cls.bob = _user(f"bob-{suffix}", "pi", "LabB") + cls.carol = _user(f"carol-{suffix}", "researcher", "LabC") + cls.olivia = _user(f"olivia-{suffix}", "equipment_owner", "Facility") + cls.admin = _user(f"admin-{suffix}", "admin", "Ops") + cls.equipment_id = f"etcher_rls_{suffix}" + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + for person in (cls.alice, cls.bob, cls.carol, cls.olivia, cls.admin): + cur.execute( + """ + INSERT INTO users (id, name, email, role, organization) + VALUES (%s, %s, %s, %s, %s) + """, + (person.id, person.name, person.email, person.role, person.organization), + ) + cur.execute( + """ + INSERT INTO equipment_metadata ( + domain_id, equipment_name, owner_id, owner_org, status + ) + VALUES (%s, 'RLS Etcher', %s, '', 'approved') + """, + (cls.equipment_id, cls.olivia.id), + ) + conn.commit() + finally: + conn.close() + + # ── helpers ────────────────────────────────────────────────────────── + + def _create_project(self, owner: PlatformUser, access: str = "private") -> str: + from data_loader_pg import create_project_pg + + project = create_project_pg( + name=f"RLS {owner.name} {uuid.uuid4().hex[:4]}", + description="rls test", + equipment_id=self.equipment_id, + equipment_name="RLS Etcher", + access_mode=access, + nanohub_user_id=owner.id, + pi_name=owner.name, + creator_role=owner.role, + ) + return project["id"] + + def _create_template(self, owner: PlatformUser, visibility: str = "private") -> str: + from metadata_pg import create_experiment_type_pg + + result = create_experiment_type_pg( + payload={ + "type_name": f"rls template {uuid.uuid4().hex[:4]}", + "equipment_id": self.equipment_id, + "visibility": visibility, + "type_parameters": [ + {"name": "pressure", "mode": "variable", "min_value": 1, "max_value": 100}, + ], + }, + user=owner, + ) + return result["type_id"] + + def _create_experiment( + self, + owner: PlatformUser, + project_id: str, + type_id: str, + *, + generate: bool = False, + client_request_id: str = "", + ) -> dict: + from metadata_pg import create_experiment_definition_pg + + return create_experiment_definition_pg( + payload={ + "experiment_name": "rls experiment", + "selected_equipment": self.equipment_id, + "project_id": project_id, + "type_id": type_id, + "planned_parameters": [{"name": "pressure", "planned_value": "10"}], + "optimization_context": {"mode": "optimization" if generate else "none"}, + "generate_proposals": generate, + "client_request_id": client_request_id, + }, + user=owner, + ) + + def _count_as(self, identity, sql: str, params: tuple = ()) -> int: + from data_loader_pg import get_pg_connection + + conn = get_pg_connection(identity) + try: + with conn.cursor() as cur: + cur.execute(sql, params) + return int(cur.fetchone()[0]) + finally: + conn.close() + + def _add_member(self, project_id: str, member: PlatformUser, by: PlatformUser) -> None: + from data_loader_pg import add_project_member_pg + + add_project_member_pg( + project_id=project_id, + invitee_email=member.email, + nanohub_user_id=by.id, + is_admin=by.role == "admin", + ) + + # ── projects and membership ───────────────────────────────────────── + + def test_private_project_is_visible_only_to_members_and_admins(self): + from data_loader_pg import get_projects_list_pg + + project_id = self._create_project(self.alice) + + def visible(person): + return { + row["id"] + for row in get_projects_list_pg( + nanohub_user_id=person.id, + is_admin=person.role == "admin", + ) + } + + self.assertIn(project_id, visible(self.alice)) + self.assertNotIn(project_id, visible(self.bob)) + self.assertIn(project_id, visible(self.admin)) + + def test_researcher_cannot_create_a_project(self): + from data_loader_pg import create_project_pg + + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + create_project_pg( + name="not allowed", + description="", + equipment_id=self.equipment_id, + equipment_name="RLS Etcher", + access_mode="private", + nanohub_user_id=self.carol.id, + pi_name=self.carol.name, + creator_role=self.carol.role, + ) + + def test_only_the_pi_can_add_members_and_nobody_can_self_enrol(self): + from data_loader_pg import get_pg_connection + + project_id = self._create_project(self.alice) + self._add_member(project_id, self.carol, by=self.alice) + self.assertEqual( + self._count_as( + self.carol, + "SELECT count(*) FROM project_members WHERE project_id = %s", + (project_id,), + ), + 2, + ) + + conn = get_pg_connection(self.bob) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + "INSERT INTO project_members (project_id, nanohub_user_id, role) " + "VALUES (%s, %s, 'member')", + (project_id, self.bob.id), + ) + conn.rollback() + with conn.cursor() as cur: + # A second PI cannot be enrolled on a project that has members. + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + "INSERT INTO project_members (project_id, nanohub_user_id, role) " + "VALUES (%s, %s, 'pi')", + (project_id, self.bob.id), + ) + finally: + conn.close() + + def test_pi_cannot_claim_an_existing_project_without_members(self): + import data_loader_pg + + seeded_project = f"seeded_{self.suffix}" + conn = data_loader_pg.get_pg_worker_connection() + try: + with conn.cursor() as cur: + # Ingestion seeds projects without members. + cur.execute( + "INSERT INTO projects (id, name, access_mode) VALUES (%s, 'seeded', 'private')", + (seeded_project,), + ) + conn.commit() + finally: + conn.close() + + conn = data_loader_pg.get_pg_connection(self.bob) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + "INSERT INTO project_members (project_id, nanohub_user_id, role) " + "VALUES (%s, %s, 'pi')", + (seeded_project, self.bob.id), + ) + finally: + conn.close() + + def test_member_cannot_delete_project_but_pi_can(self): + from data_loader_pg import delete_project_pg + + project_id = self._create_project(self.alice) + self._add_member(project_id, self.carol, by=self.alice) + with self.assertRaises(PermissionError): + delete_project_pg(project_id=project_id, nanohub_user_id=self.carol.id, is_admin=False) + # A non-member cannot even see the private project. + self.assertEqual( + delete_project_pg(project_id=project_id, nanohub_user_id=self.bob.id, is_admin=False)["status"], + "not_found", + ) + result = delete_project_pg(project_id=project_id, nanohub_user_id=self.alice.id, is_admin=False) + self.assertEqual(result["deleted_project_rows"], 1) + + # ── experiment definitions ────────────────────────────────────────── + + def test_experiment_visibility_follows_project_access(self): + from metadata_pg import list_experiment_definitions_pg + + project_id = self._create_project(self.alice) + self._add_member(project_id, self.carol, by=self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id)["experiment_id"] + + def listed(person): + return { + row["id"] + for row in list_experiment_definitions_pg( + user=person, + visible_project_ids=[project_id], + ) + } + + self.assertIn(experiment_id, listed(self.alice)) + self.assertIn(experiment_id, listed(self.carol)) + self.assertIn(experiment_id, listed(self.admin)) + # Even when the app passes the project id, RLS hides it from Bob. + self.assertNotIn(experiment_id, listed(self.bob)) + self.assertEqual( + self._count_as(self.bob, "SELECT count(*) FROM experiment_definitions WHERE id = %s", (experiment_id,)), + 0, + ) + # The private template stays readable to the project member through + # the experiment, so joined names do not disappear. + self.assertEqual( + self._count_as(self.carol, "SELECT count(*) FROM experiment_types WHERE id = %s", (type_id,)), + 1, + ) + self.assertEqual( + self._count_as(self.bob, "SELECT count(*) FROM experiment_types WHERE id = %s", (type_id,)), + 0, + ) + + def test_database_rejects_forged_owner_and_foreign_project(self): + from data_loader_pg import get_pg_connection + + project_id = self._create_project(self.alice) + conn = get_pg_connection(self.bob) + try: + for owner, target_project in ((self.alice.id, None), (self.bob.id, project_id)): + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + """ + INSERT INTO experiment_definitions (project_id, equipment_id, name, owner_id) + VALUES (%s, %s, 'forged', %s) + """, + (target_project, self.equipment_id, owner), + ) + conn.rollback() + finally: + conn.close() + + def test_hidden_template_cannot_be_used_by_id(self): + project_id = self._create_project(self.alice) + bobs_private_template = self._create_template(self.bob, visibility="private") + with self.assertRaisesRegex(ValueError, "Selected experiment was not found"): + self._create_experiment(self.alice, project_id, bobs_private_template) + + def test_hidden_template_cannot_be_forked(self): + from metadata_pg import fork_experiment_type_pg + + bobs_private_template = self._create_template(self.bob, visibility="private") + self.assertEqual( + fork_experiment_type_pg(type_id=bobs_private_template, user=self.alice), + {"status": "not_found"}, + ) + shared_template = self._create_template(self.bob, visibility="shared") + self.assertEqual( + fork_experiment_type_pg(type_id=shared_template, user=self.alice)["status"], + "forked", + ) + + def test_template_versions_follow_template_visibility(self): + from metadata_pg import get_experiment_type_versions_pg, update_experiment_type_pg + + type_id = self._create_template(self.alice, visibility="private") + update_experiment_type_pg(type_id=type_id, payload={"type_name": "renamed"}, user=self.alice) + self.assertEqual(len(get_experiment_type_versions_pg(type_id, user=self.alice)), 1) + self.assertEqual(get_experiment_type_versions_pg(type_id, user=self.bob), []) + + def test_repeated_client_request_id_returns_the_same_experiment(self): + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + request_id = f"req-{uuid.uuid4()}" + first = self._create_experiment(self.alice, project_id, type_id, client_request_id=request_id) + second = self._create_experiment(self.alice, project_id, type_id, client_request_id=request_id) + self.assertFalse(first["idempotent_replay"]) + self.assertTrue(second["idempotent_replay"]) + self.assertEqual(first["experiment_id"], second["experiment_id"]) + self.assertEqual( + self._count_as( + self.alice, + "SELECT count(*) FROM experiment_definitions WHERE client_request_id = %s", + (request_id,), + ), + 1, + ) + + def test_skipped_generation_is_recorded_atomically_with_the_experiment(self): + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + result = self._create_experiment(self.alice, project_id, type_id, generate=False) + self.assertEqual(result["proposal_generation_status"], "skipped") + + # ── proposal job queue ────────────────────────────────────────────── + + @staticmethod + def _fake_proposals(**_kwargs): + return { + "source": "rls-test", + "batches": [ + {"iteration": 1, "model": "fake", "n_train": 0, "proposals": [{"pressure": 12.0}]}, + ], + } + + def _job_state(self, experiment_id: str) -> tuple: + import data_loader_pg + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT proposal_generation_status, proposal_generation_attempts, + (SELECT count(*) FROM experiment_recipe_batches b WHERE b.experiment_id = e.id) + FROM experiment_definitions e WHERE id = %s + """, + (experiment_id,), + ) + return cur.fetchone() + finally: + conn.close() + + def test_queued_job_is_processed_and_results_follow_experiment_visibility(self): + from metadata_pg import process_experiment_proposal_jobs_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + self.assertEqual(self._job_state(experiment_id)[0], "queued") + + with patch("ml_engine.compute_proposals", side_effect=self._fake_proposals): + outcome = process_experiment_proposal_jobs_pg(experiment_id=experiment_id) + self.assertEqual(outcome["results"][0]["status"], "completed") + self.assertEqual(self._job_state(experiment_id), ("completed", 1, 1)) + self.assertEqual( + self._count_as(self.alice, "SELECT count(*) FROM experiment_recipe_proposals WHERE experiment_id = %s", (experiment_id,)), + 1, + ) + self.assertEqual( + self._count_as(self.bob, "SELECT count(*) FROM experiment_recipe_proposals WHERE experiment_id = %s", (experiment_id,)), + 0, + ) + + def test_expired_lease_is_retried_and_stale_worker_result_is_discarded(self): + import data_loader_pg + from metadata_pg import ( + claim_experiment_proposal_job_pg, + run_experiment_proposal_job_pg, + ) + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + + first_claim = claim_experiment_proposal_job_pg(experiment_id=experiment_id) + self.assertIsNotNone(first_claim) + # While the lease is live nobody else can claim it. + self.assertIsNone(claim_experiment_proposal_job_pg(experiment_id=experiment_id)) + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + "UPDATE experiment_definitions SET proposal_generation_lease_expires_at = now() - interval '1 minute' WHERE id = %s", + (experiment_id,), + ) + conn.commit() + finally: + conn.close() + + second_claim = claim_experiment_proposal_job_pg(experiment_id=experiment_id) + self.assertIsNotNone(second_claim) + self.assertEqual(second_claim["proposal_generation_attempts"], 2) + + with patch("ml_engine.compute_proposals", side_effect=self._fake_proposals): + stale = run_experiment_proposal_job_pg(first_claim) + fresh = run_experiment_proposal_job_pg(second_claim) + self.assertEqual(stale["status"], "claim_lost") + self.assertEqual(fresh["status"], "completed") + self.assertEqual(self._job_state(experiment_id), ("completed", 2, 1)) + + def test_job_that_keeps_dying_is_failed_after_max_attempts(self): + import data_loader_pg + from metadata_pg import claim_experiment_proposal_job_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE experiment_definitions + SET proposal_generation_status = 'generating', + proposal_generation_attempts = 3, + proposal_generation_lease_expires_at = now() - interval '1 minute' + WHERE id = %s + """, + (experiment_id,), + ) + conn.commit() + finally: + conn.close() + + with patch.dict(os.environ, {"PROPOSAL_JOB_MAX_ATTEMPTS": "3"}): + self.assertIsNone(claim_experiment_proposal_job_pg(experiment_id=experiment_id)) + self.assertEqual(self._job_state(experiment_id)[0], "failed") + + def test_compute_failure_marks_job_failed(self): + from metadata_pg import process_experiment_proposal_jobs_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + with patch("ml_engine.compute_proposals", side_effect=RuntimeError("model unavailable")): + outcome = process_experiment_proposal_jobs_pg(experiment_id=experiment_id) + self.assertEqual(outcome["results"][0]["status"], "failed") + self.assertEqual(self._job_state(experiment_id)[0], "failed") + + def test_recipe_proposal_status_updates_follow_project_membership(self): + from metadata_pg import process_experiment_proposal_jobs_pg, update_recipe_proposal_status_pg + import data_loader_pg + + project_id = self._create_project(self.alice) + self._add_member(project_id, self.carol, by=self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + with patch("ml_engine.compute_proposals", side_effect=self._fake_proposals): + process_experiment_proposal_jobs_pg(experiment_id=experiment_id) + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + "SELECT id::text FROM experiment_recipe_proposals WHERE experiment_id = %s", + (experiment_id,), + ) + proposal_id = cur.fetchone()[0] + finally: + conn.close() + + self.assertEqual( + update_recipe_proposal_status_pg( + proposal_id=proposal_id, status="accepted", comment="", user=self.carol + )["status"], + "accepted", + ) + self.assertEqual( + update_recipe_proposal_status_pg( + proposal_id=proposal_id, status="attempted", comment="", user=self.bob + )["status"], + "not_found", + ) + + # ── samples ───────────────────────────────────────────────────────── + + def test_samples_are_scoped_to_project_access(self): + from metadata_pg import create_sample_pg, list_samples_pg, update_sample_pg + + project_id = self._create_project(self.alice) + sample = create_sample_pg(payload={"project_id": project_id, "name": "wafer"}, user=self.alice) + self.assertIn(sample["id"], {row["id"] for row in list_samples_pg(user=self.alice)}) + self.assertNotIn(sample["id"], {row["id"] for row in list_samples_pg(user=self.bob)}) + self.assertEqual( + update_sample_pg(sample_id=sample["id"], payload={"name": "stolen"}, user=self.bob), + {"status": "not_found"}, + ) + + # ── facility execution queue ──────────────────────────────────────── + + def test_facility_manager_sees_execution_request_with_project_name(self): + from data_loader_pg import get_pg_connection + from metadata_pg import list_execution_queue_pg, update_execution_queue_status_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id)["experiment_id"] + + queue = {row["id"]: row for row in list_execution_queue_pg(self.olivia)} + self.assertIn(experiment_id, queue) + self.assertTrue(queue[experiment_id]["project_name"].startswith("RLS ")) + self.assertTrue(queue[experiment_id]["type_name"]) + + updated = update_execution_queue_status_pg( + experiment_id=experiment_id, + status="approved", + note="ok", + scheduled_for="", + user=self.olivia, + ) + self.assertEqual(updated["execution_status"], "approved") + # Facility access covers execution fields only. + conn = get_pg_connection(self.olivia) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + "UPDATE experiment_definitions SET name = 'facility edit' WHERE id = %s", + (experiment_id,), + ) + finally: + conn.close() + # Facility access does not extend to the project itself. + self.assertEqual( + self._count_as(self.olivia, "SELECT count(*) FROM projects WHERE id = %s", (project_id,)), + 0, + ) + + # ── template deletion ─────────────────────────────────────────────── + + def test_template_delete_detaches_references_in_other_projects(self): + import data_loader_pg + from metadata_pg import add_project_experiment_pg, delete_experiment_type_pg + + shared_type = self._create_template(self.alice, visibility="shared") + bobs_project = self._create_project(self.bob) + bobs_type = self._create_template(self.bob) + add_project_experiment_pg(project_id=bobs_project, type_id=bobs_type, user=self.bob) + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + "UPDATE project_experiments SET upstream_type_id = %s WHERE project_id = %s", + (shared_type, bobs_project), + ) + conn.commit() + finally: + conn.close() + + # Bob can see the shared template but cannot delete it or run the detach. + with self.assertRaises(PermissionError): + delete_experiment_type_pg(type_id=shared_type, user=self.bob) + conn = data_loader_pg.get_pg_connection(self.bob) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute("SELECT * FROM app_detach_experiment_type(%s)", (shared_type,)) + finally: + conn.close() + + result = delete_experiment_type_pg(type_id=shared_type, user=self.alice) + self.assertEqual(result["cleared_project_upstreams"], 1) + + # ── other user-scoped write paths ─────────────────────────────────── + + def test_workflow_examples_and_process_instantiation_run_as_caller(self): + from metadata_pg import ( + create_etcher_profilometer_workflow_example_pg, + create_process_definition_pg, + create_rcac_closed_loop_demo_pg, + instantiate_process_in_project_pg, + list_project_experiments_pg, + ) + + project_id = self._create_project(self.alice) + example = create_etcher_profilometer_workflow_example_pg(project_id=project_id, user=self.alice) + self.assertTrue(example) + demo = create_rcac_closed_loop_demo_pg(project_id=project_id, queue="standby", user=self.alice) + self.assertTrue(demo) + # Bob re-runs the example on his own project: the shared example + # templates Alice created are reused, not overwritten or duplicated. + bobs_project = self._create_project(self.bob) + self.assertTrue(create_etcher_profilometer_workflow_example_pg(project_id=bobs_project, user=self.bob)) + + type_id = self._create_template(self.alice, visibility="shared") + process = create_process_definition_pg( + payload={"name": f"rls process {self.suffix}", "steps": [{"type_id": type_id}]}, + user=self.alice, + ) + instantiate_process_in_project_pg( + project_id=project_id, + process_id=process["id"], + user=self.alice, + ) + attached = {row["id"] for row in list_project_experiments_pg(project_id=project_id, user=self.alice)} + self.assertIn(type_id, attached) + with self.assertRaises(PermissionError): + list_project_experiments_pg(project_id=project_id, user=self.bob) + + def test_publication_requests_follow_project_access(self): + from metadata_pg import create_publication_request_pg, list_publication_requests_pg + + project_id = self._create_project(self.alice) + created = create_publication_request_pg( + project_id=project_id, + payload={"title": "RLS dataset"}, + user=self.alice, + ) + self.assertEqual( + [row["id"] for row in list_publication_requests_pg(project_id=project_id, user=self.alice)], + [created["id"]], + ) + with self.assertRaises(PermissionError): + list_publication_requests_pg(project_id=project_id, user=self.bob) + + # ── findings from review round 1 ──────────────────────────────────── + + def test_private_template_cannot_enter_through_processes(self): + from data_loader_pg import get_pg_connection + from metadata_pg import create_process_definition_pg, instantiate_process_in_project_pg + + alices_project = self._create_project(self.alice) + bobs_private = self._create_template(self.bob, visibility="private") + with self.assertRaisesRegex(ValueError, "was not found"): + create_process_definition_pg( + payload={"name": "steal", "steps": [{"type_id": bobs_private}]}, + user=self.alice, + ) + # Bob shares a process built on his private template. + bobs_process = create_process_definition_pg( + payload={"name": f"bob shared {self.suffix}", "visibility": "published", "steps": [{"type_id": bobs_private}]}, + user=self.bob, + ) + with self.assertRaises(PermissionError): + instantiate_process_in_project_pg( + project_id=alices_project, process_id=bobs_process["id"], user=self.alice + ) + # The database refuses the attachment too. + conn = get_pg_connection(self.alice) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + "INSERT INTO project_experiments (project_id, type_id, added_by) VALUES (%s, %s, %s)", + (alices_project, bobs_private, self.alice.id), + ) + finally: + conn.close() + + def test_experiment_scope_cannot_be_moved_to_another_project(self): + from data_loader_pg import delete_project_pg, get_pg_connection + + alices_project = self._create_project(self.alice) + bobs_project = self._create_project(self.bob) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, alices_project, type_id)["experiment_id"] + + conn = get_pg_connection(self.alice) + try: + for column, value in (("project_id", bobs_project), ("owner_id", self.bob.id)): + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute( + f"UPDATE experiment_definitions SET {column} = %s WHERE id = %s", + (value, experiment_id), + ) + conn.rollback() + with conn.cursor() as cur: + cur.execute( + "UPDATE experiment_definitions SET name = 'renamed' WHERE id = %s", + (experiment_id,), + ) + self.assertEqual(cur.rowcount, 1) + conn.commit() + finally: + conn.close() + + # Deleting the project still detaches the experiment (FK SET NULL). + result = delete_project_pg(project_id=alices_project, nanohub_user_id=self.alice.id, is_admin=False) + self.assertEqual(result["deleted_project_rows"], 1) + + def test_detach_function_denies_ownerless_templates(self): + import data_loader_pg + + ownerless = f"ownerless_{self.suffix}" + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + "INSERT INTO experiment_types (id, name, owner_id, visibility) VALUES (%s, 'ownerless', NULL, 'shared')", + (ownerless,), + ) + conn.commit() + finally: + conn.close() + conn = data_loader_pg.get_pg_connection(self.bob) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute("SELECT * FROM app_detach_experiment_type(%s)", (ownerless,)) + finally: + conn.close() + + def test_reused_request_id_with_different_payload_conflicts(self): + from metadata_pg import ExperimentRequestConflictError, create_experiment_definition_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + request_id = f"req-{uuid.uuid4()}" + self._create_experiment(self.alice, project_id, type_id, client_request_id=request_id) + with self.assertRaises(ExperimentRequestConflictError): + create_experiment_definition_pg( + payload={ + "experiment_name": "a different run", + "selected_equipment": self.equipment_id, + "project_id": project_id, + "type_id": type_id, + "planned_parameters": [{"name": "pressure", "planned_value": "55"}], + "client_request_id": request_id, + }, + user=self.alice, + ) + + # ── findings from review round 2 ──────────────────────────────────── + + def test_private_template_cannot_be_attached_directly(self): + from metadata_pg import add_project_experiment_pg + + project_id = self._create_project(self.alice) + bobs_private = self._create_template(self.bob, visibility="private") + self.assertEqual( + add_project_experiment_pg(project_id=project_id, type_id=bobs_private, user=self.alice), + {"status": "not_found", "type_id": bobs_private}, + ) + + def test_process_view_hides_templates_the_viewer_cannot_read(self): + from metadata_pg import create_process_definition_pg, get_process_definition_pg, list_process_definitions_pg + + bobs_private = self._create_template(self.bob, visibility="private") + process = create_process_definition_pg( + payload={"name": f"bob process {self.suffix}", "visibility": "published", "steps": [{"type_id": bobs_private}]}, + user=self.bob, + ) + viewed = get_process_definition_pg(process_id=process["id"], user=self.alice) + self.assertEqual(len(viewed["steps"]), 1) + self.assertIsNone(viewed["steps"][0]["unit_experiment"]) + listed = next(row for row in list_process_definitions_pg(user=self.alice) if row["id"] == process["id"]) + self.assertTrue(all(not step.get("unit_experiment_name") for step in listed["steps"])) + owner_view = get_process_definition_pg(process_id=process["id"], user=self.bob) + self.assertIsNotNone(owner_view["steps"][0]["unit_experiment"]) + + def test_retry_replays_even_after_the_template_changed(self): + from metadata_pg import update_experiment_type_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + request_id = f"req-{uuid.uuid4()}" + first = self._create_experiment(self.alice, project_id, type_id, client_request_id=request_id) + # The planned parameter is no longer an input of the template. + update_experiment_type_pg(type_id=type_id, payload={"type_parameters": []}, user=self.alice) + retry = self._create_experiment(self.alice, project_id, type_id, client_request_id=request_id) + self.assertTrue(retry["idempotent_replay"]) + self.assertEqual(retry["experiment_id"], first["experiment_id"]) + + # ── findings from review round 3 ──────────────────────────────────── + + def test_example_seeding_never_touches_another_users_private_template(self): + import data_loader_pg + from metadata_pg import _ensure_workflow_example_type + + example_id = f"workflow_example_rls_{self.suffix}" + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO experiment_types (id, name, owner_id, owner_org, visibility) + VALUES (%s, 'alice private example', %s, %s, 'private') + """, + (example_id, self.alice.id, self.alice.organization), + ) + conn.commit() + finally: + conn.close() + + conn = data_loader_pg.get_pg_connection(self.bob) + try: + with conn.cursor() as cur: + with self.assertRaises(PermissionError): + _ensure_workflow_example_type( + cur, + type_id=example_id, + name="bob overwrite", + description="", + equipment_id=self.equipment_id, + parameters=[], + outputs=[], + additional_inputs=[], + user=self.bob, + ) + conn.rollback() + finally: + conn.close() + + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute("SELECT name, visibility, owner_id FROM experiment_types WHERE id = %s", (example_id,)) + self.assertEqual(cur.fetchone(), ("alice private example", "private", self.alice.id)) + finally: + conn.close() + + def test_interrupted_followup_generation_is_not_replayed_as_initial_job(self): + import data_loader_pg + from metadata_pg import claim_experiment_proposal_job_pg, process_experiment_proposal_jobs_pg + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id, generate=True)["experiment_id"] + with patch("ml_engine.compute_proposals", side_effect=self._fake_proposals): + process_experiment_proposal_jobs_pg(experiment_id=experiment_id) + + # The ingestion loop starts iteration 2 and dies with the lease held. + conn = data_loader_pg.get_pg_migration_connection() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE experiment_definitions + SET proposal_generation_status = 'generating', + proposal_generation_lease_expires_at = now() - interval '1 minute' + WHERE id = %s + """, + (experiment_id,), + ) + conn.commit() + finally: + conn.close() + + self.assertIsNone(claim_experiment_proposal_job_pg(experiment_id=experiment_id)) + status, _attempts, batches = self._job_state(experiment_id) + self.assertEqual((status, batches), ("failed", 1)) + + def test_migration_ledger_is_owner_only(self): + from data_loader_pg import get_pg_worker_connection + + conn = get_pg_worker_connection() + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute("UPDATE app_schema_migrations SET checksum = 'x'") + finally: + conn.close() + + # ── connection model ──────────────────────────────────────────────── + + def test_worker_role_is_least_privileged_but_sees_all_rows(self): + from data_loader_pg import get_pg_worker_connection + + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + experiment_id = self._create_experiment(self.alice, project_id, type_id)["experiment_id"] + + conn = get_pg_worker_connection() + try: + with conn.cursor() as cur: + cur.execute("SELECT current_user, rolsuper, rolbypassrls FROM pg_roles WHERE rolname = current_user") + self.assertEqual(cur.fetchone(), ("api_worker", False, False)) + cur.execute("SELECT count(*) FROM experiment_definitions WHERE id = %s", (experiment_id,)) + self.assertEqual(cur.fetchone()[0], 1) + with self.assertRaises(psycopg2.errors.InsufficientPrivilege): + cur.execute("CREATE TABLE worker_should_not_create (id int)") + finally: + conn.close() + + def test_connection_without_identity_sees_only_published_templates(self): + from data_loader_pg import get_pg_connection + + project_id = self._create_project(self.alice) + private_type = self._create_template(self.alice, visibility="private") + published_type = self._create_template(self.alice, visibility="published") + self._create_experiment(self.alice, project_id, private_type) + + conn = get_pg_connection() + try: + with conn.cursor() as cur: + cur.execute("SELECT count(*) FROM experiment_definitions") + self.assertEqual(cur.fetchone()[0], 0) + cur.execute( + "SELECT id FROM experiment_types WHERE id = ANY(%s)", + ([private_type, published_type],), + ) + self.assertEqual([row[0] for row in cur.fetchall()], [published_type]) + finally: + conn.close() + + def test_session_identity_survives_commit_and_rollback(self): + from data_loader_pg import get_pg_connection + + project_id = self._create_project(self.alice) + conn = get_pg_connection(self.alice) + try: + with conn.cursor() as cur: + cur.execute("SELECT 1") + conn.commit() + with conn.cursor() as cur: + cur.execute("SELECT 1") + conn.rollback() + with conn.cursor() as cur: + cur.execute("SELECT count(*) FROM projects WHERE id = %s", (project_id,)) + self.assertEqual(cur.fetchone()[0], 1) + finally: + conn.close() + + # ── HTTP endpoint ─────────────────────────────────────────────────── + + def test_experiment_endpoint_creates_and_replays_under_rls(self): + from fastapi import FastAPI + from fastapi.testclient import TestClient + from routers import equipment + + app = FastAPI() + app.include_router(equipment.router, prefix="/api") + client = TestClient(app) + project_id = self._create_project(self.alice) + type_id = self._create_template(self.alice) + headers = { + "X-System-Token": os.environ["DT_SYSTEM_TOKEN"], + "X-User-Id": self.alice.id, + "X-User-Email": self.alice.email, + "X-User-Role": self.alice.role, + "X-User-Org": self.alice.organization, + } + body = { + "experiment_name": "endpoint run", + "selected_equipment": self.equipment_id, + "project_id": project_id, + "type_id": type_id, + "planned_parameters": [{"name": "pressure", "planned_value": "20"}], + "client_request_id": f"http-{uuid.uuid4()}", + } + with patch.object(equipment, "start_experiment_proposal_job") as start_job: + first = client.post("/api/equipment/experiment", headers=headers, json=body) + second = client.post("/api/equipment/experiment", headers=headers, json=body) + self.assertEqual(first.status_code, 200, first.text) + self.assertEqual(second.status_code, 200, second.text) + self.assertEqual(first.json()["experiment_id"], second.json()["experiment_id"]) + self.assertEqual(first.json()["proposal_generation_status"], "skipped") + start_job.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/api/training/retrain.py b/api/training/retrain.py index 77412a7..853bf74 100644 --- a/api/training/retrain.py +++ b/api/training/retrain.py @@ -293,8 +293,8 @@ class _LockHeld(Exception): def _with_advisory_lock(fn): """Decorator: run `fn` only if we can grab the retrain advisory lock.""" def wrapped(*args, **kwargs): - from data_loader_pg import get_pg_superuser_connection - conn = get_pg_superuser_connection() + from data_loader_pg import get_pg_worker_connection + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute("SELECT pg_try_advisory_lock(%s)", (_ADVISORY_LOCK_KEY,)) @@ -322,7 +322,7 @@ def _current_snapshot( Records a new snapshot row if the hash is new. """ from ai_readiness import record_dataset_snapshot - from data_loader_pg import get_pg_superuser_connection + from data_loader_pg import get_pg_worker_connection if training_df is None: training_df, source = _load_training_dataset() @@ -362,7 +362,7 @@ def _current_snapshot( metadata={"training_source": source or "unknown"}, ) - conn = get_pg_superuser_connection() + conn = get_pg_worker_connection() try: with conn.cursor() as cur: cur.execute( @@ -596,6 +596,19 @@ def run_retrain_safe(**kwargs) -> dict[str, Any]: "status": "pending", "error": str(exc), } + try: + from metadata_pg import process_experiment_proposal_jobs_pg + + # Backstop for the API's in-process sweeper: run any queued or + # expired initial-proposal jobs. + result["proposal_jobs"] = process_experiment_proposal_jobs_pg(limit=20) + except Exception as exc: + logger.warning( + "Could not process pending experiment proposal jobs: %s", + exc, + exc_info=True, + ) + result["proposal_jobs"] = {"status": "pending", "error": str(exc)} return result diff --git a/geddes/k8s/01-secrets.yaml.example b/geddes/k8s/01-secrets.yaml.example index f4ee16c..f8e4364 100644 --- a/geddes/k8s/01-secrets.yaml.example +++ b/geddes/k8s/01-secrets.yaml.example @@ -50,9 +50,24 @@ stringData: PG_USER: "api_client" PG_PASS: "REPLACE_WITH_API_CLIENT_PASSWORD" - # Superuser for ingestion (bypasses RLS) + # Background jobs (ingestion, optimizer, model registry) connect as the + # least-privileged api_worker role created by the row-level security + # migration. Set its password once with: + # ALTER ROLE api_worker WITH LOGIN PASSWORD ''; + PG_WORKER_USER: "api_worker" + PG_WORKER_PASS: "REPLACE_WITH_API_WORKER_PASSWORD" + + # Table-owner credentials, used only for schema migrations. Once the + # dt-db-migrate Job (09-db-migrate-job.yaml) runs migrations, remove these + # two keys from dt-api-secret and set DB_MIGRATE_ON_STARTUP=false; until + # then the API also uses them for background jobs if PG_WORKER_PASS is unset. PG_SUPER_USER: "postgres" PG_SUPER_PASS: "REPLACE_WITH_POSTGRES_PASSWORD" + + # Row-level security for experiments and project membership: + # "enforce" turns it on, "disable" is the rollback switch, unset leaves the + # current state. See api/ROW_LEVEL_SECURITY.md. + # DB_RLS_ENFORCEMENT: "enforce" # ── Ingestion Token ── # Shared secret between Azure Function and Geddes API @@ -70,3 +85,19 @@ stringData: # ── Legacy CSV path (fallback for ML training) ── DATASET_PATH: "/app/data/full_dataset.csv" + +--- +# Credentials for the schema migration Job only (09-db-migrate-job.yaml). +apiVersion: v1 +kind: Secret +metadata: + name: dt-db-migration-secret + namespace: ncn-digitaltwins-zchen +type: Opaque +stringData: + PG_HOST: "dt-db-v2-srv.ncn-digitaltwins-zchen.svc.cluster.local" + PG_PORT: "5432" + PG_DB: "digital_twin" + PG_SUPER_USER: "postgres" + PG_SUPER_PASS: "REPLACE_WITH_POSTGRES_PASSWORD" + # DB_RLS_ENFORCEMENT: "enforce" diff --git a/geddes/k8s/02-api.yaml b/geddes/k8s/02-api.yaml index 66b45aa..8c41777 100644 --- a/geddes/k8s/02-api.yaml +++ b/geddes/k8s/02-api.yaml @@ -1,3 +1,24 @@ +# dt-api: base platform Deployment and Service. +# +# The live Deployment also carries the temporary DB3 pilot layer (the +# dt-db3-pilot-config ConfigMap in envFrom, the dt-db3-pilot-uploads volume and +# a config checksum annotation), added by geddes/k8s/pilots/db3/deploy.sh. The +# pilot is intentionally kept out of this file (see pilots/db3/README.md), so +# `kubectl diff -f` shows exactly that layer and nothing else. +# +# Do not `kubectl apply -f` this file while the pilot is live: it would remove +# the pilot's uploads volume and config, and pilots/db3/deploy.sh cannot put +# them back for a non-pilot image. Change the image or settings with +# `kubectl set image` / `kubectl set env` and record every change here. If the +# pilot layer is ever removed by mistake, restore just that layer (keeping the +# current image) with: +# kubectl -n ncn-digitaltwins-zchen patch deployment dt-api --type json -p '[ +# {"op":"add","path":"/spec/template/spec/containers/0/envFrom/-", +# "value":{"configMapRef":{"name":"dt-db3-pilot-config"}}}, +# {"op":"add","path":"/spec/template/spec/containers/0/volumeMounts", +# "value":[{"name":"dt-db3-pilot-uploads","mountPath":"/var/lib/dt-api/uploads"}]}, +# {"op":"add","path":"/spec/template/spec/volumes", +# "value":[{"name":"dt-db3-pilot-uploads","persistentVolumeClaim":{"claimName":"dt-db3-pilot-uploads"}}]}]' apiVersion: apps/v1 kind: Deployment metadata: @@ -18,7 +39,7 @@ spec: containers: - name: dt-api # Notice we are using the 'sdx' namespace on the registry now - image: geddes-registry.rcac.purdue.edu/sdx/dt-api:recipe-371da73@sha256:a9cd36d051b2aada3dd5e7f843cbaac83998711dbd205f66935357e3a0bde483 + image: geddes-registry.rcac.purdue.edu/sdx/dt-api:rls-b02e9ac@sha256:9204c93bdfba615e96e7e233990f97d2d16def03762a307161300f76fcada3f2 imagePullPolicy: Always resources: requests: @@ -30,12 +51,6 @@ spec: ports: - containerPort: 8000 env: - - name: NEXTJS_PROPOSAL_ENGINE - value: "api" - - name: ML_PROPOSAL_MODEL - value: "random_forest" - - name: ML_PROPOSAL_CANDIDATES - value: "1024" - name: ML_SHADOW_EVALUATION_ENABLED value: "false" - name: GLANCE_PRODUCTION_MAPPINGS_JSON @@ -44,6 +59,22 @@ spec: name: dt-api-secret key: GLANCE_PRODUCTION_MAPPINGS_JSON optional: true + - name: NEXTJS_PROPOSAL_ENGINE + value: "api" + - name: ML_PROPOSAL_MODEL + value: "random_forest" + - name: ML_PROPOSAL_CANDIDATES + value: "1024" + # Row-level security on experiment tables (api/ROW_LEVEL_SECURITY.md). + # Informational while DB_MIGRATE_ON_STARTUP=false: enforcement is + # changed by setting DB_RLS_ENFORCEMENT in dt-db-migration-secret and + # running the dt-db-migrate Job (09-db-migrate-job.yaml). Changing + # this value alone has no effect. + - name: DB_RLS_ENFORCEMENT + value: "enforce" + # Migrations run as a separate Job; the API has no owner credentials. + - name: DB_MIGRATE_ON_STARTUP + value: "false" envFrom: - secretRef: name: dt-api-secret diff --git a/geddes/k8s/08-retrain-cron.yaml b/geddes/k8s/08-retrain-cron.yaml index a735cfa..0d42272 100644 --- a/geddes/k8s/08-retrain-cron.yaml +++ b/geddes/k8s/08-retrain-cron.yaml @@ -49,7 +49,7 @@ spec: - name: sdx-registry-secret containers: - name: retrain - image: geddes-registry.rcac.purdue.edu/sdx/dt-api:db3-pilot-retrain-loocv-20260714@sha256:52e16ffbbdf71ba87470d8dbf686e6304955501d5b74eaf4c41ba981e14239dc + image: geddes-registry.rcac.purdue.edu/sdx/dt-api:rls-b02e9ac@sha256:9204c93bdfba615e96e7e233990f97d2d16def03762a307161300f76fcada3f2 imagePullPolicy: Always command: - python diff --git a/geddes/k8s/09-db-migrate-job.yaml b/geddes/k8s/09-db-migrate-job.yaml new file mode 100644 index 0000000..52235e9 --- /dev/null +++ b/geddes/k8s/09-db-migrate-job.yaml @@ -0,0 +1,45 @@ +# Schema migrations for the platform database. This Job is how migrations run: +# the dt-api Deployment has DB_MIGRATE_ON_STARTUP=false and no owner +# credentials, so it never migrates on its own. +# +# Before rolling out any dt-api image, set this Job's image to the same release +# and run it to completion; roll out only after it succeeds. It reads owner +# credentials and DB_RLS_ENFORCEMENT from dt-db-migration-secret, which is also +# where row-level security enforcement is switched (enforce / disable): +# +# kubectl -n ncn-digitaltwins-zchen delete job dt-db-migrate --ignore-not-found +# kubectl apply -f geddes/k8s/09-db-migrate-job.yaml +# kubectl -n ncn-digitaltwins-zchen wait --for=condition=complete job/dt-db-migrate --timeout=300s +# kubectl -n ncn-digitaltwins-zchen logs job/dt-db-migrate +# +# See api/ROW_LEVEL_SECURITY.md. +apiVersion: batch/v1 +kind: Job +metadata: + name: dt-db-migrate + namespace: ncn-digitaltwins-zchen +spec: + backoffLimit: 1 + ttlSecondsAfterFinished: 86400 + template: + metadata: + labels: + app: dt-db-migrate + spec: + restartPolicy: Never + imagePullSecrets: + - name: sdx-registry-secret + containers: + - name: migrate + image: geddes-registry.rcac.purdue.edu/sdx/dt-api:rls-b02e9ac@sha256:9204c93bdfba615e96e7e233990f97d2d16def03762a307161300f76fcada3f2 + command: ["python", "-m", "db_migrations"] + envFrom: + - secretRef: + name: dt-db-migration-secret + resources: + requests: + cpu: "50m" + memory: "256Mi" + limits: + cpu: "500m" + memory: "512Mi" diff --git a/geddes/k8s/GLANCE_INGESTION_OPERATIONS.md b/geddes/k8s/GLANCE_INGESTION_OPERATIONS.md new file mode 100644 index 0000000..db95d97 --- /dev/null +++ b/geddes/k8s/GLANCE_INGESTION_OPERATIONS.md @@ -0,0 +1,73 @@ +# GLANCE ingestion operations + +Notes for operating the `dt-glance-ingestion` CronJob +(`glance-ingestion-cronjob.yaml`, connector in `azure/common/glance_connector.py`). + +## How a failing run blocks ingestion + +When the connector cannot build a run (for example it exceeds +`GLANCE_MAX_VALUES_PER_RUN` or `GLANCE_MAX_SAMPLES_PER_RUN`), it records the +failure in the audit file and **freezes the cursor** at the last acknowledged +run, so no data is skipped silently. Every later poll re-reads the same page +(up to `GLANCE_MAX_RUNS_PER_POLL` rows in total, including the overlap rows), +re-sends the runs it can build, and never gets past the page. The failure audit only records failures; +it is not a skip list. + +## Deferring a run + +There is no built-in per-run skip. Setting the cursor to the failing run does +not work either: each poll re-reads the `GLANCE_RUN_OVERLAP_COUNT` most recent +runs at or below the cursor, which includes the cursor run itself. The +procedure that uses existing behavior only: + +1. Confirm the run is safe to defer (see the record below for the checks) and + that no proposal in the mapped project is waiting on it. +2. In `dt-api-secret`, set the tool's `initial_run_id` in + `GLANCE_PRODUCTION_MAPPINGS_JSON` to the failing run id + 1. Change only + that field and keep a copy of the original value. Do this with all + `GLANCE_BACKFILL_*` settings unset: `initial_run_id` is the lower bound for + normal polling, including the overlap rows, but a backfill minimum + overrides it and exact-id backfills bypass it. The API does not use it. +3. Let the next poll run. It ingests the following runs and acknowledges the + cursor normally. +4. Once the cursor is at least `GLANCE_RUN_OVERLAP_COUNT` real runs past the + deferred run, restore the original `initial_run_id`. +5. Record the deferral below. + +A deferred run stays in GLANCE. Import it later with an exact-id backfill +(`GLANCE_BACKFILL_RUN_IDS=`, which uses its own cursor key and leaves the +normal cursor alone) once its handling is decided. + +## Deferred runs + +### Tool 3 (`VLN-11304-CTC-PM1`), run 12822 — deferred 2026-09-25 + +- **Symptom:** ingestion failed from 2026-09-04 with "Run 12822 value count + exceeds configured limit 1000000". The backlog extended to source run id + 13188; runs beyond the frozen page (which reached id 12969) were never + ingested. +- **Run:** 2026-09-03 21:04 to 2026-09-04 08:25 UTC (11 h 21 min), 38,766 + samples at about 1 Hz, 68 parameters, 2,636,088 values. Typical runs on this + tool: median 522 samples, 99th percentile 5,186. +- **Content:** gas flows (SF6, C4F8, Ar) and RF3 forward power (up to about + 2,000 W) are active in only about 2.4% of samples (about 15 minutes, a + normal process length); the rest is idle logging with RF and gas at zero. +- **Why not ingest it whole:** `derive_trace_feature_values` averages every + sample, so wherever its trace-derived means are used they would be diluted + to about 48 W RF3 and about 6 sccm SF6 instead of the real process values. Raising the limit would also push an + untested payload size through the 1 GiB API pod. +- **Checks before deferring:** no recipe proposals exist in project + `glance_live_dse_0297fa75`; the platform already held a partial copy of the + run (first 7,000 trace samples, no inputs, not a training row) from + 2026-09-04 12:20; no other backlog run exceeded the limit. +- **Procedure:** `initial_run_id` set from 12164 to 12823 at 05:28 UTC; the + 05:30 poll moved the cursor 12822 → 13003 (20 runs, no failures); the floor + was restored to 12164 at 05:34 UTC; the 05:40 poll with the original floor + moved it 13003 → 13132 (20 runs, no failures). +- **To import later:** trim the idle tail with a validated, tool-specific + active-window rule (or accept the diluted features explicitly), then + backfill it with `GLANCE_BACKFILL_RUN_IDS=12822`. The backfill uses the + normal run builder, so importing the unchanged run needs + `GLANCE_MAX_VALUES_PER_RUN` of at least 2,636,088 (a window selected before + values are accumulated could keep the current limit). A completed backfill + list stays checkpointed; reapplying it does not re-import. diff --git a/web/app/experiment/new/page.tsx b/web/app/experiment/new/page.tsx index 17726b7..11ba6f0 100644 --- a/web/app/experiment/new/page.tsx +++ b/web/app/experiment/new/page.tsx @@ -8,6 +8,7 @@ import { useSearchParams } from "next/navigation"; import { createExperiment, createExperimentType, + newClientRequestId, getEquipmentDetail, getExperimentTypes, getEquipmentListSimple, @@ -140,10 +141,19 @@ function NewExperimentContent() { { id: 3, label: "Review & Publish", icon: FileCheck }, ]; const prefillApplied = useRef(false); + // One idempotency key per run submission, kept until it succeeds so a + // retry after a lost response cannot create a second run. + const experimentRequestKey = useRef(null); const [currentStep, setCurrentStep] = useState(projectRunMode ? 2 : 1); const [formData, setFormData] = useState(INITIAL_DATA); const [submitting, setSubmitting] = useState(false); const [submitResult, setSubmitResult] = useState<{ status: string; message: string; href?: string; linkLabel?: string } | null>(null); + + // Editing the form starts a new submission; only an unchanged retry + // reuses the idempotency key. + useEffect(() => { + experimentRequestKey.current = null; + }, [formData]); const [equipment, setEquipment] = useState([]); const [selectedEquipmentDetail, setSelectedEquipmentDetail] = useState(null); const [experimentTypes, setExperimentTypes] = useState([]); @@ -615,7 +625,11 @@ function NewExperimentContent() { return; } + if (!experimentRequestKey.current) { + experimentRequestKey.current = newClientRequestId(); + } const experimentResult = await createExperiment({ + client_request_id: experimentRequestKey.current, experiment_name: formData.experiment_name, selected_equipment: formData.selected_equipment, project_id: formData.selected_project, @@ -642,6 +656,7 @@ function NewExperimentContent() { return; } + experimentRequestKey.current = null; setSubmitResult({ status: "success", message: diff --git a/web/app/optimize/page.tsx b/web/app/optimize/page.tsx index 953a966..d3b3b5e 100644 --- a/web/app/optimize/page.tsx +++ b/web/app/optimize/page.tsx @@ -1,9 +1,11 @@ "use client"; -import { useState, useEffect } from "react"; +import { useState, useEffect, useRef } from "react"; import { cn } from "@/lib/utils"; import { createExperiment, + newClientRequestId, + type CreateExperimentPayload, getProjectExperiments, getMlRuntimeStatus, getProposalsData, @@ -36,6 +38,11 @@ interface ProposalBatch { export default function OptimizePage() { const [loading, setLoading] = useState(true); + // The exact save request per generated recipe (and target project/unit + // experiment), idempotency key included, so a double-click or retry + // resends the same request and saves one run, while recipes from a new + // optimization get fresh requests. + const saveRequests = useRef(new WeakMap>()); const [optimizing, setOptimizing] = useState(false); const [error, setError] = useState(null); const [batches, setBatches] = useState([]); @@ -182,27 +189,39 @@ export default function OptimizePage() { if (!unitExperiment) return; setSavingRecipeIndex(index); setSaveMessage(null); - const plannedParameters = Object.entries(proposal.parameters).map(([name, value]) => ({ - name, - planned_value: value === null || value === undefined ? "" : String(value), - unit: featureLabels[name]?.unit || unitExperiment.type_parameters.find((param) => param.name === name)?.unit || "", - })); - const result = await createExperiment({ - experiment_name: `ML recipe ${index + 1} - ${new Date().toLocaleDateString()}`, - selected_equipment: unitExperiment.equipment_id, - project_id: selectedProject, - type_id: unitExperiment.id, - planned_date: "", - planned_parameters: plannedParameters, - optimization_context: { - mode: "optimization", - primary_target: "AvgEtchRate", - primary_objective: "maximize", - secondary_target: "RangeEtchRate", - secondary_objective: "minimize", - notes: "Created from the Optimization page ML recommendation.", - }, - }); + let requestsForProposal = saveRequests.current.get(proposal); + if (!requestsForProposal) { + requestsForProposal = new Map(); + saveRequests.current.set(proposal, requestsForProposal); + } + const saveTarget = `${selectedProject}|${selectedUnitExperiment}`; + let request = requestsForProposal.get(saveTarget); + if (!request) { + const plannedParameters = Object.entries(proposal.parameters).map(([name, value]) => ({ + name, + planned_value: value === null || value === undefined ? "" : String(value), + unit: featureLabels[name]?.unit || unitExperiment.type_parameters.find((param) => param.name === name)?.unit || "", + })); + request = { + client_request_id: newClientRequestId(), + experiment_name: `ML recipe ${index + 1} - ${new Date().toLocaleDateString()}`, + selected_equipment: unitExperiment.equipment_id, + project_id: selectedProject, + type_id: unitExperiment.id, + planned_date: "", + planned_parameters: plannedParameters, + optimization_context: { + mode: "optimization", + primary_target: "AvgEtchRate", + primary_objective: "maximize", + secondary_target: "RangeEtchRate", + secondary_objective: "minimize", + notes: "Created from the Optimization page ML recommendation.", + }, + }; + requestsForProposal.set(saveTarget, request); + } + const result = await createExperiment(request); setSavingRecipeIndex(null); if (!result) { setSaveMessage("Could not save this ML recipe as a project run."); diff --git a/web/lib/api-client.ts b/web/lib/api-client.ts index 258d4a2..ff58d72 100644 --- a/web/lib/api-client.ts +++ b/web/lib/api-client.ts @@ -1205,11 +1205,15 @@ export interface CreateExperimentPayload { planned_parameters: PlannedExperimentParameterInput[]; optimization_context?: OptimizationContextInput; sample_ids?: string[]; + /** Reuse the same value when retrying one submission; the API then returns + * the experiment the first attempt created instead of a duplicate. */ + client_request_id?: string; } export interface CreateExperimentResult { status: string; experiment_id: string; + idempotent_replay?: boolean; execution_request_id?: string; project_id?: string; optimization_status?: "completed" | "failed" | "queued" | "generating" | "skipped" | string; @@ -1219,6 +1223,14 @@ export interface CreateExperimentResult { optimizer_result?: ProposalsData; } +/** New idempotency key for one experiment submission. */ +export function newClientRequestId(): string { + if (typeof crypto !== "undefined" && typeof crypto.randomUUID === "function") { + return crypto.randomUUID(); + } + return `req-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 12)}`; +} + export function createExperiment(payload: CreateExperimentPayload) { return apiJson("/equipment/experiment", "POST", payload); } @@ -1611,7 +1623,9 @@ export function getHealth() { system_token_configured: boolean; ingestion_token_configured: boolean; postgres_api_password_configured: boolean; + postgres_worker_credentials_configured?: boolean; postgres_superuser_password_configured: boolean; + postgres_worker_role_configured?: boolean; }; postgres: { configured: boolean; @@ -1619,6 +1633,7 @@ export function getHealth() { database?: string; user?: string; missing_tables: string[]; + row_level_security?: Record; error?: string; }; trace_id: string;