Allow advanced auth methods for dbms connection info

Author: FlorentinD

changelog.md

@@ -2,7 +2,7 @@
 
 ## Breaking changes
 
-* Drop support for PyArrow 16
+- Drop support for PyArrow 16
 
 ## New features
 
@@ -20,5 +20,6 @@
   - `DbmsConnectionInfo::from_env`
 - Retry internal functions known to be idempotent. Reduces issues such as `SessionExpiredError`.
 - Add support for PyArrow 20
+- Add support for more advanced authentication in `DbmsConnectionInfo`. Allowing to pass `auth` of type `neo4j.Auth` instead of username + password.
 
 ## Other changes

graphdatascience/graph_data_science.py

@@ -4,6 +4,7 @@
 from types import TracebackType
 from typing import Any, Optional, Type, Union
 
+import neo4j
 from neo4j import Driver
 from pandas import DataFrame
 
@@ -78,8 +79,10 @@ def __init__(
         if isinstance(endpoint, QueryRunner):
             self._query_runner = endpoint
         else:
+            if auth:
+                db_auth = neo4j.basic_auth(*auth)
             self._query_runner = Neo4jQueryRunner.create_for_db(
-                endpoint, auth, aura_ds, database, bookmarks, show_progress
+                endpoint, db_auth, aura_ds, database, bookmarks, show_progress
             )
 
         self._server_version = self._query_runner.server_version()

graphdatascience/query_runner/neo4j_query_runner.py

@@ -32,7 +32,7 @@ class Neo4jQueryRunner(QueryRunner):
     @staticmethod
     def create_for_db(
         endpoint: Union[str, neo4j.Driver],
-        auth: Optional[tuple[str, str]] = None,
+        auth: Union[tuple[str, str], neo4j.Auth, None] = None,
         aura_ds: bool = False,
         database: Optional[str] = None,
         bookmarks: Optional[Any] = None,
@@ -79,7 +79,7 @@ def create_for_db(
     @staticmethod
     def create_for_session(
         endpoint: str,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Union[tuple[str, str], neo4j.Auth, None] = None,
         show_progress: bool = True,
     ) -> Neo4jQueryRunner:
         driver_config: dict[str, Any] = {"user_agent": f"neo4j-graphdatascience-v{__version__}"}
@@ -125,7 +125,7 @@ def __init__(
         self,
         driver: neo4j.Driver,
         protocol: str,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Union[tuple[str, str], neo4j.Auth, None] = None,
         config: dict[str, Any] = {},
         database: Optional[str] = neo4j.DEFAULT_DATABASE,
         auto_close: bool = False,

graphdatascience/session/aura_graph_data_science.py

@@ -44,7 +44,7 @@ def create(
     ) -> AuraGraphDataScience:
         session_bolt_query_runner = Neo4jQueryRunner.create_for_session(
             endpoint=session_bolt_connection_info.uri,
-            auth=session_bolt_connection_info.auth(),
+            auth=session_bolt_connection_info.get_auth(),
             show_progress=show_progress,
         )
 
@@ -75,7 +75,7 @@ def create(
             else:
                 db_bolt_query_runner = Neo4jQueryRunner.create_for_db(
                     db_endpoint.uri,
-                    db_endpoint.auth(),
+                    db_endpoint.get_auth(),
                     aura_ds=True,
                     show_progress=False,
                     database=db_endpoint.database,

graphdatascience/session/dbms_connection_info.py

@@ -4,26 +4,40 @@
 from dataclasses import dataclass
 from typing import Optional
 
+from neo4j import Auth, basic_auth
+
 
 @dataclass
 class DbmsConnectionInfo:
     """
     Represents the connection information for a Neo4j DBMS, such as an AuraDB instance.
+    Supports both username/password as well as the authentication options provided by the Neo4j Python driver.
     """
 
     uri: str
-    username: str
-    password: str
+    username: Optional[str] = None
+    password: Optional[str] = None
     database: Optional[str] = None
-
-    def auth(self) -> tuple[str, str]:
+    # Optional: typed authentication, used instead of username/password. Supports for example a token. See https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods
+    auth: Optional[Auth] = None
+
+    def __post_init__(self) -> None:
+        # Validate auth fields
+        if (self.username or self.password) and self.auth:
+            raise ValueError(
+                "Cannot provide both username/password and token for authentication. "
+                "Please provide either a username/password or a token."
+            )
+
+    def get_auth(self) -> Optional[Auth]:
         """
-        Returns the username and password for authentication.
-
         Returns:
-            A tuple containing the username and password.
+            A neo4j.Auth object for authentication.
         """
-        return self.username, self.password
+        auth = self.auth
+        if self.username and self.password:
+            auth = basic_auth(self.username, self.password)
+        return auth
 
     @staticmethod
     def from_env() -> DbmsConnectionInfo:

graphdatascience/session/dedicated_sessions.py

@@ -111,7 +111,7 @@ def _create_db_runner(
     ) -> Neo4jQueryRunner:
         db_runner = Neo4jQueryRunner.create_for_db(
             endpoint=db_connection.uri,
-            auth=db_connection.auth(),
+            auth=db_connection.get_auth(),
             aura_ds=True,
             show_progress=False,
             database=db_connection.database,

graphdatascience/tests/unit/session/test_dbms_connection_info.py

@@ -0,0 +1,38 @@
+import neo4j
+from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
+
+
+def test_dbms_connection_info_username_password() -> None:
+    dci = DbmsConnectionInfo(
+        uri="foo.bar",
+        username="neo4j",
+        password="password",
+        database="neo4j",
+    )
+
+    assert dci.get_auth() == neo4j.basic_auth("neo4j", "password")
+
+
+def test_dbms_connection_info_advanced_auth() -> None:
+    advanced_auth = neo4j.kerberos_auth("foo bar")
+
+    dci = DbmsConnectionInfo(uri="foo.bar", database="neo4j", auth=advanced_auth)
+
+    assert dci.get_auth() == advanced_auth
+
+
+def test_dbms_connection_info_fail_on_auth_and_username() -> None:
+    try:
+        DbmsConnectionInfo(
+            uri="foo.bar",
+            username="neo4j",
+            password="password",
+            auth=neo4j.basic_auth("other", "other"),
+        )
+    except ValueError as e:
+        assert str(e) == (
+            "Cannot provide both username/password and token for authentication. "
+            "Please provide either a username/password or a token."
+        )
+    else:
+        assert False, "Expected ValueError was not raised"