Implement Arrow based wcc endpoints

Author: DarthMax

graphdatascience/arrow_client/authenticated_arrow_client.py

@@ -1,33 +1,41 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional, Union, Any
+from dataclasses import dataclass
+from typing import Any, Optional, Union
 
 from pyarrow import __version__ as arrow_version
 from pyarrow import flight
-from pyarrow._flight import FlightTimedOutError, FlightUnavailableError, FlightInternalError, Action
-from tenacity import retry_any, retry_if_exception_type, stop_after_delay, stop_after_attempt, wait_exponential, retry
+from pyarrow._flight import (
+    Action,
+    FlightInternalError,
+    FlightStreamReader,
+    FlightTimedOutError,
+    FlightUnavailableError,
+    Ticket,
+)
+from tenacity import retry, retry_any, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.retry_utils.retry_config import RetryConfig
-from .middleware.AuthMiddleware import AuthMiddleware, AuthFactory
-from .middleware.UserAgentMiddleware import UserAgentFactory
+
 from ..retry_utils.retry_utils import before_log
 from ..version import __version__
+from .middleware.AuthMiddleware import AuthFactory, AuthMiddleware
+from .middleware.UserAgentMiddleware import UserAgentFactory
 
 
 class AuthenticatedArrowClient:
-
     @staticmethod
     def create(
-            arrow_info: ArrowInfo,
-            auth: Optional[ArrowAuthentication] = None,
-            encrypted: bool = False,
-            disable_server_verification: bool = False,
-            tls_root_certs: Optional[bytes] = None,
-            connection_string_override: Optional[str] = None,
-            retry_config: Optional[RetryConfig] = None,
+        arrow_info: ArrowInfo,
+        auth: Optional[ArrowAuthentication] = None,
+        encrypted: bool = False,
+        disable_server_verification: bool = False,
+        tls_root_certs: Optional[bytes] = None,
+        connection_string_override: Optional[str] = None,
+        retry_config: Optional[RetryConfig] = None,
     ) -> AuthenticatedArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -59,15 +67,15 @@ def create(
         )
 
     def __init__(
-            self,
-            host: str,
-            retry_config: RetryConfig,
-            port: int = 8491,
-            auth: Optional[Union[ArrowAuthentication, tuple[str, str]]] = None,
-            encrypted: bool = False,
-            disable_server_verification: bool = False,
-            tls_root_certs: Optional[bytes] = None,
-            user_agent: Optional[str] = None,
+        self,
+        host: str,
+        retry_config: RetryConfig,
+        port: int = 8491,
+        auth: Optional[Union[ArrowAuthentication, tuple[str, str]]] = None,
+        encrypted: bool = False,
+        disable_server_verification: bool = False,
+        tls_root_certs: Optional[bytes] = None,
+        user_agent: Optional[str] = None,
     ):
         """Creates a new GdsArrowClient instance.
 
@@ -85,8 +93,6 @@ def __init__(
             A flag that disables server verification for TLS connections (default is False)
         tls_root_certs: Optional[bytes]
             PEM-encoded certificates that are used for the connection to the GDS Arrow Flight server
-        arrow_endpoint_version:
-            The version of the Arrow endpoint to use (default is ArrowEndpointVersion.V1)
         user_agent: Optional[str]
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
         retry_config: Optional[RetryConfig]
@@ -117,6 +123,48 @@ def __init__(
 
         self._flight_client = self._instantiate_flight_client()
 
+    def connection_info(self) -> ConnectionInfo:
+        """
+        Returns the host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        tuple[str, int]
+            the host and port of the GDS Arrow server
+        """
+        return ConnectionInfo(self._host, self._port, self._encrypted)
+
+    def request_token(self) -> Optional[str]:
+        """
+        Requests a token from the server and returns it.
+
+        Returns
+        -------
+        Optional[str]
+            a token from the server and returns it.
+        """
+
+        @retry(
+            reraise=True,
+            before=before_log("Request token", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def auth_with_retry() -> None:
+            client = self._flight_client
+            if self._auth:
+                auth_pair = self._auth.auth_pair()
+                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+
+        if self._auth:
+            auth_with_retry()
+            return self._auth_middleware.token()
+        else:
+            return "IGNORED"
+
+    def get_stream(self, ticket: Ticket) -> FlightStreamReader:
+        return self._flight_client.do_get(ticket)
 
     def do_action(self, endpoint: str, payload: bytes):
         return self._flight_client.do_action(Action(endpoint, payload))
@@ -155,3 +203,8 @@ def _instantiate_flight_client(self) -> flight.FlightClient:
         return flight.FlightClient(location, **client_options)
 
 
+@dataclass
+class ConnectionInfo:
+    host: str
+    port: int
+    encrypted: bool

graphdatascience/arrow_client/data_mapper.py

@@ -0,0 +1,55 @@
+import dataclasses
+import json
+from dataclasses import fields
+from typing import Any, Dict, Iterator, Type, TypeVar
+
+from pyarrow._flight import Result
+
+
+class DataMapper:
+    T = TypeVar("T")
+
+    @staticmethod
+    def deserialize_single(input_stream: Iterator[Result], cls: Type[T]) -> T:
+        rows = DataMapper.deserialize(input_stream, cls)
+
+        if len(rows) != 1:
+            raise ValueError(f"Expected exactly one row, got {len(rows)}")
+
+        return rows[0]
+
+    @staticmethod
+    def deserialize(input_stream, cls: Type[T]) -> list[T]:
+        def deserialize_row(row: Any):
+            result_dicts = json.loads(row.body.to_pybytes().decode())
+            if cls == Dict:
+                return result_dicts
+            return DataMapper.dict_to_dataclass(result_dicts, cls)
+
+        return [deserialize_row(row) for row in list(input_stream)]
+
+    @staticmethod
+    def dict_to_dataclass(data: Dict[str, Any], cls: Type[T], strict: bool = False) -> T:
+        """
+        Convert a dictionary to a dataclass instance with nested dataclass support.
+        """
+        if not dataclasses.is_dataclass(cls):
+            raise ValueError(f"{cls} is not a dataclass")
+
+        field_dict = {f.name: f for f in fields(cls)}
+        filtered_data = {}
+
+        for key, value in data.items():
+            if key in field_dict:
+                field = field_dict[key]
+                field_type = field.type
+
+                # Handle nested dataclasses
+                if dataclasses.is_dataclass(field_type) and isinstance(value, dict):
+                    filtered_data[key] = DataMapper.dict_to_dataclass(value, field_type, strict)
+                else:
+                    filtered_data[key] = value
+            elif strict:
+                raise ValueError(f"Extra field '{key}' not allowed in {cls.__name__}")
+
+        return cls(**filtered_data)

graphdatascience/arrow_client/v2/api_types.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+@dataclass(frozen=True, repr=True)
+class JobIdConfig:
+    jobId: str
+
+@dataclass(frozen=True, repr=True)
+class JobStatus:
+    jobId: str
+    status: str
+    progress: float
+
+
+@dataclass(frozen=True, repr=True)
+class MutateResult:
+    nodePropertiesWritten: int
+    relationshipsWritten: int

graphdatascience/arrow_client/v2/job_client.py

@@ -0,0 +1,64 @@
+import json
+from typing import Any, Dict
+
+from pandas import ArrowDtype, DataFrame
+from pyarrow._flight import Ticket
+
+from graphdatascience.arrow_client.authenticated_arrow_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.data_mapper import DataMapper
+from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
+
+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)
+        return DataMapper.deserialize_single(res, JobIdConfig).jobId
+
+    @staticmethod
+    def wait_for_job(client: AuthenticatedArrowClient, job_id: str):
+        while True:
+            job_id_config = {"jobId": job_id}
+            encoded_config = json.dumps(job_id_config).encode("utf-8")
+
+            arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, encoded_config)
+            job_status = DataMapper.deserialize_single(arrow_res, JobStatus)
+            if job_status.status == "Done":
+                break
+
+    @staticmethod
+    def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:
+        job_id_config = {"jobId": job_id}
+        encoded_config = json.dumps(job_id_config).encode("utf-8")
+
+        res = client.do_action_with_retry(RESULTS_SUMMARY_ENDPOINT, encoded_config)
+        return DataMapper.deserialize_single(res, Dict)
+
+    @staticmethod
+    def stream_results(client: AuthenticatedArrowClient, job_id: str) -> DataFrame:
+        job_id_config = {"jobId": job_id}
+        encoded_config = json.dumps(job_id_config).encode("utf-8")
+
+        res = client.do_action_with_retry("v2/results.stream", encoded_config)
+        export_job_id = DataMapper.deserialize_single(res, JobIdConfig).jobId
+
+        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/arrow_client/v2/mutation_client.py

@@ -0,0 +1,16 @@
+import json
+
+from graphdatascience.arrow_client.authenticated_arrow_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.data_mapper import DataMapper
+from graphdatascience.arrow_client.v2.api_types import MutateResult
+
+
+class MutationClient:
+    MUTATE_ENDPOINT = "v2/results.mutate"
+
+    @staticmethod
+    def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
+        mutate_config = {"jobId": job_id, "mutateProperty": mutate_property}
+        encoded_config = json.dumps(mutate_config).encode("utf-8")
+        mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, encoded_config)
+        return DataMapper.deserialize_single(mutate_arrow_res, MutateResult)

graphdatascience/arrow_client/v2/write_back_client.py

@@ -0,0 +1,49 @@
+import time
+from typing import Any, Optional
+
+from graphdatascience import QueryRunner
+from graphdatascience.arrow_client.authenticated_arrow_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]:
+        host, port, encrypted = self._arrow_client.connection_info()
+        token = self._arrow_client.request_token()
+        if token is None:
+            token = "IGNORED"
+        arrow_config = {"host": host, "port": port, "token": token, "encrypted": encrypted()}
+
+        return arrow_config