ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-01-25 02:08:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(attack-paths): adding stability to Neo4j driver and session (#9842)

Browse Source
This commit is contained in:
Josema Camacho
2026-01-21 12:44:31 +01:00
committed by GitHub
parent 7b839d9f9e
commit 9c76dafaa4
2 changed files with 15 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
api/src/backend/api/attack_paths/database.py
View File

@@ -42,7 +42,12 @@ def init_driver() -> neo4j.Driver:
config = settings.DATABASES["neo4j"] config = settings.DATABASES["neo4j"]
_driver = neo4j.GraphDatabase.driver( _driver = neo4j.GraphDatabase.driver(
uri, auth=(config["USER"], config["PASSWORD"]) uri,
auth=(config["USER"], config["PASSWORD"]),
keep_alive=True,
max_connection_lifetime=7200,
connection_acquisition_timeout=120,
max_connection_pool_size=50,
) )
_driver.verify_connectivity() _driver.verify_connectivity()
@@ -71,7 +76,6 @@ def get_session(database: str | None = None) -> Iterator[RetryableSession]:
try: try:
session_wrapper = RetryableSession( session_wrapper = RetryableSession(
session_factory=lambda: get_driver().session(database=database), session_factory=lambda: get_driver().session(database=database),
close_driver=close_driver, # Just to avoid circular imports
max_retries=SERVICE_UNAVAILABLE_MAX_RETRIES, max_retries=SERVICE_UNAVAILABLE_MAX_RETRIES,
) )
yield session_wrapper yield session_wrapper

16
api/src/backend/api/attack_paths/retryable_session.py
View File

@@ -17,11 +17,9 @@ class RetryableSession:
def __init__( def __init__(
self, self,
session_factory: Callable[[], neo4j.Session], session_factory: Callable[[], neo4j.Session],
close_driver: Callable[[], None], # Just to avoid circular imports
max_retries: int, max_retries: int,
) -> None: ) -> None:
self._session_factory = session_factory self._session_factory = session_factory
self._close_driver = close_driver
self._max_retries = max(0, max_retries) self._max_retries = max(0, max_retries)
self._session = self._session_factory() self._session = self._session_factory()
@@ -58,7 +56,7 @@ class RetryableSession:
def _call_with_retry(self, method_name: str, *args: Any, **kwargs: Any) -> Any: def _call_with_retry(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
attempt = 0 attempt = 0
last_exc: neo4j.exceptions.ServiceUnavailable | None = None last_exc: Exception | None = None
while attempt <= self._max_retries: while attempt <= self._max_retries:
try: try:
@@ -66,7 +64,8 @@ class RetryableSession:
return method(*args, **kwargs) return method(*args, **kwargs)
except ( except (
neo4j.exceptions.ServiceUnavailable neo4j.exceptions.ServiceUnavailable,
ConnectionResetError,
) as exc: # pragma: no cover - depends on infra ) as exc: # pragma: no cover - depends on infra
last_exc = exc last_exc = exc
attempt += 1 attempt += 1
@@ -75,7 +74,7 @@ class RetryableSession:
raise raise
logger.warning( logger.warning(
f"Neo4j session {method_name} failed with ServiceUnavailable ({attempt}/{self._max_retries} attempts). Retrying..." f"Neo4j session {method_name} failed with {type(exc).__name__} ({attempt}/{self._max_retries} attempts). Retrying..."
) )
self._refresh_session() self._refresh_session()
@@ -83,7 +82,10 @@ class RetryableSession:
def _refresh_session(self) -> None: def _refresh_session(self) -> None:
if self._session is not None: if self._session is not None:
self._session.close() try:
self._session.close()
except Exception:
# Best-effort close; failures just mean we open a new session below
pass
self._close_driver()
self._session = self._session_factory() self._session = self._session_factory()