feat: Add async run to DocumentWriter (#8962)

* add async run to DocumentWriter

* reno

Author: julian-risch

haystack/components/writers/document_writer.py

@@ -25,7 +25,8 @@ class DocumentWriter:
         Document(content="Python is a popular programming language"),
     ]
     doc_store = InMemoryDocumentStore()
-    doc_store.write_documents(docs)
+    writer = DocumentWriter(document_store=doc_store)
+    writer.run(docs)
     ```
     """
 
@@ -100,3 +101,34 @@ def run(self, documents: List[Document], policy: Optional[DuplicatePolicy] = Non
 
         documents_written = self.document_store.write_documents(documents=documents, policy=policy)
         return {"documents_written": documents_written}
+
+    @component.output_types(documents_written=int)
+    async def run_async(self, documents: List[Document], policy: Optional[DuplicatePolicy] = None):
+        """
+        Asynchronously run the DocumentWriter on the given input data.
+
+        This is the asynchronous version of the `run` method. It has the same parameters and return values
+        but can be used with `await` in async code.
+
+        :param documents:
+            A list of documents to write to the document store.
+        :param policy:
+            The policy to use when encountering duplicate documents.
+        :returns:
+            Number of documents written to the document store.
+
+        :raises ValueError:
+            If the specified document store is not found.
+        :raises TypeError:
+            If the specified document store does not implement `write_documents_async`.
+        """
+        if policy is None:
+            policy = self.policy
+
+        if not hasattr(self.document_store, "write_documents_async"):
+            raise TypeError(f"Document store {type(self.document_store).__name__} does not provide async support.")
+
+        documents_written = await self.document_store.write_documents_async(  # type: ignore
+            documents=documents, policy=policy
+        )
+        return {"documents_written": documents_written}

releasenotes/notes/async-document-writer-9458ced66c1c38c7.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Add `run_async` method to `DocumentWriter`. This method supports the same parameters as
+    the `run` method and relies on the DocumentStore to implement `write_documents_async`.
+    It returns a coroutine that can be awaited.

test/components/writers/test_document_writer.py

@@ -92,3 +92,43 @@ def test_run_skip_policy(self):
 
         result = writer.run(documents=documents)
         assert result["documents_written"] == 0
+
+    @pytest.mark.asyncio
+    async def test_run_async_invalid_docstore(self):
+        document_store = document_store_class("MockedDocumentStore")
+
+        writer = DocumentWriter(document_store)
+        documents = [
+            Document(content="This is the text of a document."),
+            Document(content="This is the text of another document."),
+        ]
+
+        with pytest.raises(TypeError, match="does not provide async support"):
+            await writer.run_async(documents=documents)
+
+    @pytest.mark.asyncio
+    async def test_run_async(self):
+        document_store = InMemoryDocumentStore()
+        writer = DocumentWriter(document_store)
+        documents = [
+            Document(content="This is the text of a document."),
+            Document(content="This is the text of another document."),
+        ]
+
+        result = await writer.run_async(documents=documents)
+        assert result["documents_written"] == 2
+
+    @pytest.mark.asyncio
+    async def test_run_async_skip_policy(self):
+        document_store = InMemoryDocumentStore()
+        writer = DocumentWriter(document_store, policy=DuplicatePolicy.SKIP)
+        documents = [
+            Document(content="This is the text of a document."),
+            Document(content="This is the text of another document."),
+        ]
+
+        result = await writer.run_async(documents=documents)
+        assert result["documents_written"] == 2
+
+        result = await writer.run_async(documents=documents)
+        assert result["documents_written"] == 0