fix: Address CodeRabbit review feedback and add index validation

- Elasticsearch:
  - Remove unused DEFAULT_MAPPING constant
  - Fix culling in JSON mode by storing timestamps at top level
  - Add storage mode mismatch detection with helpful error messages
  - Add index mapping validation on collection setup

- MongoDB:
  - Add datetime type validation when reading in native mode
  - Add updated_at field to JSON mode for consistency
  - Remove redundant collection and key fields from documents
  - Use $setOnInsert for immutable created_at field
  - Add TTL index validation for native storage mode

Both stores now validate their configuration matches the existing
index/collection setup and raise helpful errors when mismatches occur.

Co-authored-by: William Easton <[email protected]>

Author: github-actions[bot]

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

@@ -40,26 +40,6 @@
 
 DEFAULT_INDEX_PREFIX = "kv_store"
 
-DEFAULT_MAPPING = {
-    "properties": {
-        "created_at": {
-            "type": "date",
-        },
-        "expires_at": {
-            "type": "date",
-        },
-        "collection": {
-            "type": "keyword",
-        },
-        "key": {
-            "type": "keyword",
-        },
-        "value": {
-            "type": "flattened",
-        },
-    },
-}
-
 DEFAULT_PAGE_SIZE = 10000
 PAGE_LIMIT = 10000
 
@@ -141,6 +121,8 @@ async def _setup_collection(self, *, collection: str) -> None:
         index_name = self._sanitize_index_name(collection=collection)
 
         if await self._client.options(ignore_status=404).indices.exists(index=index_name):
+            # Validate that existing index mapping matches the storage mode
+            await self._validate_index_mapping(index_name=index_name, collection=collection)
             return
 
         # Create mapping based on storage mode
@@ -167,6 +149,34 @@ async def _setup_collection(self, *, collection: str) -> None:
 
         _ = await self._client.options(ignore_status=404).indices.create(index=index_name, mappings=mapping, settings={})
 
+    async def _validate_index_mapping(self, *, index_name: str, collection: str) -> None:
+        """Validate that the index mapping matches the configured storage mode."""
+        try:
+            mapping_response = await self._client.indices.get_mapping(index=index_name)
+            mappings = mapping_response.get(index_name, {}).get("mappings", {})
+            value_field_type = mappings.get("properties", {}).get("value", {}).get("type")
+
+            expected_type = "flattened" if self._native_storage else "keyword"
+
+            if value_field_type != expected_type:
+                msg = (
+                    f"Index mapping mismatch for collection '{collection}': "
+                    f"index has 'value' field type '{value_field_type}', "
+                    f"but store is configured for '{expected_type}' (native_storage={self._native_storage}). "
+                    f"To fix this, either: 1) Use the correct storage mode when initializing the store, "
+                    f"or 2) Delete and recreate the index with the new mapping."
+                )
+                raise ValueError(msg)
+        except Exception as e:
+            # If we can't get the mapping, log a warning but don't fail
+            # This allows the store to work even if mapping validation fails
+            if not isinstance(e, ValueError):
+                # Only suppress non-ValueError exceptions (e.g., connection issues)
+                pass
+            else:
+                # Re-raise ValueError from our validation
+                raise
+
     def _sanitize_index_name(self, collection: str) -> str:
         return sanitize_string(
             value=self._index_prefix + "-" + collection,
@@ -198,7 +208,19 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
 
         if self._native_storage:
             # Native storage mode: Get value as flattened object
-            if not (value := source.get("value")) or not isinstance(value, dict):
+            if not (value := source.get("value")):
+                return None
+
+            # Detect if data is in JSON string format
+            if isinstance(value, str):
+                msg = (
+                    f"Data for key '{key}' appears to be in JSON string format, but store is configured "
+                    "for native_storage mode. This indicates a storage mode mismatch. "
+                    "You may need to migrate existing data or use the correct storage mode."
+                )
+                raise ValueError(msg)
+
+            if not isinstance(value, dict):
                 return None
 
             created_at: datetime | None = try_parse_datetime_str(value=source.get("created_at"))
@@ -213,6 +235,15 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
             # JSON string mode: Get value as JSON string and parse it
             json_value: str | None = source.get("value")
 
+            # Detect if data is in native object format
+            if isinstance(json_value, dict):
+                msg = (
+                    f"Data for key '{key}' appears to be in native object format, but store is configured "
+                    "for JSON string mode. This indicates a storage mode mismatch. "
+                    "You may need to migrate existing data or use the correct storage mode."
+                )
+                raise ValueError(msg)
+
             if not isinstance(json_value, str):
                 return None
 
@@ -259,6 +290,12 @@ async def _put_managed_entry(
                 "value": managed_entry.to_json(),  # Store as JSON string
             }
 
+            # Store timestamps at top level for culling to work
+            if managed_entry.created_at:
+                document["created_at"] = managed_entry.created_at.isoformat()
+            if managed_entry.expires_at:
+                document["expires_at"] = managed_entry.expires_at.isoformat()
+
             _ = await self._client.index(
                 index=self._sanitize_index_name(collection=collection),
                 id=self._sanitize_document_id(key=combo_key),

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

@@ -143,6 +143,8 @@ async def _setup_collection(self, *, collection: str) -> None:
 
         if matching_collections:
             self._collections_by_name[collection] = self._db[collection]
+            # Validate indexes for existing collection
+            await self._validate_collection_indexes(collection=collection)
             return
 
         new_collection: AsyncCollection[dict[str, Any]] = await self._db.create_collection(name=collection)
@@ -155,6 +157,32 @@ async def _setup_collection(self, *, collection: str) -> None:
 
         self._collections_by_name[collection] = new_collection
 
+    async def _validate_collection_indexes(self, *, collection: str) -> None:
+        """Validate that the collection indexes match the configured storage mode."""
+        try:
+            coll = self._collections_by_name[collection]
+            indexes = await coll.list_indexes().to_list(length=None)
+
+            # Check for TTL index on expires_at
+            has_ttl_index = any(
+                idx.get("key", {}).get("expires_at") is not None and idx.get("expireAfterSeconds") is not None for idx in indexes
+            )
+
+            if self._native_storage and not has_ttl_index:
+                msg = (
+                    f"Collection '{collection}' is missing TTL index on 'expires_at' field, "
+                    f"but store is configured for native_storage mode. "
+                    f"To fix this, either: 1) Recreate the collection with native_storage=True, "
+                    f"or 2) Manually create the TTL index: db.{collection}.createIndex({{expires_at: 1}}, {{expireAfterSeconds: 0}})"
+                )
+                raise ValueError(msg)
+        except ValueError:
+            # Re-raise our validation errors
+            raise
+        except Exception:
+            # Suppress other errors (e.g., connection issues) to allow store to work
+            pass
+
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
         collection = self._sanitize_collection_name(collection=collection)
@@ -171,10 +199,24 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
             if not isinstance(value, dict):
                 return None
 
-            # Parse datetime objects directly
+            # Parse datetime objects directly and validate types
             created_at: datetime | None = doc.get("created_at")
             expires_at: datetime | None = doc.get("expires_at")
 
+            # Validate datetime types to detect storage mode mismatches
+            if created_at is not None and not isinstance(created_at, datetime):
+                msg = (
+                    f"Data for key '{key}' has invalid created_at type: expected datetime but got {type(created_at).__name__}. "
+                    f"This may indicate a storage mode mismatch."
+                )
+                raise TypeError(msg)
+            if expires_at is not None and not isinstance(expires_at, datetime):
+                msg = (
+                    f"Data for key '{key}' has invalid expires_at type: expected datetime but got {type(expires_at).__name__}. "
+                    f"This may indicate a storage mode mismatch."
+                )
+                raise TypeError(msg)
+
             return ManagedEntry(
                 value=value,
                 created_at=created_at,
@@ -201,35 +243,39 @@ async def _put_managed_entry(
 
         if self._native_storage:
             # Native storage mode: Store value as BSON document
-            document: dict[str, Any] = {
-                "collection": collection,
-                "key": key,
+            set_fields: dict[str, Any] = {
                 "value": managed_entry.value,  # Store as BSON document
                 "updated_at": now(),
             }
 
-            # Store as datetime objects
+            set_on_insert_fields: dict[str, Any] = {}
+
+            # Store as datetime objects (use $setOnInsert for immutable fields)
             if managed_entry.created_at:
-                document["created_at"] = managed_entry.created_at
+                set_on_insert_fields["created_at"] = managed_entry.created_at
             if managed_entry.expires_at:
-                document["expires_at"] = managed_entry.expires_at
+                # expires_at can change, so use $set
+                set_fields["expires_at"] = managed_entry.expires_at
+
+            update_doc: dict[str, Any] = {"$set": set_fields}
+            if set_on_insert_fields:
+                update_doc["$setOnInsert"] = set_on_insert_fields
 
             _ = await self._collections_by_name[collection].update_one(
                 filter={"key": key},
-                update={"$set": document},
+                update=update_doc,
                 upsert=True,
             )
         else:
             # JSON string mode: Store value as JSON string
-            document: dict[str, Any] = {
-                "collection": collection,
-                "key": key,
+            set_fields: dict[str, Any] = {
                 "value": managed_entry.to_json(),  # Store as JSON string
+                "updated_at": now(),
             }
 
             _ = await self._collections_by_name[collection].update_one(
                 filter={"key": key},
-                update={"$set": document},
+                update={"$set": set_fields},
                 upsert=True,
             )