Allow interruption during wait for job

Author: FlorentinD

graphdatascience/arrow_client/v2/job_client.py

@@ -3,12 +3,13 @@
 
 from pandas import ArrowDtype, DataFrame
 from pyarrow._flight import Ticket
-from tenacity import Retrying, retry_always, wait_exponential
+from tenacity import Retrying, retry_if_result, wait_exponential
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
 from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
 from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 
 JOB_STATUS_ENDPOINT = "v2/jobs.status"
 RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
@@ -33,11 +34,22 @@ def run_job(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, A
         single = deserialize_single(res)
         return JobIdConfig(**single).job_id
 
-    def wait_for_job(self, client: AuthenticatedArrowClient, job_id: str, show_progress: bool) -> None:
+    def wait_for_job(
+        self,
+        client: AuthenticatedArrowClient,
+        job_id: str,
+        show_progress: bool,
+        termination_flag: Optional[TerminationFlag] = None,
+    ) -> None:
         progress_bar: Optional[TqdmProgressBar] = None
 
-        for attempt in Retrying(retry=retry_always, wait=wait_exponential(min=0.1, max=5)):
+        if termination_flag is None:
+            termination_flag = TerminationFlag.create()
+
+        for attempt in Retrying(retry=retry_if_result(lambda _: True), wait=wait_exponential(min=0.1, max=5)):
             with attempt:
+                termination_flag.assert_running()
+
                 arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
                 job_status = JobStatus(**deserialize_single(arrow_res))
 

graphdatascience/arrow_client/v2/remote_write_back_client.py

@@ -8,7 +8,7 @@
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.procedure_surface.api.base_result import BaseResult
 from graphdatascience.query_runner.protocol.write_protocols import WriteProtocol
-from graphdatascience.query_runner.termination_flag import TerminationFlagNoop
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.session.dbms.protocol_resolver import ProtocolVersionResolver
 
 
@@ -53,7 +53,7 @@ def write(
             write_back_params,
             None,
             log_progress=log_progress,
-            terminationFlag=TerminationFlagNoop(),
+            terminationFlag=TerminationFlag.create(),
         ).squeeze()
         write_millis = int((time.time() - start_time) * 1000)
 

graphdatascience/tests/unit/arrow_client/V2/test_job_client.py

@@ -1,9 +1,11 @@
 from io import StringIO
 
+import pytest
 from pytest_mock import MockerFixture
 
 from graphdatascience.arrow_client.v2.api_types import UNKNOWN_PROGRESS, JobIdConfig, JobStatus
 from graphdatascience.arrow_client.v2.job_client import JobClient
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
 
 
@@ -126,6 +128,22 @@ def test_wait_for_job_waits_for_aborted(mocker: MockerFixture) -> None:
     assert mock_client.do_action_with_retry.call_count == 2
 
 
+def test_wait_for_job_stops_on_interrupt(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    job_id = "test-job-waiting"
+    mock_client.do_action_with_retry = mocker.Mock()
+
+    termination_flag = TerminationFlag.create()
+    termination_flag.set()
+
+    with pytest.raises(
+        RuntimeError, match="Closing client connection. Note, the query will be continued on the server-side"
+    ):
+        JobClient().wait_for_job(mock_client, job_id, show_progress=False, termination_flag=termination_flag)
+
+    assert mock_client.do_action_with_retry.call_count == 0
+
+
 def test_wait_for_job_progress_bar_quantive(mocker: MockerFixture) -> None:
     mock_client = mocker.Mock()
     job_id = "test-job-progress"