graphdatascience 1.17a1__tar.gz → 1.18a1__tar.gz

{graphdatascience-1.17a1/graphdatascience.egg-info → graphdatascience-1.18a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.17a1
+Version: 1.18a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -16,7 +16,6 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -25,14 +24,14 @@ Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
 Classifier: Typing :: Typed
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<3.0,>=1.0
-Requires-Dist: neo4j<6.0,>=4.4.12
-Requires-Dist: numpy<2.3
+Requires-Dist: neo4j<7.0,>=4.4.12
+Requires-Dist: numpy<2.4
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<21.0,>=17.0
+Requires-Dist: pyarrow<22.0,>=17.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0

{graphdatascience-1.17a1 → graphdatascience-1.18a1}/graphdatascience/algo/algo_proc_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC
 from typing import Any
 
@@ -19,7 +21,7 @@ class AlgoProcRunner(IllegalAttrChecker, ABC):
         return self._query_runner.call_procedure(endpoint=self._namespace, params=params, logging=with_logging)
 
     @graph_type_check
-    def estimate(self, G: Graph, **config: Any) -> "Series[Any]":
+    def estimate(self, G: Graph, **config: Any) -> Series[Any]:
         self._namespace += "." + "estimate"
         return self._run_procedure(G, config, with_logging=False).squeeze()  # type: ignore
 
@@ -30,13 +32,13 @@ class StreamModeRunner(AlgoProcRunner):
 
 
 class StandardModeRunner(AlgoProcRunner):
-    def __call__(self, G: Graph, **config: Any) -> "Series[Any]":
+    def __call__(self, G: Graph, **config: Any) -> Series[Any]:
         return self._run_procedure(G, config).squeeze()  # type: ignore
 
 
 class GraphSageRunner(AlgoProcRunner):
     @graph_type_check
-    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, "Series[Any]"]:
+    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, Series[Any]]:
         result = self._run_procedure(G, config).squeeze()
         model_name = result["modelInfo"]["modelName"]
 

graphdatascience-1.18a1/graphdatascience/arrow_client/arrow_base_model.py

@@ -0,0 +1,12 @@
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic.alias_generators import to_camel
+
+
+class ArrowBaseModel(BaseModel, alias_generator=to_camel):
+    def dump_camel(self) -> dict[str, Any]:
+        return self.model_dump(by_alias=True)
+
+    def dump_json(self) -> str:
+        return self.model_dump_json(by_alias=True)

graphdatascience-1.18a1/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+from typing import Any, Iterator
+
+from pyarrow import __version__ as arrow_version
+from pyarrow import flight
+from pyarrow._flight import (
+    Action,
+    ActionType,
+    FlightInternalError,
+    FlightStreamReader,
+    FlightTimedOutError,
+    FlightUnavailableError,
+    Result,
+    Ticket,
+)
+from tenacity import retry, retry_any, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
+
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
+from graphdatascience.retry_utils.retry_config import RetryConfig
+
+from ..retry_utils.retry_utils import before_log
+from ..version import __version__
+from .middleware.auth_middleware import AuthFactory, AuthMiddleware
+from .middleware.user_agent_middleware import UserAgentFactory
+
+
+class AuthenticatedArrowClient:
+    @staticmethod
+    def create(
+        arrow_info: ArrowInfo,
+        auth: ArrowAuthentication | None = None,
+        encrypted: bool = False,
+        arrow_client_options: dict[str, Any] | None = None,
+        connection_string_override: str | None = None,
+        retry_config: RetryConfig | None = None,
+        advertised_listen_address: tuple[str, int] | None = None,
+    ) -> AuthenticatedArrowClient:
+        connection_string: str
+        if connection_string_override is not None:
+            connection_string = connection_string_override
+        else:
+            connection_string = arrow_info.listenAddress
+
+        host, port = connection_string.split(":")
+
+        if retry_config is None:
+            retry_config = RetryConfig(
+                retry=retry_any(
+                    retry_if_exception_type(FlightTimedOutError),
+                    retry_if_exception_type(FlightUnavailableError),
+                    retry_if_exception_type(FlightInternalError),
+                ),
+                stop=(stop_after_delay(10) | stop_after_attempt(5)),
+                wait=wait_exponential(multiplier=1, min=1, max=10),
+            )
+
+        return AuthenticatedArrowClient(
+            host=host,
+            retry_config=retry_config,
+            port=int(port),
+            auth=auth,
+            encrypted=encrypted,
+            arrow_client_options=arrow_client_options,
+            advertised_listen_address=advertised_listen_address,
+        )
+
+    def __init__(
+        self,
+        host: str,
+        retry_config: RetryConfig,
+        port: int = 8491,
+        auth: ArrowAuthentication | None = None,
+        encrypted: bool = False,
+        arrow_client_options: dict[str, Any] | None = None,
+        user_agent: str | None = None,
+        advertised_listen_address: tuple[str, int] | None = None,
+    ):
+        """Creates a new GdsArrowClient instance.
+
+        Parameters
+        ----------
+        host: str
+            The host address of the GDS Arrow server
+        port: int
+            The host port of the GDS Arrow server (default is 8491)
+        auth: ArrowAuthentication | None
+            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)
+        arrow_client_options: dict[str, Any] | None
+            Additional options to be passed to the Arrow Flight client.
+        user_agent: str | None
+            The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
+        retry_config: RetryConfig | None
+            The retry configuration to use for the Arrow requests send by the client.
+        advertised_listen_address: tuple[str, int] | None
+            The advertised listen address of the GDS Arrow server. This will be used by remote projection and writeback operations.
+        """
+        self._host = host
+        self._port = port
+        self._auth = None
+        self._encrypted = encrypted
+        self._arrow_client_options = arrow_client_options
+        self._user_agent = user_agent
+        self._retry_config = retry_config
+        self._logger = logging.getLogger("gds_arrow_client")
+        self._retry_config = retry_config
+        if auth:
+            self._auth = auth
+            self._auth_middleware = AuthMiddleware(auth)
+        self.advertised_listen_address = advertised_listen_address
+
+        self._flight_client = self._instantiate_flight_client()
+
+    def connection_info(self) -> ConnectionInfo:
+        """
+        Returns the host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        tuple[str, int]
+            the host and port of the GDS Arrow server
+        """
+        return ConnectionInfo(self._host, self._port, self._encrypted)
+
+    def advertised_connection_info(self) -> ConnectionInfo:
+        """
+        Returns the advertised host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        ConnectionInfo
+            the host and port of the GDS Arrow server
+        """
+        if self.advertised_listen_address is None:
+            return self.connection_info()
+
+        h, p = self.advertised_listen_address
+        return ConnectionInfo(h, p, self._encrypted)
+
+    def request_token(self) -> str | None:
+        """
+        Requests a token from the server and returns it.
+
+        Returns
+        -------
+        str | None
+            a token from the server and returns it.
+        """
+
+        @retry(
+            reraise=True,
+            before=before_log("Request token", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def auth_with_retry() -> None:
+            client = self._flight_client
+            if self._auth:
+                auth_pair = self._auth.auth_pair()
+                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+
+        if self._auth:
+            auth_with_retry()
+            return self._auth_middleware.token()
+        else:
+            return "IGNORED"
+
+    def get_stream(self, ticket: Ticket) -> FlightStreamReader:
+        return self._flight_client.do_get(ticket)
+
+    def do_action(self, endpoint: str, payload: bytes | dict[str, Any]) -> Iterator[Result]:
+        payload_bytes = payload if isinstance(payload, bytes) else json.dumps(payload).encode("utf-8")
+
+        return self._flight_client.do_action(Action(endpoint, payload_bytes))  # type: ignore
+
+    def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
+        @retry(
+            reraise=True,
+            before=before_log("Send action", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def run_with_retry() -> list[Result]:
+            # the Flight response error code is only checked on iterator consumption
+            # we eagerly collect iterator here to trigger retry in case of an error
+            return list(self.do_action(endpoint, payload))
+
+        return run_with_retry()
+
+    def list_actions(self) -> set[ActionType]:
+        return self._flight_client.list_actions()  # type: ignore
+
+    def _instantiate_flight_client(self) -> flight.FlightClient:
+        location = (
+            flight.Location.for_grpc_tls(self._host, self._port)
+            if self._encrypted
+            else flight.Location.for_grpc_tcp(self._host, self._port)
+        )
+        client_options: dict[str, Any] = (self._arrow_client_options or {}).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),
+            ]
+
+        return flight.FlightClient(location, **client_options)
+
+
+@dataclass
+class ConnectionInfo:
+    host: str
+    port: int
+    encrypted: bool

graphdatascience-1.18a1/graphdatascience/arrow_client/middleware/auth_middleware.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import base64
+import time
+from typing import 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: str | None = None
+        self._token_timestamp = 0
+
+    def token(self) -> str | None:
+        # 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-1.18a1/graphdatascience/arrow_client/middleware/user_agent_middleware.py

@@ -0,0 +1,26 @@
+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-1.18a1/graphdatascience/arrow_client/v2/api_types.py

@@ -0,0 +1,49 @@
+from graphdatascience.arrow_client.arrow_base_model import ArrowBaseModel
+
+
+class JobIdConfig(ArrowBaseModel):
+    job_id: str
+
+
+UNKNOWN_PROGRESS = -1
+
+
+class JobStatus(ArrowBaseModel):
+    job_id: str
+    status: str
+    progress: float
+    description: str
+
+    def progress_known(self) -> bool:
+        if self.progress == UNKNOWN_PROGRESS:
+            return False
+        return True
+
+    def progress_percent(self) -> float | None:
+        if self.progress_known():
+            return self.progress * 100
+        return None
+
+    def base_task(self) -> str:
+        return self.description.split("::")[0].strip()
+
+    def sub_tasks(self) -> str | None:
+        task_split = self.description.split("::", maxsplit=1)
+        if len(task_split) > 1:
+            return task_split[1].strip()
+        return None
+
+    def aborted(self) -> bool:
+        return self.status == "Aborted"
+
+    def succeeded(self) -> bool:
+        return self.status == "Done"
+
+    def running(self) -> bool:
+        return self.status == "Running"
+
+
+class MutateResult(ArrowBaseModel):
+    mutate_millis: int
+    node_properties_written: int
+    relationships_written: int

graphdatascience-1.18a1/graphdatascience/arrow_client/v2/data_mapper_utils.py

@@ -0,0 +1,19 @@
+import json
+from typing import Any
+
+from pyarrow._flight import Result
+
+
+def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:
+    rows = deserialize(input_stream)
+    if len(rows) != 1:
+        raise ValueError(f"Expected exactly one result, got {len(rows)}")
+
+    return rows[0]
+
+
+def deserialize(input_stream: list[Result]) -> list[dict[str, Any]]:
+    def deserialize_row(row: Result):  # type:ignore
+        return json.loads(row.body.to_pybytes().decode())
+
+    return [deserialize_row(row) for row in input_stream]

graphdatascience-1.18a1/graphdatascience/arrow_client/v2/job_client.py

@@ -0,0 +1,94 @@
+import json
+from typing import Any
+
+from pandas import ArrowDtype, DataFrame
+from pyarrow._flight import Ticket
+from tenacity import Retrying, retry_if_result, wait_exponential
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
+from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
+from graphdatascience.query_runner.termination_flag import TerminationFlag
+
+JOB_STATUS_ENDPOINT = "v2/jobs.status"
+RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
+
+
+class JobClient:
+    def __init__(self, progress_bar_options: dict[str, Any] | None = None):
+        self._progress_bar_options = progress_bar_options or {}
+
+    @staticmethod
+    def run_job_and_wait(
+        client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any], show_progress: bool
+    ) -> str:
+        job_id = JobClient.run_job(client, endpoint, config)
+        JobClient().wait_for_job(client, job_id, show_progress=show_progress)
+        return job_id
+
+    @staticmethod
+    def run_job(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, Any]) -> str:
+        res = client.do_action_with_retry(endpoint, config)
+
+        single = deserialize_single(res)
+        return JobIdConfig(**single).job_id
+
+    def wait_for_job(
+        self,
+        client: AuthenticatedArrowClient,
+        job_id: str,
+        show_progress: bool,
+        termination_flag: TerminationFlag | None = None,
+    ) -> None:
+        progress_bar: TqdmProgressBar | None = None
+
+        if termination_flag is None:
+            termination_flag = TerminationFlag.create()
+
+        for attempt in Retrying(retry=retry_if_result(lambda _: True), wait=wait_exponential(min=0.1, max=5)):
+            with attempt:
+                termination_flag.assert_running()
+
+                arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
+                job_status = JobStatus(**deserialize_single(arrow_res))
+
+                if job_status.succeeded() or job_status.aborted():
+                    if progress_bar:
+                        progress_bar.finish(success=job_status.succeeded())
+                    return
+
+                if show_progress:
+                    if progress_bar is None:
+                        base_task = job_status.base_task()
+                        if base_task:
+                            progress_bar = TqdmProgressBar(
+                                task_name=base_task,
+                                relative_progress=job_status.progress_percent(),
+                                bar_options=self._progress_bar_options,
+                            )
+                    if progress_bar:
+                        progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
+
+    @staticmethod
+    def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:
+        res = client.do_action_with_retry(RESULTS_SUMMARY_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
+        return deserialize_single(res)
+
+    @staticmethod
+    def stream_results(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> DataFrame:
+        payload = {
+            "graphName": graph_name,
+            "jobId": job_id,
+        }
+
+        res = client.do_action_with_retry("v2/results.stream", payload)
+        export_job_id = JobIdConfig(**deserialize_single(res)).job_id
+
+        stream_payload = {"version": "v2", "name": export_job_id, "body": {}}
+
+        ticket = Ticket(json.dumps(stream_payload).encode("utf-8"))
+
+        get = client.get_stream(ticket)
+        arrow_table = get.read_all()
+        return arrow_table.to_pandas(types_mapper=ArrowDtype)  # type: ignore

graphdatascience-1.18a1/graphdatascience/arrow_client/v2/mutation_client.py

@@ -0,0 +1,52 @@
+import math
+import time
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.api_types import MutateResult
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
+
+
+class MutationClient:
+    MUTATE_ENDPOINT = "v2/results.mutate"
+
+    @staticmethod
+    def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
+        return MutationClient._mutate(
+            client=client,
+            job_id=job_id,
+            mutate_property=mutate_property,
+        )
+
+    @staticmethod
+    def mutate_relationship_property(
+        client: AuthenticatedArrowClient,
+        job_id: str,
+        mutate_relationship_type: str,
+        mutate_property: str | None,
+    ) -> MutateResult:
+        return MutationClient._mutate(
+            client=client,
+            job_id=job_id,
+            mutate_property=mutate_property,
+            mutate_relationship_type=mutate_relationship_type,
+        )
+
+    @staticmethod
+    def _mutate(
+        client: AuthenticatedArrowClient,
+        job_id: str,
+        mutate_property: str | None = None,
+        mutate_relationship_type: str | None = None,
+    ) -> MutateResult:
+        mutate_config = {
+            "jobId": job_id,
+        }
+        if mutate_relationship_type:
+            mutate_config["mutateRelationshipType"] = mutate_relationship_type
+        if mutate_property:
+            mutate_config["mutateProperty"] = mutate_property
+
+        start_time = time.time()
+        mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, mutate_config)
+        mutate_millis = math.ceil((time.time() - start_time) * 1000)
+        return MutateResult(mutateMillis=mutate_millis, **deserialize_single(mutate_arrow_res))