Implement logProgress

Author: DarthMax

graphdatascience/procedure_surface/api/system_endpoints.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
+
+
+class SystemEndpoints(ABC):
+    @abstractmethod
+    def list_progress(
+        self,
+        job_id: str | None = None,
+        show_completed: bool = False,
+    ) -> list[ProgressResult]:
+        pass
+
+
+class ProgressResult(BaseResult):
+    username: str
+    job_id: str
+    task_name: str
+    progress: str
+    progress_bar: str
+    status: str
+    time_started: str
+    elapsed_time: str

graphdatascience/procedure_surface/arrow/system_arrow_endpoints.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize
+from graphdatascience.procedure_surface.api.system_endpoints import ProgressResult, SystemEndpoints
+from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
+
+
+class SystemArrowEndpoints(SystemEndpoints):
+    def __init__(self, arrow_client: AuthenticatedArrowClient):
+        self._arrow_client = arrow_client
+
+    def list_progress(
+        self,
+        job_id: str | None = None,
+        show_completed: bool = False,
+    ) -> list[ProgressResult]:
+        config = ConfigConverter.convert_to_gds_config(
+            job_id=job_id,
+            show_completed=show_completed,
+        )
+
+        rows = deserialize(self._arrow_client.do_action_with_retry("v2/listProgress", config))
+        return [ProgressResult(**row) for row in rows]

graphdatascience/procedure_surface/cypher/system_cypher_endpoints.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.procedure_surface.api.system_endpoints import ProgressResult, SystemEndpoints
+from graphdatascience.query_runner.query_runner import QueryRunner
+
+
+class SystemCypherEndpoints(SystemEndpoints):
+    def __init__(self, query_runner: QueryRunner):
+        self._query_runner = query_runner
+
+    def list_progress(
+        self,
+        job_id: str | None = None,
+        show_completed: bool = False,
+    ) -> list[ProgressResult]:
+        params = CallParameters(
+            job_id=job_id if job_id else "",
+            show_completed=True,
+        )
+
+        result = self._query_runner.call_procedure(endpoint="gds.listProgress", params=params)
+        return [ProgressResult(**row.to_dict()) for _, row in result.iterrows()]

graphdatascience/session/session_v2_endpoints.py

@@ -64,6 +64,7 @@
 )
 from graphdatascience.procedure_surface.arrow.node_embedding.hashgnn_arrow_endpoints import HashGNNArrowEndpoints
 from graphdatascience.procedure_surface.arrow.node_embedding.node2vec_arrow_endpoints import Node2VecArrowEndpoints
+from graphdatascience.procedure_surface.arrow.system_arrow_endpoints import SystemArrowEndpoints
 from graphdatascience.query_runner.query_runner import QueryRunner
 
 
@@ -91,6 +92,10 @@ def graph(self) -> CatalogArrowEndpoints:
     def config(self) -> ConfigArrowEndpoints:
         return ConfigArrowEndpoints(self._arrow_client)
 
+    @property
+    def system(self) -> SystemArrowEndpoints:
+        return SystemArrowEndpoints(self._arrow_client)
+
     ## Algorithms
 
     @property

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_system_arrow_endpoints.py

@@ -0,0 +1,50 @@
+from typing import Generator
+
+import pytest
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
+from graphdatascience.procedure_surface.arrow.community.wcc_arrow_endpoints import WccArrowEndpoints
+from graphdatascience.procedure_surface.arrow.system_arrow_endpoints import SystemArrowEndpoints
+from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import create_graph
+
+graph = """
+        CREATE
+            (a1)-[:T]->(a2)
+        """
+
+
+@pytest.fixture(scope="class")
+def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
+    with create_graph(arrow_client, "g", graph) as G:
+        yield G
+
+
+@pytest.fixture(scope="class")
+def job_id(arrow_client: AuthenticatedArrowClient, sample_graph: GraphV2) -> Generator[str, None, None]:
+    job_id = "test_job_id"
+    WccArrowEndpoints(arrow_client).stats(sample_graph, job_id=job_id)
+    yield job_id
+
+
+@pytest.fixture
+def system_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[SystemArrowEndpoints, None, None]:
+    yield SystemArrowEndpoints(arrow_client)
+
+
+def test_list_progress_job_id(system_endpoints: SystemArrowEndpoints, job_id: str) -> None:
+    results = system_endpoints.list_progress(job_id=job_id, show_completed=True)
+
+    assert len(results) == 1
+
+    result = results[0]
+    assert result.username == "neo4j"
+    assert result.job_id == job_id
+    assert "WCC" in result.task_name
+    assert result.progress == "100%"
+    assert "#" in result.progress_bar
+
+
+def test_list_nothing(system_endpoints: SystemArrowEndpoints) -> None:
+    results = system_endpoints.list_progress()
+    assert len(results) == 0

graphdatascience/tests/integrationV2/procedure_surface/cypher/test_system_cypher_endpoints.py

@@ -0,0 +1,58 @@
+from typing import Generator
+
+import pytest
+
+from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
+from graphdatascience.procedure_surface.cypher.community.wcc_cypher_endpoints import WccCypherEndpoints
+from graphdatascience.procedure_surface.cypher.system_cypher_endpoints import SystemCypherEndpoints
+from graphdatascience.query_runner.query_runner import QueryRunner
+from graphdatascience.tests.integrationV2.procedure_surface.cypher.cypher_graph_helper import create_graph
+
+
+@pytest.fixture
+def system_endpoints(query_runner: QueryRunner) -> Generator[SystemCypherEndpoints, None, None]:
+    yield SystemCypherEndpoints(query_runner)
+
+
+graph = """
+        CREATE
+            (a1)-[:T]->(a2)
+        """
+
+projection_query = """
+    MATCH (n)-[r]->(m)
+    WITH gds.graph.project('g', n, m, {}) AS G
+    RETURN G
+"""
+
+
+@pytest.fixture(scope="class")
+def sample_graph(query_runner: QueryRunner) -> Generator[GraphV2, None, None]:
+    with create_graph(query_runner, "g", graph, projection_query) as G:
+        yield G
+
+
+@pytest.fixture(scope="class")
+def job_id(query_runner: QueryRunner, sample_graph: GraphV2) -> Generator[str, None, None]:
+    job_id = "test_job_id"
+    WccCypherEndpoints(query_runner).mutate(sample_graph, job_id=job_id, log_progress=True, mutate_property="wcc")
+    yield job_id
+
+
+@pytest.mark.skip(reason="Enable when we figure out how to retain jobs")
+def test_list_progress_job_id(system_endpoints: SystemCypherEndpoints, job_id: str) -> None:
+    results = system_endpoints.list_progress(show_completed=True)
+
+    assert len(results) == 1
+
+    result = results[0]
+    assert result.username == "neo4j"
+    assert result.job_id == job_id
+    assert "WCC" in result.task_name
+    assert result.progress == "100%"
+    assert "#" in result.progress_bar
+
+
+def test_list_nothing(system_endpoints: SystemCypherEndpoints) -> None:
+    results = system_endpoints.list_progress()
+    assert len(results) == 0