Introduce JobClient

Author: DarthMax

graphdatascience/arrow_client/arrow_base_model.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class ArrowBaseModel(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
+
+    def dump_camel(self) -> dict[str, Any]:
+        return self.model_dump(by_alias=True)
+
+    def dump_json(self) -> str:
+        return self.model_dump_json(by_alias=True)

graphdatascience/arrow_client/v2/api_types.py

@@ -0,0 +1,16 @@
+from graphdatascience.arrow_client.arrow_base_model import ArrowBaseModel
+
+
+class JobIdConfig(ArrowBaseModel):
+    job_id: str
+
+
+class JobStatus(ArrowBaseModel):
+    job_id: str
+    status: str
+    progress: float
+
+
+class MutateResult(ArrowBaseModel):
+    node_properties_written: int
+    relationships_written: int

graphdatascience/arrow_client/v2/data_mapper_utils.py

@@ -0,0 +1,19 @@
+import json
+from typing import Any, Iterator
+
+from pyarrow._flight import Result
+
+
+def deserialize_single(input_stream: Iterator[Result]) -> dict[str, Any]:
+    rows = deserialize(input_stream)
+    if len(rows) != 1:
+        raise ValueError(f"Expected exactly one result, got {len(rows)}")
+
+    return rows[0]
+
+
+def deserialize(input_stream: Iterator[Result]) -> list[dict[str, Any]]:
+    def deserialize_row(row: Result):  # type:ignore
+        return json.loads(row.body.to_pybytes().decode())
+
+    return [deserialize_row(row) for row in list(input_stream)]

graphdatascience/arrow_client/v2/job_client.py

@@ -0,0 +1,63 @@
+import json
+from typing import Any
+
+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
+
+JOB_STATUS_ENDPOINT = "v2/jobs.status"
+RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
+
+
+class JobClient:
+    @staticmethod
+    def run_job_and_wait(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any]) -> str:
+        job_id = JobClient.run_job(client, endpoint, config)
+        JobClient.wait_for_job(client, job_id)
+        return job_id
+
+    @staticmethod
+    def run_job(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any]) -> str:
+        encoded_config = json.dumps(config).encode("utf-8")
+        res = client.do_action_with_retry(endpoint, encoded_config)
+
+        single = deserialize_single(res)
+        return JobIdConfig(**single).job_id
+
+    @staticmethod
+    def wait_for_job(client: AuthenticatedArrowClient, job_id: str) -> None:
+        while True:
+            encoded_config = JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+
+            arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, encoded_config)
+            job_status = JobStatus(**deserialize_single(arrow_res))
+            if job_status.status == "Done":
+                break
+
+    @staticmethod
+    def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:
+        encoded_config = JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+
+        res = client.do_action_with_retry(RESULTS_SUMMARY_ENDPOINT, encoded_config)
+        return deserialize_single(res)
+
+    @staticmethod
+    def stream_results(client: AuthenticatedArrowClient, job_id: str) -> DataFrame:
+        encoded_config = JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+
+        res = client.do_action_with_retry("v2/results.stream", encoded_config)
+        export_job_id = JobIdConfig(**deserialize_single(res)).job_id
+
+        payload = {
+            "name": export_job_id,
+            "version": 1,
+        }
+
+        ticket = Ticket(json.dumps(payload).encode("utf-8"))
+        with client.get_stream(ticket) as get:
+            arrow_table = get.read_all()
+
+        return arrow_table.to_pandas(types_mapper=ArrowDtype)  # type: ignore

graphdatascience/procedure_surface/__init__.py

@@ -0,0 +1,26 @@
+from typing import Iterator
+
+import pytest
+from pyarrow._flight import Result
+
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
+from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
+
+
+def test_deserialize_single_success() -> None:
+    input_stream = iter([ArrowTestResult({"key": "value"})])
+    expected = {"key": "value"}
+    actual = deserialize_single(input_stream)
+    assert expected == actual
+
+
+def test_deserialize_single_raises_on_empty_stream() -> None:
+    input_stream: Iterator[Result] = iter([])
+    with pytest.raises(ValueError, match="Expected exactly one result, got 0"):
+        deserialize_single(input_stream)
+
+
+def test_deserialize_single_raises_on_multiple_results() -> None:
+    input_stream = iter([ArrowTestResult({"key1": "value1"}), ArrowTestResult({"key2": "value2"})])
+    with pytest.raises(ValueError, match="Expected exactly one result, got 2"):
+        deserialize_single(input_stream)

graphdatascience/procedure_surface/arrow/__init__.py

@@ -0,0 +1,112 @@
+import json
+import unittest
+from unittest.mock import MagicMock
+
+from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
+from graphdatascience.arrow_client.v2.job_client import JobClient
+from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
+
+
+class TestJobClient(unittest.TestCase):
+    def setUp(self) -> None:
+        self.mock_client = MagicMock()
+
+    def test_run_job(self) -> None:
+        # Setup
+        job_id = "test-job-123"
+        endpoint = "v2/test.endpoint"
+        config = {"param1": "value1", "param2": 42}
+
+        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult({"jobId": job_id})])
+
+        # Execute
+        result = JobClient.run_job(self.mock_client, endpoint, config)
+
+        # Verify
+        expected_config = json.dumps(config).encode("utf-8")
+        self.mock_client.do_action_with_retry.assert_called_once_with(endpoint, expected_config)
+        self.assertEqual(result, job_id)
+
+    def test_run_job_and_wait(
+        self,
+    ) -> None:
+        job_id = "test-job-456"
+        endpoint = "v2/test.endpoint"
+        config = {"param": "value"}
+
+        job_id_config = JobIdConfig(jobId=job_id)
+
+        status = JobStatus(
+            jobId=job_id,
+            progress=1.0,
+            status="Done",
+        )
+
+        do_action_with_retry = MagicMock()
+        do_action_with_retry.side_effect = [
+            iter([ArrowTestResult(job_id_config.dump_camel())]),
+            iter([ArrowTestResult(status.dump_camel())]),
+        ]
+
+        self.mock_client.do_action_with_retry = do_action_with_retry
+
+        result = JobClient.run_job_and_wait(self.mock_client, endpoint, config)
+
+        do_action_with_retry.assert_called_with("v2/jobs.status", job_id_config.dump_json().encode("utf-8"))
+        self.assertEqual(result, job_id)
+
+    def test_wait_for_job_completes_immediately(self) -> None:
+        job_id = "test-job-789"
+
+        status = JobStatus(
+            jobId=job_id,
+            progress=1.0,
+            status="Done",
+        )
+
+        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(status.dump_camel())])
+
+        JobClient.wait_for_job(self.mock_client, job_id)
+
+        self.mock_client.do_action_with_retry.assert_called_once_with(
+            "v2/jobs.status", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+        )
+
+    def test_wait_for_job_waits_for_completion(self) -> None:
+        job_id = "test-job-waiting"
+        status_running = JobStatus(
+            jobId=job_id,
+            progress=0.5,
+            status="RUNNING",
+        )
+        status_done = JobStatus(
+            jobId=job_id,
+            progress=1.0,
+            status="Done",
+        )
+
+        do_action_with_retry = MagicMock()
+        do_action_with_retry.side_effect = [
+            iter([ArrowTestResult(status_running.dump_camel())]),
+            iter([ArrowTestResult(status_done.dump_camel())]),
+        ]
+
+        self.mock_client.do_action_with_retry = do_action_with_retry
+
+        JobClient.wait_for_job(self.mock_client, job_id)
+
+        self.assertEqual(self.mock_client.do_action_with_retry.call_count, 2)
+
+    def test_get_summary(self) -> None:
+        # Setup
+        job_id = "summary-job-123"
+        expected_summary = {"nodeCount": 100, "relationshipCount": 200, "requiredMemory": "1GB"}
+
+        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(expected_summary)])
+
+        result = JobClient.get_summary(self.mock_client, job_id)
+
+        self.mock_client.do_action_with_retry.assert_called_once_with(
+            "v2/results.summary", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+        )
+        self.assertEqual(result, expected_summary)

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

@@ -0,0 +1,20 @@
+import json
+from typing import Any
+
+from pyarrow._flight import Result
+
+
+class ArrowTestResult(Result):  # type:ignore
+    def __init__(self, body: dict[str, Any]):
+        self._body = json.dumps(body).encode()
+
+    @property
+    def body(self) -> Any:
+        class MockBody:
+            def __init__(self, data: bytes):
+                self._data = data
+
+            def to_pybytes(self) -> bytes:
+                return self._data
+
+        return MockBody(self._body)