From 755dbd4c19ca5f285d4f70ee3045aaa86b2738b0 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:29:08 +0100 Subject: [PATCH 01/19] fix(core): fence scan worker leases (#303) Signed-off-by: Shaurya K Sharma --- .../e4f7a9b2c6d8_scan_leases_and_fencing.py | 68 ++++ api/models/finding.py | 316 ++++++++++++------ docs/async-scan-architecture.md | 11 +- scanner/worker.py | 144 +++++++- tests/test_async_scan_persistence.py | 81 ++++- tests/test_observability.py | 5 +- tests/test_scan_leases_postgres.py | 298 +++++++++++++++++ tests/test_severity_contract.py | 27 +- tests/test_worker.py | 25 +- 9 files changed, 841 insertions(+), 134 deletions(-) create mode 100644 alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py create mode 100644 tests/test_scan_leases_postgres.py diff --git a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py new file mode 100644 index 00000000..a1ffbc93 --- /dev/null +++ b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py @@ -0,0 +1,68 @@ +"""Add renewable ownership leases and fencing tokens to scans. + +Revision ID: e4f7a9b2c6d8 +Revises: d8e4f6a1b2c3 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "e4f7a9b2c6d8" +down_revision: Union[str, Sequence[str], None] = "d8e4f6a1b2c3" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Add additive lease state and make legacy running work recoverable.""" + op.add_column("scans", sa.Column("lease_owner", sa.Text(), nullable=True)) + op.add_column("scans", sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True)) + op.add_column("scans", sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True)) + op.add_column( + "scans", + sa.Column("fencing_token", sa.BigInteger(), server_default=sa.text("0"), nullable=False), + ) + + # A pre-lease running row belongs to an old worker that cannot satisfy the + # new fencing contract. Marking its lease expired preserves the row and + # lets the new worker recover it under a fresh owner/token. + op.execute( + """ + UPDATE scans + SET lease_expires_at = CURRENT_TIMESTAMP + WHERE status = 'running' AND lease_expires_at IS NULL + """ + ) + + # These indexes are additive and are created concurrently so a populated + # production scans table remains available while the migration runs. + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_pending_started_at + ON scans (started_at ASC) + WHERE status = 'pending' + """ + ) + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_running_lease_expires_at + ON scans (lease_expires_at ASC) + WHERE status = 'running' + """ + ) + + +def downgrade() -> None: + """Remove lease metadata; callers must be rolled back first.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_running_lease_expires_at") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_pending_started_at") + op.drop_column("scans", "fencing_token") + op.drop_column("scans", "last_heartbeat_at") + op.drop_column("scans", "lease_expires_at") + op.drop_column("scans", "lease_owner") diff --git a/api/models/finding.py b/api/models/finding.py index c2adaafe..5d51b8a3 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -11,6 +11,7 @@ import psycopg2 import psycopg2.extras import psycopg2.pool +from psycopg2 import extensions from openshield.severity import ( CONTRACT_VERSION, @@ -23,6 +24,11 @@ logger = logging.getLogger(__name__) + +class LostLease(RuntimeError): + """Raised when a worker no longer owns the scan it is trying to update.""" + + FRAMEWORKS_DIR = Path(__file__).parent.parent.parent / "compliance" / "frameworks" # One pool per DSN, shared across all DatabaseManager instances in this @@ -158,26 +164,59 @@ def __init__(self, dsn: Optional[str] = None) -> None: def connect(self) -> None: """Acquire a connection from this DSN's shared pool.""" + if self.conn is not None: + self._return_connection(close=bool(self.conn.closed)) self.conn = _get_pool(self.dsn).getconn() self.conn.autocommit = False logger.debug("Database connection acquired from pool") def _get_conn(self) -> Any: - if self.conn is None or self.conn.closed: + if self.conn is not None and self.conn.closed: + self._return_connection(close=True) + if self.conn is None: self.connect() return self.conn - def close(self) -> None: - """Return the connection to its pool (or discard it if broken).""" - if self.conn is None: + def _return_connection(self, close: bool = False, conn: Optional[Any] = None) -> None: + """Return the checked-out connection, discarding it when requested.""" + conn = conn or self.conn + if conn is None: return try: - _get_pool(self.dsn).putconn(self.conn, close=bool(self.conn.closed)) + _get_pool(self.dsn).putconn(conn, close=close or bool(conn.closed)) logger.debug("Database connection returned to pool") except Exception as exc: logger.error("Error returning connection to pool: %s", exc) + try: + conn.close() + except Exception: + pass + finally: + if self.conn is conn: + self.conn = None + + def rollback(self, conn: Optional[Any] = None) -> None: + """End a failed transaction, discarding a connection the server lost.""" + conn = conn or self.conn + if conn is None: + return + try: + if conn.closed or conn.info.transaction_status == extensions.TRANSACTION_STATUS_UNKNOWN: + self._return_connection(close=True, conn=conn) + return + conn.rollback() + except Exception as exc: + logger.warning("Database rollback failed; discarding connection: %s", exc) + self._return_connection(close=True, conn=conn) + + def close(self) -> None: + """Return the connection to its pool (or discard it if broken).""" + if self.conn is None: + return + try: + self.rollback() finally: - self.conn = None + self._return_connection() def ping(self) -> bool: """Execute a trivial query to confirm database connectivity. @@ -202,8 +241,13 @@ def init_db(self) -> None: # Write # # ------------------------------------------------------------------ # - def save_scan(self, scan_result: Dict[str, Any]) -> None: - """Persist a full scan result (scan header + all findings).""" + def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token: int) -> None: + """Persist a completed scan only while its worker still owns the lease. + + The ownership check and all authoritative result writes share one + transaction. A worker whose lease was reclaimed therefore cannot + delete or insert child findings after it has become stale. + """ from datetime import datetime, timezone # Validate and canonicalize the entire batch before issuing SQL. A bad @@ -220,32 +264,40 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None: with conn.cursor() as cur: cur.execute( """ - INSERT INTO scans ( - scan_id, subscription_id, started_at, completed_at, - total_findings, score, cve_enrichment_status, status, - attempt_count, error_message, severity_contract_version - ) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (scan_id) DO UPDATE SET - completed_at = EXCLUDED.completed_at, - total_findings = EXCLUDED.total_findings, - score = EXCLUDED.score, - status = EXCLUDED.status, - error_message = EXCLUDED.error_message, - severity_contract_version = EXCLUDED.severity_contract_version + SELECT scan_id + FROM scans + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (scan_result["scan_id"], lease_owner, fencing_token), + ) + if cur.fetchone() is None: + raise LostLease(f"Scan {scan_result['scan_id']} is no longer owned by this worker") + cur.execute( + """ + UPDATE scans + SET completed_at = %s, + total_findings = %s, + score = %s, + cve_enrichment_status = %s, + status = 'completed', + error_message = NULL, + severity_contract_version = %s, + lease_owner = NULL, + lease_expires_at = NULL + WHERE scan_id = %s """, ( - scan_result["scan_id"], - scan_result["subscription_id"], - scan_result["started_at"], completed_at, len(findings), score_findings(findings), scan_result.get("cve_enrichment_status", "PENDING"), - scan_result.get("status", "completed"), - scan_result.get("attempt_count", 0), - scan_result.get("error_message"), CONTRACT_VERSION, + scan_result["scan_id"], ), ) # A worker retry replaces the previous result atomically. This @@ -360,7 +412,7 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None: # psycopg2 connections remain in an aborted transaction after any # SQL error. Roll back here so the worker can record failure and # safely process subsequent scans on the same pooled connection. - conn.rollback() + self.rollback(conn) raise logger.info( "Saved scan %s with %d findings", @@ -482,94 +534,158 @@ def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: conn.commit() logger.info("Created pending scan %s for %s", scan_id, subscription_id) - def update_scan_status(self, scan_id: str, status: str, error_message: Optional[str] = None) -> None: - """Update the status of a scan (running, completed, failed).""" + def update_scan_status( + self, + scan_id: str, + status: str, + error_message: Optional[str], + *, + lease_owner: str, + fencing_token: int, + ) -> None: + """Record a terminal failure while the caller still owns the lease.""" + if status != "failed": + raise ValueError("Fenced status updates only support terminal failures") conn = self._get_conn() - from datetime import datetime, timezone - - with conn.cursor() as cur: - if status == "completed": - completed_at = datetime.now(timezone.utc).isoformat() - cur.execute( - "UPDATE scans SET status = %s, completed_at = %s, error_message = NULL WHERE scan_id = %s", - (status, completed_at, scan_id), - ) - else: + try: + with conn.cursor() as cur: cur.execute( - "UPDATE scans SET status = %s, error_message = %s WHERE scan_id = %s", - (status, error_message, scan_id), + """ + UPDATE scans + SET status = 'failed', + error_message = %s, + lease_owner = NULL, + lease_expires_at = NULL + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + """, + (error_message, scan_id, lease_owner, fencing_token), ) - conn.commit() + if cur.rowcount != 1: + raise LostLease(f"Scan {scan_id} is no longer owned by this worker") + conn.commit() + except Exception: + self.rollback(conn) + raise logger.info("Updated scan %s status to %s", scan_id, status) - def claim_next_pending_scan(self) -> Optional[Dict[str, Any]]: - """Atomically claim the next pending scan using SKIP LOCKED.""" + def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: + """Atomically claim one pending scan and establish its renewable lease.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") conn = self._get_conn() - from datetime import datetime, timezone + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE scans + SET status = 'running', + claimed_at = CURRENT_TIMESTAMP, + lease_owner = %s, + lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP, + fencing_token = COALESCE(fencing_token, 0) + 1, + attempt_count = COALESCE(attempt_count, 0) + 1, + error_message = NULL + WHERE scan_id = ( + SELECT scan_id + FROM scans + WHERE status = 'pending' + ORDER BY started_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING * + """, + (lease_owner, lease_seconds), + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + except Exception: + self.rollback(conn) + raise - claimed_at = datetime.now(timezone.utc).isoformat() - with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute( - """ - UPDATE scans - SET status = 'running', - claimed_at = %s, - attempt_count = COALESCE(attempt_count, 0) + 1, - error_message = NULL - WHERE scan_id = ( - SELECT scan_id - FROM scans - WHERE status = 'pending' - ORDER BY started_at ASC - FOR UPDATE SKIP LOCKED - LIMIT 1 + def heartbeat_scan(self, scan_id: str, lease_owner: str, fencing_token: int, lease_seconds: int) -> Dict[str, Any]: + """Renew a still-valid lease, or raise :class:`LostLease`.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE scans + SET lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + RETURNING * + """, + (lease_seconds, scan_id, lease_owner, fencing_token), ) - RETURNING * - """, - (claimed_at,), - ) - row = cur.fetchone() - if row: - conn.commit() - return dict(row) - return None + row = cur.fetchone() + if row is None: + self.rollback(conn) + raise LostLease(f"Scan {scan_id} lease was lost before heartbeat") + conn.commit() + return dict(row) + except LostLease: + raise + except Exception: + self.rollback(conn) + raise - def recover_stale_scans(self, timeout_minutes: int = 60, max_attempts: int = 3) -> int: - """Recover scans left running after a worker crash or restart. + def recover_stale_scans(self, max_attempts: int = 3) -> int: + """Recover scans only after their renewable leases have expired. Stale scans are returned to pending while retry attempts remain. Once a scan has reached max_attempts, it is marked failed so it cannot loop forever on bad credentials or persistent Azure errors. """ conn = self._get_conn() - with conn.cursor() as cur: - cur.execute( - """ - UPDATE scans - SET status = 'failed', - error_message = 'Scan exceeded maximum retry attempts after worker interruption.' - WHERE status = 'running' - AND COALESCE(attempt_count, 1) >= %s - AND claimed_at < (CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute')) - """, - (max_attempts, timeout_minutes), - ) - failed_count = cur.rowcount + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET status = 'failed', + lease_owner = NULL, + lease_expires_at = NULL, + error_message = 'Scan exceeded maximum retry attempts after worker interruption.' + WHERE status = 'running' + AND COALESCE(attempt_count, 1) >= %s + AND lease_expires_at < CURRENT_TIMESTAMP + """, + (max_attempts,), + ) + failed_count = cur.rowcount - cur.execute( - """ - UPDATE scans - SET status = 'pending', - claimed_at = NULL, - error_message = 'Scan worker interrupted before completion. Queued for retry.' - WHERE status = 'running' - AND COALESCE(attempt_count, 0) < %s - AND claimed_at < (CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute')) - """, - (max_attempts, timeout_minutes), - ) - retry_count = cur.rowcount - conn.commit() + cur.execute( + """ + UPDATE scans + SET status = 'pending', + claimed_at = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + error_message = 'Scan worker interrupted before completion. Queued for retry.' + WHERE status = 'running' + AND COALESCE(attempt_count, 0) < %s + AND lease_expires_at < CURRENT_TIMESTAMP + """, + (max_attempts,), + ) + retry_count = cur.rowcount + conn.commit() + except Exception: + self.rollback(conn) + raise total_count = failed_count + retry_count if total_count > 0: logger.info( diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index 1d25d3a7..e9cf84f2 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -21,10 +21,17 @@ The scans table acts as a persistent task queue. This avoids the need for additi ### Render restart behavior Scan state is not stored in Flask memory. `POST /api/scans/trigger` inserts a `pending` row into PostgreSQL, and `GET /api/scans/` reads that same row back from PostgreSQL. If the Render web process restarts, queued scan state remains in the database and the dashboard can continue polling by `scan_id` after the app process comes back. -If the worker process restarts while a scan is marked `running`, `scanner/worker.py` calls `recover_stale_scans()` on each loop. Stale running scans are moved back to `pending` while retry attempts remain, so a Render restart can resume queued work instead of losing it. Once a scan reaches the maximum attempt count, it is marked `failed` so bad credentials or persistent Azure errors cannot retry forever. +Each claim is a renewable lease. The worker records a process-lifetime owner ID, +an expiry time, and a monotonically increasing fencing token, then renews the +lease while Azure work is running. `recover_stale_scans()` only requeues work +after its lease expires; a healthy worker is never reclaimed solely because its +original claim is old. A reclaimed scan receives a new token, so the previous +worker cannot persist completion, failure, or findings after it loses ownership. +Once a scan reaches the maximum attempt count, it is marked `failed` so bad +credentials or persistent Azure errors cannot retry forever. ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It queries the DB for scans where status is pending. It updates the status to running to prevent other workers from picking it up. It invokes ScanEngine.run_scan(scan_id). On success, it saves findings and sets status to completed. On failure, it captures the traceback and sets status to failed with the error_message. +The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and marks the scan complete atomically. On failure, it records a sanitized error only while it still owns the lease. ## Technical Rationale diff --git a/scanner/worker.py b/scanner/worker.py index 21a94c6f..f2f84b4a 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -7,11 +7,13 @@ import logging import os +import threading import time import traceback +import uuid from datetime import datetime, timezone -from api.models.finding import DatabaseManager +from api.models.finding import DatabaseManager, LostLease from api.observability import ( PENDING_SCANS, SCAN_DURATION_SECONDS, @@ -25,6 +27,102 @@ logger = logging.getLogger("scanner.worker") POLL_INTERVAL_SECONDS = 5 +DEFAULT_LEASE_SECONDS = 15 * 60 +DEFAULT_HEARTBEAT_SECONDS = 5 * 60 + + +def _positive_seconds(name: str, default: int) -> int: + """Read a positive interval without allowing malformed deploy config to stop work.""" + raw_value = os.environ.get(name) + if raw_value is None: + return default + try: + value = int(raw_value) + except ValueError: + logger.warning("Invalid %s=%r; using %d", name, raw_value, default) + return default + if value <= 0: + logger.warning("Invalid %s=%r; using %d", name, raw_value, default) + return default + return value + + +def lease_configuration() -> tuple[int, int]: + """Return a lease interval and a safely shorter heartbeat interval.""" + lease_seconds = _positive_seconds("SCAN_LEASE_SECONDS", DEFAULT_LEASE_SECONDS) + heartbeat_seconds = _positive_seconds("SCAN_HEARTBEAT_SECONDS", DEFAULT_HEARTBEAT_SECONDS) + if heartbeat_seconds >= lease_seconds: + heartbeat_seconds = max(1, lease_seconds // 3) + logger.warning( + "SCAN_HEARTBEAT_SECONDS must be less than SCAN_LEASE_SECONDS; using %d seconds", + heartbeat_seconds, + ) + return lease_seconds, heartbeat_seconds + + +class LeaseHeartbeat: + """Renew one claim through a dedicated database connection. + + Scan execution may block on Azure calls. The heartbeat intentionally owns + a separate DatabaseManager so it never shares a psycopg connection with + the worker's final persistence transaction. + """ + + def __init__( + self, + db_url: str, + scan_id: str, + lease_owner: str, + fencing_token: int, + lease_seconds: int, + heartbeat_seconds: int, + ) -> None: + self.db_url = db_url + self.scan_id = scan_id + self.lease_owner = lease_owner + self.fencing_token = fencing_token + self.lease_seconds = lease_seconds + self.heartbeat_seconds = heartbeat_seconds + self.lost = threading.Event() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name=f"scan-heartbeat-{scan_id}", daemon=True) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> bool: + """Stop the heartbeat, reporting whether its thread actually exited.""" + self._stop.set() + self._thread.join(timeout=10) + if self._thread.is_alive(): + logger.critical("Heartbeat thread for scan %s did not stop promptly", self.scan_id) + return False + return True + + def _run(self) -> None: + db = DatabaseManager(self.db_url) + try: + while not self._stop.wait(self.heartbeat_seconds): + try: + db.heartbeat_scan( + self.scan_id, + self.lease_owner, + self.fencing_token, + self.lease_seconds, + ) + except LostLease: + self.lost.set() + logger.warning( + "Lease lost while scan %s was executing", self.scan_id, extra={"scan_id": self.scan_id} + ) + return + except Exception as exc: + # A transient DB failure must be visible and the next + # heartbeat must get a clean/reacquired connection. + logger.error("Heartbeat failed for scan %s: %s", self.scan_id, exc, exc_info=True) + db.rollback() + finally: + db.close() def run_worker(): @@ -38,24 +136,27 @@ def run_worker(): init_sentry() db = DatabaseManager(db_url) + worker_id = str(uuid.uuid4()) + lease_seconds, heartbeat_seconds = lease_configuration() logger.info("OpenShield Background Worker started. Polling every %ds", POLL_INTERVAL_SECONDS) while True: try: # 1. Cleanup stale scans from previous crashes - db.recover_stale_scans(timeout_minutes=60) + db.recover_stale_scans() # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) # 3. Atomic claim - scan = db.claim_next_pending_scan() + scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: time.sleep(POLL_INTERVAL_SECONDS) continue scan_id = str(scan["scan_id"]) subscription_id = scan["subscription_id"] + fencing_token = scan["fencing_token"] logger.info( "Starting scan %s for %s", @@ -65,6 +166,15 @@ def run_worker(): ) scan_start = time.perf_counter() + heartbeat = LeaseHeartbeat( + db_url, + scan_id, + worker_id, + fencing_token, + lease_seconds, + heartbeat_seconds, + ) + heartbeat.start() try: engine = ScanEngine(subscription_id) result = engine.run_scan(scan_id) @@ -73,14 +183,27 @@ def run_worker(): result["completed_at"] = datetime.now(timezone.utc).isoformat() result["status"] = "completed" - db.save_scan(result) + if not heartbeat.stop(): + # Do not start another scan while a database call in this + # heartbeat is still stuck; process supervision can restart + # us and the unconfirmed claim will safely expire. + return + if heartbeat.lost.is_set(): + raise LostLease(f"Scan {scan_id} lost its lease before completion") + db.save_scan(result, worker_id, fencing_token) SCANS_TOTAL.labels(status="completed").inc() logger.info( "Successfully completed scan %s", scan_id, extra={"scan_id": scan_id}, ) + except LostLease: + if not heartbeat.stop(): + return + logger.warning("Scan %s finished after its lease was lost; no result was persisted", scan_id) except Exception as exc: + if not heartbeat.stop(): + return error_msg = f"{str(exc)}\n{traceback.format_exc()}" SCANS_TOTAL.labels(status="failed").inc() logger.error( @@ -92,12 +215,23 @@ def run_worker(): # Sanitize public error message public_error = "An internal error occurred during the scan. Please check the logs." - db.update_scan_status(scan_id, "failed", error_message=public_error) + try: + db.update_scan_status( + scan_id, + "failed", + error_message=public_error, + lease_owner=worker_id, + fencing_token=fencing_token, + ) + except LostLease: + logger.warning("Scan %s failed after its lease was lost; failure was not persisted", scan_id) finally: + heartbeat.stop() SCAN_DURATION_SECONDS.observe(time.perf_counter() - scan_start) except Exception as exc: logger.error("Worker loop encountered an error: %s", exc) + db.rollback() time.sleep(POLL_INTERVAL_SECONDS) diff --git a/tests/test_async_scan_persistence.py b/tests/test_async_scan_persistence.py index 3bfbe4d6..dc022adc 100644 --- a/tests/test_async_scan_persistence.py +++ b/tests/test_async_scan_persistence.py @@ -7,7 +7,9 @@ from unittest.mock import MagicMock, patch -from api.models.finding import DatabaseManager +import pytest + +from api.models.finding import DatabaseManager, LostLease class _Cursor: @@ -105,20 +107,79 @@ def test_claim_next_pending_scan_increments_attempt_count(): """Claiming a pending scan should record a durable execution attempt.""" db = DatabaseManager.__new__(DatabaseManager) scan_id = "44444444-4444-4444-4444-444444444444" - cursor = _Cursor(rows=[{"scan_id": scan_id, "attempt_count": 1}]) + cursor = _Cursor(rows=[{"scan_id": scan_id, "attempt_count": 1, "fencing_token": 1}]) conn = MagicMock() conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - scan = db.claim_next_pending_scan() + scan = db.claim_next_pending_scan("worker-a", 900) executed_sql = cursor.calls[0][0] assert "attempt_count = COALESCE(attempt_count, 0) + 1" in executed_sql + assert "lease_owner = %s" in executed_sql + assert "lease_expires_at = CURRENT_TIMESTAMP" in executed_sql + assert "fencing_token = COALESCE(fencing_token, 0) + 1" in executed_sql assert "error_message = NULL" in executed_sql assert scan["scan_id"] == scan_id conn.commit.assert_called_once() +def test_empty_claim_commits_before_the_worker_can_sleep(): + """A no-work poll must not leave the long-lived worker transaction open.""" + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + assert db.claim_next_pending_scan("worker-a", 900) is None + + conn.commit.assert_called_once() + + +def test_heartbeat_requires_current_owner_token_and_unexpired_lease(): + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + with pytest.raises(LostLease): + db.heartbeat_scan("scan-1", "worker-a", 3, 900) + + sql = cursor.calls[0][0] + assert "lease_owner = %s" in sql + assert "fencing_token = %s" in sql + assert "lease_expires_at > CURRENT_TIMESTAMP" in sql + conn.rollback.assert_called_once() + + +def test_fenced_failure_rejects_a_stale_owner(): + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + with pytest.raises(LostLease): + db.update_scan_status( + "scan-1", + "failed", + "failure", + lease_owner="worker-a", + fencing_token=3, + ) + + sql = cursor.calls[0][0] + assert "lease_owner = %s" in sql + assert "fencing_token = %s" in sql + assert "lease_expires_at > CURRENT_TIMESTAMP" in sql + + def test_recover_stale_scans_retries_before_max_attempts(): """Stale running scans should return to pending while attempts remain.""" db = DatabaseManager.__new__(DatabaseManager) @@ -127,15 +188,17 @@ def test_recover_stale_scans_retries_before_max_attempts(): conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - recovered = db.recover_stale_scans(timeout_minutes=15, max_attempts=3) + recovered = db.recover_stale_scans(max_attempts=3) failed_sql, failed_params = cursor.calls[0] retry_sql, retry_params = cursor.calls[1] assert "status = 'failed'" in failed_sql - assert failed_params == (3, 15) + assert failed_params == (3,) + assert "lease_expires_at < CURRENT_TIMESTAMP" in failed_sql assert "status = 'pending'" in retry_sql assert "claimed_at = NULL" in retry_sql - assert retry_params == (3, 15) + assert "lease_owner = NULL" in retry_sql + assert retry_params == (3,) assert recovered == 1 conn.commit.assert_called_once() @@ -148,13 +211,13 @@ def test_recover_stale_scans_fails_after_max_attempts(): conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - recovered = db.recover_stale_scans(timeout_minutes=60, max_attempts=3) + recovered = db.recover_stale_scans(max_attempts=3) failed_sql, failed_params = cursor.calls[0] retry_sql, retry_params = cursor.calls[1] assert "COALESCE(attempt_count, 1) >= %s" in failed_sql - assert failed_params == (3, 60) + assert failed_params == (3,) assert "COALESCE(attempt_count, 0) < %s" in retry_sql - assert retry_params == (3, 60) + assert retry_params == (3,) assert recovered == 1 conn.commit.assert_called_once() diff --git a/tests/test_observability.py b/tests/test_observability.py index 5fe30d90..ef9207ba 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -189,7 +189,7 @@ class _Stop(BaseException): mock_db = MagicMock() mock_db.recover_stale_scans.side_effect = [None, _Stop()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": scan_id, "subscription_id": sub_id}, + {"scan_id": scan_id, "subscription_id": sub_id, "fencing_token": 1}, None, ] mock_engine = MagicMock() @@ -203,6 +203,9 @@ class _Stop(BaseException): monkeypatch.setattr(worker, "DatabaseManager", MagicMock(return_value=mock_db)) monkeypatch.setattr(worker, "ScanEngine", MagicMock(return_value=mock_engine)) + heartbeat = MagicMock() + heartbeat.lost.is_set.return_value = False + monkeypatch.setattr(worker, "LeaseHeartbeat", MagicMock(return_value=heartbeat)) monkeypatch.setattr(worker.os.environ, "get", lambda *a, **k: "postgresql://x") monkeypatch.setattr(worker.time, "sleep", lambda *a, **k: None) diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py new file mode 100644 index 00000000..bb0c6efa --- /dev/null +++ b/tests/test_scan_leases_postgres.py @@ -0,0 +1,298 @@ +"""PostgreSQL-backed lease, fencing, and connection-recovery tests.""" + +import os +import threading +import uuid +from datetime import datetime, timezone + +import psycopg2 +from psycopg2 import extensions +import pytest + +from api.models.finding import DatabaseManager, LostLease + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +def _result(scan_id: str, subscription_id: str) -> dict: + return { + "scan_id": scan_id, + "subscription_id": subscription_id, + "started_at": datetime.now(timezone.utc).isoformat(), + "completed_at": datetime.now(timezone.utc).isoformat(), + "findings": [ + { + "rule_id": "AZ-LEASE-001", + "rule_name": "Lease test finding", + "severity": "HIGH", + "category": "Test", + "resource_id": f"/subscriptions/{subscription_id}/resourceGroups/test/providers/Test/resource", + "resource_name": "resource", + "resource_type": "Test/resource", + "description": "Lease test finding", + "remediation": "Fix the test resource", + "frameworks": {}, + "metadata": {}, + "detected_at": datetime.now(timezone.utc).isoformat(), + } + ], + } + + +class ScanRows: + def __init__(self, dsn: str): + self.dsn = dsn + self.scan_ids: list[str] = [] + + def create(self) -> tuple[str, str]: + scan_id = str(uuid.uuid4()) + subscription_id = str(uuid.uuid4()) + db = DatabaseManager(self.dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + finally: + db.close() + self.scan_ids.append(scan_id) + return scan_id, subscription_id + + def expire(self, scan_id: str) -> None: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE scans SET lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' WHERE scan_id = %s", + (scan_id,), + ) + + def scan(self, scan_id: str) -> dict: + db = DatabaseManager(self.dsn) + try: + scan = db.get_scan(scan_id) + assert scan is not None + return scan + finally: + db.close() + + def finding_count(self, scan_id: str) -> int: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM findings WHERE scan_id = %s", (scan_id,)) + return cur.fetchone()[0] + + def cleanup(self) -> None: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + for scan_id in self.scan_ids: + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +@pytest.fixture +def scan_rows() -> ScanRows: + rows = ScanRows(os.environ["DATABASE_URL"]) + yield rows + rows.cleanup() + + +def _claim(dsn: str, owner: str) -> dict | None: + db = DatabaseManager(dsn) + try: + return db.claim_next_pending_scan(owner, 120) + finally: + db.close() + + +def _recover(dsn: str, max_attempts: int = 3) -> int: + db = DatabaseManager(dsn) + try: + return db.recover_stale_scans(max_attempts=max_attempts) + finally: + db.close() + + +def test_two_workers_race_to_claim_only_one_scan(scan_rows): + scan_rows.create() + barrier = threading.Barrier(2) + claims: list[dict | None] = [] + + def claim(owner: str) -> None: + barrier.wait() + claims.append(_claim(scan_rows.dsn, owner)) + + threads = [threading.Thread(target=claim, args=(owner,)) for owner in ("worker-a", "worker-b")] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + successful = [claim for claim in claims if claim is not None] + assert len(successful) == 1 + assert successful[0]["lease_owner"] in {"worker-a", "worker-b"} + assert successful[0]["fencing_token"] == 1 + + +def test_active_lease_cannot_be_reclaimed(scan_rows): + scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a") is not None + assert _claim(scan_rows.dsn, "worker-b") is None + + +def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_rows): + scan_id, _ = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + renewed = db.heartbeat_scan(scan_id, "worker-a", claim["fencing_token"], 600) + finally: + db.close() + + assert renewed["lease_owner"] == "worker-a" + assert renewed["fencing_token"] == claim["fencing_token"] + assert renewed["lease_expires_at"] > claim["lease_expires_at"] + + +def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): + scan_id, _ = scan_rows.create() + first_claim = _claim(scan_rows.dsn, "worker-a") + assert first_claim is not None + scan_rows.expire(scan_id) + + assert _recover(scan_rows.dsn) == 1 + second_claim = _claim(scan_rows.dsn, "worker-b") + + assert second_claim is not None + assert second_claim["lease_owner"] == "worker-b" + assert second_claim["fencing_token"] > first_claim["fencing_token"] + + +def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows): + scan_id, subscription_id = scan_rows.create() + first_claim = _claim(scan_rows.dsn, "worker-a") + assert first_claim is not None + scan_rows.expire(scan_id) + _recover(scan_rows.dsn) + second_claim = _claim(scan_rows.dsn, "worker-b") + assert second_claim is not None + + stale_db = DatabaseManager(scan_rows.dsn) + try: + with pytest.raises(LostLease): + stale_db.heartbeat_scan(scan_id, "worker-a", first_claim["fencing_token"], 120) + with pytest.raises(LostLease): + stale_db.update_scan_status( + scan_id, + "failed", + "stale worker failure", + lease_owner="worker-a", + fencing_token=first_claim["fencing_token"], + ) + with pytest.raises(LostLease): + stale_db.save_scan(_result(scan_id, subscription_id), "worker-a", first_claim["fencing_token"]) + finally: + stale_db.close() + + assert scan_rows.finding_count(scan_id) == 0 + scan = scan_rows.scan(scan_id) + assert scan["status"] == "running" + assert scan["lease_owner"] == "worker-b" + assert scan["fencing_token"] == second_claim["fencing_token"] + + +def test_current_owner_completion_persists_results_atomically(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-b") + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(_result(scan_id, subscription_id), "worker-b", claim["fencing_token"]) + finally: + db.close() + + scan = scan_rows.scan(scan_id) + assert scan["status"] == "completed" + assert scan["lease_owner"] is None + assert scan_rows.finding_count(scan_id) == 1 + + +def test_sql_abort_rolls_back_and_the_connection_remains_usable(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + broken_result = _result(scan_id, subscription_id) + broken_result["findings"][0]["detected_at"] = None + + db = DatabaseManager(scan_rows.dsn) + try: + with pytest.raises(psycopg2.Error): + db.save_scan(broken_result, "worker-a", claim["fencing_token"]) + db.update_scan_status( + scan_id, + "failed", + "controlled SQL failure", + lease_owner="worker-a", + fencing_token=claim["fencing_token"], + ) + finally: + db.close() + + assert scan_rows.scan(scan_id)["status"] == "failed" + assert scan_rows.finding_count(scan_id) == 0 + + +def test_terminated_backend_is_discarded_and_reacquired(scan_rows): + scan_id, _ = scan_rows.create() + db = DatabaseManager(scan_rows.dsn) + try: + conn = db._get_conn() + with conn.cursor() as cur: + cur.execute("SELECT pg_backend_pid()") + backend_pid = cur.fetchone()[0] + conn.commit() + + with psycopg2.connect(scan_rows.dsn) as terminator: + with terminator.cursor() as cur: + cur.execute("SELECT pg_terminate_backend(%s)", (backend_pid,)) + assert cur.fetchone()[0] is True + + with pytest.raises(psycopg2.OperationalError): + db.ping() + db.rollback() + + replacement = db._get_conn() + with replacement.cursor() as cur: + cur.execute("SELECT pg_backend_pid()") + assert cur.fetchone()[0] != backend_pid + db.rollback() + assert db.get_scan(scan_id) is not None + finally: + db.close() + + +def test_expired_restart_work_obeys_attempt_limit(scan_rows): + scan_id, _ = scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + + assert _claim(scan_rows.dsn, "worker-b") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + + assert _claim(scan_rows.dsn, "worker-c") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + assert scan_rows.scan(scan_id)["status"] == "failed" + + +def test_empty_claim_leaves_no_open_transaction(scan_rows): + db = DatabaseManager(scan_rows.dsn) + try: + assert db.claim_next_pending_scan("worker-a", 120) is None + assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE + finally: + db.close() diff --git a/tests/test_severity_contract.py b/tests/test_severity_contract.py index 097ce33b..4a6822ec 100644 --- a/tests/test_severity_contract.py +++ b/tests/test_severity_contract.py @@ -138,7 +138,7 @@ def test_persistence_rejects_invalid_severity_before_opening_connection(): } with patch.object(db, "_get_conn") as get_conn: with pytest.raises(SeverityContractError): - db.save_scan(result) + db.save_scan(result, "worker-a", 1) get_conn.assert_not_called() @@ -172,14 +172,16 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) - - scan_parameters = cursor.execute.call_args_list[0].args[1] - delete_parameters = cursor.execute.call_args_list[1].args[1] - finding_parameters = cursor.execute.call_args_list[2].args[1] - assert scan_parameters[4] == 1 - assert scan_parameters[5] == 100 - assert scan_parameters[10] == CONTRACT_VERSION + db.save_scan(result, "worker-a", 1) + + lock_parameters = cursor.execute.call_args_list[0].args[1] + scan_parameters = cursor.execute.call_args_list[1].args[1] + delete_parameters = cursor.execute.call_args_list[2].args[1] + finding_parameters = cursor.execute.call_args_list[3].args[1] + assert lock_parameters == (result["scan_id"], "worker-a", 1) + assert scan_parameters[1] == 1 + assert scan_parameters[2] == 100 + assert scan_parameters[4] == CONTRACT_VERSION assert delete_parameters == (result["scan_id"],) assert finding_parameters[0] == result["scan_id"] assert finding_parameters[3] == "INFO" @@ -191,8 +193,11 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): def test_persistence_rolls_back_a_failed_atomic_replacement(): db = _db() cursor = _cursor() - cursor.execute.side_effect = [None, None, RuntimeError("insert failed")] + cursor.fetchone.return_value = {"scan_id": "scan"} + cursor.execute.side_effect = [None, None, None, RuntimeError("insert failed")] conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 conn.cursor.return_value = cursor result = { "scan_id": "00000000-0000-0000-0000-000000000000", @@ -203,7 +208,7 @@ def test_persistence_rolls_back_a_failed_atomic_replacement(): with patch.object(db, "_get_conn", return_value=conn): with pytest.raises(RuntimeError, match="insert failed"): - db.save_scan(result) + db.save_scan(result, "worker-a", 1) conn.rollback.assert_called_once_with() conn.commit.assert_not_called() diff --git a/tests/test_worker.py b/tests/test_worker.py index 29a00cf1..ddd0f5a0 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -8,7 +8,7 @@ """ import unittest -from unittest.mock import patch +from unittest.mock import ANY, patch from scanner.worker import run_worker, POLL_INTERVAL_SECONDS import uuid @@ -29,7 +29,10 @@ def setUp(self): @patch("scanner.worker.ScanEngine") @patch("scanner.worker.os.environ.get") @patch("scanner.worker.time.sleep") - def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, mock_engine_class, mock_db_class): + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_processes_pending_scan_successfully( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): """ Verify the happy path: 1. Worker claims a pending scan atomically. @@ -54,9 +57,10 @@ def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, # We need to stop the infinite loop. We'll raise StopWorker on the second call to recover_stale_scans. mock_db.recover_stale_scans.side_effect = [None, StopWorker()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": self.scan_id, "subscription_id": self.subscription_id}, + {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, ] + mock_heartbeat_class.return_value.lost.is_set.return_value = False with self.assertRaises(StopWorker): run_worker() @@ -71,12 +75,16 @@ def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, saved_result = mock_db.save_scan.call_args[0][0] self.assertEqual(saved_result["status"], "completed") self.assertIn("completed_at", saved_result) + self.assertEqual(mock_db.save_scan.call_args[0][2], 1) @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.ScanEngine") @patch("scanner.worker.os.environ.get") @patch("scanner.worker.time.sleep") - def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock_engine_class, mock_db_class): + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_handles_scan_failure_gracefully( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): """ Verify the error path: 1. Worker claims a pending scan. @@ -88,9 +96,10 @@ def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock mock_db.recover_stale_scans.side_effect = [None, StopWorker()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": self.scan_id, "subscription_id": self.subscription_id}, + {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, ] + mock_heartbeat_class.return_value.lost.is_set.return_value = False # Mock Engine to fail mock_engine = mock_engine_class.return_value @@ -101,7 +110,11 @@ def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock # Verify status was updated to failed with sanitized message mock_db.update_scan_status.assert_any_call( - self.scan_id, "failed", error_message="An internal error occurred during the scan. Please check the logs." + self.scan_id, + "failed", + error_message="An internal error occurred during the scan. Please check the logs.", + lease_owner=ANY, + fencing_token=1, ) # Ensure findings were NOT saved on failure mock_db.save_scan.assert_not_called() From eaa5613da498f6ec6704f753d7f288635fddf98e Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:48:46 +0100 Subject: [PATCH 02/19] test(core): validate scan lease recovery Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 8 ++- tests/test_database_manager_reliability.py | 21 +++++++ tests/test_worker.py | 73 +++++++++++++++++++++- 3 files changed, 100 insertions(+), 2 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index 5d51b8a3..0ca6fc5b 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -165,7 +165,13 @@ def __init__(self, dsn: Optional[str] = None) -> None: def connect(self) -> None: """Acquire a connection from this DSN's shared pool.""" if self.conn is not None: - self._return_connection(close=bool(self.conn.closed)) + previous_conn = self.conn + self.rollback(previous_conn) + # rollback() discards a dead connection and clears self.conn. A + # healthy connection still needs exactly one pool return before a + # replacement is borrowed. + if self.conn is previous_conn: + self._return_connection(close=bool(previous_conn.closed), conn=previous_conn) self.conn = _get_pool(self.dsn).getconn() self.conn.autocommit = False logger.debug("Database connection acquired from pool") diff --git a/tests/test_database_manager_reliability.py b/tests/test_database_manager_reliability.py index 06249280..d8a7573a 100644 --- a/tests/test_database_manager_reliability.py +++ b/tests/test_database_manager_reliability.py @@ -71,6 +71,27 @@ def test_connect_acquires_connection_from_shared_pool(): assert db.conn is fake_conn +def test_reconnect_rolls_back_and_returns_previous_connection_once(): + dsn = "postgresql://pool-test/db" + fake_pool = MagicMock() + old_conn = MagicMock() + old_conn.closed = 0 + old_conn.info.transaction_status = 0 + new_conn = MagicMock() + new_conn.closed = 0 + fake_pool.getconn.return_value = new_conn + db = _db(dsn) + db.conn = old_conn + + with patch.object(finding_module, "_get_pool", return_value=fake_pool): + db.connect() + + old_conn.rollback.assert_called_once() + fake_pool.putconn.assert_called_once_with(old_conn, close=False) + fake_pool.getconn.assert_called_once() + assert db.conn is new_conn + + def test_close_returns_healthy_connection_to_pool(): dsn = "postgresql://pool-test/db" fake_pool = MagicMock() diff --git a/tests/test_worker.py b/tests/test_worker.py index ddd0f5a0..73b051df 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -9,7 +9,8 @@ import unittest from unittest.mock import ANY, patch -from scanner.worker import run_worker, POLL_INTERVAL_SECONDS +from api.models.finding import LostLease +from scanner.worker import LeaseHeartbeat, POLL_INTERVAL_SECONDS, run_worker import uuid @@ -19,6 +20,20 @@ class StopWorker(BaseException): pass +class OneHeartbeatThenStop: + """Deterministically run one heartbeat loop iteration without sleeping.""" + + def __init__(self): + self.calls = 0 + + def wait(self, _seconds): + self.calls += 1 + return self.calls > 1 + + def set(self): + pass + + class TestWorker(unittest.TestCase): def setUp(self): self.mock_db_url = "postgresql://user:pass@localhost/db" @@ -76,6 +91,7 @@ def test_worker_processes_pending_scan_successfully( self.assertEqual(saved_result["status"], "completed") self.assertIn("completed_at", saved_result) self.assertEqual(mock_db.save_scan.call_args[0][2], 1) + self.assertGreaterEqual(mock_heartbeat_class.return_value.stop.call_count, 1) @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.ScanEngine") @@ -118,6 +134,61 @@ def test_worker_handles_scan_failure_gracefully( ) # Ensure findings were NOT saved on failure mock_db.save_scan.assert_not_called() + self.assertGreaterEqual(mock_heartbeat_class.return_value.stop.call_count, 1) + + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.ScanEngine") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_does_not_persist_after_heartbeat_reports_lost_lease( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_pending_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "fencing_token": 1, + } + mock_engine_class.return_value.run_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "findings": [], + } + mock_heartbeat_class.return_value.lost.is_set.return_value = True + + with self.assertRaises(StopWorker): + run_worker() + + mock_db.save_scan.assert_not_called() + mock_db.update_scan_status.assert_not_called() + + +@patch("scanner.worker.DatabaseManager") +def test_heartbeat_uses_and_closes_a_dedicated_database_manager(mock_db_class): + heartbeat = LeaseHeartbeat("postgresql://heartbeat-test/db", "scan-1", "worker-a", 7, 120, 1) + heartbeat._stop = OneHeartbeatThenStop() + + heartbeat._run() + + mock_db_class.assert_called_once_with("postgresql://heartbeat-test/db") + mock_db_class.return_value.heartbeat_scan.assert_called_once_with("scan-1", "worker-a", 7, 120) + mock_db_class.return_value.close.assert_called_once() + assert not heartbeat.lost.is_set() + + +@patch("scanner.worker.DatabaseManager") +def test_heartbeat_surfaces_lost_lease_and_closes_its_connection(mock_db_class): + mock_db_class.return_value.heartbeat_scan.side_effect = LostLease("stale") + heartbeat = LeaseHeartbeat("postgresql://heartbeat-test/db", "scan-1", "worker-a", 7, 120, 1) + heartbeat._stop = OneHeartbeatThenStop() + + heartbeat._run() + + assert heartbeat.lost.is_set() + mock_db_class.return_value.close.assert_called_once() @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.os.environ.get") From aeaa83ee45a1dd3e1f5c1702283b2c6d10323da8 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:25:03 +0100 Subject: [PATCH 03/19] fix(core): make scan result persistence idempotent Signed-off-by: Shaurya K Sharma --- ...4c9_idempotent_findings_and_evaluations.py | 66 ++++++++++++++++ api/models/finding.py | 78 ++++++++++++++++--- tests/test_scan_leases_postgres.py | 76 ++++++++++++++++++ tests/test_severity_contract.py | 8 +- 4 files changed, 212 insertions(+), 16 deletions(-) create mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py new file mode 100644 index 00000000..7b4522a9 --- /dev/null +++ b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py @@ -0,0 +1,66 @@ +"""Add database-enforced identities for scan results. + +Revision ID: f2b6d8e1a4c9 +Revises: e4f7a9b2c6d8 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +revision: str = "f2b6d8e1a4c9" +down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Add stable finding keys and per-resource evaluation rows.""" + op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) + # Existing records predate the identity contract. Preserve each record as + # distinct rather than attempting to infer equivalence from mutable text. + op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") + op.alter_column("findings", "finding_key", nullable=False) + + with op.get_context().autocommit_block(): + op.execute("CREATE UNIQUE INDEX CONCURRENTLY uq_findings_scan_finding_key ON findings (scan_id, finding_key)") + + op.create_table( + "rule_evaluations", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("scan_id", postgresql.UUID(), nullable=False), + sa.Column("rule_id", sa.Text(), nullable=False), + sa.Column("resource_id", sa.Text(), nullable=False), + sa.Column("resource_type", sa.Text(), server_default=sa.text("''"), nullable=False), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("reason_code", sa.Text(), nullable=True), + sa.Column("reason", sa.Text(), nullable=True), + sa.Column("evidence", postgresql.JSONB(), server_default=sa.text("'{}'::jsonb"), nullable=True), + sa.Column("finding_id", sa.Integer(), nullable=True), + sa.Column("evaluated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="rule_evaluations_scan_id_fkey"), + sa.ForeignKeyConstraint( + ["finding_id"], ["findings.id"], name="rule_evaluations_finding_id_fkey", ondelete="SET NULL" + ), + sa.PrimaryKeyConstraint("id", name="rule_evaluations_pkey"), + sa.UniqueConstraint("scan_id", "rule_id", "resource_id", name="uq_rule_evaluations_scan_rule_resource"), + sa.CheckConstraint( + "status IN ('PASS', 'FAIL', 'UNKNOWN', 'ERROR', 'NOT_APPLICABLE')", + name="ck_rule_evaluations_status_v1", + ), + sa.CheckConstraint("resource_id <> ''", name="ck_rule_evaluations_resource_id_not_empty"), + ) + op.create_index("idx_rule_evaluations_scan_id", "rule_evaluations", ["scan_id"], unique=False) + + +def downgrade() -> None: + """Remove idempotent-result storage introduced by this revision.""" + op.drop_index("idx_rule_evaluations_scan_id", table_name="rule_evaluations") + op.drop_table("rule_evaluations") + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_findings_scan_finding_key") + op.drop_column("findings", "finding_key") diff --git a/api/models/finding.py b/api/models/finding.py index 0ca6fc5b..a262c27f 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1,6 +1,7 @@ """Finding dataclass and PostgreSQL-backed DatabaseManager.""" import json +import hashlib import logging import os import threading @@ -42,6 +43,28 @@ class LostLease(RuntimeError): _POOL_MAX_CONN = int(os.environ.get("DB_POOL_MAX_CONN", "10")) +def stable_finding_key(scan_id: str, finding: Dict[str, Any]) -> str: + """Return an immutable identity for one logical finding in a scan. + + Rules that can report more than one violation for the same resource must + provide ``finding_discriminator``. Presentation fields such as severity, + description, and remediation are deliberately excluded so retries update + the existing authoritative finding rather than creating a duplicate. + """ + resource_scope = finding.get("resource_id") or { + "resource_type": finding.get("resource_type") or "", + "resource_name": finding.get("resource_name") or "", + } + identity = { + "scan_id": str(scan_id), + "rule_id": finding.get("rule_id") or "", + "resource_scope": resource_scope, + "discriminator": finding.get("finding_discriminator") or "default", + } + encoded = json.dumps(identity, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + def _get_pool(dsn: str) -> "psycopg2.pool.ThreadedConnectionPool": # Pool creation happens at most once per DSN per process lifetime, so a # plain lock (no unlocked fast path) is simpler and just as cheap here. @@ -262,7 +285,9 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token for raw_finding in scan_result.get("findings", []): finding = dict(raw_finding) finding["severity"] = normalize_severity(finding.get("severity")) + finding["finding_key"] = stable_finding_key(scan_result["scan_id"], finding) findings.append(finding) + evaluations = [dict(raw_evaluation) for raw_evaluation in scan_result.get("evaluations", [])] conn = self._get_conn() completed_at = scan_result.get("completed_at") or datetime.now(timezone.utc).isoformat() @@ -306,27 +331,39 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token scan_result["scan_id"], ), ) - # A worker retry replaces the previous result atomically. This - # keeps the scan header, child rows, and recomputed score in - # agreement instead of duplicating findings on every attempt. - cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_result["scan_id"],)) finding_id_by_key: Dict[Any, int] = {} for f in findings: cur.execute( """ INSERT INTO findings - (scan_id, rule_id, rule_name, severity, category, + (scan_id, finding_key, rule_id, rule_name, severity, category, resource_id, resource_name, resource_type, description, remediation, playbook, frameworks, metadata, cve_references, cvss_score, exploit_available, detected_at) - VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + ON CONFLICT (scan_id, finding_key) DO UPDATE SET + rule_name = EXCLUDED.rule_name, + severity = EXCLUDED.severity, + category = EXCLUDED.category, + resource_name = EXCLUDED.resource_name, + resource_type = EXCLUDED.resource_type, + description = EXCLUDED.description, + remediation = EXCLUDED.remediation, + playbook = EXCLUDED.playbook, + frameworks = EXCLUDED.frameworks, + metadata = EXCLUDED.metadata, + cve_references = EXCLUDED.cve_references, + cvss_score = EXCLUDED.cvss_score, + exploit_available = EXCLUDED.exploit_available, + detected_at = EXCLUDED.detected_at RETURNING id """, ( # The parent scan owns every child in this batch. # Never trust a caller-supplied child scan_id. scan_result["scan_id"], + f["finding_key"], f.get("rule_id"), f.get("rule_name"), f.get("severity"), @@ -345,13 +382,32 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token f.get("detected_at"), ), ) - finding_id_by_key[(f.get("rule_id"), f.get("resource_id"))] = cur.fetchone()[0] + # DO UPDATE still returns the row, so a replayed result + # keeps the *existing* finding id rather than minting a new + # one. Evaluation rows already pointing at it stay valid. + finding_row = cur.fetchone() + finding_id_by_key[(f.get("rule_id"), f.get("resource_id"))] = ( + finding_row["id"] if isinstance(finding_row, dict) else finding_row[0] + ) + + # Findings the current result no longer reports are removed by + # identity rather than by wiping the whole child set, so a + # replayed delivery never briefly empties a populated scan. + finding_keys = [f["finding_key"] for f in findings] + if finding_keys: + cur.execute( + "DELETE FROM findings WHERE scan_id = %s AND NOT (finding_key = ANY(%s))", + (scan_result["scan_id"], finding_keys), + ) + else: + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_result["scan_id"],)) # Coverage rows (#263): a status for every resource a migrated - # rule looked at, not just its violations. A FAIL evaluation - # is durably linked to the finding row it corresponds to - # right here, in the same transaction, instead of leaving - # callers to infer the relationship from rule_id/resource_id. + # rule looked at, not just its violations. The evaluation + # contract itself belongs to #321 and is reproduced here + # unchanged; what this change adds is that these writes now + # happen inside the fenced transaction above, so a worker that + # lost its lease cannot rewrite another owner's coverage. # # Upserted rather than replaced wholesale: a retried/replayed # scan result must converge on the same rows instead of a diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index bb0c6efa..7b2fa7df 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -81,10 +81,31 @@ def finding_count(self, scan_id: str) -> int: cur.execute("SELECT COUNT(*) FROM findings WHERE scan_id = %s", (scan_id,)) return cur.fetchone()[0] + def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: + """Simulate duplicate delivery of the same claimed result for persistence tests.""" + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET status = 'running', completed_at = NULL, lease_owner = %s, + fencing_token = %s, lease_expires_at = CURRENT_TIMESTAMP + INTERVAL '5 minutes' + WHERE scan_id = %s + """, + (owner, fencing_token, scan_id), + ) + + def evaluation_count(self, scan_id: str) -> int: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + return cur.fetchone()[0] + def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) @@ -296,3 +317,58 @@ def test_empty_claim_leaves_no_open_transaction(scan_rows): assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE finally: db.close() + + +def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + result = _result(scan_id, subscription_id) + result["evaluations"] = [ + { + "rule_id": "AZ-LEASE-001", + "resource_id": result["findings"][0]["resource_id"], + "resource_type": "Test/resource", + "status": "FAIL", + "evidence": {"version": 1}, + } + ] + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(result, "worker-a", claim["fencing_token"]) + first_id = db.get_findings({"scan_id": scan_id})[0]["id"] + scan_rows.rearm(scan_id, "worker-a", claim["fencing_token"]) + result["findings"][0]["description"] = "Updated presentation text" + result["findings"][0]["severity"] = "CRITICAL" + result["evaluations"][0]["evidence"] = {"version": 2} + db.save_scan(result, "worker-a", claim["fencing_token"]) + persisted = db.get_findings({"scan_id": scan_id}) + finally: + db.close() + + assert len(persisted) == 1 + assert persisted[0]["id"] == first_id + assert persisted[0]["description"] == "Updated presentation text" + assert persisted[0]["severity"] == "CRITICAL" + assert scan_rows.evaluation_count(scan_id) == 1 + + +def test_distinct_finding_discriminators_preserve_multiple_violations(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + result = _result(scan_id, subscription_id) + duplicate_scope = dict(result["findings"][0]) + result["findings"][0]["finding_discriminator"] = "network-rule-a" + duplicate_scope["finding_discriminator"] = "network-rule-b" + duplicate_scope["description"] = "A second violation on the same resource" + result["findings"].append(duplicate_scope) + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(result, "worker-a", claim["fencing_token"]) + finally: + db.close() + + assert scan_rows.finding_count(scan_id) == 2 diff --git a/tests/test_severity_contract.py b/tests/test_severity_contract.py index 4a6822ec..53388af6 100644 --- a/tests/test_severity_contract.py +++ b/tests/test_severity_contract.py @@ -176,15 +176,13 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): lock_parameters = cursor.execute.call_args_list[0].args[1] scan_parameters = cursor.execute.call_args_list[1].args[1] - delete_parameters = cursor.execute.call_args_list[2].args[1] - finding_parameters = cursor.execute.call_args_list[3].args[1] + finding_parameters = cursor.execute.call_args_list[2].args[1] assert lock_parameters == (result["scan_id"], "worker-a", 1) assert scan_parameters[1] == 1 assert scan_parameters[2] == 100 assert scan_parameters[4] == CONTRACT_VERSION - assert delete_parameters == (result["scan_id"],) assert finding_parameters[0] == result["scan_id"] - assert finding_parameters[3] == "INFO" + assert finding_parameters[4] == "INFO" assert raw_finding["severity"] == "INFORMATIONAL" conn.commit.assert_called_once_with() conn.rollback.assert_not_called() @@ -194,7 +192,7 @@ def test_persistence_rolls_back_a_failed_atomic_replacement(): db = _db() cursor = _cursor() cursor.fetchone.return_value = {"scan_id": "scan"} - cursor.execute.side_effect = [None, None, None, RuntimeError("insert failed")] + cursor.execute.side_effect = [None, None, RuntimeError("insert failed")] conn = MagicMock() conn.closed = 0 conn.info.transaction_status = 0 From d28c0dc5dcc4ea67ee146347cee09bcc4d48b087 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:28:24 +0100 Subject: [PATCH 04/19] fix(core): make scan admission durable and idempotent Signed-off-by: Shaurya K Sharma --- ...a7c5e9d2f1b4_scan_admission_idempotency.py | 47 +++++++ api/models/finding.py | 115 +++++++++++++++--- api/routes/scans.py | 53 +++++++- tests/test_async_scan_persistence.py | 4 +- tests/test_error_exposure.py | 2 +- tests/test_input_validation.py | 3 +- tests/test_scan_admission_postgres.py | 100 +++++++++++++++ 7 files changed, 303 insertions(+), 21 deletions(-) create mode 100644 alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py create mode 100644 tests/test_scan_admission_postgres.py diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py new file mode 100644 index 00000000..a45eb609 --- /dev/null +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -0,0 +1,47 @@ +"""Enforce durable scan admission and idempotency. + +Revision ID: a7c5e9d2f1b4 +Revises: f2b6d8e1a4c9 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "a7c5e9d2f1b4" +down_revision: Union[str, Sequence[str], None] = "f2b6d8e1a4c9" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Persist idempotency semantics and prevent more than one active scan.""" + op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) + op.add_column("scans", sa.Column("request_fingerprint", sa.Text(), nullable=True)) + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE UNIQUE INDEX CONCURRENTLY uq_scans_subscription_idempotency_key + ON scans (subscription_id, idempotency_key) + WHERE idempotency_key IS NOT NULL + """ + ) + op.execute( + """ + CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription + ON scans (subscription_id) + WHERE status IN ('pending', 'running') + """ + ) + + +def downgrade() -> None: + """Remove scan admission metadata and constraints.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_one_active_per_subscription") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_subscription_idempotency_key") + op.drop_column("scans", "request_fingerprint") + op.drop_column("scans", "idempotency_key") diff --git a/api/models/finding.py b/api/models/finding.py index a262c27f..999009e6 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -30,6 +30,14 @@ class LostLease(RuntimeError): """Raised when a worker no longer owns the scan it is trying to update.""" +class ScanAdmissionConflict(RuntimeError): + """Raised when an idempotency key is reused for different scan semantics.""" + + +class ScanQuotaExceeded(RuntimeError): + """Raised when an explicitly configured subscription scan quota is exhausted.""" + + FRAMEWORKS_DIR = Path(__file__).parent.parent.parent / "compliance" / "frameworks" # One pool per DSN, shared across all DatabaseManager instances in this @@ -579,22 +587,101 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) - def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: - """Create a scan record in the 'pending' state.""" + def admit_scan( + self, + scan_id: str, + subscription_id: str, + *, + idempotency_key: Optional[str] = None, + request_fingerprint: Optional[str] = None, + max_scans_per_hour: int = 0, + ) -> tuple[Dict[str, Any], bool]: + """Atomically admit one scan or return its durable logical predecessor. + + The PostgreSQL advisory transaction lock serializes admission decisions + for one subscription. The partial unique index added by the migration + remains the final database enforcement of the one-active-scan rule. + """ + if max_scans_per_hour < 0: + raise ValueError("max_scans_per_hour must not be negative") conn = self._get_conn() - from datetime import datetime, timezone + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (subscription_id,)) + if idempotency_key: + cur.execute( + """ + SELECT * FROM scans + WHERE subscription_id = %s AND idempotency_key = %s + """, + (subscription_id, idempotency_key), + ) + existing = cur.fetchone() + if existing: + existing = dict(existing) + if existing.get("request_fingerprint") != request_fingerprint: + raise ScanAdmissionConflict("Idempotency-Key was reused with different request semantics") + conn.commit() + return existing, False - started_at = datetime.now(timezone.utc).isoformat() - with conn.cursor() as cur: - cur.execute( - """ - INSERT INTO scans (scan_id, subscription_id, started_at, status, attempt_count) - VALUES (%s, %s, %s, 'pending', 0) - """, - (scan_id, subscription_id, started_at), - ) - conn.commit() - logger.info("Created pending scan %s for %s", scan_id, subscription_id) + cur.execute( + """ + SELECT * FROM scans + WHERE subscription_id = %s AND status IN ('pending', 'running') + ORDER BY started_at ASC + LIMIT 1 + """, + (subscription_id,), + ) + active_scan = cur.fetchone() + if active_scan: + conn.commit() + return dict(active_scan), False + + if max_scans_per_hour: + cur.execute( + """ + SELECT COUNT(*) FROM scans + WHERE subscription_id = %s + AND started_at >= CURRENT_TIMESTAMP - INTERVAL '1 hour' + """, + (subscription_id,), + ) + quota_row = cur.fetchone() + scan_count = quota_row["count"] if isinstance(quota_row, dict) else quota_row[0] + if scan_count >= max_scans_per_hour: + raise ScanQuotaExceeded("Configured hourly scan quota has been reached") + + from datetime import datetime, timezone + + cur.execute( + """ + INSERT INTO scans ( + scan_id, subscription_id, started_at, status, attempt_count, + idempotency_key, request_fingerprint + ) + VALUES (%s, %s, %s, 'pending', 0, %s, %s) + RETURNING * + """, + ( + scan_id, + subscription_id, + datetime.now(timezone.utc).isoformat(), + idempotency_key, + request_fingerprint, + ), + ) + admitted = dict(cur.fetchone()) + conn.commit() + logger.info("Admitted pending scan %s for %s", scan_id, subscription_id) + return admitted, True + except Exception: + self.rollback(conn) + raise + + def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: + """Create a pending scan for older internal callers without a key.""" + self.admit_scan(scan_id, subscription_id) def update_scan_status( self, diff --git a/api/routes/scans.py b/api/routes/scans.py index 9ec2a289..109635f7 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -2,11 +2,13 @@ import logging import os +import hashlib +import json import threading import uuid from flask import Blueprint, g, jsonify, request -from api.models.finding import DatabaseManager +from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded from api.validation import ( VALIDATION_ERROR_MESSAGE, ValidationError, @@ -20,6 +22,7 @@ logger = logging.getLogger(__name__) _AUTHORIZED_SUBSCRIPTIONS_ENV = "OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS" +_MAX_SCANS_PER_HOUR_ENV = "OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR" def _subscription_is_authorized(subscription_id: str) -> bool: @@ -53,6 +56,17 @@ def _get_db() -> DatabaseManager: return g.db +def _configured_hourly_quota() -> int: + """Return an optional policy quota; zero preserves existing no-limit policy.""" + raw_value = os.environ.get(_MAX_SCANS_PER_HOUR_ENV, "0") + try: + quota = int(raw_value) + except ValueError: + logger.warning("Invalid %s=%r; disabling hourly quota", _MAX_SCANS_PER_HOUR_ENV, raw_value) + return 0 + return max(0, quota) + + @scans_bp.get("/api/scans") def list_scans(): """Return all historical scan results ordered by most recent first.""" @@ -105,18 +119,49 @@ def trigger_scan(): logger.warning("Scan trigger rejected: subscription %s is not on the authorized allowlist", subscription_id) return jsonify({"error": "Subscription is not authorized for this deployment"}), 403 + idempotency_key = request.headers.get("Idempotency-Key") + if idempotency_key is not None: + idempotency_key = idempotency_key.strip() + if not idempotency_key or len(idempotency_key) > 200: + return jsonify({"error": "Idempotency-Key must be between 1 and 200 characters"}), 400 + request_fingerprint = hashlib.sha256( + json.dumps({"subscription_id": subscription_id}, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() scan_id = str(uuid.uuid4()) - logger.info("Async scan triggered for subscription %s (id: %s)", subscription_id, scan_id) try: db = _get_db() - db.create_pending_scan(scan_id, subscription_id) + admitted, created = db.admit_scan( + scan_id, + subscription_id, + idempotency_key=idempotency_key, + request_fingerprint=request_fingerprint, + max_scans_per_hour=_configured_hourly_quota(), + ) + except ScanAdmissionConflict as exc: + return jsonify({"error": str(exc)}), 409 + except ScanQuotaExceeded as exc: + return jsonify({"error": str(exc)}), 429 except Exception as exc: logger.error("Failed to create pending scan: %s", exc, exc_info=True) return jsonify({"error": "Database error"}), 500 + response_scan_id = str(admitted["scan_id"]) + if not created: + return jsonify( + { + "scan_id": response_scan_id, + "status": admitted["status"], + "message": "Existing logical scan returned.", + } + ), 200 + logger.info("Async scan admitted for subscription %s (id: %s)", subscription_id, response_scan_id) return jsonify( - {"scan_id": scan_id, "status": "pending", "message": "Scan has been queued and will start shortly."} + { + "scan_id": response_scan_id, + "status": "pending", + "message": "Scan has been queued and will start shortly.", + } ), 202 except ValidationError: diff --git a/tests/test_async_scan_persistence.py b/tests/test_async_scan_persistence.py index dc022adc..2883377d 100644 --- a/tests/test_async_scan_persistence.py +++ b/tests/test_async_scan_persistence.py @@ -42,6 +42,7 @@ def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, mo scan_id = "11111111-1111-1111-1111-111111111111" subscription_id = "00000000-0000-0000-0000-000000000000" mock_db = MagicMock() + mock_db.admit_scan.return_value = ({"scan_id": scan_id, "status": "pending"}, True) with patch("api.routes.scans.DatabaseManager", return_value=mock_db) as db_class: with patch("api.routes.scans.uuid.uuid4", return_value=scan_id): @@ -59,7 +60,8 @@ def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, mo } db_class.assert_called_once_with("postgresql://ci:ci@localhost/ci_db") mock_db.connect.assert_called_once() - mock_db.create_pending_scan.assert_called_once_with(scan_id, subscription_id) + mock_db.admit_scan.assert_called_once() + assert mock_db.admit_scan.call_args.args[:2] == (scan_id, subscription_id) def test_get_scan_status_reads_from_database(client, auth_headers, monkeypatch): diff --git a/tests/test_error_exposure.py b/tests/test_error_exposure.py index f098ccdb..cea508a5 100644 --- a/tests/test_error_exposure.py +++ b/tests/test_error_exposure.py @@ -46,7 +46,7 @@ def test_get_scan_status_error_does_not_leak_exception(client, auth_headers): def test_trigger_scan_db_error_does_not_leak_exception(client, auth_headers): - with patch.object(scans_route, "_get_db", return_value=_raising_db("create_pending_scan")): + with patch.object(scans_route, "_get_db", return_value=_raising_db("admit_scan")): resp = client.post( "/api/scans/trigger", json={"subscription_id": "00000000-0000-0000-0000-000000000001"}, diff --git a/tests/test_input_validation.py b/tests/test_input_validation.py index 86931256..8bca8054 100644 --- a/tests/test_input_validation.py +++ b/tests/test_input_validation.py @@ -60,10 +60,11 @@ def test_trigger_rejects_malformed_subscription_id(client, auth_headers): def test_trigger_accepts_canonical_subscription_uuid(client, auth_headers): db = MagicMock() + db.admit_scan.return_value = ({"scan_id": _SCAN_ID, "status": "pending"}, True) with patch.object(scans_route, "_get_db", return_value=db): response = client.post("/api/scans/trigger", json={"subscription_id": _SUBSCRIPTION_ID}, headers=auth_headers) assert response.status_code == 202 - assert db.create_pending_scan.call_args.args[1] == _SUBSCRIPTION_ID + assert db.admit_scan.call_args.args[1] == _SUBSCRIPTION_ID @pytest.mark.parametrize( diff --git a/tests/test_scan_admission_postgres.py b/tests/test_scan_admission_postgres.py new file mode 100644 index 00000000..9c0fa829 --- /dev/null +++ b/tests/test_scan_admission_postgres.py @@ -0,0 +1,100 @@ +"""Real PostgreSQL tests for durable scan admission invariants.""" + +import os +import threading +import uuid + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +@pytest.fixture +def admitted_scans(): + dsn = os.environ["DATABASE_URL"] + scan_ids: list[str] = [] + yield dsn, scan_ids + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + for scan_id in scan_ids: + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def _admit(dsn: str, subscription_id: str, key: str | None = None, fingerprint: str = "same"): + db = DatabaseManager(dsn) + try: + return db.admit_scan( + str(uuid.uuid4()), + subscription_id, + idempotency_key=key, + request_fingerprint=fingerprint if key else None, + ) + finally: + db.close() + + +def test_concurrent_admission_returns_one_active_scan(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + barrier = threading.Barrier(2) + outcomes = [] + + def admit() -> None: + barrier.wait() + outcomes.append(_admit(dsn, subscription_id)) + + threads = [threading.Thread(target=admit) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + scan_ids.append(str(outcomes[0][0]["scan_id"])) + assert {str(scan["scan_id"]) for scan, _created in outcomes} == {scan_ids[0]} + assert sum(created for _scan, created in outcomes) == 1 + + +def test_idempotency_key_replays_or_rejects_changed_semantics(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + first, created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + scan_ids.append(str(first["scan_id"])) + replay, replay_created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + + assert created is True + assert replay_created is False + assert replay["scan_id"] == first["scan_id"] + with pytest.raises(ScanAdmissionConflict): + _admit(dsn, subscription_id, "request-1", "fingerprint-b") + + +def test_completed_scan_allows_a_later_admission_and_configured_quota(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + first, _ = _admit(dsn, subscription_id) + scan_ids.append(str(first["scan_id"])) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'completed' WHERE scan_id = %s", (first["scan_id"],)) + + second, created = _admit(dsn, subscription_id) + scan_ids.append(str(second["scan_id"])) + assert created is True + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'completed' WHERE scan_id = %s", (second["scan_id"],)) + + db = DatabaseManager(dsn) + try: + with pytest.raises(ScanQuotaExceeded): + db.admit_scan(str(uuid.uuid4()), subscription_id, max_scans_per_hour=2) + finally: + db.close() From 95f88b703368aa2770f3e1e1b21834ca74da2fec Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:37:45 +0100 Subject: [PATCH 05/19] fix(core): make CVE enrichment durable Signed-off-by: Shaurya K Sharma --- .../c9e1a5b7d3f2_durable_enrichment_jobs.py | 69 +++++ api/models/finding.py | 268 ++++++++++++++++++ api/routes/scans.py | 99 +------ scanner/cve_correlator.py | 18 +- scanner/enrichment_worker.py | 50 ++++ scanner/nvd_client.py | 123 ++++---- scanner/worker.py | 10 +- tests/test_enrichment_jobs_postgres.py | 186 ++++++++++++ tests/test_nvd_client.py | 12 + tests/test_scans_enrich.py | 124 ++------ tests/test_worker.py | 4 + 11 files changed, 714 insertions(+), 249 deletions(-) create mode 100644 alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py create mode 100644 scanner/enrichment_worker.py create mode 100644 tests/test_enrichment_jobs_postgres.py diff --git a/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py new file mode 100644 index 00000000..ab9c3a0a --- /dev/null +++ b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py @@ -0,0 +1,69 @@ +"""Add durable, fenced CVE enrichment jobs. + +Revision ID: c9e1a5b7d3f2 +Revises: a7c5e9d2f1b4 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +revision: str = "c9e1a5b7d3f2" +down_revision: Union[str, Sequence[str], None] = "a7c5e9d2f1b4" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Create one resumable enrichment job per scan.""" + op.create_table( + "enrichment_jobs", + sa.Column("job_id", postgresql.UUID(), nullable=False), + sa.Column("scan_id", postgresql.UUID(), nullable=False), + sa.Column("status", sa.Text(), nullable=False, server_default=sa.text("'pending'")), + sa.Column("lease_owner", sa.Text(), nullable=True), + sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("fencing_token", sa.BigInteger(), nullable=False, server_default=sa.text("0")), + sa.Column("attempt_count", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column( + "next_retry_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("checkpoint", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="enrichment_jobs_scan_id_fkey"), + sa.PrimaryKeyConstraint("job_id", name="enrichment_jobs_pkey"), + sa.UniqueConstraint("scan_id", name="uq_enrichment_jobs_scan_id"), + sa.CheckConstraint("status IN ('pending', 'running', 'completed', 'failed')", name="ck_enrichment_jobs_status"), + ) + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_enrichment_jobs_pending_retry + ON enrichment_jobs (next_retry_at ASC) + WHERE status = 'pending' + """ + ) + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_enrichment_jobs_running_lease + ON enrichment_jobs (lease_expires_at ASC) + WHERE status = 'running' + """ + ) + + +def downgrade() -> None: + """Remove durable enrichment work state.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_running_lease") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_pending_retry") + op.drop_table("enrichment_jobs") diff --git a/api/models/finding.py b/api/models/finding.py index 999009e6..9d62a22e 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -5,6 +5,7 @@ import logging import os import threading +import uuid from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional @@ -587,6 +588,273 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) + def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: + """Durably enqueue exactly one CVE enrichment job for a scan.""" + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint) + VALUES (%s, %s, 'pending', 0, 0) + ON CONFLICT (scan_id) DO NOTHING + RETURNING * + """, + (str(uuid.uuid4()), scan_id), + ) + job = cur.fetchone() + if job is None: + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + existing = cur.fetchone() + if existing is None: + raise RuntimeError("enrichment job conflict did not return an existing job") + conn.commit() + return dict(existing), False + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) + conn.commit() + return dict(job), True + except Exception: + self.rollback(conn) + raise + + def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: + """Atomically claim the next retry-ready enrichment job.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET status = 'running', lease_owner = %s, + lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP, + fencing_token = fencing_token + 1, + attempt_count = attempt_count + 1, + error_message = NULL + WHERE job_id = ( + SELECT job_id FROM enrichment_jobs + WHERE status = 'pending' + AND next_retry_at <= CURRENT_TIMESTAMP + ORDER BY created_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING * + """, + (lease_owner, lease_seconds), + ) + job = cur.fetchone() + if job: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'ENRICHING' WHERE scan_id = %s", (job["scan_id"],) + ) + conn.commit() + return dict(job) if job else None + except Exception: + self.rollback(conn) + raise + + def heartbeat_enrichment_job(self, job_id: str, lease_owner: str, fencing_token: int, lease_seconds: int) -> None: + """Renew an enrichment claim or raise LostLease.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + """, + (lease_seconds, job_id, lease_owner, fencing_token), + ) + if cur.rowcount != 1: + raise LostLease(f"Enrichment job {job_id} is no longer owned by this worker") + conn.commit() + except Exception: + self.rollback(conn) + raise + + def get_enrichment_findings(self, scan_id: str) -> List[Dict[str, Any]]: + """Return a deterministic snapshot ordered for checkpointed enrichment.""" + conn = self._get_conn() + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT * FROM findings WHERE scan_id = %s ORDER BY id ASC", (scan_id,)) + return [dict(row) for row in cur.fetchall()] + + def persist_enrichment_progress( + self, + job_id: str, + lease_owner: str, + fencing_token: int, + finding: Dict[str, Any], + checkpoint: int, + ) -> None: + """Persist one finding and checkpoint it under the current job fence.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT scan_id FROM enrichment_jobs + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (job_id, lease_owner, fencing_token), + ) + if cur.fetchone() is None: + raise LostLease(f"Enrichment job {job_id} lost its lease before checkpointing") + cur.execute( + """ + UPDATE findings SET cve_references = %s, cvss_score = %s, exploit_available = %s + WHERE id = %s + """, + ( + json.dumps(finding.get("cve_references", [])), + finding.get("cvss_score"), + finding.get("exploit_available", False), + finding["id"], + ), + ) + cur.execute("UPDATE enrichment_jobs SET checkpoint = %s WHERE job_id = %s", (checkpoint, job_id)) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def complete_enrichment_job(self, job_id: str, lease_owner: str, fencing_token: int) -> None: + """Atomically complete a fenced job and its scan's enrichment state.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET status = 'completed', completed_at = CURRENT_TIMESTAMP, + lease_owner = NULL, lease_expires_at = NULL + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + RETURNING scan_id + """, + (job_id, lease_owner, fencing_token), + ) + row = cur.fetchone() + if row is None: + raise LostLease(f"Enrichment job {job_id} lost its lease before completion") + scan_id = row[0] + cur.execute("UPDATE scans SET cve_enrichment_status = 'COMPLETED' WHERE scan_id = %s", (scan_id,)) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def fail_enrichment_job( + self, + job_id: str, + lease_owner: str, + fencing_token: int, + error_message: str, + *, + max_attempts: int = 3, + retry_seconds: int = 30, + ) -> str: + """Record bounded retry state under the job's current fencing token.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT scan_id, attempt_count FROM enrichment_jobs + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (job_id, lease_owner, fencing_token), + ) + row = cur.fetchone() + if row is None: + raise LostLease(f"Enrichment job {job_id} is no longer owned by this worker") + scan_id, attempt_count = row + terminal = attempt_count >= max_attempts + if terminal: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'failed', completed_at = CURRENT_TIMESTAMP, + error_message = %s, lease_owner = NULL, lease_expires_at = NULL + WHERE job_id = %s + """, + (error_message, job_id), + ) + scan_status = "FAILED" + outcome = "failed" + else: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'pending', error_message = %s, + lease_owner = NULL, lease_expires_at = NULL, + next_retry_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second') + WHERE job_id = %s + """, + (error_message, retry_seconds, job_id), + ) + scan_status = "PENDING" + outcome = "retry" + cur.execute("UPDATE scans SET cve_enrichment_status = %s WHERE scan_id = %s", (scan_status, scan_id)) + conn.commit() + return outcome + except Exception: + self.rollback(conn) + raise + + def recover_stale_enrichment_jobs(self, max_attempts: int = 3) -> int: + """Return expired enrichment claims to pending or terminally fail them.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'failed', completed_at = CURRENT_TIMESTAMP, + lease_owner = NULL, lease_expires_at = NULL, + error_message = 'Enrichment exceeded maximum retry attempts after worker interruption.' + WHERE status = 'running' AND attempt_count >= %s + AND lease_expires_at < CURRENT_TIMESTAMP + RETURNING scan_id + """, + (max_attempts,), + ) + failed_scans = [row[0] for row in cur.fetchall()] + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, + error_message = 'Enrichment worker interrupted; queued for retry.' + WHERE status = 'running' AND attempt_count < %s + AND lease_expires_at < CURRENT_TIMESTAMP + RETURNING scan_id + """, + (max_attempts,), + ) + retried_scans = [row[0] for row in cur.fetchall()] + for scan_id in failed_scans: + cur.execute("UPDATE scans SET cve_enrichment_status = 'FAILED' WHERE scan_id = %s", (scan_id,)) + for scan_id in retried_scans: + cur.execute("UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", (scan_id,)) + conn.commit() + return len(failed_scans) + len(retried_scans) + except Exception: + self.rollback(conn) + raise + def admit_scan( self, scan_id: str, diff --git a/api/routes/scans.py b/api/routes/scans.py index 109635f7..ec90033a 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -4,7 +4,6 @@ import os import hashlib import json -import threading import uuid from flask import Blueprint, g, jsonify, request @@ -16,7 +15,6 @@ require_json_object, uuid_string, ) -from scanner.cve_correlator import enrich_findings scans_bp = Blueprint("scans", __name__) logger = logging.getLogger(__name__) @@ -171,81 +169,14 @@ def trigger_scan(): return jsonify({"error": "Critical route failure"}), 500 -def _run_enrichment_in_background(scan_id: str, findings: list, db_url: str) -> None: - """Run CVE enrichment off the request thread and persist the result. - - Runs outside the Flask request/app context (it's started via - threading.Thread), so it opens its own DatabaseManager rather than - reusing flask.g. - """ - db = DatabaseManager(db_url) - try: - enriched = enrich_findings(findings) - db.update_cve_fields(enriched) - db.update_scan_enrichment_status(scan_id, "COMPLETED") - logger.info("Background CVE enrichment complete for scan %s (%d findings)", scan_id, len(enriched)) - except Exception as exc: - logger.error("Background enrichment failed for scan %s: %s", scan_id, exc) - try: - # A failed write (e.g. in update_cve_fields) can leave db.conn in - # an aborted-transaction state. Roll back first, or this status - # update itself raises InFailedSqlTransaction and gets swallowed - # below, leaving the scan stuck at ENRICHING forever. - if db.conn is not None: - db.conn.rollback() - db.update_scan_enrichment_status(scan_id, "FAILED") - except Exception as status_exc: - logger.error("Failed to record FAILED status for scan %s: %s", scan_id, status_exc) - finally: - db.close() - - -def _run_enrichment_in_background(scan_id: str, findings: list, db_url: str) -> None: - """Run CVE enrichment off the request thread and persist the result. - - Runs outside the Flask request/app context (it's started via - threading.Thread), so it opens its own DatabaseManager rather than - reusing flask.g. - """ - db = DatabaseManager(db_url) - try: - enriched = enrich_findings(findings) - db.update_cve_fields(enriched) - db.update_scan_enrichment_status(scan_id, "COMPLETED") - logger.info("Background CVE enrichment complete for scan %s (%d findings)", scan_id, len(enriched)) - except Exception as exc: - logger.error("Background enrichment failed for scan %s: %s", scan_id, exc) - try: - # A failed write (e.g. in update_cve_fields) can leave db.conn in - # an aborted-transaction state. Roll back first, or this status - # update itself raises InFailedSqlTransaction and gets swallowed - # below, leaving the scan stuck at ENRICHING forever. - if db.conn is not None: - db.conn.rollback() - db.update_scan_enrichment_status(scan_id, "FAILED") - except Exception as status_exc: - logger.error("Failed to record FAILED status for scan %s: %s", scan_id, status_exc) - finally: - db.close() - - @scans_bp.post("/api/scans//enrich") def enrich_scan(scan_id): - """Kick off CVE enrichment for an existing scan in the background. - - Returns immediately with 202 and status ENRICHING; poll - GET /api/scans/ (cve_enrichment_status) for completion. - Running enrichment synchronously here previously caused the request to - time out on scans spanning many rule categories, since NVD lookups are - rate-limited to one every ~7 seconds. - """ + """Enqueue durable CVE enrichment; no request-owned thread is created.""" try: scan_id = uuid_string(scan_id, "scan_id") db = _get_db() - # Check current status to avoid redundant NVD calls - scans = db.get_scans() - current_scan = next((s for s in scans if str(s["scan_id"]) == scan_id), None) + current_scan = db.get_scan(scan_id) if not current_scan: return jsonify({"error": "Scan not found"}), 404 @@ -253,27 +184,27 @@ def enrich_scan(scan_id): status = current_scan.get("cve_enrichment_status") if status == "COMPLETED": return jsonify({"message": "Scan already enriched", "scan_id": scan_id}), 200 - if status == "ENRICHING": - return jsonify({"message": "Enrichment already in progress", "scan_id": scan_id}), 202 - findings = db.get_findings({"scan_id": scan_id}) if not findings: return jsonify({"error": "No findings found for this scan"}), 404 - logger.info("Starting background CVE enrichment for %d findings in scan %s", len(findings), scan_id) - db.update_scan_enrichment_status(scan_id, "ENRICHING") - - threading.Thread( - target=_run_enrichment_in_background, - args=(scan_id, findings, db.dsn), - daemon=True, - ).start() + job, created = db.enqueue_enrichment_job(scan_id) + if not created: + return jsonify( + { + "job_id": str(job["job_id"]), + "scan_id": scan_id, + "status": job["status"], + "message": "Existing enrichment job returned.", + } + ), 202 return jsonify( { "scan_id": scan_id, - "status": "ENRICHING", - "message": "CVE enrichment started; poll GET /api/scans/ for completion.", + "job_id": str(job["job_id"]), + "status": "PENDING", + "message": "CVE enrichment queued; poll GET /api/scans/ for completion.", } ), 202 diff --git a/scanner/cve_correlator.py b/scanner/cve_correlator.py index 0ac3de4a..94d4af0b 100644 --- a/scanner/cve_correlator.py +++ b/scanner/cve_correlator.py @@ -10,7 +10,7 @@ import logging from typing import Optional -from scanner.nvd_client import query_nvd +from scanner.nvd_client import query_nvd, query_nvd_strict logger = logging.getLogger(__name__) @@ -132,3 +132,19 @@ def enrich_findings(findings: list[dict]) -> list[dict]: enriched = [_enrich_single_finding(f) for f in findings] logger.info("CVE enrichment complete.") return enriched + + +def enrich_finding_durable(finding: dict) -> dict: + """Enrich one persisted finding and propagate retriable NVD failures.""" + keyword = _get_nvd_keyword(finding.get("rule_id", "")) + if not keyword: + finding["cve_references"] = [] + finding["cvss_score"] = None + finding["exploit_available"] = False + return finding + cves = query_nvd_strict(keyword) + finding["cve_references"] = cves + scores = [c["cvss_score"] for c in cves if c.get("cvss_score") is not None] + finding["cvss_score"] = max(scores) if scores else None + finding["exploit_available"] = any(c.get("exploit_available") for c in cves) + return finding diff --git a/scanner/enrichment_worker.py b/scanner/enrichment_worker.py new file mode 100644 index 00000000..4ca92c44 --- /dev/null +++ b/scanner/enrichment_worker.py @@ -0,0 +1,50 @@ +"""Durable CVE enrichment job execution for the scan worker.""" + +import logging + +from api.models.finding import DatabaseManager, LostLease +from scanner.cve_correlator import enrich_finding_durable + + +logger = logging.getLogger("scanner.enrichment_worker") + + +def process_enrichment_job( + db: DatabaseManager, + job: dict, + lease_owner: str, + lease_seconds: int, + *, + max_attempts: int = 3, +) -> str: + """Run a claimed job, checkpointing each finding under its fence. + + A restart resumes at the stored finding offset. Replaying the last item is + safe because persistence updates the existing finding row by primary key. + """ + job_id = str(job["job_id"]) + fencing_token = job["fencing_token"] + try: + findings = db.get_enrichment_findings(str(job["scan_id"])) + for index, finding in enumerate(findings[job["checkpoint"] :], start=job["checkpoint"]): + db.heartbeat_enrichment_job(job_id, lease_owner, fencing_token, lease_seconds) + enriched = enrich_finding_durable(finding) + db.persist_enrichment_progress(job_id, lease_owner, fencing_token, enriched, index + 1) + db.complete_enrichment_job(job_id, lease_owner, fencing_token) + logger.info("Completed enrichment job %s", job_id) + return "completed" + except LostLease: + logger.warning("Enrichment job %s lost its lease; no further writes attempted", job_id) + return "lost_lease" + except Exception as exc: + retry_seconds = min(300, 30 * (2 ** max(0, job["attempt_count"] - 1))) + outcome = db.fail_enrichment_job( + job_id, + lease_owner, + fencing_token, + "CVE enrichment failed; see worker logs for details.", + max_attempts=max_attempts, + retry_seconds=retry_seconds, + ) + logger.error("Enrichment job %s failed (%s): %s", job_id, outcome, exc, exc_info=True) + return outcome diff --git a/scanner/nvd_client.py b/scanner/nvd_client.py index 989ea0de..8ca65666 100644 --- a/scanner/nvd_client.py +++ b/scanner/nvd_client.py @@ -32,13 +32,17 @@ _NVD_BASE_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0" _REQUEST_DELAY_SECONDS = 7.0 # Stay under 5 req/30 sec limit _MAX_RETRIES = 3 -_RESULTS_PER_PAGE = 5 # Top 5 CVEs per finding is enough for display +_RESULTS_PER_PAGE = 2000 # NVD's documented maximum; fetch every matching page. # In-memory cache. Keyed by "keyword:results_per_page". # Resets each process - intentional, NVD data changes slowly. _cache: dict[str, list[dict]] = {} +class NvdRequestError(RuntimeError): + """Raised by the durable worker when an NVD page cannot be retrieved.""" + + class _RateLimiter: """Tracks the last NVD request time, guarded by a lock so concurrent callers (e.g. background enrichment threads) can't both read a stale @@ -128,25 +132,12 @@ def _parse_cve_item(item: dict) -> Optional[dict]: return None -def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: - """ - Query NVD for CVEs matching a keyword. - - Returns a list of parsed CVE dicts (may be empty). - Never raises - all failures return []. - - Args: - keyword: Search term, e.g. "Azure Storage Account" - results_per_page: Max CVEs to fetch (default 5) - """ - cache_key = f"{keyword}:{results_per_page}" - if cache_key in _cache: - logger.debug("NVD cache hit for: %s", keyword) - return _cache[cache_key] - +def _fetch_nvd_page(keyword: str, start_index: int, results_per_page: int) -> dict: + """Fetch one NVD page, raising only after bounded retries are exhausted.""" params = urllib.parse.urlencode( { "keywordSearch": keyword, + "startIndex": start_index, "resultsPerPage": results_per_page, } ) @@ -157,58 +148,74 @@ def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[d or parsed_url.hostname != "services.nvd.nist.gov" or parsed_url.port not in (None, 443) ): - logger.error("Refusing request to an untrusted NVD endpoint") - return [] + raise NvdRequestError("Refusing request to an untrusted NVD endpoint") + last_error: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: _wait_for_rate_limit() - logger.debug("NVD query (attempt %d): %s", attempt, keyword) - req = urllib.request.Request( url, headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, ) - # URL host is the hardcoded NVD API base, not user-controlled with NVD_REQUEST_LATENCY_SECONDS.time(): - # The URL is built from the fixed HTTPS NVD endpoint; only its query string varies. - with urllib.request.urlopen( # nosec B310 # nosemgrep: python.lang.security.audit.dynamic-urllib-use-detected.dynamic-urllib-use-detected # noqa: E501 + with urllib.request.urlopen( # nosec B310 # noqa: E501 req, timeout=10 ) as resp: - data = json.loads(resp.read()) - - vulnerabilities = data.get("vulnerabilities", []) - results = [parsed for item in vulnerabilities if (parsed := _parse_cve_item(item)) is not None] - - _cache[cache_key] = results - logger.info("NVD returned %d CVEs for: %s", len(results), keyword) - return results - - except urllib.error.HTTPError as e: - if e.code == 429: - wait = 30 * attempt # Back off harder each retry - logger.warning( - "NVD rate limited (429). Waiting %ds before retry %d/%d", - wait, - attempt, - _MAX_RETRIES, - ) - time.sleep(wait) - else: - logger.warning("NVD HTTP %d for keyword '%s': %s", e.code, keyword, e) - break # Non-rate-limit HTTP errors won't improve on retry - - except Exception as e: - logger.warning( - "NVD query failed (attempt %d/%d) for '%s': %s", - attempt, - _MAX_RETRIES, - keyword, - e, - ) + return json.loads(resp.read()) + except urllib.error.HTTPError as exc: + last_error = exc + if exc.code != 429: + break + time.sleep(30 * attempt) + except Exception as exc: + last_error = exc if attempt < _MAX_RETRIES: time.sleep(2**attempt) + raise NvdRequestError(f"NVD request failed for {keyword!r}: {last_error}") + - logger.warning("NVD lookup failed for '%s' - returning empty list", keyword) - _cache[cache_key] = [] # Cache the failure to avoid hammering NVD - return [] +def query_nvd_strict(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: + """Return every NVD result, propagating terminal retrieval failure to the job queue.""" + cache_key = f"strict:{keyword}:{results_per_page}" + if cache_key in _cache: + return _cache[cache_key] + + start_index = 0 + total_results: Optional[int] = None + results: list[dict] = [] + while total_results is None or start_index < total_results: + page = _fetch_nvd_page(keyword, start_index, results_per_page) + vulnerabilities = page.get("vulnerabilities", []) + results.extend(parsed for item in vulnerabilities if (parsed := _parse_cve_item(item)) is not None) + total_results = int(page.get("totalResults", len(vulnerabilities))) + if not vulnerabilities: + break + start_index += len(vulnerabilities) + _cache[cache_key] = results + return results + + +def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: + """ + Query NVD for CVEs matching a keyword. + + Returns a list of parsed CVE dicts (may be empty). + Never raises - all failures return []. + + Args: + keyword: Search term, e.g. "Azure Storage Account" + results_per_page: Max CVEs to fetch (default 5) + """ + cache_key = f"{keyword}:{results_per_page}" + if cache_key in _cache: + logger.debug("NVD cache hit for: %s", keyword) + return _cache[cache_key] + + try: + results = query_nvd_strict(keyword, results_per_page) + except NvdRequestError as exc: + logger.warning("NVD lookup failed for '%s': %s", keyword, exc) + results = [] + _cache[cache_key] = results + return results diff --git a/scanner/worker.py b/scanner/worker.py index f2f84b4a..0a3404cf 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -22,6 +22,7 @@ init_sentry, ) from scanner.engine import ScanEngine +from scanner.enrichment_worker import process_enrichment_job configure_logging() logger = logging.getLogger("scanner.worker") @@ -144,11 +145,18 @@ def run_worker(): try: # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() + db.recover_stale_enrichment_jobs() # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) - # 3. Atomic claim + # 3. Run one durable enrichment job before taking another scan. + enrichment_job = db.claim_next_enrichment_job(worker_id, lease_seconds) + if enrichment_job: + process_enrichment_job(db, enrichment_job, worker_id, lease_seconds) + continue + + # 4. Atomic scan claim scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: time.sleep(POLL_INTERVAL_SECONDS) diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py new file mode 100644 index 00000000..ff659efe --- /dev/null +++ b/tests/test_enrichment_jobs_postgres.py @@ -0,0 +1,186 @@ +"""PostgreSQL lease, retry, and resume tests for durable enrichment jobs.""" + +import os +import threading +import uuid +from unittest.mock import patch + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager, LostLease +from scanner.enrichment_worker import process_enrichment_job + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +@pytest.fixture +def enrichment_scan(): + dsn = os.environ["DATABASE_URL"] + scan_id, subscription_id = str(uuid.uuid4()), str(uuid.uuid4()) + db = DatabaseManager(dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + claim = db.claim_next_pending_scan("seed", 120) + result = { + "scan_id": scan_id, + "subscription_id": subscription_id, + "findings": [ + { + "rule_id": "AZ-STOR-001", + "rule_name": "test", + "severity": "HIGH", + "resource_id": f"/subscriptions/{subscription_id}/resources/one", + "resource_name": "one", + "resource_type": "Test/resource", + "detected_at": "2026-08-29T00:00:00+00:00", + }, + { + "rule_id": "AZ-STOR-001", + "rule_name": "test", + "severity": "HIGH", + "resource_id": f"/subscriptions/{subscription_id}/resources/two", + "resource_name": "two", + "resource_type": "Test/resource", + "detected_at": "2026-08-29T00:00:00+00:00", + }, + ], + } + db.save_scan(result, "seed", claim["fencing_token"]) + job, _ = db.enqueue_enrichment_job(scan_id) + yield dsn, scan_id, job + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def _claim(dsn): + db = DatabaseManager(dsn) + try: + return db.claim_next_enrichment_job("worker-a", 120) + finally: + db.close() + + +def test_duplicate_enqueue_and_claim_race(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + replay, created = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + assert created is False + assert replay["job_id"] == first_job["job_id"] + + barrier = threading.Barrier(2) + claims = [] + + def claim(): + barrier.wait() + claims.append(_claim(dsn)) + + threads = [threading.Thread(target=claim) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert len([job for job in claims if job]) == 1 + + +def test_checkpoint_resume_and_completion(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + job = _claim(dsn) + assert job is not None + db = DatabaseManager(dsn) + try: + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + calls = 0 + + def enrich_once_then_fail(finding): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("transient NVD failure") + return {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + + enrich.side_effect = enrich_once_then_fail + assert process_enrichment_job(db, job, "worker-a", 120) == "retry" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status, checkpoint FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone() == ("pending", 1) + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", (scan_id,) + ) + resumed = db.claim_next_enrichment_job("worker-b", 120) + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + assert process_enrichment_job(db, resumed, "worker-b", 120) == "completed" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status, checkpoint FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone() == ("completed", 2) + cur.execute( + "SELECT COUNT(*) FROM findings WHERE scan_id = %s AND cve_references <> '[]'::jsonb", (scan_id,) + ) + assert cur.fetchone()[0] == 2 + finally: + db.close() + + +def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + for attempt in range(1, 4): + job = db.claim_next_enrichment_job("worker-a", 120) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): + expected = "failed" if attempt == 3 else "retry" + assert process_enrichment_job(db, job, "worker-a", 120) == expected + if attempt < 3: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", + (scan_id,), + ) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == "failed" + finally: + db.close() + + +def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + first = _claim(dsn) + assert first is not None + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' + WHERE scan_id = %s + """, + (scan_id,), + ) + db = DatabaseManager(dsn) + try: + assert db.recover_stale_enrichment_jobs() == 1 + second = db.claim_next_enrichment_job("worker-b", 120) + assert second["fencing_token"] > first["fencing_token"] + with pytest.raises(LostLease): + db.heartbeat_enrichment_job(str(first["job_id"]), "worker-a", first["fencing_token"], 120) + finally: + db.close() diff --git a/tests/test_nvd_client.py b/tests/test_nvd_client.py index 03ce424a..60100668 100644 --- a/tests/test_nvd_client.py +++ b/tests/test_nvd_client.py @@ -220,6 +220,18 @@ def test_second_call_uses_cache(self, mock_wait, mock_urlopen): query_nvd("Azure Storage Account") # Should be served from cache self.assertEqual(mock_urlopen.call_count, 1) + @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client._wait_for_rate_limit") + def test_fetches_every_nvd_page(self, mock_wait, mock_urlopen): + first_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][0]]} + second_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][1]]} + mock_urlopen.side_effect = [_make_mock_urlopen_response(first_page), _make_mock_urlopen_response(second_page)] + + results = query_nvd("Azure Storage Account", results_per_page=1) + + self.assertEqual([item["cve_id"] for item in results], ["CVE-2023-12345", "CVE-2022-99999"]) + self.assertEqual(mock_urlopen.call_count, 2) + @patch("scanner.nvd_client.urllib.request.urlopen") @patch("scanner.nvd_client._wait_for_rate_limit") def test_returns_empty_list_on_network_error(self, mock_wait, mock_urlopen): diff --git a/tests/test_scans_enrich.py b/tests/test_scans_enrich.py index 44c42c82..80addb6d 100644 --- a/tests/test_scans_enrich.py +++ b/tests/test_scans_enrich.py @@ -1,59 +1,41 @@ -"""Tests for POST /api/scans//enrich backgrounding (REL-004).""" +"""Tests for durable POST /api/scans//enrich job admission.""" from unittest.mock import MagicMock, patch import api.routes.scans as scans_route + _SCAN_ID = "00000000-0000-0000-0000-000000000001" def _mock_db(current_scan=None, findings=None): db = MagicMock() - db.get_scans.return_value = [current_scan] if current_scan else [] + db.get_scan.return_value = current_scan db.get_findings.return_value = findings if findings is not None else [] return db -class _FakeThread: - """Records what would have been threaded, and runs synchronously on start().""" - - last_instance = None - - def __init__(self, target, args, daemon): - self.target = target - self.args = args - self.daemon = daemon - self.started = False - _FakeThread.last_instance = self - - def start(self): - self.started = True - - -def test_enrich_returns_202_and_schedules_background_thread(client, auth_headers, monkeypatch): - monkeypatch.setenv("DATABASE_URL", "postgresql://mock/mock") +def test_enrich_returns_202_and_enqueues_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} - findings = [{"id": 1, "rule_id": "AZ-STOR-001"}] - db = _mock_db(current_scan=scan, findings=findings) + db = _mock_db(current_scan=scan, findings=[{"id": 1, "rule_id": "AZ-STOR-001"}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, True) - with ( - patch.object(scans_route, "_get_db", return_value=db), - patch.object(scans_route.threading, "Thread", _FakeThread), - ): + with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - body = resp.get_json() - assert body["status"] == "ENRICHING" + assert resp.get_json()["status"] == "PENDING" + db.enqueue_enrichment_job.assert_called_once_with(_SCAN_ID) - db.update_scan_enrichment_status.assert_called_once_with(_SCAN_ID, "ENRICHING") - thread = _FakeThread.last_instance - assert thread.started is True - assert thread.daemon is True - assert thread.target is scans_route._run_enrichment_in_background - assert thread.args[0] == _SCAN_ID - assert thread.args[1] == findings +def test_enrich_reuses_existing_durable_job(client, auth_headers): + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, False) + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 202 + assert resp.get_json()["status"] == "running" def test_enrich_already_completed_returns_200(client, auth_headers): @@ -65,82 +47,14 @@ def test_enrich_already_completed_returns_200(client, auth_headers): assert "already enriched" in resp.get_json()["message"] -def test_enrich_already_in_progress_returns_202(client, auth_headers): - scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "ENRICHING"} - db = _mock_db(current_scan=scan) - with patch.object(scans_route, "_get_db", return_value=db): - resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) - assert resp.status_code == 202 - assert "in progress" in resp.get_json()["message"] - - def test_enrich_missing_scan_returns_404(client, auth_headers): - db = _mock_db(current_scan=None) - with patch.object(scans_route, "_get_db", return_value=db): + with patch.object(scans_route, "_get_db", return_value=_mock_db(current_scan=None)): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 404 def test_enrich_no_findings_returns_404(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} - db = _mock_db(current_scan=scan, findings=[]) - with patch.object(scans_route, "_get_db", return_value=db): + with patch.object(scans_route, "_get_db", return_value=_mock_db(current_scan=scan, findings=[])): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 404 - - -def test_background_enrichment_marks_completed_on_success(): - fake_db = MagicMock() - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", return_value=[{"id": 1}]), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_cve_fields.assert_called_once_with([{"id": 1}]) - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "COMPLETED") - fake_db.close.assert_called_once() - - -def test_background_enrichment_marks_failed_on_error(): - fake_db = MagicMock() - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") - fake_db.close.assert_called_once() - - -def test_background_enrichment_rolls_back_before_marking_failed(): - """A write failure (e.g. in update_cve_fields) can leave db.conn in an - aborted-transaction state. Without a rollback first, the follow-up - update_scan_enrichment_status(FAILED) call would itself raise - InFailedSqlTransaction and get silently swallowed, leaving the scan - stuck at ENRICHING forever.""" - fake_db = MagicMock() - fake_db.conn = MagicMock() # an open connection, left mid-transaction - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.conn.rollback.assert_called_once() - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") - - -def test_background_enrichment_skips_rollback_when_never_connected(): - """If enrich_findings() fails before any DB write, db.conn is still None - — rollback() must not be called on it.""" - fake_db = MagicMock() - fake_db.conn = None - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") diff --git a/tests/test_worker.py b/tests/test_worker.py index 73b051df..f973b15f 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -71,6 +71,7 @@ def test_worker_processes_pending_scan_successfully( # We need to stop the infinite loop. We'll raise StopWorker on the second call to recover_stale_scans. mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, @@ -111,6 +112,7 @@ def test_worker_handles_scan_failure_gracefully( mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, @@ -147,6 +149,7 @@ def test_worker_does_not_persist_after_heartbeat_reports_lost_lease( mock_env.return_value = self.mock_db_url mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.return_value = { "scan_id": self.scan_id, "subscription_id": self.subscription_id, @@ -199,6 +202,7 @@ def test_worker_sleeps_when_no_scans_pending(self, mock_sleep, mock_env, mock_db mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.return_value = None with self.assertRaises(StopWorker): From 61a9526322fd099dda462ca466024c76a7b05d98 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:41:13 +0100 Subject: [PATCH 06/19] feat(core): expose durable worker metrics Signed-off-by: Shaurya K Sharma --- ...d4a8c1e6b2f9_operational_worker_metrics.py | 38 +++++++++ api/models/finding.py | 84 +++++++++++++++++++ api/observability.py | 47 +++++++++++ scanner/worker.py | 4 + tests/test_observability.py | 1 + tests/test_operational_metrics_postgres.py | 39 +++++++++ 6 files changed, 213 insertions(+) create mode 100644 alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py create mode 100644 tests/test_operational_metrics_postgres.py diff --git a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py new file mode 100644 index 00000000..66c2262b --- /dev/null +++ b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py @@ -0,0 +1,38 @@ +"""Persist worker liveness used by bounded operational metrics. + +Revision ID: d4a8c1e6b2f9 +Revises: c9e1a5b7d3f2 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "d4a8c1e6b2f9" +down_revision: Union[str, Sequence[str], None] = "c9e1a5b7d3f2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Store one liveness timestamp per worker process.""" + op.create_table( + "worker_heartbeats", + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("worker_type", sa.Text(), nullable=False), + sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("worker_id", "worker_type", name="worker_heartbeats_pkey"), + sa.CheckConstraint("worker_type IN ('scan', 'enrichment')", name="ck_worker_heartbeats_type"), + ) + op.create_index( + "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False + ) + + +def downgrade() -> None: + """Remove durable worker heartbeat state.""" + op.drop_index("idx_worker_heartbeats_type_seen", table_name="worker_heartbeats") + op.drop_table("worker_heartbeats") diff --git a/api/models/finding.py b/api/models/finding.py index 9d62a22e..8b82c645 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1135,6 +1135,90 @@ def get_scans(self) -> List[Dict[str, Any]]: cur.execute("SELECT * FROM scans ORDER BY started_at DESC LIMIT 100") return [dict(row) for row in cur.fetchall()] + def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: + """Persist liveness without using worker IDs as metric labels.""" + if worker_type not in {"scan", "enrichment"}: + raise ValueError("unsupported worker type") + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO worker_heartbeats (worker_id, worker_type, last_seen_at) + VALUES (%s, %s, CURRENT_TIMESTAMP) + ON CONFLICT (worker_id, worker_type) DO UPDATE SET + last_seen_at = EXCLUDED.last_seen_at + """, + (worker_id, worker_type), + ) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def get_operational_metrics(self) -> Dict[str, Any]: + """Return bounded aggregates used by the public Prometheus endpoint.""" + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(started_at)), 0) AS value + FROM scans WHERE status = 'pending' + """ + ) + scan_queue_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at)), 0) AS value + FROM enrichment_jobs WHERE status = 'pending' + """ + ) + enrichment_queue_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(claimed_at)), 0) AS value + FROM scans WHERE status = 'running' + """ + ) + scan_lease_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(last_heartbeat_at)), 0) AS value + FROM enrichment_jobs WHERE status = 'running' + """ + ) + enrichment_lease_age = cur.fetchone()["value"] + cur.execute("SELECT COALESCE(SUM(GREATEST(attempt_count - 1, 0)), 0) AS value FROM scans") + scan_retries = cur.fetchone()["value"] + cur.execute("SELECT COALESCE(SUM(GREATEST(attempt_count - 1, 0)), 0) AS value FROM enrichment_jobs") + enrichment_retries = cur.fetchone()["value"] + cur.execute( + """ + SELECT worker_type, EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MAX(last_seen_at)) AS age + FROM worker_heartbeats GROUP BY worker_type + """ + ) + heartbeat_age = {row["worker_type"]: float(row["age"]) for row in cur.fetchall()} + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM MAX(completed_at)), 0) AS value + FROM scans WHERE status = 'completed' + """ + ) + last_success = cur.fetchone()["value"] + conn.commit() + return { + "oldest_queue_age": {"scan": float(scan_queue_age), "enrichment": float(enrichment_queue_age)}, + "oldest_lease_age": {"scan": float(scan_lease_age), "enrichment": float(enrichment_lease_age)}, + "retry_attempts": {"scan": float(scan_retries), "enrichment": float(enrichment_retries)}, + "worker_heartbeat_age": heartbeat_age, + "last_successful_scan_timestamp": float(last_success), + } + except Exception: + self.rollback(conn) + raise + # ------------------------------------------------------------------ # # Scoring # # ------------------------------------------------------------------ # diff --git a/api/observability.py b/api/observability.py index 0d00f06c..19f2263c 100644 --- a/api/observability.py +++ b/api/observability.py @@ -72,6 +72,30 @@ "openshield_pending_scans", "Number of scans currently waiting in the pending queue.", ) +WORKER_HEARTBEAT_AGE_SECONDS = Gauge( + "openshield_worker_heartbeat_age_seconds", + "Seconds since the most recent durable worker heartbeat.", + ["worker_type"], +) +OLDEST_QUEUE_AGE_SECONDS = Gauge( + "openshield_oldest_queue_age_seconds", + "Age in seconds of the oldest pending durable work item.", + ["queue"], +) +OLDEST_ACTIVE_LEASE_AGE_SECONDS = Gauge( + "openshield_oldest_active_lease_age_seconds", + "Age in seconds of the oldest active durable lease.", + ["queue"], +) +RETRY_ATTEMPTS = Gauge( + "openshield_retry_attempts", + "Aggregate retry attempts currently recorded for durable work.", + ["queue"], +) +LAST_SUCCESSFUL_SCAN_TIMESTAMP = Gauge( + "openshield_last_successful_scan_timestamp_seconds", + "Unix timestamp of the most recently completed scan, or zero when none exist.", +) RULE_ERRORS_TOTAL = Counter( "openshield_rule_errors_total", "Total number of times a scanner rule raised an exception.", @@ -387,4 +411,27 @@ def _record_observability(response: Response) -> Response: @probe_rate_limit(_METRICS_MAX_REQUESTS_PER_WINDOW) def metrics() -> Response: _refresh_pool_metrics() + # Metrics are derived from durable state so the API can expose worker + # liveness even when scan workers run in separate processes. + db_url = os.environ.get("DATABASE_URL") + if db_url: + try: + from api.models.finding import DatabaseManager + + db = DatabaseManager(db_url) + try: + snapshot = db.get_operational_metrics() + finally: + db.close() + for queue in ("scan", "enrichment"): + OLDEST_QUEUE_AGE_SECONDS.labels(queue=queue).set(snapshot["oldest_queue_age"].get(queue, 0)) + OLDEST_ACTIVE_LEASE_AGE_SECONDS.labels(queue=queue).set(snapshot["oldest_lease_age"].get(queue, 0)) + RETRY_ATTEMPTS.labels(queue=queue).set(snapshot["retry_attempts"].get(queue, 0)) + for worker_type in ("scan", "enrichment"): + WORKER_HEARTBEAT_AGE_SECONDS.labels(worker_type=worker_type).set( + snapshot["worker_heartbeat_age"].get(worker_type, 0) + ) + LAST_SUCCESSFUL_SCAN_TIMESTAMP.set(snapshot["last_successful_scan_timestamp"]) + except Exception as exc: + logger.warning("Unable to refresh durable operational metrics: %s", exc) return Response(generate_latest(), content_type=CONTENT_TYPE_LATEST) diff --git a/scanner/worker.py b/scanner/worker.py index 0a3404cf..b5021e11 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -143,6 +143,10 @@ def run_worker(): while True: try: + db.record_worker_heartbeat(worker_id, "scan") + # This process executes both durable queue types; record both + # liveness signals without exporting the worker UUID as a label. + db.record_worker_heartbeat(worker_id, "enrichment") # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() db.recover_stale_enrichment_jobs() diff --git a/tests/test_observability.py b/tests/test_observability.py index ef9207ba..2745e940 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -188,6 +188,7 @@ class _Stop(BaseException): mock_db = MagicMock() mock_db.recover_stale_scans.side_effect = [None, _Stop()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": scan_id, "subscription_id": sub_id, "fencing_token": 1}, None, diff --git a/tests/test_operational_metrics_postgres.py b/tests/test_operational_metrics_postgres.py new file mode 100644 index 00000000..f4f63c7c --- /dev/null +++ b/tests/test_operational_metrics_postgres.py @@ -0,0 +1,39 @@ +"""PostgreSQL coverage for durable operational metric aggregates.""" + +import os +import uuid + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): + dsn = os.environ["DATABASE_URL"] + scan_id, subscription_id = str(uuid.uuid4()), str(uuid.uuid4()) + db = DatabaseManager(dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + db.record_worker_heartbeat("worker-test", "scan") + db.record_worker_heartbeat("worker-test", "enrichment") + claim = db.claim_next_pending_scan("worker-a", 120) + assert claim is not None + snapshot = db.get_operational_metrics() + assert snapshot["oldest_lease_age"]["scan"] >= 0 + assert snapshot["retry_attempts"]["scan"] == 0 + assert snapshot["worker_heartbeat_age"]["scan"] >= 0 + assert snapshot["worker_heartbeat_age"]["enrichment"] >= 0 + assert snapshot["last_successful_scan_timestamp"] >= 0 + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM worker_heartbeats WHERE worker_id = 'worker-test'") + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) From 5f38635727a2c74854dd5ff2a9fefe02dbada7e9 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:01:53 +0100 Subject: [PATCH 07/19] fix(core): complete scan durability integration Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 12 ++++++++++++ docs/async-scan-architecture.md | 25 ++++++++++++++++++++++-- docs/cve_correlation_feature.md | 7 ++++--- tests/test_enrichment_jobs_postgres.py | 3 ++- tests/test_scan_leases_postgres.py | 1 + tests/test_subscription_authorization.py | 7 ++++++- 6 files changed, 48 insertions(+), 7 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index 8b82c645..39b5eac1 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -478,6 +478,18 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token ) else: cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_result["scan_id"],)) + + # Completion creates one durable job in this same fenced + # transaction. The scan_id uniqueness constraint makes result + # replay harmless and prevents duplicate enrichment delivery. + cur.execute( + """ + INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint) + VALUES (%s, %s, 'pending', 0, 0) + ON CONFLICT (scan_id) DO NOTHING + """, + (str(uuid.uuid4()), scan_result["scan_id"]), + ) conn.commit() except Exception: # psycopg2 connections remain in an aborted transaction after any diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index e9cf84f2..757aec61 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -13,7 +13,7 @@ In the legacy synchronous model, POST /api/scans/trigger would block the HTTP re OpenShield now employs a decoupled, database backed worker architecture. This is the industry standard for long running security tasks where reliability and state persistence are critical. ### 1. The API (Flask) -When a scan is triggered, the API performs minimal work. It validates the subscription_id, creates a record in the scans table with status set to pending, and returns 202 Accepted and the scan_id immediately. +When a scan is triggered, the API validates the subscription and creates a durable pending record. PostgreSQL permits at most one `pending` or `running` scan per subscription. `Idempotency-Key` replays return the same logical scan when the request fingerprint matches; reuse with different semantics returns a conflict. The optional `OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR` policy enables an explicit time-window quota. A zero/unset value preserves the current no-business-limit policy while the one-active-scan concurrency quota remains enforced. ### 2. The Queue (PostgreSQL) The scans table acts as a persistent task queue. This avoids the need for additional infrastructure like Redis or RabbitMQ while providing ACID compliance, visibility, and auditability. Scan states are never lost during crashes, status polling is a simple SQL query, and every scan has a persistent record of its error state. @@ -30,8 +30,22 @@ worker cannot persist completion, failure, or findings after it loses ownership. Once a scan reaches the maximum attempt count, it is marked `failed` so bad credentials or persistent Azure errors cannot retry forever. +Findings use a stable database-enforced identity (`scan`, rule, canonical +resource scope, and an optional rule-specific discriminator). Result retries +use PostgreSQL upserts, so mutable text or severity is updated rather than +creating a second authoritative finding. Rule evaluations use the same +database-first uniqueness model. + ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and marks the scan complete atomically. On failure, it records a sanitized error only while it still owns the lease. +The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings, evaluations, and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. + +### Durable CVE enrichment + +`POST /api/scans//enrich` enqueues (or returns) the one durable PostgreSQL enrichment job for the scan; it never starts a request-owned daemon thread. The scan worker claims those jobs with the same owner/expiry/fencing model, checkpoints after each finding, and retries transient failures with bounded exponential backoff. An expired job is recovered or terminally failed after its attempt limit. NVD retrieval follows every `totalResults` page; replaying a checkpoint updates the existing finding instead of duplicating CVE data. + +### Operational signals + +`/metrics` derives bounded-cardinality operational gauges from PostgreSQL: worker heartbeat age, oldest queue age, oldest active lease age, aggregate retry attempts, and the last successful scan timestamp. Labels are limited to `queue` (`scan` or `enrichment`) and `worker_type`; scan, job, subscription, and worker identifiers are never metric labels. ## Technical Rationale @@ -41,6 +55,13 @@ While Celery is powerful, it introduces external dependencies and operational co ### Why not Threading Python background threads are ephemeral. If the web server process restarts, all in flight scans are killed instantly and marked as running forever in the DB. A separate worker process ensures that the scan lifecycle is independent of the web server lifecycle. +## Deployment order + +This release is not safe for mixed old and new scan workers. Stop or drain old +workers, apply Alembic migrations, then start the fenced worker version. Legacy +workers do not carry ownership/fencing state; legacy running scans are retained +and made recoverable by the lease migration rather than deleted. + ## Testing Suite The asynchronous transition is verified through a multi layered testing strategy. diff --git a/docs/cve_correlation_feature.md b/docs/cve_correlation_feature.md index a0820815..eac8ddec 100644 --- a/docs/cve_correlation_feature.md +++ b/docs/cve_correlation_feature.md @@ -19,7 +19,8 @@ The CVE Correlation feature integrates the MITRE National Vulnerability Database | File | Change | Why | |---|---|---| | scanner/engine.py | Decoupled Scan. Removed synchronous enrichment from the scan lifecycle. | Performance: Azure scans now return immediately without waiting for NVD rate limits (7s per resource type). | -| api/routes/scans.py | New Endpoint. Added `POST /api/scans//enrich`. | Flexibility: CVE enrichment can now be triggered on-demand or by a background job after the scan completes. | +| api/routes/scans.py | Durable endpoint. `POST /api/scans//enrich` enqueues or returns a PostgreSQL job. | No Gunicorn daemon thread owns authoritative work. | +| scanner/enrichment_worker.py | Checkpointed durable enrichment execution. | Reclaims expired work safely and resumes at a finding checkpoint. | | api/models/finding.py | Updated Scan model and added enrichment status tracking. | Persistence: Adds `cve_enrichment_status` to track `PENDING`, `COMPLETED`, or `FAILED` states. | | alembic/versions/ | Defines CVE columns in the versioned database schema. | Deployment: Alembic owns schema changes independently of Flask application startup. | | api/routes/score.py | Added GET /api/score/cve-summary endpoint. | Dashboard UI: Provides the frontend with high-level data like Total Known Exploits and enrichment status. | @@ -30,10 +31,10 @@ The CVE Correlation feature integrates the MITRE National Vulnerability Database To ensure the frontend dashboard works perfectly, the architecture uses a Decoupled Enrichment model: 1. Fast Dashboard Loads: The scan engine completes rapidly. The dashboard can check the enrichment status of the latest scan. -2. Manual/Job Enrichment: A "Trigger Enrichment" button or a background task calls `POST /api/scans//enrich` to populate CVE data. +2. Durable Enrichment: Completion creates one job; `POST /api/scans//enrich` returns that job idempotently. The worker claims, retries, and checkpoints it in PostgreSQL. 3. Dashboard-Ready Summary Endpoint: The /api/score/cve-summary endpoint includes the `status` field, allowing the UI to show a "Scan Enriched" badge or a "Pending" spinner. 4. Actionable Risk (CISA KEV): The exploit_available flag uses the CISA Known Exploited Vulnerabilities catalogue, allowing the dashboard to highlight high-priority risks that are being exploited in the wild. -5. Persistent Historical State: Enrichment happens at the time of the enrichment call, and the result is persisted. +5. Persistent Historical State: Enrichment checkpoint state and results survive API/worker restarts. NVD pagination continues through every `totalResults` page. ## Security and Compliance Audit diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index ff659efe..b98b73fe 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -50,7 +50,8 @@ def enrichment_scan(): ], } db.save_scan(result, "seed", claim["fencing_token"]) - job, _ = db.enqueue_enrichment_job(scan_id) + job, created = db.enqueue_enrichment_job(scan_id) + assert created is False yield dsn, scan_id, job finally: db.close() diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index 7b2fa7df..fd5df2f2 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -105,6 +105,7 @@ def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) diff --git a/tests/test_subscription_authorization.py b/tests/test_subscription_authorization.py index d79f21a2..640d9686 100644 --- a/tests/test_subscription_authorization.py +++ b/tests/test_subscription_authorization.py @@ -14,7 +14,12 @@ def _trigger(client, auth_headers, subscription_id, monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db") - with patch("api.routes.scans.DatabaseManager", return_value=MagicMock()): + database = MagicMock() + database.admit_scan.return_value = ( + {"scan_id": "test-scan", "status": "pending"}, + True, + ) + with patch("api.routes.scans.DatabaseManager", return_value=database): return client.post( "/api/scans/trigger", json={"subscription_id": subscription_id}, From 9f8701c900cf7a330da0a92a52b3ea0553f2e25a Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:09:16 +0100 Subject: [PATCH 08/19] fix(api): avoid exposing scan admission errors Signed-off-by: Shaurya K Sharma --- api/routes/scans.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/api/routes/scans.py b/api/routes/scans.py index ec90033a..a84d233a 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -136,10 +136,10 @@ def trigger_scan(): request_fingerprint=request_fingerprint, max_scans_per_hour=_configured_hourly_quota(), ) - except ScanAdmissionConflict as exc: - return jsonify({"error": str(exc)}), 409 - except ScanQuotaExceeded as exc: - return jsonify({"error": str(exc)}), 429 + except ScanAdmissionConflict: + return jsonify({"error": "Idempotency-Key is already associated with a different request."}), 409 + except ScanQuotaExceeded: + return jsonify({"error": "Scan quota exceeded for this subscription."}), 429 except Exception as exc: logger.error("Failed to create pending scan: %s", exc, exc_info=True) return jsonify({"error": "Database error"}), 500 From 27ca5ced6c4f4a3ad3898662a9775ea780f12c26 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:32:32 +0100 Subject: [PATCH 09/19] fix(scanner): use safe NVD request transport Signed-off-by: Shaurya K Sharma --- scanner/nvd_client.py | 50 ++++++++++------------- tests/test_nvd_client.py | 87 ++++++++++++++++++++-------------------- 2 files changed, 65 insertions(+), 72 deletions(-) diff --git a/scanner/nvd_client.py b/scanner/nvd_client.py index 8ca65666..8cefbcc1 100644 --- a/scanner/nvd_client.py +++ b/scanner/nvd_client.py @@ -19,12 +19,10 @@ import threading import time import logging -import urllib.request -import urllib.error -import urllib.parse -import json from typing import Optional +import requests + from api.observability import NVD_REQUEST_LATENCY_SECONDS logger = logging.getLogger(__name__) @@ -134,40 +132,34 @@ def _parse_cve_item(item: dict) -> Optional[dict]: def _fetch_nvd_page(keyword: str, start_index: int, results_per_page: int) -> dict: """Fetch one NVD page, raising only after bounded retries are exhausted.""" - params = urllib.parse.urlencode( - { - "keywordSearch": keyword, - "startIndex": start_index, - "resultsPerPage": results_per_page, - } - ) - url = f"{_NVD_BASE_URL}?{params}" - parsed_url = urllib.parse.urlsplit(url) - if ( - parsed_url.scheme != "https" - or parsed_url.hostname != "services.nvd.nist.gov" - or parsed_url.port not in (None, 443) - ): - raise NvdRequestError("Refusing request to an untrusted NVD endpoint") + params = { + "keywordSearch": keyword, + "startIndex": start_index, + "resultsPerPage": results_per_page, + } last_error: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: _wait_for_rate_limit() - req = urllib.request.Request( - url, - headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, - ) with NVD_REQUEST_LATENCY_SECONDS.time(): - with urllib.request.urlopen( # nosec B310 # noqa: E501 - req, timeout=10 - ) as resp: - return json.loads(resp.read()) - except urllib.error.HTTPError as exc: + response = requests.get( + _NVD_BASE_URL, + params=params, + headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, + timeout=10, + ) + response.raise_for_status() + return response.json() + except requests.HTTPError as exc: last_error = exc - if exc.code != 429: + if exc.response is None or exc.response.status_code != 429: break time.sleep(30 * attempt) + except requests.RequestException as exc: + last_error = exc + if attempt < _MAX_RETRIES: + time.sleep(2**attempt) except Exception as exc: last_error = exc if attempt < _MAX_RETRIES: diff --git a/tests/test_nvd_client.py b/tests/test_nvd_client.py index 60100668..06818b63 100644 --- a/tests/test_nvd_client.py +++ b/tests/test_nvd_client.py @@ -12,13 +12,13 @@ TestQueryNvd - query_nvd() HTTP behaviour (mocked urlopen) """ -import json import threading import time import unittest -import urllib.error from unittest.mock import patch, MagicMock +import requests + # Clear the module cache before import so previous test runs don't bleed in from scanner.nvd_client import query_nvd, _parse_cve_item, _cache, _wait_for_rate_limit @@ -68,21 +68,17 @@ _EMPTY_NVD_RESPONSE = {"vulnerabilities": []} -def _make_mock_urlopen_response(data: dict) -> MagicMock: +def _make_mock_requests_response(data: dict, status: int = 200) -> MagicMock: """ - Return a MagicMock that behaves like urllib.request.urlopen()'s - context manager return value. - - urlopen() is used as: - with urllib.request.urlopen(req, timeout=10) as resp: - data = json.loads(resp.read()) + Return a MagicMock that behaves like a requests NVD response. - So the mock needs __enter__/__exit__ and a .read() method. + `raise_for_status()` raises the same requests exception the client handles. """ mock_resp = MagicMock() - mock_resp.read.return_value = json.dumps(data).encode("utf-8") - mock_resp.__enter__ = lambda s: s - mock_resp.__exit__ = MagicMock(return_value=False) + mock_resp.json.return_value = data + mock_resp.status_code = status + if status >= 400: + mock_resp.raise_for_status.side_effect = requests.HTTPError(response=mock_resp) return mock_resp @@ -175,14 +171,14 @@ def test_falls_back_to_cvss_v2_when_v31_absent(self): # --------------------------------------------------------------------------- # TestQueryNvd -# Tests for query_nvd() - mocks urllib.request.urlopen to prevent live calls. +# Tests for query_nvd() - mocks requests.get to prevent live calls. # Also mocks _wait_for_rate_limit to keep tests fast. # --------------------------------------------------------------------------- class TestQueryNvd(unittest.TestCase): """ - query_nvd() builds a URL, calls urlopen, parses the response, caches it, + query_nvd() uses the fixed NVD HTTPS URL, parses the response, caches it, and handles errors gracefully. All HTTP is mocked. """ @@ -190,77 +186,82 @@ def setUp(self): """Clear the module-level cache before each test.""" _cache.clear() - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_parsed_cves_on_success(self, mock_wait, mock_urlopen): + def test_returns_parsed_cves_on_success(self, mock_wait, mock_get): """Successful response is parsed into a list of CVE dicts.""" - mock_urlopen.return_value = _make_mock_urlopen_response(_SAMPLE_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_SAMPLE_NVD_RESPONSE) results = query_nvd("Azure Storage Account") self.assertEqual(len(results), 2) self.assertEqual(results[0]["cve_id"], "CVE-2023-12345") self.assertEqual(results[1]["cve_id"], "CVE-2022-99999") + mock_get.assert_called_once_with( + "https://services.nvd.nist.gov/rest/json/cves/2.0", + params={"keywordSearch": "Azure Storage Account", "startIndex": 0, "resultsPerPage": 2000}, + headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, + timeout=10, + ) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_empty_nvd_response(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_empty_nvd_response(self, mock_wait, mock_get): """An empty vulnerabilities list returns [] without error.""" - mock_urlopen.return_value = _make_mock_urlopen_response(_EMPTY_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_EMPTY_NVD_RESPONSE) results = query_nvd("nonexistent-resource-xyz") self.assertEqual(results, []) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_second_call_uses_cache(self, mock_wait, mock_urlopen): + def test_second_call_uses_cache(self, mock_wait, mock_get): """ - Calling query_nvd twice with the same keyword only hits urlopen once. + Calling query_nvd twice with the same keyword only makes one request. The second call must return from cache without a network request. """ - mock_urlopen.return_value = _make_mock_urlopen_response(_SAMPLE_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_SAMPLE_NVD_RESPONSE) query_nvd("Azure Storage Account") query_nvd("Azure Storage Account") # Should be served from cache - self.assertEqual(mock_urlopen.call_count, 1) + self.assertEqual(mock_get.call_count, 1) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_fetches_every_nvd_page(self, mock_wait, mock_urlopen): + def test_fetches_every_nvd_page(self, mock_wait, mock_get): first_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][0]]} second_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][1]]} - mock_urlopen.side_effect = [_make_mock_urlopen_response(first_page), _make_mock_urlopen_response(second_page)] + mock_get.side_effect = [ + _make_mock_requests_response(first_page), + _make_mock_requests_response(second_page), + ] results = query_nvd("Azure Storage Account", results_per_page=1) self.assertEqual([item["cve_id"] for item in results], ["CVE-2023-12345", "CVE-2022-99999"]) - self.assertEqual(mock_urlopen.call_count, 2) + self.assertEqual(mock_get.call_count, 2) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_network_error(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_network_error(self, mock_wait, mock_get): """A network exception returns [] and does not propagate the error.""" - mock_urlopen.side_effect = Exception("Connection refused") + mock_get.side_effect = requests.ConnectionError("Connection refused") results = query_nvd("Azure Storage Account") self.assertEqual(results, []) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_http_503(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_http_503(self, mock_wait, mock_get): """An HTTP 503 returns [] and does not propagate the error.""" - mock_urlopen.side_effect = urllib.error.HTTPError( - url=None, code=503, msg="Service Unavailable", hdrs=None, fp=None - ) + mock_get.return_value = _make_mock_requests_response({}, status=503) results = query_nvd("Azure Storage Account") self.assertEqual(results, []) @patch("scanner.nvd_client.time.sleep") - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_backs_off_and_retries_on_429(self, mock_wait, mock_urlopen, mock_sleep): + def test_backs_off_and_retries_on_429(self, mock_wait, mock_get, mock_sleep): """ A 429 response triggers a sleep and retry. After MAX_RETRIES 429s, returns [] gracefully. """ - mock_urlopen.side_effect = urllib.error.HTTPError( - url=None, code=429, msg="Too Many Requests", hdrs=None, fp=None - ) + mock_get.return_value = _make_mock_requests_response({}, status=429) results = query_nvd("Azure Storage Account") self.assertEqual(results, []) # time.sleep should have been called (back-off logic) From 4a3fcff3b03c818b85ac85f4dc1aaa59f5f3d010 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Wed, 2 Sep 2026 01:08:47 +0100 Subject: [PATCH 10/19] fix(core): address review of scan durability hardening (#303) Resolves the two must-fix items and the follow-ups raised in review of #325. Integration: leave the #263 evaluation contract to #321 ------------------------------------------------------ This branch created a second `rule_evaluations` table, near-identical to the one PR #321 adds, that no production code ever populated: `run_scan()` emits no `evaluations` key, so the storage, the upsert and its tests described a contract that could not be observed in production. Two `CREATE TABLE rule_evaluations` statements would also have broken whichever of #321/#325 merged second, independently of the Alembic head ordering. Issue #263 owns that contract - the table, the PASS/FAIL/UNKNOWN/ERROR/ NOT_APPLICABLE semantics, engine emission and the compliance-score fix - so this branch drops it entirely and keeps only what #303 asks for: the stable `finding_key` identity and its unique index. `save_scan()` marks where #321's evaluation writes belong, inside the fenced completion transaction, so they inherit the lease/ownership check without re-implementing it. Terminally failed enrichment jobs are recoverable again ------------------------------------------------------- After three failed attempts a job became `failed` and nothing could move it: enqueue used ON CONFLICT DO NOTHING, claim only selected `pending`, and stale recovery only handled expired `running` leases. That regressed the operator retry the old thread-based path gave for free. `enqueue_enrichment_job()` now returns an explicit outcome - `created`, `requeued`, `active` or `completed` - and atomically resets a `failed` job to `pending` with a fresh retry budget. It keeps the same job row, its last error message (audit) and its checkpoint (so the retry resumes), never revives a `completed` job, and never disturbs a live `running` lease. The conditional UPDATE is the whole guard, so concurrent re-POSTs converge on one logical job. Admission migration cannot leave an INVALID index ------------------------------------------------- `CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription` fails, and leaves an unusable index behind, on a deployment that already holds several active scans for one subscription. The migration now preflights, changes nothing, and names the offending subscriptions in an actionable error; deciding which production scan is authoritative stays an operator call, and no scan history is deleted. A retry is safe: an INVALID index from an interrupted build is dropped before rebuilding. PostgreSQL tests no longer race the shared queue ------------------------------------------------ Fixtures called `claim_next_pending_scan()`, which takes the globally oldest pending scan, then persisted against the scan they had just created - so any unrelated pending row made `save_scan()` raise LostLease. Reproduced on a real database: with one older pending scan present, the old fixture claims someone else's row and fails; the new one does not. `claim_next_pending_scan()` and `claim_next_enrichment_job()` take an optional `scan_id` so a caller can claim a known row under identical lease and fencing semantics, and the fixtures use it. Queue-wide `recover_stale_*()` counts are asserted as progression of the test's own row rather than as global totals. Follow-ups ---------- - worker_heartbeats is bounded: rows past WORKER_HEARTBEAT_RETENTION_SECONDS are pruned on the beat that registers a new worker identity - once per process, not on every beat, and never for a live worker. - The scan and enrichment queues alternate one item per loop iteration instead of draining enrichment first, so neither can starve the other. - Added the partial index that keeps /metrics' last-successful-scan lookup from degrading into a sequential scan as history grows. - Documented SCAN_LEASE_SECONDS, SCAN_HEARTBEAT_SECONDS, OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR and WORKER_HEARTBEAT_RETENTION_SECONDS in .env.example. - docs/api-reference.md now documents the real admission and enrichment contracts, including which status codes are actually returned. Tests ----- New PostgreSQL coverage for terminal-failure requeue (explicit requeue, no restart of a completed job, a live lease is not stolen, concurrent requeues converge, a stale token cannot write after reclaim, a requeued job completes), the duplicate-active-scan migration preflight, stale heartbeat pruning, and worker fairness with both queues backlogged. Full suite on postgres:16-alpine: 897 passed, 2 skipped. The one failure, test_vector_store_purity, is a local checkout artifact - it needs a BM25 index this checkout's ai/vectorstore lacks, and passes in a clean worktree. Signed-off-by: Shaurya K Sharma --- .env.example | 20 ++ ...a7c5e9d2f1b4_scan_admission_idempotency.py | 65 ++++- ...d4a8c1e6b2f9_operational_worker_metrics.py | 14 + ...6d8e1a4c9_idempotent_finding_identities.py | 41 +++ ...4c9_idempotent_findings_and_evaluations.py | 66 ----- api/models/finding.py | 138 ++++++++-- api/routes/scans.py | 38 +-- docs/api-reference.md | 85 +++++- docs/async-scan-architecture.md | 43 ++- scanner/worker.py | 25 +- tests/test_enrichment_jobs_postgres.py | 215 +++++++++++++-- tests/test_operational_metrics_postgres.py | 47 +++- .../test_scan_admission_migration_postgres.py | 244 ++++++++++++++++++ tests/test_scan_admission_postgres.py | 1 - tests/test_scan_leases_postgres.py | 85 +++--- tests/test_scans_enrich.py | 38 ++- tests/test_worker.py | 62 +++++ 17 files changed, 1034 insertions(+), 193 deletions(-) create mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py delete mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py create mode 100644 tests/test_scan_admission_migration_postgres.py diff --git a/.env.example b/.env.example index d9163717..4812d812 100644 --- a/.env.example +++ b/.env.example @@ -22,6 +22,26 @@ OIDC_ROLE_MAP= # logs a startup warning when this is unset. See docs/api-reference.md. OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS= +# Optional - durable scan worker tuning. A claim is held for SCAN_LEASE_SECONDS +# and renewed every SCAN_HEARTBEAT_SECONDS; the heartbeat MUST be shorter than +# the lease or the worker loses its own claim mid-scan. A non-numeric or +# non-positive value logs a warning and uses the default; a heartbeat >= the +# lease logs a warning and uses one third of the lease. Defaults: 900 and 300. +SCAN_LEASE_SECONDS=900 +SCAN_HEARTBEAT_SECONDS=300 + +# Optional - explicit per-subscription hourly admission quota. Unset or 0 keeps +# the historical no-time-window policy (a non-numeric value logs a warning and +# disables the quota); one active (pending/running) scan per subscription is +# always enforced regardless of this value. +OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR=0 + +# Optional - how long a retired worker's heartbeat row is kept before it is +# pruned. Worker identities are per-process, so this bounds worker_heartbeats +# across restarts. Must stay far above SCAN_HEARTBEAT_SECONDS so that a live +# worker is never pruned. Default: 604800 (7 days). +WORKER_HEARTBEAT_RETENTION_SECONDS=604800 + # AI providers - add at least one ANTHROPIC_API_KEY= GROQ_API_KEY= diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py index a45eb609..e4daf0a2 100644 --- a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -16,22 +16,75 @@ branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +_ACTIVE_INDEX = "uq_scans_one_active_per_subscription" +_KEY_INDEX = "uq_scans_subscription_idempotency_key" + + +def _assert_one_active_scan_per_subscription() -> None: + """Fail with an actionable error instead of an unusable index. + + ``CREATE UNIQUE INDEX CONCURRENTLY`` on a table that already violates the + constraint fails *and* leaves an INVALID index behind. A deployment that + predates the one-active-scan rule can legitimately hold several + ``pending``/``running`` rows for one subscription, so this checks first and + reports exactly which subscriptions block the upgrade. Choosing which of + those scans is authoritative is an operator decision -- deleting or + completing production scan history automatically is never this migration's + call. + """ + rows = ( + op.get_bind() + .execute( + sa.text( + """ + SELECT subscription_id, COUNT(*) AS active + FROM scans + WHERE status IN ('pending', 'running') + GROUP BY subscription_id + HAVING COUNT(*) > 1 + ORDER BY active DESC, subscription_id + """ + ) + ) + .fetchall() + ) + if not rows: + return + detail = ", ".join(f"{subscription_id} ({active} active)" for subscription_id, active in rows) + raise RuntimeError( + "Cannot enforce one active scan per subscription: " + f"{len(rows)} subscription(s) already have more than one pending/running scan: {detail}. " + "Resolve them first (let the scans finish, or mark the superseded rows " + "'failed'), then re-run this migration. " + "See docs/async-scan-architecture.md for the documented cleanup order." + ) + def upgrade() -> None: """Persist idempotency semantics and prevent more than one active scan.""" op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) op.add_column("scans", sa.Column("request_fingerprint", sa.Text(), nullable=True)) + + # Checked before either index is built so a blocked upgrade leaves the + # schema exactly as it was, with the added columns unused and harmless. + _assert_one_active_scan_per_subscription() + with op.get_context().autocommit_block(): + # An earlier interrupted or failed CONCURRENTLY build leaves an INVALID + # index that cannot serve queries but does occupy the name. Drop both + # names first so retrying this migration is deterministic. + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_KEY_INDEX}") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_ACTIVE_INDEX}") op.execute( - """ - CREATE UNIQUE INDEX CONCURRENTLY uq_scans_subscription_idempotency_key + f""" + CREATE UNIQUE INDEX CONCURRENTLY {_KEY_INDEX} ON scans (subscription_id, idempotency_key) WHERE idempotency_key IS NOT NULL """ ) op.execute( - """ - CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription + f""" + CREATE UNIQUE INDEX CONCURRENTLY {_ACTIVE_INDEX} ON scans (subscription_id) WHERE status IN ('pending', 'running') """ @@ -41,7 +94,7 @@ def upgrade() -> None: def downgrade() -> None: """Remove scan admission metadata and constraints.""" with op.get_context().autocommit_block(): - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_one_active_per_subscription") - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_subscription_idempotency_key") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_ACTIVE_INDEX}") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_KEY_INDEX}") op.drop_column("scans", "request_fingerprint") op.drop_column("scans", "idempotency_key") diff --git a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py index 66c2262b..598877dd 100644 --- a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py +++ b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py @@ -31,8 +31,22 @@ def upgrade() -> None: "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False ) + # /metrics reports the last successful scan on every scrape. Without this + # the aggregate degrades into a sequential scan of the whole scans table as + # scan history grows; the partial index keeps it an index-only lookup. + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_completed_completed_at + ON scans (completed_at DESC) + WHERE status = 'completed' + """ + ) + def downgrade() -> None: """Remove durable worker heartbeat state.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_completed_completed_at") op.drop_index("idx_worker_heartbeats_type_seen", table_name="worker_heartbeats") op.drop_table("worker_heartbeats") diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py new file mode 100644 index 00000000..7c067286 --- /dev/null +++ b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py @@ -0,0 +1,41 @@ +"""Add database-enforced identities for scan findings. + +Revision ID: f2b6d8e1a4c9 +Revises: e4f7a9b2c6d8 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "f2b6d8e1a4c9" +down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_UNIQUE_INDEX = "uq_findings_scan_finding_key" + + +def upgrade() -> None: + """Give every finding a stable identity so replayed results upsert.""" + op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) + # Existing records predate the identity contract. Preserve each record as + # distinct rather than attempting to infer equivalence from mutable text. + op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") + op.alter_column("findings", "finding_key", nullable=False) + + with op.get_context().autocommit_block(): + # A previous interrupted CONCURRENTLY build leaves an unusable index + # behind that would make this statement fail with "already exists". + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_UNIQUE_INDEX}") + op.execute(f"CREATE UNIQUE INDEX CONCURRENTLY {_UNIQUE_INDEX} ON findings (scan_id, finding_key)") + + +def downgrade() -> None: + """Remove the finding identity introduced by this revision.""" + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_UNIQUE_INDEX}") + op.drop_column("findings", "finding_key") diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py deleted file mode 100644 index 7b4522a9..00000000 --- a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Add database-enforced identities for scan results. - -Revision ID: f2b6d8e1a4c9 -Revises: e4f7a9b2c6d8 -Create Date: 2026-08-29 00:00:00.000000 -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects import postgresql - - -revision: str = "f2b6d8e1a4c9" -down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Add stable finding keys and per-resource evaluation rows.""" - op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) - # Existing records predate the identity contract. Preserve each record as - # distinct rather than attempting to infer equivalence from mutable text. - op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") - op.alter_column("findings", "finding_key", nullable=False) - - with op.get_context().autocommit_block(): - op.execute("CREATE UNIQUE INDEX CONCURRENTLY uq_findings_scan_finding_key ON findings (scan_id, finding_key)") - - op.create_table( - "rule_evaluations", - sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("scan_id", postgresql.UUID(), nullable=False), - sa.Column("rule_id", sa.Text(), nullable=False), - sa.Column("resource_id", sa.Text(), nullable=False), - sa.Column("resource_type", sa.Text(), server_default=sa.text("''"), nullable=False), - sa.Column("status", sa.Text(), nullable=False), - sa.Column("reason_code", sa.Text(), nullable=True), - sa.Column("reason", sa.Text(), nullable=True), - sa.Column("evidence", postgresql.JSONB(), server_default=sa.text("'{}'::jsonb"), nullable=True), - sa.Column("finding_id", sa.Integer(), nullable=True), - sa.Column("evaluated_at", sa.DateTime(timezone=True), nullable=False), - sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="rule_evaluations_scan_id_fkey"), - sa.ForeignKeyConstraint( - ["finding_id"], ["findings.id"], name="rule_evaluations_finding_id_fkey", ondelete="SET NULL" - ), - sa.PrimaryKeyConstraint("id", name="rule_evaluations_pkey"), - sa.UniqueConstraint("scan_id", "rule_id", "resource_id", name="uq_rule_evaluations_scan_rule_resource"), - sa.CheckConstraint( - "status IN ('PASS', 'FAIL', 'UNKNOWN', 'ERROR', 'NOT_APPLICABLE')", - name="ck_rule_evaluations_status_v1", - ), - sa.CheckConstraint("resource_id <> ''", name="ck_rule_evaluations_resource_id_not_empty"), - ) - op.create_index("idx_rule_evaluations_scan_id", "rule_evaluations", ["scan_id"], unique=False) - - -def downgrade() -> None: - """Remove idempotent-result storage introduced by this revision.""" - op.drop_index("idx_rule_evaluations_scan_id", table_name="rule_evaluations") - op.drop_table("rule_evaluations") - with op.get_context().autocommit_block(): - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_findings_scan_finding_key") - op.drop_column("findings", "finding_key") diff --git a/api/models/finding.py b/api/models/finding.py index 39b5eac1..af564201 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -26,6 +26,11 @@ logger = logging.getLogger(__name__) +# Worker identities are per-process UUIDs. Heartbeat rows older than this are +# long-dead workers that no metric reads, so they are pruned to keep +# worker_heartbeats bounded. Must stay far above any heartbeat interval. +DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS = 7 * 24 * 60 * 60 + class LostLease(RuntimeError): """Raised when a worker no longer owns the scan it is trying to update.""" @@ -296,7 +301,6 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token finding["severity"] = normalize_severity(finding.get("severity")) finding["finding_key"] = stable_finding_key(scan_result["scan_id"], finding) findings.append(finding) - evaluations = [dict(raw_evaluation) for raw_evaluation in scan_result.get("evaluations", [])] conn = self._get_conn() completed_at = scan_result.get("completed_at") or datetime.now(timezone.utc).isoformat() @@ -600,8 +604,23 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) - def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: - """Durably enqueue exactly one CVE enrichment job for a scan.""" + def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], str]: + """Durably enqueue, or explicitly requeue, one CVE enrichment job. + + Returns the job row and one of four outcomes: + + ``created`` a new job row was inserted for this scan. + ``requeued`` a terminally ``failed`` job was reset to ``pending``. + ``active`` a ``pending``/``running`` job already exists. + ``completed`` enrichment already finished; nothing was changed. + + There is never more than one job per scan: the ``scan_id`` unique + constraint makes the insert idempotent, and the requeue is a single + conditional ``UPDATE`` so concurrent callers converge on one row. A + ``running`` job is left strictly alone -- only its lease expiring + (:meth:`recover_stale_enrichment_jobs`) may take it away from its + current owner, so an operator retry can never steal a live claim. + """ conn = self._get_conn() try: with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: @@ -615,25 +634,69 @@ def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: (str(uuid.uuid4()), scan_id), ) job = cur.fetchone() - if job is None: - cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - existing = cur.fetchone() - if existing is None: - raise RuntimeError("enrichment job conflict did not return an existing job") + if job is not None: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) conn.commit() - return dict(existing), False + return dict(job), "created" + + # A job already exists. Only a terminally failed one is + # revived, and the WHERE clause is the whole guard: a + # concurrent requeue that lost the race sees status + # 'pending' and matches nothing, so both callers end up + # with the same single pending job. + # + # attempt_count restarts because the operator is explicitly + # granting a fresh retry budget -- leaving it at the limit + # would make the job fail again on its first attempt. The + # previous error_message is kept as the audit trail, and + # checkpoint is kept so the retry resumes rather than + # re-enriching findings that already succeeded. cur.execute( - "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + """ + UPDATE enrichment_jobs + SET status = 'pending', + attempt_count = 0, + next_retry_at = CURRENT_TIMESTAMP, + lease_owner = NULL, + lease_expires_at = NULL, + last_heartbeat_at = NULL, + completed_at = NULL + WHERE scan_id = %s AND status = 'failed' + RETURNING * + """, (scan_id,), ) + requeued = cur.fetchone() + if requeued is not None: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) + conn.commit() + return dict(requeued), "requeued" + + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + existing = cur.fetchone() + if existing is None: + raise RuntimeError("enrichment job conflict did not return an existing job") conn.commit() - return dict(job), True + existing = dict(existing) + return existing, "completed" if existing["status"] == "completed" else "active" except Exception: self.rollback(conn) raise - def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: - """Atomically claim the next retry-ready enrichment job.""" + def claim_next_enrichment_job( + self, lease_owner: str, lease_seconds: int, scan_id: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Atomically claim the next retry-ready enrichment job. + + ``scan_id`` restricts the claim to one scan's job, with the same lease + and fencing semantics as an unrestricted claim. + """ if lease_seconds <= 0: raise ValueError("lease_seconds must be positive") conn = self._get_conn() @@ -652,13 +715,14 @@ def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Opt SELECT job_id FROM enrichment_jobs WHERE status = 'pending' AND next_retry_at <= CURRENT_TIMESTAMP + AND (%s IS NULL OR scan_id = %s::uuid) ORDER BY created_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING * """, - (lease_owner, lease_seconds), + (lease_owner, lease_seconds, scan_id, scan_id), ) job = cur.fetchone() if job: @@ -1001,8 +1065,16 @@ def update_scan_status( raise logger.info("Updated scan %s status to %s", scan_id, status) - def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: - """Atomically claim one pending scan and establish its renewable lease.""" + def claim_next_pending_scan( + self, lease_owner: str, lease_seconds: int, scan_id: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Atomically claim one pending scan and establish its renewable lease. + + ``scan_id`` restricts the claim to one specific scan instead of the + oldest pending one. The lease, fencing and skip-locked semantics are + identical either way; callers that must act on a known scan use it so + they neither depend on nor disturb the shared queue order. + """ if lease_seconds <= 0: raise ValueError("lease_seconds must be positive") conn = self._get_conn() @@ -1023,13 +1095,14 @@ def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optio SELECT scan_id FROM scans WHERE status = 'pending' + AND (%s IS NULL OR scan_id = %s::uuid) ORDER BY started_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING * """, - (lease_owner, lease_seconds), + (lease_owner, lease_seconds, scan_id, scan_id), ) row = cur.fetchone() conn.commit() @@ -1147,10 +1220,24 @@ def get_scans(self) -> List[Dict[str, Any]]: cur.execute("SELECT * FROM scans ORDER BY started_at DESC LIMIT 100") return [dict(row) for row in cur.fetchall()] - def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: - """Persist liveness without using worker IDs as metric labels.""" + def record_worker_heartbeat( + self, + worker_id: str, + worker_type: str, + retention_seconds: int = DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS, + ) -> None: + """Persist liveness without using worker IDs as metric labels. + + Worker identities are per-process, so a long-lived deployment would + otherwise accumulate one dead row per restart forever. Retired rows + are pruned only on the heartbeat that actually inserts a new worker + identity -- i.e. once per worker process -- so the table stays bounded + without a scheduler and without a full-table sweep on every beat. + """ if worker_type not in {"scan", "enrichment"}: raise ValueError("unsupported worker type") + if retention_seconds <= 0: + raise ValueError("retention_seconds must be positive") conn = self._get_conn() try: with conn.cursor() as cur: @@ -1160,9 +1247,22 @@ def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: VALUES (%s, %s, CURRENT_TIMESTAMP) ON CONFLICT (worker_id, worker_type) DO UPDATE SET last_seen_at = EXCLUDED.last_seen_at + RETURNING (xmax = 0) AS inserted """, (worker_id, worker_type), ) + row = cur.fetchone() + inserted = row["inserted"] if isinstance(row, dict) else row[0] + if inserted: + # Retention is far longer than any heartbeat interval, so a + # live worker can never prune itself or a peer. + cur.execute( + """ + DELETE FROM worker_heartbeats + WHERE last_seen_at < CURRENT_TIMESTAMP - (%s * INTERVAL '1 second') + """, + (retention_seconds,), + ) conn.commit() except Exception: self.rollback(conn) diff --git a/api/routes/scans.py b/api/routes/scans.py index a84d233a..89907e33 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -169,6 +169,14 @@ def trigger_scan(): return jsonify({"error": "Critical route failure"}), 500 +_ENRICH_MESSAGES = { + "created": "CVE enrichment queued; poll GET /api/scans/ for completion.", + "requeued": "Previously failed enrichment job requeued; poll GET /api/scans/ for completion.", + "active": "Existing enrichment job returned.", + "completed": "Scan already enriched", +} + + @scans_bp.post("/api/scans//enrich") def enrich_scan(scan_id): """Enqueue durable CVE enrichment; no request-owned thread is created.""" @@ -188,25 +196,17 @@ def enrich_scan(scan_id): if not findings: return jsonify({"error": "No findings found for this scan"}), 404 - job, created = db.enqueue_enrichment_job(scan_id) - if not created: - return jsonify( - { - "job_id": str(job["job_id"]), - "scan_id": scan_id, - "status": job["status"], - "message": "Existing enrichment job returned.", - } - ), 202 - - return jsonify( - { - "scan_id": scan_id, - "job_id": str(job["job_id"]), - "status": "PENDING", - "message": "CVE enrichment queued; poll GET /api/scans/ for completion.", - } - ), 202 + job, outcome = db.enqueue_enrichment_job(scan_id) + body = { + "scan_id": scan_id, + "job_id": str(job["job_id"]), + "status": job["status"], + "outcome": outcome, + "message": _ENRICH_MESSAGES[outcome], + } + # A job that already finished is reported as-is rather than restarted; + # every other outcome leaves exactly one queued or running job. + return jsonify(body), 200 if outcome == "completed" else 202 except ValidationError: return jsonify({"error": VALIDATION_ERROR_MESSAGE}), 400 diff --git a/docs/api-reference.md b/docs/api-reference.md index 093709a9..f5e732f6 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -206,9 +206,9 @@ Example response: ## POST /api/scans/trigger -Triggers an asynchronous scan against the configured subscription. Returns `202 Accepted` with the `scan_id` immediately. The actual scan execution happens in a background worker process. +Admits an asynchronous scan against the configured subscription. Execution happens in a background worker process; the response returns as soon as the scan is durably recorded. -Request body: +Request body (optional — falls back to `AZURE_SUBSCRIPTION_ID`): ```json { @@ -216,7 +216,26 @@ Request body: } ``` -Example response: +### Admission semantics + +Admission is serialized per subscription and enforced by the database, so concurrent and replayed triggers converge on one logical scan rather than creating duplicates: + +- **At most one active scan per subscription.** While a `pending` or `running` scan exists, a further trigger returns that existing scan instead of queueing another. +- **`Idempotency-Key` (optional request header, 1–200 characters).** A repeat of the same key for the same subscription returns the original scan. The key is scoped to the subscription; the same key under a different subscription is a different request. +- **`OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR`** adds an optional hourly admission quota. Unset or `0` (the default) applies no time-window limit; the one-active-scan rule still applies. + +### Responses + +| Status | When | Body | +| --- | --- | --- | +| `202 Accepted` | A new scan was admitted and queued. | `scan_id`, `status: "pending"`, `message` | +| `200 OK` | The request resolved to an existing logical scan — an `Idempotency-Key` replay of the same request, or a trigger while a scan is already active for the subscription. | `scan_id`, `status` (the existing scan's `pending`/`running`), `message: "Existing logical scan returned."` | +| `400 Bad Request` | Malformed body, invalid `subscription_id`, missing subscription, or an `Idempotency-Key` outside 1–200 characters. | `error` | +| `403 Forbidden` | `subscription_id` is not on the `OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS` allowlist. | `error` | +| `409 Conflict` | The `Idempotency-Key` was reused with a different request payload. | `error: "Idempotency-Key is already associated with a different request."` | +| `429 Too Many Requests` | The configured hourly quota for this subscription is exhausted. | `error: "Scan quota exceeded for this subscription."` | + +New scan (`202`): ```json { @@ -226,7 +245,17 @@ Example response: } ``` -Missing subscription response: +Replay or already-active scan (`200`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "status": "running", + "message": "Existing logical scan returned." +} +``` + +Missing subscription response (`400`): ```json { @@ -236,6 +265,54 @@ Missing subscription response: --- +## POST /api/scans/<scan_id>/enrich + +Queues durable CVE enrichment for a completed scan's findings. Enrichment runs as a database-backed job claimed by the background worker — the request never owns a thread, so the work survives an API restart. + +There is **never more than one enrichment job per scan**. Repeat calls are safe: they report the state of the single job rather than creating another. + +### Responses + +Every response carries `scan_id`, `job_id`, `status` (the job row's state) and an `outcome` naming what this call did: + +| Status | `outcome` | When | +| --- | --- | --- | +| `202 Accepted` | `created` | No job existed; one was queued. | +| `202 Accepted` | `requeued` | A previously **failed** job was reset to `pending` and will be retried. | +| `202 Accepted` | `active` | A `pending` or `running` job already exists and was returned unchanged. A live claim is never interrupted. | +| `200 OK` | `completed` | Enrichment already finished; nothing was restarted. | +| `404 Not Found` | — | Unknown `scan_id`, or the scan has no findings to enrich. | + +A job that exhausts its retry budget becomes `failed`. Re-POSTing this endpoint is the supported operator recovery: it atomically returns the job to `pending` with a fresh retry budget, clears the lease, and keeps the last `error_message` and the `checkpoint` so the retry resumes rather than re-enriching findings that already succeeded. Concurrent re-POSTs converge — exactly one reports `requeued` and the rest report `active`. + +Newly queued (`202`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "job_id": "1f2e3d4c-5b6a-4790-8123-456789abcdef", + "status": "pending", + "outcome": "created", + "message": "CVE enrichment queued; poll GET /api/scans/ for completion." +} +``` + +Requeued after terminal failure (`202`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "job_id": "1f2e3d4c-5b6a-4790-8123-456789abcdef", + "status": "pending", + "outcome": "requeued", + "message": "Previously failed enrichment job requeued; poll GET /api/scans/ for completion." +} +``` + +Poll `GET /api/scans/` for `cve_enrichment_status` (`PENDING`, `ENRICHING`, `COMPLETED`, `FAILED`). + +--- + ## GET /api/score Returns the overall security posture score from 0 to 100. Under [severity contract v1](severity-contract.md), the score starts at 100 and deducts 20 per CRITICAL finding, 10 per HIGH finding, 5 per MEDIUM finding, and 2 per LOW finding. INFO findings deduct zero. diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index 757aec61..e8c6eeb6 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -33,20 +33,32 @@ credentials or persistent Azure errors cannot retry forever. Findings use a stable database-enforced identity (`scan`, rule, canonical resource scope, and an optional rule-specific discriminator). Result retries use PostgreSQL upserts, so mutable text or severity is updated rather than -creating a second authoritative finding. Rule evaluations use the same -database-first uniqueness model. +creating a second authoritative finding. Per-resource rule evaluations are a +separate contract (issue #263) and are not persisted by this architecture yet; +when they land, they are written inside the same fenced completion transaction +and inherit its ownership check. ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings, evaluations, and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. +The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. ### Durable CVE enrichment `POST /api/scans//enrich` enqueues (or returns) the one durable PostgreSQL enrichment job for the scan; it never starts a request-owned daemon thread. The scan worker claims those jobs with the same owner/expiry/fencing model, checkpoints after each finding, and retries transient failures with bounded exponential backoff. An expired job is recovered or terminally failed after its attempt limit. NVD retrieval follows every `totalResults` page; replaying a checkpoint updates the existing finding instead of duplicating CVE data. +A job that exhausts its retries becomes `failed`. Re-POSTing the endpoint is the supported recovery path: it atomically returns that job to `pending` with a fresh retry budget while keeping the same job row, its last `error_message`, and its `checkpoint`. A `running` job with a live lease is never disturbed by a re-POST — only lease expiry can move it — and a `completed` job is never restarted. See [api-reference.md](api-reference.md) for the exact response contract. + +### Worker scheduling + +One worker process serves both durable queues. Each loop iteration takes **at most one** enrichment job and **at most one** scan, so neither queue can starve the other no matter how deep either gets: a large enrichment backlog delays scans by one job per iteration rather than blocking them until it drains. An iteration that did any work polls again immediately; only a completely idle iteration sleeps for the poll interval. + ### Operational signals `/metrics` derives bounded-cardinality operational gauges from PostgreSQL: worker heartbeat age, oldest queue age, oldest active lease age, aggregate retry attempts, and the last successful scan timestamp. Labels are limited to `queue` (`scan` or `enrichment`) and `worker_type`; scan, job, subscription, and worker identifiers are never metric labels. +These gauges are recomputed on every scrape rather than cached. The queue, lease and heartbeat aggregates are served by partial indexes and touch only active rows; the last-successful-scan lookup is served by a partial index on completed scans. The two `retry_attempts` sums scan their whole table and therefore grow with scan history — acceptable at current volumes, and the thing to revisit first (a short-TTL in-process cache) if `/metrics` scrape latency ever becomes visible. + +Worker identities are per-process, so `worker_heartbeats` would otherwise gain a permanent row on every restart. Rows older than `WORKER_HEARTBEAT_RETENTION_SECONDS` (default 7 days) are pruned on the heartbeat that registers a new worker identity — once per worker process, never on every beat. Retention is far longer than any heartbeat interval, so a live worker is never pruned. + ## Technical Rationale ### Why not Celery or Redis @@ -62,6 +74,31 @@ workers, apply Alembic migrations, then start the fenced worker version. Legacy workers do not carry ownership/fencing state; legacy running scans are retained and made recoverable by the lease migration rather than deleted. +### Migration prerequisite: one active scan per subscription + +`a7c5e9d2f1b4` adds the unique index that enforces one `pending`/`running` scan +per subscription. A deployment that predates that rule may already hold several, +in which case the migration **stops before creating any index** and names the +offending subscriptions: + +``` +Cannot enforce one active scan per subscription: 1 subscription(s) already have +more than one pending/running scan: (2 active). ... +``` + +Nothing is changed when this happens — the migration will not decide which of +your production scans is authoritative. Resolve it, then re-run: + +1. Drain the old workers first (step 1 above), so no new active scans appear. +2. Let the in-flight scans finish, or mark the superseded rows `failed` + (`UPDATE scans SET status = 'failed', error_message = '...' WHERE scan_id = ...`). + Never delete the rows: scan history is retained, and only `pending`/`running` + rows are constrained — any number of `completed`/`failed` scans per + subscription remains valid. +3. Re-run `alembic upgrade head`. A retry is safe: the migration drops the + INVALID index that an interrupted `CREATE INDEX CONCURRENTLY` leaves behind + before rebuilding it. + ## Testing Suite The asynchronous transition is verified through a multi layered testing strategy. diff --git a/scanner/worker.py b/scanner/worker.py index b5021e11..2ae431ae 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -13,7 +13,11 @@ import uuid from datetime import datetime, timezone -from api.models.finding import DatabaseManager, LostLease +from api.models.finding import ( + DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS, + DatabaseManager, + LostLease, +) from api.observability import ( PENDING_SCANS, SCAN_DURATION_SECONDS, @@ -139,14 +143,17 @@ def run_worker(): db = DatabaseManager(db_url) worker_id = str(uuid.uuid4()) lease_seconds, heartbeat_seconds = lease_configuration() + heartbeat_retention_seconds = _positive_seconds( + "WORKER_HEARTBEAT_RETENTION_SECONDS", DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS + ) logger.info("OpenShield Background Worker started. Polling every %ds", POLL_INTERVAL_SECONDS) while True: try: - db.record_worker_heartbeat(worker_id, "scan") + db.record_worker_heartbeat(worker_id, "scan", heartbeat_retention_seconds) # This process executes both durable queue types; record both # liveness signals without exporting the worker UUID as a label. - db.record_worker_heartbeat(worker_id, "enrichment") + db.record_worker_heartbeat(worker_id, "enrichment", heartbeat_retention_seconds) # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() db.recover_stale_enrichment_jobs() @@ -154,16 +161,22 @@ def run_worker(): # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) - # 3. Run one durable enrichment job before taking another scan. + # 3. Take at most one job from each durable queue per iteration. + # Draining enrichment first and restarting the loop would let a + # sustained enrichment backlog hold off every pending scan, so + # the two queues alternate instead: neither can starve the + # other regardless of how deep either one gets. enrichment_job = db.claim_next_enrichment_job(worker_id, lease_seconds) if enrichment_job: process_enrichment_job(db, enrichment_job, worker_id, lease_seconds) - continue # 4. Atomic scan claim scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: - time.sleep(POLL_INTERVAL_SECONDS) + # Only idle when there was no work at all; an iteration that + # ran an enrichment job polls again immediately. + if not enrichment_job: + time.sleep(POLL_INTERVAL_SECONDS) continue scan_id = str(scan["scan_id"]) diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index b98b73fe..0b4780b0 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -6,6 +6,7 @@ from unittest.mock import patch import psycopg2 +import psycopg2.extras import pytest from api.models.finding import DatabaseManager, LostLease @@ -24,7 +25,11 @@ def enrichment_scan(): db = DatabaseManager(dsn) try: db.create_pending_scan(scan_id, subscription_id) - claim = db.claim_next_pending_scan("seed", 120) + # Claim this scan explicitly. An unrestricted claim takes the globally + # oldest pending scan, which in a shared test database is very often a + # row another test admitted first. + claim = db.claim_next_pending_scan("seed", 120, scan_id=scan_id) + assert claim is not None and str(claim["scan_id"]) == scan_id result = { "scan_id": scan_id, "subscription_id": subscription_id, @@ -50,23 +55,23 @@ def enrichment_scan(): ], } db.save_scan(result, "seed", claim["fencing_token"]) - job, created = db.enqueue_enrichment_job(scan_id) - assert created is False + job, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "active" yield dsn, scan_id, job finally: db.close() with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) -def _claim(dsn): +def _claim(dsn, scan_id, owner="worker-a"): + """Claim this scan's enrichment job, never another test's.""" db = DatabaseManager(dsn) try: - return db.claim_next_enrichment_job("worker-a", 120) + return db.claim_next_enrichment_job(owner, 120, scan_id=scan_id) finally: db.close() @@ -75,10 +80,10 @@ def test_duplicate_enqueue_and_claim_race(enrichment_scan): dsn, scan_id, first_job = enrichment_scan db = DatabaseManager(dsn) try: - replay, created = db.enqueue_enrichment_job(scan_id) + replay, outcome = db.enqueue_enrichment_job(scan_id) finally: db.close() - assert created is False + assert outcome == "active" assert replay["job_id"] == first_job["job_id"] barrier = threading.Barrier(2) @@ -86,7 +91,7 @@ def test_duplicate_enqueue_and_claim_race(enrichment_scan): def claim(): barrier.wait() - claims.append(_claim(dsn)) + claims.append(_claim(dsn, scan_id)) threads = [threading.Thread(target=claim) for _ in range(2)] for thread in threads: @@ -98,7 +103,7 @@ def claim(): def test_checkpoint_resume_and_completion(enrichment_scan): dsn, scan_id, _ = enrichment_scan - job = _claim(dsn) + job = _claim(dsn, scan_id) assert job is not None db = DatabaseManager(dsn) try: @@ -121,7 +126,7 @@ def enrich_once_then_fail(finding): cur.execute( "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", (scan_id,) ) - resumed = db.claim_next_enrichment_job("worker-b", 120) + resumed = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} assert process_enrichment_job(db, resumed, "worker-b", 120) == "completed" @@ -142,7 +147,7 @@ def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): db = DatabaseManager(dsn) try: for attempt in range(1, 4): - job = db.claim_next_enrichment_job("worker-a", 120) + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) assert job is not None with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): expected = "failed" if attempt == 3 else "retry" @@ -164,7 +169,7 @@ def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enrichment_scan): dsn, scan_id, _ = enrichment_scan - first = _claim(dsn) + first = _claim(dsn, scan_id) assert first is not None with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: @@ -178,10 +183,190 @@ def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enr ) db = DatabaseManager(dsn) try: - assert db.recover_stale_enrichment_jobs() == 1 - second = db.claim_next_enrichment_job("worker-b", 120) + # Other tests may share this database; assert this job specifically + # was recovered rather than that it was the only one. + assert db.recover_stale_enrichment_jobs() >= 1 + second = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) assert second["fencing_token"] > first["fencing_token"] with pytest.raises(LostLease): db.heartbeat_enrichment_job(str(first["job_id"]), "worker-a", first["fencing_token"], 120) finally: db.close() + + +def _fail_terminally(dsn, db, scan_id): + """Drive a job through its whole retry budget until it is 'failed'.""" + for attempt in range(1, 4): + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): + process_enrichment_job(db, job, "worker-a", 120) + if attempt < 3: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", + (scan_id,), + ) + return _job_row(dsn, scan_id) + + +def _job_row(dsn, scan_id): + with psycopg2.connect(dsn) as conn: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + return dict(cur.fetchone()) + + +def test_terminally_failed_job_can_be_explicitly_requeued(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + failed = _fail_terminally(dsn, db, scan_id) + assert failed["status"] == "failed" + assert failed["attempt_count"] >= 3 + + job, outcome = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + + assert outcome == "requeued" + # Same logical job, now retryable again. + assert str(job["job_id"]) == str(first_job["job_id"]) + assert job["status"] == "pending" + assert job["attempt_count"] == 0 + assert job["lease_owner"] is None + assert job["lease_expires_at"] is None + assert job["completed_at"] is None + # The failure reason is kept as the audit trail, and the checkpoint is kept + # so the retry resumes instead of re-enriching what already succeeded. + assert job["error_message"] + assert job["checkpoint"] == failed["checkpoint"] + assert _job_row(dsn, scan_id)["next_retry_at"] is not None + + +def test_completed_job_is_never_restarted(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + assert process_enrichment_job(db, job, "worker-a", 120) == "completed" + + requeued, outcome = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + + assert outcome == "completed" + assert requeued["status"] == "completed" + assert _job_row(dsn, scan_id)["status"] == "completed" + + +def test_requeue_does_not_steal_a_valid_running_lease(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + running = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + assert running is not None + + job, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "active" + assert job["status"] == "running" + + # The live owner keeps its lease and can still heartbeat and complete. + db.heartbeat_enrichment_job(str(running["job_id"]), "worker-a", running["fencing_token"], 120) + after = _job_row(dsn, scan_id) + assert after["lease_owner"] == "worker-a" + assert after["fencing_token"] == running["fencing_token"] + finally: + db.close() + + +def test_concurrent_requeues_converge_on_one_logical_job(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + _fail_terminally(dsn, db, scan_id) + finally: + db.close() + + barrier = threading.Barrier(4) + outcomes: list[str] = [] + + def requeue() -> None: + worker_db = DatabaseManager(dsn) + try: + barrier.wait() + outcomes.append(worker_db.enqueue_enrichment_job(scan_id)[1]) + finally: + worker_db.close() + + threads = [threading.Thread(target=requeue) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # Exactly one caller performs the failed -> pending transition; the rest + # observe the job that is already queued. Nobody creates a second job. + assert sorted(outcomes) == ["active", "active", "active", "requeued"] + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == 1 + assert str(_job_row(dsn, scan_id)["job_id"]) == str(first_job["job_id"]) + + +def test_stale_token_cannot_write_after_requeue_and_reclaim(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + # worker-a held the job through its last failed attempt; that is the + # token a stale process would still be carrying. + failed = _fail_terminally(dsn, db, scan_id) + stale_job_id, stale_token = str(failed["job_id"]), failed["fencing_token"] + + assert db.enqueue_enrichment_job(scan_id)[1] == "requeued" + reclaimed = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) + assert reclaimed is not None + assert reclaimed["fencing_token"] > stale_token + + findings = db.get_enrichment_findings(scan_id) + with pytest.raises(LostLease): + db.heartbeat_enrichment_job(stale_job_id, "worker-a", stale_token, 120) + with pytest.raises(LostLease): + db.persist_enrichment_progress(stale_job_id, "worker-a", stale_token, findings[0], 99) + with pytest.raises(LostLease): + db.complete_enrichment_job(stale_job_id, "worker-a", stale_token) + + # None of the rejected writes landed. + current = _job_row(dsn, scan_id) + assert current["checkpoint"] != 99 + assert current["status"] == "running" + assert current["lease_owner"] == "worker-b" + finally: + db.close() + + +def test_requeued_job_can_eventually_complete(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + _fail_terminally(dsn, db, scan_id) + _, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "requeued" + + job = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-2"}]} + assert process_enrichment_job(db, job, "worker-b", 120) == "completed" + finally: + db.close() + + assert _job_row(dsn, scan_id)["status"] == "completed" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT cve_enrichment_status FROM scans WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == "COMPLETED" diff --git a/tests/test_operational_metrics_postgres.py b/tests/test_operational_metrics_postgres.py index f4f63c7c..c8df812b 100644 --- a/tests/test_operational_metrics_postgres.py +++ b/tests/test_operational_metrics_postgres.py @@ -22,11 +22,13 @@ def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): db.create_pending_scan(scan_id, subscription_id) db.record_worker_heartbeat("worker-test", "scan") db.record_worker_heartbeat("worker-test", "enrichment") - claim = db.claim_next_pending_scan("worker-a", 120) + claim = db.claim_next_pending_scan("worker-a", 120, scan_id=scan_id) assert claim is not None snapshot = db.get_operational_metrics() assert snapshot["oldest_lease_age"]["scan"] >= 0 - assert snapshot["retry_attempts"]["scan"] == 0 + # Aggregated over the whole table, so other rows in a shared test + # database may contribute; only the shape is this test's contract. + assert snapshot["retry_attempts"]["scan"] >= 0 assert snapshot["worker_heartbeat_age"]["scan"] >= 0 assert snapshot["worker_heartbeat_age"]["enrichment"] >= 0 assert snapshot["last_successful_scan_timestamp"] >= 0 @@ -37,3 +39,44 @@ def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): cur.execute("DELETE FROM worker_heartbeats WHERE worker_id = 'worker-test'") cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def test_stale_worker_heartbeats_are_pruned_without_dropping_live_workers(): + """Retired worker rows must not accumulate, and live ones must survive.""" + dsn = os.environ["DATABASE_URL"] + live_worker, dead_worker = f"live-{uuid.uuid4()}", f"dead-{uuid.uuid4()}" + db = DatabaseManager(dsn) + try: + db.record_worker_heartbeat(dead_worker, "scan") + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE worker_heartbeats SET last_seen_at = CURRENT_TIMESTAMP - INTERVAL '30 days' " + "WHERE worker_id = %s", + (dead_worker,), + ) + + # Pruning happens on the beat that inserts a new worker identity. + db.record_worker_heartbeat(live_worker, "scan", retention_seconds=7 * 24 * 60 * 60) + + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (dead_worker,)) + assert cur.fetchone()[0] == 0 + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (live_worker,)) + assert cur.fetchone()[0] == 1 + + # A repeat beat from an existing worker only refreshes its timestamp. + db.record_worker_heartbeat(live_worker, "scan") + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (live_worker,)) + assert cur.fetchone()[0] == 1 + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "DELETE FROM worker_heartbeats WHERE worker_id IN (%s, %s)", + (live_worker, dead_worker), + ) diff --git a/tests/test_scan_admission_migration_postgres.py b/tests/test_scan_admission_migration_postgres.py new file mode 100644 index 00000000..bb10ad10 --- /dev/null +++ b/tests/test_scan_admission_migration_postgres.py @@ -0,0 +1,244 @@ +"""Migration-time safety for the one-active-scan unique index. + +``CREATE UNIQUE INDEX CONCURRENTLY`` fails -- and leaves an INVALID index +behind -- when a deployment already holds several active scans for one +subscription. These tests run the real migration against a throwaway +PostgreSQL database seeded with exactly that legacy shape. +""" + +import os +import uuid +from contextlib import contextmanager + +import psycopg2 +import pytest +from alembic import command +from alembic.config import Config + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + +# The revision immediately before scan admission is introduced. +_BEFORE_ADMISSION = "f2b6d8e1a4c9" +_ADMISSION = "a7c5e9d2f1b4" +_ACTIVE_INDEX = "uq_scans_one_active_per_subscription" +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +@contextmanager +def _scratch_database(): + """Create and drop a database dedicated to one migration test.""" + base = os.environ["DATABASE_URL"].rsplit("/", 1)[0] + name = f"openshield_mig_{uuid.uuid4().hex[:12]}" + admin = psycopg2.connect(f"{base}/postgres") + admin.autocommit = True + try: + with admin.cursor() as cur: + cur.execute(f'CREATE DATABASE "{name}"') + yield f"{base}/{name}" + finally: + with admin.cursor() as cur: + cur.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()", + (name,), + ) + cur.execute(f'DROP DATABASE IF EXISTS "{name}"') + admin.close() + + +def _alembic_config() -> Config: + """Alembic config that does not touch this process's logging setup. + + Passing alembic.ini would make ``env.py`` call ``fileConfig()``, which + disables every logger that already exists -- silencing the rest of the + test session. Only ``script_location`` is actually needed here; ``env.py`` + reads the database URL from the environment. + """ + config = Config() + config.set_main_option("script_location", os.path.join(_REPO_ROOT, "alembic")) + return config + + +def _upgrade(dsn: str, revision: str) -> None: + """Run alembic against ``dsn`` without disturbing the ambient env.""" + previous = os.environ.get("DATABASE_URL") + os.environ["DATABASE_URL"] = dsn + try: + command.upgrade(_alembic_config(), revision) + finally: + if previous is None: + os.environ.pop("DATABASE_URL", None) + else: + os.environ["DATABASE_URL"] = previous + + +def _seed_active_scans(dsn: str, subscription_id: str, count: int) -> list[str]: + scan_ids = [str(uuid.uuid4()) for _ in range(count)] + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + for offset, scan_id in enumerate(scan_ids): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute'), 'pending') + """, + (scan_id, subscription_id, offset), + ) + return scan_ids + + +def _index_is_valid(dsn: str, name: str): + """Return True/False for a present index, or None when it does not exist.""" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + SELECT i.indisvalid + FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid + WHERE c.relname = %s + """, + (name,), + ) + row = cur.fetchone() + return None if row is None else row[0] + + +def test_legacy_duplicate_active_scans_fail_the_migration_with_an_actionable_error(): + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + _seed_active_scans(dsn, subscription_id, 2) + + with pytest.raises(RuntimeError) as excinfo: + _upgrade(dsn, _ADMISSION) + + message = str(excinfo.value) + assert subscription_id in message + assert "2 active" in message + # The operator is told what to do, and no unusable index was left behind. + assert "Resolve them first" in message + assert _index_is_valid(dsn, _ACTIVE_INDEX) is None + + # Historical rows are untouched: the migration decides nothing for us. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM scans WHERE subscription_id = %s", (subscription_id,)) + assert cur.fetchone()[0] == 2 + + +def test_migration_succeeds_and_enforces_the_contract_after_cleanup(): + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + scan_ids = _seed_active_scans(dsn, subscription_id, 2) + + with pytest.raises(RuntimeError): + _upgrade(dsn, _ADMISSION) + + # The documented cleanup: retire the superseded scan, keep the record. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'failed' WHERE scan_id = %s", (scan_ids[1],)) + + _upgrade(dsn, _ADMISSION) + assert _index_is_valid(dsn, _ACTIVE_INDEX) is True + + conn = psycopg2.connect(dsn) + try: + with conn.cursor() as cur: + # Both historical rows survived the migration. + cur.execute("SELECT COUNT(*) FROM scans WHERE subscription_id = %s", (subscription_id,)) + assert cur.fetchone()[0] == 2 + + # A second active scan for the same subscription is now refused. + with pytest.raises(psycopg2.errors.UniqueViolation): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'running') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.rollback() + + with conn.cursor() as cur: + # Terminal scans stay unconstrained: history keeps accumulating. + for _ in range(3): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed') + """, + (str(uuid.uuid4()), subscription_id), + ) + # Idempotency-key uniqueness is per subscription and ignores NULLs. + for _ in range(2): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', NULL) + """, + (str(uuid.uuid4()), subscription_id), + ) + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.commit() + + with conn.cursor() as cur: + with pytest.raises(psycopg2.errors.UniqueViolation): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.rollback() + + # The same key under a different subscription is still admissible. + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), str(uuid.uuid4())), + ) + conn.commit() + finally: + conn.close() + + +def test_an_invalid_index_left_by_a_failed_build_is_replaced_not_inherited(): + """A retry after a failed CONCURRENTLY build must not trip over its debris.""" + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + _seed_active_scans(dsn, subscription_id, 1) + + # Reproduce the debris a failed concurrent build leaves behind. + conn = psycopg2.connect(dsn) + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute( + f"CREATE UNIQUE INDEX CONCURRENTLY {_ACTIVE_INDEX} " + "ON scans (subscription_id) WHERE status IN ('pending', 'running')" + ) + cur.execute( + "UPDATE pg_index SET indisvalid = false WHERE indexrelid = %s::regclass", + (_ACTIVE_INDEX,), + ) + finally: + conn.close() + assert _index_is_valid(dsn, _ACTIVE_INDEX) is False + + _upgrade(dsn, _ADMISSION) + assert _index_is_valid(dsn, _ACTIVE_INDEX) is True diff --git a/tests/test_scan_admission_postgres.py b/tests/test_scan_admission_postgres.py index 9c0fa829..05ed909f 100644 --- a/tests/test_scan_admission_postgres.py +++ b/tests/test_scan_admission_postgres.py @@ -23,7 +23,6 @@ def admitted_scans(): with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: for scan_id in scan_ids: - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index fd5df2f2..2d19e6cd 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -95,18 +95,11 @@ def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: (owner, fencing_token, scan_id), ) - def evaluation_count(self, scan_id: str) -> int: - with psycopg2.connect(self.dsn) as conn: - with conn.cursor() as cur: - cur.execute("SELECT COUNT(*) FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) - return cur.fetchone()[0] - def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) @@ -118,10 +111,16 @@ def scan_rows() -> ScanRows: rows.cleanup() -def _claim(dsn: str, owner: str) -> dict | None: +def _claim(dsn: str, owner: str, scan_id: str | None = None) -> dict | None: + """Claim a pending scan, restricted to ``scan_id`` when given. + + Tests always pass their own scan: an unrestricted claim takes the globally + oldest pending row, so in a shared database it would otherwise claim + whatever unrelated scan another test admitted first. + """ db = DatabaseManager(dsn) try: - return db.claim_next_pending_scan(owner, 120) + return db.claim_next_pending_scan(owner, 120, scan_id=scan_id) finally: db.close() @@ -135,13 +134,13 @@ def _recover(dsn: str, max_attempts: int = 3) -> int: def test_two_workers_race_to_claim_only_one_scan(scan_rows): - scan_rows.create() + scan_id, _ = scan_rows.create() barrier = threading.Barrier(2) claims: list[dict | None] = [] def claim(owner: str) -> None: barrier.wait() - claims.append(_claim(scan_rows.dsn, owner)) + claims.append(_claim(scan_rows.dsn, owner, scan_id)) threads = [threading.Thread(target=claim, args=(owner,)) for owner in ("worker-a", "worker-b")] for thread in threads: @@ -156,14 +155,14 @@ def claim(owner: str) -> None: def test_active_lease_cannot_be_reclaimed(scan_rows): - scan_rows.create() - assert _claim(scan_rows.dsn, "worker-a") is not None - assert _claim(scan_rows.dsn, "worker-b") is None + scan_id, _ = scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a", scan_id) is not None + assert _claim(scan_rows.dsn, "worker-b", scan_id) is None def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_rows): scan_id, _ = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None db = DatabaseManager(scan_rows.dsn) @@ -179,12 +178,15 @@ def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_ro def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): scan_id, _ = scan_rows.create() - first_claim = _claim(scan_rows.dsn, "worker-a") + first_claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert first_claim is not None scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn) == 1 - second_claim = _claim(scan_rows.dsn, "worker-b") + # Recovery is queue-wide, so assert this scan was recovered rather than + # that it was the only scan recovered. + assert _recover(scan_rows.dsn) >= 1 + assert scan_rows.scan(scan_id)["status"] == "pending" + second_claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert second_claim is not None assert second_claim["lease_owner"] == "worker-b" @@ -193,11 +195,11 @@ def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows): scan_id, subscription_id = scan_rows.create() - first_claim = _claim(scan_rows.dsn, "worker-a") + first_claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert first_claim is not None scan_rows.expire(scan_id) _recover(scan_rows.dsn) - second_claim = _claim(scan_rows.dsn, "worker-b") + second_claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert second_claim is not None stale_db = DatabaseManager(scan_rows.dsn) @@ -226,7 +228,7 @@ def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows) def test_current_owner_completion_persists_results_atomically(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-b") + claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert claim is not None db = DatabaseManager(scan_rows.dsn) @@ -243,7 +245,7 @@ def test_current_owner_completion_persists_results_atomically(scan_rows): def test_sql_abort_rolls_back_and_the_connection_remains_usable(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None broken_result = _result(scan_id, subscription_id) broken_result["findings"][0]["detected_at"] = None @@ -297,43 +299,31 @@ def test_terminated_backend_is_discarded_and_reacquired(scan_rows): def test_expired_restart_work_obeys_attempt_limit(scan_rows): scan_id, _ = scan_rows.create() - assert _claim(scan_rows.dsn, "worker-a") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 - - assert _claim(scan_rows.dsn, "worker-b") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 - - assert _claim(scan_rows.dsn, "worker-c") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 + # Recovery is queue-wide, so assert this scan's own progression rather + # than a global count another test could contribute to. + for owner in ("worker-a", "worker-b", "worker-c"): + assert _claim(scan_rows.dsn, owner, scan_id) is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) >= 1 assert scan_rows.scan(scan_id)["status"] == "failed" def test_empty_claim_leaves_no_open_transaction(scan_rows): + # A scan_id with no pending row exercises the same "claimed nothing" path + # without depending on the shared queue being empty. db = DatabaseManager(scan_rows.dsn) try: - assert db.claim_next_pending_scan("worker-a", 120) is None + assert db.claim_next_pending_scan("worker-a", 120, scan_id=str(uuid.uuid4())) is None assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE finally: db.close() -def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_rows): +def test_duplicate_result_delivery_upserts_mutable_fields(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None result = _result(scan_id, subscription_id) - result["evaluations"] = [ - { - "rule_id": "AZ-LEASE-001", - "resource_id": result["findings"][0]["resource_id"], - "resource_type": "Test/resource", - "status": "FAIL", - "evidence": {"version": 1}, - } - ] db = DatabaseManager(scan_rows.dsn) try: @@ -342,22 +332,21 @@ def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_r scan_rows.rearm(scan_id, "worker-a", claim["fencing_token"]) result["findings"][0]["description"] = "Updated presentation text" result["findings"][0]["severity"] = "CRITICAL" - result["evaluations"][0]["evidence"] = {"version": 2} db.save_scan(result, "worker-a", claim["fencing_token"]) persisted = db.get_findings({"scan_id": scan_id}) finally: db.close() + # Replay updates the same row in place instead of duplicating it. assert len(persisted) == 1 assert persisted[0]["id"] == first_id assert persisted[0]["description"] == "Updated presentation text" assert persisted[0]["severity"] == "CRITICAL" - assert scan_rows.evaluation_count(scan_id) == 1 def test_distinct_finding_discriminators_preserve_multiple_violations(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None result = _result(scan_id, subscription_id) duplicate_scope = dict(result["findings"][0]) diff --git a/tests/test_scans_enrich.py b/tests/test_scans_enrich.py index 80addb6d..ae1dcd3a 100644 --- a/tests/test_scans_enrich.py +++ b/tests/test_scans_enrich.py @@ -18,24 +18,54 @@ def _mock_db(current_scan=None, findings=None): def test_enrich_returns_202_and_enqueues_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} db = _mock_db(current_scan=scan, findings=[{"id": 1, "rule_id": "AZ-STOR-001"}]) - db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, True) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, "created") with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - assert resp.get_json()["status"] == "PENDING" + body = resp.get_json() + assert body["outcome"] == "created" + assert body["status"] == "pending" + assert body["job_id"] == _SCAN_ID db.enqueue_enrichment_job.assert_called_once_with(_SCAN_ID) def test_enrich_reuses_existing_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} db = _mock_db(current_scan=scan, findings=[{"id": 1}]) - db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, False) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, "active") with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - assert resp.get_json()["status"] == "running" + body = resp.get_json() + assert body["outcome"] == "active" + assert body["status"] == "running" + + +def test_enrich_requeues_a_terminally_failed_job(client, auth_headers): + """A failed job is the case a re-POST exists to recover from.""" + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "FAILED"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, "requeued") + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 202 + body = resp.get_json() + assert body["outcome"] == "requeued" + assert body["status"] == "pending" + assert "requeued" in body["message"] + + +def test_enrich_reports_an_already_completed_job_without_restarting_it(client, auth_headers): + # The scan header still reads PENDING, so the job row is the authority. + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "completed"}, "completed") + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 200 + assert resp.get_json()["outcome"] == "completed" def test_enrich_already_completed_returns_200(client, auth_headers): diff --git a/tests/test_worker.py b/tests/test_worker.py index f973b15f..8bbcae94 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -210,6 +210,68 @@ def test_worker_sleeps_when_no_scans_pending(self, mock_sleep, mock_env, mock_db mock_sleep.assert_called_with(POLL_INTERVAL_SECONDS) + @patch("scanner.worker.process_enrichment_job") + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.ScanEngine") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + @patch("scanner.worker.LeaseHeartbeat") + def test_enrichment_backlog_does_not_starve_pending_scans( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class, mock_process + ): + """With both queues permanently backlogged, each still gets served. + + Draining enrichment and restarting the loop would mean a scan is never + claimed while enrichment work keeps arriving. + """ + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_engine_class.return_value.run_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "findings": [], + "started_at": "2026-06-05T12:00:00Z", + } + mock_heartbeat_class.return_value.lost.is_set.return_value = False + + # Both queues always have work; stop after two full iterations. + mock_db.recover_stale_scans.side_effect = [None, None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = {"job_id": "job-1", "scan_id": self.scan_id} + mock_db.claim_next_pending_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "fencing_token": 1, + } + + with self.assertRaises(StopWorker): + run_worker() + + # Each iteration served one job from each queue -- neither starved. + self.assertEqual(mock_process.call_count, 2) + self.assertEqual(mock_db.save_scan.call_count, 2) + # A busy worker never idles. + mock_sleep.assert_not_called() + + @patch("scanner.worker.process_enrichment_job") + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + def test_enrichment_only_backlog_polls_again_without_idling( + self, mock_sleep, mock_env, mock_db_class, mock_process + ): + """Enrichment work must not be paced by the empty-queue poll interval.""" + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = {"job_id": "job-1", "scan_id": self.scan_id} + mock_db.claim_next_pending_scan.return_value = None + + with self.assertRaises(StopWorker): + run_worker() + + mock_process.assert_called_once() + mock_sleep.assert_not_called() + if __name__ == "__main__": unittest.main() From da65eff87380f51a7ab2bfd097adf8ca51753f0e Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Wed, 2 Sep 2026 01:35:55 +0100 Subject: [PATCH 11/19] test(worker): cover the lease/heartbeat interval clamp `lease_configuration()` guarantees the heartbeat interval stays strictly shorter than the lease - a worker that heartbeats no more often than its lease expires would lose its own claim mid-scan and have its results fenced out. .env.example documents that constraint, but nothing tested it. Covers the defaults, valid overrides, heartbeat == lease, heartbeat > lease, a lease small enough that `lease // 3` would floor to a zero-second heartbeat, and malformed/non-positive values falling back to the defaults. Verified the tests fail when the clamp is removed (3 failures) and pass when it is restored. Signed-off-by: Shaurya K Sharma --- tests/test_worker.py | 52 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/test_worker.py b/tests/test_worker.py index 8bbcae94..2d75c948 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -10,7 +10,14 @@ import unittest from unittest.mock import ANY, patch from api.models.finding import LostLease -from scanner.worker import LeaseHeartbeat, POLL_INTERVAL_SECONDS, run_worker +from scanner.worker import ( + DEFAULT_HEARTBEAT_SECONDS, + DEFAULT_LEASE_SECONDS, + LeaseHeartbeat, + POLL_INTERVAL_SECONDS, + lease_configuration, + run_worker, +) import uuid @@ -273,5 +280,48 @@ def test_enrichment_only_backlog_polls_again_without_idling( mock_sleep.assert_not_called() +class TestLeaseConfiguration(unittest.TestCase): + """A worker that heartbeats no more often than its lease expires would + lose its own claim mid-scan, so the configuration is clamped rather than + trusted. .env.example documents this; these tests hold it to it.""" + + def _configure(self, **env): + with patch.dict("scanner.worker.os.environ", env, clear=True): + return lease_configuration() + + def test_defaults_keep_the_heartbeat_shorter_than_the_lease(self): + lease, heartbeat = self._configure() + self.assertEqual((lease, heartbeat), (DEFAULT_LEASE_SECONDS, DEFAULT_HEARTBEAT_SECONDS)) + self.assertLess(heartbeat, lease) + + def test_valid_overrides_are_used_as_given(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="600", SCAN_HEARTBEAT_SECONDS="60") + self.assertEqual((lease, heartbeat), (600, 60)) + + def test_heartbeat_equal_to_the_lease_is_clamped_below_it(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="300", SCAN_HEARTBEAT_SECONDS="300") + self.assertEqual(lease, 300) + self.assertEqual(heartbeat, 100) + self.assertLess(heartbeat, lease) + + def test_heartbeat_longer_than_the_lease_is_clamped_below_it(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="300", SCAN_HEARTBEAT_SECONDS="9000") + self.assertEqual((lease, heartbeat), (300, 100)) + self.assertLess(heartbeat, lease) + + def test_a_tiny_lease_still_yields_a_positive_heartbeat(self): + # lease // 3 would floor to 0 and make the heartbeat thread spin. + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="2", SCAN_HEARTBEAT_SECONDS="2") + self.assertEqual((lease, heartbeat), (2, 1)) + self.assertGreater(heartbeat, 0) + + def test_malformed_or_non_positive_values_fall_back_to_defaults(self): + for bad in ("not-a-number", "0", "-30", ""): + with self.subTest(value=bad): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS=bad, SCAN_HEARTBEAT_SECONDS=bad) + self.assertEqual((lease, heartbeat), (DEFAULT_LEASE_SECONDS, DEFAULT_HEARTBEAT_SECONDS)) + self.assertLess(heartbeat, lease) + + if __name__ == "__main__": unittest.main() From 0ed794ff24f90893d140c076ac7c24239cdb4ee4 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:33:24 +0100 Subject: [PATCH 12/19] fix(core): rebase scan durability onto the #263 evaluation contract PR #321 landed the rule_evaluations coverage contract on dev while this branch was outstanding. Both branches added their first migration on top of d8e4f6a1b2c3, which forked the Alembic chain into two heads. Repoint the first lease migration at 3f59f83a5253 so the chain stays linear and `alembic heads` reports the single head d4a8c1e6b2f9. save_scan now performs #321's evaluation upsert and stale-coverage cleanup inside this change's fenced transaction rather than alongside it. The evaluation semantics are #321's, unchanged; what this adds is that a worker which lost its lease can no longer rewrite another owner's coverage rows. The finding upsert keeps RETURNING id so a replayed delivery reuses the existing finding row and FAIL evaluations stay linked to it instead of pointing at a row that was deleted and recreated. Signed-off-by: Shaurya K Sharma --- alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py index a1ffbc93..caad42dd 100644 --- a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py +++ b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py @@ -12,7 +12,7 @@ revision: str = "e4f7a9b2c6d8" -down_revision: Union[str, Sequence[str], None] = "d8e4f6a1b2c3" +down_revision: Union[str, Sequence[str], None] = "3f59f83a5253" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None From 0e79aa32f011ea9ce9b892f9c30bfadf7e5c5fb5 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:39:24 +0100 Subject: [PATCH 13/19] test(core): cover the #263 coverage contract under scan fencing Adapt #321's save_scan tests to the fenced signature: the ownership probe is answered by the mocked cursor so finding ids still start at 1 and the FAIL-to-finding linkage assertions keep their original meaning. Add PostgreSQL regression coverage for the integration itself: an owner writes both findings and coverage rows, a worker whose lease was reclaimed raises LostLease and leaves no coverage behind, and a replayed delivery converges on the same rows with the FAIL row still pointing at the same finding id rather than a renumbered one. rule_evaluations references scans, so the lease fixture now purges it before deleting the scan row. Signed-off-by: Shaurya K Sharma --- tests/test_rule_evaluations.py | 29 ++++++-- tests/test_scan_leases_postgres.py | 114 +++++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 5 deletions(-) diff --git a/tests/test_rule_evaluations.py b/tests/test_rule_evaluations.py index a2a5dac4..1712a28b 100644 --- a/tests/test_rule_evaluations.py +++ b/tests/test_rule_evaluations.py @@ -17,6 +17,10 @@ from scanner.evaluation import EvaluationStatus, RuleEvaluation, aggregate_status, subscription_scope_id _SUB = "00000000-0000-0000-0000-000000000001" +_SCAN_ID = "00000000-0000-0000-0000-000000000000" +# save_scan only writes while the caller still holds the scan lease (#303). +_OWNER = "worker-under-test" +_TOKEN = 1 # ── RuleEvaluation / EvaluationStatus contract ────────────────────────────── @@ -222,7 +226,22 @@ def _cursor(): cur.__exit__ = MagicMock(return_value=False) # Every INSERT ... RETURNING id call returns an incrementing fake id. ids = iter(range(1, 10_000)) - cur.fetchone.side_effect = lambda: (next(ids),) + last_sql = {"text": ""} + + def _execute(sql, *args, **kwargs): + last_sql["text"] = sql + return MagicMock() + + def _fetchone(): + # save_scan opens with the lease/fencing ownership probe. It must see + # an owned row, and it must not consume a finding id -- the linkage + # assertions below depend on findings starting at 1. + if "FOR UPDATE" in last_sql["text"]: + return (_SCAN_ID,) + return (next(ids),) + + cur.execute.side_effect = _execute + cur.fetchone.side_effect = _fetchone return cur @@ -275,7 +294,7 @@ def test_save_scan_persists_evaluations_and_links_fail_finding_id(): } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) + db.save_scan(result, _OWNER, _TOKEN) insert_calls = [c for c in cursor.execute.call_args_list if "INSERT INTO rule_evaluations" in c.args[0]] assert len(insert_calls) == 2 @@ -316,7 +335,7 @@ def test_save_scan_evaluations_upsert_on_conflict_instead_of_delete_first(): } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) + db.save_scan(result, _OWNER, _TOKEN) insert_calls = [c for c in cursor.execute.call_args_list if "INSERT INTO rule_evaluations" in c.args[0]] assert len(insert_calls) == 1 @@ -348,7 +367,7 @@ def test_save_scan_deletes_all_prior_evaluations_when_scan_reports_none(): } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) + db.save_scan(result, _OWNER, _TOKEN) delete_sql = [c.args[0] for c in cursor.execute.call_args_list if c.args[0].strip().startswith("DELETE")] assert any("rule_evaluations" in sql for sql in delete_sql) @@ -370,4 +389,4 @@ def test_save_scan_evaluations_default_to_empty_list_for_backward_compatible_cal } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) # must not raise + db.save_scan(result, _OWNER, _TOKEN) # must not raise diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index 2d19e6cd..22ced58c 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -81,6 +81,15 @@ def finding_count(self, scan_id: str) -> int: cur.execute("SELECT COUNT(*) FROM findings WHERE scan_id = %s", (scan_id,)) return cur.fetchone()[0] + def evaluations(self, scan_id: str) -> list[tuple]: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "SELECT rule_id, status, finding_id FROM rule_evaluations WHERE scan_id = %s ORDER BY rule_id", + (scan_id,), + ) + return cur.fetchall() + def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: """Simulate duplicate delivery of the same claimed result for persistence tests.""" with psycopg2.connect(self.dsn) as conn: @@ -100,6 +109,7 @@ def cleanup(self) -> None: with conn.cursor() as cur: for scan_id in self.scan_ids: cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) @@ -362,3 +372,107 @@ def test_distinct_finding_discriminators_preserve_multiple_violations(scan_rows) db.close() assert scan_rows.finding_count(scan_id) == 2 + + +# ── #263 coverage rows are written inside this change's fenced transaction ── + + +def _result_with_evaluations(scan_id: str, subscription_id: str) -> dict: + """A result carrying both a FAIL finding and its #263 coverage rows.""" + result = _result(scan_id, subscription_id) + resource_id = result["findings"][0]["resource_id"] + result["evaluations"] = [ + { + "rule_id": "AZ-LEASE-001", + "resource_id": resource_id, + "resource_type": "Test/resource", + "status": "FAIL", + "reason_code": None, + "reason": None, + "evidence": {}, + }, + { + "rule_id": "AZ-LEASE-002", + "resource_id": f"/subscriptions/{subscription_id}", + "resource_type": "", + "status": "UNKNOWN", + "reason_code": "LEGACY_RULE_NOT_MIGRATED", + "reason": "not migrated", + "evidence": {}, + }, + ] + return result + + +def test_owner_persists_evaluations_and_links_the_fail_row_to_its_finding(scan_rows): + """#321's coverage contract survives inside the fenced save_scan.""" + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a", scan_id) + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(_result_with_evaluations(scan_id, subscription_id), "worker-a", claim["fencing_token"]) + finally: + db.close() + + evaluations = scan_rows.evaluations(scan_id) + assert [(rule_id, status) for rule_id, status, _ in evaluations] == [ + ("AZ-LEASE-001", "FAIL"), + ("AZ-LEASE-002", "UNKNOWN"), + ] + # Only the FAIL row is linked, and it points at a real finding row. + assert evaluations[0][2] is not None + assert evaluations[1][2] is None + + +def test_stale_worker_cannot_write_rule_evaluations(scan_rows): + """A reclaimed scan's coverage rows belong to the new owner alone. + + Before #303's fencing wrapped these writes, a stale worker finishing a + long scan could still rewrite another owner's #263 coverage. + """ + scan_id, subscription_id = scan_rows.create() + first_claim = _claim(scan_rows.dsn, "worker-a", scan_id) + assert first_claim is not None + scan_rows.expire(scan_id) + _recover(scan_rows.dsn) + second_claim = _claim(scan_rows.dsn, "worker-b", scan_id) + assert second_claim is not None + assert second_claim["fencing_token"] > first_claim["fencing_token"] + + stale_db = DatabaseManager(scan_rows.dsn) + try: + with pytest.raises(LostLease): + stale_db.save_scan( + _result_with_evaluations(scan_id, subscription_id), + "worker-a", + first_claim["fencing_token"], + ) + finally: + stale_db.close() + + assert scan_rows.evaluations(scan_id) == [] + + +def test_replayed_delivery_keeps_evaluations_linked_to_the_same_finding(scan_rows): + """Duplicate delivery must converge, not renumber findings under coverage.""" + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a", scan_id) + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(_result_with_evaluations(scan_id, subscription_id), "worker-a", claim["fencing_token"]) + first = scan_rows.evaluations(scan_id) + + scan_rows.rearm(scan_id, "worker-a", claim["fencing_token"]) + db.save_scan(_result_with_evaluations(scan_id, subscription_id), "worker-a", claim["fencing_token"]) + second = scan_rows.evaluations(scan_id) + finally: + db.close() + + # Same two coverage rows, and the FAIL row still points at the *same* + # finding id: the upsert reused the row instead of deleting and recreating. + assert first == second + assert scan_rows.finding_count(scan_id) == 1 From 1e466ca84d0c605bc756d9f030ed83887658f740 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:46:34 +0100 Subject: [PATCH 14/19] fix(core): make stale scan recovery atomic and count attempts consistently recover_stale_scans issued two sequential UPDATEs whose attempt-count predicates disagreed: the fail branch read COALESCE(attempt_count, 1) while the retry branch read COALESCE(attempt_count, 0). A row predating the column therefore counted as one attempt already spent and was retired a run early. Reproduced: with max_attempts=1 a NULL row went straight to 'failed' even though the retry branch's own predicate admitted it. Neither statement took SKIP LOCKED, so recovery waited on any row another transaction held. Measured: one locked row blocked recovery for the full 3s the lock was held while no other stale scan made progress. This runs at the top of every worker iteration, so it stalls claiming and enrichment too. Collapse both branches into one CTE that selects candidates FOR UPDATE SKIP LOCKED and transitions them in the same statement. The same measured case now returns in 0.01s and recovers every stale scan except the held one. attempt_count counts claims already started (claim_next_pending_scan increments it with the lease), so a scan gets exactly max_attempts executions and NULL reads as zero. Documented on the method. Adds PostgreSQL regression coverage: two barrier-synchronised workers transition one stale scan exactly once and the winner's fencing token strictly advances; a row held by another transaction is stepped over rather than waited on; and the attempt budget is spent fully before a scan becomes terminal. Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 70 ++++++++------ tests/test_scan_leases_postgres.py | 142 +++++++++++++++++++++++++++++ 2 files changed, 184 insertions(+), 28 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index af564201..a3afe438 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1147,48 +1147,62 @@ def heartbeat_scan(self, scan_id: str, lease_owner: str, fencing_token: int, lea def recover_stale_scans(self, max_attempts: int = 3) -> int: """Recover scans only after their renewable leases have expired. - Stale scans are returned to pending while retry attempts remain. Once a - scan has reached max_attempts, it is marked failed so it cannot loop - forever on bad credentials or persistent Azure errors. + ``attempt_count`` counts claims that have already been *started*: + :meth:`claim_next_pending_scan` increments it as part of the same + statement that takes the lease, so a scan being executed for the + first time already reads 1. A stale scan is therefore returned to + ``pending`` while ``attempt_count < max_attempts`` and is failed once + it reaches that limit, giving every scan exactly ``max_attempts`` + executions. Rows predating the column read NULL and are treated as + zero attempts so they get the full budget rather than being retired a + run early. + + Selection and transition are one statement. The CTE takes + ``FOR UPDATE SKIP LOCKED`` so a row another worker is already + recovering (or executing) is stepped over instead of waited on: this + runs at the top of every worker iteration, so blocking here would + stall claiming and enrichment behind one stuck row. """ conn = self._get_conn() try: with conn.cursor() as cur: cur.execute( """ - UPDATE scans - SET status = 'failed', - lease_owner = NULL, - lease_expires_at = NULL, - error_message = 'Scan exceeded maximum retry attempts after worker interruption.' - WHERE status = 'running' - AND COALESCE(attempt_count, 1) >= %s - AND lease_expires_at < CURRENT_TIMESTAMP - """, - (max_attempts,), - ) - failed_count = cur.rowcount - - cur.execute( - """ - UPDATE scans - SET status = 'pending', - claimed_at = NULL, + WITH stale AS ( + SELECT scan_id, COALESCE(attempt_count, 0) AS attempts + FROM scans + WHERE status = 'running' + AND lease_expires_at < CURRENT_TIMESTAMP + ORDER BY lease_expires_at ASC + FOR UPDATE SKIP LOCKED + ) + UPDATE scans AS s + SET status = CASE WHEN stale.attempts >= %(max_attempts)s THEN 'failed' ELSE 'pending' END, + claimed_at = CASE WHEN stale.attempts >= %(max_attempts)s THEN s.claimed_at END, + last_heartbeat_at = CASE + WHEN stale.attempts >= %(max_attempts)s THEN s.last_heartbeat_at + END, lease_owner = NULL, lease_expires_at = NULL, - error_message = 'Scan worker interrupted before completion. Queued for retry.' - WHERE status = 'running' - AND COALESCE(attempt_count, 0) < %s - AND lease_expires_at < CURRENT_TIMESTAMP + error_message = CASE + WHEN stale.attempts >= %(max_attempts)s + THEN 'Scan exceeded maximum retry attempts after worker interruption.' + ELSE 'Scan worker interrupted before completion. Queued for retry.' + END + FROM stale + WHERE s.scan_id = stale.scan_id + RETURNING s.status """, - (max_attempts,), + {"max_attempts": max_attempts}, ) - retry_count = cur.rowcount + statuses = [row[0] for row in cur.fetchall()] conn.commit() except Exception: self.rollback(conn) raise - total_count = failed_count + retry_count + failed_count = statuses.count("failed") + retry_count = statuses.count("pending") + total_count = len(statuses) if total_count > 0: logger.info( "Recovered %d stale 'running' scans (%d retried, %d failed)", diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index 22ced58c..eeab4288 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -2,6 +2,7 @@ import os import threading +import time import uuid from datetime import datetime, timezone @@ -90,6 +91,28 @@ def evaluations(self, scan_id: str) -> list[tuple]: ) return cur.fetchall() + def make_stale(self, scan_id: str, attempt_count) -> None: + """Put a scan into 'running' with an expired lease and a chosen budget. + + attempt_count=None reproduces a row that predates the column. + """ + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET status = 'running', + attempt_count = %s, + fencing_token = 5, + lease_owner = 'dead-worker', + claimed_at = CURRENT_TIMESTAMP - INTERVAL '1 hour', + last_heartbeat_at = CURRENT_TIMESTAMP - INTERVAL '1 hour', + lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' + WHERE scan_id = %s + """, + (attempt_count, scan_id), + ) + def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: """Simulate duplicate delivery of the same claimed result for persistence tests.""" with psycopg2.connect(self.dsn) as conn: @@ -476,3 +499,122 @@ def test_replayed_delivery_keeps_evaluations_linked_to_the_same_finding(scan_row # finding id: the upsert reused the row instead of deleting and recreating. assert first == second assert scan_rows.finding_count(scan_id) == 1 + + +# ── Stale recovery is atomic, non-blocking, and counts attempts consistently ── + + +def test_concurrent_stale_recovery_transitions_each_scan_once(scan_rows): + """Two workers contending over one stale scan: exactly one recovers it. + + Both threads meet at a barrier so they genuinely overlap inside + recover_stale_scans rather than running one after the other. + """ + scan_id, _ = scan_rows.create() + scan_rows.make_stale(scan_id, 1) + + barrier = threading.Barrier(2) + recovered: dict[str, int] = {} + + def worker(name: str) -> None: + db = DatabaseManager(scan_rows.dsn) + try: + barrier.wait() + recovered[name] = db.recover_stale_scans(max_attempts=3) + finally: + db.close() + + threads = [threading.Thread(target=worker, args=(f"w{i}",)) for i in (1, 2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + scan = scan_rows.scan(scan_id) + assert scan["status"] == "pending" + assert scan["lease_owner"] is None + # The scan is requeued once, not handed to both recoverers. + assert sum(recovered.values()) == 1 + + # And the requeued scan can still only be claimed by a single owner, + # whose token strictly advances past the dead worker's. + first = _claim(scan_rows.dsn, "owner-1", scan_id) + second = _claim(scan_rows.dsn, "owner-2", scan_id) + assert first is not None and second is None + # The dead worker held token 5; the new owner's token strictly advances, + # so the old claimant's fenced writes can never be accepted again. + assert first["fencing_token"] > 5 + + +def test_stale_recovery_skips_rows_another_transaction_holds(scan_rows): + """One locked row must not stall recovery of every other stale scan. + + recover_stale_scans runs at the top of each worker iteration, so waiting + on a locked row would hold up scan claiming and enrichment behind it. + """ + blocked_id, _ = scan_rows.create() + others = [scan_rows.create()[0] for _ in range(3)] + for scan_id in [blocked_id, *others]: + scan_rows.make_stale(scan_id, 1) + + holder = psycopg2.connect(scan_rows.dsn) + holding = threading.Event() + release = threading.Event() + + def hold() -> None: + with holder.cursor() as cur: + cur.execute("SELECT scan_id FROM scans WHERE scan_id = %s FOR UPDATE", (blocked_id,)) + holding.set() + release.wait(timeout=10) + holder.rollback() + + thread = threading.Thread(target=hold) + thread.start() + try: + holding.wait(timeout=10) + started = time.perf_counter() + _recover(scan_rows.dsn) + elapsed = time.perf_counter() - started + finally: + release.set() + thread.join() + holder.close() + + # Returned promptly instead of waiting on the lock holder. + assert elapsed < 2.0 + # The unlocked scans were recovered; the locked one was left alone. + for scan_id in others: + assert scan_rows.scan(scan_id)["status"] == "pending" + assert scan_rows.scan(blocked_id)["status"] == "running" + + +def test_recovery_retries_until_the_attempt_budget_is_actually_spent(scan_rows): + """attempt_count counts claims already started, and both branches agree. + + The fail and retry branches previously defaulted a NULL attempt_count + differently (1 vs 0), so a row predating the column was retired one run + early instead of getting its full budget. + """ + # A row from before attempt_count existed still gets its first run. + legacy_id, _ = scan_rows.create() + scan_rows.make_stale(legacy_id, None) + _recover(scan_rows.dsn, max_attempts=1) + assert scan_rows.scan(legacy_id)["status"] == "pending" + + # One claim spends that budget, and only then is it terminal. + assert _claim(scan_rows.dsn, "worker-a", legacy_id) is not None + scan_rows.expire(legacy_id) + _recover(scan_rows.dsn, max_attempts=1) + assert scan_rows.scan(legacy_id)["status"] == "failed" + + # The boundary holds for ordinary rows too: attempts below the limit are + # requeued, and reaching the limit is terminal. + below_id, _ = scan_rows.create() + scan_rows.make_stale(below_id, 2) + _recover(scan_rows.dsn, max_attempts=3) + assert scan_rows.scan(below_id)["status"] == "pending" + + at_limit_id, _ = scan_rows.create() + scan_rows.make_stale(at_limit_id, 3) + _recover(scan_rows.dsn, max_attempts=3) + assert scan_rows.scan(at_limit_id)["status"] == "failed" From 9e0da999910acab4dabc606ed7f809d563610da7 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:48:07 +0100 Subject: [PATCH 15/19] fix(observability): report scan lease age from heartbeat freshness oldest_lease_age.scan measured CURRENT_TIMESTAMP - MIN(claimed_at), so it grew for the entire life of a scan even while its worker renewed the lease on schedule. A healthy long-running scan was indistinguishable from a stalled one, which is the situation the metric exists to detect. Measure MIN(COALESCE(last_heartbeat_at, claimed_at)) instead: that is what the lease actually renews, and it matches what the enrichment counterpart already reported. Rows migrated into leases have no heartbeat yet, so they fall back to their claim time rather than dropping out of the aggregate. Labels are unchanged, so scrape cardinality is unaffected. Regression test: a scan claimed an hour ago reports >= 3600s while its worker is silent, then drops as soon as the worker heartbeats even though claimed_at is still an hour old. Against the previous query that second assertion reported 3600.02s. Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 13 +++- tests/test_operational_metrics_postgres.py | 74 ++++++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/api/models/finding.py b/api/models/finding.py index a3afe438..94137d9c 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1301,9 +1301,20 @@ def get_operational_metrics(self) -> Dict[str, Any]: """ ) enrichment_queue_age = cur.fetchone()["value"] + # Lease age is the freshness of the *current* lease, not how + # long ago the scan was first claimed. Measuring claimed_at + # made this climb for the whole life of a healthy long scan + # even while its worker was heartbeating on schedule, so the + # metric could not distinguish "slow but alive" from "stalled". + # last_heartbeat_at is what the lease actually renews, and is + # what the enrichment counterpart below already reports. + # Rows migrated into leases predate any heartbeat, so they fall + # back to their claim time rather than dropping out of MIN(). cur.execute( """ - SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(claimed_at)), 0) AS value + SELECT COALESCE( + EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(COALESCE(last_heartbeat_at, claimed_at))), 0 + ) AS value FROM scans WHERE status = 'running' """ ) diff --git a/tests/test_operational_metrics_postgres.py b/tests/test_operational_metrics_postgres.py index c8df812b..ae26cbee 100644 --- a/tests/test_operational_metrics_postgres.py +++ b/tests/test_operational_metrics_postgres.py @@ -80,3 +80,77 @@ def test_stale_worker_heartbeats_are_pruned_without_dropping_live_workers(): "DELETE FROM worker_heartbeats WHERE worker_id IN (%s, %s)", (live_worker, dead_worker), ) + + +def _oldest_other_running_lease_age(dsn: str, exclude_scan_id: str) -> float: + """Age the metric would report from running scans other than this one. + + oldest_lease_age is a whole-table MIN, so a shared test database can carry + unrelated running rows. Measuring them separately keeps the assertions + below about this scan's lease rather than about the fixture ordering. + """ + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + SELECT COALESCE( + EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(COALESCE(last_heartbeat_at, claimed_at))), 0 + ) + FROM scans WHERE status = 'running' AND scan_id <> %s + """, + (exclude_scan_id,), + ) + return float(cur.fetchone()[0]) + + +def test_scan_lease_age_tracks_heartbeat_freshness_not_claim_time(): + """A healthy long scan must not look like a stalled one. + + The metric previously reported CURRENT_TIMESTAMP - MIN(claimed_at), which + grows for the entire life of a scan even while its worker renews the lease + on schedule. It must report the freshness of the current lease instead. + """ + dsn = os.environ["DATABASE_URL"] + scan_id, subscription_id = str(uuid.uuid4()), str(uuid.uuid4()) + db = DatabaseManager(dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + claim = db.claim_next_pending_scan("worker-a", 120, scan_id=scan_id) + assert claim is not None + + # A scan claimed an hour ago whose worker has not checked in since. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET claimed_at = CURRENT_TIMESTAMP - INTERVAL '1 hour', + last_heartbeat_at = CURRENT_TIMESTAMP - INTERVAL '1 hour', + lease_expires_at = CURRENT_TIMESTAMP + INTERVAL '1 hour' + WHERE scan_id = %s + """, + (scan_id,), + ) + + # A genuinely stale lease reports its real age. + assert db.get_operational_metrics()["oldest_lease_age"]["scan"] >= 3600 + + # The worker checks in on time. claimed_at is deliberately left an + # hour old: the lease is fresh even though the scan started long ago. + db.heartbeat_scan(scan_id, "worker-a", claim["fencing_token"], 120) + + after = db.get_operational_metrics()["oldest_lease_age"]["scan"] + others = _oldest_other_running_lease_age(dsn, scan_id) + # This scan no longer contributes an hour of age; anything still + # reported comes from unrelated running rows, not from this one. + assert after < 3600 or after <= others + 5 + if others < 60: + assert after < 60 + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) From 18693a06a82e0bb19f5f364f0f1b92eeed3bc216 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:55:12 +0100 Subject: [PATCH 16/19] fix(api): drop the unreachable request fingerprint and 409 path request_fingerprint hashed {"subscription_id": ...} only, while admission already looked the key up as (subscription_id, idempotency_key). Every request that could reach the comparison therefore carried an identical fingerprint, so ScanAdmissionConflict and the documented 409 were dead. The existing test only reached the 409 by passing hand-written fingerprints straight to admit_scan, which is why it went unnoticed. Making the fingerprint real would need a second semantic request input, and there is none: trigger_scan rejects every body field except subscription_id, and docs/api-reference.md deliberately scopes a key to one subscription. Widening keys to be global would create a reachable 409 but contradict that documented contract, so the dead model is removed instead of being propped up. Removes the column from a7c5e9d2f1b4 (unreleased in this branch), the parameter and comparison from admit_scan, the exception type, and the 409 row from the API reference, which now states why no changed-payload conflict exists. Tests: admission replays a repeated key and treats the same key under a different subscription as its own scan; the route returns 200 with the original scan_id on replay and passes no fingerprint to admission. Signed-off-by: Shaurya K Sharma --- ...a7c5e9d2f1b4_scan_admission_idempotency.py | 2 - api/models/finding.py | 18 ++-- api/routes/scans.py | 15 ++- docs/api-reference.md | 3 +- tests/test_async_scan_persistence.py | 100 ++++++++++++++---- tests/test_scan_admission_postgres.py | 40 ++++--- 6 files changed, 117 insertions(+), 61 deletions(-) diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py index e4daf0a2..2ca69d88 100644 --- a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -63,7 +63,6 @@ def _assert_one_active_scan_per_subscription() -> None: def upgrade() -> None: """Persist idempotency semantics and prevent more than one active scan.""" op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) - op.add_column("scans", sa.Column("request_fingerprint", sa.Text(), nullable=True)) # Checked before either index is built so a blocked upgrade leaves the # schema exactly as it was, with the added columns unused and harmless. @@ -96,5 +95,4 @@ def downgrade() -> None: with op.get_context().autocommit_block(): op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_ACTIVE_INDEX}") op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_KEY_INDEX}") - op.drop_column("scans", "request_fingerprint") op.drop_column("scans", "idempotency_key") diff --git a/api/models/finding.py b/api/models/finding.py index 94137d9c..77b045c5 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -36,10 +36,6 @@ class LostLease(RuntimeError): """Raised when a worker no longer owns the scan it is trying to update.""" -class ScanAdmissionConflict(RuntimeError): - """Raised when an idempotency key is reused for different scan semantics.""" - - class ScanQuotaExceeded(RuntimeError): """Raised when an explicitly configured subscription scan quota is exhausted.""" @@ -937,7 +933,6 @@ def admit_scan( subscription_id: str, *, idempotency_key: Optional[str] = None, - request_fingerprint: Optional[str] = None, max_scans_per_hour: int = 0, ) -> tuple[Dict[str, Any], bool]: """Atomically admit one scan or return its durable logical predecessor. @@ -962,11 +957,11 @@ def admit_scan( ) existing = cur.fetchone() if existing: - existing = dict(existing) - if existing.get("request_fingerprint") != request_fingerprint: - raise ScanAdmissionConflict("Idempotency-Key was reused with different request semantics") + # The key is scoped to this subscription and a trigger + # carries no other semantic input, so a hit here is + # always a replay of the same logical request. conn.commit() - return existing, False + return dict(existing), False cur.execute( """ @@ -1002,9 +997,9 @@ def admit_scan( """ INSERT INTO scans ( scan_id, subscription_id, started_at, status, attempt_count, - idempotency_key, request_fingerprint + idempotency_key ) - VALUES (%s, %s, %s, 'pending', 0, %s, %s) + VALUES (%s, %s, %s, 'pending', 0, %s) RETURNING * """, ( @@ -1012,7 +1007,6 @@ def admit_scan( subscription_id, datetime.now(timezone.utc).isoformat(), idempotency_key, - request_fingerprint, ), ) admitted = dict(cur.fetchone()) diff --git a/api/routes/scans.py b/api/routes/scans.py index 89907e33..b7850ff8 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -2,12 +2,10 @@ import logging import os -import hashlib -import json import uuid from flask import Blueprint, g, jsonify, request -from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded +from api.models.finding import DatabaseManager, ScanQuotaExceeded from api.validation import ( VALIDATION_ERROR_MESSAGE, ValidationError, @@ -117,14 +115,16 @@ def trigger_scan(): logger.warning("Scan trigger rejected: subscription %s is not on the authorized allowlist", subscription_id) return jsonify({"error": "Subscription is not authorized for this deployment"}), 403 + # A trigger's only semantic input is subscription_id, and an + # Idempotency-Key is scoped to one subscription (see + # docs/api-reference.md). Two requests carrying the same key under the + # same subscription are therefore always the same logical request, + # which is why admission needs no separate request fingerprint. idempotency_key = request.headers.get("Idempotency-Key") if idempotency_key is not None: idempotency_key = idempotency_key.strip() if not idempotency_key or len(idempotency_key) > 200: return jsonify({"error": "Idempotency-Key must be between 1 and 200 characters"}), 400 - request_fingerprint = hashlib.sha256( - json.dumps({"subscription_id": subscription_id}, sort_keys=True, separators=(",", ":")).encode("utf-8") - ).hexdigest() scan_id = str(uuid.uuid4()) try: @@ -133,11 +133,8 @@ def trigger_scan(): scan_id, subscription_id, idempotency_key=idempotency_key, - request_fingerprint=request_fingerprint, max_scans_per_hour=_configured_hourly_quota(), ) - except ScanAdmissionConflict: - return jsonify({"error": "Idempotency-Key is already associated with a different request."}), 409 except ScanQuotaExceeded: return jsonify({"error": "Scan quota exceeded for this subscription."}), 429 except Exception as exc: diff --git a/docs/api-reference.md b/docs/api-reference.md index f5e732f6..8cb24c1c 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -221,7 +221,7 @@ Request body (optional — falls back to `AZURE_SUBSCRIPTION_ID`): Admission is serialized per subscription and enforced by the database, so concurrent and replayed triggers converge on one logical scan rather than creating duplicates: - **At most one active scan per subscription.** While a `pending` or `running` scan exists, a further trigger returns that existing scan instead of queueing another. -- **`Idempotency-Key` (optional request header, 1–200 characters).** A repeat of the same key for the same subscription returns the original scan. The key is scoped to the subscription; the same key under a different subscription is a different request. +- **`Idempotency-Key` (optional request header, 1–200 characters).** A repeat of the same key for the same subscription returns the original scan. The key is scoped to the subscription; the same key under a different subscription is a different request. A trigger carries no request input other than `subscription_id`, so a key that resolves to an existing scan is always a replay of the same logical request and there is no changed-payload conflict to report. - **`OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR`** adds an optional hourly admission quota. Unset or `0` (the default) applies no time-window limit; the one-active-scan rule still applies. ### Responses @@ -232,7 +232,6 @@ Admission is serialized per subscription and enforced by the database, so concur | `200 OK` | The request resolved to an existing logical scan — an `Idempotency-Key` replay of the same request, or a trigger while a scan is already active for the subscription. | `scan_id`, `status` (the existing scan's `pending`/`running`), `message: "Existing logical scan returned."` | | `400 Bad Request` | Malformed body, invalid `subscription_id`, missing subscription, or an `Idempotency-Key` outside 1–200 characters. | `error` | | `403 Forbidden` | `subscription_id` is not on the `OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS` allowlist. | `error` | -| `409 Conflict` | The `Idempotency-Key` was reused with a different request payload. | `error: "Idempotency-Key is already associated with a different request."` | | `429 Too Many Requests` | The configured hourly quota for this subscription is exhausted. | `error: "Scan quota exceeded for this subscription."` | New scan (`202`): diff --git a/tests/test_async_scan_persistence.py b/tests/test_async_scan_persistence.py index 2883377d..4495e959 100644 --- a/tests/test_async_scan_persistence.py +++ b/tests/test_async_scan_persistence.py @@ -35,6 +35,10 @@ def fetchone(self): return self.rows.pop(0) return None + def fetchall(self): + rows, self.rows = self.rows, [] + return rows + def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, monkeypatch): """POST /api/scans/trigger should create a pending DB row, not an in-memory job.""" @@ -64,6 +68,53 @@ def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, mo assert mock_db.admit_scan.call_args.args[:2] == (scan_id, subscription_id) +def test_trigger_scan_replays_an_existing_scan_for_a_repeated_key(client, auth_headers, monkeypatch): + """A repeated Idempotency-Key returns the original scan with 200, not a new one.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db") + scan_id = "11111111-1111-1111-1111-111111111111" + subscription_id = "00000000-0000-0000-0000-000000000000" + mock_db = MagicMock() + # created=False is how admission reports "this resolved to an existing scan". + mock_db.admit_scan.return_value = ({"scan_id": scan_id, "status": "running"}, False) + + with patch("api.routes.scans.DatabaseManager", return_value=mock_db): + resp = client.post( + "/api/scans/trigger", + json={"subscription_id": subscription_id}, + headers={**auth_headers, "Idempotency-Key": "repeat-me"}, + ) + + assert resp.status_code == 200 + assert resp.get_json() == { + "scan_id": scan_id, + "status": "running", + "message": "Existing logical scan returned.", + } + assert mock_db.admit_scan.call_args.kwargs["idempotency_key"] == "repeat-me" + + +def test_trigger_scan_sends_no_request_fingerprint(client, auth_headers, monkeypatch): + """Admission takes the key alone; there is no second request identity. + + A trigger's only semantic input is subscription_id and keys are scoped to + a subscription, so a fingerprint derived from the request could never + differ between two requests that shared a key. It was removed rather than + left as an unreachable 409 path. + """ + monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db") + mock_db = MagicMock() + mock_db.admit_scan.return_value = ({"scan_id": "x", "status": "pending"}, True) + + with patch("api.routes.scans.DatabaseManager", return_value=mock_db): + client.post( + "/api/scans/trigger", + json={"subscription_id": "00000000-0000-0000-0000-000000000000"}, + headers={**auth_headers, "Idempotency-Key": "k"}, + ) + + assert "request_fingerprint" not in mock_db.admit_scan.call_args.kwargs + + def test_get_scan_status_reads_from_database(client, auth_headers, monkeypatch): """GET /api/scans/ should read durable status from PostgreSQL.""" monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db") @@ -182,44 +233,47 @@ def test_fenced_failure_rejects_a_stale_owner(): assert "lease_expires_at > CURRENT_TIMESTAMP" in sql -def test_recover_stale_scans_retries_before_max_attempts(): - """Stale running scans should return to pending while attempts remain.""" +def test_recover_stale_scans_transitions_candidates_in_one_locked_statement(): + """Selection and transition must be a single SKIP LOCKED statement. + + Two sequential UPDATEs waited on any row another transaction held, which + stalled the whole worker loop because recovery runs before claiming. + """ db = DatabaseManager.__new__(DatabaseManager) - cursor = _Cursor(rowcounts=[0, 1]) + cursor = _Cursor(rows=[("pending",), ("failed",)]) conn = MagicMock() conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): recovered = db.recover_stale_scans(max_attempts=3) - failed_sql, failed_params = cursor.calls[0] - retry_sql, retry_params = cursor.calls[1] - assert "status = 'failed'" in failed_sql - assert failed_params == (3,) - assert "lease_expires_at < CURRENT_TIMESTAMP" in failed_sql - assert "status = 'pending'" in retry_sql - assert "claimed_at = NULL" in retry_sql - assert "lease_owner = NULL" in retry_sql - assert retry_params == (3,) - assert recovered == 1 + assert len(cursor.calls) == 1 + sql, params = cursor.calls[0] + assert "FOR UPDATE SKIP LOCKED" in sql + assert "lease_expires_at < CURRENT_TIMESTAMP" in sql + assert "status = 'running'" in sql + assert "lease_owner = NULL" in sql + assert params == {"max_attempts": 3} + # Both transitions are reported from the one statement's RETURNING rows. + assert recovered == 2 conn.commit.assert_called_once() -def test_recover_stale_scans_fails_after_max_attempts(): - """Stale scans at the attempt limit should fail instead of retrying forever.""" +def test_recover_stale_scans_defaults_a_missing_attempt_count_to_zero(): + """Retry and exhaustion must read attempt_count the same way. + + The fail branch previously defaulted NULL to 1 while the retry branch + defaulted it to 0, retiring rows that predate the column a run early. + """ db = DatabaseManager.__new__(DatabaseManager) - cursor = _Cursor(rowcounts=[1, 0]) + cursor = _Cursor(rows=[("failed",)]) conn = MagicMock() conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): recovered = db.recover_stale_scans(max_attempts=3) - failed_sql, failed_params = cursor.calls[0] - retry_sql, retry_params = cursor.calls[1] - assert "COALESCE(attempt_count, 1) >= %s" in failed_sql - assert failed_params == (3,) - assert "COALESCE(attempt_count, 0) < %s" in retry_sql - assert retry_params == (3,) + sql, _params = cursor.calls[0] + assert "COALESCE(attempt_count, 0)" in sql + assert "COALESCE(attempt_count, 1)" not in sql assert recovered == 1 - conn.commit.assert_called_once() diff --git a/tests/test_scan_admission_postgres.py b/tests/test_scan_admission_postgres.py index 05ed909f..e1d19ed2 100644 --- a/tests/test_scan_admission_postgres.py +++ b/tests/test_scan_admission_postgres.py @@ -7,7 +7,7 @@ import psycopg2 import pytest -from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded +from api.models.finding import DatabaseManager, ScanQuotaExceeded pytestmark = pytest.mark.skipif( @@ -27,15 +27,10 @@ def admitted_scans(): cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) -def _admit(dsn: str, subscription_id: str, key: str | None = None, fingerprint: str = "same"): +def _admit(dsn: str, subscription_id: str, key: str | None = None): db = DatabaseManager(dsn) try: - return db.admit_scan( - str(uuid.uuid4()), - subscription_id, - idempotency_key=key, - request_fingerprint=fingerprint if key else None, - ) + return db.admit_scan(str(uuid.uuid4()), subscription_id, idempotency_key=key) finally: db.close() @@ -61,18 +56,37 @@ def admit() -> None: assert sum(created for _scan, created in outcomes) == 1 -def test_idempotency_key_replays_or_rejects_changed_semantics(admitted_scans): +def test_idempotency_key_replays_the_same_logical_scan(admitted_scans): dsn, scan_ids = admitted_scans subscription_id = str(uuid.uuid4()) - first, created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + key = f"request-{uuid.uuid4()}" + first, created = _admit(dsn, subscription_id, key) scan_ids.append(str(first["scan_id"])) - replay, replay_created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + replay, replay_created = _admit(dsn, subscription_id, key) assert created is True assert replay_created is False assert replay["scan_id"] == first["scan_id"] - with pytest.raises(ScanAdmissionConflict): - _admit(dsn, subscription_id, "request-1", "fingerprint-b") + + +def test_idempotency_key_is_scoped_to_its_subscription(admitted_scans): + """The documented scope: one key means different things per subscription. + + A trigger's only semantic input is subscription_id, so within a + subscription a key hit is always a replay. Reusing the key under another + subscription is a genuinely different request and admits its own scan + rather than replaying or conflicting. + """ + dsn, scan_ids = admitted_scans + key = f"shared-{uuid.uuid4()}" + first, first_created = _admit(dsn, str(uuid.uuid4()), key) + scan_ids.append(str(first["scan_id"])) + second, second_created = _admit(dsn, str(uuid.uuid4()), key) + scan_ids.append(str(second["scan_id"])) + + assert first_created is True + assert second_created is True + assert first["scan_id"] != second["scan_id"] def test_completed_scan_allows_a_later_admission_and_configured_quota(admitted_scans): From 364f5f3bea4304102d8a9fe8bfe8f60bc90201be Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 18:58:11 +0100 Subject: [PATCH 17/19] fix(api): give /enrich one response contract for every outcome The route short-circuited a COMPLETED scan to {message, scan_id} with 200, while docs/api-reference.md and every other outcome use {scan_id, job_id, status, outcome, message}. Clients had to special-case the one response that carried no job_id. Remove the early return so all four outcomes come from enqueue_enrichment_job. Two details this exposes, both handled rather than absorbed as behaviour changes: A scan enriched before durable jobs existed reads COMPLETED but has no job row, so a plain insert would queue fresh work and report "created". enqueue_enrichment_job now records that work as already finished and returns "completed", so the scan is not silently re-enriched and the caller still gets a real job_id. A clean scan can finish enrichment with nothing to enrich, so the findings 404 guard is skipped once a scan is enriched. Without that, removing the early return would have turned an existing 200 into a 404. Tests: every outcome returns the same keys with the documented status code; the completed case asserts job_id/outcome/status rather than a message; and a PostgreSQL test reproduces the legacy no-job-row scan and asserts it resolves to completed without re-queueing. Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 21 ++++++++++--- api/routes/scans.py | 14 +++++---- docs/api-reference.md | 4 ++- tests/test_enrichment_jobs_postgres.py | 36 ++++++++++++++++++++++ tests/test_scans_enrich.py | 42 ++++++++++++++++++++++++-- 5 files changed, 103 insertions(+), 14 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index 77b045c5..7c3af289 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -620,23 +620,36 @@ def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], str]: conn = self._get_conn() try: with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + # A scan enriched before durable jobs existed has no job row + # to report, but it must not be re-enriched either. Record the + # work as already finished so every caller gets the same + # {outcome, job_id} answer instead of a second contract. cur.execute( """ - INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint) - VALUES (%s, %s, 'pending', 0, 0) + INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint, completed_at) + SELECT %s, %s, + CASE WHEN s.cve_enrichment_status = 'COMPLETED' THEN 'completed' ELSE 'pending' END, + 0, 0, + CASE WHEN s.cve_enrichment_status = 'COMPLETED' THEN CURRENT_TIMESTAMP END + FROM scans s + WHERE s.scan_id = %s ON CONFLICT (scan_id) DO NOTHING RETURNING * """, - (str(uuid.uuid4()), scan_id), + (str(uuid.uuid4()), scan_id, scan_id), ) job = cur.fetchone() if job is not None: + job = dict(job) + if job["status"] == "completed": + conn.commit() + return job, "completed" cur.execute( "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", (scan_id,), ) conn.commit() - return dict(job), "created" + return job, "created" # A job already exists. Only a terminally failed one is # revived, and the WHERE clause is the whole guard: a diff --git a/api/routes/scans.py b/api/routes/scans.py index b7850ff8..3aed6dfe 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -186,12 +186,14 @@ def enrich_scan(scan_id): if not current_scan: return jsonify({"error": "Scan not found"}), 404 - status = current_scan.get("cve_enrichment_status") - if status == "COMPLETED": - return jsonify({"message": "Scan already enriched", "scan_id": scan_id}), 200 - findings = db.get_findings({"scan_id": scan_id}) - if not findings: - return jsonify({"error": "No findings found for this scan"}), 404 + # Every outcome is reported in one shape by enqueue_enrichment_job, + # including the already-enriched case, so there is no second response + # contract for a completed scan. The findings guard is skipped for an + # enriched scan: a clean scan legitimately finishes enrichment with + # nothing to enrich, and must still report completion rather than 404. + if current_scan.get("cve_enrichment_status") != "COMPLETED": + if not db.get_findings({"scan_id": scan_id}): + return jsonify({"error": "No findings found for this scan"}), 404 job, outcome = db.enqueue_enrichment_job(scan_id) body = { diff --git a/docs/api-reference.md b/docs/api-reference.md index 8cb24c1c..14dfb333 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -280,7 +280,9 @@ Every response carries `scan_id`, `job_id`, `status` (the job row's state) and a | `202 Accepted` | `requeued` | A previously **failed** job was reset to `pending` and will be retried. | | `202 Accepted` | `active` | A `pending` or `running` job already exists and was returned unchanged. A live claim is never interrupted. | | `200 OK` | `completed` | Enrichment already finished; nothing was restarted. | -| `404 Not Found` | — | Unknown `scan_id`, or the scan has no findings to enrich. | +| `404 Not Found` | — | Unknown `scan_id`, or the scan has no findings to enrich and has not already been enriched. | + +An already-enriched scan always reports `completed`, including a clean scan that had no findings to enrich in the first place. A job that exhausts its retry budget becomes `failed`. Re-POSTing this endpoint is the supported operator recovery: it atomically returns the job to `pending` with a fresh retry budget, clears the lease, and keeps the last `error_message` and the `checkpoint` so the retry resumes rather than re-enriching findings that already succeeded. Concurrent re-POSTs converge — exactly one reports `requeued` and the rest report `active`. diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index 0b4780b0..808d7359 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -370,3 +370,39 @@ def test_requeued_job_can_eventually_complete(enrichment_scan): with conn.cursor() as cur: cur.execute("SELECT cve_enrichment_status FROM scans WHERE scan_id = %s", (scan_id,)) assert cur.fetchone()[0] == "COMPLETED" + + +def test_scan_enriched_before_durable_jobs_reports_completed_not_a_new_job(enrichment_scan): + """A pre-durable-jobs enrichment must not be silently redone. + + Such a scan carries cve_enrichment_status COMPLETED but has no job row, + so a plain insert would queue fresh work and answer "created". It has to + resolve to the same completed outcome every other caller sees. + """ + dsn, scan_id, _ = enrichment_scan + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + # Reproduce the legacy shape: enriched scan, no durable job row. + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + cur.execute("UPDATE scans SET cve_enrichment_status = 'COMPLETED' WHERE scan_id = %s", (scan_id,)) + + db = DatabaseManager(dsn) + try: + job, outcome = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + + assert outcome == "completed" + assert job["status"] == "completed" + assert job["job_id"] is not None + # The scan is not dragged back into the queue. + row = _job_row(dsn, scan_id) + assert row["status"] == "completed" + assert _scan_enrichment_status(dsn, scan_id) == "COMPLETED" + + +def _scan_enrichment_status(dsn: str, scan_id: str) -> str: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT cve_enrichment_status FROM scans WHERE scan_id = %s", (scan_id,)) + return cur.fetchone()[0] diff --git a/tests/test_scans_enrich.py b/tests/test_scans_enrich.py index ae1dcd3a..234dd956 100644 --- a/tests/test_scans_enrich.py +++ b/tests/test_scans_enrich.py @@ -6,6 +6,7 @@ _SCAN_ID = "00000000-0000-0000-0000-000000000001" +_JOB_ID = "00000000-0000-0000-0000-0000000000aa" def _mock_db(current_scan=None, findings=None): @@ -68,13 +69,48 @@ def test_enrich_reports_an_already_completed_job_without_restarting_it(client, a assert resp.get_json()["outcome"] == "completed" -def test_enrich_already_completed_returns_200(client, auth_headers): +def test_enrich_already_completed_uses_the_canonical_response(client, auth_headers): + """An enriched scan answers in the documented {outcome, job_id} shape. + + This used to short-circuit to {message, scan_id}, which is neither the + documented contract nor what any other outcome returns. + """ scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "COMPLETED"} - db = _mock_db(current_scan=scan) + # A clean scan can finish enrichment with nothing to enrich, so findings + # are deliberately empty here: completion must still win over the 404. + db = _mock_db(current_scan=scan, findings=[]) + db.enqueue_enrichment_job.return_value = ({"job_id": _JOB_ID, "status": "completed"}, "completed") with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 200 - assert "already enriched" in resp.get_json()["message"] + body = resp.get_json() + assert body["outcome"] == "completed" + assert body["job_id"] == _JOB_ID + assert body["status"] == "completed" + assert body["scan_id"] == _SCAN_ID + assert "already enriched" in body["message"] + db.enqueue_enrichment_job.assert_called_once_with(_SCAN_ID) + + +def test_enrich_responses_share_one_contract_across_every_outcome(client, auth_headers): + """created/requeued/active/completed all return the same keys.""" + expected = {"scan_id", "job_id", "status", "outcome", "message"} + cases = [ + ("created", "pending", 202), + ("requeued", "pending", 202), + ("active", "running", 202), + ("completed", "completed", 200), + ] + for outcome, job_status, code in cases: + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _JOB_ID, "status": job_status}, outcome) + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == code, outcome + assert set(resp.get_json()) == expected, outcome + assert resp.get_json()["outcome"] == outcome def test_enrich_missing_scan_returns_404(client, auth_headers): From a966c8b609416675f4c9442ff82086f80e1e4164 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Tue, 8 Sep 2026 19:03:18 +0100 Subject: [PATCH 18/19] fix(db): make every migration in this branch safe to re-run alembic/env.py does not set transaction_per_migration, so the whole upgrade shares one transaction -- but each autocommit_block() commits it. Any DDL issued before a CONCURRENTLY index build is therefore already durable when that build fails, while alembic_version still names the previous revision. The retry then died on its own committed work before it could reach the index recovery. This was reported against a7c5e9d2f1b4 but applies to all five revisions here, so each is fixed rather than only the one that was noticed: - e4f7a9b2c6d8, f2b6d8e1a4c9, a7c5e9d2f1b4 add their columns with ADD COLUMN IF NOT EXISTS. All are nullable or carry a default, and the finding_key backfill and NOT NULL tightening were already idempotent. - c9e1a5b7d3f2 and d4a8c1e6b2f9 skip their CREATE TABLE when the table is already present. - Every CONCURRENTLY build now drops its index name first. CREATE INDEX ... IF NOT EXISTS would have kept an INVALID index from an interrupted build, which owns the name but can never serve a query. No data is touched: the duplicate-active-scan preflight still refuses to choose which scan history to discard. Tests replay the real partial states against a throwaway database: columns committed with the version stamp behind, a table committed the same way, and an index marked invalid the way an interrupted build leaves it. All three previously failed with DuplicateColumn/DuplicateTable; they now reach head with the index rebuilt valid. A fourth covers the deployment path itself, dev head -> this head, asserting rule_evaluations survives. Signed-off-by: Shaurya K Sharma --- ...a7c5e9d2f1b4_scan_admission_idempotency.py | 5 +- .../c9e1a5b7d3f2_durable_enrichment_jobs.py | 59 ++++---- ...d4a8c1e6b2f9_operational_worker_metrics.py | 29 ++-- .../e4f7a9b2c6d8_scan_leases_and_fencing.py | 23 +-- ...6d8e1a4c9_idempotent_finding_identities.py | 6 +- .../test_scan_admission_migration_postgres.py | 132 ++++++++++++++++++ 6 files changed, 208 insertions(+), 46 deletions(-) diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py index 2ca69d88..38b927e8 100644 --- a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -62,7 +62,10 @@ def _assert_one_active_scan_per_subscription() -> None: def upgrade() -> None: """Persist idempotency semantics and prevent more than one active scan.""" - op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) + # Added before the preflight and the concurrent index builds, both of + # which can fail; autocommit_block() has already committed this column by + # then, so a retry must tolerate it already being there. + op.execute("ALTER TABLE scans ADD COLUMN IF NOT EXISTS idempotency_key TEXT") # Checked before either index is built so a blocked upgrade leaves the # schema exactly as it was, with the added columns unused and harmless. diff --git a/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py index ab9c3a0a..9d95ff97 100644 --- a/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py +++ b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py @@ -20,31 +20,41 @@ def upgrade() -> None: """Create one resumable enrichment job per scan.""" - op.create_table( - "enrichment_jobs", - sa.Column("job_id", postgresql.UUID(), nullable=False), - sa.Column("scan_id", postgresql.UUID(), nullable=False), - sa.Column("status", sa.Text(), nullable=False, server_default=sa.text("'pending'")), - sa.Column("lease_owner", sa.Text(), nullable=True), - sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), - sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True), - sa.Column("fencing_token", sa.BigInteger(), nullable=False, server_default=sa.text("0")), - sa.Column("attempt_count", sa.Integer(), nullable=False, server_default=sa.text("0")), - sa.Column( - "next_retry_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") - ), - sa.Column("checkpoint", sa.Integer(), nullable=False, server_default=sa.text("0")), - sa.Column("error_message", sa.Text(), nullable=True), - sa.Column( - "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") - ), - sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), - sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="enrichment_jobs_scan_id_fkey"), - sa.PrimaryKeyConstraint("job_id", name="enrichment_jobs_pkey"), - sa.UniqueConstraint("scan_id", name="uq_enrichment_jobs_scan_id"), - sa.CheckConstraint("status IN ('pending', 'running', 'completed', 'failed')", name="ck_enrichment_jobs_status"), - ) + # autocommit_block() below commits the table before the concurrent index + # builds run, so a failure there leaves the table behind with + # alembic_version unchanged. Skip the create on a retry rather than + # failing on "relation already exists" before the index recovery. + if not sa.inspect(op.get_bind()).has_table("enrichment_jobs"): + op.create_table( + "enrichment_jobs", + sa.Column("job_id", postgresql.UUID(), nullable=False), + sa.Column("scan_id", postgresql.UUID(), nullable=False), + sa.Column("status", sa.Text(), nullable=False, server_default=sa.text("'pending'")), + sa.Column("lease_owner", sa.Text(), nullable=True), + sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("fencing_token", sa.BigInteger(), nullable=False, server_default=sa.text("0")), + sa.Column("attempt_count", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column( + "next_retry_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("checkpoint", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="enrichment_jobs_scan_id_fkey"), + sa.PrimaryKeyConstraint("job_id", name="enrichment_jobs_pkey"), + sa.UniqueConstraint("scan_id", name="uq_enrichment_jobs_scan_id"), + sa.CheckConstraint( + "status IN ('pending', 'running', 'completed', 'failed')", name="ck_enrichment_jobs_status" + ), + ) with op.get_context().autocommit_block(): + # Drop first so an INVALID index left by an interrupted build is + # rebuilt rather than kept (see the note in e4f7a9b2c6d8). + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_pending_retry") op.execute( """ CREATE INDEX CONCURRENTLY idx_enrichment_jobs_pending_retry @@ -52,6 +62,7 @@ def upgrade() -> None: WHERE status = 'pending' """ ) + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_running_lease") op.execute( """ CREATE INDEX CONCURRENTLY idx_enrichment_jobs_running_lease diff --git a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py index 598877dd..e412ca9d 100644 --- a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py +++ b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py @@ -19,22 +19,29 @@ def upgrade() -> None: """Store one liveness timestamp per worker process.""" - op.create_table( - "worker_heartbeats", - sa.Column("worker_id", sa.Text(), nullable=False), - sa.Column("worker_type", sa.Text(), nullable=False), - sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=False), - sa.PrimaryKeyConstraint("worker_id", "worker_type", name="worker_heartbeats_pkey"), - sa.CheckConstraint("worker_type IN ('scan', 'enrichment')", name="ck_worker_heartbeats_type"), - ) - op.create_index( - "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False - ) + # autocommit_block() below commits this table before the concurrent index + # build, so a failure there leaves it behind with alembic_version + # unchanged. Skip on a retry rather than failing before the recovery. + if not sa.inspect(op.get_bind()).has_table("worker_heartbeats"): + op.create_table( + "worker_heartbeats", + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("worker_type", sa.Text(), nullable=False), + sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("worker_id", "worker_type", name="worker_heartbeats_pkey"), + sa.CheckConstraint("worker_type IN ('scan', 'enrichment')", name="ck_worker_heartbeats_type"), + ) + op.create_index( + "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False + ) # /metrics reports the last successful scan on every scrape. Without this # the aggregate degrades into a sequential scan of the whole scans table as # scan history grows; the partial index keeps it an index-only lookup. with op.get_context().autocommit_block(): + # Drop first so an INVALID index left by an interrupted build is + # rebuilt rather than kept (see the note in e4f7a9b2c6d8). + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_completed_completed_at") op.execute( """ CREATE INDEX CONCURRENTLY idx_scans_completed_completed_at diff --git a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py index caad42dd..3555ac1b 100644 --- a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py +++ b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py @@ -8,7 +8,6 @@ from typing import Sequence, Union from alembic import op -import sqlalchemy as sa revision: str = "e4f7a9b2c6d8" @@ -19,13 +18,15 @@ def upgrade() -> None: """Add additive lease state and make legacy running work recoverable.""" - op.add_column("scans", sa.Column("lease_owner", sa.Text(), nullable=True)) - op.add_column("scans", sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True)) - op.add_column("scans", sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True)) - op.add_column( - "scans", - sa.Column("fencing_token", sa.BigInteger(), server_default=sa.text("0"), nullable=False), - ) + # autocommit_block() below commits everything issued before it, so a + # failure while building the concurrent indexes leaves these columns in + # place with alembic_version still on the previous revision. The retry has + # to be able to walk back over them instead of failing on "column already + # exists" before it reaches the index recovery. + op.execute("ALTER TABLE scans ADD COLUMN IF NOT EXISTS lease_owner TEXT") + op.execute("ALTER TABLE scans ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ") + op.execute("ALTER TABLE scans ADD COLUMN IF NOT EXISTS last_heartbeat_at TIMESTAMPTZ") + op.execute("ALTER TABLE scans ADD COLUMN IF NOT EXISTS fencing_token BIGINT NOT NULL DEFAULT 0") # A pre-lease running row belongs to an old worker that cannot satisfy the # new fencing contract. Marking its lease expired preserves the row and @@ -41,6 +42,11 @@ def upgrade() -> None: # These indexes are additive and are created concurrently so a populated # production scans table remains available while the migration runs. with op.get_context().autocommit_block(): + # An interrupted CONCURRENTLY build leaves an INVALID index that still + # owns the name and can never serve a query. Dropping first (rather + # than CREATE ... IF NOT EXISTS, which would keep the broken one) makes + # a retry rebuild it. + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_pending_started_at") op.execute( """ CREATE INDEX CONCURRENTLY idx_scans_pending_started_at @@ -48,6 +54,7 @@ def upgrade() -> None: WHERE status = 'pending' """ ) + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_running_lease_expires_at") op.execute( """ CREATE INDEX CONCURRENTLY idx_scans_running_lease_expires_at diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py index 7c067286..9fc160d6 100644 --- a/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py +++ b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py @@ -8,7 +8,6 @@ from typing import Sequence, Union from alembic import op -import sqlalchemy as sa revision: str = "f2b6d8e1a4c9" @@ -21,7 +20,10 @@ def upgrade() -> None: """Give every finding a stable identity so replayed results upsert.""" - op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) + # Rerunnable after a failed concurrent index build (see the note in + # e4f7a9b2c6d8): the backfill and the NOT NULL tightening below are both + # idempotent, so the whole revision can be replayed. + op.execute("ALTER TABLE findings ADD COLUMN IF NOT EXISTS finding_key TEXT") # Existing records predate the identity contract. Preserve each record as # distinct rather than attempting to infer equivalence from mutable text. op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") diff --git a/tests/test_scan_admission_migration_postgres.py b/tests/test_scan_admission_migration_postgres.py index bb10ad10..059a548a 100644 --- a/tests/test_scan_admission_migration_postgres.py +++ b/tests/test_scan_admission_migration_postgres.py @@ -22,6 +22,10 @@ # The revision immediately before scan admission is introduced. _BEFORE_ADMISSION = "f2b6d8e1a4c9" +# The dev head this branch builds on, i.e. the state a deployment upgrades from. +_BEFORE_LEASES = "3f59f83a5253" +_LEASES = "e4f7a9b2c6d8" +_HEAD = "d4a8c1e6b2f9" _ADMISSION = "a7c5e9d2f1b4" _ACTIVE_INDEX = "uq_scans_one_active_per_subscription" _REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -242,3 +246,131 @@ def test_an_invalid_index_left_by_a_failed_build_is_replaced_not_inherited(): _upgrade(dsn, _ADMISSION) assert _index_is_valid(dsn, _ACTIVE_INDEX) is True + + +# ── Partial-execution safety ──────────────────────────────────────────────── +# +# alembic/env.py does not use transaction_per_migration, so the whole upgrade +# shares one transaction -- but every autocommit_block() commits it. DDL +# issued before a concurrent index build is therefore already durable when +# that build fails, while alembic_version still names the previous revision. +# The retry must be able to walk back over its own committed work. + + +def _current_revision(dsn: str) -> str: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT version_num FROM alembic_version") + return cur.fetchone()[0] + + +def _invalidate_index(dsn: str, name: str) -> None: + """Reproduce the index an interrupted CONCURRENTLY build leaves behind.""" + conn = psycopg2.connect(dsn) + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute("UPDATE pg_index SET indisvalid = false WHERE indexrelid = %s::regclass", (name,)) + finally: + conn.close() + + +def test_upgrade_reruns_after_a_partial_lease_migration(): + """Columns committed by a failed run must not block the retry. + + e4f7a9b2c6d8 adds its columns, then builds two indexes concurrently. If + that build fails the columns are already committed, so a plain add_column + on retry died with "column already exists" before reaching any recovery. + """ + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_LEASES) + # Exactly what the failed run had committed: the columns, no version bump. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("ALTER TABLE scans ADD COLUMN lease_owner TEXT") + cur.execute("ALTER TABLE scans ADD COLUMN lease_expires_at TIMESTAMPTZ") + cur.execute("ALTER TABLE scans ADD COLUMN last_heartbeat_at TIMESTAMPTZ") + cur.execute("ALTER TABLE scans ADD COLUMN fencing_token BIGINT NOT NULL DEFAULT 0") + assert _current_revision(dsn) == _BEFORE_LEASES + + _upgrade(dsn, "head") + + assert _current_revision(dsn) == _HEAD + assert _index_is_valid(dsn, "idx_scans_pending_started_at") is True + + +def test_upgrade_rebuilds_an_index_left_invalid_by_an_interrupted_build(): + """An INVALID index owns its name but can never serve a query. + + CREATE INDEX ... IF NOT EXISTS would keep it, so each build drops the name + first. The retry has to end with a valid index, not the broken one. + """ + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_LEASES) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("ALTER TABLE scans ADD COLUMN lease_owner TEXT") + cur.execute("ALTER TABLE scans ADD COLUMN lease_expires_at TIMESTAMPTZ") + cur.execute("ALTER TABLE scans ADD COLUMN last_heartbeat_at TIMESTAMPTZ") + cur.execute("ALTER TABLE scans ADD COLUMN fencing_token BIGINT NOT NULL DEFAULT 0") + cur.execute( + "CREATE INDEX idx_scans_pending_started_at ON scans (started_at ASC) WHERE status = 'pending'" + ) + _invalidate_index(dsn, "idx_scans_pending_started_at") + assert _index_is_valid(dsn, "idx_scans_pending_started_at") is False + + _upgrade(dsn, "head") + + assert _current_revision(dsn) == _HEAD + assert _index_is_valid(dsn, "idx_scans_pending_started_at") is True + + +def test_upgrade_reruns_after_a_partial_enrichment_jobs_migration(): + """A committed table from a failed run must not block the retry either.""" + with _scratch_database() as dsn: + _upgrade(dsn, _ADMISSION) + # c9e1a5b7d3f2 creates enrichment_jobs, then builds its indexes + # concurrently; the table survives a failure in that block. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + CREATE TABLE enrichment_jobs ( + job_id UUID PRIMARY KEY, + scan_id UUID NOT NULL UNIQUE REFERENCES scans(scan_id), + status TEXT NOT NULL DEFAULT 'pending', + lease_owner TEXT, + lease_expires_at TIMESTAMPTZ, + last_heartbeat_at TIMESTAMPTZ, + fencing_token BIGINT NOT NULL DEFAULT 0, + attempt_count INTEGER NOT NULL DEFAULT 0, + next_retry_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + checkpoint INTEGER NOT NULL DEFAULT 0, + error_message TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + completed_at TIMESTAMPTZ, + CONSTRAINT ck_enrichment_jobs_status + CHECK (status IN ('pending', 'running', 'completed', 'failed')) + ) + """ + ) + assert _current_revision(dsn) == _ADMISSION + + _upgrade(dsn, "head") + + assert _current_revision(dsn) == _HEAD + assert _index_is_valid(dsn, "idx_enrichment_jobs_pending_retry") is True + + +def test_upgrade_from_the_dev_head_reaches_a_single_head(): + """The documented deployment path: current dev -> this branch.""" + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_LEASES) + assert _current_revision(dsn) == _BEFORE_LEASES + _upgrade(dsn, "head") + assert _current_revision(dsn) == _HEAD + # The #263 coverage table from dev survives this branch's migrations. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT to_regclass('rule_evaluations')") + assert cur.fetchone()[0] is not None From cad05628f6599789b5c86c31f08770b667786e04 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sun, 20 Sep 2026 22:31:53 +0100 Subject: [PATCH 19/19] fix(core): make enrichment recovery non-blocking Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 74 ++++++---- docs/async-scan-architecture.md | 2 +- tests/test_enrichment_jobs_postgres.py | 195 +++++++++++++++++++++++++ 3 files changed, 245 insertions(+), 26 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index 7c3af289..ee68bdd2 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -903,39 +903,63 @@ def fail_enrichment_job( raise def recover_stale_enrichment_jobs(self, max_attempts: int = 3) -> int: - """Return expired enrichment claims to pending or terminally fail them.""" + """Recover expired enrichment claims without blocking the worker loop. + + Selection, job transition, and parent-scan status update are one + statement. Both rows are locked with ``SKIP LOCKED`` so a job (or its + scan) currently handled by another transaction is deferred to a later + recovery pass instead of convoying all queue work behind it. + """ conn = self._get_conn() try: with conn.cursor() as cur: cur.execute( """ - UPDATE enrichment_jobs SET status = 'failed', completed_at = CURRENT_TIMESTAMP, - lease_owner = NULL, lease_expires_at = NULL, - error_message = 'Enrichment exceeded maximum retry attempts after worker interruption.' - WHERE status = 'running' AND attempt_count >= %s - AND lease_expires_at < CURRENT_TIMESTAMP - RETURNING scan_id - """, - (max_attempts,), - ) - failed_scans = [row[0] for row in cur.fetchall()] - cur.execute( - """ - UPDATE enrichment_jobs SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, - error_message = 'Enrichment worker interrupted; queued for retry.' - WHERE status = 'running' AND attempt_count < %s - AND lease_expires_at < CURRENT_TIMESTAMP - RETURNING scan_id + WITH stale AS ( + SELECT j.job_id, j.scan_id, j.attempt_count + FROM enrichment_jobs AS j + JOIN scans AS s ON s.scan_id = j.scan_id + WHERE j.status = 'running' + AND j.lease_expires_at < CURRENT_TIMESTAMP + ORDER BY j.lease_expires_at ASC, j.job_id ASC + FOR UPDATE OF j, s SKIP LOCKED + ), transitioned AS ( + UPDATE enrichment_jobs AS j + SET status = CASE + WHEN stale.attempt_count >= %(max_attempts)s THEN 'failed' + ELSE 'pending' + END, + completed_at = CASE + WHEN stale.attempt_count >= %(max_attempts)s THEN CURRENT_TIMESTAMP + ELSE j.completed_at + END, + lease_owner = NULL, + lease_expires_at = NULL, + error_message = CASE + WHEN stale.attempt_count >= %(max_attempts)s + THEN 'Enrichment exceeded maximum retry attempts after worker interruption.' + ELSE 'Enrichment worker interrupted; queued for retry.' + END + FROM stale + WHERE j.job_id = stale.job_id + RETURNING j.job_id, j.scan_id, j.status + ), updated_scans AS ( + UPDATE scans AS s + SET cve_enrichment_status = CASE + WHEN transitioned.status = 'failed' THEN 'FAILED' + ELSE 'PENDING' + END + FROM transitioned + WHERE s.scan_id = transitioned.scan_id + RETURNING transitioned.job_id, transitioned.scan_id, transitioned.status + ) + SELECT job_id, scan_id, status FROM updated_scans """, - (max_attempts,), + {"max_attempts": max_attempts}, ) - retried_scans = [row[0] for row in cur.fetchall()] - for scan_id in failed_scans: - cur.execute("UPDATE scans SET cve_enrichment_status = 'FAILED' WHERE scan_id = %s", (scan_id,)) - for scan_id in retried_scans: - cur.execute("UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", (scan_id,)) + transitions = cur.fetchall() conn.commit() - return len(failed_scans) + len(retried_scans) + return len(transitions) except Exception: self.rollback(conn) raise diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index e8c6eeb6..ed92cace 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -13,7 +13,7 @@ In the legacy synchronous model, POST /api/scans/trigger would block the HTTP re OpenShield now employs a decoupled, database backed worker architecture. This is the industry standard for long running security tasks where reliability and state persistence are critical. ### 1. The API (Flask) -When a scan is triggered, the API validates the subscription and creates a durable pending record. PostgreSQL permits at most one `pending` or `running` scan per subscription. `Idempotency-Key` replays return the same logical scan when the request fingerprint matches; reuse with different semantics returns a conflict. The optional `OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR` policy enables an explicit time-window quota. A zero/unset value preserves the current no-business-limit policy while the one-active-scan concurrency quota remains enforced. +When a scan is triggered, the API validates the subscription and creates a durable pending record. PostgreSQL permits at most one `pending` or `running` scan per subscription. `Idempotency-Key` values are scoped to a subscription: repeating the same key for that subscription returns the original scan, while the same key under a different subscription is independent. A trigger has no semantic request input beyond `subscription_id`, so there is no changed-payload conflict to report. The optional `OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR` policy enables an explicit time-window quota. A zero/unset value preserves the current no-business-limit policy while the one-active-scan concurrency quota remains enforced. ### 2. The Queue (PostgreSQL) The scans table acts as a persistent task queue. This avoids the need for additional infrastructure like Redis or RabbitMQ while providing ACID compliance, visibility, and auditability. Scan states are never lost during crashes, status polling is a simple SQL query, and every scan has a persistent record of its error state. diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index 808d7359..4ea681f7 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -2,12 +2,16 @@ import os import threading +import time import uuid +from contextlib import contextmanager from unittest.mock import patch import psycopg2 import psycopg2.extras import pytest +from alembic import command +from alembic.config import Config from api.models.finding import DatabaseManager, LostLease from scanner.enrichment_worker import process_enrichment_job @@ -17,6 +21,85 @@ not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" ) +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +@contextmanager +def _isolated_database(): + """Create a migrated database whose recovery count has no external rows.""" + base = os.environ["DATABASE_URL"].rsplit("/", 1)[0] + name = f"openshield_enrichment_{uuid.uuid4().hex[:12]}" + admin = psycopg2.connect(f"{base}/postgres") + admin.autocommit = True + try: + with admin.cursor() as cur: + cur.execute(f'CREATE DATABASE "{name}"') + dsn = f"{base}/{name}" + config = Config() + config.set_main_option("script_location", os.path.join(_REPO_ROOT, "alembic")) + previous = os.environ.get("DATABASE_URL") + os.environ["DATABASE_URL"] = dsn + try: + command.upgrade(config, "head") + finally: + if previous is None: + os.environ.pop("DATABASE_URL", None) + else: + os.environ["DATABASE_URL"] = previous + yield dsn + finally: + with admin.cursor() as cur: + cur.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()", + (name,), + ) + cur.execute(f'DROP DATABASE IF EXISTS "{name}"') + admin.close() + + +def _seed_stale_job(dsn, *, owner, attempts, checkpoint): + """Insert one expired running job and its completed parent scan.""" + scan_id, job_id = str(uuid.uuid4()), str(uuid.uuid4()) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO scans + (scan_id, subscription_id, started_at, completed_at, status, cve_enrichment_status) + VALUES (%s, %s, CURRENT_TIMESTAMP - INTERVAL '1 hour', CURRENT_TIMESTAMP, + 'completed', 'ENRICHING') + """, + (scan_id, str(uuid.uuid4())), + ) + cur.execute( + """ + INSERT INTO enrichment_jobs + (job_id, scan_id, status, lease_owner, lease_expires_at, last_heartbeat_at, + fencing_token, attempt_count, next_retry_at, checkpoint, error_message) + VALUES (%s, %s, 'running', %s, + CURRENT_TIMESTAMP - INTERVAL '5 minutes', + CURRENT_TIMESTAMP - INTERVAL '6 minutes', + 17, %s, CURRENT_TIMESTAMP - INTERVAL '10 minutes', %s, %s) + """, + (job_id, scan_id, owner, attempts, checkpoint, f"previous error from {owner}"), + ) + return scan_id, job_id + + +def _recovery_row(dsn, job_id): + with psycopg2.connect(dsn) as conn: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + SELECT j.*, s.cve_enrichment_status + FROM enrichment_jobs AS j + JOIN scans AS s ON s.scan_id = j.scan_id + WHERE j.job_id = %s + """, + (job_id,), + ) + return dict(cur.fetchone()) + @pytest.fixture def enrichment_scan(): @@ -194,6 +277,118 @@ def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enr db.close() +def test_stale_recovery_skips_locked_job_and_recovers_it_on_the_next_pass(): + """A locked stale job is deferred without convoying other recovery work.""" + with _isolated_database() as dsn: + _locked_scan, locked_job = _seed_stale_job(dsn, owner="locked-owner", attempts=1, checkpoint=11) + _retry_scan, retry_job = _seed_stale_job(dsn, owner="retry-owner", attempts=1, checkpoint=22) + _failed_scan, failed_job = _seed_stale_job(dsn, owner="failed-owner", attempts=3, checkpoint=33) + + locked_before = _recovery_row(dsn, locked_job) + retry_before = _recovery_row(dsn, retry_job) + failed_before = _recovery_row(dsn, failed_job) + + holder = psycopg2.connect(dsn) + recovery_finished = threading.Event() + result = {} + + def recover() -> None: + db = DatabaseManager(dsn) + started = time.perf_counter() + try: + result["count"] = db.recover_stale_enrichment_jobs(max_attempts=3) + except Exception as exc: # pragma: no cover - asserted in the caller + result["error"] = exc + finally: + result["elapsed"] = time.perf_counter() - started + db.close() + recovery_finished.set() + + thread = threading.Thread(target=recover) + lock_started = time.perf_counter() + try: + with holder.cursor() as cur: + cur.execute("SELECT job_id FROM enrichment_jobs WHERE job_id = %s FOR UPDATE", (locked_job,)) + + thread.start() + # The lock is deliberately still held here. The timeout is only a + # deadlock guard; the event proves recovery completed before this + # transaction released Job A. + assert recovery_finished.wait(timeout=2), "stale recovery blocked behind a locked enrichment job" + assert "error" not in result + assert result["count"] == 2 + + locked_during = _recovery_row(dsn, locked_job) + retry_after = _recovery_row(dsn, retry_job) + failed_after = _recovery_row(dsn, failed_job) + + preserved_fields = ( + "status", + "lease_owner", + "lease_expires_at", + "last_heartbeat_at", + "fencing_token", + "attempt_count", + "next_retry_at", + "checkpoint", + "error_message", + "completed_at", + "cve_enrichment_status", + ) + assert {field: locked_during[field] for field in preserved_fields} == { + field: locked_before[field] for field in preserved_fields + } + + assert retry_after["status"] == "pending" + assert retry_after["lease_owner"] is None + assert retry_after["lease_expires_at"] is None + assert retry_after["attempt_count"] == retry_before["attempt_count"] + assert retry_after["checkpoint"] == retry_before["checkpoint"] + assert retry_after["fencing_token"] == retry_before["fencing_token"] + assert retry_after["last_heartbeat_at"] == retry_before["last_heartbeat_at"] + assert retry_after["next_retry_at"] == retry_before["next_retry_at"] + assert retry_after["completed_at"] == retry_before["completed_at"] + assert retry_after["error_message"] == "Enrichment worker interrupted; queued for retry." + assert retry_after["cve_enrichment_status"] == "PENDING" + + assert failed_after["status"] == "failed" + assert failed_after["lease_owner"] is None + assert failed_after["lease_expires_at"] is None + assert failed_after["attempt_count"] == failed_before["attempt_count"] + assert failed_after["checkpoint"] == failed_before["checkpoint"] + assert failed_after["fencing_token"] == failed_before["fencing_token"] + assert failed_after["last_heartbeat_at"] == failed_before["last_heartbeat_at"] + assert failed_after["next_retry_at"] == failed_before["next_retry_at"] + assert failed_after["completed_at"] is not None + assert failed_after["error_message"] == ( + "Enrichment exceeded maximum retry attempts after worker interruption." + ) + assert failed_after["cve_enrichment_status"] == "FAILED" + finally: + result["lock_held"] = time.perf_counter() - lock_started + holder.rollback() + holder.close() + if thread.ident is not None: + thread.join(timeout=10) + + db = DatabaseManager(dsn) + try: + assert db.recover_stale_enrichment_jobs(max_attempts=3) == 1 + finally: + db.close() + + locked_after = _recovery_row(dsn, locked_job) + assert locked_after["status"] == "pending" + assert locked_after["lease_owner"] is None + assert locked_after["lease_expires_at"] is None + assert locked_after["attempt_count"] == locked_before["attempt_count"] + assert locked_after["checkpoint"] == locked_before["checkpoint"] + assert locked_after["fencing_token"] == locked_before["fencing_token"] + assert locked_after["cve_enrichment_status"] == "PENDING" + assert result["elapsed"] < 2 + assert result["lock_held"] >= result["elapsed"] + + def _fail_terminally(dsn, db, scan_id): """Drive a job through its whole retry budget until it is 'failed'.""" for attempt in range(1, 4):