feat: add native storage mode for MongoDB

Add support for storing dictionaries as native BSON objects in MongoDB,
similar to the native storage feature implemented for Elasticsearch in #149.

Changes:
- Add native_storage parameter to MongoDBStore (defaults to True)
- Use two-field approach: value.dict for native BSON, value.string for legacy JSON strings
- Support reading both formats for backward compatibility and migration
- Update serialization/deserialization functions to handle both modes
- Add comprehensive tests for native mode, legacy mode, and migration scenarios
- Update sync library via codegen

The native storage mode stores values as BSON documents instead of JSON strings,
enabling better query support and more efficient storage in MongoDB.

Resolves #159

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: William Easton <strawgate@users.noreply.github.com>

Author: github-actions[bot]

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

@@ -37,36 +37,66 @@ def document_to_managed_entry(document: dict[str, Any]) -> ManagedEntry:
     """Convert a MongoDB document back to a ManagedEntry.
 
     This function deserializes a MongoDB document (created by `managed_entry_to_document`) back to a
-    ManagedEntry object, parsing the stringified value field and preserving all metadata.
+    ManagedEntry object. It supports both native BSON storage (dict in value.dict field) and legacy
+    JSON string storage (string in value.string field) for migration support.
 
     Args:
         document: The MongoDB document to convert.
 
     Returns:
         A ManagedEntry object reconstructed from the document.
     """
+    # Check if we have the new two-field structure
+    value_field = document.get("value")
+
+    if isinstance(value_field, dict):
+        # New format: check for dict or string subfields
+        if "dict" in value_field:
+            # Native storage mode - value is already a dict
+            return ManagedEntry.from_dict(
+                data={"value": value_field["dict"], **{k: v for k, v in document.items() if k != "value"}}, stringified_value=False
+            )
+        if "string" in value_field:
+            # Legacy storage mode - value is a JSON string
+            return ManagedEntry.from_dict(
+                data={"value": value_field["string"], **{k: v for k, v in document.items() if k != "value"}}, stringified_value=True
+            )
+
+    # Old format: value field is directly a string
     return ManagedEntry.from_dict(data=document, stringified_value=True)
 
 
-def managed_entry_to_document(key: str, managed_entry: ManagedEntry) -> dict[str, Any]:
+def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_storage: bool = True) -> dict[str, Any]:
     """Convert a ManagedEntry to a MongoDB document for storage.
 
     This function serializes a ManagedEntry to a MongoDB document format, including the key and all
-    metadata (TTL, creation, and expiration timestamps). The value is stringified to ensure proper
-    storage in MongoDB. The serialization is designed to preserve all entry information for round-trip
-    conversion back to a ManagedEntry.
+    metadata (TTL, creation, and expiration timestamps). The value storage format depends on the
+    native_storage parameter.
 
     Args:
         key: The key associated with this entry.
         managed_entry: The ManagedEntry to serialize.
+        native_storage: If True (default), store value as native BSON dict in value.dict field.
+                       If False, store as JSON string in value.string field for backward compatibility.
 
     Returns:
         A MongoDB document dict containing the key, value, and all metadata.
     """
-    return {
-        "key": key,
-        **managed_entry.to_dict(include_metadata=True, include_expiration=True, include_creation=True, stringify_value=True),
-    }
+    document: dict[str, Any] = {"key": key, "value": {}}
+
+    # Store in appropriate field based on mode
+    if native_storage:
+        document["value"]["dict"] = managed_entry.value_as_dict
+    else:
+        document["value"]["string"] = managed_entry.value_as_json
+
+    # Add metadata fields
+    if managed_entry.created_at:
+        document["created_at"] = managed_entry.created_at
+    if managed_entry.expires_at:
+        document["expires_at"] = managed_entry.expires_at
+
+    return document
 
 
 class MongoDBStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore):
@@ -75,6 +105,7 @@ class MongoDBStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, Ba
     _client: AsyncMongoClient[dict[str, Any]]
     _db: AsyncDatabase[dict[str, Any]]
     _collections_by_name: dict[str, AsyncCollection[dict[str, Any]]]
+    _native_storage: bool
 
     @overload
     def __init__(
@@ -83,6 +114,7 @@ def __init__(
         client: AsyncMongoClient[dict[str, Any]],
         db_name: str | None = None,
         coll_name: str | None = None,
+        native_storage: bool = True,
         default_collection: str | None = None,
     ) -> None:
         """Initialize the MongoDB store.
@@ -91,19 +123,27 @@ def __init__(
             client: The MongoDB client to use.
             db_name: The name of the MongoDB database.
             coll_name: The name of the MongoDB collection.
+            native_storage: Whether to use native BSON storage (True, default) or JSON string storage (False).
             default_collection: The default collection to use if no collection is provided.
         """
 
     @overload
     def __init__(
-        self, *, url: str, db_name: str | None = None, coll_name: str | None = None, default_collection: str | None = None
+        self,
+        *,
+        url: str,
+        db_name: str | None = None,
+        coll_name: str | None = None,
+        native_storage: bool = True,
+        default_collection: str | None = None,
     ) -> None:
         """Initialize the MongoDB store.
 
         Args:
             url: The url of the MongoDB cluster.
             db_name: The name of the MongoDB database.
             coll_name: The name of the MongoDB collection.
+            native_storage: Whether to use native BSON storage (True, default) or JSON string storage (False).
             default_collection: The default collection to use if no collection is provided.
         """
 
@@ -114,9 +154,21 @@ def __init__(
         url: str | None = None,
         db_name: str | None = None,
         coll_name: str | None = None,
+        native_storage: bool = True,
         default_collection: str | None = None,
     ) -> None:
-        """Initialize the MongoDB store."""
+        """Initialize the MongoDB store.
+
+        Args:
+            client: The MongoDB client to use (mutually exclusive with url).
+            url: The url of the MongoDB cluster (mutually exclusive with client).
+            db_name: The name of the MongoDB database.
+            coll_name: The name of the MongoDB collection.
+            native_storage: Whether to use native BSON storage (True, default) or JSON string storage (False).
+                           Native storage stores values as BSON dicts for better query support.
+                           Legacy mode stores values as JSON strings for backward compatibility.
+            default_collection: The default collection to use if no collection is provided.
+        """
 
         if client:
             self._client = client
@@ -131,6 +183,7 @@ def __init__(
 
         self._db = self._client[db_name]
         self._collections_by_name = {}
+        self._native_storage = native_storage
 
         super().__init__(default_collection=default_collection)
 
@@ -217,7 +270,7 @@ async def _put_managed_entry(
         collection: str,
         managed_entry: ManagedEntry,
     ) -> None:
-        mongo_doc: dict[str, Any] = managed_entry_to_document(key=key, managed_entry=managed_entry)
+        mongo_doc: dict[str, Any] = managed_entry_to_document(key=key, managed_entry=managed_entry, native_storage=self._native_storage)
 
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
@@ -248,7 +301,7 @@ async def _put_managed_entries(
 
         operations: list[UpdateOne] = []
         for key, managed_entry in zip(keys, managed_entries, strict=True):
-            mongo_doc: dict[str, Any] = managed_entry_to_document(key=key, managed_entry=managed_entry)
+            mongo_doc: dict[str, Any] = managed_entry_to_document(key=key, managed_entry=managed_entry, native_storage=self._native_storage)
 
             operations.append(
                 UpdateOne(

key-value/key-value-aio/tests/stores/mongodb/test_mongodb.py

@@ -44,19 +44,19 @@ class MongoDBFailedToStartError(Exception):
     pass
 
 
-def test_managed_entry_document_conversion():
+def test_managed_entry_document_conversion_native_mode():
     created_at = datetime(year=2025, month=1, day=1, hour=0, minute=0, second=0, tzinfo=timezone.utc)
     expires_at = created_at + timedelta(seconds=10)
 
     managed_entry = ManagedEntry(value={"test": "test"}, created_at=created_at, expires_at=expires_at)
-    document = managed_entry_to_document(key="test", managed_entry=managed_entry)
+    document = managed_entry_to_document(key="test", managed_entry=managed_entry, native_storage=True)
 
     assert document == snapshot(
         {
             "key": "test",
-            "value": '{"test": "test"}',
-            "created_at": "2025-01-01T00:00:00+00:00",
-            "expires_at": "2025-01-01T00:00:10+00:00",
+            "value": {"dict": {"test": "test"}},
+            "created_at": datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc),
+            "expires_at": datetime(2025, 1, 1, 0, 0, 10, tzinfo=timezone.utc),
         }
     )
 
@@ -68,8 +68,55 @@ def test_managed_entry_document_conversion():
     assert round_trip_managed_entry.expires_at == expires_at
 
 
+def test_managed_entry_document_conversion_legacy_mode():
+    created_at = datetime(year=2025, month=1, day=1, hour=0, minute=0, second=0, tzinfo=timezone.utc)
+    expires_at = created_at + timedelta(seconds=10)
+
+    managed_entry = ManagedEntry(value={"test": "test"}, created_at=created_at, expires_at=expires_at)
+    document = managed_entry_to_document(key="test", managed_entry=managed_entry, native_storage=False)
+
+    assert document == snapshot(
+        {
+            "key": "test",
+            "value": {"string": '{"test": "test"}'},
+            "created_at": datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc),
+            "expires_at": datetime(2025, 1, 1, 0, 0, 10, tzinfo=timezone.utc),
+        }
+    )
+
+    round_trip_managed_entry = document_to_managed_entry(document=document)
+
+    assert round_trip_managed_entry.value == managed_entry.value
+    assert round_trip_managed_entry.created_at == created_at
+    assert round_trip_managed_entry.ttl == IsFloat(lt=0)
+    assert round_trip_managed_entry.expires_at == expires_at
+
+
+def test_managed_entry_document_conversion_old_format():
+    """Test backward compatibility with old format where value is directly a string."""
+    created_at = datetime(year=2025, month=1, day=1, hour=0, minute=0, second=0, tzinfo=timezone.utc)
+    expires_at = created_at + timedelta(seconds=10)
+
+    # Simulate old document format
+    old_document = {
+        "key": "test",
+        "value": '{"test": "test"}',
+        "created_at": "2025-01-01T00:00:00+00:00",
+        "expires_at": "2025-01-01T00:00:10+00:00",
+    }
+
+    round_trip_managed_entry = document_to_managed_entry(document=old_document)
+
+    assert round_trip_managed_entry.value == {"test": "test"}
+    assert round_trip_managed_entry.created_at == created_at
+    assert round_trip_managed_entry.ttl == IsFloat(lt=0)
+    assert round_trip_managed_entry.expires_at == expires_at
+
+
 @pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not available")
 class TestMongoDBStore(ContextManagerStoreTestMixin, BaseStoreTests):
+    """Test MongoDBStore with native_storage=False (legacy mode) for backward compatibility."""
+
     @pytest.fixture(autouse=True, scope="session", params=MONGODB_VERSIONS_TO_TEST)
     async def setup_mongodb(self, request: pytest.FixtureRequest) -> AsyncGenerator[None, None]:
         version = request.param
@@ -84,7 +131,8 @@ async def setup_mongodb(self, request: pytest.FixtureRequest) -> AsyncGenerator[
     @override
     @pytest.fixture
     async def store(self, setup_mongodb: None) -> MongoDBStore:
-        store = MongoDBStore(url=f"mongodb://{MONGODB_HOST}:{MONGODB_HOST_PORT}", db_name=MONGODB_TEST_DB)
+        # Use legacy mode (native_storage=False) to test backward compatibility
+        store = MongoDBStore(url=f"mongodb://{MONGODB_HOST}:{MONGODB_HOST_PORT}", db_name=MONGODB_TEST_DB, native_storage=False)
         # Ensure a clean db by dropping our default test collection if it exists
         with contextlib.suppress(Exception):
             _ = await store._client.drop_database(name_or_database=MONGODB_TEST_DB)  # pyright: ignore[reportPrivateUsage]
@@ -106,3 +154,89 @@ async def test_mongodb_collection_name_sanitization(self, mongodb_store: MongoDB
 
         collections = await mongodb_store.collections()
         assert collections == snapshot(["test_collection_-daf4a2ec"])
+
+
+@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not available")
+class TestMongoDBStoreNativeMode(ContextManagerStoreTestMixin, BaseStoreTests):
+    """Test MongoDBStore with native_storage=True (default)."""
+
+    @pytest.fixture(autouse=True, scope="session", params=MONGODB_VERSIONS_TO_TEST)
+    async def setup_mongodb(self, request: pytest.FixtureRequest) -> AsyncGenerator[None, None]:
+        version = request.param
+
+        with docker_container(f"mongodb-test-native-{version}", f"mongo:{version}", {str(MONGODB_HOST_PORT): MONGODB_HOST_PORT}):
+            if not await async_wait_for_true(bool_fn=ping_mongodb, tries=WAIT_FOR_MONGODB_TIMEOUT, wait_time=1):
+                msg = f"MongoDB {version} failed to start"
+                raise MongoDBFailedToStartError(msg)
+
+            yield
+
+    @override
+    @pytest.fixture
+    async def store(self, setup_mongodb: None) -> MongoDBStore:
+        store = MongoDBStore(url=f"mongodb://{MONGODB_HOST}:{MONGODB_HOST_PORT}", db_name=f"{MONGODB_TEST_DB}-native", native_storage=True)
+        # Ensure a clean db by dropping our default test collection if it exists
+        with contextlib.suppress(Exception):
+            _ = await store._client.drop_database(name_or_database=f"{MONGODB_TEST_DB}-native")  # pyright: ignore[reportPrivateUsage]
+
+        return store
+
+    @pytest.fixture
+    async def mongodb_store(self, store: MongoDBStore) -> MongoDBStore:
+        return store
+
+    @pytest.mark.skip(reason="Distributed Caches are unbounded")
+    @override
+    async def test_not_unbounded(self, store: BaseStore): ...
+
+    async def test_value_stored_as_bson_dict(self, mongodb_store: MongoDBStore):
+        """Verify values are stored as BSON dicts, not JSON strings."""
+        await mongodb_store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
+
+        # Get the raw MongoDB document
+        await mongodb_store._setup_collection(collection="test")  # pyright: ignore[reportPrivateUsage]
+        sanitized_collection = mongodb_store._sanitize_collection_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        collection = mongodb_store._collections_by_name[sanitized_collection]  # pyright: ignore[reportPrivateUsage]
+        doc = await collection.find_one({"key": "test_key"})
+
+        # In native mode, value should be a dict with "dict" subfield
+        assert doc is not None
+        assert isinstance(doc["value"], dict)
+        assert "dict" in doc["value"]
+        assert doc["value"]["dict"] == {"name": "Alice", "age": 30}
+
+    async def test_migration_from_legacy_mode(self, mongodb_store: MongoDBStore):
+        """Verify native mode can read legacy JSON string data."""
+        # Manually insert a legacy document with JSON string value in the new format
+        await mongodb_store._setup_collection(collection="test")  # pyright: ignore[reportPrivateUsage]
+        sanitized_collection = mongodb_store._sanitize_collection_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        collection = mongodb_store._collections_by_name[sanitized_collection]  # pyright: ignore[reportPrivateUsage]
+
+        await collection.insert_one(
+            {
+                "key": "legacy_key",
+                "value": {"string": '{"legacy": "data"}'},  # New format with JSON string
+            }
+        )
+
+        # Should be able to read it in native mode
+        result = await mongodb_store.get(collection="test", key="legacy_key")
+        assert result == {"legacy": "data"}
+
+    async def test_migration_from_old_format(self, mongodb_store: MongoDBStore):
+        """Verify native mode can read old format where value is directly a string."""
+        # Manually insert an old document with value directly as JSON string
+        await mongodb_store._setup_collection(collection="test")  # pyright: ignore[reportPrivateUsage]
+        sanitized_collection = mongodb_store._sanitize_collection_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        collection = mongodb_store._collections_by_name[sanitized_collection]  # pyright: ignore[reportPrivateUsage]
+
+        await collection.insert_one(
+            {
+                "key": "old_key",
+                "value": '{"old": "format"}',  # Old format: value directly as JSON string
+            }
+        )
+
+        # Should be able to read it in native mode
+        result = await mongodb_store.get(collection="test", key="old_key")
+        assert result == {"old": "format"}