Python: adds SUNION command (valkey-io#1583)

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))

glide-core/src/client/value_conversion.rs

@@ -864,7 +864,7 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
         | 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),

glide-core/src/protobuf/redis_request.proto

@@ -223,6 +223,7 @@ enum RequestType {
     Watch = 183;
     UnWatch = 184;
     GeoSearchStore = 185;
+    SUnion = 186;
 }
 
 message Command {

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<ProtobufRequestType>> 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")),

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,

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,

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"):

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"):