Use new GdsArrowClient internally

FlorentinD · FlorentinD · commit 80ff4d1c09b2 · 2025-11-20T16:50:59.000+01:00
new one based on AuthenticatedArrowClient
diff --git a/graphdatascience/arrow_client/authenticated_flight_client.py b/graphdatascience/arrow_client/authenticated_flight_client.py
@@ -124,7 +124,7 @@ def connection_info(self) -> ConnectionInfo:
 
         Returns
         -------
-        tuple[str, int]
+        ConnectionInfo
             the host and port of the GDS Arrow server
         """
         return ConnectionInfo(self._host, self._port, self._encrypted)
diff --git a/graphdatascience/arrow_client/v1/gds_arrow_client.py b/graphdatascience/arrow_client/v1/gds_arrow_client.py
@@ -14,7 +14,7 @@
 from pydantic import BaseModel
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
-from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_single
 
 from ...semantic_version.semantic_version import SemanticVersion
@@ -451,6 +451,29 @@ def upload_triplets(
         """
         self._upload_data(graph_name, "triplet", triplet_data, batch_size, progress_callback)
 
+    def advertised_connection_info(self) -> ConnectionInfo:
+        """
+        Returns the host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        ConnectionInfo
+            the host and port of the GDS Arrow server
+        """
+        return self._flight_client.advertised_connection_info()
+
+    def request_token(self) -> str | None:
+        """
+        Requests a token from the server and returns it.
+
+        Returns
+        -------
+        str | None
+            a token from the server and returns it.
+        """
+
+        return self._flight_client.request_token()
+
     def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str, Any]:
         action_type = f"{ArrowEndpointVersion.V1.prefix()}{action_type}"
         raw_result = self._flight_client.do_action_with_retry(action_type, meta_data)
diff --git a/graphdatascience/procedure_surface/cypher/catalog/node_properties_cypher_endpoints.py b/graphdatascience/procedure_surface/cypher/catalog/node_properties_cypher_endpoints.py
@@ -1,5 +1,6 @@
 from pandas import DataFrame
 
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.catalog.node_properties_endpoints import (
@@ -11,7 +12,6 @@
 from graphdatascience.procedure_surface.api.default_values import ALL_LABELS
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 from graphdatascience.procedure_surface.utils.result_utils import join_db_node_properties, transpose_property_columns
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.query_runner import QueryRunner
 
 
diff --git a/graphdatascience/procedure_surface/cypher/catalog/relationship_cypher_endpoints.py b/graphdatascience/procedure_surface/cypher/catalog/relationship_cypher_endpoints.py
@@ -1,6 +1,7 @@
 from pandas import DataFrame
 
 from graphdatascience import QueryRunner
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.catalog.relationships_endpoints import (
@@ -14,7 +15,6 @@
 )
 from graphdatascience.procedure_surface.api.default_values import ALL_TYPES
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 
 
 class RelationshipCypherEndpoints(RelationshipsEndpoints):
diff --git a/graphdatascience/query_runner/arrow_graph_constructor.py b/graphdatascience/query_runner/arrow_graph_constructor.py
@@ -11,7 +11,8 @@
 from pandas import DataFrame
 from tqdm.auto import tqdm
 
-from .gds_arrow_client import GdsArrowClient
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
+
 from .graph_constructor import GraphConstructor
 
 
diff --git a/graphdatascience/query_runner/arrow_query_runner.py b/graphdatascience/query_runner/arrow_query_runner.py
@@ -5,15 +5,16 @@
 
 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.query_runner.arrow_authentication import ArrowAuthentication
 from graphdatascience.query_runner.query_mode import QueryMode
-from graphdatascience.retry_utils.retry_config import RetryConfig
+from graphdatascience.retry_utils.retry_config import RetryConfigV2
 
 from ..call_parameters import CallParameters
 from ..query_runner.arrow_info import ArrowInfo
 from ..server_version.server_version import ServerVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
-from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
 
@@ -27,20 +28,22 @@ def create(
         encrypted: bool = False,
         arrow_client_options: dict[str, Any] | None = None,
         connection_string_override: str | None = None,
-        retry_config: RetryConfig | None = None,
+        retry_config: RetryConfigV2 | None = None,
     ) -> ArrowQueryRunner:
         if not arrow_info.enabled:
             raise ValueError("Arrow is not enabled on the server")
 
-        gds_arrow_client = GdsArrowClient.create(
+        arrow_client = AuthenticatedArrowClient.create(
             arrow_info=arrow_info,
             auth=arrow_authentication,
             encrypted=encrypted,
+            arrow_client_options=arrow_client_options,
             connection_string_override=connection_string_override,
             retry_config=retry_config,
-            arrow_client_options=arrow_client_options,
         )
 
+        gds_arrow_client = GdsArrowClient(flight_client=arrow_client)
+
         return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())
 
     def __init__(
diff --git a/graphdatascience/query_runner/session_query_runner.py b/graphdatascience/query_runner/session_query_runner.py
@@ -6,6 +6,7 @@
 
 from pandas import DataFrame
 
+from graphdatascience.arrow_client.v1.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
@@ -14,7 +15,6 @@
 
 from ..call_parameters import CallParameters
 from ..session.dbms.protocol_resolver import ProtocolVersionResolver
-from .gds_arrow_client import GdsArrowClient
 from .protocol.project_protocols import ProjectProtocol
 from .protocol.write_protocols import WriteProtocol
 from .query_runner import QueryRunner
@@ -263,12 +263,12 @@ def _resolve_show_progress(self, show_progress: bool) -> bool:
         return self._show_progress and show_progress
 
     def _inject_arrow_config(self, params: dict[str, Any]) -> None:
-        host, port = self._gds_arrow_client.connection_info()
+        connection_info = self._gds_arrow_client.advertised_connection_info()
         token = self._gds_arrow_client.request_token()
         if token is None:
             token = "IGNORED"
 
-        params["host"] = host
-        params["port"] = port
+        params["host"] = connection_info.host
+        params["port"] = connection_info.port
         params["token"] = token
-        params["encrypted"] = self._gds_query_runner.encrypted()
+        params["encrypted"] = connection_info.encrypted
diff --git a/graphdatascience/session/aura_graph_data_science.py b/graphdatascience/session/aura_graph_data_science.py
@@ -6,6 +6,7 @@
 
 from graphdatascience import QueryRunner, ServerVersion
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.call_builder import IndirectCallBuilder
 from graphdatascience.endpoints import (
     AlphaRemoteEndpoints,
@@ -17,7 +18,6 @@
 from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
 from graphdatascience.query_runner.arrow_info import ArrowInfo
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.query_runner.session_query_runner import SessionQueryRunner
@@ -59,23 +59,17 @@ def create(
             arrow_client_options=arrow_client_options,
         )
 
-        # TODO: merge with the gds_arrow_client created inside ArrowQueryRunner
-        session_arrow_client = GdsArrowClient.create(
-            arrow_info,
-            arrow_authentication,
-            session_bolt_query_runner.encrypted(),
-            arrow_client_options=arrow_client_options,
-        )
-
-        gds_version = session_bolt_query_runner.server_version()
-
         session_auth_arrow_client = AuthenticatedArrowClient.create(
             arrow_info=arrow_info,
             auth=arrow_authentication,
             encrypted=session_bolt_query_runner.encrypted(),
             arrow_client_options=arrow_client_options,
         )
 
+        session_arrow_client = GdsArrowClient(flight_client=session_auth_arrow_client)
+
+        gds_version = session_bolt_query_runner.server_version()
+
         if db_endpoint is not None:
             if isinstance(db_endpoint, Neo4jQueryRunner):
                 db_bolt_query_runner = db_endpoint
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/cypher/catalog/test_node_properties_cypher_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/cypher/catalog/test_node_properties_cypher_endpoints.py
@@ -3,11 +3,11 @@
 import pytest
 
 from graphdatascience import QueryRunner
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.cypher.catalog.node_properties_cypher_endpoints import (
     NodePropertiesCypherEndpoints,
 )
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.tests.integrationV2.procedure_surface.cypher.cypher_graph_helper import create_graph
 
 
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/cypher/catalog/test_relationship_cypher_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/cypher/catalog/test_relationship_cypher_endpoints.py
@@ -3,12 +3,12 @@
 import pytest
 
 from graphdatascience import QueryRunner
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.catalog.relationships_endpoints import Aggregation
 from graphdatascience.procedure_surface.cypher.catalog.relationship_cypher_endpoints import (
     RelationshipCypherEndpoints,
 )
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.tests.integrationV2.procedure_surface.cypher.cypher_graph_helper import create_graph
 
 
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/cypher/conftest.py b/graphdatascience/tests/integrationV2/procedure_surface/cypher/conftest.py
@@ -9,7 +9,11 @@
 from testcontainers.neo4j import Neo4jContainer
 
 from graphdatascience import QueryRunner
-from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
+from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
+from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.tests.integrationV2.conftest import inside_ci
 
@@ -84,14 +88,18 @@ def query_runner(gds_plugin_container: DockerContainer) -> Generator[QueryRunner
 
 
 @pytest.fixture(scope="package")
-def gds_arrow_client(
-    gds_plugin_container: DockerContainer, query_runner: QueryRunner
-) -> Generator[GdsArrowClient, None, None]:
+def gds_arrow_client(gds_plugin_container: DockerContainer) -> Generator[GdsArrowClient, None, None]:
     arrow_port = int(gds_plugin_container.get_exposed_port(8491))
     with GdsArrowClient(
-        gds_plugin_container.get_container_host_ip(),
-        arrow_port,
-        ("neo4j", "password"),
-        encrypted=False,
+        flight_client=AuthenticatedArrowClient.create(
+            arrow_info=ArrowInfo(
+                listenAddress=f"{gds_plugin_container.get_container_host_ip()}:{arrow_port}",
+                enabled=True,
+                running=True,
+                versions=[ArrowEndpointVersion.V2.version()],
+            ),
+            auth=UsernamePasswordAuthentication("neo4j", "password"),
+            encrypted=False,
+        )
     ) as client:
         yield client
diff --git a/graphdatascience/tests/unit/test_arrow_runner.py b/graphdatascience/tests/unit/test_arrow_runner.py
@@ -1,25 +1,21 @@
 import pytest
 from pyarrow.flight import FlightUnavailableError
-from tenacity import retry_any, stop_after_attempt, wait_fixed
 
 from graphdatascience.query_runner.arrow_info import ArrowInfo
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
-from graphdatascience.retry_utils.retry_config import RetryConfig
-from graphdatascience.server_version.server_version import ServerVersion
+from graphdatascience.retry_utils.retry_config import RetryConfigV2, StopConfig
 
 from ...arrow_client.arrow_endpoint_version import ArrowEndpointVersion
 from .conftest import CollectingQueryRunner
 
 
-@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
 def test_create(runner: CollectingQueryRunner) -> None:
     arrow_info = ArrowInfo(
         listenAddress="localhost:1234", enabled=True, running=True, versions=[ArrowEndpointVersion.V1.version()]
     )
-    retry_config = RetryConfig(
-        retry=retry_any(),
-        stop=(stop_after_attempt(1)),
-        wait=wait_fixed(0),
+    retry_config = RetryConfigV2(
+        retryable_exceptions=[],
+        stop_config=StopConfig(after_attempt=1),
     )
 
     arrow_runner = ArrowQueryRunner.create(runner, arrow_info, retry_config=retry_config)
@@ -30,15 +26,13 @@ def test_create(runner: CollectingQueryRunner) -> None:
         arrow_runner._gds_arrow_client._send_action("TEST", {})
 
 
-@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
 def test_return_fallback_when_arrow_is_not_enabled(runner: CollectingQueryRunner) -> None:
     arrow_info = ArrowInfo(listenAddress="localhost:1234", enabled=False, running=False, versions=[])
 
     with pytest.raises(ValueError, match="Arrow is not enabled"):
         ArrowQueryRunner.create(runner, arrow_info)
 
 
-@pytest.mark.parametrize("server_version", [ServerVersion(2, 6, 0)])
 def test_create_with_provided_connection(runner: CollectingQueryRunner) -> None:
     arrow_info = ArrowInfo(
         listenAddress="localhost:1234", enabled=True, running=True, versions=[ArrowEndpointVersion.V1.version()]
