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'",
+    "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",
@@ -42,11 +43,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",
     "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",
@@ -56,6 +60,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,24 @@
+/*!
+ * 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 {
+  COMPONENT = 'component',
+  PG_HOST = 'pg.host',
+  PG_PORT = 'pg.port',
+  PG_TEXT = 'pg.text',
+  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,160 @@
+/*!
+ * 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 } from '@opentelemetry/types';
+import { AttributeNames } from './enums';
+import { PostgresCallback, PostgresPluginOptions } from './types';
+import * as path from 'path';
+import * as pgTypes from 'pg';
+import * as shimmer from 'shimmer';
+
+export class PostgresPlugin extends BasePlugin<typeof pgTypes> {
+  static readonly component = 'pg';
+  readonly supportedVersions = ['^7.12.1'];
+  protected _config: PostgresPluginOptions;
+
+  constructor(readonly moduleName: string, readonly version: 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, ...args: unknown[]) {
+        // setup span
+        let callbackProvided: boolean =
+          args.length > 1 && typeof args[args.length - 1] === 'function';
+        const span = plugin._tracer.startSpan(
+          `${PostgresPlugin.component}.query`,
+          {
+            kind: SpanKind.CLIENT,
+            parent: plugin._tracer.getCurrentSpan() || undefined,
+            attributes: {
+              [AttributeNames.COMPONENT]: PostgresPlugin.component,
+              [AttributeNames.PG_HOST]: (this as any).connectionParameters.host,
+              [AttributeNames.PG_PORT]: (this as any).connectionParameters.port,
+            },
+          }
+        );
+
+        try {
+          if (typeof args[0] === 'string') {
+            span.setAttribute(AttributeNames.PG_TEXT, args[0]);
+            if (args[1] instanceof Array) {
+              span.setAttribute(AttributeNames.PG_VALUES, args[1]);
+              if (callbackProvided) {
+                args[2] = plugin._tracer.bind(args[2]);
+              }
+            } else {
+              if (callbackProvided) {
+                args[1] = plugin._tracer.bind(args[1]);
+              }
+            }
+          } else {
+            const config = args[0] as pgTypes.QueryConfig & {
+              callback?: PostgresCallback;
+            };
+            if (typeof config.name === 'string') {
+              span.setAttribute(AttributeNames.PG_PLAN, config.name);
+            } else {
+              if (typeof config.text === 'string') {
+                span.setAttribute(AttributeNames.PG_TEXT, config.text);
+              }
+              if (config.values instanceof Array) {
+                span.setAttribute(AttributeNames.PG_VALUES, config.values);
+              }
+            }
+
+            if (callbackProvided) {
+              if (typeof args[1] === 'function') {
+                args[1] = plugin._tracer.bind(args[1]);
+              } else if (typeof args[2] === 'function') {
+                args[2] = plugin._tracer.bind(args[2]);
+              }
+            } else if (
+              config.callback &&
+              typeof config.callback === 'function'
+            ) {
+              callbackProvided = true;
+              config.callback = plugin._tracer.bind(config.callback);
+            }
+          }
+        } catch (e) {
+          plugin._logger.warn(
+            `pg Plugin failed to trace query: error: ${e.message}`
+          );
+          const result = original.apply(this, arguments as any);
+          span.end();
+          return result;
+        }
+
+        const queryResult = original.apply(this, args as any);
+
+        // No callback was provided, return a promise instead (new as of [email protected])
+        if (!callbackProvided) {
+          const queryResultPromise = (queryResult as unknown) as Promise<
+            unknown
+          >;
+          return plugin._tracer.bind(
+            queryResultPromise
+              .then((result: any) => {
+                // Return a pass-along promise which ends the span and then goes to user's orig resolvers
+                return new Promise((resolve, _) => {
+                  span.end();
+                  resolve(result);
+                });
+              })
+              .catch((error: Error) => {
+                return new Promise((_, reject) => {
+                  span.end();
+                  reject(error);
+                });
+              })
+          );
+        }
+
+        // Else a callback was provided, so just return the result
+        span.end();
+        return queryResult;
+      };
+    };
+  }
+}
+
+const basedir = path.dirname(require.resolve('pg'));
+const version = require(path.join(basedir, '../', 'package.json')).version;
+export const plugin = new PostgresPlugin(PostgresPlugin.component, version);

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

@@ -0,0 +1,19 @@
+/*!
+ * 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 interface PostgresPluginOptions {}
+
+export type PostgresCallback = (err: Error, res: object) => unknown;

packages/opentelemetry-plugin-postgres/test/assertionUtils.ts

@@ -0,0 +1,90 @@
+/*!
+ * 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 {
+  SpanKind,
+  Attributes,
+  Event,
+  Span,
+  TimedEvent,
+} from '@opentelemetry/types';
+import * as assert from 'assert';
+import { PostgresPlugin } from '../src';
+import { ReadableSpan } from '@opentelemetry/tracing';
+import {
+  hrTimeToMilliseconds,
+  hrTimeToMicroseconds,
+} from '@opentelemetry/core';
+import { AttributeNames } from '../src/enums';
+
+export const assertSpan = (
+  span: ReadableSpan,
+  kind: SpanKind,
+  attributes: Attributes,
+  events: Event[]
+) => {
+  assert.strictEqual(span.spanContext.traceId.length, 32);
+  assert.strictEqual(span.spanContext.spanId.length, 16);
+  assert.strictEqual(span.kind, kind);
+
+  assert.strictEqual(
+    span.attributes[AttributeNames.COMPONENT],
+    PostgresPlugin.component
+  );
+  assert.ok(span.endTime);
+  assert.strictEqual(span.links.length, 0);
+
+  assert.ok(
+    hrTimeToMicroseconds(span.startTime) < hrTimeToMicroseconds(span.endTime)
+  );
+  assert.ok(hrTimeToMilliseconds(span.endTime) > 0);
+
+  // attributes
+  assert.strictEqual(
+    Object.keys(span.attributes).length,
+    Object.keys(attributes).length,
+    'Should contain same number of attributes'
+  );
+  Object.keys(span.attributes).forEach(attribute => {
+    assert.deepStrictEqual(span.attributes[attribute], attributes[attribute]);
+  });
+
+  // events
+  assert.strictEqual(
+    span.events.length,
+    events.length,
+    'Should contain same number of events'
+  );
+  span.events.forEach((_: TimedEvent, index: number) => {
+    assert.deepStrictEqual(span.events[index], events[index]);
+  });
+};
+
+// Check if sourceSpan was propagated to targetSpan
+export const assertPropagation = (
+  childSpan: ReadableSpan,
+  parentSpan: Span
+) => {
+  const targetSpanContext = childSpan.spanContext;
+  const sourceSpanContext = parentSpan.context();
+  assert.strictEqual(targetSpanContext.traceId, sourceSpanContext.traceId);
+  assert.strictEqual(childSpan.parentSpanId, sourceSpanContext.spanId);
+  assert.strictEqual(
+    targetSpanContext.traceFlags,
+    sourceSpanContext.traceFlags
+  );
+  assert.notStrictEqual(targetSpanContext.spanId, sourceSpanContext.spanId);
+};