fix(attack-paths): recover graph_data_ready when scan fails during graph swap (#10354)

Author: josema-xyz

api/CHANGELOG.md

@@ -13,6 +13,10 @@ All notable changes to the **Prowler API** are documented in this file.
 - Attack Paths: Complete migration to private graph labels and properties, removing deprecated dual-write support [(#10268)](https://github.com/prowler-cloud/prowler/pull/10268)
 - Attack Paths: Added tenant and provider related labels to the nodes so they can be easily filtered on custom queries [(#10308)](https://github.com/prowler-cloud/prowler/pull/10308)
 
+### 🐞 Fixed
+
+- Attack Paths: Recover `graph_data_ready` flag when scan fails during graph swap, preventing query endpoints from staying blocked until the next successful scan [(#10354)](https://github.com/prowler-cloud/prowler/pull/10354)
+
 ### 🔐 Security
 
 - Use `psycopg2.sql` to safely compose DDL in `PostgresEnumMigration`, preventing SQL injection via f-string interpolation [(#10166)](https://github.com/prowler-cloud/prowler/pull/10166)

api/src/backend/api/attack_paths/database.py

@@ -1,26 +1,22 @@
 import atexit
 import logging
 import threading
-
-from typing import Any
-
 from contextlib import contextmanager
-from typing import Iterator
+from typing import Any, Iterator
 from uuid import UUID
 
 import neo4j
 import neo4j.exceptions
-
-from django.conf import settings
-
-from api.attack_paths.retryable_session import RetryableSession
 from config.env import env
+from django.conf import settings
 from tasks.jobs.attack_paths.config import (
     BATCH_SIZE,
     PROVIDER_ID_PROPERTY,
     PROVIDER_RESOURCE_LABEL,
 )
 
+from api.attack_paths.retryable_session import RetryableSession
+
 # Without this Celery goes crazy with Neo4j logging
 logging.getLogger("neo4j").setLevel(logging.ERROR)
 logging.getLogger("neo4j").propagate = False
@@ -197,6 +193,29 @@ def drop_subgraph(database: str, provider_id: str) -> int:
     return deleted_nodes
 
 
+def has_provider_data(database: str, provider_id: str) -> bool:
+    """
+    Check if any ProviderResource node exists for this provider.
+
+    Returns `False` if the database doesn't exist.
+    """
+    query = (
+        f"MATCH (n:{PROVIDER_RESOURCE_LABEL} "
+        f"{{{PROVIDER_ID_PROPERTY}: $provider_id}}) "
+        "RETURN 1 LIMIT 1"
+    )
+
+    try:
+        with get_session(database, default_access_mode=neo4j.READ_ACCESS) as session:
+            result = session.run(query, {"provider_id": provider_id})
+            return result.single() is not None
+
+    except GraphDatabaseQueryException as exc:
+        if exc.code == "Neo.ClientError.Database.DatabaseNotFound":
+            return False
+        raise
+
+
 def clear_cache(database: str) -> None:
     query = "CALL db.clearQueryCaches()"
 

api/src/backend/api/tests/test_attack_paths_database.py

@@ -442,3 +442,78 @@ def call_init():
         # All threads got the same driver instance
         assert all(r is mock_driver for r in results)
         assert len(results) == 10
+
+
+class TestHasProviderData:
+    """Test has_provider_data helper for checking provider nodes in Neo4j."""
+
+    def test_returns_true_when_nodes_exist(self):
+        import api.attack_paths.database as db_module
+
+        mock_session = MagicMock()
+        mock_result = MagicMock()
+        mock_result.single.return_value = MagicMock()  # non-None record
+        mock_session.run.return_value = mock_result
+
+        session_ctx = MagicMock()
+        session_ctx.__enter__.return_value = mock_session
+        session_ctx.__exit__.return_value = False
+
+        with patch(
+            "api.attack_paths.database.get_session",
+            return_value=session_ctx,
+        ):
+            assert db_module.has_provider_data("db-tenant-abc", "provider-123") is True
+
+        mock_session.run.assert_called_once()
+
+    def test_returns_false_when_no_nodes(self):
+        import api.attack_paths.database as db_module
+
+        mock_session = MagicMock()
+        mock_result = MagicMock()
+        mock_result.single.return_value = None
+        mock_session.run.return_value = mock_result
+
+        session_ctx = MagicMock()
+        session_ctx.__enter__.return_value = mock_session
+        session_ctx.__exit__.return_value = False
+
+        with patch(
+            "api.attack_paths.database.get_session",
+            return_value=session_ctx,
+        ):
+            assert db_module.has_provider_data("db-tenant-abc", "provider-123") is False
+
+    def test_returns_false_when_database_not_found(self):
+        import api.attack_paths.database as db_module
+
+        session_ctx = MagicMock()
+        session_ctx.__enter__.side_effect = db_module.GraphDatabaseQueryException(
+            message="Database does not exist",
+            code="Neo.ClientError.Database.DatabaseNotFound",
+        )
+
+        with patch(
+            "api.attack_paths.database.get_session",
+            return_value=session_ctx,
+        ):
+            assert (
+                db_module.has_provider_data("db-tenant-gone", "provider-123") is False
+            )
+
+    def test_raises_on_other_errors(self):
+        import api.attack_paths.database as db_module
+
+        session_ctx = MagicMock()
+        session_ctx.__enter__.side_effect = db_module.GraphDatabaseQueryException(
+            message="Connection refused",
+            code="Neo.TransientError.General.UnknownError",
+        )
+
+        with patch(
+            "api.attack_paths.database.get_session",
+            return_value=session_ctx,
+        ):
+            with pytest.raises(db_module.GraphDatabaseQueryException):
+                db_module.has_provider_data("db-tenant-abc", "provider-123")

api/src/backend/tasks/jobs/attack_paths/db_utils.py

@@ -3,15 +3,13 @@
 
 from cartography.config import Config as CartographyConfig
 from celery.utils.log import get_task_logger
+from tasks.jobs.attack_paths.config import is_provider_available
 
 from api.attack_paths import database as graph_database
 from api.db_utils import rls_transaction
-from api.models import (
-    AttackPathsScan as ProwlerAPIAttackPathsScan,
-    Provider as ProwlerAPIProvider,
-    StateChoices,
-)
-from tasks.jobs.attack_paths.config import is_provider_available
+from api.models import AttackPathsScan as ProwlerAPIAttackPathsScan
+from api.models import Provider as ProwlerAPIProvider
+from api.models import StateChoices
 
 logger = get_task_logger(__name__)
 
@@ -155,6 +153,37 @@ def set_provider_graph_data_ready(
         attack_paths_scan.refresh_from_db(fields=["graph_data_ready"])
 
 
+def recover_graph_data_ready(
+    attack_paths_scan: ProwlerAPIAttackPathsScan,
+) -> None:
+    """
+    Best-effort recovery of `graph_data_ready` after a scan failure.
+
+    Queries Neo4j to check if the provider still has data in the tenant
+    database. If data exists, restores `graph_data_ready=True` for all scans
+    of this provider. Never raises.
+
+    Trade-off: if the worker crashed mid-sync, partial data may exist and
+    this will re-enable queries against it. We accept that because leaving
+    `graph_data_ready=False` permanently (blocking all queries until the
+    next successful scan) is a worse outcome for the user.
+    """
+    try:
+        tenant_db = graph_database.get_database_name(attack_paths_scan.tenant_id)
+        if graph_database.has_provider_data(
+            tenant_db, str(attack_paths_scan.provider_id)
+        ):
+            set_provider_graph_data_ready(attack_paths_scan, True)
+            logger.info(
+                f"Recovered `graph_data_ready` for provider {attack_paths_scan.provider_id}"
+            )
+
+    except Exception:
+        logger.exception(
+            f"Failed to recover `graph_data_ready` for provider {attack_paths_scan.provider_id}"
+        )
+
+
 def fail_attack_paths_scan(
     tenant_id: str,
     scan_id: str,
@@ -185,3 +214,5 @@ def fail_attack_paths_scan(
             StateChoices.FAILED,
             {"global_error": error},
         )
+
+        recover_graph_data_ready(attack_paths_scan)

api/src/backend/tasks/jobs/attack_paths/scan.py

@@ -55,24 +55,21 @@
 
 import logging
 import time
-
 from typing import Any
 
 from cartography.config import Config as CartographyConfig
 from cartography.intel import analysis as cartography_analysis
 from cartography.intel import create_indexes as cartography_create_indexes
 from cartography.intel import ontology as cartography_ontology
 from celery.utils.log import get_task_logger
+from tasks.jobs.attack_paths import db_utils, findings, internet, sync, utils
+from tasks.jobs.attack_paths.config import get_cartography_ingestion_function
 
 from api.attack_paths import database as graph_database
 from api.db_utils import rls_transaction
-from api.models import (
-    Provider as ProwlerAPIProvider,
-    StateChoices,
-)
+from api.models import Provider as ProwlerAPIProvider
+from api.models import StateChoices
 from api.utils import initialize_prowler_provider
-from tasks.jobs.attack_paths import db_utils, findings, internet, sync, utils
-from tasks.jobs.attack_paths.config import get_cartography_ingestion_function
 
 # Without this Celery goes crazy with Cartography logging
 logging.getLogger("cartography").setLevel(logging.ERROR)
@@ -147,6 +144,10 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
         attack_paths_scan, task_id, tenant_cartography_config
     )
 
+    subgraph_dropped = False
+    sync_completed = False
+    provider_gated = False
+
     try:
         logger.info(
             f"Creating Neo4j database {tmp_cartography_config.neo4j_database} for tenant {prowler_api_provider.tenant_id}"
@@ -225,10 +226,12 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
 
         logger.info(f"Deleting existing provider graph in {tenant_database_name}")
         db_utils.set_provider_graph_data_ready(attack_paths_scan, False)
+        provider_gated = True
         graph_database.drop_subgraph(
             database=tenant_database_name,
             provider_id=str(prowler_api_provider.id),
         )
+        subgraph_dropped = True
         db_utils.update_attack_paths_scan_progress(attack_paths_scan, 98)
 
         logger.info(
@@ -240,6 +243,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
             tenant_id=str(prowler_api_provider.tenant_id),
             provider_id=str(prowler_api_provider.id),
         )
+        sync_completed = True
         db_utils.set_graph_data_ready(attack_paths_scan, True)
         db_utils.update_attack_paths_scan_progress(attack_paths_scan, 99)
 
@@ -264,23 +268,39 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
         logger.exception(exception_message)
         ingestion_exceptions["global_error"] = exception_message
 
-        # Handling databases changes
+        # Recover graph_data_ready based on how far the swap got.
+        # Partial drop (mid-batch failure) may leave `subgraph_dropped=False`
+        # with data partially deleted, so we prefer that over permanently blocked queries.
+        try:
+            if sync_completed:
+                db_utils.set_graph_data_ready(attack_paths_scan, True)
+            elif provider_gated and not subgraph_dropped:
+                db_utils.set_provider_graph_data_ready(attack_paths_scan, True)
+
+        except Exception:
+            logger.error(
+                f"Failed to recover `graph_data_ready` for provider {attack_paths_scan.provider_id}",
+                exc_info=True,
+            )
+
+        # Dropping the temporary database if it still exists
         try:
             graph_database.drop_database(tmp_cartography_config.neo4j_database)
 
         except Exception as e:
             logger.error(
-                f"Failed to drop temporary Neo4j database {tmp_cartography_config.neo4j_database} during cleanup: {e}",
+                f"Failed to drop temporary Neo4j database `{tmp_cartography_config.neo4j_database}` during cleanup: {e}",
                 exc_info=True,
             )
 
+        # Set Attack Paths scan state to FAILED
         try:
             db_utils.finish_attack_paths_scan(
                 attack_paths_scan, StateChoices.FAILED, ingestion_exceptions
             )
         except Exception as e:
             logger.error(
-                f"Could not mark attack paths scan {attack_paths_scan.id} as FAILED (row may have been deleted): {e}",
+                f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` (row may have been deleted): {e}",
                 exc_info=True,
             )