make ZUNION, ZUNIONSTORE to avoid command_keys as well

redis · Jun 30, 2023 · ca93c1c · ca93c1c
1 parent b011b6b
commit ca93c1c
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git 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 <SUM | MIN | MAX>]
+            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 <SUM | MIN | MAX>] [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
@@ -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),