Merge pull request #376 from soerenreichardt/arrow-remote-write-back

Arrow remote write back

Author: soerenreichardt

graphdatascience/query_runner/aura_db_arrow_query_runner.py

@@ -66,6 +66,17 @@ def run_query(
             params["host"] = aura_db_arrow_endpoint
             params["config"] = {"useEncryption": False}
 
+        elif ".write" in query and "config" in params and "remote" in params["config"] and params["config"]["remote"]:
+            token, aura_db_arrow_endpoint = self._get_or_request_auth_pair()
+            host, port_string = aura_db_arrow_endpoint.split(":")
+            del params["config"]["remote"]
+            params["config"]["arrowConnectionInfo"] = {
+                "hostname": host,
+                "port": int(port_string),
+                "bearerToken": token,
+                "useEncryption": False,
+            }
+
         return self._fallback_query_runner.run_query(query, params, database, custom_error)
 
     def set_database(self, database: str) -> None:

graphdatascience/tests/integration/test_remote_graph_ops.py

@@ -43,3 +43,15 @@ def test_remote_projection(gds_with_cloud_setup: GraphDataScience) -> None:
 
     assert G.name() == GRAPH_NAME
     assert result["nodeCount"] == 3
+
+
+@pytest.mark.cloud_architecture
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 4, 0))
+def test_remote_write_back(gds_with_cloud_setup: GraphDataScience) -> None:
+    G, result = gds_with_cloud_setup.alpha.graph.project.remote(
+        GRAPH_NAME, "MATCH (n)-->(m) RETURN n as sourceNode, m as targetNode", "neo4j"
+    )
+
+    result = gds_with_cloud_setup.pageRank.write(G, writeProperty="score", remote=True)
+
+    assert result["nodePropertiesWritten"] == 3