From 39e64ed83bc5a7059d21fac20683501f1301c11b Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Mon, 19 Dec 2022 16:21:52 -0500
Subject: [PATCH] fix(populate): handle virtual populate underneath document
 array with justOne=true and sort set where 1 element has only 1 result

Fix #12730
Re: #10552
---
 .../populate/assignRawDocsToIdStructure.js    |   5 +-
 test/model.populate.test.js                   | 122 ++++++++++++++++++
 2 files changed, 126 insertions(+), 1 deletion(-)

diff --git a/lib/helpers/populate/assignRawDocsToIdStructure.js b/lib/helpers/populate/assignRawDocsToIdStructure.js
index fbff9d10480..be39c370dd8 100644
--- a/lib/helpers/populate/assignRawDocsToIdStructure.js
+++ b/lib/helpers/populate/assignRawDocsToIdStructure.js
@@ -43,6 +43,8 @@ function assignRawDocsToIdStructure(rawIds, resultDocs, resultOrder, options, re
 
   let i = 0;
   const len = rawIds.length;
+  const hasResultArrays = Object.values(resultOrder).find(o => Array.isArray(o));
+
   for (i = 0; i < len; ++i) {
     id = rawIds[i];
 
@@ -77,7 +79,8 @@ function assignRawDocsToIdStructure(rawIds, resultDocs, resultOrder, options, re
       if (doc) {
         if (sorting) {
           const _resultOrder = resultOrder[sid];
-          if (Array.isArray(_resultOrder) && Array.isArray(doc) && _resultOrder.length === doc.length) {
+          if (hasResultArrays) {
+            // If result arrays, rely on the MongoDB server response for ordering
             newOrder.push(doc);
           } else {
             newOrder[_resultOrder] = doc;
diff --git a/test/model.populate.test.js b/test/model.populate.test.js
index c26b2741006..f37e523b1c7 100644
--- a/test/model.populate.test.js
+++ b/test/model.populate.test.js
@@ -4429,6 +4429,55 @@ describe('model: populate:', function() {
         assert.deepEqual(app.modules[1].menu.map(i => i.title), ['Redo', 'Undo']);
       });
 
+      it('in embedded array with sort and one result (gh-10552)', async function() {
+        const AppMenuItemSchema = new Schema({
+          appId: 'ObjectId',
+          moduleId: Number,
+          title: String,
+          parent: {
+            type: mongoose.ObjectId,
+            ref: 'AppMenuItem'
+          },
+          order: Number
+        });
+
+        const moduleSchema = new Schema({
+          _id: Number,
+          title: { type: String },
+          hidden: { type: Boolean }
+        });
+
+        moduleSchema.virtual('menu', {
+          ref: 'Test1',
+          localField: '_id',
+          foreignField: 'moduleId',
+          options: { sort: { title: 1 } }
+        });
+
+        const appSchema = new Schema({
+          modules: [moduleSchema]
+        });
+
+        const App = db.model('Test', appSchema);
+        const AppMenuItem = db.model('Test1', AppMenuItemSchema);
+
+        let app = await App.create({ modules: [{ _id: 1, title: 'File' }, { _id: 2, title: 'Preferences' }] });
+        await AppMenuItem.create([
+          { title: 'Save', moduleId: 1 },
+          { title: 'Save As', moduleId: 1 },
+          // { title: 'Undo', moduleId: 2 },
+          { title: 'Redo', moduleId: 2 }
+        ]);
+
+        app = await App.findById(app).populate('modules.menu');
+        app = app.toObject({ virtuals: true });
+
+        assert.equal(app.modules.length, 2);
+        assert.equal(app.modules[0].menu.length, 2);
+        assert.deepEqual(app.modules[0].menu.map(i => i.title), ['Save', 'Save As']);
+        assert.deepEqual(app.modules[1].menu.map(i => i.title), ['Redo']);
+      });
+
       it('justOne option (gh-4263)', function(done) {
         const PersonSchema = new Schema({
           name: String,
@@ -10793,6 +10842,79 @@ describe('model: populate:', function() {
     assert.equal(row.values.get(createList._id.toString()).valueObject.name, 'test');
   });
 
+  it('handles virtual populate with `justOne` underneath document array and sort (gh-12730) (gh-10552)', async function() {
+    const shiftSchema = new mongoose.Schema({
+      employeeId: mongoose.Types.ObjectId,
+      startedAt: Date,
+      endedAt: Date,
+      name: String
+    });
+
+    const Shift = db.model('Child', shiftSchema);
+
+    const employeeSchema = new mongoose.Schema({
+      name: String
+    }, { toJSON: { virtuals: true }, toObject: { virtuals: true } });
+
+    employeeSchema.virtual('mostRecentShift', {
+      ref: Shift,
+      localField: '_id',
+      foreignField: 'employeeId',
+      options: {
+        sort: { startedAt: -1 }
+      },
+      justOne: true
+    });
+
+    const storeSchema = new mongoose.Schema({
+      location: String,
+      employees: [employeeSchema]
+    });
+
+    const Store = db.model('Parent', storeSchema);
+
+    const store = await Store.create({
+      location: 'Tashbaan',
+      employees: [
+        { name: 'Aravis' },
+        { name: 'Shasta' }
+      ]
+    });
+
+    const employeeAravis = store.employees.find(({ name }) => name === 'Aravis');
+    const employeeShasta = store.employees.find(({ name }) => name === 'Shasta');
+
+    await Shift.insertMany([
+      {
+        employeeId: employeeAravis._id,
+        startedAt: new Date(Date.now() - 57600000),
+        endedAt: new Date(Date.now() - 43200000),
+        name: 'shift1'
+      },
+      {
+        employeeId: employeeAravis._id,
+        startedAt: new Date(Date.now() - 28800000),
+        endedAt: new Date(Date.now() - 14400000),
+        name: 'shift2'
+      },
+      {
+        employeeId: employeeShasta._id,
+        startedAt: new Date(Date.now() - 14400000),
+        endedAt: new Date(),
+        name: 'shift3'
+      }
+    ]);
+
+    const storeWithMostRecentShifts = await Store.
+      findOne({ location: 'Tashbaan' }).
+      populate('employees.mostRecentShift');
+
+    assert.deepStrictEqual(
+      storeWithMostRecentShifts.employees.map(e => e.mostRecentShift.name),
+      ['shift2', 'shift3']
+    );
+  });
+
   describe('strictPopulate', function() {
     it('reports full path when throwing `strictPopulate` error with deep populate (gh-10923)', async function() {
       const L2 = db.model('Test', new Schema({ name: String }));