feat: add native storage mode for Elasticsearch using flattened fields

Add optional `native_storage` parameter to ElasticsearchStore that stores
values as Elasticsearch flattened objects instead of JSON strings.

Key changes:
- Add `native_storage` parameter to constructor (defaults to False)
- Update mapping to include `value_flattened` field (flattened type)
- Update `managed_entry_to_document` to store in appropriate field
- Update `source_to_managed_entry` to read from both fields for migration
- Add comprehensive tests for native storage mode and migration

Benefits:
- More efficient storage and retrieval
- Better aligned with Elasticsearch's document model
- Enables basic querying on nested fields

The dual-field approach (value + value_flattened) enables backward
compatibility and gradual migration from legacy JSON string mode.

Related to #87

Co-authored-by: William Easton <[email protected]>

Author: github-actions[bot]

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

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Any, overload
+from typing import Any, cast, overload
 
 from elastic_transport import ObjectApiResponse  # noqa: TC002
 from key_value.shared.errors import DeserializationError
@@ -60,6 +60,9 @@
             "doc_values": False,
             "ignore_above": 256,
         },
+        "value_flattened": {
+            "type": "flattened",
+        },
     },
 }
 
@@ -73,13 +76,18 @@
 ALLOWED_INDEX_CHARACTERS: str = LOWERCASE_ALPHABET + NUMBERS + "_" + "-" + "."
 
 
-def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedEntry) -> dict[str, Any]:
+def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedEntry, *, native_storage: bool = False) -> dict[str, Any]:
     document: dict[str, Any] = {
         "collection": collection,
         "key": key,
-        "value": managed_entry.to_json(include_metadata=False),
     }
 
+    # Store in appropriate field based on mode
+    if native_storage:
+        document["value_flattened"] = dict(managed_entry.value)
+    else:
+        document["value"] = managed_entry.to_json(include_metadata=False)
+
     if managed_entry.created_at:
         document["created_at"] = managed_entry.created_at.isoformat()
     if managed_entry.expires_at:
@@ -89,15 +97,26 @@ def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedE
 
 
 def source_to_managed_entry(source: dict[str, Any]) -> ManagedEntry:
-    if not (value_str := source.get("value")) or not isinstance(value_str, str):
-        msg = "Value is not a string"
+    # Try flattened field first, fall back to string field
+    value_flattened = source.get("value_flattened")
+    value_str = source.get("value")
+
+    value: dict[str, Any]
+    if value_flattened and isinstance(value_flattened, dict):
+        # Native storage mode - cast to the correct type
+        value = cast(dict[str, Any], value_flattened)
+    elif value_str and isinstance(value_str, str):
+        # Legacy JSON string mode
+        value = load_from_json(value_str)
+    else:
+        msg = "Value field not found or invalid type"
         raise DeserializationError(msg)
 
     created_at: datetime | None = try_parse_datetime_str(value=source.get("created_at"))
     expires_at: datetime | None = try_parse_datetime_str(value=source.get("expires_at"))
 
     return ManagedEntry(
-        value=load_from_json(value_str),
+        value=value,
         created_at=created_at,
         expires_at=expires_at,
     )
@@ -114,11 +133,28 @@ class ElasticsearchStore(
 
     _index_prefix: str
 
+    _native_storage: bool
+
     @overload
-    def __init__(self, *, elasticsearch_client: AsyncElasticsearch, index_prefix: str, default_collection: str | None = None) -> None: ...
+    def __init__(
+        self,
+        *,
+        elasticsearch_client: AsyncElasticsearch,
+        index_prefix: str,
+        native_storage: bool = False,
+        default_collection: str | None = None,
+    ) -> None: ...
 
     @overload
-    def __init__(self, *, url: str, api_key: str | None = None, index_prefix: str, default_collection: str | None = None) -> None: ...
+    def __init__(
+        self,
+        *,
+        url: str,
+        api_key: str | None = None,
+        index_prefix: str,
+        native_storage: bool = False,
+        default_collection: str | None = None,
+    ) -> None: ...
 
     def __init__(
         self,
@@ -127,6 +163,7 @@ def __init__(
         url: str | None = None,
         api_key: str | None = None,
         index_prefix: str,
+        native_storage: bool = False,
         default_collection: str | None = None,
     ) -> None:
         """Initialize the elasticsearch store.
@@ -136,6 +173,8 @@ def __init__(
             url: The url of the elasticsearch cluster.
             api_key: The api key to use.
             index_prefix: The index prefix to use. Collections will be prefixed with this prefix.
+            native_storage: Whether to use native storage mode (flattened field type) for values.
+                Defaults to False for backward compatibility.
             default_collection: The default collection to use if no collection is provided.
         """
         if elasticsearch_client is None and url is None:
@@ -153,6 +192,7 @@ def __init__(
             raise ValueError(msg)
 
         self._index_prefix = index_prefix
+        self._native_storage = native_storage
         self._is_serverless = False
 
         super().__init__(default_collection=default_collection)
@@ -205,18 +245,11 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
         if not (source := get_source_from_body(body=body)):
             return None
 
-        if not (value_str := source.get("value")) or not isinstance(value_str, str):
+        try:
+            return source_to_managed_entry(source=source)
+        except DeserializationError:
             return None
 
-        created_at: datetime | None = try_parse_datetime_str(value=source.get("created_at"))
-        expires_at: datetime | None = try_parse_datetime_str(value=source.get("expires_at"))
-
-        return ManagedEntry(
-            value=load_from_json(value_str),
-            created_at=created_at,
-            expires_at=expires_at,
-        )
-
     @override
     async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) -> list[ManagedEntry | None]:
         if not keys:
@@ -265,7 +298,9 @@ async def _put_managed_entry(
         index_name: str = self._sanitize_index_name(collection=collection)
         document_id: str = self._sanitize_document_id(key=key)
 
-        document: dict[str, Any] = managed_entry_to_document(collection=collection, key=key, managed_entry=managed_entry)
+        document: dict[str, Any] = managed_entry_to_document(
+            collection=collection, key=key, managed_entry=managed_entry, native_storage=self._native_storage
+        )
 
         _ = await self._client.index(
             index=index_name,
@@ -297,7 +332,9 @@ async def _put_managed_entries(
 
             index_action: dict[str, Any] = new_bulk_action(action="index", index=index_name, document_id=document_id)
 
-            document: dict[str, Any] = managed_entry_to_document(collection=collection, key=key, managed_entry=managed_entry)
+            document: dict[str, Any] = managed_entry_to_document(
+                collection=collection, key=key, managed_entry=managed_entry, native_storage=self._native_storage
+            )
 
             operations.extend([index_action, document])
 

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

@@ -69,6 +69,31 @@ def test_managed_entry_document_conversion():
     assert round_trip_managed_entry.expires_at == expires_at
 
 
+def test_managed_entry_document_conversion_native_storage():
+    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(collection="test_collection", key="test_key", managed_entry=managed_entry, native_storage=True)
+
+    assert document == snapshot(
+        {
+            "collection": "test_collection",
+            "key": "test_key",
+            "value_flattened": {"test": "test"},
+            "created_at": "2025-01-01T00:00:00+00:00",
+            "expires_at": "2025-01-01T00:00:10+00:00",
+        }
+    )
+
+    round_trip_managed_entry = source_to_managed_entry(source=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
+
+
 @pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
 class TestElasticsearchStore(ContextManagerStoreTestMixin, BaseStoreTests):
     @pytest.fixture(autouse=True, scope="session", params=ELASTICSEARCH_VERSIONS_TO_TEST)
@@ -121,3 +146,89 @@ async def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: A
         assert len(indices.body) == 2
         assert "kv-store-e2e-test-test_collection" in indices
         assert "kv-store-e2e-test-test_collection_2" in indices
+
+
+@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
+class TestElasticsearchStoreNativeMode(ContextManagerStoreTestMixin, BaseStoreTests):
+    """Test ElasticsearchStore with native_storage=True"""
+
+    @pytest.fixture(autouse=True, scope="session", params=ELASTICSEARCH_VERSIONS_TO_TEST)
+    async def setup_elasticsearch(self, request: pytest.FixtureRequest) -> AsyncGenerator[None, None]:
+        version = request.param
+        es_image = f"docker.elastic.co/elasticsearch/elasticsearch:{version}"
+
+        with docker_container(
+            f"elasticsearch-test-native-{version}",
+            es_image,
+            {str(ES_CONTAINER_PORT): ES_PORT},
+            {"discovery.type": "single-node", "xpack.security.enabled": "false"},
+        ):
+            if not await async_wait_for_true(bool_fn=ping_elasticsearch, tries=WAIT_FOR_ELASTICSEARCH_TIMEOUT, wait_time=2):
+                msg = f"Elasticsearch {version} failed to start"
+                raise ElasticsearchFailedToStartError(msg)
+
+            yield
+
+    @pytest.fixture
+    async def es_client(self) -> AsyncGenerator[AsyncElasticsearch, None]:
+        async with AsyncElasticsearch(hosts=[ES_URL]) as es_client:
+            yield es_client
+
+    @override
+    @pytest.fixture
+    async def store(self) -> AsyncGenerator[ElasticsearchStore, None]:
+        async with get_elasticsearch_client() as es_client:
+            indices = await es_client.options(ignore_status=404).indices.get(index="kv-store-native-test-*")
+            for index in indices:
+                _ = await es_client.options(ignore_status=404).indices.delete(index=index)
+        async with ElasticsearchStore(url=ES_URL, index_prefix="kv-store-native-test", native_storage=True) as store:
+            yield store
+
+    @pytest.mark.skip(reason="Distributed Caches are unbounded")
+    @override
+    async def test_not_unbounded(self, store: BaseStore): ...
+
+    @pytest.mark.skip(reason="Skip concurrent tests on distributed caches")
+    @override
+    async def test_concurrent_operations(self, store: BaseStore): ...
+
+    async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
+        """Verify values are stored as flattened objects, not JSON strings"""
+        await store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
+
+        # Check raw Elasticsearch document using public sanitization methods
+        # Note: We need to access these internal methods for testing the storage format
+        index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        doc_id = store._sanitize_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
+
+        response = await es_client.get(index=index_name, id=doc_id)
+        source = response.body["_source"]
+
+        # Should have value_flattened field as dict
+        assert "value_flattened" in source
+        assert isinstance(source["value_flattened"], dict)
+        assert source["value_flattened"] == {"name": "Alice", "age": 30}
+
+        # Should NOT have value field (JSON string)
+        assert "value" not in source or source.get("value") is None
+
+    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
+        index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        doc_id = store._sanitize_document_id(key="legacy_key")  # pyright: ignore[reportPrivateUsage]
+
+        await es_client.index(
+            index=index_name,
+            id=doc_id,
+            body={
+                "collection": "test",
+                "key": "legacy_key",
+                "value": '{"legacy": "data"}',  # JSON string
+            },
+        )
+        await es_client.indices.refresh(index=index_name)
+
+        # Should be able to read it in native mode
+        result = await store.get(collection="test", key="legacy_key")
+        assert result == {"legacy": "data"}