Move arrow client related code into its own package

Author: DarthMax

graphdatascience/arrow_client/__init__.py

@@ -2,8 +2,8 @@
 
 from dataclasses import dataclass
 
-from ..query_runner.query_runner import QueryRunner
-from ..server_version.server_version import ServerVersion
+from graphdatascience.query_runner.query_runner import QueryRunner
+from graphdatascience.server_version.server_version import ServerVersion
 
 
 @dataclass(frozen=True)

graphdatascience/query_runner/arrow_authentication.py renamed to graphdatascience/arrow_client/arrow_authentication.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import base64
+import time
+from typing import Optional, Any
+
+from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
+
+
+class AuthFactory(ClientMiddlewareFactory):  # type: ignore
+    def __init__(self, middleware: AuthMiddleware, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._middleware = middleware
+
+    def start_call(self, info: Any) -> AuthMiddleware:
+        return self._middleware
+
+
+class AuthMiddleware(ClientMiddleware):  # type: ignore
+    def __init__(self, auth: ArrowAuthentication, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._auth = auth
+        self._token: Optional[str] = None
+        self._token_timestamp = 0
+
+    def token(self) -> Optional[str]:
+        # check whether the token is older than 10 minutes. If so, reset it.
+        if self._token and int(time.time()) - self._token_timestamp > 600:
+            self._token = None
+
+        return self._token
+
+    def _set_token(self, token: str) -> None:
+        self._token = token
+        self._token_timestamp = int(time.time())
+
+    def received_headers(self, headers: dict[str, Any]) -> None:
+        auth_header = headers.get("authorization", None)
+        if not auth_header:
+            return
+
+        # the result is always a list
+        header_value = auth_header[0]
+
+        if not isinstance(header_value, str):
+            raise ValueError(f"Incompatible header value received from server: `{header_value}`")
+
+        auth_type, token = header_value.split(" ", 1)
+        if auth_type == "Bearer":
+            self._set_token(token)
+
+    def sending_headers(self) -> dict[str, str]:
+        token = self.token()
+        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}

graphdatascience/query_runner/arrow_info.py renamed to graphdatascience/arrow_client/arrow_info.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+
+class UserAgentFactory(ClientMiddlewareFactory):  # type: ignore
+    def __init__(self, useragent: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._middleware = UserAgentMiddleware(useragent)
+
+    def start_call(self, info: Any) -> ClientMiddleware:
+        return self._middleware
+
+class UserAgentMiddleware(ClientMiddleware):  # type: ignore
+    def __init__(self, useragent: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._useragent = useragent
+
+    def sending_headers(self) -> dict[str, str]:
+        return {"x-gds-user-agent": self._useragent}
+
+    def received_headers(self, headers: dict[str, Any]) -> None:
+        pass

graphdatascience/arrow_client/middleware/AuthMiddleware.py

@@ -8,15 +8,15 @@
 from neo4j import Driver
 from pandas import DataFrame
 
-from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
+from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.procedure_surface.api.wcc_endpoints import WccEndpoints
 from graphdatascience.procedure_surface.cypher.wcc_proc_runner import WccCypherEndpoints
 
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
 from .error.uncallable_namespace import UncallableNamespace
 from .graph.graph_proc_runner import GraphProcRunner
-from .query_runner.arrow_info import ArrowInfo
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from .query_runner.arrow_query_runner import ArrowQueryRunner
 from .query_runner.neo4j_query_runner import Neo4jQueryRunner
 from .query_runner.query_runner import QueryRunner

graphdatascience/arrow_client/middleware/UserAgentMiddleware.py

@@ -5,12 +5,12 @@
 
 from pandas import DataFrame
 
-from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.retry_utils.retry_config import RetryConfig
 
 from ..call_parameters import CallParameters
-from ..query_runner.arrow_info import ArrowInfo
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from ..server_version.server_version import ServerVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
 from .gds_arrow_client import GdsArrowClient

graphdatascience/graph_data_science.py

@@ -35,14 +35,14 @@
     wait_exponential,
 )
 
-from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication, UsernamePasswordAuthentication
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication, UsernamePasswordAuthentication
 from graphdatascience.retry_utils.retry_config import RetryConfig
 from graphdatascience.retry_utils.retry_utils import before_log
 
 from ..semantic_version.semantic_version import SemanticVersion
 from ..version import __version__
 from .arrow_endpoint_version import ArrowEndpointVersion
-from .arrow_info import ArrowInfo
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
 
 
 class GdsArrowClient:

graphdatascience/query_runner/arrow_query_runner.py

@@ -1,4 +1,4 @@
-from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.session.aura_api import AuraApi
 
 

graphdatascience/query_runner/gds_arrow_client.py

@@ -13,8 +13,8 @@
 )
 from graphdatascience.error.uncallable_namespace import UncallableNamespace
 from graphdatascience.graph.graph_remote_proc_runner import GraphRemoteProcRunner
-from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
-from graphdatascience.query_runner.arrow_info import ArrowInfo
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
 from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner