Expose mode and retryable on user-facing run_cypher

seen in our e2e we now need to manually instantiate a neo4j driver to run retryable cypher queries such as for cleanup.

Author: FlorentinD

changelog.md

@@ -6,8 +6,12 @@
 
 ## Bug fixes
 
-* Fix reporting error based on http responses from the Aura-API with an invalid JSON body. Earlier the client would report JSONDecodeError instead of showing the actual issue.
+- Fix reporting error based on http responses from the Aura-API with an invalid JSON body. Earlier the client would report JSONDecodeError instead of showing the actual issue.
 
 ## Improvements
 
+- `GraphDataScience::run_query` now supports setting the `mode` of the query to be used for routing. Previously queries would always route the leader of the cluster, assuming write mode.
+- `GraphDataScience::run_query` now support setting `retryable` to enable a retry-mechanism for appropriate errors. This requires `neo4j>=5.5.0`.
+
+
 ## Other changes

graphdatascience/graph/graph_cypher_runner.py

@@ -6,6 +6,8 @@
 
 from pandas import Series
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..caller_base import CallerBase
 from ..query_runner.query_runner import QueryRunner
 from ..server_version.server_version import ServerVersion
@@ -46,7 +48,7 @@ def project(
         GraphCypherRunner._verify_query_ends_with_return_clause(self._namespace, query)
 
         result: Optional[dict[str, Any]] = self._query_runner.run_retryable_cypher(
-            query, params, database, custom_error=False
+            query, params, database, custom_error=False, mode=QueryMode.READ
         ).squeeze()
 
         if not result:

graphdatascience/graph/graph_entity_ops_runner.py

@@ -5,6 +5,8 @@
 import pandas as pd
 from pandas import DataFrame, Series
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..call_parameters import CallParameters
 from ..error.cypher_warning_handler import (
     filter_id_func_deprecation_warning,
@@ -164,7 +166,9 @@ def _process_result(
             unique_node_ids = result["nodeId"].drop_duplicates().tolist()
 
             db_properties_df = query_runner.run_retryable_cypher(
-                GraphNodePropertiesRunner._build_query(db_node_properties), params={"ids": unique_node_ids}
+                GraphNodePropertiesRunner._build_query(db_node_properties),
+                params={"ids": unique_node_ids},
+                mode=QueryMode.READ,
             )
 
             if "propertyValue" not in result.keys():

graphdatascience/graph_data_science.py

@@ -9,6 +9,7 @@
 from pandas import DataFrame
 
 from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
+from graphdatascience.query_runner.query_mode import QueryMode
 
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
@@ -199,7 +200,12 @@ def last_bookmarks(self) -> Optional[Any]:
         return self._query_runner.last_bookmarks()
 
     def run_cypher(
-        self, query: str, params: Optional[dict[str, Any]] = None, database: Optional[str] = None
+        self,
+        query: str,
+        params: Optional[dict[str, Any]] = None,
+        database: Optional[str] = None,
+        retryable: bool = False,
+        mode: QueryMode = QueryMode.WRITE,
     ) -> DataFrame:
         """
         Run a Cypher query
@@ -212,6 +218,10 @@ def run_cypher(
             parameters to the query
         database: str
             the database on which to run the query
+        retryable: bool
+            whether the query can be automatically retried. Make sure the query is idempotent if set to True.
+        mode: QueryMode
+            the query mode to use (READ or WRITE). Set based on the operation performed in the query.
 
         Returns:
             The query result as a DataFrame
@@ -222,8 +232,10 @@ def run_cypher(
         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)
+        if retryable:
+            return qr.run_retryable_cypher(query, params, database, custom_error=False, mode=mode)
+        else:
+            return qr.run_cypher(query, params, database, custom_error=False, mode=mode)
 
     def driver_config(self) -> dict[str, Any]:
         """

graphdatascience/query_runner/arrow_query_runner.py

@@ -65,18 +65,22 @@ def run_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._fallback_query_runner.run_cypher(query, params, database, custom_error)
+        return self._fallback_query_runner.run_cypher(query, params, database, mode, custom_error=custom_error)
 
     def run_retryable_cypher(
         self,
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._fallback_query_runner.run_retryable_cypher(query, params, database, custom_error=custom_error)
+        return self._fallback_query_runner.run_retryable_cypher(
+            query, params, database, mode, custom_error=custom_error
+        )
 
     def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
         return self._fallback_query_runner.call_function(endpoint, params)

graphdatascience/query_runner/cypher_graph_constructor.py

@@ -7,6 +7,8 @@
 
 from pandas import DataFrame, concat
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..server_version.server_version import ServerVersion
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
@@ -105,7 +107,9 @@ def graph_construct_error_multidf(element: str) -> str:
     def _should_warn_about_arrow_missing(self) -> bool:
         try:
             license: str = self._query_runner.run_retryable_cypher(
-                "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value", custom_error=False
+                "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value",
+                custom_error=False,
+                mode=QueryMode.READ,
             ).squeeze()
             should_warn = license == "Licensed"
         except Exception as e:

graphdatascience/query_runner/neo4j_query_runner.py

@@ -161,20 +161,28 @@ def run_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
         connectivity_retry_config: Optional[ConnectivityRetriesConfig] = None,
     ) -> DataFrame:
         if params is None:
             params = {}
 
+        if mode is None:
+            mode = QueryMode.WRITE
+
         if database is None:
             database = self._database
 
         if connectivity_retry_config is None:
             connectivity_retry_config = Neo4jQueryRunner.ConnectivityRetriesConfig()
         self._verify_connectivity(database=database, retry_config=connectivity_retry_config)
 
-        with self._driver.session(database=database, bookmarks=self.bookmarks()) as session:
+        with self._driver.session(
+            database=database,
+            bookmarks=self.bookmarks(),
+            default_access_mode=mode.neo4j_access_mode(),
+        ) as session:
             try:
                 result = session.run(query, params)
             except Exception as e:
@@ -205,18 +213,18 @@ def run_retryable_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
-        custom_error: bool = True,
         mode: Optional[QueryMode] = None,
+        custom_error: bool = True,
         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)
+            return self.run_cypher(query, params, database, mode, custom_error, connectivity_retry_config)
 
         if not mode:
-            routing = neo4j.RoutingControl.READ
+            routing = neo4j.RoutingControl.WRITE
         else:
             routing = mode.neo4j_routing()
 
@@ -263,9 +271,9 @@ def call_procedure(
 
         def run_cypher_query() -> DataFrame:
             if retryable:
-                return self.run_retryable_cypher(query, params, database, custom_error, mode=mode)
+                return self.run_retryable_cypher(query, params, database, custom_error=custom_error, mode=mode)
             else:
-                return self.run_cypher(query, params, database, custom_error)
+                return self.run_cypher(query, params, database, custom_error=custom_error)
 
         job_id = None if not params else params.get_job_id()
         if self._resolve_show_progress(logging) and job_id:

graphdatascience/query_runner/query_mode.py

@@ -14,3 +14,11 @@ def neo4j_routing(self) -> "neo4j.RoutingControl":
             return neo4j.RoutingControl.WRITE
         else:
             raise ValueError(f"Unknown query mode: {self}")
+
+    def neo4j_access_mode(self) -> str:
+        if self == QueryMode.READ:
+            return neo4j.READ_ACCESS
+        elif self == QueryMode.WRITE:
+            return neo4j.WRITE_ACCESS
+        else:
+            raise ValueError(f"Unknown query mode: {self}")

graphdatascience/query_runner/query_runner.py

@@ -35,6 +35,7 @@ def run_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
         pass
@@ -45,6 +46,7 @@ def run_retryable_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
         pass

graphdatascience/query_runner/session_query_runner.py

@@ -54,18 +54,20 @@ def run_cypher(
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._db_query_runner.run_cypher(query, params, database, custom_error)
+        return self._db_query_runner.run_cypher(query, params, database, mode, custom_error)
 
     def run_retryable_cypher(
         self,
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._db_query_runner.run_retryable_cypher(query, params, database, custom_error=custom_error)
+        return self._db_query_runner.run_retryable_cypher(query, params, database, mode=mode, custom_error=custom_error)
 
     def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
         return self._gds_query_runner.call_function(endpoint, params)