Azure · Rakshith Bhyravabhotla (rakshith91) · Aug 16, 2022 · Aug 15, 2022 · Aug 15, 2022 · Aug 15, 2022
@@ -3,6 +3,7 @@
 ## 11.3.0b9 (Unreleased)
 
 ### Features Added
+- Added support for other national clouds.
 
 ### Breaking Changes
 

@@ -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
 

@@ -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,

@@ -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,

@@ -3,18 +3,24 @@
 # 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):
     # type: (int) -> bool
     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
@@ -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,

@@ -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,

@@ -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,

@@ -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,

@@ -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,

@@ -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,