Optionally override arrow address

Co-Authored-By: Paul Horn <paul.horn@neo4j.com>

Author: jjaderberg

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,12 +31,16 @@ def create(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
+        listen_address_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"])
+        if listen_address_override is not None:
+            listen_address = listen_address_override
+        else:
+            listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
         arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
 
         if arrow_info["running"]: