Skip to content

Commit

Permalink
psycopg(2) instrument_connection support enable_attr_commenter
Browse files Browse the repository at this point in the history
  • Loading branch information
@tammy-baylis-swi
tammy-baylis-swi committed Jan 14, 2025
1 parent 90051b7 commit 3587846
Show file tree
Hide file tree
Showing 4 changed files with 132 additions and 0 deletions.
6 changes: 6 additions & 0 deletions ...entelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down
60 changes: 60 additions & 0 deletions instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand All @@ -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")
Expand All @@ -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(
Expand Down
6 changes: 6 additions & 0 deletions ...telemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down
60 changes: 60 additions & 0 deletions instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand All @@ -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")
Expand All @@ -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(
Expand Down

0 comments on commit 3587846

Please sign in to comment.