fix: Correctly handle pg.Cursor in pg query method (#3567)

getsentry · May 19, 2021 · c3b1bb5 · c3b1bb5
1 parent 95fe018
commit c3b1bb5
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/packages/tracing/src/integrations/postgres.ts b/packages/tracing/src/integrations/postgres.ts
@@ -1,6 +1,6 @@
 import { Hub } from '@sentry/hub';
 import { EventProcessor, Integration } from '@sentry/types';
-import { fill, loadModule, logger } from '@sentry/utils';
+import { fill, isThenable, loadModule, logger } from '@sentry/utils';
 
 interface PgClient {
   prototype: {
@@ -36,6 +36,7 @@ export class Postgres implements Integration {
      * function (query, params, callback) => void
      * function (query) => Promise
      * function (query, params) => Promise
+     * function (pg.Cursor) => pg.Cursor
      */
     fill(pkg.Client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
       return function(this: unknown, config: unknown, values: unknown, callback: unknown) {
@@ -60,10 +61,17 @@ export class Postgres implements Integration {
           });
         }
 
-        return (orig.call(this, config, values) as Promise<unknown>).then((res: unknown) => {
-          span?.finish();
-          return res;
-        });
+        const rv = typeof values !== 'undefined' ? orig.call(this, config, values) : orig.call(this, config);
+
+        if (isThenable(rv)) {
+          return (rv as Promise<unknown>).then((res: unknown) => {
+            span?.finish();
+            return res;
+          });
+        }
+
+        span?.finish();
+        return rv;
       };
     });
   }