Add policy (#6379)

* add opencensus impl

* dont need to check for noop tracer

* get rid of span and trace id

* fix span instance

* added documentation

* added documentation and refactor names of a few variables

* write test for opencensus wrapper

* put opencensus in the dev requirements

* initial common

* only import the wrapper when necessary

* add check for the exporter

* rework logic and fix some settings

* added initial decorator

* some mroe documentation

* added decorators

* small change

* some minor fixes

* test patch happening

* fix space

* share a function

* clearer logic for setting span context

* better logic

* better logic

* fix environ variable

* test the way opencensus does it

* middle of tests

* only load if opencensus has already been imported

* fix spelling mistake

* temp

* finish writing tests for common

* charles fixes

* fix tests

* fix test settings

* to header should not take a dict

* from header should be class method

* initial tests

* dont create trace and get rid of end_tracer

* dont need to save the trace

* more little fixes

* some intermediatary changes

* fix type annotations

* rst fix types

* add :class:annotations

* fix line wrapping

* added tests for decorator

* rename opencensus wrapper

* intermediate changes

* use spans the right way

* some formatting

* some grammar

* restructure settings and make tests pass

* rename get_parent

* fix typings

* use protocol and from_headers becomes links

* ramifications of opencensus wrapper being a protocol

* add tests for link

* added async tests

* delete the unit test thing

* added add_attribute

* added add_attribute

* added tests for add attributes

* add initial policy

* remove unused import

* added docstrings

* minor docstring formatting

* fix pylint errors

* don't rely on opencensus children to check

* use exporter to not rely on parent.children

* added documentation and span attributes

* added test tracing policy

* test should only propagate

* made test tracing helper

* decrease flakiness of test

* simplify get parent

* calling a decorator decorator is redundant

* middle of writing tests

* fix settings

* add tests

* test propogation also happens

* more elegant code

* await async stuff

* add await for async

* should only have to wait a 1/1000 of a second

* fix tests spans too short

* accidentally deleted setup.cfg

* add component

* add set_http_attributes

* fix span network name

* fix http request types

* fix more types

* bryan fixes

* more efficient tests

* non flakey tests

* make tracing only use my context

* test user agent on exception

* pylint formatting

* delete unused import

* fix spelling

* pylint

sdk/core/azure-core/azure/core/pipeline/policies/distributed_tracing.py

@@ -0,0 +1,106 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+"""Traces network calls using the implementation library from the settings."""
+
+from azure.core.pipeline import PipelineRequest, PipelineResponse
+from azure.core.tracing.context import tracing_context
+from azure.core.tracing.abstract_span import AbstractSpan
+from azure.core.tracing.common import set_span_contexts
+from azure.core.pipeline.policies import SansIOHTTPPolicy
+from azure.core.settings import settings
+
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse
+
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from azure.core.pipeline.transport import HttpRequest, HttpResponse
+
+
+class DistributedTracingPolicy(SansIOHTTPPolicy):
+    """The policy to create spans for Azure Calls"""
+
+    def __init__(self):
+        # type: (str, str, str) -> None
+        self.parent_span_dict = {}
+        self._request_id = "x-ms-request-id"
+
+    def set_header(self, request, span):
+        # type: (PipelineRequest[HttpRequest], Any) -> None
+        """
+        Sets the header information on the span.
+        """
+        headers = span.to_header()
+        request.http_request.headers.update(headers)
+
+    def on_request(self, request, **kwargs):
+        # type: (PipelineRequest[HttpRequest], Any) -> None
+        parent_span = tracing_context.current_span.get()  # type: AbstractSpan
+
+        if parent_span is None:
+            return
+
+        only_propagate = settings.tracing_should_only_propagate()
+        if only_propagate:
+            self.set_header(request, parent_span)
+            return
+
+        path = urlparse(request.http_request.url).path
+        if not path:
+            path = "/"
+        child = parent_span.span(name=path)
+        child.start()
+
+        set_span_contexts(child)
+        self.parent_span_dict[child] = parent_span
+        self.set_header(request, child)
+
+    def end_span(self, request, response=None):
+        # type: (HttpRequest, Optional[HttpResponse]) -> None
+        """Ends the span that is tracing the network and updates its status."""
+        span = tracing_context.current_span.get()   # type: AbstractSpan
+        only_propagate = settings.tracing_should_only_propagate()
+        if span and not only_propagate:
+            span.set_http_attributes(request, response=response)
+            if response and self._request_id in response.headers:
+                span.add_attribute(self._request_id, response.headers[self._request_id])
+            span.finish()
+            set_span_contexts(self.parent_span_dict[span])
+
+    def on_response(self, request, response, **kwargs):
+        # type: (PipelineRequest[HttpRequest], PipelineResponse[HttpRequest, HttpResponse], Any) -> None
+        self.end_span(request.http_request, response=response.http_response)
+
+    def on_exception(self, _request, **kwargs):  # pylint: disable=unused-argument
+        # type: (PipelineRequest[HttpRequest], Any) -> bool
+        self.end_span(_request.http_request)

sdk/core/azure-core/azure/core/tracing/abstract_span.py

@@ -12,6 +12,8 @@
 if TYPE_CHECKING:
     from typing import Any, Dict, Optional, Union
 
+    from azure.core.pipeline.transport import HttpRequest, HttpResponse
+
 try:
     from typing_extensions import Protocol
 except ImportError:
@@ -21,7 +23,7 @@
 class AbstractSpan(Protocol):
     """Wraps a span from a distributed tracing implementation."""
 
-    def __init__(self, span=None, name=None): # pylint: disable=super-init-not-called
+    def __init__(self, span=None, name=None):  # pylint: disable=super-init-not-called
         # type: (Optional[Any], Optional[str]) -> None
         """
         If a span is given wraps the span. Else a new span is created.
@@ -66,6 +68,18 @@ def add_attribute(self, key, value):
         """
         pass
 
+    def set_http_attributes(self, request, response=None):
+        # type: (HttpRequest, HttpResponse) -> None
+        """
+        Add correct attributes for a http client span.
+
+        :param request: The request made
+        :type request: HttpRequest
+        :param response: The response received by the server. Is None if no response received.
+        :type response: HttpResponse
+        """
+        pass
+
     @property
     def span_instance(self):
         # type: () -> Any

sdk/core/azure-core/azure/core/tracing/ext/opencensus_span.py

@@ -5,6 +5,7 @@
 """Implements azure.core.tracing.AbstractSpan to wrap opencensus spans."""
 
 from opencensus.trace import Span, execution_context
+from opencensus.trace.span import SpanKind
 from opencensus.trace.link import Link
 from opencensus.trace.propagation import trace_context_http_header_format
 
@@ -14,7 +15,9 @@
     TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from typing import Dict, Optional, Union
+    from typing import Dict, Optional, Union, TypeVar
+
+    from azure.core.pipeline.transport import HttpRequest, HttpResponse
 
 
 class OpenCensusSpan(object):
@@ -33,8 +36,13 @@ def __init__(self, span=None, name="span"):
         """
         if not span:
             tracer = self.get_current_tracer()
-            span = tracer.span(name=name) # type: Span
+            span = tracer.span(name=name)  # type: Span
         self._span_instance = span
+        self._span_component = "component"
+        self._http_user_agent = "http.user_agent"
+        self._http_method = "http.method"
+        self._http_url = "http.url"
+        self._http_status_code = "http.status_code"
 
     @property
     def span_instance(self):
@@ -89,6 +97,28 @@ def add_attribute(self, key, value):
         """
         self.span_instance.add_attribute(key, value)
 
+    def set_http_attributes(self, request, response=None):
+        # type: (HttpRequest, Optional[HttpResponse]) -> None
+        """
+        Add correct attributes for a http client span.
+
+        :param request: The request made
+        :type request: HttpRequest
+        :param response: The response received by the server. Is None if no response received.
+        :type response: HttpResponse
+        """
+        self._span_instance.span_kind = SpanKind.CLIENT
+        self.span_instance.add_attribute(self._span_component, "http")
+        self.span_instance.add_attribute(self._http_method, request.method)
+        self.span_instance.add_attribute(self._http_url, request.url)
+        user_agent = request.headers.get("User-Agent")
+        if user_agent:
+            self.span_instance.add_attribute(self._http_user_agent, user_agent)
+        if response:
+            self._span_instance.add_attribute(self._http_status_code, response.status_code)
+        else:
+            self._span_instance.add_attribute(self._http_status_code, 504)
+
     @classmethod
     def link(cls, headers):
         # type: (Dict[str, str]) -> None

sdk/core/azure-core/tests/test_tracing_implementations.py

@@ -13,65 +13,14 @@
 
 from azure.core.tracing.ext.opencensus_span import OpenCensusSpan
 from opencensus.trace import tracer as tracer_module
+from opencensus.trace.span import SpanKind
 from opencensus.trace.samplers import AlwaysOnSampler
 from opencensus.trace.base_exporter import Exporter
 from opencensus.common.utils import timestamp_to_microseconds
+from tracing_common import MockExporter, ContextHelper
 import os
 
 
-class Node:
-    def __init__(self, span_data):
-        self.span_data = span_data  # type: SpanData
-        self.parent = None
-        self.children = []
-
-
-class MockExporter(Exporter):
-    def __init__(self):
-        self.root = None
-        self._all_nodes = []
-
-    def export(self, span_datas):
-        # type: (List[SpanData]) -> None
-        sp = span_datas[0]  # type: SpanData
-        node = Node(sp)
-        if not node.span_data.parent_span_id:
-            self.root = node
-        self._all_nodes.append(node)
-
-    def build_tree(self):
-        parent_dict = {}
-        for node in self._all_nodes:
-            parent_span_id = node.span_data.parent_span_id
-            if parent_span_id not in parent_dict:
-                parent_dict[parent_span_id] = []
-            parent_dict[parent_span_id].append(node)
-
-        for node in self._all_nodes:
-            if node.span_data.span_id in parent_dict:
-                node.children = sorted(
-                    parent_dict[node.span_data.span_id], key=lambda x: timestamp_to_microseconds(x.span_data.start_time)
-                )
-
-
-class ContextHelper(object):
-    def __init__(self, environ={}):
-        self.orig_tracer = OpenCensusSpan.get_current_tracer()
-        self.orig_current_span = OpenCensusSpan.get_current_span()
-        self.os_env = mock.patch.dict(os.environ, environ)
-
-    def __enter__(self):
-        self.orig_tracer = OpenCensusSpan.get_current_tracer()
-        self.orig_current_span = OpenCensusSpan.get_current_span()
-        self.os_env.start()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        OpenCensusSpan.set_current_tracer(self.orig_tracer)
-        OpenCensusSpan.set_current_span(self.orig_current_span)
-        self.os_env.stop()
-
-
 class TestOpencensusWrapper(unittest.TestCase):
     def test_span_passed_in(self):
         with ContextHelper():
@@ -162,3 +111,26 @@ def test_add_attribute(self):
             wrapped_class.add_attribute("test", "test2")
             assert wrapped_class.span_instance.attributes["test"] == "test2"
             assert parent.attributes["test"] == "test2"
+
+    def test_set_http_attributes(self):
+        with ContextHelper():
+            trace = tracer_module.Tracer(sampler=AlwaysOnSampler())
+            parent = trace.start_span()
+            wrapped_class = OpenCensusSpan(span=parent)
+            request = mock.Mock()
+            setattr(request, "method", "GET")
+            setattr(request, "url", "some url")
+            response = mock.Mock()
+            setattr(request, "headers", {})
+            setattr(response, "status_code", 200)
+            wrapped_class.set_http_attributes(request)
+            assert wrapped_class.span_instance.span_kind == SpanKind.CLIENT
+            assert wrapped_class.span_instance.attributes.get("http.method") == request.method
+            assert wrapped_class.span_instance.attributes.get("component") == "http"
+            assert wrapped_class.span_instance.attributes.get("http.url") == request.url
+            assert wrapped_class.span_instance.attributes.get("http.status_code") == 504
+            assert wrapped_class.span_instance.attributes.get("http.user_agent") is None
+            request.headers["User-Agent"] = "some user agent"
+            wrapped_class.set_http_attributes(request, response)
+            assert wrapped_class.span_instance.attributes.get("http.status_code") == response.status_code
+            assert wrapped_class.span_instance.attributes.get("http.user_agent") == request.headers.get("User-Agent")