DBM: Add option to SQL comment propagation

docs/test.ts

@@ -370,6 +370,7 @@ tracer.use('paperplane', httpServerOptions);
 tracer.use('playwright');
 tracer.use('pg');
 tracer.use('pg', { service: params => `${params.host}-${params.database}` });
+tracer.use('pg', { appendComment: true });
 tracer.use('pino');
 tracer.use('protobufjs');
 tracer.use('redis');

index.d.ts

@@ -1768,6 +1768,10 @@ declare namespace tracer {
        * The database monitoring propagation mode to be used for this plugin.
        */
       dbmPropagationMode?: string;
+      /**
+       * Appends the SQL comment propagation to the query string. Prepends the comment if `false`. For long query strings, the appended propagation comment might be truncated, causing loss of correlation between the query and trace.
+       */
+      appendComment?: boolean;
     }
 
     /**

packages/datadog-plugin-pg/test/index.spec.js

@@ -691,6 +691,56 @@ describe('Plugin', () => {
           queryText = client.queryQueue[0].text
         })
       })
+
+      describe.skip('with DBM propagation enabled with append comment configurations', () => {
+        before(() => {
+          return agent.load('pg', [{ dbmPropagationMode: 'service', service: () => 'serviced', appendComment: true }])
+        })
+
+        after(() => {
+          return agent.close({ ritmReset: false })
+        })
+
+        beforeEach(done => {
+          pg = require(`../../../versions/pg@${version}`).get()
+
+          client = new pg.Client({
+            host: '127.0.0.1',
+            user: 'postgres',
+            password: 'postgres',
+            database: 'postgres'
+          })
+          client.connect(err => done(err))
+        })
+
+        it('should append comment in query text', done => {
+          const client = new pg.Client({
+            host: '127.0.0.1',
+            user: 'postgres',
+            password: 'postgres',
+            database: 'postgres'
+          })
+
+          client.connect(err => done(err))
+
+          client.query('SELECT $1::text as message', ['Hello world!'], (err, result) => {
+            if (err) return done(err)
+
+            client.end((err) => {
+              if (err) return done(err)
+            })
+          })
+          if (client.queryQueue[0] !== undefined) {
+            try {
+              expect(client.queryQueue[0].text).to.equal(
+                'SELECT $1::text as message /*dddb=\'postgres\',dddbs=\'serviced\',dde=\'tester\',' +
+                `ddh='127.0.0.1',ddps='test',ddpv='${ddpv}'*/ `)
+            } catch (e) {
+              done(e)
+            }
+          }
+        })
+      })
     })
   })
 })

packages/dd-trace/src/plugins/database.js

@@ -63,6 +63,12 @@ class DatabasePlugin extends StoragePlugin {
     return tracerService
   }
 
+  addComment (query, comment) {
+    return this.config.appendComment
+      ? `${query} /*${comment}*/`
+      : `/*${comment}*/ ${query}`
+  }
+
   injectDbmQuery (span, query, serviceName, isPreparedStatement = false) {
     const mode = this.config.dbmPropagationMode
     const dbmService = this.getDbmServiceName(span, serviceName)
@@ -74,11 +80,11 @@ class DatabasePlugin extends StoragePlugin {
     const servicePropagation = this.createDBMPropagationCommentService(dbmService, span)
 
     if (isPreparedStatement || mode === 'service') {
-      return `/*${servicePropagation}*/ ${query}`
+      return this.addComment(query, servicePropagation)
     } else if (mode === 'full') {
       span.setTag('_dd.dbm_trace_injected', 'true')
       const traceparent = span._spanContext.toTraceparent()
-      return `/*${servicePropagation},traceparent='${traceparent}'*/ ${query}`
+      return this.addComment(query, `${servicePropagation},traceparent='${traceparent}'`)
     }
   }