From 99ea3df3e02a3d83fea0bd3a2a041b95f25d7e44 Mon Sep 17 00:00:00 2001
From: Rakshith Bhyravabhotla <rakshith.bhyravabhotla@gmail.com>
Date: Mon, 15 Aug 2022 12:39:20 -0700
Subject: [PATCH 1/4] Support national clouds in search

---
 .../azure-search-documents/CHANGELOG.md       |  1 +
 sdk/search/azure-search-documents/README.md   | 20 +++++++++++++++++++
 .../azure/search/documents/_search_client.py  |  6 +++++-
 .../_search_indexing_buffered_sender.py       |  6 +++++-
 .../azure/search/documents/_utils.py          | 16 ++++++++++-----
 .../documents/aio/_search_client_async.py     |  8 ++++++--
 .../_search_indexing_buffered_sender_async.py |  9 ++++++---
 .../search/documents/aio/_utils_async.py      | 15 --------------
 .../documents/indexes/_search_index_client.py |  7 +++++--
 .../indexes/_search_indexer_client.py         |  7 +++++--
 .../indexes/aio/_search_index_client.py       |  7 +++++--
 .../indexes/aio/_search_indexer_client.py     |  7 +++++--
 12 files changed, 74 insertions(+), 35 deletions(-)
 delete mode 100644 sdk/search/azure-search-documents/azure/search/documents/aio/_utils_async.py

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..9c8bf672aad7 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
+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..6abc33307b62 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
@@ -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_async_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..91c71cf7a3ba 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
@@ -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_async_authentication_policy(credential, audience=audience, is_async=True)
             self._client = _SearchServiceClient(
                 endpoint=endpoint,
                 authentication_policy=authentication_policy,

From ab02544f549eadc72dab71c48eab3190f83718c2 Mon Sep 17 00:00:00 2001
From: Rakshith Bhyravabhotla <rakshith.bhyravabhotla@gmail.com>
Date: Mon, 15 Aug 2022 13:38:24 -0700
Subject: [PATCH 2/4] Update sdk/search/azure-search-documents/README.md

---
 sdk/search/azure-search-documents/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdk/search/azure-search-documents/README.md b/sdk/search/azure-search-documents/README.md
index 9c8bf672aad7..aa57eafbe7e2 100644
--- a/sdk/search/azure-search-documents/README.md
+++ b/sdk/search/azure-search-documents/README.md
@@ -285,7 +285,7 @@ To authenticate in a [National Cloud](https://docs.microsoft.com/azure/active-di
 ```python
 # Create a SearchClient that will authenticate through AAD in the China national cloud.
 import os
-from azure.identity import DefaultAzureCredential
+from azure.identity import DefaultAzureCredential, AzureAuthorityHosts
 from azure.search.documents import SearchClient
 
 index_name = "hotels"

From 1b86e747f079bead22621a7d56a460c447242136 Mon Sep 17 00:00:00 2001
From: Rakshith Bhyravabhotla <rakshith.bhyravabhotla@gmail.com>
Date: Mon, 15 Aug 2022 15:13:22 -0700
Subject: [PATCH 3/4] lint

---
 .../search/documents/indexes/aio/_search_index_client.py      | 4 ++--
 .../search/documents/indexes/aio/_search_indexer_client.py    | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

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 6abc33307b62..554392b250a0 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,10 +12,10 @@
 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,
+    get_authentication_policy,
 )
 from ..._api_versions import DEFAULT_VERSION
 from ..._headers_mixin import HeadersMixin
@@ -69,7 +69,7 @@ def __init__(
             )  # type: _SearchServiceClient
         else:
             self._aad = True
-            authentication_policy = get_async_authentication_policy(credential, audience=audience, is_async=True)
+            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 91c71cf7a3ba..341ab4ef030a 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,6 +14,7 @@
 from .._utils import (
     get_access_conditions,
     normalize_endpoint,
+    get_authentication_policy,
 )
 from ..models import (
     SearchIndexerDataSourceConnection,
@@ -21,7 +22,6 @@
 from ..._api_versions import DEFAULT_VERSION
 from ..._headers_mixin import HeadersMixin
 from ..._version import SDK_MONIKER
-from ...aio._utils_async import get_async_authentication_policy
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
@@ -65,7 +65,7 @@ def __init__(
             )  # type: _SearchServiceClient
         else:
             self._aad = True
-            authentication_policy = get_async_authentication_policy(credential, audience=audience, is_async=True)
+            authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True)
             self._client = _SearchServiceClient(
                 endpoint=endpoint,
                 authentication_policy=authentication_policy,

From 075e4df9d940606786fee4eae211cd1e7b557dcd Mon Sep 17 00:00:00 2001
From: Rakshith Bhyravabhotla <rakshith.bhyravabhotla@gmail.com>
Date: Tue, 16 Aug 2022 13:33:52 -0700
Subject: [PATCH 4/4] lint

---
 .../search/documents/indexes/aio/_search_index_client.py      | 2 +-
 .../search/documents/indexes/aio/_search_indexer_client.py    | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

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 554392b250a0..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
@@ -15,10 +15,10 @@
 from .._utils import (
     get_access_conditions,
     normalize_endpoint,
-    get_authentication_policy,
 )
 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,
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 341ab4ef030a..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,13 +14,13 @@
 from .._utils import (
     get_access_conditions,
     normalize_endpoint,
-    get_authentication_policy,
-)
+    )
 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
 
 if TYPE_CHECKING: