Skip to content

Commit

Permalink
Python: adds SUNION command
Browse files Browse the repository at this point in the history
  • Loading branch information
shohamazon committed Jun 16, 2024
1 parent 4ef40e1 commit dab998e
Show file tree
Hide file tree
Showing 8 changed files with 75 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,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))
Expand Down
2 changes: 1 addition & 1 deletion glide-core/src/client/value_conversion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -865,7 +865,7 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
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),
Expand Down
1 change: 1 addition & 0 deletions glide-core/src/protobuf/redis_request.proto
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ enum RequestType {
LPos = 180;
LCS = 181;
GeoSearch = 182;
SUnion = 183;
}

message Command {
Expand Down
3 changes: 3 additions & 0 deletions glide-core/src/request_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,7 @@ pub enum RequestType {
LPos = 180,
LCS = 181,
GeoSearch = 182,
SUnion = 183,
}

fn get_two_word_command(first: &str, second: &str) -> Cmd {
Expand Down Expand Up @@ -381,6 +382,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
ProtobufRequestType::LPos => RequestType::LPos,
ProtobufRequestType::LCS => RequestType::LCS,
ProtobufRequestType::GeoSearch => RequestType::GeoSearch,
ProtobufRequestType::SUnion => RequestType::SUnion,
}
}
}
Expand Down Expand Up @@ -569,6 +571,7 @@ impl RequestType {
RequestType::LPos => Some(cmd("LPOS")),
RequestType::LCS => Some(cmd("LCS")),
RequestType::GeoSearch => Some(cmd("GEOSEARCH")),
RequestType::SUnion => Some(cmd("SUNION")),
}
}
}
26 changes: 26 additions & 0 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1849,6 +1849,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,
Expand Down
15 changes: 15 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -1225,6 +1225,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,
Expand Down
26 changes: 26 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1474,6 +1474,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):
Expand Down Expand Up @@ -4491,6 +4516,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.suion(["def", "ghi"]),
]

if not await check_if_server_version_lt(redis_client, "7.0.0"):
Expand Down
2 changes: 2 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,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"):
Expand Down

0 comments on commit dab998e

Please sign in to comment.