feat: Add support for Azure TokenCredentials to AzureAISearchDocumentStore (#3014)

* feat: Add support for Azure TokenCredentials

* Fix mypy error

* Address review comments and add unit tests

---------

Co-authored-by: bogdankostic <bogdankostic@web.de>

Author: maxdswain

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

@@ -6,7 +6,7 @@
 from datetime import datetime
 from typing import Any
 
-from azure.core.credentials import AzureKeyCredential
+from azure.core.credentials import AzureKeyCredential, TokenCredential
 from azure.core.exceptions import (
     ClientAuthenticationError,
     HttpResponseError,
@@ -120,6 +120,7 @@ def __init__(
         metadata_fields: dict[str, SearchField | type] | None = None,
         vector_search_configuration: VectorSearch | None = None,
         include_search_metadata: bool = False,
+        azure_token_credential: TokenCredential | None = None,
         **index_creation_kwargs: Any,
     ) -> None:
         """
@@ -154,6 +155,8 @@ def __init__(
             in the returned documents. When set to True, the `meta` field of the returned
             documents will contain the @search.score, @search.reranker_score, @search.highlights,
             @search.captions, and other fields returned by Azure AI Search.
+        :param azure_token_credential: An Azure `TokenCredential` instance used to authenticate requests.
+            When provided, this takes priority over `api_key`.
         :param index_creation_kwargs: Optional keyword parameters to be passed to `SearchIndex` class
             during index creation. Some of the supported parameters:
                 - `semantic_search`: Defines semantic configuration of the search index. This parameter is needed
@@ -175,6 +178,7 @@ def __init__(
         self._metadata_fields = AzureAISearchDocumentStore._normalize_metadata_index_fields(metadata_fields)
         self._vector_search_configuration = vector_search_configuration or DEFAULT_VECTOR_SEARCH
         self._include_search_metadata = include_search_metadata
+        self._azure_token_credential = azure_token_credential
         self._index_creation_kwargs = index_creation_kwargs
 
     @property
@@ -183,7 +187,12 @@ def client(self) -> SearchClient:
         resolved_endpoint = self._azure_endpoint.resolve_value()
         resolved_key = self._api_key.resolve_value()
 
-        credential = AzureKeyCredential(resolved_key) if resolved_key else DefaultAzureCredential()
+        if self._azure_token_credential is not None:
+            credential: TokenCredential | AzureKeyCredential = self._azure_token_credential
+        elif resolved_key:
+            credential = AzureKeyCredential(resolved_key)
+        else:
+            credential = DefaultAzureCredential()
 
         # build a UserAgentPolicy to be used for the request
         ua_policy = UserAgentPolicy(user_agent=USER_AGENT)
@@ -316,6 +325,13 @@ def to_dict(self) -> dict[str, Any]:
         :returns:
             Dictionary with serialized data.
         """
+        if self._azure_token_credential:
+            logger.warning(
+                "AzureAISearchDocumentStore was initialized with `azure_token_credential`, "
+                "which cannot be serialized. It will be excluded from the serialized output "
+                "and must be provided again when deserializing."
+            )
+
         return default_to_dict(
             self,
             azure_endpoint=(self._azure_endpoint.to_dict() if self._azure_endpoint else None),

integrations/azure_ai_search/tests/test_document_store.py

@@ -2,12 +2,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 import os
 import random
 from datetime import datetime, timezone
 from unittest.mock import Mock, patch
 
 import pytest
+from azure.core.credentials import TokenCredential
 from azure.search.documents.indexes.models import (
     CustomAnalyzer,
     SearchableField,
@@ -135,6 +137,24 @@ def test_to_dict_with_params(monkeypatch):
     }
 
 
+def test_to_dict_emits_warning_when_token_credential_is_used(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
+    monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
+
+    mock_token_credential = Mock(spec=TokenCredential)
+    document_store = AzureAISearchDocumentStore(azure_token_credential=mock_token_credential)
+
+    with caplog.at_level(logging.WARNING):
+        result = document_store.to_dict()
+
+    assert "`azure_token_credential`, which cannot be serialized." in caplog.text
+
+    # token credential should not appear in the serialized output
+    assert "azure_token_credential" not in result["init_parameters"]
+
+
 def test_from_dict(monkeypatch):
     monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
     monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
@@ -246,6 +266,28 @@ def test_init():
     assert document_store._vector_search_configuration == DEFAULT_VECTOR_SEARCH
 
 
+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")
+
+    mock_token_credential = Mock(spec=TokenCredential)
+    document_store = AzureAISearchDocumentStore(azure_token_credential=mock_token_credential)
+
+    with patch(
+        "haystack_integrations.document_stores.azure_ai_search.document_store.SearchIndexClient"
+    ) as mock_index_client_cls:
+        mock_index_client = Mock()
+        mock_index_client.list_index_names.return_value = ["default"]
+        mock_index_client.get_index.return_value = Mock(fields=[])
+        mock_index_client.get_search_client.return_value = Mock()
+        mock_index_client_cls.return_value = mock_index_client
+
+        _ = document_store.client
+
+        _, kwargs = mock_index_client_cls.call_args
+        assert kwargs["credential"] is mock_token_credential
+
+
 def _build_mock_document_store_with_schema(index_fields):
     store = AzureAISearchDocumentStore(
         api_key=Secret.from_token("fake-api-key"),