Log on retries

Author: FlorentinD

graphdatascience/query_runner/gds_arrow_client.py

@@ -2,6 +2,7 @@
 
 import base64
 import json
+import logging
 import re
 import time
 import warnings
@@ -34,11 +35,15 @@
     wait_fixed,
 )
 
+from graphdatascience.retry_utils.retry_utils import before_log
+
 from ..semantic_version.semantic_version import SemanticVersion
 from ..version import __version__
 from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_info import ArrowInfo
 
+_arrow_client_logger = logging.getLogger("gds_arrow_client")
+
 
 class GdsArrowClient:
     @staticmethod
@@ -148,6 +153,7 @@ def connection_info(self) -> tuple[str, int]:
         return self._host, self._port
 
     @retry(
+        before=before_log("Request token", _arrow_client_logger, logging.DEBUG),
         retry=retry_any(retry_if_exception_type(FlightTimedOutError), retry_if_exception_type(FlightUnavailableError)),
         stop=stop_after_attempt(3),
         wait=wait_fixed(1),
@@ -593,6 +599,7 @@ def _client(self) -> flight.FlightClient:
         return self._flight_client
 
     @retry(
+        before=before_log("Send action", _arrow_client_logger, logging.DEBUG),
         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),
@@ -614,6 +621,7 @@ def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str,
             raise e  # unreachable
 
     @retry(
+        before=before_log("Do put", _arrow_client_logger, logging.DEBUG),
         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),
@@ -644,6 +652,7 @@ def _upload_data(
         put_stream, ack_stream = self._safe_do_put(upload_descriptor, batches[0].schema)
 
         @retry(
+            before=before_log("Upload batch", _arrow_client_logger, logging.DEBUG),
             stop=(stop_after_delay(10) | stop_after_attempt(5)),
             wait=wait_exponential(multiplier=1, min=1, max=10),
             retry=(
@@ -665,6 +674,7 @@ def upload_batch(p: RecordBatch) -> None:
             GdsArrowClient.handle_flight_error(e)
 
     @retry(
+        before=before_log("Do get", _arrow_client_logger, logging.DEBUG),
         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),

graphdatascience/query_runner/protocol/project_protocols.py

@@ -7,9 +7,9 @@
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 

graphdatascience/query_runner/protocol/write_protocols.py

@@ -7,9 +7,9 @@
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 

graphdatascience/retry_utils/__init__.py

@@ -11,9 +11,10 @@ def before_log(
     sec_format: str = "%0.3f",
 ) -> typing.Callable[[RetryCallState], None]:
     def log_it(retry_state: RetryCallState) -> None:
-        logger.log(
-            log_level,
-            f"Retry of '{fn_name}', " f"attempt: {retry_state.attempt_number}",
-        )
+        if retry_state.attempt_number > 1:  # only log on actual retry
+            logger.log(
+                log_level,
+                f"Retry of '{fn_name}', " f"attempt: {retry_state.attempt_number}",
+            )
 
     return log_it