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 d8db967f47..6f9ca5d0dd 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py @@ -71,6 +71,7 @@ def trace_integration( capture_parameters: bool = False, enable_commenter: bool = False, db_api_integration_factory=None, + enable_attribute_commenter: bool = False, ): """Integrate with DB API library. https://www.python.org/dev/peps/pep-0249/ @@ -88,6 +89,7 @@ def trace_integration( enable_commenter: Flag to enable/disable sqlcommenter. db_api_integration_factory: The `DatabaseApiIntegration` to use. If none is passed the default one is used. + enable_attribute_commenter: Flag to enable/disable sqlcomment inclusion in `db.statement` span attribute. Only available if enable_commenter=True. """ wrap_connect( __name__, @@ -100,6 +102,7 @@ def trace_integration( capture_parameters=capture_parameters, enable_commenter=enable_commenter, db_api_integration_factory=db_api_integration_factory, + enable_attribute_commenter=enable_attribute_commenter, ) @@ -115,6 +118,7 @@ def wrap_connect( enable_commenter: bool = False, db_api_integration_factory=None, commenter_options: dict = None, + enable_attribute_commenter: bool = False, ): """Integrate with DB API library. https://www.python.org/dev/peps/pep-0249/ @@ -133,6 +137,7 @@ def wrap_connect( db_api_integration_factory: The `DatabaseApiIntegration` to use. If none is passed the default one is used. commenter_options: Configurations for tags to be appended at the sql query. + enable_attribute_commenter: Flag to enable/disable sqlcomment inclusion in `db.statement` span attribute. Only available if enable_commenter=True. """ db_api_integration_factory = ( @@ -156,6 +161,7 @@ def wrap_connect_( enable_commenter=enable_commenter, commenter_options=commenter_options, connect_module=connect_module, + enable_attribute_commenter=enable_attribute_commenter, ) return db_integration.wrapped_connection(wrapped, args, kwargs) @@ -191,6 +197,7 @@ def instrument_connection( enable_commenter: bool = False, commenter_options: dict = None, connect_module: typing.Callable[..., typing.Any] = None, + enable_attribute_commenter: bool = False, ): """Enable instrumentation in a database connection. @@ -206,6 +213,7 @@ def instrument_connection( enable_commenter: Flag to enable/disable sqlcommenter. commenter_options: Configurations for tags to be appended at the sql query. connect_module: Module name where connect method is available. + enable_attribute_commenter: Flag to enable/disable sqlcomment inclusion in `db.statement` span attribute. Only available if enable_commenter=True. Returns: An instrumented connection. @@ -224,6 +232,7 @@ def instrument_connection( enable_commenter=enable_commenter, commenter_options=commenter_options, connect_module=connect_module, + enable_attribute_commenter=enable_attribute_commenter, ) db_integration.get_connection_attributes(connection) return get_traced_connection_proxy(connection, db_integration) @@ -257,6 +266,7 @@ def __init__( enable_commenter: bool = False, commenter_options: dict = None, connect_module: typing.Callable[..., typing.Any] = None, + enable_attribute_commenter: bool = False, ): self.connection_attributes = connection_attributes if self.connection_attributes is None: @@ -277,6 +287,7 @@ def __init__( self.capture_parameters = capture_parameters self.enable_commenter = enable_commenter self.commenter_options = commenter_options + self.enable_attribute_commenter = enable_attribute_commenter self.database_system = database_system self.connection_props = {} self.span_attributes = {} @@ -434,6 +445,9 @@ def __init__(self, db_api_integration: DatabaseApiIntegration) -> None: if self._db_api_integration.commenter_options else {} ) + self._enable_attribute_commenter = ( + self._db_api_integration.enable_attribute_commenter + ) self._connect_module = self._db_api_integration.connect_module self._leading_comment_remover = re.compile(r"^/\*.*?\*/") @@ -497,51 +511,109 @@ def traced_execution( ) as span: if span.is_recording(): if args and self._commenter_enabled: - try: - args_list = list(args) - - # lazy capture of mysql-connector client version using cursor - if ( - self._db_api_integration.database_system == "mysql" - and self._db_api_integration.connect_module.__name__ - == "mysql.connector" - and not self._db_api_integration.commenter_data[ - "mysql_client_version" - ] - ): - self._db_api_integration.commenter_data[ - "mysql_client_version" - ] = cursor._cnx._cmysql.get_client_info() - - commenter_data = dict( - self._db_api_integration.commenter_data - ) - if self._commenter_options.get( - "opentelemetry_values", True - ): - commenter_data.update( - **_get_opentelemetry_values() + # sqlcomment in query and span attribute + if self._enable_attribute_commenter: + try: + args_list = list(args) + + # lazy capture of mysql-connector client version using cursor + if ( + self._db_api_integration.database_system + == "mysql" + and self._db_api_integration.connect_module.__name__ + == "mysql.connector" + and not self._db_api_integration.commenter_data[ + "mysql_client_version" + ] + ): + self._db_api_integration.commenter_data[ + "mysql_client_version" + ] = cursor._cnx._cmysql.get_client_info() + + commenter_data = dict( + self._db_api_integration.commenter_data + ) + 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, + ) + + self._populate_span(span, cursor, *args) + + # sqlcomment in query only + else: + self._populate_span(span, cursor, *args) + + try: + args_list = list(args) + + # lazy capture of mysql-connector client version using cursor + if ( + self._db_api_integration.database_system + == "mysql" + and self._db_api_integration.connect_module.__name__ + == "mysql.connector" + and not self._db_api_integration.commenter_data[ + "mysql_client_version" + ] + ): + self._db_api_integration.commenter_data[ + "mysql_client_version" + ] = cursor._cnx._cmysql.get_client_info() + + commenter_data = dict( + self._db_api_integration.commenter_data + ) + 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, ) - # 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 - ) - - self._populate_span(span, cursor, *args) + # No sqlcommenting + else: + self._populate_span(span, cursor, *args) return query_method(*args, **kwargs) diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py b/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py index 2ffa2f3d5b..794bcd6efb 100644 --- a/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py +++ b/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py @@ -277,6 +277,47 @@ def test_executemany_comment(self): cursor.query, 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}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_executemany_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "test" + 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", + "postgresql", + enable_commenter=True, + commenter_options={"db_driver": False, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + 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}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + 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_executemany_comment_non_pep_249_compliant(self): class MockConnectModule: @@ -306,8 +347,54 @@ def __getattr__(self, name): cursor.query, r"Select 1 /\*dbapi_level='1.0',dbapi_threadsafety='unknown',driver_paramstyle='unknown',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) - def test_executemany_comment_matches_db_statement_attribute(self): + def test_executemany_comment_non_pep_249_compliant_stmt_enabled(self): + class MockConnectModule: + def __getattr__(self, name): + if name == "__name__": + return "test" + if name == "__version__": + return mock.MagicMock() + if name == "__libpq_version__": + return 123 + raise AttributeError("attribute missing") + + connect_module = MockConnectModule() + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "postgresql", + enable_commenter=True, + connect_module=connect_module, + commenter_options={"db_driver": False}, + enable_attribute_commenter=True, + ) + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*dbapi_level='1.0',dbapi_threadsafety='unknown',driver_paramstyle='unknown',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*dbapi_level='1.0',dbapi_threadsafety='unknown',driver_paramstyle='unknown',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + + def test_executemany_comment_stmt_enabled_matches_db_statement_attribute( + self, + ): connect_module = mock.MagicMock() connect_module.__version__ = mock.MagicMock() connect_module.__libpq_version__ = 123 @@ -321,6 +408,7 @@ def test_executemany_comment_matches_db_statement_attribute(self): enable_commenter=True, commenter_options={"db_driver": False, "dbapi_level": False}, connect_module=connect_module, + enable_attribute_commenter=True, ) mock_connection = db_integration.wrapped_connection( mock_connect, {}, {} @@ -371,6 +459,50 @@ def test_compatible_build_version_psycopg_psycopg2_libpq(self): cursor.query, 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}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_compatible_build_version_psycopg_psycopg2_libpq_stmt_enabled( + self, + ): + connect_module = mock.MagicMock() + connect_module.__name__ = "test" + connect_module.__version__ = mock.MagicMock() + connect_module.pq = mock.MagicMock() + connect_module.pq.__build_version__ = 123 + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "postgresql", + enable_commenter=True, + commenter_options={"db_driver": False, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + 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}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + 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_executemany_psycopg2_integration_comment(self): connect_module = mock.MagicMock() @@ -397,6 +529,47 @@ def test_executemany_psycopg2_integration_comment(self): cursor.query, r"Select 1 /\*db_driver='psycopg2%%3A1.2.3',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}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_executemany_psycopg2_integration_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "psycopg2" + connect_module.__version__ = "1.2.3" + connect_module.__libpq_version__ = 123 + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "postgresql", + enable_commenter=True, + commenter_options={"db_driver": True, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*db_driver='psycopg2%%3A1.2.3',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}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*db_driver='psycopg2%%3A1.2.3',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_executemany_psycopg_integration_comment(self): connect_module = mock.MagicMock() @@ -424,6 +597,48 @@ def test_executemany_psycopg_integration_comment(self): cursor.query, r"Select 1 /\*db_driver='psycopg%%3A1.2.3',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}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_executemany_psycopg_integration_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "psycopg" + connect_module.__version__ = "1.2.3" + connect_module.pq = mock.MagicMock() + connect_module.pq.__build_version__ = 123 + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "postgresql", + enable_commenter=True, + commenter_options={"db_driver": True, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*db_driver='psycopg%%3A1.2.3',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}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*db_driver='psycopg%%3A1.2.3',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_executemany_mysqlconnector_integration_comment(self): connect_module = mock.MagicMock() @@ -450,6 +665,47 @@ def test_executemany_mysqlconnector_integration_comment(self): cursor.query, r"Select 1 /\*db_driver='mysql.connector%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='1.2.3',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_executemany_mysqlconnector_integration_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "mysql.connector" + connect_module.__version__ = "1.2.3" + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "mysql", + enable_commenter=True, + commenter_options={"db_driver": True, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*db_driver='mysql.connector%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='1.2.3',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*db_driver='mysql.connector%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='1.2.3',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) @mock.patch("opentelemetry.instrumentation.dbapi.util_version") def test_executemany_mysqlclient_integration_comment( @@ -485,6 +741,56 @@ def test_executemany_mysqlclient_integration_comment( cursor.query, r"Select 1 /\*db_driver='MySQLdb%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + @mock.patch("opentelemetry.instrumentation.dbapi.util_version") + def test_executemany_mysqlclient_integration_comment_stmt_enabled( + self, + mock_dbapi_util_version, + ): + mock_dbapi_util_version.return_value = "1.2.3" + connect_module = mock.MagicMock() + connect_module.__name__ = "MySQLdb" + connect_module.__version__ = "1.2.3" + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + connect_module._mysql = mock.MagicMock() + connect_module._mysql.get_client_info = mock.MagicMock( + return_value="123" + ) + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "mysql", + enable_commenter=True, + commenter_options={"db_driver": True, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*db_driver='MySQLdb%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*db_driver='MySQLdb%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) def test_executemany_pymysql_integration_comment(self): connect_module = mock.MagicMock() @@ -512,6 +818,48 @@ def test_executemany_pymysql_integration_comment(self): cursor.query, r"Select 1 /\*db_driver='pymysql%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + def test_executemany_pymysql_integration_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "pymysql" + connect_module.__version__ = "1.2.3" + connect_module.apilevel = 123 + connect_module.threadsafety = 123 + connect_module.paramstyle = "test" + connect_module.get_client_info = mock.MagicMock(return_value="123") + + db_integration = dbapi.DatabaseApiIntegration( + "testname", + "mysql", + enable_commenter=True, + commenter_options={"db_driver": True, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*db_driver='pymysql%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*db_driver='pymysql%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) def test_executemany_flask_integration_comment(self): connect_module = mock.MagicMock() @@ -544,6 +892,58 @@ def test_executemany_flask_integration_comment(self): cursor.query, r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',flask=1,libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1;", + ) + + clear_context = context.set_value( + "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}, current_context + ) + context.attach(clear_context) + + def test_executemany_flask_integration_comment_stmt_enabled(self): + connect_module = mock.MagicMock() + connect_module.__name__ = "test" + 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", + "postgresql", + enable_commenter=True, + commenter_options={"db_driver": False, "dbapi_level": False}, + connect_module=connect_module, + enable_attribute_commenter=True, + ) + current_context = context.get_current() + sqlcommenter_context = context.set_value( + "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context + ) + context.attach(sqlcommenter_context) + + mock_connection = db_integration.wrapped_connection( + mock_connect, {}, {} + ) + cursor = mock_connection.cursor() + cursor.executemany("Select 1;") + self.assertRegex( + cursor.query, + r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',flask=1,libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) + spans_list = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans_list), 1) + span = spans_list[0] + self.assertRegex( + span.attributes[SpanAttributes.DB_STATEMENT], + r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',flask=1,libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;", + ) clear_context = context.set_value( "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}, current_context @@ -603,6 +1003,7 @@ def test_instrument_connection_kwargs_defaults(self, mock_dbapiint): self.assertEqual(kwargs["enable_commenter"], False) self.assertEqual(kwargs["commenter_options"], None) self.assertEqual(kwargs["connect_module"], None) + self.assertEqual(kwargs["enable_attribute_commenter"], False) @mock.patch("opentelemetry.instrumentation.dbapi.DatabaseApiIntegration") def test_instrument_connection_kwargs_provided(self, mock_dbapiint): @@ -619,6 +1020,7 @@ def test_instrument_connection_kwargs_provided(self, mock_dbapiint): enable_commenter=True, commenter_options={"foo": "bar"}, connect_module=mock_connect_module, + enable_attribute_commenter=True, ) kwargs = mock_dbapiint.call_args[1] self.assertEqual(kwargs["connection_attributes"], {"foo": "bar"}) @@ -628,6 +1030,7 @@ def test_instrument_connection_kwargs_provided(self, mock_dbapiint): self.assertEqual(kwargs["enable_commenter"], True) self.assertEqual(kwargs["commenter_options"], {"foo": "bar"}) self.assertIs(kwargs["connect_module"], mock_connect_module) + self.assertEqual(kwargs["enable_attribute_commenter"], True) def test_uninstrument_connection(self): connection = mock.Mock()