fix: avoid fake .then() method on mysql2 Query class (#501)

* fix: avoid fake .then() method on mysql2 Query class

* test: add unit test for fake .then() method on mysql2 Query class

Author: mickdekkers

packages/mysql/lib/mysql_p.js

@@ -33,13 +33,26 @@ module.exports = function captureMySQL(mysql) {
   return mysql;
 };
 
+function isPromise(maybePromise) {
+  if (maybePromise != null && maybePromise.then instanceof Function) {
+    // mysql2 has a `Query` class with a `then` method which always throws an error when called.
+    // We want to avoid calling this, so we need to check for more than just the presence of a `then` method.
+    // See https://github.com/sidorares/node-mysql2/blob/dbb344e89a1cc8bb457b24e67b07cdb3013fe844/lib/commands/query.js#L38-L44
+    // Since it's highly unlikely that any Promise implementation would name their class `Query`,
+    // we can safely use this to determine whether or not this is actually a Promise.
+    const constructorName = maybePromise.constructor != null ? maybePromise.constructor.name : undefined;
+    return constructorName !== 'Query';
+  }
+  return false;
+}
+
 function patchCreateConnection(mysql) {
   var baseFcn = '__createConnection';
   mysql[baseFcn] = mysql['createConnection'];
 
   mysql['createConnection'] = function patchedCreateConnection() {
     var connection = mysql[baseFcn].apply(connection, arguments);
-    if (connection && connection.then instanceof Function) {
+    if (isPromise(connection)) {
       connection = connection.then((result) => {
         patchObject(result.connection);
         return result;
@@ -57,7 +70,7 @@ function patchCreatePool(mysql) {
 
   mysql['createPool'] = function patchedCreatePool() {
     var pool = mysql[baseFcn].apply(pool, arguments);
-    if (pool && pool.then instanceof Function) {
+    if (isPromise(pool)) {
       pool = pool.then((result) => {
         patchObject(result.pool);
         return result;
@@ -112,7 +125,7 @@ function patchGetConnection(pool) {
     }
 
     var result = pool[baseFcn].apply(pool, args);
-    if (result && result.then instanceof Function) {
+    if (isPromise(result)) {
       return result.then(patchObject);
     } else {
       return result;
@@ -232,7 +245,7 @@ function captureOperation(name) {
         }
       };
 
-      if (command.then instanceof Function) {
+      if (isPromise(command)) {
         command.then(() => {
           subsegment.close();
         });

packages/mysql/test/unit/mysql_p.test.js

@@ -4,6 +4,7 @@ var sinon = require('sinon');
 var sinonChai = require('sinon-chai');
 
 var AWSXRay = require('aws-xray-sdk-core');
+var EventEmitter = require('events').EventEmitter;
 
 var captureMySQL = require('../../lib/mysql_p');
 var Segment = AWSXRay.Segment;
@@ -34,6 +35,55 @@ describe('captureMySQL', function() {
       assert.equal(patched.createPool.name, 'patchedCreatePool');
       assert.equal(patched.createPoolCluster.name, 'patchedCreatePoolCluster');
     });
+
+    describe('handles mysql2 quirks', function() {
+      var conn, connectionObj, mysql,  sandbox, segment;
+
+      // See https://github.com/sidorares/node-mysql2/blob/dbb344e89a1cc8bb457b24e67b07cdb3013fe844/lib/commands/query.js#L38-L44
+      class Query extends EventEmitter {
+        then() {
+          throw new Error('You have tried to call .then(), .catch(), or invoked await on the result of query that is not a promise [...]');
+        }
+      }
+
+      before(function() {
+        conn = {
+          config: {
+            user: 'mcmuls',
+            host: 'database.location',
+            port: '8080',
+            database: 'myTestDb'
+          },
+          query: function() {
+            return new Query();
+          }
+        };
+
+        mysql = { createConnection: function() {
+          return conn;
+        } };
+        mysql = captureMySQL(mysql);
+        connectionObj = mysql.createConnection();
+      });
+
+      beforeEach(function() {
+        sandbox = sinon.createSandbox();
+        segment = new Segment('test');
+        segment.addNewSubsegment('testSub');
+
+        sandbox = sinon.createSandbox();
+        sandbox.stub(AWSXRay, 'getSegment').returns(segment);
+        sandbox.stub(AWSXRay, 'isAutomaticMode').returns(true);
+      });
+
+      afterEach(function() {
+        sandbox.restore();
+      });
+
+      it('should not call fake .then() on mysql2 Query class', function() {
+        (() => connectionObj.query()).should.not.throw();
+      });
+    });
   });
 
   describe('#captureQuery', function() {