fix(mysql): instrument pool.getConnection properly

Author: Shadowbeetle

lib/instrumentations/trace-instrumentation-mysql.js

@@ -23,39 +23,53 @@ module.exports = function (mysql, agent, pkg) {
       var callback = args[last]
 
       if (name === 'query') {
-        var config = moduleName === 'pool'
-          ? this.config.connectionConfig
-          : this.config
+        wrapQuery.call(this, moduleName, original, args, callback)
+      } else if (name === 'getConnection') {
+        wrapGetConnection.call(this, original, callback)
+      } else {
+        return original.apply(this, args)
+      }
+    }
+  }
 
-        var user = config.user || 'unknown'
-        var host = config.host || 'unknown'
-        var port = config.port ? ':' + config.port : ''
-        var database = config.database ? '/' + config.database : ''
+  function wrapQuery (moduleName, original, args, callback) {
+    var config = moduleName === 'pool'
+      ? this.config.connectionConfig
+      : this.config
 
-        var connectionURI = 'mysql://' +
-          user + '@' +
-          host +
-          port +
-          database
+    var user = config.user || 'unknown'
+    var host = config.host || 'unknown'
+    var port = config.port ? ':' + config.port : ''
+    var database = config.database ? '/' + config.database : ''
 
-        var hostString = host +
-          port +
-          database
+    var connectionURI = 'mysql://' +
+      user + '@' +
+      host +
+      port +
+      database
 
-        return utils.wrapQuery.call(this,
-          original,
-          args,
-          agent, {
-            continuationMethod: typeof callback === 'function' ? 'callback' : 'eventEmitter',
-            protocol: consts.PROTOCOLS.MYSQL,
-            url: connectionURI,
-            host: hostString,
-            method: utils.tryParseSql(args[0]) || 'unknown'
-          })
-      } else {
-        return original.apply(this, args)
-      }
+    var hostString = host +
+      port +
+      database
+
+    return utils.wrapQuery.call(this,
+      original,
+      args,
+      agent, {
+        continuationMethod: typeof callback === 'function' ? 'callback' : 'eventEmitter',
+        protocol: consts.PROTOCOLS.MYSQL,
+        url: connectionURI,
+        host: hostString,
+        method: utils.tryParseSql(args[0]) || 'unknown'
+      })
+  }
+
+  function wrapGetConnection (original, callback) {
+    var wrappedCallback = function (err, connection) {
+      Shimmer.wrap(connection, CONNECTION_OPERATIONS, wrapOperations.bind(this, 'connection'))
+      return callback.call(this, err, connection)
     }
+    return original.call(this, wrappedCallback)
   }
 
   Shimmer.wrap(mysql, 'createConnection', function (original) {

test/instrumentations/mysql/mysql.spec.js

@@ -5,6 +5,7 @@ var wrapper = require('../../../lib/instrumentations/trace-instrumentation-mysql
 var expect = require('chai').expect
 var Shimmer = require('../../../lib/utils/shimmer')
 var utils = require('../../../lib/instrumentations/utils')
+var sinon = require('sinon')
 
 describe('The mysql wrapper', function () {
   var CONNECTION_OPERATIONS = [
@@ -19,6 +20,44 @@ describe('The mysql wrapper', function () {
     'destroy'
   ]
 
+  var sandbox = sinon.sandbox.create()
+  var fakeAgent = {
+    incomingEdgeMetrics: {
+      report: sandbox.spy()
+    },
+    tracer: {
+      collector: {
+        mustCollectSeverity: 9,
+        defaultSeverity: 0,
+        clientRecv: sandbox.spy(),
+        clientSend: sandbox.stub().returns({
+          event: {
+            p: 'dfasdfs'
+          },
+          duffelBag: {
+            timestamp: 0
+          }
+        }),
+        systemError: sandbox.spy()
+      },
+      send: sandbox.spy()
+    },
+    storage: {
+      get: sandbox.spy()
+    },
+    externalEdgeMetrics: {
+      report: sandbox.spy(),
+      EDGE_STATUS: {
+        OK: 1,
+        NOT_OK: 0
+      }
+    }
+  }
+
+  afterEach(function () {
+    sandbox.restore()
+  })
+
   describe('Connection', function () {
     it('should wrap var connection = mysql.createConnection and var connection = mysql.createPool', function () {
       var shimmerWrapStub = this.sandbox.stub(Shimmer, 'wrap')
@@ -44,7 +83,6 @@ describe('The mysql wrapper', function () {
     it('should use wrapQuery with expected arguments to wrap connection.query', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')
 
@@ -87,7 +125,6 @@ describe('The mysql wrapper', function () {
     it('should use wrapQuery with expected arguments to wrap pool.query', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')
 
@@ -124,10 +161,39 @@ describe('The mysql wrapper', function () {
       expect(shimmerWrapStub).to.have.been.called
     })
 
+    it('should wrap connection returned by pool.getConnection properly', function (done) {
+      require.cache = {}
+      var mysql = require('mysql')
+
+      wrapper(mysql, fakeAgent)
+
+      // This is tricky. We have to stub exactly after wrapping, and before
+      // createConnection to catch the wrapping of the query operation
+
+      var pool = mysql.createPool({
+        host: process.env.MYSQL_HOST || 'localhost',
+        user: 'root',
+        password: '',
+        database: 'information_schema'
+      })
+
+      pool.getConnection(function (err, conn) {
+        if (err) throw err
+        conn.query('SELECT 1 + 1 AS solution', function (error, results, fields) {
+          if (error) throw error
+          expect(fakeAgent.tracer.collector.clientSend).to.have.been.calledOnce
+          expect(fakeAgent.tracer.collector.clientRecv).to.have.been.calledWith({
+            protocol: 'mysql',
+            status: 'ok'
+          })
+          done()
+        })
+      })
+    })
+
     it('should not leak credentials', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')
 
@@ -171,7 +237,6 @@ describe('The mysql wrapper', function () {
     it('should use wrapQuery with expected arguments to wrap connection.query', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')
 
@@ -215,7 +280,6 @@ describe('The mysql wrapper', function () {
     it('should use wrapQuery with expected arguments to wrap pool.query', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')
 
@@ -255,7 +319,6 @@ describe('The mysql wrapper', function () {
     it('should not leak credentials', function (done) {
       var sandbox = this.sandbox
       var fakeWrapQuery = sandbox.stub(utils, 'wrapQuery')
-      var fakeAgent = { clearly: 'a mock' }
 
       var mysql = require('mysql')