diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f6223fb8d94..6ef99f5ffd1 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -143,6 +143,7 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( _LUA_COUNT: Final = TypeAdapter(int) _PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") +_PIPELINE_COMMAND_NUMBER: Final = re.compile(r"Command # ([0-9]+) \(") _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) @@ -153,6 +154,13 @@ def _as_incrbyfloat_writes(value: float) -> str: return "0" if text == "-0" else text +def _failed_after_the_first_command(error: Exception) -> bool: + """Whether redis-py names a pipeline command after the first as the one that failed, so the first one applied.""" + message: Final = error.args[0] if error.args else None + number: Final = _PIPELINE_COMMAND_NUMBER.match(message) if isinstance(message, str) else None + return number is not None and int(number.group(1)) > 1 + + def _plain_float(reply: object) -> float | None: """The finite float a stored value spells in plain decimal or exponent form, or None for anything else.""" text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply @@ -1494,8 +1502,9 @@ class RedisCache(BaseCache): async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: """SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or - None when the server refused a command rather than failed to answer, or the key holds a value this does not - read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing.""" + None when the server refused SET NX GET rather than failed to answer, or the key holds a value this does not + read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing. A refused TTL command + raises, as it does in main's add.""" from redis.exceptions import ClusterError, RedisError from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError @@ -1510,7 +1519,8 @@ class RedisCache(BaseCache): pipe.expire(key, ttl) replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics except RedisError as e: - if isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)): + failed_to_answer: Final = isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)) + if failed_to_answer or _failed_after_the_first_command(e): raise return None if replies[0] is None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2a2ac5f7063..e673f16e5a6 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1727,8 +1727,6 @@ class _KeyspacePipeline(_RaisingPipeline): self.commands.append(("SETNX", name, value)) def _write(self, command: str) -> None: - if command in self.keyspace.refused: - raise NoPermissionError(f"this user has no permissions to run the '{command.lower()}' command") if self.writes_left is not None: if self.writes_left == 0: raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.") @@ -1737,6 +1735,8 @@ class _KeyspacePipeline(_RaisingPipeline): def _run(self, command: tuple[object, ...]) -> object: name: Final = str(command[1]) held: Final = self.keyspace.values.get(name) + if command[0] in self.keyspace.refused: + raise NoPermissionError(f"this user has no permissions to run the '{str(command[0]).lower()}' command") if command[0] == "EXPIRE": return held is not None self._write(str(command[0])) @@ -1756,10 +1756,12 @@ class _KeyspacePipeline(_RaisingPipeline): if self.lost: raise RedisConnectionError("Connection closed by server.") replies: Final[list[object]] = [] - for command in self.commands: + for number, command in enumerate(self.commands, start=1): try: replies.append(self._run(command)) except RedisError as e: + words: str = " ".join(str(part) for part in command) + e.args = (f"Command # {number} ({words}) of pipeline caused error: {e.args}",) replies.append(e) first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) if first_error is not None: @@ -1801,6 +1803,25 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redi added.assert_called_once_with() +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_raises_a_refused_ttl_command_as_mains_add_does_without_adding_again( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("EXPIRE",))) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(NoPermissionError): + await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert keyspace.values == {"ns:spend:key:k": "6"} + assert keyspace.pipelines == 1 + added.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"]) async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add(