Fix Elasticsearch and Mongo collection handling

Author: strawgate

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

@@ -211,9 +211,15 @@ def __init__(
         self._native_storage = native_storage
         self._is_serverless = False
 
+        # We have 240 characters to work with
+        # We need to account for the index prefix and the hyphen.
+        max_index_length = MAX_INDEX_LENGTH - (len(self._index_prefix) + 1)
+
         self._serializer = ElasticsearchSerializationAdapter(native_storage=native_storage)
         self._collection_sanitization = HybridSanitizationStrategy(
-            replacement_character="_", max_length=MAX_INDEX_LENGTH, allowed_characters=ALLOWED_INDEX_CHARACTERS
+            replacement_character="_",
+            max_length=max_index_length,
+            allowed_characters=ALLOWED_INDEX_CHARACTERS,
         )
         self._key_sanitization = AlwaysHashStrategy()
 

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

@@ -205,7 +205,7 @@ async def _setup_collection(self, *, collection: str) -> None:
             self._collections_by_name[collection] = self._db[sanitized_collection]
             return
 
-        new_collection: AsyncCollection[dict[str, Any]] = await self._db.create_collection(name=collection)
+        new_collection: AsyncCollection[dict[str, Any]] = await self._db.create_collection(name=sanitized_collection)
 
         # Index for efficient key lookups
         _ = await new_collection.create_index(keys="key")
@@ -217,9 +217,7 @@ async def _setup_collection(self, *, collection: str) -> None:
 
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
-        if doc := await self._collections_by_name[sanitized_collection].find_one(filter={"key": key}):
+        if doc := await self._collections_by_name[collection].find_one(filter={"key": key}):
             try:
                 return self._adapter.load_dict(data=doc)
             except DeserializationError:
@@ -232,10 +230,8 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) ->
         if not keys:
             return []
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         # Use find with $in operator to get multiple documents at once
-        cursor = self._collections_by_name[sanitized_collection].find(filter={"key": {"$in": keys}})
+        cursor = self._collections_by_name[collection].find(filter={"key": {"$in": keys}})
 
         managed_entries_by_key: dict[str, ManagedEntry | None] = dict.fromkeys(keys)
 
@@ -258,12 +254,10 @@ async def _put_managed_entry(
     ) -> None:
         mongo_doc = self._adapter.dump_dict(entry=managed_entry)
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         try:
             # Ensure that the value is serializable to JSON
             _ = managed_entry.value_as_json
-            _ = await self._collections_by_name[sanitized_collection].update_one(
+            _ = await self._collections_by_name[collection].update_one(
                 filter={"key": key},
                 update={
                     "$set": {
@@ -291,8 +285,6 @@ async def _put_managed_entries(
         if not keys:
             return
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         operations: list[UpdateOne] = []
         for key, managed_entry in zip(keys, managed_entries, strict=True):
             mongo_doc = self._adapter.dump_dict(entry=managed_entry)
@@ -311,13 +303,11 @@ async def _put_managed_entries(
                 )
             )
 
-        _ = await self._collections_by_name[sanitized_collection].bulk_write(operations)  # pyright: ignore[reportUnknownMemberType]
+        _ = await self._collections_by_name[collection].bulk_write(operations)  # pyright: ignore[reportUnknownMemberType]
 
     @override
     async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
-        result: DeleteResult = await self._collections_by_name[sanitized_collection].delete_one(filter={"key": key})
+        result: DeleteResult = await self._collections_by_name[collection].delete_one(filter={"key": key})
         return bool(result.deleted_count)
 
     @override

key-value/key-value-sync/src/key_value/sync/code_gen/stores/elasticsearch/store.py

@@ -180,9 +180,13 @@ def __init__(
         self._native_storage = native_storage
         self._is_serverless = False
 
+        # We have 240 characters to work with
+        # We need to account for the index prefix and the hyphen.
+        max_index_length = MAX_INDEX_LENGTH - (len(self._index_prefix) + 1)
+
         self._serializer = ElasticsearchSerializationAdapter(native_storage=native_storage)
         self._collection_sanitization = HybridSanitizationStrategy(
-            replacement_character="_", max_length=MAX_INDEX_LENGTH, allowed_characters=ALLOWED_INDEX_CHARACTERS
+            replacement_character="_", max_length=max_index_length, allowed_characters=ALLOWED_INDEX_CHARACTERS
         )
         self._key_sanitization = AlwaysHashStrategy()
 

key-value/key-value-sync/src/key_value/sync/code_gen/stores/mongodb/store.py

@@ -212,7 +212,7 @@ def _setup_collection(self, *, collection: str) -> None:
             self._collections_by_name[collection] = self._db[sanitized_collection]
             return
 
-        new_collection: Collection[dict[str, Any]] = self._db.create_collection(name=collection)
+        new_collection: Collection[dict[str, Any]] = self._db.create_collection(name=sanitized_collection)
 
         # Index for efficient key lookups
         _ = new_collection.create_index(keys="key")
@@ -224,9 +224,7 @@ def _setup_collection(self, *, collection: str) -> None:
 
     @override
     def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
-        if doc := self._collections_by_name[sanitized_collection].find_one(filter={"key": key}):
+        if doc := self._collections_by_name[collection].find_one(filter={"key": key}):
             try:
                 return self._adapter.load_dict(data=doc)
             except DeserializationError:
@@ -239,10 +237,8 @@ def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) -> list[
         if not keys:
             return []
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         # Use find with $in operator to get multiple documents at once
-        cursor = self._collections_by_name[sanitized_collection].find(filter={"key": {"$in": keys}})
+        cursor = self._collections_by_name[collection].find(filter={"key": {"$in": keys}})
 
         managed_entries_by_key: dict[str, ManagedEntry | None] = dict.fromkeys(keys)
 
@@ -259,12 +255,10 @@ def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) -> list[
     def _put_managed_entry(self, *, key: str, collection: str, managed_entry: ManagedEntry) -> None:
         mongo_doc = self._adapter.dump_dict(entry=managed_entry)
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         try:
             # Ensure that the value is serializable to JSON
             _ = managed_entry.value_as_json
-            _ = self._collections_by_name[sanitized_collection].update_one(
+            _ = self._collections_by_name[collection].update_one(
                 filter={"key": key}, update={"$set": {"key": key, **mongo_doc}}, upsert=True
             )
         except InvalidDocument as e:
@@ -285,8 +279,6 @@ def _put_managed_entries(
         if not keys:
             return
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         operations: list[UpdateOne] = []
         for key, managed_entry in zip(keys, managed_entries, strict=True):
             mongo_doc = self._adapter.dump_dict(entry=managed_entry)
@@ -295,13 +287,11 @@ def _put_managed_entries(
                 UpdateOne(filter={"key": key}, update={"$set": {"collection": collection, "key": key, **mongo_doc}}, upsert=True)
             )
 
-        _ = self._collections_by_name[sanitized_collection].bulk_write(operations)  # pyright: ignore[reportUnknownMemberType]
+        _ = self._collections_by_name[collection].bulk_write(operations)  # pyright: ignore[reportUnknownMemberType]
 
     @override
     def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
-        result: DeleteResult = self._collections_by_name[sanitized_collection].delete_one(filter={"key": key})
+        result: DeleteResult = self._collections_by_name[collection].delete_one(filter={"key": key})
         return bool(result.deleted_count)
 
     @override