feature(plugin): implement postgres plugin

.circleci/config.yml

@@ -1,5 +1,11 @@
 version: 2
 
+test_env: &test_env
+  TEST_POSTGRES: 1
+
+postgres_service: &postgres_service
+  image: postgres:alpine
+
 node_unit_tests: &node_unit_tests
   steps:
     - checkout
@@ -71,18 +77,26 @@ jobs:
   node8:
     docker:
       - image: node:8
+        environment: *test_env
+      - *postgres_service
     <<: *node_unit_tests
   node10:
     docker:
       - image: node:10
+        environment: *test_env
+      - *postgres_service
     <<: *node_unit_tests
   node11:
     docker:
       - image: node:11
+        environment: *test_env
+      - *postgres_service
     <<: *node_unit_tests
   node12:
     docker:
       - image: node:12
+        environment: *test_env
+      - *postgres_service
     <<: *node_unit_tests
   node12-browsers:
     docker:

packages/opentelemetry-plugin-postgres/package.json

@@ -7,7 +7,8 @@
   "types": "build/src/index.d.ts",
   "repository": "open-telemetry/opentelemetry-js",
   "scripts": {
-    "test": "nyc ts-mocha -p tsconfig.json 'test/**/*.ts'",
+    "test": "nyc ts-mocha -p tsconfig.json 'test/**/*.test.ts'",
+    "test:debug": "ts-mocha --inspect-brk --no-timeouts -p tsconfig.json 'test/**/*.test.ts'",
     "tdd": "yarn test -- --watch-extensions ts --watch",
     "clean": "rimraf build/*",
     "check": "gts check",
@@ -43,11 +44,14 @@
   "devDependencies": {
     "@types/mocha": "^5.2.7",
     "@types/node": "^12.6.9",
+    "@types/pg": "^7.11.2",
+    "@types/shimmer": "^1.0.1",
     "codecov": "^3.5.0",
-    "gts": "^1.1.0",
+    "gts": "^1.0.0",
     "mocha": "^6.2.0",
     "nyc": "^14.1.1",
     "rimraf": "^3.0.0",
+    "pg": "^7.12.1",
     "tslint-microsoft-contrib": "^6.2.0",
     "tslint-consistent-codestyle": "^1.15.1",
     "ts-mocha": "^6.0.0",
@@ -57,6 +61,8 @@
   "dependencies": {
     "@opentelemetry/core": "^0.1.0",
     "@opentelemetry/node": "^0.1.0",
-    "@opentelemetry/types": "^0.1.0"
+    "@opentelemetry/tracing": "^0.1.0",
+    "@opentelemetry/types": "^0.1.0",
+    "shimmer": "^1.2.1"
   }
 }

packages/opentelemetry-plugin-postgres/src/enums.ts

@@ -0,0 +1,36 @@
+/*!
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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.
+ */
+
+export enum AttributeNames {
+  // required by https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/data-semantic-conventions.md#databases-client-calls
+  COMPONENT = 'component',
+  DB_TYPE = 'db.type',
+  DB_INSTANCE = 'db.instance',
+  DB_STATEMENT = 'db.statement',
+  PEER_ADDRESS = 'peer.address',
+  PEER_HOSTNAME = 'peer.host',
+
+  // optional
+  DB_USER = 'db.user',
+  PEER_PORT = 'peer.port',
+  PEER_IPV4 = 'peer.ipv4',
+  PEER_IPV6 = 'peer.ipv6',
+  PEER_SERVICE = 'peer.service',
+
+  // PG specific -- not specified by spec
+  PG_VALUES = 'pg.values',
+  PG_PLAN = 'pg.plan',
+}

packages/opentelemetry-plugin-postgres/src/index.ts

@@ -13,3 +13,5 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
+export * from './pg';

packages/opentelemetry-plugin-postgres/src/pg.ts

@@ -0,0 +1,178 @@
+/*!
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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.
+ */
+
+import { BasePlugin } from '@opentelemetry/core';
+import { SpanKind, CanonicalCode } from '@opentelemetry/types';
+import { AttributeNames } from './enums';
+import {
+  PostgresPluginOptions,
+  PgClientExtended,
+  PgPluginQueryConfig,
+  PostgresCallback,
+} from './types';
+import * as pgTypes from 'pg';
+import * as shimmer from 'shimmer';
+import * as utils from './utils';
+
+export class PostgresPlugin extends BasePlugin<typeof pgTypes> {
+  protected _config: PostgresPluginOptions;
+
+  static readonly COMPONENT = 'pg';
+  static readonly DB_TYPE = 'sql';
+
+  static readonly BASE_SPAN_NAME = PostgresPlugin.COMPONENT + '.query';
+
+  readonly supportedVersions = ['^7.12.1'];
+
+  constructor(readonly moduleName: string) {
+    super();
+    this._config = {};
+  }
+
+  protected patch(): typeof pgTypes {
+    if (this._moduleExports.Client.prototype.query) {
+      shimmer.wrap(
+        this._moduleExports.Client.prototype,
+        'query',
+        this._getClientQueryPatch() as never
+      );
+    }
+    return this._moduleExports;
+  }
+
+  protected unpatch(): void {
+    if (this._moduleExports.Client.prototype.query) {
+      shimmer.unwrap(this._moduleExports.Client.prototype, 'query');
+    }
+  }
+
+  private _getClientQueryPatch() {
+    const plugin = this;
+    return (original: typeof pgTypes.Client.prototype.query) => {
+      plugin._logger.debug(
+        `Patching ${PostgresPlugin.COMPONENT}.Client.prototype.query`
+      );
+      return function query(
+        this: pgTypes.Client & PgClientExtended,
+        ...args: unknown[]
+      ) {
+        let callbackProvided = false;
+        const span = plugin._pgStartSpan(this);
+
+        // Handle different client.query(...) signatures
+        if (typeof args[0] === 'string') {
+          if (args.length > 1 && args[1] instanceof Array) {
+            utils._handleParameterizedQuery.call(this, span, ...args);
+          } else {
+            utils._handleTextQuery.call(this, span, ...args);
+          }
+        } else if (typeof args[0] === 'object') {
+          utils._handleConfigQuery.call(this, span, ...args);
+        }
+
+        // Bind callback to parent span
+        if (args.length > 0) {
+          const parentSpan = plugin._tracer.getCurrentSpan();
+          if (typeof args[args.length - 1] === 'function') {
+            // Patch ParameterQuery callback
+            args[args.length - 1] = utils._patchCallback(span, args[
+              args.length - 1
+            ] as PostgresCallback);
+            // If a parent span exists, bind the callback
+            if (parentSpan) {
+              args[args.length - 1] = plugin._tracer.bind(
+                args[args.length - 1]
+              );
+            }
+            callbackProvided = true;
+          } else if (
+            typeof (args[0] as PgPluginQueryConfig).callback === 'function'
+          ) {
+            // Patch ConfigQuery callback
+            let callback = utils._patchCallback(
+              span,
+              (args[0] as PgPluginQueryConfig).callback!
+            );
+            // If a parent span existed, bind the callback
+            if (parentSpan) {
+              callback = plugin._tracer.bind(callback);
+            }
+
+            // Copy the callback instead of writing to args.callback so that we don't modify user's
+            // original callback reference
+            args[0] = { ...args[0], callback };
+            callbackProvided = true;
+          }
+        }
+
+        // Perform the original query
+        const result: unknown = original.apply(this, args as never);
+
+        // Bind promise to parent span and end the span
+        if (result instanceof Promise) {
+          return result
+            .then((result: unknown) => {
+              // Return a pass-along promise which ends the span and then goes to user's orig resolvers
+              return new Promise((resolve, _) => {
+                span.setStatus({ code: CanonicalCode.OK });
+                span.end();
+                resolve(result);
+              });
+            })
+            .catch((error: Error) => {
+              return new Promise((_, reject) => {
+                span.setStatus({ code: CanonicalCode.UNKNOWN });
+                span.end();
+                reject(error);
+              });
+            });
+        }
+
+        // If a promise was not returned and no callback is provided, we recieved invalid args
+        if (!callbackProvided) {
+          span.setStatus({
+            code: CanonicalCode.INVALID_ARGUMENT,
+            message: 'Invalid query provided to the driver',
+          });
+          span.end();
+        }
+
+        // else returns void
+        return result; // void
+      };
+    };
+  }
+
+  // Private helper function to start a span
+  private _pgStartSpan(client: pgTypes.Client & PgClientExtended) {
+    const jdbcString = utils._getJDBCString(client.connectionParameters);
+    return this._tracer.startSpan(PostgresPlugin.BASE_SPAN_NAME, {
+      kind: SpanKind.CLIENT,
+      parent: this._tracer.getCurrentSpan() || undefined,
+      attributes: {
+        [AttributeNames.COMPONENT]: PostgresPlugin.COMPONENT, // required
+        [AttributeNames.DB_INSTANCE]: client.connectionParameters.database, // required
+        [AttributeNames.DB_TYPE]: PostgresPlugin.DB_TYPE, // required
+        [AttributeNames.PEER_ADDRESS]: jdbcString, // required
+        [AttributeNames.PEER_HOSTNAME]: client.connectionParameters.host, // required
+        [AttributeNames.PEER_PORT]: client.connectionParameters.port,
+        [AttributeNames.DB_USER]: client.connectionParameters.user,
+      },
+    });
+  }
+}
+
+export const plugin = new PostgresPlugin(PostgresPlugin.COMPONENT);

packages/opentelemetry-plugin-postgres/src/types.ts

@@ -0,0 +1,38 @@
+/*!
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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.
+ */
+
+import * as pgTypes from 'pg';
+
+export interface PostgresPluginOptions {}
+
+export type PostgresCallback = (err: Error, res: object) => unknown;
+
+// These are not included in @types/pg, so manually define them.
+// https://github.com/brianc/node-postgres/blob/fde5ec586e49258dfc4a2fcd861fcdecb4794fc3/lib/client.js#L25
+export interface PgClientConnectionParams {
+  database: string;
+  host: string;
+  port: number;
+  user: string;
+}
+
+export interface PgClientExtended {
+  connectionParameters: PgClientConnectionParams;
+}
+
+export interface PgPluginQueryConfig extends pgTypes.QueryConfig {
+  callback?: PostgresCallback;
+}