diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py index 1295e1fddf..b0a982ebf7 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -243,6 +243,7 @@ def instrument_connection( tracer_provider: typing.Optional[trace_api.TracerProvider] = None, enable_commenter: bool = False, commenter_options: dict = None, + enable_attribute_commenter=None, ): """Enable instrumentation of a Psycopg connection. @@ -256,6 +257,8 @@ def instrument_connection( Optional flag to enable/disable sqlcommenter (default False). commenter_options: dict, optional Optional configurations for tags to be appended at the sql query. + enable_attribute_commenter: + Optional flag to enable/disable addition of sqlcomment to span attribute (default False). Requires enable_commenter=True. Returns: An instrumented psycopg connection object. @@ -271,6 +274,7 @@ def instrument_connection( tracer_provider=tracer_provider, enable_commenter=enable_commenter, commenter_options=commenter_options, + enable_attribute_commenter=enable_attribute_commenter, ) connection._is_instrumented_by_opentelemetry = True else: @@ -360,6 +364,7 @@ def _new_cursor_factory( tracer_provider: typing.Optional[trace_api.TracerProvider] = None, enable_commenter: bool = False, commenter_options: dict = None, + enable_attribute_commenter: bool = False, ): if not db_api: db_api = DatabaseApiIntegration( @@ -371,6 +376,7 @@ def _new_cursor_factory( enable_commenter=enable_commenter, commenter_options=commenter_options, connect_module=psycopg, + enable_attribute_commenter=enable_attribute_commenter, ) base_factory = base_factory or pg_cursor diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py index 4081070aeb..e8ad82a5f1 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py @@ -20,6 +20,7 @@ import opentelemetry.instrumentation.psycopg from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor from opentelemetry.sdk import resources +from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.test.test_base import TestBase @@ -411,6 +412,49 @@ def test_sqlcommenter_enabled_instrument_connection_defaults(self): MockCursor.execute.call_args[0][0], f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) + + def test_sqlcommenter_enabled_instrument_connection_stmt_enabled(self): + with mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.__version__", + "foobar", + ), mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__", + "foobaz", + ), mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.threadsafety", + "123", + ), mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.apilevel", + "123", + ), mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.paramstyle", + "test", + ): + cnx = psycopg.connect(database="test") + cnx = PsycopgInstrumentor().instrument_connection( + cnx, + enable_commenter=True, + enable_attribute_commenter=True, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + 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( + MockCursor.execute.call_args[0][0], + f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", + ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", + ) def test_sqlcommenter_enabled_instrument_connection_with_options(self): with mock.patch( @@ -445,6 +489,10 @@ def test_sqlcommenter_enabled_instrument_connection_with_options(self): MockCursor.execute.call_args[0][0], f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") def test_sqlcommenter_disabled(self, event_mocked): @@ -468,6 +516,12 @@ def test_sqlcommenter_disabled_default_instrument_connection(self): MockCursor.execute.call_args[0][0], "Select 1", ) + spans_list = self.memory_exporter.get_finished_spans() + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) def test_sqlcommenter_disabled_explicit_instrument_connection(self): cnx = psycopg.connect(database="test") @@ -482,6 +536,12 @@ def test_sqlcommenter_disabled_explicit_instrument_connection(self): MockCursor.execute.call_args[0][0], "Select 1", ) + spans_list = self.memory_exporter.get_finished_spans() + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) class TestPostgresqlIntegrationAsync( 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 36a8714222..6f996f0115 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -208,6 +208,7 @@ def instrument_connection( tracer_provider: typing.Optional[trace_api.TracerProvider] = None, enable_commenter: bool = False, commenter_options: dict = None, + enable_attribute_commenter=None, ): """Enable instrumentation of a Psycopg2 connection. @@ -221,6 +222,8 @@ def instrument_connection( Optional flag to enable/disable sqlcommenter (default False). commenter_options: dict, optional Optional configurations for tags to be appended at the sql query. + enable_attribute_commenter: + Optional flag to enable/disable addition of sqlcomment to span attribute (default False). Requires enable_commenter=True. Returns: An instrumented psycopg2 connection object. @@ -236,6 +239,7 @@ def instrument_connection( tracer_provider=tracer_provider, enable_commenter=enable_commenter, commenter_options=commenter_options, + enable_attribute_commenter=enable_attribute_commenter, ) connection._is_instrumented_by_opentelemetry = True else: @@ -304,6 +308,7 @@ def _new_cursor_factory( tracer_provider: typing.Optional[trace_api.TracerProvider] = None, enable_commenter: bool = False, commenter_options: dict = None, + enable_attribute_commenter: bool = False, ): if not db_api: db_api = DatabaseApiIntegration( @@ -315,6 +320,7 @@ def _new_cursor_factory( enable_commenter=enable_commenter, commenter_options=commenter_options, connect_module=psycopg2, + enable_attribute_commenter=enable_attribute_commenter, ) base_factory = base_factory or pg_cursor diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py index 6d38b1c34c..0ed104539a 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py @@ -21,6 +21,7 @@ from opentelemetry import trace from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor from opentelemetry.sdk import resources +from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.test.test_base import TestBase @@ -294,6 +295,49 @@ def test_sqlcommenter_enabled_instrument_connection_defaults(self): MockCursor.execute.call_args[0][0], f"Select 1 /*db_driver='psycopg2%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) + + def test_sqlcommenter_enabled_instrument_connection_stmt_enabled(self): + with mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.__version__", + "foobar", + ), mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__", + "foobaz", + ), mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety", + "123", + ), mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.apilevel", + "123", + ), mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.paramstyle", + "test", + ): + cnx = psycopg2.connect(database="test") + cnx = Psycopg2Instrumentor().instrument_connection( + cnx, + enable_commenter=True, + enable_attribute_commenter=True, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + 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( + MockCursor.execute.call_args[0][0], + f"Select 1 /*db_driver='psycopg2%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", + ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + f"Select 1 /*db_driver='psycopg2%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", + ) def test_sqlcommenter_enabled_instrument_connection_with_options(self): with mock.patch( @@ -328,6 +372,10 @@ def test_sqlcommenter_enabled_instrument_connection_with_options(self): MockCursor.execute.call_args[0][0], f"Select 1 /*db_driver='psycopg2%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/", ) + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") def test_sqlcommenter_disabled(self, event_mocked): @@ -351,6 +399,12 @@ def test_sqlcommenter_disabled_default_instrument_connection(self): MockCursor.execute.call_args[0][0], "Select 1", ) + spans_list = self.memory_exporter.get_finished_spans() + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) def test_sqlcommenter_disabled_explicit_instrument_connection(self): cnx = psycopg2.connect(database="test") @@ -365,6 +419,12 @@ def test_sqlcommenter_disabled_explicit_instrument_connection(self): MockCursor.execute.call_args[0][0], "Select 1", ) + spans_list = self.memory_exporter.get_finished_spans() + span = spans_list[0] + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + "Select 1", + ) def test_no_op_tracer_provider(self): Psycopg2Instrumentor().instrument(