Support call_functions in QueryRunner

Allows to route to GDS session

Author: FlorentinD

graphdatascience/query_runner/arrow_query_runner.py

@@ -64,6 +64,9 @@ def run_cypher(
     ) -> DataFrame:
         return self._fallback_query_runner.run_cypher(query, params, database, custom_error)
 
+    def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
+        return self._fallback_query_runner.call_function(endpoint, params)
+
     def call_procedure(
         self,
         endpoint: str,

graphdatascience/query_runner/aura_db_query_runner.py

@@ -35,6 +35,9 @@ def run_cypher(
     ) -> DataFrame:
         return self._db_query_runner.run_cypher(query, params, database, custom_error)
 
+    def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
+        return self._gds_query_runner.call_function(endpoint, params)
+
     def call_procedure(
         self,
         endpoint: str,

graphdatascience/query_runner/neo4j_query_runner.py

@@ -131,6 +131,13 @@ def run_cypher(
 
             return df
 
+    def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
+        if params is None:
+            params = CallParameters()
+        query = f"RETURN {endpoint}({params.placeholder_str()})"
+
+        return self.run_cypher(query, params).squeeze()
+
     def call_procedure(
         self,
         endpoint: str,

graphdatascience/query_runner/query_runner.py

@@ -21,6 +21,10 @@ def call_procedure(
     ) -> DataFrame:
         pass
 
+    @abstractmethod
+    def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
+        pass
+
     @abstractmethod
     def run_cypher(
         self,

graphdatascience/session/aura_graph_data_science.py

@@ -3,7 +3,11 @@
 from pandas import DataFrame
 
 from graphdatascience.call_builder import IndirectCallBuilder
-from graphdatascience.endpoints import AlphaRemoteEndpoints, BetaEndpoints, DirectEndpoints
+from graphdatascience.endpoints import (
+    AlphaRemoteEndpoints,
+    BetaEndpoints,
+    DirectEndpoints,
+)
 from graphdatascience.error.uncallable_namespace import UncallableNamespace
 from graphdatascience.graph.graph_remote_proc_runner import GraphRemoteProcRunner
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
@@ -12,6 +16,7 @@
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.server_version.server_version import ServerVersion
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
+from graphdatascience.utils.util_remote_proc_runner import UtilRemoteProcRunner
 
 
 class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
@@ -75,7 +80,7 @@ def __init__(
 
         self._delete_fn = delete_fn
 
-        super().__init__(self._query_runner, "gds", self._server_version)
+        super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
 
     def run_cypher(
         self, query: str, params: Optional[Dict[str, Any]] = None, database: Optional[str] = None
@@ -102,6 +107,10 @@ def run_cypher(
     def graph(self) -> GraphRemoteProcRunner:
         return GraphRemoteProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
 
+    @property
+    def util(self) -> UtilRemoteProcRunner:
+        return UtilRemoteProcRunner(self._query_runner, f"{self._namespace}.util", self._server_version)
+
     @property
     def alpha(self) -> AlphaRemoteEndpoints:
         return AlphaRemoteEndpoints(self._query_runner, "gds.alpha", self._server_version)