fix: address CodeRabbit review feedback

- Elasticsearch: disable doc_values for JSON mode, enhance mapping validation, fix empty dict handling
- MongoDB: add inverse TTL validation, explicit type errors, fix expires_at=None bug
- Fix line length issues and apply linting suggestions

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

Author: github-actions[bot]

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

@@ -64,10 +64,25 @@ class ElasticsearchStore(
     _native_storage: bool
 
     @overload
-    def __init__(self, *, elasticsearch_client: AsyncElasticsearch, index_prefix: str, default_collection: str | None = None, native_storage: bool = False) -> None: ...
+    def __init__(
+        self,
+        *,
+        elasticsearch_client: AsyncElasticsearch,
+        index_prefix: str,
+        default_collection: str | None = None,
+        native_storage: bool = False,
+    ) -> None: ...
 
     @overload
-    def __init__(self, *, url: str, api_key: str | None = None, index_prefix: str, default_collection: str | None = None, native_storage: bool = False) -> None: ...
+    def __init__(
+        self,
+        *,
+        url: str,
+        api_key: str | None = None,
+        index_prefix: str,
+        default_collection: str | None = None,
+        native_storage: bool = False,
+    ) -> None: ...
 
     def __init__(
         self,
@@ -142,7 +157,7 @@ async def _setup_collection(self, *, collection: str) -> None:
                 },
                 "value": {
                     "type": "flattened" if self._native_storage else "keyword",
-                    **({"index": False} if not self._native_storage else {}),
+                    **({"index": False, "doc_values": False} if not self._native_storage else {}),
                 },
             },
         }
@@ -154,7 +169,10 @@ async def _validate_index_mapping(self, *, index_name: str, collection: str) ->
         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")
+            props = mappings.get("properties", {})
+            value_field_type = props.get("value", {}).get("type")
+            created_type = props.get("created_at", {}).get("type")
+            expires_type = props.get("expires_at", {}).get("type")
 
             expected_type = "flattened" if self._native_storage else "keyword"
 
@@ -166,16 +184,24 @@ async def _validate_index_mapping(self, *, index_name: str, collection: str) ->
                     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
+                raise ValueError(msg)  # noqa: TRY301
+
+            # Enforce date types for timestamps (both modes)
+            for field_name, field_type in (("created_at", created_type), ("expires_at", expires_type)):
+                if field_type not in ("date", None):  # None => not yet created; will be added on first write
+                    msg = (
+                        f"Index mapping mismatch for collection '{collection}': "
+                        f"'{field_name}' is mapped as '{field_type}', expected 'date'. "
+                        f"Delete and recreate the index or fix the mapping."
+                    )
+                    raise ValueError(msg)  # noqa: TRY301
+        except ValueError:
+            raise
+        except Exception:
+            # Log a warning but do not fail hard (keep behavior)
+            import logging
+
+            logging.getLogger(__name__).warning("Failed to validate mapping for index '%s' (collection '%s')", index_name, collection)
 
     def _sanitize_index_name(self, collection: str) -> str:
         return sanitize_string(
@@ -208,8 +234,9 @@ 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")):
+            if "value" not in source:
                 return None
+            value = source["value"]
 
             # Detect if data is in JSON string format
             if isinstance(value, str):
@@ -218,7 +245,7 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
                     "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)
+                raise TypeError(msg)
 
             if not isinstance(value, dict):
                 return None
@@ -231,23 +258,22 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
                 created_at=created_at,
                 expires_at=expires_at,
             )
-        else:
-            # 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)
+        # 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 TypeError(msg)
 
-            if not isinstance(json_value, str):
-                return None
+        if not isinstance(json_value, str):
+            return None
 
-            return ManagedEntry.from_json(json_str=json_value)
+        return ManagedEntry.from_json(json_str=json_value)
 
     @property
     def _should_refresh_on_put(self) -> bool:

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

@@ -71,7 +71,13 @@ def __init__(
 
     @overload
     def __init__(
-        self, *, url: str, db_name: str | None = None, coll_name: str | None = None, default_collection: str | None = None, native_storage: bool = False
+        self,
+        *,
+        url: str,
+        db_name: str | None = None,
+        coll_name: str | None = None,
+        default_collection: str | None = None,
+        native_storage: bool = False,
     ) -> None:
         """Initialize the MongoDB store.
 
@@ -175,11 +181,19 @@ async def _validate_collection_indexes(self, *, collection: str) -> None:
                     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)
+                raise ValueError(msg)  # noqa: TRY301
+            if not self._native_storage and has_ttl_index:
+                msg = (
+                    f"Collection '{collection}' has a TTL index on 'expires_at' field, "
+                    f"but store is configured for JSON string mode (native_storage=False). "
+                    f"This may cause unexpected behavior. Consider either: "
+                    f"1) Using native_storage=True, or 2) Dropping the TTL index."
+                )
+                raise ValueError(msg)  # noqa: TRY301
         except ValueError:
             # Re-raise our validation errors
             raise
-        except Exception:
+        except Exception:  # noqa: S110
             # Suppress other errors (e.g., connection issues) to allow store to work
             pass
 
@@ -197,7 +211,11 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
             value: dict[str, Any] | None = doc.get("value")
 
             if not isinstance(value, dict):
-                return None
+                msg = (
+                    f"Data for key '{key}' has invalid value type: expected dict but got {type(value).__name__}. "
+                    f"This may indicate the collection contains JSON-mode data but native_storage=True."
+                )
+                raise TypeError(msg)
 
             # Parse datetime objects directly and validate types
             created_at: datetime | None = doc.get("created_at")
@@ -222,14 +240,17 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
                 created_at=created_at,
                 expires_at=expires_at,
             )
-        else:
-            # JSON string mode: Get value as JSON string and parse it
-            json_value: str | None = doc.get("value")
+        # JSON string mode: Get value as JSON string and parse it
+        json_value: str | None = doc.get("value")
 
-            if not isinstance(json_value, str):
-                return None
+        if not isinstance(json_value, str):
+            msg = (
+                f"Data for key '{key}' has invalid value type: expected str but got {type(json_value).__name__}. "
+                f"This may indicate the collection contains native-mode data but native_storage=False."
+            )
+            raise TypeError(msg)
 
-            return ManagedEntry.from_json(json_str=json_value)
+        return ManagedEntry.from_json(json_str=json_value)
 
     @override
     async def _put_managed_entry(
@@ -253,14 +274,19 @@ async def _put_managed_entry(
             # Store as datetime objects (use $setOnInsert for immutable fields)
             if managed_entry.created_at:
                 set_on_insert_fields["created_at"] = managed_entry.created_at
-            if managed_entry.expires_at:
-                # expires_at can change, so use $set
-                set_fields["expires_at"] = managed_entry.expires_at
 
+            # Build update document
             update_doc: dict[str, Any] = {"$set": set_fields}
             if set_on_insert_fields:
                 update_doc["$setOnInsert"] = set_on_insert_fields
 
+            # Always handle expires_at to support removing expiration
+            if managed_entry.expires_at is not None:
+                set_fields["expires_at"] = managed_entry.expires_at
+            else:
+                # Use $unset to remove the field when expires_at is None
+                update_doc["$unset"] = {"expires_at": ""}
+
             _ = await self._collections_by_name[collection].update_one(
                 filter={"key": key},
                 update=update_doc,