Span decorator to pass INTERNAL by default

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

@@ -30,6 +30,7 @@
 from typing import Callable, Any, TypeVar, overload
 from typing_extensions import ParamSpec
 from .common import change_context, get_function_and_class_name
+from . import SpanKind
 from ..settings import settings
 
 
@@ -57,10 +58,14 @@ def distributed_trace(  # pylint:disable=function-redefined
     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
+    :keyword name_of_span: The span name to replace func name if necessary
+    :paramtype name_of_span: str
+    :keyword kind: The kind of the span. INTERNAL by default.
+    :paramtype kind: ~azure.core.tracing.SpanKind
     """
     name_of_span = kwargs.pop("name_of_span", None)
     tracing_attributes = kwargs.pop("tracing_attributes", {})
+    kind = kwargs.pop("kind", SpanKind.INTERNAL)
 
     def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
@@ -78,7 +83,7 @@ def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T:
 
             with change_context(passed_in_parent):
                 name = name_of_span or get_function_and_class_name(func, *args)
-                with span_impl_type(name=name) as span:
+                with span_impl_type(name=name, kind=kind) as span:
                     for key, value in tracing_attributes.items():
                         span.add_attribute(key, value)
                     return func(*args, **kwargs)

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

@@ -30,6 +30,7 @@
 from typing import Awaitable, Callable, Any, TypeVar, overload
 from typing_extensions import ParamSpec
 from .common import change_context, get_function_and_class_name
+from . import SpanKind
 from ..settings import settings
 
 P = ParamSpec("P")
@@ -58,10 +59,14 @@ def distributed_trace_async(  # pylint:disable=function-redefined
     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
+    :keyword name_of_span: The span name to replace func name if necessary
+    :paramtype name_of_span: str
+    :keyword kind: The kind of the span. INTERNAL by default.
+    :paramtype kind: ~azure.core.tracing.SpanKind
     """
     name_of_span = kwargs.pop("name_of_span", None)
     tracing_attributes = kwargs.pop("tracing_attributes", {})
+    kind = kwargs.pop("kind", SpanKind.INTERNAL)
 
     def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
         @functools.wraps(func)
@@ -79,7 +84,7 @@ async def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T:
 
             with change_context(passed_in_parent):
                 name = name_of_span or get_function_and_class_name(func, *args)
-                with span_impl_type(name=name) as span:
+                with span_impl_type(name=name, kind=kind) as span:
                     for key, value in tracing_attributes.items():
                         span.add_attribute(key, value)
                     return await func(*args, **kwargs)

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

@@ -17,6 +17,7 @@
 from azure.core.pipeline.policies import HTTPPolicy
 from azure.core.pipeline.transport import HttpTransport
 from azure.core.settings import settings
+from azure.core.tracing import SpanKind
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing.decorator_async import distributed_trace_async
 from tracing_common import FakeSpan
@@ -80,6 +81,10 @@ async def check_name_is_different(self):
     async def tracing_attr(self):
         time.sleep(0.001)
 
+    @distributed_trace_async(kind=SpanKind.PRODUCER)
+    async def kind_override(self):
+        time.sleep(0.001)
+
     @distributed_trace_async
     async def raising_exception(self):
         raise ValueError("Something went horribly wrong here")
@@ -98,6 +103,7 @@ async def test_decorator_tracing_attr(self, http_request):
         assert len(parent.children) == 2
         assert parent.children[0].name == "MockClient.__init__"
         assert parent.children[1].name == "MockClient.tracing_attr"
+        assert parent.children[1].kind == SpanKind.INTERNAL
         assert parent.children[1].attributes == {'foo': 'bar'}
 
 
@@ -110,7 +116,19 @@ async def test_decorator_has_different_name(self, http_request):
         assert len(parent.children) == 2
         assert parent.children[0].name == "MockClient.__init__"
         assert parent.children[1].name == "different name"
+        assert parent.children[1].kind == SpanKind.INTERNAL
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+    async def test_kind_override(self, http_request):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient(http_request)
+            await client.kind_override()
+
+        assert len(parent.children) == 2
+        assert parent.children[0].name == "MockClient.__init__"
+        assert parent.children[1].name == "MockClient.kind_override"
+        assert parent.children[1].kind == SpanKind.PRODUCER
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("http_request", HTTP_REQUESTS)

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

@@ -16,7 +16,7 @@
 from azure.core.pipeline.policies import HTTPPolicy
 from azure.core.pipeline.transport import HttpTransport
 from azure.core.settings import settings
-from azure.core.tracing import common
+from azure.core.tracing import common, SpanKind
 from azure.core.tracing.decorator import distributed_trace
 from tracing_common import FakeSpan
 from utils import HTTP_REQUESTS
@@ -78,6 +78,10 @@ def check_name_is_different(self):
     def tracing_attr(self):
         time.sleep(0.001)
 
+    @distributed_trace(kind=SpanKind.PRODUCER)
+    def kind_override(self):
+        time.sleep(0.001)
+
     @distributed_trace
     def raising_exception(self):
         raise ValueError("Something went horribly wrong here")
@@ -106,6 +110,7 @@ def test_decorator_tracing_attr(self, http_request):
         assert len(parent.children) == 2
         assert parent.children[0].name == "MockClient.__init__"
         assert parent.children[1].name == "MockClient.tracing_attr"
+        assert parent.children[1].kind == SpanKind.INTERNAL
         assert parent.children[1].attributes == {'foo': 'bar'}
 
     @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
@@ -117,6 +122,18 @@ def test_decorator_has_different_name(self, http_request):
         assert len(parent.children) == 2
         assert parent.children[0].name == "MockClient.__init__"
         assert parent.children[1].name == "different name"
+        assert parent.children[1].kind == SpanKind.INTERNAL
+
+    @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+    def test_kind_override(self, http_request):
+        with FakeSpan(name="parent") as parent:
+            client = MockClient(http_request)
+            client.kind_override()
+
+        assert len(parent.children) == 2
+        assert parent.children[0].name == "MockClient.__init__"
+        assert parent.children[1].name == "MockClient.kind_override"
+        assert parent.children[1].kind == SpanKind.PRODUCER
 
     @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
     def test_used(self, http_request):