diff --git a/README.md b/README.md index 35291871..6f998bbd 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,26 @@ 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, + ) +``` + +Wrappers always satisfy this `isinstance` check, since they forward the call +if the underlying store supports it. If it doesn't, `put_if_absent()` raises +`NotImplementedError` instead of the check being `False`. + ### Stores The library provides multiple store implementations organized into three diff --git a/docs/api/protocols.md b/docs/api/protocols.md index 50e80ca6..d8b0b218 100644 --- a/docs/api/protocols.md +++ b/docs/api/protocols.md @@ -11,3 +11,31 @@ 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, + ) +``` + +Wrappers always satisfy this `isinstance` check, since they forward the call +if the underlying store supports it. If it doesn't, `put_if_absent()` raises +`NotImplementedError` instead of the check being `False`. + +::: 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..d6ad225b 100644 --- a/src/key_value/aio/protocols/__init__.py +++ b/src/key_value/aio/protocols/__init__.py @@ -1 +1,7 @@ +from key_value.aio.protocols.key_value import AsyncCullProtocol as AsyncCullProtocol +from key_value.aio.protocols.key_value import AsyncDestroyCollectionProtocol as AsyncDestroyCollectionProtocol +from key_value.aio.protocols.key_value import AsyncDestroyStoreProtocol as AsyncDestroyStoreProtocol +from key_value.aio.protocols.key_value import AsyncEnumerateCollectionsProtocol as AsyncEnumerateCollectionsProtocol +from key_value.aio.protocols.key_value import AsyncEnumerateKeysProtocol as AsyncEnumerateKeysProtocol 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..26583eef 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]]] @@ -315,19 +316,26 @@ async def _put_managed_entries( managed_entry=managed_entry, ) + async def _prepare_write( + self, *, collection: str | None, value: Mapping[str, Any], ttl: SupportsFloat | None + ) -> tuple[str, ManagedEntry]: + """Resolve the collection, ensure it's set up, and build a ManagedEntry for a new write.""" + resolved_collection = collection or self.default_collection + await self.setup_collection(collection=resolved_collection) + + created_at, _, expires_at = prepare_entry_timestamps(ttl=ttl) + managed_entry = ManagedEntry(value=value, created_at=created_at, expires_at=expires_at) + + return resolved_collection, managed_entry + @bear_enforce @override async def put(self, key: str, value: Mapping[str, Any], *, collection: str | None = None, ttl: SupportsFloat | None = None) -> None: """Store a key-value pair in the specified collection with optional TTL.""" - collection = collection or self.default_collection - await self.setup_collection(collection=collection) - - created_at, _, expires_at = prepare_entry_timestamps(ttl=ttl) - - managed_entry: ManagedEntry = ManagedEntry(value=value, created_at=created_at, expires_at=expires_at) + resolved_collection, managed_entry = await self._prepare_write(collection=collection, value=value, ttl=ttl) await self._put_managed_entry( - collection=collection, + collection=resolved_collection, key=key, managed_entry=managed_entry, ) @@ -407,6 +415,40 @@ 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.""" + resolved_collection, managed_entry = await self._prepare_write(collection=collection, value=value, ttl=ttl) + + return await self._put_managed_entry_if_absent( + collection=resolved_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..355fcd4e 100644 --- a/src/key_value/aio/stores/memory/store.py +++ b/src/key_value/aio/stores/memory/store.py @@ -1,18 +1,22 @@ import sys +import time from dataclasses import dataclass from datetime import datetime +from threading import RLock from typing import Any from typing_extensions import override from key_value.aio._utils.managed_entry import ManagedEntry from key_value.aio._utils.serialization import BasicSerializationAdapter +from key_value.aio._utils.time_to_live import now from key_value.aio.stores.base import ( SEED_DATA_TYPE, BaseDestroyCollectionStore, BaseDestroyStore, BaseEnumerateCollectionsStore, BaseEnumerateKeysStore, + BasePutIfAbsentStore, ) try: @@ -32,7 +36,13 @@ class MemoryCacheEntry: def _memory_cache_ttu(_key: Any, value: MemoryCacheEntry, _now: float) -> float: - """Calculate time-to-use for cache entries based on their expiration time.""" + """Calculate time-to-use for cache entries based on their expiration time. + + Returns a wall-clock epoch timestamp, so the cache must be constructed with a + matching `timer=time.time` -- TLRUCache's default `timer` is `time.monotonic`, + whose epoch is arbitrary (e.g. time since boot), which would compare a wall-clock + expires_at against a monotonic "now" and never actually expire anything. + """ if value.expires_at is None: return float(sys.maxsize) @@ -60,34 +70,52 @@ def __init__(self, max_entries: int | None = None): self._cache = TLRUCache[str, MemoryCacheEntry]( maxsize=max_entries if max_entries is not None else sys.maxsize, ttu=_memory_cache_ttu, + timer=time.time, getsizeof=_memory_cache_getsizeof, ) 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_entry: MemoryCacheEntry | None = self._cache.get(key) + if existing_entry is not None and (existing_entry.expires_at is None or existing_entry.expires_at > now()): + return False + + json_str: str = self._serialization_adapter.dump_json(entry=value) + self._cache[key] = MemoryCacheEntry(json_str=json_str, expires_at=value.expires_at) + 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 +201,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..0fad0a4e 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 @@ -9,8 +10,14 @@ from key_value.aio._utils.compound import compound_key, get_keys_from_compound_keys 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.errors import DeserializationError, InvalidTTLError +from key_value.aio.stores.base import ( + BaseContextManagerStore, + BaseDestroyStore, + BaseEnumerateKeysStore, + BasePutIfAbsentStore, + BaseStore, +) try: from redis.asyncio import Redis @@ -153,14 +160,29 @@ 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.""" + if ttl is None: + return None + if not math.isfinite(ttl): + raise InvalidTTLError(ttl=ttl) + return max(math.ceil(ttl * 1000), 1) -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( + client: Redis, + name: str, + value: str, + ttl: float | None, +) -> bool: + """Set a value atomically when its key does not exist.""" + result = await client.set(name=name, value=value, nx=True, px=_ttl_to_milliseconds(ttl)) + return bool(result) async def _redis_pipeline_execute(pipeline: Any) -> None: @@ -183,7 +205,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 @@ -339,6 +361,11 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) -> return entries + def _combo_key_and_json_value(self, *, key: str, collection: str, managed_entry: ManagedEntry) -> tuple[str, str]: + combo_key: str = compound_key(collection=collection, key=key) + json_value: str = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection) + return combo_key, json_value + @override async def _put_managed_entry( self, @@ -347,17 +374,20 @@ async def _put_managed_entry( collection: str, managed_entry: ManagedEntry, ) -> None: - combo_key: str = compound_key(collection=collection, key=key) - - json_value: str = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection) + combo_key, json_value = self._combo_key_and_json_value(key=key, collection=collection, managed_entry=managed_entry) - if managed_entry.ttl is not None: - # Redis does not support <= 0 TTLs - ttl = max(int(managed_entry.ttl), 1) + await _redis_set(self._client, combo_key, json_value, managed_entry.ttl) - await _redis_setex(self._client, combo_key, ttl, json_value) - 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, json_value = self._combo_key_and_json_value(key=key, collection=collection, managed_entry=managed_entry) + return await _redis_set_if_absent(self._client, combo_key, json_value, managed_entry.ttl) @override async def _put_managed_entries( @@ -384,8 +414,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() @@ -394,7 +423,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/src/key_value/aio/wrappers/base.py b/src/key_value/aio/wrappers/base.py index 29e5abcf..83262e0d 100644 --- a/src/key_value/aio/wrappers/base.py +++ b/src/key_value/aio/wrappers/base.py @@ -4,10 +4,10 @@ from typing_extensions import override from key_value.aio._utils.beartype import bear_enforce -from key_value.aio.protocols.key_value import AsyncKeyValue +from key_value.aio.protocols.key_value import AsyncKeyValue, AsyncPutIfAbsentProtocol -class BaseWrapper(AsyncKeyValue): +class BaseWrapper(AsyncKeyValue, AsyncPutIfAbsentProtocol): """A base wrapper for KVStore implementations that passes through to the underlying store. This class implements the passthrough pattern where all operations are delegated to the wrapped @@ -75,3 +75,20 @@ async def delete(self, key: str, *, collection: str | None = None) -> bool: @override async def delete_many(self, keys: Sequence[str], *, collection: str | None = None) -> int: return await self.key_value.delete_many(keys=keys, collection=collection) + + @bear_enforce + @override + async def put_if_absent( + self, key: str, value: Mapping[str, Any], *, collection: str | None = None, ttl: SupportsFloat | None = None + ) -> bool: + """Forward to the wrapped store, if it supports atomic conditional writes. + + `isinstance(wrapper, AsyncPutIfAbsentProtocol)` can't reflect whether the *wrapped* store + supports this -- Python's runtime-checkable protocols only check for the method's presence + on the wrapper's own class, not the object it delegates to -- so this raises clearly instead + of letting an isinstance-guarded caller believe the write was atomic when it wasn't attempted. + """ + if not isinstance(self.key_value, AsyncPutIfAbsentProtocol): + msg = f"{type(self.key_value).__name__} does not support put_if_absent" + raise NotImplementedError(msg) + return await self.key_value.put_if_absent(key=key, value=value, collection=collection, ttl=ttl) 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..9f0c3705 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,102 @@ 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}, + ) + + 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..f9097218 100644 --- a/tests/stores/memory/test_memory.py +++ b/tests/stores/memory/test_memory.py @@ -1,11 +1,13 @@ +import asyncio + import pytest 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: @@ -14,3 +16,21 @@ async def store(self) -> MemoryStore: async def test_seed(self): store = MemoryStore(max_entries_per_collection=500, seed={"test_collection": {"test_key": {"obj_key": "obj_value"}}}) assert await store.get(key="test_key", collection="test_collection") == {"obj_key": "obj_value"} + + async def test_expired_entries_are_evicted_from_the_underlying_cache(self, store: MemoryStore): + """get() filtering expired entries isn't enough on its own -- the raw cache entry has to + actually go away too, or a write-once/never-reread key leaks forever. TLRUCache's default + timer (monotonic) doesn't match the wall-clock `expires_at` this store hands it, so its + own eviction silently never fired even though get() correctly returned None. + """ + await store.put(collection="test", key="short_lived", value={"data": "value"}, ttl=0.2) + + for _ in range(20): + if await store.get(collection="test", key="short_lived") is None: + break + await asyncio.sleep(0.05) + else: + pytest.fail("entry never expired") + + raw_cache = store._cache["test"]._cache + assert len(raw_cache) == 0 diff --git a/tests/stores/redis/test_redis.py b/tests/stores/redis/test_redis.py index 9ef7ea68..37c0701b 100644 --- a/tests/stores/redis/test_redis.py +++ b/tests/stores/redis/test_redis.py @@ -1,18 +1,24 @@ 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 from typing_extensions import override from key_value.aio._utils.wait import async_wait_for_true +from key_value.aio.errors import InvalidTTLError from key_value.aio.stores.base import BaseStore from key_value.aio.stores.redis import RedisStore +from key_value.aio.stores.redis.store import _ttl_to_milliseconds 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 @@ -33,8 +39,24 @@ def get_client_from_store(store: RedisStore) -> Redis: return store._client +class TestTtlToMilliseconds: + """Unit tests for _ttl_to_milliseconds; no Docker/Redis needed.""" + + def test_none_ttl_means_no_expiry(self): + assert _ttl_to_milliseconds(None) is None + + @pytest.mark.parametrize("ttl", [float("inf"), float("nan"), float("-inf")]) + def test_non_finite_ttl_raises_invalid_ttl_error(self, ttl: float): + with pytest.raises(InvalidTTLError): + _ttl_to_milliseconds(ttl) + + @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 @@ -81,6 +103,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.5]) + @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}" 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..4282a9e1 --- /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}', + 0.5, + ) + + assert stored is True + set_mock.assert_awaited_once_with( + name="collection::key", + value='{"value": 1}', + nx=True, + px=500, + ) diff --git a/tests/stores/wrappers/test_base.py b/tests/stores/wrappers/test_base.py new file mode 100644 index 00000000..d67b28e3 --- /dev/null +++ b/tests/stores/wrappers/test_base.py @@ -0,0 +1,33 @@ +"""Tests for BaseWrapper's optional-capability forwarding.""" + +from pathlib import Path + +import pytest + +from key_value.aio.protocols import AsyncPutIfAbsentProtocol +from key_value.aio.stores.disk import DiskStore +from key_value.aio.stores.memory import MemoryStore +from key_value.aio.wrappers.logging import LoggingWrapper + + +class TestBaseWrapperPutIfAbsent: + """BaseWrapper should forward put_if_absent when the wrapped store supports it.""" + + async def test_wrapper_forwards_put_if_absent_when_supported(self): + wrapped = LoggingWrapper(key_value=MemoryStore()) + + assert isinstance(wrapped, AsyncPutIfAbsentProtocol) + assert await wrapped.put_if_absent(collection="test", key="k", value={"data": "first"}) is True + assert await wrapped.put_if_absent(collection="test", key="k", value={"data": "second"}) is False + assert await wrapped.get(collection="test", key="k") == {"data": "first"} + + async def test_wrapper_raises_when_wrapped_store_does_not_support_it(self, tmp_path: Path): + wrapped = LoggingWrapper(key_value=DiskStore(directory=tmp_path)) + + # isinstance can't reflect the wrapped store's actual capability (Python's runtime-checkable + # protocols only check the wrapper's own class), so the wrapper always satisfies the check -- + # the failure has to surface at call time instead. + assert isinstance(wrapped, AsyncPutIfAbsentProtocol) + + with pytest.raises(NotImplementedError): + await wrapped.put_if_absent(collection="test", key="k", value={"data": "value"})