From 187b2ba42839b3b06f528488b2ef244852a7179b Mon Sep 17 00:00:00 2001 From: Yi-Pin Chen Date: Thu, 27 Jun 2024 15:03:45 -0700 Subject: [PATCH] Python: add FUNCTION LOAD command (#1699) * Python: Added FUNCTION LOAD command * Addressed review comments * Fixed CI/CD build error --------- Co-authored-by: Shoham Elias Co-authored-by: Andrew Carbonetto --- CHANGELOG.md | 1 + .../glide/async_commands/cluster_commands.py | 34 +++++++ .../async_commands/standalone_commands.py | 29 ++++++ .../glide/async_commands/transaction.py | 23 +++++ python/python/tests/test_async_client.py | 97 +++++++++++++++++++ python/python/tests/test_transaction.py | 15 ++- python/python/tests/utils/utils.py | 15 +++ 7 files changed, 213 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bcb350c6e..6bdce0d49d 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/python/python/glide/async_commands/cluster_commands.py b/python/python/glide/async_commands/cluster_commands.py index f2f17a01c6..9294c1c8ba 100644 --- a/python/python/glide/async_commands/cluster_commands.py +++ b/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. diff --git a/python/python/glide/async_commands/standalone_commands.py b/python/python/glide/async_commands/standalone_commands.py index cd5518d417..497fbdd901 100644 --- a/python/python/glide/async_commands/standalone_commands.py +++ b/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. diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 71ca2f2e18..5add3172d6 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/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, diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 11155191ea..a85ead46f3 100644 --- a/python/python/tests/test_async_client.py +++ b/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): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 89673c7dd3..15638826a3 100644 --- a/python/python/tests/test_transaction.py +++ b/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) diff --git a/python/python/tests/utils/utils.py b/python/python/tests/utils/utils.py index 6bbde86f86..f97bcc14e6 100644 --- a/python/python/tests/utils/utils.py +++ b/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