Retry on Arrow unavailable and timeout exceptions

Author: FlorentinD

changelog.md

@@ -15,6 +15,7 @@
 * Reduce calls to the Aura API during GdsSessions::get_or_create.
 * Improve error message when a query is interrupted by a signal (SIGINT or SIGTERM).
 * Improve error message if session is expired.
+* Improve robustness of Arrow client against connection errors such as `FlightUnavailableError` and `FlightTimedOutError`.
 
 
 ## Other changes

graphdatascience/query_runner/gds_arrow_client.py

@@ -12,11 +12,27 @@
 import pandas
 import pyarrow
 from neo4j.exceptions import ClientError
-from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Table, chunked_array, flight
+from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Schema, Table, chunked_array, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import (
+    ClientMiddleware,
+    ClientMiddlewareFactory,
+    FlightDescriptor,
+    FlightMetadataReader,
+    FlightStreamWriter,
+    FlightTimedOutError,
+    FlightUnavailableError,
+)
 from pyarrow.types import is_dictionary
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
+from tenacity import (
+    retry,
+    retry_any,
+    retry_if_exception_type,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_exponential,
+    wait_fixed,
+)
 
 from ..semantic_version.semantic_version import SemanticVersion
 from ..version import __version__
@@ -131,6 +147,11 @@ def connection_info(self) -> tuple[str, int]:
         """
         return self._host, self._port
 
+    @retry(
+        retry=retry_any(retry_if_exception_type(FlightTimedOutError), retry_if_exception_type(FlightUnavailableError)),
+        stop=stop_after_attempt(3),
+        wait=wait_fixed(1),
+    )
     def request_token(self) -> Optional[str]:
         """
         Requests a token from the server and returns it.
@@ -571,6 +592,11 @@ def _client(self) -> flight.FlightClient:
             self._flight_client = self._instantiate_flight_client()
         return self._flight_client
 
+    @retry(
+        retry=retry_any(retry_if_exception_type(FlightTimedOutError), retry_if_exception_type(FlightUnavailableError)),
+        stop=(stop_after_delay(10) | stop_after_attempt(5)),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+    )
     def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str, Any]:
         action_type = self._versioned_action_type(action_type)
 
@@ -587,6 +613,16 @@ def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str,
             self.handle_flight_error(e)
             raise e  # unreachable
 
+    @retry(
+        retry=retry_any(retry_if_exception_type(FlightTimedOutError), retry_if_exception_type(FlightUnavailableError)),
+        stop=(stop_after_delay(10) | stop_after_attempt(5)),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+    )
+    def _safe_do_put(
+        self, upload_descriptor: FlightDescriptor, schema: Schema
+    ) -> tuple[FlightStreamWriter, FlightMetadataReader]:
+        return self._client().do_put(upload_descriptor, schema)  # type: ignore
+
     def _upload_data(
         self,
         graph_name: str,
@@ -605,8 +641,7 @@ def _upload_data(
         flight_descriptor = self._versioned_flight_descriptor({"name": graph_name, "entity_type": entity_type})
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 
-        client = self._client()
-        put_stream, ack_stream = client.do_put(upload_descriptor, batches[0].schema)
+        put_stream, ack_stream = self._safe_do_put(upload_descriptor, batches[0].schema)
 
         @retry(
             stop=(stop_after_delay(10) | stop_after_attempt(5)),
@@ -629,6 +664,11 @@ def upload_batch(p: RecordBatch) -> None:
         except Exception as e:
             GdsArrowClient.handle_flight_error(e)
 
+    @retry(
+        retry=retry_any(retry_if_exception_type(FlightTimedOutError), retry_if_exception_type(FlightUnavailableError)),
+        stop=(stop_after_delay(10) | stop_after_attempt(5)),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+    )
     def _do_get(
         self,
         database: str,

graphdatascience/tests/unit/test_gds_arrow_client.py

@@ -4,16 +4,22 @@
 
 import pyarrow as pa
 import pytest
-from pyarrow import flight
 from pyarrow._flight import GeneratorStream
-from pyarrow.flight import Action, Ticket
+from pyarrow.flight import (
+    Action,
+    FlightServerBase,
+    FlightServerError,
+    FlightTimedOutError,
+    FlightUnavailableError,
+    Ticket,
+)
 
 from graphdatascience.query_runner.gds_arrow_client import AuthMiddleware, GdsArrowClient
 
 ActionParam = Union[str, tuple[str, Any], Action]
 
 
-class FlightServer(flight.FlightServerBase):  # type: ignore
+class FlightServer(FlightServerBase):  # type: ignore
     def __init__(self, location: str = "grpc://0.0.0.0:0", **kwargs: dict[str, Any]) -> None:
         super(FlightServer, self).__init__(location, **kwargs)
         self._location: str = location
@@ -49,18 +55,81 @@ def do_action(self, context: Any, action: ActionParam) -> list[bytes]:
         return [json.dumps(response).encode("utf-8")]
 
 
+class FlakyFlightServer(FlightServerBase):  # type: ignore
+    def __init__(self, location: str = "grpc://0.0.0.0:0", **kwargs: dict[str, Any]) -> None:
+        super(FlakyFlightServer, self).__init__(location, **kwargs)
+        self._location: str = location
+        self._actions: list[ActionParam] = []
+        self._tickets: list[Ticket] = []
+        self._expected_failures = [
+            FlightUnavailableError("Flight server is unavailable", "some reason"),
+            FlightTimedOutError("Time out for some reason", "still timed out"),
+        ]
+        self._expected_retries = len(self._expected_failures) + 1
+
+    def expected_retries(self) -> int:
+        return self._expected_retries
+
+    def do_get(self, context: Any, ticket: Ticket) -> GeneratorStream:
+        self._tickets.append(ticket)
+
+        if len(self._expected_failures) > 0:
+            raise self._expected_failures.pop()
+
+        table = pa.Table.from_pydict({"ids": [42, 1337, 1234]})
+        return GeneratorStream(schema=table.schema, generator=table.to_batches())
+
+    def do_action(self, context: Any, action: ActionParam) -> list[bytes]:
+        self._actions.append(action)
+
+        if len(self._expected_failures) > 0:
+            raise self._expected_failures.pop()
+
+        if isinstance(action, Action):
+            actionType = action.type
+        elif isinstance(action, tuple):
+            actionType = action[0]
+        elif isinstance(action, str):
+            actionType = action
+
+        response: dict[str, Any] = {}
+        if "CREATE" in actionType:
+            response = {"name": "g"}
+        elif "NODE_LOAD_DONE" in actionType:
+            response = {"name": "g", "node_count": 42}
+        elif "RELATIONSHIP_LOAD_DONE" in actionType:
+            response = {"name": "g", "relationship_count": 42}
+        elif "TRIPLET_LOAD_DONE" in actionType:
+            response = {"name": "g", "node_count": 42, "relationship_count": 1337}
+        else:
+            response = {}
+        return [json.dumps(response).encode("utf-8")]
+
+
 @pytest.fixture()
 def flight_server() -> Generator[None, FlightServer, None]:
     with FlightServer() as server:
         yield server
 
 
+@pytest.fixture()
+def flaky_flight_server() -> Generator[None, FlakyFlightServer, None]:
+    with FlakyFlightServer() as server:
+        yield server
+
+
 @pytest.fixture()
 def flight_client(flight_server: FlightServer) -> Generator[GdsArrowClient, None, None]:
     with GdsArrowClient("localhost", flight_server.port) as client:
         yield client
 
 
+@pytest.fixture()
+def flaky_flight_client(flaky_flight_server: FlakyFlightServer) -> Generator[GdsArrowClient, None, None]:
+    with GdsArrowClient("localhost", flaky_flight_server.port) as client:
+        yield client
+
+
 def test_create_graph_with_defaults(flight_server: FlightServer, flight_client: GdsArrowClient) -> None:
     flight_client.create_graph("g", "DB")
     actions = flight_server._actions
@@ -87,6 +156,15 @@ def test_create_graph_with_options(flight_server: FlightServer, flight_client: G
     )
 
 
+def test_create_graph_with_flaky_server(
+    flaky_flight_server: FlakyFlightServer, flaky_flight_client: GdsArrowClient
+) -> None:
+    flaky_flight_client.create_graph("g", "DB")
+    actions = flaky_flight_server._actions
+    assert len(actions) == flaky_flight_server.expected_retries()
+    assert_action(actions[0], "v1/CREATE_GRAPH", {"name": "g", "database_name": "DB"})
+
+
 def test_create_graph_from_triplets_with_defaults(flight_server: FlightServer, flight_client: GdsArrowClient) -> None:
     flight_client.create_graph_from_triplets("g", "DB")
     actions = flight_server._actions
@@ -202,6 +280,22 @@ def test_get_node_property(flight_server: FlightServer, flight_client: GdsArrowC
     )
 
 
+def test_flakey_get_node_property(flaky_flight_server: FlakyFlightServer, flaky_flight_client: GdsArrowClient) -> None:
+    flaky_flight_client.get_node_properties("g", "db", "id", ["Person"], concurrency=42)
+    tickets = flaky_flight_server._tickets
+    assert len(tickets) == flaky_flight_server.expected_retries()
+    assert_ticket(
+        tickets[0],
+        {
+            "concurrency": 42,
+            "configuration": {"list_node_labels": False, "node_labels": ["Person"], "node_property": "id"},
+            "database_name": "db",
+            "graph_name": "g",
+            "procedure_name": "gds.graph.nodeProperty.stream",
+        },
+    )
+
+
 def test_get_node_properties(flight_server: FlightServer, flight_client: GdsArrowClient) -> None:
     flight_client.get_node_properties("g", "db", ["foo", "bar"], ["Person"], list_node_labels=True, concurrency=42)
     tickets = flight_server._tickets
@@ -314,21 +408,21 @@ def test_auth_middleware_bad_headers() -> None:
 
 def test_handle_flight_error() -> None:
     with pytest.raises(
-        flight.FlightServerError,
+        FlightServerError,
         match="FlightServerError: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database.",
     ):
         GdsArrowClient.handle_flight_error(
-            flight.FlightServerError(
+            FlightServerError(
                 'FlightServerError: Flight RPC failed with message: org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database.. gRPC client debug context: UNKNOWN:Error received from peer ipv4:35.241.177.75:8491 {created_time:"2024-08-29T15:59:03.828903999+02:00", grpc_status:2, grpc_message:"org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database."}. Client context: IOError: Server never sent a data message. Detail: Internal'
             )
         )
 
     with pytest.raises(
-        flight.FlightServerError,
+        FlightServerError,
         match=re.escape("FlightServerError: UNKNOWN: Unexpected configuration key(s): [undirectedRelationshipTypes]"),
     ):
         GdsArrowClient.handle_flight_error(
-            flight.FlightServerError(
+            FlightServerError(
                 "FlightServerError: Flight returned internal error, with message: org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Unexpected configuration key(s): [undirectedRelationshipTypes]"
             )
         )