graphdatascience 1.15a2__tar.gz → 1.15.1__tar.gz

{graphdatascience-1.15a2/graphdatascience.egg-info → graphdatascience-1.15.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.15a2
+Version: 1.15.1
 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.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import concurrent
+import logging
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
@@ -33,6 +34,7 @@ class ArrowGraphConstructor(GraphConstructor):
         )
         self._chunk_size = chunk_size
         self._min_batch_size = chunk_size * 10
+        self._logger = logging.getLogger()
 
     def run(self, node_dfs: list[DataFrame], relationship_dfs: list[DataFrame]) -> None:
         try:
@@ -60,8 +62,11 @@ class ArrowGraphConstructor(GraphConstructor):
 
             self._client.relationship_load_done(self._graph_name)
         except (Exception, KeyboardInterrupt) as e:
-            self._client.abort(self._graph_name)
-
+            try:
+                self._client.abort(self._graph_name)
+            except Exception as abort_exception:
+                if "No arrow process" not in str(abort_exception):
+                    self._logger.warning(f"error aborting graph creation: {abort_exception}")
             raise e
 
     def _partition_dfs(self, dfs: list[DataFrame]) -> list[DataFrame]:

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/arrow_query_runner.py

@@ -211,6 +211,13 @@ class ArrowQueryRunner(QueryRunner):
         self._fallback_query_runner.close()
         self._gds_arrow_client.close()
 
+    def clone(self, host: str, port: int) -> "QueryRunner":
+        return ArrowQueryRunner(
+            self._gds_arrow_client,
+            self._fallback_query_runner.clone(host, port),
+            self._server_version,
+        )
+
     def fallback_query_runner(self) -> QueryRunner:
         return self._fallback_query_runner
 

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/gds_arrow_client.py

@@ -640,7 +640,10 @@ class GdsArrowClient:
         )
         def send_with_retry() -> dict[str, Any]:
             try:
-                result = client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
+                result = client.do_action(
+                    action=flight.Action(action_type, json.dumps(meta_data).encode("utf-8")),
+                    options=flight.FlightCallOptions(timeout=20.0),
+                )
 
                 # Consume result fully to sanity check and avoid cancelled streams
                 collected_result = list(result)

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/neo4j_query_runner.py

@@ -46,6 +46,8 @@ class Neo4jQueryRunner(QueryRunner):
 
             query_runner = Neo4jQueryRunner(
                 driver,
+                Neo4jQueryRunner.parse_protocol(endpoint),
+                auth,
                 auto_close=True,
                 bookmarks=bookmarks,
                 config=config,
@@ -54,8 +56,14 @@ class Neo4jQueryRunner(QueryRunner):
             )
 
         elif isinstance(endpoint, neo4j.Driver):
+            protocol = "neo4j+s" if endpoint.encrypted else "bolt"
             query_runner = Neo4jQueryRunner(
-                endpoint, auto_close=False, bookmarks=bookmarks, database=database, show_progress=show_progress
+                endpoint,
+                protocol,
+                auto_close=False,
+                bookmarks=bookmarks,
+                database=database,
+                show_progress=show_progress,
             )
         else:
             raise ValueError(f"Invalid endpoint type: {type(endpoint)}")
@@ -76,6 +84,8 @@ class Neo4jQueryRunner(QueryRunner):
 
         query_runner = Neo4jQueryRunner(
             driver,
+            Neo4jQueryRunner.parse_protocol(endpoint),
+            auth,
             auto_close=True,
             show_progress=show_progress,
             bookmarks=None,
@@ -94,9 +104,18 @@ class Neo4jQueryRunner(QueryRunner):
         config["keep_alive"] = True
         config["max_connection_pool_size"] = 50
 
+    @staticmethod
+    def parse_protocol(endpoint: str) -> str:
+        protocol_match = re.match(r"^([^:]+)://", endpoint)
+        if not protocol_match:
+            raise ValueError(f"Invalid endpoint URI format: {endpoint}")
+        return protocol_match.group(1)
+
     def __init__(
         self,
         driver: neo4j.Driver,
+        protocol: str,
+        auth: Optional[tuple[str, str]] = None,
         config: dict[str, Any] = {},
         database: Optional[str] = neo4j.DEFAULT_DATABASE,
         auto_close: bool = False,
@@ -105,6 +124,8 @@ class Neo4jQueryRunner(QueryRunner):
         instance_description: str = "Neo4j DBMS",
     ):
         self._driver = driver
+        self._protocol = protocol
+        self._auth = auth
         self._config = config
         self._auto_close = auto_close
         self._database = database
@@ -279,6 +300,22 @@ class Neo4jQueryRunner(QueryRunner):
     def set_show_progress(self, show_progress: bool) -> None:
         self._show_progress = show_progress
 
+    def clone(self, host: str, port: int) -> QueryRunner:
+        endpoint = "{}://{}:{}".format(self._protocol, host, port)
+        driver = neo4j.GraphDatabase.driver(endpoint, auth=self._auth, **self.driver_config())
+
+        return Neo4jQueryRunner(
+            driver,
+            self._protocol,
+            self._auth,
+            self._config,
+            self._database,
+            self._auto_close,
+            self._bookmarks,
+            self._show_progress,
+            self._instance_description,
+        )
+
     @staticmethod
     def handle_driver_exception(session: neo4j.Session, e: Exception) -> None:
         reg_gds_hit = re.search(

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/protocol/project_protocols.py

@@ -121,7 +121,7 @@ class ProjectProtocolV3(ProjectProtocol):
         query_runner: QueryRunner,
         endpoint: str,
         params: CallParameters,
-        terminationFlag: TerminationFlag,
+        termination_flag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -132,6 +132,14 @@ class ProjectProtocolV3(ProjectProtocol):
 
         logger = getLogger()
 
+        # We need to pin the driver to a specific cluster member
+        response = query_runner.call_procedure(
+            ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
+        ).squeeze()
+        member_host = response["host"]
+        member_port = response["port"] if ("port" in response.index) else 7687
+        projection_query_runner = query_runner.clone(member_host, member_port)
+
         @retry(
             reraise=True,
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
@@ -139,9 +147,13 @@ class ProjectProtocolV3(ProjectProtocol):
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
         )
         def project_fn() -> DataFrame:
-            terminationFlag.assert_running()
-            return query_runner.call_procedure(
+            termination_flag.assert_running()
+            return projection_query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
             )
 
-        return project_fn()
+        projection_result = project_fn()
+
+        projection_query_runner.close()
+
+        return projection_result

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/query_runner.py

@@ -80,5 +80,9 @@ class QueryRunner(ABC):
     def set_show_progress(self, show_progress: bool) -> None:
         pass
 
+    @abstractmethod
+    def clone(self, host: str, port: int) -> "QueryRunner":
+        pass
+
     def set_server_version(self, _: ServerVersion) -> None:
         pass

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/session_query_runner.py

@@ -120,6 +120,14 @@ class SessionQueryRunner(QueryRunner):
         self._show_progress = show_progress
         self._gds_query_runner.set_show_progress(show_progress)
 
+    def clone(self, host: str, port: int) -> QueryRunner:
+        return SessionQueryRunner(
+            self._gds_query_runner,
+            self._db_query_runner.clone(host, port),
+            self._gds_arrow_client,
+            self._show_progress,
+        )
+
     def close(self) -> None:
         self._gds_arrow_client.close()
         self._gds_query_runner.close()
@@ -202,12 +210,13 @@ class SessionQueryRunner(QueryRunner):
             return write_protocol.run_write_back(self._db_query_runner, write_back_params, yields, terminationFlag)
 
         try:
-            if self._resolve_show_progress(logging):
-                database_write_result = self._progress_logger.run_with_progress_logging(
-                    run_write_back, job_id, database
-                )
-            else:
-                database_write_result = run_write_back()
+            # Skipping progress for now as export has a different jobId
+            # if self._resolve_show_progress(logging):
+            #     database_write_result = self._progress_logger.run_with_progress_logging(
+            #         run_write_back, job_id, database
+            #     )
+            # else:
+            database_write_result = run_write_back()
         except Exception as e:
             # catch the case nothing was needed to write-back (empty graph)
             # once we have the Arrow Endpoints V2, we could catch by first checking the jobs summary

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/query_runner/standalone_session_query_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Optional
 
 from pandas import DataFrame
@@ -74,3 +76,6 @@ class StandaloneSessionQueryRunner(QueryRunner):
 
     def set_server_version(self, _: ServerVersion) -> None:
         super().set_server_version(_)
+
+    def clone(self, host: str, port: int) -> QueryRunner:
+        return self

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/session/aura_api.py

@@ -352,7 +352,7 @@ class AuraApi:
     def _check_status_code(self, resp: requests.Response) -> None:
         if resp.status_code >= 400:
             raise AuraApiError(
-                f"Request for {resp.url} failed with status code {resp.status_code} - {resp.reason}: {resp.text}",
+                f"Request for {resp.url} failed with status code {resp.status_code} - {resp.reason}: {resp.json()}",
                 status_code=resp.status_code,
             )
 

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/session/session_sizes.py

@@ -57,6 +57,7 @@ class SessionMemory(Enum):
     m_192GB = SessionMemoryValue("192GB")
     m_256GB = SessionMemoryValue("256GB")
     m_384GB = SessionMemoryValue("384GB")
+    m_512GB = SessionMemoryValue("512GB")
 
     @classmethod
     def all_values(cls) -> list[SessionMemoryValue]:

{graphdatascience-1.15a2 → graphdatascience-1.15.1}/graphdatascience/version.py

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

{graphdatascience-1.15a2 → graphdatascience-1.15.1/graphdatascience.egg-info}/PKG-INFO

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