Use tracing policy (Azure#6390)

SuyogSoti · Jul 24, 2019 · eefbc79 · eefbc79
1 parent 9b40e0a
commit eefbc79
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 10 deletions.
diff --git a/sdk/core/azure-core/azure/core/pipeline_client.py b/sdk/core/azure-core/azure/core/pipeline_client.py
@@ -26,16 +26,35 @@
 
 import logging
 
-from typing import List, Any, Dict, Union, IO, Tuple, Optional, Callable, Iterator, cast, TYPE_CHECKING  # pylint: disable=unused-import
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import (
+        List,
+        Any,
+        Dict,
+        Union,
+        IO,
+        Tuple,
+        Optional,
+        Callable,
+        Iterator,
+        cast,
+    )  # pylint: disable=unused-import
 
 from .pipeline import Pipeline
 from .pipeline.transport.base import PipelineClientBase
 from .pipeline.policies import ContentDecodePolicy
 from .pipeline.transport import RequestsTransport
+from .pipeline.policies.distributed_tracing import DistributedTracingPolicy
 
 
 _LOGGER = logging.getLogger(__name__)
 
+
 class PipelineClient(PipelineClientBase):
     """Service client core methods.
 
@@ -62,14 +81,15 @@ class PipelineClient(PipelineClientBase):
             :dedent: 4
             :caption: Builds the pipeline client.
     """
+
     def __init__(self, base_url, config, **kwargs):
         super(PipelineClient, self).__init__(base_url)
         if config is None:
             raise ValueError("Config is a required parameter")
         self._config = config
         self._base_url = base_url
-        if kwargs.get('pipeline'):
-            self._pipeline = kwargs['pipeline']
+        if kwargs.get("pipeline"):
+            self._pipeline = kwargs["pipeline"]
         else:
             self._pipeline = self._build_pipeline(config, **kwargs)
 
@@ -97,12 +117,10 @@ def _build_pipeline(self, config, **kwargs): # pylint: disable=no-self-use
                 config.retry_policy,
                 config.custom_hook_policy,
                 config.logging_policy,
+                DistributedTracingPolicy(),
             ]
 
         if not transport:
             transport = RequestsTransport(**kwargs)
 
-        return Pipeline(
-            transport,
-            policies
-        )
+        return Pipeline(transport, policies)
diff --git a/sdk/core/azure-core/azure/core/pipeline_client_async.py b/sdk/core/azure-core/azure/core/pipeline_client_async.py
@@ -26,16 +26,35 @@
 
 import logging
 
-from typing import List, Any, Dict, Union, IO, Tuple, Optional, Callable, Iterator, cast, TYPE_CHECKING  # pylint: disable=unused-import
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import (
+        List,
+        Any,
+        Dict,
+        Union,
+        IO,
+        Tuple,
+        Optional,
+        Callable,
+        Iterator,
+        cast,
+    )  # pylint: disable=unused-import
 
 from .pipeline import AsyncPipeline
 from .pipeline.transport.base import PipelineClientBase
 from .pipeline.policies import ContentDecodePolicy
 from .pipeline.transport import AioHttpTransport
+from .pipeline.policies.distributed_tracing import DistributedTracingPolicy
 
 
 _LOGGER = logging.getLogger(__name__)
 
+
 class AsyncPipelineClient(PipelineClientBase):
     """Service client core methods.
 
@@ -62,14 +81,15 @@ class AsyncPipelineClient(PipelineClientBase):
             :dedent: 4
             :caption: Builds the async pipeline client.
     """
+
     def __init__(self, base_url, config, **kwargs):
         super(AsyncPipelineClient, self).__init__(base_url)
         if config is None:
             raise ValueError("Config is a required parameter")
         self._config = config
         self._base_url = base_url
-        if kwargs.get('pipeline'):
-            self._pipeline = kwargs['pipeline']
+        if kwargs.get("pipeline"):
+            self._pipeline = kwargs["pipeline"]
         else:
             self._pipeline = self._build_pipeline(config, **kwargs)
 
@@ -97,6 +117,7 @@ def _build_pipeline(self, config, **kwargs): # pylint: disable=no-self-use
                 config.retry_policy,
                 config.custom_hook_policy,
                 config.logging_policy,
+                DistributedTracingPolicy(),
             ]
 
         if not transport: