Skip to content

Commit

Permalink
fix enable_commenter functionality (#1440)
Browse files Browse the repository at this point in the history
  • Loading branch information
avzis authored Dec 6, 2022
1 parent 99e0b42 commit cfd017e
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 8 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#1461](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1461))
- Add grpc.aio instrumentation to package entry points
([#1442](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1442))
- Fix a bug in SQLAlchemy instrumentation - support disabling enable_commenter variable
([#1440](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1440))

## Version 1.14.0/0.35b0 (2022-11-03)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -136,11 +136,16 @@ def _instrument(self, **kwargs):
An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
"""
tracer_provider = kwargs.get("tracer_provider")
_w("sqlalchemy", "create_engine", _wrap_create_engine(tracer_provider))
enable_commenter = kwargs.get("enable_commenter", False)
_w(
"sqlalchemy",
"create_engine",
_wrap_create_engine(tracer_provider, enable_commenter),
)
_w(
"sqlalchemy.engine",
"create_engine",
_wrap_create_engine(tracer_provider),
_wrap_create_engine(tracer_provider, enable_commenter),
)
_w(
"sqlalchemy.engine.base",
Expand All @@ -151,7 +156,7 @@ def _instrument(self, **kwargs):
_w(
"sqlalchemy.ext.asyncio",
"create_async_engine",
_wrap_create_async_engine(tracer_provider),
_wrap_create_async_engine(tracer_provider, enable_commenter),
)
if kwargs.get("engine") is not None:
return EngineTracer(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,27 +49,29 @@ def _get_tracer(tracer_provider=None):
)


def _wrap_create_async_engine(tracer_provider=None):
def _wrap_create_async_engine(tracer_provider=None, enable_commenter=False):
# pylint: disable=unused-argument
def _wrap_create_async_engine_internal(func, module, args, kwargs):
"""Trace the SQLAlchemy engine, creating an `EngineTracer`
object that will listen to SQLAlchemy events.
"""
engine = func(*args, **kwargs)
EngineTracer(_get_tracer(tracer_provider), engine.sync_engine)
EngineTracer(
_get_tracer(tracer_provider), engine.sync_engine, enable_commenter
)
return engine

return _wrap_create_async_engine_internal


def _wrap_create_engine(tracer_provider=None):
def _wrap_create_engine(tracer_provider=None, enable_commenter=False):
# pylint: disable=unused-argument
def _wrap_create_engine_internal(func, module, args, kwargs):
"""Trace the SQLAlchemy engine, creating an `EngineTracer`
object that will listen to SQLAlchemy events.
"""
engine = func(*args, **kwargs)
EngineTracer(_get_tracer(tracer_provider), engine)
EngineTracer(_get_tracer(tracer_provider), engine, enable_commenter)
return engine

return _wrap_create_engine_internal
Expand All @@ -94,7 +96,7 @@ def _wrap_connect_internal(func, module, args, kwargs):

class EngineTracer:
def __init__(
self, tracer, engine, enable_commenter=True, commenter_options=None
self, tracer, engine, enable_commenter=False, commenter_options=None
):
self.tracer = tracer
self.engine = engine
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,30 @@ def test_sqlcommenter_flask_integration(self):
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

def test_sqlcommenter_enabled_create_engine_after_instrumentation(self):
SQLAlchemyInstrumentor().instrument(
tracer_provider=self.tracer_provider,
enable_commenter=True,
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute("SELECT 1;").fetchall()
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

def test_sqlcommenter_disabled_create_engine_after_instrumentation(self):
SQLAlchemyInstrumentor().instrument(
tracer_provider=self.tracer_provider,
enable_commenter=False,
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute("SELECT 1;").fetchall()
self.assertEqual(self.caplog.records[-2].getMessage(), "SELECT 1;")

0 comments on commit cfd017e

Please sign in to comment.