Skip to content

Commit

Permalink
Add PyMySQL instrumentor support for sqlcommenting (open-telemetry#2942)
Browse files Browse the repository at this point in the history 
* WIP

* Add _DB_DRIVER_ALIASES

* Add mysql_client_version to sqlcomment

* lint

* Fix existing tests

* lint test

* Add PyMySQL dbapi commenter case

* Add test

* Add test

* Add test

* Add tests

* Changelog

* calculate_commenter_data at init of DatabaseApiIntegration

* try-except if NoneType module

* Add pymysql sqlcomment support

* Add unit tests

* Update docstring

* Changelog

* pymysql instrument_connection specifies connect_module

* lint

* Add tests

* Fix doc

---------

Co-authored-by: Riccardo Magliocchetti <[email protected]>
  • Loading branch information
@tammy-baylis-swi @xrmx
tammy-baylis-swi and xrmx authored Nov 21, 2024
1 parent beff723 commit 7b1554b
Show file tree
Hide file tree
Showing 3 changed files with 327 additions and 5 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#2635](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2635))
- `opentelemetry-instrumentation` Add support for string based dotted module paths in unwrap
([#2919](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2919))
- `opentelemetry-instrumentation-pymysql` Add sqlcommenter support
([#2942](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2942))

### Fixed

Expand Down
91 changes: 86 additions & 5 deletions ...entelemetry-instrumentation-pymysql/src/opentelemetry/instrumentation/pymysql/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
import pymysql
from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor
PyMySQLInstrumentor().instrument()
cnx = pymysql.connect(database="MySQL_Database")
Expand All @@ -36,6 +35,76 @@
cursor.close()
cnx.close()
SQLCOMMENTER
*****************************************
You can optionally configure PyMySQL instrumentation to enable sqlcommenter which enriches
the query with contextual information.
Usage
-----
.. code:: python
import pymysql
from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor
PyMySQLInstrumentor().instrument(enable_commenter=True, commenter_options={})
cnx = pymysql.connect(database="MySQL_Database")
cursor = cnx.cursor()
cursor.execute("INSERT INTO test (testField) VALUES (123)"
cnx.commit()
cursor.close()
cnx.close()
For example,
::
Invoking cursor.execute("INSERT INTO test (testField) VALUES (123)") will lead to sql query "INSERT INTO test (testField) VALUES (123)" but when SQLCommenter is enabled
the query will get appended with some configurable tags like "INSERT INTO test (testField) VALUES (123) /*tag=value*/;"
SQLCommenter Configurations
***************************
We can configure the tags to be appended to the sqlquery log by adding configuration inside commenter_options(default:{}) keyword
db_driver = True(Default) or False
For example,
::
Enabling this flag will add pymysql and its version, e.g. /*pymysql%%3A1.2.3*/
dbapi_threadsafety = True(Default) or False
For example,
::
Enabling this flag will add threadsafety /*dbapi_threadsafety=2*/
dbapi_level = True(Default) or False
For example,
::
Enabling this flag will add dbapi_level /*dbapi_level='2.0'*/
mysql_client_version = True(Default) or False
For example,
::
Enabling this flag will add mysql_client_version /*mysql_client_version='123'*/
driver_paramstyle = True(Default) or False
For example,
::
Enabling this flag will add driver_paramstyle /*driver_paramstyle='pyformat'*/
opentelemetry_values = True(Default) or False
For example,
::
Enabling this flag will add traceparent values /*traceparent='00-03afa25236b8cd948fa853d67038ac79-405ff022e8247c46-01'*/
API
---
"""
Expand All @@ -59,14 +128,16 @@


class PyMySQLInstrumentor(BaseInstrumentor):
def instrumentation_dependencies(self) -> Collection[str]:
def instrumentation_dependencies(self) -> Collection[str]: # pylint: disable=no-self-use
return _instruments

def _instrument(self, **kwargs):
def _instrument(self, **kwargs): # pylint: disable=no-self-use
"""Integrate with the PyMySQL library.
https://github.com/PyMySQL/PyMySQL/
"""
tracer_provider = kwargs.get("tracer_provider")
enable_sqlcommenter = kwargs.get("enable_commenter", False)
commenter_options = kwargs.get("commenter_options", {})

dbapi.wrap_connect(
__name__,
Expand All @@ -76,14 +147,21 @@ def _instrument(self, **kwargs):
_CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
enable_commenter=enable_sqlcommenter,
commenter_options=commenter_options,
)

def _uninstrument(self, **kwargs):
def _uninstrument(self, **kwargs): # pylint: disable=no-self-use
""" "Disable PyMySQL instrumentation"""
dbapi.unwrap_connect(pymysql, "connect")

@staticmethod
def instrument_connection(connection, tracer_provider=None):
def instrument_connection(
connection,
tracer_provider=None,
enable_commenter=None,
commenter_options=None,
):
"""Enable instrumentation in a PyMySQL connection.
Args:
Expand All @@ -102,6 +180,9 @@ def instrument_connection(connection, tracer_provider=None):
_CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
enable_commenter=enable_commenter,
commenter_options=commenter_options,
connect_module=pymysql,
)

@staticmethod
Expand Down
239 changes: 239 additions & 0 deletions instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@


class TestPyMysqlIntegration(TestBase):
# pylint: disable=invalid-name
def tearDown(self):
super().tearDown()
with self.disable_logging():
Expand Down Expand Up @@ -111,6 +112,244 @@ def test_instrument_connection(self, mock_connect):
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)

@mock.patch("opentelemetry.instrumentation.dbapi.instrument_connection")
@mock.patch("pymysql.connect")
# pylint: disable=unused-argument
def test_instrument_connection_enable_commenter_dbapi_kwargs(
self,
mock_connect,
mock_instrument_connection,
):
cnx = pymysql.connect(database="test")
cnx = PyMySQLInstrumentor().instrument_connection(
cnx,
enable_commenter=True,
commenter_options={"foo": True},
)
cursor = cnx.cursor()
cursor.execute("SELECT * FROM test")
kwargs = mock_instrument_connection.call_args[1]
self.assertEqual(kwargs["enable_commenter"], True)
self.assertEqual(kwargs["commenter_options"], {"foo": True})

def test_instrument_connection_with_dbapi_sqlcomment_enabled(self):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
cnx_proxy = PyMySQLInstrumentor().instrument_connection(
mock_connection,
enable_commenter=True,
)
cnx_proxy.cursor().execute("Select 1;")

spans_list = self.memory_exporter.get_finished_spans()
span = spans_list[0]
span_id = format(span.get_span_context().span_id, "016x")
trace_id = format(span.get_span_context().trace_id, "032x")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
)

def test_instrument_connection_with_dbapi_sqlcomment_enabled_with_options(
self,
):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
cnx_proxy = PyMySQLInstrumentor().instrument_connection(
mock_connection,
enable_commenter=True,
commenter_options={
"dbapi_level": False,
"dbapi_threadsafety": True,
"driver_paramstyle": False,
},
)
cnx_proxy.cursor().execute("Select 1;")

spans_list = self.memory_exporter.get_finished_spans()
span = spans_list[0]
span_id = format(span.get_span_context().span_id, "016x")
trace_id = format(span.get_span_context().trace_id, "032x")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_threadsafety='123',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
)

def test_instrument_connection_with_dbapi_sqlcomment_not_enabled_default(
self,
):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
cnx_proxy = PyMySQLInstrumentor().instrument_connection(
mock_connection,
)
cnx_proxy.cursor().execute("Select 1;")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
"Select 1;",
)

@mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
@mock.patch("pymysql.connect")
# pylint: disable=unused-argument
def test_instrument_enable_commenter_dbapi_kwargs(
self,
mock_connect,
mock_wrap_connect,
):
PyMySQLInstrumentor()._instrument(
enable_commenter=True,
commenter_options={"foo": True},
)
kwargs = mock_wrap_connect.call_args[1]
self.assertEqual(kwargs["enable_commenter"], True)
self.assertEqual(kwargs["commenter_options"], {"foo": True})

def test_instrument_with_dbapi_sqlcomment_enabled(
self,
):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
PyMySQLInstrumentor()._instrument(
enable_commenter=True,
)
cnx = mock_connect_module.connect(database="test")
cursor = cnx.cursor()
cursor.execute("Select 1;")

spans_list = self.memory_exporter.get_finished_spans()
span = spans_list[0]
span_id = format(span.get_span_context().span_id, "016x")
trace_id = format(span.get_span_context().trace_id, "032x")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
)

def test_instrument_with_dbapi_sqlcomment_enabled_with_options(
self,
):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
PyMySQLInstrumentor()._instrument(
enable_commenter=True,
commenter_options={
"dbapi_level": False,
"dbapi_threadsafety": True,
"driver_paramstyle": False,
},
)
cnx = mock_connect_module.connect(database="test")
cursor = cnx.cursor()
cursor.execute("Select 1;")

spans_list = self.memory_exporter.get_finished_spans()
span = spans_list[0]
span_id = format(span.get_span_context().span_id, "016x")
trace_id = format(span.get_span_context().trace_id, "032x")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_threadsafety='123',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
)

def test_instrument_with_dbapi_sqlcomment_not_enabled_default(
self,
):
mock_connect_module = mock.MagicMock(
__name__="pymysql",
__version__="foobar",
threadsafety="123",
apilevel="123",
paramstyle="test",
)
mock_connect_module.get_client_info.return_value = "foobaz"
mock_cursor = mock_connect_module.connect().cursor()
mock_connection = mock.MagicMock()
mock_connection.cursor.return_value = mock_cursor

with mock.patch(
"opentelemetry.instrumentation.pymysql.pymysql",
mock_connect_module,
):
PyMySQLInstrumentor()._instrument()
cnx = mock_connect_module.connect(database="test")
cursor = cnx.cursor()
cursor.execute("Select 1;")
self.assertEqual(
mock_cursor.execute.call_args[0][0],
"Select 1;",
)

@mock.patch("pymysql.connect")
# pylint: disable=unused-argument
def test_uninstrument_connection(self, mock_connect):
Expand Down

0 comments on commit 7b1554b

Please sign in to comment.