fix(shell-api): implement own map/forEach for cursors MONGOSH-936 (#1062)

Because the driver expects `.map()` and `.forEach()` callbacks to work
synchronously, but mongosh users can use pseudo-synchronous async
functions, the driver cannot properly handle the functions we pass
to it.
In order to fix this, if we don’t jump through a lot of extra hoops,
it’s easiest to just implement `.map()` and `.forEach()` ourselves.

Author: addaleax

packages/cli-repl/test/e2e.spec.ts

@@ -446,6 +446,59 @@ describe('e2e', function() {
         });
       });
     });
+
+    describe('cursor transform operations', () => {
+      beforeEach(async() => {
+        await shell.executeLine(`use ${dbName}`);
+        await shell.executeLine('for (let i = 0; i < 3; i++) db.coll.insertOne({i})');
+      });
+
+      it('works with .map() with immediate .toArray() iteration', async() => {
+        const result = await shell.executeLine(`const cs = db.coll.find().map((doc) => {
+          print('mapped');
+          return db.coll.find({_id:doc._id}).toArray()
+        }); print('after'); cs.toArray()`);
+        expect(result).to.include('after');
+        expect(result).to.include('mapped');
+        expect(result).to.include('i: 1');
+      });
+
+      it('works with .map() with later .toArray() iteration', async() => {
+        const before = await shell.executeLine(`const cs = db.coll.find().map((doc) => {
+          print('mapped');
+          return db.coll.find({_id:doc._id}).toArray()
+        }); print('after');`);
+        expect(before).to.include('after');
+        expect(before).not.to.include('mapped');
+        const result = await shell.executeLine('cs.toArray()');
+        expect(result).to.include('mapped');
+        expect(result).to.include('i: 1');
+      });
+
+      it('works with .map() with implicit iteration', async() => {
+        const before = await shell.executeLine(`const cs = db.coll.find().map((doc) => {
+          print('mapped');
+          return db.coll.findOne({_id:doc._id});
+        }); print('after');`);
+        expect(before).to.include('after');
+        expect(before).not.to.include('mapped');
+        const result = await shell.executeLine('cs');
+        expect(result).to.include('mapped');
+        expect(result).to.include('i: 1');
+      });
+
+      it('works with .forEach() iteration', async() => {
+        await shell.executeLine('out = [];');
+        const before = await shell.executeLine(`db.coll.find().forEach((doc) => {
+          print('enter forEach');
+          out.push(db.coll.findOne({_id:doc._id}));
+          print('leave forEach');
+        }); print('after');`);
+        expect(before).to.match(/(enter forEach\r?\nleave forEach\r?\n){3}after/);
+        const result = await shell.executeLine('out[1]');
+        expect(result).to.include('i: 1');
+      });
+    });
   });
 
   describe('with --host', () => {

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/BaseMongoIterableHelper.kt

@@ -15,12 +15,6 @@ internal abstract class BaseMongoIterableHelper<T : MongoIterable<*>>(val iterab
     abstract val converters: Map<String, (T, Any?) -> Either<T>>
     abstract val defaultConverter: (T, String, Any?) -> Either<T>
 
-    fun map(function: Value): BaseMongoIterableHelper<*> {
-        return helper(iterable.map { v ->
-            converter.toJava(function.execute(converter.toJs(v))).value
-        }, converter)
-    }
-
     fun itcount(): Int {
         return iterable.count()
     }
@@ -202,4 +196,4 @@ internal fun helper(iterable: MongoIterable<out Any?>, converter: MongoShellConv
         is AggregateIterable -> AggregateIterableHelper(iterable, converter, Document(), null)
         else -> MongoIterableHelper(iterable, converter, Document())
     }
-}
+}

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/Cursor.kt

@@ -103,16 +103,6 @@ internal class Cursor(private var helper: BaseMongoIterableHelper<*>, private va
         return converter.toJs(helper.explain(verbosity))
     }
 
-    @HostAccess.Export
-    override fun forEach(func: Value) {
-        if (!func.canExecute()) {
-            throw IllegalArgumentException("Expected one argument of type function. Got: $func")
-        }
-        getOrCreateIterator().forEach { v ->
-            func.execute(converter.toJs(v))
-        }
-    }
-
     @HostAccess.Export
     override fun hasNext(): Boolean = getOrCreateIterator().hasNext()
 
@@ -148,16 +138,6 @@ internal class Cursor(private var helper: BaseMongoIterableHelper<*>, private va
         return this
     }
 
-    @HostAccess.Export
-    override fun map(func: Value): Cursor {
-        checkQueryNotExecuted()
-        if (!func.canExecute()) {
-            throw IllegalArgumentException("Expected one argument of type function. Got: $func")
-        }
-        helper = helper.map(func)
-        return this
-    }
-
     @HostAccess.Export
     override fun max(v: Value): Cursor {
         checkQueryNotExecuted()

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/ServiceProviderCursor.kt

@@ -13,13 +13,11 @@ interface ServiceProviderCursor {
   fun collation(v: Value): ServiceProviderCursor
   fun comment(v: String): ServiceProviderCursor
   fun count(): Long
-  fun forEach(func: Value)
   fun hasNext(): Boolean
   fun hint(v: Value): ServiceProviderCursor
   fun isExhausted(): Boolean
   fun itcount(): Int
   fun limit(v: Int): ServiceProviderCursor
-  fun map(func: Value): ServiceProviderCursor
   fun max(v: Value): ServiceProviderCursor
   fun maxTimeMS(v: Long): ServiceProviderCursor
   fun maxAwaitTimeMS(value: Int): ServiceProviderCursor

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

@@ -21,39 +21,15 @@ import { iterate, validateExplainableVerbosity, markAsExplainOutput } from './he
 export abstract class AbstractCursor<CursorType extends ServiceProviderAggregationCursor | ServiceProviderCursor> extends ShellApiWithMongoClass {
   _mongo: Mongo;
   _cursor: CursorType;
+  _transform: ((doc: any) => any) | null;
 
   _currentIterationResult: CursorIterationResult | null = null;
-  _mapError: Error | null = null;
 
   constructor(mongo: Mongo, cursor: CursorType) {
     super();
     this._mongo = mongo;
     this._cursor = cursor;
-  }
-
-  // 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;
-    }
-    // eslint-disable-next-line @typescript-eslint/await-thenable
-    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;
+    this._transform = null;
   }
 
   /**
@@ -82,32 +58,27 @@ export abstract class AbstractCursor<CursorType extends ServiceProviderAggregati
   }
 
   @returnsPromise
-  async forEach(f: (doc: Document) => void): Promise<void> {
-    // 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.
+  async forEach(f: (doc: Document) => void | boolean | Promise<void> | Promise<boolean>): Promise<void> {
+    // Do not use the driver method because it does not have Promise support.
+    for await (const doc of this) {
+      if ((await f(doc)) === false) {
+        break;
       }
-    };
-    await this._cursor.forEach(wrapped);
-    if (exception) {
-      throw exception;
     }
   }
 
   @returnsPromise
   async hasNext(): Promise<boolean> {
-    return this._withCheckMapError(() => this._cursor.hasNext());
+    return this._cursor.hasNext();
   }
 
   @returnsPromise
   async tryNext(): Promise<Document | null> {
-    return this._withCheckMapError(() => this._cursor.tryNext());
+    let result = await this._cursor.tryNext();
+    if (result !== null && this._transform !== null) {
+      result = await this._transform(result);
+    }
+    return result;
   }
 
   async* [Symbol.asyncIterator]() {
@@ -136,7 +107,11 @@ export abstract class AbstractCursor<CursorType extends ServiceProviderAggregati
 
   @returnsPromise
   async toArray(): Promise<Document[]> {
-    return this._withCheckMapError(() => this._cursor.toArray());
+    const result = [];
+    for await (const doc of this) {
+      result.push(doc);
+    }
+    return result;
   }
 
   @returnType('this')
@@ -146,20 +121,12 @@ export abstract class AbstractCursor<CursorType extends ServiceProviderAggregati
 
   @returnType('this')
   map(f: (doc: Document) => Document): this {
-    // 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);
+    if (this._transform === null) {
+      this._transform = f;
+    } else {
+      const g = this._transform;
+      this._transform = (doc: any) => f(g(doc));
+    }
     return this;
   }
 
@@ -171,7 +138,11 @@ export abstract class AbstractCursor<CursorType extends ServiceProviderAggregati
 
   @returnsPromise
   async next(): Promise<Document | null> {
-    return this._withCheckMapError(() => this._cursor.next());
+    let result = await this._cursor.next();
+    if (result !== null && this._transform !== null) {
+      result = await this._transform(result);
+    }
+    return result;
   }
 
   @returnType('this')

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

@@ -65,12 +65,6 @@ describe('AggregationCursor', () => {
     it('pretty returns the same cursor', () => {
       expect(cursor.pretty()).to.equal(cursor);
     });
-
-    it('calls wrappee.map with arguments', () => {
-      const arg = {};
-      cursor.map(arg);
-      expect(wrappee.map).to.have.callCount(1);
-    });
   });
 
   describe('Cursor Internals', () => {

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

@@ -180,15 +180,16 @@ describe('Collection', () => {
       });
 
       it('returns an AggregationCursor that wraps the service provider one', async() => {
-        const toArrayResult = [];
-        serviceProviderCursor.toArray.resolves(toArrayResult);
+        const toArrayResult = [{ foo: 'bar' }];
+        serviceProviderCursor.tryNext.onFirstCall().resolves({ foo: 'bar' });
+        serviceProviderCursor.tryNext.onSecondCall().resolves(null);
         serviceProvider.aggregate.returns(serviceProviderCursor);
 
         const cursor = await collection.aggregate([{
           $piplelineStage: {}
         }]);
 
-        expect(await (cursor as AggregationCursor).toArray()).to.equal(toArrayResult);
+        expect(await (cursor as AggregationCursor).toArray()).to.deep.equal(toArrayResult);
       });
 
       it('throws if serviceProvider.aggregate rejects', async() => {

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

@@ -70,12 +70,6 @@ describe('Cursor', () => {
       expect(cursor.pretty()).to.equal(cursor);
     });
 
-    it('calls wrappee.map with arguments', () => {
-      const arg = {};
-      cursor.map(arg);
-      expect(wrappee.map).to.have.callCount(1);
-    });
-
     it('has the correct metadata', () => {
       expect(cursor.collation.serverVersions).to.deep.equal(['3.4.0', ServerVersions.latest]);
     });

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

@@ -297,12 +297,13 @@ describe('Database', () => {
       });
 
       it('returns an AggregationCursor that wraps the service provider one', async() => {
-        const toArrayResult = [];
-        serviceProviderCursor.toArray.resolves(toArrayResult);
+        const toArrayResult = [{ foo: 'bar' }];
+        serviceProviderCursor.tryNext.onFirstCall().resolves({ foo: 'bar' });
+        serviceProviderCursor.tryNext.onSecondCall().resolves(null);
         serviceProvider.aggregateDb.returns(serviceProviderCursor);
 
         const cursor = await database.aggregate([{ $piplelineStage: {} }]);
-        expect(await cursor.toArray()).to.equal(toArrayResult);
+        expect(await cursor.toArray()).to.deep.equal(toArrayResult);
       });
 
       it('throws if serviceProvider.aggregateDb rejects', async() => {

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

@@ -55,12 +55,6 @@ describe('ExplainableCursor', () => {
       expect(eCursor.map()).to.equal(eCursor);
     });
 
-    it('calls wrappee.map with arguments', () => {
-      const arg = () => {};
-      eCursor.map(arg);
-      expect(wrappee.map).to.have.callCount(1);
-    });
-
     it('has the correct metadata', () => {
       expect(eCursor.collation.serverVersions).to.deep.equal(['3.4.0', ServerVersions.latest]);
     });