diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index 40c47bffbc..6e3cac8875 100644 --- a/python/python/glide/__init__.py +++ b/python/python/glide/__init__.py @@ -1,8 +1,9 @@ from glide.async_commands.core import ( - ConditionalSet, + ConditionalChange, ExpireOptions, ExpirySet, ExpiryType, + UpdateOptions, ) from glide.async_commands.transaction import ClusterTransaction, Transaction from glide.config import ( @@ -37,10 +38,11 @@ "BaseClientConfiguration", "ClusterClientConfiguration", "RedisClientConfiguration", - "ConditionalSet", + "ConditionalChange", "ExpireOptions", "ExpirySet", "ExpiryType", + "UpdateOptions", "Logger", "LogLevel", "OK", diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index ef0cdef8d8..90d0674717 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -19,14 +19,15 @@ from glide.routes import Route -class ConditionalSet(Enum): - """SET option: A condition to the "SET" command. - - ONLY_IF_EXISTS - Only set the key if it already exist. Equivalent to `XX` in the Redis API - - ONLY_IF_DOES_NOT_EXIST - Only set the key if it does not already exist. Equivalent to `NX` in the Redis API +class ConditionalChange(Enum): + """ + A condition to the "SET" and "ZADD" commands. + - ONLY_IF_EXISTS - Only update key / elements that already exist. Equivalent to `XX` in the Redis API + - ONLY_IF_DOES_NOT_EXIST - Only set key / add elements that does not already exist. Equivalent to `NX` in the Redis API """ - ONLY_IF_EXISTS = 0 # Equivalent to `XX` in the Redis API - ONLY_IF_DOES_NOT_EXIST = 1 # Equivalent to `NX` in the Redis API + ONLY_IF_EXISTS = "XX" + ONLY_IF_DOES_NOT_EXIST = "NX" class ExpiryType(Enum): @@ -106,6 +107,18 @@ class ExpireOptions(Enum): NewExpiryLessThanCurrent = "LT" +class UpdateOptions(Enum): + """ + Options for updating elements of a sorted set key. + + - LESS_THAN: Only update existing elements if the new score is less than the current score. + - GREATER_THAN: Only update existing elements if the new score is greater than the current score. + """ + + LESS_THAN = "LT" + GREATER_THAN = "GT" + + class ExpirySet: """SET option: Represents the expiry type and value to be executed with "SET" command.""" @@ -179,7 +192,7 @@ async def set( self, key: str, value: str, - conditional_set: Union[ConditionalSet, None] = None, + conditional_set: Union[ConditionalChange, None] = None, expiry: Union[ExpirySet, None] = None, return_old_value: bool = False, ) -> TResult: @@ -188,11 +201,11 @@ async def set( @example - Set "foo" to "bar" only if "foo" already exists, and set the key expiration to 5 seconds: - connection.set("foo", "bar", conditional_set=ConditionalSet.ONLY_IF_EXISTS, expiry=Expiry(ExpiryType.SEC, 5)) + connection.set("foo", "bar", conditional_set=ConditionalChange.ONLY_IF_EXISTS, expiry=Expiry(ExpiryType.SEC, 5)) Args: key (str): the key to store. value (str): the value to store with the given key. - conditional_set (Union[ConditionalSet, None], optional): set the key only if the given condition is met. + conditional_set (Union[ConditionalChange, None], optional): set the key only if the given condition is met. Equivalent to [`XX` | `NX`] in the Redis API. Defaults to None. expiry (Union[Expiry, None], optional): set expiriation to the given key. Equivalent to [`EX` | `PX` | `EXAT` | `PXAT` | `KEEPTTL`] in the Redis API. Defaults to None. @@ -207,10 +220,7 @@ async def set( """ args = [key, value] if conditional_set: - if conditional_set == ConditionalSet.ONLY_IF_EXISTS: - args.append("XX") - if conditional_set == ConditionalSet.ONLY_IF_DOES_NOT_EXIST: - args.append("NX") + args.append(conditional_set.value) if return_old_value: args.append("GET") if expiry is not None: @@ -1019,3 +1029,122 @@ async def ttl(self, key: str) -> int: -2 # Returns -2 for a non-existing key. """ return cast(int, await self._execute_command(RequestType.TTL, [key])) + + async def zadd( + self, + key: str, + members_scores: Mapping[str, float], + existing_options: Optional[ConditionalChange] = None, + update_condition: Optional[UpdateOptions] = None, + changed: bool = False, + ) -> int: + """ + Adds members with their scores to the sorted set stored at `key`. + If a member is already a part of the sorted set, its score is updated. + + See https://redis.io/commands/zadd/ for more details. + + Args: + key (str): The key of the sorted set. + members_scores (Mapping[str, float]): A mapping of members to their corresponding scores. + existing_options (Optional[ConditionalChange]): Options for handling existing members. + - NX: Only add new elements. + - XX: Only update existing elements. + update_condition (Optional[UpdateOptions]): Options for updating scores. + - GT: Only update scores greater than the current values. + - LT: Only update scores less than the current values. + changed (bool): Modify the return value to return the number of changed elements, instead of the number of new elements added. + + Returns: + int: The number of elements added to the sorted set. + If `changed` is set, returns the number of elements updated in the sorted set. + + Examples: + >>> await zadd("my_sorted_set", {"member1": 10.5, "member2": 8.2}) + 2 # Indicates that two elements have been added or updated in the sorted set "my_sorted_set." + >>> await zadd("existing_sorted_set", {"member1": 15.0, "member2": 5.5}, existing_options=ConditionalChange.XX) + 2 # Updates the scores of two existing members in the sorted set "existing_sorted_set." + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if update_condition: + args.append(update_condition.value) + + if changed: + args.append("CH") + + if existing_options and update_condition: + if existing_options == ConditionalChange.ONLY_IF_DOES_NOT_EXIST: + raise ValueError( + "The GT, LT and NX options are mutually exclusive. " + f"Cannot choose both {update_condition.value} and NX." + ) + + members_scores_list = [ + str(item) for pair in members_scores.items() for item in pair[::-1] + ] + args += members_scores_list + + return cast( + int, + await self._execute_command(RequestType.Zadd, args), + ) + + async def zadd_incr( + self, + key: str, + member: str, + increment: float, + existing_options: Optional[ConditionalChange] = None, + update_condition: Optional[UpdateOptions] = None, + ) -> Optional[float]: + """ + Increments the score of member in the sorted set stored at `key` by `increment`. + If `member` does not exist in the sorted set, it is added with `increment` as its score (as if its previous score was 0.0). + If `key` does not exist, a new sorted set with the specified member as its sole member is created. + + See https://redis.io/commands/zadd/ for more details. + + Args: + key (str): The key of the sorted set. + member (str): A member in the sorted set to increment. + increment (float): The score to increment the member. + existing_options (Optional[ConditionalChange]): Options for handling the member's existence. + - NX: Only increment a member that doesn't exist. + - XX: Only increment an existing member. + update_condition (Optional[UpdateOptions]): Options for updating the score. + - GT: Only increment the score of the member if the new score will be greater than the current score. + - LT: Only increment (decrement) the score of the member if the new score will be less than the current score. + + Returns: + Optional[float]: The score of the member. + If there was a conflict with choosing the XX/NX/LT/GT options, the operation aborts and null is returned. + Examples: + >>> await zaddIncr("my_sorted_set", member , 5.0) + 5.0 + >>> await zaddIncr("existing_sorted_set", member , "3.0" , UpdateOptions.LESS_THAN) + None + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if update_condition: + args.append(update_condition.value) + + args.append("INCR") + + if existing_options and update_condition: + if existing_options == ConditionalChange.ONLY_IF_DOES_NOT_EXIST: + raise ValueError( + "The GT, LT and NX options are mutually exclusive. " + f"Cannot choose both {update_condition.value} and NX." + ) + + args += [str(increment), member] + return cast( + Optional[float], + await self._execute_command(RequestType.Zadd, args), + ) diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index b4a30723e6..67b31c026c 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -2,7 +2,7 @@ from typing import List, Mapping, Optional, Tuple, Union from glide.async_commands.core import ( - ConditionalSet, + ConditionalChange, ExpireOptions, ExpirySet, InfoSection, @@ -48,15 +48,15 @@ def set( self, key: str, value: str, - conditional_set: Union[ConditionalSet, None] = None, + conditional_set: Union[ConditionalChange, None] = None, expiry: Union[ExpirySet, None] = None, return_old_value: bool = False, ): args = [key, value] if conditional_set: - if conditional_set == ConditionalSet.ONLY_IF_EXISTS: + if conditional_set == ConditionalChange.ONLY_IF_EXISTS: args.append("XX") - if conditional_set == ConditionalSet.ONLY_IF_DOES_NOT_EXIST: + if conditional_set == ConditionalChange.ONLY_IF_DOES_NOT_EXIST: args.append("NX") if return_old_value: args.append("GET") diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index a6482c65f6..ca57978f1e 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -10,11 +10,12 @@ import pytest from glide import ClosingError, RequestError, TimeoutError from glide.async_commands.core import ( - ConditionalSet, + ConditionalChange, ExpireOptions, ExpirySet, ExpiryType, InfoSection, + UpdateOptions, ) from glide.config import ProtocolVersion, RedisCredentials from glide.constants import OK @@ -288,16 +289,16 @@ async def test_conditional_set(self, redis_client: TRedisClient): key = get_random_string(10) value = get_random_string(10) res = await redis_client.set( - key, value, conditional_set=ConditionalSet.ONLY_IF_EXISTS + key, value, conditional_set=ConditionalChange.ONLY_IF_EXISTS ) assert res is None res = await redis_client.set( - key, value, conditional_set=ConditionalSet.ONLY_IF_DOES_NOT_EXIST + key, value, conditional_set=ConditionalChange.ONLY_IF_DOES_NOT_EXIST ) assert res == OK assert await redis_client.get(key) == value res = await redis_client.set( - key, "foobar", conditional_set=ConditionalSet.ONLY_IF_DOES_NOT_EXIST + key, "foobar", conditional_set=ConditionalChange.ONLY_IF_DOES_NOT_EXIST ) assert res is None assert await redis_client.get(key) == value @@ -949,6 +950,100 @@ async def test_expire_pexpire_expireAt_pexpireAt_ttl_non_existing_key( ) assert await redis_client.ttl(key) == -2 + @pytest.mark.parametrize("cluster_mode", [True, False]) + async def test_zadd_zaddincr(self, redis_client: TRedisClient): + key = get_random_string(10) + members_scores = {"one": 1, "two": 2, "three": 3} + assert await redis_client.zadd(key, members_scores=members_scores) == 3 + assert await redis_client.zadd_incr(key, member="one", increment=2) == 3.0 + + @pytest.mark.parametrize("cluster_mode", [True, False]) + async def test_zadd_nx_xx(self, redis_client: TRedisClient): + key = get_random_string(10) + members_scores = {"one": 1, "two": 2, "three": 3} + assert ( + await redis_client.zadd( + key, + members_scores=members_scores, + existing_options=ConditionalChange.ONLY_IF_EXISTS, + ) + == 0 + ) + assert ( + await redis_client.zadd( + key, + members_scores=members_scores, + existing_options=ConditionalChange.ONLY_IF_DOES_NOT_EXIST, + ) + == 3 + ) + + assert ( + await redis_client.zadd_incr( + key, + member="one", + increment=5.0, + existing_options=ConditionalChange.ONLY_IF_DOES_NOT_EXIST, + ) + == None + ) + + assert ( + await redis_client.zadd_incr( + key, + member="one", + increment=5.0, + existing_options=ConditionalChange.ONLY_IF_EXISTS, + ) + == 6.0 + ) + + @pytest.mark.parametrize("cluster_mode", [True, False]) + async def test_zadd_gt_lt(self, redis_client: TRedisClient): + key = get_random_string(10) + members_scores = {"one": -3, "two": 2, "three": 3} + assert await redis_client.zadd(key, members_scores=members_scores) == 3 + members_scores["one"] = 10 + assert ( + await redis_client.zadd( + key, + members_scores=members_scores, + update_condition=UpdateOptions.GREATER_THAN, + changed=True, + ) + == 1 + ) + + assert ( + await redis_client.zadd( + key, + members_scores=members_scores, + update_condition=UpdateOptions.LESS_THAN, + changed=True, + ) + == 0 + ) + + assert ( + await redis_client.zadd_incr( + key, + member="one", + increment=-3.0, + update_condition=UpdateOptions.LESS_THAN, + ) + == 7.0 + ) + + assert ( + await redis_client.zadd_incr( + key, + member="one", + increment=-3.0, + update_condition=UpdateOptions.GREATER_THAN, + ) + == None + ) + class TestCommandsUnitTests: def test_expiry_cmd_args(self):