Python: add FUNCTION LOAD command (#1699)

* Python: Added FUNCTION LOAD command

* Addressed review comments

* Fixed CI/CD build error

---------

Co-authored-by: Shoham Elias <shohame@amazon.com>
Co-authored-by: Andrew Carbonetto <andrew.carbonetto@improving.com>

CHANGELOG.md

@@ -59,6 +59,7 @@
 * Python: Added XACK command ([#1681](https://github.com/aws/glide-for-redis/pull/1681))
 * Python: Added FLUSHDB command ([#1680](https://github.com/aws/glide-for-redis/pull/1680))
 * Python: Added XGROUP SETID command ([#1683](https://github.com/aws/glide-for-redis/pull/1683))
+* Python: Added FUNCTION LOAD command ([#1699](https://github.com/aws/glide-for-redis/pull/1699))
 
 ### Breaking Changes
 * Node: Update XREAD to return a Map of Map ([#1494](https://github.com/aws/glide-for-redis/pull/1494))

python/python/glide/async_commands/cluster_commands.py

@@ -319,6 +319,40 @@ async def echo(
             await self._execute_command(RequestType.Echo, [message], route),
         )
 
+    async def function_load(
+        self, library_code: str, replace: bool = False, route: Optional[Route] = None
+    ) -> str:
+        """
+        Loads a library to Redis.
+
+        See https://valkey.io/docs/latest/commands/function-load/ for more details.
+
+        Args:
+            library_code (str): The source code that implements the library.
+            replace (bool): Whether the given library should overwrite a library with the same name if
+                it already exists.
+            route (Optional[Route]): The command will be routed to all primaries, unless `route` is provided,
+                in which case the client will route the command to the nodes defined by `route`.
+
+        Returns:
+            str: The library name that was loaded.
+
+        Examples:
+            >>> code = "#!lua name=mylib \n redis.register_function('myfunc', function(keys, args) return args[1] end)"
+            >>> await client.function_load(code, True, RandomNode())
+                "mylib"
+
+        Since: Redis 7.0.0.
+        """
+        return cast(
+            str,
+            await self._execute_command(
+                RequestType.FunctionLoad,
+                ["REPLACE", library_code] if replace else [library_code],
+                route,
+            ),
+        )
+
     async def time(self, route: Optional[Route] = None) -> TClusterResponse[List[str]]:
         """
         Returns the server time.

python/python/glide/async_commands/standalone_commands.py

@@ -233,6 +233,35 @@ async def echo(self, message: str) -> str:
         """
         return cast(str, await self._execute_command(RequestType.Echo, [message]))
 
+    async def function_load(self, library_code: str, replace: bool = False) -> str:
+        """
+        Loads a library to Redis.
+
+        See https://valkey.io/docs/latest/commands/function-load/ for more details.
+
+        Args:
+            library_code (str): The source code that implements the library.
+            replace (bool): Whether the given library should overwrite a library with the same name if
+                it already exists.
+
+        Returns:
+            str: The library name that was loaded.
+
+        Examples:
+            >>> code = "#!lua name=mylib \n redis.register_function('myfunc', function(keys, args) return args[1] end)"
+            >>> await client.function_load(code, True)
+                "mylib"
+
+        Since: Redis 7.0.0.
+        """
+        return cast(
+            str,
+            await self._execute_command(
+                RequestType.FunctionLoad,
+                ["REPLACE", library_code] if replace else [library_code],
+            ),
+        )
+
     async def time(self) -> List[str]:
         """
         Returns the server time.

python/python/glide/async_commands/transaction.py

@@ -1776,6 +1776,29 @@ def type(self: TTransaction, key: str) -> TTransaction:
         """
         return self.append_command(RequestType.Type, [key])
 
+    def function_load(
+        self: TTransaction, library_code: str, replace: bool = False
+    ) -> TTransaction:
+        """
+        Loads a library to Redis.
+
+        See https://valkey.io/docs/latest/commands/function-load/ for more details.
+
+        Args:
+            library_code (str): The source code that implements the library.
+            replace (bool): Whether the given library should overwrite a library with the same name if
+                it already exists.
+
+        Commands response:
+            str: The library name that was loaded.
+
+        Since: Redis 7.0.0.
+        """
+        return self.append_command(
+            RequestType.FunctionLoad,
+            ["REPLACE", library_code] if replace else [library_code],
+        )
+
     def xadd(
         self: TTransaction,
         key: str,

python/python/tests/test_async_client.py

@@ -89,6 +89,7 @@
 from tests.utils.utils import (
     check_if_server_version_lt,
     compare_maps,
+    generate_lua_lib_code,
     get_first_result,
     get_random_string,
     is_single_response,
@@ -6299,6 +6300,102 @@ async def test_object_refcount(self, redis_client: TGlideClient):
         refcount = await redis_client.object_refcount(string_key)
         assert refcount is not None and refcount >= 0
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_function_load(self, redis_client: TGlideClient):
+        # TODO: Test function with FCALL
+        # TODO: Test with FUNCTION LIST
+        min_version = "7.0.0"
+        if await check_if_server_version_lt(redis_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        lib_name = f"mylib1C{get_random_string(5)}"
+        func_name = f"myfunc1c{get_random_string(5)}"
+        code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, True)
+
+        assert await redis_client.function_load(code) == lib_name
+
+        # TODO: change when FCALL, FCALL_RO is implemented
+        assert (
+            await redis_client.custom_command(["FCALL", func_name, "0", "one", "two"])
+            == "one"
+        )
+        assert (
+            await redis_client.custom_command(
+                ["FCALL_RO", func_name, "0", "one", "two"]
+            )
+            == "one"
+        )
+
+        # TODO: add FUNCTION LIST once implemented
+
+        # re-load library without replace
+        with pytest.raises(RequestError) as e:
+            await redis_client.function_load(code)
+        assert "Library '" + lib_name + "' already exists" in str(e)
+
+        # re-load library with replace
+        assert await redis_client.function_load(code, True) == lib_name
+
+        func2_name = f"myfunc2c{get_random_string(5)}"
+        new_code = generate_lua_lib_code(
+            lib_name, {func_name: "return args[1]", func2_name: "return #args"}, True
+        )
+
+        assert await redis_client.function_load(new_code, True) == lib_name
+
+    @pytest.mark.parametrize("cluster_mode", [True])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    @pytest.mark.parametrize("single_route", [True, False])
+    async def test_function_load_cluster_with_route(
+        self, redis_client: GlideClusterClient, single_route: bool
+    ):
+        # TODO: Test function with FCALL
+        # TODO: Test with FUNCTION LIST
+        min_version = "7.0.0"
+        if await check_if_server_version_lt(redis_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        lib_name = f"mylib1C{get_random_string(5)}"
+        func_name = f"myfunc1c{get_random_string(5)}"
+        code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, True)
+        route = SlotKeyRoute(SlotType.PRIMARY, "1") if single_route else AllPrimaries()
+
+        assert await redis_client.function_load(code, False, route) == lib_name
+
+        # TODO: change when FCALL, FCALL_RO is implemented.
+        assert (
+            await redis_client.custom_command(
+                ["FCALL", func_name, "0", "one", "two"],
+                SlotKeyRoute(SlotType.PRIMARY, "1"),
+            )
+            == "one"
+        )
+        assert (
+            await redis_client.custom_command(
+                ["FCALL_RO", func_name, "0", "one", "two"],
+                SlotKeyRoute(SlotType.PRIMARY, "1"),
+            )
+            == "one"
+        )
+
+        # TODO: add FUNCTION LIST once implemented
+
+        # re-load library without replace
+        with pytest.raises(RequestError) as e:
+            await redis_client.function_load(code, False, route)
+        assert "Library '" + lib_name + "' already exists" in str(e)
+
+        # re-load library with replace
+        assert await redis_client.function_load(code, True, route) == lib_name
+
+        func2_name = f"myfunc2c{get_random_string(5)}"
+        new_code = generate_lua_lib_code(
+            lib_name, {func_name: "return args[1]", func2_name: "return #args"}, True
+        )
+
+        assert await redis_client.function_load(new_code, True, route) == lib_name
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_srandmember(self, redis_client: TGlideClient):

python/python/tests/test_transaction.py

@@ -53,7 +53,11 @@
 from glide.constants import OK, TResult
 from glide.glide_client import GlideClient, GlideClusterClient, TGlideClient
 from tests.conftest import create_client
-from tests.utils.utils import check_if_server_version_lt, get_random_string
+from tests.utils.utils import (
+    check_if_server_version_lt,
+    generate_lua_lib_code,
+    get_random_string,
+)
 
 
 async def transaction_test(
@@ -86,8 +90,17 @@ async def transaction_test(
     value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S")
     value2 = get_random_string(5)
     value3 = get_random_string(5)
+    lib_name = f"mylib1C{get_random_string(5)}"
+    func_name = f"myfunc1c{get_random_string(5)}"
+    code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, True)
     args: List[TResult] = []
 
+    if not await check_if_server_version_lt(redis_client, "7.0.0"):
+        transaction.function_load(code)
+        args.append(lib_name)
+        transaction.function_load(code, True)
+        args.append(lib_name)
+
     transaction.dbsize()
     args.append(0)
 

python/python/tests/utils/utils.py

@@ -113,3 +113,18 @@ def compare_maps(
     if map1 is None or map2 is None:
         return False
     return json.dumps(map1) == json.dumps(map2)
+
+
+def generate_lua_lib_code(
+    lib_name: str, functions: Mapping[str, str], readonly: bool
+) -> str:
+    code = f"#!lua name={lib_name}\n"
+    for function_name, function_body in functions.items():
+        code += (
+            f"redis.register_function{{ function_name = '{function_name}', callback = function(keys, args) "
+            f"{function_body} end"
+        )
+        if readonly:
+            code += ", flags = { 'no-writes' }"
+        code += " }\n"
+    return code