feat: add key, collection, and version fields to serialized documents

github-actions[bot] · strawgate · github-actions[bot] · commit 339f258dd9cc · 2025-11-07T18:03:00.000Z
This change adds support for storing unsanitized key and collection names along with a version number in serialized ManagedEntry documents. This enables proper enumeration support in stores that perform sanitization or hashing of keys and collection names. Changes: - Updated SerializationAdapter.dump_dict() to accept optional key, collection, and version parameters - Updated SerializationAdapter.dump_json() to pass through these parameters - Added version field (default: 1) to all serialized documents - Updated all store implementations to pass key and collection when serializing entries - Updated Elasticsearch mapping to include version field (integer type) - Fixed MongoDB _setup_collection to use sanitized collection name in filter - Fixed MongoDB _delete_collection to properly remove collection from cache Stores updated: - Elasticsearch, MongoDB, Keyring, WindowsRegistry, Memcached (sanitize keys/collections) - DynamoDB, Redis, RocksDB, Simple, Valkey, Vault, Disk (consistency) Resolves #203 Co-authored-by: William Easton <strawgate@users.noreply.github.com>
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/disk/multi_store.py b/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:
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/disk/store.py b/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:
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/dynamodb/store.py b/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},
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/elasticsearch/store.py b/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])
 
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/keyring/store.py b/key-value/key-value-aio/src/key_value/aio/stores/keyring/store.py
@@ -105,7 +105,7 @@ 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)
 
         keyring.set_password(service_name=self._service_name, username=combo_key, password=json_str)
 
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/memcached/store.py b/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"),
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/mongodb/store.py b/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
 
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/redis/store.py b/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)
 
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/rocksdb/store.py b/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)
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/simple/store.py b/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,
         )
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/valkey/store.py b/key-value/key-value-aio/src/key_value/aio/stores/valkey/store.py
@@ -126,7 +126,7 @@ async def _put_managed_entry(
     ) -> None:
         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)
 
         expiry: ExpirySet | None = ExpirySet(expiry_type=ExpiryType.SEC, value=int(managed_entry.ttl)) if managed_entry.ttl else None
 
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/vault/store.py b/key-value/key-value-aio/src/key_value/aio/stores/vault/store.py
@@ -108,7 +108,7 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
     async def _put_managed_entry(self, *, key: str, collection: str, managed_entry: ManagedEntry) -> None:
         combo_key: str = compound_key(collection=collection, key=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)
 
         # Store the JSON string in a 'value' field
         secret_data = {"value": json_str}
diff --git a/key-value/key-value-aio/src/key_value/aio/stores/windows_registry/store.py b/key-value/key-value-aio/src/key_value/aio/stores/windows_registry/store.py
@@ -102,7 +102,7 @@ async def _put_managed_entry(self, *, key: str, collection: str, managed_entry:
         sanitized_key = self._sanitize_key(key=key)
         registry_path = self._get_registry_path(collection=collection)
 
-        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)
 
         set_reg_sz_value(hive=self._hive, sub_key=registry_path, value_name=sanitized_key, value=json_str)
 
diff --git a/key-value/key-value-shared/src/key_value/shared/utils/serialization.py b/key-value/key-value-shared/src/key_value/shared/utils/serialization.py
@@ -111,13 +111,39 @@ def prepare_dump(self, data: dict[str, Any]) -> dict[str, Any]:
 
         This method is used by subclasses to handle any required transformations before dumping the data to a dictionary."""
 
-    def dump_dict(self, entry: ManagedEntry, exclude_none: bool = True) -> dict[str, Any]:
-        """Convert a ManagedEntry to a dictionary."""
+    def dump_dict(
+        self,
+        entry: ManagedEntry,
+        exclude_none: bool = True,
+        *,
+        key: str | None = None,
+        collection: str | None = None,
+        version: int = 1,
+    ) -> dict[str, Any]:
+        """Convert a ManagedEntry to a dictionary.
+
+        Args:
+            entry: The ManagedEntry to serialize
+            exclude_none: Whether to exclude None values from the output
+            key: Optional unsanitized key name to include in the document
+            collection: Optional unsanitized collection name to include in the document
+            version: Document schema version (default: 1)
+
+        Returns:
+            A dictionary representation of the ManagedEntry with optional metadata fields
+        """
 
         data: dict[str, Any] = {
+            "version": version,
             "value": entry.value_as_dict if self._value_format == "dict" else entry.value_as_json,
         }
 
+        if key is not None:
+            data["key"] = key
+
+        if collection is not None:
+            data["collection"] = collection
+
         if self._date_format == "isoformat":
             data["created_at"] = entry.created_at_isoformat
             data["expires_at"] = entry.expires_at_isoformat
@@ -131,12 +157,31 @@ def dump_dict(self, entry: ManagedEntry, exclude_none: bool = True) -> dict[str,
 
         return self.prepare_dump(data=data)
 
-    def dump_json(self, entry: ManagedEntry, exclude_none: bool = True) -> str:
-        """Convert a ManagedEntry to a JSON string."""
+    def dump_json(
+        self,
+        entry: ManagedEntry,
+        exclude_none: bool = True,
+        *,
+        key: str | None = None,
+        collection: str | None = None,
+        version: int = 1,
+    ) -> str:
+        """Convert a ManagedEntry to a JSON string.
+
+        Args:
+            entry: The ManagedEntry to serialize
+            exclude_none: Whether to exclude None values from the output
+            key: Optional unsanitized key name to include in the document
+            collection: Optional unsanitized collection name to include in the document
+            version: Document schema version (default: 1)
+
+        Returns:
+            A JSON string representation of the ManagedEntry with optional metadata fields
+        """
         if self._date_format == "datetime":
             msg = 'dump_json is incompatible with date_format="datetime"; use date_format="isoformat" or dump_dict().'
             raise SerializationError(msg)
-        return dump_to_json(obj=self.dump_dict(entry=entry, exclude_none=exclude_none))
+        return dump_to_json(obj=self.dump_dict(entry=entry, exclude_none=exclude_none, key=key, collection=collection, version=version))
 
 
 class BasicSerializationAdapter(SerializationAdapter):
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/disk/multi_store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/disk/multi_store.py
@@ -129,7 +129,11 @@ def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | Non
 
     @override
     def _put_managed_entry(self, *, key: str, 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
     def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/disk/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/disk/store.py
@@ -104,7 +104,11 @@ def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | Non
     def _put_managed_entry(self, *, key: str, collection: str, managed_entry: ManagedEntry) -> 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
     def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/elasticsearch/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/elasticsearch/store.py
@@ -51,6 +51,7 @@
         "expires_at": {"type": "date"},
         "collection": {"type": "keyword"},
         "key": {"type": "keyword"},
+        "version": {"type": "integer"},
         "value": {"properties": {"flattened": {"type": "flattened"}}},
     }
 }
@@ -321,7 +322,7 @@ def _put_managed_entry(self, *, key: str, collection: str, managed_entry: Manage
         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:
             _ = self._client.index(index=index_name, id=document_id, body=document, refresh=self._should_refresh_on_put)
@@ -354,7 +355,7 @@ 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])
 
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/keyring/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/keyring/store.py
@@ -104,7 +104,7 @@ def _put_managed_entry(self, *, key: str, collection: str, managed_entry: Manage
 
         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)
 
         keyring.set_password(service_name=self._service_name, username=combo_key, password=json_str)
 
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/mongodb/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/mongodb/store.py
@@ -214,7 +214,7 @@ 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] = self._db.list_collection_names(filter=collection_filter)
 
         if matching_collections:
@@ -262,7 +262,7 @@ def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) -> list[
 
     @override
     def _put_managed_entry(self, *, key: str, 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
@@ -290,7 +290,7 @@ 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(filter={"key": key}, update={"$set": {"collection": collection, "key": key, **mongo_doc}}, upsert=True)
@@ -318,8 +318,7 @@ def _delete_collection(self, *, collection: str) -> bool:
 
         _ = 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
 
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/redis/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/redis/store.py
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/rocksdb/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/rocksdb/store.py
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/simple/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/simple/store.py
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/valkey/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/valkey/store.py
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/vault/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/vault/store.py
diff --git a/key-value/key-value-sync/src/key_value/sync/code_gen/stores/windows_registry/store.py b/key-value/key-value-sync/src/key_value/sync/code_gen/stores/windows_registry/store.py

Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ async def _put_managed_entry(self, *, key: str, collection: str, managed_entry:`
`69`	`69`	`_ = self._data.pop(next(iter(self._data)))`
`70`	`70`
`71`	`71`	`self._data[combo_key] = SimpleStoreEntry(`
`72`		`- json_str=self._serialization_adapter.dump_json(entry=managed_entry),`
	`72`	`+ json_str=self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection),`
`73`	`73`	`expires_at=managed_entry.expires_at,`
`74`	`74`	`created_at=managed_entry.created_at,`
`75`	`75`	`)`