diff --git a/CHANGELOG.md b/CHANGELOG.md index 80d63d0395..d59c11bb8d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,6 +60,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ([#2942](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2942)) - `opentelemetry-instrumentation-click`: new instrumentation to trace click commands ([#2994](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2994)) +- `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 d547e806a3..9274aeef8d 100644 --- a/docs-requirements.txt +++ b/docs-requirements.txt @@ -32,6 +32,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 8233fccb15..3dca3d05f5 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -122,6 +122,8 @@ "https://opentelemetry-python.readthedocs.io/en/latest/", None, ), + "psycopg": ("https://www.psycopg.org/psycopg3/docs/", None), + "psycopg2": ("https://www.psycopg.org/docs/", None), } # http://www.sphinx-doc.org/en/master/config.html#confval-nitpicky 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 8c608b7655..f4523dd8c8 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -16,7 +16,7 @@ The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by using ``PsycopgInstrumentor``. -.. _Psycopg: http://initd.org/psycopg/ +.. _Psycopg: https://www.psycopg.org/psycopg3/docs/ SQLCOMMENTER ***************************************** @@ -130,6 +130,7 @@ ) from psycopg.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.psycopg.package import _instruments @@ -154,7 +155,7 @@ def instrumentation_dependencies(self) -> Collection[str]: def _instrument(self, **kwargs): """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) @@ -211,8 +212,13 @@ 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 psycopg connection. + def instrument_connection( + connection: psycopg.Connection, + tracer_provider: typing.Optional[trace_api.TracerProvider] = None, + enable_commenter: bool = False, + commenter_options: dict = None, + ): + """Enable instrumentation of a Psycopg connection. Args: connection: psycopg.Connection @@ -220,6 +226,10 @@ def instrument_connection(connection, tracer_provider=None): 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. Returns: An instrumented psycopg connection object. @@ -232,7 +242,9 @@ 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, ) connection._is_instrumented_by_opentelemetry = True else: @@ -316,7 +328,13 @@ 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, +): if not db_api: db_api = DatabaseApiIntegration( __name__, @@ -324,6 +342,9 @@ def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None): connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES, version=__version__, tracer_provider=tracer_provider, + enable_commenter=enable_commenter, + commenter_options=commenter_options, + connect_module=psycopg, ) 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 4ddaad9174..4081070aeb 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py @@ -378,6 +378,74 @@ 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.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'*/", + ) + + 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.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'*/", + ) + @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") def test_sqlcommenter_disabled(self, event_mocked): cnx = psycopg.connect(database="test") @@ -388,6 +456,33 @@ 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", + ) + + 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", + ) + 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 a811f4285a..341ddb1b7d 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -13,10 +13,10 @@ # limitations under the License. """ -The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by +The integration with PostgreSQL supports the `Psycopg2`_ library, it can be enabled by using ``Psycopg2Instrumentor``. -.. _Psycopg: http://initd.org/psycopg/ +.. _Psycopg2: https://www.psycopg.org/docs/ SQLCOMMENTER ***************************************** @@ -122,11 +122,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 _instruments @@ -150,8 +154,8 @@ def instrumentation_dependencies(self) -> Collection[str]: return _instruments 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) @@ -175,20 +179,28 @@ 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 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. Returns: An instrumented psycopg2 connection object. """ - if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -197,7 +209,9 @@ 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, ) connection._is_instrumented_by_opentelemetry = True else: @@ -260,7 +274,13 @@ 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, +): if not db_api: db_api = DatabaseApiIntegration( __name__, @@ -268,6 +288,9 @@ 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, ) 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..6d38b1c34c 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py @@ -261,6 +261,74 @@ 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'*/", + ) + + 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'*/", + ) + @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") def test_sqlcommenter_disabled(self, event_mocked): cnx = psycopg2.connect(database="test") @@ -271,6 +339,33 @@ 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", + ) + + 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", + ) + def test_no_op_tracer_provider(self): Psycopg2Instrumentor().instrument( tracer_provider=trace.NoOpTracerProvider()