diff --git a/CHANGELOG.md b/CHANGELOG.md index 52bc006893..3c13223d51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ([#4141](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4141)) - `opentelemetry-instrumentation-pyramid`: pass request attributes at span creation ([#4139](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4139)) +- `opentelemetry-instrumentation-psycopg2`, `opentelemetry-instrumentation-psycopg`: Add sqlcommenter support for `instrument_connection` + ([#3071](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3071)) ### Fixed diff --git a/docs-requirements.txt b/docs-requirements.txt index 9d5e42fe74..da5e9b6b25 100644 --- a/docs-requirements.txt +++ b/docs-requirements.txt @@ -39,6 +39,7 @@ mysqlclient~=2.1.1 openai >= 1.26.0 psutil>=5 psycopg~=3.1.17 +psycopg2~=2.9.9 pika>=0.12.0 pymongo~=4.6.3 PyMySQL~=1.1.1 diff --git a/docs/conf.py b/docs/conf.py index e3e8ae145b..223b4d35a5 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -129,6 +129,8 @@ "https://opentelemetry-python.readthedocs.io/en/latest/", None, ), + "psycopg": ("https://www.psycopg.org/psycopg3/docs/", None), + "psycopg2": ("https://www.psycopg.org/docs/", None), "redis": ("https://redis.readthedocs.io/en/latest/", None), "fsspec": ("https://filesystem-spec.readthedocs.io/en/latest/", None), } 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 28896be138..01964f8986 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -72,6 +72,13 @@ from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor PsycopgInstrumentor().instrument(enable_commenter=True) + # OR with specific connection + cnx = psycopg.connect(database='Database') + instrumented_cnx = PsycopgInstrumentor().instrument_connection( + cnx, + enable_commenter=True, + ) + SQLCommenter with commenter_options @@ -144,11 +151,12 @@ from __future__ import annotations import logging -from typing import Any, Callable, Collection, TypeVar +from typing import Any, Callable, Collection, Optional, TypeVar import psycopg # pylint: disable=import-self from psycopg.sql import Composable # pylint: disable=no-name-in-module +from opentelemetry import trace as trace_api from opentelemetry.instrumentation import dbapi from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.psycopg.package import _instruments @@ -179,7 +187,7 @@ def instrumentation_dependencies(self) -> Collection[str]: def _instrument(self, **kwargs: Any): """Integrate with PostgreSQL Psycopg library. - Psycopg: http://initd.org/psycopg/ + Psycopg: https://www.psycopg.org/psycopg3/docs/ """ tracer_provider = kwargs.get("tracer_provider") enable_sqlcommenter = kwargs.get("enable_commenter", False) @@ -247,9 +255,13 @@ def _uninstrument(self, **kwargs: Any): # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod def instrument_connection( - connection: ConnectionT, tracer_provider: TracerProvider | None = None - ) -> ConnectionT: - """Enable instrumentation in a psycopg connection. + connection: psycopg.Connection, + tracer_provider: Optional[trace_api.TracerProvider] = None, + enable_commenter: bool = False, + commenter_options: dict = None, + enable_attribute_commenter: bool = False, + ): + """Enable instrumentation of a Psycopg connection. Args: connection: psycopg.Connection @@ -257,6 +269,12 @@ def instrument_connection( tracer_provider: opentelemetry.trace.TracerProvider, optional The TracerProvider to use for instrumentation. If not provided, the global TracerProvider will be used. + enable_commenter: bool, optional + 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. @@ -270,11 +288,17 @@ def instrument_connection( ) if isinstance(connection, psycopg.AsyncConnection): connection.cursor_factory = _new_cursor_async_factory( - tracer_provider=tracer_provider + tracer_provider=tracer_provider, + enable_commenter=enable_commenter, + commenter_options=commenter_options, + enable_attribute_commenter=enable_attribute_commenter, ) else: connection.cursor_factory = _new_cursor_factory( - tracer_provider=tracer_provider + 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: @@ -359,9 +383,12 @@ def get_statement(self, cursor: CursorT, args: list[Any]) -> str: def _new_cursor_factory( - db_api: DatabaseApiIntegration | None = None, - base_factory: type[psycopg.Cursor] | None = None, - tracer_provider: TracerProvider | None = None, + db_api: dbapi.DatabaseApiIntegration = None, + base_factory: psycopg.Cursor = None, + tracer_provider: 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( @@ -370,6 +397,10 @@ def _new_cursor_factory( connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES, version=__version__, tracer_provider=tracer_provider, + enable_commenter=enable_commenter, + commenter_options=commenter_options, + connect_module=psycopg, + enable_attribute_commenter=enable_attribute_commenter, ) base_factory = base_factory or psycopg.Cursor @@ -398,6 +429,9 @@ def _new_cursor_async_factory( db_api: DatabaseApiAsyncIntegration | None = None, base_factory: type[psycopg.AsyncCursor] | None = None, tracer_provider: TracerProvider | None = None, + enable_commenter: bool = False, + commenter_options: dict = None, + enable_attribute_commenter: bool = False, ): if not db_api: db_api = DatabaseApiAsyncIntegration( @@ -406,6 +440,10 @@ def _new_cursor_async_factory( connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES, version=__version__, tracer_provider=tracer_provider, + enable_commenter=enable_commenter, + commenter_options=commenter_options, + connect_module=psycopg, + enable_attribute_commenter=enable_attribute_commenter, ) base_factory = base_factory or psycopg.AsyncCursor _cursor_tracer = CursorTracer(db_api) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py index 84d6709bbb..c287891088 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py @@ -22,6 +22,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 @@ -451,6 +452,149 @@ def test_sqlcommenter_enabled(self, event_mocked): kwargs = event_mocked.call_args[1] self.assertEqual(kwargs["enable_commenter"], True) + def test_sqlcommenter_enabled_instrument_connection_defaults(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.pq.version", + new=lambda: "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, + ) + 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], + "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.pq.version", + new=lambda: "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( + "opentelemetry.instrumentation.psycopg.psycopg.__version__", + "foobar", + ), + mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__", + "foobaz", + ), + mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.pq.version", + new=lambda: "foobaz", + ), + mock.patch( + "opentelemetry.instrumentation.psycopg.psycopg.threadsafety", + "123", + ), + ): + cnx = psycopg.connect(database="test") + cnx = PsycopgInstrumentor().instrument_connection( + cnx, + enable_commenter=True, + commenter_options={ + "dbapi_level": False, + "dbapi_threadsafety": True, + "driver_paramstyle": False, + "foo": "ignored", + }, + ) + 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_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): cnx = psycopg.connect(database="test") @@ -461,6 +605,45 @@ def test_sqlcommenter_disabled(self, event_mocked): kwargs = event_mocked.call_args[1] self.assertEqual(kwargs["enable_commenter"], False) + def test_sqlcommenter_disabled_default_instrument_connection(self): + cnx = psycopg.connect(database="test") + cnx = PsycopgInstrumentor().instrument_connection( + cnx, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + self.assertEqual( + 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") + cnx = PsycopgInstrumentor().instrument_connection( + cnx, + enable_commenter=False, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + self.assertEqual( + 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( PostgresqlIntegrationTestMixin, TestBase, IsolatedAsyncioTestCase 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 f5e113a47e..2d4eb72a8f 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -72,6 +72,13 @@ Psycopg2Instrumentor().instrument(enable_commenter=True) + # OR with specific connection + cnx = psycopg2.connect(database='Database') + instrumented_cnx = Psycopg2Instrumentor().instrument_connection( + cnx, + enable_commenter=True, + ) + SQLCommenter with commenter_options *********************************** @@ -146,11 +153,15 @@ from typing import Collection import psycopg2 +from psycopg2.extensions import ( + connection as pg_connection, # pylint: disable=no-name-in-module +) from psycopg2.extensions import ( cursor as pg_cursor, # pylint: disable=no-name-in-module ) from psycopg2.sql import Composed # pylint: disable=no-name-in-module +from opentelemetry import trace as trace_api from opentelemetry.instrumentation import dbapi from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.psycopg2.package import ( @@ -193,8 +204,8 @@ def instrumentation_dependencies(self) -> Collection[str]: return _instruments_any def _instrument(self, **kwargs): - """Integrate with PostgreSQL Psycopg library. - Psycopg: http://initd.org/psycopg/ + """Integrate with PostgreSQL Psycopg2 library. + Psycopg2: https://www.psycopg.org/docs/ """ tracer_provider = kwargs.get("tracer_provider") enable_sqlcommenter = kwargs.get("enable_commenter", False) @@ -222,20 +233,31 @@ def _uninstrument(self, **kwargs): # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql @staticmethod - def instrument_connection(connection, tracer_provider=None): - """Enable instrumentation in a psycopg2 connection. + def instrument_connection( + connection: pg_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. Args: connection: psycopg2.extensions.connection The psycopg2 connection object to be instrumented. tracer_provider: opentelemetry.trace.TracerProvider, optional - The TracerProvider to use for instrumentation. If not specified, + The TracerProvider to use for instrumentation. If not provided, the global TracerProvider will be used. + enable_commenter: bool, optional + 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. """ - if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -244,7 +266,10 @@ def instrument_connection(connection, tracer_provider=None): connection, _OTEL_CURSOR_FACTORY_KEY, connection.cursor_factory ) connection.cursor_factory = _new_cursor_factory( - tracer_provider=tracer_provider + 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: @@ -307,7 +332,14 @@ def get_statement(self, cursor, args): return statement -def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None): +def _new_cursor_factory( + db_api: dbapi.DatabaseApiIntegration = None, + base_factory: pg_cursor = None, + 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( __name__, @@ -315,6 +347,10 @@ def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None): connection_attributes=Psycopg2Instrumentor._CONNECTION_ATTRIBUTES, version=__version__, tracer_provider=tracer_provider, + 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 9a6a5ff2fa..ac4fe939e7 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 @@ -261,6 +262,137 @@ def test_sqlcommenter_enabled(self, event_mocked): kwargs = event_mocked.call_args[1] self.assertEqual(kwargs["enable_commenter"], True) + def test_sqlcommenter_enabled_instrument_connection_defaults(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, + ) + 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], + "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( + "opentelemetry.instrumentation.psycopg2.psycopg2.__version__", + "foobar", + ), + mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__", + "foobaz", + ), + mock.patch( + "opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety", + "123", + ), + ): + cnx = psycopg2.connect(database="test") + cnx = Psycopg2Instrumentor().instrument_connection( + cnx, + enable_commenter=True, + commenter_options={ + "dbapi_level": False, + "dbapi_threadsafety": True, + "driver_paramstyle": False, + "foo": "ignored", + }, + ) + 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_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): cnx = psycopg2.connect(database="test") @@ -271,6 +403,45 @@ def test_sqlcommenter_disabled(self, event_mocked): kwargs = event_mocked.call_args[1] self.assertEqual(kwargs["enable_commenter"], False) + def test_sqlcommenter_disabled_default_instrument_connection(self): + cnx = psycopg2.connect(database="test") + cnx = Psycopg2Instrumentor().instrument_connection( + cnx, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + self.assertEqual( + 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") + cnx = Psycopg2Instrumentor().instrument_connection( + cnx, + enable_commenter=False, + ) + query = "Select 1" + cursor = cnx.cursor() + cursor.execute(query) + self.assertEqual( + 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( tracer_provider=trace.NoOpTracerProvider()