do retries and add tests

Author: zprobst

nodestream_plugin_neo4j/neo4j_database.py

@@ -1,19 +1,19 @@
+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,
-    TransientError,
     ServiceUnavailable,
     SessionExpired,
+    TransientError,
 )
 from nodestream.file_io import LazyLoadedArgument
 
 from .query import Query
 
-
 RETRYABLE_EXCEPTIONS = (TransientError, ServiceUnavailable, SessionExpired, AuthError)
 
 
@@ -52,21 +52,27 @@ def from_configuration(
         password: Union[str, LazyLoadedArgument],
         database_name: str = "neo4j",
         max_retry_attempts: int = 3,
+        retry_factor: int = 1,
         **driver_kwargs,
     ):
         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)
+        return cls(driver_factory, database_name, max_retry_attempts, retry_factor)
 
     def __init__(
-        self, driver_factory, database_name: str, max_retry_attempts: int
+        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
 
     def acquire_driver(self) -> AsyncDriver:
@@ -97,7 +103,10 @@ def log_record(self, record: Record):
         )
 
     async def _execute_query(
-        self, query: Query, log_result: bool = False, routing_=RoutingControl.WRITE
+        self,
+        query: Query,
+        log_result: bool = False,
+        routing_=RoutingControl.WRITE,
     ) -> Record:
         result = await self.driver.execute_query(
             query.query_statement,
@@ -113,7 +122,10 @@ async def _execute_query(
         return records
 
     async def execute(
-        self, query: Query, log_result: bool = False, routing_=RoutingControl.WRITE
+        self,
+        query: Query,
+        log_result: bool = False,
+        routing_=RoutingControl.WRITE,
     ) -> Iterable[Record]:
         self.log_query_start(query)
         attempts = 0
@@ -126,7 +138,7 @@ async def execute(
                     f"Error executing query, retrying. Attempt {attempts + 1}",
                     exc_info=e,
                 )
+                await asyncio.sleep(self.retry_factor * attempts)
                 self.acquire_driver()
                 if attempts >= self.max_retry_attempts:
-                    message = f"Failed to execute after {self.max_retry_attempts} tries"
-                    raise Exception(message) from e
+                    raise e

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", 3)
+    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()