fix(elasticsearch): fix async write/delete bugs and restore comments

Fix three async behavior deviations from the standard DocumentStore interface:

- write_documents_async with DuplicatePolicy.FAIL now raises
  DuplicateDocumentError instead of DocumentStoreError (outer try/except
  was swallowing the inner DuplicateDocumentError)
- write_documents_async with DuplicatePolicy.SKIP now returns 0 when
  skipping an existing document (errors are now categorized like the sync
  version: version_conflict_engine_exception with SKIP policy is silently
  ignored)
- delete_documents_async no longer raises on non-existent IDs (adds
  raise_on_error=False to match the sync version behaviour)

Also restore docstrings and inline comments that were accidentally removed
in the previous refactor commit.

Author: SyedShahmeerAli12

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

@@ -710,29 +710,38 @@ async def write_documents_async(
             }
             actions.append(action)
 
-        try:
-            success, failed = await helpers.async_bulk(
-                client=self.async_client,
-                actions=actions,
-                index=self._index,
-                refresh=refresh,
-                raise_on_error=False,
-                stats_only=False,
-            )
-            if failed:
-                # with stats_only=False, failed is guaranteed to be a list of dicts
-                assert isinstance(failed, list)
-                if policy == DuplicatePolicy.FAIL:
-                    for error in failed:
-                        if "create" in error and error["create"]["status"] == DOC_ALREADY_EXISTS:
-                            msg = f"ID '{error['create']['_id']}' already exists in the document store"
-                            raise DuplicateDocumentError(msg)
-                msg = f"Failed to write documents to Elasticsearch. Errors:\n{failed}"
+        documents_written, errors = await helpers.async_bulk(
+            client=self.async_client,
+            actions=actions,
+            index=self._index,
+            refresh=refresh,
+            raise_on_error=False,
+            stats_only=False,
+        )
+
+        if errors:
+            # with stats_only=False, errors is guaranteed to be a list of dicts
+            assert isinstance(errors, list)
+            duplicate_errors_ids = []
+            other_errors = []
+            for e in errors:
+                error_type = e["create"]["error"]["type"]
+                if policy == DuplicatePolicy.FAIL and error_type == "version_conflict_engine_exception":
+                    duplicate_errors_ids.append(e["create"]["_id"])
+                elif policy == DuplicatePolicy.SKIP and error_type == "version_conflict_engine_exception":
+                    continue
+                else:
+                    other_errors.append(e)
+
+            if len(duplicate_errors_ids) > 0:
+                msg = f"IDs '{', '.join(duplicate_errors_ids)}' already exist in the document store."
+                raise DuplicateDocumentError(msg)
+
+            if len(other_errors) > 0:
+                msg = f"Failed to write documents to Elasticsearch. Errors:\n{other_errors}"
                 raise DocumentStoreError(msg)
-            return success
-        except Exception as e:
-            msg = f"Failed to write documents to Elasticsearch: {e!s}"
-            raise DocumentStoreError(msg) from e
+
+        return documents_written
 
     def delete_documents(self, document_ids: list[str], refresh: Literal["wait_for", True, False] = "wait_for") -> None:
         """
@@ -776,16 +785,13 @@ async def delete_documents_async(
         """
         self._ensure_initialized()
 
-        try:
-            await helpers.async_bulk(
-                client=self.async_client,
-                actions=({"_op_type": "delete", "_id": id_} for id_ in document_ids),
-                index=self._index,
-                refresh=refresh,
-            )
-        except Exception as e:
-            msg = f"Failed to delete documents from Elasticsearch: {e!s}"
-            raise DocumentStoreError(msg) from e
+        await helpers.async_bulk(
+            client=self.async_client,
+            actions=({"_op_type": "delete", "_id": id_} for id_ in document_ids),
+            index=self._index,
+            refresh=refresh,
+            raise_on_error=False,
+        )
 
     def delete_all_documents(self, recreate_index: bool = False, refresh: bool = True) -> None:
         """

integrations/elasticsearch/tests/test_document_store_async.py

@@ -43,12 +43,17 @@ class TestElasticsearchDocumentStoreAsync(
 ):
     @pytest_asyncio.fixture
     async def document_store(self, request):
+        """
+        Basic fixture providing a document store instance for async tests
+        """
         hosts = ["http://localhost:9200"]
+        # Use a different index for each test so we can run them in parallel
         index = f"{request.node.name}"
 
         store = ElasticsearchDocumentStore(hosts=hosts, index=index)
         yield store
         store.client.options(ignore_status=[400, 404]).indices.delete(index=index)
+
         await store.async_client.close()
 
     def assert_documents_are_equal(self, received: list[Document], expected: list[Document]):
@@ -63,37 +68,6 @@ async def test_count_not_empty_async(self, document_store):
         )
         assert await document_store.count_documents_async() == 3
 
-    @pytest.mark.asyncio
-    async def test_delete_documents_empty_document_store_async(self, document_store):
-        # Elasticsearch raises DocumentStoreError when deleting a non-existent document
-        # rather than silently ignoring it, so we override the mixin test here.
-        with pytest.raises(DocumentStoreError):
-            await document_store.delete_documents_async(["non_existing_id"])
-
-    @pytest.mark.asyncio
-    async def test_delete_documents_non_existing_document_async(self, document_store):
-        # Same as above: Elasticsearch raises on missing IDs rather than a no-op.
-        doc = Document(content="test doc")
-        await document_store.write_documents_async([doc])
-        assert await document_store.count_documents_async() == 1
-        with pytest.raises(DocumentStoreError):
-            await document_store.delete_documents_async(["non_existing_id"])
-
-    @pytest.mark.asyncio
-    async def test_write_documents_duplicate_fail_async(self, document_store):
-        # Elasticsearch raises DocumentStoreError instead of DuplicateDocumentError on duplicate FAIL
-        doc = Document(content="test doc")
-        assert await document_store.write_documents_async([doc], policy=DuplicatePolicy.FAIL) == 1
-        with pytest.raises(DocumentStoreError):
-            await document_store.write_documents_async(documents=[doc], policy=DuplicatePolicy.FAIL)
-
-    @pytest.mark.asyncio
-    async def test_write_documents_duplicate_skip_async(self, document_store):
-        # Elasticsearch returns 1 (not 0) when skipping an already-existing document
-        doc = Document(content="test doc")
-        assert await document_store.write_documents_async([doc], policy=DuplicatePolicy.SKIP) == 1
-        assert await document_store.write_documents_async(documents=[doc], policy=DuplicatePolicy.SKIP) == 1
-
     @pytest.mark.asyncio
     async def test_write_documents_async(self, document_store):
         docs = [Document(id="1", content="test")]
@@ -104,12 +78,14 @@ async def test_write_documents_async(self, document_store):
 
     @pytest.mark.asyncio
     async def test_write_documents_async_invalid_document_type(self, document_store):
+        """Test write_documents with invalid document type"""
         invalid_docs = [{"id": "1", "content": "test"}]
         with pytest.raises(ValueError, match="param 'documents' must contain a list of objects of type Document"):
             await document_store.write_documents_async(invalid_docs)
 
     @pytest.mark.asyncio
     async def test_write_documents_async_with_sparse_embedding_warning(self, document_store, caplog):
+        """Test write_documents with document containing sparse_embedding field"""
         doc = Document(id="1", content="test", sparse_embedding=SparseEmbedding(indices=[0, 1], values=[0.5, 0.5]))
         await document_store.write_documents_async([doc])
         assert "but `sparse_vector_field` is not configured" in caplog.text
@@ -121,6 +97,7 @@ async def test_write_documents_async_with_sparse_embedding_warning(self, documen
 
     @pytest.mark.asyncio
     async def test_write_documents_async_with_sparse_vectors(self):
+        """Test write_documents with document containing sparse_embedding field"""
         store = ElasticsearchDocumentStore(
             hosts=["http://localhost:9200"], index="test_async_sparse", sparse_vector_field="sparse_vec"
         )
@@ -129,9 +106,11 @@ async def test_write_documents_async_with_sparse_vectors(self):
         doc = Document(id="1", content="test", sparse_embedding=SparseEmbedding(indices=[0, 1], values=[0.5, 0.5]))
         await store.write_documents_async([doc])
 
+        # check ES natively
         raw_doc = await store.async_client.get(index="test_async_sparse", id="1")
         assert raw_doc["_source"]["sparse_vec"] == {"0": 0.5, "1": 0.5}
 
+        # check retrieval
         results = await store.filter_documents_async()
         assert len(results) == 1
         assert results[0].sparse_embedding is not None