Sync row usage with app deletion

packages/server/src/middleware/usageQuota.js

@@ -1,5 +1,6 @@
 const CouchDB = require("../db")
 const usageQuota = require("../utilities/usageQuota")
+const { getUniqueRows } = require("../utilities/usageQuota/rows")
 const {
   isExternalTable,
   isRowId: isExternalRowId,
@@ -74,9 +75,81 @@ module.exports = async (ctx, next) => {
   }
   try {
     await quotaMigration.runIfRequired()
-    await usageQuota.update(property, usage)
-    return next()
+    await performRequest(ctx, next, property, usage)
   } catch (err) {
     ctx.throw(400, err)
   }
 }
+
+const performRequest = async (ctx, next, property, usage) => {
+  const usageContext = {
+    skipNext: false,
+    skipUsage: false,
+    [usageQuota.Properties.APPS]: {},
+  }
+
+  if (usage === -1) {
+    if (PRE_DELETE[property]) {
+      await PRE_DELETE[property](ctx, usageContext)
+    }
+  } else {
+    if (PRE_CREATE[property]) {
+      await PRE_CREATE[property](ctx, usageContext)
+    }
+  }
+
+  // run the request
+  if (!usageContext.skipNext) {
+    await usageQuota.update(property, usage, { dryRun: true })
+    await next()
+  }
+
+  if (usage === -1) {
+    if (POST_DELETE[property]) {
+      await POST_DELETE[property](ctx, usageContext)
+    }
+  } else {
+    if (POST_CREATE[property]) {
+      await POST_CREATE[property](ctx, usageContext)
+    }
+  }
+
+  // update the usage
+  if (!usageContext.skipUsage) {
+    await usageQuota.update(property, usage)
+  }
+}
+
+const appPreDelete = async (ctx, usageContext) => {
+  if (ctx.query.unpublish) {
+    // don't run usage decrement for unpublish
+    usageContext.skipUsage = true
+    return
+  }
+
+  // store the row count to delete
+  const rows = await getUniqueRows([ctx.appId])
+  if (rows.size) {
+    usageContext[usageQuota.Properties.APPS] = { rowCount: rows.size }
+  }
+}
+
+const appPostDelete = async (ctx, usageContext) => {
+  // delete the app rows from usage
+  const rowCount = usageContext[usageQuota.Properties.APPS].rowCount
+  if (rowCount) {
+    await usageQuota.update(usageQuota.Properties.ROW, -rowCount)
+  }
+}
+
+const PRE_DELETE = {
+  [usageQuota.Properties.APPS]: appPreDelete,
+}
+
+const POST_DELETE = {
+  [usageQuota.Properties.APPS]: appPostDelete,
+}
+
+const PRE_CREATE = {}
+
+const POST_CREATE = {}

packages/server/src/migrations/sync_app_and_reset_rows_quotas.js

@@ -7,44 +7,7 @@ const { getGlobalDB } = require("@budibase/backend-core/tenancy")
 const { getAllApps } = require("@budibase/backend-core/db")
 const CouchDB = require("../db")
 const { getUsageQuotaDoc, useQuotas } = require("../utilities/usageQuota")
-const { getRowParams } = require("../db/utils")
-
-/**
- * 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
-  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))
-    } catch (e) {
-      // don't error out if we can't count the app rows, just continue
-    }
-  }
-
-  return allRows
-}
-
-/**
- * Get all rows in the given app ids.
- *
- * The returned rows will be unique, duplicated rows across
- * production and dev apps will be removed.
- */
-const getUniqueRows = async appIds => {
-  const allRows = await getAllRows(appIds)
-  return new Set(allRows)
-}
+const { getUniqueRows } = require("../utilities/usageQuota/rows")
 
 const syncRowsQuota = async db => {
   // get all rows in all apps

packages/server/src/utilities/tests/usageQuota/usageQuota.spec.js

@@ -3,7 +3,7 @@ jest.mock("@budibase/backend-core/tenancy", () => ({
   getTenantId
 }))
 const usageQuota = require("../usageQuota")
-const env = require("../../environment")
+const env = require("../../../environment")
 
 class TestConfiguration {
   constructor() {

packages/server/src/utilities/usageQuota/index.js

@@ -1,4 +1,4 @@
-const env = require("../environment")
+const env = require("../../environment")
 const { getGlobalDB, getTenantId } = require("@budibase/backend-core/tenancy")
 const {
   StaticDatabases,
@@ -55,7 +55,7 @@ exports.getUsageQuotaDoc = async db => {
  * @returns {Promise<void>} When this completes the API key will now be up to date - the quota period may have
  * also been reset after this call.
  */
-exports.update = async (property, usage) => {
+exports.update = async (property, usage, opts = { dryRun: false }) => {
   if (!exports.useQuotas()) {
     return
   }
@@ -67,14 +67,24 @@ exports.update = async (property, usage) => {
     // increment the quota
     quota.usageQuota[property] += usage
 
-    if (quota.usageQuota[property] > quota.usageLimits[property]) {
+    if (
+      quota.usageQuota[property] > quota.usageLimits[property] &&
+      usage > 0 // allow for decrementing usage when the quota is already exceeded
+    ) {
       throw new Error(
         `You have exceeded your usage quota of ${quota.usageLimits[property]} ${property}.`
       )
     }
 
+    if (quota.usageQuota[property] < 0) {
+      // never go negative if the quota has previously been exceeded
+      quota.usageQuota[property] = 0
+    }
+
     // update the usage quotas
-    await db.put(quota)
+    if (!opts.dryRun) {
+      await db.put(quota)
+    }
   } catch (err) {
     console.error(`Error updating usage quotas for ${property}`, err)
     throw err

packages/server/src/utilities/usageQuota/rows.js

@@ -0,0 +1,52 @@
+const { getRowParams, USER_METDATA_PREFIX } = require("../../db/utils")
+const CouchDB = require("../../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
+  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
+    }
+  }
+
+  return allRows
+}
+
+/**
+ * Get all rows in the given app ids.
+ *
+ * The returned rows will be unique, duplicated rows across
+ * production and dev apps will be removed.
+ */
+exports.getUniqueRows = async appIds => {
+  const allRows = await getAllRows(appIds)
+  return new Set(allRows)
+}