Use the right protocol when cloning the query runner for remote projections

(cherry picked from commit 3285d74)

Author: DarthMax

graphdatascience/query_runner/arrow_query_runner.py

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

graphdatascience/query_runner/neo4j_query_runner.py

@@ -46,6 +46,7 @@ def create_for_db(
 
             query_runner = Neo4jQueryRunner(
                 driver,
+                Neo4jQueryRunner.parse_protocol(endpoint),
                 auth,
                 auto_close=True,
                 bookmarks=bookmarks,
@@ -55,8 +56,14 @@ def create_for_db(
             )
 
         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)}")
@@ -77,6 +84,7 @@ def create_for_session(
 
         query_runner = Neo4jQueryRunner(
             driver,
+            Neo4jQueryRunner.parse_protocol(endpoint),
             auth,
             auto_close=True,
             show_progress=show_progress,
@@ -96,9 +104,17 @@ def _configure_aura(config: dict[str, Any]) -> None:
         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,
@@ -108,6 +124,7 @@ def __init__(
         instance_description: str = "Neo4j DBMS",
     ):
         self._driver = driver
+        self._protocol = protocol
         self._auth = auth
         self._config = config
         self._auto_close = auto_close
@@ -283,10 +300,13 @@ def create_graph_constructor(
     def set_show_progress(self, show_progress: bool) -> None:
         self._show_progress = show_progress
 
-    def clone(self, endpoint: str) -> QueryRunner:
+    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,

graphdatascience/query_runner/protocol/project_protocols.py

@@ -136,7 +136,8 @@ def is_not_done(result: DataFrame) -> bool:
             ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
         ).squeeze()["host"]
 
-        projection_query_runner = query_runner.clone(f"neo4j+s://{member_address}:7687")
+        # TODO: retrieve the port from the server
+        projection_query_runner = query_runner.clone(member_address, 7687)
 
         @retry(
             reraise=True,

graphdatascience/query_runner/query_runner.py

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

graphdatascience/query_runner/session_query_runner.py

@@ -120,10 +120,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, endpoint: str) -> QueryRunner:
+    def clone(self, host: str, port: int) -> QueryRunner:
         return SessionQueryRunner(
-            self._gds_query_runner.clone(endpoint),
-            self._db_query_runner.clone(endpoint),
+            self._gds_query_runner,
+            self._db_query_runner.clone(host, port),
             self._gds_arrow_client,
             self._show_progress,
         )

graphdatascience/query_runner/standalone_session_query_runner.py

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

graphdatascience/tests/unit/conftest.py

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