dbm: tedious (sql server) service mode (#5375)

Author: tlhunter

packages/datadog-instrumentations/src/tedious.js

@@ -16,7 +16,7 @@ addHook({ name: 'tedious', versions: ['>=1.0.0'] }, tedious => {
       return makeRequest.apply(this, arguments)
     }
 
-    const queryOrProcedure = getQueryOrProcedure(request)
+    const [queryOrProcedure, queryParent, queryField] = getQueryOrProcedure(request)
 
     if (!queryOrProcedure) {
       return makeRequest.apply(this, arguments)
@@ -28,7 +28,9 @@ addHook({ name: 'tedious', versions: ['>=1.0.0'] }, tedious => {
     const connectionConfig = this.config
 
     return asyncResource.runInAsyncScope(() => {
-      startCh.publish({ queryOrProcedure, connectionConfig })
+      const payload = { queryOrProcedure, connectionConfig }
+      startCh.publish(payload)
+      queryParent[queryField] = payload.sql
 
       const cb = callbackResource.bind(request.callback, request)
       request.callback = asyncResource.bind(function (error) {
@@ -53,14 +55,15 @@ addHook({ name: 'tedious', versions: ['>=1.0.0'] }, tedious => {
   return tedious
 })
 
+// returns [queryOrProcedure, parentObjectToSet, propertyNameToSet]
 function getQueryOrProcedure (request) {
-  if (!request.parameters) return
-
-  const statement = request.parametersByName.statement || request.parametersByName.stmt
-
-  if (!statement) {
-    return request.sqlTextOrProcedure
+  if (!request.parameters) return [null]
+
+  if (request.parametersByName.statement) {
+    return [request.parametersByName.statement.value, request.parametersByName.statement, 'value']
+  } else if (request.parametersByName.stmt) {
+    return [request.parametersByName.stmt.value, request.parametersByName.stmt, 'value']
+  } else {
+    return [request.sqlTextOrProcedure, request, 'sqlTextOrProcedure']
   }
-
-  return statement.value
 }

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

@@ -8,22 +8,27 @@ class TediousPlugin extends DatabasePlugin {
   static get operation () { return 'request' } // TODO: change to match other database plugins
   static get system () { return 'mssql' }
 
-  start ({ queryOrProcedure, connectionConfig }) {
-    this.startSpan(this.operationName(), {
-      service: this.serviceName({ pluginConfig: this.config, system: this.system }),
-      resource: queryOrProcedure,
+  start (payload) {
+    const service = this.serviceName({ pluginConfig: this.config, system: this.system })
+    const span = this.startSpan(this.operationName(), {
+      service,
+      resource: payload.queryOrProcedure,
       type: 'sql',
       kind: 'client',
       meta: {
         'db.type': 'mssql',
         component: 'tedious',
-        'out.host': connectionConfig.server,
-        [CLIENT_PORT_KEY]: connectionConfig.options.port,
-        'db.user': connectionConfig.userName || connectionConfig.authentication.options.userName,
-        'db.name': connectionConfig.options.database,
-        'db.instance': connectionConfig.options.instanceName
+        'out.host': payload.connectionConfig.server,
+        [CLIENT_PORT_KEY]: payload.connectionConfig.options.port,
+        'db.user': payload.connectionConfig.userName || payload.connectionConfig.authentication.options.userName,
+        'db.name': payload.connectionConfig.options.database,
+        'db.instance': payload.connectionConfig.options.instanceName
       }
     })
+
+    // SQL Server includes comments when caching queries
+    // For that reason we allow service mode but not full mode
+    payload.sql = this.injectDbmQuery(span, payload.queryOrProcedure, service, true)
   }
 }
 

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

@@ -12,7 +12,6 @@ describe('Plugin', () => {
   let tds
   let tracer
   let connection
-  let connectionIsClosed
 
   withVersions('tedious', 'tedious', version => {
     beforeEach(() => {
@@ -53,7 +52,6 @@ describe('Plugin', () => {
           config.password = MSSQL_PASSWORD
         }
 
-        connectionIsClosed = false
         connection = new tds.Connection(config)
           .on('connect', done)
 
@@ -64,12 +62,8 @@ describe('Plugin', () => {
       })
 
       afterEach(function (done) {
-        if (connectionIsClosed) {
-          done()
-        } else {
-          connection.on('end', () => done())
-          connection.close()
-        }
+        connection.on('end', () => done())
+        connection.close()
       })
 
       withNamingSchema(
@@ -405,5 +399,68 @@ describe('Plugin', () => {
         })
       }
     })
+
+    // it's a pretty old version with a different enough API that I don't think it's worth supporting
+    const testDbm = semver.intersects(version, '<10') ? describe.skip : describe
+    testDbm('with configuration and DBM enabled', () => {
+      let config
+      let tds
+      let connection
+
+      beforeEach(() => {
+        return agent.load('tedious', { dbmPropagationMode: 'service', service: 'custom' }).then(() => {
+          tds = require(`../../../versions/tedious@${version}`).get()
+        })
+      })
+
+      afterEach(() => {
+        return agent.close({ ritmReset: false })
+      })
+
+      beforeEach((done) => {
+        config = {
+          server: 'localhost',
+          options: {
+            database: 'master',
+            trustServerCertificate: true
+          },
+          authentication: {
+            options: {
+              userName: MSSQL_USERNAME,
+              password: MSSQL_PASSWORD
+            },
+            type: 'default'
+          }
+        }
+
+        connection = new tds.Connection(config)
+          .on('connect', done)
+
+        connection.connect()
+      })
+
+      afterEach(function (done) {
+        connection.on('end', () => done())
+        connection.close()
+      })
+
+      it('should inject the correct DBM comment into query but not into trace', done => {
+        const query = 'SELECT 1 + 1 AS solution'
+
+        const request = new tds.Request(query, (err) => {
+          if (err) return done(err)
+          promise.then(done, done)
+        })
+
+        const promise = agent
+          .use(traces => {
+            expect(traces[0][0]).to.have.property('resource', 'SELECT 1 + 1 AS solution')
+            expect(request.sqlTextOrProcedure).to.equal("/*dddb='master',dddbs='custom',dde='tester'," +
+              "ddh='localhost',ddps='test',ddpv='10.8.2'*/ SELECT 1 + 1 AS solution")
+          })
+
+        connection.execSql(request)
+      })
+    })
   })
 })

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

@@ -63,7 +63,7 @@ class DatabasePlugin extends StoragePlugin {
     return tracerService
   }
 
-  createDbmComment (span, serviceName, isPreparedStatement = false) {
+  createDbmComment (span, serviceName, disableFullMode = false) {
     const mode = this.config.dbmPropagationMode
     const dbmService = this.getDbmServiceName(span, serviceName)
 
@@ -73,7 +73,7 @@ class DatabasePlugin extends StoragePlugin {
 
     const servicePropagation = this.createDBMPropagationCommentService(dbmService, span)
 
-    if (isPreparedStatement || mode === 'service') {
+    if (disableFullMode || mode === 'service') {
       return servicePropagation
     } else if (mode === 'full') {
       span.setTag('_dd.dbm_trace_injected', 'true')
@@ -82,8 +82,8 @@ class DatabasePlugin extends StoragePlugin {
     }
   }
 
-  injectDbmQuery (span, query, serviceName, isPreparedStatement = false) {
-    const dbmTraceComment = this.createDbmComment(span, serviceName, isPreparedStatement)
+  injectDbmQuery (span, query, serviceName, disableFullMode = false) {
+    const dbmTraceComment = this.createDbmComment(span, serviceName, disableFullMode)
 
     if (!dbmTraceComment) {
       return query