From 039a64ea9b5c426a0ec4fdd9335743fa5f40f015 Mon Sep 17 00:00:00 2001
From: Matt Loring <mattloring@google.com>
Date: Thu, 24 Mar 2016 13:59:02 -0700
Subject: [PATCH] Patch mysql connection pool

If the mysql connection pool is accessed during a request we drop
context. By patching getConnection we can avoid this.

Fixes the mysql component of #221
---
 lib/hooks/userspace/hook-mysql.js  | 15 +++++++
 package.json                       |  1 +
 test/standalone/test-mysql-pool.js | 67 ++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+)
 create mode 100644 test/standalone/test-mysql-pool.js

diff --git a/lib/hooks/userspace/hook-mysql.js b/lib/hooks/userspace/hook-mysql.js
index 5b003beba..aeb7a5f25 100644
--- a/lib/hooks/userspace/hook-mysql.js
+++ b/lib/hooks/userspace/hook-mysql.js
@@ -72,6 +72,12 @@ function wrapCallback(span, done) {
   return cls.getNamespace().bind(fn);
 }
 
+function wrapGetConnection(getConnection) {
+  return function getConnection_trace(cb) {
+    return getConnection.call(this, cls.getNamespace().bind(cb));
+  };
+}
+
 module.exports = function(version_, agent_) {
   if (!semver.satisfies(version_, SUPPORTED_VERSIONS)) {
     agent_.logger.info('Mysql: unsupported version ' + version_ + ' loaded');
@@ -87,6 +93,15 @@ module.exports = function(version_, agent_) {
         shimmer.unwrap(Connection, 'createQuery');
         agent_.logger.info('Mysql: unpatched');
       }
+    },
+    'lib/Pool.js': {
+      patch: function(Pool) {
+        shimmer.wrap(Pool.prototype, 'getConnection', wrapGetConnection);
+      },
+      unpatch: function(Pool) {
+        shimmer.unwrap(Pool.prototype, 'getConnection');
+        agent_.logger.info('Mysql connection pool: unpatched');
+      }
     }
   };
 };
diff --git a/package.json b/package.json
index 64a80cd58..adb1bede3 100644
--- a/package.json
+++ b/package.json
@@ -38,6 +38,7 @@
     "closure-npc": "*",
     "coveralls": "^2.11.8",
     "glob": "^7.0.3",
+    "hapi-plugin-mysql": "^3.1.3",
     "istanbul": "^0.4.2",
     "jshint": "^2.9.1",
     "mocha": "^2.2.4",
diff --git a/test/standalone/test-mysql-pool.js b/test/standalone/test-mysql-pool.js
new file mode 100644
index 000000000..b43634487
--- /dev/null
+++ b/test/standalone/test-mysql-pool.js
@@ -0,0 +1,67 @@
+/**
+ * Copyright 2015 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+'use strict';
+
+var common = require('../hooks/common.js');
+var assert = require('assert');
+var http = require('http');
+var semver = require('semver');
+
+// hapi 13 and hapi-plugin-mysql uses const
+if (semver.satisfies(process.version, '>=4')) {
+  var Hapi = require('../hooks/fixtures/hapi13');
+  describe('test-trace-mysql', function() {
+    it('should work with connection pool access', function(done) {
+      var server = new Hapi.Server();
+      server.connection({ port: common.serverPort });
+      server.register({
+        register: require('hapi-plugin-mysql'),
+        options: {
+          host     : 'localhost',
+          user     : 'travis',
+          password : '',
+          database : 'test'
+        }
+      }, function (err) {
+        assert(!err);
+        server.route({ 
+          method: 'GET', 
+          path: '/', 
+          handler: function (request, reply) {
+            request.app.db.query('SELECT * FROM t', function(err, res) {
+              return reply(common.serverRes);
+            }); 
+          } 
+        });
+        server.start(function(err) {
+          assert(!err);
+          http.get({port: common.serverPort}, function(res) {
+            var result = '';
+            res.on('data', function(data) { result += data; });
+            res.on('end', function() {
+              var spans = common.getMatchingSpans(function (span) {
+                return span.name === 'mysql-query';
+              });
+              assert.equal(spans.length, 1);
+              assert.equal(spans[0].labels.sql, 'SELECT * FROM t');
+              done();
+            });
+          });
+        });
+      });
+    });
+  });
+}