Skip COMMAND GETKEYS for movablekeys commands with complete key specs

redis/_parsers/commands.py

@@ -77,6 +77,63 @@ def _get_pubsub_keys(self, *args):
     def parse_subcommand(self, command, **options):
         return _parse_subcommand(command)
 
+    def _get_keys_from_specs(self, command, args):
+        """
+        Find the keys of a movablekeys command from its key specs, the way the
+        server does, which saves the COMMAND GETKEYS round trip. Returns None
+        when the specs are not enough and the server has to be asked.
+        """
+        specs = command.get("key_specifications")
+        if not specs:
+            return None
+        args = args[0].split() + list(args[1:])
+        argc = len(args)
+        keys = []
+        for spec in specs:
+            spec = _spec_dict(spec)
+            begin = _spec_dict(spec.get("begin_search", ()))
+            find = _spec_dict(spec.get("find_keys", ()))
+            if (
+                "incomplete" in {str_if_bytes(f) for f in spec.get("flags", ())}
+                or str_if_bytes(begin.get("type")) != "index"
+            ):
+                return None
+            first = _spec_dict(begin["spec"])["index"]
+            find_type = str_if_bytes(find.get("type"))
+            find_spec = _spec_dict(find.get("spec", ()))
+            if find_type == "range":
+                lastkey = find_spec["lastkey"]
+                limit = find_spec["limit"]
+                if lastkey >= 0:
+                    last = first + lastkey
+                elif not limit:
+                    last = argc + lastkey
+                else:
+                    last = first + (argc - first) // limit + lastkey
+            elif find_type == "keynum":
+                try:
+                    numkeys = int(args[first + find_spec["keynumidx"]])
+                except (IndexError, TypeError, ValueError):
+                    return None
+                if numkeys < 0:
+                    return None
+                first += find_spec["firstkey"]
+                last = first + (numkeys - 1) * find_spec["keystep"]
+            else:
+                return None
+            step = find_spec["keystep"]
+            if step < 1 or first < 1 or last >= argc:
+                return None
+            keys.extend(args[first : last + 1 : step])
+        return keys
+
+
+def _spec_dict(spec):
+    # RESP2 sends a key spec as a flat [name, value, ...] list, RESP3 as a map.
+    if isinstance(spec, dict):
+        return {str_if_bytes(k): v for k, v in spec.items()}
+    return {str_if_bytes(k): v for k, v in zip(spec[::2], spec[1::2])}
+
 
 class CommandsParser(AbstractCommandsParser):
     """
@@ -140,7 +197,9 @@ def get_keys(self, redis_conn, *args):
 
         command = self.commands.get(cmd_name)
         if "movablekeys" in command["flags"]:
-            keys = self._get_moveable_keys(redis_conn, *args)
+            keys = self._get_keys_from_specs(command, args)
+            if keys is None:
+                keys = self._get_moveable_keys(redis_conn, *args)
         elif "pubsub" in command["flags"] or command["name"] == "pubsub":
             keys = self._get_pubsub_keys(*args)
         else:
@@ -316,7 +375,9 @@ async def get_keys(self, *args: Any) -> Optional[Tuple[str, ...]]:
 
         command = self.commands.get(cmd_name)
         if "movablekeys" in command["flags"]:
-            keys = await self._get_moveable_keys(*args)
+            keys = self._get_keys_from_specs(command, args)
+            if keys is None:
+                keys = await self._get_moveable_keys(*args)
         elif "pubsub" in command["flags"] or command["name"] == "pubsub":
             keys = self._get_pubsub_keys(*args)
         else:

redis/crc.py

@@ -15,6 +15,9 @@ def key_slot(key: EncodedT, bucket: int = REDIS_CLUSTER_HASH_SLOTS) -> int:
     :param key - bytes
     :param bucket - int
     """
+    if isinstance(key, memoryview):
+        # The encoder passes memoryviews through, and they have no find().
+        key = key.tobytes()
     start = key.find(b"{")
     if start > -1:
         end = key.find(b"}", start + 1)

tests/test_command_parser.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 import pytest
 from redis._parsers import CommandsParser
 from redis.commands.metadata import RequestPolicy, ResponsePolicy
@@ -57,24 +59,14 @@ def test_get_moveable_keys(self, r):
         ]
         args7 = ["MIGRATE", "192.168.1.34", 6379, "key1", 0, 5000]
 
-        assert_resp_response(
-            r,
-            sorted(commands_parser.get_keys(r, *args1)),
-            ["key1", "key2"],
-            [b"key1", b"key2"],
-        )
+        assert commands_parser.get_keys(r, *args1) == ["key1", "key2"]
         assert_resp_response(
             r,
             sorted(commands_parser.get_keys(r, *args2)),
             ["mystream", "writers"],
             [b"mystream", b"writers"],
         )
-        assert_resp_response(
-            r,
-            sorted(commands_parser.get_keys(r, *args3)),
-            ["out", "zset1", "zset2"],
-            [b"out", b"zset1", b"zset2"],
-        )
+        assert commands_parser.get_keys(r, *args3) == ["out", "zset1", "zset2"]
         assert_resp_response(
             r,
             sorted(commands_parser.get_keys(r, *args4)),
@@ -92,6 +84,29 @@ def test_get_moveable_keys(self, r):
             r, sorted(commands_parser.get_keys(r, *args7)), ["key1"], [b"key1"]
         )
 
+    @skip_if_server_version_lt("7.0.0")
+    def test_get_moveable_keys_from_key_specs(self, r):
+        commands_parser = CommandsParser(r)
+        cases = [
+            (["ZUNIONSTORE", "out", 2, "a", "b", "WEIGHTS", 2, 3], ["out", "a", "b"]),
+            (["LMPOP", 2, "a", "b", "LEFT"], ["a", "b"]),
+            (["BZMPOP", 0, 1, "a", "MIN"], ["a"]),
+            (["SINTERCARD", 2, "a", "b", "LIMIT", 1], ["a", "b"]),
+            (["FCALL", "f", 0, "arg"], []),
+        ]
+        with mock.patch.object(commands_parser, "_get_moveable_keys") as server:
+            for args, keys in cases:
+                assert commands_parser.get_keys(r, *args) == keys
+        server.assert_not_called()
+
+    @skip_if_server_version_lt("8.4.0")
+    def test_get_moveable_keys_from_key_specs_with_keystep(self, r):
+        commands_parser = CommandsParser(r)
+        args = ["MSETEX", 2, "k1", "v1", "k2", "v2", "EX", 10]
+        with mock.patch.object(commands_parser, "_get_moveable_keys") as server:
+            assert commands_parser.get_keys(r, *args) == ["k1", "k2"]
+        server.assert_not_called()
+
     # A bug in redis<7.0 causes this to fail: https://github.com/redis/redis/issues/9493
     @skip_if_server_version_lt("7.0.0")
     def test_get_eval_keys_with_0_keys(self, r):

tests/test_crc.py

@@ -57,3 +57,7 @@ def test_unclosed_brace_falls_back_to_whole_key(self):
     def test_custom_bucket_is_honored(self):
         assert key_slot(b"foo", 100) == crc_hqx(b"foo", 0) % 100
         assert 0 <= key_slot(b"123456789", 100) < 100
+
+    def test_memoryview_key_matches_bytes(self):
+        assert key_slot(memoryview(b"foo")) == key_slot(b"foo")
+        assert key_slot(memoryview(b"{user1000}.x")) == key_slot(b"{user1000}")