Allow passing in root certificates and disable server verification

Co-authored-by: Max Kießling <max.kiessling@neotechnology.com>

Author: Mats-SX

changelog.md

@@ -4,6 +4,10 @@
 
 ## 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
+
+
 ## 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.

graphdatascience/session/dedicated_sessions.py

@@ -63,6 +63,8 @@ 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,
     ) -> AuraGraphDataScience:
         if db_connection is None:
             if not cloud_location:
@@ -104,6 +106,8 @@ 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,
         )
 
     def _create_db_runner(
@@ -209,10 +213,14 @@ def _construct_client(
         session_bolt_connection_info: DbmsConnectionInfo,
         arrow_authentication: ArrowAuthentication,
         db_runner: Optional[Neo4jQueryRunner],
+        tls_root_certs: Optional[bytes],
+        disable_server_verification: bool,
     ) -> 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,
         )

graphdatascience/session/gds_sessions.py

@@ -106,6 +106,8 @@ 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,
     ) -> AuraGraphDataScience:
         """
         Retrieves an existing session with the given session name and database connection,
@@ -122,6 +124,8 @@ 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.
         Returns:
             AuraGraphDataScience: The session.
         """
@@ -133,6 +137,8 @@ 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,
         )
 
     def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str] = None) -> bool:

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -365,6 +365,60 @@ def test_create_attached_session(mocker: MockerFixture, aura_api: AuraApi) -> No
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "ffff0-ffff1",
+        "disable_server_verification": False,
+        "tls_root_certs": None,
+    }
+
+    assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
+
+    assert len(sessions.list()) == 1
+    actual_session = sessions.list()[0]
+
+    assert actual_session.name == "my-session"
+    assert actual_session.user_id == "user-1"
+    assert actual_session.ttl == ttl
+
+
+def test_create_attached_session_passthrough_tls_settings(mocker: MockerFixture, aura_api: AuraApi) -> None:
+    _setup_db_instance(aura_api)
+
+    sessions = DedicatedSessions(aura_api)
+
+    patch_construct_client(mocker)
+    patch_neo4j_query_runner(mocker)
+
+    ttl = timedelta(hours=42)
+    fake_certs = bytes(1)
+    gds_parameters = sessions.get_or_create(
+        "my-session",
+        SessionMemory.m_8GB,
+        DbmsConnectionInfo("neo4j+s://ffff0.databases.neo4j.io", "dbuser", "db_pw"),
+        ttl=ttl,
+        disable_server_verification=True,
+        tls_root_certs=fake_certs,
+    )
+
+    arrow_authentication = gds_parameters["arrow_authentication"]  # type: ignore
+    del gds_parameters["arrow_authentication"]
+
+    dbms_authentication = gds_parameters["db_runner"].pop("auth")  # type: ignore
+
+    assert (dbms_authentication.principal, dbms_authentication.credentials) == ("dbuser", "db_pw")
+
+    assert gds_parameters == {  # type: ignore
+        "db_runner": {
+            "endpoint": "neo4j+s://ffff0.databases.neo4j.io",
+            "aura_ds": True,
+            "database": None,
+            "show_progress": False,
+            "config": None,
+        },
+        "session_bolt_connection_info": DbmsConnectionInfo(
+            uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
+        ),
+        "session_id": "ffff0-ffff1",
+        "disable_server_verification": True,
+        "tls_root_certs": fake_certs,
     }
 
     assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
@@ -412,6 +466,8 @@ def test_create_standalone_session(mocker: MockerFixture, aura_api: AuraApi) ->
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "None-ffff0",
+        "disable_server_verification": False,
+        "tls_root_certs": None,
     }
 
     assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
@@ -465,6 +521,8 @@ def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "ffff0-ffff1",
+        "disable_server_verification": False,
+        "tls_root_certs": None,
     }
 
     assert gds_args1 == gds_args2