20
20
)
21
21
22
22
from opentelemetry import trace
23
- from opentelemetry .instrumentation .sqlalchemy .package import (
24
- _instrumenting_module_name ,
25
- )
26
23
from opentelemetry .instrumentation .sqlalchemy .version import __version__
27
24
from opentelemetry .instrumentation .sqlcommenter_utils import _add_sql_comment
28
25
from opentelemetry .instrumentation .utils import _get_opentelemetry_values
@@ -44,49 +41,36 @@ def _normalize_vendor(vendor):
44
41
return vendor
45
42
46
43
47
- def _get_tracer (tracer_provider = None ):
48
- return trace .get_tracer (
49
- _instrumenting_module_name ,
50
- __version__ ,
51
- tracer_provider = tracer_provider ,
52
- )
53
-
54
-
55
- def _wrap_create_async_engine (tracer_provider = None , enable_commenter = False ):
44
+ def _wrap_create_async_engine (
45
+ tracer , connections_usage , enable_commenter = False
46
+ ):
56
47
# pylint: disable=unused-argument
57
48
def _wrap_create_async_engine_internal (func , module , args , kwargs ):
58
49
"""Trace the SQLAlchemy engine, creating an `EngineTracer`
59
50
object that will listen to SQLAlchemy events.
60
51
"""
61
52
engine = func (* args , ** kwargs )
62
53
EngineTracer (
63
- _get_tracer ( tracer_provider ) , engine .sync_engine , enable_commenter
54
+ tracer , engine .sync_engine , connections_usage , enable_commenter
64
55
)
65
56
return engine
66
57
67
58
return _wrap_create_async_engine_internal
68
59
69
60
70
- def _wrap_create_engine (tracer_provider = None , enable_commenter = False ):
71
- # pylint: disable=unused-argument
72
- def _wrap_create_engine_internal (func , module , args , kwargs ):
61
+ def _wrap_create_engine (tracer , connections_usage , enable_commenter = False ):
62
+ def _wrap_create_engine_internal (func , _module , args , kwargs ):
73
63
"""Trace the SQLAlchemy engine, creating an `EngineTracer`
74
64
object that will listen to SQLAlchemy events.
75
65
"""
76
66
engine = func (* args , ** kwargs )
77
- EngineTracer (_get_tracer ( tracer_provider ) , engine , enable_commenter )
67
+ EngineTracer (tracer , engine , connections_usage , enable_commenter )
78
68
return engine
79
69
80
70
return _wrap_create_engine_internal
81
71
82
72
83
- def _wrap_connect (tracer_provider = None ):
84
- tracer = trace .get_tracer (
85
- _instrumenting_module_name ,
86
- __version__ ,
87
- tracer_provider = tracer_provider ,
88
- )
89
-
73
+ def _wrap_connect (tracer ):
90
74
# pylint: disable=unused-argument
91
75
def _wrap_connect_internal (func , module , args , kwargs ):
92
76
with tracer .start_as_current_span (
@@ -107,10 +91,16 @@ class EngineTracer:
107
91
_remove_event_listener_params = []
108
92
109
93
def __init__ (
110
- self , tracer , engine , enable_commenter = False , commenter_options = None
94
+ self ,
95
+ tracer ,
96
+ engine ,
97
+ connections_usage ,
98
+ enable_commenter = False ,
99
+ commenter_options = None ,
111
100
):
112
101
self .tracer = tracer
113
102
self .engine = engine
103
+ self .connections_usage = connections_usage
114
104
self .vendor = _normalize_vendor (engine .name )
115
105
self .enable_commenter = enable_commenter
116
106
self .commenter_options = commenter_options if commenter_options else {}
@@ -123,6 +113,49 @@ def __init__(
123
113
engine , "after_cursor_execute" , _after_cur_exec
124
114
)
125
115
self ._register_event_listener (engine , "handle_error" , _handle_error )
116
+ self ._register_event_listener (engine , "connect" , self ._pool_connect )
117
+ self ._register_event_listener (engine , "close" , self ._pool_close )
118
+ self ._register_event_listener (engine , "checkin" , self ._pool_checkin )
119
+ self ._register_event_listener (engine , "checkout" , self ._pool_checkout )
120
+
121
+ def _get_pool_name (self ):
122
+ return self .engine .pool .logging_name or ""
123
+
124
+ def _add_idle_to_connection_usage (self , value ):
125
+ self .connections_usage .add (
126
+ value ,
127
+ attributes = {
128
+ "pool.name" : self ._get_pool_name (),
129
+ "state" : "idle" ,
130
+ },
131
+ )
132
+
133
+ def _add_used_to_connection_usage (self , value ):
134
+ self .connections_usage .add (
135
+ value ,
136
+ attributes = {
137
+ "pool.name" : self ._get_pool_name (),
138
+ "state" : "used" ,
139
+ },
140
+ )
141
+
142
+ def _pool_connect (self , _dbapi_connection , _connection_record ):
143
+ self ._add_idle_to_connection_usage (1 )
144
+
145
+ def _pool_close (self , _dbapi_connection , _connection_record ):
146
+ self ._add_idle_to_connection_usage (- 1 )
147
+
148
+ # Called when a connection returns to the pool.
149
+ def _pool_checkin (self , _dbapi_connection , _connection_record ):
150
+ self ._add_used_to_connection_usage (- 1 )
151
+ self ._add_idle_to_connection_usage (1 )
152
+
153
+ # Called when a connection is retrieved from the Pool.
154
+ def _pool_checkout (
155
+ self , _dbapi_connection , _connection_record , _connection_proxy
156
+ ):
157
+ self ._add_idle_to_connection_usage (- 1 )
158
+ self ._add_used_to_connection_usage (1 )
126
159
127
160
@classmethod
128
161
def _register_event_listener (cls , target , identifier , func , * args , ** kw ):
@@ -153,9 +186,8 @@ def _operation_name(self, db_name, statement):
153
186
return self .vendor
154
187
return " " .join (parts )
155
188
156
- # pylint: disable=unused-argument
157
189
def _before_cur_exec (
158
- self , conn , cursor , statement , params , context , executemany
190
+ self , conn , cursor , statement , params , context , _executemany
159
191
):
160
192
attrs , found = _get_attributes_from_url (conn .engine .url )
161
193
if not found :
0 commit comments