Experimental postgres plugin

.travis.yml

@@ -11,10 +11,13 @@ node_js:
 services:
   - mongodb
   - redis-server
+  - postgresql
 
 before_install:
   - mysql -e "create database IF NOT EXISTS test;" -uroot
   - echo "USE mysql;\nUPDATE user SET password=PASSWORD('Password12!') WHERE user='root';\nFLUSH PRIVILEGES;\n" | mysql -u root
+  - psql -c 'create database test;' -U postgres
+  - psql -c "alter user postgres with password 'Password12!';" -U postgres
 
 env:
   - GCLOUD_PROJECT=0 CXX=g++-4.8

appveyor.yml

@@ -9,6 +9,7 @@ environment:
 services:
   - mongodb
   - mysql
+  - postgresql
 
 # Install scripts. (runs after repo cloning)
 install:
@@ -28,6 +29,10 @@ install:
 before_test:
   - SET PATH=C:\Program Files\MySql\MySQL Server 5.7\bin;%PATH%
   - mysqladmin --host=localhost --user=root --password=Password12! create test
+  - SET PGUSER=postgres
+  - SET PGPASSWORD=Password12!
+  - PATH=C:\Program Files\PostgreSQL\9.4\bin\;%PATH%
+  - createdb test
 
 # Post-install test scripts.
 test_script:

bin/docker-trace.sh

@@ -5,11 +5,13 @@ if [ ! -z $1 ]; then
     docker run --name trace-test-mongo -p 127.0.0.1:27017:27017 -d mongo &&\
     docker run --name trace-test-redis -p 127.0.0.1:6379:6379 -d redis &&\
     docker run --name trace-test-mysql -p 127.0.0.1:3306:3306 -e MYSQL_ROOT_PASSWORD='Password12!' -e MYSQL_DATABASE=test -d mysql
+    docker run --name trace-test-postgres -p 127.0.0.1:5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD='Password12!' -e POSTGRES_DB=test -d postgres
     exit $?
   elif [ $COMMAND = 'stop' ]; then
     docker stop trace-test-mongo
     docker stop trace-test-redis
     docker stop trace-test-mysql
+    docker stop trace-test-postgres
     exit $?
   fi
 fi

config.js

@@ -52,6 +52,7 @@ module.exports = {
       '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'),
+      'pg': path.join(__dirname, 'src/plugins/plugin-pg.js'),
       'redis': path.join(__dirname, 'src/plugins/plugin-redis.js'),
       'restify': path.join(__dirname, 'src/plugins/plugin-restify.js')
     },

src/plugins/plugin-pg.js

@@ -0,0 +1,71 @@
+/**
+ * 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 SUPPORTED_VERSIONS = '^6.x';
+
+module.exports = [
+  {
+    file: 'lib/client.js',
+    versions: SUPPORTED_VERSIONS,
+    patch: function(Client, api) {
+      function queryWrap(query) {
+        return function query_trace() {
+          var span = api.createChildSpan({
+            name: 'pg-query'
+          });
+          var pgQuery = query.apply(this, arguments);
+          if (!span) {
+            return pgQuery;
+          }
+          if (api.enhancedDatabaseReportingEnabled()) {
+            span.addLabel('query', pgQuery.text);
+            if (pgQuery.values) {
+              span.addLabel('values', pgQuery.values);
+            }
+          }
+          api.wrapEmitter(pgQuery);
+          var done = pgQuery.callback;
+          pgQuery.callback = api.wrap(function(err, res) {
+            if (api.enhancedDatabaseReportingEnabled()) {
+              if (err) {
+                span.addLabel('error', err);
+              }
+              if (res) {
+                span.addLabel('row_count', res.rowCount);
+                span.addLabel('oid', res.oid);
+                span.addLabel('rows', res.rows);
+                span.addLabel('fields', res.fields);
+              }
+            }
+            span.endSpan();
+            if (done) {
+              done(err, res);
+            }
+          });
+          return pgQuery;
+        };
+      }
+
+      shimmer.wrap(Client.prototype, 'query', queryWrap);
+    },
+    unpatch: function(Client) {
+      shimmer.unwrap(Client.prototype, 'query');
+    }
+  }
+];

test/pg-config.js

@@ -0,0 +1,26 @@
+/**
+ * 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';
+
+module.exports = (process.env.CIRCLECI === 'true') ? {
+  host: 'localhost',
+  user: 'ubuntu',
+  database: 'circle_test'
+} : {
+  user: 'postgres',
+  password: 'Password12!',
+  database: 'test'
+};

test/plugins/fixtures/pg6/index.js

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

test/plugins/fixtures/pg6/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "pg6",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "pg": "^6.1.2"
+  }
+}

test/plugins/test-trace-pg.js

@@ -0,0 +1,129 @@
+/**
+ * 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 common = require('./common.js');
+var traceLabels = require('../../src/trace-labels.js');
+var assert = require('assert');
+
+describe('test-trace-pg', function() {
+  var traceApi;
+  var pool;
+  var client;
+  var releaseClient;
+  before(function() {
+    traceApi = require('../..').start({
+      samplingRate: 0,
+      enhancedDatabaseReporting: true
+    });
+    var pg = require('./fixtures/pg6');
+    pool = new pg.Pool(require('../pg-config.js'));
+  });
+
+  beforeEach(function(done) {
+    pool.connect(function(err, c, release) {
+      client = c;
+      releaseClient = release;
+      assert(!err);
+      client.query('CREATE TABLE t (name text NOT NULL, id text NOT NULL)', [],
+          function(err, res) {
+        assert(!err);
+        common.cleanTraces(traceApi);
+        done();
+      });
+    });
+  });
+
+  afterEach(function(done) {
+    client.query('DROP TABLE t', [], function(err, res) {
+      assert(!err);
+      releaseClient();
+      common.cleanTraces(traceApi);
+      done();
+    });
+  });
+
+  it('should perform basic operations', function(done) {
+    common.runInTransaction(traceApi, function(endRootSpan) {
+      client.query('INSERT INTO t (name, id) VALUES($1, $2)',
+          ['test_name', 'test_id'], function(err, res) {
+        endRootSpan();
+        assert(!err);
+        var span = common.getMatchingSpan(traceApi, function (span) {
+          return span.name === 'pg-query';
+        });
+        assert.equal(span.labels.query, 'INSERT INTO t (name, id) VALUES($1, $2)');
+        assert.equal(span.labels.values, '[ \'test_name\', \'test_id\' ]');
+        assert.equal(span.labels.row_count, '1');
+        assert.equal(span.labels.oid, '0');
+        assert.equal(span.labels.rows, '[]');
+        assert.equal(span.labels.fields, '[]');
+        done();
+      });
+    });
+  });
+
+  it('should remove trace frames from stack', function(done) {
+    common.runInTransaction(traceApi, function(endRootSpan) {
+      client.query('SELECT $1::int AS number', [1], function(err, res) {
+        endRootSpan();
+        assert(!err);
+        var span = common.getMatchingSpan(traceApi, function (span) {
+          return span.name === 'pg-query';
+        });
+        var labels = span.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('query_trace') !== -1);
+        done();
+      });
+    });
+  });
+
+  it('should work with events', function(done) {
+    common.runInTransaction(traceApi, function(endRootSpan) {
+      var query = client.query('SELECT $1::int AS number', [1]);
+      query.on('row', function(row) {
+        assert.strictEqual(row.number, 1);
+      });
+      query.on('end', function() {
+        endRootSpan();
+        var span = common.getMatchingSpan(traceApi, function (span) {
+          return span.name === 'pg-query';
+        });
+        assert.equal(span.labels.query, 'SELECT $1::int AS number');
+        assert.equal(span.labels.values, '[ \'1\' ]');
+        done();
+      });
+    });
+  });
+
+  it('should work without events or callback', function(done) {
+    common.runInTransaction(traceApi, function(endRootSpan) {
+      client.query('SELECT $1::int AS number', [1]);
+      setTimeout(function() {
+        endRootSpan();
+        var span = common.getMatchingSpan(traceApi, function (span) {
+          return span.name === 'pg-query';
+        });
+        assert.equal(span.labels.query, 'SELECT $1::int AS number');
+        assert.equal(span.labels.values, '[ \'1\' ]');
+        done();
+      }, 50);
+    });
+  });
+});

test/test-config-plugins.js

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