From fc54aaa6430c065c2a18b5e298da89b3894b5d26 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Sep 2025 01:25:18 +0000 Subject: [PATCH 1/3] Initial plan From 50b19f8759eb04337909db8ba139c9b3ce694de0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Sep 2025 01:37:19 +0000 Subject: [PATCH 2/3] Refactor wrappers to use KVStoreProtocol and move to new location Co-authored-by: strawgate <6384545+strawgate@users.noreply.github.com> --- README.md | 2 +- src/kv_store_adapter/wrappers/__init__.py | 15 ++ src/kv_store_adapter/wrappers/clamp_ttl.py | 42 ++++++ .../wrappers/passthrough_cache.py | 43 ++++++ .../wrappers/prefix_collection.py | 36 +++++ src/kv_store_adapter/wrappers/prefix_key.py | 36 +++++ .../wrappers/single_collection.py | 31 ++++ src/kv_store_adapter/wrappers/statistics.py | 133 ++++++++++++++++++ tests/stores/wrappers/conftest.py | 80 +++++++++++ tests/stores/wrappers/test_clamp_ttl.py | 16 ++- tests/stores/wrappers/test_passthrough.py | 6 +- .../stores/wrappers/test_prefix_collection.py | 6 +- tests/stores/wrappers/test_prefix_key.py | 6 +- .../stores/wrappers/test_single_collection.py | 23 +-- tests/stores/wrappers/test_statistics.py | 81 +++++++++++ 15 files changed, 520 insertions(+), 36 deletions(-) create mode 100644 src/kv_store_adapter/wrappers/__init__.py create mode 100644 src/kv_store_adapter/wrappers/clamp_ttl.py create mode 100644 src/kv_store_adapter/wrappers/passthrough_cache.py create mode 100644 src/kv_store_adapter/wrappers/prefix_collection.py create mode 100644 src/kv_store_adapter/wrappers/prefix_key.py create mode 100644 src/kv_store_adapter/wrappers/single_collection.py create mode 100644 src/kv_store_adapter/wrappers/statistics.py create mode 100644 tests/stores/wrappers/conftest.py create mode 100644 tests/stores/wrappers/test_statistics.py diff --git a/README.md b/README.md index 78f7f767..cd69402b 100644 --- a/README.md +++ b/README.md @@ -161,7 +161,7 @@ Imagine you have a service where you want to cache 3 pydantic models in a single import asyncio from kv_store_adapter.adapters.pydantic import PydanticAdapter -from kv_store_adapter.stores.wrappers.single_collection import SingleCollectionWrapper +from kv_store_adapter.wrappers.single_collection import SingleCollectionWrapper from kv_store_adapter.stores.memory import MemoryStore from pydantic import BaseModel diff --git a/src/kv_store_adapter/wrappers/__init__.py b/src/kv_store_adapter/wrappers/__init__.py new file mode 100644 index 00000000..de3f0dc6 --- /dev/null +++ b/src/kv_store_adapter/wrappers/__init__.py @@ -0,0 +1,15 @@ +from .clamp_ttl import TTLClampWrapper +from .passthrough_cache import PassthroughCacheWrapper +from .prefix_collection import PrefixCollectionWrapper +from .prefix_key import PrefixKeyWrapper +from .single_collection import SingleCollectionWrapper +from .statistics import StatisticsWrapper + +__all__ = [ + "TTLClampWrapper", + "PassthroughCacheWrapper", + "PrefixCollectionWrapper", + "PrefixKeyWrapper", + "SingleCollectionWrapper", + "StatisticsWrapper" +] \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/clamp_ttl.py b/src/kv_store_adapter/wrappers/clamp_ttl.py new file mode 100644 index 00000000..e6966483 --- /dev/null +++ b/src/kv_store_adapter/wrappers/clamp_ttl.py @@ -0,0 +1,42 @@ +from typing import Any + +from kv_store_adapter.types import KVStoreProtocol + + +class TTLClampWrapper: + """Wrapper that enforces a maximum TTL for puts into the store.""" + + def __init__(self, store: KVStoreProtocol, min_ttl: float, max_ttl: float, missing_ttl: float | None = None) -> None: + """Initialize the TTL clamp wrapper. + + Args: + store: The store to wrap. + min_ttl: The minimum TTL for puts into the store. + max_ttl: The maximum TTL for puts into the store. + missing_ttl: The TTL to use for entries that do not have a TTL. Defaults to None. + """ + self.store: KVStoreProtocol = store + self.min_ttl: float = min_ttl + self.max_ttl: float = max_ttl + self.missing_ttl: float | None = missing_ttl + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + return await self.store.get(collection=collection, key=key) + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + if ttl is None and self.missing_ttl: + ttl = self.missing_ttl + + if ttl and ttl < self.min_ttl: + ttl = self.min_ttl + + if ttl and ttl > self.max_ttl: + ttl = self.max_ttl + + await self.store.put(collection=collection, key=key, value=value, ttl=ttl) + + async def delete(self, collection: str, key: str) -> bool: + return await self.store.delete(collection=collection, key=key) + + async def exists(self, collection: str, key: str) -> bool: + return await self.store.exists(collection=collection, key=key) \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/passthrough_cache.py b/src/kv_store_adapter/wrappers/passthrough_cache.py new file mode 100644 index 00000000..9c2a813f --- /dev/null +++ b/src/kv_store_adapter/wrappers/passthrough_cache.py @@ -0,0 +1,43 @@ +from typing import Any + +from kv_store_adapter.types import KVStoreProtocol + + +class PassthroughCacheWrapper: + """Wrapper that uses two stores, ideal for combining a local and distributed store.""" + + def __init__(self, primary_store: KVStoreProtocol, cache_store: KVStoreProtocol) -> None: + """Initialize the passthrough cache wrapper. Items are first checked in the primary store and if not found, are + checked in the secondary store. Operations are performed on both stores but are not atomic. + + Note: This wrapper only implements the core KVStoreProtocol operations. Operations like expiry culling + against the primary store will not be reflected in the cache store if the underlying stores support such operations. + + Args: + primary_store: The primary store the data will live in. + cache_store: The write-through (likely ephemeral) cache to use. + """ + self.cache_store: KVStoreProtocol = cache_store + self.primary_store: KVStoreProtocol = primary_store + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + if cache_store_value := await self.cache_store.get(collection=collection, key=key): + return cache_store_value + + if primary_store_value := await self.primary_store.get(collection=collection, key=key): + # Cache the value from primary store (without TTL since we can't get it from protocol) + await self.cache_store.put(collection=collection, key=key, value=primary_store_value) + return primary_store_value + return None + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + _ = await self.cache_store.delete(collection=collection, key=key) + await self.primary_store.put(collection=collection, key=key, value=value, ttl=ttl) + + async def delete(self, collection: str, key: str) -> bool: + deleted = await self.primary_store.delete(collection=collection, key=key) + _ = await self.cache_store.delete(collection=collection, key=key) + return deleted + + async def exists(self, collection: str, key: str) -> bool: + return await self.get(collection=collection, key=key) is not None \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/prefix_collection.py b/src/kv_store_adapter/wrappers/prefix_collection.py new file mode 100644 index 00000000..07ac6f8f --- /dev/null +++ b/src/kv_store_adapter/wrappers/prefix_collection.py @@ -0,0 +1,36 @@ +from typing import Any + +from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_collection +from kv_store_adapter.types import KVStoreProtocol + + +class PrefixCollectionWrapper: + """Wrapper that prefixes all collections with a given prefix.""" + + def __init__(self, store: KVStoreProtocol, prefix: str, separator: str | None = None) -> None: + """Initialize the prefix collection wrapper. + + Args: + store: The store to wrap. + prefix: The prefix to add to all collections. + separator: The separator to use between the prefix and the collection. Defaults to "__". + """ + self.store: KVStoreProtocol = store + self.prefix: str = prefix + self.separator: str = separator or DEFAULT_PREFIX_SEPARATOR + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) + return await self.store.get(collection=prefixed_collection, key=key) + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) + await self.store.put(collection=prefixed_collection, key=key, value=value, ttl=ttl) + + async def delete(self, collection: str, key: str) -> bool: + prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) + return await self.store.delete(collection=prefixed_collection, key=key) + + async def exists(self, collection: str, key: str) -> bool: + prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) + return await self.store.exists(collection=prefixed_collection, key=key) \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/prefix_key.py b/src/kv_store_adapter/wrappers/prefix_key.py new file mode 100644 index 00000000..a5e9002c --- /dev/null +++ b/src/kv_store_adapter/wrappers/prefix_key.py @@ -0,0 +1,36 @@ +from typing import Any + +from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_key +from kv_store_adapter.types import KVStoreProtocol + + +class PrefixKeyWrapper: + """Wrapper that prefixes all keys with a given prefix.""" + + def __init__(self, store: KVStoreProtocol, prefix: str, separator: str | None = None) -> None: + """Initialize the prefix key wrapper. + + Args: + store: The store to wrap. + prefix: The prefix to add to all keys. + separator: The separator to use between the prefix and the key. Defaults to "__". + """ + self.store: KVStoreProtocol = store + self.prefix: str = prefix + self.separator: str = separator or DEFAULT_PREFIX_SEPARATOR + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) + return await self.store.get(collection=collection, key=prefixed_key) + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) + await self.store.put(collection=collection, key=prefixed_key, value=value, ttl=ttl) + + async def delete(self, collection: str, key: str) -> bool: + prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) + return await self.store.delete(collection=collection, key=prefixed_key) + + async def exists(self, collection: str, key: str) -> bool: + prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) + return await self.store.exists(collection=collection, key=prefixed_key) \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/single_collection.py b/src/kv_store_adapter/wrappers/single_collection.py new file mode 100644 index 00000000..5759d3c8 --- /dev/null +++ b/src/kv_store_adapter/wrappers/single_collection.py @@ -0,0 +1,31 @@ +from typing import Any + +from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_key +from kv_store_adapter.types import KVStoreProtocol + + +class SingleCollectionWrapper: + """Wrapper that forces all requests into a single collection, prefixes the keys with the original collection name. + + The single collection wrapper does not support management operations.""" + + def __init__(self, store: KVStoreProtocol, collection: str, prefix_separator: str | None = None) -> None: + self.collection: str = collection + self.prefix_separator: str = prefix_separator or DEFAULT_PREFIX_SEPARATOR + self.store: KVStoreProtocol = store + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) + return await self.store.get(collection=self.collection, key=prefixed_key) + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) + await self.store.put(collection=self.collection, key=prefixed_key, value=value, ttl=ttl) + + async def delete(self, collection: str, key: str) -> bool: + prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) + return await self.store.delete(collection=self.collection, key=prefixed_key) + + async def exists(self, collection: str, key: str) -> bool: + prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) + return await self.store.exists(collection=self.collection, key=prefixed_key) \ No newline at end of file diff --git a/src/kv_store_adapter/wrappers/statistics.py b/src/kv_store_adapter/wrappers/statistics.py new file mode 100644 index 00000000..36375d7e --- /dev/null +++ b/src/kv_store_adapter/wrappers/statistics.py @@ -0,0 +1,133 @@ +from dataclasses import dataclass, field +from typing import Any + +from kv_store_adapter.types import KVStoreProtocol + + +@dataclass +class BaseStatistics: + """Base statistics container with operation counting.""" + + count: int = field(default=0) + """The number of operations.""" + + def increment(self) -> None: + self.count += 1 + + +@dataclass +class BaseHitMissStatistics(BaseStatistics): + """Statistics container with hit/miss tracking for cache-like operations.""" + + hit: int = field(default=0) + """The number of hits.""" + miss: int = field(default=0) + """The number of misses.""" + + def increment_hit(self) -> None: + self.increment() + self.hit += 1 + + def increment_miss(self) -> None: + self.increment() + self.miss += 1 + + +@dataclass +class GetStatistics(BaseHitMissStatistics): + """A class for statistics about GET operations.""" + + +@dataclass +class SetStatistics(BaseStatistics): + """A class for statistics about PUT operations.""" + + +@dataclass +class DeleteStatistics(BaseHitMissStatistics): + """A class for statistics about DELETE operations.""" + + +@dataclass +class ExistsStatistics(BaseHitMissStatistics): + """A class for statistics about EXISTS operations.""" + + +@dataclass +class KVStoreCollectionStatistics(BaseStatistics): + """A class for statistics about a KV Store collection.""" + + get: GetStatistics = field(default_factory=GetStatistics) + """The statistics for the get operation.""" + + set: SetStatistics = field(default_factory=SetStatistics) + """The statistics for the set operation.""" + + delete: DeleteStatistics = field(default_factory=DeleteStatistics) + """The statistics for the delete operation.""" + + exists: ExistsStatistics = field(default_factory=ExistsStatistics) + """The statistics for the exists operation.""" + + +@dataclass +class KVStoreStatistics: + """Statistics container for a KV Store.""" + + collections: dict[str, KVStoreCollectionStatistics] = field(default_factory=dict) + + def get_collection(self, collection: str) -> KVStoreCollectionStatistics: + if collection not in self.collections: + self.collections[collection] = KVStoreCollectionStatistics() + return self.collections[collection] + + +class StatisticsWrapper: + """Statistics wrapper around a KV Store that tracks operation statistics.""" + + def __init__(self, store: KVStoreProtocol, track_statistics: bool = True) -> None: + self.store: KVStoreProtocol = store + self._statistics: KVStoreStatistics | None = KVStoreStatistics() if track_statistics else None + + @property + def statistics(self) -> KVStoreStatistics | None: + return self._statistics + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + if value := await self.store.get(collection=collection, key=key): + if self.statistics: + self.statistics.get_collection(collection).get.increment_hit() + return value + + if self.statistics: + self.statistics.get_collection(collection).get.increment_miss() + + return None + + async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: + await self.store.put(collection=collection, key=key, value=value, ttl=ttl) + + if self.statistics: + self.statistics.get_collection(collection).set.increment() + + async def delete(self, collection: str, key: str) -> bool: + if await self.store.delete(collection=collection, key=key): + if self.statistics: + self.statistics.get_collection(collection).delete.increment_hit() + return True + + if self.statistics: + self.statistics.get_collection(collection).delete.increment_miss() + + return False + + async def exists(self, collection: str, key: str) -> bool: + if await self.store.exists(collection=collection, key=key): + if self.statistics: + self.statistics.get_collection(collection).exists.increment_hit() + return True + + if self.statistics: + self.statistics.get_collection(collection).exists.increment_miss() + + return False \ No newline at end of file diff --git a/tests/stores/wrappers/conftest.py b/tests/stores/wrappers/conftest.py new file mode 100644 index 00000000..bbd89be2 --- /dev/null +++ b/tests/stores/wrappers/conftest.py @@ -0,0 +1,80 @@ +"""Test configuration for protocol-only wrappers.""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from kv_store_adapter.types import KVStoreProtocol + + +class BaseProtocolTests(ABC): + """Base test class for KVStoreProtocol implementations.""" + + @pytest.fixture + @abstractmethod + async def store(self) -> "KVStoreProtocol": ... + + async def test_empty_get(self, store: "KVStoreProtocol"): + """Tests that the get method returns None from an empty store.""" + assert await store.get(collection="test", key="test") is None + + async def test_empty_put(self, store: "KVStoreProtocol"): + """Tests that the put method does not raise an exception when called on a new store.""" + await store.put(collection="test", key="test", value={"test": "test"}) + + async def test_empty_exists(self, store: "KVStoreProtocol"): + """Tests that the exists method returns False from an empty store.""" + assert await store.exists(collection="test", key="test") is False + + async def test_empty_delete(self, store: "KVStoreProtocol"): + """Tests that the delete method returns False from an empty store.""" + assert await store.delete(collection="test", key="test") is False + + async def test_get_put_get_put_delete_get(self, store: "KVStoreProtocol"): + """Tests that the get, put, get, put, delete, and get methods work together.""" + await store.put(collection="test", key="test", value={"test": "test"}) + assert await store.get(collection="test", key="test") == {"test": "test"} + + await store.put(collection="test", key="test", value={"test": "test_2"}) + + assert await store.get(collection="test", key="test") == {"test": "test_2"} + assert await store.delete(collection="test", key="test") + assert await store.get(collection="test", key="test") is None + + async def test_exists_functionality(self, store: "KVStoreProtocol"): + """Tests that the exists method works correctly.""" + assert await store.exists(collection="test", key="test") is False + + await store.put(collection="test", key="test", value={"test": "test"}) + assert await store.exists(collection="test", key="test") is True + + await store.delete(collection="test", key="test") + assert await store.exists(collection="test", key="test") is False + + async def test_multiple_collections(self, store: "KVStoreProtocol"): + """Tests that multiple collections work independently.""" + await store.put(collection="test_one", key="test", value={"test": "test_one"}) + await store.put(collection="test_two", key="test", value={"test": "test_two"}) + + assert await store.get(collection="test_one", key="test") == {"test": "test_one"} + assert await store.get(collection="test_two", key="test") == {"test": "test_two"} + + assert await store.exists(collection="test_one", key="test") is True + assert await store.exists(collection="test_two", key="test") is True + + async def test_multiple_keys(self, store: "KVStoreProtocol"): + """Tests that multiple keys work independently in the same collection.""" + await store.put(collection="test", key="key_one", value={"test": "value_one"}) + await store.put(collection="test", key="key_two", value={"test": "value_two"}) + + assert await store.get(collection="test", key="key_one") == {"test": "value_one"} + assert await store.get(collection="test", key="key_two") == {"test": "value_two"} + + assert await store.exists(collection="test", key="key_one") is True + assert await store.exists(collection="test", key="key_two") is True + + assert await store.delete(collection="test", key="key_one") is True + assert await store.exists(collection="test", key="key_one") is False + assert await store.exists(collection="test", key="key_two") is True \ No newline at end of file diff --git a/tests/stores/wrappers/test_clamp_ttl.py b/tests/stores/wrappers/test_clamp_ttl.py index 65c4c8bc..1a9f1164 100644 --- a/tests/stores/wrappers/test_clamp_ttl.py +++ b/tests/stores/wrappers/test_clamp_ttl.py @@ -5,14 +5,15 @@ from typing_extensions import override from kv_store_adapter.stores.memory.store import MemoryStore -from kv_store_adapter.stores.wrappers.clamp_ttl import TTLClampWrapper -from tests.stores.conftest import BaseStoreTests, now, now_plus +from kv_store_adapter.wrappers.clamp_ttl import TTLClampWrapper +from tests.stores.conftest import now, now_plus +from tests.stores.wrappers.conftest import BaseProtocolTests if TYPE_CHECKING: from kv_store_adapter.types import TTLInfo -class TestTTLClampWrapper(BaseStoreTests): +class TestTTLClampWrapper(BaseProtocolTests): @pytest.fixture async def memory_store(self) -> MemoryStore: return MemoryStore() @@ -28,7 +29,8 @@ async def test_put_below_min_ttl(self, memory_store: MemoryStore): await ttl_clamp_store.put(collection="test", key="test", value={"test": "test"}, ttl=5) assert await ttl_clamp_store.get(collection="test", key="test") is not None - ttl_info: TTLInfo | None = await ttl_clamp_store.ttl(collection="test", key="test") + # Check the TTL through the underlying store since wrapper doesn't expose ttl() method + ttl_info: TTLInfo | None = await memory_store.ttl(collection="test", key="test") assert ttl_info is not None assert ttl_info.ttl == 50 @@ -44,7 +46,8 @@ async def test_put_above_max_ttl(self, memory_store: MemoryStore): await ttl_clamp_store.put(collection="test", key="test", value={"test": "test"}, ttl=1000) assert await ttl_clamp_store.get(collection="test", key="test") is not None - ttl_info: TTLInfo | None = await ttl_clamp_store.ttl(collection="test", key="test") + # Check the TTL through the underlying store since wrapper doesn't expose ttl() method + ttl_info: TTLInfo | None = await memory_store.ttl(collection="test", key="test") assert ttl_info is not None assert ttl_info.ttl == 100 @@ -60,7 +63,8 @@ async def test_put_missing_ttl(self, memory_store: MemoryStore): await ttl_clamp_store.put(collection="test", key="test", value={"test": "test"}, ttl=None) assert await ttl_clamp_store.get(collection="test", key="test") is not None - ttl_info: TTLInfo | None = await ttl_clamp_store.ttl(collection="test", key="test") + # Check the TTL through the underlying store since wrapper doesn't expose ttl() method + ttl_info: TTLInfo | None = await memory_store.ttl(collection="test", key="test") assert ttl_info is not None assert ttl_info.ttl == 50 diff --git a/tests/stores/wrappers/test_passthrough.py b/tests/stores/wrappers/test_passthrough.py index 051fa859..33260382 100644 --- a/tests/stores/wrappers/test_passthrough.py +++ b/tests/stores/wrappers/test_passthrough.py @@ -6,13 +6,13 @@ from kv_store_adapter.stores.disk.store import DiskStore from kv_store_adapter.stores.memory.store import MemoryStore -from kv_store_adapter.stores.wrappers.passthrough_cache import PassthroughCacheWrapper -from tests.stores.conftest import BaseStoreTests +from kv_store_adapter.wrappers.passthrough_cache import PassthroughCacheWrapper +from tests.stores.wrappers.conftest import BaseProtocolTests DISK_STORE_SIZE_LIMIT = 1 * 1024 * 1024 # 1MB -class TestPrefixCollectionWrapper(BaseStoreTests): +class TestPassthroughCacheWrapper(BaseProtocolTests): @pytest.fixture async def primary_store(self) -> AsyncGenerator[DiskStore, None]: with tempfile.TemporaryDirectory() as temp_dir: diff --git a/tests/stores/wrappers/test_prefix_collection.py b/tests/stores/wrappers/test_prefix_collection.py index 15fc19f3..ad4e6e4b 100644 --- a/tests/stores/wrappers/test_prefix_collection.py +++ b/tests/stores/wrappers/test_prefix_collection.py @@ -2,11 +2,11 @@ from typing_extensions import override from kv_store_adapter.stores.memory.store import MemoryStore -from kv_store_adapter.stores.wrappers.prefix_collection import PrefixCollectionWrapper -from tests.stores.conftest import BaseStoreTests +from kv_store_adapter.wrappers.prefix_collection import PrefixCollectionWrapper +from tests.stores.wrappers.conftest import BaseProtocolTests -class TestPrefixCollectionWrapper(BaseStoreTests): +class TestPrefixCollectionWrapper(BaseProtocolTests): @override @pytest.fixture async def store(self) -> PrefixCollectionWrapper: diff --git a/tests/stores/wrappers/test_prefix_key.py b/tests/stores/wrappers/test_prefix_key.py index 3868e420..918d768c 100644 --- a/tests/stores/wrappers/test_prefix_key.py +++ b/tests/stores/wrappers/test_prefix_key.py @@ -2,11 +2,11 @@ from typing_extensions import override from kv_store_adapter.stores.memory.store import MemoryStore -from kv_store_adapter.stores.wrappers.prefix_key import PrefixKeyWrapper -from tests.stores.conftest import BaseStoreTests +from kv_store_adapter.wrappers.prefix_key import PrefixKeyWrapper +from tests.stores.wrappers.conftest import BaseProtocolTests -class TestPrefixKeyWrapper(BaseStoreTests): +class TestPrefixKeyWrapper(BaseProtocolTests): @override @pytest.fixture async def store(self) -> PrefixKeyWrapper: diff --git a/tests/stores/wrappers/test_single_collection.py b/tests/stores/wrappers/test_single_collection.py index 963a4f29..7880ce71 100644 --- a/tests/stores/wrappers/test_single_collection.py +++ b/tests/stores/wrappers/test_single_collection.py @@ -1,31 +1,14 @@ import pytest from typing_extensions import override -from kv_store_adapter.stores.base.unmanaged import BaseKVStore from kv_store_adapter.stores.memory.store import MemoryStore -from kv_store_adapter.stores.wrappers.single_collection import SingleCollectionWrapper -from tests.stores.conftest import BaseStoreTests +from kv_store_adapter.wrappers.single_collection import SingleCollectionWrapper +from tests.stores.wrappers.conftest import BaseProtocolTests -class TestSingleCollectionWrapper(BaseStoreTests): +class TestSingleCollectionWrapper(BaseProtocolTests): @override @pytest.fixture async def store(self) -> SingleCollectionWrapper: memory_store: MemoryStore = MemoryStore() return SingleCollectionWrapper(store=memory_store, collection="test") - - @pytest.mark.skip(reason="SingleCollectionWrapper does not support collection operations") - @override - async def test_empty_clear_collection(self, store: BaseKVStore): ... - - @pytest.mark.skip(reason="SingleCollectionWrapper does not support collection operations") - @override - async def test_empty_list_collections(self, store: BaseKVStore): ... - - @pytest.mark.skip(reason="SingleCollectionWrapper does not support collection operations") - @override - async def test_list_collections(self, store: BaseKVStore): ... - - @pytest.mark.skip(reason="SingleCollectionWrapper does not support collection operations") - @override - async def test_set_set_list_collections(self, store: BaseKVStore): ... diff --git a/tests/stores/wrappers/test_statistics.py b/tests/stores/wrappers/test_statistics.py new file mode 100644 index 00000000..7a5cc871 --- /dev/null +++ b/tests/stores/wrappers/test_statistics.py @@ -0,0 +1,81 @@ +import pytest +from typing_extensions import override + +from kv_store_adapter.stores.memory.store import MemoryStore +from kv_store_adapter.wrappers.statistics import StatisticsWrapper +from tests.stores.wrappers.conftest import BaseProtocolTests + + +class TestStatisticsWrapper(BaseProtocolTests): + @override + @pytest.fixture + async def store(self) -> StatisticsWrapper: + memory_store: MemoryStore = MemoryStore() + return StatisticsWrapper(store=memory_store, track_statistics=True) + + async def test_statistics_tracking(self): + memory_store: MemoryStore = MemoryStore() + stats_wrapper = StatisticsWrapper(store=memory_store, track_statistics=True) + + # Initially no statistics + assert stats_wrapper.statistics is not None + assert len(stats_wrapper.statistics.collections) == 0 + + # Test GET miss + result = await stats_wrapper.get(collection="test", key="key1") + assert result is None + + collection_stats = stats_wrapper.statistics.get_collection("test") + assert collection_stats.get.count == 1 + assert collection_stats.get.miss == 1 + assert collection_stats.get.hit == 0 + + # Test PUT + await stats_wrapper.put(collection="test", key="key1", value={"data": "value1"}) + assert collection_stats.set.count == 1 + + # Test GET hit + result = await stats_wrapper.get(collection="test", key="key1") + assert result == {"data": "value1"} + assert collection_stats.get.count == 2 + assert collection_stats.get.miss == 1 + assert collection_stats.get.hit == 1 + + # Test EXISTS hit + exists = await stats_wrapper.exists(collection="test", key="key1") + assert exists is True + assert collection_stats.exists.count == 1 + assert collection_stats.exists.hit == 1 + assert collection_stats.exists.miss == 0 + + # Test DELETE hit + deleted = await stats_wrapper.delete(collection="test", key="key1") + assert deleted is True + assert collection_stats.delete.count == 1 + assert collection_stats.delete.hit == 1 + assert collection_stats.delete.miss == 0 + + # Test EXISTS miss + exists = await stats_wrapper.exists(collection="test", key="key1") + assert exists is False + assert collection_stats.exists.count == 2 + assert collection_stats.exists.hit == 1 + assert collection_stats.exists.miss == 1 + + # Test DELETE miss + deleted = await stats_wrapper.delete(collection="test", key="key1") + assert deleted is False + assert collection_stats.delete.count == 2 + assert collection_stats.delete.hit == 1 + assert collection_stats.delete.miss == 1 + + async def test_statistics_disabled(self): + memory_store: MemoryStore = MemoryStore() + stats_wrapper = StatisticsWrapper(store=memory_store, track_statistics=False) + + assert stats_wrapper.statistics is None + + # Operations should still work + await stats_wrapper.put(collection="test", key="key1", value={"data": "value1"}) + result = await stats_wrapper.get(collection="test", key="key1") + assert result == {"data": "value1"} \ No newline at end of file From 56e8a1fd7d3942ae4958e4b46996b313875324c2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 25 Sep 2025 01:39:47 +0000 Subject: [PATCH 3/3] Remove old wrapper files from stores directory Co-authored-by: strawgate <6384545+strawgate@users.noreply.github.com> --- .../stores/wrappers/__init__.py | 6 - .../stores/wrappers/clamp_ttl.py | 69 ------ .../stores/wrappers/passthrough_cache.py | 81 ------- .../stores/wrappers/prefix_collection.py | 76 ------- .../stores/wrappers/prefix_key.py | 69 ------ .../stores/wrappers/single_collection.py | 68 ------ .../stores/wrappers/statistics.py | 197 ------------------ 7 files changed, 566 deletions(-) delete mode 100644 src/kv_store_adapter/stores/wrappers/__init__.py delete mode 100644 src/kv_store_adapter/stores/wrappers/clamp_ttl.py delete mode 100644 src/kv_store_adapter/stores/wrappers/passthrough_cache.py delete mode 100644 src/kv_store_adapter/stores/wrappers/prefix_collection.py delete mode 100644 src/kv_store_adapter/stores/wrappers/prefix_key.py delete mode 100644 src/kv_store_adapter/stores/wrappers/single_collection.py delete mode 100644 src/kv_store_adapter/stores/wrappers/statistics.py diff --git a/src/kv_store_adapter/stores/wrappers/__init__.py b/src/kv_store_adapter/stores/wrappers/__init__.py deleted file mode 100644 index b7e7f975..00000000 --- a/src/kv_store_adapter/stores/wrappers/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .prefix_collection import PrefixCollectionWrapper -from .prefix_key import PrefixKeyWrapper -from .single_collection import SingleCollectionWrapper -from .statistics import StatisticsWrapper - -__all__ = ["PrefixCollectionWrapper", "PrefixKeyWrapper", "SingleCollectionWrapper", "StatisticsWrapper"] diff --git a/src/kv_store_adapter/stores/wrappers/clamp_ttl.py b/src/kv_store_adapter/stores/wrappers/clamp_ttl.py deleted file mode 100644 index 59ec24b4..00000000 --- a/src/kv_store_adapter/stores/wrappers/clamp_ttl.py +++ /dev/null @@ -1,69 +0,0 @@ -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.types import TTLInfo - - -class TTLClampWrapper(BaseKVStore): - """Wrapper that enforces a maximum TTL for puts into the store.""" - - def __init__(self, store: BaseKVStore, min_ttl: float, max_ttl: float, missing_ttl: float | None = None) -> None: - """Initialize the TTL clamp wrapper. - - Args: - store: The store to wrap. - min_ttl: The minimum TTL for puts into the store. - max_ttl: The maximum TTL for puts into the store. - missing_ttl: The TTL to use for entries that do not have a TTL. Defaults to None. - """ - self.store: BaseKVStore = store - self.min_ttl: float = min_ttl - self.max_ttl: float = max_ttl - self.missing_ttl: float | None = missing_ttl - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - return await self.store.get(collection=collection, key=key) - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - if ttl is None and self.missing_ttl: - ttl = self.missing_ttl - - if ttl and ttl < self.min_ttl: - ttl = self.min_ttl - - if ttl and ttl > self.max_ttl: - ttl = self.max_ttl - - await self.store.put(collection=collection, key=key, value=value, ttl=ttl) - - @override - async def delete(self, collection: str, key: str) -> bool: - return await self.store.delete(collection=collection, key=key) - - @override - async def exists(self, collection: str, key: str) -> bool: - return await self.store.exists(collection=collection, key=key) - - @override - async def keys(self, collection: str) -> list[str]: - return await self.store.keys(collection=collection) - - @override - async def clear_collection(self, collection: str) -> int: - return await self.store.clear_collection(collection=collection) - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - return await self.store.ttl(collection=collection, key=key) - - @override - async def list_collections(self) -> list[str]: - return await self.store.list_collections() - - @override - async def cull(self) -> None: - await self.store.cull() diff --git a/src/kv_store_adapter/stores/wrappers/passthrough_cache.py b/src/kv_store_adapter/stores/wrappers/passthrough_cache.py deleted file mode 100644 index 713bf037..00000000 --- a/src/kv_store_adapter/stores/wrappers/passthrough_cache.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.types import TTLInfo - - -class PassthroughCacheWrapper(BaseKVStore): - """Wrapper that users two stores, ideal for combining a local and distributed store.""" - - def __init__(self, primary_store: BaseKVStore, cache_store: BaseKVStore) -> None: - """Initialize the passthrough cache wrapper. Items are first checked in the primary store and if not found, are - checked in the secondary store. Operations are performed on both stores but are not atomic. - - Operations like expiry culling against the primary store will not be reflected in the cache store. This may - lead to stale data in the cache store. One way to combat this is to use a TTLClampWrapper on the cache store to - enforce a lower TTL on the cache store than the primary store. - - Args: - primary_store: The primary store the data will live in. - cache_store: The write-through (likely ephemeral) cache to use. - """ - self.cache_store: BaseKVStore = cache_store - self.primary_store: BaseKVStore = primary_store - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - if cache_store_value := await self.cache_store.get(collection=collection, key=key): - return cache_store_value - - if primary_store_value := await self.primary_store.get(collection=collection, key=key): - ttl_info: TTLInfo | None = await self.primary_store.ttl(collection=collection, key=key) - - await self.cache_store.put(collection=collection, key=key, value=primary_store_value, ttl=ttl_info.ttl if ttl_info else None) - - return primary_store_value - return None - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - _ = await self.cache_store.delete(collection=collection, key=key) - await self.primary_store.put(collection=collection, key=key, value=value, ttl=ttl) - - @override - async def delete(self, collection: str, key: str) -> bool: - deleted = await self.primary_store.delete(collection=collection, key=key) - _ = await self.cache_store.delete(collection=collection, key=key) - return deleted - - @override - async def exists(self, collection: str, key: str) -> bool: - return await self.get(collection=collection, key=key) is not None - - @override - async def keys(self, collection: str) -> list[str]: - return await self.primary_store.keys(collection=collection) - - @override - async def clear_collection(self, collection: str) -> int: - removed: int = await self.primary_store.clear_collection(collection=collection) - _ = await self.cache_store.clear_collection(collection=collection) - return removed - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - if ttl_info := await self.cache_store.ttl(collection=collection, key=key): - return ttl_info - - return await self.primary_store.ttl(collection=collection, key=key) - - @override - async def list_collections(self) -> list[str]: - collections: list[str] = await self.primary_store.list_collections() - - return collections - - @override - async def cull(self) -> None: - await self.primary_store.cull() - await self.cache_store.cull() diff --git a/src/kv_store_adapter/stores/wrappers/prefix_collection.py b/src/kv_store_adapter/stores/wrappers/prefix_collection.py deleted file mode 100644 index 6488e611..00000000 --- a/src/kv_store_adapter/stores/wrappers/prefix_collection.py +++ /dev/null @@ -1,76 +0,0 @@ -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_collection, unprefix_collection -from kv_store_adapter.types import TTLInfo - - -class PrefixCollectionWrapper(BaseKVStore): - """Wrapper that prefixes all collections with a given prefix.""" - - def __init__(self, store: BaseKVStore, prefix: str, separator: str | None = None) -> None: - """Initialize the prefix collection wrapper. - - Args: - store: The store to wrap. - prefix: The prefix to add to all collections. - separator: The separator to use between the prefix and the collection. Defaults to "__". - """ - self.store: BaseKVStore = store - self.prefix: str = prefix - self.separator: str = separator or DEFAULT_PREFIX_SEPARATOR - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - return await self.store.get(collection=prefixed_collection, key=key) - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - await self.store.put(collection=prefixed_collection, key=key, value=value, ttl=ttl) - - @override - async def delete(self, collection: str, key: str) -> bool: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - return await self.store.delete(collection=prefixed_collection, key=key) - - @override - async def exists(self, collection: str, key: str) -> bool: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - return await self.store.exists(collection=prefixed_collection, key=key) - - @override - async def keys(self, collection: str) -> list[str]: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - return await self.store.keys(collection=prefixed_collection) - - @override - async def clear_collection(self, collection: str) -> int: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - return await self.store.clear_collection(collection=prefixed_collection) - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - prefixed_collection: str = prefix_collection(collection=collection, prefix=self.prefix, separator=self.separator) - ttl_info: TTLInfo | None = await self.store.ttl(collection=prefixed_collection, key=key) - if ttl_info: - ttl_info.collection = collection - ttl_info.key = key - return ttl_info - - @override - async def list_collections(self) -> list[str]: - collections: list[str] = await self.store.list_collections() - - return [ - unprefix_collection(collection=collection, separator=self.separator) - for collection in collections - if collection.startswith(self.prefix) - ] - - @override - async def cull(self) -> None: - await self.store.cull() diff --git a/src/kv_store_adapter/stores/wrappers/prefix_key.py b/src/kv_store_adapter/stores/wrappers/prefix_key.py deleted file mode 100644 index a7c43fe2..00000000 --- a/src/kv_store_adapter/stores/wrappers/prefix_key.py +++ /dev/null @@ -1,69 +0,0 @@ -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_key, unprefix_key -from kv_store_adapter.types import TTLInfo - - -class PrefixKeyWrapper(BaseKVStore): - """Wrapper that prefixes all keys with a given prefix.""" - - def __init__(self, store: BaseKVStore, prefix: str, separator: str | None = None) -> None: - """Initialize the prefix key wrapper. - - Args: - store: The store to wrap. - prefix: The prefix to add to all keys. - separator: The separator to use between the prefix and the key. Defaults to "__". - """ - self.store: BaseKVStore = store - self.prefix: str = prefix - self.separator: str = separator or DEFAULT_PREFIX_SEPARATOR - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) - return await self.store.get(collection=collection, key=prefixed_key) - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) - await self.store.put(collection=collection, key=prefixed_key, value=value, ttl=ttl) - - @override - async def delete(self, collection: str, key: str) -> bool: - prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) - return await self.store.delete(collection=collection, key=prefixed_key) - - @override - async def exists(self, collection: str, key: str) -> bool: - prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) - return await self.store.exists(collection=collection, key=prefixed_key) - - @override - async def keys(self, collection: str) -> list[str]: - keys: list[str] = await self.store.keys(collection=collection) - return [unprefix_key(key=key, separator=self.separator) for key in keys] - - @override - async def clear_collection(self, collection: str) -> int: - return await self.store.clear_collection(collection=collection) - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - prefixed_key: str = prefix_key(key=key, prefix=self.prefix, separator=self.separator) - ttl_info: TTLInfo | None = await self.store.ttl(collection=collection, key=prefixed_key) - if ttl_info: - ttl_info.collection = collection - ttl_info.key = key - return ttl_info - - @override - async def list_collections(self) -> list[str]: - return await self.store.list_collections() - - @override - async def cull(self) -> None: - await self.store.cull() diff --git a/src/kv_store_adapter/stores/wrappers/single_collection.py b/src/kv_store_adapter/stores/wrappers/single_collection.py deleted file mode 100644 index 6806a6cc..00000000 --- a/src/kv_store_adapter/stores/wrappers/single_collection.py +++ /dev/null @@ -1,68 +0,0 @@ -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.stores.utils.compound import DEFAULT_PREFIX_SEPARATOR, prefix_key, unprefix_key -from kv_store_adapter.types import TTLInfo - - -class SingleCollectionWrapper(BaseKVStore): - """Wrapper that forces all requests into a single collection, prefixes the keys with the original collection name. - - The single collection wrapper does not support collection operations.""" - - def __init__(self, store: BaseKVStore, collection: str, prefix_separator: str | None = None) -> None: - self.collection: str = collection - self.prefix_separator: str = prefix_separator or DEFAULT_PREFIX_SEPARATOR - self.store: BaseKVStore = store - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) - return await self.store.get(collection=self.collection, key=prefixed_key) - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) - await self.store.put(collection=self.collection, key=prefixed_key, value=value, ttl=ttl) - - @override - async def delete(self, collection: str, key: str) -> bool: - prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) - return await self.store.delete(collection=self.collection, key=prefixed_key) - - @override - async def exists(self, collection: str, key: str) -> bool: - prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) - return await self.store.exists(collection=self.collection, key=prefixed_key) - - @override - async def keys(self, collection: str) -> list[str]: - keys: list[str] = await self.store.keys(collection=collection) - return [unprefix_key(key=key, separator=self.prefix_separator) for key in keys] - - @override - async def clear_collection(self, collection: str) -> int: - msg = "Clearing a collection is not supported for SingleCollectionWrapper" - raise NotImplementedError(msg) - - # return await self.store.clear_collection(collection=self.collection) - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - prefixed_key: str = prefix_key(key=key, prefix=collection, separator=self.prefix_separator) - ttl: TTLInfo | None = await self.store.ttl(collection=self.collection, key=prefixed_key) - if ttl: - ttl.collection = collection - ttl.key = key - return ttl - - @override - async def list_collections(self) -> list[str]: - msg = "Listing collections is not supported for SingleCollectionWrapper" - raise NotImplementedError(msg) - - @override - async def cull(self) -> None: - await self.store.cull() diff --git a/src/kv_store_adapter/stores/wrappers/statistics.py b/src/kv_store_adapter/stores/wrappers/statistics.py deleted file mode 100644 index 5163808f..00000000 --- a/src/kv_store_adapter/stores/wrappers/statistics.py +++ /dev/null @@ -1,197 +0,0 @@ -from dataclasses import dataclass, field -from typing import Any - -from typing_extensions import override - -from kv_store_adapter.stores.base.unmanaged import BaseKVStore -from kv_store_adapter.types import TTLInfo - - -@dataclass -class BaseStatistics: - """Base statistics container with operation counting.""" - - count: int = field(default=0) - """The number of operations.""" - - def increment(self) -> None: - self.count += 1 - - -@dataclass -class BaseHitMissStatistics(BaseStatistics): - """Statistics container with hit/miss tracking for cache-like operations.""" - - hit: int = field(default=0) - """The number of hits.""" - miss: int = field(default=0) - """The number of misses.""" - - def increment_hit(self) -> None: - self.increment() - self.hit += 1 - - def increment_miss(self) -> None: - self.increment() - self.miss += 1 - - -@dataclass -class GetStatistics(BaseHitMissStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class SetStatistics(BaseStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class DeleteStatistics(BaseHitMissStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class ExistsStatistics(BaseHitMissStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class KeysStatistics(BaseStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class ClearCollectionStatistics(BaseHitMissStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class ListCollectionsStatistics(BaseStatistics): - """A class for statistics about a KV Store collection.""" - - -@dataclass -class KVStoreCollectionStatistics(BaseStatistics): - """A class for statistics about a KV Store collection.""" - - get: GetStatistics = field(default_factory=GetStatistics) - """The statistics for the get operation.""" - - set: SetStatistics = field(default_factory=SetStatistics) - """The statistics for the set operation.""" - - delete: DeleteStatistics = field(default_factory=DeleteStatistics) - """The statistics for the delete operation.""" - - exists: ExistsStatistics = field(default_factory=ExistsStatistics) - """The statistics for the exists operation.""" - - keys: KeysStatistics = field(default_factory=KeysStatistics) - """The statistics for the keys operation.""" - - clear_collection: ClearCollectionStatistics = field(default_factory=ClearCollectionStatistics) - """The statistics for the clear collection operation.""" - - list_collections: ListCollectionsStatistics = field(default_factory=ListCollectionsStatistics) - """The statistics for the list collections operation.""" - - -@dataclass -class KVStoreStatistics: - """Statistics container for a KV Store.""" - - collections: dict[str, KVStoreCollectionStatistics] = field(default_factory=dict) - - def get_collection(self, collection: str) -> KVStoreCollectionStatistics: - if collection not in self.collections: - self.collections[collection] = KVStoreCollectionStatistics() - return self.collections[collection] - - -class StatisticsWrapper(BaseKVStore): - """Statistics wrapper around a KV Store that tracks operation statistics.""" - - def __init__(self, store: BaseKVStore, track_statistics: bool = True) -> None: - self.store: BaseKVStore = store - self._statistics: KVStoreStatistics | None = KVStoreStatistics() if track_statistics else None - - @property - def statistics(self) -> KVStoreStatistics | None: - return self._statistics - - @override - async def get(self, collection: str, key: str) -> dict[str, Any] | None: - if value := await self.store.get(collection=collection, key=key): - if self.statistics: - self.statistics.get_collection(collection).get.increment_hit() - return value - - if self.statistics: - self.statistics.get_collection(collection).get.increment_miss() - - return None - - @override - async def put(self, collection: str, key: str, value: dict[str, Any], *, ttl: float | None = None) -> None: - await self.store.put(collection=collection, key=key, value=value, ttl=ttl) - - if self.statistics: - self.statistics.get_collection(collection).set.increment() - - @override - async def delete(self, collection: str, key: str) -> bool: - if await self.store.delete(collection=collection, key=key): - if self.statistics: - self.statistics.get_collection(collection).delete.increment_hit() - return True - - if self.statistics: - self.statistics.get_collection(collection).delete.increment_miss() - - return False - - @override - async def exists(self, collection: str, key: str) -> bool: - if await self.store.exists(collection=collection, key=key): - if self.statistics: - self.statistics.get_collection(collection).exists.increment_hit() - return True - - if self.statistics: - self.statistics.get_collection(collection).exists.increment_miss() - - return False - - @override - async def keys(self, collection: str) -> list[str]: - keys: list[str] = await self.store.keys(collection) - - if self.statistics: - self.statistics.get_collection(collection).keys.increment() - - return keys - - @override - async def clear_collection(self, collection: str) -> int: - if count := await self.store.clear_collection(collection): - if self.statistics: - self.statistics.get_collection(collection).clear_collection.increment_hit() - return count - - if self.statistics: - self.statistics.get_collection(collection).clear_collection.increment_miss() - - return 0 - - @override - async def ttl(self, collection: str, key: str) -> TTLInfo | None: - return await self.store.ttl(collection=collection, key=key) - - @override - async def list_collections(self) -> list[str]: - return await self.store.list_collections() - - @override - async def cull(self) -> None: - await self.store.cull()