Update Mongo storage

Author: strawgate

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

@@ -127,7 +127,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
 
 
 class MongoDBStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore):
-    """MongoDB-based key-value store using Motor (async MongoDB driver)."""
+    """MongoDB-based key-value store using pymongo."""
 
     _client: AsyncMongoClient[dict[str, Any]]
     _db: AsyncDatabase[dict[str, Any]]

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

@@ -4,11 +4,11 @@
 from typing import Any
 
 import pytest
+from bson import ObjectId
 from dirty_equals import IsDatetime, IsFloat, IsInstance
 from inline_snapshot import snapshot
 from key_value.shared.stores.wait import async_wait_for_true
 from key_value.shared.utils.managed_entry import ManagedEntry
-from bson import ObjectId
 from pymongo import AsyncMongoClient
 from typing_extensions import override
 
@@ -133,9 +133,8 @@ class TestMongoDBStoreNativeMode(BaseMongoDBStoreTests):
     @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]
+
+        await clean_mongodb_database(store=store)
 
         return store
 
@@ -172,7 +171,6 @@ async def test_migration_from_legacy_mode(self, store: MongoDBStore):
             }
         )
 
-        # Should be able to read it in native mode
         result = await store.get(collection="test", key="legacy_key")
         assert result == {"legacy": "data"}
 
@@ -210,7 +208,7 @@ async def test_value_stored_as_json(self, store: MongoDBStore):
         )
 
     async def test_migration_from_native_mode(self, store: MongoDBStore):
-        """Verify legacy mode can read native-mode object data."""
+        """Verify non-native mode can read native mode data."""
         # Manually insert a legacy document with JSON string value in the new format
         await store._setup_collection(collection="test")  # pyright: ignore[reportPrivateUsage]
         sanitized_collection = store._sanitize_collection_name(collection="test")  # pyright: ignore[reportPrivateUsage]
@@ -223,6 +221,5 @@ async def test_migration_from_native_mode(self, store: MongoDBStore):
             }
         )
 
-        # Should be able to read it in native mode
         result = await store.get(collection="test", key="legacy_key")
         assert result == {"name": "Alice", "age": 30}

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

@@ -24,7 +24,7 @@
     from pymongo.database import Database
     from pymongo.results import DeleteResult  # noqa: TC002
 except ImportError as e:
-    msg = "MongoDBStore requires py-key-value-sync[mongodb]"
+    msg = "MongoDBStore requires py-key-value-aio[mongodb]"
     raise ImportError(msg) from e
 
 DEFAULT_DB = "kv-store-adapter"
@@ -134,7 +134,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
 
 
 class MongoDBStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore):
-    """MongoDB-based key-value store using PyMongo (sync MongoDB driver)."""
+    """MongoDB-based key-value store using pymongo."""
 
     _client: MongoClient[dict[str, Any]]
     _db: Database[dict[str, Any]]

key-value/key-value-sync/src/key_value/sync/code_gen/test.py

@@ -0,0 +1,18 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file 'test.py'
+# DO NOT CHANGE! Change the original file instead.
+import os
+
+from cryptography.fernet import Fernet
+from key_value.sync.code_gen.stores.dynamodb.store import DynamoDBStore
+
+from key_value.sync.code_gen.wrappers.encryption.fernet import FernetEncryptionWrapper
+
+PERSISTENCE_TABLE_NAME = os.getenv("PERSISTENCE_TABLE_NAME", "test-key-value")
+STORAGE_ENCRYPTION_KEY = os.getenv("STORAGE_ENCRYPTION_KEY")
+
+wrapped_client_storage: DynamoDBStore
+
+client_storage = DynamoDBStore(table_name=PERSISTENCE_TABLE_NAME)
+if STORAGE_ENCRYPTION_KEY:
+    wrapped_client_storage = FernetEncryptionWrapper(key_value=client_storage, fernet=Fernet(STORAGE_ENCRYPTION_KEY))

key-value/key-value-sync/tests/code_gen/stores/mongodb/test_mongodb.py

@@ -7,11 +7,11 @@
 from typing import Any
 
 import pytest
+from bson import ObjectId
 from dirty_equals import IsDatetime, IsFloat, IsInstance
 from inline_snapshot import snapshot
 from key_value.shared.stores.wait import wait_for_true
 from key_value.shared.utils.managed_entry import ManagedEntry
-from bson import ObjectId
 from pymongo import MongoClient
 from typing_extensions import override
 
@@ -134,9 +134,8 @@ class TestMongoDBStoreNativeMode(BaseMongoDBStoreTests):
     @pytest.fixture
     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):
-            _ = store._client.drop_database(name_or_database=f"{MONGODB_TEST_DB}-native")  # pyright: ignore[reportPrivateUsage]
+
+        clean_mongodb_database(store=store)
 
         return store
 
@@ -169,7 +168,6 @@ def test_migration_from_legacy_mode(self, store: MongoDBStore):
         # New format with JSON string
         collection.insert_one({"key": "legacy_key", "value": {"string": '{"legacy": "data"}'}})
 
-        # Should be able to read it in native mode
         result = store.get(collection="test", key="legacy_key")
         assert result == {"legacy": "data"}
 
@@ -207,14 +205,13 @@ def test_value_stored_as_json(self, store: MongoDBStore):
         )
 
     def test_migration_from_native_mode(self, store: MongoDBStore):
-        """Verify legacy mode can read native-mode object data."""
+        """Verify non-native mode can read native mode data."""
         # Manually insert a legacy document with JSON string value in the new format
         store._setup_collection(collection="test")  # pyright: ignore[reportPrivateUsage]
         sanitized_collection = store._sanitize_collection_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         collection = store._collections_by_name[sanitized_collection]  # pyright: ignore[reportPrivateUsage]
 
         collection.insert_one({"key": "legacy_key", "value": {"object": {"name": "Alice", "age": 30}}})
 
-        # Should be able to read it in native mode
         result = store.get(collection="test", key="legacy_key")
         assert result == {"name": "Alice", "age": 30}