Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

sqlalchemy wrap_create_engine now accepts sqlcommenter options #1873

1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#1870](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1870))
- Update falcon instrumentation to follow semantic conventions
([#1824](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1824))
- Fix sqlalchemy instrumentation wrap methods to accept sqlcommenter options([#1873](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1873))

### Added

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,9 @@ def _instrument(self, **kwargs):
``engine``: a SQLAlchemy engine instance
``engines``: a list of SQLAlchemy engine instances
``tracer_provider``: a TracerProvider, defaults to global
``meter_provider``: a MeterProvider, defaults to global
``enable_commenter``: bool to enable sqlcommenter, defaults to False
``commenter_options``: dict of sqlcommenter config, defaults to None

Returns:
An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
Expand All @@ -151,16 +154,21 @@ def _instrument(self, **kwargs):
)

enable_commenter = kwargs.get("enable_commenter", False)
commenter_options = kwargs.get("commenter_options", {})

_w(
"sqlalchemy",
"create_engine",
_wrap_create_engine(tracer, connections_usage, enable_commenter),
_wrap_create_engine(
tracer, connections_usage, enable_commenter, commenter_options
),
)
_w(
"sqlalchemy.engine",
"create_engine",
_wrap_create_engine(tracer, connections_usage, enable_commenter),
_wrap_create_engine(
tracer, connections_usage, enable_commenter, commenter_options
),
)
_w(
"sqlalchemy.engine.base",
Expand All @@ -172,7 +180,10 @@ def _instrument(self, **kwargs):
"sqlalchemy.ext.asyncio",
"create_async_engine",
_wrap_create_async_engine(
tracer, connections_usage, enable_commenter
tracer,
connections_usage,
enable_commenter,
commenter_options,
),
)
if kwargs.get("engine") is not None:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def _normalize_vendor(vendor):


def _wrap_create_async_engine(
tracer, connections_usage, enable_commenter=False
tracer, connections_usage, enable_commenter=False, commenter_options=None
):
# pylint: disable=unused-argument
def _wrap_create_async_engine_internal(func, module, args, kwargs):
Expand All @@ -52,20 +52,32 @@ def _wrap_create_async_engine_internal(func, module, args, kwargs):
"""
engine = func(*args, **kwargs)
EngineTracer(
tracer, engine.sync_engine, connections_usage, enable_commenter
tracer,
engine.sync_engine,
connections_usage,
enable_commenter,
commenter_options,
)
return engine

return _wrap_create_async_engine_internal


def _wrap_create_engine(tracer, connections_usage, enable_commenter=False):
def _wrap_create_engine(
tracer, connections_usage, enable_commenter=False, commenter_options=None
):
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(tracer, engine, connections_usage, enable_commenter)
EngineTracer(
tracer,
engine,
connections_usage,
enable_commenter,
commenter_options,
)
return engine

return _wrap_create_engine_internal
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
import logging
from unittest import mock

import pytest
Expand Down Expand Up @@ -176,6 +177,43 @@ def test_create_engine_wrapper(self):
"opentelemetry.instrumentation.sqlalchemy",
)

def test_create_engine_wrapper_enable_commenter(self):
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={"db_framework": False},
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute("SELECT 1;").fetchall()
# sqlcommenter
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_create_engine_wrapper_enable_commenter_otel_values_false(self):
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute("SELECT 1;").fetchall()
# sqlcommenter
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)

def test_custom_tracer_provider(self):
provider = TracerProvider(
resource=Resource.create(
Expand Down Expand Up @@ -242,6 +280,65 @@ async def run():

asyncio.get_event_loop().run_until_complete(run())

@pytest.mark.skipif(
not sqlalchemy.__version__.startswith("1.4"),
reason="only run async tests for 1.4",
)
def test_create_async_engine_wrapper_enable_commenter(self):
async def run():
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
},
)
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
create_async_engine,
)

engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.connect() as cnx:
await cnx.execute(sqlalchemy.text("SELECT 1;"))
# sqlcommenter
self.assertRegex(
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

asyncio.get_event_loop().run_until_complete(run())

@pytest.mark.skipif(
not sqlalchemy.__version__.startswith("1.4"),
reason="only run async tests for 1.4",
)
def test_create_async_engine_wrapper_enable_commenter_otel_values_false(
self,
):
async def run():
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
)
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
create_async_engine,
)

engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.connect() as cnx:
await cnx.execute(sqlalchemy.text("SELECT 1;"))
# sqlcommenter
self.assertRegex(
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)

asyncio.get_event_loop().run_until_complete(run())

def test_uninstrument(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,24 @@ def test_sqlcommenter_enabled(self):
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_enabled_otel_values_false(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
)
cnx = engine.connect()
cnx.execute("SELECT 1;").fetchall()
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)

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