Add query mode to call_procedure and call_function

Author: FlorentinD

graphdatascience/query_runner/arrow_query_runner.py

@@ -6,6 +6,7 @@
 from pandas import DataFrame
 
 from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
+from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.retry_utils.retry_config import RetryConfig
 
 from ..call_parameters import CallParameters
@@ -86,6 +87,7 @@ def call_procedure(
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
         retryable: bool = False,
         custom_error: bool = True,

graphdatascience/query_runner/neo4j_query_runner.py

@@ -9,6 +9,8 @@
 import neo4j
 from pandas import DataFrame
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..call_parameters import CallParameters
 from ..error.endpoint_suggester import generate_suggestive_error_message
 from ..error.gds_not_installed import GdsNotFound
@@ -238,14 +240,17 @@ def call_function(self, endpoint: str, params: Optional[CallParameters] = None,
         query = f"RETURN {endpoint}({params.placeholder_str()})"
 
         # 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()
+        return self.run_retryable_cypher(
+            query, params, custom_error=custom_error, routing=neo4j.RoutingControl.READ
+        ).squeeze()
 
     def call_procedure(
         self,
         endpoint: str,
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
         retryable: bool = False,
         custom_error: bool = True,
@@ -258,8 +263,7 @@ def call_procedure(
 
         def run_cypher_query() -> DataFrame:
             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)
+                return self.run_retryable_cypher(query, params, database, custom_error, routing=mode.neo4j_mode())
             else:
                 return self.run_cypher(query, params, database, custom_error)
 

graphdatascience/query_runner/protocol/project_protocols.py

@@ -138,7 +138,12 @@ def is_not_done(result: DataFrame) -> bool:
 
         # 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
+            ProtocolVersion.V3.versioned_procedure_name(endpoint),
+            params,
+            yields,
+            database,
+            logging=logging,
+            custom_error=False,
         ).squeeze()
         member_host = response["host"]
         member_port = response["port"] if ("port" in response.index) else 7687

graphdatascience/query_runner/protocol/write_protocols.py

@@ -7,6 +7,7 @@
 
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.query_runner.protocol.status import Status
+from graphdatascience.query_runner.query_mode import QueryMode
 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
@@ -76,6 +77,7 @@ def run_write_back(
             retryable=False,
             database=None,
             logging=False,
+            mode=QueryMode.WRITE,
             custom_error=False,
         )
 
@@ -115,6 +117,7 @@ def run_write_back(
             retryable=False,
             database=None,
             logging=False,
+            mode=QueryMode.WRITE,
             custom_error=False,
         )
 
@@ -161,6 +164,7 @@ def write_fn() -> DataFrame:
                 yields,
                 retryable=True,
                 logging=False,
+                mode=QueryMode.WRITE,
                 custom_error=False,
             )
 

graphdatascience/query_runner/query_mode.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import neo4j
+import neo4j.routing
+
+
+class QueryMode(str, Enum):
+    READ = "read"
+    WRITE = "write"
+
+    def neo4j_mode(self) -> neo4j.RoutingControl:
+        if self == QueryMode.READ:
+            return neo4j.RoutingControl.READ
+        elif self == QueryMode.WRITE:
+            return neo4j.RoutingControl.WRITE
+        else:
+            raise ValueError(f"Unknown query mode: {self}")

graphdatascience/query_runner/query_runner.py

@@ -3,6 +3,8 @@
 
 from pandas import DataFrame
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
 from .graph_constructor import GraphConstructor
@@ -16,6 +18,7 @@ def call_procedure(
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
         retryable: bool = False,
         custom_error: bool = True,

graphdatascience/query_runner/session_query_runner.py

@@ -8,6 +8,7 @@
 
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
 from graphdatascience.query_runner.progress.query_progress_logger import QueryProgressLogger
+from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.server_version.server_version import ServerVersion
 
@@ -75,6 +76,7 @@ def call_procedure(
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
         retryable: bool = False,
         custom_error: bool = True,
@@ -206,7 +208,7 @@ def _remote_write_back(
         config["writeToResultStore"] = True
 
         gds_write_result = self._gds_query_runner.call_procedure(
-            endpoint, params, yields, database, logging, custom_error
+            endpoint, params, yields, database=database, logging=logging, custom_error=custom_error
         )
         terminationFlag.assert_running()
 

graphdatascience/query_runner/standalone_session_query_runner.py

@@ -7,6 +7,7 @@
 from graphdatascience import QueryRunner, ServerVersion
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
+from graphdatascience.query_runner.query_mode import QueryMode
 
 
 class StandaloneSessionQueryRunner(QueryRunner):
@@ -19,6 +20,7 @@ def call_procedure(
         params: Optional[CallParameters] = None,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
+        mode: QueryMode = QueryMode.READ,
         logging: bool = False,
         retryable: bool = False,
         custom_error: bool = True,

graphdatascience/tests/integration/test_remote_graph_ops.py

@@ -36,7 +36,7 @@ def run_around_tests(gds_with_cloud_setup: AuraGraphDataScience) -> Generator[No
         gds_with_cloud_setup.graph.get(graph_name).drop(failIfMissing=True)
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_projection(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
@@ -87,7 +87,7 @@ def test_remote_projection_and_writeback_custom_database_name(gds_with_cloud_set
         gds_with_cloud_setup.set_database("neo4j")
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_projection_with_small_batch_size(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -98,7 +98,7 @@ def test_remote_projection_with_small_batch_size(gds_with_cloud_setup: AuraGraph
     assert result["nodeCount"] == 3
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_page_rank(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
@@ -108,7 +108,7 @@ def test_remote_write_back_page_rank(gds_with_cloud_setup: AuraGraphDataScience)
     assert result["nodePropertiesWritten"] == 3
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_node_similarity(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
@@ -120,7 +120,7 @@ def test_remote_write_back_node_similarity(gds_with_cloud_setup: AuraGraphDataSc
     assert result["relationshipsWritten"] == 2
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_node_properties(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
@@ -130,7 +130,7 @@ def test_remote_write_back_node_properties(gds_with_cloud_setup: AuraGraphDataSc
     assert result["propertiesWritten"] == 3
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_node_properties_with_multiple_labels(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -143,7 +143,7 @@ def test_remote_write_back_node_properties_with_multiple_labels(gds_with_cloud_s
     assert result["nodePropertiesWritten"] == 3
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_node_properties_with_select_labels(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -157,7 +157,7 @@ def test_remote_write_back_node_properties_with_select_labels(gds_with_cloud_set
     assert result["propertiesWritten"] == 1
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_node_label(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
@@ -166,7 +166,7 @@ def test_remote_write_back_node_label(gds_with_cloud_setup: AuraGraphDataScience
     assert result["nodeLabelsWritten"] == 3
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_relationship_topology(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -177,7 +177,7 @@ def test_remote_write_back_relationship_topology(gds_with_cloud_setup: AuraGraph
     assert result["relationshipsWritten"] == 4
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_relationship_property(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -190,7 +190,7 @@ def test_remote_write_back_relationship_property(gds_with_cloud_setup: AuraGraph
     assert result["relationshipsWritten"] == 4
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_relationship_properties(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     G, result = gds_with_cloud_setup.graph.project(
@@ -207,7 +207,7 @@ def test_remote_write_back_relationship_properties(gds_with_cloud_setup: AuraGra
     assert result["relationshipsWritten"] == 4
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_write_back_relationship_property_from_pathfinding_algo(
     gds_with_cloud_setup: AuraGraphDataScience,
@@ -223,7 +223,7 @@ def test_remote_write_back_relationship_property_from_pathfinding_algo(
     assert result["relationshipsWritten"] == 1
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_empty_graph_write_back(
     gds_with_cloud_setup: AuraGraphDataScience,

graphdatascience/tests/integration/test_remote_util_ops.py

@@ -50,15 +50,15 @@ def G(gds_with_cloud_setup: AuraGraphDataScience) -> Generator[Graph, None, None
     gds_with_cloud_setup.run_cypher("MATCH (n) DETACH DELETE n")
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_util_as_node(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     id = gds_with_cloud_setup.find_node_id(["Location"], {"name": "A"})
     result = gds_with_cloud_setup.util.asNode(id)
     assert result["name"] == "A"
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_remote_util_as_nodes(gds_with_cloud_setup: AuraGraphDataScience) -> None:
     ids = [
@@ -69,7 +69,7 @@ def test_remote_util_as_nodes(gds_with_cloud_setup: AuraGraphDataScience) -> Non
     assert len(result) == 2
 
 
-@pytest.mark.cloud_architecture
+# @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 7, 0))
 def test_util_nodeProperty(gds_with_cloud_setup: AuraGraphDataScience, G: Graph) -> None:
     id = gds_with_cloud_setup.find_node_id(["Location"], {"name": "A"})