fix(shell-api): handle errors in cursor.map/.forEach MONGOSH-703 (#816)

Handle exceptions from callbacks passed to `cursor.map()` and
`cursor.forEach()`.

Since `cursor.forEach()` returns a Promise, use that to propagate
the error to the user.

For `cursor.map()`, this is a bit trickier, since the method itself
does not return any results to the user. Therefore, we add checks
before and after accessing cursor methods that read data which forward
the exception if one has occurred.

Both of these feel like situations that the driver should take care of
explicitly, so NODE tickets have been opened for both cases.

(One alternative would have been to add async context/Node.js domain
propagation to these methods. However, I decided against that, because
that solution would be a) Node.js-specific, and b) because it would
still lead to undesirable results, like the Promise returned by
`.forEach()` never resolving.)

Author: addaleax

packages/java-shell/src/test/resources/cursor/forEach.expected.txt

@@ -1,2 +1 @@
-null
-[ { "_id": 1, "name": "Vasya" }, { "_id": 2, "name": "Petya" }, { "_id": 3, "name": "Lyusya" } ]
+[ { "_id": 1, "name": "Vasya" }, { "_id": 2, "name": "Petya" }, { "_id": 3, "name": "Lyusya" } ]

packages/shell-api/src/abstract-cursor.ts

@@ -24,12 +24,37 @@ export abstract class AbstractCursor extends ShellApiClass {
   abstract _cursor: ServiceProviderAggregationCursor | ServiceProviderCursor;
   _currentIterationResult: CursorIterationResult | null = null;
   _batchSize: number | null = null;
+  _mapError: Error | null = null;
 
   constructor(mongo: Mongo) {
     super();
     this._mongo = mongo;
   }
 
+  // Wrap a function with checks before and after that verify whether a .map()
+  // callback has resulted in an exception. Such an error would otherwise result
+  // in an uncaught exception, bringing the whole process down.
+  // The downside to this is that errors will not actually be visible until
+  // the caller tries to interact with this cursor in a way that triggers
+  // these checks. Since that is also the behavior for errors coming from the
+  // database server, it makes sense to match that.
+  // Ideally, this kind of code could be lifted into the driver (NODE-3231 and
+  // NODE-3232 are the tickets for that).
+  async _withCheckMapError<Ret>(fn: () => Ret): Promise<Ret> {
+    if (this._mapError) {
+      // If an error has already occurred, we don't want to call the function
+      // at all.
+      throw this._mapError;
+    }
+    const ret = await fn();
+    if (this._mapError) {
+      // If an error occurred during the function, we don't want to forward its
+      // results.
+      throw this._mapError;
+    }
+    return ret;
+  }
+
   /**
    * Internal method to determine what is printed for this class.
    */
@@ -39,7 +64,7 @@ export abstract class AbstractCursor extends ShellApiClass {
 
   async _it(): Promise<CursorIterationResult> {
     const results = this._currentIterationResult = new CursorIterationResult();
-    await iterate(results, this._cursor, this._batchSize ?? await this._mongo._batchSize());
+    await iterate(results, this, this._batchSize ?? await this._mongo._batchSize());
     results.cursorHasMore = !this.isExhausted();
     return results;
   }
@@ -57,17 +82,31 @@ export abstract class AbstractCursor extends ShellApiClass {
 
   @returnsPromise
   async forEach(f: (doc: Document) => void): Promise<void> {
-    return this._cursor.forEach(f);
+    // Work around https://jira.mongodb.org/browse/NODE-3231
+    let exception;
+    const wrapped = (doc: Document): boolean | undefined => {
+      try {
+        f(doc);
+        return undefined;
+      } catch (err) {
+        exception = err;
+        return false; // Stop iteration.
+      }
+    };
+    await this._cursor.forEach(wrapped);
+    if (exception) {
+      throw exception;
+    }
   }
 
   @returnsPromise
   async hasNext(): Promise<boolean> {
-    return this._cursor.hasNext();
+    return this._withCheckMapError(() => this._cursor.hasNext());
   }
 
   @returnsPromise
   async tryNext(): Promise<Document | null> {
-    return this._cursor.tryNext();
+    return this._withCheckMapError(() => this._cursor.tryNext());
   }
 
   async* [Symbol.asyncIterator]() {
@@ -96,7 +135,7 @@ export abstract class AbstractCursor extends ShellApiClass {
 
   @returnsPromise
   async toArray(): Promise<Document[]> {
-    return this._cursor.toArray();
+    return this._withCheckMapError(() => this._cursor.toArray());
   }
 
   @returnType('this')
@@ -106,7 +145,20 @@ export abstract class AbstractCursor extends ShellApiClass {
 
   @returnType('this')
   map(f: (doc: Document) => Document): this {
-    this._cursor.map(f);
+    // Work around https://jira.mongodb.org/browse/NODE-3232
+    const wrapped = (doc: Document): Document => {
+      if (this._mapError) {
+        // These errors should never become visible to the user.
+        return { __errored: true };
+      }
+      try {
+        return f(doc);
+      } catch (err) {
+        this._mapError = err;
+        return { __errored: true };
+      }
+    };
+    this._cursor.map(wrapped);
     return this;
   }
 
@@ -118,7 +170,7 @@ export abstract class AbstractCursor extends ShellApiClass {
 
   @returnsPromise
   async next(): Promise<Document | null> {
-    return this._cursor.next();
+    return this._withCheckMapError(() => this._cursor.next());
   }
 
   @returnType('this')

packages/shell-api/src/aggregation-cursor.spec.ts

@@ -67,7 +67,7 @@ describe('AggregationCursor', () => {
     it('calls wrappee.map with arguments', () => {
       const arg = {};
       cursor.map(arg);
-      expect(wrappee.map.calledWith(arg)).to.equal(true);
+      expect(wrappee.map).to.have.callCount(1);
     });
   });
 

packages/shell-api/src/change-stream-cursor.ts

@@ -43,7 +43,7 @@ export default class ChangeStreamCursor extends ShellApiClass {
       throw new MongoshRuntimeError('ChangeStreamCursor is closed');
     }
     const result = this._currentIterationResult = new CursorIterationResult();
-    return iterate(result, this._cursor, this._batchSize ?? await this._mongo._batchSize());
+    return iterate(result, this, this._batchSize ?? await this._mongo._batchSize());
   }
 
   /**

packages/shell-api/src/cursor.spec.ts

@@ -71,7 +71,7 @@ describe('Cursor', () => {
     it('calls wrappee.map with arguments', () => {
       const arg = {};
       cursor.map(arg);
-      expect(wrappee.map.calledWith(arg)).to.equal(true);
+      expect(wrappee.map).to.have.callCount(1);
     });
 
     it('has the correct metadata', () => {

packages/shell-api/src/explainable-cursor.spec.ts

@@ -56,7 +56,7 @@ describe('ExplainableCursor', () => {
     it('calls wrappee.map with arguments', () => {
       const arg = () => {};
       eCursor.map(arg);
-      expect(wrappee.map.calledWith(arg)).to.equal(true);
+      expect(wrappee.map).to.have.callCount(1);
     });
 
     it('has the correct metadata', () => {

packages/shell-api/src/helpers.ts

@@ -3,12 +3,9 @@ import type {
   DbOptions,
   Document,
   ExplainVerbosityLike,
-  FindCursor,
-  AggregationCursor as SPAggregationCursor,
   FindAndModifyOptions,
   DeleteOptions,
   MapReduceOptions,
-  ChangeStream,
   KMSProviders,
   ExplainOptions
 } from '@mongosh/service-provider-core';
@@ -22,6 +19,8 @@ import { BinaryType, ReplPlatform } from '@mongosh/service-provider-core';
 import { ClientSideFieldLevelEncryptionOptions } from './field-level-encryption';
 import { AutoEncryptionOptions } from 'mongodb';
 import { shellApiType } from './enums';
+import type { AbstractCursor } from './abstract-cursor';
+import type ChangeStreamCursor from './change-stream-cursor';
 
 /**
  * Helper method to adapt aggregation pipeline options.
@@ -525,9 +524,9 @@ export function addHiddenDataProperty<T = any>(target: T, key: string|symbol, va
 
 export async function iterate(
   results: CursorIterationResult,
-  cursor: FindCursor | SPAggregationCursor | ChangeStream,
+  cursor: AbstractCursor | ChangeStreamCursor,
   batchSize: number): Promise<CursorIterationResult> {
-  if (cursor.closed) {
+  if (cursor.isClosed()) {
     return results;
   }
 

packages/shell-api/src/integration.spec.ts

@@ -1959,4 +1959,52 @@ describe('Shell API (integration)', function() {
       expect((await collection.find().batchSize(50)._it()).documents).to.have.lengthOf(50);
     });
   });
+
+  describe('cursor map/forEach', () => {
+    beforeEach(async() => {
+      await collection.insertMany([...Array(10).keys()].map(i => ({ i })));
+    });
+
+    it('forEach() iterates over input but does not return anything', async() => {
+      let value = 0;
+      const result = await collection.find().forEach(({ i }) => { value += i; });
+      expect(result).to.equal(undefined);
+      expect(value).to.equal(45);
+    });
+
+    it('map() iterates over input and changes documents in-place', async() => {
+      const cursor = collection.find();
+      cursor.map(({ i }) => ({ j: i }));
+      expect((await cursor._it()).documents[0]).to.deep.equal({ j: 0 });
+    });
+
+    it('forEach() errors lead to a rejected promise', async() => {
+      const error = new Error();
+      let calls = 0;
+      try {
+        await collection.find().forEach(() => { calls++; throw error; });
+        expect.fail('missed exception');
+      } catch (err) {
+        expect(err).to.equal(error);
+      }
+      expect(calls).to.equal(1);
+    });
+
+    it('map() errors show up when reading the cursor', async() => {
+      const error = new Error();
+      const cursor = collection.find();
+      let calls = 0;
+      cursor.map(() => { calls++; throw error; });
+      for (let i = 0; i < 2; i++) {
+        // Try reading twice to make sure .map() is not called again for the second attempt.
+        try {
+          await cursor.tryNext();
+          expect.fail('missed exception');
+        } catch (err) {
+          expect(err).to.equal(error);
+        }
+      }
+      expect(calls).to.equal(1);
+    });
+  });
 });