Merge pull request #603 from neo4j/robinhood

Support overriding Arrow connection string

Author: FlorentinD

graphdatascience/graph_data_science.py

@@ -28,7 +28,7 @@ def __init__(
         auth: Optional[Tuple[str, str]] = None,
         aura_ds: bool = False,
         database: Optional[str] = None,
-        arrow: bool = True,
+        arrow: Union[str, bool] = True,
         arrow_disable_server_verification: bool = True,
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
@@ -47,9 +47,10 @@ def __init__(
             to a Neo4j Aura instance.
         database: Optional[str], default None
             The Neo4j database to query against.
-        arrow : bool, default True
-            A flag that indicates that the client should use Apache Arrow
-            for data streaming if it is available on the server.
+        arrow : Union[str, bool], default True
+            Arrow connection information. Either a flag that indicates whether the client should use Apache Arrow
+            for data streaming if it is available on the server. True means discover the connection URI from the server.
+            A connection URI (str) can also be provided.
         arrow_disable_server_verification : bool, default True
             A flag that indicates that, if the flight client is connecting with
             TLS, that it skips server verification. If this is enabled, all
@@ -77,6 +78,7 @@ def __init__(
                 self._query_runner.encrypted(),
                 arrow_disable_server_verification,
                 arrow_tls_root_certs,
+                None if arrow is True else arrow,
             )
 
         super().__init__(self._query_runner, "gds", self._server_version)

graphdatascience/query_runner/arrow_query_runner.py

@@ -31,17 +31,22 @@ def create(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
+        connection_string_override: Optional[str] = None,
     ) -> QueryRunner:
         arrow_info = (
             fallback_query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
         )
         server_version = fallback_query_runner.server_version()
-        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
+        connection_string: str
+        if connection_string_override is not None:
+            connection_string = connection_string_override
+        else:
+            connection_string = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
         arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
 
         if arrow_info["running"]:
             return ArrowQueryRunner(
-                listen_address,
+                connection_string,
                 fallback_query_runner,
                 server_version,
                 auth,

graphdatascience/tests/unit/test_arrow_runner.py

@@ -0,0 +1,40 @@
+import pytest
+from pandas import DataFrame
+from pyarrow.flight import FlightUnavailableError
+
+from .conftest import CollectingQueryRunner
+from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
+from graphdatascience.server_version.server_version import ServerVersion
+
+
+@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
+def test_create(runner: CollectingQueryRunner) -> None:
+    runner.set__mock_result(DataFrame([{"running": True, "listenAddress": "localhost:1234"}]))
+
+    arrow_runner = ArrowQueryRunner.create(runner)
+
+    assert isinstance(arrow_runner, ArrowQueryRunner)
+
+    with pytest.raises(FlightUnavailableError, match=".+ failed to connect .+ ipv4:127.0.0.1:1234: .+"):
+        arrow_runner._flight_client.list_actions()
+
+
+@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
+def test_return_fallback_when_arrow_is_not_running(runner: CollectingQueryRunner) -> None:
+    runner.set__mock_result(DataFrame([{"running": False, "listenAddress": "localhost:1234"}]))
+
+    arrow_runner = ArrowQueryRunner.create(runner)
+
+    assert arrow_runner is runner
+
+
+@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
+def test_create_with_provided_connection(runner: CollectingQueryRunner) -> None:
+    runner.set__mock_result(DataFrame([{"running": True, "listenAddress": "localhost:1234"}]))
+
+    arrow_runner = ArrowQueryRunner.create(runner, connection_string_override="localhost:4321")
+
+    assert isinstance(arrow_runner, ArrowQueryRunner)
+
+    with pytest.raises(FlightUnavailableError, match=".+ failed to connect .+ ipv4:127.0.0.1:4321: .+"):
+        arrow_runner._flight_client.list_actions()