Support Knex

config.js

@@ -49,6 +49,7 @@ module.exports = {
       'grpc': path.join(__dirname, 'src/plugins/plugin-grpc.js'),
       'hapi': path.join(__dirname, 'src/plugins/plugin-hapi.js'),
       'http': path.join(__dirname, 'src/plugins/plugin-http.js'),
+      'knex': path.join(__dirname, 'src/plugins/plugin-knex.js'),
       'koa': path.join(__dirname, 'src/plugins/plugin-koa.js'),
       'mongodb-core': path.join(__dirname, 'src/plugins/plugin-mongodb-core.js'),
       'mysql': path.join(__dirname, 'src/plugins/plugin-mysql.js'),

src/plugins/plugin-knex.js

@@ -0,0 +1,58 @@
+/**
+ * Copyright 2017 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 shimmer = require('shimmer');
+var util = require('util');
+
+// knex 0.10.x and 0.11.x do not need patching
+var VERSIONS = '>=0.12 <=0.13';
+
+function patchClient(Client, api) {
+  shimmer.wrap(Client.prototype, 'runner', function(original) {
+    return function() {
+      var runner = original.apply(this, arguments);
+      runner.query = api.wrap(runner.query);
+      return runner;
+    };
+  });
+}
+
+function unpatchClient(Client) {
+  shimmer.unwrap(Client.prototype, 'runner');
+}
+
+function interceptTransaction(Transaction, api) {
+  function WrappedTransaction(client, container, config, outerTx) {
+    Transaction.call(this, client, api.wrap(container), config, outerTx);
+  }
+  util.inherits(WrappedTransaction, Transaction);
+  return WrappedTransaction;
+}
+
+module.exports = [
+  {
+    file: 'lib/transaction.js',
+    versions: VERSIONS,
+    intercept: interceptTransaction
+  },
+  {
+    file: 'lib/client.js',
+    versions: VERSIONS,
+    patch: patchClient,
+    unpatch: unpatchClient
+  }
+];

test/plugins/fixtures/knex0.10/index.js

@@ -0,0 +1 @@
+module.exports = require('knex');

test/plugins/fixtures/knex0.10/package.json

@@ -0,0 +1,9 @@
+{
+  "name": "knex0.10",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "knex": "^0.10.0",
+    "mysql": "^2.13.0"
+  }
+}

test/plugins/fixtures/knex0.11/index.js

@@ -0,0 +1 @@
+module.exports = require('knex');

test/plugins/fixtures/knex0.11/package.json

@@ -0,0 +1,9 @@
+{
+  "name": "knex0.11",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "knex": "^0.11.0",
+    "mysql": "^2.13.0"
+  }
+}

test/plugins/fixtures/knex0.12/index.js

@@ -0,0 +1 @@
+module.exports = require('knex');

test/plugins/fixtures/knex0.12/package.json

@@ -0,0 +1,9 @@
+{
+  "name": "knex0.12",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "knex": "^0.12.9",
+    "mysql": "^2.13.0"
+  }
+}

test/plugins/fixtures/knex0.13/index.js

@@ -0,0 +1 @@
+module.exports = require('knex');

test/plugins/fixtures/knex0.13/package.json

@@ -0,0 +1,9 @@
+{
+  "name": "knex0.13",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "knex": "^0.13.0",
+    "mysql": "^2.13.0"
+  }
+}

test/plugins/test-trace-knex.js

@@ -0,0 +1,267 @@
+/**
+ * Copyright 2017 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 assert = require('assert');
+
+var common = require('./common.js');
+var traceLabels = require('../../src/trace-labels.js');
+
+var RESULT_SIZE = 5;
+var TABLE_NAME = 't';
+
+var obj = {
+  k: 1,
+  v: 'obj'
+};
+
+var versions = {
+  knex10: './fixtures/knex0.10',
+  knex11: './fixtures/knex0.11',
+  knex12: './fixtures/knex0.12',
+  knex13: './fixtures/knex0.13'
+};
+
+describe('test-trace-knex', function() {
+  var agent;
+  before(function() {
+    agent = require('../..').start({
+      logLevel: 2,
+      flushDelaySeconds: 1,
+      enhancedDatabaseReporting: true,
+      databaseResultReportingSize: RESULT_SIZE
+    });
+  });
+
+  Object.keys(versions).forEach(function(version) {
+    describe('with mysql and ' + version, function() {
+      var knex;
+      before(function() {
+        knex = require(versions[version])({
+          client: 'mysql',
+          connection: require('../mysql-config')
+        });
+      });
+
+      beforeEach(function(done) {
+        knex.schema.createTable(TABLE_NAME, function(table) {
+          table.integer('k', 10);
+          table.string('v', 10);
+        }).then(function(result) {
+          assert.ok(result);
+          knex.insert(obj).into(TABLE_NAME).then(function(result) {
+            assert.ok(result);
+            common.cleanTraces(agent);
+            done();
+          });
+        });
+      });
+
+      afterEach(function(done) {
+        knex.schema.dropTable(TABLE_NAME).then(function(result) {
+          assert.ok(result);
+          common.cleanTraces(agent);
+          done();
+        });
+      });
+
+      it('should perform basic operations using ' + version, function(done) {
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex(TABLE_NAME).select().then(function(res) {
+            endRootSpan();
+            assert(res);
+            assert.equal(res.length, 1);
+            assert.equal(res[0].k, 1);
+            assert.equal(res[0].v, 'obj');
+            var spans = common.getMatchingSpans(agent, function (span) {
+              return span.name === 'mysql-query';
+            });
+            if (version === 'knex11') {
+              assert.equal(spans.length, 2);
+              assert.equal(spans[0].labels.sql, 'SELECT 1');
+              assert.equal(spans[1].labels.sql, 'select * from `t`');
+            }
+            else {
+              assert.equal(spans.length, 1);
+              assert.equal(spans[0].labels.sql, 'select * from `t`');
+            }
+            done();
+          });
+        });
+      });
+
+      it('should propagate context using ' + version, function(done) {
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex.select().from(TABLE_NAME).then(function(res) {
+            assert.ok(common.hasContext());
+            endRootSpan();
+            done();
+          }).catch(function(e) {
+            assert.ifError(e);
+          });
+        });
+      });
+
+      it('should remove trace frames from stack using ' + version, function(done) {
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex.select().from(TABLE_NAME).then(function(res) {
+            endRootSpan();
+            var spans = common.getMatchingSpans(agent, function (span) {
+              return span.name === 'mysql-query';
+            });
+            var labels = spans[0].labels;
+            var stackTrace = JSON.parse(labels[traceLabels.STACK_TRACE_DETAILS_KEY]);
+            // Ensure that our patch is on top of the stack
+            assert(
+              stackTrace.stack_frame[0].method_name.indexOf('createQuery_trace') !== -1);
+            done();
+          }).catch(function(e) {
+            assert.ifError(e);
+          });
+        });
+      });
+
+      it('should work with events using ' + version, function(done) {
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex.select().from(TABLE_NAME).on('query-response', function(response, obj, builder) {
+            var row = response[0];
+            assert.ok(row);
+            assert.equal(row.k, 1);
+            assert.equal(row.v, 'obj');
+          }).on('query-error', function(err, obj) {
+            assert.ifError(err);
+          }).then(function(res) {
+            endRootSpan();
+            var spans = common.getMatchingSpans(agent, function (span) {
+              return span.name === 'mysql-query';
+            });
+            if (version === 'knex11') {
+              assert.equal(spans.length, 2);
+              assert.equal(spans[0].labels.sql, 'SELECT 1');
+              assert.equal(spans[1].labels.sql, 'select * from `t`');
+            }
+            else {
+              assert.equal(spans.length, 1);
+              assert.equal(spans[0].labels.sql, 'select * from `t`');
+            }
+            done();
+          }).catch(function(e) {
+            assert.ifError(e);
+          });
+        });
+      });
+
+      it('should work without events or callback using ' + version, function(done) {
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex.select().from(TABLE_NAME).then(function(result) {
+            setTimeout(function() {
+              endRootSpan();
+              var spans = common.getMatchingSpans(agent, function (span) {
+                return span.name === 'mysql-query';
+              });
+              if (version === 'knex11') {
+                assert.equal(spans.length, 2);
+                assert.equal(spans[0].labels.sql, 'SELECT 1');
+                assert.equal(spans[1].labels.sql, 'select * from `t`');
+              }
+              else {
+                assert.equal(spans.length, 1);
+                assert.equal(spans[0].labels.sql, 'select * from `t`');
+              }
+              done();
+            }, 50);
+          });
+        });
+      });
+
+      it('should perform basic transaction using ' + version, function(done) {
+        var obj2 = {
+          k: 2,
+          v: 'obj2'
+        };
+        common.runInTransaction(agent, function(endRootSpan) {
+          knex.transaction(function(trx) {
+            knex.insert(obj2)
+                .into(TABLE_NAME)
+                .transacting(trx)
+                .then(function(res) {
+                  return trx.select()
+                     .from(TABLE_NAME)
+                     .then(function(res) {
+                       assert.equal(res.length, 2);
+                       assert.equal(res[0].k, 1);
+                       assert.equal(res[0].v, 'obj');
+                       assert.equal(res[1].k, 2);
+                       assert.equal(res[1].v, 'obj2');
+                     }).catch(function(err) {
+                       assert.ifError(err);
+                     });
+                })
+                .then(function() {
+                  trx.rollback(new Error('Rolling back'));
+                }).catch(function(err) {
+                  assert.ifError(err);
+                });
+          }).catch(function(err) {
+            assert.ok(err);
+            assert.strictEqual(err.message, 'Rolling back');
+            knex.select()
+              .from(TABLE_NAME)
+              .then(function(res) {
+                endRootSpan();
+                assert.equal(res.length, 1);
+                assert.equal(res[0].k, 1);
+                assert.equal(res[0].v, 'obj');
+                var spans = common.getMatchingSpans(agent, function (span) {
+                  return span.name === 'mysql-query';
+                });
+                var expectedCmds;
+                if (version === 'knex10') {
+                  expectedCmds = ['BEGIN;',
+                                  'insert into `t` (`k`, `v`) values (?, ?)',
+                                  'select * from `t`',
+                                  'ROLLBACK;',
+                                  'select * from `t`'];
+                }
+                else if (version === 'knex11') {
+                  expectedCmds = ['SELECT 1',
+                                  'BEGIN;',
+                                  'insert into `t` (`k`, `v`) values (?, ?)',
+                                  'select * from `t`',
+                                  'ROLLBACK;',
+                                  'SELECT 1',
+                                  'select * from `t`'];
+                }
+                else {
+                  expectedCmds = ['insert into `t` (`k`, `v`) values (?, ?)',
+                                  'select * from `t`',
+                                  'ROLLBACK;',
+                                  'select * from `t`'];
+                }
+                assert.equal(expectedCmds.length, spans.length);
+                for (var i = 0; i < spans.length; i++) {
+                  assert.equal(spans[i].labels.sql, expectedCmds[i]);
+                }
+                done();
+              }).catch(function(err) {
+                assert.ifError(err);
+              });
+          });
+        });
+      });
+    });
+  });
+});

test/test-config-plugins.js

@@ -21,8 +21,8 @@ var trace = require('..');
 
 var common = require('./plugins/common.js');
 
-var instrumentedModules = ['connect', 'express', 'generic-pool', 'grpc',
-  'hapi', 'http', 'koa', 'mongodb-core', 'mysql', 'pg', 'redis', 'restify'];
+var instrumentedModules = ['connect', 'express', 'generic-pool', 'grpc', 'hapi',
+  'http', 'knex', 'koa', 'mongodb-core', 'mysql', 'pg', 'redis', 'restify'];
 
 describe('plugin configuration', function() {
   it('should have correct defaults', function() {