Trace identity (Azure#6455)

* make sure its int. Need it to pass the test * added policy and decorator * dont decorate authn * loggin policy behind retry policy * managed should not be decoratored * dont need my workaround * dont trace get_token * dont need this extra indent
danigm · Jul 25, 2019 · 8b2d345 · 8b2d345
1 parent f249119
commit 8b2d345
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/sdk/identity/azure-identity/azure/identity/_authn_client.py b/sdk/identity/azure-identity/azure/identity/_authn_client.py
@@ -10,6 +10,7 @@
 from azure.core.exceptions import ClientAuthenticationError
 from azure.core.pipeline import Pipeline
 from azure.core.pipeline.policies import ContentDecodePolicy, NetworkTraceLoggingPolicy, RetryPolicy
+from azure.core.pipeline.policies.distributed_tracing import DistributedTracingPolicy
 from azure.core.pipeline.transport import HttpTransport, RequestsTransport
 from msal import TokenCache
 
@@ -119,7 +120,7 @@ class AuthnClient(AuthnClientBase):
     def __init__(self, auth_url, config=None, policies=None, transport=None, **kwargs):
         # type: (str, Optional[Configuration], Optional[Iterable[HTTPPolicy]], Optional[HttpTransport], Mapping[str, Any]) -> None
         config = config or self.create_config(**kwargs)
-        policies = policies or [ContentDecodePolicy(), config.logging_policy, config.retry_policy]
+        policies = policies or [ContentDecodePolicy(), config.retry_policy, config.logging_policy, DistributedTracingPolicy()]
         if not transport:
             transport = RequestsTransport(**kwargs)
         self._pipeline = Pipeline(transport=transport, policies=policies)

diff --git a/sdk/identity/azure-identity/azure/identity/aio/_authn_client.py b/sdk/identity/azure-identity/azure/identity/aio/_authn_client.py
@@ -8,6 +8,7 @@
 from azure.core import Configuration
 from azure.core.credentials import AccessToken
 from azure.core.pipeline import AsyncPipeline
+from azure.core.pipeline.policies.distributed_tracing import DistributedTracingPolicy
 from azure.core.pipeline.policies import AsyncRetryPolicy, ContentDecodePolicy, HTTPPolicy, NetworkTraceLoggingPolicy
 from azure.core.pipeline.transport import AsyncHttpTransport
 from azure.core.pipeline.transport.requests_asyncio import AsyncioRequestsTransport
@@ -27,7 +28,12 @@ def __init__(
         **kwargs: Mapping[str, Any]
     ) -> None:
         config = config or self.create_config(**kwargs)
-        policies = policies or [ContentDecodePolicy(), config.logging_policy, config.retry_policy]
+        policies = policies or [
+            ContentDecodePolicy(),
+            config.retry_policy,
+            config.logging_policy,
+            DistributedTracingPolicy(),
+        ]
         if not transport:
             transport = AsyncioRequestsTransport(**kwargs)
         self._pipeline = AsyncPipeline(transport=transport, policies=policies)