feat(azure_functions): add tracing support for service bus triggers (#13324)

This PR adds tracing support for Service Bus triggers (queues and
topics) in Azure Functions.

### Additional Notes:

- Span Kind set to `consumer` for Service Bus triggers
- Refactors Azure Functions integration with
`wrap_function_with_tracing` method so additional trigger
[decorators](https://learn.microsoft.com/en-us/python/api/azure-functions/azure.functions.decorators?view=azure-python)
can be added with the same approach
- Make [manual
request](https://learn.microsoft.com/en-us/azure/azure-functions/functions-manually-run-non-http?tabs=azure-portal)
to service bus functions in test


## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: duncanpharvey

ddtrace/_trace/trace_handlers.py

@@ -823,7 +823,7 @@ def _set_span_pointer(span: "Span", span_pointer_description: _SpanPointerDescri
     )
 
 
-def _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind=SpanKind.INTERNAL):
+def _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind):
     span.set_tag_str(COMPONENT, azure_functions_config.integration_name)
     span.set_tag_str(SPAN_KIND, span_kind)
     span.set_tag_str("aas.function.name", function_name)  # codespell:ignore
@@ -857,9 +857,9 @@ def _on_azure_functions_start_response(ctx, azure_functions_config, res, functio
     )
 
 
-def _on_azure_functions_trigger_span_modifier(ctx, azure_functions_config, function_name, trigger):
+def _on_azure_functions_trigger_span_modifier(ctx, azure_functions_config, function_name, trigger, span_kind):
     span = ctx.get_item("trigger_span")
-    _set_azure_function_tags(span, azure_functions_config, function_name, trigger)
+    _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind)
 
 
 def listen():
@@ -968,6 +968,7 @@ def listen():
         "rq.job.perform",
         "rq.job.fetch_many",
         "azure.functions.patched_route_request",
+        "azure.functions.patched_service_bus",
         "azure.functions.patched_timer",
     ):
         core.on(f"context.started.start_span.{context_name}", _start_span)

ddtrace/contrib/internal/azure_functions/patch.py

@@ -1,19 +1,18 @@
-import functools
-import inspect
 import os
 
 import azure.functions as azure_functions
 from wrapt import wrap_function_wrapper as _w
 
 from ddtrace import config
 from ddtrace.contrib.internal.trace_utils import unwrap as _u
-from ddtrace.internal import core
+from ddtrace.ext import SpanKind
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.internal.utils.formats import asbool
 from ddtrace.trace import Pin
 
 from .utils import create_context
 from .utils import get_function_name
+from .utils import wrap_function_with_tracing
 
 
 config._add(
@@ -42,6 +41,8 @@ def patch():
     Pin().onto(azure_functions.FunctionApp)
     _w("azure.functions", "FunctionApp.function_name", _patched_function_name)
     _w("azure.functions", "FunctionApp.route", _patched_route)
+    _w("azure.functions", "FunctionApp.service_bus_queue_trigger", _patched_service_bus_trigger)
+    _w("azure.functions", "FunctionApp.service_bus_topic_trigger", _patched_service_bus_trigger)
     _w("azure.functions", "FunctionApp.timer_trigger", _patched_timer_trigger)
 
 
@@ -61,39 +62,49 @@ def _patched_route(wrapped, instance, args, kwargs):
     def _wrapper(func):
         function_name = get_function_name(pin, instance, func)
 
-        if inspect.iscoroutinefunction(func):
-
-            @functools.wraps(func)
-            async def async_wrap_function(*args, **kwargs):
-                req = kwargs.get(trigger_arg_name)
-                with create_context("azure.functions.patched_route_request", pin, headers=req.headers) as ctx, ctx.span:
-                    ctx.set_item("req_span", ctx.span)
-                    core.dispatch("azure.functions.request_call_modifier", (ctx, config.azure_functions, req))
-                    res = None
-                    try:
-                        res = await func(*args, **kwargs)
-                        return res
-                    finally:
-                        core.dispatch(
-                            "azure.functions.start_response", (ctx, config.azure_functions, res, function_name, trigger)
-                        )
-
-            return wrapped(*args, **kwargs)(async_wrap_function)
-
-        @functools.wraps(func)
-        def wrap_function(*args, **kwargs):
+        def context_factory(kwargs):
             req = kwargs.get(trigger_arg_name)
-            with create_context("azure.functions.patched_route_request", pin, headers=req.headers) as ctx, ctx.span:
-                ctx.set_item("req_span", ctx.span)
-                core.dispatch("azure.functions.request_call_modifier", (ctx, config.azure_functions, req))
-                res = None
-                try:
-                    res = func(*args, **kwargs)
-                    return res
-                finally:
-                    core.dispatch(
-                        "azure.functions.start_response", (ctx, config.azure_functions, res, function_name, trigger)
-                    )
+            return create_context("azure.functions.patched_route_request", pin, headers=req.headers)
+
+        def pre_dispatch(ctx, kwargs):
+            req = kwargs.get(trigger_arg_name)
+            ctx.set_item("req_span", ctx.span)
+            return ("azure.functions.request_call_modifier", (ctx, config.azure_functions, req))
+
+        def post_dispatch(ctx, res):
+            return ("azure.functions.start_response", (ctx, config.azure_functions, res, function_name, trigger))
+
+        wrap_function = wrap_function_with_tracing(
+            func, context_factory, pre_dispatch=pre_dispatch, post_dispatch=post_dispatch
+        )
+
+        return wrapped(*args, **kwargs)(wrap_function)
+
+    return _wrapper
+
+
+def _patched_service_bus_trigger(wrapped, instance, args, kwargs):
+    trigger = "ServiceBus"
+
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return wrapped(*args, **kwargs)
+
+    def _wrapper(func):
+        function_name = get_function_name(pin, instance, func)
+
+        def context_factory(kwargs):
+            resource_name = f"{trigger} {function_name}"
+            return create_context("azure.functions.patched_service_bus", pin, resource_name)
+
+        def pre_dispatch(ctx, kwargs):
+            ctx.set_item("trigger_span", ctx.span)
+            return (
+                "azure.functions.trigger_call_modifier",
+                (ctx, config.azure_functions, function_name, trigger, SpanKind.CONSUMER),
+            )
+
+        wrap_function = wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch)
 
         return wrapped(*args, **kwargs)(wrap_function)
 
@@ -109,31 +120,19 @@ def _patched_timer_trigger(wrapped, instance, args, kwargs):
 
     def _wrapper(func):
         function_name = get_function_name(pin, instance, func)
-        resource_name = f"{trigger} {function_name}"
-
-        if inspect.iscoroutinefunction(func):
-
-            @functools.wraps(func)
-            async def async_wrap_function(*args, **kwargs):
-                with create_context("azure.functions.patched_timer", pin, resource_name) as ctx, ctx.span:
-                    ctx.set_item("trigger_span", ctx.span)
-                    core.dispatch(
-                        "azure.functions.trigger_call_modifier",
-                        (ctx, config.azure_functions, function_name, trigger),
-                    )
-                    await func(*args, **kwargs)
-
-            return wrapped(*args, **kwargs)(async_wrap_function)
-
-        @functools.wraps(func)
-        def wrap_function(*args, **kwargs):
-            with create_context("azure.functions.patched_timer", pin, resource_name) as ctx, ctx.span:
-                ctx.set_item("trigger_span", ctx.span)
-                core.dispatch(
-                    "azure.functions.trigger_call_modifier",
-                    (ctx, config.azure_functions, function_name, trigger),
-                )
-                func(*args, **kwargs)
+
+        def context_factory(kwargs):
+            resource_name = f"{trigger} {function_name}"
+            return create_context("azure.functions.patched_timer", pin, resource_name)
+
+        def pre_dispatch(ctx, kwargs):
+            ctx.set_item("trigger_span", ctx.span)
+            return (
+                "azure.functions.trigger_call_modifier",
+                (ctx, config.azure_functions, function_name, trigger, SpanKind.INTERNAL),
+            )
+
+        wrap_function = wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch)
 
         return wrapped(*args, **kwargs)(wrap_function)
 

ddtrace/contrib/internal/azure_functions/utils.py

@@ -1,3 +1,6 @@
+import functools
+import inspect
+
 from ddtrace import config
 from ddtrace.contrib.internal.trace_utils import int_service
 from ddtrace.ext import SpanTypes
@@ -30,3 +33,39 @@ def get_function_name(pin, instance, func):
     else:
         function_name = func.__name__
     return function_name
+
+
+def wrap_function_with_tracing(func, context_factory, pre_dispatch=None, post_dispatch=None):
+    if inspect.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            with context_factory(kwargs) as ctx, ctx.span:
+                if pre_dispatch:
+                    core.dispatch(*pre_dispatch(ctx, kwargs))
+
+                res = None
+                try:
+                    res = await func(*args, **kwargs)
+                    return res
+                finally:
+                    if post_dispatch:
+                        core.dispatch(*post_dispatch(ctx, res))
+
+        return async_wrapper
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        with context_factory(kwargs) as ctx, ctx.span:
+            if pre_dispatch:
+                core.dispatch(*pre_dispatch(ctx, kwargs))
+
+            res = None
+            try:
+                res = func(*args, **kwargs)
+                return res
+            finally:
+                if post_dispatch:
+                    core.dispatch(*post_dispatch(ctx, res))
+
+    return wrapper

releasenotes/notes/feat-azure-functions-service-bus-trigger-cf816fdbb710f970.yaml

@@ -0,0 +1,3 @@
+features:
+  - |
+    azure_functions: This introduces tracing support for service bus triggers.

tests/contrib/azure_functions/azure_function_app/function_app.py

@@ -64,6 +64,20 @@ def http_get_child(req: func.HttpRequest) -> func.HttpResponse:
     return func.HttpResponse("Hello Datadog!")
 
 
+@app.function_name(name="servicebusqueue")
+@app.service_bus_queue_trigger(arg_name="msg", queue_name="queue.1", connection="CONNECTION_SETTING")
+def service_bus_queue(msg: func.ServiceBusMessage):
+    pass
+
+
+@app.function_name(name="servicebustopic")
+@app.service_bus_topic_trigger(
+    arg_name="msg", topic_name="topic.1", connection="CONNECTION_SETTING", subscription_name="subscription.1"
+)
+def service_bus_topic(msg: func.ServiceBusMessage):
+    pass
+
+
 @app.timer_trigger(schedule="0 0 0 1 1 *", arg_name="timer")
 def timer(timer: func.TimerRequest) -> None:
     pass

tests/contrib/azure_functions/azure_function_app/local.settings.json

@@ -5,6 +5,7 @@
     "FUNCTIONS_EXTENSION_VERSION": "~4",
     "AzureWebJobsFeatureFlags": "EnableWorkerIndexing",
     "AzureWebJobsStorage": "UseDevelopmentStorage=true",
+    "CONNECTION_SETTING": "Endpoint=sb://localhost;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=SAS_KEY_VALUE;UseDevelopmentEmulator=true;",
     "WEBSITE_SITE_NAME": "test-func"
   }
 }

tests/contrib/azure_functions/test_azure_functions_snapshot.py

@@ -118,6 +118,30 @@ def test_http_get_distributed_tracing(azure_functions_client: Client) -> None:
     assert azure_functions_client.get("/api/httpgetroot", headers=DEFAULT_HEADERS).status_code == 200
 
 
+@pytest.mark.snapshot
+def test_service_bus_queue(azure_functions_client: Client) -> None:
+    assert (
+        azure_functions_client.post(
+            "/admin/functions/servicebusqueue",
+            headers={"User-Agent": "python-httpx/x.xx.x", "Content-Type": "application/json"},
+            data=json.dumps({"input": '{"msg": "test message"}'}),
+        ).status_code
+        == 202
+    )
+
+
+@pytest.mark.snapshot
+def test_service_bus_topic(azure_functions_client: Client) -> None:
+    assert (
+        azure_functions_client.post(
+            "/admin/functions/servicebustopic",
+            headers={"User-Agent": "python-httpx/x.xx.x", "Content-Type": "application/json"},
+            data=json.dumps({"input": '{"msg": "test message"}'}),
+        ).status_code
+        == 202
+    )
+
+
 @pytest.mark.snapshot
 def test_timer(azure_functions_client: Client) -> None:
     assert (

tests/snapshots/tests.contrib.azure_functions.test_azure_functions_snapshot.test_service_bus_queue.json

@@ -0,0 +1,28 @@
+[[
+  {
+    "name": "azure.functions.invoke",
+    "service": "test-func",
+    "resource": "ServiceBus servicebusqueue",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "serverless",
+    "meta": {
+      "_dd.p.dm": "-0",
+      "_dd.p.tid": "6815246300000000",
+      "aas.function.name": "servicebusqueue",
+      "aas.function.trigger": "ServiceBus",
+      "component": "azure_functions",
+      "language": "python",
+      "runtime-id": "0a56574a66cd46eb8288f2856b00d420",
+      "span.kind": "consumer"
+    },
+    "metrics": {
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 46647
+    },
+    "duration": 195625,
+    "start": 1746216035448955969
+  }]]

tests/snapshots/tests.contrib.azure_functions.test_azure_functions_snapshot.test_service_bus_topic.json

@@ -0,0 +1,28 @@
+[[
+  {
+    "name": "azure.functions.invoke",
+    "service": "test-func",
+    "resource": "ServiceBus servicebustopic",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "serverless",
+    "meta": {
+      "_dd.p.dm": "-0",
+      "_dd.p.tid": "6815246c00000000",
+      "aas.function.name": "servicebustopic",
+      "aas.function.trigger": "ServiceBus",
+      "component": "azure_functions",
+      "language": "python",
+      "runtime-id": "39a73bcd6b804b0c9fc28bd839d213f1",
+      "span.kind": "consumer"
+    },
+    "metrics": {
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 47763
+    },
+    "duration": 220333,
+    "start": 1746216044422172751
+  }]]