Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Python: adds SUNION command #1583

Merged
merged 2 commits into from
Jun 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
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 @@ -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),
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 @@ -223,6 +223,7 @@ enum RequestType {
Watch = 183;
UnWatch = 184;
GeoSearchStore = 185;
SUnion = 186;
}

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 @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")),
Expand Down
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 @@ -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,
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 @@ -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,
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 @@ -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):
Expand Down Expand Up @@ -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"):
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 @@ -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"):
Expand Down
Loading