Skip to content

Commit

Permalink
Fix sqlalchemy uninstrument (#1581)
Browse files Browse the repository at this point in the history
  • Loading branch information
shalevr committed Jan 29, 2023
1 parent a300d65 commit bbe7578
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 4 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Fixed

- Fix SQLAlchemy uninstrumentation
([#1581](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1581))
- `opentelemetry-instrumentation-grpc` Fix code()/details() of _OpentelemetryServicerContext.
([#1578](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1578))
- Fix aiopg instrumentation to work with aiopg < 2.0.0
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -186,3 +186,4 @@ def _uninstrument(self, **kwargs):
unwrap(Engine, "connect")
if parse_version(sqlalchemy.__version__).release >= (1, 4):
unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
EngineTracer.remove_all_event_listeners()
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@
import os
import re

from sqlalchemy.event import listen # pylint: disable=no-name-in-module
from sqlalchemy.event import ( # pylint: disable=no-name-in-module
listen,
remove,
)

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


class EngineTracer:
_remove_event_listener_params = []

def __init__(
self, tracer, engine, enable_commenter=False, commenter_options=None
):
Expand All @@ -105,11 +110,24 @@ def __init__(
self.commenter_options = commenter_options if commenter_options else {}
self._leading_comment_remover = re.compile(r"^/\*.*?\*/")

listen(
self._register_event_listener(
engine, "before_cursor_execute", self._before_cur_exec, retval=True
)
listen(engine, "after_cursor_execute", _after_cur_exec)
listen(engine, "handle_error", _handle_error)
self._register_event_listener(
engine, "after_cursor_execute", _after_cur_exec
)
self._register_event_listener(engine, "handle_error", _handle_error)

@classmethod
def _register_event_listener(cls, target, identifier, func, *args, **kw):
listen(target, identifier, func, *args, **kw)
cls._remove_event_listener_params.append((target, identifier, func))

@classmethod
def remove_all_event_listeners(cls):
for remove_params in cls._remove_event_listener_params:
remove(*remove_params)
cls._remove_event_listener_params.clear()

def _operation_name(self, db_name, statement):
parts = []
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -248,12 +248,41 @@ def test_uninstrument(self):

self.memory_exporter.clear()
SQLAlchemyInstrumentor().uninstrument()
cnx.execute("SELECT 1 + 1;").fetchall()
engine2 = create_engine("sqlite:///:memory:")
cnx2 = engine2.connect()
cnx2.execute("SELECT 2 + 2;").fetchall()
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 0)

SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
)
cnx = engine.connect()
cnx.execute("SELECT 1 + 1;").fetchall()
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)

def test_uninstrument_without_engine(self):
SQLAlchemyInstrumentor().instrument(
tracer_provider=self.tracer_provider
)
from sqlalchemy import create_engine

engine = create_engine("sqlite:///:memory:")

cnx = engine.connect()
cnx.execute("SELECT 1 + 1;").fetchall()
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)

self.memory_exporter.clear()
SQLAlchemyInstrumentor().uninstrument()
cnx.execute("SELECT 1 + 1;").fetchall()
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 0)

def test_no_op_tracer_provider(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
Expand Down

0 comments on commit bbe7578

Please sign in to comment.