postgres: DBM full service fallback w/ prepared statements

Author: tlhunter

packages/datadog-plugin-pg/src/index.js

@@ -27,7 +27,7 @@ class PGPlugin extends DatabasePlugin {
       }
     })
 
-    query.text = this.injectDbmQuery(query.text, service)
+    query.text = this.injectDbmQuery(query.text, service, !!query.name)
   }
 }
 

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

@@ -271,6 +271,7 @@ describe('Plugin', () => {
           })
         })
       })
+
       describe('with DBM propagation enabled with service using plugin configurations', () => {
         before(() => {
           return agent.load('pg', [{ dbmPropagationMode: 'service', service: () => 'serviced' }])
@@ -318,6 +319,7 @@ describe('Plugin', () => {
             }
           }
         })
+
         it('trace query resource should not be changed when propagation is enabled', done => {
           agent.use(traces => {
             expect(traces[0][0]).to.have.property('resource', 'SELECT $1::text as message')
@@ -331,15 +333,18 @@ describe('Plugin', () => {
           })
         })
       })
+
       describe('DBM propagation should handle special characters', () => {
         let clientDBM
+
         before(() => {
           return agent.load('pg', [{ dbmPropagationMode: 'service', service: '~!@#$%^&*()_+|??/<>' }])
         })
 
         after(() => {
           return agent.close({ ritmReset: false })
         })
+
         beforeEach(done => {
           pg = require(`../../../versions/pg@${version}`).get()
 
@@ -352,6 +357,7 @@ describe('Plugin', () => {
 
           clientDBM.connect(err => done(err))
         })
+
         it('DBM propagation should handle special characters', done => {
           clientDBM.query('SELECT $1::text as message', ['Hello world!'], (err, result) => {
             if (err) return done(err)
@@ -373,15 +379,18 @@ 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')
         })
+
         beforeEach(done => {
           pg = require(`../../../versions/pg@${version}`).get()
 
@@ -409,9 +418,11 @@ describe('Plugin', () => {
             return originalWrite.apply(this, arguments)
           }
         })
+
         afterEach(() => {
           net.Socket.prototype.write = originalWrite
         })
+
         it('query text should contain traceparent', done => {
           agent.use(traces => {
             const traceId = traces[0][0].trace_id.toString(16).padStart(32, '0')
@@ -428,6 +439,7 @@ describe('Plugin', () => {
             })
           })
         })
+
         it('query should inject _dd.dbm_trace_injected into span', done => {
           agent.use(traces => {
             expect(traces[0][0].meta).to.have.property('_dd.dbm_trace_injected', 'true')
@@ -442,6 +454,7 @@ describe('Plugin', () => {
             })
           })
         })
+
         it('service should default to tracer service name', done => {
           tracer
           agent.use(traces => {
@@ -458,12 +471,14 @@ describe('Plugin', () => {
           })
         })
       })
+
       describe('DBM propagation enabled with full should handle query config objects', () => {
         const tracer = require('../../dd-trace')
 
         before(() => {
           return agent.load('pg')
         })
+
         beforeEach(done => {
           pg = require(`../../../versions/pg@${version}`).get()
 
@@ -484,10 +499,12 @@ describe('Plugin', () => {
 
         it('query config objects should be handled', done => {
           let queryText = ''
+
           const query = {
             name: 'pgSelectQuery',
             text: 'SELECT $1::text as message'
           }
+
           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')
@@ -496,6 +513,7 @@ describe('Plugin', () => {
               `/*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(query, ['Hello world!'], (err) => {
             if (err) return done(err)
 
@@ -505,23 +523,51 @@ describe('Plugin', () => {
           })
           queryText = client.queryQueue[0].text
         })
+
         it('query config object should persist when comment is injected', done => {
           const query = {
             name: 'pgSelectQuery',
             text: 'SELECT $1::text as message'
           }
+
           client.query(query, ['Hello world!'], (err) => {
             if (err) return done(err)
 
             client.end((err) => {
               if (err) return done(err)
             })
           })
+
           agent.use(traces => {
             expect(query).to.have.property(
               'name', 'pgSelectQuery')
           }).then(done, done)
         })
+
+        it('falls back to service with prepared statements', done => {
+          let queryText = ''
+
+          const query = {
+            name: 'pgSelectQuery',
+            text: 'SELECT $1::text as message',
+            name: 'getText'
+          }
+
+          agent.use(traces => {
+            expect(queryText).to.equal(
+              `/*dddbs='post',dde='tester',ddps='test',ddpv='8.4.0'` +
+              `*/ SELECT $1::text as message`)
+          }).then(done, done)
+
+          client.query(query, ['Hello world!'], (err) => {
+            if (err) return done(err)
+
+            client.end((err) => {
+              if (err) return done(err)
+            })
+          })
+          queryText = client.queryQueue[0].text
+        })
       })
     })
   })

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

@@ -37,12 +37,12 @@ class DatabasePlugin extends StoragePlugin {
     `ddps='${encodedDdps}',ddpv='${encodedDdpv}'`
   }
 
-  injectDbmQuery (query, serviceName) {
+  injectDbmQuery (query, serviceName, isPreparedStatement = false) {
     if (this.config.dbmPropagationMode === 'disabled') {
       return query
     }
     const servicePropagation = this.createDBMPropagationCommentService(serviceName)
-    if (this.config.dbmPropagationMode === 'service') {
+    if (isPreparedStatement || this.config.dbmPropagationMode === 'service') {
       return `/*${servicePropagation}*/ ${query}`
     } else if (this.config.dbmPropagationMode === 'full') {
       this.activeSpan.setTag('_dd.dbm_trace_injected', 'true')