feat: add key, collection, and version fields to serialized documents (#204)

Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: William Easton <strawgate@users.noreply.github.com>

Author: strawgate

key-value/key-value-aio/src/key_value/aio/stores/disk/multi_store.py

@@ -132,7 +132,11 @@ async def _put_managed_entry(
         collection: str,
         managed_entry: ManagedEntry,
     ) -> None:
-        _ = self._cache[collection].set(key=key, value=self._serialization_adapter.dump_json(entry=managed_entry), expire=managed_entry.ttl)
+        _ = self._cache[collection].set(
+            key=key,
+            value=self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection),
+            expire=managed_entry.ttl,
+        )
 
     @override
     async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:

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

@@ -107,7 +107,11 @@ async def _put_managed_entry(
     ) -> None:
         combo_key: str = compound_key(collection=collection, key=key)
 
-        _ = self._cache.set(key=combo_key, value=self._serialization_adapter.dump_json(entry=managed_entry), expire=managed_entry.ttl)
+        _ = self._cache.set(
+            key=combo_key,
+            value=self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection),
+            expire=managed_entry.ttl,
+        )
 
     @override
     async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:

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

@@ -219,7 +219,7 @@ async def _put_managed_entry(
         managed_entry: ManagedEntry,
     ) -> None:
         """Store a managed entry in DynamoDB."""
-        json_value = self._serialization_adapter.dump_json(entry=managed_entry)
+        json_value = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection)
 
         item: dict[str, Any] = {
             "collection": {"S": collection},

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

@@ -67,6 +67,9 @@
         "key": {
             "type": "keyword",
         },
+        "version": {
+            "type": "integer",
+        },
         "value": {
             "properties": {
                 "flattened": {
@@ -357,7 +360,7 @@ async def _put_managed_entry(
         index_name: str = self._get_index_name(collection=collection)
         document_id: str = self._get_document_id(key=key)
 
-        document: dict[str, Any] = self._serializer.dump_dict(entry=managed_entry)
+        document: dict[str, Any] = self._serializer.dump_dict(entry=managed_entry, key=key, collection=collection)
 
         try:
             _ = await self._client.index(
@@ -395,7 +398,7 @@ async def _put_managed_entries(
 
             index_action: dict[str, Any] = new_bulk_action(action="index", index=index_name, document_id=document_id)
 
-            document: dict[str, Any] = self._serializer.dump_dict(entry=managed_entry)
+            document: dict[str, Any] = self._serializer.dump_dict(entry=managed_entry, key=key, collection=collection)
 
             operations.extend([index_action, document])
 

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

@@ -1,5 +1,8 @@
 """Python keyring-based key-value store."""
 
+import os
+
+from key_value.shared.errors.key_value import ValueTooLargeError
 from key_value.shared.utils.compound import compound_key
 from key_value.shared.utils.managed_entry import ManagedEntry
 from key_value.shared.utils.sanitization import HybridSanitizationStrategy, SanitizationStrategy
@@ -17,6 +20,16 @@
 
 DEFAULT_KEYCHAIN_SERVICE = "py-key-value"
 
+
+def is_value_too_large(value: bytes) -> bool:
+    value_length = len(value)
+    if os.name == "nt":
+        return value_length > WINDOWS_MAX_VALUE_LENGTH
+    return False
+
+
+WINDOWS_MAX_VALUE_LENGTH = 2560  # bytes
+
 MAX_KEY_COLLECTION_LENGTH = 256
 ALLOWED_KEY_COLLECTION_CHARACTERS: str = ALPHANUMERIC_CHARACTERS
 
@@ -105,7 +118,11 @@ async def _put_managed_entry(self, *, key: str, collection: str, managed_entry:
 
         combo_key: str = compound_key(collection=sanitized_collection, key=sanitized_key)
 
-        json_str: str = self._serialization_adapter.dump_json(entry=managed_entry)
+        json_str: str = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection)
+        encoded_json_bytes: bytes = json_str.encode(encoding="utf-8")
+
+        if is_value_too_large(value=encoded_json_bytes):
+            raise ValueTooLargeError(size=len(encoded_json_bytes), max_size=2560, collection=sanitized_collection, key=sanitized_key)
 
         keyring.set_password(service_name=self._service_name, username=combo_key, password=json_str)
 

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

@@ -125,7 +125,7 @@ async def _put_managed_entry(
         else:
             exptime = max(int(managed_entry.ttl), 1)
 
-        json_value: str = self._serialization_adapter.dump_json(entry=managed_entry)
+        json_value: str = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection)
 
         _ = await self._client.set(
             key=combo_key.encode(encoding="utf-8"),

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

@@ -219,7 +219,7 @@ async def _setup_collection(self, *, collection: str) -> None:
         # Ensure index on the unique combo key and supporting queries
         sanitized_collection = self._sanitize_collection(collection=collection)
 
-        collection_filter: dict[str, str] = {"name": collection}
+        collection_filter: dict[str, str] = {"name": sanitized_collection}
         matching_collections: list[str] = await self._db.list_collection_names(filter=collection_filter)
 
         if matching_collections:
@@ -273,7 +273,7 @@ async def _put_managed_entry(
         collection: str,
         managed_entry: ManagedEntry,
     ) -> None:
-        mongo_doc = self._adapter.dump_dict(entry=managed_entry)
+        mongo_doc = self._adapter.dump_dict(entry=managed_entry, key=key, collection=collection)
 
         try:
             # Ensure that the value is serializable to JSON
@@ -308,7 +308,7 @@ async def _put_managed_entries(
 
         operations: list[UpdateOne] = []
         for key, managed_entry in zip(keys, managed_entries, strict=True):
-            mongo_doc = self._adapter.dump_dict(entry=managed_entry)
+            mongo_doc = self._adapter.dump_dict(entry=managed_entry, key=key, collection=collection)
 
             operations.append(
                 UpdateOne(
@@ -346,8 +346,7 @@ async def _delete_collection(self, *, collection: str) -> bool:
 
         _ = await self._db.drop_collection(name_or_collection=collection_name)
 
-        if collection_name in self._collections_by_name:
-            del self._collections_by_name[collection]
+        self._collections_by_name.pop(collection, None)
 
         return True
 

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

@@ -132,7 +132,7 @@ async def _put_managed_entry(
     ) -> None:
         combo_key: str = compound_key(collection=collection, key=key)
 
-        json_value: str = self._adapter.dump_json(entry=managed_entry)
+        json_value: str = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection)
 
         if managed_entry.ttl is not None:
             # Redis does not support <= 0 TTLs
@@ -160,7 +160,7 @@ async def _put_managed_entries(
             # If there is no TTL, we can just do a simple mset
             mapping: dict[str, str] = {}
             for key, managed_entry in zip(keys, managed_entries, strict=True):
-                json_value = self._adapter.dump_json(entry=managed_entry)
+                json_value = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection)
                 mapping[compound_key(collection=collection, key=key)] = json_value
 
             await self._client.mset(mapping=mapping)
@@ -175,7 +175,7 @@ async def _put_managed_entries(
 
         for key, managed_entry in zip(keys, managed_entries, strict=True):
             combo_key: str = compound_key(collection=collection, key=key)
-            json_value = self._adapter.dump_json(entry=managed_entry)
+            json_value = self._adapter.dump_json(entry=managed_entry, key=key, collection=collection)
 
             pipeline.setex(name=combo_key, time=ttl_seconds, value=json_value)
 

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

@@ -124,7 +124,7 @@ async def _put_managed_entry(
         self._fail_on_closed_store()
 
         combo_key: str = compound_key(collection=collection, key=key)
-        json_value: str = self._serialization_adapter.dump_json(entry=managed_entry)
+        json_value: str = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection)
 
         self._db[combo_key] = json_value.encode("utf-8")
 
@@ -147,7 +147,7 @@ async def _put_managed_entries(
         batch = WriteBatch()
         for key, managed_entry in zip(keys, managed_entries, strict=True):
             combo_key: str = compound_key(collection=collection, key=key)
-            json_value: str = self._serialization_adapter.dump_json(entry=managed_entry)
+            json_value: str = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection)
             batch.put(combo_key, json_value.encode("utf-8"))
 
         self._db.write(batch)

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

@@ -69,7 +69,7 @@ async def _put_managed_entry(self, *, key: str, collection: str, managed_entry:
             _ = self._data.pop(next(iter(self._data)))
 
         self._data[combo_key] = SimpleStoreEntry(
-            json_str=self._serialization_adapter.dump_json(entry=managed_entry),
+            json_str=self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection),
             expires_at=managed_entry.expires_at,
             created_at=managed_entry.created_at,
         )