Centralize credential pipeline construction (Azure#19864)

Author: chlowell

sdk/identity/azure-identity/CHANGELOG.md

@@ -16,6 +16,11 @@
   ([#19943](https://github.com/Azure/azure-sdk-for-python/issues/19943))
 
 ### Other Changes
+- Added `CustomHookPolicy` to credential HTTP pipelines. This allows applications
+  to initialize credentials with `raw_request_hook` and `raw_response_hook`
+  keyword arguments. The value of these arguments should be a callback taking a
+  `PipelineRequest` and `PipelineResponse`, respectively. For example:
+  `ManagedIdentityCredential(raw_request_hook=lambda request: print(request.http_request.url))`
 - Reduced redundant `ChainedTokenCredential` and `DefaultAzureCredential`
   logging. On Python 3.7+, credentials invoked by these classes now log debug
   rather than info messages.
@@ -25,6 +30,7 @@
   fails
   ([#19989](https://github.com/Azure/azure-sdk-for-python/issues/19989))
 
+
 ## 1.7.0b2 (2021-07-08)
 ### Features Added
 - `InteractiveBrowserCredential` keyword argument `login_hint` enables

sdk/identity/azure-identity/azure/identity/_credentials/azure_arc.py

@@ -8,24 +8,16 @@
 
 from azure.core.exceptions import ClientAuthenticationError
 from azure.core.pipeline.transport import HttpRequest
-from azure.core.pipeline.policies import (
-    DistributedTracingPolicy,
-    HttpLoggingPolicy,
-    HTTPPolicy,
-    UserAgentPolicy,
-    NetworkTraceLoggingPolicy,
-)
+from azure.core.pipeline.policies import HTTPPolicy
 
 from .. import CredentialUnavailableError
 from .._constants import EnvironmentVariables
-from .._internal.managed_identity_client import ManagedIdentityClient, _get_configuration
+from .._internal.managed_identity_client import ManagedIdentityClient
 from .._internal.get_token_mixin import GetTokenMixin
-from .._internal.user_agent import USER_AGENT
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, List, Optional, Union
-    from azure.core.configuration import Configuration
+    from typing import Any, Optional, Union
     from azure.core.credentials import AccessToken
     from azure.core.pipeline import PipelineRequest, PipelineResponse
     from azure.core.pipeline.policies import SansIOHTTPPolicy
@@ -42,10 +34,8 @@ def __init__(self, **kwargs):
         imds = os.environ.get(EnvironmentVariables.IMDS_ENDPOINT)
         self._available = url and imds
         if self._available:
-            config = _get_configuration()
-
             self._client = ManagedIdentityClient(
-                policies=_get_policies(config),
+                _per_retry_policies=[ArcChallengeAuthPolicy()],
                 request_factory=functools.partial(_get_request, url),
                 **kwargs
             )
@@ -67,19 +57,6 @@ def _request_token(self, *scopes, **kwargs):
         return self._client.request_token(*scopes, **kwargs)
 
 
-def _get_policies(config, **kwargs):
-    # type: (Configuration, **Any) -> List[PolicyType]
-    return [
-        UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs),
-        config.proxy_policy,
-        config.retry_policy,
-        ArcChallengeAuthPolicy(),
-        NetworkTraceLoggingPolicy(**kwargs),
-        DistributedTracingPolicy(**kwargs),
-        HttpLoggingPolicy(**kwargs),
-    ]
-
-
 def _get_request(url, scope, identity_config):
     # type: (str, str, dict) -> HttpRequest
     if identity_config:

sdk/identity/azure-identity/azure/identity/_internal/aad_client.py

@@ -5,30 +5,16 @@
 import time
 from typing import TYPE_CHECKING
 
-from azure.core.configuration import Configuration
-from azure.core.pipeline import Pipeline
-from azure.core.pipeline.policies import (
-    NetworkTraceLoggingPolicy,
-    RetryPolicy,
-    ProxyPolicy,
-    UserAgentPolicy,
-    DistributedTracingPolicy,
-    HttpLoggingPolicy,
-)
-
 from .aad_client_base import AadClientBase
-from .user_agent import USER_AGENT
+from .._internal.pipeline import build_pipeline
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Iterable, List, Optional, Union
+    from typing import Any, Iterable, Optional
     from azure.core.credentials import AccessToken
-    from azure.core.pipeline.policies import HTTPPolicy, SansIOHTTPPolicy
-    from azure.core.pipeline.transport import HttpTransport
+    from azure.core.pipeline import Pipeline
     from .._internal import AadClientCertificate
 
-    Policy = Union[HTTPPolicy, SansIOHTTPPolicy]
-
 
 class AadClient(AadClientBase):
     def obtain_token_by_authorization_code(self, scopes, code, redirect_uri, client_secret=None, **kwargs):
@@ -62,30 +48,6 @@ def obtain_token_by_refresh_token(self, scopes, refresh_token, **kwargs):
         return self._process_response(response, now)
 
     # pylint:disable=no-self-use
-    def _build_pipeline(self, config=None, policies=None, transport=None, **kwargs):
-        # type: (Optional[Configuration], Optional[List[Policy]], Optional[HttpTransport], **Any) -> Pipeline
-        config = config or _create_config(**kwargs)
-        policies = policies or [
-            config.user_agent_policy,
-            config.proxy_policy,
-            config.retry_policy,
-            config.logging_policy,
-            DistributedTracingPolicy(**kwargs),
-            HttpLoggingPolicy(**kwargs),
-        ]
-        if not transport:
-            from azure.core.pipeline.transport import RequestsTransport
-
-            transport = RequestsTransport(**kwargs)
-
-        return Pipeline(transport=transport, policies=policies)
-
-
-def _create_config(**kwargs):
-    # type: (**Any) -> Configuration
-    config = Configuration(**kwargs)
-    config.logging_policy = NetworkTraceLoggingPolicy(**kwargs)
-    config.retry_policy = RetryPolicy(**kwargs)
-    config.proxy_policy = ProxyPolicy(**kwargs)
-    config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs)
-    return config
+    def _build_pipeline(self, **kwargs):
+        # type: (**Any) -> Pipeline
+        return build_pipeline(**kwargs)

sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py

@@ -93,12 +93,14 @@ def obtain_token_by_refresh_token(self, scopes, refresh_token, **kwargs):
         pass
 
     @abc.abstractmethod
-    def _build_pipeline(self, config=None, policies=None, transport=None, **kwargs):
+    def _build_pipeline(self, **kwargs):
         pass
 
     def _process_response(self, response, request_time):
         # type: (PipelineResponse, int) -> AccessToken
-        content = ContentDecodePolicy.deserialize_from_http_generics(response.http_response)
+        content = response.context.get(
+            ContentDecodePolicy.CONTEXT_NAME
+        ) or ContentDecodePolicy.deserialize_from_http_generics(response.http_response)
 
         if response.http_request.body.get("grant_type") == "refresh_token":
             if content.get("error") == "invalid_grant":

sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py

@@ -9,22 +9,11 @@
 from msal import TokenCache
 import six
 
-from azure.core.configuration import Configuration
 from azure.core.credentials import AccessToken
 from azure.core.exceptions import ClientAuthenticationError, DecodeError
-from azure.core.pipeline import Pipeline
-from azure.core.pipeline.policies import (
-    ContentDecodePolicy,
-    DistributedTracingPolicy,
-    HeadersPolicy,
-    HttpLoggingPolicy,
-    UserAgentPolicy,
-    RetryPolicy,
-    NetworkTraceLoggingPolicy,
-)
-from azure.identity._internal import _scopes_to_resource
-
-from .user_agent import USER_AGENT
+from azure.core.pipeline.policies import ContentDecodePolicy
+from .._internal import _scopes_to_resource
+from .._internal.pipeline import build_pipeline
 
 try:
     ABC = abc.ABC
@@ -33,10 +22,10 @@
 
 if TYPE_CHECKING:
     # pylint:disable=ungrouped-imports
-    from typing import Any, Callable, Dict, List, Optional, Union
+    from typing import Any, Callable, Dict, Optional, Union
     from azure.core.pipeline import PipelineResponse
     from azure.core.pipeline.policies import HTTPPolicy, SansIOHTTPPolicy
-    from azure.core.pipeline.transport import HttpTransport, HttpRequest
+    from azure.core.pipeline.transport import HttpRequest
 
     PolicyType = Union[HTTPPolicy, SansIOHTTPPolicy]
 
@@ -50,27 +39,27 @@ def __init__(self, request_factory, client_id=None, identity_config=None, **kwar
         self._identity_config = identity_config or {}
         if client_id:
             self._identity_config["client_id"] = client_id
-
-        config = kwargs.pop("_config", None) or _get_configuration(**kwargs)
-        self._pipeline = self._build_pipeline(config, **kwargs)
-
+        self._pipeline = self._build_pipeline(**kwargs)
         self._request_factory = request_factory
 
     def _process_response(self, response, request_time):
         # type: (PipelineResponse, int) -> AccessToken
 
-        try:
-            content = ContentDecodePolicy.deserialize_from_text(
-                response.http_response.text(), mime_type="application/json"
-            )
-            if not content:
-                raise ClientAuthenticationError(message="No token received.", response=response.http_response)
-        except DecodeError as ex:
-            if response.http_response.content_type.startswith("application/json"):
-                message = "Failed to deserialize JSON from response"
-            else:
-                message = 'Unexpected content type "{}"'.format(response.http_response.content_type)
-            six.raise_from(ClientAuthenticationError(message=message, response=response.http_response), ex)
+        content = response.context.get(ContentDecodePolicy.CONTEXT_NAME)
+        if not content:
+            try:
+                content = ContentDecodePolicy.deserialize_from_text(
+                    response.http_response.text(), mime_type="application/json"
+                )
+            except DecodeError as ex:
+                if response.http_response.content_type.startswith("application/json"):
+                    message = "Failed to deserialize JSON from response"
+                else:
+                    message = 'Unexpected content type "{}"'.format(response.http_response.content_type)
+                six.raise_from(ClientAuthenticationError(message=message, response=response.http_response), ex)
+
+        if not content:
+            raise ClientAuthenticationError(message="No token received.", response=response.http_response)
 
         if "access_token" not in content or not ("expires_in" in content or "expires_on" in content):
             if content and "access_token" in content:
@@ -110,7 +99,7 @@ def request_token(self, *scopes, **kwargs):
         pass
 
     @abc.abstractmethod
-    def _build_pipeline(self, config, policies=None, transport=None, **kwargs):
+    def _build_pipeline(self, **kwargs):
         pass
 
 
@@ -124,32 +113,5 @@ def request_token(self, *scopes, **kwargs):
         token = self._process_response(response, request_time)
         return token
 
-    def _build_pipeline(self, config, policies=None, transport=None, **kwargs):  # pylint:disable=no-self-use
-        # type: (Configuration, Optional[List[PolicyType]], Optional[HttpTransport], **Any) -> Pipeline
-        if policies is None:  # [] is a valid policy list
-            policies = _get_policies(config, **kwargs)
-        if not transport:
-            from azure.core.pipeline.transport import RequestsTransport
-
-            transport = RequestsTransport(**kwargs)
-
-        return Pipeline(transport=transport, policies=policies)
-
-
-def _get_policies(config, **kwargs):
-    return [
-        HeadersPolicy(**kwargs),
-        UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs),
-        config.proxy_policy,
-        config.retry_policy,
-        NetworkTraceLoggingPolicy(**kwargs),
-        DistributedTracingPolicy(**kwargs),
-        HttpLoggingPolicy(**kwargs),
-    ]
-
-
-def _get_configuration(**kwargs):
-    # type: (**Any) -> Configuration
-    config = Configuration()
-    config.retry_policy = RetryPolicy(**kwargs)
-    return config
+    def _build_pipeline(self, **kwargs):
+        return build_pipeline(**kwargs)

sdk/identity/azure-identity/azure/identity/_internal/msal_client.py

@@ -6,21 +6,11 @@
 
 import six
 
-from azure.core.configuration import Configuration
 from azure.core.exceptions import ClientAuthenticationError
-from azure.core.pipeline import Pipeline
-from azure.core.pipeline.policies import (
-    ContentDecodePolicy,
-    DistributedTracingPolicy,
-    HttpLoggingPolicy,
-    NetworkTraceLoggingPolicy,
-    ProxyPolicy,
-    RetryPolicy,
-    UserAgentPolicy,
-)
-from azure.core.pipeline.transport import HttpRequest, RequestsTransport
-
-from .user_agent import USER_AGENT
+from azure.core.pipeline.policies import ContentDecodePolicy
+from azure.core.pipeline.transport import HttpRequest
+
+from .pipeline import build_pipeline
 
 try:
     from typing import TYPE_CHECKING
@@ -29,12 +19,10 @@
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Dict, List, Optional, Union
+    from typing import Any, Dict, Optional, Union
     from azure.core.pipeline import PipelineResponse
-    from azure.core.pipeline.policies import HTTPPolicy, SansIOHTTPPolicy
-    from azure.core.pipeline.transport import HttpResponse, HttpTransport
+    from azure.core.pipeline.transport import HttpResponse
 
-    PolicyList = List[Union[HTTPPolicy, SansIOHTTPPolicy]]
     RequestData = Union[Dict[str, str], str]
 
 
@@ -83,7 +71,7 @@ class MsalClient(object):
     def __init__(self, **kwargs):  # pylint:disable=missing-client-constructor-parameter-credential
         # type: (**Any) -> None
         self._local = threading.local()
-        self._pipeline = _build_pipeline(**kwargs)
+        self._pipeline = build_pipeline(**kwargs)
 
     def post(self, url, params=None, data=None, headers=None, **kwargs):  # pylint:disable=unused-argument
         # type: (str, Optional[Dict[str, str]], RequestData, Optional[Dict[str, str]], **Any) -> MsalResponse
@@ -129,34 +117,3 @@ def _store_auth_error(self, response):
             content = response.context.get(ContentDecodePolicy.CONTEXT_NAME)
             if content and "error" in content:
                 self._local.error = (content["error"], response.http_response)
-
-
-def _create_config(**kwargs):
-    # type: (Any) -> Configuration
-    config = Configuration(**kwargs)
-    config.logging_policy = NetworkTraceLoggingPolicy(**kwargs)
-    config.retry_policy = RetryPolicy(**kwargs)
-    config.proxy_policy = ProxyPolicy(**kwargs)
-    config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs)
-    return config
-
-
-def _build_pipeline(config=None, policies=None, transport=None, **kwargs):
-    # type: (Optional[Configuration], Optional[PolicyList], Optional[HttpTransport], **Any) -> Pipeline
-    config = config or _create_config(**kwargs)
-
-    if policies is None:  # [] is a valid policy list
-        policies = [
-            ContentDecodePolicy(),
-            config.user_agent_policy,
-            config.proxy_policy,
-            config.retry_policy,
-            config.logging_policy,
-            DistributedTracingPolicy(**kwargs),
-            HttpLoggingPolicy(**kwargs),
-        ]
-
-    if not transport:
-        transport = RequestsTransport(**kwargs)
-
-    return Pipeline(transport=transport, policies=policies)