From d2da8cda4e1c9f12f8b9ca787f0817abc30fa6d8 Mon Sep 17 00:00:00 2001 From: tammy-baylis-swi Date: Thu, 5 Dec 2024 15:27:18 -0800 Subject: [PATCH] enable_commenter for instrument_connection psycopg(2) --- .../instrumentation/psycopg/__init__.py | 23 ++++++++++++++++--- .../instrumentation/psycopg2/__init__.py | 23 ++++++++++++++++--- 2 files changed, 40 insertions(+), 6 deletions(-) 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 e986ec0d46..50913e2703 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py @@ -114,6 +114,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 @@ -195,7 +196,12 @@ 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): + def instrument_connection( + connection, + tracer_provider=None, + enable_commenter: bool = False, + commenter_options: dict = None, + ): if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -204,7 +210,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: @@ -288,7 +296,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__, @@ -296,6 +310,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-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py index de2e49f4c3..600616a0b7 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -111,6 +111,7 @@ ) 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 @@ -159,7 +160,12 @@ 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): + def instrument_connection( + connection, + tracer_provider=None, + enable_commenter: bool = False, + commenter_options: dict = None, + ): if not hasattr(connection, "_is_instrumented_by_opentelemetry"): connection._is_instrumented_by_opentelemetry = False @@ -168,7 +174,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: @@ -231,7 +239,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__, @@ -239,6 +253,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