feat(tracing): track Cursor.toArray method for Mongo

packages/tracing/src/integrations/node/mongo.ts

@@ -1,7 +1,9 @@
 import { Hub } from '@sentry/hub';
-import { EventProcessor, Integration, SpanContext } from '@sentry/types';
+import { EventProcessor, Integration, Span, SpanContext } from '@sentry/types';
 import { fill, isThenable, loadModule, logger } from '@sentry/utils';
 
+type Callback = (...args: unknown[]) => void;
+
 // 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" | ... )
 // and not just a string[])
@@ -82,12 +84,97 @@ interface MongoCollection {
   };
 }
 
+export interface MongoCursor {
+  // count,
+  // explain,
+  // hasNext,
+  // next,
+  // tryNext,
+  // forEach,
+  // close,
+
+  toArray(): Promise<unknown[]>;
+  toArray(callback: Callback): void;
+}
+
 interface MongoOptions {
   operations?: Operation[];
   describeOperations?: boolean | Operation[];
   useMongoose?: boolean;
 }
 
+function measurePromiseOrCb(
+  orig: (...args: unknown[]) => void | Promise<unknown>,
+  args: unknown[],
+  span: Span | undefined,
+  cb: Callback | undefined,
+): unknown {
+  if (cb) {
+    return orig(...args.slice(0, -1), function (...cbArgs: unknown[]) {
+      span?.finish();
+      cb(...cbArgs);
+    });
+  }
+
+  const maybePromise = orig(...args) as Promise<unknown>;
+
+  if (isThenable(maybePromise)) {
+    return maybePromise.then((res: unknown) => {
+      span?.finish();
+      return res;
+    });
+  } else {
+    span?.finish();
+    return maybePromise;
+  }
+}
+
+function instrumentCursor(cursor: MongoCursor, parentSpan: Span | undefined): MongoCursor {
+  fill(cursor, 'toArray', (orig: () => void | Promise<unknown>) => {
+    return function (this: MongoCursor, ...args: unknown[]) {
+      const lastArg = args[args.length - 1];
+
+      const span = parentSpan?.startChild({
+        op: 'db',
+        description: 'Cursor.toArray',
+      });
+
+      return measurePromiseOrCb(
+        orig.bind(this),
+        args,
+        span,
+        typeof lastArg === 'function' ? (lastArg as Callback) : undefined,
+      );
+    };
+  });
+
+  return cursor;
+}
+
+const serializeArg = (arg: unknown): string => {
+  if (typeof arg === 'function') {
+    return arg.name || '<anonymous>';
+  }
+
+  if (typeof arg === 'string') {
+    return arg;
+  }
+
+  return JSON.stringify(arg);
+};
+
+function getOperationCallbackFromArgsOrUndefined(operation: string, args: unknown[]): Callback | undefined {
+  const lastArg = args[args.length - 1];
+
+  // Check if the operation was passed a callback. (mapReduce requires a different check, as
+  // its (non-callback) arguments can also be functions.)
+  if (typeof lastArg !== 'function' || (operation === 'mapReduce' && args.length === 2)) {
+    return undefined;
+  }
+
+  return lastArg as Callback;
+}
+
 /** Tracing integration for mongo package */
 export class Mongo implements Integration {
   /**
@@ -141,36 +228,27 @@ export class Mongo implements Integration {
   private _patchOperation(collection: MongoCollection, operation: Operation, getCurrentHub: () => Hub): void {
     if (!(operation in collection.prototype)) return;
 
-    const getSpanContext = this._getSpanContextFromOperationArguments.bind(this);
+    const getSpanContext: Mongo['_getSpanContextFromOperationArguments'] =
+      this._getSpanContextFromOperationArguments.bind(this);
 
     fill(collection.prototype, operation, function (orig: () => void | Promise<unknown>) {
-      return function (this: unknown, ...args: unknown[]) {
-        const lastArg = args[args.length - 1];
+      return function (this: MongoCollection, ...args: unknown[]) {
         const scope = getCurrentHub().getScope();
         const parentSpan = scope?.getSpan();
+        const spanContext = getSpanContext(this, operation, args);
 
-        // Check if the operation was passed a callback. (mapReduce requires a different check, as
-        // 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));
-          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 maybePromise;
-          }
+        // `find` returns a cursor, wrap cursor methods
+        if (operation === 'find') {
+          const cursor = orig.call(this, ...args); // => FindCursor
+          return instrumentCursor(cursor, parentSpan?.startChild(spanContext));
         }
 
-        const span = parentSpan?.startChild(getSpanContext(this, operation, args.slice(0, -1)));
-        return orig.call(this, ...args.slice(0, -1), function (err: Error, result: unknown) {
-          span?.finish();
-          lastArg(err, result);
-        });
+        return measurePromiseOrCb(
+          orig.bind(this),
+          args,
+          parentSpan?.startChild(spanContext),
+          getOperationCallbackFromArgsOrUndefined(operation, args),
+        );
       };
     });
   }
@@ -206,15 +284,8 @@ export class Mongo implements Integration {
     }
 
     try {
-      // Special case for `mapReduce`, as the only one accepting functions as arguments.
-      if (operation === 'mapReduce') {
-        const [map, reduce] = args as { name?: string }[];
-        data[signature[0]] = typeof map === 'string' ? map : map.name || '<anonymous>';
-        data[signature[1]] = typeof reduce === 'string' ? reduce : reduce.name || '<anonymous>';
-      } else {
-        for (let i = 0; i < signature.length; i++) {
-          data[signature[i]] = JSON.stringify(args[i]);
-        }
+      for (let i = 0; i < signature.length; i++) {
+        data[signature[i]] = serializeArg(args[i]);
       }
     } catch (_oO) {
       // no-empty

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

@@ -1,9 +1,12 @@
 /* eslint-disable @typescript-eslint/unbound-method */
 import { Hub, Scope } from '@sentry/hub';
 
-import { Mongo } from '../../../src/integrations/node/mongo';
+import { Mongo, MongoCursor } from '../../../src/integrations/node/mongo';
 import { Span } from '../../../src/span';
 
+type MapReduceArg = () => void | string;
+type Callback<T = any> = (error?: any, result?: T) => void;
+
 class Collection {
   public collectionName: string = 'mockedCollectionName';
   public dbName: string = 'mockedDbName';
@@ -17,10 +20,40 @@ class Collection {
     }
     return Promise.resolve();
   }
+
   // Method that has no callback as last argument, and doesnt return promise.
   public initializeOrderedBulkOp() {
     return {};
   }
+
+  mapReduce(map: MapReduceArg, reduce: MapReduceArg): Promise<unknown>;
+  mapReduce(map: MapReduceArg, reduce: MapReduceArg, callback: Callback<unknown>): void;
+  mapReduce(map: MapReduceArg, reduce: MapReduceArg, options: Record<string, any>, callback?: Callback<unknown>): void;
+
+  mapReduce(...args: any[]): Promise<unknown> | void {
+    const lastArg = args[args.length - 1];
+
+    // (map, reduce) => promise
+    // (map, reduce, options) => promise
+    if (args.length === 2 || (args.length === 3 && typeof lastArg !== 'function')) {
+      return Promise.resolve();
+    }
+
+    // (map, reduce, cb) => void
+    // (map, reduce, options, cb) => void
+    if (typeof lastArg === 'function') {
+      lastArg();
+      return;
+    }
+  }
+
+  find(_query: any): MongoCursor {
+    return {
+      async toArray(): Promise<unknown[]> {
+        return [];
+      },
+    };
+  }
 }
 
 jest.mock('@sentry/utils', () => {
@@ -47,7 +80,7 @@ describe('patchOperation()', () => {
 
   beforeAll(() => {
     new Mongo({
-      operations: ['insertOne', 'initializeOrderedBulkOp'],
+      operations: ['insertOne', 'mapReduce', 'find', 'initializeOrderedBulkOp'],
     }).setupOnce(
       () => undefined,
       () => new Hub(undefined, scope),
@@ -63,6 +96,19 @@ describe('patchOperation()', () => {
     jest.spyOn(childSpan, 'finish');
   });
 
+  it('should call orig methods with origin context (this) and params', () => {
+    const spy = jest.spyOn(collection, 'insertOne');
+
+    const cb = function () {};
+    const options = {};
+
+    collection.insertOne(doc, options, cb) as void;
+    expect(spy.mock.instances[0]).toBe(collection);
+    expect(spy).toBeCalledWith(doc, options, cb);
+
+    spy.mockRestore();
+  });
+
   it('should wrap method accepting callback as the last argument', done => {
     collection.insertOne(doc, {}, function () {
       expect(scope.getSpan).toBeCalled();
@@ -111,4 +157,114 @@ describe('patchOperation()', () => {
     });
     expect(childSpan.finish).toBeCalled();
   });
+
+  describe('mapReduce operation', () => {
+    describe('variable arguments', () => {
+      const expectToBeTracked = () => {
+        expect(scope.getSpan).toBeCalled();
+        expect(parentSpan.startChild).toBeCalledWith({
+          data: {
+            collectionName: 'mockedCollectionName',
+            dbName: 'mockedDbName',
+            namespace: 'mockedNamespace',
+            map: '<anonymous>',
+            reduce: '<anonymous>',
+          },
+          op: `db`,
+          description: 'mapReduce',
+        });
+        expect(childSpan.finish).toBeCalled();
+      };
+
+      it('should work when (map, reduce, cb)', done => {
+        collection.mapReduce(
+          () => {},
+          () => {},
+          function () {
+            expectToBeTracked();
+            done();
+          },
+        );
+      });
+
+      it('should work when (map, reduce, options, cb)', done => {
+        collection.mapReduce(
+          () => {},
+          () => {},
+          {},
+          function () {
+            expectToBeTracked();
+            done();
+          },
+        );
+      });
+
+      it('should work when (map, reduce) => Promise', async () => {
+        await collection.mapReduce(
+          () => {},
+          () => {},
+        );
+        expectToBeTracked();
+      });
+
+      it('should work when (map, reduce, options) => Promise', async () => {
+        await collection.mapReduce(
+          () => {},
+          () => {},
+          {},
+        );
+        expectToBeTracked();
+      });
+    });
+
+    it('Should store function names', async () => {
+      await collection.mapReduce(
+        function TestMapFn() {},
+        function TestReduceFn() {},
+      );
+      expect(scope.getSpan).toBeCalled();
+      expect(parentSpan.startChild).toBeCalledWith({
+        data: {
+          collectionName: 'mockedCollectionName',
+          dbName: 'mockedDbName',
+          namespace: 'mockedNamespace',
+          map: 'TestMapFn',
+          reduce: 'TestReduceFn',
+        },
+        op: `db`,
+        description: 'mapReduce',
+      });
+      expect(childSpan.finish).toBeCalled();
+    });
+  });
+
+  describe('Cursor', () => {
+    it('Should instrument cursor', async () => {
+      const cursorSpan = new Span();
+
+      jest.spyOn(cursorSpan, 'finish');
+      jest.spyOn(childSpan, 'startChild').mockReturnValueOnce(cursorSpan);
+
+      const cursor = collection.find({ _id: '1234567' });
+      await cursor.toArray();
+
+      expect(parentSpan.startChild).toBeCalledWith({
+        data: {
+          collectionName: 'mockedCollectionName',
+          dbName: 'mockedDbName',
+          namespace: 'mockedNamespace',
+          query: '{"_id":"1234567"}',
+        },
+        op: `db`,
+        description: 'find',
+      });
+
+      expect(childSpan.startChild).toBeCalledWith({
+        op: `db`,
+        description: 'Cursor.toArray',
+      });
+
+      expect(cursorSpan.finish).toBeCalled();
+    });
+  });
 });