feat(tracing): allow direct pg module to enable esbuild support (#9227)

Author: aldenquimby

packages/tracing-internal/src/node/integrations/postgres.ts

@@ -5,9 +5,15 @@ import { fill, isThenable, loadModule, logger } from '@sentry/utils';
 import type { LazyLoadedIntegration } from './lazy';
 import { shouldDisableAutoInstrumentation } from './utils/node-utils';
 
+type PgClientQuery = (
+  config: unknown,
+  values?: unknown,
+  callback?: (err: unknown, result: unknown) => void,
+) => void | Promise<unknown>;
+
 interface PgClient {
   prototype: {
-    query: () => void | Promise<unknown>;
+    query: PgClientQuery;
   };
 }
 
@@ -20,9 +26,23 @@ interface PgClientThis {
 
 interface PgOptions {
   usePgNative?: boolean;
+  /**
+   * Supply your postgres module directly, instead of having Sentry attempt automatic resolution.
+   * Use this if you (a) use a module that's not `pg`, or (b) use a bundler that breaks resolution (e.g. esbuild).
+   *
+   * Usage:
+   * ```
+   * import pg from 'pg';
+   *
+   * Sentry.init({
+   *   integrations: [new Sentry.Integrations.Postgres({ module: pg })],
+   * });
+   * ```
+   */
+  module?: PGModule;
 }
 
-type PGModule = { Client: PgClient; native: { Client: PgClient } };
+type PGModule = { Client: PgClient; native: { Client: PgClient } | null };
 
 /** Tracing integration for node-postgres package */
 export class Postgres implements LazyLoadedIntegration<PGModule> {
@@ -43,6 +63,7 @@ export class Postgres implements LazyLoadedIntegration<PGModule> {
   public constructor(options: PgOptions = {}) {
     this.name = Postgres.id;
     this._usePgNative = !!options.usePgNative;
+    this._module = options.module;
   }
 
   /** @inheritdoc */
@@ -66,21 +87,21 @@ export class Postgres implements LazyLoadedIntegration<PGModule> {
       return;
     }
 
-    if (this._usePgNative && !pkg.native?.Client) {
+    const Client = this._usePgNative ? pkg.native?.Client : pkg.Client;
+
+    if (!Client) {
       __DEBUG_BUILD__ && logger.error("Postgres Integration was unable to access 'pg-native' bindings.");
       return;
     }
 
-    const { Client } = this._usePgNative ? pkg.native : pkg;
-
     /**
      * function (query, callback) => void
      * function (query, params, callback) => void
      * function (query) => Promise
      * function (query, params) => Promise
      * function (pg.Cursor) => pg.Cursor
      */
-    fill(Client.prototype, 'query', function (orig: () => void | Promise<unknown>) {
+    fill(Client.prototype, 'query', function (orig: PgClientQuery) {
       return function (this: PgClientThis, config: unknown, values: unknown, callback: unknown) {
         const scope = getCurrentHub().getScope();
         const parentSpan = scope.getSpan();

packages/tracing/test/integrations/node/postgres.test.ts

@@ -1,16 +1,17 @@
 /* eslint-disable deprecation/deprecation */
 /* eslint-disable @typescript-eslint/unbound-method */
 import { Hub, Scope } from '@sentry/core';
-import { logger } from '@sentry/utils';
+import { loadModule, logger } from '@sentry/utils';
+import pg from 'pg';
 
 import { Integrations, Span } from '../../../src';
 import { getTestClient } from '../../testutils';
 
 class PgClient {
   // https://node-postgres.com/api/client#clientquery
-  public query(_text: unknown, values: unknown, callback?: () => void) {
+  public query(_text: unknown, values: unknown, callback?: (err: unknown, result: unknown) => void) {
     if (typeof callback === 'function') {
-      callback();
+      callback(null, null);
       return;
     }
 
@@ -25,25 +26,28 @@ class PgClient {
 
 // Jest mocks get hoisted. vars starting with `mock` are hoisted before imports.
 /* eslint-disable no-var */
-var mockClient = PgClient;
+var mockModule = {
+  Client: PgClient,
+  native: {
+    Client: PgClient,
+  },
+};
 
 // mock for 'pg' / 'pg-native' package
 jest.mock('@sentry/utils', () => {
   const actual = jest.requireActual('@sentry/utils');
   return {
     ...actual,
-    loadModule() {
-      return {
-        Client: mockClient,
-        native: {
-          Client: mockClient,
-        },
-      };
-    },
+    loadModule: jest.fn(() => mockModule),
   };
 });
 
 describe('setupOnce', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    jest.resetAllMocks();
+  });
+
   ['pg', 'pg-native'].forEach(pgApi => {
     const Client: PgClient = new PgClient();
     let scope = new Scope();
@@ -127,4 +131,23 @@ describe('setupOnce', () => {
 
     expect(loggerLogSpy).toBeCalledWith('Postgres Integration is skipped because of instrumenter configuration.');
   });
+
+  it('does not attempt resolution when module is passed directly', async () => {
+    const scope = new Scope();
+    jest.spyOn(scope, 'getSpan').mockReturnValueOnce(new Span());
+
+    new Integrations.Postgres({ module: mockModule }).setupOnce(
+      () => undefined,
+      () => new Hub(undefined, scope),
+    );
+
+    await new PgClient().query('SELECT NOW()', null);
+
+    expect(loadModule).not.toBeCalled();
+    expect(scope.getSpan).toBeCalled();
+  });
+
+  it('has valid module type', () => {
+    expect(() => new Integrations.Postgres({ module: pg })).not.toThrow();
+  });
 });