Add progress bar to job client

FlorentinD · FlorentinD · commit 6152a2482282 · 2025-09-30T09:34:11.000+02:00
diff --git a/graphdatascience/arrow_client/v2/api_types.py b/graphdatascience/arrow_client/v2/api_types.py
@@ -5,10 +5,42 @@ class JobIdConfig(ArrowBaseModel):
     job_id: str
 
 
+UNKNOWN_PROGRESS = -1
+
+
 class JobStatus(ArrowBaseModel):
     job_id: str
     status: str
     progress: float
+    description: str
+
+    def progress_known(self) -> bool:
+        if self.progress == UNKNOWN_PROGRESS:
+            return False
+        return True
+
+    def progress_percent(self) -> float | None:
+        if self.progress_known():
+            return self.progress * 100
+        return None
+
+    def base_task(self) -> str:
+        return self.description.split("::")[0].strip()
+
+    def sub_tasks(self) -> str | None:
+        task_split = self.description.split("::", maxsplit=1)
+        if len(task_split) > 1:
+            return task_split[1].strip()
+        return None
+
+    def aborted(self) -> bool:
+        return self.status == "Aborted"
+
+    def succeeded(self) -> bool:
+        return self.status == "Done"
+
+    def running(self) -> bool:
+        return self.status == "Running"
 
 
 class MutateResult(ArrowBaseModel):
diff --git a/graphdatascience/arrow_client/v2/job_client.py b/graphdatascience/arrow_client/v2/job_client.py
@@ -1,22 +1,28 @@
 import json
-from typing import Any
+from typing import Any, Optional
 
 from pandas import ArrowDtype, DataFrame
 from pyarrow._flight import Ticket
 
 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
 
 JOB_STATUS_ENDPOINT = "v2/jobs.status"
 RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
 
 
 class JobClient:
+    def __init__(self, progress_bar_options: dict[str, Any] | None = None):
+        self._progress_bar_options = progress_bar_options or {}
+
     @staticmethod
-    def run_job_and_wait(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any]) -> str:
+    def run_job_and_wait(
+        client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any], show_progress: bool
+    ) -> str:
         job_id = JobClient.run_job(client, endpoint, config)
-        JobClient.wait_for_job(client, job_id)
+        JobClient().wait_for_job(client, job_id, show_progress=show_progress)
         return job_id
 
     @staticmethod
@@ -26,14 +32,30 @@ def run_job(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, A
         single = deserialize_single(res)
         return JobIdConfig(**single).job_id
 
-    @staticmethod
-    def wait_for_job(client: AuthenticatedArrowClient, job_id: str) -> None:
+    def wait_for_job(self, client: AuthenticatedArrowClient, job_id: str, show_progress: bool) -> 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.status == "Done":
+
+            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())
+
     @staticmethod
     def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:
         res = client.do_action_with_retry(RESULTS_SUMMARY_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
diff --git a/graphdatascience/procedure_surface/arrow/catalog/node_label_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/catalog/node_label_arrow_endpoints.py
@@ -19,6 +19,7 @@ def __init__(
     ):
         self._arrow_client = arrow_client
         self._node_property_endpoints = NodePropertyEndpoints(arrow_client, write_back_client)
+        self._show_progress = False  # TODO add option to show progress
 
     def mutate(
         self,
@@ -45,7 +46,10 @@ def mutate(
             job_id=job_id,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.nodeLabel.mutate", config)
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.nodeLabel.mutate", config, show_progress=show_progress
+        )
         return NodeLabelMutateResult(**JobClient.get_summary(self._arrow_client, job_id))
 
     def write(
diff --git a/graphdatascience/procedure_surface/arrow/catalog/relationship_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/catalog/relationship_arrow_endpoints.py
@@ -19,9 +19,15 @@
 
 
 class RelationshipArrowEndpoints(RelationshipsEndpoints):
-    def __init__(self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[RemoteWriteBackClient]):
+    def __init__(
+        self,
+        arrow_client: AuthenticatedArrowClient,
+        write_back_client: Optional[RemoteWriteBackClient],
+        show_progress: bool = False,
+    ):
         self._arrow_client = arrow_client
         self._write_back_client = write_back_client
+        self._show_progress = show_progress
 
     def stream(
         self,
@@ -154,7 +160,10 @@ def index_inverse(
             job_id=job_id,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.relationships.indexInverse", config)
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.relationships.indexInverse", config, show_progress=show_progress
+        )
         result = JobClient.get_summary(self._arrow_client, job_id)
         return RelationshipsInverseIndexResult(**result)
 
@@ -182,7 +191,10 @@ def to_undirected(
             username=username,
             job_id=job_id,
         )
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.relationships.toUndirected", config)
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.relationships.toUndirected", config, show_progress=show_progress
+        )
         result = JobClient.get_summary(self._arrow_client, job_id)
         return RelationshipsToUndirectedResult(**result)
diff --git a/graphdatascience/procedure_surface/arrow/catalog_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/catalog_arrow_endpoints.py
@@ -37,10 +37,16 @@
 class CatalogArrowEndpoints(CatalogEndpoints):
     GDS_REMOTE_PROJECTION_PROC_NAME = "gds.arrow.project"
 
-    def __init__(self, arrow_client: AuthenticatedArrowClient, query_runner: Optional[QueryRunner] = None):
+    def __init__(
+        self,
+        arrow_client: AuthenticatedArrowClient,
+        query_runner: Optional[QueryRunner] = None,
+        show_progress: bool = False,
+    ):
         self._arrow_client = arrow_client
         self._query_runner = query_runner
         self._graph_backend = GraphOpsArrow(arrow_client)
+        self._show_progress = show_progress
         if query_runner is not None:
             protocol_version = ProtocolVersionResolver(query_runner).resolve()
             self._project_protocol = ProjectProtocol.select(protocol_version)
@@ -152,7 +158,9 @@ def filter(
             job_id=job_id,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.project.filter", config)
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.project.filter", config, show_progress=self._show_progress
+        )
 
         return GraphWithFilterResult(
             get_graph(graph_name, self._arrow_client),
@@ -192,7 +200,10 @@ def generate(
             username=username,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.generate", config)
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.generate", config, show_progress=show_progress
+        )
 
         return GraphWithGenerationStats(
             get_graph(graph_name, self._arrow_client),
@@ -201,7 +212,7 @@ def generate(
 
     @property
     def sample(self) -> GraphSamplingEndpoints:
-        return GraphSamplingArrowEndpoints(self._arrow_client)
+        return GraphSamplingArrowEndpoints(self._arrow_client, show_progress=self._show_progress)
 
     @property
     def node_labels(self) -> NodeLabelArrowEndpoints:
@@ -218,6 +229,7 @@ def relationships(self) -> RelationshipArrowEndpoints:
         return RelationshipArrowEndpoints(
             self._arrow_client,
             RemoteWriteBackClient(self._arrow_client, self._query_runner) if self._query_runner else None,
+            show_progress=self._show_progress,
         )
 
     def _arrow_config(self) -> dict[str, Any]:
diff --git a/graphdatascience/procedure_surface/arrow/graph_sampling_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/graph_sampling_arrow_endpoints.py
@@ -15,8 +15,9 @@
 
 
 class GraphSamplingArrowEndpoints(GraphSamplingEndpoints):
-    def __init__(self, arrow_client: AuthenticatedArrowClient):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, show_progress: bool = False):
         self._arrow_client = arrow_client
+        self._show_progress = show_progress
 
     def rwr(
         self,
@@ -52,7 +53,10 @@ def rwr(
             job_id=job_id,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.sample.rwr", config)
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.sample.rwr", config, show_progress=show_progress
+        )
 
         return GraphWithSamplingResult(
             get_graph(graph_name, self._arrow_client),
@@ -93,7 +97,10 @@ def cnarw(
             job_id=job_id,
         )
 
-        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.sample.cnarw", config)
+        show_progress = self._show_progress and log_progress if log_progress is not None else self._show_progress
+        job_id = JobClient.run_job_and_wait(
+            self._arrow_client, "v2/graph.sample.cnarw", config, show_progress=show_progress
+        )
 
         return GraphWithSamplingResult(
             get_graph(graph_name, self._arrow_client),
diff --git a/graphdatascience/procedure_surface/arrow/node_property_endpoints.py b/graphdatascience/procedure_surface/arrow/node_property_endpoints.py
@@ -20,21 +20,26 @@ class NodePropertyEndpoints:
     """
 
     def __init__(
-        self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[RemoteWriteBackClient] = None
+        self,
+        arrow_client: AuthenticatedArrowClient,
+        write_back_client: Optional[RemoteWriteBackClient] = None,
+        show_progress: bool = True,
     ):
         self._arrow_client = arrow_client
         self._write_back_client = write_back_client
+        self._show_progress = show_progress
 
     def run_job_and_get_summary(self, endpoint: str, G: GraphV2, config: Dict[str, Any]) -> Dict[str, Any]:
         """Run a job and return the computation summary."""
-        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
+        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config, self._show_progress)
         return JobClient.get_summary(self._arrow_client, job_id)
 
+    # TODO expose show_progress option to endpoints
     def run_job_and_mutate(
         self, endpoint: str, G: GraphV2, config: Dict[str, Any], mutate_property: str
     ) -> Dict[str, Any]:
         """Run a job, mutate node properties, and return summary with mutation result."""
-        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
+        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config, self._show_progress)
         mutate_result = MutationClient.mutate_node_property(self._arrow_client, job_id, mutate_property)
         computation_result = JobClient.get_summary(self._arrow_client, job_id)
 
@@ -51,11 +56,13 @@ def run_job_and_mutate(
 
         return computation_result
 
+    # TODO expose show_progress option to endpoints
     def run_job_and_stream(self, endpoint: str, G: GraphV2, config: Dict[str, Any]) -> DataFrame:
         """Run a job and return streamed results."""
-        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
+        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config, show_progress=self._show_progress)
         return JobClient.stream_results(self._arrow_client, G.name(), job_id)
 
+    # TODO expose show_progress option to endpoints
     def run_job_and_write(
         self,
         endpoint: str,
@@ -66,7 +73,7 @@ def run_job_and_write(
         property_overwrites: Optional[Union[str, dict[str, str]]] = None,
     ) -> Dict[str, Any]:
         """Run a job, write results, and return summary with write time."""
-        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
+        job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config, show_progress=self._show_progress)
         computation_result = JobClient.get_summary(self._arrow_client, job_id)
 
         if self._write_back_client is None:
diff --git a/graphdatascience/procedure_surface/arrow/pagerank_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/pagerank_arrow_endpoints.py
@@ -18,9 +18,14 @@
 
 class PageRankArrowEndpoints(PageRankEndpoints):
     def __init__(
-        self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[RemoteWriteBackClient] = None
+        self,
+        arrow_client: AuthenticatedArrowClient,
+        write_back_client: Optional[RemoteWriteBackClient] = None,
+        show_progress: bool = False,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(arrow_client, write_back_client)
+        self._node_property_endpoints = NodePropertyEndpoints(
+            arrow_client, write_back_client, show_progress=show_progress
+        )
 
     def mutate(
         self,
diff --git a/graphdatascience/query_runner/progress/progress_bar.py b/graphdatascience/query_runner/progress/progress_bar.py
@@ -6,6 +6,7 @@
 
 
 class TqdmProgressBar:
+    # TODO helper method for creating for a test with obserable progress
     def __init__(self, task_name: str, relative_progress: Optional[float], bar_options: dict[str, Any] = {}):
         root_task_name = task_name
         if relative_progress is None:  # Qualitative progress report
@@ -21,6 +22,7 @@ def __init__(self, task_name: str, relative_progress: Optional[float], bar_optio
                 total=100,
                 unit="%",
                 desc=root_task_name,
+                initial=relative_progress,
                 **bar_options,
             )
 
diff --git a/graphdatascience/session/aura_graph_data_science.py b/graphdatascience/session/aura_graph_data_science.py
@@ -96,15 +96,17 @@ def create(
                 query_runner=session_query_runner,
                 delete_fn=delete_fn,
                 gds_version=gds_version,
-                v2_endpoints=SessionV2Endpoints(session_auth_arrow_client, db_bolt_query_runner),
+                v2_endpoints=SessionV2Endpoints(
+                    session_auth_arrow_client, db_bolt_query_runner, show_progress=show_progress
+                ),
             )
         else:
             standalone_query_runner = StandaloneSessionQueryRunner(session_arrow_query_runner)
             return cls(
                 query_runner=standalone_query_runner,
                 delete_fn=delete_fn,
                 gds_version=gds_version,
-                v2_endpoints=SessionV2Endpoints(session_auth_arrow_client, None),
+                v2_endpoints=SessionV2Endpoints(session_auth_arrow_client, None, show_progress=show_progress),
             )
 
     def __init__(
diff --git a/graphdatascience/session/session_v2_endpoints.py b/graphdatascience/session/session_v2_endpoints.py
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/arrow/graph_creation_helper.py b/graphdatascience/tests/integrationV2/procedure_surface/arrow/graph_creation_helper.py
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/arrow/test_pagerank_arrow_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/arrow/test_pagerank_arrow_endpoints.py
diff --git a/graphdatascience/tests/unit/arrow_client/V2/test_api_types.py b/graphdatascience/tests/unit/arrow_client/V2/test_api_types.py
diff --git a/graphdatascience/tests/unit/arrow_client/V2/test_job_client.py b/graphdatascience/tests/unit/arrow_client/V2/test_job_client.py
