@@ -261,6 +261,39 @@ def test_sqlcommenter_enabled(self, event_mocked):
261261 kwargs = event_mocked .call_args [1 ]
262262 self .assertEqual (kwargs ["enable_commenter" ], True )
263263
264+ def test_sqlcommenter_enabled_instrument_connection (self ):
265+ with mock .patch (
266+ "opentelemetry.instrumentation.psycopg2.psycopg2.__version__" ,
267+ "foobar" ,
268+ ), mock .patch (
269+ "opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__" ,
270+ "foobaz" ,
271+ ), mock .patch (
272+ "opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety" ,
273+ "123" ,
274+ ), mock .patch (
275+ "opentelemetry.instrumentation.psycopg2.psycopg2.apilevel" ,
276+ "123" ,
277+ ), mock .patch (
278+ "opentelemetry.instrumentation.psycopg2.psycopg2.paramstyle" ,
279+ "test" ,
280+ ):
281+ cnx = psycopg2 .connect (database = "test" )
282+ cnx = Psycopg2Instrumentor ().instrument_connection (
283+ cnx , enable_commenter = True , commenter_options = {"foo" : "bar" }
284+ )
285+ query = "Select 1"
286+ cursor = cnx .cursor ()
287+ cursor .execute (query )
288+ spans_list = self .memory_exporter .get_finished_spans ()
289+ span = spans_list [0 ]
290+ span_id = format (span .get_span_context ().span_id , "016x" )
291+ trace_id = format (span .get_span_context ().trace_id , "032x" )
292+ self .assertEqual (
293+ MockCursor .execute .call_args [0 ][0 ],
294+ 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'*/" ,
295+ )
296+
264297 @mock .patch ("opentelemetry.instrumentation.dbapi.wrap_connect" )
265298 def test_sqlcommenter_disabled (self , event_mocked ):
266299 cnx = psycopg2 .connect (database = "test" )
0 commit comments