diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md index e8312e2c08bc..833d4268aca6 100644 --- a/sdk/core/azure-core/CHANGELOG.md +++ b/sdk/core/azure-core/CHANGELOG.md @@ -6,6 +6,10 @@ - Added a default implementation to handle token challenges in `BearerTokenCredentialPolicy` and `AsyncBearerTokenCredentialPolicy`. +### Bugs Fixed + +- Fixed an issue where the `tracing_attributes` keyword argument wasn't being handled at the request/method level. #38164 + ### Other Changes - Log "x-vss-e2eid" and "x-msedge-ref" headers in `HttpLoggingPolicy`. diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py b/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py index 1206311b44e8..d049881dd27c 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py +++ b/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py @@ -93,10 +93,11 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None: return namer = ctxt.pop("network_span_namer", self._network_span_namer) + tracing_attributes = ctxt.pop("tracing_attributes", self._tracing_attributes) span_name = namer(request.http_request) span = span_impl_type(name=span_name, kind=SpanKind.CLIENT) - for attr, value in self._tracing_attributes.items(): + for attr, value in tracing_attributes.items(): span.add_attribute(attr, value) span.start() diff --git a/sdk/core/azure-core/azure/core/tracing/decorator.py b/sdk/core/azure-core/azure/core/tracing/decorator.py index e13f13d3068c..adca3aff3356 100644 --- a/sdk/core/azure-core/azure/core/tracing/decorator.py +++ b/sdk/core/azure-core/azure/core/tracing/decorator.py @@ -97,6 +97,9 @@ def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T: merge_span = kwargs.pop("merge_span", False) passed_in_parent = kwargs.pop("parent_span", None) + # Assume this will be popped in DistributedTracingPolicy. + func_tracing_attributes = kwargs.pop("tracing_attributes", tracing_attributes) + span_impl_type = settings.tracing_implementation() if span_impl_type is None: return func(*args, **kwargs) @@ -108,7 +111,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, kind=kind) as span: - for key, value in tracing_attributes.items(): + for key, value in func_tracing_attributes.items(): span.add_attribute(key, value) return func(*args, **kwargs) diff --git a/sdk/core/azure-core/azure/core/tracing/decorator_async.py b/sdk/core/azure-core/azure/core/tracing/decorator_async.py index 6ea8ffc352e9..f17081d1ad48 100644 --- a/sdk/core/azure-core/azure/core/tracing/decorator_async.py +++ b/sdk/core/azure-core/azure/core/tracing/decorator_async.py @@ -106,6 +106,9 @@ async def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T: merge_span = kwargs.pop("merge_span", False) passed_in_parent = kwargs.pop("parent_span", None) + # Assume this will be popped in DistributedTracingPolicy. + func_tracing_attributes = kwargs.get("tracing_attributes", tracing_attributes) + span_impl_type = settings.tracing_implementation() if span_impl_type is None: return await func(*args, **kwargs) @@ -117,7 +120,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, kind=kind) as span: - for key, value in tracing_attributes.items(): + for key, value in func_tracing_attributes.items(): span.add_attribute(key, value) return await func(*args, **kwargs) diff --git a/sdk/core/azure-core/tests/async_tests/test_tracing_decorator_async.py b/sdk/core/azure-core/tests/async_tests/test_tracing_decorator_async.py index b668e49ce000..f897ce9c1688 100644 --- a/sdk/core/azure-core/tests/async_tests/test_tracing_decorator_async.py +++ b/sdk/core/azure-core/tests/async_tests/test_tracing_decorator_async.py @@ -78,7 +78,7 @@ async def check_name_is_different(self): time.sleep(0.001) @distributed_trace_async(tracing_attributes={"foo": "bar"}) - async def tracing_attr(self): + async def tracing_attr(self, **kwargs): time.sleep(0.001) @distributed_trace_async(kind=SpanKind.PRODUCER) @@ -105,6 +105,19 @@ async def test_decorator_tracing_attr(self, http_request): assert parent.children[1].kind == SpanKind.INTERNAL assert parent.children[1].attributes == {"foo": "bar"} + @pytest.mark.asyncio + @pytest.mark.parametrize("http_request", HTTP_REQUESTS) + async def test_decorator_tracing_attr_custom(self, http_request): + with FakeSpan(name="parent") as parent: + client = MockClient(http_request) + await client.tracing_attr(tracing_attributes={"biz": "baz"}) + + 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 == {"biz": "baz"} + @pytest.mark.asyncio @pytest.mark.parametrize("http_request", HTTP_REQUESTS) async def test_decorator_has_different_name(self, http_request): diff --git a/sdk/core/azure-core/tests/test_tracing_decorator.py b/sdk/core/azure-core/tests/test_tracing_decorator.py index 4ee7ee621d96..01060a24faba 100644 --- a/sdk/core/azure-core/tests/test_tracing_decorator.py +++ b/sdk/core/azure-core/tests/test_tracing_decorator.py @@ -76,7 +76,7 @@ def check_name_is_different(self): time.sleep(0.001) @distributed_trace(tracing_attributes={"foo": "bar"}) - def tracing_attr(self): + def tracing_attr(self, **kwargs): time.sleep(0.001) @distributed_trace(kind=SpanKind.PRODUCER) @@ -113,6 +113,18 @@ def test_decorator_tracing_attr(self, http_request): assert parent.children[1].kind == SpanKind.INTERNAL assert parent.children[1].attributes == {"foo": "bar"} + @pytest.mark.parametrize("http_request", HTTP_REQUESTS) + def test_decorator_tracing_attr_custom(self, http_request): + with FakeSpan(name="parent") as parent: + client = MockClient(http_request) + client.tracing_attr(tracing_attributes={"biz": "baz"}) + + 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 == {"biz": "baz"} + @pytest.mark.parametrize("http_request", HTTP_REQUESTS) def test_decorator_has_different_name(self, http_request): with FakeSpan(name="parent") as parent: diff --git a/sdk/core/azure-core/tests/test_tracing_policy.py b/sdk/core/azure-core/tests/test_tracing_policy.py index 425b00bc94dc..75ae6c14813b 100644 --- a/sdk/core/azure-core/tests/test_tracing_policy.py +++ b/sdk/core/azure-core/tests/test_tracing_policy.py @@ -118,6 +118,29 @@ def test_distributed_tracing_policy_attributes(http_request, http_response): assert network_span.attributes.get("myattr") == "myvalue" +@pytest.mark.parametrize("http_request,http_response", request_and_responses_product(HTTP_RESPONSES)) +def test_distributed_tracing_policy_attributes_per_operation(http_request, http_response): + """Test policy with no other policy and happy path""" + settings.tracing_implementation.set_value(FakeSpan) + with FakeSpan(name="parent") as root_span: + policy = DistributedTracingPolicy(tracing_attributes={"myattr": "myvalue"}) + + request = http_request("GET", "http://localhost/temp?query=query") + + pipeline_request = PipelineRequest(request, PipelineContext(None, tracing_attributes={"foo": "bar"})) + policy.on_request(pipeline_request) + + response = create_http_response(http_response, request, None) + response.headers = request.headers + response.status_code = 202 + + policy.on_response(pipeline_request, PipelineResponse(request, response, PipelineContext(None))) + + # Check on_response + network_span = root_span.children[0] + assert network_span.attributes.get("foo") == "bar" + + @pytest.mark.parametrize("http_request,http_response", request_and_responses_product(HTTP_RESPONSES)) def test_distributed_tracing_policy_badurl(caplog, http_request, http_response): """Test policy with a bad url that will throw, and be sure policy ignores it"""