Merge pull request #912 from neo4j/tls-session

DarthMax · web-flow · commit 2a954451e4a5 · 2025-07-15T11:50:05.000+02:00
Allow passing in root certificates and disable server verification
diff --git a/changelog.md b/changelog.md
@@ -4,6 +4,9 @@
 
 ## New features
 
+- `sessions.get_or_create()` now supports passing additional configuration options for the Arrow Flight Client
+
+
 ## Bug fixes
 
 - Fix reporting error based on http responses from the Aura-API with an invalid JSON body. Earlier the client would report JSONDecodeError instead of showing the actual issue.
diff --git a/doc/modules/ROOT/pages/graph-analytics-serverless.adoc b/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,19 +96,23 @@ sessions.get_or_create(
     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,
 ): AuraGraphDataScience
 ----
 
 .Parameters:
 [opts="header",cols="3m,1m,1,1m,6", role="no-break"]
 |===
-| Name          | Type               | Optional | Default | Description
-| session_name  | str                | no       | -       | Name of the session. Must be unique within the project.
-| memory        | https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/session_memory[SessionMemory]      | no       | -       | Amount of memory available to the session.
-| db_connection |  https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/DbmsConnectionInfo[DbmsConnectionInfo] | yes      | None    | Bolt server URL, username, and password to a Neo4j DBMS. Required for the Attached and Self-managed types. Alternatively to username and password, you can provide a `neo4j.Auth` https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods[object].
-| 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.
+| Name                        | Type               | Optional | Default | Description
+| session_name                | str                | no       | -       | Name of the session. Must be unique within the project.
+| memory                      | https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/session_memory[SessionMemory]      | no       | -       | Amount of memory available to the session.
+| db_connection               |  https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/DbmsConnectionInfo[DbmsConnectionInfo] | yes      | None    | Bolt server URL, username, and password to a Neo4j DBMS. Required for the Attached and Self-managed types. Alternatively to username and password, you can provide a `neo4j.Auth` https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods[object].
+| 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.
+| neo4j_driver_options        | dict[str, any]     | yes      | None    | Additional options passed to the Neo4j driver to the Neo4j DBMS. Only relevant if `db_connection` is specified.
+| arrow_client_options        | dict[str, any]     | yes      | None    | Additional options passed to the Arrow Flight Client used to connect to the Session.
 |===
 
 
diff --git a/graphdatascience/graph/graph_object.py b/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
diff --git a/graphdatascience/graph_data_science.py b/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)
diff --git a/graphdatascience/query_runner/arrow_query_runner.py b/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())
diff --git a/graphdatascience/query_runner/gds_arrow_client.py b/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,39 @@ 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 = arrow_client_options if arrow_client_options is not None else {}
+
+        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 +164,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),
+                ]
+            )
+
         return flight.FlightClient(location, **client_options)
 
     def connection_info(self) -> tuple[str, int]:
diff --git a/graphdatascience/session/aura_graph_data_science.py b/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()
diff --git a/graphdatascience/session/dedicated_sessions.py b/graphdatascience/session/dedicated_sessions.py
@@ -63,6 +63,7 @@ def get_or_create(
         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,
     ) -> AuraGraphDataScience:
         if db_connection is None:
             if not cloud_location:
@@ -104,6 +105,7 @@ def get_or_create(
             session_bolt_connection_info=session_bolt_connection_info,
             arrow_authentication=arrow_authentication,
             db_runner=db_runner,
+            arrow_client_options=arrow_client_options,
         )
 
     def _create_db_runner(
@@ -209,10 +211,12 @@ def _construct_client(
         session_bolt_connection_info: DbmsConnectionInfo,
         arrow_authentication: ArrowAuthentication,
         db_runner: Optional[Neo4jQueryRunner],
+        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_client_options=arrow_client_options,
         )
diff --git a/graphdatascience/session/gds_sessions.py b/graphdatascience/session/gds_sessions.py
@@ -106,6 +106,7 @@ def get_or_create(
         cloud_location: Optional[CloudLocation] = None,
         timeout: Optional[int] = None,
         neo4j_driver_config: Optional[dict[str, Any]] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> AuraGraphDataScience:
         """
         Retrieves an existing session with the given session name and database connection,
@@ -121,7 +122,8 @@ def get_or_create(
             ttl: (Optional[timedelta]): The sessions time to live after inactivity in seconds.
             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.
+           neo4j_driver_config (Optional[dict[str, Any]]): Optional configuration for the Neo4j driver to the Neo4j DBMS. Only relevant if `db_connection` is specified..
+            arrow_client_options (Optional[dict[str, Any]]): Optional configuration for the Arrow Flight client.
         Returns:
             AuraGraphDataScience: The session.
         """
@@ -133,6 +135,7 @@ def get_or_create(
             cloud_location=cloud_location,
             timeout=timeout,
             neo4j_driver_options=neo4j_driver_config,
+            arrow_client_options=arrow_client_options,
         )
 
     def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str] = None) -> bool:
diff --git a/graphdatascience/tests/unit/test_dedicated_sessions.py b/graphdatascience/tests/unit/test_dedicated_sessions.py
