@@ -32,73 +32,61 @@ def _collect_kvs(self, span, sql):
32
32
return span
33
33
34
34
def execute (self , sql , params = None ):
35
- try :
36
- span = None
37
- context = internal_tracer .current_context ()
35
+ parent_span = internal_tracer .active_span
38
36
39
- # If we're not tracing, just return
40
- if context is None :
41
- return self .__wrapped__ .execute (sql , params )
37
+ # If we're not tracing, just return
38
+ if parent_span is None :
39
+ return self .__wrapped__ .execute (sql , params )
42
40
43
- span = internal_tracer .start_span (self ._module_name , child_of = context )
44
- span = self ._collect_kvs (span , sql )
41
+ with internal_tracer .start_active_span (self ._module_name , child_of = parent_span ) as scope :
42
+ try :
43
+ self ._collect_kvs (scope .span , sql )
45
44
46
- result = self .__wrapped__ .execute (sql , params )
47
- except Exception as e :
48
- if span :
49
- span .log_exception (e )
50
- raise
51
- else :
52
- return result
53
- finally :
54
- if span :
55
- span .finish ()
45
+ result = self .__wrapped__ .execute (sql , params )
46
+ except Exception as e :
47
+ if scope .span :
48
+ scope .span .log_exception (e )
49
+ raise
50
+ else :
51
+ return result
56
52
57
53
def executemany (self , sql , seq_of_parameters ):
58
- try :
59
- span = None
60
- context = internal_tracer .current_context ()
54
+ parent_span = internal_tracer .active_span
61
55
62
- # If we're not tracing, just return
63
- if context is None :
64
- return self .__wrapped__ .executemany (sql , seq_of_parameters )
56
+ # If we're not tracing, just return
57
+ if parent_span is None :
58
+ return self .__wrapped__ .executemany (sql , seq_of_parameters )
65
59
66
- span = internal_tracer .start_span (self ._module_name , child_of = context )
67
- span = self ._collect_kvs (span , sql )
60
+ with internal_tracer .start_active_span (self ._module_name , child_of = parent_span ) as scope :
61
+ try :
62
+ self ._collect_kvs (scope .span , sql )
68
63
69
- result = self .__wrapped__ .executemany (sql , seq_of_parameters )
70
- except Exception as e :
71
- if span :
72
- span .log_exception (e )
73
- raise
74
- else :
75
- return result
76
- finally :
77
- if span :
78
- span .finish ()
64
+ result = self .__wrapped__ .executemany (sql , seq_of_parameters )
65
+ except Exception as e :
66
+ if scope .span :
67
+ scope .span .log_exception (e )
68
+ raise
69
+ else :
70
+ return result
79
71
80
72
def callproc (self , proc_name , params ):
81
- try :
82
- span = None
83
- context = internal_tracer .current_context ()
84
-
85
- # If we're not tracing, just return
86
- if context is None :
87
- return self .__wrapped__ .execute (proc_name , params )
88
-
89
- span = internal_tracer .start_span (self ._module_name , child_of = context )
90
- span = self ._collect_kvs (span , proc_name )
91
-
92
- result = self .__wrapped__ .callproc (proc_name , params )
93
- except Exception as e :
94
- if span :
95
- span .log_exception (e )
96
- raise
97
- else :
98
- return result
99
- finally :
100
- if span :
101
- span .finish ()
73
+ parent_span = internal_tracer .active_span
74
+
75
+ # If we're not tracing, just return
76
+ if parent_span is None :
77
+ return self .__wrapped__ .execute (proc_name , params )
78
+
79
+ with internal_tracer .start_active_span (self ._module_name , child_of = parent_span ) as scope :
80
+ try :
81
+ self ._collect_kvs (scope .span , proc_name )
82
+
83
+ result = self .__wrapped__ .callproc (proc_name , params )
84
+ except Exception as e :
85
+ if scope .span :
86
+ scope .span .log_exception (e )
87
+ raise
88
+ else :
89
+ return result
102
90
103
91
104
92
class ConnectionWrapper (wrapt .ObjectProxy ):
0 commit comments