From c6e8207ce33fb4311aeb789f50b095baf14e1fe0 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:51:01 -0400 Subject: [PATCH 1/3] Add atomic conditional writes --- README.md | 20 +++- docs/api/protocols.md | 24 +++++ docs/stores.md | 2 + src/key_value/aio/protocols/__init__.py | 1 + src/key_value/aio/protocols/key_value.py | 31 ++++++ src/key_value/aio/stores/base.py | 42 ++++++++ src/key_value/aio/stores/memory/store.py | 50 +++++++-- src/key_value/aio/stores/redis/store.py | 43 +++++++- tests/protocols/test_types.py | 15 ++- tests/stores/base.py | 102 +++++++++++++++++- tests/stores/memory/test_memory.py | 4 +- tests/stores/redis/test_redis.py | 12 ++- .../stores/redis/test_redis_put_if_absent.py | 27 +++++ 13 files changed, 352 insertions(+), 21 deletions(-) create mode 100644 tests/stores/redis/test_redis_put_if_absent.py diff --git a/README.md b/README.md index 35291871..9d95fadb 100644 --- a/README.md +++ b/README.md @@ -174,8 +174,8 @@ asyncio.run(main()) - **Async**: `key_value.aio.protocols.AsyncKeyValue` — async `get/put/delete/ttl` and bulk variants; optional protocol segments for - culling, destroying stores/collections, and enumerating keys/collections - implemented by capable stores. + conditional writes, culling, destroying stores/collections, and enumerating + keys/collections implemented by capable stores. The protocols offer a simple interface for your application to interact with the store: @@ -194,6 +194,22 @@ ttl(key: str, collection: str | None = None) -> tuple[dict[str, Any] | None, flo ttl_many(keys: list[str], collection: str | None = None) -> list[tuple[dict[str, Any] | None, float | None]]: ``` +Stores with native atomic conditional writes implement +`AsyncPutIfAbsentProtocol`. Use a runtime check before calling the optional +method: + +```python +from key_value.aio.protocols import AsyncPutIfAbsentProtocol + +if isinstance(store, AsyncPutIfAbsentProtocol): + stored = await store.put_if_absent( + key="request-123", + value={"status": "started"}, + collection="idempotency", + ttl=300, + ) +``` + ### Stores The library provides multiple store implementations organized into three diff --git a/docs/api/protocols.md b/docs/api/protocols.md index 50e80ca6..0901a68e 100644 --- a/docs/api/protocols.md +++ b/docs/api/protocols.md @@ -11,3 +11,27 @@ composability. show_source: true members: true show_root_heading: true + +## Optional Atomic Conditional Writes + +`AsyncPutIfAbsentProtocol` is implemented only by stores that can atomically +check for a missing key and write it. Check the capability at runtime before +calling `put_if_absent()`. + +```python +from key_value.aio.protocols import AsyncPutIfAbsentProtocol + +if isinstance(store, AsyncPutIfAbsentProtocol): + stored = await store.put_if_absent( + key="request-123", + value={"status": "started"}, + collection="idempotency", + ttl=300, + ) +``` + +::: key_value.aio.protocols.key_value.AsyncPutIfAbsentProtocol + options: + show_source: true + members: true + show_root_heading: true diff --git a/docs/stores.md b/docs/stores.md index 5912c9d2..537bec10 100644 --- a/docs/stores.md +++ b/docs/stores.md @@ -71,6 +71,7 @@ pip install py-key-value-aio[memory] - Extremely fast - No external dependencies - Thread-safe +- Atomic `put_if_absent()` support --- @@ -595,6 +596,7 @@ pip install py-key-value-aio[redis] - Production-ready - Rich feature set - Horizontal scaling support +- Atomic `put_if_absent()` support - SSL/TLS and mutual TLS connection options - **Stable storage format** diff --git a/src/key_value/aio/protocols/__init__.py b/src/key_value/aio/protocols/__init__.py index 1314fc2c..20860d68 100644 --- a/src/key_value/aio/protocols/__init__.py +++ b/src/key_value/aio/protocols/__init__.py @@ -1 +1,2 @@ from key_value.aio.protocols.key_value import AsyncKeyValue as AsyncKeyValue +from key_value.aio.protocols.key_value import AsyncPutIfAbsentProtocol as AsyncPutIfAbsentProtocol diff --git a/src/key_value/aio/protocols/key_value.py b/src/key_value/aio/protocols/key_value.py index 5e8a1426..e8b72703 100644 --- a/src/key_value/aio/protocols/key_value.py +++ b/src/key_value/aio/protocols/key_value.py @@ -124,6 +124,37 @@ async def delete_many(self, keys: Sequence[str], *, collection: str | None = Non ... +@runtime_checkable +class AsyncPutIfAbsentProtocol(Protocol): + """Protocol segment for atomic conditional writes.""" + + async def put_if_absent( + self, + key: str, + value: Mapping[str, Any], + *, + collection: str | None = None, + ttl: SupportsFloat | None = None, + ) -> bool: + """Store a value only when the key does not already exist. + + The existence check and write must be one atomic operation. Expired + entries are treated as absent. + + Args: + key: The key to store the value under. + value: The value to store. + collection: The collection to store the value in. If no collection + is provided, the default collection is used. + ttl: Optional time-to-live in seconds. + + Returns: + True when the value was stored, or False when an unexpired value + already exists. + """ + ... + + @runtime_checkable class AsyncCullProtocol(Protocol): async def cull(self) -> None: diff --git a/src/key_value/aio/stores/base.py b/src/key_value/aio/stores/base.py index da565962..04ef17ce 100644 --- a/src/key_value/aio/stores/base.py +++ b/src/key_value/aio/stores/base.py @@ -27,6 +27,7 @@ AsyncEnumerateCollectionsProtocol, AsyncEnumerateKeysProtocol, AsyncKeyValueProtocol, + AsyncPutIfAbsentProtocol, ) SEED_DATA_TYPE = Mapping[str, Mapping[str, Mapping[str, Any]]] @@ -407,6 +408,47 @@ def _warn_about_stability(self) -> None: ) +class BasePutIfAbsentStore(BaseStore, AsyncPutIfAbsentProtocol, ABC): + """Base class for stores with native atomic conditional writes.""" + + @abstractmethod + async def _put_managed_entry_if_absent( + self, + *, + collection: str, + key: str, + managed_entry: ManagedEntry, + ) -> bool: + """Atomically store a managed entry only when its key is absent.""" + ... + + @bear_enforce + @override + async def put_if_absent( + self, + key: str, + value: Mapping[str, Any], + *, + collection: str | None = None, + ttl: SupportsFloat | None = None, + ) -> bool: + """Store a value only when the key does not already exist.""" + collection = collection or self.default_collection + await self.setup_collection(collection=collection) + + created_at, _, expires_at = prepare_entry_timestamps(ttl=ttl) + managed_entry = ManagedEntry( + value=value, + created_at=created_at, + expires_at=expires_at, + ) + return await self._put_managed_entry_if_absent( + collection=collection, + key=key, + managed_entry=managed_entry, + ) + + class BaseEnumerateKeysStore(BaseStore, AsyncEnumerateKeysProtocol, ABC): """An abstract base class for enumerate key-value stores. diff --git a/src/key_value/aio/stores/memory/store.py b/src/key_value/aio/stores/memory/store.py index 4272a0f5..189c08bd 100644 --- a/src/key_value/aio/stores/memory/store.py +++ b/src/key_value/aio/stores/memory/store.py @@ -1,6 +1,7 @@ import sys from dataclasses import dataclass from datetime import datetime +from threading import RLock from typing import Any from typing_extensions import override @@ -13,6 +14,7 @@ BaseDestroyStore, BaseEnumerateCollectionsStore, BaseEnumerateKeysStore, + BasePutIfAbsentStore, ) try: @@ -64,30 +66,45 @@ def __init__(self, max_entries: int | None = None): ) self._serialization_adapter = BasicSerializationAdapter() + self._lock = RLock() def get(self, key: str) -> ManagedEntry | None: - managed_entry_str: MemoryCacheEntry | None = self._cache.get(key) + with self._lock: + managed_entry_str: MemoryCacheEntry | None = self._cache.get(key) - if managed_entry_str is None: - return None + if managed_entry_str is None: + return None - managed_entry: ManagedEntry = self._serialization_adapter.load_json(json_str=managed_entry_str.json_str) + managed_entry: ManagedEntry = self._serialization_adapter.load_json(json_str=managed_entry_str.json_str) - return managed_entry + return managed_entry def put(self, key: str, value: ManagedEntry) -> None: - json_str: str = self._serialization_adapter.dump_json(entry=value) - self._cache[key] = MemoryCacheEntry(json_str=json_str, expires_at=value.expires_at) + with self._lock: + json_str: str = self._serialization_adapter.dump_json(entry=value) + self._cache[key] = MemoryCacheEntry(json_str=json_str, expires_at=value.expires_at) + + def put_if_absent(self, key: str, value: ManagedEntry) -> bool: + with self._lock: + existing = self.get(key) + if existing is not None and not existing.is_expired: + return False + self.put(key, value) + return True def delete(self, key: str) -> bool: - return self._cache.pop(key, None) is not None + with self._lock: + return self._cache.pop(key, None) is not None def keys(self, *, limit: int | None = None) -> list[str]: - limit = min(limit or DEFAULT_PAGE_SIZE, PAGE_LIMIT) - return list(self._cache.keys())[:limit] + with self._lock: + limit = min(limit or DEFAULT_PAGE_SIZE, PAGE_LIMIT) + return list(self._cache.keys())[:limit] -class MemoryStore(BaseDestroyStore, BaseDestroyCollectionStore, BaseEnumerateCollectionsStore, BaseEnumerateKeysStore): +class MemoryStore( + BasePutIfAbsentStore, BaseDestroyStore, BaseDestroyCollectionStore, BaseEnumerateCollectionsStore, BaseEnumerateKeysStore +): """A fixed-size in-memory key-value store using TLRU (Time-aware Least Recently Used) cache.""" max_entries_per_collection: int @@ -173,6 +190,17 @@ async def _put_managed_entry( collection_cache = self._get_collection_or_raise(collection) collection_cache.put(key=key, value=managed_entry) + @override + async def _put_managed_entry_if_absent( + self, + *, + key: str, + collection: str, + managed_entry: ManagedEntry, + ) -> bool: + collection_cache = self._get_collection_or_raise(collection) + return collection_cache.put_if_absent(key=key, value=managed_entry) + @override async def _delete_managed_entry(self, *, key: str, collection: str) -> bool: collection_cache = self._get_collection_or_raise(collection) diff --git a/src/key_value/aio/stores/redis/store.py b/src/key_value/aio/stores/redis/store.py index 324c8cc6..002db135 100644 --- a/src/key_value/aio/stores/redis/store.py +++ b/src/key_value/aio/stores/redis/store.py @@ -10,7 +10,13 @@ from key_value.aio._utils.managed_entry import ManagedEntry from key_value.aio._utils.serialization import BasicSerializationAdapter, SerializationAdapter from key_value.aio.errors import DeserializationError -from key_value.aio.stores.base import BaseContextManagerStore, BaseDestroyStore, BaseEnumerateKeysStore, BaseStore +from key_value.aio.stores.base import ( + BaseContextManagerStore, + BaseDestroyStore, + BaseEnumerateKeysStore, + BasePutIfAbsentStore, + BaseStore, +) try: from redis.asyncio import Redis @@ -163,6 +169,17 @@ async def _redis_setex(client: Redis, name: str, time: int, value: str) -> None: _ = await client.setex(name=name, time=time, value=value) +async def _redis_set_if_absent( + client: Redis, + name: str, + value: str, + ttl: int | None, +) -> bool: + """Set a value atomically when its key does not exist.""" + result = await client.set(name=name, value=value, nx=True, ex=ttl) + return bool(result) + + async def _redis_pipeline_execute(pipeline: Any) -> None: """Execute a Redis pipeline.""" await pipeline.execute() @@ -183,7 +200,7 @@ async def _redis_flushdb(client: Redis) -> bool: return await client.flushdb() # pyright: ignore[reportUnknownMemberType] -class RedisStore(BaseDestroyStore, BaseEnumerateKeysStore, BaseContextManagerStore, BaseStore): +class RedisStore(BasePutIfAbsentStore, BaseDestroyStore, BaseEnumerateKeysStore, BaseContextManagerStore, BaseStore): """Redis-based key-value store.""" _client: Redis @@ -359,6 +376,28 @@ async def _put_managed_entry( else: await _redis_set(self._client, combo_key, json_value) + @override + async def _put_managed_entry_if_absent( + self, + *, + key: str, + collection: str, + managed_entry: ManagedEntry, + ) -> bool: + combo_key = compound_key(collection=collection, key=key) + json_value = self._adapter.dump_json( + entry=managed_entry, + key=key, + collection=collection, + ) + ttl = max(int(managed_entry.ttl), 1) if managed_entry.ttl is not None else None + return await _redis_set_if_absent( + self._client, + combo_key, + json_value, + ttl, + ) + @override async def _put_managed_entries( self, diff --git a/tests/protocols/test_types.py b/tests/protocols/test_types.py index 9df69a27..bc4d0200 100644 --- a/tests/protocols/test_types.py +++ b/tests/protocols/test_types.py @@ -1,5 +1,11 @@ -from key_value.aio.protocols.key_value import AsyncKeyValue +import pytest + +from key_value.aio.protocols.key_value import ( + AsyncKeyValue, + AsyncPutIfAbsentProtocol, +) from key_value.aio.stores.memory import MemoryStore +from key_value.aio.stores.null import NullStore async def test_key_value_protocol(): @@ -15,3 +21,10 @@ async def test_protocol(key_value: AsyncKeyValue): assert await memory_store.get(collection="test", key="test") is None assert await memory_store.get(collection="test", key="test_2") == {"test": "test"} + + +def test_put_if_absent_is_an_optional_protocol(): + assert isinstance(MemoryStore(), AsyncPutIfAbsentProtocol) + with pytest.warns(UserWarning, match="configured store is unstable"): + null_store = NullStore() + assert not isinstance(null_store, AsyncPutIfAbsentProtocol) diff --git a/tests/stores/base.py b/tests/stores/base.py index 733c11dc..f0ea668f 100644 --- a/tests/stores/base.py +++ b/tests/stores/base.py @@ -12,7 +12,10 @@ from pydantic import AnyHttpUrl from key_value.aio.errors import InvalidTTLError, SerializationError -from key_value.aio.protocols.key_value import AsyncKeyValueProtocol +from key_value.aio.protocols.key_value import ( + AsyncKeyValueProtocol, + AsyncPutIfAbsentProtocol, +) from key_value.aio.stores.base import BaseContextManagerStore, BaseStore from tests.conftest import async_running_in_event_loop from tests.shared.cases import ( @@ -271,6 +274,103 @@ async def test_minimum_put_many_delete_many_performance(self, store: BaseStore): assert await store.delete_many(collection="test_collection", keys=keys) == 10 +class PutIfAbsentStoreTestMixin: + async def test_put_if_absent_stores_new_value(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + + stored = await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 1}, + ) + + assert stored is True + assert await store.get(collection="test", key="conditional") == {"winner": 1} + + async def test_put_if_absent_preserves_existing_value(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + await store.put( + collection="test", + key="conditional", + value={"winner": 1}, + ) + + stored = await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 2}, + ) + + assert stored is False + assert await store.get(collection="test", key="conditional") == {"winner": 1} + + async def test_put_if_absent_applies_ttl(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + + stored = await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 1}, + ttl=100, + ) + value, ttl = await store.ttl(collection="test", key="conditional") + + assert stored is True + assert value == {"winner": 1} + assert ttl == IsFloat(approx=100, delta=2) + + async def test_put_if_absent_rejects_invalid_ttl(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + + with pytest.raises(InvalidTTLError): + await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 1}, + ttl=-1, + ) + + async def test_put_if_absent_accepts_expired_key(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + assert await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 1}, + ttl=1, + ) + + for _ in range(20): + if await store.get(collection="test", key="conditional") is None: + break + await asyncio.sleep(0.1) + else: + pytest.fail("conditional entry did not expire") + + assert await store.put_if_absent( + collection="test", + key="conditional", + value={"winner": 2}, + ) + + @pytest.mark.skipif(condition=not async_running_in_event_loop(), reason="Cannot run concurrent operations outside of event loop") + async def test_put_if_absent_is_atomic(self, store: BaseStore): + assert isinstance(store, AsyncPutIfAbsentProtocol) + + results = await asyncio.gather( + *( + store.put_if_absent( + collection="test", + key="conditional", + value={"candidate": candidate}, + ) + for candidate in range(20) + ) + ) + + assert results.count(True) == 1 + assert results.count(False) == 19 + + class ContextManagerStoreTestMixin: @pytest.fixture(params=[True, False], ids=["with_ctx_manager", "no_ctx_manager"], autouse=True) async def enter_exit_store( diff --git a/tests/stores/memory/test_memory.py b/tests/stores/memory/test_memory.py index a98a2a31..7b6f0f57 100644 --- a/tests/stores/memory/test_memory.py +++ b/tests/stores/memory/test_memory.py @@ -2,10 +2,10 @@ from typing_extensions import override from key_value.aio.stores.memory.store import MemoryStore -from tests.stores.base import BaseStoreTests +from tests.stores.base import BaseStoreTests, PutIfAbsentStoreTestMixin -class TestMemoryStore(BaseStoreTests): +class TestMemoryStore(PutIfAbsentStoreTestMixin, BaseStoreTests): @override @pytest.fixture async def store(self) -> MemoryStore: diff --git a/tests/stores/redis/test_redis.py b/tests/stores/redis/test_redis.py index 9ef7ea68..8f91bca0 100644 --- a/tests/stores/redis/test_redis.py +++ b/tests/stores/redis/test_redis.py @@ -12,7 +12,11 @@ from key_value.aio.stores.base import BaseStore from key_value.aio.stores.redis import RedisStore from tests.conftest import should_skip_docker_tests -from tests.stores.base import BaseStoreTests, ContextManagerStoreTestMixin +from tests.stores.base import ( + BaseStoreTests, + ContextManagerStoreTestMixin, + PutIfAbsentStoreTestMixin, +) # Redis test configuration REDIS_DB = 15 # Use a separate database for tests @@ -34,7 +38,11 @@ def get_client_from_store(store: RedisStore) -> Redis: @pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running") -class TestRedisStore(ContextManagerStoreTestMixin, BaseStoreTests): +class TestRedisStore( + ContextManagerStoreTestMixin, + PutIfAbsentStoreTestMixin, + BaseStoreTests, +): @pytest.fixture(autouse=True, scope="module", params=REDIS_VERSIONS_TO_TEST) def redis_container(self, request: pytest.FixtureRequest): version = request.param diff --git a/tests/stores/redis/test_redis_put_if_absent.py b/tests/stores/redis/test_redis_put_if_absent.py new file mode 100644 index 00000000..57feb2b0 --- /dev/null +++ b/tests/stores/redis/test_redis_put_if_absent.py @@ -0,0 +1,27 @@ +from typing import cast +from unittest.mock import AsyncMock + +from redis.asyncio import Redis + +from key_value.aio.stores.redis.store import _redis_set_if_absent + + +async def test_redis_set_if_absent_uses_atomic_set() -> None: + client = AsyncMock(spec=Redis) + set_mock = AsyncMock(return_value=True) + client.configure_mock(set=set_mock) + + stored = await _redis_set_if_absent( + cast("Redis", client), + "collection::key", + '{"value": 1}', + 300, + ) + + assert stored is True + set_mock.assert_awaited_once_with( + name="collection::key", + value='{"value": 1}', + nx=True, + ex=300, + ) From 1bd742e2ba758fd301cd97fc287cdd715df399ea Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:06:25 -0400 Subject: [PATCH 2/3] Fix conditional write edge cases --- src/key_value/aio/stores/redis/store.py | 9 +++++---- tests/stores/base.py | 1 - tests/stores/redis/test_redis_put_if_absent.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/key_value/aio/stores/redis/store.py b/src/key_value/aio/stores/redis/store.py index 002db135..d57f26ce 100644 --- a/src/key_value/aio/stores/redis/store.py +++ b/src/key_value/aio/stores/redis/store.py @@ -1,3 +1,4 @@ +import math from collections.abc import Sequence from datetime import datetime from typing import Any, Literal, overload @@ -173,10 +174,11 @@ async def _redis_set_if_absent( client: Redis, name: str, value: str, - ttl: int | None, + ttl: float | None, ) -> bool: """Set a value atomically when its key does not exist.""" - result = await client.set(name=name, value=value, nx=True, ex=ttl) + ttl_ms = max(math.ceil(ttl * 1000), 1) if ttl is not None else None + result = await client.set(name=name, value=value, nx=True, px=ttl_ms) return bool(result) @@ -390,12 +392,11 @@ async def _put_managed_entry_if_absent( key=key, collection=collection, ) - ttl = max(int(managed_entry.ttl), 1) if managed_entry.ttl is not None else None return await _redis_set_if_absent( self._client, combo_key, json_value, - ttl, + managed_entry.ttl, ) @override diff --git a/tests/stores/base.py b/tests/stores/base.py index f0ea668f..9f0c3705 100644 --- a/tests/stores/base.py +++ b/tests/stores/base.py @@ -352,7 +352,6 @@ async def test_put_if_absent_accepts_expired_key(self, store: BaseStore): value={"winner": 2}, ) - @pytest.mark.skipif(condition=not async_running_in_event_loop(), reason="Cannot run concurrent operations outside of event loop") async def test_put_if_absent_is_atomic(self, store: BaseStore): assert isinstance(store, AsyncPutIfAbsentProtocol) diff --git a/tests/stores/redis/test_redis_put_if_absent.py b/tests/stores/redis/test_redis_put_if_absent.py index 57feb2b0..4282a9e1 100644 --- a/tests/stores/redis/test_redis_put_if_absent.py +++ b/tests/stores/redis/test_redis_put_if_absent.py @@ -15,7 +15,7 @@ async def test_redis_set_if_absent_uses_atomic_set() -> None: cast("Redis", client), "collection::key", '{"value": 1}', - 300, + 0.5, ) assert stored is True @@ -23,5 +23,5 @@ async def test_redis_set_if_absent_uses_atomic_set() -> None: name="collection::key", value='{"value": 1}', nx=True, - ex=300, + px=500, ) From e9e3dae17c11aed906b3ed4bfe8b8e28935bd9d2 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 5 Sep 2026 10:42:12 -0400 Subject: [PATCH 3/3] Preserve Redis TTL precision across writes --- src/key_value/aio/stores/redis/store.py | 28 +++++++++---------------- tests/stores/redis/test_redis.py | 27 ++++++++++++++++++++++-- 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/src/key_value/aio/stores/redis/store.py b/src/key_value/aio/stores/redis/store.py index d57f26ce..9505b314 100644 --- a/src/key_value/aio/stores/redis/store.py +++ b/src/key_value/aio/stores/redis/store.py @@ -160,14 +160,14 @@ async def _redis_mget(client: Redis, keys: list[str]) -> list[Any]: return await client.mget(keys=keys) -async def _redis_set(client: Redis, name: str, value: str) -> None: - """Set a value in Redis without TTL.""" - _ = await client.set(name=name, value=value) +def _ttl_to_milliseconds(ttl: float | None) -> int | None: + """Preserve TTL precision while keeping Redis expiry positive.""" + return max(math.ceil(ttl * 1000), 1) if ttl is not None else None -async def _redis_setex(client: Redis, name: str, time: int, value: str) -> None: - """Set a value in Redis with TTL.""" - _ = await client.setex(name=name, time=time, value=value) +async def _redis_set(client: Redis, name: str, value: str, ttl: float | None = None) -> None: + """Set a value in Redis with an optional TTL.""" + _ = await client.set(name=name, value=value, px=_ttl_to_milliseconds(ttl)) async def _redis_set_if_absent( @@ -177,8 +177,7 @@ async def _redis_set_if_absent( ttl: float | None, ) -> bool: """Set a value atomically when its key does not exist.""" - ttl_ms = max(math.ceil(ttl * 1000), 1) if ttl is not None else None - result = await client.set(name=name, value=value, nx=True, px=ttl_ms) + result = await client.set(name=name, value=value, nx=True, px=_ttl_to_milliseconds(ttl)) return bool(result) @@ -370,13 +369,7 @@ async def _put_managed_entry( json_value: str = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection) - if managed_entry.ttl is not None: - # Redis does not support <= 0 TTLs - ttl = max(int(managed_entry.ttl), 1) - - await _redis_setex(self._client, combo_key, ttl, json_value) - else: - await _redis_set(self._client, combo_key, json_value) + await _redis_set(self._client, combo_key, json_value, managed_entry.ttl) @override async def _put_managed_entry_if_absent( @@ -424,8 +417,7 @@ async def _put_managed_entries( return - # Convert TTL to integer seconds for Redis - ttl_seconds: int = max(int(ttl), 1) + ttl_ms = _ttl_to_milliseconds(ttl) # Use pipeline for bulk operations pipeline = self._client.pipeline() @@ -434,7 +426,7 @@ async def _put_managed_entries( combo_key: str = compound_key(collection=collection, key=key) json_value = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection) - pipeline.setex(name=combo_key, time=ttl_seconds, value=json_value) + pipeline.set(name=combo_key, value=json_value, px=ttl_ms) await _redis_pipeline_execute(pipeline) diff --git a/tests/stores/redis/test_redis.py b/tests/stores/redis/test_redis.py index 8f91bca0..58e671ed 100644 --- a/tests/stores/redis/test_redis.py +++ b/tests/stores/redis/test_redis.py @@ -1,8 +1,8 @@ import json -from typing import Any +from typing import Any, Literal import pytest -from dirty_equals import IsDatetime +from dirty_equals import IsDatetime, IsInt from inline_snapshot import snapshot from redis.asyncio.client import Redis from testcontainers.core.container import DockerContainer @@ -89,6 +89,29 @@ async def store(self, setup_redis: None, redis_host: str, redis_port: int) -> Re def redis_client(self, store: RedisStore) -> Redis: return get_client_from_store(store=store) + @pytest.mark.parametrize("ttl", [None, 0.5, 1.9, 2.0]) + @pytest.mark.parametrize("write_method", ["put", "put_many", "put_if_absent"]) + async def test_writes_preserve_ttl_precision( + self, + store: RedisStore, + redis_client: Redis, + ttl: float | None, + write_method: Literal["put", "put_many", "put_if_absent"], + ): + if write_method == "put_many": + await store.put_many(collection="test", keys=["ttl_precision"], values=[{"value": 1}], ttl=ttl) + elif write_method == "put": + await store.put(collection="test", key="ttl_precision", value={"value": 1}, ttl=ttl) + else: + assert await store.put_if_absent(collection="test", key="ttl_precision", value={"value": 1}, ttl=ttl) + + remaining_ms = await redis_client.pttl("test::ttl_precision") + + if ttl is None: + assert remaining_ms == -1 + else: + assert remaining_ms == IsInt(approx=int(ttl * 1000), delta=100) + async def test_redis_url_connection(self, setup_redis: None, redis_host: str, redis_port: int): """Test Redis store creation with URL.""" redis_url = f"redis://{redis_host}:{redis_port}/{REDIS_DB}"