PR Feedback

Author: strawgate

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

@@ -23,7 +23,7 @@
     BaseEnumerateKeysStore,
     BaseStore,
 )
-from key_value.aio.stores.elasticsearch.utils import LessCapableJsonSerializer, NdjsonSerializer, new_bulk_action
+from key_value.aio.stores.elasticsearch.utils import LessCapableJsonSerializer, LessCapableNdjsonSerializer, new_bulk_action
 
 try:
     from elasticsearch import AsyncElasticsearch
@@ -99,14 +99,16 @@ def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedE
 def source_to_managed_entry(source: dict[str, Any]) -> ManagedEntry:
     value: dict[str, Any] = {}
 
+    raw_value = source.get("value")
+
     # Try flattened field first, fall back to string field
-    if not (value := source.get("value")) or not isinstance(value, dict):
+    if not raw_value or not isinstance(raw_value, dict):
         msg = "Value field not found or invalid type"
         raise DeserializationError(msg)
 
-    if value_flattened := value.get("flattened"):
+    if value_flattened := raw_value.get("flattened"):  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
         value = verify_dict(obj=value_flattened)
-    elif value_str := value.get("string"):
+    elif value_str := raw_value.get("string"):  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
         if not isinstance(value_str, str):
             msg = "Value in `value` field is not a string"
             raise DeserializationError(msg)
@@ -196,7 +198,7 @@ def __init__(
 
         LessCapableJsonSerializer.install_serializer(client=self._client)
         LessCapableJsonSerializer.install_default_serializer(client=self._client)
-        NdjsonSerializer.install_serializer(client=self._client)
+        LessCapableNdjsonSerializer.install_serializer(client=self._client)
 
         self._index_prefix = index_prefix
         self._native_storage = native_storage

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

@@ -1,6 +1,11 @@
 from typing import Any, ClassVar, TypeVar, cast
 
-from elastic_transport import JsonSerializer, NdjsonSerializer, ObjectApiResponse, SerializationError
+from elastic_transport import (
+    JsonSerializer,
+    NdjsonSerializer,
+    ObjectApiResponse,
+    SerializationError,
+)
 
 from elasticsearch import AsyncElasticsearch
 
@@ -29,7 +34,10 @@ def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
     if not (aggregations := body.get("aggregations")):
         return {}
 
-    if not isinstance(aggregations, dict) or not all(isinstance(key, str) for key in aggregations):  # pyright: ignore[reportUnknownVariableType]
+    if not isinstance(aggregations, dict) or not all(
+        isinstance(key, str)
+        for key in aggregations  # pyright: ignore[reportUnknownVariableType]
+    ):  # pyright: ignore[reportUnknownVariableType]
         return {}
 
     return cast("dict[str, Any]", aggregations)
@@ -113,44 +121,46 @@ def new_bulk_action(action: str, index: str, document_id: str) -> dict[str, Any]
     return {action: {"_index": index, "_id": document_id}}
 
 
-class InstallSerializerMixin:
-    """A mixin that installs the serializer into the transport."""
+class LessCapableJsonSerializer(JsonSerializer):
+    """A JSON Serializer that doesnt try to be smart with datetime, floats, etc."""
 
-    mimetype: ClassVar[str]
-    compatibility_mimetype: ClassVar[str]
+    mimetype: ClassVar[str] = "application/json"
+    compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+json"
 
-    @classmethod
-    def install_serializer(cls, client: AsyncElasticsearch) -> None:
-        client.transport.serializers.serializers.update(
-            {
-                cls.mimetype: cls(),
-                cls.compatibility_mimetype: cls(),
-            }
+    def default(self, data: Any) -> Any:
+        raise SerializationError(
+            message=f"Unable to serialize to JSON: {data!r} (type: {type(data).__name__})",
         )
 
     @classmethod
     def install_default_serializer(cls, client: AsyncElasticsearch) -> None:
         cls.install_serializer(client=client)
         client.transport.serializers.default_serializer = cls()
 
-
-class LessCapableJsonSerializer(InstallSerializerMixin, JsonSerializer):
-    """A JSON Serializer that doesnt try to be smart with datetime, floats, etc."""
-
-    mimetype: ClassVar[str] = "application/json"
-    compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+json"
-
-    def default(self, data: Any) -> Any:
-        raise SerializationError(
-            message=f"Unable to serialize to JSON: {data!r} (type: {type(data).__name__})",
+    @classmethod
+    def install_serializer(cls, client: AsyncElasticsearch) -> None:
+        client.transport.serializers.serializers.update(
+            {
+                cls.mimetype: cls(),
+                cls.compatibility_mimetype: cls(),
+            }
         )
 
 
-class NdjsonSerializer(InstallSerializerMixin, NdjsonSerializer):
+class LessCapableNdjsonSerializer(NdjsonSerializer):
     """A NDJSON Serializer that doesnt try to be smart with datetime, floats, etc."""
 
     mimetype: ClassVar[str] = "application/x-ndjson"
     compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+x-ndjson"
 
     def default(self, data: Any) -> Any:
-        return LessCapableJsonSerializer.default(self, data)
+        return LessCapableJsonSerializer.default(data=data)  # pyright: ignore[reportCallIssue, reportUnknownVariableType]
+
+    @classmethod
+    def install_serializer(cls, client: AsyncElasticsearch) -> None:
+        client.transport.serializers.serializers.update(
+            {
+                cls.mimetype: cls(),
+                cls.compatibility_mimetype: cls(),
+            }
+        )

key-value/key-value-aio/tests/stores/elasticsearch/test_elasticsearch.py

@@ -185,9 +185,8 @@ async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore,
             }
         )
 
-    async def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
-        """Verify native mode can read legacy JSON string data"""
-        # Manually insert a legacy document with JSON string value
+    async def test_migration_from_non_native_mode(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
+        """Verify native mode can read a document with stringified data"""
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="legacy_key")  # pyright: ignore[reportPrivateUsage]
 
@@ -197,14 +196,15 @@ async def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_cl
             body={
                 "collection": "test",
                 "key": "legacy_key",
-                "value": '{"legacy": "data"}',  # JSON string
+                "value": {
+                    "string": '{"legacy": "data"}',
+                },
             },
         )
         await es_client.indices.refresh(index=index_name)
 
-        # Should be able to read stringified values too
         result = await store.get(collection="test", key="legacy_key")
-        assert result == snapshot(None)
+        assert result == snapshot({"legacy": "data"})
 
 
 class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):

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

@@ -21,7 +21,7 @@
     BaseEnumerateKeysStore,
     BaseStore,
 )
-from key_value.sync.code_gen.stores.elasticsearch.utils import LessCapableJsonSerializer, NdjsonSerializer, new_bulk_action
+from key_value.sync.code_gen.stores.elasticsearch.utils import LessCapableJsonSerializer, LessCapableNdjsonSerializer, new_bulk_action
 
 try:
     from elasticsearch import Elasticsearch
@@ -79,14 +79,16 @@ def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedE
 def source_to_managed_entry(source: dict[str, Any]) -> ManagedEntry:
     value: dict[str, Any] = {}
 
+    raw_value = source.get("value")
+
     # Try flattened field first, fall back to string field
-    if not (value := source.get("value")) or not isinstance(value, dict):
+    if not raw_value or not isinstance(raw_value, dict):
         msg = "Value field not found or invalid type"
         raise DeserializationError(msg)
 
-    if value_flattened := value.get("flattened"):
+    if value_flattened := raw_value.get("flattened"):  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
         value = verify_dict(obj=value_flattened)
-    elif value_str := value.get("string"):
+    elif value_str := raw_value.get("string"):  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
         if not isinstance(value_str, str):
             msg = "Value in `value` field is not a string"
             raise DeserializationError(msg)
@@ -161,7 +163,7 @@ def __init__(
 
         LessCapableJsonSerializer.install_serializer(client=self._client)
         LessCapableJsonSerializer.install_default_serializer(client=self._client)
-        NdjsonSerializer.install_serializer(client=self._client)
+        LessCapableNdjsonSerializer.install_serializer(client=self._client)
 
         self._index_prefix = index_prefix
         self._native_storage = native_storage

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

@@ -32,6 +32,7 @@ def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
         return {}
 
     if not isinstance(aggregations, dict) or not all(isinstance(key, str) for key in aggregations):  # pyright: ignore[reportUnknownVariableType]
+        # pyright: ignore[reportUnknownVariableType]
         return {}
 
     return cast("dict[str, Any]", aggregations)
@@ -115,37 +116,34 @@ def new_bulk_action(action: str, index: str, document_id: str) -> dict[str, Any]
     return {action: {"_index": index, "_id": document_id}}
 
 
-class InstallSerializerMixin:
-    """A mixin that installs the serializer into the transport."""
+class LessCapableJsonSerializer(JsonSerializer):
+    """A JSON Serializer that doesnt try to be smart with datetime, floats, etc."""
 
-    mimetype: ClassVar[str]
-    compatibility_mimetype: ClassVar[str]
+    mimetype: ClassVar[str] = "application/json"
+    compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+json"
 
-    @classmethod
-    def install_serializer(cls, client: Elasticsearch) -> None:
-        client.transport.serializers.serializers.update({cls.mimetype: cls(), cls.compatibility_mimetype: cls()})
+    def default(self, data: Any) -> Any:
+        raise SerializationError(message=f"Unable to serialize to JSON: {data!r} (type: {type(data).__name__})")
 
     @classmethod
     def install_default_serializer(cls, client: Elasticsearch) -> None:
         cls.install_serializer(client=client)
         client.transport.serializers.default_serializer = cls()
 
-
-class LessCapableJsonSerializer(InstallSerializerMixin, JsonSerializer):
-    """A JSON Serializer that doesnt try to be smart with datetime, floats, etc."""
-
-    mimetype: ClassVar[str] = "application/json"
-    compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+json"
-
-    def default(self, data: Any) -> Any:
-        raise SerializationError(message=f"Unable to serialize to JSON: {data!r} (type: {type(data).__name__})")
+    @classmethod
+    def install_serializer(cls, client: Elasticsearch) -> None:
+        client.transport.serializers.serializers.update({cls.mimetype: cls(), cls.compatibility_mimetype: cls()})
 
 
-class NdjsonSerializer(InstallSerializerMixin, NdjsonSerializer):
+class LessCapableNdjsonSerializer(NdjsonSerializer):
     """A NDJSON Serializer that doesnt try to be smart with datetime, floats, etc."""
 
     mimetype: ClassVar[str] = "application/x-ndjson"
     compatibility_mimetype: ClassVar[str] = "application/vnd.elasticsearch+x-ndjson"
 
     def default(self, data: Any) -> Any:
-        return LessCapableJsonSerializer.default(self, data)
+        return LessCapableJsonSerializer.default(data=data)  # pyright: ignore[reportCallIssue, reportUnknownVariableType]
+
+    @classmethod
+    def install_serializer(cls, client: Elasticsearch) -> None:
+        client.transport.serializers.serializers.update({cls.mimetype: cls(), cls.compatibility_mimetype: cls()})

key-value/key-value-sync/tests/code_gen/stores/elasticsearch/test_elasticsearch.py

@@ -186,19 +186,18 @@ def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_cl
             }
         )
 
-    def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_client: Elasticsearch):
-        """Verify native mode can read legacy JSON string data"""
-        # Manually insert a legacy document with JSON string value
+    def test_migration_from_non_native_mode(self, store: ElasticsearchStore, es_client: Elasticsearch):
+        """Verify native mode can read a document with stringified data"""
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="legacy_key")  # pyright: ignore[reportPrivateUsage]
 
-        # JSON string
-        es_client.index(index=index_name, id=doc_id, body={"collection": "test", "key": "legacy_key", "value": '{"legacy": "data"}'})
+        es_client.index(
+            index=index_name, id=doc_id, body={"collection": "test", "key": "legacy_key", "value": {"string": '{"legacy": "data"}'}}
+        )
         es_client.indices.refresh(index=index_name)
 
-        # Should be able to read stringified values too
         result = store.get(collection="test", key="legacy_key")
-        assert result == snapshot(None)
+        assert result == snapshot({"legacy": "data"})
 
 
 class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):

pyproject.toml

@@ -99,6 +99,9 @@ typeCheckingMode = "strict"
 reportExplicitAny = false
 reportMissingTypeStubs = false
 #reportUnnecessaryTypeIgnoreComment = "error"
+include = [
+    "key-value"
+]
 exclude = [
     "**/playground/**",
     "**/node_modules/**",