Fix to DBM propagation - comment is injected before query is sent (#2938)

* adding query to async resource to ensure we add the comment before query.apply is called

* changing full propagation test to check for comment before call to db

Author: crysmags

packages/datadog-instrumentations/src/pg.js

@@ -27,34 +27,38 @@ function wrapQuery (query) {
       return query.apply(this, arguments)
     }
 
-    const retval = query.apply(this, arguments)
-
-    const queryQueue = this.queryQueue || this._queryQueue
-    const activeQuery = this.activeQuery || this._activeQuery
-    const pgQuery = queryQueue[queryQueue.length - 1] || activeQuery
-
-    if (!pgQuery) {
-      return retval
-    }
-
     const callbackResource = new AsyncResource('bound-anonymous-fn')
     const asyncResource = new AsyncResource('bound-anonymous-fn')
     const processId = this.processID
+    let pgQuery = {}
+    pgQuery.text = arguments[0]
+
     return asyncResource.runInAsyncScope(() => {
       startCh.publish({
         params: this.connectionParameters,
-        originalQuery: pgQuery.text,
         query: pgQuery,
         processId
       })
 
+      arguments[0] = pgQuery.text
+
       const finish = asyncResource.bind(function (error) {
         if (error) {
           errorCh.publish(error)
         }
         finishCh.publish()
       })
 
+      const retval = query.apply(this, arguments)
+      const queryQueue = this.queryQueue || this._queryQueue
+      const activeQuery = this.activeQuery || this._activeQuery
+
+      pgQuery = queryQueue[queryQueue.length - 1] || activeQuery
+
+      if (!pgQuery) {
+        return retval
+      }
+
       if (pgQuery.callback) {
         const originalCallback = callbackResource.bind(pgQuery.callback)
         pgQuery.callback = function (err, res) {

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

@@ -4,6 +4,7 @@ const { expect } = require('chai')
 const semver = require('semver')
 const agent = require('../../dd-trace/test/plugins/agent')
 const { ERROR_MESSAGE, ERROR_TYPE, ERROR_STACK } = require('../../dd-trace/src/constants')
+const net = require('net')
 
 const clients = {
   pg: pg => pg.Client
@@ -374,6 +375,10 @@ describe('Plugin', () => {
       })
       describe('with DBM propagation enabled with full using tracer configurations', () => {
         const tracer = require('../../dd-trace')
+        let seenTraceParent
+        let seenTraceId
+        let seenSpanId
+        let originalWrite
         before(() => {
           return agent.load('pg')
         })
@@ -393,27 +398,36 @@ describe('Plugin', () => {
             database: 'postgres'
           })
           client.connect(err => done(err))
+          originalWrite = net.Socket.prototype.write
+          net.Socket.prototype.write = function (buffer) {
+            let strBuf = buffer.toString()
+            if (strBuf.includes('traceparent=\'')) {
+              strBuf = strBuf.split('-')
+              seenTraceParent = true
+              seenTraceId = strBuf[1]
+              seenSpanId = strBuf[2]
+            }
+            return originalWrite.apply(this, arguments)
+          }
+        })
+        afterEach(() => {
+          net.Socket.prototype.write = originalWrite
         })
-
         it('query text should contain traceparent', done => {
-          let queryText = ''
           agent.use(traces => {
             const traceId = traces[0][0].trace_id.toString(16).padStart(32, '0')
             const spanId = traces[0][0].span_id.toString(16).padStart(16, '0')
-
-            expect(queryText).to.equal(
-              `/*dddbs='post',dde='tester',ddps='test',ddpv='8.4.0',` +
-              `traceparent='00-${traceId}-${spanId}-00'*/ SELECT $1::text as message`)
+            expect(seenTraceId).to.equal(traceId)
+            expect(seenSpanId).to.equal(spanId)
           }).then(done, done)
 
           client.query('SELECT $1::text as message', ['Hello World!'], (err, result) => {
             if (err) return done(err)
-
+            expect(seenTraceParent).to.be.true
             client.end((err) => {
               if (err) return done(err)
             })
           })
-          queryText = client.queryQueue[0].text
         })
         it('query should inject _dd.dbm_trace_injected into span', done => {
           agent.use(traces => {