2020)
2121
2222from opentelemetry import trace
23- from opentelemetry .instrumentation .sqlalchemy .package import (
24- _instrumenting_module_name ,
25- )
2623from opentelemetry .instrumentation .sqlalchemy .version import __version__
2724from opentelemetry .instrumentation .sqlcommenter_utils import _add_sql_comment
2825from opentelemetry .instrumentation .utils import _get_opentelemetry_values
@@ -44,49 +41,36 @@ def _normalize_vendor(vendor):
4441 return vendor
4542
4643
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+ ):
5647 # pylint: disable=unused-argument
5748 def _wrap_create_async_engine_internal (func , module , args , kwargs ):
5849 """Trace the SQLAlchemy engine, creating an `EngineTracer`
5950 object that will listen to SQLAlchemy events.
6051 """
6152 engine = func (* args , ** kwargs )
6253 EngineTracer (
63- _get_tracer ( tracer_provider ) , engine .sync_engine , enable_commenter
54+ tracer , engine .sync_engine , connections_usage , enable_commenter
6455 )
6556 return engine
6657
6758 return _wrap_create_async_engine_internal
6859
6960
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 ):
7363 """Trace the SQLAlchemy engine, creating an `EngineTracer`
7464 object that will listen to SQLAlchemy events.
7565 """
7666 engine = func (* args , ** kwargs )
77- EngineTracer (_get_tracer ( tracer_provider ) , engine , enable_commenter )
67+ EngineTracer (tracer , engine , connections_usage , enable_commenter )
7868 return engine
7969
8070 return _wrap_create_engine_internal
8171
8272
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 ):
9074 # pylint: disable=unused-argument
9175 def _wrap_connect_internal (func , module , args , kwargs ):
9276 with tracer .start_as_current_span (
@@ -107,10 +91,16 @@ class EngineTracer:
10791 _remove_event_listener_params = []
10892
10993 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 ,
111100 ):
112101 self .tracer = tracer
113102 self .engine = engine
103+ self .connections_usage = connections_usage
114104 self .vendor = _normalize_vendor (engine .name )
115105 self .enable_commenter = enable_commenter
116106 self .commenter_options = commenter_options if commenter_options else {}
@@ -123,6 +113,49 @@ def __init__(
123113 engine , "after_cursor_execute" , _after_cur_exec
124114 )
125115 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 )
126159
127160 @classmethod
128161 def _register_event_listener (cls , target , identifier , func , * args , ** kw ):
@@ -153,9 +186,8 @@ def _operation_name(self, db_name, statement):
153186 return self .vendor
154187 return " " .join (parts )
155188
156- # pylint: disable=unused-argument
157189 def _before_cur_exec (
158- self , conn , cursor , statement , params , context , executemany
190+ self , conn , cursor , statement , params , context , _executemany
159191 ):
160192 attrs , found = _get_attributes_from_url (conn .engine .url )
161193 if not found :
0 commit comments