Skip to content

Commit 3587846

Browse files
tammy-baylis-switammy-baylis-swi
committed
psycopg(2) instrument_connection support enable_attr_commenter
1 parent 90051b7 commit 3587846

File tree

4 files changed

+132
-0
lines changed

4 files changed

+132
-0
lines changed

‎instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@ def instrument_connection(
243243
tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
244244
enable_commenter: bool = False,
245245
commenter_options: dict = None,
246+
enable_attribute_commenter=None,
246247
):
247248
"""Enable instrumentation of a Psycopg connection.
248249
@@ -256,6 +257,8 @@ def instrument_connection(
256257
Optional flag to enable/disable sqlcommenter (default False).
257258
commenter_options: dict, optional
258259
Optional configurations for tags to be appended at the sql query.
260+
enable_attribute_commenter:
261+
Optional flag to enable/disable addition of sqlcomment to span attribute (default False). Requires enable_commenter=True.
259262
260263
Returns:
261264
An instrumented psycopg connection object.
@@ -271,6 +274,7 @@ def instrument_connection(
271274
tracer_provider=tracer_provider,
272275
enable_commenter=enable_commenter,
273276
commenter_options=commenter_options,
277+
enable_attribute_commenter=enable_attribute_commenter,
274278
)
275279
connection._is_instrumented_by_opentelemetry = True
276280
else:
@@ -360,6 +364,7 @@ def _new_cursor_factory(
360364
tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
361365
enable_commenter: bool = False,
362366
commenter_options: dict = None,
367+
enable_attribute_commenter: bool = False,
363368
):
364369
if not db_api:
365370
db_api = DatabaseApiIntegration(
@@ -371,6 +376,7 @@ def _new_cursor_factory(
371376
enable_commenter=enable_commenter,
372377
commenter_options=commenter_options,
373378
connect_module=psycopg,
379+
enable_attribute_commenter=enable_attribute_commenter,
374380
)
375381

376382
base_factory = base_factory or pg_cursor

‎instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import opentelemetry.instrumentation.psycopg
2121
from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
2222
from opentelemetry.sdk import resources
23+
from opentelemetry.semconv.trace import SpanAttributes
2324
from opentelemetry.test.test_base import TestBase
2425

2526

@@ -411,6 +412,49 @@ def test_sqlcommenter_enabled_instrument_connection_defaults(self):
411412
MockCursor.execute.call_args[0][0],
412413
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'*/",
413414
)
415+
self.assertEqual(
416+
span.attributes[SpanAttributes.DB_STATEMENT],
417+
"Select 1",
418+
)
419+
420+
def test_sqlcommenter_enabled_instrument_connection_stmt_enabled(self):
421+
with mock.patch(
422+
"opentelemetry.instrumentation.psycopg.psycopg.__version__",
423+
"foobar",
424+
), mock.patch(
425+
"opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
426+
"foobaz",
427+
), mock.patch(
428+
"opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
429+
"123",
430+
), mock.patch(
431+
"opentelemetry.instrumentation.psycopg.psycopg.apilevel",
432+
"123",
433+
), mock.patch(
434+
"opentelemetry.instrumentation.psycopg.psycopg.paramstyle",
435+
"test",
436+
):
437+
cnx = psycopg.connect(database="test")
438+
cnx = PsycopgInstrumentor().instrument_connection(
439+
cnx,
440+
enable_commenter=True,
441+
enable_attribute_commenter=True,
442+
)
443+
query = "Select 1"
444+
cursor = cnx.cursor()
445+
cursor.execute(query)
446+
spans_list = self.memory_exporter.get_finished_spans()
447+
span = spans_list[0]
448+
span_id = format(span.get_span_context().span_id, "016x")
449+
trace_id = format(span.get_span_context().trace_id, "032x")
450+
self.assertEqual(
451+
MockCursor.execute.call_args[0][0],
452+
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'*/",
453+
)
454+
self.assertEqual(
455+
span.attributes[SpanAttributes.DB_STATEMENT],
456+
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'*/",
457+
)
414458

415459
def test_sqlcommenter_enabled_instrument_connection_with_options(self):
416460
with mock.patch(
@@ -445,6 +489,10 @@ def test_sqlcommenter_enabled_instrument_connection_with_options(self):
445489
MockCursor.execute.call_args[0][0],
446490
f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
447491
)
492+
self.assertEqual(
493+
span.attributes[SpanAttributes.DB_STATEMENT],
494+
"Select 1",
495+
)
448496

449497
@mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
450498
def test_sqlcommenter_disabled(self, event_mocked):
@@ -468,6 +516,12 @@ def test_sqlcommenter_disabled_default_instrument_connection(self):
468516
MockCursor.execute.call_args[0][0],
469517
"Select 1",
470518
)
519+
spans_list = self.memory_exporter.get_finished_spans()
520+
span = spans_list[0]
521+
self.assertEqual(
522+
span.attributes[SpanAttributes.DB_STATEMENT],
523+
"Select 1",
524+
)
471525

472526
def test_sqlcommenter_disabled_explicit_instrument_connection(self):
473527
cnx = psycopg.connect(database="test")
@@ -482,6 +536,12 @@ def test_sqlcommenter_disabled_explicit_instrument_connection(self):
482536
MockCursor.execute.call_args[0][0],
483537
"Select 1",
484538
)
539+
spans_list = self.memory_exporter.get_finished_spans()
540+
span = spans_list[0]
541+
self.assertEqual(
542+
span.attributes[SpanAttributes.DB_STATEMENT],
543+
"Select 1",
544+
)
485545

486546

487547
class TestPostgresqlIntegrationAsync(

‎instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ def instrument_connection(
208208
tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
209209
enable_commenter: bool = False,
210210
commenter_options: dict = None,
211+
enable_attribute_commenter=None,
211212
):
212213
"""Enable instrumentation of a Psycopg2 connection.
213214
@@ -221,6 +222,8 @@ def instrument_connection(
221222
Optional flag to enable/disable sqlcommenter (default False).
222223
commenter_options: dict, optional
223224
Optional configurations for tags to be appended at the sql query.
225+
enable_attribute_commenter:
226+
Optional flag to enable/disable addition of sqlcomment to span attribute (default False). Requires enable_commenter=True.
224227
225228
Returns:
226229
An instrumented psycopg2 connection object.
@@ -236,6 +239,7 @@ def instrument_connection(
236239
tracer_provider=tracer_provider,
237240
enable_commenter=enable_commenter,
238241
commenter_options=commenter_options,
242+
enable_attribute_commenter=enable_attribute_commenter,
239243
)
240244
connection._is_instrumented_by_opentelemetry = True
241245
else:
@@ -304,6 +308,7 @@ def _new_cursor_factory(
304308
tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
305309
enable_commenter: bool = False,
306310
commenter_options: dict = None,
311+
enable_attribute_commenter: bool = False,
307312
):
308313
if not db_api:
309314
db_api = DatabaseApiIntegration(
@@ -315,6 +320,7 @@ def _new_cursor_factory(
315320
enable_commenter=enable_commenter,
316321
commenter_options=commenter_options,
317322
connect_module=psycopg2,
323+
enable_attribute_commenter=enable_attribute_commenter,
318324
)
319325

320326
base_factory = base_factory or pg_cursor

‎instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from opentelemetry import trace
2222
from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
2323
from opentelemetry.sdk import resources
24+
from opentelemetry.semconv.trace import SpanAttributes
2425
from opentelemetry.test.test_base import TestBase
2526

2627

@@ -294,6 +295,49 @@ def test_sqlcommenter_enabled_instrument_connection_defaults(self):
294295
MockCursor.execute.call_args[0][0],
295296
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'*/",
296297
)
298+
self.assertEqual(
299+
span.attributes[SpanAttributes.DB_STATEMENT],
300+
"Select 1",
301+
)
302+
303+
def test_sqlcommenter_enabled_instrument_connection_stmt_enabled(self):
304+
with mock.patch(
305+
"opentelemetry.instrumentation.psycopg2.psycopg2.__version__",
306+
"foobar",
307+
), mock.patch(
308+
"opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__",
309+
"foobaz",
310+
), mock.patch(
311+
"opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety",
312+
"123",
313+
), mock.patch(
314+
"opentelemetry.instrumentation.psycopg2.psycopg2.apilevel",
315+
"123",
316+
), mock.patch(
317+
"opentelemetry.instrumentation.psycopg2.psycopg2.paramstyle",
318+
"test",
319+
):
320+
cnx = psycopg2.connect(database="test")
321+
cnx = Psycopg2Instrumentor().instrument_connection(
322+
cnx,
323+
enable_commenter=True,
324+
enable_attribute_commenter=True,
325+
)
326+
query = "Select 1"
327+
cursor = cnx.cursor()
328+
cursor.execute(query)
329+
spans_list = self.memory_exporter.get_finished_spans()
330+
span = spans_list[0]
331+
span_id = format(span.get_span_context().span_id, "016x")
332+
trace_id = format(span.get_span_context().trace_id, "032x")
333+
self.assertEqual(
334+
MockCursor.execute.call_args[0][0],
335+
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'*/",
336+
)
337+
self.assertEqual(
338+
span.attributes[SpanAttributes.DB_STATEMENT],
339+
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'*/",
340+
)
297341

298342
def test_sqlcommenter_enabled_instrument_connection_with_options(self):
299343
with mock.patch(
@@ -328,6 +372,10 @@ def test_sqlcommenter_enabled_instrument_connection_with_options(self):
328372
MockCursor.execute.call_args[0][0],
329373
f"Select 1 /*db_driver='psycopg2%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
330374
)
375+
self.assertEqual(
376+
span.attributes[SpanAttributes.DB_STATEMENT],
377+
"Select 1",
378+
)
331379

332380
@mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
333381
def test_sqlcommenter_disabled(self, event_mocked):
@@ -351,6 +399,12 @@ def test_sqlcommenter_disabled_default_instrument_connection(self):
351399
MockCursor.execute.call_args[0][0],
352400
"Select 1",
353401
)
402+
spans_list = self.memory_exporter.get_finished_spans()
403+
span = spans_list[0]
404+
self.assertEqual(
405+
span.attributes[SpanAttributes.DB_STATEMENT],
406+
"Select 1",
407+
)
354408

355409
def test_sqlcommenter_disabled_explicit_instrument_connection(self):
356410
cnx = psycopg2.connect(database="test")
@@ -365,6 +419,12 @@ def test_sqlcommenter_disabled_explicit_instrument_connection(self):
365419
MockCursor.execute.call_args[0][0],
366420
"Select 1",
367421
)
422+
spans_list = self.memory_exporter.get_finished_spans()
423+
span = spans_list[0]
424+
self.assertEqual(
425+
span.attributes[SpanAttributes.DB_STATEMENT],
426+
"Select 1",
427+
)
368428

369429
def test_no_op_tracer_provider(self):
370430
Psycopg2Instrumentor().instrument(

0 commit comments

Comments
 (0)