Use execute query for retrying on retryable Neo4j Exceptions

Author: FlorentinD

changelog.md

@@ -1,21 +1,20 @@
 # Changes in 1.16
 
-
 ## Breaking changes
 
 ## New features
 
 ## Bug fixes
 
-* Fixed a bug where remote projections would fail when the database is clustered
+- Fixed a bug where remote projections would fail when the database is clustered
 
 ## Improvements
 
-* Allow creating sessions of size `512GB`.
-* Allow passing additional parameters for the Neo4j driver connection to `GdsSessions.get_or_create(neo4j_driver_config={..})`
-* Add helper functions to create config objects from environment variables
-    * `AuraApiCredentials::from_env`
-    * `DbmsConnectionInfo::from_env`
-
+- Allow creating sessions of size `512GB`.
+- Allow passing additional parameters for the Neo4j driver connection to `GdsSessions.get_or_create(neo4j_driver_config={..})`
+- Add helper functions to create config objects from environment variables
+  - `AuraApiCredentials::from_env`
+  - `DbmsConnectionInfo::from_env`
+- Retry internal functions known to be idempotent. Reduces issues such as `SessionExpiredError`.
 
 ## Other changes

graphdatascience/graph/graph_cypher_runner.py

@@ -45,7 +45,9 @@ def project(
 
         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, False
+        ).squeeze()
 
         if not result:
             raise ValueError("Projected graph cannot be empty.")

graphdatascience/graph/graph_entity_ops_runner.py

@@ -162,7 +162,8 @@ def _process_result(
                 )
 
             unique_node_ids = result["nodeId"].drop_duplicates().tolist()
-            db_properties_df = query_runner.run_cypher(
+
+            db_properties_df = query_runner.run_retryable_cypher(
                 GraphNodePropertiesRunner._build_query(db_node_properties), {"ids": unique_node_ids}
             )
 

graphdatascience/graph_data_science.py

@@ -218,6 +218,7 @@ def run_cypher(
         if isinstance(self._query_runner, ArrowQueryRunner):
             qr = self._query_runner.fallback_query_runner()
 
+        # not using qr.execute_query 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/query_runner/arrow_query_runner.py

@@ -68,6 +68,15 @@ def run_cypher(
     ) -> 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)
 
@@ -78,6 +87,7 @@ def call_procedure(
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
+        retryable: bool = False,
         custom_error: bool = True,
     ) -> DataFrame:
         if params is None:
@@ -171,7 +181,9 @@ def call_procedure(
                 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()

graphdatascience/query_runner/cypher_graph_constructor.py

@@ -104,6 +104,7 @@ def graph_construct_error_multidf(element: str) -> str:
 
     def _should_warn_about_arrow_missing(self) -> bool:
         try:
+            # TOOD use execute_query
             license: str = self._query_runner.run_cypher(
                 "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value", custom_error=False
             ).squeeze()
@@ -210,6 +211,7 @@ def run(self, node_dfs: list[DataFrame], relationship_dfs: list[DataFrame]) -> N
                 "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/query_runner/neo4j_query_runner.py

@@ -150,8 +150,10 @@ def __init__(
     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 execute_query 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,
@@ -195,12 +197,47 @@ def run_cypher(
 
             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,
+        routing: Optional[neo4j.RoutingControl] = 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 routing:
+            routing = neo4j.RoutingControl.READ
+
+        try:
+            return self._driver.execute_query(
+                query_=query,
+                parameters_=params,
+                database=database,
+                result_transformer_=neo4j.Result.to_df,
+                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).squeeze()
 
     def call_procedure(
         self,
@@ -209,6 +246,7 @@ def call_procedure(
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
+        retryable: bool = False,
         custom_error: bool = True,
     ) -> DataFrame:
         if params is None:
@@ -218,7 +256,11 @@ def call_procedure(
         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:
+                routing = neo4j.RoutingControl.WRITE if "write" in endpoint else neo4j.RoutingControl.READ
+                return self.run_retryable_cypher(query, params, database, custom_error, routing=routing)
+            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:
@@ -234,7 +276,7 @@ def server_version(self) -> ServerVersion:
             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
@@ -325,7 +367,7 @@ def clone(self, host: str, port: int) -> QueryRunner:
         )
 
     @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),
@@ -335,8 +377,16 @@ def handle_driver_exception(session: neo4j.Session, e: Exception) -> None:
 
         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/query_runner/protocol/project_protocols.py

@@ -68,7 +68,9 @@ def run_projection(
         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=True, custom_error=False
+        )
 
 
 class ProjectProtocolV2(ProjectProtocol):
@@ -97,7 +99,9 @@ def run_projection(
         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=True, custom_error=False
+        )
 
 
 class ProjectProtocolV3(ProjectProtocol):
@@ -149,7 +153,14 @@ def is_not_done(result: DataFrame) -> bool:
         def project_fn() -> DataFrame:
             termination_flag.assert_running()
             return projection_query_runner.call_procedure(
-                ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
+                ProtocolVersion.V3.versioned_procedure_name(endpoint),
+                params,
+                yields,
+                database,
+                database=database,
+                logging=logging,
+                retryable=True,
+                custom_error=False,
             )
 
         projection_result = project_fn()

graphdatascience/query_runner/protocol/write_protocols.py

@@ -5,9 +5,9 @@
 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_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 +73,10 @@ def run_write_back(
             ProtocolVersion.V1.versioned_procedure_name("gds.arrow.write"),
             parameters,
             yields,
-            None,
-            False,
-            False,
+            retryable=True,
+            database=None,
+            logging=False,
+            custom_error=False,
         )
 
 
@@ -111,9 +112,10 @@ def run_write_back(
             ProtocolVersion.V2.versioned_procedure_name("gds.arrow.write"),
             parameters,
             yields,
-            None,
-            False,
-            False,
+            retryable=True,
+            database=None,
+            logging=False,
+            custom_error=False,
         )
 
 
@@ -157,9 +159,10 @@ def write_fn() -> DataFrame:
                 ProtocolVersion.V3.versioned_procedure_name("gds.arrow.write"),
                 parameters,
                 yields,
-                None,
-                False,
-                False,
+                retryable=True,
+                database=None,
+                logging=False,
+                custom_error=False,
             )
 
         return write_fn()

graphdatascience/query_runner/query_runner.py

@@ -17,6 +17,7 @@ def call_procedure(
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
+        retryable: bool = False,
         custom_error: bool = True,
     ) -> DataFrame:
         pass
@@ -35,6 +36,16 @@ def run_cypher(
     ) -> DataFrame:
         pass
 
+    @abstractmethod
+    def run_retryable_cypher(
+        self,
+        query: str,
+        params: Optional[dict[str, Any]] = None,
+        database: Optional[str] = None,
+        custom_error: bool = True,
+    ) -> DataFrame:
+        pass
+
     @abstractmethod
     def server_version(self) -> ServerVersion:
         pass