Merge pull request #887 from DarthMax/session_remote_projection_fix

Session remote projection fix

Author: DarthMax

changelog.md

@@ -7,6 +7,8 @@
 
 ## Bug fixes
 
+* Fixed a bug where remote projections would fail when the database is clustered
+
 ## Improvements
 
 * Allow creating sessions of size `512GB`.

graphdatascience/query_runner/arrow_query_runner.py

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

graphdatascience/query_runner/neo4j_query_runner.py

@@ -46,6 +46,8 @@ def create_for_db(
 
             query_runner = Neo4jQueryRunner(
                 driver,
+                Neo4jQueryRunner.parse_protocol(endpoint),
+                auth,
                 auto_close=True,
                 bookmarks=bookmarks,
                 config=config,
@@ -54,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)}")
@@ -76,6 +84,8 @@ def create_for_session(
 
         query_runner = Neo4jQueryRunner(
             driver,
+            Neo4jQueryRunner.parse_protocol(endpoint),
+            auth,
             auto_close=True,
             show_progress=show_progress,
             bookmarks=None,
@@ -94,9 +104,18 @@ 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,
         auto_close: bool = False,
@@ -105,6 +124,8 @@ def __init__(
         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
@@ -279,6 +300,22 @@ 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)
+        driver = neo4j.GraphDatabase.driver(endpoint, auth=self._auth, **self.driver_config())
+
+        return Neo4jQueryRunner(
+            driver,
+            self._protocol,
+            self._auth,
+            self._config,
+            self._database,
+            self._auto_close,
+            self._bookmarks,
+            self._show_progress,
+            self._instance_description,
+        )
+
     @staticmethod
     def handle_driver_exception(session: neo4j.Session, e: Exception) -> None:
         reg_gds_hit = re.search(

graphdatascience/query_runner/protocol/project_protocols.py

@@ -121,7 +121,7 @@ def run_projection(
         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,16 +132,28 @@ def is_not_done(result: DataFrame) -> bool:
 
         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, False
+        ).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)
+
         @retry(
             reraise=True,
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
             retry=retry_if_result(is_not_done),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
         )
         def project_fn() -> DataFrame:
-            terminationFlag.assert_running()
-            return query_runner.call_procedure(
+            termination_flag.assert_running()
+            return projection_query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
             )
 
-        return project_fn()
+        projection_result = project_fn()
+
+        projection_query_runner.close()
+
+        return projection_result

graphdatascience/query_runner/query_runner.py

@@ -80,5 +80,9 @@ def last_bookmarks(self) -> Optional[Any]:
     def set_show_progress(self, show_progress: bool) -> None:
         pass
 
+    @abstractmethod
+    def clone(self, host: str, port: int) -> "QueryRunner":
+        pass
+
     def set_server_version(self, _: ServerVersion) -> None:
         pass

graphdatascience/query_runner/session_query_runner.py

@@ -120,6 +120,14 @@ 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:
+        return SessionQueryRunner(
+            self._gds_query_runner,
+            self._db_query_runner.clone(host, port),
+            self._gds_arrow_client,
+            self._show_progress,
+        )
+
     def close(self) -> None:
         self._gds_arrow_client.close()
         self._gds_query_runner.close()

graphdatascience/query_runner/standalone_session_query_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Optional
 
 from pandas import DataFrame
@@ -74,3 +76,6 @@ 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:
+        return self

graphdatascience/tests/unit/conftest.py

@@ -126,6 +126,9 @@ def create_graph_constructor(
             self, graph_name, concurrency, undirected_relationship_types, self._server_version
         )
 
+    def clone(self, host: str, port: int) -> QueryRunner:
+        return self
+
     def set__mock_result(self, result: DataFrame) -> None:
         self._result_map.clear()
         self._result_map[""] = result