feat: Azure AI Search - add closing methods (#3662)

Author: anakin87

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/bm25_retriever.py

@@ -96,6 +96,12 @@ def from_dict(cls, data: dict[str, Any]) -> "AzureAISearchBM25Retriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(data["init_parameters"]["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()
+
     @component.output_types(documents=list[Document])
     def run(
         self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/embedding_retriever.py

@@ -93,6 +93,12 @@ def from_dict(cls, data: dict[str, Any]) -> "AzureAISearchEmbeddingRetriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(data["init_parameters"]["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()
+
     @component.output_types(documents=list[Document])
     def run(
         self, query_embedding: list[float], filters: dict[str, Any] | None = None, top_k: int | None = None

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/hybrid_retriever.py

@@ -96,6 +96,12 @@ def from_dict(cls, data: dict[str, Any]) -> "AzureAISearchHybridRetriever":
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(data["init_parameters"]["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()
+
     @component.output_types(documents=list[Document])
     def run(
         self,

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

@@ -4,6 +4,7 @@
 
 import logging as python_logging
 from collections.abc import Mapping, Sequence
+from contextlib import suppress
 from datetime import datetime
 from typing import Any
 
@@ -387,6 +388,19 @@ def from_dict(cls, data: dict[str, Any]) -> "AzureAISearchDocumentStore":
             data["init_parameters"]["vector_search_configuration"] = VectorSearch(vector_search_configuration)
         return default_from_dict(cls, data)
 
+    def close(self) -> None:
+        """
+        Release the associated synchronous resources.
+        """
+        if self._client is not None:
+            with suppress(Exception):
+                self._client.close()
+            self._client = None
+        if self._index_client is not None:
+            with suppress(Exception):
+                self._index_client.close()
+            self._index_client = None
+
     def count_documents(self) -> int:
         """
         Returns how many documents are present in the search index.

integrations/azure_ai_search/tests/test_bm25_retriever.py

@@ -97,6 +97,16 @@ def test_from_dict():
     assert retriever._filter_policy == FilterPolicy.REPLACE
 
 
+def test_close():
+    mock_store = Mock(spec=AzureAISearchDocumentStore)
+    retriever = AzureAISearchBM25Retriever(document_store=mock_store)
+
+    retriever.close()
+
+    mock_store.close.assert_called_once()
+    assert retriever._document_store is mock_store
+
+
 def test_run():
     mock_store = Mock(spec=AzureAISearchDocumentStore)
     mock_store._bm25_retrieval.return_value = [Document(content="Test doc")]

integrations/azure_ai_search/tests/test_document_store.py

@@ -261,6 +261,47 @@ def test_init():
     assert document_store._vector_search_configuration == DEFAULT_VECTOR_SEARCH
 
 
+def test_close():
+    store = AzureAISearchDocumentStore(
+        api_key=Secret.from_token("fake-api-key"),
+        azure_endpoint=Secret.from_token("fake-endpoint"),
+    )
+    client = Mock()
+    index_client = Mock()
+    store._client = client
+    store._index_client = index_client
+
+    store.close()
+
+    client.close.assert_called_once()
+    index_client.close.assert_called_once()
+    assert store._client is None
+    assert store._index_client is None
+
+    store.close()
+
+    client.close.assert_called_once()
+    index_client.close.assert_called_once()
+
+
+def test_close_is_exception_safe():
+    store = AzureAISearchDocumentStore(
+        api_key=Secret.from_token("fake-api-key"),
+        azure_endpoint=Secret.from_token("fake-endpoint"),
+    )
+    client = Mock()
+    client.close.side_effect = RuntimeError("boom")
+    index_client = Mock()
+    index_client.close.side_effect = RuntimeError("boom")
+    store._client = client
+    store._index_client = index_client
+
+    store.close()
+
+    assert store._client is None
+    assert store._index_client is None
+
+
 def test_token_credential_takes_priority_over_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
     monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
@@ -496,6 +537,13 @@ class TestDocumentStore(
     def assert_documents_are_equal(self, received: list[Document], expected: list[Document]):
         _assert_documents_are_equal(received, expected)
 
+    def test_close_and_reopen(self, document_store: AzureAISearchDocumentStore):
+        assert document_store.count_documents() == 0
+        document_store.close()
+        assert document_store._client is None
+        assert document_store._index_client is None
+        assert document_store.count_documents() == 0
+
     def test_write_documents(self, document_store: AzureAISearchDocumentStore):
         docs = [Document(id="1")]
         assert document_store.write_documents(docs) == 1

integrations/azure_ai_search/tests/test_embedding_retriever.py

@@ -100,6 +100,16 @@ def test_from_dict():
     assert retriever._filter_policy == FilterPolicy.REPLACE
 
 
+def test_close():
+    mock_store = Mock(spec=AzureAISearchDocumentStore)
+    retriever = AzureAISearchEmbeddingRetriever(document_store=mock_store)
+
+    retriever.close()
+
+    mock_store.close.assert_called_once()
+    assert retriever._document_store is mock_store
+
+
 def test_run():
     mock_store = Mock(spec=AzureAISearchDocumentStore)
     mock_store._embedding_retrieval.return_value = [Document(content="Test doc", embedding=[0.1, 0.2])]

integrations/azure_ai_search/tests/test_hybrid_retriever.py

@@ -100,6 +100,16 @@ def test_from_dict():
     assert retriever._filter_policy == FilterPolicy.REPLACE
 
 
+def test_close():
+    mock_store = Mock(spec=AzureAISearchDocumentStore)
+    retriever = AzureAISearchHybridRetriever(document_store=mock_store)
+
+    retriever.close()
+
+    mock_store.close.assert_called_once()
+    assert retriever._document_store is mock_store
+
+
 def test_run():
     mock_store = Mock(spec=AzureAISearchDocumentStore)
     mock_store._hybrid_retrieval.return_value = [Document(content="Test doc", embedding=[0.1, 0.2])]