Allow arrow authentication via aura api token

Always pass api token from DedicatedSessions

Author: soerenreichardt

graphdatascience/graph_data_science.py

@@ -16,6 +16,7 @@
 from .query_runner.neo4j_query_runner import Neo4jQueryRunner
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
+from .session.arrow_authentication import UsernamePasswordAuthentication
 from .utils.util_proc_runner import UtilProcRunner
 from .version import __min_server_version__
 
@@ -93,10 +94,15 @@ def __init__(
 
         arrow_info = ArrowInfo.create(self._query_runner)
         if arrow and arrow_info.enabled and self._server_version >= ServerVersion(2, 1, 0):
+            arrow_auth = None
+            if auth is not None:
+                username, password = auth
+                arrow_auth = UsernamePasswordAuthentication(username, password)
+
             self._query_runner = ArrowQueryRunner.create(
                 self._query_runner,
                 arrow_info,
-                auth,
+                arrow_auth,
                 self._query_runner.encrypted(),
                 arrow_disable_server_verification,
                 arrow_tls_root_certs,

graphdatascience/query_runner/arrow_query_runner.py

@@ -14,14 +14,15 @@
 from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
+from ..session.arrow_authentication import ArrowAuthentication
 
 
 class ArrowQueryRunner(QueryRunner):
     @staticmethod
     def create(
         fallback_query_runner: QueryRunner,
         arrow_info: ArrowInfo,
-        auth: Optional[tuple[str, str]] = None,
+        arrow_authentication: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
@@ -33,7 +34,7 @@ def create(
 
         gds_arrow_client = GdsArrowClient.create(
             arrow_info,
-            auth,
+            arrow_authentication,
             encrypted,
             disable_server_verification,
             tls_root_certs,

graphdatascience/query_runner/gds_arrow_client.py

@@ -39,6 +39,7 @@
 from graphdatascience.retry_utils.retry_utils import before_log
 
 from ..semantic_version.semantic_version import SemanticVersion
+from ..session.arrow_authentication import ArrowAuthentication
 from ..version import __version__
 from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_info import ArrowInfo
@@ -48,7 +49,7 @@ class GdsArrowClient:
     @staticmethod
     def create(
         arrow_info: ArrowInfo,
-        auth: Optional[tuple[str, str]] = None,
+        arrow_authentication: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
@@ -80,7 +81,7 @@ def create(
             host,
             retry_config,
             int(port),
-            auth,
+            arrow_authentication,
             encrypted,
             disable_server_verification,
             tls_root_certs,
@@ -92,7 +93,7 @@ def __init__(
         host: str,
         retry_config: RetryConfig,
         port: int = 8491,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
@@ -107,8 +108,8 @@ def __init__(
             The host address of the GDS Arrow server
         port: int
             The host port of the GDS Arrow server (default is 8491)
-        auth: Optional[tuple[str, str]]
-            A tuple containing the username and password for authentication
+        auth: Optional[ArrowAuthentication]
+            An implementation of ArrowAuthentication providing a pair to be used for basic authentication
         encrypted: bool
             A flag that indicates whether the connection should be encrypted (default is False)
         disable_server_verification: bool
@@ -189,7 +190,8 @@ def request_token(self) -> Optional[str]:
         def auth_with_retry() -> None:
             client = self._client()
             if self._auth:
-                client.authenticate_basic_token(self._auth[0], self._auth[1])
+                auth_pair = self._auth.auth_pair()
+                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
 
         if self._auth:
             auth_with_retry()
@@ -884,7 +886,7 @@ def start_call(self, info: Any) -> AuthMiddleware:
 
 
 class AuthMiddleware(ClientMiddleware):  # type: ignore
-    def __init__(self, auth: tuple[str, str], *args: Any, **kwargs: Any) -> None:
+    def __init__(self, auth: ArrowAuthentication, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._auth = auth
         self._token: Optional[str] = None
@@ -918,15 +920,15 @@ def received_headers(self, headers: dict[str, Any]) -> None:
 
     def sending_headers(self) -> dict[str, str]:
         token = self.token()
-        if not token:
-            username, password = self._auth
-            auth_token = f"{username}:{password}"
-            auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
-            # There seems to be a bug, `authorization` must be lower key
-            return {"authorization": auth_token}
-        else:
+        if token is not None:
             return {"authorization": "Bearer " + token}
 
+        auth_pair = self._auth.auth_pair()
+        auth_token = f"{auth_pair[0]}:{auth_pair[1]}"
+        auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
+        # There seems to be a bug, `authorization` must be lower key
+        return {"authorization": auth_token}
+
 
 @dataclass(repr=True, frozen=True)
 class NodeLoadDoneResult:

graphdatascience/session/arrow_authentication.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from typing import Callable
+
+from graphdatascience.session.aura_api import AuraApi
+
+
+class ArrowAuthentication(ABC):
+    type AuthTokenFn = Callable[[], str]
+
+    @abstractmethod
+    def auth_pair(self) -> tuple[str, str]:
+        """Returns the auth pair used for authentication."""
+        pass
+
+
+class UsernamePasswordAuthentication(ArrowAuthentication):
+    def __init__(self, username: str, password: str):
+        self._username = username
+        self._password = password
+
+    def auth_pair(self) -> tuple[str, str]:
+        return self._username, self._password
+
+
+class AuraApiTokenAuthentication(ArrowAuthentication):
+    def __init__(self, aura_api: AuraApi):
+        self._aura_api = aura_api
+
+    def auth_pair(self) -> tuple[str, str]:
+        return "", self._aura_api._auth._auth_token()

graphdatascience/session/aura_graph_data_science.py

@@ -19,6 +19,7 @@
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.query_runner.session_query_runner import SessionQueryRunner
 from graphdatascience.query_runner.standalone_session_query_runner import StandaloneSessionQueryRunner
+from graphdatascience.session.arrow_authentication import ArrowAuthentication
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.utils.util_remote_proc_runner import UtilRemoteProcRunner
 
@@ -32,7 +33,8 @@ class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
     @classmethod
     def create(
         cls,
-        gds_session_connection_info: DbmsConnectionInfo,
+        session_bolt_connection_info: DbmsConnectionInfo,
+        arrow_authentication: Optional[ArrowAuthentication],
         db_endpoint: Optional[Union[Neo4jQueryRunner, DbmsConnectionInfo]],
         delete_fn: Callable[[], bool],
         arrow_disable_server_verification: bool = False,
@@ -41,16 +43,16 @@ def create(
         show_progress: bool = True,
     ) -> AuraGraphDataScience:
         session_bolt_query_runner = Neo4jQueryRunner.create_for_session(
-            endpoint=gds_session_connection_info.uri,
-            auth=gds_session_connection_info.auth(),
+            endpoint=session_bolt_connection_info.uri,
+            auth=session_bolt_connection_info.auth(),
             show_progress=show_progress,
         )
 
         arrow_info = ArrowInfo.create(session_bolt_query_runner)
         session_arrow_query_runner = ArrowQueryRunner.create(
             fallback_query_runner=session_bolt_query_runner,
             arrow_info=arrow_info,
-            auth=gds_session_connection_info.auth(),
+            arrow_authentication=arrow_authentication,
             encrypted=session_bolt_query_runner.encrypted(),
             disable_server_verification=arrow_disable_server_verification,
             tls_root_certs=arrow_tls_root_certs,
@@ -59,7 +61,7 @@ def create(
         # TODO: merge with the gds_arrow_client created inside ArrowQueryRunner
         session_arrow_client = GdsArrowClient.create(
             arrow_info,
-            gds_session_connection_info.auth(),
+            arrow_authentication,
             session_bolt_query_runner.encrypted(),
             arrow_disable_server_verification,
             arrow_tls_root_certs,

graphdatascience/session/dedicated_sessions.py

@@ -7,6 +7,7 @@
 
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.session.algorithm_category import AlgorithmCategory
+from graphdatascience.session.arrow_authentication import AuraApiTokenAuthentication, ArrowAuthentication
 from graphdatascience.session.aura_api import AuraApi
 from graphdatascience.session.aura_api_responses import SessionDetails
 from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
@@ -88,15 +89,18 @@ def get_or_create(
 
         self._await_session_running(session_details, timeout)
 
-        session_connection = DbmsConnectionInfo(
+        session_bolt_connection_info = DbmsConnectionInfo(
             uri=session_details.bolt_connection_url(),
-            username="",
-            password=self._aura_api._auth._auth_token(),
+            username=self._aura_api._credentials[0],
+            password=self._aura_api._credentials[1],
         )
 
+        arrow_authentication = AuraApiTokenAuthentication(self._aura_api)
+
         return self._construct_client(
             session_id=session_details.id,
-            session_connection=session_connection,
+            session_bolt_connection_info=session_bolt_connection_info,
+            arrow_authentication=arrow_authentication,
             db_runner=db_runner,
         )
 
@@ -197,11 +201,13 @@ def _get_or_create_self_managed_session(
     def _construct_client(
         self,
         session_id: str,
-        session_connection: DbmsConnectionInfo,
+        session_bolt_connection_info: DbmsConnectionInfo,
+        arrow_authentication: ArrowAuthentication,
         db_runner: Optional[Neo4jQueryRunner],
     ) -> AuraGraphDataScience:
         return AuraGraphDataScience.create(
-            gds_session_connection_info=session_connection,
+            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),
         )

graphdatascience/tests/integration/conftest.py

@@ -9,6 +9,7 @@
 from graphdatascience.graph_data_science import GraphDataScience
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.server_version.server_version import ServerVersion
+from graphdatascience.session.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 
@@ -92,7 +93,8 @@ def gds_without_arrow() -> Generator[GraphDataScience, None, None]:
 @pytest.fixture(scope="package", autouse=False)
 def gds_with_cloud_setup(request: pytest.FixtureRequest) -> Generator[AuraGraphDataScience, None, None]:
     _gds = AuraGraphDataScience.create(
-        gds_session_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
+        session_bolt_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
+        arrow_authentication=UsernamePasswordAuthentication(AUTH[0], AUTH[1]),
         db_endpoint=DbmsConnectionInfo(AURA_DB_URI, AURA_DB_AUTH[0], AURA_DB_AUTH[1]),
         delete_fn=lambda: True,
     )
@@ -106,7 +108,8 @@ def gds_with_cloud_setup(request: pytest.FixtureRequest) -> Generator[AuraGraphD
 @pytest.fixture(scope="package", autouse=False)
 def standalone_aura_gds() -> Generator[AuraGraphDataScience, None, None]:
     _gds = AuraGraphDataScience.create(
-        gds_session_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
+        session_bolt_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
+        arrow_authentication=UsernamePasswordAuthentication(AUTH[0], AUTH[1]),
         db_endpoint=None,
         delete_fn=lambda: True,
     )

graphdatascience/tests/unit/conftest.py

@@ -16,6 +16,7 @@
 )
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
 from graphdatascience.server_version.server_version import ServerVersion
+from graphdatascience.session.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 
@@ -179,7 +180,8 @@ def aura_gds(runner: CollectingQueryRunner, mocker: MockerFixture) -> Generator[
     mocker.patch("graphdatascience.query_runner.gds_arrow_client.GdsArrowClient.create", return_value=None)
 
     aura_gds = AuraGraphDataScience.create(
-        gds_session_connection_info=DbmsConnectionInfo("address", "some", "auth"),
+        session_bolt_connection_info=DbmsConnectionInfo("address", "some", "auth"),
+        arrow_authentication=UsernamePasswordAuthentication("some", "auth"),
         db_endpoint=DbmsConnectionInfo("address", "some", "auth"),
         delete_fn=lambda: True,
     )

graphdatascience/tests/unit/test_aura_api.py

@@ -807,14 +807,14 @@ def test_auth_token(requests_mock: Mocker) -> None:
         json={"access_token": "very_short_token", "expires_in": 0, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth._auth_token() == "very_short_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "very_short_token"  # type: ignore
 
     requests_mock.post(
         "https://api.neo4j.io/oauth/token",
         json={"access_token": "longer_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth._auth_token() == "longer_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "longer_token"  # type: ignore
 
 
 def test_auth_token_reused(requests_mock: Mocker) -> None:
@@ -825,15 +825,15 @@ def test_auth_token_reused(requests_mock: Mocker) -> None:
         json={"access_token": "one_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
 
     requests_mock.post(
         "https://api.neo4j.io/oauth/token",
         json={"access_token": "new_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
     # no new token requested
-    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
 
 
 def test_auth_token_use_short_token(requests_mock: Mocker) -> None:
@@ -844,8 +844,8 @@ def test_auth_token_use_short_token(requests_mock: Mocker) -> None:
         json={"access_token": "one_token", "expires_in": 10, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
-    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
 
 
 def test_derive_tenant(requests_mock: Mocker) -> None:

graphdatascience/tests/unit/test_gds_arrow_client.py

@@ -16,6 +16,7 @@
 
 from graphdatascience.query_runner.arrow_info import ArrowInfo
 from graphdatascience.query_runner.gds_arrow_client import AuthMiddleware, GdsArrowClient
+from graphdatascience.session.session_connection_info import UsernamePasswordAuthentication
 
 ActionParam = Union[str, tuple[str, Any], Action]
 
@@ -382,7 +383,7 @@ def test_get_relationship_topologys(flight_server: FlightServer, flight_client:
 
 
 def test_auth_middleware() -> None:
-    middleware = AuthMiddleware(("user", "password"))
+    middleware = AuthMiddleware(UsernamePasswordAuthentication("user", "password"))
 
     first_header = middleware.sending_headers()
     assert first_header == {"authorization": "Basic dXNlcjpwYXNzd29yZA=="}
@@ -401,7 +402,7 @@ def test_auth_middleware() -> None:
 
 
 def test_auth_middleware_bad_headers() -> None:
-    middleware = AuthMiddleware(("user", "password"))
+    middleware = AuthMiddleware(UsernamePasswordAuthentication("user", "password"))
 
     with pytest.raises(ValueError, match="Incompatible header value received from server: `12342`"):
         middleware.received_headers({"authorization": [12342]})