From a09bbfb49297a470cfe81eebb4936fcfdb057de3 Mon Sep 17 00:00:00 2001
From: Rory Powell <rory.codes@gmail.com>
Date: Tue, 18 Jan 2022 12:48:54 +0000
Subject: [PATCH] Account for duplicate row ids caused by app import

---
 packages/server/src/middleware/usageQuota.js  |  4 +-
 .../sync_app_and_reset_rows_quotas.js         |  2 +-
 .../server/src/utilities/usageQuota/rows.js   | 95 +++++++++++--------
 3 files changed, 61 insertions(+), 40 deletions(-)

diff --git a/packages/server/src/middleware/usageQuota.js b/packages/server/src/middleware/usageQuota.js
index 2d55aac61d..c42442affe 100644
--- a/packages/server/src/middleware/usageQuota.js
+++ b/packages/server/src/middleware/usageQuota.js
@@ -129,8 +129,8 @@ const appPreDelete = async (ctx, usageContext) => {
 
   // store the row count to delete
   const rows = await getUniqueRows([ctx.appId])
-  if (rows.size) {
-    usageContext[usageQuota.Properties.APPS] = { rowCount: rows.size }
+  if (rows.length) {
+    usageContext[usageQuota.Properties.APPS] = { rowCount: rows.length }
   }
 }
 
diff --git a/packages/server/src/migrations/sync_app_and_reset_rows_quotas.js b/packages/server/src/migrations/sync_app_and_reset_rows_quotas.js
index 5445919f26..4ab2e413f3 100644
--- a/packages/server/src/migrations/sync_app_and_reset_rows_quotas.js
+++ b/packages/server/src/migrations/sync_app_and_reset_rows_quotas.js
@@ -17,7 +17,7 @@ const syncRowsQuota = async db => {
 
   // sync row count
   const usageDoc = await getUsageQuotaDoc(db)
-  usageDoc.usageQuota.rows = rows.size
+  usageDoc.usageQuota.rows = rows.length
   await db.put(usageDoc)
 }
 
diff --git a/packages/server/src/utilities/usageQuota/rows.js b/packages/server/src/utilities/usageQuota/rows.js
index 3a5602958b..19c3e5b8bf 100644
--- a/packages/server/src/utilities/usageQuota/rows.js
+++ b/packages/server/src/utilities/usageQuota/rows.js
@@ -1,52 +1,73 @@
 const { getRowParams, USER_METDATA_PREFIX } = require("../../db/utils")
 const CouchDB = require("../../db")
+const { isDevAppID, getDevelopmentAppID } = require("@budibase/backend-core/db")
 
 const ROW_EXCLUSIONS = [USER_METDATA_PREFIX]
 
-/**
- * Get all rows in the given app ids.
- *
- * The returned rows may contan duplicates if there
- * is a production and dev app.
- */
-const getAllRows = async appIds => {
-  const allRows = []
-  let appDb
+const getAppPairs = appIds => {
+  // collect the app ids into dev / prod pairs
+  // keyed by the dev app id
+  const pairs = {}
   for (let appId of appIds) {
-    try {
-      appDb = new CouchDB(appId)
-      const response = await appDb.allDocs(
-        getRowParams(null, null, {
-          include_docs: false,
-        })
-      )
-      allRows.push(
-        ...response.rows
-          .map(r => r.id)
-          .filter(id => {
-            for (let exclusion of ROW_EXCLUSIONS) {
-              if (id.startsWith(exclusion)) {
-                return false
-              }
-            }
-            return true
-          })
-      )
-    } catch (e) {
-      // don't error out if we can't count the app rows, just continue
+    const devId = getDevelopmentAppID(appId)
+    if (!pairs[devId]) {
+      pairs[devId] = {}
+    }
+    if (isDevAppID(appId)) {
+      pairs[devId].devId = appId
+    } else {
+      pairs[devId].prodId = appId
     }
   }
+  return pairs
+}
 
-  return allRows
+const getAppRows = async appId => {
+  const appDb = new CouchDB(appId)
+  const response = await appDb.allDocs(
+    getRowParams(null, null, {
+      include_docs: false,
+    })
+  )
+  return response.rows
+    .map(r => r.id)
+    .filter(id => {
+      for (let exclusion of ROW_EXCLUSIONS) {
+        if (id.startsWith(exclusion)) {
+          return false
+        }
+      }
+      return true
+    })
 }
 
 /**
- * Get all rows in the given app ids.
- *
- * The returned rows will be unique, duplicated rows across
- * production and dev apps will be removed.
+ * Return a set of all rows in the given app ids.
+ * The returned rows will be unique on a per dev/prod app basis.
+ * Rows duplicates may exist across apps due to data import so they are not filtered out.
  */
 exports.getUniqueRows = async appIds => {
-  const allRows = await getAllRows(appIds)
-  return new Set(allRows)
+  let uniqueRows = []
+  const pairs = getAppPairs(appIds)
+
+  for (let pair of Object.values(pairs)) {
+    let appRows = []
+    for (let appId of [pair.devId, pair.prodId]) {
+      if (!appId) {
+        continue
+      }
+      try {
+        appRows.push(await getAppRows(appId))
+      } catch (e) {
+        // don't error out if we can't count the app rows, just continue
+      }
+    }
+
+    // ensure uniqueness on a per app pair basis
+    // this can't be done on all rows because app import results in
+    // duplicate row ids across apps
+    uniqueRows = uniqueRows.concat(...new Set(appRows))
+  }
+
+  return uniqueRows
 }