Add tracing_attributes to tracing decorator

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

@@ -36,38 +36,49 @@
     TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from typing import Callable, Any
+    from typing import Callable, Dict, Optional, Any, cast
 
 
-def distributed_trace(func=None, name_of_span=None):
-    # type: (Callable, str) -> Callable[[Any], Any]
+def distributed_trace(_func=None, name_of_span=None, **kwargs):
+    # type: (Callable, Optional[str], Optional[Dict[str, Any]]) -> Callable
     """Decorator to apply to function to get traced automatically.
 
     Span will use the func name or "name_of_span".
 
     :param callable func: A function to decorate
     :param str name_of_span: The span name to replace func name if necessary
     """
-    if func is None:
-        return functools.partial(distributed_trace, name_of_span=name_of_span)
+    tracing_attributes = kwargs.get('tracing_attributes')
+    # https://github.com/python/mypy/issues/2608
+    if _func is None:
+        return functools.partial(
+            distributed_trace,
+            name_of_span=name_of_span,
+            tracing_attributes=tracing_attributes,
+        )
+    func = _func  # mypy is happy now
+
+    not_none_tracing_attributes = tracing_attributes if tracing_attributes else {}
 
     @functools.wraps(func)
     def wrapper_use_tracer(*args, **kwargs):
         # type: (Any, Any) -> Any
-        merge_span = kwargs.pop('merge_span', False)
+        merge_span = kwargs.pop("merge_span", False)
         passed_in_parent = kwargs.pop("parent_span", None)
 
         span_impl_type = settings.tracing_implementation()
         if span_impl_type is None:
-            return func(*args, **kwargs) # type: ignore
+            return func(*args, **kwargs)
 
         # Merge span is parameter is set, but only if no explicit parent are passed
         if merge_span and not passed_in_parent:
-            return func(*args, **kwargs) # type: ignore
+            return func(*args, **kwargs)
 
         with change_context(passed_in_parent):
-            name = name_of_span or get_function_and_class_name(func, *args)  # type: ignore
-            with span_impl_type(name=name):
-                return func(*args, **kwargs) # type: ignore
+            name = name_of_span or get_function_and_class_name(func, *args)
+            with span_impl_type(name=name) as span:
+                for key, value in not_none_tracing_attributes.items():
+                    span.add_attribute(key, value)
+                return func(*args, **kwargs)
 
     return wrapper_use_tracer

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

@@ -36,38 +36,48 @@
     TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from typing import Callable, Any
+    from typing import Callable, Dict, Optional, Any
 
 
-def distributed_trace_async(func=None, name_of_span=None):
-    # type: (Callable, str) -> Callable[[Any], Any]
+def distributed_trace_async(_func=None, name_of_span=None, *, tracing_attributes=None):
+    # type: (Callable, Optional[str], Optional[Dict[str, Any]]) -> Callable
     """Decorator to apply to async function to get traced automatically.
 
     Span will use the func name or "name_of_span".
 
     :param callable func: A function to decorate
     :param str name_of_span: The span name to replace func name if necessary
     """
-    if func is None:
-        return functools.partial(distributed_trace_async, name_of_span=name_of_span)
+    # https://github.com/python/mypy/issues/2608
+    if _func is None:
+        return functools.partial(
+            distributed_trace_async,
+            name_of_span=name_of_span,
+            tracing_attributes=tracing_attributes,
+        )
+    func = _func  # mypy is happy now
+
+    not_none_tracing_attributes = tracing_attributes if tracing_attributes else {}
 
     @functools.wraps(func)
     async def wrapper_use_tracer(*args, **kwargs):
         # type: (Any, Any) -> Any
-        merge_span = kwargs.pop('merge_span', False)
+        merge_span = kwargs.pop("merge_span", False)
         passed_in_parent = kwargs.pop("parent_span", None)
 
         span_impl_type = settings.tracing_implementation()
         if span_impl_type is None:
-            return await func(*args, **kwargs) # type: ignore
+            return await func(*args, **kwargs)
 
         # Merge span is parameter is set, but only if no explicit parent are passed
         if merge_span and not passed_in_parent:
-            return await func(*args, **kwargs) # type: ignore
+            return await func(*args, **kwargs)
 
         with change_context(passed_in_parent):
-            name = name_of_span or get_function_and_class_name(func, *args)  # type: ignore
-            with span_impl_type(name=name):
-                return await func(*args, **kwargs)  # type: ignore
+            name = name_of_span or get_function_and_class_name(func, *args)
+            with span_impl_type(name=name) as span:
+                for key, value in not_none_tracing_attributes.items():
+                    span.add_attribute(key, value)
+                return await func(*args, **kwargs)
 
     return wrapper_use_tracer

sdk/core/azure-core/tests/azure_core_asynctests/test_tracing_decorator_async.py

@@ -0,0 +1,182 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+"""The tests for decorators_async.py"""
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+import sys
+import time
+
+import pytest
+from azure.core.pipeline import Pipeline, PipelineResponse
+from azure.core.pipeline.policies import HTTPPolicy
+from azure.core.pipeline.transport import HttpTransport, HttpRequest
+from azure.core.settings import settings
+from azure.core.tracing.decorator import distributed_trace
+from azure.core.tracing.decorator_async import distributed_trace_async
+from tracing_common import FakeSpan
+
+
+@pytest.fixture(scope="module")
+def fake_span():
+    settings.tracing_implementation.set_value(FakeSpan)
+
+
+class MockClient:
+    @distributed_trace
+    def __init__(self, policies=None, assert_current_span=False):
+        time.sleep(0.001)
+        self.request = HttpRequest("GET", "https://bing.com")
+        if policies is None:
+            policies = []
+        policies.append(mock.Mock(spec=HTTPPolicy, send=self.verify_request))
+        self.policies = policies
+        self.transport = mock.Mock(spec=HttpTransport)
+        self.pipeline = Pipeline(self.transport, policies=policies)
+
+        self.expected_response = mock.Mock(spec=PipelineResponse)
+        self.assert_current_span = assert_current_span
+
+    def verify_request(self, request):
+        if self.assert_current_span:
+            assert execution_context.get_current_span() is not None
+        return self.expected_response
+
+    @distributed_trace_async
+    async def make_request(self, numb_times, **kwargs):
+        time.sleep(0.001)
+        if numb_times < 1:
+            return None
+        response = self.pipeline.run(self.request, **kwargs)
+        await self.get_foo(merge_span=True)
+        kwargs['merge_span'] = True
+        await self.make_request(numb_times - 1, **kwargs)
+        return response
+
+    @distributed_trace_async
+    async def merge_span_method(self):
+        return await self.get_foo(merge_span=True)
+
+    @distributed_trace_async
+    async def no_merge_span_method(self):
+        return await self.get_foo()
+
+    @distributed_trace_async
+    async def get_foo(self):
+        time.sleep(0.001)
+        return 5
+
+    @distributed_trace_async(name_of_span="different name")
+    async def check_name_is_different(self):
+        time.sleep(0.001)
+
+    @distributed_trace_async(tracing_attributes={'foo': 'bar'})
+    async def tracing_attr(self):
+        time.sleep(0.001)
+
+    @distributed_trace_async
+    async def raising_exception(self):
+        raise ValueError("Something went horribly wrong here")
+
+
+@pytest.mark.usefixtures("fake_span")
+class TestAsyncDecorator(object):
+
+    @pytest.mark.asyncio
+    async def test_decorator_tracing_attr(self):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient()
+            await client.tracing_attr()
+
+        assert len(parent.children) == 2
+        assert parent.children[0].name == "MockClient.__init__"
+        assert parent.children[1].name == "MockClient.tracing_attr"
+        assert parent.children[1].attributes == {'foo': 'bar'}
+
+
+    @pytest.mark.asyncio
+    async def test_decorator_has_different_name(self):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient()
+            await client.check_name_is_different()
+        assert len(parent.children) == 2
+        assert parent.children[0].name == "MockClient.__init__"
+        assert parent.children[1].name == "different name"
+
+
+    @pytest.mark.asyncio
+    async def test_used(self):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient(policies=[])
+            await client.get_foo(parent_span=parent)
+            await client.get_foo()
+
+            assert len(parent.children) == 3
+            assert parent.children[0].name == "MockClient.__init__"
+            assert not parent.children[0].children
+            assert parent.children[1].name == "MockClient.get_foo"
+            assert not parent.children[1].children
+            assert parent.children[2].name == "MockClient.get_foo"
+            assert not parent.children[2].children
+
+
+    @pytest.mark.asyncio
+    async def test_span_merge_span(self):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient()
+            await client.merge_span_method()
+            await client.no_merge_span_method()
+
+        assert len(parent.children) == 3
+        assert parent.children[0].name == "MockClient.__init__"
+        assert not parent.children[0].children
+        assert parent.children[1].name == "MockClient.merge_span_method"
+        assert not parent.children[1].children
+        assert parent.children[2].name == "MockClient.no_merge_span_method"
+        assert parent.children[2].children[0].name == "MockClient.get_foo"
+
+
+    @pytest.mark.asyncio
+    async def test_span_complicated(self):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient()
+            await client.make_request(2)
+            with parent.span("child") as child:
+                time.sleep(0.001)
+                await client.make_request(2, parent_span=parent)
+                assert FakeSpan.get_current_span() == child
+                await client.make_request(2)
+
+        assert len(parent.children) == 4
+        assert parent.children[0].name == "MockClient.__init__"
+        assert not parent.children[0].children
+        assert parent.children[1].name == "MockClient.make_request"
+        assert not parent.children[1].children
+        assert parent.children[2].name == "child"
+        assert parent.children[2].children[0].name == "MockClient.make_request"
+        assert parent.children[3].name == "MockClient.make_request"
+        assert not parent.children[3].children
+
+    @pytest.mark.asyncio
+    async def test_span_with_exception(self):
+        """Assert that if an exception is raised, the next sibling method is actually a sibling span.
+        """
+        with FakeSpan(name="parent") as parent:
+            client = MockClient()
+            try:
+                await client.raising_exception()
+            except:
+                pass
+            await client.get_foo()
+
+        assert len(parent.children) == 3
+        assert parent.children[0].name == "MockClient.__init__"
+        assert parent.children[1].name == "MockClient.raising_exception"
+        # Exception should propagate status for Opencensus
+        assert parent.children[1].status == 'Something went horribly wrong here'
+        assert parent.children[2].name == "MockClient.get_foo"

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

@@ -29,3 +29,13 @@
 collect_ignore = []
 if sys.version_info < (3, 5):
     collect_ignore.append("azure_core_asynctests")
+
+
+# If opencensus is loadable while doing these tests, register an empty tracer to avoid this:
+# https://github.com/census-instrumentation/opencensus-python/issues/442
+try:
+    from azure.core.tracing.ext.opencensus_span import OpenCensusSpan
+    from opencensus.trace.tracer import Tracer
+    Tracer()
+except ImportError:
+    pass

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

@@ -6,7 +6,10 @@
 from azure.core.pipeline.policies import RequestIdPolicy
 from azure.core.pipeline.transport import HttpRequest
 from azure.core.pipeline import PipelineRequest, PipelineContext
-from mock import patch
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 from itertools import product
 import pytest
 
@@ -31,7 +34,7 @@ def test_request_id_policy(auto_request_id, request_id_init, request_id_set, req
     pipeline_request = PipelineRequest(request, PipelineContext(None))
     if request_id_req != "_unset":
         pipeline_request.context.options['request_id'] = request_id_req
-    with patch('uuid.uuid1', return_value="VALUE"):
+    with mock.patch('uuid.uuid1', return_value="VALUE"):
         request_id_policy.on_request(pipeline_request)
 
     if request_id_req != "_unset":