fix: add support for pg 7 changes (#702)

PR-URL: #702

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",

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<pg_7.QueryResult>;
+
+// 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<typeof done>.
+          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<typeof pg_6.Client>,
+  {
+    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<pg.Result>
+          // where [*] is one of:
+          // - ...[query: { text: string, values?: Array<any> }]
+          // - ...[text: string, values?: Array<any>]
+          // 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<typeof pg_7.Client>
 ];
 
-export default {};
+export = plugin;

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<any>));
+
+  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
 };

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"