Merge pull request #15430 from Automattic/vkarpov15/mongoose-lean-getters-44

fix(queryCursor): add _transformForAsyncIterator transform after user-defined transforms

Author: vkarpov15

lib/cursor/queryCursor.js

@@ -83,6 +83,9 @@ function QueryCursor(query) {
       // Max out the number of documents we'll populate in parallel at 5000.
       this.options._populateBatchSize = Math.min(this.options.batchSize, 5000);
     }
+    if (query._mongooseOptions._asyncIterator) {
+      this._mongooseOptions._asyncIterator = true;
+    }
 
     if (model.collection._shouldBufferCommands() && model.collection.buffer) {
       model.collection.queue.push([
@@ -379,29 +382,6 @@ QueryCursor.prototype.addCursorFlag = function(flag, value) {
   return this;
 };
 
-/*!
- * ignore
- */
-
-QueryCursor.prototype.transformNull = function(val) {
-  if (arguments.length === 0) {
-    val = true;
-  }
-  this._mongooseOptions.transformNull = val;
-  return this;
-};
-
-/*!
- * ignore
- */
-
-QueryCursor.prototype._transformForAsyncIterator = function() {
-  if (this._transforms.indexOf(_transformForAsyncIterator) === -1) {
-    this.map(_transformForAsyncIterator);
-  }
-  return this;
-};
-
 /**
  * Returns an asyncIterator for use with [`for/await/of` loops](https://thecodebarbarian.com/getting-started-with-async-iterators-in-node-js).
  * You do not need to call this function explicitly, the JavaScript runtime
@@ -433,19 +413,13 @@ QueryCursor.prototype._transformForAsyncIterator = function() {
  */
 
 if (Symbol.asyncIterator != null) {
-  QueryCursor.prototype[Symbol.asyncIterator] = function() {
-    return this.transformNull()._transformForAsyncIterator();
+  QueryCursor.prototype[Symbol.asyncIterator] = function queryCursorAsyncIterator() {
+    // Set so QueryCursor knows it should transform results for async iterators into `{ value, done }` syntax
+    this._mongooseOptions._asyncIterator = true;
+    return this;
   };
 }
 
-/*!
- * ignore
- */
-
-function _transformForAsyncIterator(doc) {
-  return doc == null ? { done: true } : { value: doc, done: false };
-}
-
 /**
  * Get the next doc from the underlying cursor and mongooseify it
  * (populate, etc.)
@@ -456,16 +430,38 @@ function _transformForAsyncIterator(doc) {
 
 function _next(ctx, cb) {
   let callback = cb;
-  if (ctx._transforms.length) {
-    callback = function(err, doc) {
-      if (err || (doc === null && !ctx._mongooseOptions.transformNull)) {
-        return cb(err, doc);
+
+  // Create a custom callback to handle transforms, async iterator, and transformNull
+  callback = function(err, doc) {
+    if (err) {
+      return cb(err);
+    }
+
+    // Handle null documents - if asyncIterator, we need to return `done: true`, otherwise just
+    // skip. In either case, avoid transforms.
+    if (doc === null) {
+      if (ctx._mongooseOptions._asyncIterator) {
+        return cb(null, { done: true });
+      } else {
+        return cb(null, null);
       }
-      cb(err, ctx._transforms.reduce(function(doc, fn) {
+    }
+
+    // Apply transforms
+    if (ctx._transforms.length && doc !== null) {
+      doc = ctx._transforms.reduce(function(doc, fn) {
         return fn.call(ctx, doc);
-      }, doc));
-    };
-  }
+      }, doc);
+    }
+
+    // This option is set in `Symbol.asyncIterator` code paths.
+    // For async iterator, we need to convert to {value, done} format
+    if (ctx._mongooseOptions._asyncIterator) {
+      return cb(null, { value: doc, done: false });
+    }
+
+    return cb(null, doc);
+  };
 
   if (ctx._error) {
     return immediate(function() {

lib/query.js

@@ -5476,8 +5476,10 @@ Query.prototype.nearSphere = function() {
  */
 
 if (Symbol.asyncIterator != null) {
-  Query.prototype[Symbol.asyncIterator] = function() {
-    return this.cursor().transformNull()._transformForAsyncIterator();
+  Query.prototype[Symbol.asyncIterator] = function queryAsyncIterator() {
+    // Set so QueryCursor knows it should transform results for async iterators into `{ value, done }` syntax
+    this._mongooseOptions._asyncIterator = true;
+    return this.cursor();
   };
 }
 

test/query.test.js

@@ -2354,6 +2354,19 @@ describe('Query', function() {
     assert.strictEqual(called, 1);
   });
 
+  it('transform with for/await and cursor', async function() {
+    const Model = db.model('Test', new Schema({ name: String }));
+
+    await Model.create({ name: 'test' });
+    const cursor = Model.find().transform(doc => doc.name.toUpperCase()).cursor();
+    const names = [];
+    for await (const name of cursor) {
+      names.push(name);
+    }
+
+    assert.deepStrictEqual(names, ['TEST']);
+  });
+
   describe('orFail (gh-6841)', function() {
     let Model;