remove COMMAND GETKEYS call from ZUNION and ZUNIONSTORE (#4)

* make EVALSHA and EVALSHA_RO to avoid command_keys as well

* make ZUNION, ZUNIONSTORE to avoid command_keys as well

Author: zach-iee

redis/cluster.py

@@ -957,7 +957,7 @@ def determine_slot(self, *args):
         # redis server to parse the keys. Besides, there is a bug in redis<7.0
         # where `self._get_command_keys()` fails anyway. So, we special case
         # EVAL/EVALSHA.
-        if command in ("EVAL", "EVALSHA"):
+        if command in ("EVAL", "EVALSHA", "EVAL_RO", "EVALSHA_RO"):
             # command syntax: EVAL "script body" num_keys ...
             if len(args) <= 2:
                 raise RedisClusterException(f"Invalid args in command: {args}")
@@ -968,6 +968,14 @@ 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://github.com/redis/redis/pull/12380
+            # 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]
         else:
             keys = self._get_command_keys(*args)
             if keys is None or len(keys) == 0:

tests/test_cluster.py

@@ -2006,10 +2006,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),