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

Fix Automattic#15048

Author: vkarpov15

lib/document.js

@@ -805,6 +805,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,26 @@ 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/);
+      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);
+    });
   });
 });