From 39095138e61491f17d1b5a32158e86cd634a8ca4 Mon Sep 17 00:00:00 2001
From: Matt Meigs <matt.meigs@mongodb.com>
Date: Mon, 18 Sep 2023 13:03:05 -0400
Subject: [PATCH] persistence module repos_branches migrate to docsets

---
 .../services/metadata/repos_branches/index.ts | 61 ++++++++++++++-----
 1 file changed, 46 insertions(+), 15 deletions(-)

diff --git a/modules/persistence/src/services/metadata/repos_branches/index.ts b/modules/persistence/src/services/metadata/repos_branches/index.ts
index 3de8086a2..bad39e56c 100644
--- a/modules/persistence/src/services/metadata/repos_branches/index.ts
+++ b/modules/persistence/src/services/metadata/repos_branches/index.ts
@@ -27,6 +27,40 @@ export interface ReposBranchesDocument extends WithId<Document> {
 
 const internals: { [key: project]: ReposBranchesDocument } = {};
 
+const getAggregationPipeline = (matchCondition: any) => {
+  return [
+    // Stage 1: Unwind the repos array to create multiple documents for each referenced repo
+    {
+      $unwind: '$repos',
+    },
+    // Stage 2: Lookup to join with the repos_branches collection
+    {
+      $lookup: {
+        from: 'repos_branches',
+        localField: 'repos',
+        foreignField: '_id',
+        as: 'repo',
+      },
+    },
+    // Stage 3: Match documents based on given field(s)
+    {
+      $match: matchCondition,
+    },
+    // Stage 4: Merge/flatten repo into docset
+    {
+      $replaceRoot: { newRoot: { $mergeObjects: [{ $arrayElemAt: ['$repo', 0] }, '$$ROOT'] } },
+    },
+    // Stage 5: Exclude fields
+    {
+      $project: {
+        _id: 0,
+        repos: 0,
+        repo: 0,
+      },
+    },
+  ];
+};
+
 // Queries pool*.repos_branches for all entries for associated_products in a shared metadata entry
 export const getAllAssociatedRepoBranchesEntries = async (metadata: Metadata) => {
   const { associated_products = [] } = metadata;
@@ -48,14 +82,13 @@ export const getAllAssociatedRepoBranchesEntries = async (metadata: Metadata) =>
 
   try {
     const db = await pool();
-    await db
-      .collection('repos_branches')
-      .find({ project: { $in: fetch } })
-      .forEach((doc: ReposBranchesDocument) => {
-        // TODO: store in cache
-        internals[doc['project']] = doc;
-        res.push(doc);
-      });
+    const aggregationPipeline = getAggregationPipeline({ project: { $in: fetch } });
+    const res = await db.collection('docsets').aggregate(aggregationPipeline).toArray();
+    res.forEach((doc: ReposBranchesDocument) => {
+      // TODO: store in cache
+      internals[doc['project']] = doc;
+      res.push(doc);
+    });
     return res;
   } catch (e) {
     console.error(`Error while getting associated repo branches: ${e}`);
@@ -80,15 +113,13 @@ export const getRepoBranchesEntry = async (project: project, branch = ''): Promi
   // get from DB if not cached
   try {
     const db = await pool();
-    const query = {
-      project,
-    };
+    const matchCondition = { project };
     if (branch) {
-      query['branches'] = {
-        $elemMatch: { gitBranchName: branch },
-      };
+      matchCondition['branches'] = { branches: { $elemMatch: { gitBranchName: branch } } };
     }
-    const res = (await db.collection('repos_branches').findOne(query)) as unknown as ReposBranchesDocument;
+    const aggregationPipeline = getAggregationPipeline(matchCondition);
+
+    const res = (await db.collection('docsets').aggregate(aggregationPipeline)) as unknown as ReposBranchesDocument;
     // if not already set, set cache value for repo_branches
     if (!internals[project]) {
       internals[project] = res;