diff --git a/redis/cluster.py b/redis/cluster.py index f19408ad8d..d6a38f0d09 100644 --- a/redis/cluster.py +++ b/redis/cluster.py @@ -968,6 +968,22 @@ def determine_slot(self, *args): if len(eval_keys) == 0: return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS) keys = eval_keys + elif command == "ZUNIONSTORE": + # https://sendbird.slack.com/archives/GFLAQKVDZ/p1688047132982659?thread_ts=1688046729.152689&cid=GFLAQKVDZ + # command syntax: ZUNIONSTORE destination numkeys key [key ...] + # [WEIGHTS weight [weight ...]] [AGGREGATE ] + if len(args) <= 4: + raise RedisClusterException(f"Invalid args in ZUNIONSTORE: {args}") + num_actual_keys = args[2] + keys = (args[1],) + args[3 : 3 + num_actual_keys] + elif command == "ZUNION": + # https://sendbird.slack.com/archives/GFLAQKVDZ/p1688047132982659?thread_ts=1688046729.152689&cid=GFLAQKVDZ + # command syntax: ZUNION numkeys key [key ...] [WEIGHTS weight [weight ...]] + # [AGGREGATE ] [WITHSCORES] + if len(args) <= 3: + raise RedisClusterException(f"Invalid args in ZUNION: {args}") + num_actual_keys = args[1] + keys = args[2 : 2 + num_actual_keys] else: keys = self._get_command_keys(*args) if keys is None or len(keys) == 0: diff --git a/tests/test_cluster.py b/tests/test_cluster.py index e3b43fc4be..0610d7adbe 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -1981,6 +1981,14 @@ def test_cluster_zunion(self, r): r.zadd("{foo}a", {"a1": 1, "a2": 1, "a3": 1}) r.zadd("{foo}b", {"a1": 2, "a2": 2, "a3": 2}) r.zadd("{foo}c", {"a1": 6, "a3": 5, "a4": 4}) + + failed_keys = ["{foo1}a", "{foo}b", "{foo}c"] + with pytest.raises( + RedisClusterException, + match="ZUNION - all keys must map to the same key slot" + ): + r.zunion(failed_keys) + # sum assert r.zunion(["{foo}a", "{foo}b", "{foo}c"]) == [b"a2", b"a4", b"a3", b"a1"] assert r.zunion(["{foo}a", "{foo}b", "{foo}c"], withscores=True) == [ @@ -2009,6 +2017,23 @@ def test_cluster_zunionstore_sum(self, r): r.zadd("{foo}a", {"a1": 1, "a2": 1, "a3": 1}) r.zadd("{foo}b", {"a1": 2, "a2": 2, "a3": 2}) r.zadd("{foo}c", {"a1": 6, "a3": 5, "a4": 4}) + + result_key = "{foo}d" + failed_keys = ["{foo1}a", "{foo}b", "{foo}c"] + with pytest.raises( + RedisClusterException, + match="ZUNIONSTORE - all keys must map to the same key slot" + ): + r.zunionstore(result_key, failed_keys) + + result_key = "{foo1}d" + failed_keys = ["{foo}a", "{foo}b"] + with pytest.raises( + RedisClusterException, + match="ZUNIONSTORE - all keys must map to the same key slot" + ) as e: + r.zunionstore(result_key, failed_keys) + assert r.zunionstore("{foo}d", ["{foo}a", "{foo}b", "{foo}c"]) == 4 assert r.zrange("{foo}d", 0, -1, withscores=True) == [ (b"a2", 3),