From 798ed70f8a32fb874d7db84c7261ac6083928d7a Mon Sep 17 00:00:00 2001
From: Onur Temizkan <onur@narval.co.uk>
Date: Mon, 16 Aug 2021 10:21:01 +0000
Subject: [PATCH] Add `pg-native` support to Postgres integration.

---
 packages/tracing/src/integrations/postgres.ts | 21 +++-
 .../test/integrations/postgres.test.ts        | 97 +++++++++++++++++++
 2 files changed, 116 insertions(+), 2 deletions(-)
 create mode 100644 packages/tracing/test/integrations/postgres.test.ts

diff --git a/packages/tracing/src/integrations/postgres.ts b/packages/tracing/src/integrations/postgres.ts
index 1a8e4fad2ab7..b345f8a2876b 100644
--- a/packages/tracing/src/integrations/postgres.ts
+++ b/packages/tracing/src/integrations/postgres.ts
@@ -8,6 +8,10 @@ interface PgClient {
   };
 }
 
+interface PgOptions {
+  usePgNative?: boolean;
+}
+
 /** Tracing integration for node-postgres package */
 export class Postgres implements Integration {
   /**
@@ -20,17 +24,30 @@ export class Postgres implements Integration {
    */
   public name: string = Postgres.id;
 
+  private _usePgNative: boolean;
+
+  public constructor(options: PgOptions = {}) {
+    this._usePgNative = !!options.usePgNative;
+  }
+
   /**
    * @inheritDoc
    */
   public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
-    const pkg = loadModule<{ Client: PgClient }>('pg');
+    const pkg = loadModule<{ Client: PgClient; native: { Client: PgClient } }>('pg');
 
     if (!pkg) {
       logger.error('Postgres Integration was unable to require `pg` package.');
       return;
     }
 
+    if (this._usePgNative && !pkg.native?.Client) {
+      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
@@ -38,7 +55,7 @@ export class Postgres implements Integration {
      * function (query, params) => Promise
      * function (pg.Cursor) => pg.Cursor
      */
-    fill(pkg.Client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
+    fill(Client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
       return function(this: unknown, config: unknown, values: unknown, callback: unknown) {
         const scope = getCurrentHub().getScope();
         const parentSpan = scope?.getSpan();
diff --git a/packages/tracing/test/integrations/postgres.test.ts b/packages/tracing/test/integrations/postgres.test.ts
new file mode 100644
index 000000000000..98bb45357783
--- /dev/null
+++ b/packages/tracing/test/integrations/postgres.test.ts
@@ -0,0 +1,97 @@
+/* eslint-disable @typescript-eslint/unbound-method */
+import { Hub, Scope } from '@sentry/hub';
+
+import { Postgres } from '../../src/integrations/postgres';
+import { Span } from '../../src/span';
+
+class PgClient {
+  // https://node-postgres.com/api/client#clientquery
+  public query(_text: unknown, values: unknown, callback?: () => void) {
+    if (typeof callback === 'function') {
+      callback();
+      return;
+    }
+
+    if (typeof values === 'function') {
+      values();
+      return;
+    }
+
+    return Promise.resolve();
+  }
+}
+
+// mock for 'pg' / 'pg-native' package
+jest.mock('@sentry/utils', () => {
+  const actual = jest.requireActual('@sentry/utils');
+  return {
+    ...actual,
+    loadModule() {
+      return {
+        Client: PgClient,
+        native: {
+          Client: PgClient,
+        },
+      };
+    },
+  };
+});
+
+describe('setupOnce', () => {
+  ['pg', 'pg-native'].forEach(pgApi => {
+    const Client: PgClient = new PgClient();
+    let scope = new Scope();
+    let parentSpan: Span;
+    let childSpan: Span;
+
+    beforeAll(() => {
+      (pgApi === 'pg' ? new Postgres() : new Postgres({ usePgNative: true })).setupOnce(
+        () => undefined,
+        () => new Hub(undefined, scope),
+      );
+    });
+
+    beforeEach(() => {
+      scope = new Scope();
+      parentSpan = new Span();
+      childSpan = parentSpan.startChild();
+      jest.spyOn(scope, 'getSpan').mockReturnValueOnce(parentSpan);
+      jest.spyOn(parentSpan, 'startChild').mockReturnValueOnce(childSpan);
+      jest.spyOn(childSpan, 'finish');
+    });
+
+    it(`should wrap ${pgApi}'s query method accepting callback as the last argument`, done => {
+      Client.query('SELECT NOW()', {}, function() {
+        expect(scope.getSpan).toBeCalled();
+        expect(parentSpan.startChild).toBeCalledWith({
+          description: 'SELECT NOW()',
+          op: 'db',
+        });
+        expect(childSpan.finish).toBeCalled();
+        done();
+      }) as void;
+    });
+
+    it(`should wrap ${pgApi}'s query method accepting callback as the second argument`, done => {
+      Client.query('SELECT NOW()', function() {
+        expect(scope.getSpan).toBeCalled();
+        expect(parentSpan.startChild).toBeCalledWith({
+          description: 'SELECT NOW()',
+          op: 'db',
+        });
+        expect(childSpan.finish).toBeCalled();
+        done();
+      }) as void;
+    });
+
+    it(`should wrap ${pgApi}'s query method accepting no callback as the last argument but returning promise`, async () => {
+      await Client.query('SELECT NOW()', null);
+      expect(scope.getSpan).toBeCalled();
+      expect(parentSpan.startChild).toBeCalledWith({
+        description: 'SELECT NOW()',
+        op: 'db',
+      });
+      expect(childSpan.finish).toBeCalled();
+    });
+  });
+});