From f070636eb4ed9adb3f59c55e354913915f495f1e Mon Sep 17 00:00:00 2001 From: Kelvin Jin Date: Thu, 29 Mar 2018 08:47:18 -0700 Subject: [PATCH] fix: add support for pg 7 changes (#702) PR-URL: #702 --- package-lock.json | 26 ++++ package.json | 1 + src/plugins/plugin-pg.ts | 215 ++++++++++++++++++++----- src/plugins/types/index.d.ts | 29 +++- test/fixtures/plugin-fixtures.json | 5 + test/plugins/test-trace-pg.ts | 242 ++++++++++++++++++----------- tsconfig.json | 1 + 7 files changed, 386 insertions(+), 133 deletions(-) diff --git a/package-lock.json b/package-lock.json index 276ae623d..058aa1d14 100644 --- a/package-lock.json +++ b/package-lock.json @@ -298,6 +298,26 @@ "integrity": "sha512-cnEvTAVVRqF6OQg/4SLnbxQ0slZJHqZQDve5BzGhcIQtuMpPv8T5QNS2cBPa/W0jTxciqwn7bmJAIGe/bOJ5Kw==", "dev": true }, + "@types/pg": { + "version": "7.4.5", + "resolved": "https://registry.npmjs.org/@types/pg/-/pg-7.4.5.tgz", + "integrity": "sha512-DV9A1X9duAnZrF+ANT9i7Z3k+49Dfl96hJlmpz8KCZtBaB7ck3eaAX/37P/vOtpb1VBS5C7xfYI1oRnAfL71DQ==", + "dev": true, + "requires": { + "@types/events": "1.2.0", + "@types/node": "9.4.6", + "@types/pg-types": "1.11.4" + } + }, + "@types/pg-types": { + "version": "1.11.4", + "resolved": "https://registry.npmjs.org/@types/pg-types/-/pg-types-1.11.4.tgz", + "integrity": "sha512-WdIiQmE347LGc1Vq3Ki8sk3iyCuLgnccqVzgxek6gEHp2H0p3MQ3jniIHt+bRODXKju4kNQ+mp53lmP5+/9moQ==", + "dev": true, + "requires": { + "moment": "2.21.0" + } + }, "@types/pify": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/@types/pify/-/pify-3.0.0.tgz", @@ -4227,6 +4247,12 @@ "resolved": "https://registry.npmjs.org/module-details-from-path/-/module-details-from-path-1.0.3.tgz", "integrity": "sha1-EUyUlnPiqKNenTV4hSeqN7Z52is=" }, + "moment": { + "version": "2.21.0", + "resolved": "https://registry.npmjs.org/moment/-/moment-2.21.0.tgz", + "integrity": "sha512-TCZ36BjURTeFTM/CwRcViQlfkMvL1/vFISuNLO5GkcVm1+QHfbSiNqZuWeMFjj1/3+uAjXswgRk30j1kkLYJBQ==", + "dev": true + }, "ms": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", diff --git a/package.json b/package.json index 7b46f7310..6818544a7 100644 --- a/package.json +++ b/package.json @@ -64,6 +64,7 @@ "@types/nock": "^9.1.2", "@types/node": "^9.4.6", "@types/once": "^1.4.0", + "@types/pg": "^7.4.5", "@types/pify": "^3.0.0", "@types/proxyquire": "^1.3.28", "@types/request": "^2.0.8", diff --git a/src/plugins/plugin-pg.ts b/src/plugins/plugin-pg.ts index bf5035cc8..d9b909888 100644 --- a/src/plugins/plugin-pg.ts +++ b/src/plugins/plugin-pg.ts @@ -13,61 +13,196 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -'use strict'; -var shimmer = require('shimmer'); +import {EventEmitter} from 'events'; +import * as shimmer from 'shimmer'; +import {Readable} from 'stream'; -var SUPPORTED_VERSIONS = '^6.x || ^7.x'; +import {Patch, Plugin, SpanData} from '../plugin-types'; -module.exports = [ +import {pg_6, pg_7} from './types'; + +// TS: Client#query also accepts a callback as a last argument, but TS cannot +// detect this as it's a dependent type. So we don't specify it here. +type ClientQueryArguments = + [{submit?: Function} & pg_7.QueryConfig]|[string]|[string, {}]; +type PG7QueryReturnValue = (pg_7.QueryConfig&({submit: Function}&EventEmitter)| + pg_7.Query)|Promise; + +// tslint:disable-next-line:no-any +function isSubmittable(obj: any): obj is {submit: Function} { + return typeof obj.submit === 'function'; +} + +const noOp = () => {}; + +function populateLabelsFromInputs(span: SpanData, args: ClientQueryArguments) { + const queryObj = args[0]; + if (typeof queryObj === 'object') { + if (queryObj.text) { + span.addLabel('query', queryObj.text); + } + if (queryObj.values) { + span.addLabel('values', queryObj.values); + } + } else if (typeof queryObj === 'string') { + span.addLabel('query', queryObj); + if (args.length >= 2 && typeof args[1] !== 'function') { + span.addLabel('values', args[1]); + } + } +} + +function populateLabelsFromOutputs( + span: SpanData, err: Error|null, res?: pg_7.QueryResult) { + 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); + } +} + +const plugin: Plugin = [ { 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); + versions: '^6.x', + // TS: Client is a class name. + // tslint:disable-next-line:variable-name + patch: (Client, api) => { + const maybePopulateLabelsFromInputs = + api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs : + noOp; + const maybePopulateLabelsFromOutputs = + api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs : + noOp; + shimmer.wrap(Client.prototype, 'query', (query) => { + return function query_trace(this: pg_6.Client) { + const span = api.createChildSpan({name: 'pg-query'}); if (!api.isRealSpan(span)) { - return pgQuery; + return query.apply(this, arguments); } - if (api.enhancedDatabaseReportingEnabled()) { - span.addLabel('query', pgQuery.text); - if (pgQuery.values) { - span.addLabel('values', pgQuery.values); - } + const argLength = arguments.length; + if (argLength >= 1) { + const args: ClientQueryArguments = + Array.prototype.slice.call(arguments, 0); + // Extract query text and values, if needed. + maybePopulateLabelsFromInputs(span, args); } + const pgQuery: pg_6.QueryReturnValue = query.apply(this, arguments); 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); - } + const done = pgQuery.callback; + // TODO(kjin): Clean up this line a little bit by casting the function + // passed to api.wrap as a NonNullable. + pgQuery.callback = + api.wrap((err: Error|null, res?: pg_7.QueryResult) => { + maybePopulateLabelsFromOutputs(span, err, res); + span.endSpan(); + if (done) { + done(err, res); + } + }); + return pgQuery; + }; + }); + }, + // TS: Client is a class name. + // tslint:disable-next-line:variable-name + unpatch(Client) { + shimmer.unwrap(Client.prototype, 'query'); + } + } as Patch, + { + file: 'lib/client.js', + versions: '^7.x', + // TS: Client is a class name. + // tslint:disable-next-line:variable-name + patch: (Client, api) => { + const maybePopulateLabelsFromInputs = + api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs : + noOp; + const maybePopulateLabelsFromOutputs = + api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs : + noOp; + shimmer.wrap(Client.prototype, 'query', (query) => { + return function query_trace(this: pg_7.Client) { + const span = api.createChildSpan({name: 'pg-query'}); + if (!api.isRealSpan(span)) { + return query.apply(this, arguments); + } + + let pgQuery: PG7QueryReturnValue; + // In 7.x, the value of pgQuery depends on how the query() was called. + // It can be one of: + // - (query: pg.Submittable) => EventEmitter + // - Note: return value is the same as the argument. + // - ([*], callback: (err, res: pg.Result) => void) => void + // - ([*]) => Promise + // where [*] is one of: + // - ...[query: { text: string, values?: Array }] + // - ...[text: string, values?: Array] + // See: https://node-postgres.com/guides/upgrading + const argLength = arguments.length; + if (argLength >= 1) { + const args: ClientQueryArguments = + Array.prototype.slice.call(arguments, 0); + + // Extract query text and values, if needed. + maybePopulateLabelsFromInputs(span, args); + + // If we received a callback, bind it to the current context, + // optionally adding labels as well. + const callback = args[args.length - 1]; + if (typeof callback === 'function') { + args[args.length - 1] = + api.wrap((err: Error|null, res?: pg_7.QueryResult) => { + maybePopulateLabelsFromOutputs(span, err, res); + span.endSpan(); + // TS: Type cast is safe as we know that callback is a + // Function. + (callback as (err: Error|null, res?: pg_7.QueryResult) => + void)(err, res); + }); + pgQuery = query.apply(this, args); + } else { + pgQuery = query.apply(this, arguments); } - span.endSpan(); - if (done) { - done(err, res); + } else { + pgQuery = query.apply(this, arguments); + } + + if (pgQuery) { + if (pgQuery instanceof EventEmitter) { + api.wrapEmitter(pgQuery); + } else if (typeof pgQuery.then === 'function') { + // Ensure that the span is ended, optionally adding labels as + // well. + pgQuery = pgQuery.then( + (res) => { + maybePopulateLabelsFromOutputs(span, null, res); + span.endSpan(); + return res; + }, + (err) => { + maybePopulateLabelsFromOutputs(span, err); + span.endSpan(); + throw err; + }); } - }); + } return pgQuery; }; - } - - shimmer.wrap(Client.prototype, 'query', queryWrap); + }); }, - unpatch: function(Client) { + // TS: Client is a class name. + // tslint:disable-next-line:variable-name + unpatch(Client) { shimmer.unwrap(Client.prototype, 'query'); } - } + } as Patch ]; -export default {}; +export = plugin; diff --git a/src/plugins/types/index.d.ts b/src/plugins/types/index.d.ts index 649f7542d..383f7146b 100644 --- a/src/plugins/types/index.d.ts +++ b/src/plugins/types/index.d.ts @@ -2,12 +2,14 @@ import * as connect_3 from 'connect'; // connect@3 import * as express_4 from 'express'; // express@4 import * as hapi_16 from 'hapi'; // hapi@16 import * as koa_2 from 'koa'; // koa@2 +import * as pg_7 from 'pg'; // pg@7 import * as restify_5 from 'restify'; // restify@5 -//---koa@1---// - import { EventEmitter } from 'events'; import { Server } from 'http'; +import { Readable } from 'stream'; + +//---koa@1---// declare class koa_1 extends EventEmitter { use(middleware: koa_1.Middleware): this; @@ -28,6 +30,27 @@ declare namespace koa_1 { interface Context extends koa_2.Context {} } +//---pg@6---// + +declare namespace pg_6 { + // PG 6's method signature for Client#query differs from that of PG 7 in that + // the return value is either a Submittable if one was passed in, or a + // pg.Query object instead. (In PG 6, pg.Query is PromiseLike and contains + // values passed in as the query configuration.) + // + // References: + // https://node-postgres.com/guides/upgrading#client-query-on + // https://github.com/brianc/node-postgres/blob/v6.4.2/lib/client.js#L355 + type QueryReturnValue = ( + pg_7.QueryConfig & + { callback?: (err: Error|null, res?: pg_7.QueryResult) => void } + ) & (({ submit: Function } & Readable) | (pg_7.Query & PromiseLike)); + + class Client { + query(...args: any[]): QueryReturnValue; + } +} + //---exports---// export { @@ -36,5 +59,7 @@ export { hapi_16, koa_1, koa_2, + pg_6, + pg_7, restify_5 }; diff --git a/test/fixtures/plugin-fixtures.json b/test/fixtures/plugin-fixtures.json index 74de24373..b8c02c235 100644 --- a/test/fixtures/plugin-fixtures.json +++ b/test/fixtures/plugin-fixtures.json @@ -152,6 +152,11 @@ "pg": "^6.1.2" } }, + "pg7": { + "dependencies": { + "pg": "^7.4.1" + } + }, "redis0.12": { "dependencies": { "redis": "^0.12.1" diff --git a/test/plugins/test-trace-pg.ts b/test/plugins/test-trace-pg.ts index 4175d9794..3e45c8361 100644 --- a/test/plugins/test-trace-pg.ts +++ b/test/plugins/test-trace-pg.ts @@ -16,127 +16,187 @@ 'use strict'; import { TraceLabels } from '../../src/trace-labels'; +import { FORCE_NEW } from '../../src/util'; var common = require('./common'/*.js*/); var assert = require('assert'); -describe('test-trace-pg', function() { - var traceApi; - var pool; - var client; - var releaseClient; - before(function() { - traceApi = require('../../..').start({ - projectId: '0', - samplingRate: 0, - enhancedDatabaseReporting: true +var pgVersions = ['6', '7']; + +pgVersions.forEach(pgVersion => { + describe(`test-trace-pg (v${pgVersion})`, function() { + var pg; + var traceApi; + var pool; + var client; + var releaseClient; + before(function() { + traceApi = require('../../..').start({ + projectId: '0', + samplingRate: 0, + enhancedDatabaseReporting: true, + [FORCE_NEW]: true + }); + pg = require(`./fixtures/pg${pgVersion}`); + 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(); + done(); + }); + }); }); - 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) { + afterEach(function(done) { + client.query('DROP TABLE t', [], function(err, res) { assert(!err); + releaseClient(); common.cleanTraces(); done(); }); }); - }); - afterEach(function(done) { - client.query('DROP TABLE t', [], function(err, res) { - assert(!err); - releaseClient(); - common.cleanTraces(); - done(); + it('should perform basic operations', function(done) { + common.runInTransaction(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(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 perform basic operations', function(done) { - common.runInTransaction(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(function (span) { - return span.name === 'pg-query'; + it('should perform basic operations with promises', function(done) { + common.runInTransaction(function(endRootSpan) { + client.query('INSERT INTO t (name, id) VALUES($1, $2)', + ['test_name', 'test_id']) + .then((res) => { + endRootSpan(); + var span = common.getMatchingSpan(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(); + }, (err) => { + assert.fail('Error not expected'); + }); + }); + }); + + it('should propagate context', function(done) { + common.runInTransaction(function(endRootSpan) { + client.query('INSERT INTO t (name, id) VALUES($1, $2)', + ['test_name', 'test_id'], function(err, res) { + assert.ok(common.hasContext()); + endRootSpan(); + done(); }); - 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 propagate context', function(done) { - common.runInTransaction(function(endRootSpan) { - client.query('INSERT INTO t (name, id) VALUES($1, $2)', - ['test_name', 'test_id'], function(err, res) { - assert.ok(common.hasContext()); - endRootSpan(); - done(); + it('should propagate context with promises', function(done) { + common.runInTransaction(function(endRootSpan) { + client.query('INSERT INTO t (name, id) VALUES($1, $2)', + ['test_name', 'test_id']) + .then((res) => { + assert.ok(common.hasContext()); + endRootSpan(); + done(); + }); }); }); - }); - it('should remove trace frames from stack', function(done) { - common.runInTransaction(function(endRootSpan) { - client.query('SELECT $1::int AS number', [1], function(err, res) { - endRootSpan(); - assert(!err); - var span = common.getMatchingSpan(function (span) { - return span.name === 'pg-query'; + it('should remove trace frames from stack', function(done) { + common.runInTransaction(function(endRootSpan) { + client.query('SELECT $1::int AS number', [1], function(err, res) { + endRootSpan(); + assert(!err); + var span = common.getMatchingSpan(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(); }); - 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(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(function (span) { - return span.name === 'pg-query'; + it('should work with events', function(done) { + common.runInTransaction(function(endRootSpan) { + var query = client.query(new pg.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(function (span) { + return span.name === 'pg-query'; + }); + assert.equal(span.labels.query, 'SELECT $1::int AS number'); + assert.equal(span.labels.values, '[ 1 ]'); + done(); }); - 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(function(endRootSpan) { - client.query('SELECT $1::int AS number', [1]); - setTimeout(function() { - endRootSpan(); - var span = common.getMatchingSpan(function (span) { - return span.name === 'pg-query'; + it('should work with generic Submittables', function(done) { + common.runInTransaction(function(endRootSpan) { + let submitCalled = false; + client.query({ + submit: (connection) => { + // Indicate that the next item may be processed. + connection.emit('readyForQuery'); + submitCalled = true; + endRootSpan(); + common.getMatchingSpan(function (span) { + return span.name === 'pg-query'; + }); + done(); + }, + handleReadyForQuery: () => {} }); - assert.equal(span.labels.query, 'SELECT $1::int AS number'); - assert.equal(span.labels.values, '[ \'1\' ]'); - done(); - }, 50); + }); + }); + + it('should work without events or callback', function(done) { + common.runInTransaction(function(endRootSpan) { + client.query('SELECT $1::int AS number', [1]); + setTimeout(function() { + endRootSpan(); + var span = common.getMatchingSpan(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); + }); }); }); }); diff --git a/tsconfig.json b/tsconfig.json index 6fa33d958..7bce0c978 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -15,6 +15,7 @@ "src/plugins/plugin-http2.ts", "src/plugins/plugin-https.ts", "src/plugins/plugin-koa.ts", + "src/plugins/plugin-pg.ts", "src/plugins/plugin-restify.ts", "test/plugins/test-trace-google-gax.ts", "test/plugins/test-trace-http2.ts",