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

This reverts commit 840f390.

Author: simon-id

packages/datadog-instrumentations/src/pg.js

@@ -27,38 +27,34 @@ 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,7 +4,6 @@ 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
@@ -375,10 +374,6 @@ 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')
         })
@@ -398,36 +393,27 @@ 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(seenTraceId).to.equal(traceId)
-            expect(seenSpanId).to.equal(spanId)
+
+            expect(queryText).to.equal(
+              `/*dddbs='post',dde='tester',ddps='test',ddpv='8.4.0',` +
+              `traceparent='00-${traceId}-${spanId}-00'*/ SELECT $1::text as message`)
           }).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 => {