sqlalchemy wrap_create_engine now accepts sqlcommenter options (#1873)

* sqlalchemy wrap_create_engine accepts sqlcommenter options

* Changelog

* Lint

* Fix default val

* Add sqlalchemy tests

* Change a default in _instrument get

* Lint

* More lint

* Update default

Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>

* Update args doc

* lintttt

---------

Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>
This commit is contained in:
Tammy Baylis
2023-06-27 03:43:35 -07:00
committed by GitHub
parent 2e49ba1af8
commit 79d62b3bcd
5 changed files with 146 additions and 7 deletions

View File

@ -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)) ([#1870](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1870))
- Update falcon instrumentation to follow semantic conventions - Update falcon instrumentation to follow semantic conventions
([#1824](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1824)) ([#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 ### Added

View File

@ -134,6 +134,9 @@ class SQLAlchemyInstrumentor(BaseInstrumentor):
``engine``: a SQLAlchemy engine instance ``engine``: a SQLAlchemy engine instance
``engines``: a list of SQLAlchemy engine instances ``engines``: a list of SQLAlchemy engine instances
``tracer_provider``: a TracerProvider, defaults to global ``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: Returns:
An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise. An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
@ -151,16 +154,21 @@ class SQLAlchemyInstrumentor(BaseInstrumentor):
) )
enable_commenter = kwargs.get("enable_commenter", False) enable_commenter = kwargs.get("enable_commenter", False)
commenter_options = kwargs.get("commenter_options", {})
_w( _w(
"sqlalchemy", "sqlalchemy",
"create_engine", "create_engine",
_wrap_create_engine(tracer, connections_usage, enable_commenter), _wrap_create_engine(
tracer, connections_usage, enable_commenter, commenter_options
),
) )
_w( _w(
"sqlalchemy.engine", "sqlalchemy.engine",
"create_engine", "create_engine",
_wrap_create_engine(tracer, connections_usage, enable_commenter), _wrap_create_engine(
tracer, connections_usage, enable_commenter, commenter_options
),
) )
_w( _w(
"sqlalchemy.engine.base", "sqlalchemy.engine.base",
@ -172,7 +180,10 @@ class SQLAlchemyInstrumentor(BaseInstrumentor):
"sqlalchemy.ext.asyncio", "sqlalchemy.ext.asyncio",
"create_async_engine", "create_async_engine",
_wrap_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: if kwargs.get("engine") is not None:

View File

@ -43,7 +43,7 @@ def _normalize_vendor(vendor):
def _wrap_create_async_engine( def _wrap_create_async_engine(
tracer, connections_usage, enable_commenter=False tracer, connections_usage, enable_commenter=False, commenter_options=None
): ):
# pylint: disable=unused-argument # pylint: disable=unused-argument
def _wrap_create_async_engine_internal(func, module, args, kwargs): def _wrap_create_async_engine_internal(func, module, args, kwargs):
@ -52,20 +52,32 @@ def _wrap_create_async_engine(
""" """
engine = func(*args, **kwargs) engine = func(*args, **kwargs)
EngineTracer( EngineTracer(
tracer, engine.sync_engine, connections_usage, enable_commenter tracer,
engine.sync_engine,
connections_usage,
enable_commenter,
commenter_options,
) )
return engine return engine
return _wrap_create_async_engine_internal 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): def _wrap_create_engine_internal(func, _module, args, kwargs):
"""Trace the SQLAlchemy engine, creating an `EngineTracer` """Trace the SQLAlchemy engine, creating an `EngineTracer`
object that will listen to SQLAlchemy events. object that will listen to SQLAlchemy events.
""" """
engine = func(*args, **kwargs) engine = func(*args, **kwargs)
EngineTracer(tracer, engine, connections_usage, enable_commenter) EngineTracer(
tracer,
engine,
connections_usage,
enable_commenter,
commenter_options,
)
return engine return engine
return _wrap_create_engine_internal return _wrap_create_engine_internal

View File

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import asyncio import asyncio
import logging
from unittest import mock from unittest import mock
import pytest import pytest
@ -176,6 +177,43 @@ class TestSqlalchemyInstrumentation(TestBase):
"opentelemetry.instrumentation.sqlalchemy", "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): def test_custom_tracer_provider(self):
provider = TracerProvider( provider = TracerProvider(
resource=Resource.create( resource=Resource.create(
@ -242,6 +280,65 @@ class TestSqlalchemyInstrumentation(TestBase):
asyncio.get_event_loop().run_until_complete(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): def test_uninstrument(self):
engine = create_engine("sqlite:///:memory:") engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument( SQLAlchemyInstrumentor().instrument(

View File

@ -56,6 +56,24 @@ class TestSqlalchemyInstrumentationWithSQLCommenter(TestBase):
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", 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): def test_sqlcommenter_flask_integration(self):
engine = create_engine("sqlite:///:memory:") engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument( SQLAlchemyInstrumentor().instrument(