diff --git a/sdk/search/azure-search-documents/CHANGELOG.md b/sdk/search/azure-search-documents/CHANGELOG.md index 67ab923a2ca9..84b2f2aca1c3 100644 --- a/sdk/search/azure-search-documents/CHANGELOG.md +++ b/sdk/search/azure-search-documents/CHANGELOG.md @@ -3,6 +3,7 @@ ## 11.3.0b9 (Unreleased) ### Features Added +- Added support for other national clouds. ### Breaking Changes diff --git a/sdk/search/azure-search-documents/README.md b/sdk/search/azure-search-documents/README.md index f2f9e3152c7b..aa57eafbe7e2 100644 --- a/sdk/search/azure-search-documents/README.md +++ b/sdk/search/azure-search-documents/README.md @@ -275,6 +275,26 @@ result = search_client.upload_documents(documents=[DOCUMENT]) print("Upload of new document succeeded: {}".format(result[0].succeeded)) ``` +### Authenticate in a National Cloud + +To authenticate in a [National Cloud](https://docs.microsoft.com/azure/active-directory/develop/authentication-national-cloud), you will need to make the following additions to your client configuration: + +- Set the `AuthorityHost` in the credential options or via the `AZURE_AUTHORITY_HOST` environment variable +- Set the `audience` in `SearchClient`, `SearchIndexClient`, or `SearchIndexerClient` + +```python +# Create a SearchClient that will authenticate through AAD in the China national cloud. +import os +from azure.identity import DefaultAzureCredential, AzureAuthorityHosts +from azure.search.documents import SearchClient + +index_name = "hotels" +endpoint = os.environ["SEARCH_ENDPOINT"] +key = os.environ["SEARCH_API_KEY"] +credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_CHINA) + +search_client = SearchClient(endpoint, index_name, crdential=credential, audience="https://search.azure.cn") +``` ### Retrieving a specific document from your index diff --git a/sdk/search/azure-search-documents/azure/search/documents/_search_client.py b/sdk/search/azure-search-documents/azure/search/documents/_search_client.py index 09e2bbee6f4e..caf6f90f352b 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/_search_client.py +++ b/sdk/search/azure-search-documents/azure/search/documents/_search_client.py @@ -63,6 +63,9 @@ class SearchClient(HeadersMixin): :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.TokenCredential :keyword str api_version: The Search API version to use for requests. + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. .. admonition:: Example: @@ -83,6 +86,7 @@ def __init__(self, endpoint, index_name, credential, **kwargs): self._endpoint = endpoint # type: str self._index_name = index_name # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = SearchIndexClient( @@ -94,7 +98,7 @@ def __init__(self, endpoint, index_name, credential, **kwargs): ) # type: SearchIndexClient else: self._aad = True - authentication_policy = get_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience) self._client = SearchIndexClient( endpoint=endpoint, index_name=index_name, diff --git a/sdk/search/azure-search-documents/azure/search/documents/_search_indexing_buffered_sender.py b/sdk/search/azure-search-documents/azure/search/documents/_search_indexing_buffered_sender.py index caa03eb76e88..8fba34a89ebd 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/_search_indexing_buffered_sender.py +++ b/sdk/search/azure-search-documents/azure/search/documents/_search_indexing_buffered_sender.py @@ -50,6 +50,9 @@ class SearchIndexingBufferedSender(SearchIndexingBufferedSenderBase, HeadersMixi is a IndexAction removed from the queue (succeeds or fails). This may be called from main thread or a worker thread. :keyword str api_version: The Search API version to use for requests. + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ # pylint: disable=too-many-instance-attributes @@ -60,6 +63,7 @@ def __init__(self, endpoint, index_name, credential, **kwargs): endpoint=endpoint, index_name=index_name, credential=credential, **kwargs ) self._index_documents_batch = IndexDocumentsBatch() + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = SearchIndexClient( @@ -71,7 +75,7 @@ def __init__(self, endpoint, index_name, credential, **kwargs): ) # type: SearchIndexClient else: self._aad = True - authentication_policy = get_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience) self._client = SearchIndexClient( endpoint=endpoint, index_name=index_name, diff --git a/sdk/search/azure-search-documents/azure/search/documents/_utils.py b/sdk/search/azure-search-documents/azure/search/documents/_utils.py index bf9cad202f9a..c7c8aa783c3a 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/_utils.py +++ b/sdk/search/azure-search-documents/azure/search/documents/_utils.py @@ -3,9 +3,9 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- -from azure.core.pipeline import policies +from azure.core.pipeline.policies import BearerTokenCredentialPolicy, AsyncBearerTokenCredentialPolicy -CREDENTIAL_SCOPES = ["https://search.azure.com/.default"] +DEFAULT_AUDIENCE = "https://search.azure.com" def is_retryable_status_code(status_code): @@ -13,8 +13,14 @@ def is_retryable_status_code(status_code): return status_code in [422, 409, 503] -def get_authentication_policy(credential): - authentication_policy = policies.BearerTokenCredentialPolicy( - credential, *CREDENTIAL_SCOPES +def get_authentication_policy(credential, **kwargs): + audience = kwargs.get('audience', None) + is_async = kwargs.get('is_async', False) + if not audience: + audience = DEFAULT_AUDIENCE + scope = audience.rstrip('/') + '/.default' + _policy = BearerTokenCredentialPolicy if not is_async else AsyncBearerTokenCredentialPolicy + authentication_policy = _policy( + credential, scope ) return authentication_policy diff --git a/sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py b/sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py index d3c4af8cc4f4..8f715b576e89 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py +++ b/sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py @@ -9,7 +9,7 @@ from azure.core.credentials import AzureKeyCredential from azure.core.tracing.decorator_async import distributed_trace_async from ._paging import AsyncSearchItemPaged, AsyncSearchPageIterator -from ._utils_async import get_async_authentication_policy +from .._utils import get_authentication_policy from .._generated.aio import SearchClient as SearchIndexClient from .._generated.models import IndexingResult from .._search_documents_error import RequestEntityTooLargeError @@ -35,6 +35,9 @@ class SearchClient(HeadersMixin): :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials_async.AsyncTokenCredential :keyword str api_version: The Search API version to use for requests. + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. .. admonition:: Example: @@ -60,6 +63,7 @@ def __init__( self._endpoint = endpoint # type: str self._index_name = index_name # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = SearchIndexClient( @@ -71,7 +75,7 @@ def __init__( ) # type: SearchIndexClient else: self._aad = True - authentication_policy = get_async_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True) self._client = SearchIndexClient( endpoint=endpoint, index_name=index_name, diff --git a/sdk/search/azure-search-documents/azure/search/documents/aio/_search_indexing_buffered_sender_async.py b/sdk/search/azure-search-documents/azure/search/documents/aio/_search_indexing_buffered_sender_async.py index 4cae7a0152f5..ab434a8709c4 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/aio/_search_indexing_buffered_sender_async.py +++ b/sdk/search/azure-search-documents/azure/search/documents/aio/_search_indexing_buffered_sender_async.py @@ -10,8 +10,7 @@ from azure.core.tracing.decorator_async import distributed_trace_async from azure.core.exceptions import ServiceResponseTimeoutError from ._timer import Timer -from ._utils_async import get_async_authentication_policy -from .._utils import is_retryable_status_code +from .._utils import is_retryable_status_code, get_authentication_policy from .._search_indexing_buffered_sender_base import SearchIndexingBufferedSenderBase from .._generated.aio import SearchClient as SearchIndexClient from .._generated.models import IndexingResult @@ -49,6 +48,9 @@ class SearchIndexingBufferedSender(SearchIndexingBufferedSenderBase, HeadersMixi :keyword callable on_remove: If it is set, the client will call corresponding methods when there is a IndexAction removed from the queue (succeeds or fails). :keyword str api_version: The Search API version to use for requests. + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ # pylint: disable=too-many-instance-attributes @@ -64,6 +66,7 @@ def __init__( endpoint=endpoint, index_name=index_name, credential=credential, **kwargs ) self._index_documents_batch = IndexDocumentsBatch() + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = SearchIndexClient( @@ -75,7 +78,7 @@ def __init__( ) # type: SearchIndexClient else: self._aad = True - authentication_policy = get_async_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True) self._client = SearchIndexClient( endpoint=endpoint, index_name=index_name, diff --git a/sdk/search/azure-search-documents/azure/search/documents/aio/_utils_async.py b/sdk/search/azure-search-documents/azure/search/documents/aio/_utils_async.py deleted file mode 100644 index 697aae0f9adb..000000000000 --- a/sdk/search/azure-search-documents/azure/search/documents/aio/_utils_async.py +++ /dev/null @@ -1,15 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See License.txt in the project root for -# license information. -# -------------------------------------------------------------------------- -from azure.core.pipeline import policies - -from .._utils import CREDENTIAL_SCOPES - - -def get_async_authentication_policy(credential): - authentication_policy = policies.AsyncBearerTokenCredentialPolicy( - credential, *CREDENTIAL_SCOPES - ) - return authentication_policy diff --git a/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_index_client.py b/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_index_client.py index 0299ed345ac1..7078d029d072 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_index_client.py +++ b/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_index_client.py @@ -37,7 +37,9 @@ class SearchIndexClient(HeadersMixin): # pylint:disable=too-many-public-methods :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.TokenCredential :keyword str api_version: The Search API version to use for requests. - + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ _ODATA_ACCEPT = "application/json;odata.metadata=minimal" # type: str @@ -48,6 +50,7 @@ def __init__(self, endpoint, credential, **kwargs): self._api_version = kwargs.pop("api_version", DEFAULT_VERSION) self._endpoint = normalize_endpoint(endpoint) # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = _SearchServiceClient( @@ -58,7 +61,7 @@ def __init__(self, endpoint, credential, **kwargs): ) # type: _SearchServiceClient else: self._aad = True - authentication_policy = get_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience) self._client = _SearchServiceClient( endpoint=endpoint, authentication_policy=authentication_policy, diff --git a/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_indexer_client.py b/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_indexer_client.py index ac21a5cf7e3b..44196185f631 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_indexer_client.py +++ b/sdk/search/azure-search-documents/azure/search/documents/indexes/_search_indexer_client.py @@ -40,7 +40,9 @@ class SearchIndexerClient(HeadersMixin): # pylint: disable=R0904 :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.TokenCredential :keyword str api_version: The Search API version to use for requests. - + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ _ODATA_ACCEPT = "application/json;odata.metadata=minimal" # type: str @@ -51,6 +53,7 @@ def __init__(self, endpoint, credential, **kwargs): self._api_version = kwargs.pop("api_version", DEFAULT_VERSION) self._endpoint = normalize_endpoint(endpoint) # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = _SearchServiceClient( @@ -61,7 +64,7 @@ def __init__(self, endpoint, credential, **kwargs): ) # type: _SearchServiceClient else: self._aad = True - authentication_policy = get_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience) self._client = _SearchServiceClient( endpoint=endpoint, authentication_policy=authentication_policy, diff --git a/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_index_client.py b/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_index_client.py index 6ecc8182f794..3b3de7cd5def 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_index_client.py +++ b/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_index_client.py @@ -12,13 +12,13 @@ from azure.core.async_paging import AsyncItemPaged from .._generated.aio import SearchClient as _SearchServiceClient from ...aio._search_client_async import SearchClient -from ...aio._utils_async import get_async_authentication_policy from .._utils import ( get_access_conditions, normalize_endpoint, ) from ..._api_versions import DEFAULT_VERSION from ..._headers_mixin import HeadersMixin +from ..._utils import get_authentication_policy from ..._version import SDK_MONIKER from ..models import ( SearchIndex, @@ -42,7 +42,9 @@ class SearchIndexClient(HeadersMixin): # pylint:disable=too-many-public-methods :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials_async.AsyncTokenCredential :keyword str api_version: The Search API version to use for requests. - + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ _ODATA_ACCEPT = "application/json;odata.metadata=minimal" # type: str @@ -56,6 +58,7 @@ def __init__( self._api_version = kwargs.pop("api_version", DEFAULT_VERSION) self._endpoint = normalize_endpoint(endpoint) # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = _SearchServiceClient( @@ -66,7 +69,7 @@ def __init__( ) # type: _SearchServiceClient else: self._aad = True - authentication_policy = get_async_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True) self._client = _SearchServiceClient( endpoint=endpoint, authentication_policy=authentication_policy, diff --git a/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_indexer_client.py b/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_indexer_client.py index 89f451b5bbca..9ed924cb83dd 100644 --- a/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_indexer_client.py +++ b/sdk/search/azure-search-documents/azure/search/documents/indexes/aio/_search_indexer_client.py @@ -14,14 +14,14 @@ from .._utils import ( get_access_conditions, normalize_endpoint, -) + ) from ..models import ( SearchIndexerDataSourceConnection, ) from ..._api_versions import DEFAULT_VERSION from ..._headers_mixin import HeadersMixin +from ..._utils import get_authentication_policy from ..._version import SDK_MONIKER -from ...aio._utils_async import get_async_authentication_policy if TYPE_CHECKING: # pylint:disable=unused-import,ungrouped-imports @@ -38,7 +38,9 @@ class SearchIndexerClient(HeadersMixin): # pylint: disable=R0904 :param credential: A credential to authorize search client requests :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials_async.AsyncTokenCredential :keyword str api_version: The Search API version to use for requests. - + :keyword str audience: sets the Audience to use for authentication with Azure Active Directory (AAD). The + audience is not considered when using a shared key. If audience is not provided, the public cloud audience + will be assumed. """ _ODATA_ACCEPT = "application/json;odata.metadata=minimal" # type: str @@ -52,6 +54,7 @@ def __init__( self._api_version = kwargs.pop("api_version", DEFAULT_VERSION) self._endpoint = normalize_endpoint(endpoint) # type: str self._credential = credential + audience = kwargs.pop("audience", None) if isinstance(credential, AzureKeyCredential): self._aad = False self._client = _SearchServiceClient( @@ -62,7 +65,7 @@ def __init__( ) # type: _SearchServiceClient else: self._aad = True - authentication_policy = get_async_authentication_policy(credential) + authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True) self._client = _SearchServiceClient( endpoint=endpoint, authentication_policy=authentication_policy,