fix: OpenSearch async client initialization (#2645)

* offload init to a new thread in async environment

* Add _ensure_initialized_async method

* Fix method call

* Fix lint

* Fix format

* Update tests teardown

* Fix tests fixtures

* Update test treardown

* Update conftest

* Update conftest

---------

Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: Nitro963

integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py

@@ -262,7 +262,7 @@ def from_dict(cls, data: dict[str, Any]) -> "OpenSearchDocumentStore":
 
     def _ensure_initialized(self):
         # Ideally, we have a warm-up stage for document stores as well as components.
-        if not self._initialized:
+        if not self._client:
             self._client = OpenSearch(
                 hosts=self._hosts,
                 http_auth=self._http_auth,
@@ -271,6 +271,12 @@ def _ensure_initialized(self):
                 timeout=self._timeout,
                 **self._kwargs,
             )
+            self._initialized = True
+
+            self._ensure_index_exists()
+
+    async def _ensure_initialized_async(self):
+        if not self._async_client:
             async_http_auth = AsyncAWSAuth(self._http_auth) if isinstance(self._http_auth, AWSAuth) else self._http_auth
             self._async_client = AsyncOpenSearch(
                 hosts=self._hosts,
@@ -283,10 +289,22 @@ def _ensure_initialized(self):
                 connection_class=AsyncHttpConnection,
                 **self._kwargs,
             )
-
             self._initialized = True
+            await self._ensure_index_exists_async()
 
-            self._ensure_index_exists()
+    async def _ensure_index_exists_async(self):
+        assert self._async_client is not None
+
+        if await self._async_client.indices.exists(index=self._index):
+            logger.debug(
+                "The index '{index}' already exists. The `embedding_dim`, `method`, `mappings`, and "
+                "`settings` values will be ignored.",
+                index=self._index,
+            )
+        elif self._create_index:
+            # Create the index if it doesn't exist
+            body = {"mappings": self._mappings, "settings": self._settings}
+            await self._async_client.indices.create(index=self._index, body=body)
 
     def _ensure_index_exists(self):
         assert self._client is not None
@@ -315,7 +333,7 @@ async def count_documents_async(self) -> int:
         """
         Asynchronously returns the total number of documents in the document store.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
 
         assert self._async_client is not None
         return (await self._async_client.count(index=self._index))["count"]
@@ -376,7 +394,8 @@ async def filter_documents_async(self, filters: Optional[dict[str, Any]] = None)
         :param filters: The filters to apply to the document list.
         :returns: A list of Documents that match the given filters.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
+
         return await self._search_documents_async(self._prepare_filter_search_request(filters))
 
     def _prepare_bulk_write_request(
@@ -477,7 +496,8 @@ async def write_documents_async(
         :param policy: The duplicate policy to use when writing documents.
         :returns: The number of documents written to the document store.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
+        assert self._async_client is not None
         bulk_params = self._prepare_bulk_write_request(documents=documents, policy=policy, is_async=True)
         documents_written, errors = await async_bulk(**bulk_params)
         # since we call async_bulk with stats_only=False, errors is guaranteed to be a list (not int)
@@ -525,7 +545,8 @@ async def delete_documents_async(self, document_ids: list[str]) -> None:
 
         :param document_ids: the document ids to delete
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
+        assert self._async_client is not None
 
         await async_bulk(**self._prepare_bulk_delete_request(document_ids=document_ids, is_async=True))
 
@@ -583,7 +604,7 @@ async def delete_all_documents_async(self, recreate_index: bool = False) -> None
         :param recreate_index: If True, the index will be deleted and recreated with the original mappings and
             settings. If False, all documents will be deleted using the `delete_by_query` API.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
         assert self._async_client is not None
 
         try:
@@ -643,7 +664,7 @@ async def delete_by_filter_async(self, filters: dict[str, Any]) -> int:
             For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
         :returns: The number of documents deleted.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
         assert self._async_client is not None
 
         try:
@@ -707,7 +728,7 @@ async def update_by_filter_async(self, filters: dict[str, Any], meta: dict[str,
         :param meta: The metadata fields to update.
         :returns: The number of documents updated.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
         assert self._async_client is not None
 
         try:
@@ -863,7 +884,8 @@ async def _bm25_retrieval_async(
         See `OpenSearchBM25Retriever` for more information.
         """
 
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
+        assert self._async_client is not None
 
         search_params = self._prepare_bm25_search_request(
             query=query,
@@ -982,7 +1004,8 @@ async def _embedding_retrieval_async(
 
         See `OpenSearchEmbeddingRetriever` for more information.
         """
-        self._ensure_initialized()
+        await self._ensure_initialized_async()
+        assert self._async_client is not None
 
         search_params = self._prepare_embedding_search_request(
             query_embedding=query_embedding,

integrations/opensearch/tests/conftest.py

@@ -25,10 +25,12 @@ def document_store(request):
         return_embedding=True,
         method={"space_type": "cosinesimil", "engine": "nmslib", "name": "hnsw"},
     )
+    store._ensure_initialized()
     yield store
 
-    store._ensure_initialized()
+    asyncio.run(store._ensure_initialized_async())
     assert store._client
+    assert store._async_client
     store._client.indices.delete(index=index, params={"ignore": [400, 404]})
     asyncio.run(store._async_client.close())
 
@@ -51,8 +53,11 @@ def document_store_2(request):
 
     # Cleanup
     store._ensure_initialized()
+    asyncio.run(store._ensure_initialized_async())
     assert store._client
+    assert store._async_client
     store._client.indices.delete(index=index, params={"ignore": [400, 404]})
+    asyncio.run(store._async_client.close())
 
 
 @pytest.fixture
@@ -74,13 +79,14 @@ def document_store_readonly(request):
         create_index=False,
     )
     store._ensure_initialized()
+    asyncio.run(store._ensure_initialized_async())
     assert store._client
+    assert store._async_client
     store._client.cluster.put_settings(body={"transient": {"action.auto_create_index": False}})
     yield store
 
     store._client.cluster.put_settings(body={"transient": {"action.auto_create_index": True}})
     store._client.indices.delete(index=index, params={"ignore": [400, 404]})
-    asyncio.run(store._async_client.close())
 
 
 @pytest.fixture
@@ -104,7 +110,6 @@ def document_store_embedding_dim_4_no_emb_returned(request):
     yield store
 
     store._client.indices.delete(index=index, params={"ignore": [400, 404]})
-    asyncio.run(store._async_client.close())
 
 
 @pytest.fixture
@@ -129,7 +134,6 @@ def document_store_embedding_dim_4_no_emb_returned_faiss(request):
     yield store
 
     store._client.indices.delete(index=index, params={"ignore": [400, 404]})
-    asyncio.run(store._async_client.close())
 
 
 @pytest.fixture

integrations/opensearch/tests/test_document_store_async.py

@@ -6,6 +6,7 @@
 
 import pytest
 from haystack.dataclasses import Document
+from haystack.document_stores.types import DuplicatePolicy
 
 from haystack_integrations.document_stores.opensearch.document_store import OpenSearchDocumentStore
 
@@ -14,7 +15,7 @@
 class TestDocumentStoreAsync:
     @pytest.mark.asyncio
     async def test_write_documents(self, document_store: OpenSearchDocumentStore):
-        assert await document_store.write_documents_async([Document(id="1")]) == 1
+        assert await document_store.write_documents_async([Document(id="1")], policy=DuplicatePolicy.OVERWRITE) == 1
 
     @pytest.mark.asyncio
     async def test_bm25_retrieval(self, document_store: OpenSearchDocumentStore, test_documents: list[Document]):