Updates based on code reviews

Author: aavasthy

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

@@ -5,19 +5,24 @@
 # -------------------------------------------------------------------------
 from typing import TypeVar, Any, MutableMapping, cast, Optional
 
+from azure.cosmos import _constants as Constants
 from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.core.pipeline.transport import HttpRequest as LegacyHttpRequest
 from azure.core.rest import HttpRequest
 from azure.core.credentials import AccessToken
+from azure.core.exceptions import HttpResponseError
 
 from .http_constants import HttpHeaders
 
 HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
 
-
+# NOTE: This class accesses protected members (_scopes, _token) of the parent class
+# to implement fallback and scope-switching logic not exposed by the public API.
+# Composition was considered, but still required accessing protected members, so inheritance is retained
+# for seamless Azure SDK pipeline integration.
 class CosmosBearerTokenCredentialPolicy(BearerTokenCredentialPolicy):
-    AadDefaultScope = "https://cosmos.azure.com/.default"
+    AadDefaultScope = Constants._Constants.AAD_DEFAULT_SCOPE
 
     def __init__(self, credential, account_scope: str, override_scope: Optional[str] = None):
         self._account_scope = account_scope
@@ -48,7 +53,7 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
                 # 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:
+            except HttpResponseError as ex:
                 # Only fallback if not using override, not already tried, and error is AADSTS500011
                 if (
                         not self._override_scope and

sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py

@@ -56,6 +56,7 @@ class _Constants:
     CIRCUIT_BREAKER_ENABLED_CONFIG: str =  "AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"
     CIRCUIT_BREAKER_ENABLED_CONFIG_DEFAULT: str = "False"
     AAD_SCOPE_OVERRIDE: str = "AZURE_COSMOS_AAD_SCOPE_OVERRIDE"
+    AAD_DEFAULT_SCOPE: str = "https://cosmos.azure.com/.default"
 
     # Database Account Retry Policy constants
     AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES: str = "AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES"

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -118,6 +118,7 @@ def __init__( # pylint: disable=too-many-statements
         auth: CredentialDict,
         connection_policy: Optional[ConnectionPolicy] = None,
         consistency_level: Optional[str] = None,
+        audience: Optional[str] = None,
         **kwargs: Any
     ) -> None:
         """
@@ -132,7 +133,8 @@ def __init__( # pylint: disable=too-many-statements
             The connection policy for the client.
         :param documents.ConsistencyLevel consistency_level:
             The default consistency policy for client operations.
-
+        :param str audience:
+            The overridden scope value.
         """
         self.client_id = str(uuid.uuid4())
         self.url_connection = url_connection
@@ -204,7 +206,7 @@ def __init__( # pylint: disable=too-many-statements
 
         credentials_policy = None
         if self.aad_credentials:
-            scope_override = os.environ.get(Constants.AAD_SCOPE_OVERRIDE, "")
+            scope_override = audience or os.environ.get(Constants.AAD_SCOPE_OVERRIDE, "")
             account_scope = base.create_scope_from_url(self.url_connection)
             credentials_policy = CosmosBearerTokenCredentialPolicy(
                 self.aad_credentials,

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

@@ -6,19 +6,24 @@
 
 from typing import Any, MutableMapping, TypeVar, cast, Optional
 
+from azure.cosmos import _constants as Constants
 from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
 from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.transport import HttpRequest as LegacyHttpRequest
 from azure.core.rest import HttpRequest
 from azure.core.credentials import AccessToken
+from azure.core.exceptions import HttpResponseError
 
 from ..http_constants import HttpHeaders
 
 HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
 
-
+# NOTE: This class accesses protected members (_scopes, _token) of the parent class
+# to implement fallback and scope-switching logic not exposed by the public API.
+# Composition was considered, but still required accessing protected members, so inheritance is retained
+# for seamless Azure SDK pipeline integration.
 class AsyncCosmosBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
-    AadDefaultScope = "https://cosmos.azure.com/.default"
+    AadDefaultScope = Constants._Constants.AAD_DEFAULT_SCOPE
 
     def __init__(self, credential, account_scope: str, override_scope: Optional[str] = None):
         self._account_scope = account_scope
@@ -49,7 +54,7 @@ async def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
                 # 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:
+            except HttpResponseError as ex:
                 # Only fallback if not using override, not already tried, and error is AADSTS500011
                 if (
                         not self._override_scope and

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

@@ -124,6 +124,7 @@ def __init__( # pylint: disable=too-many-statements
             auth: CredentialDict,
             connection_policy: Optional[ConnectionPolicy] = None,
             consistency_level: Optional[str] = None,
+            audience: Optional[str] = None,
             **kwargs: Any
     ) -> None:
         """
@@ -138,7 +139,8 @@ def __init__( # pylint: disable=too-many-statements
             The connection policy for the client.
         :param documents.ConsistencyLevel consistency_level:
             The default consistency policy for client operations.
-
+        :param str audience:
+            The overridden scope value.
         """
         self.client_id = str(uuid.uuid4())
         self.url_connection = url_connection
@@ -212,7 +214,7 @@ def __init__( # pylint: disable=too-many-statements
 
         credentials_policy = None
         if self.aad_credentials:
-            scope_override = os.environ.get(Constants.AAD_SCOPE_OVERRIDE, "")
+            scope_override = audience or os.environ.get(Constants.AAD_SCOPE_OVERRIDE, "")
             account_scope = base.create_scope_from_url(self.url_connection)
             credentials_policy = AsyncCosmosBearerTokenCredentialPolicy(
                 self.aad_credentials,

sdk/cosmos/azure-cosmos/tests/test_aad.py

@@ -14,7 +14,7 @@
 import azure.cosmos.cosmos_client as cosmos_client
 import test_config
 from azure.cosmos import DatabaseProxy, ContainerProxy, exceptions
-
+from azure.core.exceptions import HttpResponseError
 
 def _remove_padding(encoded_string):
     while encoded_string.endswith("="):
@@ -216,7 +216,7 @@ def __init__(self):
             def get_token(self, *scopes, **kwargs):
                 self.call_count += 1
                 if self.call_count == 1:
-                    raise Exception("AADSTS500011: Simulated error for fallback")
+                    raise HttpResponseError(message="AADSTS500011: Simulated error for fallback")
                 return super().get_token(*scopes, **kwargs)
 
         def action(scopes_captured):

sdk/cosmos/azure-cosmos/tests/test_aad_async.py

@@ -14,7 +14,7 @@
 import test_config
 from azure.cosmos import exceptions
 from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
-
+from azure.core.exceptions import HttpResponseError
 
 def _remove_padding(encoded_string):
     while encoded_string.endswith("="):
@@ -243,7 +243,7 @@ def __init__(self):
             async def get_token(self, *scopes, **kwargs):
                 self.call_count += 1
                 if self.call_count == 1:
-                    raise Exception("AADSTS500011: Simulated error for fallback")
+                    raise HttpResponseError(message="AADSTS500011: Simulated error for fallback")
                 return await super().get_token(*scopes, **kwargs)
 
         async def action(scopes_captured):