graphdatascience 1.10__tar.gz → 1.11a1__tar.gz

{graphdatascience-1.10/graphdatascience.egg-info → graphdatascience-1.11a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.10
+Version: 1.11a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j

{graphdatascience-1.10 → graphdatascience-1.11a1}/graphdatascience/graph/base_graph_proc_runner.py

@@ -517,8 +517,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         ).squeeze()
 
     @multimethod
-    def removeNodeProperties(self) -> None:
-        ...
+    def removeNodeProperties(self) -> None: ...
 
     @removeNodeProperties.register
     @graph_type_check

{graphdatascience-1.10 → graphdatascience-1.11a1}/graphdatascience/graph_data_science.py

@@ -44,20 +44,20 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
             A username, password pair for database authentication.
         aura_ds : bool, default False
             A flag that indicates that that the client is used to connect
-            to a Neo4j Aura instance.
+            to a Neo4j AuraDS instance.
         database: Optional[str], default None
             The Neo4j database to query against.
         arrow : Union[str, bool], default True
-            Arrow connection information. Either a flag that indicates whether the client should use Apache Arrow
-            for data streaming if it is available on the server. True means discover the connection URI from the server.
-            A connection URI (str) can also be provided.
+            Arrow connection information. This is either a bool or a string.
+            If it is a string, it will be interpreted as a connection URL to a GDS Arrow Server.
+            If it is a bool,
+                True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint,
+                while False will make the client use Bolt for all operations.
         arrow_disable_server_verification : bool, default True
-            A flag that indicates that, if the flight client is connecting with
-            TLS, that it skips server verification. If this is enabled, all
-            other TLS settings are overridden.
+            A flag that overrides other TLS settings and disables server verification for TLS connections.
         arrow_tls_root_certs : Optional[bytes], default None
-            PEM-encoded certificates that are used for the connecting to the
-            Arrow Flight server.
+            PEM-encoded certificates that are used for the connection to the
+            GDS Arrow Flight server.
         bookmarks : Optional[Any], default None
             The Neo4j bookmarks to require a certain state before the next query gets executed.
         """

{graphdatascience-1.10 → graphdatascience-1.11a1}/graphdatascience/query_runner/arrow_query_runner.py

@@ -282,7 +282,7 @@ class ArrowQueryRunner(QueryRunner):
 
         if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
             payload = {
-                "name": "GET_MESSAGE",
+                "name": "GET_COMMAND",
                 "version": ArrowEndpointVersion.V1.version(),
                 "body": payload,
             }
@@ -377,10 +377,17 @@ class AuthMiddleware(ClientMiddleware):  # type: ignore
         self._factory = factory
 
     def received_headers(self, headers: Dict[str, Any]) -> None:
-        auth_header: str = headers.get("Authorization", None)
+        auth_header = headers.get("authorization", None)
         if not auth_header:
             return
-        [auth_type, token] = auth_header.split(" ", 1)
+
+        # 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._factory.set_token(token)
 

graphdatascience-1.11a1/graphdatascience/session/__init__.py

@@ -0,0 +1,16 @@
+from .algorithm_category import AlgorithmCategory
+from .dbms_connection_info import DbmsConnectionInfo
+from .gds_sessions import AuraAPICredentials, GdsSessions
+from .schema import GdsPropertyTypes
+from .session_info import SessionInfo
+from .session_sizes import SessionMemory
+
+__all__ = [
+    "GdsSessions",
+    "SessionInfo",
+    "DbmsConnectionInfo",
+    "AuraAPICredentials",
+    "SessionMemory",
+    "GdsPropertyTypes",
+    "AlgorithmCategory",
+]

graphdatascience-1.11a1/graphdatascience/session/algorithm_category.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class AlgorithmCategory(Enum):
+    """
+    Enumeration of supported algorithm categories used for size estimation.
+    """
+
+    CENTRALITY = "centrality"
+    COMMUNITY_DETECTION = "community-detection"
+    MACHINE_LEARNING = "machine-learning"
+    NODE_EMBEDDING = "node-embedding"
+    PATH_FINDING = "path-finding"
+    SIMILARITY = "similarity"

{graphdatascience-1.10 → graphdatascience-1.11a1}/graphdatascience/session/aura_api.py

@@ -1,121 +1,27 @@
 from __future__ import annotations
 
-import dataclasses
 import logging
 import os
 import time
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, List, NamedTuple, Optional, Set
+from typing import Any, List, Optional
 from urllib.parse import urlparse
 
 import requests as req
 from requests import HTTPError
 
+from graphdatascience.session.algorithm_category import AlgorithmCategory
+from graphdatascience.session.aura_api_responses import (
+    EstimationDetails,
+    InstanceCreateDetails,
+    InstanceDetails,
+    InstanceSpecificDetails,
+    SessionDetails,
+    TenantDetails,
+    WaitResult,
+)
 from graphdatascience.version import __version__
 
 
-@dataclass(repr=True, frozen=True)
-class InstanceDetails:
-    id: str
-    name: str
-    tenant_id: str
-    cloud_provider: str
-
-    @classmethod
-    def fromJson(cls, json: dict[str, Any]) -> InstanceDetails:
-        return cls(
-            id=json["id"],
-            name=json["name"],
-            tenant_id=json["tenant_id"],
-            cloud_provider=json["cloud_provider"],
-        )
-
-
-@dataclass(repr=True, frozen=True)
-class InstanceSpecificDetails(InstanceDetails):
-    status: str
-    connection_url: str
-    memory: str
-    type: str
-    region: str
-
-    @classmethod
-    def fromJson(cls, json: dict[str, Any]) -> InstanceSpecificDetails:
-        return cls(
-            id=json["id"],
-            name=json["name"],
-            tenant_id=json["tenant_id"],
-            cloud_provider=json["cloud_provider"],
-            status=json["status"],
-            connection_url=json.get("connection_url", ""),
-            memory=json.get("memory", ""),
-            type=json["type"],
-            region=json["region"],
-        )
-
-
-@dataclass(repr=True, frozen=True)
-class InstanceCreateDetails:
-    id: str
-    username: str
-    password: str
-    connection_url: str
-
-    @classmethod
-    def from_json(cls, json: dict[str, Any]) -> InstanceCreateDetails:
-        fields = dataclasses.fields(cls)
-        if any(f.name not in json for f in fields):
-            raise RuntimeError(f"Missing required field. Expected `{[f.name for f in fields]}` but got `{json}`")
-
-        return cls(**{f.name: json[f.name] for f in fields})
-
-
-class WaitResult(NamedTuple):
-    connection_url: str
-    error: str
-
-    @classmethod
-    def from_error(cls, error: str) -> WaitResult:
-        return cls(connection_url="", error=error)
-
-    @classmethod
-    def from_connection_url(cls, connection_url: str) -> WaitResult:
-        return cls(connection_url=connection_url, error="")
-
-
-@dataclass(repr=True, frozen=True)
-class TenantDetails:
-    id: str
-    ds_type: str
-    regions_per_provider: dict[str, Set[str]]
-
-    @classmethod
-    def from_json(cls, json: dict[str, Any]) -> TenantDetails:
-        regions_per_provider = defaultdict(set)
-        instance_types = set()
-        ds_type = None
-
-        for configs in json["instance_configurations"]:
-            type = configs["type"]
-            if type.split("-")[1] == "ds":
-                regions_per_provider[configs["cloud_provider"]].add(configs["region"])
-                ds_type = type
-            instance_types.add(configs["type"])
-
-        id = json["id"]
-        if not ds_type:
-            raise RuntimeError(
-                f"Tenant with id `{id}` cannot create DS instances. Available instances are `{instance_types}`."
-            )
-
-        return cls(
-            id=id,
-            ds_type=ds_type,
-            regions_per_provider=regions_per_provider,
-        )
-
-
 class AuraApi:
     class AuraAuthToken:
         access_token: str
@@ -156,6 +62,79 @@ class AuraApi:
 
         return host.split(".")[0].split("-")[0]
 
+    def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> SessionDetails:
+        response = req.post(
+            f"{self._base_uri}/v1beta5/data-science/sessions",
+            headers=self._build_header(),
+            json={"name": name, "instance_id": dbid, "password": pwd, "memory": memory},
+        )
+
+        response.raise_for_status()
+
+        return SessionDetails.fromJson(response.json())
+
+    def list_session(self, session_id: str, dbid: str) -> Optional[SessionDetails]:
+        response = req.get(
+            f"{self._base_uri}/v1beta5/data-science/sessions/{session_id}?instanceId={dbid}",
+            headers=self._build_header(),
+        )
+
+        if response.status_code == 404:
+            return None
+
+        response.raise_for_status()
+
+        return SessionDetails.fromJson(response.json())
+
+    def list_sessions(self, dbid: str) -> List[SessionDetails]:
+        response = req.get(
+            f"{self._base_uri}/v1beta5/data-science/sessions?instanceId={dbid}",
+            headers=self._build_header(),
+        )
+
+        response.raise_for_status()
+
+        return [SessionDetails.fromJson(s) for s in response.json()]
+
+    def wait_for_session_running(
+        self, session_id: str, dbid: str, sleep_time: float = 0.2, max_sleep_time: float = 300
+    ) -> WaitResult:
+        waited_time = 0.0
+        while waited_time <= max_sleep_time:
+            session = self.list_session(session_id, dbid)
+            if session is None:
+                return WaitResult.from_error(f"Session `{session_id}` for database `{dbid}` not found -- please retry")
+            elif session.status == "Ready" and session.host:  # check host needed until dns based routing
+                return WaitResult.from_connection_url(session.bolt_connection_url())
+            else:
+                self._logger.debug(
+                    f"Session `{session_id}` is not yet running. "
+                    f"Current status: {session.status} Host: {session.host}. "
+                    f"Retrying in {sleep_time} seconds..."
+                )
+            waited_time += sleep_time
+            time.sleep(sleep_time)
+
+        return WaitResult.from_error(
+            f"Session `{session_id}` for database `{dbid}` is not running after {waited_time} seconds"
+        )
+
+    def delete_session(self, session_id: str, dbid: str) -> bool:
+        response = req.delete(
+            f"{self._base_uri}/v1beta5/data-science/sessions/{session_id}",
+            headers=self._build_header(),
+            json={"instance_id": dbid},
+        )
+
+        if response.status_code == 404:
+            return False
+        elif response.status_code == 202:
+            return True
+
+        response.raise_for_status()
+
+        return False
+
     def create_instance(self, name: str, memory: str, cloud_provider: str, region: str) -> InstanceCreateDetails:
         tenant_details = self.tenant_details()
 
@@ -247,6 +226,21 @@ class AuraApi:
 
         return WaitResult.from_error(f"Instance is not running after waiting for {waited_time} seconds")
 
+    def estimate_size(
+        self, node_count: int, relationship_count: int, algorithm_categories: List[AlgorithmCategory]
+    ) -> EstimationDetails:
+        data = {
+            "node_count": node_count,
+            "relationship_count": relationship_count,
+            "algorithm_categories": [i.value for i in algorithm_categories],
+            "instance_type": "dsenterprise",
+        }
+
+        response = req.post(f"{self._base_uri}/v1/instances/sizing", headers=self._build_header(), json=data)
+        response.raise_for_status()
+
+        return EstimationDetails.from_json(response.json()["data"])
+
     def _get_tenant_id(self) -> str:
         response = req.get(
             f"{self._base_uri}/v1/tenants",
@@ -257,8 +251,9 @@ class AuraApi:
         raw_data = response.json()["data"]
 
         if len(raw_data) != 1:
+            tenants_dict = {d["id"]: d["name"] for d in raw_data}
             raise RuntimeError(
-                f"This account has access to multiple tenants `{raw_data}`. Please specify which one to use."
+                f"This account has access to multiple tenants: `{tenants_dict}`. Please specify which one to use."
             )
 
         return raw_data[0]["id"]  # type: ignore

graphdatascience-1.11a1/graphdatascience/session/aura_api_responses.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import dataclasses
+import sys
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, NamedTuple, Optional, Set
+
+
+@dataclass(repr=True, frozen=True)
+class SessionDetails:
+    id: str
+    name: str
+    instance_id: str
+    memory: str
+    status: str
+    host: Optional[str]
+    expiry_date: Optional[datetime]
+    created_at: datetime
+
+    @classmethod
+    def fromJson(cls, json: dict[str, Any]) -> SessionDetails:
+        expiry_date = json.get("expiry_date")
+
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            instance_id=json["instance_id"],
+            memory=json["memory"],
+            status=json["status"],
+            host=json.get("host"),
+            expiry_date=TimeParser.fromisoformat(expiry_date) if expiry_date else None,
+            created_at=TimeParser.fromisoformat(json["created_at"]),
+        )
+
+    def bolt_connection_url(self) -> str:
+        return f"neo4j+ssc://{self.host}"  # TODO use neo4j+s
+
+
+@dataclass(repr=True, frozen=True)
+class InstanceDetails:
+    id: str
+    name: str
+    tenant_id: str
+    cloud_provider: str
+
+    @classmethod
+    def fromJson(cls, json: dict[str, Any]) -> InstanceDetails:
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            tenant_id=json["tenant_id"],
+            cloud_provider=json["cloud_provider"],
+        )
+
+
+@dataclass(repr=True, frozen=True)
+class InstanceSpecificDetails(InstanceDetails):
+    status: str
+    connection_url: str
+    memory: str
+    type: str
+    region: str
+
+    @classmethod
+    def fromJson(cls, json: dict[str, Any]) -> InstanceSpecificDetails:
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            tenant_id=json["tenant_id"],
+            cloud_provider=json["cloud_provider"],
+            status=json["status"],
+            connection_url=json.get("connection_url", ""),
+            memory=json.get("memory", ""),
+            type=json["type"],
+            region=json["region"],
+        )
+
+
+@dataclass(repr=True, frozen=True)
+class InstanceCreateDetails:
+    id: str
+    username: str
+    password: str
+    connection_url: str
+
+    @classmethod
+    def from_json(cls, json: dict[str, Any]) -> InstanceCreateDetails:
+        fields = dataclasses.fields(cls)
+        if any(f.name not in json for f in fields):
+            raise RuntimeError(f"Missing required field. Expected `{[f.name for f in fields]}` but got `{json}`")
+
+        return cls(**{f.name: json[f.name] for f in fields})
+
+
+@dataclass(repr=True, frozen=True)
+class EstimationDetails:
+    min_required_memory: str
+    recommended_size: str
+    did_exceed_maximum: bool
+
+    @classmethod
+    def from_json(cls, json: dict[str, Any]) -> EstimationDetails:
+        fields = dataclasses.fields(cls)
+        if any(f.name not in json for f in fields):
+            raise RuntimeError(f"Missing required field. Expected `{[f.name for f in fields]}` but got `{json}`")
+
+        return cls(**{f.name: json[f.name] for f in fields})
+
+
+class WaitResult(NamedTuple):
+    connection_url: str
+    error: str
+
+    @classmethod
+    def from_error(cls, error: str) -> WaitResult:
+        return cls(connection_url="", error=error)
+
+    @classmethod
+    def from_connection_url(cls, connection_url: str) -> WaitResult:
+        return cls(connection_url=connection_url, error="")
+
+
+@dataclass(repr=True, frozen=True)
+class TenantDetails:
+    id: str
+    ds_type: str
+    regions_per_provider: dict[str, Set[str]]
+
+    @classmethod
+    def from_json(cls, json: dict[str, Any]) -> TenantDetails:
+        regions_per_provider = defaultdict(set)
+        instance_types = set()
+        ds_type = None
+
+        for configs in json["instance_configurations"]:
+            type = configs["type"]
+            if type.split("-")[1] == "ds":
+                regions_per_provider[configs["cloud_provider"]].add(configs["region"])
+                ds_type = type
+            instance_types.add(configs["type"])
+
+        id = json["id"]
+        if not ds_type:
+            raise RuntimeError(
+                f"Tenant with id `{id}` cannot create DS instances. Available instances are `{instance_types}`."
+            )
+
+        return cls(
+            id=id,
+            ds_type=ds_type,
+            regions_per_provider=regions_per_provider,
+        )
+
+
+# datetime.fromisoformat only works with Python version > 3.9
+class TimeParser:
+
+    @staticmethod
+    def fromisoformat(date: str) -> datetime:
+        if sys.version_info >= (3, 11):
+            return datetime.fromisoformat(date)
+        else:
+            # Aura API example: 1970-01-01T00:00:00Z
+            return datetime.strptime(date, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)