From fbac4245159e51877efa86581724995bd1415a18 Mon Sep 17 00:00:00 2001 From: Thiyagu55 <64461612+Thiyagu55@users.noreply.github.com> Date: Sat, 30 Jul 2022 05:14:47 +0530 Subject: [PATCH] Add psycopg2 native tags to sqlcommenter (#1203) --- CHANGELOG.md | 2 + .../instrumentation/dbapi/__init__.py | 46 +++++++++++-- .../tests/test_dbapi_integration.py | 16 ++++- .../instrumentation/psycopg2/__init__.py | 64 +++++++++++++++++++ .../postgres/test_psycopg_sqlcommenter.py | 2 +- 5 files changed, 123 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 793568e433..f252ee5336 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ([#1187](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1187)) - SQLCommenter semicolon bug fix ([#1200](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1200/files)) +- Add psycopg2 native tags to sqlcommenter + ([#1203](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1203)) ### Added - `opentelemetry-instrumentation-redis` add support to instrument RedisCluster clients diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py index 397a80b2bd..1afeb5499c 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py @@ -105,6 +105,7 @@ def wrap_connect( capture_parameters: bool = False, enable_commenter: bool = False, db_api_integration_factory=None, + commenter_options: dict = None, ): """Integrate with DB API library. https://www.python.org/dev/peps/pep-0249/ @@ -119,6 +120,8 @@ def wrap_connect( tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to use. If omitted the current configured one is used. capture_parameters: Configure if db.statement.parameters should be captured. + enable_commenter: Flag to enable/disable sqlcommenter. + commenter_options: Configurations for tags to be appended at the sql query. """ db_api_integration_factory = ( @@ -140,6 +143,8 @@ def wrap_connect_( tracer_provider=tracer_provider, capture_parameters=capture_parameters, enable_commenter=enable_commenter, + commenter_options=commenter_options, + connect_module=connect_module, ) return db_integration.wrapped_connection(wrapped, args, kwargs) @@ -173,6 +178,7 @@ def instrument_connection( tracer_provider: typing.Optional[TracerProvider] = None, capture_parameters: bool = False, enable_commenter: bool = False, + commenter_options: dict = None, ): """Enable instrumentation in a database connection. @@ -185,6 +191,9 @@ def instrument_connection( tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to use. If omitted the current configured one is used. capture_parameters: Configure if db.statement.parameters should be captured. + enable_commenter: Flag to enable/disable sqlcommenter. + commenter_options: Configurations for tags to be appended at the sql query. + Returns: An instrumented connection. """ @@ -200,6 +209,7 @@ def instrument_connection( tracer_provider=tracer_provider, capture_parameters=capture_parameters, enable_commenter=enable_commenter, + commenter_options=commenter_options, ) db_integration.get_connection_attributes(connection) return get_traced_connection_proxy(connection, db_integration) @@ -231,6 +241,8 @@ def __init__( tracer_provider: typing.Optional[TracerProvider] = None, capture_parameters: bool = False, enable_commenter: bool = False, + commenter_options: dict = None, + connect_module: typing.Callable[..., typing.Any] = None, ): self.connection_attributes = connection_attributes if self.connection_attributes is None: @@ -249,11 +261,13 @@ def __init__( ) self.capture_parameters = capture_parameters self.enable_commenter = enable_commenter + self.commenter_options = commenter_options self.database_system = database_system self.connection_props = {} self.span_attributes = {} self.name = "" self.database = "" + self.connect_module = connect_module def wrapped_connection( self, @@ -335,12 +349,18 @@ class CursorTracer: def __init__(self, db_api_integration: DatabaseApiIntegration) -> None: self._db_api_integration = db_api_integration self._commenter_enabled = self._db_api_integration.enable_commenter + self._commenter_options = ( + self._db_api_integration.commenter_options + if self._db_api_integration.commenter_options + else {} + ) + self._connect_module = self._db_api_integration.connect_module def _populate_span( self, span: trace_api.Span, cursor, - *args: typing.Tuple[typing.Any, typing.Any] + *args: typing.Tuple[typing.Any, typing.Any], ): if not span.is_recording(): return @@ -380,7 +400,7 @@ def traced_execution( cursor, query_method: typing.Callable[..., typing.Any], *args: typing.Tuple[typing.Any, typing.Any], - **kwargs: typing.Dict[typing.Any, typing.Any] + **kwargs: typing.Dict[typing.Any, typing.Any], ): name = self.get_operation_name(cursor, args) if not name: @@ -397,14 +417,32 @@ def traced_execution( if args and self._commenter_enabled: try: args_list = list(args) - commenter_data = {} - commenter_data.update(_get_opentelemetry_values()) + commenter_data = dict( + # Psycopg2/framework information + db_driver=f"psycopg2:{self._connect_module.__version__.split(' ')[0]}", + dbapi_threadsafety=self._connect_module.threadsafety, + dbapi_level=self._connect_module.apilevel, + libpq_version=self._connect_module.__libpq_version__, + driver_paramstyle=self._connect_module.paramstyle, + ) + if self._commenter_options.get( + "opentelemetry_values", True + ): + commenter_data.update(**_get_opentelemetry_values()) + + # Filter down to just the requested attributes. + commenter_data = { + k: v + for k, v in commenter_data.items() + if self._commenter_options.get(k, True) + } statement = _add_sql_comment( args_list[0], **commenter_data ) args_list[0] = statement args = tuple(args_list) + except Exception as exc: # pylint: disable=broad-except _logger.exception( "Exception while generating sql comment: %s", exc diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py b/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py index 1ec5cd9df2..af2cbeea15 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py +++ b/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py @@ -229,8 +229,20 @@ def test_executemany(self): ) def test_executemany_comment(self): + + connect_module = mock.MagicMock() + connect_module.__version__ = mock.MagicMock() + connect_module.__libpq_version__ = 123 + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + db_integration = dbapi.DatabaseApiIntegration( - "testname", "testcomponent", enable_commenter=True + "testname", + "testcomponent", + enable_commenter=True, + commenter_options={"db_driver": False, "dbapi_level": False}, + connect_module=connect_module, ) mock_connection = db_integration.wrapped_connection( mock_connect, {}, {} @@ -239,7 +251,7 @@ def test_executemany_comment(self): cursor.executemany("Select 1;") self.assertRegex( cursor.query, - r"Select 1 /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) def test_callproc(self): diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py index 4b7ff23ea6..ebad55cf2a 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -18,6 +18,68 @@ .. _Psycopg: http://initd.org/psycopg/ +SQLCOMMENTER +***************************************** +You can optionally configure Psycopg2 instrumentation to enable sqlcommenter which enriches +the query with contextual information. + +Usage +----- + +.. code:: python + + from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + + Psycopg2Instrumentor().instrument(enable_commenter=True, commenter_options={}) + + +For example, +:: + + Invoking cursor.execute("select * from auth_users") will lead to sql query "select * from auth_users" but when SQLCommenter is enabled + the query will get appended with some configurable tags like "select * from auth_users /*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 psycopg2 and it's version which is /*psycopg2%%3A2.9.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'*/ + +libpq_version = True(Default) or False + +For example, +:: +Enabling this flag will add libpq_version /*libpq_version=140001*/ + +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'*/ + Usage ----- @@ -77,6 +139,7 @@ def _instrument(self, **kwargs): """ tracer_provider = kwargs.get("tracer_provider") enable_sqlcommenter = kwargs.get("enable_commenter", False) + commenter_options = kwargs.get("commenter_options", {}) dbapi.wrap_connect( __name__, psycopg2, @@ -87,6 +150,7 @@ def _instrument(self, **kwargs): tracer_provider=tracer_provider, db_api_integration_factory=DatabaseApiIntegration, enable_commenter=enable_sqlcommenter, + commenter_options=commenter_options, ) def _uninstrument(self, **kwargs): diff --git a/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py index 4f5679a56b..0dee17f865 100644 --- a/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py +++ b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py @@ -50,5 +50,5 @@ def test_commenter_enabled(self): self._cursor.execute("SELECT 1;") self.assertRegex( self._cursor.query.decode("ascii"), - r"SELECT 1 /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + r"SELECT 1 /\*db_driver='psycopg2(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", )