Updates for PR Feedback

Author: strawgate

key-value/key-value-aio/src/key_value/aio/stores/base.py

@@ -24,6 +24,8 @@
     AsyncKeyValueProtocol,
 )
 
+SEED_DATA_TYPE = Mapping[str, Mapping[str, Mapping[str, Any]]]
+
 
 class BaseStore(AsyncKeyValueProtocol, ABC):
     """An opinionated Abstract base class for managed key-value stores using ManagedEntry objects.
@@ -43,21 +45,29 @@ class BaseStore(AsyncKeyValueProtocol, ABC):
     _setup_collection_locks: defaultdict[str, Lock]
     _setup_collection_complete: defaultdict[str, bool]
 
+    _seed: SEED_DATA_TYPE
+
     default_collection: str
 
-    def __init__(self, *, default_collection: str | None = None) -> None:
+    def __init__(self, *, default_collection: str | None = None, seed: SEED_DATA_TYPE | None = None) -> None:
         """Initialize the managed key-value store.
 
         Args:
             default_collection: The default collection to use if no collection is provided.
                 Defaults to "default_collection".
+            seed: Optional seed data to pre-populate the store. Format: {collection: {key: {field: value, ...}}}.
+                Each value must be a mapping (dict) that will be stored as the entry's value.
+                Seeding occurs during every store setup (when the store is entered or when the first operation
+                is performed on the store).
         """
 
         self._setup_complete = False
         self._setup_lock = Lock()
         self._setup_collection_locks = defaultdict(Lock)
         self._setup_collection_complete = defaultdict(bool)
 
+        self._seed = seed or {}
+
         self.default_collection = default_collection or DEFAULT_COLLECTION_NAME
 
         if not hasattr(self, "_stable_api"):
@@ -74,6 +84,13 @@ async def _setup(self) -> None:
     async def _setup_collection(self, *, collection: str) -> None:
         """Initialize the collection (called once before first use of the collection)."""
 
+    async def _seed_store(self) -> None:
+        """Seed the store with the data from the seed."""
+        for collection, items in self._seed.items():
+            await self.setup_collection(collection=collection)
+            for key, value in items.items():
+                await self.put(key=key, value=value, collection=collection)
+
     async def setup(self) -> None:
         if not self._setup_complete:
             async with self._setup_lock:
@@ -86,6 +103,8 @@ async def setup(self) -> None:
                         ) from e
                     self._setup_complete = True
 
+                await self._seed_store()
+
     async def setup_collection(self, *, collection: str) -> None:
         await self.setup()
 

key-value/key-value-aio/src/key_value/aio/stores/memory/store.py

@@ -1,5 +1,4 @@
 import sys
-from collections.abc import Mapping
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any, SupportsFloat
@@ -9,6 +8,7 @@
 from typing_extensions import Self, override
 
 from key_value.aio.stores.base import (
+    SEED_DATA_TYPE,
     BaseDestroyCollectionStore,
     BaseDestroyStore,
     BaseEnumerateCollectionsStore,
@@ -109,14 +109,13 @@ class MemoryStore(BaseDestroyStore, BaseDestroyCollectionStore, BaseEnumerateCol
     max_entries_per_collection: int
 
     _cache: dict[str, MemoryCollection]
-    _seed: Mapping[str, Mapping[str, Mapping[str, Any]]]
 
     def __init__(
         self,
         *,
         max_entries_per_collection: int = DEFAULT_MAX_ENTRIES_PER_COLLECTION,
         default_collection: str | None = None,
-        seed: Mapping[str, Mapping[str, Mapping[str, Any]]] | None = None,
+        seed: SEED_DATA_TYPE | None = None,
     ):
         """Initialize a fixed-size in-memory store.
 
@@ -131,24 +130,15 @@ def __init__(
         self.max_entries_per_collection = max_entries_per_collection
 
         self._cache = {}
-        self._seed = seed or {}
 
         self._stable_api = True
 
-        super().__init__(default_collection=default_collection)
+        super().__init__(default_collection=default_collection, seed=seed)
 
-    def _create_collection(self, collection: str) -> MemoryCollection:
-        """Create a new collection.
-
-        Args:
-            collection: The collection name.
-
-        Returns:
-            The created MemoryCollection instance.
-        """
-        collection_cache = MemoryCollection(max_entries=self.max_entries_per_collection)
-        self._cache[collection] = collection_cache
-        return collection_cache
+    @override
+    async def _setup(self) -> None:
+        for collection in self._seed:
+            await self._setup_collection(collection=collection)
 
     @override
     async def _setup_collection(self, *, collection: str) -> None:
@@ -157,15 +147,11 @@ async def _setup_collection(self, *, collection: str) -> None:
         Args:
             collection: The collection name.
         """
-        # Create the collection
-        collection_cache = self._create_collection(collection)
-
-        # Seed the collection if seed data is available for it
-        if collection in self._seed:
-            items = self._seed[collection]
-            for key, value in items.items():
-                managed_entry = ManagedEntry(value=value)
-                collection_cache.put(key=key, value=managed_entry)
+        if collection in self._cache:
+            return
+
+        collection_cache = MemoryCollection(max_entries=self.max_entries_per_collection)
+        self._cache[collection] = collection_cache
 
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:

key-value/key-value-aio/src/key_value/aio/wrappers/default_value/wrapper.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping, Sequence
-from typing import Any
+from typing import Any, SupportsFloat
 
+from key_value.shared.utils.managed_entry import dump_to_json, load_from_json
 from typing_extensions import override
 
 from key_value.aio.protocols.key_value import AsyncKeyValue
@@ -12,28 +13,35 @@ class DefaultValueWrapper(BaseWrapper):
 
     This wrapper provides dict.get(key, default) behavior for the key-value store,
     allowing you to specify a default value to return instead of None when a key doesn't exist.
+
+    It does not store the default value in the underlying key-value store and the TTL returned with the default
+    value is hard-coded based on the default_ttl parameter. Picking a default_ttl requires careful consideration
+    of how the value will be used and if any other wrappers will be used that may rely on the TTL.
     """
 
-    _key_value: AsyncKeyValue
     key_value: AsyncKeyValue  # Alias for BaseWrapper compatibility
+    _default_ttl: float | None
+    _default_value_json: str
 
     def __init__(
         self,
         key_value: AsyncKeyValue,
         default_value: Mapping[str, Any],
-        default_ttl: float | None = None,
+        default_ttl: SupportsFloat | None = None,
     ) -> None:
         """Initialize the DefaultValueWrapper.
 
         Args:
             key_value: The underlying key-value store to wrap.
             default_value: The default value to return when a key is not found.
-            default_ttl: The TTL to return for default values. Defaults to None.
+            default_ttl: The TTL to return to the caller for default values. Defaults to None.
         """
-        self._key_value = key_value
-        self.key_value = key_value  # Alias for BaseWrapper compatibility
-        self._default_value = default_value
-        self._default_ttl = default_ttl
+        self.key_value = key_value
+        self._default_value_json = dump_to_json(obj=dict(default_value))
+        self._default_ttl = None if default_ttl is None else float(default_ttl)
+
+    def _new_default_value(self) -> dict[str, Any]:
+        return load_from_json(json_str=self._default_value_json)
 
     @override
     async def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
@@ -46,8 +54,8 @@ async def get(self, key: str, *, collection: str | None = None) -> dict[str, Any
         Returns:
             The value associated with the key, or the default value if not found.
         """
-        result = await self._key_value.get(key=key, collection=collection)
-        return result if result is not None else dict(self._default_value)
+        result = await self.key_value.get(key=key, collection=collection)
+        return result if result is not None else self._new_default_value()
 
     @override
     async def get_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[dict[str, Any] | None]:
@@ -60,8 +68,8 @@ async def get_many(self, keys: Sequence[str], *, collection: str | None = None)
         Returns:
             A list of values, with default values for missing keys.
         """
-        results = await self._key_value.get_many(keys=keys, collection=collection)
-        return [result if result is not None else dict(self._default_value) for result in results]
+        results = await self.key_value.get_many(keys=keys, collection=collection)
+        return [result if result is not None else self._new_default_value() for result in results]
 
     @override
     async def ttl(self, key: str, *, collection: str | None = None) -> tuple[dict[str, Any] | None, float | None]:
@@ -74,9 +82,9 @@ async def ttl(self, key: str, *, collection: str | None = None) -> tuple[dict[st
         Returns:
             A tuple of (value, ttl), with default value and default TTL if not found.
         """
-        result, ttl_value = await self._key_value.ttl(key=key, collection=collection)
+        result, ttl_value = await self.key_value.ttl(key=key, collection=collection)
         if result is None:
-            return (dict(self._default_value), self._default_ttl)
+            return (self._new_default_value(), self._default_ttl)
         return (result, ttl_value)
 
     @override
@@ -90,7 +98,7 @@ async def ttl_many(self, keys: Sequence[str], *, collection: str | None = None)
         Returns:
             A list of (value, ttl) tuples, with default values and default TTL for missing keys.
         """
-        results = await self._key_value.ttl_many(keys=keys, collection=collection)
+        results = await self.key_value.ttl_many(keys=keys, collection=collection)
         return [
-            (result, ttl_value) if result is not None else (dict(self._default_value), self._default_ttl) for result, ttl_value in results
+            (result, ttl_value) if result is not None else (self._new_default_value(), self._default_ttl) for result, ttl_value in results
         ]

key-value/key-value-aio/src/key_value/aio/wrappers/pydantic_json/__init__.py

@@ -10,3 +10,7 @@ class TestMemoryStore(BaseStoreTests):
     @pytest.fixture
     async def store(self) -> MemoryStore:
         return MemoryStore(max_entries_per_collection=500)
+
+    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"}