graphdatascience 1.14__tar.gz → 1.15a1__tar.gz

{graphdatascience-1.14/graphdatascience.egg-info → graphdatascience-1.15a1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.14
+Version: 1.15a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -51,6 +51,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: project-url
 Dynamic: provides-extra
 Dynamic: requires-dist

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/query_runner/arrow_query_runner.py

@@ -5,6 +5,8 @@ from typing import Any, Optional
 
 from pandas import DataFrame
 
+from graphdatascience.retry_utils.retry_config import RetryConfig
+
 from ..call_parameters import CallParameters
 from ..query_runner.arrow_info import ArrowInfo
 from ..server_version.server_version import ServerVersion
@@ -24,6 +26,7 @@ class ArrowQueryRunner(QueryRunner):
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
+        retry_config: Optional[RetryConfig] = None,
     ) -> ArrowQueryRunner:
         if not arrow_info.enabled:
             raise ValueError("Arrow is not enabled on the server")
@@ -35,6 +38,7 @@ class ArrowQueryRunner(QueryRunner):
             disable_server_verification,
             tls_root_certs,
             connection_string_override,
+            retry_config=retry_config,
         )
 
         return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/query_runner/gds_arrow_client.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import base64
 import json
+import logging
 import re
 import time
 import warnings
@@ -12,11 +13,30 @@ from typing import Any, Callable, Dict, Iterable, Optional, Type, Union
 import pandas
 import pyarrow
 from neo4j.exceptions import ClientError
-from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Table, chunked_array, flight
+from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Schema, Table, chunked_array, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import (
+    ClientMiddleware,
+    ClientMiddlewareFactory,
+    FlightDescriptor,
+    FlightInternalError,
+    FlightMetadataReader,
+    FlightStreamWriter,
+    FlightTimedOutError,
+    FlightUnavailableError,
+)
 from pyarrow.types import is_dictionary
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
+from tenacity import (
+    retry,
+    retry_any,
+    retry_if_exception_type,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_exponential,
+)
+
+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__
@@ -33,6 +53,7 @@ class GdsArrowClient:
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
+        retry_config: Optional[RetryConfig] = None,
     ) -> GdsArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -44,8 +65,20 @@ class GdsArrowClient:
 
         arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.versions)
 
+        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 GdsArrowClient(
             host,
+            retry_config,
             int(port),
             auth,
             encrypted,
@@ -57,6 +90,7 @@ class GdsArrowClient:
     def __init__(
         self,
         host: str,
+        retry_config: RetryConfig,
         port: int = 8491,
         auth: Optional[tuple[str, str]] = None,
         encrypted: bool = False,
@@ -85,6 +119,8 @@ class GdsArrowClient:
             The version of the Arrow endpoint to use (default is ArrowEndpointVersion.V1)
         user_agent: Optional[str]
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
+        retry_config: Optional[RetryConfig]
+            The retry configuration to use for the Arrow requests send by the client.
         """
         self._arrow_endpoint_version = arrow_endpoint_version
         self._host = host
@@ -94,6 +130,8 @@ class GdsArrowClient:
         self._disable_server_verification = disable_server_verification
         self._tls_root_certs = tls_root_certs
         self._user_agent = user_agent
+        self._retry_config = retry_config
+        self._logger = logging.getLogger("gds_arrow_client")
 
         if auth:
             self._auth_middleware = AuthMiddleware(auth)
@@ -140,9 +178,21 @@ class GdsArrowClient:
         Optional[str]
             a token from the server and returns it.
         """
-        if self._auth:
+
+        @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._client()
-            client.authenticate_basic_token(self._auth[0], self._auth[1])
+            if self._auth:
+                client.authenticate_basic_token(self._auth[0], self._auth[1])
+
+        if self._auth:
+            auth_with_retry()
             return self._auth_middleware.token()
         else:
             return "IGNORED"
@@ -193,7 +243,7 @@ class GdsArrowClient:
         if node_labels:
             config["node_labels"] = node_labels
 
-        return self._do_get(database, graph_name, proc, concurrency, config)
+        return self._do_get_with_retry(database, graph_name, proc, concurrency, config)
 
     def get_node_labels(self, graph_name: str, database: str, concurrency: Optional[int] = None) -> pandas.DataFrame:
         """
@@ -213,7 +263,7 @@ class GdsArrowClient:
         DataFrame
             The requested nodes as a DataFrame
         """
-        return self._do_get(database, graph_name, "gds.graph.nodeLabels.stream", concurrency, {})
+        return self._do_get_with_retry(database, graph_name, "gds.graph.nodeLabels.stream", concurrency, {})
 
     def get_relationships(
         self, graph_name: str, database: str, relationship_types: list[str], concurrency: Optional[int] = None
@@ -237,7 +287,7 @@ class GdsArrowClient:
         DataFrame
             The requested relationships as a DataFrame
         """
-        return self._do_get(
+        return self._do_get_with_retry(
             database,
             graph_name,
             "gds.graph.relationships.stream",
@@ -285,7 +335,7 @@ class GdsArrowClient:
         if relationship_types:
             config["relationship_types"] = relationship_types
 
-        return self._do_get(database, graph_name, proc, concurrency, config)
+        return self._do_get_with_retry(database, graph_name, proc, concurrency, config)
 
     def create_graph(
         self,
@@ -573,19 +623,29 @@ class GdsArrowClient:
 
     def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str, Any]:
         action_type = self._versioned_action_type(action_type)
+        client = self._client()
 
-        try:
-            client = self._client()
-            result = client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
+        @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 send_with_retry() -> dict[str, Any]:
+            try:
+                result = client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
 
-            # Consume result fully to sanity check and avoid cancelled streams
-            collected_result = list(result)
-            assert len(collected_result) == 1
+                # Consume result fully to sanity check and avoid cancelled streams
+                collected_result = list(result)
+                assert len(collected_result) == 1
 
-            return json.loads(collected_result[0].body.to_pybytes().decode())  # type: ignore
-        except Exception as e:
-            self.handle_flight_error(e)
-            raise e  # unreachable
+                return json.loads(collected_result[0].body.to_pybytes().decode())  # type: ignore
+            except Exception as e:
+                self.handle_flight_error(e)
+                raise e  # unreachable
+
+        return send_with_retry()
 
     def _upload_data(
         self,
@@ -605,17 +665,26 @@ class GdsArrowClient:
         flight_descriptor = self._versioned_flight_descriptor({"name": graph_name, "entity_type": entity_type})
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 
-        client = self._client()
-        put_stream, ack_stream = client.do_put(upload_descriptor, batches[0].schema)
+        @retry(
+            reraise=True,
+            before=before_log("Do put", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def safe_do_put(
+            upload_descriptor: FlightDescriptor, schema: Schema
+        ) -> tuple[FlightStreamWriter, FlightMetadataReader]:
+            return self._client().do_put(upload_descriptor, schema)  # type: ignore
+
+        put_stream, ack_stream = safe_do_put(upload_descriptor, batches[0].schema)
 
         @retry(
-            stop=(stop_after_delay(10) | stop_after_attempt(5)),
-            wait=wait_exponential(multiplier=1, min=1, max=10),
-            retry=(
-                retry_if_exception_type(flight.FlightUnavailableError)
-                | retry_if_exception_type(flight.FlightTimedOutError)
-                | retry_if_exception_type(flight.FlightInternalError)
-            ),
+            reraise=True,
+            before=before_log("Upload batch", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
         )
         def upload_batch(p: RecordBatch) -> None:
             put_stream.write_batch(p)
@@ -629,6 +698,26 @@ class GdsArrowClient:
         except Exception as e:
             GdsArrowClient.handle_flight_error(e)
 
+    def _do_get_with_retry(
+        self,
+        database: str,
+        graph_name: str,
+        procedure_name: str,
+        concurrency: Optional[int],
+        configuration: dict[str, Any],
+    ) -> pandas.DataFrame:
+        @retry(
+            reraise=True,
+            before=before_log("Do get", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def safe_do_get() -> pandas.DataFrame:
+            return self._do_get(database, graph_name, procedure_name, concurrency, configuration)
+
+        return safe_do_get()
+
     def _do_get(
         self,
         database: str,

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/query_runner/protocol/project_protocols.py

@@ -7,9 +7,9 @@ from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -133,6 +133,7 @@ class ProjectProtocolV3(ProjectProtocol):
         logger = getLogger()
 
         @retry(
+            reraise=True,
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
             retry=retry_if_result(is_not_done),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/query_runner/protocol/write_protocols.py

@@ -7,9 +7,9 @@ from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -142,6 +142,7 @@ class RemoteWriteBackV3(WriteProtocol):
         logger = logging.getLogger()
 
         @retry(
+            reraise=True,
             retry=retry_if_result(is_not_completed),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
             before=before_log(

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/query_runner/termination_flag.py

@@ -18,8 +18,7 @@ class TerminationFlag(ABC):
         pass
 
     def assert_running(self) -> None:
-        if self.is_set():
-            raise RuntimeError("Query has been terminated")
+        pass
 
     @staticmethod
     def create(signals: Optional[list[signal.Signals]] = None) -> TerminationFlag:
@@ -38,7 +37,7 @@ class TerminationFlagImpl(TerminationFlag):
         self._event = threading.Event()
 
         def receive_signal(sig: int, frame: Optional[FrameType]) -> None:
-            logging.debug(f"Received signal {sig}. Interrupting query.")
+            logging.info(f"Received signal {sig}.")
             self._event.set()
 
         for sig in signals:
@@ -52,7 +51,7 @@ class TerminationFlagImpl(TerminationFlag):
 
     def assert_running(self) -> None:
         if self.is_set():
-            raise RuntimeError("Query has been terminated")
+            raise RuntimeError("Closing client connection. Note, the query will be continued on the server-side")
 
 
 class TerminationFlagNoop(TerminationFlag):

graphdatascience-1.15a1/graphdatascience/retry_utils/retry_config.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+from tenacity.retry import retry_base
+from tenacity.stop import stop_base
+from tenacity.wait import wait_base
+
+
+@dataclass(frozen=True, repr=True)
+class RetryConfig:
+    stop: stop_base
+    wait: wait_base
+    retry: retry_base

{graphdatascience-1.14/graphdatascience/query_runner/protocol → graphdatascience-1.15a1/graphdatascience/retry_utils}/retry_utils.py

@@ -11,9 +11,10 @@ def before_log(
     sec_format: str = "%0.3f",
 ) -> typing.Callable[[RetryCallState], None]:
     def log_it(retry_state: RetryCallState) -> None:
-        logger.log(
-            log_level,
-            f"Retry of '{fn_name}', " f"attempt: {retry_state.attempt_number}",
-        )
+        if retry_state.attempt_number > 1:  # only log on actual retry
+            logger.log(
+                log_level,
+                f"Retry of '{fn_name}', " f"attempt: {retry_state.attempt_number}",
+            )
 
     return log_it

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/session/aura_api.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 import math
 import time
+import uuid
 import warnings
 from collections import defaultdict
 from datetime import timedelta
@@ -44,8 +45,9 @@ class AuraApi:
         self, client_id: str, client_secret: str, tenant_id: Optional[str] = None, aura_env: Optional[str] = None
     ) -> None:
         self._base_uri = AuraApi.base_uri(aura_env)
+        self._credentials = (client_id, client_secret)
 
-        self._request_session = self._init_request_session((client_id, client_secret))
+        self._request_session = self._init_request_session(self._credentials)
         self._logger = logging.getLogger()
 
         self._tenant_id = tenant_id if tenant_id else self._get_tenant_id()
@@ -99,12 +101,12 @@ class AuraApi:
     def get_or_create_session(
         self,
         name: str,
-        pwd: str,
         memory: SessionMemoryValue,
         dbid: Optional[str] = None,
         ttl: Optional[timedelta] = None,
         cloud_location: Optional[CloudLocation] = None,
     ) -> SessionDetails:
+        pwd = str(uuid.uuid4())  # password wont be used and will go away in v1 endpoints
         json = {"name": name, "password": pwd, "memory": memory.value, "tenant_id": self._tenant_id}
 
         if dbid:
@@ -164,7 +166,7 @@ class AuraApi:
     def wait_for_session_running(
         self,
         session_id: str,
-        sleep_time: float = 0.2,
+        sleep_time: float = 1.0,
         max_sleep_time: float = 10,
         max_wait_time: float = math.inf,
     ) -> WaitResult:
@@ -173,10 +175,16 @@ class AuraApi:
             session = self.get_session(session_id)
             if session is None:
                 return WaitResult.from_error(f"Session `{session_id}` not found -- please retry")
-            elif session.status == "Ready" and session.host:  # check host needed until dns based routing
+            elif session.status == "Ready":
                 return WaitResult.from_connection_url(session.bolt_connection_url())
             elif session.status == "Failed":
-                return WaitResult.from_error(f"Session `{session_id}` failed due to: {session.errors}")
+                return WaitResult.from_error(
+                    f"Session `{session_id}` with name `{session.name}` failed due to: {session.errors}"
+                )
+            elif session.is_expired():
+                return WaitResult.from_error(
+                    f"Session `{session_id}` with name `{session.name}` is expired. Expired due to: {session.errors}"
+                )
             else:
                 self._logger.debug(
                     f"Session `{session_id}` is not yet running. "

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/session/dedicated_sessions.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import hashlib
 import math
 import warnings
 from datetime import datetime, timedelta, timezone
@@ -74,10 +73,7 @@ class DedicatedSessions:
 
         dbid = AuraApi.extract_id(db_connection.uri)
 
-        # hashing the password to avoid storing the actual db password in Aura
-        password = hashlib.sha256(db_connection.password.encode()).hexdigest()
-
-        session_details = self._get_or_create_session(session_name, dbid, password, memory.value, ttl, cloud_location)
+        session_details = self._get_or_create_session(session_name, dbid, memory.value, ttl, cloud_location)
 
         if session_details.expiry_date:
             until_expiry: timedelta = session_details.expiry_date - datetime.now(timezone.utc)
@@ -95,8 +91,8 @@ class DedicatedSessions:
 
         session_connection = DbmsConnectionInfo(
             uri=connection_url,
-            username=DedicatedSessions.GDS_SESSION_USER,
-            password=password,
+            username=self._aura_api._credentials[0],
+            password=self._aura_api._credentials[1],
         )
 
         return self._construct_client(
@@ -154,7 +150,6 @@ class DedicatedSessions:
         self,
         session_name: str,
         dbid: str,
-        pwd: str,
         memory: SessionMemoryValue,
         ttl: Optional[timedelta] = None,
         cloud_location: Optional[CloudLocation] = None,
@@ -171,13 +166,12 @@ class DedicatedSessions:
         if cloud_location:
             return self._aura_api.get_or_create_session(
                 name=session_name,
-                pwd=pwd,
                 memory=memory,
                 ttl=ttl,
                 cloud_location=cloud_location,
             )
         else:
-            return self._aura_api.get_or_create_session(name=session_name, dbid=dbid, pwd=pwd, memory=memory, ttl=ttl)
+            return self._aura_api.get_or_create_session(name=session_name, dbid=dbid, memory=memory, ttl=ttl)
 
     def _construct_client(
         self,

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience/version.py

@@ -1,2 +1,2 @@
-__version__ = "1.14"
+__version__ = "1.15a1"
 __min_server_version__ = "2.6.0"  # matches installation.adoc

{graphdatascience-1.14 → graphdatascience-1.15a1/graphdatascience.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.14
+Version: 1.15a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -51,6 +51,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: project-url
 Dynamic: provides-extra
 Dynamic: requires-dist

{graphdatascience-1.14 → graphdatascience-1.15a1}/graphdatascience.egg-info/SOURCES.txt

@@ -93,7 +93,6 @@ graphdatascience/query_runner/progress/query_progress_provider.py
 graphdatascience/query_runner/progress/static_progress_provider.py
 graphdatascience/query_runner/protocol/__init__.py
 graphdatascience/query_runner/protocol/project_protocols.py
-graphdatascience/query_runner/protocol/retry_utils.py
 graphdatascience/query_runner/protocol/status.py
 graphdatascience/query_runner/protocol/write_protocols.py
 graphdatascience/resources/__init__.py
@@ -118,6 +117,9 @@ graphdatascience/resources/lastfm/user_friend_df_directed.parquet.gzip
 graphdatascience/resources/lastfm/user_listen_artist_rels.parquet.gzip
 graphdatascience/resources/lastfm/user_nodes.parquet.gzip
 graphdatascience/resources/lastfm/user_tag_artist_rels.parquet.gzip
+graphdatascience/retry_utils/__init__.py
+graphdatascience/retry_utils/retry_config.py
+graphdatascience/retry_utils/retry_utils.py
 graphdatascience/semantic_version/__init__.py
 graphdatascience/semantic_version/semantic_version.py
 graphdatascience/server_version/__init__.py