Check for signals during projection and write-back

Author: FlorentinD

graphdatascience/query_runner/protocol/project_protocols.py

@@ -1,3 +1,4 @@
+import signal
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
@@ -6,6 +7,7 @@
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
+from graphdatascience.query_runner.protocol.retry_utils import retry_unless_signal
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
@@ -123,7 +125,10 @@ def is_not_done(result: DataFrame) -> bool:
             status: str = result.squeeze()["status"]
             return status != Status.DONE.name
 
-        @retry(retry=retry_if_result(is_not_done), wait=wait_incrementing(start=0.2, increment=0.2, max=2))
+        @retry(
+            retry=retry_if_result(is_not_done) and retry_unless_signal([signal.SIGTERM, signal.SIGINT]),
+            wait=wait_incrementing(start=0.2, increment=0.2, max=2),
+        )
         def project_fn() -> DataFrame:
             return query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False

graphdatascience/query_runner/protocol/retry_utils.py

@@ -0,0 +1,21 @@
+import signal
+from types import FrameType
+from typing import Optional
+
+from tenacity import RetryCallState, retry_base
+
+
+class retry_unless_signal(retry_base):
+    """Retries unless one of the given signals is raised."""
+
+    def __init__(self, signals: list[signal.Signals]) -> None:
+        self.signal_received = False
+
+        def receive_signal(sig: int, frame: Optional[FrameType]) -> None:
+            self.signal_received = True
+
+        for sig in signals:
+            signal.signal(sig, receive_signal)
+
+    def __call__(self, retry_state: RetryCallState) -> bool:
+        return not self.signal_received

graphdatascience/query_runner/protocol/write_protocols.py

@@ -1,11 +1,14 @@
+import logging
+import signal
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
 from pandas import DataFrame
-from tenacity import retry, retry_if_result, wait_incrementing
+from tenacity import after_log, retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
+from graphdatascience.query_runner.protocol.retry_utils import retry_unless_signal
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
@@ -120,7 +123,13 @@ def is_not_completed(result: DataFrame) -> bool:
             status: str = result.squeeze()["status"]
             return status != Status.COMPLETED.name
 
-        @retry(retry=retry_if_result(is_not_completed), wait=wait_incrementing(start=0.2, increment=0.2, max=2))
+        logger = logging.getLogger()
+
+        @retry(
+            retry=retry_if_result(is_not_completed) and retry_unless_signal([signal.SIGTERM, signal.SIGINT]),
+            wait=wait_incrementing(start=0.2, increment=0.2, max=2),
+            after=after_log(logger, logging.WARN),
+        )
         def write_fn() -> DataFrame:
             return query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name("gds.arrow.write"),