feat(plugin): pg-pool plugin implementation (#501)

* feat: pg-pool plugin implementation

* feat: pg-pool plugin implementation

* fix: linting

* fix: add attributes for span & add tests for pool.query()

* fix: add span.setStatus

* chore: address comments

* fix: linting

Co-authored-by: Mayur Kale <[email protected]>

packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/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-POOL specific -- not specified by spec
+  IDLE_TIMEOUT_MILLIS = 'idle.timeout.millis',
+  MAX_CLIENT = 'max',
+}

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

@@ -0,0 +1,17 @@
+/*!
+ * 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 * from './pg-pool';

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

@@ -13,3 +13,120 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
+import { BasePlugin } from '@opentelemetry/core';
+import { CanonicalCode, SpanKind } from '@opentelemetry/types';
+import { AttributeNames } from './enums';
+import * as shimmer from 'shimmer';
+import * as pgPoolTypes from 'pg-pool';
+import {
+  PostgresPoolPluginOptions,
+  PgPoolCallback,
+  PgPoolExtended,
+} from './types';
+import * as utils from './utils';
+
+export class PostgresPoolPlugin extends BasePlugin<typeof pgPoolTypes> {
+  protected _config: PostgresPoolPluginOptions;
+
+  static readonly COMPONENT = 'pg-pool';
+  static readonly DB_TYPE = 'sql';
+
+  readonly supportedVersions = ['2.*'];
+
+  constructor(readonly moduleName: string) {
+    super();
+    this._config = {};
+  }
+
+  protected patch(): typeof pgPoolTypes {
+    shimmer.wrap(
+      this._moduleExports.prototype,
+      'connect',
+      this._getPoolConnectPatch() as never
+    );
+
+    return this._moduleExports;
+  }
+
+  protected unpatch(): void {
+    shimmer.unwrap(this._moduleExports.prototype, 'connect');
+  }
+
+  private _getPoolConnectPatch() {
+    const plugin = this;
+    return (originalConnect: typeof pgPoolTypes.prototype.connect) => {
+      plugin._logger.debug(
+        `Patching ${PostgresPoolPlugin.COMPONENT}.prototype.connect`
+      );
+      return function connect(this: PgPoolExtended, callback?: PgPoolCallback) {
+        const jdbcString = utils.getJDBCString(this.options);
+        // setup span
+        const span = plugin._tracer.startSpan(
+          `${PostgresPoolPlugin.COMPONENT}.connect`,
+          {
+            kind: SpanKind.CLIENT,
+            parent: plugin._tracer.getCurrentSpan() || undefined,
+            attributes: {
+              [AttributeNames.COMPONENT]: PostgresPoolPlugin.COMPONENT, // required
+              [AttributeNames.DB_TYPE]: PostgresPoolPlugin.DB_TYPE, // required
+              [AttributeNames.DB_INSTANCE]: this.options.database, // required
+              [AttributeNames.PEER_HOSTNAME]: this.options.host, // required
+              [AttributeNames.PEER_ADDRESS]: jdbcString, // required
+              [AttributeNames.PEER_PORT]: this.options.port,
+              [AttributeNames.DB_USER]: this.options.user,
+              [AttributeNames.IDLE_TIMEOUT_MILLIS]: this.options
+                .idleTimeoutMillis,
+              [AttributeNames.MAX_CLIENT]: this.options.maxClient,
+            },
+          }
+        );
+
+        if (callback) {
+          const parentSpan = plugin._tracer.getCurrentSpan();
+          callback = utils.patchCallback(span, callback) as PgPoolCallback;
+          // If a parent span exists, bind the callback
+          if (parentSpan) {
+            callback = plugin._tracer.bind(callback);
+          }
+        }
+
+        const connectResult: unknown = originalConnect.call(
+          this,
+          callback as never
+        );
+
+        // No callback was provided, return a promise instead
+        if (connectResult instanceof Promise) {
+          const connectResultPromise = connectResult as Promise<unknown>;
+          return plugin._tracer.bind(
+            connectResultPromise
+              .then((result: any) => {
+                // Resturn 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,
+                    message: error.message,
+                  });
+                  span.end();
+                  reject(error);
+                });
+              })
+          );
+        }
+
+        // Else a callback was provided, so just return the result
+        return connectResult;
+      };
+    };
+  }
+}
+
+export const plugin = new PostgresPoolPlugin(PostgresPoolPlugin.COMPONENT);

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

@@ -0,0 +1,39 @@
+/*!
+ * 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';
+import * as pgPoolTypes from 'pg-pool';
+
+export interface PostgresPoolPluginOptions {}
+
+export type PgPoolCallback = (
+  err: Error,
+  client: any,
+  done: (release?: any) => void
+) => void;
+
+export interface PgPoolOptionsParams {
+  database: string;
+  host: string;
+  port: number;
+  user: string;
+  idleTimeoutMillis: number; // the minimum amount of time that an object may sit idle in the pool before it is eligible for eviction due to idle time
+  maxClient: number; // maximum size of the pool
+}
+
+export interface PgPoolExtended extends pgPoolTypes<pgTypes.Client> {
+  options: PgPoolOptionsParams;
+}

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

@@ -0,0 +1,45 @@
+/*!
+ * 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 { Span, CanonicalCode } from '@opentelemetry/types';
+import { PgPoolOptionsParams, PgPoolCallback, PgPoolExtended } from './types';
+
+export function getJDBCString(params: PgPoolOptionsParams) {
+  const host = params.host || 'localhost'; // postgres defaults to localhost
+  const port = params.port || 5432; // postgres defaults to port 5432
+  const database = params.database || '';
+  return `jdbc:postgresql://${host}:${port}/${database}`;
+}
+
+export function patchCallback(span: Span, cb: PgPoolCallback): PgPoolCallback {
+  return function patchedCallback(
+    this: PgPoolExtended,
+    err: Error,
+    res: object,
+    done: any
+  ) {
+    if (err) {
+      span.setStatus({
+        code: CanonicalCode.UNKNOWN,
+        message: err.message,
+      });
+    } else if (res) {
+      span.setStatus({ code: CanonicalCode.OK });
+    }
+    span.end();
+    cb.call(this, err, res, done);
+  };
+}

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

@@ -0,0 +1,79 @@
+/*!
+ * 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 { ReadableSpan } from '@opentelemetry/tracing';
+import {
+  hrTimeToMilliseconds,
+  hrTimeToMicroseconds,
+} from '@opentelemetry/core';
+
+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);
+
+  // check all the AttributeNames fields
+  Object.keys(span.attributes).forEach(key => {
+    assert.deepStrictEqual(span.attributes[key], attributes[key]);
+  });
+
+  assert.ok(span.endTime);
+  assert.strictEqual(span.links.length, 0);
+
+  assert.ok(
+    hrTimeToMicroseconds(span.startTime) < hrTimeToMicroseconds(span.endTime)
+  );
+  assert.ok(hrTimeToMilliseconds(span.endTime) > 0);
+
+  // 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);
+};