mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(redis): add native seconds transport expiry commands
This commit is contained in:
Devin AI
parent
d373b499df
commit
5b7f7125c2
5 changed files with 213 additions and 2 deletions
|
|
@ -37,7 +37,9 @@ class AsyncRedisCommands(Protocol):
|
|||
|
||||
def mget(self, keys: Sequence[str]) -> Awaitable[object]: ...
|
||||
|
||||
def set(self, name: str, value: Arg, *, nx: bool, px: int | None) -> Awaitable[object]: ...
|
||||
def set(
|
||||
self, name: str, value: Arg, *, nx: bool, px: int | None = None, ex: timedelta | None = None
|
||||
) -> Awaitable[object]: ...
|
||||
|
||||
def delete(self, *names: str) -> Awaitable[object]: ...
|
||||
|
||||
|
|
@ -47,6 +49,8 @@ class AsyncRedisCommands(Protocol):
|
|||
|
||||
def pexpire(self, name: str, time: int) -> Awaitable[object]: ...
|
||||
|
||||
def expire(self, name: str, time: int) -> Awaitable[object]: ...
|
||||
|
||||
def ttl(self, name: str) -> Awaitable[object]: ...
|
||||
|
||||
def rpush(self, name: str, *values: Arg) -> Awaitable[object]: ...
|
||||
|
|
|
|||
|
|
@ -53,12 +53,18 @@ class RedisCommandBatch(Protocol):
|
|||
|
||||
def delete(self, *keys: str) -> BatchReply[int | RedisFailure]: ...
|
||||
|
||||
def set_seconds(
|
||||
self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False
|
||||
) -> BatchReply[bool | RedisFailure]: ...
|
||||
|
||||
def incrby(self, key: str, amount: int) -> BatchReply[int | RedisFailure]: ...
|
||||
|
||||
def incrbyfloat(self, key: str, amount: float) -> BatchReply[float | RedisFailure]: ...
|
||||
|
||||
def expire(self, key: str, ttl: timedelta) -> BatchReply[bool | RedisFailure]: ...
|
||||
|
||||
def expire_seconds(self, key: str, ttl: timedelta) -> BatchReply[bool | RedisFailure]: ...
|
||||
|
||||
def ttl(self, key: str) -> BatchReply[int | RedisFailure]: ...
|
||||
|
||||
def rpush(self, key: str, *values: Arg) -> BatchReply[int | RedisFailure]: ...
|
||||
|
|
@ -89,12 +95,18 @@ class RedisTransport(Protocol):
|
|||
|
||||
async def delete(self, *keys: str) -> int | RedisFailure: ...
|
||||
|
||||
async def set_seconds(
|
||||
self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False
|
||||
) -> bool | RedisFailure: ...
|
||||
|
||||
async def incrby(self, key: str, amount: int) -> int | RedisFailure: ...
|
||||
|
||||
async def incrbyfloat(self, key: str, amount: float) -> float | RedisFailure: ...
|
||||
|
||||
async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
|
||||
async def expire_seconds(self, key: str, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
|
||||
async def ttl(self, key: str) -> int | RedisFailure: ...
|
||||
|
||||
async def rpush(self, key: str, *values: Arg) -> int | RedisFailure: ...
|
||||
|
|
|
|||
|
|
@ -61,6 +61,13 @@ class RedisPyTransport:
|
|||
px: Final = None if ttl is None else ttl_ms(ttl)
|
||||
return _bool_reply(await self._run("set", lambda: self.client().set(self._key(key), value, px=px, nx=nx)))
|
||||
|
||||
async def set_seconds(
|
||||
self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False
|
||||
) -> bool | RedisFailure:
|
||||
return _set_reply(
|
||||
await self._run("set_seconds", lambda: self.client().set(self._key(key), value, ex=ttl, nx=nx))
|
||||
)
|
||||
|
||||
async def delete(self, *keys: str) -> int | RedisFailure:
|
||||
namespaced: Final = tuple(self._key(k) for k in keys)
|
||||
return _int_reply(await self._run("delete", lambda: self.client().delete(*namespaced)))
|
||||
|
|
@ -74,6 +81,11 @@ class RedisPyTransport:
|
|||
async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure:
|
||||
return _bool_reply(await self._run("expire", lambda: self.client().pexpire(self._key(key), ttl_ms(ttl))))
|
||||
|
||||
async def expire_seconds(self, key: str, ttl: timedelta) -> bool | RedisFailure:
|
||||
return _expire_reply(
|
||||
await self._run("expire_seconds", lambda: self.client().expire(self._key(key), int(ttl.total_seconds())))
|
||||
)
|
||||
|
||||
async def ttl(self, key: str) -> int | RedisFailure:
|
||||
return _int_reply(await self._run("ttl", lambda: self.client().ttl(self._key(key))))
|
||||
|
||||
|
|
@ -228,6 +240,17 @@ class RedisPyCommandBatch:
|
|||
lambda: self.transport.delete(*keys),
|
||||
)
|
||||
|
||||
def set_seconds(
|
||||
self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False
|
||||
) -> BatchResult[bool | RedisFailure]:
|
||||
name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own
|
||||
self.engine.note_present(key)
|
||||
return self._declare(
|
||||
lambda pipe: _one(pipe.set(name, value, ex=ttl, nx=nx)),
|
||||
lambda replies: _set_reply(replies[0]),
|
||||
lambda: self.transport.set_seconds(key, value, ttl=ttl, nx=nx),
|
||||
)
|
||||
|
||||
def incrby(self, key: str, amount: int) -> BatchResult[int | RedisFailure]:
|
||||
name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own
|
||||
self.engine.note_present(key)
|
||||
|
|
@ -260,6 +283,14 @@ class RedisPyCommandBatch:
|
|||
lambda pipe: _one(pipe.ttl(name)), lambda replies: _int_reply(replies[0]), lambda: self.transport.ttl(key)
|
||||
)
|
||||
|
||||
def expire_seconds(self, key: str, ttl: timedelta) -> BatchResult[bool | RedisFailure]:
|
||||
name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own
|
||||
return self._declare(
|
||||
lambda pipe: _one(pipe.expire(name, ttl)),
|
||||
lambda replies: _expire_reply(replies[0]),
|
||||
lambda: self.transport.expire_seconds(key, ttl),
|
||||
)
|
||||
|
||||
def rpush(self, key: str, *values: Arg) -> BatchResult[int | RedisFailure]:
|
||||
name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own
|
||||
self.engine.note_present(key)
|
||||
|
|
@ -491,6 +522,24 @@ def _bool_reply(reply: object) -> bool | RedisFailure:
|
|||
return reply is not None and reply is not False and reply != 0
|
||||
|
||||
|
||||
def _expire_reply(reply: object) -> bool | RedisFailure:
|
||||
if isinstance(reply, (RedisFailure, bool)):
|
||||
return reply
|
||||
if type(reply) is int and reply in (0, 1):
|
||||
return reply == 1
|
||||
return _decode_failure(reply)
|
||||
|
||||
|
||||
def _set_reply(reply: object) -> bool | RedisFailure:
|
||||
if isinstance(reply, RedisFailure):
|
||||
return reply
|
||||
if reply is None or reply is False:
|
||||
return False
|
||||
if reply is True or reply in (b"OK", "OK"):
|
||||
return True
|
||||
return _decode_failure(reply)
|
||||
|
||||
|
||||
def _int_reply(reply: object) -> int | RedisFailure:
|
||||
if isinstance(reply, RedisFailure):
|
||||
return reply
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ class ExpiredLockTransport:
|
|||
async def set(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> bool | RedisFailure:
|
||||
return not nx
|
||||
|
||||
async def set_seconds(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> NoReturn:
|
||||
raise AssertionError("unexpected set_seconds")
|
||||
|
||||
async def mget(self, keys: Sequence[str]) -> tuple[bytes | None, ...] | RedisFailure:
|
||||
raise AssertionError("unexpected mget")
|
||||
|
||||
|
|
@ -35,6 +38,9 @@ class ExpiredLockTransport:
|
|||
async def expire(self, key: str, ttl: timedelta) -> NoReturn:
|
||||
raise AssertionError("unexpected expire")
|
||||
|
||||
async def expire_seconds(self, key: str, ttl: timedelta) -> NoReturn:
|
||||
raise AssertionError("unexpected expire_seconds")
|
||||
|
||||
async def ttl(self, key: str) -> NoReturn:
|
||||
raise AssertionError("unexpected ttl")
|
||||
|
||||
|
|
|
|||
|
|
@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Sequence
|
|||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
from redis.asyncio.client import Pipeline, Redis
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
import litellm
|
||||
|
|
@ -389,6 +391,8 @@ _WRONG_TYPE_COMMANDS: Final[tuple[_Command, ...]] = (
|
|||
_OUTAGE_COMMANDS: Final[tuple[_Command, ...]] = (
|
||||
lambda t: t.incrby("k", 1),
|
||||
lambda t: t.expire("k", timedelta(seconds=1)),
|
||||
lambda t: t.expire_seconds("k", timedelta(seconds=1)),
|
||||
lambda t: t.set_seconds("k", "v", ttl=timedelta(seconds=1)),
|
||||
lambda t: t.ttl("k"),
|
||||
lambda t: t.lpop("k"),
|
||||
lambda t: t.scan("k", limit=1),
|
||||
|
|
@ -428,6 +432,139 @@ async def test_expire_sets_the_ttl_in_milliseconds_and_ttl_reads_it_back(
|
|||
assert await transport.ttl("absent") == -2
|
||||
|
||||
|
||||
def _typed_pipeline(client: "Redis[bytes]") -> "Pipeline[bytes]":
|
||||
return client.pipeline(transaction=False)
|
||||
|
||||
|
||||
def _seconds_commands(transport: RedisPyTransport, mode: str) -> RedisPyTransport | RedisPyCommandBatch:
|
||||
if mode == "direct":
|
||||
return transport
|
||||
if mode == "alone":
|
||||
backend: Final = replace(transport.batch().engine.backend, runs_alone=True)
|
||||
return RedisPyCommandBatch(transport, PipelineBatch(backend, name="seconds"))
|
||||
return transport.batch(transaction=mode == "transaction")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone"))
|
||||
@pytest.mark.parametrize("seconds", (2, 1.9, 0.5, 0, -0.5, -1))
|
||||
async def test_native_expire_truncates_seconds_and_deletes_nonpositive_expiry(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, seconds: float
|
||||
) -> None:
|
||||
transport: Final = fake_redis_transport(raw, failure_threshold=1, namespace="ns")
|
||||
commands: Final = _seconds_commands(transport, mode)
|
||||
with patch("time.time", return_value=1700000000.0):
|
||||
await raw.set("ns:k", "v")
|
||||
assert await commands.expire_seconds("ns:k", timedelta(seconds=seconds)) is True
|
||||
assert await commands.expire_seconds("absent", timedelta(seconds=seconds)) is False
|
||||
if seconds >= 1:
|
||||
assert await raw.pttl("ns:k") == int(seconds) * 1000
|
||||
assert await raw.get("ns:k") == b"v"
|
||||
else:
|
||||
with patch("time.time", return_value=1700000000.001):
|
||||
assert await raw.get("ns:k") is None
|
||||
assert await raw.exists("ns:ns:k") == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone"))
|
||||
async def test_native_set_preserves_seconds_expiry_persistence_and_nx(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str
|
||||
) -> None:
|
||||
commands: Final = _seconds_commands(fake_redis_transport(raw, failure_threshold=1, namespace="ns"), mode)
|
||||
with patch("time.time", return_value=1700000000.0):
|
||||
assert await commands.set_seconds("k", b"1.25", ttl=timedelta(seconds=1.9), nx=True) is True
|
||||
assert await commands.set_seconds("ns:k", "rejected", ttl=None, nx=True) is False
|
||||
assert await raw.get("ns:k") == b"1.25"
|
||||
assert await raw.pttl("ns:k") == 1000
|
||||
assert await commands.set_seconds("ns:k", "persist", ttl=None) is True
|
||||
assert await raw.get("ns:k") == b"persist"
|
||||
assert await raw.pttl("ns:k") == -1
|
||||
assert await raw.keys("*") == [b"ns:k"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone"))
|
||||
@pytest.mark.parametrize("seconds", (0.5, 0, -0.5, -1))
|
||||
async def test_native_set_rejection_keeps_value_and_connection_and_emits_one_event(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, seconds: float
|
||||
) -> None:
|
||||
events: Final = _RecordingEvents()
|
||||
transport: Final = fake_redis_transport(raw, failure_threshold=1, events=events, namespace="ns")
|
||||
commands: Final = _seconds_commands(transport, mode)
|
||||
await raw.set("ns:k", "before")
|
||||
with service_target("seconds_test"):
|
||||
rejected: Final = await commands.set_seconds("k", "after", ttl=timedelta(seconds=seconds))
|
||||
assert isinstance(rejected, RedisFailure) and rejected.kind == "server"
|
||||
assert await raw.get("ns:k") == b"before"
|
||||
assert len(events.seen) == 1
|
||||
assert events.seen[0].outcome == ("failure" if mode in ("direct", "alone") else "success")
|
||||
assert events.seen[0].target == "seconds_test"
|
||||
assert await commands.expire_seconds("k", timedelta(seconds=30)) is True
|
||||
await _settle()
|
||||
assert len(events.seen) == 2
|
||||
assert events.seen[1].outcome == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("direct", "batch"))
|
||||
async def test_native_seconds_outage_opens_breaker_without_duplicate_events(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str
|
||||
) -> None:
|
||||
events: Final = _RecordingEvents()
|
||||
commands: Final = _seconds_commands(fake_redis_transport(raw, failure_threshold=1, events=events), mode)
|
||||
server.connected = False
|
||||
failed: Final = await commands.expire_seconds("k", timedelta(seconds=30))
|
||||
assert isinstance(failed, RedisFailure) and failed.kind == "transport"
|
||||
server.connected = True
|
||||
refused: Final = await commands.set_seconds("k", "v", ttl=None)
|
||||
assert isinstance(refused, RedisFailure) and refused.kind == "breaker_open"
|
||||
await _settle()
|
||||
assert tuple(event.outcome for event in events.seen) == (
|
||||
("failure",) if mode == "direct" else ("failure", "failure")
|
||||
)
|
||||
assert await raw.get("k") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("direct", "batch"))
|
||||
@pytest.mark.parametrize("reply", (b"malformed", 2, 1.0, (True,), [True], {"unexpected": True}))
|
||||
async def test_native_seconds_rejects_malformed_command_replies(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, reply: object
|
||||
) -> None:
|
||||
transport: Final = fake_redis_transport(raw, failure_threshold=1)
|
||||
pipeline: Final = _typed_pipeline(raw)
|
||||
with (
|
||||
patch.object(raw, "expire", new=AsyncMock(return_value=reply)),
|
||||
patch.object(raw, "set", new=AsyncMock(return_value=reply)),
|
||||
patch.object(raw, "pipeline", return_value=pipeline),
|
||||
patch.object(pipeline, "execute", return_value=[reply]),
|
||||
):
|
||||
commands: Final = _seconds_commands(transport, mode)
|
||||
expired: Final = await commands.expire_seconds("k", timedelta(seconds=30))
|
||||
written: Final = await commands.set_seconds("k", "v", ttl=None)
|
||||
assert isinstance(expired, RedisFailure) and expired.kind == "decode"
|
||||
assert isinstance(written, RedisFailure) and written.kind == "decode"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_expire_lost_reply_is_unknown_applied_without_retry(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
) -> None:
|
||||
transport: Final = fake_redis_transport(raw, failure_threshold=1)
|
||||
expire: Final = raw.expire
|
||||
await raw.set("k", "v")
|
||||
|
||||
async def lose_reply(name: str, time: int) -> bool:
|
||||
await expire(name, time)
|
||||
raise RedisConnectionError("reply lost after expiry applied")
|
||||
|
||||
with patch("time.time", return_value=1700000000.0), patch.object(raw, "expire", side_effect=lose_reply):
|
||||
result: Final = await transport.expire_seconds("k", timedelta(seconds=30))
|
||||
assert isinstance(result, RedisFailure) and result.kind == "transport"
|
||||
assert await raw.pttl("k") == 30000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_commands_push_trim_and_pop_from_the_head(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
|
|
@ -654,10 +791,13 @@ async def test_request_batch_commands_share_the_round_trip_of_redis_batch_caller
|
|||
old_read = legacy.mget(["legacy"])
|
||||
new_read = commands.get("moved")
|
||||
missing = commands.get("absent")
|
||||
written = commands.set_seconds("seconds", "v", ttl=None)
|
||||
expired = commands.expire_seconds("seconds", timedelta(seconds=30))
|
||||
|
||||
assert await old_read == {"legacy": "old"}
|
||||
assert await new_read == b"new"
|
||||
assert await missing is None
|
||||
assert await written is True and await expired is True
|
||||
assert legacy.flushes == 1
|
||||
assert legacy.read_as_missing("absent") is True
|
||||
assert legacy.read_as_missing("moved") is False
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue