From 3326b4fc1b402e9fed9a828c87a72fce0aedf849 Mon Sep 17 00:00:00 2001 From: Aaron <69273634+aaron-congo@users.noreply.github.com> Date: Mon, 6 May 2024 15:48:24 -0700 Subject: [PATCH] Python: add ZRANGESTORE command (#1377) * Python: add ZRANGESTORE command (#258) * Update with PR link --- CHANGELOG.md | 1 + python/python/glide/async_commands/core.py | 41 +++ .../python/glide/async_commands/sorted_set.py | 26 ++ .../glide/async_commands/transaction.py | 33 +++ python/python/tests/test_async_client.py | 269 ++++++++++++++++++ python/python/tests/test_transaction.py | 2 + 6 files changed, 372 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e0ef8ea07..92ee6b4da9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ * Python: Added ZMSCORE command ([#1357](https://github.com/aws/glide-for-redis/pull/1357)) * Python: Added HRANDFIELD command ([#1334](https://github.com/aws/glide-for-redis/pull/1334)) * Python: Added XADD, XTRIM commands ([#1320](https://github.com/aws/glide-for-redis/pull/1320)) +* Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377)) #### Fixes * Python: Fix typing error "‘type’ object is not subscriptable" ([#1203](https://github.com/aws/glide-for-redis/pull/1203)) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index e8667c0696..0bfe51680f 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -24,6 +24,7 @@ RangeByScore, ScoreBoundary, _create_zrange_args, + _create_zrangestore_args, ) from glide.constants import TOK, TResult from glide.protobuf.redis_request_pb2 import RequestType @@ -2338,6 +2339,46 @@ async def zrange_withscores( Mapping[str, float], await self._execute_command(RequestType.Zrange, args) ) + async def zrangestore( + self, + destination: str, + source: str, + range_query: Union[RangeByIndex, RangeByLex, RangeByScore], + reverse: bool = False, + ) -> int: + """ + Stores a specified range of elements from the sorted set at `source`, into a new sorted set at `destination`. If + `destination` doesn't exist, a new sorted set is created; if it exists, it's overwritten. + + When in Cluster mode, all keys must map to the same hash slot. + + ZRANGESTORE can perform different types of range queries: by index (rank), by the score, or by lexicographical + order. + + See https://valkey.io/commands/zrangestore for more details. + + Args: + destination (str): The key for the destination sorted set. + source (str): The key of the source sorted set. + range_query (Union[RangeByIndex, RangeByLex, RangeByScore]): The range query object representing the type of range query to perform. + - For range queries by index (rank), use RangeByIndex. + - For range queries by lexicographical order, use RangeByLex. + - For range queries by score, use RangeByScore. + reverse (bool): If True, reverses the sorted set, with index 0 as the element with the highest score. + + Returns: + int: The number of elements in the resulting sorted set. + + Examples: + >>> await client.zrangestore("destination_key", "my_sorted_set", RangeByIndex(0, 2), True) + 3 # The 3 members with the highest scores from "my_sorted_set" were stored in the sorted set at "destination_key". + >>> await client.zrangestore("destination_key", "my_sorted_set", RangeByScore(InfBound.NEG_INF, ScoreBoundary(3))) + 2 # The 2 members with scores between negative infinity and 3 (inclusive) from "my_sorted_set" were stored in the sorted set at "destination_key". + """ + args = _create_zrangestore_args(destination, source, range_query, reverse) + + return cast(int, await self._execute_command(RequestType.ZRangeStore, args)) + async def zrank( self, key: str, diff --git a/python/python/glide/async_commands/sorted_set.py b/python/python/glide/async_commands/sorted_set.py index 83c6037341..f5198741f2 100644 --- a/python/python/glide/async_commands/sorted_set.py +++ b/python/python/glide/async_commands/sorted_set.py @@ -158,3 +158,29 @@ def _create_zrange_args( args.append("WITHSCORES") return args + + +def _create_zrangestore_args( + destination: str, + source: str, + range_query: Union[RangeByLex, RangeByScore, RangeByIndex], + reverse: bool, +) -> List[str]: + args = [destination, source, str(range_query.start), str(range_query.stop)] + + if isinstance(range_query, RangeByScore): + args.append("BYSCORE") + elif isinstance(range_query, RangeByLex): + args.append("BYLEX") + if reverse: + args.append("REV") + if hasattr(range_query, "limit") and range_query.limit is not None: + args.extend( + [ + "LIMIT", + str(range_query.limit.offset), + str(range_query.limit.count), + ] + ) + + return args diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index e91a52cce4..2867661a49 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -23,6 +23,7 @@ RangeByScore, ScoreBoundary, _create_zrange_args, + _create_zrangestore_args, ) from glide.protobuf.redis_request_pb2 import RequestType @@ -1664,6 +1665,38 @@ def zrange_withscores( return self.append_command(RequestType.Zrange, args) + def zrangestore( + self: TTransaction, + destination: str, + source: str, + range_query: Union[RangeByIndex, RangeByLex, RangeByScore], + reverse: bool = False, + ) -> TTransaction: + """ + Stores a specified range of elements from the sorted set at `source`, into a new sorted set at `destination`. If + `destination` doesn't exist, a new sorted set is created; if it exists, it's overwritten. + + ZRANGESTORE can perform different types of range queries: by index (rank), by the score, or by lexicographical + order. + + See https://valkey.io/commands/zrangestore for more details. + + Args: + destination (str): The key for the destination sorted set. + source (str): The key of the source sorted set. + range_query (Union[RangeByIndex, RangeByLex, RangeByScore]): The range query object representing the type of range query to perform. + - For range queries by index (rank), use RangeByIndex. + - For range queries by lexicographical order, use RangeByLex. + - For range queries by score, use RangeByScore. + reverse (bool): If True, reverses the sorted set, with index 0 as the element with the highest score. + + Command response: + int: The number of elements in the resulting sorted set. + """ + args = _create_zrangestore_args(destination, source, range_query, reverse) + + return self.append_command(RequestType.ZRangeStore, args) + def zrank( self: TTransaction, key: str, diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 9d3efdf43c..95285233d4 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -2022,6 +2022,275 @@ async def test_zrange_different_types_of_keys(self, redis_client: TRedisClient): with pytest.raises(RequestError): await redis_client.zrange_withscores(key, RangeByIndex(start=0, stop=1)) + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrangestore_by_index(self, redis_client: TRedisClient): + destination = f"{{testKey}}:{get_random_string(10)}" + source = f"{{testKey}}:{get_random_string(10)}" + string_key = f"{{testKey}}:{get_random_string(10)}" + non_existing_key = f"{{testKey}}:{get_random_string(10)}" + + member_scores = {"one": 1.0, "two": 2.0, "three": 3.0} + assert await redis_client.zadd(source, member_scores) == 3 + + # full range + assert ( + await redis_client.zrangestore(destination, source, RangeByIndex(0, -1)) + == 3 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"one": 1.0, "two": 2.0, "three": 3.0} + + # range from rank 0 to 1, from highest to lowest score + assert ( + await redis_client.zrangestore( + destination, source, RangeByIndex(0, 1), True + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"three": 3.0, "two": 2.0} + + # incorrect range, as start > stop + assert ( + await redis_client.zrangestore(destination, source, RangeByIndex(3, 1)) == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # non-existing source + assert ( + await redis_client.zrangestore( + destination, non_existing_key, RangeByIndex(0, -1) + ) + == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrangestore(destination, string_key, RangeByIndex(0, -1)) + + # same-slot requirement + if isinstance(redis_client, RedisClusterClient): + with pytest.raises(RequestError) as e: + await redis_client.zrangestore("abc", "def", RangeByIndex(0, -1)) + assert "CrossSlot" in str(e) + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrangestore_by_score(self, redis_client: TRedisClient): + destination = f"{{testKey}}:{get_random_string(10)}" + source = f"{{testKey}}:{get_random_string(10)}" + string_key = f"{{testKey}}:{get_random_string(10)}" + non_existing_key = f"{{testKey}}:{get_random_string(10)}" + + member_scores = {"one": 1.0, "two": 2.0, "three": 3.0} + assert await redis_client.zadd(source, member_scores) == 3 + + # range from negative infinity to 3 (exclusive) + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByScore(InfBound.NEG_INF, ScoreBoundary(3, False)), + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"one": 1.0, "two": 2.0} + + # range from 1 (inclusive) to positive infinity + assert ( + await redis_client.zrangestore( + destination, source, RangeByScore(ScoreBoundary(1), InfBound.POS_INF) + ) + == 3 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"one": 1.0, "two": 2.0, "three": 3.0} + + # range from negative to positive infinity, limited to ranks 1 to 2 + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByScore(InfBound.NEG_INF, InfBound.POS_INF, Limit(1, 2)), + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"two": 2.0, "three": 3.0} + + # range from positive to negative infinity reversed, limited to ranks 1 to 2 + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByScore(InfBound.POS_INF, InfBound.NEG_INF, Limit(1, 2)), + True, + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"two": 2.0, "one": 1.0} + + # incorrect range as start > stop + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByScore(ScoreBoundary(3, False), InfBound.NEG_INF), + ) + == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # non-existing source + assert ( + await redis_client.zrangestore( + destination, + non_existing_key, + RangeByScore(InfBound.NEG_INF, ScoreBoundary(3, False)), + ) + == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrangestore( + destination, + string_key, + RangeByScore(ScoreBoundary(0), ScoreBoundary(3)), + ) + + # same-slot requirement + if isinstance(redis_client, RedisClusterClient): + with pytest.raises(RequestError) as e: + await redis_client.zrangestore( + "abc", "def", RangeByScore(ScoreBoundary(0), ScoreBoundary(3)) + ) + assert "CrossSlot" in str(e) + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrangestore_by_lex(self, redis_client: TRedisClient): + destination = f"{{testKey}}:{get_random_string(10)}" + source = f"{{testKey}}:{get_random_string(10)}" + string_key = f"{{testKey}}:4-{get_random_string(10)}" + non_existing_key = f"{{testKey}}:5-{get_random_string(10)}" + + member_scores = {"a": 1.0, "b": 2.0, "c": 3.0} + assert await redis_client.zadd(source, member_scores) == 3 + + # range from negative infinity to "c" (exclusive) + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByLex(InfBound.NEG_INF, LexBoundary("c", False)), + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"a": 1.0, "b": 2.0} + + # range from "a" (inclusive) to positive infinity + assert ( + await redis_client.zrangestore( + destination, source, RangeByLex(LexBoundary("a"), InfBound.POS_INF) + ) + == 3 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"a": 1.0, "b": 2.0, "c": 3.0} + + # range from negative to positive infinity, limited to ranks 1 to 2 + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByLex(InfBound.NEG_INF, InfBound.POS_INF, Limit(1, 2)), + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"b": 2.0, "c": 3.0} + + # range from positive to negative infinity reversed, limited to ranks 1 to 2 + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByLex(InfBound.POS_INF, InfBound.NEG_INF, Limit(1, 2)), + True, + ) + == 2 + ) + assert await redis_client.zrange_withscores( + destination, RangeByIndex(0, -1) + ) == {"b": 2.0, "a": 1.0} + + # incorrect range as start > stop + assert ( + await redis_client.zrangestore( + destination, + source, + RangeByLex(LexBoundary("c", False), InfBound.NEG_INF), + ) + == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # non-existing source + assert ( + await redis_client.zrangestore( + destination, + non_existing_key, + RangeByLex(InfBound.NEG_INF, InfBound.POS_INF), + ) + == 0 + ) + assert ( + await redis_client.zrange_withscores(destination, RangeByIndex(0, -1)) == {} + ) + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrangestore( + destination, string_key, RangeByLex(InfBound.NEG_INF, InfBound.POS_INF) + ) + + # same-slot requirement + if isinstance(redis_client, RedisClusterClient): + with pytest.raises(RequestError) as e: + await redis_client.zrangestore( + "abc", "def", RangeByLex(InfBound.NEG_INF, InfBound.POS_INF) + ) + assert "CrossSlot" in str(e) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_zrank(self, redis_client: TRedisClient): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index bf6227c5bf..a462c4b2da 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -217,6 +217,8 @@ async def transaction_test( args.append({"two": 2, "three": 3, "four": 4}) transaction.zmscore(key8, ["two", "three"]) args.append([2.0, 3.0]) + transaction.zrangestore(key8, key8, RangeByIndex(0, -1)) + args.append(3) transaction.zpopmin(key8) args.append({"two": 2.0}) transaction.zpopmax(key8)