Merge pull request #11814 from Automattic/gh-10423

Gh 10423 adds a transform option to lean queries
Automattic · Jun 3, 2022 · f5dc4f8 · f5dc4f8
2 parents 8c5263d + 1c7696e
commit f5dc4f8
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 5 deletions.
diff --git a/lib/query.js b/lib/query.js
@@ -27,6 +27,7 @@ const immediate = require('./helpers/immediate');
 const isExclusive = require('./helpers/projection/isExclusive');
 const isInclusive = require('./helpers/projection/isInclusive');
 const isSubpath = require('./helpers/projection/isSubpath');
+const mpath = require('mpath');
 const mquery = require('mquery');
 const parseProjection = require('./helpers/projection/parseProjection');
 const removeUnusedArrayFilters = require('./helpers/update/removeUnusedArrayFilters');
@@ -2217,7 +2218,8 @@ Query.prototype._find = wrapThunk(function(callback) {
   // Separate options to pass down to `completeMany()` in case we need to
   // set a session on the document
   const completeManyOptions = Object.assign({}, {
-    session: this && this.options && this.options.session || null
+    session: this && this.options && this.options.session || null,
+    lean: mongooseOptions.lean || null
   });
 
   const cb = (err, docs) => {
@@ -2241,7 +2243,9 @@ Query.prototype._find = wrapThunk(function(callback) {
         });
       }
       return mongooseOptions.lean ?
-        callback(null, docs) :
+      // call _completeManyLean here?
+        _completeManyLean(_this.model.schema, docs, null, completeManyOptions, callback) :
+        // callback(null, docs) :
         completeMany(_this.model, docs, fields, userProvidedFields, completeManyOptions, callback);
     }
 
@@ -2415,6 +2419,9 @@ Query.prototype._completeOne = function(doc, res, callback) {
   const mongooseOptions = this._mongooseOptions;
   // `rawResult`
   const options = this.options;
+  if (!options.lean && mongooseOptions.lean) {
+    options.lean = mongooseOptions.lean;
+  }
 
   if (options.explain) {
     return callback(null, doc);
@@ -2428,7 +2435,7 @@ Query.prototype._completeOne = function(doc, res, callback) {
       }
     }
     return mongooseOptions.lean ?
-      _completeOneLean(doc, res, options, callback) :
+      _completeOneLean(model.schema, doc, null, res, options, callback) :
       completeOne(model, doc, res, options, projection, userProvidedFields,
         null, callback);
   }
@@ -2439,7 +2446,7 @@ Query.prototype._completeOne = function(doc, res, callback) {
       if (err != null) {
         return callback(err);
       }
-      _completeOneLean(doc, res, options, callback);
+      _completeOneLean(model.schema, doc, null, res, options, callback);
     });
   }
 
@@ -4006,13 +4013,62 @@ Query.prototype._findAndModify = function(type, callback) {
  * ignore
  */
 
-function _completeOneLean(doc, res, opts, callback) {
+function _completeOneLean(schema, doc, path, res, opts, callback) {
+  if (opts.lean && opts.lean.transform) {
+    for (let i = 0; i < schema.childSchemas.length; i++) {
+      const childPath = path ? path + '.' + schema.childSchemas[i].model.path : schema.childSchemas[i].model.path;
+      const _schema = schema.childSchemas[i].schema;
+      const obj = mpath.get(childPath, doc);
+      if (obj == null) {
+        continue;
+      }
+      if (Array.isArray(obj)) {
+        for (let i = 0; i < obj.length; i++) {
+          opts.lean.transform(obj[i]);
+        }
+      } else {
+        opts.lean.transform(obj);
+      }
+      _completeOneLean(_schema, obj, childPath, res, opts);
+    }
+    if (callback) {
+      return callback(null, doc);
+    } else {
+      return;
+    }
+  }
   if (opts.rawResult) {
     return callback(null, res);
   }
   return callback(null, doc);
 }
 
+/*!
+ * ignore
+ */
+
+function _completeManyLean(schema, docs, path, opts, callback) {
+  if (opts.lean && opts.lean.transform) {
+    for (let i = 0; i < schema.childSchemas.length; i++) {
+      const childPath = path ? path + '.' + schema.childSchemas[i].model.path : schema.childSchemas[i].model.path;
+      const _schema = schema.childSchemas[i].schema;
+      let doc = mpath.get(childPath, docs);
+      if (doc == null) {
+        continue;
+      }
+      doc = doc.flat();
+      for (let i = 0; i < doc.length; i++) {
+        opts.lean.transform(doc[i]);
+      }
+      _completeManyLean(_schema, doc, childPath, opts);
+    }
+  }
+
+  if (!callback) {
+    return;
+  }
+  return callback(null, docs);
+}
 /*!
  * Override mquery.prototype._mergeUpdate to handle mongoose objects in
  * updates.

diff --git a/test/query.test.js b/test/query.test.js
@@ -3962,4 +3962,35 @@ describe('Query', function() {
     assert.equal(result.length, 1);
     assert.equal(result[0].name, '@foo.com');
   });
+  it('allows a transform option for lean on a query gh-10423', async function() {
+    const arraySchema = new mongoose.Schema({
+      sub: String
+    });
+    const subDoc = new mongoose.Schema({
+      nickName: String
+    });
+    const testSchema = new mongoose.Schema({
+      name: String,
+      foo: [arraySchema],
+      otherName: subDoc
+    });
+    const Test = db.model('gh10423', testSchema);
+    await Test.create({ name: 'foo', foo: [{ sub: 'Test' }, { sub: 'Testerson' }], otherName: { nickName: 'Bar' } });
+    const result = await Test.find().lean({ transform: (doc) => {
+      delete doc._id;
+      return doc;
+    } });
+    assert(result[0]._id);
+    assert.equal(result[0].otherName._id, undefined);
+    assert.equal(result[0].foo[0]._id, undefined);
+    assert.equal(result[0].foo[1]._id, undefined);
+    const single = await Test.findOne().lean({ transform: (doc) => {
+      delete doc._id;
+      return doc;
+    } });
+    assert(single._id);
+    assert.equal(single.otherName._id, undefined);
+    assert.equal(single.foo[0]._id, undefined);
+    assert.equal(single.foo[0]._id, undefined);
+  });
 });