feat: add closing methods for Qdrant (#3671)

Author: anakin87

integrations/qdrant/src/haystack_integrations/components/retrievers/qdrant/retriever.py

@@ -130,6 +130,18 @@ def from_dict(cls, data: dict[str, Any]) -> "QdrantEmbeddingRetriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
         return default_from_dict(cls, data)
 
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        self._document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        await self._document_store.close_async()
+
     @component.output_types(documents=list[Document])
     def run(
         self,
@@ -358,6 +370,18 @@ def from_dict(cls, data: dict[str, Any]) -> "QdrantSparseEmbeddingRetriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
         return default_from_dict(cls, data)
 
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        self._document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        await self._document_store.close_async()
+
     @component.output_types(documents=list[Document])
     def run(
         self,
@@ -596,6 +620,18 @@ def from_dict(cls, data: dict[str, Any]) -> "QdrantHybridRetriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
         return default_from_dict(cls, data)
 
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        self._document_store.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the asynchronous resources of the underlying Document Store.
+        """
+        await self._document_store.close_async()
+
     @component.output_types(documents=list[Document])
     def run(
         self,

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

@@ -1,5 +1,6 @@
 import inspect
 from collections.abc import AsyncGenerator, Generator
+from contextlib import suppress
 from dataclasses import replace
 from itertools import islice
 from typing import Any, ClassVar, cast
@@ -301,6 +302,24 @@ async def _initialize_async_client(self) -> None:
                 self.payload_fields_to_index,
             )
 
+    def close(self) -> None:
+        """
+        Release the associated synchronous resources.
+        """
+        if self._client is not None:
+            with suppress(Exception):
+                self._client.close()
+            self._client = None
+
+    async def close_async(self) -> None:
+        """
+        Release the associated asynchronous resources.
+        """
+        if self._async_client is not None:
+            with suppress(Exception):
+                await self._async_client.close()
+            self._async_client = None
+
     def count_documents(self) -> int:
         """
         Returns the number of documents present in the Document Store.

integrations/qdrant/tests/test_document_store.py

@@ -370,6 +370,29 @@ def test_metadata_methods_swallow_client_errors(self, method_name, args, expecte
         ):
             assert getattr(document_store, method_name)(*args) == expected
 
+    def test_close(self):
+        document_store = QdrantDocumentStore(location=":memory:")
+        mock_client = MagicMock()
+        document_store._client = mock_client
+
+        document_store.close()
+
+        mock_client.close.assert_called_once()
+        assert document_store._client is None
+
+        document_store.close()
+        mock_client.close.assert_called_once()
+
+    def test_close_is_exception_safe(self):
+        document_store = QdrantDocumentStore(location=":memory:")
+        mock_client = MagicMock()
+        mock_client.close.side_effect = RuntimeError("boom")
+        document_store._client = mock_client
+
+        document_store.close()
+
+        assert document_store._client is None
+
 
 @pytest.mark.integration
 class TestQdrantDocumentStore(
@@ -409,6 +432,16 @@ def assert_documents_are_equal(self, received: list[Document], expected: list[Do
         # Check that the sets are equal, meaning the content and IDs match regardless of order
         assert {doc.id for doc in received} == {doc.id for doc in expected}
 
+    def test_close_and_reopen(self, document_store: QdrantDocumentStore):
+        assert document_store.count_documents() == 0
+        assert document_store._client is not None
+
+        document_store.close()
+
+        assert document_store._client is None
+        assert document_store.count_documents() == 0
+        assert document_store._client is not None
+
     def test_prepare_client_params_no_mutability(self):
         metadata = {"key": "value"}
         doc_store = QdrantDocumentStore(

integrations/qdrant/tests/test_document_store_async.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 import pytest_asyncio
@@ -162,6 +162,29 @@ async def test_metadata_methods_async_absorb_client_errors(self, method_name, ar
         ):
             assert await getattr(document_store, method_name)(*args) == expected
 
+    async def test_close_async(self):
+        document_store = QdrantDocumentStore(location=":memory:")
+        mock_client = AsyncMock()
+        document_store._async_client = mock_client
+
+        await document_store.close_async()
+
+        mock_client.close.assert_awaited_once()
+        assert document_store._async_client is None
+
+        await document_store.close_async()
+        mock_client.close.assert_awaited_once()
+
+    async def test_close_async_is_exception_safe(self):
+        document_store = QdrantDocumentStore(location=":memory:")
+        mock_client = AsyncMock()
+        mock_client.close.side_effect = RuntimeError("boom")
+        document_store._async_client = mock_client
+
+        await document_store.close_async()
+
+        assert document_store._async_client is None
+
 
 @pytest.mark.integration
 @pytest.mark.asyncio
@@ -195,6 +218,16 @@ def assert_documents_are_equal(self, received: list[Document], expected: list[Do
         assert len(received) == len(expected)
         assert {doc.id for doc in received} == {doc.id for doc in expected}
 
+    async def test_close_async_and_reopen(self, document_store: QdrantDocumentStore):
+        assert await document_store.count_documents_async() == 0
+        assert document_store._async_client is not None
+
+        await document_store.close_async()
+
+        assert document_store._async_client is None
+        assert await document_store.count_documents_async() == 0
+        assert document_store._async_client is not None
+
     @pytest.mark.asyncio
     async def test_write_documents_async(self, document_store: QdrantDocumentStore):
         docs = [Document(id="1")]

integrations/qdrant/tests/test_embedding_retriever.py

@@ -146,6 +146,26 @@ def test_from_dict_no_filter_policy(self):
         retriever = QdrantEmbeddingRetriever.from_dict(data)
         assert retriever._filter_policy == FilterPolicy.REPLACE  # defaults to REPLACE
 
+    def test_close(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        retriever = QdrantEmbeddingRetriever(document_store=mock_store)
+
+        retriever.close()
+
+        mock_store.close.assert_called_once_with()
+        assert retriever._document_store is mock_store
+
+    @pytest.mark.asyncio
+    async def test_close_async(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        mock_store.close_async = AsyncMock()
+        retriever = QdrantEmbeddingRetriever(document_store=mock_store)
+
+        await retriever.close_async()
+
+        mock_store.close_async.assert_awaited_once_with()
+        assert retriever._document_store is mock_store
+
     def test_run(self):
         mock_store = Mock(spec=QdrantDocumentStore)
         mock_store._query_by_embedding.return_value = [Document(content="doc", embedding=[0.1, 0.2])]

integrations/qdrant/tests/test_hybrid_retriever.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from haystack.dataclasses import Document, SparseEmbedding
@@ -146,6 +146,26 @@ def test_from_dict_no_filter_policy(self):
         assert retriever._group_by is None
         assert retriever._group_size is None
 
+    def test_close(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        retriever = QdrantHybridRetriever(document_store=mock_store)
+
+        retriever.close()
+
+        mock_store.close.assert_called_once_with()
+        assert retriever._document_store is mock_store
+
+    @pytest.mark.asyncio
+    async def test_close_async(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        mock_store.close_async = AsyncMock()
+        retriever = QdrantHybridRetriever(document_store=mock_store)
+
+        await retriever.close_async()
+
+        mock_store.close_async.assert_awaited_once_with()
+        assert retriever._document_store is mock_store
+
     def test_run(self):
         mock_store = Mock(spec=QdrantDocumentStore)
         sparse_embedding = SparseEmbedding(indices=[0, 1, 2, 3], values=[0.1, 0.8, 0.05, 0.33])

integrations/qdrant/tests/test_sparse_embedding_retriever.py

@@ -154,6 +154,26 @@ def test_from_dict_no_filter_policy(self):
         assert retriever._group_by is None
         assert retriever._group_size is None
 
+    def test_close(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        retriever = QdrantSparseEmbeddingRetriever(document_store=mock_store)
+
+        retriever.close()
+
+        mock_store.close.assert_called_once_with()
+        assert retriever._document_store is mock_store
+
+    @pytest.mark.asyncio
+    async def test_close_async(self):
+        mock_store = Mock(spec=QdrantDocumentStore)
+        mock_store.close_async = AsyncMock()
+        retriever = QdrantSparseEmbeddingRetriever(document_store=mock_store)
+
+        await retriever.close_async()
+
+        mock_store.close_async.assert_awaited_once_with()
+        assert retriever._document_store is mock_store
+
     def test_run(self):
         mock_store = Mock(spec=QdrantDocumentStore)
         sparse = SparseEmbedding(indices=[0, 5], values=[0.1, 0.7])