Merge pull request #22 from nodestream-proj/handle-spurious-errors

Handle spurious errors

Author: zprobst

nodestream_plugin_neo4j/neo4j_database.py

@@ -1,12 +1,21 @@
+import asyncio
 from logging import getLogger
 from typing import Awaitable, Iterable, Tuple, Union
 
 from neo4j import AsyncDriver, AsyncGraphDatabase, AsyncSession, Record, RoutingControl
 from neo4j.auth_management import AsyncAuthManagers
+from neo4j.exceptions import (
+    AuthError,
+    ServiceUnavailable,
+    SessionExpired,
+    TransientError,
+)
 from nodestream.file_io import LazyLoadedArgument
 
 from .query import Query
 
+RETRYABLE_EXCEPTIONS = (TransientError, ServiceUnavailable, SessionExpired, AuthError)
+
 
 def auth_provider_factory(
     username: Union[str, LazyLoadedArgument],
@@ -42,20 +51,40 @@ def from_configuration(
         username: Union[str, LazyLoadedArgument],
         password: Union[str, LazyLoadedArgument],
         database_name: str = "neo4j",
-        **driver_kwargs
+        max_retry_attempts: int = 3,
+        retry_factor: int = 1,
+        **driver_kwargs,
     ):
-        auth = AsyncAuthManagers.basic(auth_provider_factory(username, password))
-        driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs)
-        return cls(driver, database_name)
+        def driver_factory():
+            auth = AsyncAuthManagers.basic(auth_provider_factory(username, password))
+            return AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs)
+
+        return cls(driver_factory, database_name, max_retry_attempts, retry_factor)
 
-    def __init__(self, driver: AsyncDriver, database_name: str) -> None:
-        self.driver = driver
+    def __init__(
+        self,
+        driver_factory,
+        database_name: str,
+        max_retry_attempts: int = 3,
+        retry_factor: float = 1,
+    ) -> None:
+        self.driver_factory = driver_factory
         self.database_name = database_name
         self.logger = getLogger(self.__class__.__name__)
+        self.max_retry_attempts = max_retry_attempts
+        self.retry_factor = retry_factor
+        self._driver = None
 
-    async def execute(
-        self, query: Query, log_result: bool = False, routing_=RoutingControl.WRITE
-    ) -> Iterable[Record]:
+    def acquire_driver(self) -> AsyncDriver:
+        self._driver = self.driver_factory()
+
+    @property
+    def driver(self):
+        if self._driver is None:
+            self.acquire_driver()
+        return self._driver
+
+    def log_query_start(self, query: Query):
         self.logger.info(
             "Executing Cypher Query to Neo4j",
             extra={
@@ -64,23 +93,52 @@ async def execute(
             },
         )
 
+    def log_record(self, record: Record):
+        self.logger.info(
+            "Gathered Query Results",
+            extra=dict(**record, uri=self.driver._pool.address.host),
+        )
+
+    async def _execute_query(
+        self,
+        query: Query,
+        log_result: bool = False,
+        routing_=RoutingControl.WRITE,
+    ) -> Record:
         result = await self.driver.execute_query(
             query.query_statement,
             query.parameters,
             database_=self.database_name,
             routing_=routing_,
         )
+        records = result.records
         if log_result:
-            for record in result.records:
-                self.logger.info(
-                    "Gathered Query Results",
-                    extra=dict(
-                        **record,
-                        query=query.query_statement,
-                        uri=self.driver._pool.address.host
-                    ),
+            for record in records:
+                self.log_record(record)
+
+        return records
+
+    async def execute(
+        self,
+        query: Query,
+        log_result: bool = False,
+        routing_=RoutingControl.WRITE,
+    ) -> Iterable[Record]:
+        self.log_query_start(query)
+        attempts = 0
+        while True:
+            attempts += 1
+            try:
+                return await self._execute_query(query, log_result, routing_)
+            except RETRYABLE_EXCEPTIONS as e:
+                self.logger.warning(
+                    f"Error executing query, retrying. Attempt {attempts + 1}",
+                    exc_info=e,
                 )
-        return result.records
+                await asyncio.sleep(self.retry_factor * attempts)
+                self.acquire_driver()
+                if attempts >= self.max_retry_attempts:
+                    raise e
 
     def session(self) -> AsyncSession:
         return self.driver.session(database=self.database_name)

tests/unit/test_neo4j_database.py

@@ -1,6 +1,7 @@
 import pytest
 from hamcrest import assert_that, equal_to
 from neo4j import AsyncDriver, RoutingControl
+from neo4j.exceptions import TransientError
 from nodestream.file_io import LazyLoadedArgument
 
 from nodestream_plugin_neo4j.neo4j_database import (
@@ -9,30 +10,58 @@
 )
 from nodestream_plugin_neo4j.query import Query
 
+A_QUERY = Query("MATCH (n) RETURN n LIMIT $limit", {"limit": 2})
+SOME_RECORDS = [
+    {"n": {"name": "foo"}},
+    {"n": {"name": "bar"}},
+]
+
 
 @pytest.fixture
 def database_connection(mocker):
-    return Neo4jDatabaseConnection(mocker.AsyncMock(AsyncDriver), "neo4j")
+    return Neo4jDatabaseConnection(
+        lambda: mocker.AsyncMock(AsyncDriver), "neo4j", 2, 0.1
+    )
 
 
 @pytest.mark.asyncio
 async def test_execute(database_connection):
-    query = Query("MATCH (n) RETURN n LIMIT $limit", {"limit": 2})
-    records = [
-        {"n": {"name": "foo"}},
-        {"n": {"name": "bar"}},
-    ]
-    database_connection.driver.execute_query.return_value.records = records
-    result = await database_connection.execute(query, log_result=True)
-    assert_that(result, equal_to(records))
+    database_connection.driver.execute_query.return_value.records = SOME_RECORDS
+    result = await database_connection.execute(A_QUERY, log_result=True)
+    assert_that(result, equal_to(SOME_RECORDS))
     database_connection.driver.execute_query.assert_called_once_with(
-        query.query_statement,
-        query.parameters,
+        A_QUERY.query_statement,
+        A_QUERY.parameters,
         database_="neo4j",
         routing_=RoutingControl.WRITE,
     )
 
 
+@pytest.mark.asyncio
+async def test_execute_fail_and_then_succeed(database_connection, mocker):
+    database_connection.acquire_driver = mocker.Mock(
+        wraps=database_connection.acquire_driver
+    )
+    database_connection.driver.execute_query.side_effect = [
+        TransientError("Failed to execute query"),
+        SOME_RECORDS,
+    ]
+    await database_connection.execute(A_QUERY)
+    assert_that(database_connection.acquire_driver.call_count, equal_to(2))
+
+
+@pytest.mark.asyncio
+async def test_execute_fail_and_then_fail(database_connection, mocker):
+    def driver_factory():
+        driver = mocker.AsyncMock(AsyncDriver)
+        driver.execute_query.side_effect = TransientError("Failed to execute query")
+        return driver
+
+    database_connection.driver_factory = driver_factory
+    with pytest.raises(TransientError):
+        await database_connection.execute(A_QUERY)
+
+
 @pytest.mark.asyncio
 async def test_session(database_connection):
     session = database_connection.session()