Add MutationClient

Author: DarthMax

graphdatascience/arrow_client/v2/write_back_client.py

@@ -0,0 +1,54 @@
+import time
+from typing import Any, Optional
+
+from graphdatascience import QueryRunner
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.query_runner.protocol.write_protocols import WriteProtocol
+from graphdatascience.query_runner.termination_flag import TerminationFlagNoop
+from graphdatascience.session.dbms.protocol_resolver import ProtocolVersionResolver
+
+
+class WriteBackClient:
+    def __init__(self, arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner):
+        self._arrow_client = arrow_client
+        self._query_runner = query_runner
+
+        protocol_version = ProtocolVersionResolver(query_runner).resolve()
+        self._write_protocol = WriteProtocol.select(protocol_version)
+
+    # TODO: Add progress logging
+    # TODO: Support setting custom writeProperties and relationshipTypes
+    def write(self, graph_name: str, job_id: str, concurrency: Optional[int]) -> int:
+        arrow_config = self._arrow_configuration()
+
+        configuration = {}
+        if concurrency is not None:
+            configuration["concurrency"] = concurrency
+
+        write_back_params = CallParameters(
+            graphName=graph_name,
+            jobId=job_id,
+            arrowConfiguration=arrow_config,
+            configuration=configuration,
+        )
+
+        start_time = time.time()
+
+        self._write_protocol.run_write_back(self._query_runner, write_back_params, None, TerminationFlagNoop())
+
+        return int((time.time() - start_time) * 1000)
+
+    def _arrow_configuration(self) -> dict[str, Any]:
+        connection_info = self._arrow_client.connection_info()
+        token = self._arrow_client.request_token()
+        if token is None:
+            token = "IGNORED"
+        arrow_config = {
+            "host": connection_info.host,
+            "port": connection_info.port,
+            "token": token,
+            "encrypted": connection_info.encrypted,
+        }
+
+        return arrow_config

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

@@ -0,0 +1,27 @@
+import unittest
+from unittest.mock import MagicMock
+
+from graphdatascience.arrow_client.v2.api_types import MutateResult
+from graphdatascience.arrow_client.v2.mutation_client import MutationClient
+from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
+
+
+class TestMutationClient(unittest.TestCase):
+    def setUp(self) -> None:
+        self.mock_client = MagicMock()
+
+    def test_mutate_node_property_success(self) -> None:
+        job_id = "test-job-123"
+        expected_mutation_result = MutateResult(nodePropertiesWritten=42, relationshipsWritten=1337)
+
+        self.mock_client.do_action_with_retry.return_value = iter(
+            [ArrowTestResult(expected_mutation_result.dump_camel())]
+        )
+
+        result = MutationClient.mutate_node_property(self.mock_client, job_id, "propertyName")
+
+        assert result == expected_mutation_result
+
+        self.mock_client.do_action_with_retry.assert_called_once_with(
+            MutationClient.MUTATE_ENDPOINT, b'{"jobId": "test-job-123", "mutateProperty": "propertyName"}'
+        )