Update placement of fallback logic.

Author: aavasthy

sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policy.py

@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # -------------------------------------------------------------------------
-from typing import TypeVar, Any, MutableMapping, cast
+from typing import TypeVar, Any, MutableMapping, cast, Optional
 
 from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
@@ -19,7 +19,7 @@
 class CosmosBearerTokenCredentialPolicy(BearerTokenCredentialPolicy):
     AadDefaultScope = "https://cosmos.azure.com/.default"
 
-    def __init__(self, credential, account_scope: str, override_scope: str = None):
+    def __init__(self, credential, account_scope: str, override_scope: Optional[str] = None):
         self._account_scope = account_scope
         self._override_scope = override_scope
         self._current_scope = override_scope or account_scope
@@ -41,25 +41,11 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
 
         :param ~azure.core.pipeline.PipelineRequest request: the request
         """
-        self.authorize_request(request)
-        super().on_request(request)
-        # The None-check for self._token is done in the parent on_request
-        self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)
-
-    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
-        """Acquire a token from the credential and authorize the request with it.
-
-        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
-        authorize future requests.
-
-        :param ~azure.core.pipeline.PipelineRequest request: the request
-        :param str scopes: required scopes of authentication
-        """
         tried_fallback = False
         while True:
             try:
-                super().authorize_request(request, self._current_scope, **kwargs)
-                # The None-check for self._token is done in the parent authorize_request
+                super().on_request(request)
+                # The None-check for self._token is done in the parent on_request
                 self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)
                 break
             except Exception as ex:
@@ -70,7 +56,22 @@ def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes:
                         self._current_scope != self.AadDefaultScope and
                         "AADSTS500011" in str(ex)
                 ):
+                    self._scopes = (self.AadDefaultScope,)
                     self._current_scope = self.AadDefaultScope
                     tried_fallback = True
                     continue
                 raise
+
+    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
+        """Acquire a token from the credential and authorize the request with it.
+
+        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
+        authorize future requests.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        :param str scopes: required scopes of authentication
+        """
+
+        super().authorize_request(request, self._current_scope, **kwargs)
+        # The None-check for self._token is done in the parent authorize_request
+        self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py

@@ -4,7 +4,7 @@
 # license information.
 # -------------------------------------------------------------------------
 
-from typing import Any, MutableMapping, TypeVar, cast
+from typing import Any, MutableMapping, TypeVar, cast, Optional
 
 from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
 from azure.core.pipeline import PipelineRequest
@@ -20,7 +20,7 @@
 class AsyncCosmosBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
     AadDefaultScope = "https://cosmos.azure.com/.default"
 
-    def __init__(self, credential, account_scope: str, override_scope: str):
+    def __init__(self, credential, account_scope: str, override_scope: Optional[str] = None):
         self._account_scope = account_scope
         self._override_scope = override_scope
         self._current_scope = override_scope or account_scope
@@ -42,25 +42,11 @@ async def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         :type request: ~azure.core.pipeline.PipelineRequest
         :raises: :class:`~azure.core.exceptions.ServiceRequestError`
         """
-        await self.authorize_request(request)
-        await super().on_request(request)
-        # The None-check for self._token is done in the parent on_request
-        self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)
-
-    async def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
-        """Acquire a token from the credential and authorize the request with it.
-
-        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
-        authorize future requests.
-
-        :param ~azure.core.pipeline.PipelineRequest request: the request
-        :param str scopes: required scopes of authentication
-        """
         tried_fallback = False
         while True:
             try:
-                await super().authorize_request(request, self._current_scope, **kwargs)
-                # The None-check for self._token is done in the parent authorize_request
+                await super().on_request(request)
+                # The None-check for self._token is done in the parent on_request
                 self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)
                 break
             except Exception as ex:
@@ -71,7 +57,22 @@ async def authorize_request(self, request: PipelineRequest[HTTPRequestType], *sc
                         self._current_scope != self.AadDefaultScope and
                         "AADSTS500011" in str(ex)
                 ):
+                    self._scopes = (self.AadDefaultScope,)
                     self._current_scope = self.AadDefaultScope
                     tried_fallback = True
                     continue
                 raise
+
+    async def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
+        """Acquire a token from the credential and authorize the request with it.
+
+        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
+        authorize future requests.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        :param str scopes: required scopes of authentication
+        """
+
+        await super().authorize_request(request, self._current_scope, **kwargs)
+        # The None-check for self._token is done in the parent authorize_request
+        self._update_headers(request.http_request.headers, cast(AccessToken, self._token).token)