graphdatascience 1.15a2__tar.gz → 1.16__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.15a2
+Version: 1.16
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -32,7 +32,7 @@ Requires-Dist: multimethod<3.0,>=1.0
 Requires-Dist: neo4j<6.0,>=4.4.12
 Requires-Dist: numpy<2.3
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<20.0,>=16.0
+Requires-Dist: pyarrow<21.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.15a2 → graphdatascience-1.16}/graphdatascience/error/cypher_warning_handler.py

@@ -19,15 +19,12 @@ def filter_id_func_deprecation_warning() -> Callable[[F], F]:
                 message=r"^The query used a deprecated function: `id`\.",
             )
 
+            # previously      The query used a deprecated function. ('id' is no longer supported)
+            # since 2025.04.0 The query used a deprecated function. ('id' has been replaced by 'elementId or an application-generated id')
+            # since 2025.06   The query used a deprecated function. ('id' has been replaced by 'elementId or consider using an application-generated id')
             warnings.filterwarnings(
                 "ignore",
-                message=r"^The query used a deprecated function. \('id' is no longer supported\)",
-            )
-
-            # since 2025.04.0
-            warnings.filterwarnings(
-                "ignore",
-                message=r"^The query used a deprecated function. \('id' has been replaced by 'elementId or an application-generated id'\)",
+                message=r"The query used a deprecated function. \('id'.*",
             )
 
             return func(self, *args, **kwargs)

{graphdatascience-1.15a2 → graphdatascience-1.16}/graphdatascience/graph/graph_cypher_runner.py

@@ -45,7 +45,9 @@ class GraphCypherRunner(CallerBase):
 
         GraphCypherRunner._verify_query_ends_with_return_clause(self._namespace, query)
 
-        result: Optional[dict[str, Any]] = self._query_runner.run_cypher(query, params, database, False).squeeze()
+        result: Optional[dict[str, Any]] = self._query_runner.run_retryable_cypher(
+            query, params, database, custom_error=False
+        ).squeeze()
 
         if not result:
             raise ValueError("Projected graph cannot be empty.")
@@ -101,7 +103,7 @@ class GraphCypherRunner(CallerBase):
                             at_end = True
                             break
 
-            if query_token == "RETURN":
+            if query_token.upper() == "RETURN":
                 # State 1: We found the start of a `RETURN` clause.
                 # Check if it is the `RETURN gds.graph.project` call.
                 # We split tokens on `__separators` and flatten the nested iters.

{graphdatascience-1.15a2 → graphdatascience-1.16}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -158,12 +158,13 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
             duplicate_properties = set(db_node_properties).intersection(set(node_properties))
             if duplicate_properties:
                 raise ValueError(
-                    f"Duplicate property keys '{duplicate_properties}' in db_node_properties and " f"node_properties."
+                    f"Duplicate property keys '{duplicate_properties}' in db_node_properties and node_properties."
                 )
 
             unique_node_ids = result["nodeId"].drop_duplicates().tolist()
-            db_properties_df = query_runner.run_cypher(
-                GraphNodePropertiesRunner._build_query(db_node_properties), {"ids": unique_node_ids}
+
+            db_properties_df = query_runner.run_retryable_cypher(
+                GraphNodePropertiesRunner._build_query(db_node_properties), params={"ids": unique_node_ids}
             )
 
             if "propertyValue" not in result.keys():
@@ -342,7 +343,7 @@ class GraphRelationshipsRunner(GraphEntityOpsBaseRunner):
     def stream(self, G: Graph, relationship_types: list[str] = ["*"], **config: Any) -> TopologyDataFrame:
         self._namespace += ".stream"
         params = CallParameters(graph_name=G.name(), relationship_types=relationship_types, config=config)
-        result = self._query_runner.call_procedure(endpoint=self._namespace, params=params)
+        result = self._query_runner.call_procedure(endpoint=self._namespace, params=params, retryable=True)
 
         return TopologyDataFrame(result)
 
@@ -360,7 +361,9 @@ class GraphRelationshipsBetaRunner(GraphEntityOpsBaseRunner):
         self._namespace += ".stream"
         params = CallParameters(graph_name=G.name(), relationship_types=relationship_types, config=config)
 
-        return TopologyDataFrame(self._query_runner.call_procedure(endpoint=self._namespace, params=params))
+        return TopologyDataFrame(
+            self._query_runner.call_procedure(endpoint=self._namespace, params=params, retryable=True)
+        )
 
     @property
     @compatible_with("toUndirected", min_inclusive=ServerVersion(2, 3, 0))
@@ -381,7 +384,7 @@ class GraphPropertyRunner(UncallableNamespace, IllegalAttrChecker):
         self._namespace += ".stream"
         params = CallParameters(graph_name=G.name(), graph_property=graph_property, config=config)
 
-        return self._query_runner.call_procedure(endpoint=self._namespace, params=params)
+        return self._query_runner.call_procedure(endpoint=self._namespace, params=params, retryable=True)
 
     @compatible_with("drop", min_inclusive=ServerVersion(2, 2, 0))
     @graph_type_check

{graphdatascience-1.15a2 → graphdatascience-1.16}/graphdatascience/graph_data_science.py

@@ -4,6 +4,7 @@ import warnings
 from types import TracebackType
 from typing import Any, Optional, Type, Union
 
+import neo4j
 from neo4j import Driver
 from pandas import DataFrame
 
@@ -78,8 +79,11 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         if isinstance(endpoint, QueryRunner):
             self._query_runner = endpoint
         else:
+            db_auth = None
+            if auth:
+                db_auth = neo4j.basic_auth(*auth)
             self._query_runner = Neo4jQueryRunner.create_for_db(
-                endpoint, auth, aura_ds, database, bookmarks, show_progress
+                endpoint, db_auth, aura_ds, database, bookmarks, show_progress
             )
 
         self._server_version = self._query_runner.server_version()
@@ -218,6 +222,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         if isinstance(self._query_runner, ArrowQueryRunner):
             qr = self._query_runner.fallback_query_runner()
 
+        # not using qr.run_retryable_cypher as we dont know if it can be retried
         return qr.run_cypher(query, params, database, False)
 
     def driver_config(self) -> dict[str, Any]:

{graphdatascience-1.15a2 → graphdatascience-1.16}/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.16}/graphdatascience/query_runner/arrow_query_runner.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 from pandas import DataFrame
 
 from graphdatascience.query_runner.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
@@ -68,6 +69,15 @@ class ArrowQueryRunner(QueryRunner):
     ) -> DataFrame:
         return self._fallback_query_runner.run_cypher(query, params, database, custom_error)
 
+    def run_retryable_cypher(
+        self,
+        query: str,
+        params: Optional[dict[str, Any]] = None,
+        database: Optional[str] = None,
+        custom_error: bool = True,
+    ) -> DataFrame:
+        return self._fallback_query_runner.run_retryable_cypher(query, params, database, custom_error=custom_error)
+
     def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
         return self._fallback_query_runner.call_function(endpoint, params)
 
@@ -77,7 +87,9 @@ class ArrowQueryRunner(QueryRunner):
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
+        retryable: bool = False,
         custom_error: bool = True,
     ) -> DataFrame:
         if params is None:
@@ -171,7 +183,9 @@ class ArrowQueryRunner(QueryRunner):
                 graph_name, self._database_or_throw(), relationship_types, concurrency
             )
 
-        return self._fallback_query_runner.call_procedure(endpoint, params, yields, database, logging, custom_error)
+        return self._fallback_query_runner.call_procedure(
+            endpoint, params, yields, database, logging=logging, retryable=retryable, custom_error=custom_error
+        )
 
     def server_version(self) -> ServerVersion:
         return self._fallback_query_runner.server_version()
@@ -211,6 +225,13 @@ class ArrowQueryRunner(QueryRunner):
         self._fallback_query_runner.close()
         self._gds_arrow_client.close()
 
+    def cloneWithoutRouting(self, host: str, port: int) -> "QueryRunner":
+        return ArrowQueryRunner(
+            self._gds_arrow_client,
+            self._fallback_query_runner.cloneWithoutRouting(host, port),
+            self._server_version,
+        )
+
     def fallback_query_runner(self) -> QueryRunner:
         return self._fallback_query_runner
 

{graphdatascience-1.15a2 → graphdatascience-1.16}/graphdatascience/query_runner/cypher_graph_constructor.py

@@ -104,15 +104,14 @@ class CypherGraphConstructor(GraphConstructor):
 
     def _should_warn_about_arrow_missing(self) -> bool:
         try:
-            license: str = self._query_runner.run_cypher(
+            license: str = self._query_runner.run_retryable_cypher(
                 "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value", custom_error=False
             ).squeeze()
             should_warn = license == "Licensed"
         except Exception as e:
             # It's not a user's concern whether Arrow is set up or not in AuraDS.
-            if (
-                "There is no procedure with the name `gds.debug.sysInfo` "
-                "registered for this database instance." in str(e)
+            if "There is no procedure with the name `gds.debug.sysInfo` registered for this database instance." in str(
+                e
             ):
                 should_warn = False
             else:
@@ -210,6 +209,7 @@ class CypherGraphConstructor(GraphConstructor):
                 "undirectedRelationshipTypes": self._undirected_relationship_types,
             }
 
+            # not using retryable here as gds.graph.project adds a graph to the gds graph catalog
             self._query_runner.run_cypher(
                 query,
                 {

{graphdatascience-1.15a2 → graphdatascience-1.16}/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.16}/graphdatascience/query_runner/neo4j_query_runner.py

@@ -9,6 +9,8 @@ from typing import Any, NamedTuple, Optional, Union
 import neo4j
 from pandas import DataFrame
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..call_parameters import CallParameters
 from ..error.endpoint_suggester import generate_suggestive_error_message
 from ..error.gds_not_installed import GdsNotFound
@@ -30,14 +32,18 @@ class Neo4jQueryRunner(QueryRunner):
     @staticmethod
     def create_for_db(
         endpoint: Union[str, neo4j.Driver],
-        auth: Optional[tuple[str, str]] = None,
+        auth: Union[tuple[str, str], neo4j.Auth, None] = None,
         aura_ds: bool = False,
         database: Optional[str] = None,
         bookmarks: Optional[Any] = None,
         show_progress: bool = True,
+        config: Optional[dict[str, Any]] = None,
     ) -> Neo4jQueryRunner:
         if isinstance(endpoint, str):
-            config: dict[str, Any] = {"user_agent": f"neo4j-graphdatascience-v{__version__}"}
+            if config is None:
+                config = {}
+
+            config["user_agent"] = f"neo4j-graphdatascience-v{__version__}"
 
             if aura_ds:
                 Neo4jQueryRunner._configure_aura(config)
@@ -46,6 +52,8 @@ class Neo4jQueryRunner(QueryRunner):
 
             query_runner = Neo4jQueryRunner(
                 driver,
+                Neo4jQueryRunner.parse_protocol(endpoint),
+                auth,
                 auto_close=True,
                 bookmarks=bookmarks,
                 config=config,
@@ -54,8 +62,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)}")
@@ -65,7 +79,7 @@ class Neo4jQueryRunner(QueryRunner):
     @staticmethod
     def create_for_session(
         endpoint: str,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Union[tuple[str, str], neo4j.Auth, None] = None,
         show_progress: bool = True,
     ) -> Neo4jQueryRunner:
         driver_config: dict[str, Any] = {"user_agent": f"neo4j-graphdatascience-v{__version__}"}
@@ -76,6 +90,8 @@ class Neo4jQueryRunner(QueryRunner):
 
         query_runner = Neo4jQueryRunner(
             driver,
+            Neo4jQueryRunner.parse_protocol(endpoint),
+            auth,
             auto_close=True,
             show_progress=show_progress,
             bookmarks=None,
@@ -90,13 +106,26 @@ class Neo4jQueryRunner(QueryRunner):
 
     @staticmethod
     def _configure_aura(config: dict[str, Any]) -> None:
-        config["max_connection_lifetime"] = 60 * 8  # 8 minutes
-        config["keep_alive"] = True
-        config["max_connection_pool_size"] = 50
+        # defaults as documented in https://support.neo4j.com/s/article/1500001173021-How-to-handle-Session-Expired-Errors-while-connecting-to-Neo4j-Aura
+        config.setdefault("max_connection_lifetime", 60 * 50)  # 50 minutes
+        config.setdefault("keep_alive", True)
+        config.setdefault("max_connection_pool_size", 50)
+
+        if Neo4jQueryRunner._NEO4J_DRIVER_VERSION >= SemanticVersion(5, 16, 0):
+            config.setdefault("liveness_check_timeout", 60 * 5)  # 5 minutes
+
+    @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: Union[tuple[str, str], neo4j.Auth, None] = None,
         config: dict[str, Any] = {},
         database: Optional[str] = neo4j.DEFAULT_DATABASE,
         auto_close: bool = False,
@@ -105,6 +134,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
@@ -121,8 +152,10 @@ class Neo4jQueryRunner(QueryRunner):
     def __run_cypher_simplified_for_query_progress_logger(self, query: str, database: Optional[str]) -> DataFrame:
         # progress logging should not retry a lot as it perodically fetches the latest progress anyway
         connectivity_retry_config = Neo4jQueryRunner.ConnectivityRetriesConfig(max_retries=2)
+        # not using retryable cypher as failing is okay
         return self.run_cypher(query=query, database=database, connectivity_retry_config=connectivity_retry_config)
 
+    # only use for user defined queries
     def run_cypher(
         self,
         query: str,
@@ -166,12 +199,50 @@ class Neo4jQueryRunner(QueryRunner):
 
             return df
 
-    def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
+    # better retry mechanism than run_cypher. The neo4j driver handles retryable errors internally
+    def run_retryable_cypher(
+        self,
+        query: str,
+        params: Optional[dict[str, Any]] = None,
+        database: Optional[str] = None,
+        custom_error: bool = True,
+        mode: Optional[QueryMode] = None,
+        connectivity_retry_config: Optional[ConnectivityRetriesConfig] = None,
+    ) -> DataFrame:
+        if not database:
+            database = self._database
+
+        if self._NEO4J_DRIVER_VERSION < SemanticVersion(5, 5, 0):
+            return self.run_cypher(query, params, database, custom_error, connectivity_retry_config)
+
+        if not mode:
+            routing = neo4j.RoutingControl.READ
+        else:
+            routing = mode.neo4j_routing()
+
+        try:
+            return self._driver.execute_query(
+                query_=query,
+                parameters_=params,
+                database_=database,
+                result_transformer_=neo4j.Result.to_df,
+                bookmark_manager_=self.bookmarks(),
+                routing_=routing,
+            )
+        except Exception as e:
+            if custom_error:
+                Neo4jQueryRunner.handle_driver_exception(self._driver, e)
+                raise e
+            else:
+                raise e
+
+    def call_function(self, endpoint: str, params: Optional[CallParameters] = None, custom_error: bool = True) -> Any:
         if params is None:
             params = CallParameters()
         query = f"RETURN {endpoint}({params.placeholder_str()})"
 
-        return self.run_cypher(query, params).squeeze()
+        # we can use retryable cypher as we expect all gds functions to be idempotent
+        return self.run_retryable_cypher(query, params, custom_error=custom_error, mode=QueryMode.READ).squeeze()
 
     def call_procedure(
         self,
@@ -179,7 +250,9 @@ class Neo4jQueryRunner(QueryRunner):
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
+        retryable: bool = False,
         custom_error: bool = True,
     ) -> DataFrame:
         if params is None:
@@ -189,7 +262,10 @@ class Neo4jQueryRunner(QueryRunner):
         query = f"CALL {endpoint}({params.placeholder_str()}){yields_clause}"
 
         def run_cypher_query() -> DataFrame:
-            return self.run_cypher(query, params, database, custom_error)
+            if retryable:
+                return self.run_retryable_cypher(query, params, database, custom_error, mode=mode)
+            else:
+                return self.run_cypher(query, params, database, custom_error)
 
         job_id = None if not params else params.get_job_id()
         if self._resolve_show_progress(logging) and job_id:
@@ -205,7 +281,7 @@ class Neo4jQueryRunner(QueryRunner):
             return self._server_version
 
         try:
-            server_version_string = self.run_cypher("RETURN gds.version()", custom_error=False).squeeze()
+            server_version_string = self.call_function("gds.version", custom_error=False)
             server_version = ServerVersion.from_string(server_version_string)
             self._server_version = server_version
             return server_version
@@ -279,8 +355,25 @@ class Neo4jQueryRunner(QueryRunner):
     def set_show_progress(self, show_progress: bool) -> None:
         self._show_progress = show_progress
 
+    def cloneWithoutRouting(self, host: str, port: int) -> QueryRunner:
+        protocol = self._protocol.replace("neo4j", "bolt")
+        endpoint = "{}://{}:{}".format(protocol, host, port)
+        driver = neo4j.GraphDatabase.driver(endpoint, auth=self._auth, **self.driver_config())
+
+        return Neo4jQueryRunner(
+            driver=driver,
+            protocol=protocol,
+            auth=self._auth,
+            config=self._config,
+            database=self._database,
+            auto_close=self._auto_close,
+            bookmarks=self._bookmarks,
+            show_progress=self._show_progress,
+            instance_description=self._instance_description,
+        )
+
     @staticmethod
-    def handle_driver_exception(session: neo4j.Session, e: Exception) -> None:
+    def handle_driver_exception(cypher_executor: Union[neo4j.Session, neo4j.Driver], e: Exception) -> None:
         reg_gds_hit = re.search(
             r"There is no procedure with the name `(gds(?:\.\w+)+)` registered for this database instance",
             str(e),
@@ -290,8 +383,16 @@ class Neo4jQueryRunner(QueryRunner):
 
         requested_endpoint = reg_gds_hit.group(1)
 
-        list_result = session.run("CALL gds.list() YIELD name")
-        all_endpoints = list_result.to_df()["name"].tolist()
+        if isinstance(cypher_executor, neo4j.Session):
+            list_result = cypher_executor.run("CALL gds.list() YIELD name")
+            all_endpoints = list_result.to_df()["name"].tolist()
+        elif isinstance(cypher_executor, neo4j.Driver):
+            result = cypher_executor.execute_query("CALL gds.list() YIELD name", result_transformer_=neo4j.Result.to_df)
+            all_endpoints = result["name"].tolist()
+        else:
+            raise TypeError(
+                f"Expected cypher_executor to be a neo4j.Session or neo4j.Driver, got {type(cypher_executor)}"
+            )
 
         raise SyntaxError(generate_suggestive_error_message(requested_endpoint, all_endpoints)) from e
 

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

@@ -68,7 +68,9 @@ class ProjectProtocolV1(ProjectProtocol):
         logging: bool = False,
     ) -> DataFrame:
         versioned_endpoint = ProtocolVersion.V1.versioned_procedure_name(endpoint)
-        return query_runner.call_procedure(versioned_endpoint, params, yields, database, logging, False)
+        return query_runner.call_procedure(
+            versioned_endpoint, params, yields, database=database, logging=logging, retryable=False, custom_error=False
+        )
 
 
 class ProjectProtocolV2(ProjectProtocol):
@@ -97,7 +99,9 @@ class ProjectProtocolV2(ProjectProtocol):
         logging: bool = False,
     ) -> DataFrame:
         versioned_endpoint = ProtocolVersion.V2.versioned_procedure_name(endpoint)
-        return query_runner.call_procedure(versioned_endpoint, params, yields, database, logging, False)
+        return query_runner.call_procedure(
+            versioned_endpoint, params, yields, database=database, logging=logging, retryable=False, custom_error=False
+        )
 
 
 class ProjectProtocolV3(ProjectProtocol):
@@ -121,7 +125,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 +136,20 @@ 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=logging,
+            custom_error=False,
+            retryable=True,
+        ).squeeze()
+        member_host = response["host"]
+        member_port = response["port"] if ("port" in response.index) else 7687
+        projection_query_runner = query_runner.cloneWithoutRouting(member_host, member_port)
+
         @retry(
             reraise=True,
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
@@ -139,9 +157,19 @@ 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(
-                ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
+            termination_flag.assert_running()
+            return projection_query_runner.call_procedure(
+                ProtocolVersion.V3.versioned_procedure_name(endpoint),
+                params,
+                yields,
+                database=database,
+                logging=logging,
+                retryable=True,
+                custom_error=False,
             )
 
-        return project_fn()
+        projection_result = project_fn()
+
+        projection_query_runner.close()
+
+        return projection_result

{graphdatascience-1.15a2 → graphdatascience-1.16}/graphdatascience/query_runner/protocol/write_protocols.py

@@ -5,9 +5,10 @@ from typing import Any, Optional
 from pandas import DataFrame
 from tenacity import retry, retry_if_result, wait_incrementing
 
-from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.query_runner.protocol.status import Status
+from graphdatascience.query_runner.query_mode import QueryMode
+from graphdatascience.query_runner.query_runner import QueryRunner
 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
@@ -73,9 +74,11 @@ class RemoteWriteBackV1(WriteProtocol):
             ProtocolVersion.V1.versioned_procedure_name("gds.arrow.write"),
             parameters,
             yields,
-            None,
-            False,
-            False,
+            retryable=False,
+            database=None,
+            logging=False,
+            mode=QueryMode.WRITE,
+            custom_error=False,
         )
 
 
@@ -111,9 +114,11 @@ class RemoteWriteBackV2(WriteProtocol):
             ProtocolVersion.V2.versioned_procedure_name("gds.arrow.write"),
             parameters,
             yields,
-            None,
-            False,
-            False,
+            retryable=False,
+            database=None,
+            logging=False,
+            mode=QueryMode.WRITE,
+            custom_error=False,
         )
 
 
@@ -157,9 +162,10 @@ class RemoteWriteBackV3(WriteProtocol):
                 ProtocolVersion.V3.versioned_procedure_name("gds.arrow.write"),
                 parameters,
                 yields,
-                None,
-                False,
-                False,
+                retryable=True,
+                logging=False,
+                mode=QueryMode.WRITE,
+                custom_error=False,
             )
 
         return write_fn()

graphdatascience-1.16/graphdatascience/query_runner/query_mode.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+import neo4j
+
+
+class QueryMode(str, Enum):
+    READ = "read"
+    WRITE = "write"
+
+    def neo4j_routing(self) -> "neo4j.RoutingControl":
+        if self == QueryMode.READ:
+            return neo4j.RoutingControl.READ
+        elif self == QueryMode.WRITE:
+            return neo4j.RoutingControl.WRITE
+        else:
+            raise ValueError(f"Unknown query mode: {self}")