feat: add closing methods to components that hold Document Stores (#12169)

Author: anakin87

haystack/components/caching/cache_checker.py

@@ -121,3 +121,17 @@ async def run_async(self, items: list[Any]) -> dict[str, Any]:
             else:
                 misses.append(item)
         return {"hits": found_documents, "misses": misses}
+
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close"):
+            self.document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close_async"):
+            await self.document_store.close_async()

haystack/components/retrievers/auto_merging_retriever.py

@@ -224,3 +224,17 @@ async def _try_merge_level(docs_to_merge: list[Document], docs_to_return: list[D
             return await _try_merge_level(merged_docs, docs_to_return)
 
         return {"documents": await _try_merge_level(documents, [])}
+
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close"):
+            self.document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close_async"):
+            await self.document_store.close_async()

haystack/components/retrievers/filter_retriever.py

@@ -102,3 +102,17 @@ async def run_async(self, filters: dict[str, Any] | None = None) -> dict[str, An
         # 'ignore' since filter_documents_async is not defined in the Protocol but exists in the implementations
         out_documents = await self.document_store.filter_documents_async(filters=filters or self.filters)  # type: ignore[attr-defined]
         return {"documents": out_documents}
+
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close"):
+            self.document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close_async"):
+            await self.document_store.close_async()

haystack/components/retrievers/sentence_window_retriever.py

@@ -319,3 +319,17 @@ def _build_filter_conditions(self, split_id: int, window_size: int, source_ids:
             *source_id_filters,
         ]
         return {"operator": "AND", "conditions": conditions}
+
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close"):
+            self.document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close_async"):
+            await self.document_store.close_async()

haystack/components/writers/document_writer.py

@@ -125,3 +125,17 @@ async def run_async(self, documents: list[Document], policy: DuplicatePolicy | N
 
         documents_written = await self.document_store.write_documents_async(documents=documents, policy=policy)
         return {"documents_written": documents_written}
+
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close"):
+            self.document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        if hasattr(self.document_store, "close_async"):
+            await self.document_store.close_async()

releasenotes/notes/closing-methods-for-ds-holding-comps-e2dcfc4e5afddd0f.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    Haystack components that use a document store now provide ``close`` and ``close_async`` methods for releasing
+    resources. These methods are available on: ``AutoMergingRetriever``, ``CacheChecker``, ``DocumentWriter``,
+    ``FilterRetriever``, and ``SentenceWindowRetriever``. If the underlying Document Store does not implement the
+    corresponding method, calling ``close`` or ``close_async`` has no effect.

…onents/caching/test_url_cache_checker.py …components/caching/test_cache_checker.pytest/components/caching/test_url_cache_checker.py renamed to test/components/caching/test_cache_checker.py test/components/caching/test_url_cache_checker.py renamed to test/components/caching/test_cache_checker.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 
@@ -92,3 +92,14 @@ def test_filters_syntax(self):
         checker.run(items=["https://example.com/1"])
         valid_filters_syntax = {"field": "url", "operator": "==", "value": "https://example.com/1"}
         mocked_docstore_class.filter_documents.assert_any_call(filters=valid_filters_syntax)
+
+    def test_close(self):
+        closable_document_store = Mock(spec=["close"])
+        checker = CacheChecker(document_store=closable_document_store, cache_field="url")
+        checker.close()
+        closable_document_store.close.assert_called_once_with()
+
+        nonclosable_document_store = Mock(spec=[])
+        checker = CacheChecker(document_store=nonclosable_document_store, cache_field="url")
+        checker.close()
+        assert nonclosable_document_store.mock_calls == []

test/components/caching/test_cache_checker_async.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, Mock
 
 import pytest
 
@@ -60,3 +60,16 @@ async def test_run_async_filters_syntax(self):
         await checker.run_async(items=["https://example.com/1"])
         expected_filters = {"field": "url", "operator": "==", "value": "https://example.com/1"}
         mock_store.filter_documents_async.assert_awaited_once_with(filters=expected_filters)
+
+    @pytest.mark.asyncio
+    async def test_close_async(self):
+        closable_document_store = Mock(spec=["close_async"])
+        closable_document_store.close_async = AsyncMock()
+        checker = CacheChecker(document_store=closable_document_store, cache_field="url")
+        await checker.close_async()
+        closable_document_store.close_async.assert_awaited_once_with()
+
+        nonclosable_document_store = Mock(spec=[])
+        checker = CacheChecker(document_store=nonclosable_document_store, cache_field="url")
+        await checker.close_async()
+        assert nonclosable_document_store.mock_calls == []

test/components/retrievers/test_auto_merging_retriever.py

@@ -2,6 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from unittest.mock import Mock
+
 import pytest
 
 from haystack import Document, Pipeline
@@ -252,3 +254,14 @@ def test_run_go_up_hierarchy_multiple_levels_hit_root_document(self, in_memory_d
 
         assert len(result["documents"]) == 1
         assert result["documents"][0].meta["__level"] == 0  # hit root document
+
+    def test_close(self):
+        closable_document_store = Mock(spec=["close"])
+        retriever = AutoMergingRetriever(document_store=closable_document_store)
+        retriever.close()
+        closable_document_store.close.assert_called_once_with()
+
+        nonclosable_document_store = Mock(spec=[])
+        retriever = AutoMergingRetriever(document_store=nonclosable_document_store)
+        retriever.close()
+        assert nonclosable_document_store.mock_calls == []

test/components/retrievers/test_auto_merging_retriever_async.py

@@ -2,6 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from unittest.mock import AsyncMock, Mock
+
 import pytest
 
 from haystack import Document
@@ -206,3 +208,16 @@ async def test_run_go_up_hierarchy_multiple_levels_hit_root_document(self, in_me
 
         assert len(result["documents"]) == 1
         assert result["documents"][0].meta["__level"] == 0  # hit root document
+
+    @pytest.mark.asyncio
+    async def test_close_async(self):
+        closable_document_store = Mock(spec=["close_async"])
+        closable_document_store.close_async = AsyncMock()
+        retriever = AutoMergingRetriever(document_store=closable_document_store)
+        await retriever.close_async()
+        closable_document_store.close_async.assert_awaited_once_with()
+
+        nonclosable_document_store = Mock(spec=[])
+        retriever = AutoMergingRetriever(document_store=nonclosable_document_store)
+        await retriever.close_async()
+        assert nonclosable_document_store.mock_calls == []