@@ -32,6 +32,7 @@ import {
3232} from '@opentelemetry/sdk-trace-base' ;
3333import * as assert from 'assert' ;
3434import type * as pg from 'pg' ;
35+ import * as sinon from 'sinon' ;
3536import {
3637 PgInstrumentation ,
3738 PgInstrumentationConfig ,
@@ -44,6 +45,7 @@ import {
4445 DbSystemValues ,
4546} from '@opentelemetry/semantic-conventions' ;
4647import { isSupported } from './utils' ;
48+ import { addSqlCommenterComment } from '../src/utils' ;
4749
4850const pgVersion = require ( 'pg/package.json' ) . version ;
4951const nodeVersion = process . versions . node ;
@@ -110,6 +112,12 @@ describe('pg', () => {
110112 const testPostgresLocally = process . env . RUN_POSTGRES_TESTS_LOCAL ; // For local: spins up local postgres db via docker
111113 const shouldTest = testPostgres || testPostgresLocally ; // Skips these tests if false (default)
112114
115+ function getExecutedQueries ( ) {
116+ return ( client as any ) . queryQueue . push . args . flat ( ) as ( pg . Query & {
117+ text ?: string ;
118+ } ) [ ] ;
119+ }
120+
113121 before ( async function ( ) {
114122 const skipForUnsupported =
115123 process . env . IN_TAV && ! isSupported ( nodeVersion , pgVersion ) ;
@@ -156,11 +164,16 @@ describe('pg', () => {
156164 beforeEach ( ( ) => {
157165 contextManager = new AsyncHooksContextManager ( ) . enable ( ) ;
158166 context . setGlobalContextManager ( contextManager ) ;
167+
168+ // Add a spy on the underlying client's internal query queue so that
169+ // we could assert on what the final queries are that are executed
170+ sinon . spy ( ( client as any ) . queryQueue , 'push' ) ;
159171 } ) ;
160172
161173 afterEach ( ( ) => {
162174 memoryExporter . reset ( ) ;
163175 context . disable ( ) ;
176+ sinon . restore ( ) ;
164177 } ) ;
165178
166179 it ( 'should return an instrumentation' , ( ) => {
@@ -649,6 +662,116 @@ describe('pg', () => {
649662 } ) ;
650663 } ) ;
651664
665+ it ( 'should not add sqlcommenter comment when flag is not specified' , async ( ) => {
666+ const span = tracer . startSpan ( 'test span' ) ;
667+ await context . with ( trace . setSpan ( context . active ( ) , span ) , async ( ) => {
668+ try {
669+ const query = 'SELECT NOW()' ;
670+ const resPromise = await client . query ( query ) ;
671+ assert . ok ( resPromise ) ;
672+
673+ const [ span ] = memoryExporter . getFinishedSpans ( ) ;
674+ assert . ok ( span ) ;
675+
676+ const commentedQuery = addSqlCommenterComment (
677+ trace . wrapSpanContext ( span . spanContext ( ) ) ,
678+ query
679+ ) ;
680+
681+ const executedQueries = getExecutedQueries ( ) ;
682+ assert . equal ( executedQueries . length , 1 ) ;
683+ assert . equal ( executedQueries [ 0 ] . text , query ) ;
684+ assert . notEqual ( query , commentedQuery ) ;
685+ } catch ( e ) {
686+ assert . ok ( false , e . message ) ;
687+ }
688+ } ) ;
689+ } ) ;
690+
691+ it ( 'should not add sqlcommenter comment with client.query({text, callback}) when flag is not specified' , done => {
692+ const span = tracer . startSpan ( 'test span' ) ;
693+ context . with ( trace . setSpan ( context . active ( ) , span ) , ( ) => {
694+ const query = 'SELECT NOW()' ;
695+ client . query ( {
696+ text : query ,
697+ callback : ( err : Error , res : pg . QueryResult ) => {
698+ assert . strictEqual ( err , null ) ;
699+ assert . ok ( res ) ;
700+
701+ const [ span ] = memoryExporter . getFinishedSpans ( ) ;
702+ const commentedQuery = addSqlCommenterComment (
703+ trace . wrapSpanContext ( span . spanContext ( ) ) ,
704+ query
705+ ) ;
706+
707+ const executedQueries = getExecutedQueries ( ) ;
708+ assert . equal ( executedQueries . length , 1 ) ;
709+ assert . equal ( executedQueries [ 0 ] . text , query ) ;
710+ assert . notEqual ( query , commentedQuery ) ;
711+ done ( ) ;
712+ } ,
713+ } as pg . QueryConfig ) ;
714+ } ) ;
715+ } ) ;
716+
717+ it ( 'should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified' , async ( ) => {
718+ instrumentation . setConfig ( {
719+ addSqlCommenterCommentToQueries : true ,
720+ } ) ;
721+
722+ const span = tracer . startSpan ( 'test span' ) ;
723+ await context . with ( trace . setSpan ( context . active ( ) , span ) , async ( ) => {
724+ try {
725+ const query = 'SELECT NOW()' ;
726+ const resPromise = await client . query ( query ) ;
727+ assert . ok ( resPromise ) ;
728+
729+ const [ span ] = memoryExporter . getFinishedSpans ( ) ;
730+ const commentedQuery = addSqlCommenterComment (
731+ trace . wrapSpanContext ( span . spanContext ( ) ) ,
732+ query
733+ ) ;
734+
735+ const executedQueries = getExecutedQueries ( ) ;
736+ assert . equal ( executedQueries . length , 1 ) ;
737+ assert . equal ( executedQueries [ 0 ] . text , commentedQuery ) ;
738+ assert . notEqual ( query , commentedQuery ) ;
739+ } catch ( e ) {
740+ assert . ok ( false , e . message ) ;
741+ }
742+ } ) ;
743+ } ) ;
744+
745+ it ( 'should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified with client.query({text, callback})' , done => {
746+ instrumentation . setConfig ( {
747+ addSqlCommenterCommentToQueries : true ,
748+ } ) ;
749+
750+ const span = tracer . startSpan ( 'test span' ) ;
751+ context . with ( trace . setSpan ( context . active ( ) , span ) , ( ) => {
752+ const query = 'SELECT NOW()' ;
753+ client . query ( {
754+ text : query ,
755+ callback : ( err : Error , res : pg . QueryResult ) => {
756+ assert . strictEqual ( err , null ) ;
757+ assert . ok ( res ) ;
758+
759+ const [ span ] = memoryExporter . getFinishedSpans ( ) ;
760+ const commentedQuery = addSqlCommenterComment (
761+ trace . wrapSpanContext ( span . spanContext ( ) ) ,
762+ query
763+ ) ;
764+
765+ const executedQueries = getExecutedQueries ( ) ;
766+ assert . equal ( executedQueries . length , 1 ) ;
767+ assert . equal ( executedQueries [ 0 ] . text , commentedQuery ) ;
768+ assert . notEqual ( query , commentedQuery ) ;
769+ done ( ) ;
770+ } ,
771+ } as pg . QueryConfig ) ;
772+ } ) ;
773+ } ) ;
774+
652775 it ( 'should not generate traces for client.query() when requireParentSpan=true is specified' , done => {
653776 instrumentation . setConfig ( {
654777 requireParentSpan : true ,
0 commit comments