[Core] Allow operation-level tracing attributes (Azure#38164)

Our documentation advertises that the `tracing_attributes` keyword
argument can be passed in at both the client constructor level and the
operation/method level. This makes this actually the case.

Per-operation customization of tracing attributes can help users mark/identify
spans from specific operations when spans are handled through custom span
processors.

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

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`.

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()
 

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)
 

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)
 

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):

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:

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"""