From 62a6e61f8ad0713682ae8370002eb2fd3ad464c2 Mon Sep 17 00:00:00 2001 From: Shadowbeetle Date: Fri, 26 May 2017 10:42:47 +0200 Subject: [PATCH] fix(mysql): instrument pool.getConnection properly --- .../trace-instrumentation-mysql.js | 70 ++++++++++------- test/instrumentations/mysql/mysql.spec.js | 75 +++++++++++++++++-- 2 files changed, 111 insertions(+), 34 deletions(-) diff --git a/lib/instrumentations/trace-instrumentation-mysql.js b/lib/instrumentations/trace-instrumentation-mysql.js index 09ebf11..969bea7 100644 --- a/lib/instrumentations/trace-instrumentation-mysql.js +++ b/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) { diff --git a/test/instrumentations/mysql/mysql.spec.js b/test/instrumentations/mysql/mysql.spec.js index 8adb1ac..c1d7244 100644 --- a/test/instrumentations/mysql/mysql.spec.js +++ b/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')