Use retries but disable routing by manipulating protocol

Co-authored-by: Florentin Dörre <florentin.dorre@neo4j.com>
Co-authored-by: Max Kießling <max.kiessling@neotechnology.com>

Author: 3 people

graphdatascience/query_runner/arrow_query_runner.py

@@ -225,10 +225,10 @@ def close(self) -> None:
         self._fallback_query_runner.close()
         self._gds_arrow_client.close()
 
-    def clone(self, host: str, port: int) -> "QueryRunner":
+    def cloneWithoutRouting(self, host: str, port: int) -> "QueryRunner":
         return ArrowQueryRunner(
             self._gds_arrow_client,
-            self._fallback_query_runner.clone(host, port),
+            self._fallback_query_runner.cloneWithoutRouting(host, port),
             self._server_version,
         )
 

graphdatascience/query_runner/neo4j_query_runner.py

@@ -355,13 +355,14 @@ def create_graph_constructor(
     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)
+    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=self._protocol,
+            protocol=protocol,
             auth=self._auth,
             config=self._config,
             database=self._database,

graphdatascience/query_runner/protocol/project_protocols.py

@@ -144,10 +144,11 @@ def is_not_done(result: DataFrame) -> bool:
             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.clone(member_host, member_port)
+        projection_query_runner = query_runner.cloneWithoutRouting(member_host, member_port)
 
         @retry(
             reraise=True,
@@ -163,7 +164,7 @@ def project_fn() -> DataFrame:
                 yields,
                 database=database,
                 logging=logging,
-                retryable=False,
+                retryable=True,
                 custom_error=False,
             )
 

graphdatascience/query_runner/query_runner.py

@@ -95,7 +95,7 @@ def set_show_progress(self, show_progress: bool) -> None:
         pass
 
     @abstractmethod
-    def clone(self, host: str, port: int) -> "QueryRunner":
+    def cloneWithoutRouting(self, host: str, port: int) -> "QueryRunner":
         pass
 
     def set_server_version(self, _: ServerVersion) -> None:

graphdatascience/query_runner/session_query_runner.py

@@ -137,10 +137,10 @@ def set_show_progress(self, show_progress: bool) -> None:
         self._show_progress = show_progress
         self._gds_query_runner.set_show_progress(show_progress)
 
-    def clone(self, host: str, port: int) -> QueryRunner:
+    def cloneWithoutRouting(self, host: str, port: int) -> QueryRunner:
         return SessionQueryRunner(
             self._gds_query_runner,
-            self._db_query_runner.clone(host, port),
+            self._db_query_runner.cloneWithoutRouting(host, port),
             self._gds_arrow_client,
             self._show_progress,
         )

graphdatascience/query_runner/standalone_session_query_runner.py

@@ -91,5 +91,5 @@ def last_bookmarks(self) -> Optional[Any]:
     def set_server_version(self, _: ServerVersion) -> None:
         super().set_server_version(_)
 
-    def clone(self, host: str, port: int) -> QueryRunner:
+    def cloneWithoutRouting(self, host: str, port: int) -> QueryRunner:
         return self

graphdatascience/tests/unit/conftest.py

@@ -138,7 +138,7 @@ def create_graph_constructor(
             self, graph_name, concurrency, undirected_relationship_types, self._server_version
         )
 
-    def clone(self, host: str, port: int) -> QueryRunner:
+    def cloneWithoutRouting(self, host: str, port: int) -> QueryRunner:
         return self
 
     def set__mock_result(self, result: DataFrame) -> None: