fix: address EVALSHA ClusterPipeline review feedback

Author: ahmed5145

docs/lua_scripting.rst

@@ -114,8 +114,11 @@ The following commands are not supported:
 
 ``EVALSHA`` can be used inside a ``ClusterPipeline``. Keys must still map to
 the same hash slot (or ``numkeys`` may be ``0``, in which case the command is
-routed to a random primary). The script must already be loaded on the target
-node, for example via ``SCRIPT LOAD`` on the cluster client, which loads the
-script on all primaries. Other scripting helpers such as ``EVAL``,
-``load_scripts``, and ``script_load_for_pipeline`` remain unsupported on
-cluster pipelines.
+routed to a random primary). In a transactional pipeline
+(``pipeline(transaction=True)``), zero-key ``EVALSHA`` reuses the
+transaction's existing slot when one is already chosen so multiple zero-key
+scripts (or a mix with keyed commands) stay single-slot. The script must
+already be loaded on the target node, for example via ``SCRIPT LOAD`` on the
+cluster client, which loads the script on all primaries. Other scripting
+helpers such as ``EVAL``, ``load_scripts``, and ``script_load_for_pipeline``
+remain unsupported on cluster pipelines.

redis/asyncio/cluster.py

@@ -75,6 +75,7 @@
     LoadBalancingStrategy,
     block_pipeline_command,
     get_node_name,
+    is_zero_key_eval_command,
     parse_cluster_shards,
     parse_cluster_shards_unified,
     parse_cluster_shards_with_str_keys,
@@ -2917,6 +2918,8 @@ def __init__(self, pipe: ClusterPipeline) -> None:
         self._explicit_transaction = False
         self._watching = False
         self._pipeline_slots: Set[int] = set()
+        # True once a keyed (non-slot-agnostic) command has fixed the slot
+        self._transaction_has_keyed_slot = False
         self._transaction_node: Optional[ClusterNode] = None
         self._transaction_connection: Optional[Connection] = None
         self._executing = False
@@ -2925,6 +2928,34 @@ def __init__(self, pipe: ClusterPipeline) -> None:
             RedisCluster.ERRORS_ALLOW_RETRY + self.SLOT_REDIRECT_ERRORS
         )
 
+    async def _resolve_transaction_slot(self, *args) -> Optional[int]:
+        """
+        Pick a slot for a transactional pipeline command.
+
+        Zero-key EVAL/EVALSHA can run on any primary. Reuse an existing
+        transaction slot when present so multiple zero-key scripts (or a
+        mix with keyed commands) stay single-slot.
+        """
+        if args[0] in self.NO_SLOTS_COMMANDS:
+            return None
+
+        if is_zero_key_eval_command(*args):
+            if self._pipeline_slots:
+                return next(iter(self._pipeline_slots))
+            return await self._pipe.cluster_client._determine_slot(*args)
+
+        slot_number = await self._pipe.cluster_client._determine_slot(*args)
+        if (
+            slot_number is not None
+            and self._pipeline_slots
+            and slot_number not in self._pipeline_slots
+            and not self._transaction_has_keyed_slot
+        ):
+            # Prior slots came only from zero-key EVAL/EVALSHA; retarget.
+            self._pipeline_slots.clear()
+        self._transaction_has_keyed_slot = True
+        return slot_number
+
     def _get_client_and_connection_for_transaction(
         self,
     ) -> Tuple[ClusterNode, Connection]:
@@ -2982,7 +3013,7 @@ async def _execute_command(
 
         slot_number: Optional[int] = None
         if args[0] not in self.NO_SLOTS_COMMANDS:
-            slot_number = await self._pipe.cluster_client._determine_slot(*args)
+            slot_number = await self._resolve_transaction_slot(*args)
 
         if (
             self._watching or args[0] in self.IMMEDIATE_EXECUTE_COMMANDS
@@ -3327,6 +3358,7 @@ async def reset(self):
             self._watching = False
             self._explicit_transaction = False
             self._pipeline_slots = set()
+            self._transaction_has_keyed_slot = False
             self._executing = False
 
     def multi(self):

redis/cluster.py

@@ -3623,6 +3623,20 @@ def inner(*args, **kwargs):
     return inner
 
 
+def is_zero_key_eval_command(*args) -> bool:
+    """
+    True for EVAL/EVALSHA with numkeys=0 (may run on any primary).
+    """
+    if len(args) < 3:
+        return False
+    if str(args[0]).upper() not in ("EVAL", "EVALSHA"):
+        return False
+    try:
+        return int(args[2]) == 0
+    except (TypeError, ValueError):
+        return False
+
+
 # Blocked pipeline commands
 PIPELINE_BLOCKED_COMMANDS = (
     "BGREWRITEAOF",
@@ -4112,20 +4126,26 @@ def _send_cluster_commands(
                         command_flag = self.command_flags.get(command)
                         if not command_flag:
                             # Fallback to default policy.
-                            # Use determine_slot (not _get_command_keys) so EVAL /
-                            # EVALSHA with numkeys=0 work on Redis <7, matching
-                            # the async ClusterPipeline path.
-                            if not self._pipe.get_default_node():
-                                slot = None
-                            else:
-                                slot = self._pipe.determine_slot(*c.args)
-                            if slot is None:
-                                command_policies = CommandPolicies()
-                            else:
+                            # EVAL/EVALSHA must not use _get_command_keys(): Redis
+                            # <7 breaks on COMMAND GETKEYS when numkeys is 0.
+                            # Other unflagged commands keep the keyless fallback.
+                            if command in ("EVAL", "EVALSHA"):
                                 command_policies = CommandPolicies(
                                     request_policy=RequestPolicy.DEFAULT_KEYED,
                                     response_policy=ResponsePolicy.DEFAULT_KEYED,
                                 )
+                            else:
+                                if not self._pipe.get_default_node():
+                                    keys = None
+                                else:
+                                    keys = self._pipe._get_command_keys(*c.args)
+                                if not keys or len(keys) == 0:
+                                    command_policies = CommandPolicies()
+                                else:
+                                    command_policies = CommandPolicies(
+                                        request_policy=RequestPolicy.DEFAULT_KEYED,
+                                        response_policy=ResponsePolicy.DEFAULT_KEYED,
+                                    )
                         else:
                             if command_flag in self._pipe._command_flags_mapping:
                                 command_policies = CommandPolicies(
@@ -4419,13 +4439,43 @@ def __init__(self, pipe: ClusterPipeline):
         self._explicit_transaction = False
         self._watching = False
         self._pipeline_slots: Set[int] = set()
+        # True once a keyed (non-slot-agnostic) command has fixed the slot
+        self._transaction_has_keyed_slot = False
         self._transaction_connection: Optional[Connection] = None
         self._executing = False
         self._retry = copy(self._pipe.retry)
         self._retry.update_supported_errors(
             RedisCluster.ERRORS_ALLOW_RETRY + self.SLOT_REDIRECT_ERRORS
         )
 
+    def _resolve_transaction_slot(self, *args) -> Optional[int]:
+        """
+        Pick a slot for a transactional pipeline command.
+
+        Zero-key EVAL/EVALSHA can run on any primary. Reuse an existing
+        transaction slot when present so multiple zero-key scripts (or a
+        mix with keyed commands) stay single-slot.
+        """
+        if args[0] in ClusterPipeline.NO_SLOTS_COMMANDS:
+            return None
+
+        if is_zero_key_eval_command(*args):
+            if self._pipeline_slots:
+                return next(iter(self._pipeline_slots))
+            return self._pipe.determine_slot(*args)
+
+        slot_number = self._pipe.determine_slot(*args)
+        if (
+            slot_number is not None
+            and self._pipeline_slots
+            and slot_number not in self._pipeline_slots
+            and not self._transaction_has_keyed_slot
+        ):
+            # Prior slots came only from zero-key EVAL/EVALSHA; retarget.
+            self._pipeline_slots.clear()
+        self._transaction_has_keyed_slot = True
+        return slot_number
+
     def _get_client_and_connection_for_transaction(self) -> Tuple[Redis, Connection]:
         """
         Find a connection for a pipeline transaction.
@@ -4462,7 +4512,7 @@ def _get_client_and_connection_for_transaction(self) -> Tuple[Redis, Connection]
     def execute_command(self, *args, **kwargs):
         slot_number: Optional[int] = None
         if args[0] not in ClusterPipeline.NO_SLOTS_COMMANDS:
-            slot_number = self._pipe.determine_slot(*args)
+            slot_number = self._resolve_transaction_slot(*args)
 
         if (
             self._watching or args[0] in self.IMMEDIATE_EXECUTE_COMMANDS
@@ -4787,6 +4837,7 @@ def reset(self):
         self._watching = False
         self._explicit_transaction = False
         self._pipeline_slots = set()
+        self._transaction_has_keyed_slot = False
         self._executing = False
 
     def send_cluster_commands(

tests/test_asyncio/test_cluster.py

@@ -3863,6 +3863,74 @@ async def test_blocked_methods(self, r: RedisCluster) -> None:
                 "when running redis in cluster mode..."
             )
 
+    async def test_evalsha_not_blocked_on_cluster_pipeline(self) -> None:
+        """EVALSHA must be usable on async ClusterPipeline (see #2914)."""
+        assert "EVALSHA" not in PIPELINE_BLOCKED_COMMANDS
+        r = await get_mocked_redis_client(host=default_host, port=default_port)
+        try:
+            pipe = r.pipeline()
+            sha = "a" * 40
+            returned = pipe.evalsha(sha, 1, "foo", "bar")
+            assert returned is pipe
+            queue = pipe._execution_strategy._command_queue
+            assert len(queue) == 1
+            assert queue[0].args == ("EVALSHA", sha, 1, "foo", "bar")
+        finally:
+            await r.aclose()
+
+    async def test_evalsha_zero_keys_pipeline_execute(self) -> None:
+        """Async ClusterPipeline must execute EVALSHA with numkeys=0."""
+        r = await get_mocked_redis_client(host=default_host, port=default_port)
+        try:
+            mock_all_nodes_resp(r, 42)
+            sha = "a" * 40
+            async with r.pipeline() as pipe:
+                pipe.evalsha(sha, 0)
+                assert await pipe.execute() == [42]
+        finally:
+            await r.aclose()
+
+    async def test_evalsha_zero_keys_reuses_slot_in_transaction(self) -> None:
+        """Zero-key EVALSHA in a transactional async pipeline reuses one slot."""
+        r = await get_mocked_redis_client(host=default_host, port=default_port)
+        try:
+            sha = "a" * 40
+            async with r.pipeline(transaction=True) as pipe:
+                pipe.evalsha(sha, 0)
+                pipe.evalsha(sha, 0)
+                slots = pipe._execution_strategy._pipeline_slots
+                assert len(slots) == 1
+
+                keyed_slot = key_slot(b"foo")
+
+                async def _fake_determine_slot(*_args, **_kwargs):
+                    return keyed_slot
+
+                with mock.patch.object(
+                    pipe.cluster_client,
+                    "_determine_slot",
+                    side_effect=_fake_determine_slot,
+                ):
+                    pipe.set("foo", "bar")
+                assert pipe._execution_strategy._pipeline_slots == {keyed_slot}
+                assert pipe._execution_strategy._transaction_has_keyed_slot is True
+        finally:
+            await r.aclose()
+
+    async def test_evalsha_zero_keys_follows_keyed_slot_in_transaction(self) -> None:
+        """Zero-key EVALSHA after a keyed slot is fixed reuses that slot."""
+        r = await get_mocked_redis_client(host=default_host, port=default_port)
+        try:
+            sha = "a" * 40
+            async with r.pipeline(transaction=True) as pipe:
+                strategy = pipe._execution_strategy
+                strategy._pipeline_slots = {key_slot(b"foo")}
+                strategy._transaction_has_keyed_slot = True
+                pipe.evalsha(sha, 0)
+                assert strategy._pipeline_slots == {key_slot(b"foo")}
+        finally:
+            await r.aclose()
+
     async def test_empty_stack(self, r: RedisCluster) -> None:
         """If a pipeline is executed with no commands it should return a empty list."""
         p = r.pipeline()

tests/test_cluster.py

@@ -3886,6 +3886,45 @@ def test_evalsha_zero_keys_pipeline_execute_skips_get_command_keys():
         r.close()
 
 
+def test_evalsha_zero_keys_reuses_slot_in_transaction():
+    """
+    Zero-key EVALSHA in a transactional pipeline must reuse one slot so
+    execute() does not raise CrossSlotTransactionError.
+    """
+    r = get_mocked_redis_client(host=default_host, port=default_port)
+    try:
+        sha = "a" * 40
+        with r.pipeline(transaction=True) as pipe:
+            pipe.evalsha(sha, 0)
+            pipe.evalsha(sha, 0)
+            slots = pipe._execution_strategy._pipeline_slots
+            assert len(slots) == 1
+
+            # Keyed follow-up must retarget without needing a live COMMAND map.
+            keyed_slot = key_slot(b"foo")
+            with patch.object(pipe, "determine_slot", return_value=keyed_slot):
+                pipe.set("foo", "bar")
+            assert pipe._execution_strategy._pipeline_slots == {keyed_slot}
+            assert pipe._execution_strategy._transaction_has_keyed_slot is True
+    finally:
+        r.close()
+
+
+def test_evalsha_zero_keys_follows_keyed_slot_in_transaction():
+    """Zero-key EVALSHA after a keyed slot is fixed reuses that slot."""
+    r = get_mocked_redis_client(host=default_host, port=default_port)
+    try:
+        sha = "a" * 40
+        with r.pipeline(transaction=True) as pipe:
+            strategy = pipe._execution_strategy
+            strategy._pipeline_slots = {key_slot(b"foo")}
+            strategy._transaction_has_keyed_slot = True
+            pipe.evalsha(sha, 0)
+            assert strategy._pipeline_slots == {key_slot(b"foo")}
+    finally:
+        r.close()
+
+
 @pytest.mark.onlycluster
 class TestClusterPipeline:
     """
@@ -3944,6 +3983,17 @@ def test_evalsha_zero_keys_in_pipeline(self, r):
             pipe.evalsha(sha, 0)
             assert pipe.execute() == [42]
 
+    def test_evalsha_zero_keys_in_transaction_pipeline(self, r):
+        """
+        Multiple zero-key EVALSHA calls in a transactional pipeline share
+        one slot and execute successfully.
+        """
+        sha = r.script_load("return 42")
+        with r.pipeline(transaction=True) as pipe:
+            pipe.evalsha(sha, 0)
+            pipe.evalsha(sha, 0)
+            assert pipe.execute() == [42, 42]
+
     def test_redis_cluster_pipeline(self, r):
         """
         Test that we can use a pipeline with the RedisCluster class