diff --git a/.env b/.env index 8294fac91c..49a03a7c69 100644 --- a/.env +++ b/.env @@ -72,8 +72,8 @@ NEO4J_APOC_IMPORT_FILE_ENABLED=false NEO4J_APOC_IMPORT_FILE_USE_NEO4J_CONFIG=true NEO4J_APOC_TRIGGER_ENABLED=false NEO4J_DBMS_CONNECTOR_BOLT_LISTEN_ADDRESS=0.0.0.0:7687 -# Neo4j Prowler settings -ATTACK_PATHS_BATCH_SIZE=1000 +# Attack Paths graph settings +ATTACK_PATHS_GRAPH_MUTATION_BATCH_SIZE=1000 ATTACK_PATHS_SERVICE_UNAVAILABLE_MAX_RETRIES=3 ATTACK_PATHS_READ_QUERY_TIMEOUT_SECONDS=30 ATTACK_PATHS_MAX_CUSTOM_QUERY_NODES=250 diff --git a/api/changelog.d/attack-paths-scan-reliability.fixed.md b/api/changelog.d/attack-paths-scan-reliability.fixed.md new file mode 100644 index 0000000000..1532943dc8 --- /dev/null +++ b/api/changelog.d/attack-paths-scan-reliability.fixed.md @@ -0,0 +1 @@ +Attack Paths scans handle provider deletion races cleanly, detect stale tasks after 16 hours, use backend-specific graph synchronization batches, and report exhausted Neptune write retries with the original database error diff --git a/api/src/backend/api/attack_paths/database.py b/api/src/backend/api/attack_paths/database.py index 3a33b964b7..3ef55b7eca 100644 --- a/api/src/backend/api/attack_paths/database.py +++ b/api/src/backend/api/attack_paths/database.py @@ -27,6 +27,7 @@ from django.conf import ( MAX_CUSTOM_QUERY_NODES = env.int("ATTACK_PATHS_MAX_CUSTOM_QUERY_NODES", default=250) TEMP_DB_PREFIX = "db-tmp-scan-" +DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound" # Exceptions @@ -44,6 +45,10 @@ class GraphDatabaseQueryException(Exception): return self.message +class NeptuneWriteRetryExhaustedException(GraphDatabaseQueryException): + pass + + class WriteQueryNotAllowedException(GraphDatabaseQueryException): pass diff --git a/api/src/backend/api/attack_paths/retryable_session.py b/api/src/backend/api/attack_paths/retryable_session.py index 2cd32f54ee..e7e78e9798 100644 --- a/api/src/backend/api/attack_paths/retryable_session.py +++ b/api/src/backend/api/attack_paths/retryable_session.py @@ -10,6 +10,28 @@ import neo4j.exceptions logger = logging.getLogger(__name__) +class RetryExhaustedError(Exception): + def __init__( + self, + *, + retry_context: str, + method_name: str, + attempts: int, + elapsed_seconds: float, + last_error: Exception, + ) -> None: + self.retry_context = retry_context + self.method_name = method_name + self.attempts = attempts + self.elapsed_seconds = elapsed_seconds + self.last_error = last_error + last_message = getattr(last_error, "message", None) or str(last_error) + super().__init__( + f"{retry_context} {method_name} failed after {attempts} attempts over " + f"{elapsed_seconds:.3f}s. Last error: {last_message}" + ) + + class RetryableSession: """Wrapper around ``neo4j.Session`` with a refreshable retry policy.""" @@ -19,11 +41,13 @@ class RetryableSession: max_retries: int, retry_if: Callable[[Exception], bool] | None = None, initial_retry_delay_seconds: float = 0, + retry_context: str | None = None, ) -> None: self._session_factory = session_factory self._max_retries = max(0, max_retries) self._retry_if = retry_if self._initial_retry_delay_seconds = max(0.0, initial_retry_delay_seconds) + self._retry_context = retry_context self._session = self._session_factory() def close(self) -> None: @@ -54,6 +78,7 @@ class RetryableSession: def _call_with_retry(self, method_name: str, *args: Any, **kwargs: Any) -> Any: attempt = 0 last_exc: Exception | None = None + started_at = time.monotonic() while attempt <= self._max_retries: try: @@ -68,17 +93,38 @@ class RetryableSession: attempt += 1 if attempt > self._max_retries: + if self._retry_context is not None: + raise RetryExhaustedError( + retry_context=self._retry_context, + method_name=method_name, + attempts=attempt, + elapsed_seconds=time.monotonic() - started_at, + last_error=exc, + ) from exc raise delay = self._retry_delay(attempt) - logger.warning( - "Graph session %s failed with %s; retry %s/%s in %.3fs", - method_name, - type(exc).__name__, - attempt, - self._max_retries, - delay, - ) + if self._retry_context is not None: + error_message = getattr(exc, "message", None) or str(exc) + logger.warning( + "%s %s failed with %s: %s; retry %s/%s in %.3fs", + self._retry_context, + method_name, + type(exc).__name__, + error_message, + attempt, + self._max_retries, + delay, + ) + else: + logger.warning( + "Graph session %s failed with %s; retry %s/%s in %.3fs", + method_name, + type(exc).__name__, + attempt, + self._max_retries, + delay, + ) self._refresh_session() if delay: time.sleep(delay) diff --git a/api/src/backend/api/attack_paths/sink/base.py b/api/src/backend/api/attack_paths/sink/base.py index 0ba4737f5e..0134e0a41f 100644 --- a/api/src/backend/api/attack_paths/sink/base.py +++ b/api/src/backend/api/attack_paths/sink/base.py @@ -15,6 +15,8 @@ class SinkDatabase(Protocol): has a single graph, and isolation is label-based). """ + sync_batch_size: int + def init(self) -> None: ... def close(self) -> None: ... diff --git a/api/src/backend/api/attack_paths/sink/neo4j.py b/api/src/backend/api/attack_paths/sink/neo4j.py index cb7d4889b0..c820ed9d93 100644 --- a/api/src/backend/api/attack_paths/sink/neo4j.py +++ b/api/src/backend/api/attack_paths/sink/neo4j.py @@ -54,6 +54,8 @@ DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound" class Neo4jSink(SinkDatabase): """Neo4j-backed sink. Multi-database cluster; tenant isolation is physical.""" + sync_batch_size = env.int("ATTACK_PATHS_NEO4J_SYNC_BATCH_SIZE", default=1000) + def __init__(self) -> None: self._driver: neo4j.Driver | None = None self._lock = threading.Lock() @@ -203,7 +205,7 @@ class Neo4jSink(SinkDatabase): """ from api.attack_paths.database import GraphDatabaseQueryException from tasks.jobs.attack_paths.config import ( - BATCH_SIZE, + GRAPH_MUTATION_BATCH_SIZE, PROVIDER_RESOURCE_LABEL, get_provider_label, ) @@ -251,7 +253,7 @@ class Neo4jSink(SinkDatabase): total_key="rels", deleted_key="deleted_rels", initial_total=deleted_relationships, - batch_size=BATCH_SIZE, + batch_size=GRAPH_MUTATION_BATCH_SIZE, drop_t0=drop_t0, ) relationship_batches += phase_batches @@ -270,7 +272,7 @@ class Neo4jSink(SinkDatabase): total_key="nodes", deleted_key="deleted_nodes", initial_total=0, - batch_size=BATCH_SIZE, + batch_size=GRAPH_MUTATION_BATCH_SIZE, drop_t0=drop_t0, ) diff --git a/api/src/backend/api/attack_paths/sink/neptune.py b/api/src/backend/api/attack_paths/sink/neptune.py index c68b6f8277..022e3c8669 100644 --- a/api/src/backend/api/attack_paths/sink/neptune.py +++ b/api/src/backend/api/attack_paths/sink/neptune.py @@ -25,7 +25,7 @@ from urllib.parse import urlsplit import neo4j import neo4j.exceptions -from api.attack_paths.retryable_session import RetryableSession +from api.attack_paths.retryable_session import RetryableSession, RetryExhaustedError from api.attack_paths.sink.base import SinkDatabase from api.attack_paths.sink.drop import ( NODE_DELETE_QUERY_TEMPLATE, @@ -85,6 +85,8 @@ def _is_retryable_write_error(exc: Exception) -> bool: class NeptuneSink(SinkDatabase): """Neptune-backed sink. Single database; isolation is label-based.""" + sync_batch_size = env.int("ATTACK_PATHS_NEPTUNE_SYNC_BATCH_SIZE", default=500) + def __init__(self) -> None: self._writer: neo4j.Driver | None = None self._reader: neo4j.Driver | None = None @@ -206,6 +208,7 @@ class NeptuneSink(SinkDatabase): from api.attack_paths.database import ( ClientStatementException, GraphDatabaseQueryException, + NeptuneWriteRetryExhaustedException, WriteQueryNotAllowedException, ) @@ -227,9 +230,17 @@ class NeptuneSink(SinkDatabase): initial_retry_delay_seconds=( NEPTUNE_WRITE_RETRY_DELAY_SECONDS if is_write_session else 0 ), + retry_context="Neptune write" if is_write_session else None, ) yield session_wrapper + except RetryExhaustedError as exc: + last_error = exc.last_error + raise NeptuneWriteRetryExhaustedException( + message=str(exc), + code=getattr(last_error, "code", None), + ) from last_error + except neo4j.exceptions.Neo4jError as exc: if ( default_access_mode == neo4j.READ_ACCESS @@ -291,7 +302,7 @@ class NeptuneSink(SinkDatabase): graph's branching factor. """ from tasks.jobs.attack_paths.config import ( - BATCH_SIZE, + GRAPH_MUTATION_BATCH_SIZE, PROVIDER_RESOURCE_LABEL, get_provider_label, ) @@ -330,7 +341,7 @@ class NeptuneSink(SinkDatabase): total_key="rels", deleted_key="deleted_rels", initial_total=deleted_relationships, - batch_size=BATCH_SIZE, + batch_size=GRAPH_MUTATION_BATCH_SIZE, drop_t0=drop_t0, ) relationship_batches += phase_batches @@ -349,7 +360,7 @@ class NeptuneSink(SinkDatabase): total_key="nodes", deleted_key="deleted_nodes", initial_total=0, - batch_size=BATCH_SIZE, + batch_size=GRAPH_MUTATION_BATCH_SIZE, drop_t0=drop_t0, ) diff --git a/api/src/backend/api/decorators.py b/api/src/backend/api/decorators.py index a055b2252f..2dd2d5fea4 100644 --- a/api/src/backend/api/decorators.py +++ b/api/src/backend/api/decorators.py @@ -1,12 +1,13 @@ import uuid from functools import wraps +from api.attack_paths.database import GraphDatabaseQueryException from api.db_router import READ_REPLICA_ALIAS from api.db_utils import POSTGRES_TENANT_VAR, SET_CONFIG_QUERY, rls_transaction from api.exceptions import ProviderDeletedException -from api.models import Provider, Scan +from api.models import Membership, Provider, Scan, Tenant from django.core.exceptions import ObjectDoesNotExist -from django.db import DatabaseError, connection, transaction +from django.db import DEFAULT_DB_ALIAS, DatabaseError, connection, transaction from rest_framework_json_api.serializers import ValidationError @@ -75,9 +76,11 @@ def handle_provider_deletion(func): """ Decorator that raises `ProviderDeletedException` if provider was deleted during execution. - Catches `ObjectDoesNotExist` and `DatabaseError` (including `IntegrityError`), checks if - provider still exists, and raises `ProviderDeletedException` if not. Otherwise, - re-raises original exception. + Catches `ObjectDoesNotExist`, `DatabaseError` (including `IntegrityError`), and + `GraphDatabaseQueryException`, checks if provider still exists, and raises + `ProviderDeletedException` if not. Graph database errors also check whether the + tenant still exists and has memberships. Otherwise, re-raises the original + exception. Requires `tenant_id` and `provider_id` in kwargs. @@ -92,11 +95,16 @@ def handle_provider_deletion(func): def wrapper(*args, **kwargs): try: return func(*args, **kwargs) - except (ObjectDoesNotExist, DatabaseError): + except (ObjectDoesNotExist, DatabaseError, GraphDatabaseQueryException) as exc: tenant_id = kwargs.get("tenant_id") provider_id = kwargs.get("provider_id") + database_alias = ( + DEFAULT_DB_ALIAS + if isinstance(exc, GraphDatabaseQueryException) + else READ_REPLICA_ALIAS + ) - with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS): + with rls_transaction(tenant_id, using=database_alias): if provider_id is None: scan_id = kwargs.get("scan_id") if scan_id is None: @@ -113,6 +121,13 @@ def handle_provider_deletion(func): raise ProviderDeletedException( f"Provider '{provider_id}' was deleted during the scan" ) from None + if isinstance(exc, GraphDatabaseQueryException) and ( + not Tenant.objects.filter(pk=tenant_id).exists() + or not Membership.objects.filter(tenant_id=tenant_id).exists() + ): + raise ProviderDeletedException( + f"Tenant '{tenant_id}' was deleted during the scan" + ) from None raise return wrapper diff --git a/api/src/backend/api/tests/test_decorators.py b/api/src/backend/api/tests/test_decorators.py index 5c2897730f..0bf58340a1 100644 --- a/api/src/backend/api/tests/test_decorators.py +++ b/api/src/backend/api/tests/test_decorators.py @@ -2,11 +2,12 @@ import uuid from unittest.mock import call, patch import pytest +from api.attack_paths.database import GraphDatabaseQueryException from api.db_utils import POSTGRES_TENANT_VAR, SET_CONFIG_QUERY from api.decorators import handle_provider_deletion, set_tenant from api.exceptions import ProviderDeletedException from django.core.exceptions import ObjectDoesNotExist -from django.db import DatabaseError, IntegrityError +from django.db import DEFAULT_DB_ALIAS, DatabaseError, IntegrityError @pytest.mark.django_db @@ -204,6 +205,106 @@ class TestHandleProviderDeletionDecorator: with pytest.raises(DatabaseError): task_func(tenant_id=str(tenant.id), provider_id=str(provider.id)) + @patch("api.decorators.rls_transaction") + @patch("api.decorators.Provider.objects.filter") + def test_graph_database_error_provider_missing_or_soft_deleted( + self, mock_provider_filter, mock_rls, tenants_fixture + ): + tenant = tenants_fixture[0] + provider_id = str(uuid.uuid4()) + + mock_rls.return_value.__enter__ = lambda s: None + mock_rls.return_value.__exit__ = lambda s, *args: None + mock_provider_filter.return_value.exists.return_value = False + + @handle_provider_deletion + def task_func(**kwargs): + raise GraphDatabaseQueryException("Temporary database not found") + + with pytest.raises(ProviderDeletedException): + task_func(tenant_id=str(tenant.id), provider_id=provider_id) + + @patch("api.decorators.rls_transaction") + @patch("api.decorators.Tenant.objects.filter") + @patch("api.decorators.Provider.objects.filter") + def test_graph_database_error_tenant_missing( + self, mock_provider_filter, mock_tenant_filter, mock_rls, tenants_fixture + ): + tenant = tenants_fixture[0] + provider_id = str(uuid.uuid4()) + + mock_rls.return_value.__enter__ = lambda s: None + mock_rls.return_value.__exit__ = lambda s, *args: None + mock_provider_filter.return_value.exists.return_value = True + mock_tenant_filter.return_value.exists.return_value = False + + @handle_provider_deletion + def task_func(**kwargs): + raise GraphDatabaseQueryException("Temporary database not found") + + with pytest.raises(ProviderDeletedException): + task_func(tenant_id=str(tenant.id), provider_id=provider_id) + + @patch("api.decorators.rls_transaction") + @patch("api.decorators.Membership.objects.filter") + @patch("api.decorators.Tenant.objects.filter") + @patch("api.decorators.Provider.objects.filter") + def test_graph_database_error_tenant_without_memberships( + self, + mock_provider_filter, + mock_tenant_filter, + mock_membership_filter, + mock_rls, + tenants_fixture, + ): + tenant = tenants_fixture[0] + provider_id = str(uuid.uuid4()) + + mock_rls.return_value.__enter__ = lambda s: None + mock_rls.return_value.__exit__ = lambda s, *args: None + mock_provider_filter.return_value.exists.return_value = True + mock_tenant_filter.return_value.exists.return_value = True + mock_membership_filter.return_value.exists.return_value = False + + @handle_provider_deletion + def task_func(**kwargs): + raise GraphDatabaseQueryException("Temporary database not found") + + with pytest.raises(ProviderDeletedException): + task_func(tenant_id=str(tenant.id), provider_id=provider_id) + + @patch("api.decorators.rls_transaction") + @patch("api.decorators.Membership.objects.filter") + @patch("api.decorators.Tenant.objects.filter") + @patch("api.decorators.Provider.objects.filter") + def test_graph_database_error_active_provider_and_tenant_reraises( + self, + mock_provider_filter, + mock_tenant_filter, + mock_membership_filter, + mock_rls, + tenants_fixture, + ): + tenant = tenants_fixture[0] + provider_id = str(uuid.uuid4()) + graph_error = GraphDatabaseQueryException("Temporary database not found") + + mock_rls.return_value.__enter__ = lambda s: None + mock_rls.return_value.__exit__ = lambda s, *args: None + mock_provider_filter.return_value.exists.return_value = True + mock_tenant_filter.return_value.exists.return_value = True + mock_membership_filter.return_value.exists.return_value = True + + @handle_provider_deletion + def task_func(**kwargs): + raise graph_error + + with pytest.raises(GraphDatabaseQueryException) as exc_info: + task_func(tenant_id=str(tenant.id), provider_id=provider_id) + + assert exc_info.value is graph_error + mock_rls.assert_called_once_with(str(tenant.id), using=DEFAULT_DB_ALIAS) + def test_missing_provider_and_scan_raises_assertion(self, tenants_fixture): """Raises AssertionError when neither provider_id nor scan_id in kwargs.""" diff --git a/api/src/backend/api/tests/test_retryable_session.py b/api/src/backend/api/tests/test_retryable_session.py index 8bb457b8b6..09b184c223 100644 --- a/api/src/backend/api/tests/test_retryable_session.py +++ b/api/src/backend/api/tests/test_retryable_session.py @@ -1,7 +1,7 @@ from unittest.mock import MagicMock, patch import pytest -from api.attack_paths.retryable_session import RetryableSession +from api.attack_paths.retryable_session import RetryableSession, RetryExhaustedError from neo4j.exceptions import ServiceUnavailable @@ -24,6 +24,7 @@ class TestRetryableSession: max_retries=3, retry_if=lambda exc: exc is retryable_error, initial_retry_delay_seconds=2, + retry_context="Neptune write", ) assert session.execute_write(work) == "success" @@ -54,6 +55,7 @@ class TestRetryableSession: max_retries=3, retry_if=lambda _: False, initial_retry_delay_seconds=2, + retry_context="Neptune write", ) with pytest.raises(RuntimeError) as exc_info: @@ -83,3 +85,81 @@ class TestRetryableSession: driver_sessions[0].close.assert_called_once_with() driver_sessions[1].close.assert_called_once_with() driver_sessions[2].close.assert_not_called() + + def test_retry_exhaustion_with_context_reports_attempts_and_elapsed_time(self): + error = RuntimeError("still retryable") + driver_sessions = [MagicMock() for _ in range(3)] + for driver_session in driver_sessions: + driver_session.execute_write.side_effect = error + session = RetryableSession( + session_factory=MagicMock(side_effect=driver_sessions), + max_retries=2, + retry_if=lambda _: True, + retry_context="Neptune write", + ) + + with ( + patch( + "api.attack_paths.retryable_session.time.monotonic", + side_effect=[100.0, 127.1234], + ), + pytest.raises(RetryExhaustedError) as exc_info, + ): + session.execute_write(MagicMock()) + + assert exc_info.value.method_name == "execute_write" + assert exc_info.value.attempts == 3 + assert exc_info.value.elapsed_seconds == pytest.approx(27.1234) + assert exc_info.value.last_error is error + assert exc_info.value.__cause__ is error + assert str(exc_info.value) == ( + "Neptune write execute_write failed after 3 attempts over 27.123s. " + "Last error: still retryable" + ) + + def test_retry_exhaustion_with_zero_retries_reports_one_attempt(self): + error = ServiceUnavailable("still unavailable") + driver_session = MagicMock() + driver_session.execute_write.side_effect = error + session = RetryableSession( + session_factory=MagicMock(return_value=driver_session), + max_retries=0, + retry_context="Neptune write", + ) + + with pytest.raises(RetryExhaustedError) as exc_info: + session.execute_write(MagicMock()) + + assert exc_info.value.attempts == 1 + + @patch("api.attack_paths.retryable_session.time.sleep") + @patch("api.attack_paths.retryable_session.random.uniform", return_value=3.0) + def test_contextual_retry_warning_includes_original_error( + self, _mock_uniform, _mock_sleep + ): + error = RuntimeError("retryable detail") + first_session = MagicMock() + first_session.execute_write.side_effect = error + second_session = MagicMock() + second_session.execute_write.return_value = "success" + session = RetryableSession( + session_factory=MagicMock(side_effect=[first_session, second_session]), + max_retries=1, + retry_if=lambda _: True, + initial_retry_delay_seconds=2, + retry_context="Neptune write", + ) + + with patch("api.attack_paths.retryable_session.logger.warning") as mock_warning: + assert session.execute_write(MagicMock()) == "success" + + mock_warning.assert_called_once_with( + "%s %s failed with %s: %s; retry %s/%s in %.3fs", + "Neptune write", + "execute_write", + "RuntimeError", + "retryable detail", + 1, + 1, + 3.0, + ) diff --git a/api/src/backend/api/tests/test_sentry.py b/api/src/backend/api/tests/test_sentry.py index 082f563808..67ba1ee01e 100644 --- a/api/src/backend/api/tests/test_sentry.py +++ b/api/src/backend/api/tests/test_sentry.py @@ -1,6 +1,7 @@ import logging from unittest.mock import MagicMock, patch +import pytest from config.settings import sentry as sentry_settings from config.settings.sentry import before_send @@ -82,6 +83,45 @@ def test_before_send_passes_through_non_ignored_log(): assert result == event +def test_before_send_ignores_cartography_missing_temporary_database_log(): + log_record = _make_log_record( + msg="Cartography job failed with %s for database %s", + name="cartography.graph.job", + args=( + "Neo.ClientError.Database.DatabaseNotFound", + "db-tmp-scan-12345678", + ), + ) + + event = MagicMock() + + assert before_send(event, {"log_record": log_record}) is None + + +@pytest.mark.parametrize( + ("logger_name", "message"), + [ + ( + "cartography.graph.job.worker", + "Neo.ClientError.Database.DatabaseNotFound for db-tmp-scan-12345678", + ), + ( + "cartography.graph.job", + "DatabaseNotFound for db-tmp-scan-12345678", + ), + ( + "cartography.graph.job", + "Neo.ClientError.Database.DatabaseNotFound for db-tenant-12345678", + ), + ], +) +def test_before_send_passes_through_similar_cartography_logs(logger_name, message): + log_record = _make_log_record(msg=message, name=logger_name) + event = MagicMock() + + assert before_send(event, {"log_record": log_record}) is event + + def test_before_send_passes_through_non_ignored_exception(): """Test that before_send passes through exceptions that don't contain ignored exceptions.""" exc_info = (Exception, Exception("Some other error message"), None) diff --git a/api/src/backend/api/tests/test_sink.py b/api/src/backend/api/tests/test_sink.py index 487519923e..c778626b62 100644 --- a/api/src/backend/api/tests/test_sink.py +++ b/api/src/backend/api/tests/test_sink.py @@ -11,7 +11,11 @@ from unittest.mock import MagicMock, patch import neo4j import pytest from api.attack_paths import sink as sink_module -from api.attack_paths.database import GraphDatabaseQueryException +from api.attack_paths.database import ( + GraphDatabaseQueryException, + NeptuneWriteRetryExhaustedException, +) +from api.attack_paths.retryable_session import RetryExhaustedError from api.attack_paths.sink import factory from api.attack_paths.sink.neo4j import DATABASE_NOT_FOUND_CODE, Neo4jSink from api.attack_paths.sink.neptune import ( @@ -123,6 +127,14 @@ class TestSinkFactory: assert mock_driver.call_count == 1 +def test_neo4j_sync_batch_size_defaults_to_1000(): + assert Neo4jSink.sync_batch_size == 1000 + + +def test_neptune_sync_batch_size_defaults_to_500(): + assert NeptuneSink.sync_batch_size == 500 + + class TestGetBackendForScan: """``get_backend_for_scan`` routes by the row's recorded sink backend.""" @@ -372,6 +384,7 @@ class TestNeptuneRetryPolicy: assert ( kwargs["initial_retry_delay_seconds"] == NEPTUNE_WRITE_RETRY_DELAY_SECONDS ) + assert kwargs["retry_context"] == "Neptune write" @patch("api.attack_paths.sink.neptune.RetryableSession") def test_reader_session_does_not_enable_write_retry_policy(self, retryable_session): @@ -384,6 +397,48 @@ class TestNeptuneRetryPolicy: kwargs = retryable_session.call_args.kwargs assert kwargs["retry_if"] is None assert kwargs["initial_retry_delay_seconds"] == 0 + assert kwargs["retry_context"] is None + + def test_writer_retry_exhaustion_preserves_neptune_error_details(self): + message = ( + "Unexpected server exception 'Operation failed due to conflicting " + "concurrent operations (please retry), 0 transactions are currently " + "rolling back.'" + ) + error = neo4j.exceptions.Neo4jError._hydrate_neo4j( + code="BoltProtocol.unexpectedException", + message=message, + ) + retry_error = RetryExhaustedError( + retry_context="Neptune write", + method_name="execute_write", + attempts=4, + elapsed_seconds=27.1234, + last_error=error, + ) + sink = NeptuneSink() + driver = MagicMock() + retryable_session = MagicMock() + retryable_session.execute_write.side_effect = retry_error + + with ( + patch.object(sink, "_get_writer", return_value=driver), + patch( + "api.attack_paths.sink.neptune.RetryableSession", + return_value=retryable_session, + ), + pytest.raises(NeptuneWriteRetryExhaustedException) as exc_info, + ): + with sink.get_session() as session: + session.execute_write(MagicMock()) + + assert exc_info.value.code == "BoltProtocol.unexpectedException" + assert str(exc_info.value) == ( + "BoltProtocol.unexpectedException: Neptune write execute_write failed " + "after 4 attempts over 27.123s. Last error: " + f"{message}" + ) + assert exc_info.value.__cause__ is error class TestNeptuneSinkDropSubgraph: diff --git a/api/src/backend/config/django/base.py b/api/src/backend/config/django/base.py index 04251b4479..a079942600 100644 --- a/api/src/backend/config/django/base.py +++ b/api/src/backend/config/django/base.py @@ -312,8 +312,8 @@ ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES = env.int( "ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES", 30 ) ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES = env.int( - "ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES", 2880 -) # 48h + "ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES", 960 +) # 16h # Selects where the persistent attack-paths graph is stored. The scan # temporary database is always Neo4j; only the sink is configurable. diff --git a/api/src/backend/config/settings/sentry.py b/api/src/backend/config/settings/sentry.py index d75f9360e5..a3acbe37bb 100644 --- a/api/src/backend/config/settings/sentry.py +++ b/api/src/backend/config/settings/sentry.py @@ -91,6 +91,13 @@ def before_send(event, hint): log_msg = log_record.getMessage() log_lvl = log_record.levelno + if ( + getattr(log_record, "name", "") == "cartography.graph.job" + and "Neo.ClientError.Database.DatabaseNotFound" in log_msg + and "db-tmp-scan-" in log_msg + ): + return None + # The Neo4j driver logs transient connection errors (defunct # connections, resets) at ERROR level via the `neo4j.io` logger. # `RetryableSession` handles these with retries. If all retries diff --git a/api/src/backend/tasks/jobs/attack_paths/aws.py b/api/src/backend/tasks/jobs/attack_paths/aws.py index 15ecd86a19..4b2b96dd1b 100644 --- a/api/src/backend/tasks/jobs/attack_paths/aws.py +++ b/api/src/backend/tasks/jobs/attack_paths/aws.py @@ -8,6 +8,8 @@ import aioboto3 import boto3 import botocore import neo4j +import neo4j.exceptions +from api.attack_paths.database import DATABASE_NOT_FOUND_CODE from api.models import ( AttackPathsScan as ProwlerAPIAttackPathsScan, ) @@ -347,6 +349,12 @@ def sync_aws_account( ) except Exception as e: + if ( + isinstance(e, neo4j.exceptions.Neo4jError) + and e.code == DATABASE_NOT_FOUND_CODE + ): + raise + logger.info( f"Synced function {func_name} for AWS account {prowler_api_provider.uid} in {time.perf_counter() - func_t0:.3f}s (FAILED)" ) diff --git a/api/src/backend/tasks/jobs/attack_paths/config.py b/api/src/backend/tasks/jobs/attack_paths/config.py index d8ed63a8fc..122fc8a9e0 100644 --- a/api/src/backend/tasks/jobs/attack_paths/config.py +++ b/api/src/backend/tasks/jobs/attack_paths/config.py @@ -10,13 +10,10 @@ NormalizedList = _provider_config.NormalizedList PROVIDER_CONFIGS = _provider_config.PROVIDER_CONFIGS ProviderConfig = _provider_config.ProviderConfig -# Batch size for Neo4j write operations (resource labeling, cleanup) -BATCH_SIZE = env.int("ATTACK_PATHS_BATCH_SIZE", 1000) +# Batch size for graph mutation operations (resource labeling and subgraph deletion) +GRAPH_MUTATION_BATCH_SIZE = env.int("ATTACK_PATHS_GRAPH_MUTATION_BATCH_SIZE", 1000) # Batch size for Postgres findings fetch (keyset pagination page size) FINDINGS_BATCH_SIZE = env.int("ATTACK_PATHS_FINDINGS_BATCH_SIZE", 1000) -# Batch size for temp-to-tenant graph sync (nodes and relationships per cursor page) -SYNC_BATCH_SIZE = env.int("ATTACK_PATHS_SYNC_BATCH_SIZE", 1000) - # Neo4j internal labels (Prowler-specific, not provider-specific) # - `Internet`: Singleton node representing external internet access for exposed-resource queries # - `ProwlerFinding`: Label for finding nodes created by Prowler and linked to cloud resources diff --git a/api/src/backend/tasks/jobs/attack_paths/findings.py b/api/src/backend/tasks/jobs/attack_paths/findings.py index 6cc7ddb2e0..c47c9f1149 100644 --- a/api/src/backend/tasks/jobs/attack_paths/findings.py +++ b/api/src/backend/tasks/jobs/attack_paths/findings.py @@ -21,8 +21,8 @@ from cartography.config import Config as CartographyConfig from celery.utils.log import get_task_logger from prowler.config import config as ProwlerConfig from tasks.jobs.attack_paths.config import ( - BATCH_SIZE, FINDINGS_BATCH_SIZE, + GRAPH_MUTATION_BATCH_SIZE, get_node_uid_field, get_provider_resource_label, get_root_node_label, @@ -135,7 +135,7 @@ def add_resource_label( while labeled_count > 0: result = neo4j_session.run( query, - {"provider_uid": provider_uid, "batch_size": BATCH_SIZE}, + {"provider_uid": provider_uid, "batch_size": GRAPH_MUTATION_BATCH_SIZE}, ) labeled_count = result.single().get("labeled_count", 0) total_labeled += labeled_count diff --git a/api/src/backend/tasks/jobs/attack_paths/scan.py b/api/src/backend/tasks/jobs/attack_paths/scan.py index 32337c6832..e161c9eb8d 100644 --- a/api/src/backend/tasks/jobs/attack_paths/scan.py +++ b/api/src/backend/tasks/jobs/attack_paths/scan.py @@ -372,7 +372,19 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]: except Exception as e: exception_message = utils.stringify_exception(e, "Attack Paths scan failed") - logger.exception(exception_message) + temporary_database_missing = ( + isinstance(e, graph_database.GraphDatabaseQueryException) + and e.code == graph_database.DATABASE_NOT_FOUND_CODE + and tmp_database_name in str(e) + ) + if temporary_database_missing: + logger.warning(exception_message) + else: + logger.exception(exception_message) + cleanup_log_level = ( + logging.WARNING if temporary_database_missing else logging.ERROR + ) + cleanup_exc_info = not temporary_database_missing ingestion_exceptions["global_error"] = exception_message # Recover `graph_data_ready` based on how far the swap got @@ -387,19 +399,24 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]: ) except Exception: - logger.error( - f"Failed to recover `graph_data_ready` for provider {attack_paths_scan.provider_id}", - exc_info=True, + logger.log( + cleanup_log_level, + "Failed to recover `graph_data_ready` for provider " + f"{attack_paths_scan.provider_id}", + exc_info=cleanup_exc_info, ) # 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}", - exc_info=True, + except Exception as cleanup_error: + logger.log( + cleanup_log_level, + "Failed to drop temporary Neo4j database " + f"`{tmp_cartography_config.neo4j_database}` during cleanup: " + f"{cleanup_error}", + exc_info=cleanup_exc_info, ) # Set Attack Paths scan state to FAILED @@ -407,10 +424,12 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]: 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}", - exc_info=True, + except Exception as cleanup_error: + logger.log( + cleanup_log_level, + f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` " + f"(row may have been deleted): {cleanup_error}", + exc_info=cleanup_exc_info, ) raise diff --git a/api/src/backend/tasks/jobs/attack_paths/sync.py b/api/src/backend/tasks/jobs/attack_paths/sync.py index 00c2c585c7..98a52bd48b 100644 --- a/api/src/backend/tasks/jobs/attack_paths/sync.py +++ b/api/src/backend/tasks/jobs/attack_paths/sync.py @@ -30,7 +30,6 @@ from tasks.jobs.attack_paths.config import ( PROVIDER_CONFIGS, PROVIDER_ISOLATION_PROPERTIES, PROVIDER_RESOURCE_LABEL, - SYNC_BATCH_SIZE, NormalizedList, get_provider_label, get_tenant_label, @@ -116,6 +115,7 @@ def sync_nodes( Source and target sessions are opened sequentially per batch to avoid holding two Bolt connections simultaneously for the entire sync duration. """ + batch_size = sink.sync_batch_size t0 = time.perf_counter() last_id = -1 parents_synced = 0 @@ -137,7 +137,7 @@ def sync_nodes( with graph_database.get_session(source_database) as source_session: result = source_session.run( NODE_FETCH_QUERY, - {"last_id": last_id, "batch_size": SYNC_BATCH_SIZE}, + {"last_id": last_id, "batch_size": batch_size}, ) for record in result: batch_count += 1 @@ -156,17 +156,17 @@ def sync_nodes( for labels, batch in parent_groups.items(): rendered_labels = _render_labels(labels, extra_labels) - for sink_batch in _iter_sink_batches(batch): + for sink_batch in _iter_sink_batches(batch, batch_size): sink.write_nodes(target_database, rendered_labels, sink_batch) for child_label, batch in child_groups.items(): rendered_labels = _render_labels((child_label,), extra_labels) - for sink_batch in _iter_sink_batches(batch): + for sink_batch in _iter_sink_batches(batch, batch_size): sink.write_nodes(target_database, rendered_labels, sink_batch) children_synced += len(batch) for rel_type, batch in rel_groups.items(): - for sink_batch in _iter_sink_batches(batch): + for sink_batch in _iter_sink_batches(batch, batch_size): sink.write_relationships( target_database, rel_type, provider_id, sink_batch ) @@ -205,6 +205,7 @@ def sync_relationships( Source and target sessions are opened sequentially per batch to avoid holding two Bolt connections simultaneously for the entire sync duration. """ + batch_size = sink.sync_batch_size t0 = time.perf_counter() last_id = -1 total_synced = 0 @@ -217,7 +218,7 @@ def sync_relationships( with graph_database.get_session(source_database) as source_session: result = source_session.run( RELATIONSHIPS_FETCH_QUERY, - {"last_id": last_id, "batch_size": SYNC_BATCH_SIZE}, + {"last_id": last_id, "batch_size": batch_size}, ) for record in result: batch_count += 1 @@ -229,7 +230,7 @@ def sync_relationships( break for rel_type, batch in grouped.items(): - for sink_batch in _iter_sink_batches(batch): + for sink_batch in _iter_sink_batches(batch, batch_size): sink.write_relationships( target_database, rel_type, provider_id, sink_batch ) @@ -247,10 +248,9 @@ def sync_relationships( def _iter_sink_batches( rows: list[dict[str, Any]], - batch_size: int | None = None, + batch_size: int, ) -> Iterator[list[dict[str, Any]]]: """Yield final sink write batches after source rows have been transformed.""" - batch_size = SYNC_BATCH_SIZE if batch_size is None else batch_size if batch_size <= 0: raise ValueError("Sink batch size must be greater than zero") diff --git a/api/src/backend/tasks/tasks.py b/api/src/backend/tasks/tasks.py index 7a1ff54131..a91ff85c01 100644 --- a/api/src/backend/tasks/tasks.py +++ b/api/src/backend/tasks/tasks.py @@ -11,6 +11,7 @@ from api.compliance import ( from api.db_router import READ_REPLICA_ALIAS from api.db_utils import delete_related_daily_task, rls_transaction from api.decorators import handle_provider_deletion, set_tenant +from api.exceptions import ProviderDeletedException from api.models import ( Finding, Integration, @@ -666,7 +667,13 @@ class AttackPathsScanRLSTask(RLSTask): scan_id = kwargs.get("scan_id") if tenant_id and scan_id: - logger.error(f"Attack paths scan task {task_id} failed: {exc}") + if isinstance(exc, ProviderDeletedException): + logger.warning( + f"Attack paths scan task {task_id} stopped because its provider " + f"or tenant was deleted: {exc}" + ) + else: + logger.error(f"Attack paths scan task {task_id} failed: {exc}") attack_paths_db_utils.fail_attack_paths_scan(tenant_id, scan_id, str(exc)) diff --git a/api/src/backend/tasks/tests/test_attack_paths_aws.py b/api/src/backend/tasks/tests/test_attack_paths_aws.py new file mode 100644 index 0000000000..dc2c59d614 --- /dev/null +++ b/api/src/backend/tasks/tests/test_attack_paths_aws.py @@ -0,0 +1,102 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import neo4j.exceptions +import pytest +from tasks.jobs.attack_paths import aws + +DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound" + + +def _make_neo4j_error(code: str) -> neo4j.exceptions.Neo4jError: + return neo4j.exceptions.Neo4jError._hydrate_neo4j( + code=code, + message="graph query failed", + ) + + +def _resource_functions(failing_sync, following_sync): + return { + "failing_sync": failing_sync, + "following_sync": following_sync, + "permission_relationships": MagicMock(), + "resourcegroupstaggingapi": MagicMock(), + } + + +def test_sync_aws_account_reraises_database_not_found_immediately(): + error = _make_neo4j_error(DATABASE_NOT_FOUND_CODE) + failing_sync = MagicMock(side_effect=error) + following_sync = MagicMock() + + with ( + patch.object( + aws.cartography_aws, + "RESOURCE_FUNCTIONS", + _resource_functions(failing_sync, following_sync), + ), + patch.object(aws.db_utils, "update_attack_paths_scan_progress"), + patch.object(aws.utils, "stringify_exception") as stringify_exception, + patch.object(aws.logger, "warning") as warning, + pytest.raises(neo4j.exceptions.Neo4jError) as exc_info, + ): + aws.sync_aws_account( + SimpleNamespace(uid="123456789012"), + [ + "failing_sync", + "following_sync", + "permission_relationships", + "resourcegroupstaggingapi", + ], + {}, + MagicMock(), + ) + + assert exc_info.value is error + following_sync.assert_not_called() + stringify_exception.assert_not_called() + warning.assert_not_called() + + +@pytest.mark.parametrize( + "error", + [ + _make_neo4j_error("Neo.ClientError.Statement.SyntaxError"), + RuntimeError("resource sync failed"), + ], + ids=["different-neo4j-error", "non-neo4j-error"], +) +def test_sync_aws_account_warns_and_continues_for_other_exceptions(error): + failing_sync = MagicMock(side_effect=error) + following_sync = MagicMock() + + with ( + patch.object( + aws.cartography_aws, + "RESOURCE_FUNCTIONS", + _resource_functions(failing_sync, following_sync), + ), + patch.object(aws.db_utils, "update_attack_paths_scan_progress"), + patch.object( + aws.utils, + "stringify_exception", + return_value="formatted failure", + ), + patch.object(aws.logger, "warning") as warning, + ): + failed_syncs = aws.sync_aws_account( + SimpleNamespace(uid="123456789012"), + [ + "failing_sync", + "following_sync", + "permission_relationships", + "resourcegroupstaggingapi", + ], + {}, + MagicMock(), + ) + + assert failed_syncs == {"failing_sync": "formatted failure"} + following_sync.assert_called_once_with() + warning.assert_called_once() + assert "Continuing to the next AWS sync function" in warning.call_args.args[0] diff --git a/api/src/backend/tasks/tests/test_attack_paths_scan.py b/api/src/backend/tasks/tests/test_attack_paths_scan.py index 3be885a7fc..4409e5f19d 100644 --- a/api/src/backend/tasks/tests/test_attack_paths_scan.py +++ b/api/src/backend/tasks/tests/test_attack_paths_scan.py @@ -1,3 +1,4 @@ +import logging from contextlib import nullcontext from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -5,7 +6,9 @@ from unittest.mock import MagicMock, call, patch from uuid import uuid4 import pytest +from api.attack_paths.database import GraphDatabaseQueryException from api.db_utils import rls_transaction +from api.exceptions import ProviderDeletedException from api.models import ( AttackPathsScan, Finding, @@ -250,6 +253,32 @@ class TestAttackPathsRun: mock_starting.assert_not_called() mock_create_db.assert_not_called() + @pytest.mark.parametrize( + ("ingestion_error", "temporary_database_missing"), + [ + (RuntimeError("ingestion boom"), False), + ( + GraphDatabaseQueryException( + message="Graph not found: db-scan-id", + code="Neo.ClientError.Database.DatabaseNotFound", + ), + True, + ), + ( + GraphDatabaseQueryException( + message="Graph not found: db-tenant-id", + code="Neo.ClientError.Database.DatabaseNotFound", + ), + False, + ), + ], + ids=[ + "regular-error", + "temporary-database-missing", + "sink-database-missing", + ], + ) + @patch("tasks.jobs.attack_paths.scan.logger") @patch( "tasks.jobs.attack_paths.scan.utils.stringify_exception", return_value="Cartography failed: ingestion boom", @@ -302,6 +331,9 @@ class TestAttackPathsRun: mock_drop_db, mock_event_loop, mock_stringify, + mock_logger, + ingestion_error, + temporary_database_missing, tenants_fixture, aws_provider, scans_fixture, @@ -321,7 +353,11 @@ class TestAttackPathsRun: session_ctx = MagicMock() session_ctx.__enter__.return_value = mock_session session_ctx.__exit__.return_value = False - ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom")) + ingestion_fn = MagicMock(side_effect=ingestion_error) + if temporary_database_missing: + mock_finish.side_effect = DatabaseError( + "Save with update_fields did not affect any rows" + ) with ( patch( @@ -337,13 +373,28 @@ class TestAttackPathsRun: return_value=ingestion_fn, ), ): - with pytest.raises(RuntimeError, match="ingestion boom"): + with pytest.raises(type(ingestion_error)): attack_paths_run(str(tenant.id), str(scan.id), "task-456") failure_args = mock_finish.call_args[0] assert failure_args[0] is attack_paths_scan assert failure_args[1] == StateChoices.FAILED assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"} + mock_drop_db.assert_called_once_with("db-scan-id") + if temporary_database_missing: + mock_logger.warning.assert_any_call("Cartography failed: ingestion boom") + mock_logger.exception.assert_not_called() + mock_logger.log.assert_called_once_with( + logging.WARNING, + f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` " + "(row may have been deleted): Save with update_fields did not affect " + "any rows", + exc_info=False, + ) + else: + mock_logger.exception.assert_called_once_with( + "Cartography failed: ingestion boom" + ) @patch( "tasks.jobs.attack_paths.scan.utils.stringify_exception", @@ -1265,6 +1316,33 @@ class TestAttackPathsScanRLSTaskOnFailure: mock_fail.assert_called_once_with("t-1", "s-1", "boom") + def test_on_failure_logs_provider_deletion_as_warning(self): + from tasks.tasks import AttackPathsScanRLSTask + + task = AttackPathsScanRLSTask() + error = ProviderDeletedException("provider deleted") + + with ( + patch("tasks.tasks.logger") as mock_logger, + patch( + "tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan" + ) as mock_fail, + ): + task.on_failure( + exc=error, + task_id="task-abc", + args=(), + kwargs={"tenant_id": "t-1", "scan_id": "s-1"}, + _einfo=None, + ) + + mock_logger.warning.assert_called_once_with( + "Attack paths scan task task-abc stopped because its provider or tenant " + "was deleted: provider deleted" + ) + mock_logger.error.assert_not_called() + mock_fail.assert_called_once_with("t-1", "s-1", "provider deleted") + def test_on_failure_skips_when_missing_kwargs(self): from tasks.tasks import AttackPathsScanRLSTask @@ -1896,7 +1974,7 @@ class TestSyncNodes: mock_source_1.run.return_value = [row] mock_source_2 = MagicMock() mock_source_2.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1000) with patch( "tasks.jobs.attack_paths.sync.graph_database.get_session", @@ -1933,7 +2011,7 @@ class TestSyncNodes: src_1.run.return_value = [row] src_2 = MagicMock() src_2.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1000) sink.write_nodes.side_effect = lambda *_a, **_kw: call_order.append( "sink:write" ) @@ -1969,18 +2047,15 @@ class TestSyncNodes: src_2.run.return_value = [row_b] src_3 = MagicMock() src_3.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1) - with ( - patch( - "tasks.jobs.attack_paths.sync.graph_database.get_session", - side_effect=[ - _make_session_ctx(src_1), - _make_session_ctx(src_2), - _make_session_ctx(src_3), - ], - ), - patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1), + with patch( + "tasks.jobs.attack_paths.sync.graph_database.get_session", + side_effect=[ + _make_session_ctx(src_1), + _make_session_ctx(src_2), + _make_session_ctx(src_3), + ], ): result = sync_module.sync_nodes("src", "tgt", "t-1", "p-1", sink, []) @@ -2009,17 +2084,14 @@ class TestSyncNodes: src_1.run.return_value = [row] src_2 = MagicMock() src_2.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=2) - with ( - patch( - "tasks.jobs.attack_paths.sync.graph_database.get_session", - side_effect=[ - _make_session_ctx(src_1), - _make_session_ctx(src_2), - ], - ), - patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 2), + with patch( + "tasks.jobs.attack_paths.sync.graph_database.get_session", + side_effect=[ + _make_session_ctx(src_1), + _make_session_ctx(src_2), + ], ): result = sync_module.sync_nodes( "src", "tgt", "t-1", "p-1", sink, normalized_lists @@ -2037,7 +2109,7 @@ class TestSyncNodes: def test_sync_nodes_empty_source_returns_zero(self): src = MagicMock() src.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1000) with patch( "tasks.jobs.attack_paths.sync.graph_database.get_session", @@ -2066,7 +2138,7 @@ class TestSyncRelationships: src_1.run.return_value = [row] src_2 = MagicMock() src_2.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1000) sink.write_relationships.side_effect = lambda *_a, **_kw: call_order.append( "sink:write" ) @@ -2104,18 +2176,15 @@ class TestSyncRelationships: src_2.run.return_value = [row_b] src_3 = MagicMock() src_3.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1) - with ( - patch( - "tasks.jobs.attack_paths.sync.graph_database.get_session", - side_effect=[ - _make_session_ctx(src_1), - _make_session_ctx(src_2), - _make_session_ctx(src_3), - ], - ), - patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1), + with patch( + "tasks.jobs.attack_paths.sync.graph_database.get_session", + side_effect=[ + _make_session_ctx(src_1), + _make_session_ctx(src_2), + _make_session_ctx(src_3), + ], ): total = sync_module.sync_relationships("src", "tgt", "p-1", sink) @@ -2140,17 +2209,14 @@ class TestSyncRelationships: src_1.run.return_value = rows src_2 = MagicMock() src_2.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=2) - with ( - patch( - "tasks.jobs.attack_paths.sync.graph_database.get_session", - side_effect=[ - _make_session_ctx(src_1), - _make_session_ctx(src_2), - ], - ), - patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 2), + with patch( + "tasks.jobs.attack_paths.sync.graph_database.get_session", + side_effect=[ + _make_session_ctx(src_1), + _make_session_ctx(src_2), + ], ): total = sync_module.sync_relationships("src", "tgt", "p-1", sink) @@ -2163,7 +2229,7 @@ class TestSyncRelationships: def test_sync_relationships_empty_source_returns_zero(self): src = MagicMock() src.run.return_value = [] - sink = MagicMock() + sink = MagicMock(sync_batch_size=1000) with patch( "tasks.jobs.attack_paths.sync.graph_database.get_session", @@ -3056,6 +3122,61 @@ class TestCleanupStaleAttackPathsScans: ap_scan.refresh_from_db() assert ap_scan.state == StateChoices.FAILED + @pytest.mark.parametrize( + ("age_seconds", "should_clean"), + [ + (960 * 60 - 1, False), + (960 * 60, False), + (960 * 60 + 1, True), + ], + ) + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_stale_threshold_boundary_is_strict( + self, + mock_ping, + mock_revoke, + mock_drop_db, + mock_recover, + age_seconds, + should_clean, + tenants_fixture, + aws_provider, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + now = datetime.now(tz=UTC) + ap_scan, task_result = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + started_at=now - timedelta(seconds=age_seconds), + worker="live-worker@host", + ) + mock_ping.return_value = ({"live-worker@host"}, set()) + + with patch("tasks.jobs.attack_paths.cleanup.datetime") as mock_datetime: + mock_datetime.now.return_value = now + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == int(should_clean) + ap_scan.refresh_from_db() + expected_state = StateChoices.FAILED if should_clean else StateChoices.EXECUTING + assert ap_scan.state == expected_state + if should_clean: + mock_revoke.assert_called_once_with(task_result, terminate=True) + mock_drop_db.assert_called_once() + mock_recover.assert_called_once() + else: + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") @patch(