Support serializing authenticated flight client

Author: FlorentinD

graphdatascience/arrow_client/authenticated_flight_client.py

@@ -4,11 +4,12 @@
 import logging
 import platform
 from dataclasses import dataclass
-from typing import Any, Iterator
+from types import TracebackType
+from typing import Any, Iterator, Type
 
 import certifi
+from pyarrow import Schema, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow import flight, Schema
 from pyarrow._flight import (
     Action,
     ActionType,
@@ -19,13 +20,11 @@
     Result,
     Ticket,
 )
-from tenacity import retry, retry_any, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
-from graphdatascience.retry_utils.retry_config import RetryConfig
+from graphdatascience.retry_utils.retry_config import ExponentialWaitConfig, RetryConfigV2, StopConfig
 
-from ..retry_utils.retry_utils import before_log
 from ..version import __version__
 from .middleware.auth_middleware import AuthFactory, AuthMiddleware
 from .middleware.user_agent_middleware import UserAgentFactory
@@ -39,7 +38,7 @@ def create(
         encrypted: bool = False,
         arrow_client_options: dict[str, Any] | None = None,
         connection_string_override: str | None = None,
-        retry_config: RetryConfig | None = None,
+        retry_config: RetryConfigV2 | None = None,
         advertised_listen_address: tuple[str, int] | None = None,
     ) -> AuthenticatedArrowClient:
         connection_string: str
@@ -51,14 +50,14 @@ def create(
         host, port = connection_string.split(":")
 
         if retry_config is None:
-            retry_config = RetryConfig(
-                retry=retry_any(
-                    retry_if_exception_type(FlightTimedOutError),
-                    retry_if_exception_type(FlightUnavailableError),
-                    retry_if_exception_type(FlightInternalError),
-                ),
-                stop=(stop_after_delay(10) | stop_after_attempt(5)),
-                wait=wait_exponential(multiplier=1, min=1, max=10),
+            retry_config = RetryConfigV2(
+                retryable_exceptions=[
+                    FlightTimedOutError,
+                    FlightUnavailableError,
+                    FlightInternalError,
+                ],
+                stop=StopConfig(after_delay=10, after_attempt=5),
+                wait=ExponentialWaitConfig(multiplier=1, min=1, max=10),
             )
 
         return AuthenticatedArrowClient(
@@ -74,7 +73,7 @@ def create(
     def __init__(
         self,
         host: str,
-        retry_config: RetryConfig,
+        retry_config: RetryConfigV2,
         port: int = 8491,
         auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
@@ -117,7 +116,7 @@ def __init__(
             self._auth_middleware = AuthMiddleware(auth)
         self.advertised_listen_address = advertised_listen_address
 
-        self._flight_client = self._instantiate_flight_client()
+        self._flight_client: flight.FlightClient = self._instantiate_flight_client()
 
     def connection_info(self) -> ConnectionInfo:
         """
@@ -155,13 +154,7 @@ def request_token(self) -> str | None:
             a token from the server and returns it.
         """
 
-        @retry(
-            reraise=True,
-            before=before_log("Request token", self._logger, logging.DEBUG),
-            retry=self._retry_config.retry,
-            stop=self._retry_config.stop,
-            wait=self._retry_config.wait,
-        )
+        @self._retry_config.decorator(operation_name="Request token", logger=self._logger)
         def auth_with_retry() -> None:
             client = self._flight_client
             if self._auth:
@@ -183,13 +176,7 @@ def do_action(self, endpoint: str, payload: bytes | dict[str, Any]) -> Iterator[
         return self._flight_client.do_action(Action(endpoint, payload_bytes))  # type: ignore
 
     def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
-        @retry(
-            reraise=True,
-            before=before_log("Send action", self._logger, logging.DEBUG),
-            retry=self._retry_config.retry,
-            stop=self._retry_config.stop,
-            wait=self._retry_config.wait,
-        )
+        @self._retry_config.decorator(operation_name="Send action", logger=self._logger)
         def run_with_retry() -> list[Result]:
             # the Flight response error code is only checked on iterator consumption
             # we eagerly collect iterator here to trigger retry in case of an error
@@ -203,18 +190,27 @@ def list_actions(self) -> set[ActionType]:
     def do_put_with_retry(
         self, descriptor: flight.FlightDescriptor, schema: Schema
     ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
-        @retry(
-            reraise=True,
-            before=before_log("Do put", self._logger, logging.DEBUG),
-            retry=self._retry_config.retry,
-            stop=self._retry_config.stop,
-            wait=self._retry_config.wait,
-        )
-        def run_with_retry() -> list[Result]:
-            return self._client().do_put(descriptor, schema)
+        @self._retry_config.decorator(operation_name="Do put", logger=self._logger)
+        def run_with_retry() -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
+            return self._flight_client.do_put(descriptor, schema)  # type: ignore
 
         return run_with_retry()
 
+    def __enter__(self) -> AuthenticatedArrowClient:
+        return self
+
+    def __exit__(
+        self,
+        exception_type: Type[BaseException] | None,
+        exception_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._flight_client:
+            self._flight_client.close()
+
     def _instantiate_flight_client(self) -> flight.FlightClient:
         location = (
             flight.Location.for_grpc_tls(self._host, self._port)

graphdatascience/retry_utils/retry_config.py

@@ -1,12 +1,142 @@
+from __future__ import annotations
+
+import logging
+import sys
 from dataclasses import dataclass
+from typing import Any, Callable
+
+from pydantic import BaseModel
+from tenacity import WrappedFn, retry
+from tenacity.retry import retry_always, retry_any, retry_base, retry_if_exception_type
+from tenacity.stop import stop_after_attempt, stop_after_delay, stop_any, stop_base
+from tenacity.wait import wait_base, wait_exponential
 
-from tenacity.retry import retry_base
-from tenacity.stop import stop_base
-from tenacity.wait import wait_base
+from graphdatascience.retry_utils.retry_utils import before_log
 
 
 @dataclass(frozen=True, repr=True)
 class RetryConfig:
     stop: stop_base
     wait: wait_base
     retry: retry_base
+
+
+class RetryConfigV2(BaseModel):
+    """Retry configuration which can be serialized/deserialized."""
+
+    retryable_exceptions: list[type[BaseException]] | None = None
+    stop_config: StopConfig | None = None
+    wait_config: ExponentialWaitConfig | None = None
+
+    def decorator(
+        self,
+        logger: logging.Logger | None = None,
+        log_level: int = logging.DEBUG,
+        operation_name: str = "Operation",
+    ) -> Callable[[WrappedFn], WrappedFn]:
+        """
+        Create a tenacity retry decorator configured with this retry config.
+
+        Parameters
+        ----------
+        logger
+            Logger to use for retry logging. If None, no logging is performed.
+        log_level
+            Logging level to use (e.g., logging.DEBUG, logging.INFO).
+        operation_name
+            Name of the operation being retried, used in log messages.
+
+        Returns
+        -------
+        Callable
+            A decorator that can be applied to functions to add retry behavior.
+
+        Examples
+        --------
+        >>> config = RetryConfigV2(
+        ...     retryable_exceptions=[ValueError, KeyError],
+        ...     stop_config=StopConfig(after_attempt=3),
+        ...     wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10)
+        ... )
+        >>> @config.create_retry_decorator()
+        ... def my_function():
+        ...     # function logic here
+        ...     pass
+        """
+        kwargs: dict[str, Any] = {}
+
+        # Add retry condition
+        retry_condition = self.tenacity_retry()
+        if retry_condition is not None:
+            kwargs["retry"] = retry_condition
+
+        # Add stop condition
+        stop_condition = self.tenacity_stop()
+        if stop_condition is not None:
+            kwargs["stop"] = stop_condition
+
+        # Add wait strategy
+        wait_strategy = self.tenacity_wait()
+        if wait_strategy is not None:
+            kwargs["wait"] = wait_strategy
+
+        # Add logging if logger is provided
+        if logger is not None:
+            kwargs["before"] = before_log(operation_name, logger, log_level)
+
+        return retry(**kwargs, reraise=True)
+
+    def tenacity_wait(self) -> wait_base | None:
+        if self.wait_config is None:
+            return None
+
+        return self.wait_config.to_tenacity()
+
+    def tenacity_stop(self) -> stop_base | None:
+        if self.stop_config is None:
+            return None
+
+        return self.stop_config.to_tenacity()
+
+    def tenacity_retry(self) -> retry_base | None:
+        if self.retryable_exceptions is None:
+            return retry_always
+
+        retries = [retry_if_exception_type(exc) for exc in self.retryable_exceptions]
+        if len(retries) == 1:
+            return retries[0]
+
+        return retry_any(*retries)
+
+
+class StopConfig(BaseModel):
+    after_attempt: int | None = None
+    after_delay: int | None = None
+
+    def to_tenacity(self) -> stop_base | None:
+        stops: list[stop_base] = []
+        if self.after_attempt is not None:
+            stops.append(stop_after_attempt(self.after_attempt))
+        if self.after_delay is not None:
+            stops.append(stop_after_delay(self.after_delay))
+
+        if not stops:
+            return None
+        if len(stops) == 1:
+            return stops[0]
+
+        return stop_any(*stops)
+
+
+class ExponentialWaitConfig(BaseModel):
+    multiplier: float = 1
+    min: int = 0
+    max: float | int = sys.maxsize / 2
+    exp_base: int = 2
+
+    def to_tenacity(self) -> wait_base | None:
+        return wait_exponential(
+            multiplier=self.multiplier,
+            min=self.min,
+            max=self.max,
+        )

graphdatascience/tests/unit/arrow_client/test_authenticated_flight_client.py

@@ -1,12 +1,11 @@
 # graphdatascience/tests/test_authenticated_flight_client.py
 import pytest
 from pyarrow._flight import FlightInternalError, FlightTimedOutError, FlightUnavailableError
-from tenacity import retry_any, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
-from graphdatascience.retry_utils.retry_config import RetryConfig
+from graphdatascience.retry_utils.retry_config import ExponentialWaitConfig, RetryConfigV2, StopConfig
 
 
 @pytest.fixture
@@ -15,15 +14,15 @@ def arrow_info() -> ArrowInfo:
 
 
 @pytest.fixture
-def retry_config() -> RetryConfig:
-    return RetryConfig(
-        retry=retry_any(
-            retry_if_exception_type(FlightTimedOutError),
-            retry_if_exception_type(FlightUnavailableError),
-            retry_if_exception_type(FlightInternalError),
-        ),
-        stop=(stop_after_delay(10) | stop_after_attempt(5)),
-        wait=wait_exponential(multiplier=1, min=1, max=10),
+def retry_config() -> RetryConfigV2:
+    return RetryConfigV2(
+        retryable_exceptions=[
+            FlightTimedOutError,
+            FlightUnavailableError,
+            FlightInternalError,
+        ],
+        stop_config=StopConfig(after_delay=10, after_attempt=5),
+        wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
     )
 
 
@@ -37,7 +36,7 @@ def auth_pair(self) -> tuple[str, str]:
 
 
 def test_create_authenticated_arrow_client(
-    arrow_info: ArrowInfo, retry_config: RetryConfig, mock_auth: ArrowAuthentication
+    arrow_info: ArrowInfo, retry_config: RetryConfigV2, mock_auth: ArrowAuthentication
 ) -> None:
     client = AuthenticatedArrowClient.create(
         arrow_info=arrow_info, auth=mock_auth, encrypted=True, retry_config=retry_config
@@ -46,7 +45,18 @@ def test_create_authenticated_arrow_client(
     assert client.connection_info() == ConnectionInfo("localhost", 8491, encrypted=True)
 
 
-def test_connection_info(arrow_info: ArrowInfo, retry_config: RetryConfig) -> None:
+def test_connection_info(arrow_info: ArrowInfo, retry_config: RetryConfigV2) -> None:
     client = AuthenticatedArrowClient(host="localhost", port=8491, retry_config=retry_config)
     connection_info = client.connection_info()
     assert connection_info == ConnectionInfo("localhost", 8491, encrypted=False)
+
+
+def test_pickle_roundtrip(arrow_info: ArrowInfo, retry_config: RetryConfigV2) -> None:
+    client = AuthenticatedArrowClient(host="localhost", port=8491, retry_config=retry_config)
+    import pickle
+
+    pickled_client = pickle.dumps(client)
+    unpickled_client = pickle.loads(pickled_client)
+    assert isinstance(unpickled_client, AuthenticatedArrowClient)
+    assert unpickled_client.connection_info() == client.connection_info()
+    assert unpickled_client._retry_config == client._retry_config