diff --git a/lib/cursor/queryCursor.js b/lib/cursor/queryCursor.js index f25a06f2fd1..f157b71df4d 100644 --- a/lib/cursor/queryCursor.js +++ b/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() { diff --git a/lib/query.js b/lib/query.js index 32fc1811b17..ad539e3518e 100644 --- a/lib/query.js +++ b/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(); }; } diff --git a/test/query.test.js b/test/query.test.js index bd7cf01d3ba..fb779b82fc6 100644 --- a/test/query.test.js +++ b/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;