Merge pull request #974 from FlorentinD/v2-terminationflag-GDSA-208

Allow interruption during wait for job

Author: Mats-SX

graphdatascience/arrow_client/v2/job_client.py

@@ -3,11 +3,13 @@
 
 from pandas import ArrowDtype, DataFrame
 from pyarrow._flight import Ticket
+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"
@@ -32,28 +34,41 @@ 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
-        while True:
-            arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
-            job_status = JobStatus(**deserialize_single(arrow_res))
-
-            if job_status.succeeded() or job_status.aborted():
-                if progress_bar:
-                    progress_bar.finish(success=job_status.succeeded())
-                break
-
-            if show_progress:
-                if progress_bar is None:
-                    base_task = job_status.base_task()
-                    if base_task:
-                        progress_bar = TqdmProgressBar(
-                            task_name=base_task,
-                            relative_progress=job_status.progress_percent(),
-                            bar_options=self._progress_bar_options,
-                        )
-                if progress_bar:
-                    progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
+
+        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))
+
+                if job_status.succeeded() or job_status.aborted():
+                    if progress_bar:
+                        progress_bar.finish(success=job_status.succeeded())
+                    return
+
+                if show_progress:
+                    if progress_bar is None:
+                        base_task = job_status.base_task()
+                        if base_task:
+                            progress_bar = TqdmProgressBar(
+                                task_name=base_task,
+                                relative_progress=job_status.progress_percent(),
+                                bar_options=self._progress_bar_options,
+                            )
+                    if progress_bar:
+                        progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
 
     @staticmethod
     def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:

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"
@@ -173,8 +191,8 @@ def test_wait_for_job_progress_bar_qualitative(mocker: MockerFixture) -> None:
 
         progress_output = pbarOutputStream.getvalue().split("\r")
         assert "Algo [elapsed: 00:00 ]" in progress_output
-        assert "Algo [elapsed: 00:00 , status: RUNNING, task: Halfway there]" in progress_output
-        assert any("Algo [elapsed: 00:00 , status: FINISHED]" in line for line in progress_output)
+        assert "Algo [elapsed: 00:01 , status: RUNNING, task: Halfway there]" in progress_output
+        assert any("Algo [elapsed: 00:03 , status: FINISHED]" in line for line in progress_output)
 
 
 def test_get_summary(mocker: MockerFixture) -> None: