Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sdk/core/azure-core/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

### Bugs Fixed

- Declare method level span as INTERNAL by default #24492
- Fixed type hints for `azure.core.paging.ItemPaged` #24548

### Other Changes
Expand Down
9 changes: 7 additions & 2 deletions sdk/core/azure-core/azure/core/tracing/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
9 changes: 7 additions & 2 deletions sdk/core/azure-core/azure/core/tracing/decorator_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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'}


Expand All @@ -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)
Expand Down
19 changes: 18 additions & 1 deletion sdk/core/azure-core/tests/test_tracing_decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand All @@ -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):
Expand Down