Skip to content

Commit bbe7578

Browse files
authored
Fix sqlalchemy uninstrument (#1581)
1 parent a300d65 commit bbe7578

File tree

4 files changed

+54
-4
lines changed

4 files changed

+54
-4
lines changed

CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1616

1717
### Fixed
1818

19+
- Fix SQLAlchemy uninstrumentation
20+
([#1581](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1581))
1921
- `opentelemetry-instrumentation-grpc` Fix code()/details() of _OpentelemetryServicerContext.
2022
([#1578](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1578))
2123
- Fix aiopg instrumentation to work with aiopg < 2.0.0

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -186,3 +186,4 @@ def _uninstrument(self, **kwargs):
186186
unwrap(Engine, "connect")
187187
if parse_version(sqlalchemy.__version__).release >= (1, 4):
188188
unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
189+
EngineTracer.remove_all_event_listeners()

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

+22-4
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@
1414
import os
1515
import re
1616

17-
from sqlalchemy.event import listen # pylint: disable=no-name-in-module
17+
from sqlalchemy.event import ( # pylint: disable=no-name-in-module
18+
listen,
19+
remove,
20+
)
1821

1922
from opentelemetry import trace
2023
from opentelemetry.instrumentation.sqlalchemy.package import (
@@ -95,6 +98,8 @@ def _wrap_connect_internal(func, module, args, kwargs):
9598

9699

97100
class EngineTracer:
101+
_remove_event_listener_params = []
102+
98103
def __init__(
99104
self, tracer, engine, enable_commenter=False, commenter_options=None
100105
):
@@ -105,11 +110,24 @@ def __init__(
105110
self.commenter_options = commenter_options if commenter_options else {}
106111
self._leading_comment_remover = re.compile(r"^/\*.*?\*/")
107112

108-
listen(
113+
self._register_event_listener(
109114
engine, "before_cursor_execute", self._before_cur_exec, retval=True
110115
)
111-
listen(engine, "after_cursor_execute", _after_cur_exec)
112-
listen(engine, "handle_error", _handle_error)
116+
self._register_event_listener(
117+
engine, "after_cursor_execute", _after_cur_exec
118+
)
119+
self._register_event_listener(engine, "handle_error", _handle_error)
120+
121+
@classmethod
122+
def _register_event_listener(cls, target, identifier, func, *args, **kw):
123+
listen(target, identifier, func, *args, **kw)
124+
cls._remove_event_listener_params.append((target, identifier, func))
125+
126+
@classmethod
127+
def remove_all_event_listeners(cls):
128+
for remove_params in cls._remove_event_listener_params:
129+
remove(*remove_params)
130+
cls._remove_event_listener_params.clear()
113131

114132
def _operation_name(self, db_name, statement):
115133
parts = []

instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py

+29
Original file line numberDiff line numberDiff line change
@@ -248,12 +248,41 @@ def test_uninstrument(self):
248248

249249
self.memory_exporter.clear()
250250
SQLAlchemyInstrumentor().uninstrument()
251+
cnx.execute("SELECT 1 + 1;").fetchall()
251252
engine2 = create_engine("sqlite:///:memory:")
252253
cnx2 = engine2.connect()
253254
cnx2.execute("SELECT 2 + 2;").fetchall()
254255
spans = self.memory_exporter.get_finished_spans()
255256
self.assertEqual(len(spans), 0)
256257

258+
SQLAlchemyInstrumentor().instrument(
259+
engine=engine,
260+
tracer_provider=self.tracer_provider,
261+
)
262+
cnx = engine.connect()
263+
cnx.execute("SELECT 1 + 1;").fetchall()
264+
spans = self.memory_exporter.get_finished_spans()
265+
self.assertEqual(len(spans), 2)
266+
267+
def test_uninstrument_without_engine(self):
268+
SQLAlchemyInstrumentor().instrument(
269+
tracer_provider=self.tracer_provider
270+
)
271+
from sqlalchemy import create_engine
272+
273+
engine = create_engine("sqlite:///:memory:")
274+
275+
cnx = engine.connect()
276+
cnx.execute("SELECT 1 + 1;").fetchall()
277+
spans = self.memory_exporter.get_finished_spans()
278+
self.assertEqual(len(spans), 2)
279+
280+
self.memory_exporter.clear()
281+
SQLAlchemyInstrumentor().uninstrument()
282+
cnx.execute("SELECT 1 + 1;").fetchall()
283+
spans = self.memory_exporter.get_finished_spans()
284+
self.assertEqual(len(spans), 0)
285+
257286
def test_no_op_tracer_provider(self):
258287
engine = create_engine("sqlite:///:memory:")
259288
SQLAlchemyInstrumentor().instrument(

0 commit comments

Comments
 (0)