Fix async scan state recovery (#168)

api/models/finding.py

@@ -110,15 +110,15 @@ def close(self) -> None:
     def init_db(self) -> None:
         """Alias for run_migrations. Called by startup.sh on every boot.
 
-        Calling this is always safe — run_migrations() handles both fresh
+        Calling this is always safe; run_migrations() handles both fresh
         databases and existing ones via IF NOT EXISTS guards throughout.
         """
         self.run_migrations()
 
     def create_tables(self) -> None:
         """Create the findings, scans, and rules tables if they do not exist.
 
-        Includes all columns — including CVE columns — so fresh databases
+        Includes all columns, including CVE columns, so fresh databases
         never need the ALTER TABLE path in run_migrations().
         """
         conn = self._get_conn()
@@ -134,6 +134,7 @@ def create_tables(self) -> None:
                     score           INTEGER DEFAULT NULL,
                     cve_enrichment_status TEXT DEFAULT 'PENDING',
                     status          TEXT DEFAULT 'pending',
+                    attempt_count   INTEGER DEFAULT 0,
                     error_message   TEXT
                 );
             """)
@@ -173,7 +174,7 @@ def create_tables(self) -> None:
     def run_migrations(self) -> None:
         """Ensure the schema is fully current. Safe to call on every startup.
 
-        Calls create_tables() first so the call order never matters — this
+        Calls create_tables() first so the call order never matters; this
         method is safe whether the database is brand new or has existing data.
 
         On a fresh database:
@@ -187,7 +188,7 @@ def run_migrations(self) -> None:
         Concurrent startup safety:
             Both CREATE TABLE IF NOT EXISTS and ALTER TABLE ADD COLUMN IF NOT
             EXISTS are atomic at the PostgreSQL catalog level. Two Render
-            instances racing at boot will not error — the second call silently
+            instances racing at boot will not error; the second call silently
             no-ops on whichever statement the first already completed.
         """
         self.create_tables()
@@ -205,6 +206,7 @@ def run_migrations(self) -> None:
                     ALTER TABLE scans
                         ADD COLUMN IF NOT EXISTS cve_enrichment_status TEXT DEFAULT 'COMPLETED',
                         ADD COLUMN IF NOT EXISTS status                TEXT DEFAULT 'completed',
+                        ADD COLUMN IF NOT EXISTS attempt_count         INTEGER DEFAULT 0,
                         ADD COLUMN IF NOT EXISTS error_message         TEXT,
                         ADD COLUMN IF NOT EXISTS claimed_at            TIMESTAMPTZ
                 """)
@@ -239,9 +241,10 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None:
                 """
                 INSERT INTO scans (
                     scan_id, subscription_id, started_at, completed_at,
-                    total_findings, score, cve_enrichment_status, status, error_message
+                    total_findings, score, cve_enrichment_status, status,
+                    attempt_count, error_message
                 )
-                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
+                VALUES (%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,
@@ -258,6 +261,7 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None:
                     scan_result.get("score"),
                     scan_result.get("cve_enrichment_status", "PENDING"),
                     scan_result.get("status", "completed"),
+                    scan_result.get("attempt_count", 0),
                     scan_result.get("error_message"),
                 ),
             )
@@ -395,8 +399,8 @@ def create_pending_scan(self, scan_id: str, subscription_id: str) -> None:
         with conn.cursor() as cur:
             cur.execute(
                 """
-                INSERT INTO scans (scan_id, subscription_id, started_at, status)
-                VALUES (%s, %s, %s, 'pending')
+                INSERT INTO scans (scan_id, subscription_id, started_at, status, attempt_count)
+                VALUES (%s, %s, %s, 'pending', 0)
                 """,
                 (scan_id, subscription_id, started_at),
             )
@@ -412,7 +416,7 @@ def update_scan_status(self, scan_id: str, status: str, error_message: Optional[
             if status == "completed":
                 completed_at = datetime.now(timezone.utc).isoformat()
                 cur.execute(
-                    "UPDATE scans SET status = %s, completed_at = %s WHERE scan_id = %s",
+                    "UPDATE scans SET status = %s, completed_at = %s, error_message = NULL WHERE scan_id = %s",
                     (status, completed_at, scan_id),
                 )
             else:
@@ -433,7 +437,10 @@ def claim_next_pending_scan(self) -> Optional[Dict[str, Any]]:
             cur.execute(
                 """
                 UPDATE scans
-                SET status = 'running', claimed_at = %s
+                SET status = 'running',
+                    claimed_at = %s,
+                    attempt_count = COALESCE(attempt_count, 0) + 1,
+                    error_message = NULL
                 WHERE scan_id = (
                     SELECT scan_id
                     FROM scans
@@ -452,25 +459,51 @@ def claim_next_pending_scan(self) -> Optional[Dict[str, Any]]:
                 return dict(row)
             return None
 
-    def recover_stale_scans(self, timeout_minutes: int = 60) -> int:
-        """Mark scans that have been 'running' for too long as 'failed'."""
+    def recover_stale_scans(self, timeout_minutes: int = 60, max_attempts: int = 3) -> int:
+        """Recover scans left running after a worker crash or restart.
+
+        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 timed out after remaining in running state for too long.'
+                    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
+
+            cur.execute(
+                """
+                UPDATE scans
+                SET status = 'pending',
+                    claimed_at = NULL,
+                    error_message = 'Scan worker interrupted before completion. Queued for retry.'
                 WHERE status = 'running'
-                  AND claimed_at < (CURRENT_TIMESTAMP - INTERVAL '%s minutes')
+                  AND COALESCE(attempt_count, 0) < %s
+                  AND claimed_at < (CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute'))
                 """,
-                (timeout_minutes,),
+                (max_attempts, timeout_minutes),
             )
-            count = cur.rowcount
+            retry_count = cur.rowcount
         conn.commit()
-        if count > 0:
-            logger.info("Recovered %d stale 'running' scans", count)
-        return count
+        total_count = failed_count + retry_count
+        if total_count > 0:
+            logger.info(
+                "Recovered %d stale 'running' scans (%d retried, %d failed)",
+                total_count,
+                retry_count,
+                failed_count,
+            )
+        return total_count
 
     def get_pending_scans(self) -> List[Dict[str, Any]]:
         """Return all scans in the 'pending' state."""

docs/async-scan-architecture.md

@@ -18,6 +18,11 @@ When a scan is triggered, the API performs minimal work. It validates the subscr
 ### 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.
 
+### 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/<scan_id>` 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.
+
 ### 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.
 

tests/test_async_scan_persistence.py

@@ -0,0 +1,160 @@
+"""Regression tests for DB-backed async scan state.
+
+These tests protect against reintroducing an in-memory scan job store. The API
+must persist queued scan state through DatabaseManager so status survives web
+process restarts.
+"""
+
+from unittest.mock import MagicMock, patch
+
+from api.models.finding import DatabaseManager
+
+
+class _Cursor:
+    def __init__(self, rows=None, rowcounts=None):
+        self.rows = rows or []
+        self.rowcounts = rowcounts or []
+        self.calls = []
+        self.rowcount = 0
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+    def execute(self, sql, params=None):
+        self.calls.append((sql, params))
+        if self.rowcounts:
+            self.rowcount = self.rowcounts.pop(0)
+
+    def fetchone(self):
+        if self.rows:
+            return self.rows.pop(0)
+        return None
+
+
+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."""
+    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()
+
+    with patch("api.routes.scans.DatabaseManager", return_value=mock_db) as db_class:
+        with patch("api.routes.scans.uuid.uuid4", return_value=scan_id):
+            resp = client.post(
+                "/api/scans/trigger",
+                json={"subscription_id": subscription_id},
+                headers=auth_headers,
+            )
+
+    assert resp.status_code == 202
+    assert resp.get_json() == {
+        "scan_id": scan_id,
+        "status": "pending",
+        "message": "Scan has been queued and will start shortly.",
+    }
+    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)
+
+
+def test_get_scan_status_reads_from_database(client, auth_headers, monkeypatch):
+    """GET /api/scans/<scan_id> should read durable status from PostgreSQL."""
+    monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db")
+    scan_id = "22222222-2222-2222-2222-222222222222"
+    mock_db = MagicMock()
+    mock_db.get_scan.return_value = {
+        "scan_id": scan_id,
+        "subscription_id": "00000000-0000-0000-0000-000000000000",
+        "status": "running",
+        "started_at": "2026-07-07T12:00:00Z",
+        "completed_at": None,
+        "total_findings": 0,
+        "score": None,
+        "error_message": None,
+    }
+
+    with patch("api.routes.scans.DatabaseManager", return_value=mock_db) as db_class:
+        resp = client.get(f"/api/scans/{scan_id}", headers=auth_headers)
+
+    assert resp.status_code == 200
+    assert resp.get_json()["status"] == "running"
+    db_class.assert_called_once_with("postgresql://ci:ci@localhost/ci_db")
+    mock_db.connect.assert_called_once()
+    mock_db.get_scan.assert_called_once_with(scan_id)
+
+
+def test_get_scan_status_returns_not_found_for_missing_database_row(client, auth_headers, monkeypatch):
+    """Missing persisted scan state should return 404 instead of consulting memory."""
+    monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db")
+    scan_id = "33333333-3333-3333-3333-333333333333"
+    mock_db = MagicMock()
+    mock_db.get_scan.return_value = None
+
+    with patch("api.routes.scans.DatabaseManager", return_value=mock_db):
+        resp = client.get(f"/api/scans/{scan_id}", headers=auth_headers)
+
+    assert resp.status_code == 404
+    assert resp.get_json() == {"error": "Scan not found"}
+    mock_db.get_scan.assert_called_once_with(scan_id)
+
+
+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}])
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    with patch.object(db, "_get_conn", return_value=conn):
+        scan = db.claim_next_pending_scan()
+
+    executed_sql = cursor.calls[0][0]
+    assert "attempt_count = COALESCE(attempt_count, 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_recover_stale_scans_retries_before_max_attempts():
+    """Stale running scans should return to pending while attempts remain."""
+    db = DatabaseManager.__new__(DatabaseManager)
+    cursor = _Cursor(rowcounts=[0, 1])
+    conn = MagicMock()
+    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)
+
+    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 "status = 'pending'" in retry_sql
+    assert "claimed_at = NULL" in retry_sql
+    assert retry_params == (3, 15)
+    assert recovered == 1
+    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."""
+    db = DatabaseManager.__new__(DatabaseManager)
+    cursor = _Cursor(rowcounts=[1, 0])
+    conn = MagicMock()
+    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)
+
+    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 "COALESCE(attempt_count, 0) < %s" in retry_sql
+    assert retry_params == (3, 60)
+    assert recovered == 1
+    conn.commit.assert_called_once()