WIP

Author: Mats-SX

graphdatascience/graph/graph_remote_project_runner.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any
+
 from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..query_runner.session_query_runner import SessionQueryRunner
@@ -21,6 +23,7 @@ def __call__(
         inverse_indexed_relationship_types: list[str] | None = None,
         batch_size: int | None = None,
         logging: bool = True,
+        query_parameters: dict[str, Any] = None,
     ) -> GraphCreateResult:
         if inverse_indexed_relationship_types is None:
             inverse_indexed_relationship_types = []

graphdatascience/query_runner/protocol/project_protocols.py

@@ -113,6 +113,7 @@ def project_params(
         config = {
             "undirectedRelationshipTypes": params["undirected_relationship_types"],
             "inverseIndexedRelationshipTypes": params["inverse_indexed_relationship_types"],
+            "queryParameters": params["query_parameters"],
         }
         if "concurrency" in params:
             config["concurrency"] = params["concurrency"]

graphdatascience/tests/unit/session/test_aura_graph_data_science.py

@@ -70,6 +70,35 @@ def test_remote_projection_defaults(mocker: MockerFixture) -> None:
     }
 
 
+def test_remote_projection_query_parameters(mocker: MockerFixture) -> None:
+    v = ServerVersion(9, 9, 9)
+    query_runner = CollectingQueryRunner(v)
+    gds = AuraGraphDataScience(
+        query_runner=query_runner,
+        delete_fn=lambda: True,
+        gds_version=v,
+        v2_endpoints=mocker.Mock(),
+    )
+
+    g = gds.graph.project("foo", "RETURN gds.graph.project($one, $two)", query_parameters={"one": 1, "two": 2})
+
+    assert g.graph.name() == "foo"
+    assert (
+        query_runner.last_query()
+        == "CALL gds.arrow.project($graph_name, $query, $job_id, $concurrency, $undirected_relationship_types, $inverse_indexed_relationship_types, $arrow_configuration)"
+    )
+    assert query_runner.last_params() == {
+        "graph_name": "foo",
+        "query": "RETURN gds.graph.project($one, $two)",
+        "job_id": None,
+        "concurrency": 4,
+        "undirected_relationship_types": [],
+        "inverse_indexed_relationship_types": [],
+        "arrow_configuration": {},
+        "query_parameters": {"one": 1, "two": 2},
+    }
+
+
 def test_remote_algo_write(mocker: MockerFixture) -> None:
     v = ServerVersion(9, 9, 9)
     query_runner = CollectingQueryRunner(v)