Merge pull request #141 from FalkorDB/propagate-udf

propagate udf commands

Author: swilly22

falkordb/asyncio/falkordb.py

@@ -1,9 +1,10 @@
 import redis.asyncio as redis
 from .cluster import *
 from .graph import AsyncGraph
-from typing import List, Union
+from typing import List, Union, Optional
 
 # config command
+UDF_CMD    = "GRAPH.UDF"
 LIST_CMD   = "GRAPH.LIST"
 CONFIG_CMD = "GRAPH.CONFIG"
 
@@ -197,3 +198,93 @@ async def config_set(self, name: str, value=None) -> None:
         """
 
         return await self.connection.execute_command(CONFIG_CMD, "SET", name, value)
+
+    # GRAPH.UDF LOAD [REPLACE] <lib> <script>
+    async def udf_load(self, name: str, script: str, replace: bool = False):
+        """
+        Load a User Defined Function (UDF) library.
+
+        Args:
+            name (str): The name of the library to load.
+            script (str): The UDF script contents.
+            replace (bool, optional): If True, replace an existing library with the same name.
+                                      Defaults to False.
+        """
+
+        # prep arguments
+        args = [UDF_CMD, "LOAD"]
+        if replace:
+            args.append("REPLACE")
+        args.extend([name, script])
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = await client.execute_command(*args)
+        else:
+            resp = await self.connection.execute_command(*args)
+
+        return resp
+
+    # GRAPH.UDF LIST [LIBRARYNAME] [WITHCODE]
+    async def udf_list(self, lib: Optional[str] = None, with_code: bool = False):
+        """
+        List User Defined Function (UDF) libraries.
+
+        Args:
+            lib (str, optional): If provided, filter the list to this specific library.
+            with_code (bool, optional): If True, include the library source code in the result.
+                                        Defaults to False.
+
+        Returns:
+            list: A list of UDF libraries and their metadata.
+        """
+
+        args = [UDF_CMD, "LIST"]
+        if lib is not None:
+            args.append(lib)
+
+        if with_code:
+            args.append("WITHCODE")
+
+        return await self.connection.execute_command(*args)
+
+    # GRAPH.UDF FLUSH
+    async def udf_flush(self):
+        """
+        Flush (remove) all User Defined Function (UDF) libraries.
+        """
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = await client.execute_command(UDF_CMD, "FLUSH")
+        else:
+            resp = await self.connection.execute_command(UDF_CMD, "FLUSH")
+
+        return resp
+
+    # GRAPH.UDF DELETE <lib>
+    async def udf_delete(self, lib: str):
+        """
+        Delete a User Defined Function (UDF) library.
+
+        Args:
+            lib (str): The name of the library to delete.
+        """
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = await client.execute_command(UDF_CMD, "DELETE", lib)
+        else:
+            resp = await self.connection.execute_command(UDF_CMD, "DELETE", lib)
+
+        return resp
+

falkordb/falkordb.py

@@ -5,9 +5,9 @@
 from typing import List, Union, Optional
 
 # config commands
+UDF_CMD    = "GRAPH.UDF"
 LIST_CMD   = "GRAPH.LIST"
 CONFIG_CMD = "GRAPH.CONFIG"
-UDF_CMD    = "GRAPH.UDF"
 
 class FalkorDB:
     """
@@ -249,11 +249,23 @@ def udf_load(self, name: str, script: str, replace: bool = False):
             replace (bool, optional): If True, replace an existing library with the same name.
                                       Defaults to False.
         """
+
+        # prep arguments
         args = [UDF_CMD, "LOAD"]
         if replace:
             args.append("REPLACE")
         args.extend([name, script])
-        return self.connection.execute_command(*args)
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = client.execute_command(*args)
+        else:
+            resp = self.connection.execute_command(*args)
+
+        return resp
 
     # GRAPH.UDF LIST [LIBRARYNAME] [WITHCODE]
     def udf_list(self, lib: Optional[str] = None, with_code: bool = False):
@@ -268,19 +280,32 @@ def udf_list(self, lib: Optional[str] = None, with_code: bool = False):
         Returns:
             list: A list of UDF libraries and their metadata.
         """
+
         args = [UDF_CMD, "LIST"]
         if lib is not None:
             args.append(lib)
+
         if with_code:
             args.append("WITHCODE")
+
         return self.connection.execute_command(*args)
 
     # GRAPH.UDF FLUSH
     def udf_flush(self):
         """
         Flush (remove) all User Defined Function (UDF) libraries.
         """
-        return self.connection.execute_command(UDF_CMD, "FLUSH")
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = client.execute_command(UDF_CMD, "FLUSH")
+        else:
+            resp = self.connection.execute_command(UDF_CMD, "FLUSH")
+
+        return resp
 
     # GRAPH.UDF DELETE <lib>
     def udf_delete(self, lib: str):
@@ -290,5 +315,15 @@ def udf_delete(self, lib: str):
         Args:
             lib (str): The name of the library to delete.
         """
-        return self.connection.execute_command(UDF_CMD, "DELETE", lib)
+
+        # propagate command in cluster mode
+        if Is_Cluster(self.connection):
+            for node in self.connection.get_primaries():
+                # create a direct connection to this node
+                client = self.connection.get_redis_connection(node)
+                resp = client.execute_command(UDF_CMD, "DELETE", lib)
+        else:
+            resp = self.connection.execute_command(UDF_CMD, "DELETE", lib)
+
+        return resp
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "FalkorDB"
-version = "1.2.1"
+version = "1.2.2"
 description = "Python client for interacting with FalkorDB database"
 authors = ["FalkorDB inc <info@falkordb.com>"]