From f222008f372a5d1529f6be6f2a75fe00bdea31b5 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 29 Apr 2024 17:54:10 -0700 Subject: [PATCH 1/5] Python: add ZDIFFSTORE command --- CHANGELOG.md | 1 + python/python/glide/async_commands/core.py | 27 +++++++++++ .../glide/async_commands/transaction.py | 19 ++++++++ python/python/tests/test_async_client.py | 46 +++++++++++++++++++ python/python/tests/test_transaction.py | 2 + 5 files changed, 95 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5239d482d7..4b156d9f45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ * Python: Added GEOPOS command ([#1301](https://github.com/aws/glide-for-redis/pull/1301)) * Node: Added PFADD command ([#1317](https://github.com/aws/glide-for-redis/pull/1317)) * Python: Added PFADD command ([#1315](https://github.com/aws/glide-for-redis/pull/1315)) +* Python: Added ZDIFFSTORE command (TODO: add PR link) #### 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 a2818b7174..3251ecaa99 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2326,6 +2326,33 @@ async def zscore(self, key: str, member: str) -> Optional[float]: await self._execute_command(RequestType.ZScore, [key, member]), ) + async def zdiffstore(self, destination: str, keys: List[str]) -> int: + """ + Calculates the difference between the first sorted set and all the successive sorted sets at keys and stores the + difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are treated + as empty sets. + When in Cluster mode, all keys must map to the same hash slot. + + See https://valkey.io/commands/zdiffstore for more details. + + Args: + destination (str): The key for the resulting sorted set. + keys (List[str]): The keys of the sorted sets to compare. + + Returns: + The number of members in the resulting sorted set stored at `destination`. + + Examples: + >>> await client.zdiffstore("my_sorted_set", ["key1", "key2"]) + 1 # One member exists in "key1" but not "key2", and this member was stored in my_sorted_set. + """ + return cast( + int, + await self._execute_command( + RequestType.ZDiffStore, [destination, str(len(keys))] + keys + ), + ) + async def invoke_script( self, script: Script, diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index e90b339086..7389b0109c 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1751,6 +1751,25 @@ def zscore(self: TTransaction, key: str, member: str) -> TTransaction: """ return self.append_command(RequestType.ZScore, [key, member]) + def zdiffstore(self: TTransaction, destination: str, keys: List[str]) -> TTransaction: + """ + Calculates the difference between the first sorted set and all the successive sorted sets at keys and stores the + difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are treated + as empty sets. + + See https://valkey.io/commands/zdiffstore for more details. + + Args: + destination (str): The key for the resulting sorted set. + keys (List[str]): The keys of the sorted sets to compare. + + Command response: + The number of members in the resulting sorted set stored at `destination`. + """ + return self.append_command( + RequestType.ZDiffStore, [destination, str(len(keys))] + keys + ) + def dbsize(self: TTransaction) -> TTransaction: """ Returns the number of keys in the currently selected database. diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 849425ab8a..df56db5f8f 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -1944,6 +1944,52 @@ async def test_zrank(self, redis_client: TRedisClient): with pytest.raises(RequestError): await redis_client.zrank(key, "one") + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zdiffstore(self, redis_client: TRedisClient): + key1 = f"{{testKey}}:1-{get_random_string(10)}" + key2 = f"{{testKey}}:2-{get_random_string(10)}" + key3 = f"{{testKey}}:3-{get_random_string(10)}" + key4 = f"{{testKey}}:4-{get_random_string(10)}" + string_key = f"{{testKey}}:4-{get_random_string(10)}" + non_existing_key = f"{{testKey}}:5-{get_random_string(10)}" + + member_scores1 = {"one": 1.0, "two": 2.0, "three": 3.0} + member_scores2 = {"two": 2.0} + member_scores3 = {"one": 1.0, "two": 2.0, "three": 3.0, "four": 4.0} + + assert await redis_client.zadd(key1, member_scores1) == 3 + assert await redis_client.zadd(key2, member_scores2) == 1 + assert await redis_client.zadd(key3, member_scores3) == 4 + + assert await redis_client.zdiffstore(key4, [key1, key2]) == 2 + assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == { + "one": 1.0, + "three": 3.0, + } + + assert await redis_client.zdiffstore(key4, [key3, key2, key1]) == 1 + assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == { + "four": 4.0 + } + + assert await redis_client.zdiffstore(key4, [key1, key3]) == 0 + assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {} + + assert await redis_client.zdiffstore(key4, [non_existing_key, key1]) == 0 + assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {} + + # key exists, but it is not a sorted set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zdiffstore(key4, [string_key, key1]) + + # same-slot requirement + if isinstance(redis_client, RedisClusterClient): + with pytest.raises(RequestError) as e: + await redis_client.zdiffstore("abc", ["zxy", "lkn"]) + assert "CrossSlot" in str(e) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_type(self, redis_client: TRedisClient): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 983d10bb12..1a1535ce19 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -211,6 +211,8 @@ async def transaction_test( args.append(1) transaction.zremrangebylex(key8, InfBound.NEG_INF, InfBound.POS_INF) args.append(0) + transaction.zdiffstore(key8, [key8, key8]) + args.append(0) transaction.pfadd(key10, ["a", "b", "c"]) args.append(1) From e3d19f3789b9896b32e8099b62f7ee1b6a6b4c5f Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 29 Apr 2024 17:56:31 -0700 Subject: [PATCH 2/5] Minor formatting improvement --- python/python/glide/async_commands/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 3251ecaa99..ced7ed0931 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2344,7 +2344,7 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int: Examples: >>> await client.zdiffstore("my_sorted_set", ["key1", "key2"]) - 1 # One member exists in "key1" but not "key2", and this member was stored in my_sorted_set. + 1 # One member exists in "key1" but not "key2", and this member was stored in "my_sorted_set". """ return cast( int, From c4f7348764693d4cc332e30af9cd4acb89fe820b Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 29 Apr 2024 17:58:26 -0700 Subject: [PATCH 3/5] Fix formatting --- python/python/glide/async_commands/transaction.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 7389b0109c..7202c613e8 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1751,7 +1751,9 @@ def zscore(self: TTransaction, key: str, member: str) -> TTransaction: """ return self.append_command(RequestType.ZScore, [key, member]) - def zdiffstore(self: TTransaction, destination: str, keys: List[str]) -> TTransaction: + def zdiffstore( + self: TTransaction, destination: str, keys: List[str] + ) -> TTransaction: """ Calculates the difference between the first sorted set and all the successive sorted sets at keys and stores the difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are treated From 4f588bd10c800938a0cde3cf106b02f75837882a Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 30 Apr 2024 09:43:05 -0700 Subject: [PATCH 4/5] PR suggestions --- python/python/glide/async_commands/core.py | 2 +- python/python/glide/async_commands/transaction.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index ced7ed0931..1092945ee9 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2340,7 +2340,7 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int: keys (List[str]): The keys of the sorted sets to compare. Returns: - The number of members in the resulting sorted set stored at `destination`. + int: The number of members in the resulting sorted set stored at `destination`. Examples: >>> await client.zdiffstore("my_sorted_set", ["key1", "key2"]) diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 7202c613e8..61379682bb 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1766,7 +1766,7 @@ def zdiffstore( keys (List[str]): The keys of the sorted sets to compare. Command response: - The number of members in the resulting sorted set stored at `destination`. + int: The number of members in the resulting sorted set stored at `destination`. """ return self.append_command( RequestType.ZDiffStore, [destination, str(len(keys))] + keys From 4eb6d36fb31206601a52a963775d2d36c69871bd Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 2 May 2024 12:33:55 -0700 Subject: [PATCH 5/5] Adjust docs to align with existing doc patterns --- python/python/glide/async_commands/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 1092945ee9..c9eef3e291 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2331,6 +2331,7 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int: Calculates the difference between the first sorted set and all the successive sorted sets at keys and stores the difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are treated as empty sets. + When in Cluster mode, all keys must map to the same hash slot. See https://valkey.io/commands/zdiffstore for more details.