Add cluster support for scripting (#1937)

* Add cluster support for scripting

* Fall back to connection_pool.get_encoder if necessary

* Add documentation for cluster-based scripting

* Add test for flush response

Co-authored-by: dvora-h <dvora.heller@redis.com>

1 files changed, 69 insertions, 18 deletions

diff --git a/redis/cluster.py b/redis/cluster.py
index 7151cfe..b8d6b19 100644
--- a/redis/cluster.py
+++ b/redis/cluster.py
@@ -289,6 +289,9 @@ class RedisCluster(RedisClusterCommands):
             [
                 "FLUSHALL",
                 "FLUSHDB",
+                "SCRIPT EXISTS",
+                "SCRIPT FLUSH",
+                "SCRIPT LOAD",
             ],
             PRIMARIES,
         ),
@@ -379,6 +382,24 @@ class RedisCluster(RedisClusterCommands):
             ],
             parse_scan_result,
         ),
+        list_keys_to_dict(
+            [
+                "SCRIPT LOAD",
+            ],
+            lambda command, res: list(res.values()).pop(),
+        ),
+        list_keys_to_dict(
+            [
+                "SCRIPT EXISTS",
+            ],
+            lambda command, res: [all(k) for k in zip(*res.values())],
+        ),
+        list_keys_to_dict(
+            [
+                "SCRIPT FLUSH",
+            ],
+            lambda command, res: all(res.values()),
+        ),
     )
 
     ERRORS_ALLOW_RETRY = (
@@ -778,40 +799,70 @@ class RedisCluster(RedisClusterCommands):
         """
         Get the keys in the command. If the command has no keys in in, None is
         returned.
+
+        NOTE: Due to a bug in redis<7.0, this function does not work properly
+        for EVAL or EVALSHA when the `numkeys` arg is 0.
+         - issue: https://github.com/redis/redis/issues/9493
+         - fix: https://github.com/redis/redis/pull/9733
+
+        So, don't use this function with EVAL or EVALSHA.
         """
         redis_conn = self.get_default_node().redis_connection
         return self.commands_parser.get_keys(redis_conn, *args)
 
     def determine_slot(self, *args):
         """
-        Figure out what slot based on command and args
+        Figure out what slot to use based on args.
+
+        Raises a RedisClusterException if there's a missing key and we can't
+            determine what slots to map the command to; or, if the keys don't
+            all map to the same key slot.
         """
-        if self.command_flags.get(args[0]) == SLOT_ID:
+        command = args[0]
+        if self.command_flags.get(command) == SLOT_ID:
             # The command contains the slot ID
             return args[1]
 
         # Get the keys in the command
-        keys = self._get_command_keys(*args)
-        if keys is None or len(keys) == 0:
-            raise RedisClusterException(
-                "No way to dispatch this command to Redis Cluster. "
-                "Missing key.\nYou can execute the command by specifying "
-                f"target nodes.\nCommand: {args}"
-            )
 
-        if len(keys) > 1:
-            # multi-key command, we need to make sure all keys are mapped to
-            # the same slot
-            slots = {self.keyslot(key) for key in keys}
-            if len(slots) != 1:
+        # EVAL and EVALSHA are common enough that it's wasteful to go to the
+        # 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"):
+            # command syntax: EVAL "script body" num_keys ...
+            if len(args) <= 2:
+                raise RedisClusterException(f"Invalid args in command: {args}")
+            num_actual_keys = args[2]
+            eval_keys = args[3 : 3 + num_actual_keys]
+            # if there are 0 keys, that means the script can be run on any node
+            # so we can just return a random slot
+            if len(eval_keys) == 0:
+                return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
+            keys = eval_keys
+        else:
+            keys = self._get_command_keys(*args)
+            if keys is None or len(keys) == 0:
                 raise RedisClusterException(
-                    f"{args[0]} - all keys must map to the same key slot"
+                    "No way to dispatch this command to Redis Cluster. "
+                    "Missing key.\nYou can execute the command by specifying "
+                    f"target nodes.\nCommand: {args}"
                 )
-            return slots.pop()
-        else:
-            # single key command
+
+        # single key command
+        if len(keys) == 1:
             return self.keyslot(keys[0])
 
+        # multi-key command; we need to make sure all keys are mapped to
+        # the same slot
+        slots = {self.keyslot(key) for key in keys}
+        if len(slots) != 1:
+            raise RedisClusterException(
+                f"{command} - all keys must map to the same key slot"
+            )
+
+        return slots.pop()
+
     def reinitialize_caches(self):
         self.nodes_manager.initialize()