|
11 | 11 | from ddtrace.trace import Pin
|
12 | 12 |
|
13 | 13 | from .utils import create_context
|
14 |
| -from .utils import get_function_name |
15 | 14 | from .utils import wrap_function_with_tracing
|
16 | 15 |
|
17 | 16 |
|
@@ -39,109 +38,92 @@ def patch():
|
39 | 38 | azure_functions._datadog_patch = True
|
40 | 39 |
|
41 | 40 | Pin().onto(azure_functions.FunctionApp)
|
42 |
| - _w("azure.functions", "FunctionApp.function_name", _patched_function_name) |
43 |
| - _w("azure.functions", "FunctionApp.route", _patched_route) |
44 |
| - _w("azure.functions", "FunctionApp.service_bus_queue_trigger", _patched_service_bus_trigger) |
45 |
| - _w("azure.functions", "FunctionApp.service_bus_topic_trigger", _patched_service_bus_trigger) |
46 |
| - _w("azure.functions", "FunctionApp.timer_trigger", _patched_timer_trigger) |
| 41 | + _w("azure.functions", "FunctionApp.get_functions", _patched_get_functions) |
47 | 42 |
|
48 | 43 |
|
49 |
| -def _patched_function_name(wrapped, instance, args, kwargs): |
50 |
| - Pin.override(instance, tags={"function_name": kwargs.get("name")}) |
51 |
| - return wrapped(*args, **kwargs) |
52 |
| - |
53 |
| - |
54 |
| -def _patched_route(wrapped, instance, args, kwargs): |
55 |
| - trigger = "Http" |
56 |
| - trigger_arg_name = kwargs.get("trigger_arg_name", "req") |
57 |
| - |
| 44 | +def _patched_get_functions(wrapped, instance, args, kwargs): |
58 | 45 | pin = Pin.get_from(instance)
|
59 | 46 | if not pin or not pin.enabled():
|
60 | 47 | return wrapped(*args, **kwargs)
|
61 | 48 |
|
62 |
| - def _wrapper(func): |
63 |
| - function_name = get_function_name(pin, instance, func) |
64 |
| - |
65 |
| - def context_factory(kwargs): |
66 |
| - req = kwargs.get(trigger_arg_name) |
67 |
| - return create_context("azure.functions.patched_route_request", pin, headers=req.headers) |
| 49 | + functions = wrapped(*args, **kwargs) |
68 | 50 |
|
69 |
| - def pre_dispatch(ctx, kwargs): |
70 |
| - req = kwargs.get(trigger_arg_name) |
71 |
| - ctx.set_item("req_span", ctx.span) |
72 |
| - return ("azure.functions.request_call_modifier", (ctx, config.azure_functions, req)) |
| 51 | + for function in functions: |
| 52 | + trigger = function.get_trigger() |
| 53 | + if not trigger: |
| 54 | + continue |
73 | 55 |
|
74 |
| - def post_dispatch(ctx, res): |
75 |
| - return ("azure.functions.start_response", (ctx, config.azure_functions, res, function_name, trigger)) |
| 56 | + trigger_type = trigger.get_binding_name() |
| 57 | + trigger_arg_name = trigger.name |
76 | 58 |
|
77 |
| - wrap_function = wrap_function_with_tracing( |
78 |
| - func, context_factory, pre_dispatch=pre_dispatch, post_dispatch=post_dispatch |
79 |
| - ) |
| 59 | + function_name = function.get_function_name() |
| 60 | + func = function.get_user_function() |
80 | 61 |
|
81 |
| - return wrapped(*args, **kwargs)(wrap_function) |
| 62 | + if trigger_type == "httpTrigger": |
| 63 | + function._func = _wrap_http_trigger(pin, func, function_name, trigger_arg_name) |
| 64 | + elif trigger_type == "timerTrigger": |
| 65 | + function._func = _wrap_timer_trigger(pin, func, function_name) |
| 66 | + elif trigger_type == "serviceBusTrigger": |
| 67 | + function._func = _wrap_service_bus_trigger(pin, func, function_name) |
82 | 68 |
|
83 |
| - return _wrapper |
| 69 | + return functions |
84 | 70 |
|
85 | 71 |
|
86 |
| -def _patched_service_bus_trigger(wrapped, instance, args, kwargs): |
87 |
| - trigger = "ServiceBus" |
| 72 | +def _wrap_http_trigger(pin, func, function_name, trigger_arg_name): |
| 73 | + trigger_type = "Http" |
88 | 74 |
|
89 |
| - pin = Pin.get_from(instance) |
90 |
| - if not pin or not pin.enabled(): |
91 |
| - return wrapped(*args, **kwargs) |
| 75 | + def context_factory(kwargs): |
| 76 | + req = kwargs.get(trigger_arg_name) |
| 77 | + return create_context("azure.functions.patched_route_request", pin, headers=req.headers) |
92 | 78 |
|
93 |
| - def _wrapper(func): |
94 |
| - function_name = get_function_name(pin, instance, func) |
| 79 | + def pre_dispatch(ctx, kwargs): |
| 80 | + req = kwargs.get(trigger_arg_name) |
| 81 | + ctx.set_item("req_span", ctx.span) |
| 82 | + return ("azure.functions.request_call_modifier", (ctx, config.azure_functions, req)) |
95 | 83 |
|
96 |
| - def context_factory(kwargs): |
97 |
| - resource_name = f"{trigger} {function_name}" |
98 |
| - return create_context("azure.functions.patched_service_bus", pin, resource_name) |
| 84 | + def post_dispatch(ctx, res): |
| 85 | + return ("azure.functions.start_response", (ctx, config.azure_functions, res, function_name, trigger_type)) |
99 | 86 |
|
100 |
| - def pre_dispatch(ctx, kwargs): |
101 |
| - ctx.set_item("trigger_span", ctx.span) |
102 |
| - return ( |
103 |
| - "azure.functions.trigger_call_modifier", |
104 |
| - (ctx, config.azure_functions, function_name, trigger, SpanKind.CONSUMER), |
105 |
| - ) |
| 87 | + return wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch, post_dispatch=post_dispatch) |
106 | 88 |
|
107 |
| - wrap_function = wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch) |
108 | 89 |
|
109 |
| - return wrapped(*args, **kwargs)(wrap_function) |
| 90 | +def _wrap_service_bus_trigger(pin, func, function_name): |
| 91 | + trigger_type = "ServiceBus" |
110 | 92 |
|
111 |
| - return _wrapper |
| 93 | + def context_factory(kwargs): |
| 94 | + resource_name = f"{trigger_type} {function_name}" |
| 95 | + return create_context("azure.functions.patched_service_bus", pin, resource_name) |
112 | 96 |
|
| 97 | + def pre_dispatch(ctx, kwargs): |
| 98 | + ctx.set_item("trigger_span", ctx.span) |
| 99 | + return ( |
| 100 | + "azure.functions.trigger_call_modifier", |
| 101 | + (ctx, config.azure_functions, function_name, trigger_type, SpanKind.CONSUMER), |
| 102 | + ) |
113 | 103 |
|
114 |
| -def _patched_timer_trigger(wrapped, instance, args, kwargs): |
115 |
| - trigger = "Timer" |
116 |
| - |
117 |
| - pin = Pin.get_from(instance) |
118 |
| - if not pin or not pin.enabled(): |
119 |
| - return wrapped(*args, **kwargs) |
120 |
| - |
121 |
| - def _wrapper(func): |
122 |
| - function_name = get_function_name(pin, instance, func) |
| 104 | + return wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch) |
123 | 105 |
|
124 |
| - def context_factory(kwargs): |
125 |
| - resource_name = f"{trigger} {function_name}" |
126 |
| - return create_context("azure.functions.patched_timer", pin, resource_name) |
127 | 106 |
|
128 |
| - def pre_dispatch(ctx, kwargs): |
129 |
| - ctx.set_item("trigger_span", ctx.span) |
130 |
| - return ( |
131 |
| - "azure.functions.trigger_call_modifier", |
132 |
| - (ctx, config.azure_functions, function_name, trigger, SpanKind.INTERNAL), |
133 |
| - ) |
| 107 | +def _wrap_timer_trigger(pin, func, function_name): |
| 108 | + trigger_type = "Timer" |
134 | 109 |
|
135 |
| - wrap_function = wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch) |
| 110 | + def context_factory(kwargs): |
| 111 | + resource_name = f"{trigger_type} {function_name}" |
| 112 | + return create_context("azure.functions.patched_timer", pin, resource_name) |
136 | 113 |
|
137 |
| - return wrapped(*args, **kwargs)(wrap_function) |
| 114 | + def pre_dispatch(ctx, kwargs): |
| 115 | + ctx.set_item("trigger_span", ctx.span) |
| 116 | + return ( |
| 117 | + "azure.functions.trigger_call_modifier", |
| 118 | + (ctx, config.azure_functions, function_name, trigger_type, SpanKind.INTERNAL), |
| 119 | + ) |
138 | 120 |
|
139 |
| - return _wrapper |
| 121 | + return wrap_function_with_tracing(func, context_factory, pre_dispatch=pre_dispatch) |
140 | 122 |
|
141 | 123 |
|
142 | 124 | def unpatch():
|
143 | 125 | if not getattr(azure_functions, "_datadog_patch", False):
|
144 | 126 | return
|
145 | 127 | azure_functions._datadog_patch = False
|
146 | 128 |
|
147 |
| - _u(azure_functions.FunctionApp, "route") |
| 129 | + _u(azure_functions.FunctionApp, "get_functions") |
0 commit comments