From 2ea5d6caddd3db02baa51e46a5f37836cf6ff61e Mon Sep 17 00:00:00 2001 From: Adan Wattad Date: Thu, 1 Feb 2024 13:07:02 +0200 Subject: [PATCH] Added Script API in python. (#860) * Added Script API in python. * Change the name of the execute functions --- node/src/BaseClient.ts | 2 +- python/python/glide/__init__.py | 3 ++ .../glide/async_commands/cluster_commands.py | 2 +- python/python/glide/async_commands/core.py | 40 ++++++++++++++++++- .../async_commands/standalone_commands.py | 2 +- python/python/glide/glide.pyi | 5 +++ python/python/glide/redis_client.py | 31 ++++++++++---- python/python/tests/test_async_client.py | 27 ++++++++++++- python/src/lib.rs | 23 +++++++++++ 9 files changed, 123 insertions(+), 12 deletions(-) diff --git a/node/src/BaseClient.ts b/node/src/BaseClient.ts index 2f2b71353e..a5e909f12f 100644 --- a/node/src/BaseClient.ts +++ b/node/src/BaseClient.ts @@ -958,7 +958,7 @@ export class BaseClient { * @returns a value that depends on the script that was executed. * * @example - * const luaScript = "return \{ KEYS[1], ARGV[1] \}"; + * const luaScript = new Script("return \{ KEYS[1], ARGV[1] \}"); * const scriptOptions = \{ * keys: ["foo"], * args: ["bar"], diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index 4d2d4cafaf..eb187a2c1c 100644 --- a/python/python/glide/__init__.py +++ b/python/python/glide/__init__.py @@ -37,6 +37,8 @@ SlotType, ) +from .glide import Script + __all__ = [ "BaseClientConfiguration", "ClusterClientConfiguration", @@ -54,6 +56,7 @@ "RedisClient", "RedisClusterClient", "RedisCredentials", + "Script", "NodeAddress", "Transaction", "ClusterTransaction", diff --git a/python/python/glide/async_commands/cluster_commands.py b/python/python/glide/async_commands/cluster_commands.py index c43efd48cd..00bb1376d4 100644 --- a/python/python/glide/async_commands/cluster_commands.py +++ b/python/python/glide/async_commands/cluster_commands.py @@ -80,7 +80,7 @@ async def exec( If the transaction failed due to a WATCH command, `exec` will return `None`. """ commands = transaction.commands[:] - return await self.execute_transaction(commands, route) + return await self._execute_transaction(commands, route) async def config_resetstat( self, diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 5a27d79b8a..a699fd475f 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -20,6 +20,8 @@ from glide.protobuf.redis_request_pb2 import RequestType from glide.routes import Route +from ..glide import Script + class ConditionalChange(Enum): """ @@ -182,12 +184,20 @@ async def _execute_command( route: Optional[Route] = ..., ) -> TResult: ... - async def execute_transaction( + async def _execute_transaction( self, commands: List[Tuple[RequestType.ValueType, List[str]]], route: Optional[Route] = None, ) -> List[TResult]: ... + async def _execute_script( + self, + hash: str, + keys: Optional[List[str]] = None, + args: Optional[List[str]] = None, + route: Optional[Route] = None, + ) -> TResult: ... + async def set( self, key: str, @@ -1234,3 +1244,31 @@ async def zrem( int, await self._execute_command(RequestType.Zrem, [key] + members), ) + + async def invoke_script( + self, + script: Script, + keys: Optional[List[str]] = None, + args: Optional[List[str]] = None, + ) -> TResult: + """ + Invokes a Lua script with its keys and arguments. + This method simplifies the process of invoking scripts on a Redis server by using an object that represents a Lua script. + The script loading, argument preparation, and execution will all be handled internally. + If the script has not already been loaded, it will be loaded automatically using the Redis `SCRIPT LOAD` command. + After that, it will be invoked using the Redis `EVALSHA` command. + + Args: + script (Script): The Lua script to execute. + keys (List[str]): The keys that are used in the script. + args (List[str]): The arguments for the script. + + Returns: + TResult: a value that depends on the script that was executed. + + Examples: + >>> lua_script = Script("return { KEYS[1], ARGV[1] }") + >>> await invoke_script(lua_script, keys=["foo"], args=["bar"] ); + ["foo", "bar"] + """ + return await self._execute_script(script.get_hash(), keys, args) diff --git a/python/python/glide/async_commands/standalone_commands.py b/python/python/glide/async_commands/standalone_commands.py index 907cd60e67..77c731e1cc 100644 --- a/python/python/glide/async_commands/standalone_commands.py +++ b/python/python/glide/async_commands/standalone_commands.py @@ -61,7 +61,7 @@ async def exec( If the transaction failed due to a WATCH command, `exec` will return `None`. """ commands = transaction.commands[:] - return await self.execute_transaction(commands) + return await self._execute_transaction(commands) async def select(self, index: int) -> TOK: """Change the currently selected Redis database. diff --git a/python/python/glide/glide.pyi b/python/python/glide/glide.pyi index d3ff5a7d53..d155757bbd 100644 --- a/python/python/glide/glide.pyi +++ b/python/python/glide/glide.pyi @@ -15,6 +15,11 @@ class Level(Enum): def is_lower(self, level: Level) -> bool: ... +class Script: + def __init__(self, code: str) -> None: ... + def get_hash(self) -> str: ... + def __del__(self) -> None: ... + def start_socket_listener_external(init_callback: Callable) -> None: ... def value_from_pointer(pointer: int) -> TResult: ... def create_leaked_value(message: str) -> int: ... diff --git a/python/python/glide/redis_client.py b/python/python/glide/redis_client.py index b50dd050ba..989cb6b495 100644 --- a/python/python/glide/redis_client.py +++ b/python/python/glide/redis_client.py @@ -197,14 +197,9 @@ async def _execute_command( request.single_command.request_type = request_type request.single_command.args_array.args[:] = args # TODO - use arg pointer set_protobuf_route(request, route) - # Create a response future for this request and add it to the available - # futures map - response_future = self._get_future(request.callback_idx) - self._create_write_task(request) - await response_future - return response_future.result() + return await self._write_request_await_response(request) - async def execute_transaction( + async def _execute_transaction( self, commands: List[Tuple[RequestType.ValueType, List[str]]], route: Optional[Route] = None, @@ -223,6 +218,28 @@ async def execute_transaction( transaction_commands.append(command) request.transaction.commands.extend(transaction_commands) set_protobuf_route(request, route) + return await self._write_request_await_response(request) + + async def _execute_script( + self, + hash: str, + keys: Optional[List[str]] = None, + args: Optional[List[str]] = None, + route: Optional[Route] = None, + ) -> TResult: + if self._is_closed: + raise ClosingError( + "Unable to execute requests; the client is closed. Please create a new client." + ) + request = RedisRequest() + request.callback_idx = self._get_callback_index() + request.script_invocation.hash = hash + request.script_invocation.args[:] = args if args is not None else [] + request.script_invocation.keys[:] = keys if keys is not None else [] + set_protobuf_route(request, route) + return await self._write_request_await_response(request) + + async def _write_request_await_response(self, request: RedisRequest): # Create a response future for this request and add it to the available # futures map response_future = self._get_future(request.callback_idx) diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 77737bfdc2..173ce6dd9d 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -10,7 +10,7 @@ from typing import Dict, List, TypeVar, Union, cast import pytest -from glide import ClosingError, RequestError, TimeoutError +from glide import ClosingError, RequestError, Script, TimeoutError from glide.async_commands.core import ( ConditionalChange, ExpireOptions, @@ -1229,3 +1229,28 @@ async def test_timeout_exception_with_blpop(self, redis_client: TRedisClient): key = get_random_string(10) with pytest.raises(TimeoutError) as e: await redis_client.custom_command(["BLPOP", key, "1"]) + + +@pytest.mark.asyncio +class TestScripts: + @pytest.mark.smoke_test + @pytest.mark.parametrize("cluster_mode", [True, False]) + async def test_script(self, redis_client: TRedisClient): + key1 = get_random_string(10) + key2 = get_random_string(10) + script = Script("return 'Hello'") + assert await redis_client.invoke_script(script) == "Hello" + + script = Script("return redis.call('SET', KEYS[1], ARGV[1])") + assert ( + await redis_client.invoke_script(script, keys=[key1], args=["value1"]) + == "OK" + ) + # Reuse the same script with different parameters. + assert ( + await redis_client.invoke_script(script, keys=[key2], args=["value2"]) + == "OK" + ) + script = Script("return redis.call('GET', KEYS[1])") + assert await redis_client.invoke_script(script, keys=[key1]) == "value1" + assert await redis_client.invoke_script(script, keys=[key2]) == "value2" diff --git a/python/src/lib.rs b/python/src/lib.rs index da738690a9..fd05126c17 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -30,10 +30,33 @@ impl Level { } } +#[pyclass] +pub struct Script { + hash: String, +} + +#[pymethods] +impl Script { + #[new] + fn new(code: String) -> Self { + let hash = glide_core::scripts_container::add_script(&code); + Script { hash } + } + + fn get_hash(&self) -> String { + self.hash.clone() + } + + fn __del__(&mut self) { + glide_core::scripts_container::remove_script(&self.hash); + } +} + /// A Python module implemented in Rust. #[pymodule] fn glide(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; + m.add_class::