Patches pg-pool to trace with dd-trace

pgPoolTracing.ts

import tracer, { Span } from "dd-trace";
import pg from "pg";

/**
 * Patches pg.Pool to add tracing.
 */
export function patchPgPool(serviceName: string) {
  const name = "pg-pool.connect";
  const service = `${serviceName}-postgres`;

  class TracedPool extends pg.Pool {
    connect(): Promise<pg.PoolClient>;

    connect(
      callback?: Parameters<typeof pg.Pool.prototype.connect>[0],
    ): Promise<pg.PoolClient> | void {
      const { waitingCount, idleCount } = this;
      // Pool.connect has two overloads: callback and async
      if (callback) {
        return tracer.trace(name, { service }, (span, spanDone) => {
          if (span) this.updateSpanBefore(span, waitingCount, idleCount);
          return super.connect((err, client, done) => {
            if (span && spanDone) {
              this.updateSpanAfter(span, client);
              spanDone(err);
            }
            return callback(err, client, done);
          });
        });
      } else {
        return tracer.trace(name, { service }, async (span) => {
          if (span) this.updateSpanBefore(span, waitingCount, idleCount);
          const client = await super.connect();
          if (span) this.updateSpanAfter(span, client);
          return client;
        });
      }
    }

    updateSpanBefore(span: Span, waitingCount: number, idleCount: number) {
      span.addTags({
        component: "pg",
        "db.type": "postgres",
        "db.poolWaitingCount": waitingCount,
        "db.poolIdleCount": idleCount,
      });
    }

    updateSpanAfter(span: Span, client: pg.PoolClient) {
      span.addTags({
        "db.pid": (client as any).processID, // internal property
      });
    }
  }

  pg.Pool = TracedPool;
}