fix: Make sure that mongo method is thenable before calling it (#3173)

* fix: Make sure that mongo method is thenable before calling it

* test: Add test suite for MongoDB tracing integration

Author: kamilogorek

packages/tracing/src/integrations/mongo.ts

@@ -1,6 +1,6 @@
 import { Hub } from '@sentry/hub';
 import { EventProcessor, Integration, SpanContext } from '@sentry/types';
-import { dynamicRequire, fill, logger } from '@sentry/utils';
+import { dynamicRequire, fill, isThenable, logger } from '@sentry/utils';
 
 // This allows us to use the same array for both defaults options and the type itself.
 // (note `as const` at the end to make it a union of string literal types (i.e. "a" | "b" | ... )
@@ -152,10 +152,17 @@ export class Mongo implements Integration {
         // its (non-callback) arguments can also be functions.)
         if (typeof lastArg !== 'function' || (operation === 'mapReduce' && args.length === 2)) {
           const span = parentSpan?.startChild(getSpanContext(this, operation, args));
-          return (orig.call(this, ...args) as Promise<unknown>).then((res: unknown) => {
+          const maybePromise = orig.call(this, ...args) as Promise<unknown>;
+
+          if (isThenable(maybePromise)) {
+            return maybePromise.then((res: unknown) => {
+              span?.finish();
+              return res;
+            });
+          } else {
             span?.finish();
-            return res;
-          });
+            return maybePromise;
+          }
         }
 
         const span = parentSpan?.startChild(getSpanContext(this, operation, args.slice(0, -1)));

packages/tracing/test/integrations/mongo.test.ts

@@ -0,0 +1,114 @@
+/* eslint-disable @typescript-eslint/unbound-method */
+import { Hub, Scope } from '@sentry/hub';
+
+import { Mongo } from '../../src/integrations/mongo';
+import { Span } from '../../src/span';
+
+class Collection {
+  public collectionName: string = 'mockedCollectionName';
+  public dbName: string = 'mockedDbName';
+  public namespace: string = 'mockedNamespace';
+
+  // Method that can have a callback as last argument, or return a promise otherwise.
+  insertOne(_doc: unknown, _options: unknown, callback?: () => void) {
+    if (typeof callback === 'function') {
+      callback();
+      return;
+    }
+    return Promise.resolve();
+  }
+  // Method that has no callback as last argument, and doesnt return promise.
+  initializeOrderedBulkOp() {
+    return {};
+  }
+}
+
+jest.mock('@sentry/utils', () => {
+  const actual = jest.requireActual('@sentry/utils');
+  return {
+    ...actual,
+    dynamicRequire() {
+      return {
+        Collection,
+      };
+    },
+  };
+});
+
+describe('patchOperation()', () => {
+  const doc = {
+    name: 'PickleRick',
+    answer: 42,
+  };
+  const collection: Collection = new Collection();
+  let scope = new Scope();
+  let parentSpan: Span;
+  let childSpan: Span;
+
+  beforeAll(() => {
+    new Mongo({
+      operations: ['insertOne', 'initializeOrderedBulkOp'],
+    }).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 method accepting callback as the last argument', done => {
+    collection.insertOne(doc, {}, function() {
+      expect(scope.getSpan).toBeCalled();
+      expect(parentSpan.startChild).toBeCalledWith({
+        data: {
+          collectionName: 'mockedCollectionName',
+          dbName: 'mockedDbName',
+          doc: JSON.stringify(doc),
+          namespace: 'mockedNamespace',
+        },
+        op: `db`,
+        description: 'insertOne',
+      });
+      expect(childSpan.finish).toBeCalled();
+      done();
+    }) as void;
+  });
+
+  it('should wrap method accepting no callback as the last argument but returning promise', async () => {
+    await collection.insertOne(doc, {});
+    expect(scope.getSpan).toBeCalled();
+    expect(parentSpan.startChild).toBeCalledWith({
+      data: {
+        collectionName: 'mockedCollectionName',
+        dbName: 'mockedDbName',
+        doc: JSON.stringify(doc),
+        namespace: 'mockedNamespace',
+      },
+      op: `db`,
+      description: 'insertOne',
+    });
+    expect(childSpan.finish).toBeCalled();
+  });
+
+  it('should wrap method accepting no callback as the last argument and not returning promise', () => {
+    collection.initializeOrderedBulkOp();
+    expect(scope.getSpan).toBeCalled();
+    expect(parentSpan.startChild).toBeCalledWith({
+      data: {
+        collectionName: 'mockedCollectionName',
+        dbName: 'mockedDbName',
+        namespace: 'mockedNamespace',
+      },
+      op: `db`,
+      description: 'initializeOrderedBulkOp',
+    });
+    expect(childSpan.finish).toBeCalled();
+  });
+});