revert: remove Elasticsearch changes from OpenSearch PR

These changes were unrelated to OpenSearch support and should be in a separate PR.

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

Author: github-actions[bot]

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

@@ -270,8 +270,6 @@ async def _setup_collection(self, *, collection: str) -> None:
         except BadRequestError as e:
             if "index_already_exists_exception" in str(e).lower():
                 return
-            if "resource_already_exists_exception" in str(e).lower():
-                return
             raise
 
     def _get_index_name(self, collection: str) -> str:

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

@@ -1,5 +1,6 @@
 from collections.abc import AsyncGenerator
 from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from dirty_equals import IsFloat, IsStr
@@ -19,11 +20,14 @@
 from tests.conftest import docker_container, should_skip_docker_tests
 from tests.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
+if TYPE_CHECKING:
+    from elastic_transport._response import ObjectApiResponse
+
 TEST_SIZE_LIMIT = 1 * 1024 * 1024  # 1MB
 ES_HOST = "localhost"
-CONTAINER_PORT = 9200
-HOST_PORT = 19200
-ES_URL = f"http://{ES_HOST}:{HOST_PORT}"
+ES_PORT = 9200
+ES_URL = f"http://{ES_HOST}:{ES_PORT}"
+ES_CONTAINER_PORT = 9200
 
 WAIT_FOR_ELASTICSEARCH_TIMEOUT = 30
 
@@ -41,7 +45,12 @@ async def ping_elasticsearch() -> bool:
     es_client: AsyncElasticsearch = get_elasticsearch_client()
 
     async with es_client:
-        return await es_client.ping()
+        if not await es_client.ping():
+            return False
+
+        status: ObjectApiResponse[dict[str, Any]] = await es_client.options(ignore_status=404).cluster.health(wait_for_status="green")
+
+        return status.body.get("status") == "green"
 
 
 async def cleanup_elasticsearch_indices(elasticsearch_client: AsyncElasticsearch):
@@ -90,7 +99,7 @@ async def setup_elasticsearch(self, request: pytest.FixtureRequest) -> AsyncGene
         with docker_container(
             f"elasticsearch-test-{version}",
             es_image,
-            {str(CONTAINER_PORT): HOST_PORT},
+            {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):
@@ -99,6 +108,11 @@ async def setup_elasticsearch(self, request: pytest.FixtureRequest) -> AsyncGene
 
             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) -> ElasticsearchStore:
@@ -114,11 +128,10 @@ async def sanitizing_store(self) -> ElasticsearchStore:
         )
 
     @pytest.fixture(autouse=True)
-    async def cleanup_elasticsearch(self):
-        async with get_elasticsearch_client() as es_client:
-            await cleanup_elasticsearch_indices(elasticsearch_client=es_client)
-            yield
-            await cleanup_elasticsearch_indices(elasticsearch_client=es_client)
+    async def cleanup_elasticsearch_indices(self, es_client: AsyncElasticsearch):
+        await cleanup_elasticsearch_indices(elasticsearch_client=es_client)
+        yield
+        await cleanup_elasticsearch_indices(elasticsearch_client=es_client)
 
     @pytest.mark.skip(reason="Distributed Caches are unbounded")
     @override
@@ -145,19 +158,18 @@ async def test_long_key_name(self, store: ElasticsearchStore, sanitizing_store:
         await sanitizing_store.put(collection="test_collection", key="test_key" * 100, value={"test": "test"})
         assert await sanitizing_store.get(collection="test_collection", key="test_key" * 100) == {"test": "test"}
 
-    async def test_put_put_two_indices(self, store: ElasticsearchStore):
+    async def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
         await store.put(collection="test_collection", key="test_key", value={"test": "test"})
         await store.put(collection="test_collection_2", key="test_key", value={"test": "test"})
         assert await store.get(collection="test_collection", key="test_key") == {"test": "test"}
         assert await store.get(collection="test_collection_2", key="test_key") == {"test": "test"}
 
-        async with get_elasticsearch_client() as es_client:
-            indices = await es_client.options(ignore_status=404).indices.get(index="kv-store-e2e-test-*")
-            assert len(indices.body) == 2
-            index_names: list[str] = [str(key) for key in indices]
-            assert index_names == snapshot(["kv-store-e2e-test-test_collection", "kv-store-e2e-test-test_collection_2"])
+        indices = await es_client.options(ignore_status=404).indices.get(index="kv-store-e2e-test-*")
+        assert len(indices.body) == 2
+        index_names: list[str] = [str(key) for key in indices]
+        assert index_names == snapshot(["kv-store-e2e-test-test_collection", "kv-store-e2e-test-test_collection_2"])
 
-    async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore):
+    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})
 
@@ -166,32 +178,30 @@ async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore)
         index_name = store._get_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._get_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
 
-        async with get_elasticsearch_client() as es_client:
-            response = await es_client.get(index=index_name, id=doc_id)
-            assert response.body["_source"] == snapshot(
-                {
-                    "version": 1,
-                    "key": "test_key",
-                    "collection": "test",
-                    "value": {"flattened": {"name": "Alice", "age": 30}},
-                    "created_at": IsStr(min_length=20, max_length=40),
-                }
-            )
-
-            # Test with TTL
-            await store.put(collection="test", key="test_key", value={"name": "Bob", "age": 25}, ttl=10)
-
-            response = await es_client.get(index=index_name, id=doc_id)
-            assert response.body["_source"] == snapshot(
-                {
-                    "version": 1,
-                    "key": "test_key",
-                    "collection": "test",
-                    "value": {"flattened": {"name": "Bob", "age": 25}},
-                    "created_at": IsStr(min_length=20, max_length=40),
-                    "expires_at": IsStr(min_length=20, max_length=40),
-                }
-            )
+        response = await es_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "version": 1,
+                "key": "test_key",
+                "collection": "test",
+                "value": {"flattened": {"name": "Alice", "age": 30}},
+                "created_at": IsStr(min_length=20, max_length=40),
+            }
+        )
+
+        # Test with TTL
+        await store.put(collection="test", key="test_key", value={"name": "Bob", "age": 25}, ttl=10)
+        response = await es_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "version": 1,
+                "key": "test_key",
+                "collection": "test",
+                "value": {"flattened": {"name": "Bob", "age": 25}},
+                "created_at": IsStr(min_length=20, max_length=40),
+                "expires_at": IsStr(min_length=20, max_length=40),
+            }
+        )
 
     @override
     async def test_special_characters_in_collection_name(self, store: ElasticsearchStore, sanitizing_store: ElasticsearchStore):  # pyright: ignore[reportIncompatibleMethodOverride]

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

@@ -38,7 +38,7 @@
         get_source_from_body,
     )
 except ImportError as e:
-    msg = "ElasticsearchStore requires py-key-value-sync[elasticsearch]"
+    msg = "ElasticsearchStore requires py-key-value-aio[elasticsearch]"
     raise ImportError(msg) from e
 
 logger = logging.getLogger(__name__)
@@ -242,8 +242,6 @@ def _setup_collection(self, *, collection: str) -> None:
         except BadRequestError as e:
             if "index_already_exists_exception" in str(e).lower():
                 return
-            if "resource_already_exists_exception" in str(e).lower():
-                return
             raise
 
     def _get_index_name(self, collection: str) -> str:

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

@@ -3,6 +3,7 @@
 # DO NOT CHANGE! Change the original file instead.
 from collections.abc import Generator
 from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from dirty_equals import IsFloat, IsStr
@@ -22,11 +23,14 @@
 from tests.code_gen.conftest import docker_container, should_skip_docker_tests
 from tests.code_gen.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
+if TYPE_CHECKING:
+    from elastic_transport._response import ObjectApiResponse
+
 TEST_SIZE_LIMIT = 1 * 1024 * 1024  # 1MB
 ES_HOST = "localhost"
-CONTAINER_PORT = 9200
-HOST_PORT = 19200
-ES_URL = f"http://{ES_HOST}:{HOST_PORT}"
+ES_PORT = 9200
+ES_URL = f"http://{ES_HOST}:{ES_PORT}"
+ES_CONTAINER_PORT = 9200
 
 WAIT_FOR_ELASTICSEARCH_TIMEOUT = 30
 # Released Apr 2025
@@ -42,7 +46,12 @@ def ping_elasticsearch() -> bool:
     es_client: Elasticsearch = get_elasticsearch_client()
 
     with es_client:
-        return es_client.ping()
+        if not es_client.ping():
+            return False
+
+        status: ObjectApiResponse[dict[str, Any]] = es_client.options(ignore_status=404).cluster.health(wait_for_status="green")
+
+        return status.body.get("status") == "green"
 
 
 def cleanup_elasticsearch_indices(elasticsearch_client: Elasticsearch):
@@ -91,7 +100,7 @@ def setup_elasticsearch(self, request: pytest.FixtureRequest) -> Generator[None,
         with docker_container(
             f"elasticsearch-test-{version}",
             es_image,
-            {str(CONTAINER_PORT): HOST_PORT},
+            {str(ES_CONTAINER_PORT): ES_PORT},
             {"discovery.type": "single-node", "xpack.security.enabled": "false"},
         ):
             if not wait_for_true(bool_fn=ping_elasticsearch, tries=WAIT_FOR_ELASTICSEARCH_TIMEOUT, wait_time=2):
@@ -100,6 +109,11 @@ def setup_elasticsearch(self, request: pytest.FixtureRequest) -> Generator[None,
 
             yield
 
+    @pytest.fixture
+    def es_client(self) -> Generator[Elasticsearch, None, None]:
+        with Elasticsearch(hosts=[ES_URL]) as es_client:
+            yield es_client
+
     @override
     @pytest.fixture
     def store(self) -> ElasticsearchStore:
@@ -115,11 +129,10 @@ def sanitizing_store(self) -> ElasticsearchStore:
         )
 
     @pytest.fixture(autouse=True)
-    def cleanup_elasticsearch(self):
-        with get_elasticsearch_client() as es_client:
-            cleanup_elasticsearch_indices(elasticsearch_client=es_client)
-            yield
-            cleanup_elasticsearch_indices(elasticsearch_client=es_client)
+    def cleanup_elasticsearch_indices(self, es_client: Elasticsearch):
+        cleanup_elasticsearch_indices(elasticsearch_client=es_client)
+        yield
+        cleanup_elasticsearch_indices(elasticsearch_client=es_client)
 
     @pytest.mark.skip(reason="Distributed Caches are unbounded")
     @override
@@ -146,19 +159,18 @@ def test_long_key_name(self, store: ElasticsearchStore, sanitizing_store: Elasti
         sanitizing_store.put(collection="test_collection", key="test_key" * 100, value={"test": "test"})
         assert sanitizing_store.get(collection="test_collection", key="test_key" * 100) == {"test": "test"}
 
-    def test_put_put_two_indices(self, store: ElasticsearchStore):
+    def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: Elasticsearch):
         store.put(collection="test_collection", key="test_key", value={"test": "test"})
         store.put(collection="test_collection_2", key="test_key", value={"test": "test"})
         assert store.get(collection="test_collection", key="test_key") == {"test": "test"}
         assert store.get(collection="test_collection_2", key="test_key") == {"test": "test"}
 
-        with get_elasticsearch_client() as es_client:
-            indices = es_client.options(ignore_status=404).indices.get(index="kv-store-e2e-test-*")
-            assert len(indices.body) == 2
-            index_names: list[str] = [str(key) for key in indices]
-            assert index_names == snapshot(["kv-store-e2e-test-test_collection", "kv-store-e2e-test-test_collection_2"])
+        indices = es_client.options(ignore_status=404).indices.get(index="kv-store-e2e-test-*")
+        assert len(indices.body) == 2
+        index_names: list[str] = [str(key) for key in indices]
+        assert index_names == snapshot(["kv-store-e2e-test-test_collection", "kv-store-e2e-test-test_collection_2"])
 
-    def test_value_stored_as_flattened_object(self, store: ElasticsearchStore):
+    def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_client: Elasticsearch):
         """Verify values are stored as flattened objects, not JSON strings"""
         store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
 
@@ -167,32 +179,30 @@ def test_value_stored_as_flattened_object(self, store: ElasticsearchStore):
         index_name = store._get_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._get_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
 
-        with get_elasticsearch_client() as es_client:
-            response = es_client.get(index=index_name, id=doc_id)
-            assert response.body["_source"] == snapshot(
-                {
-                    "version": 1,
-                    "key": "test_key",
-                    "collection": "test",
-                    "value": {"flattened": {"name": "Alice", "age": 30}},
-                    "created_at": IsStr(min_length=20, max_length=40),
-                }
-            )
-
-            # Test with TTL
-            store.put(collection="test", key="test_key", value={"name": "Bob", "age": 25}, ttl=10)
-
-            response = es_client.get(index=index_name, id=doc_id)
-            assert response.body["_source"] == snapshot(
-                {
-                    "version": 1,
-                    "key": "test_key",
-                    "collection": "test",
-                    "value": {"flattened": {"name": "Bob", "age": 25}},
-                    "created_at": IsStr(min_length=20, max_length=40),
-                    "expires_at": IsStr(min_length=20, max_length=40),
-                }
-            )
+        response = es_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "version": 1,
+                "key": "test_key",
+                "collection": "test",
+                "value": {"flattened": {"name": "Alice", "age": 30}},
+                "created_at": IsStr(min_length=20, max_length=40),
+            }
+        )
+
+        # Test with TTL
+        store.put(collection="test", key="test_key", value={"name": "Bob", "age": 25}, ttl=10)
+        response = es_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "version": 1,
+                "key": "test_key",
+                "collection": "test",
+                "value": {"flattened": {"name": "Bob", "age": 25}},
+                "created_at": IsStr(min_length=20, max_length=40),
+                "expires_at": IsStr(min_length=20, max_length=40),
+            }
+        )
 
     @override
     def test_special_characters_in_collection_name(self, store: ElasticsearchStore, sanitizing_store: ElasticsearchStore):  # pyright: ignore[reportIncompatibleMethodOverride]