Expose arrow_client method in AuraGraphDataScience

Author: DarthMax

src/graphdatascience/arrow_client/v1/gds_arrow_client.py

@@ -2,13 +2,11 @@
 
 import json
 import logging
-import re
 from types import TracebackType
 from typing import Any, Callable, Iterable, Type
 
 import pandas
 import pyarrow
-from neo4j.exceptions import ClientError
 from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Table, chunked_array, flight
 from pyarrow.types import is_dictionary
 from pydantic import BaseModel
@@ -18,6 +16,7 @@
 from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_single
 
 from ...semantic_version.semantic_version import SemanticVersion
+from ..error_handler import handle_flight_error
 
 
 class GdsArrowClient:
@@ -515,7 +514,7 @@ def upload_batch(p: RecordBatch) -> None:
                     ack_stream.read()
                     progress_callback(partition.num_rows)
         except Exception as e:
-            GdsArrowClient.handle_flight_error(e)
+            handle_flight_error(e)
 
     def _get_data(
         self,
@@ -560,7 +559,7 @@ def _fetch_get_result(self, get: flight.FlightStreamReader) -> pandas.DataFrame:
         try:
             arrow_table = get.read_all()
         except Exception as e:
-            GdsArrowClient.handle_flight_error(e)
+            handle_flight_error(e)
         arrow_table = self._sanitize_arrow_table(arrow_table)
         if SemanticVersion.from_string(pandas.__version__) >= SemanticVersion(2, 0, 0):
             return arrow_table.to_pandas(types_mapper=pandas.ArrowDtype)  # type: ignore
@@ -615,26 +614,6 @@ def _decode_pyarrow_array(array: Array) -> Array:
         else:
             return array
 
-    @staticmethod
-    def handle_flight_error(e: Exception) -> None:
-        if isinstance(e, flight.FlightServerError | flight.FlightInternalError | ClientError):
-            original_message = e.args[0] if len(e.args) > 0 else e.message
-            improved_message = original_message.replace(
-                "Flight RPC failed with message: org.apache.arrow.flight.FlightRuntimeException: ", ""
-            )
-            improved_message = improved_message.replace(
-                "Flight returned internal error, with message: org.apache.arrow.flight.FlightRuntimeException: ", ""
-            )
-            improved_message = improved_message.replace(
-                "Failed to invoke procedure `gds.arrow.project`: Caused by: org.apache.arrow.flight.FlightRuntimeException: ",
-                "",
-            )
-            improved_message = re.sub(r"(\. )?gRPC client debug context: .+$", "", improved_message)
-
-            raise flight.FlightServerError(improved_message)
-        else:
-            raise e
-
 
 class NodeLoadDoneResult(BaseModel):
     name: str

src/graphdatascience/arrow_client/v2/api_types.py

@@ -37,10 +37,10 @@ def aborted(self) -> bool:
         return self.status == "Aborted"
 
     def succeeded(self) -> bool:
-        return self.status == "Done"
+        return self.status.lower() == "done"
 
     def running(self) -> bool:
-        return self.status == "Running"
+        return self.status.lower() == "running"
 
 
 class MutateResult(ArrowBaseModel):

src/graphdatascience/query_runner/session_query_runner.py

@@ -6,13 +6,14 @@
 
 from pandas import DataFrame
 
-from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
+from graphdatascience.arrow_client.v2.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
 from graphdatascience.query_runner.progress.query_progress_logger import QueryProgressLogger
 from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.server_version.server_version import ServerVersion
 
+from ..arrow_client.error_handler import handle_flight_error
 from ..call_parameters import CallParameters
 from ..session.dbms.protocol_resolver import ProtocolVersionResolver
 from .protocol.project_protocols import ProjectProtocol
@@ -183,7 +184,7 @@ def run_projection() -> DataFrame:
             else:
                 return run_projection()
         except Exception as e:
-            GdsArrowClient.handle_flight_error(e)
+            handle_flight_error(e)
             raise e  # above should already raise
 
     def _remote_write_back(

src/graphdatascience/session/aura_graph_data_science.py

@@ -5,7 +5,7 @@
 from pandas import DataFrame
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
+from graphdatascience.arrow_client.v2.gds_arrow_client import GdsArrowClient
 from graphdatascience.call_builder import IndirectCallBuilder
 from graphdatascience.endpoints import (
     AlphaRemoteEndpoints,
@@ -94,6 +94,7 @@ def create(
                 v2_endpoints=SessionV2Endpoints(
                     session_auth_arrow_client, db_bolt_query_runner, show_progress=show_progress
                 ),
+                authenticated_arrow_client=session_auth_arrow_client,
             )
         else:
             standalone_query_runner = StandaloneSessionQueryRunner(session_arrow_query_runner)
@@ -102,6 +103,7 @@ def create(
                 delete_fn=delete_fn,
                 gds_version=gds_version,
                 v2_endpoints=SessionV2Endpoints(session_auth_arrow_client, None, show_progress=show_progress),
+                authenticated_arrow_client=session_auth_arrow_client,
             )
 
     def __init__(
@@ -110,11 +112,13 @@ def __init__(
         delete_fn: Callable[[], bool],
         gds_version: ServerVersion,
         v2_endpoints: SessionV2Endpoints,
+        authenticated_arrow_client: AuthenticatedArrowClient,
     ):
         self._query_runner = query_runner
         self._delete_fn = delete_fn
         self._server_version = gds_version
         self._v2_endpoints = v2_endpoints
+        self._authenticated_arrow_client = authenticated_arrow_client
 
         super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
 
@@ -177,6 +181,18 @@ def v2(self) -> SessionV2Endpoints:
     def __getattr__(self, attr: str) -> IndirectCallBuilder:
         return IndirectCallBuilder(self._query_runner, f"gds.{attr}", self._server_version)
 
+    def arrow_client(self) -> GdsArrowClient:
+        """
+        Returns a GdsArrowClient that is authenticated to communicate with the Aura Graph Analytics Session.
+        This client can be used to get direct access to the sessions Arrow Flight server.
+
+        Returns:
+            A GdsArrowClient
+        -------
+
+        """
+        return GdsArrowClient(self._authenticated_arrow_client)
+
     def set_database(self, database: str) -> None:
         """
         Set the database which cypher queries are run against.

tests/integrationV2/procedure_surface/session/test_walking_skeleton.py

@@ -15,6 +15,7 @@ def gds(arrow_client: AuthenticatedArrowClient, db_query_runner: QueryRunner) ->
         delete_fn=lambda: True,
         gds_version=ServerVersion.from_string("1.2.3"),
         v2_endpoints=SessionV2Endpoints(arrow_client, db_query_runner),
+        authenticated_arrow_client=arrow_client,
     )
 
 

tests/unit/arrow_client/V1/test_gds_arrow_client.py

@@ -15,6 +15,7 @@
 )
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.error_handler import handle_flight_error
 from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.query_runner.arrow_info import ArrowInfo
@@ -397,7 +398,7 @@ def test_handle_flight_error() -> None:
         FlightServerError,
         match="FlightServerError: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database.",
     ):
-        GdsArrowClient.handle_flight_error(
+        handle_flight_error(
             FlightServerError(
                 'FlightServerError: Flight RPC failed with message: org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database.. gRPC client debug context: UNKNOWN:Error received from peer ipv4:35.241.177.75:8491 {created_time:"2024-08-29T15:59:03.828903999+02:00", grpc_status:2, grpc_message:"org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Graph with name `people-and-fruits` does not exist on database `neo4j`. It might exist on another database."}. Client context: IOError: Server never sent a data message. Detail: Internal'
             )
@@ -407,7 +408,7 @@ def test_handle_flight_error() -> None:
         FlightServerError,
         match=re.escape("FlightServerError: UNKNOWN: Unexpected configuration key(s): [undirectedRelationshipTypes]"),
     ):
-        GdsArrowClient.handle_flight_error(
+        handle_flight_error(
             FlightServerError(
                 "FlightServerError: Flight returned internal error, with message: org.apache.arrow.flight.FlightRuntimeException: UNKNOWN: Unexpected configuration key(s): [undirectedRelationshipTypes]"
             )

tests/unit/session/test_aura_graph_data_science.py

@@ -14,6 +14,7 @@ def test_remote_projection_configuration(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     g = gds.graph.project(
@@ -50,6 +51,7 @@ def test_remote_projection_defaults(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     g = gds.graph.project("foo", "RETURN gds.graph.project(0, 1)")
@@ -78,6 +80,7 @@ def test_remote_algo_write(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     G, _ = gds.graph.project("foo", "RETURN gds.graph.project(0, 1)")
@@ -99,6 +102,7 @@ def test_remote_algo_write_configuration(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     G, _ = gds.graph.project("foo", "RETURN gds.graph.project(0, 1)")
@@ -125,6 +129,7 @@ def test_remote_graph_write(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     G, _ = gds.graph.project("foo", "RETURN gds.graph.project(0, 1)")
@@ -149,6 +154,7 @@ def test_remote_graph_write_configuration(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     G, _ = gds.graph.project("foo", "RETURN gds.graph.project(0, 1)")
@@ -176,6 +182,7 @@ def test_run_cypher_write(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     gds.run_cypher("RETURN 1", params={"foo": 1}, mode=QueryMode.WRITE, database="bar", retryable=True)
@@ -193,6 +200,7 @@ def test_run_cypher_read(mocker: MockerFixture) -> None:
         delete_fn=lambda: True,
         gds_version=v,
         v2_endpoints=mocker.Mock(),
+        authenticated_arrow_client=mocker.Mock(),
     )
 
     gds.run_cypher("RETURN 1", params={"foo": 1}, mode=QueryMode.READ, retryable=False)