Skip to content

Commit 9ef3a28

Browse files
strawgatestrawgate
committed
Mongo PR Cleanup
1 parent c612c08 commit 9ef3a28

File tree

5 files changed

+106
-90
lines changed

5 files changed

+106
-90
lines changed

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

Lines changed: 41 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@
22
from datetime import datetime
33
from typing import Any, overload
44

5-
from key_value.shared.utils.managed_entry import ManagedEntry
5+
from key_value.shared.errors import DeserializationError
6+
from key_value.shared.utils.managed_entry import ManagedEntry, verify_dict
67
from key_value.shared.utils.sanitize import ALPHANUMERIC_CHARACTERS, sanitize_string
8+
from key_value.shared.utils.time_to_live import timezone
79
from typing_extensions import Self, override
810

911
from key_value.aio.stores.base import BaseContextManagerStore, BaseDestroyCollectionStore, BaseEnumerateCollectionsStore, BaseStore
@@ -46,24 +48,42 @@ def document_to_managed_entry(document: dict[str, Any]) -> ManagedEntry:
4648
Returns:
4749
A ManagedEntry object reconstructed from the document.
4850
"""
49-
# Check if we have the new two-field structure
50-
value_field = document.get("value")
51-
52-
if isinstance(value_field, dict):
53-
# New format: check for dict or string subfields
54-
if "dict" in value_field:
55-
# Native storage mode - value is already a dict
56-
return ManagedEntry.from_dict(
57-
data={"value": value_field["dict"], **{k: v for k, v in document.items() if k != "value"}}, stringified_value=False
58-
)
59-
if "string" in value_field:
60-
# Legacy storage mode - value is a JSON string
61-
return ManagedEntry.from_dict(
62-
data={"value": value_field["string"], **{k: v for k, v in document.items() if k != "value"}}, stringified_value=True
63-
)
51+
if not (value_field := document.get("value")):
52+
msg = "Value field not found"
53+
raise DeserializationError(msg)
54+
55+
if not isinstance(value_field, dict):
56+
msg = "Expected `value` field to be an object"
57+
raise DeserializationError(msg)
58+
59+
value_holder: dict[str, Any] = verify_dict(obj=value_field)
60+
61+
data: dict[str, Any] = {}
62+
63+
# The Value field is an object with two possible fields: `object` and `string`
64+
# - `object`: The value is a native BSON dict
65+
# - `string`: The value is a JSON string
66+
# Mongo stores datetimes without timezones as UTC so we mark them as UTC
67+
68+
if created_at_datetime := document.get("created_at"):
69+
if not isinstance(created_at_datetime, datetime):
70+
msg = "Expected `created_at` field to be a datetime"
71+
raise DeserializationError(msg)
72+
data["created_at"] = created_at_datetime.replace(tzinfo=timezone.utc)
73+
if expires_at_datetime := document.get("expires_at"):
74+
if not isinstance(expires_at_datetime, datetime):
75+
msg = "Expected `expires_at` field to be a datetime"
76+
raise DeserializationError(msg)
77+
data["expires_at"] = expires_at_datetime.replace(tzinfo=timezone.utc)
78+
79+
if value_object := value_holder.get("object"):
80+
return ManagedEntry.from_dict(data={"value": value_object, **data})
81+
82+
if value_string := value_holder.get("string"):
83+
return ManagedEntry.from_dict(data={"value": value_string, **data}, stringified_value=True)
6484

65-
# Old format: value field is directly a string
66-
return ManagedEntry.from_dict(data=document, stringified_value=True)
85+
msg = "Expected `value` field to be an object with `object` or `string` subfield"
86+
raise DeserializationError(msg)
6787

6888

6989
def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_storage: bool = True) -> dict[str, Any]:
@@ -76,7 +96,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
7696
Args:
7797
key: The key associated with this entry.
7898
managed_entry: The ManagedEntry to serialize.
79-
native_storage: If True (default), store value as native BSON dict in value.dict field.
99+
native_storage: If True (default), store value as native BSON dict in value.object field.
80100
If False, store as JSON string in value.string field for backward compatibility.
81101
82102
Returns:
@@ -86,7 +106,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
86106

87107
# Store in appropriate field based on mode
88108
if native_storage:
89-
document["value"]["dict"] = managed_entry.value_as_dict
109+
document["value"]["object"] = managed_entry.value_as_dict
90110
else:
91111
document["value"]["string"] = managed_entry.value_as_json
92112

@@ -240,7 +260,7 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
240260
sanitized_collection = self._sanitize_collection_name(collection=collection)
241261

242262
if doc := await self._collections_by_name[sanitized_collection].find_one(filter={"key": key}):
243-
return ManagedEntry.from_dict(data=doc, stringified_value=True)
263+
return document_to_managed_entry(document=doc)
244264

245265
return None
246266

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

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def test_managed_entry_document_conversion_native_mode():
5454
assert document == snapshot(
5555
{
5656
"key": "test",
57-
"value": {"dict": {"test": "test"}},
57+
"value": {"object": {"test": "test"}},
5858
"created_at": datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc),
5959
"expires_at": datetime(2025, 1, 1, 0, 0, 10, tzinfo=timezone.utc),
6060
}
@@ -92,27 +92,6 @@ def test_managed_entry_document_conversion_legacy_mode():
9292
assert round_trip_managed_entry.expires_at == expires_at
9393

9494

95-
def test_managed_entry_document_conversion_old_format():
96-
"""Test backward compatibility with old format where value is directly a string."""
97-
created_at = datetime(year=2025, month=1, day=1, hour=0, minute=0, second=0, tzinfo=timezone.utc)
98-
expires_at = created_at + timedelta(seconds=10)
99-
100-
# Simulate old document format
101-
old_document = {
102-
"key": "test",
103-
"value": '{"test": "test"}',
104-
"created_at": "2025-01-01T00:00:00+00:00",
105-
"expires_at": "2025-01-01T00:00:10+00:00",
106-
}
107-
108-
round_trip_managed_entry = document_to_managed_entry(document=old_document)
109-
110-
assert round_trip_managed_entry.value == {"test": "test"}
111-
assert round_trip_managed_entry.created_at == created_at
112-
assert round_trip_managed_entry.ttl == IsFloat(lt=0)
113-
assert round_trip_managed_entry.expires_at == expires_at
114-
115-
11695
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not available")
11796
class TestMongoDBStore(ContextManagerStoreTestMixin, BaseStoreTests):
11897
"""Test MongoDBStore with native_storage=False (legacy mode) for backward compatibility."""

‎key-value/key-value-shared/src/key_value/shared/utils/managed_entry.py‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,16 +79,34 @@ def to_json(
7979
)
8080

8181
@classmethod
82-
def from_dict(
82+
def from_dict( # noqa: PLR0912
8383
cls, data: dict[str, Any], includes_metadata: bool = True, ttl: SupportsFloat | None = None, stringified_value: bool = False
8484
) -> Self:
8585
if not includes_metadata:
8686
return cls(
8787
value=data,
8888
)
8989

90-
created_at: datetime | None = try_parse_datetime_str(value=data.get("created_at"))
91-
expires_at: datetime | None = try_parse_datetime_str(value=data.get("expires_at"))
90+
created_at: datetime | None = None
91+
expires_at: datetime | None = None
92+
93+
if created_at_value := data.get("created_at"):
94+
if isinstance(created_at_value, str):
95+
created_at = try_parse_datetime_str(value=created_at_value)
96+
elif isinstance(created_at_value, datetime):
97+
created_at = created_at_value
98+
else:
99+
msg = "Expected `created_at` field to be a string or datetime"
100+
raise DeserializationError(msg)
101+
102+
if expires_at_value := data.get("expires_at"):
103+
if isinstance(expires_at_value, str):
104+
expires_at = try_parse_datetime_str(value=expires_at_value)
105+
elif isinstance(expires_at_value, datetime):
106+
expires_at = expires_at_value
107+
else:
108+
msg = "Expected `expires_at` field to be a string or datetime"
109+
raise DeserializationError(msg)
92110

93111
if not (raw_value := data.get("value")):
94112
msg = "Value is None"

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

Lines changed: 42 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@
55
from datetime import datetime
66
from typing import Any, overload
77

8-
from key_value.shared.utils.managed_entry import ManagedEntry
8+
from key_value.shared.errors import DeserializationError
9+
from key_value.shared.utils.managed_entry import ManagedEntry, verify_dict
910
from key_value.shared.utils.sanitize import ALPHANUMERIC_CHARACTERS, sanitize_string
11+
from key_value.shared.utils.time_to_live import timezone
1012
from typing_extensions import Self, override
1113

1214
from key_value.sync.code_gen.stores.base import (
@@ -53,24 +55,42 @@ def document_to_managed_entry(document: dict[str, Any]) -> ManagedEntry:
5355
Returns:
5456
A ManagedEntry object reconstructed from the document.
5557
"""
56-
# Check if we have the new two-field structure
57-
value_field = document.get("value")
58-
59-
if isinstance(value_field, dict):
60-
# New format: check for dict or string subfields
61-
if "dict" in value_field:
62-
# Native storage mode - value is already a dict
63-
return ManagedEntry.from_dict(
64-
data={"value": value_field["dict"], **{k: v for (k, v) in document.items() if k != "value"}}, stringified_value=False
65-
)
66-
if "string" in value_field:
67-
# Legacy storage mode - value is a JSON string
68-
return ManagedEntry.from_dict(
69-
data={"value": value_field["string"], **{k: v for (k, v) in document.items() if k != "value"}}, stringified_value=True
70-
)
71-
72-
# Old format: value field is directly a string
73-
return ManagedEntry.from_dict(data=document, stringified_value=True)
58+
if not (value_field := document.get("value")):
59+
msg = "Value field not found"
60+
raise DeserializationError(msg)
61+
62+
if not isinstance(value_field, dict):
63+
msg = "Expected `value` field to be an object"
64+
raise DeserializationError(msg)
65+
66+
value_holder: dict[str, Any] = verify_dict(obj=value_field)
67+
68+
data: dict[str, Any] = {}
69+
70+
# The Value field is an object with two possible fields: `object` and `string`
71+
# - `object`: The value is a native BSON dict
72+
# - `string`: The value is a JSON string
73+
# Mongo stores datetimes without timezones as UTC so we mark them as UTC
74+
75+
if created_at_datetime := document.get("created_at"):
76+
if not isinstance(created_at_datetime, datetime):
77+
msg = "Expected `created_at` field to be a datetime"
78+
raise DeserializationError(msg)
79+
data["created_at"] = created_at_datetime.replace(tzinfo=timezone.utc)
80+
if expires_at_datetime := document.get("expires_at"):
81+
if not isinstance(expires_at_datetime, datetime):
82+
msg = "Expected `expires_at` field to be a datetime"
83+
raise DeserializationError(msg)
84+
data["expires_at"] = expires_at_datetime.replace(tzinfo=timezone.utc)
85+
86+
if value_object := value_holder.get("object"):
87+
return ManagedEntry.from_dict(data={"value": value_object, **data})
88+
89+
if value_string := value_holder.get("string"):
90+
return ManagedEntry.from_dict(data={"value": value_string, **data}, stringified_value=True)
91+
92+
msg = "Expected `value` field to be an object with `object` or `string` subfield"
93+
raise DeserializationError(msg)
7494

7595

7696
def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_storage: bool = True) -> dict[str, Any]:
@@ -83,7 +103,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
83103
Args:
84104
key: The key associated with this entry.
85105
managed_entry: The ManagedEntry to serialize.
86-
native_storage: If True (default), store value as native BSON dict in value.dict field.
106+
native_storage: If True (default), store value as native BSON dict in value.object field.
87107
If False, store as JSON string in value.string field for backward compatibility.
88108
89109
Returns:
@@ -93,7 +113,7 @@ def managed_entry_to_document(key: str, managed_entry: ManagedEntry, *, native_s
93113

94114
# Store in appropriate field based on mode
95115
if native_storage:
96-
document["value"]["dict"] = managed_entry.value_as_dict
116+
document["value"]["object"] = managed_entry.value_as_dict
97117
else:
98118
document["value"]["string"] = managed_entry.value_as_json
99119

@@ -247,7 +267,7 @@ def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | Non
247267
sanitized_collection = self._sanitize_collection_name(collection=collection)
248268

249269
if doc := self._collections_by_name[sanitized_collection].find_one(filter={"key": key}):
250-
return ManagedEntry.from_dict(data=doc, stringified_value=True)
270+
return document_to_managed_entry(document=doc)
251271

252272
return None
253273

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

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ def test_managed_entry_document_conversion_native_mode():
5555
assert document == snapshot(
5656
{
5757
"key": "test",
58-
"value": {"dict": {"test": "test"}},
58+
"value": {"object": {"test": "test"}},
5959
"created_at": datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc),
6060
"expires_at": datetime(2025, 1, 1, 0, 0, 10, tzinfo=timezone.utc),
6161
}
@@ -93,27 +93,6 @@ def test_managed_entry_document_conversion_legacy_mode():
9393
assert round_trip_managed_entry.expires_at == expires_at
9494

9595

96-
def test_managed_entry_document_conversion_old_format():
97-
"""Test backward compatibility with old format where value is directly a string."""
98-
created_at = datetime(year=2025, month=1, day=1, hour=0, minute=0, second=0, tzinfo=timezone.utc)
99-
expires_at = created_at + timedelta(seconds=10)
100-
101-
# Simulate old document format
102-
old_document = {
103-
"key": "test",
104-
"value": '{"test": "test"}',
105-
"created_at": "2025-01-01T00:00:00+00:00",
106-
"expires_at": "2025-01-01T00:00:10+00:00",
107-
}
108-
109-
round_trip_managed_entry = document_to_managed_entry(document=old_document)
110-
111-
assert round_trip_managed_entry.value == {"test": "test"}
112-
assert round_trip_managed_entry.created_at == created_at
113-
assert round_trip_managed_entry.ttl == IsFloat(lt=0)
114-
assert round_trip_managed_entry.expires_at == expires_at
115-
116-
11796
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not available")
11897
class TestMongoDBStore(ContextManagerStoreTestMixin, BaseStoreTests):
11998
"""Test MongoDBStore with native_storage=False (legacy mode) for backward compatibility."""

0 commit comments

Comments
 (0)