From 4f1d102327609fbd4c4c671e99318fa39c756a40 Mon Sep 17 00:00:00 2001 From: Shoham Elias <116083498+shohamazon@users.noreply.github.com> Date: Tue, 18 Jun 2024 11:59:33 +0300 Subject: [PATCH] Python: adds SUNION command (#1583) --- CHANGELOG.md | 1 + glide-core/src/client/value_conversion.rs | 2 +- glide-core/src/protobuf/redis_request.proto | 1 + glide-core/src/request_type.rs | 3 +++ python/python/glide/async_commands/core.py | 26 +++++++++++++++++++ .../glide/async_commands/transaction.py | 15 +++++++++++ python/python/tests/test_async_client.py | 26 +++++++++++++++++++ python/python/tests/test_transaction.py | 2 ++ 8 files changed, 75 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ab4b77f2af..284c13b628 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ * Node: Added OBJECT ENCODING command ([#1518](https://github.com/aws/glide-for-redis/pull/1518), [#1559](https://github.com/aws/glide-for-redis/pull/1559)) * Python: Added LMOVE and BLMOVE commands ([#1536](https://github.com/aws/glide-for-redis/pull/1536)) * Node: Added SUNIONSTORE command ([#1549](https://github.com/aws/glide-for-redis/pull/1549)) +* Python: Added SUNION command ([#1583](https://github.com/aws/glide-for-redis/pull/1583)) * Node: Added PFCOUNT command ([#1545](https://github.com/aws/glide-for-redis/pull/1545)) * Node: Added OBJECT FREQ command ([#1542](https://github.com/aws/glide-for-redis/pull/1542), [#1559](https://github.com/aws/glide-for-redis/pull/1559)) * Node: Added LINSERT command ([#1544](https://github.com/aws/glide-for-redis/pull/1544)) diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs index c44604ce34..b45b9637fc 100644 --- a/glide-core/src/client/value_conversion.rs +++ b/glide-core/src/client/value_conversion.rs @@ -864,7 +864,7 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option { | b"SISMEMBER" | b"PERSIST" | b"SMOVE" | b"RENAMENX" | b"MOVE" | b"COPY" | b"XGROUP DESTROY" | b"MSETNX" => Some(ExpectedReturnType::Boolean), b"SMISMEMBER" => Some(ExpectedReturnType::ArrayOfBools), - b"SMEMBERS" | b"SINTER" | b"SDIFF" => Some(ExpectedReturnType::Set), + b"SMEMBERS" | b"SINTER" | b"SDIFF" | b"SUNION" => Some(ExpectedReturnType::Set), b"ZSCORE" | b"GEODIST" => Some(ExpectedReturnType::DoubleOrNull), b"ZMSCORE" => Some(ExpectedReturnType::ArrayOfDoubleOrNull), b"ZPOPMIN" | b"ZPOPMAX" => Some(ExpectedReturnType::MapOfStringToDouble), diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto index 2d68131274..984484d210 100644 --- a/glide-core/src/protobuf/redis_request.proto +++ b/glide-core/src/protobuf/redis_request.proto @@ -223,6 +223,7 @@ enum RequestType { Watch = 183; UnWatch = 184; GeoSearchStore = 185; + SUnion = 186; } message Command { diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs index 947236672e..19487468b0 100644 --- a/glide-core/src/request_type.rs +++ b/glide-core/src/request_type.rs @@ -193,6 +193,7 @@ pub enum RequestType { Watch = 183, UnWatch = 184, GeoSearchStore = 185, + SUnion = 186, } fn get_two_word_command(first: &str, second: &str) -> Cmd { @@ -386,6 +387,7 @@ impl From<::protobuf::EnumOrUnknown> for RequestType { ProtobufRequestType::LPos => RequestType::LPos, ProtobufRequestType::LCS => RequestType::LCS, ProtobufRequestType::GeoSearch => RequestType::GeoSearch, + ProtobufRequestType::SUnion => RequestType::SUnion, ProtobufRequestType::Watch => RequestType::Watch, ProtobufRequestType::UnWatch => RequestType::UnWatch, ProtobufRequestType::GeoSearchStore => RequestType::GeoSearchStore, @@ -578,6 +580,7 @@ impl RequestType { RequestType::LPos => Some(cmd("LPOS")), RequestType::LCS => Some(cmd("LCS")), RequestType::GeoSearch => Some(cmd("GEOSEARCH")), + RequestType::SUnion => Some(cmd("SUNION")), RequestType::Watch => Some(cmd("WATCH")), RequestType::UnWatch => Some(cmd("UNWATCH")), RequestType::GeoSearchStore => Some(cmd("GEOSEARCHSTORE")), diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 39e06735dd..52b4e447e0 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -1896,6 +1896,32 @@ async def smove( ), ) + async def sunion(self, keys: List[str]) -> Set[str]: + """ + Gets the union of all the given sets. + + See https://valkey.io/commands/sunion for more details. + + Note: + When in cluster mode, all `keys` must map to the same hash slot. + + Args: + keys (List[str]): The keys of the sets. + + Returns: + Set[str]: A set of members which are present in at least one of the given sets. + If none of the sets exist, an empty set will be returned. + + Examples: + >>> await client.sadd("my_set1", ["member1", "member2"]) + >>> await client.sadd("my_set2", ["member2", "member3"]) + >>> await client.sunion(["my_set1", "my_set2"]) + {"member1", "member2", "member3"} # sets "my_set1" and "my_set2" have three unique members + >>> await client.sunion(["my_set1", "non_existing_set"]) + {"member1", "member2"} + """ + return cast(Set[str], await self._execute_command(RequestType.SUnion, keys)) + async def sunionstore( self, destination: str, diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index fd9462bf70..2dfc3c662f 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1255,6 +1255,21 @@ def smove( """ return self.append_command(RequestType.SMove, [source, destination, member]) + def sunion(self: TTransaction, keys: List[str]) -> TTransaction: + """ + Gets the union of all the given sets. + + See https://valkey.io/commands/sunion for more details. + + Args: + keys (List[str]): The keys of the sets. + + Commands response: + Set[str]: A set of members which are present in at least one of the given sets. + If none of the sets exist, an empty set will be returned. + """ + return self.append_command(RequestType.SUnion, keys) + def sunionstore( self: TTransaction, destination: str, diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 8b25b4573c..78b43f239a 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -1502,6 +1502,31 @@ async def test_smove(self, redis_client: TRedisClient): with pytest.raises(RequestError): await redis_client.smove(string_key, key1, "_") + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_sunion(self, redis_client: TRedisClient): + key1 = f"{{testKey}}:{get_random_string(10)}" + key2 = f"{{testKey}}:{get_random_string(10)}" + non_existing_key = f"{{testKey}}:non_existing_key" + member1_list = ["a", "b", "c"] + member2_list = ["b", "c", "d", "e"] + + assert await redis_client.sadd(key1, member1_list) == 3 + assert await redis_client.sadd(key2, member2_list) == 4 + assert await redis_client.sunion([key1, key2]) == {"a", "b", "c", "d", "e"} + + # invalid argument - key list must not be empty + with pytest.raises(RequestError): + await redis_client.sunion([]) + + # non-existing key returns the set of existing keys + assert await redis_client.sunion([key1, non_existing_key]) == set(member1_list) + + # non-set key + assert await redis_client.set(key2, "value") == OK + with pytest.raises(RequestError) as e: + await redis_client.sunion([key2]) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_sunionstore(self, redis_client: TRedisClient): @@ -4910,6 +4935,7 @@ async def test_multi_key_command_returns_cross_slot_error( "abc", "zxy", ListDirection.LEFT, ListDirection.LEFT, 1 ), redis_client.msetnx({"abc": "abc", "zxy": "zyx"}), + redis_client.sunion(["def", "ghi"]), ] if not await check_if_server_version_lt(redis_client, "6.2.0"): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index bdd30ac449..1d6827cafd 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -247,6 +247,8 @@ async def transaction_test( args.append(2) transaction.sinter([key7, key7]) args.append({"foo", "bar"}) + transaction.sunion([key7, key7]) + args.append({"foo", "bar"}) transaction.sinterstore(key7, [key7, key7]) args.append(2) if not await check_if_server_version_lt(redis_client, "7.0.0"):