Merge pull request #14965 from Automattic/vkarpov15/refpath-ids

fix(populate): handle array of ids with parent refPath
Automattic · Oct 17, 2024 · 562aabd · 562aabd
2 parents de211db + 304c791
commit 562aabd
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 7 deletions.
diff --git a/lib/helpers/populate/getModelsMapForPopulate.js b/lib/helpers/populate/getModelsMapForPopulate.js
@@ -478,9 +478,10 @@ function addModelNamesToMap(model, map, available, modelNames, options, data, re
     return;
   }
 
-  let k = modelNames.length;
+  const flatModelNames = utils.array.flatten(modelNames);
+  let k = flatModelNames.length;
   while (k--) {
-    let modelName = modelNames[k];
+    let modelName = flatModelNames[k];
     if (modelName == null) {
       continue;
     }
@@ -503,11 +504,10 @@ function addModelNamesToMap(model, map, available, modelNames, options, data, re
     }
 
     let ids = ret;
-    const flat = Array.isArray(ret) ? utils.array.flatten(ret) : [];
 
     const modelNamesForRefPath = data.modelNamesInOrder ? data.modelNamesInOrder : modelNames;
-    if (data.isRefPath && Array.isArray(ret) && flat.length === modelNamesForRefPath.length) {
-      ids = flat.filter((val, i) => modelNamesForRefPath[i] === modelName);
+    if (data.isRefPath && Array.isArray(ret) && ret.length === modelNamesForRefPath.length) {
+      ids = matchIdsToRefPaths(ret, modelNamesForRefPath, modelName);
     }
 
     const perDocumentLimit = options.perDocumentLimit == null ?
@@ -569,6 +569,20 @@ function _getModelFromConn(conn, modelName) {
   return conn.model(modelName);
 }
 
+function matchIdsToRefPaths(ids, refPaths, refPathToFind) {
+  if (!Array.isArray(refPaths)) {
+    return refPaths === refPathToFind
+      ? Array.isArray(ids)
+        ? utils.array.flatten(ids)
+        : [ids]
+      : [];
+  }
+  if (Array.isArray(ids) && Array.isArray(refPaths)) {
+    return ids.flatMap((id, index) => matchIdsToRefPaths(id, refPaths[index], refPathToFind));
+  }
+  return [];
+}
+
 /*!
  * ignore
  */

diff --git a/lib/helpers/populate/modelNamesFromRefPath.js b/lib/helpers/populate/modelNamesFromRefPath.js
@@ -62,7 +62,5 @@ module.exports = function modelNamesFromRefPath(refPath, doc, populatedPath, mod
     modelNames = Array.isArray(refValue) ? refValue : [refValue];
   }
 
-  modelNames = utils.array.flatten(modelNames);
-
   return modelNames;
 };
diff --git a/test/model.populate.test.js b/test/model.populate.test.js
@@ -2582,6 +2582,91 @@ describe('model: populate:', function() {
       assert.strictEqual(doc.parts[0].contents[1].item.url, 'https://youtube.com');
     });
 
+    it('with refPath and array of ids with parent refPath', async function() {
+      const Child = db.model(
+        'Child',
+        new mongoose.Schema({
+          fetched: Boolean
+        })
+      );
+
+      const Parent = db.model(
+        'Parent',
+        new mongoose.Schema({
+          docArray: [
+            {
+              type: {
+                type: String,
+                enum: ['Child', 'OtherModel']
+              },
+              ids: [
+                {
+                  type: mongoose.Schema.ObjectId,
+                  refPath: 'docArray.type'
+                }
+              ]
+            }
+          ]
+        })
+      );
+      await Child.insertMany([
+        { _id: new mongoose.Types.ObjectId('6671a008596112f0729c2045'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('667195f3596112f0728abe24'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('6671bd39596112f072cda69c'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('6672c351596112f072868565'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('66734edd596112f0727304a2'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('66726eff596112f072f8e834'), fetched: true },
+        { _id: new mongoose.Types.ObjectId('667267ff596112f072ed56b1'), fetched: true }
+      ]);
+      const { _id } = await Parent.create(
+        {
+          docArray: [
+            {},
+            {
+              type: 'Child',
+              ids: [
+                new mongoose.Types.ObjectId('6671a008596112f0729c2045'),
+                new mongoose.Types.ObjectId('667195f3596112f0728abe24'),
+                new mongoose.Types.ObjectId('6671bd39596112f072cda69c'),
+                new mongoose.Types.ObjectId('6672c351596112f072868565')
+              ]
+            },
+            {
+              type: 'Child',
+              ids: [new mongoose.Types.ObjectId('66734edd596112f0727304a2')]
+            },
+            {},
+            {
+              type: 'Child',
+              ids: [new mongoose.Types.ObjectId('66726eff596112f072f8e834')]
+            },
+            {},
+            {
+              type: 'Child',
+              ids: [new mongoose.Types.ObjectId('667267ff596112f072ed56b1')]
+            }
+          ]
+        }
+      );
+
+      const doc = await Parent.findById(_id).populate('docArray.ids').orFail();
+      assert.strictEqual(doc.docArray.length, 7);
+      assert.strictEqual(doc.docArray[0].ids.length, 0);
+      assert.strictEqual(doc.docArray[1].ids.length, 4);
+      assert.ok(doc.docArray[1].ids[0].fetched);
+      assert.ok(doc.docArray[1].ids[1].fetched);
+      assert.ok(doc.docArray[1].ids[2].fetched);
+      assert.ok(doc.docArray[1].ids[3].fetched);
+      assert.strictEqual(doc.docArray[2].ids.length, 1);
+      assert.ok(doc.docArray[2].ids[0].fetched);
+      assert.strictEqual(doc.docArray[3].ids.length, 0);
+      assert.strictEqual(doc.docArray[4].ids.length, 1);
+      assert.ok(doc.docArray[4].ids[0].fetched);
+      assert.strictEqual(doc.docArray[5].ids.length, 0);
+      assert.strictEqual(doc.docArray[6].ids.length, 1);
+      assert.ok(doc.docArray[6].ids[0].fetched);
+    });
+
     it('with nested nonexistant refPath (gh-6457)', async function() {
       const CommentSchema = new Schema({
         text: String,