make ZUNION, ZUNIONSTORE to avoid command_keys as well

Author: zach-iee

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 <SUM | MIN | MAX>]
+            if len(args) <= 3:
+                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 <SUM | MIN | MAX>] [WITHSCORES]
+            if len(args) <= 2:
+                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:

tests/test_cluster.py

@@ -1978,10 +1978,21 @@ def test_cluster_zrangestore(self, r):
 
     @skip_if_server_version_lt("6.2.0")
     def test_cluster_zunion(self, r):
+        assert r.zunion(["{foo}" + str(i) for i in range(0, 256)]) == []
+
         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"]) == [b"a1", b"a2", b"a3"]
         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) == [
             (b"a2", 3),
@@ -2006,10 +2017,30 @@ def test_cluster_zunion(self, r):
         ]
 
     def test_cluster_zunionstore_sum(self, r):
+        assert r.zunionstore("{foo}d", ["{foo}" + str(i) for i in range(0, 256)]) == 0
+
         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",
+        ):
+            r.zunionstore(result_key, failed_keys)
+
         assert r.zunionstore("{foo}d", ["{foo}a", "{foo}b", "{foo}c"]) == 4
+        assert r.zunionstore("{foo}e", ["{foo}a"]) == 3
         assert r.zrange("{foo}d", 0, -1, withscores=True) == [
             (b"a2", 3),
             (b"a4", 4),