Add more generic option arrow_client_options to allow user specific arrow client configuration

Author: DarthMax

changelog.md

@@ -4,8 +4,7 @@
 
 ## New features
 
-- `sessions.get_or_create()` now supports passing in a manual selection of root certificates for verifying server certificate
-- `sessions.get_or_create()` now supports disabling server certificate verification
+- `sessions.get_or_create()` now supports passing additional configuration options for the Arrow Flight Client
 
 
 ## Bug fixes

doc/modules/ROOT/pages/graph-analytics-serverless.adoc

@@ -77,6 +77,16 @@ This is a hard limit and cannot be changed.
 
 ==== Syntax
 
+        self,
+        session_name: str,
+        memory: SessionMemory,
+        db_connection: Optional[DbmsConnectionInfo] = None,
+        ttl: Optional[timedelta] = None,
+        cloud_location: Optional[CloudLocation] = None,
+        timeout: Optional[int] = None,
+        neo4j_driver_options: Optional[dict[str, Any]] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
+
 [source, role=no-test]
 ----
 sessions.get_or_create(
@@ -86,8 +96,8 @@ sessions.get_or_create(
     ttl: Optional[timedelta] = None,
     cloud_location: Optional[CloudLocation] = None,
     timeout: Optional[int] = None,
-   tls_root_certs: Optional[bytes] = None,
-   disable_server_verification: bool = False,
+    neo4j_driver_options: Optional[dict[str, Any]] = None,
+    arrow_client_options: Optional[dict[str, Any]] = None,
 ): AuraGraphDataScience
 ----
 
@@ -101,8 +111,8 @@ sessions.get_or_create(
 | ttl                         | datetime.timedelta | yes      | 1h      | Time-to-live for the session.
 | cloud_location              | https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/cloud_location[CloudLocation]      | yes      | None    | Aura-supported cloud provider and region where the GDS Session will run. Required for the Self-managed and Standalone types.
 | timeout                     | int                | yes      | None    | Seconds to wait for the session to enter Ready state. If the time is exceeded, an error will be returned.
-| tls_root_certs              | bytes              | yes      | None    | PEM-encoded root certificates used for verifying server certificate. If not specified, platform-specific default root certificates will be used.
-| disable_server_verification | bool               | yes      | False   | Set to True to disable server TLS certificate verification. Use with caution.
+| neo4j_driver_options        | dict[str, any]     | yes      | None    | Additional options passed to the Neo4j driver
+| arrow_client_options        | dict[str, any]     | yes      | None    | Additional options passed to the Arrow Flight Client
 |===
 
 

graphdatascience/graph/graph_object.py

@@ -68,7 +68,7 @@ def database(self) -> str:
         """
         return self._graph_info(["database"])  # type: ignore
 
-    def configuration(self) -> "Series[Any]":
+    def configuration(self) -> Series[Any]:
         """
         Returns:
             the configuration of the graph

graphdatascience/graph_data_science.py

@@ -42,6 +42,7 @@ def __init__(
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
         show_progress: bool = True,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ):
         """
         Construct a new GraphDataScience object.
@@ -65,14 +66,20 @@ def __init__(
                 - True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint.
                 - False will make the client use Bolt for all operations.
         arrow_disable_server_verification : bool, default True
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             A flag that overrides other TLS settings and disables server verification for TLS connections.
         arrow_tls_root_certs : Optional[bytes], default None
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             PEM-encoded certificates that are used for the connection to the
             GDS Arrow Flight server.
         bookmarks : Optional[Any], default None
             The Neo4j bookmarks to require a certain state before the next query gets executed.
         show_progress : bool, default True
             A flag to indicate whether to show progress bars for running procedures.
+        arrow_client_options : Optional[dict[str, Any]], default None
+            Additional options to be passed to the Arrow Flight client.
         """
         if aura_ds:
             GraphDataScience._validate_endpoint(endpoint)
@@ -105,14 +112,19 @@ def __init__(
                 username, password = auth
                 arrow_auth = UsernamePasswordAuthentication(username, password)
 
+            if arrow_client_options is None:
+                arrow_client_options = {}
+            if arrow_disable_server_verification:
+                arrow_client_options["disable_server_verification"] = True
+            if arrow_tls_root_certs is not None:
+                arrow_client_options["tls_root_certs"] = arrow_tls_root_certs
             self._query_runner = ArrowQueryRunner.create(
                 self._query_runner,
-                arrow_info,
-                arrow_auth,
-                self._query_runner.encrypted(),
-                arrow_disable_server_verification,
-                arrow_tls_root_certs,
-                None if arrow is True else arrow,
+                arrow_info=arrow_info,
+                arrow_authentication=arrow_auth,
+                encrypted=self._query_runner.encrypted(),
+                arrow_client_options=arrow_client_options,
+                connection_string_override=None if arrow is True else arrow,
             )
 
         self._query_runner.set_show_progress(show_progress)

graphdatascience/query_runner/arrow_query_runner.py

@@ -25,22 +25,20 @@ def create(
         arrow_info: ArrowInfo,
         arrow_authentication: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
-        disable_server_verification: bool = False,
-        tls_root_certs: Optional[bytes] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
         connection_string_override: Optional[str] = None,
         retry_config: Optional[RetryConfig] = None,
     ) -> ArrowQueryRunner:
         if not arrow_info.enabled:
             raise ValueError("Arrow is not enabled on the server")
 
         gds_arrow_client = GdsArrowClient.create(
-            arrow_info,
-            arrow_authentication,
-            encrypted,
-            disable_server_verification,
-            tls_root_certs,
-            connection_string_override,
+            arrow_info=arrow_info,
+            auth=arrow_authentication,
+            encrypted=encrypted,
+            connection_string_override=connection_string_override,
             retry_config=retry_config,
+            arrow_client_options=arrow_client_options,
         )
 
         return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())

graphdatascience/query_runner/gds_arrow_client.py

@@ -55,6 +55,7 @@ def create(
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
         retry_config: Optional[RetryConfig] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> GdsArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -78,14 +79,15 @@ def create(
             )
 
         return GdsArrowClient(
-            host,
-            retry_config,
-            int(port),
-            auth,
-            encrypted,
-            disable_server_verification,
-            tls_root_certs,
-            arrow_endpoint_version,
+            host=host,
+            retry_config=retry_config,
+            port=int(port),
+            auth=auth,
+            encrypted=encrypted,
+            disable_server_verification=disable_server_verification,
+            tls_root_certs=tls_root_certs,
+            arrow_endpoint_version=arrow_endpoint_version,
+            arrow_client_options=arrow_client_options,
         )
 
     def __init__(
@@ -99,6 +101,7 @@ def __init__(
         tls_root_certs: Optional[bytes] = None,
         arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.V1,
         user_agent: Optional[str] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ):
         """Creates a new GdsArrowClient instance.
 
@@ -113,27 +116,41 @@ def __init__(
         encrypted: bool
             A flag that indicates whether the connection should be encrypted (default is False)
         disable_server_verification: bool
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             A flag that disables server verification for TLS connections (default is False)
         tls_root_certs: Optional[bytes]
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             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]
             The retry configuration to use for the Arrow requests send by the client.
+        arrow_client_options: Optional[dict[str, Any]]
+            Additional configuration for the Arrow flight client.
+
         """
         self._arrow_endpoint_version = arrow_endpoint_version
         self._host = host
         self._port = port
         self._auth = None
         self._encrypted = encrypted
-        self._disable_server_verification = disable_server_verification
-        self._tls_root_certs = tls_root_certs
         self._user_agent = user_agent
         self._retry_config = retry_config
         self._logger = logging.getLogger("gds_arrow_client")
 
+        self._arrow_client_options: dict[str, Any] = arrow_client_options  # type: ignore
+        if self._arrow_client_options is None:
+            self._arrow_client_options = {}
+
+        if disable_server_verification:
+            self._arrow_client_options["disable_server_verification"] = True
+        if tls_root_certs is not None:
+            self._arrow_client_options["tls_root_certs"] = tls_root_certs
+
         if auth:
             if not isinstance(auth, ArrowAuthentication):
                 username, password = auth
@@ -149,18 +166,27 @@ def _instantiate_flight_client(self) -> flight.FlightClient:
             if self._encrypted
             else flight.Location.for_grpc_tcp(self._host, self._port)
         )
-        client_options: dict[str, Any] = {"disable_server_verification": self._disable_server_verification}
+
+        client_options = self._arrow_client_options.copy()
+
         if self._auth:
             user_agent = f"neo4j-graphdatascience-v{__version__} pyarrow-v{arrow_version}"
             if self._user_agent:
                 user_agent = self._user_agent
 
-            client_options["middleware"] = [
-                AuthFactory(self._auth_middleware),
-                UserAgentFactory(useragent=user_agent),
-            ]
-        if self._tls_root_certs:
-            client_options["tls_root_certs"] = self._tls_root_certs
+            if "middleware" in client_options:
+                if not isinstance(client_options["middleware"], list):
+                    raise TypeError("client_options['middleware'] must be a list")
+            else:
+                client_options["middleware"] = []
+
+            client_options["middleware"].extend(
+                [
+                    AuthFactory(self._auth_middleware),
+                    UserAgentFactory(useragent=user_agent),
+                ]
+            )  # type: ignore
+
         return flight.FlightClient(location, **client_options)
 
     def connection_info(self) -> tuple[str, int]:

graphdatascience/session/aura_graph_data_science.py

@@ -38,8 +38,7 @@ def create(
         arrow_authentication: Optional[ArrowAuthentication],
         db_endpoint: Optional[Union[Neo4jQueryRunner, DbmsConnectionInfo]],
         delete_fn: Callable[[], bool],
-        arrow_disable_server_verification: bool = False,
-        arrow_tls_root_certs: Optional[bytes] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
         bookmarks: Optional[Any] = None,
         show_progress: bool = True,
     ) -> AuraGraphDataScience:
@@ -55,17 +54,15 @@ def create(
             arrow_info=arrow_info,
             arrow_authentication=arrow_authentication,
             encrypted=session_bolt_query_runner.encrypted(),
-            disable_server_verification=arrow_disable_server_verification,
-            tls_root_certs=arrow_tls_root_certs,
+            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_disable_server_verification,
-            arrow_tls_root_certs,
+            arrow_client_options=arrow_client_options,
         )
 
         gds_version = session_bolt_query_runner.server_version()

graphdatascience/session/dedicated_sessions.py

@@ -63,8 +63,7 @@ def get_or_create(
         cloud_location: Optional[CloudLocation] = None,
         timeout: Optional[int] = None,
         neo4j_driver_options: Optional[dict[str, Any]] = None,
-        tls_root_certs: Optional[bytes] = None,
-        disable_server_verification: bool = False,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> AuraGraphDataScience:
         if db_connection is None:
             if not cloud_location:
@@ -106,8 +105,7 @@ def get_or_create(
             session_bolt_connection_info=session_bolt_connection_info,
             arrow_authentication=arrow_authentication,
             db_runner=db_runner,
-            tls_root_certs=tls_root_certs,
-            disable_server_verification=disable_server_verification,
+            arrow_client_options=arrow_client_options,
         )
 
     def _create_db_runner(
@@ -213,14 +211,12 @@ def _construct_client(
         session_bolt_connection_info: DbmsConnectionInfo,
         arrow_authentication: ArrowAuthentication,
         db_runner: Optional[Neo4jQueryRunner],
-        tls_root_certs: Optional[bytes],
-        disable_server_verification: bool,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> AuraGraphDataScience:
         return AuraGraphDataScience.create(
             session_bolt_connection_info=session_bolt_connection_info,
             arrow_authentication=arrow_authentication,
             db_endpoint=db_runner,
             delete_fn=lambda: self._aura_api.delete_session(session_id=session_id),
-            arrow_tls_root_certs=tls_root_certs,
-            arrow_disable_server_verification=disable_server_verification,
+            arrow_client_options=arrow_client_options,
         )

graphdatascience/session/gds_sessions.py

@@ -106,8 +106,7 @@ def get_or_create(
         cloud_location: Optional[CloudLocation] = None,
         timeout: Optional[int] = None,
         neo4j_driver_config: Optional[dict[str, Any]] = None,
-        tls_root_certs: Optional[bytes] = None,
-        disable_server_verification: bool = False,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> AuraGraphDataScience:
         """
         Retrieves an existing session with the given session name and database connection,
@@ -124,8 +123,7 @@ def get_or_create(
             cloud_location (Optional[CloudLocation]): The cloud location. Required if the GDS session is for a self-managed database.
             timeout (Optional[int]): Optional timeout (in seconds) when waiting for session to become ready. If unset the method will wait forever. If set and session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation.
             neo4j_driver_config (Optional[dict[str, Any]]): Optional configuration for the Neo4j driver.
-            tls_root_certs (Optional[bytes]): Manually specify PEM-encoded root certificates used for verifying server certificate. If not specified, platform-specific default root certificates will be used.
-            disable_server_verification (bool): Set to True to disable server certificate verification. Use with caution.
+            arrow_client_options (Optional[dict[str, Any]]): Optional configuration for the Arrow Flight client.
         Returns:
             AuraGraphDataScience: The session.
         """
@@ -137,8 +135,7 @@ def get_or_create(
             cloud_location=cloud_location,
             timeout=timeout,
             neo4j_driver_options=neo4j_driver_config,
-            tls_root_certs=tls_root_certs,
-            disable_server_verification=disable_server_verification,
+            arrow_client_options=arrow_client_options,
         )
 
     def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str] = None) -> bool: