Merge pull request #15080 from Automattic/vkarpov15/gh-15048

fix: mark documents that are populated using hydratedPopulatedDocs option as populated in top-level doc

lib/document.js

@@ -624,16 +624,17 @@ Document.prototype.toBSON = function() {
 };
 
 /**
- * Initializes the document without setters or marking anything modified.
+ * Hydrates this document with the data in `doc`. Does not run setters or mark any paths modified.
  *
- * Called internally after a document is returned from mongodb. Normally,
+ * Called internally after a document is returned from MongoDB. Normally,
  * you do **not** need to call this function on your own.
  *
  * This function triggers `init` [middleware](https://mongoosejs.com/docs/middleware.html).
  * Note that `init` hooks are [synchronous](https://mongoosejs.com/docs/middleware.html#synchronous).
  *
- * @param {Object} doc document returned by mongo
+ * @param {Object} doc raw document returned by mongo
  * @param {Object} [opts]
+ * @param {Boolean} [opts.hydratedPopulatedDocs=false] If true, hydrate and mark as populated any paths that are populated in the raw document
  * @param {Function} [fn]
  * @api public
  * @memberOf Document
@@ -805,6 +806,14 @@ function init(self, obj, doc, opts, prefix) {
               reason: e
             }));
           }
+        } else if (opts.hydratedPopulatedDocs) {
+          doc[i] = schemaType.cast(value, self, true);
+
+          if (doc[i] && doc[i].$__ && doc[i].$__.wasPopulated) {
+            self.$populated(path, doc[i].$__.wasPopulated.value, doc[i].$__.wasPopulated.options);
+          } else if (Array.isArray(doc[i]) && doc[i].length && doc[i][0]?.$__?.wasPopulated) {
+            self.$populated(path, doc[i].map(populatedDoc => populatedDoc?.$__?.wasPopulated?.value).filter(val => val != null), doc[i][0].$__.wasPopulated.options);
+          }
         } else {
           doc[i] = value;
         }

lib/schemaType.js

@@ -1567,8 +1567,9 @@ SchemaType.prototype._castRef = function _castRef(value, doc, init) {
     !doc.$__.populated[path].options ||
     !doc.$__.populated[path].options.options ||
     !doc.$__.populated[path].options.options.lean) {
-    ret = new pop.options[populateModelSymbol](value);
-    ret.$__.wasPopulated = { value: ret._doc._id };
+    const PopulatedModel = pop ? pop.options[populateModelSymbol] : doc.constructor.db.model(this.options.ref);
+    ret = new PopulatedModel(value);
+    ret.$__.wasPopulated = { value: ret._doc._id, options: { [populateModelSymbol]: PopulatedModel } };
   }
 
   return ret;

test/model.hydrate.test.js

@@ -108,8 +108,10 @@ describe('model', function() {
         users: [{ ref: 'User', type: Schema.Types.ObjectId }]
       });
 
-      db.model('UserTestHydrate', userSchema);
-      const Company = db.model('CompanyTestHyrdrate', companySchema);
+      db.deleteModel(/User/);
+      db.deleteModel(/Company/);
+      db.model('User', userSchema);
+      const Company = db.model('Company', companySchema);
 
       const users = [{ _id: new mongoose.Types.ObjectId(), name: 'Val' }];
       const company = { _id: new mongoose.Types.ObjectId(), name: 'Booster', users: [users[0]] };
@@ -144,6 +146,7 @@ describe('model', function() {
         count: true
       });
 
+      db.deleteModel(/User/);
       const User = db.model('User', UserSchema);
       const Story = db.model('Story', StorySchema);
 
@@ -173,5 +176,27 @@ describe('model', function() {
 
       assert.strictEqual(hydrated.storiesCount, 2);
     });
+
+    it('sets hydrated docs as populated (gh-15048)', async function() {
+      const userSchema = new Schema({
+        name: String
+      });
+      const companySchema = new Schema({
+        name: String,
+        users: [{ ref: 'User', type: Schema.Types.ObjectId }]
+      });
+
+      db.deleteModel(/User/);
+      db.deleteModel(/Company/);
+      const User = db.model('User', userSchema);
+      const Company = db.model('Company', companySchema);
+
+      const users = [{ _id: new mongoose.Types.ObjectId(), name: 'Val' }];
+      const company = { _id: new mongoose.Types.ObjectId(), name: 'Acme', users: [users[0]] };
+
+      const c = Company.hydrate(company, null, { hydratedPopulatedDocs: true });
+      assert.ok(c.populated('users'));
+      assert.ok(c.users[0] instanceof User);
+    });
   });
 });