1
1
import asyncio
2
2
from functools import wraps
3
-
4
3
from opentelemetry import trace
5
4
6
- # Get the _tracer instance (You can set your own _tracer name)
7
5
tracer = trace .get_tracer (__name__ )
8
6
9
7
@@ -16,41 +14,47 @@ def trace_function(trace_attributes: bool = True, trace_result: bool = True):
16
14
- trace_result (bool): If False, disables adding the function's result to the span.
17
15
"""
18
16
17
+ def set_span_attributes (span , args , kwargs , result = None ):
18
+ """Helper to set function arguments and results as span attributes."""
19
+ if trace_attributes :
20
+ span .set_attribute ("function.args" , str (args ))
21
+ span .set_attribute ("function.kwargs" , str (kwargs ))
22
+ if trace_result and result is not None :
23
+ span .set_attribute ("function.result" , str (result ))
24
+
25
+ def record_exception (span , exception ):
26
+ """Helper to handle exception recording in a span."""
27
+ span .record_exception (exception )
28
+ span .set_status (trace .status .Status (trace .status .StatusCode .ERROR ))
29
+
30
+ async def handle_async (span , func , * args , ** kwargs ):
31
+ """Handle asynchronous functions."""
32
+ try :
33
+ result = await func (* args , ** kwargs )
34
+ set_span_attributes (span , args , kwargs , result )
35
+ return result
36
+ except Exception as e :
37
+ record_exception (span , e )
38
+ raise
39
+
40
+ def handle_sync (span , func , * args , ** kwargs ):
41
+ """Handle synchronous functions."""
42
+ try :
43
+ result = func (* args , ** kwargs )
44
+ set_span_attributes (span , args , kwargs , result )
45
+ return result
46
+ except Exception as e :
47
+ record_exception (span , e )
48
+ raise
49
+
19
50
def decorator (func ):
20
51
@wraps (func )
21
- def sync_or_async_wrapper (* args , ** kwargs ):
52
+ def wrapper (* args , ** kwargs ):
22
53
with tracer .start_as_current_span (func .__name__ ) as span :
23
- try :
24
- # Set function arguments as attributes
25
- if trace_attributes :
26
- span .set_attribute ("function.args" , str (args ))
27
- span .set_attribute ("function.kwargs" , str (kwargs ))
28
-
29
- async def async_handler ():
30
- result = await func (* args , ** kwargs )
31
- # Add result to span
32
- if trace_result :
33
- span .set_attribute ("function.result" , str (result ))
34
- return result
35
-
36
- def sync_handler ():
37
- result = func (* args , ** kwargs )
38
- # Add result to span
39
- if trace_result :
40
- span .set_attribute ("function.result" , str (result ))
41
- return result
42
-
43
- if asyncio .iscoroutinefunction (func ):
44
- return async_handler ()
45
- else :
46
- return sync_handler ()
47
-
48
- except Exception as e :
49
- # Record the exception in the span
50
- span .record_exception (e )
51
- span .set_status (trace .status .Status (trace .status .StatusCode .ERROR ))
52
- raise
53
-
54
- return sync_or_async_wrapper
54
+ if asyncio .iscoroutinefunction (func ):
55
+ return handle_async (span , func , * args , ** kwargs )
56
+ return handle_sync (span , func , * args , ** kwargs )
57
+
58
+ return wrapper
55
59
56
60
return decorator
0 commit comments