Calculate total rows in migration, centralise quota enabled logic and tidy to use env vars only

charts/budibase/templates/app-service-deployment.yaml

@@ -99,7 +99,9 @@ spec:
         - name: PLATFORM_URL
           value: {{ .Values.globals.platformUrl | quote }}
         - name: USE_QUOTAS
-          value: "1" 
+          value: {{ .Values.globals.useQuotas | quote }}
+        - name: EXCLUDE_QUOTAS_TENANTS
+          value: {{ .Values.globals.excludeQuotasTenants | quote }}
         - name: ACCOUNT_PORTAL_URL
           value: {{ .Values.globals.accountPortalUrl | quote }}
         - name: ACCOUNT_PORTAL_API_KEY

charts/budibase/values.yaml

@@ -93,6 +93,8 @@ globals:
   logLevel: info
   selfHosted: "1" # set to 0 for budibase cloud environment, set to 1 for self-hosted setup
   multiTenancy: "0" # set to 0 to disable multiple orgs, set to 1 to enable multiple orgs
+  useQuotas: "0"
+  excludeQuotasTenants: "" # comma seperated list of tenants to exclude from quotas
   accountPortalUrl: ""
   accountPortalApiKey: ""
   cookieDomain: ""
@@ -239,7 +241,8 @@ couchdb:
     hosts:
       - chart-example.local
     path: /
-    annotations: []
+    annotations:
+      []
       # kubernetes.io/ingress.class: nginx
       # kubernetes.io/tls-acme: "true"
     tls:

packages/server/src/automations/steps/createRow.js

@@ -1,6 +1,5 @@
 const rowController = require("../../api/controllers/row")
 const automationUtils = require("../automationUtils")
-const env = require("../../environment")
 const usage = require("../../utilities/usageQuota")
 const { buildCtx } = require("./utils")
 
@@ -83,9 +82,7 @@ exports.run = async function ({ inputs, appId, emitter }) {
       inputs.row.tableId,
       inputs.row
     )
-    if (env.USE_QUOTAS) {
-      await usage.update(usage.Properties.ROW, 1)
-    }
+    await usage.update(usage.Properties.ROW, 1)
     await rowController.save(ctx)
     return {
       row: inputs.row,

packages/server/src/environment.js

@@ -38,6 +38,7 @@ module.exports = {
   MINIO_ACCESS_KEY: process.env.MINIO_ACCESS_KEY,
   MINIO_SECRET_KEY: process.env.MINIO_SECRET_KEY,
   USE_QUOTAS: process.env.USE_QUOTAS,
+  EXCLUDE_QUOTAS_TENANTS: process.env.EXCLUDE_QUOTAS_TENANTS,
   REDIS_URL: process.env.REDIS_URL,
   REDIS_PASSWORD: process.env.REDIS_PASSWORD,
   INTERNAL_API_KEY: process.env.INTERNAL_API_KEY,

packages/server/src/integrations/s3.ts

@@ -38,7 +38,7 @@ module S3Module {
       signatureVersion: {
         type: "string",
         required: false,
-        default: "v4"
+        default: "v4",
       },
     },
     query: {

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

@@ -1,11 +1,5 @@
 jest.mock("../../db")
 jest.mock("../../utilities/usageQuota")
-jest.mock("../../environment", () => ({
-  isTest: () => true,
-  isProd: () => false,
-  isDev: () => true,
-  _set: () => {},
-}))
 jest.mock("@budibase/backend-core/tenancy", () => ({
   getTenantId: () => "testing123"
 }))
@@ -32,6 +26,7 @@ class TestConfiguration {
         url: "/applications"
       }
     }
+    usageQuota.useQuotas = () => true
   }
 
   executeMiddleware() {
@@ -113,12 +108,10 @@ describe("usageQuota middleware", () => {
 
   it("calculates and persists the correct usage quota for the relevant action", async () => {
     config.setUrl("/rows")
-    config.setProd(true)
 
     await config.executeMiddleware()
 
-    // expect(usageQuota.update).toHaveBeenCalledWith("rows", 1)
-    expect(usageQuota.update).not.toHaveBeenCalledWith("rows", 1)
+    expect(usageQuota.update).toHaveBeenCalledWith("rows", 1)
     expect(config.next).toHaveBeenCalled()
   })
 

packages/server/src/middleware/usageQuota.js

@@ -1,18 +1,11 @@
 const CouchDB = require("../db")
 const usageQuota = require("../utilities/usageQuota")
-const env = require("../environment")
-const { getTenantId } = require("@budibase/backend-core/tenancy")
 const {
   isExternalTable,
   isRowId: isExternalRowId,
 } = require("../integrations/utils")
 const quotaMigration = require("../migrations/sync_app_and_reset_rows_quotas")
 
-const testing = false
-
-// tenants without limits
-const EXCLUDED_TENANTS = ["bb", "default", "bbtest", "bbstaging"]
-
 // currently only counting new writes and deletes
 const METHOD_MAP = {
   POST: 1,
@@ -20,7 +13,7 @@ const METHOD_MAP = {
 }
 
 const DOMAIN_MAP = {
-  // rows: usageQuota.Properties.ROW,                 // works - disabled
+  rows: usageQuota.Properties.ROW,
   // upload: usageQuota.Properties.UPLOAD,            // doesn't work yet
   // views: usageQuota.Properties.VIEW,               // doesn't work yet
   // users: usageQuota.Properties.USER,               // doesn't work yet
@@ -39,13 +32,7 @@ function getProperty(url) {
 }
 
 module.exports = async (ctx, next) => {
-  const tenantId = getTenantId()
-
-  // if in development or a self hosted cloud usage quotas should not be executed
-  if (
-    (env.isDev() || env.SELF_HOSTED || EXCLUDED_TENANTS.includes(tenantId)) &&
-    !testing
-  ) {
+  if (!usageQuota.useQuotas()) {
     return next()
   }
 

packages/server/src/migrations/sync_app_and_reset_rows_quotas.js

@@ -6,25 +6,80 @@ const {
 const { getGlobalDB } = require("@budibase/backend-core/tenancy")
 const { getAllApps } = require("@budibase/backend-core/db")
 const CouchDB = require("../db")
-const { getUsageQuotaDoc } = require("../utilities/usageQuota")
+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 syncRowsQuota = async db => {
+  // get all rows in all apps
+  const allApps = await getAllApps(CouchDB, { all: true })
+  const appIds = allApps ? allApps.map(app => app.appId) : []
+  const rows = await getUniqueRows(appIds)
+
+  // sync row count
+  const usageDoc = await getUsageQuotaDoc(db)
+  usageDoc.usageQuota.rows = rows.size
+  await db.put(usageDoc)
+}
+
+const syncAppsQuota = async db => {
+  // get app count
+  const devApps = await getAllApps(CouchDB, { dev: true })
+  const appCount = devApps ? devApps.length : 0
+
+  // sync app count
+  const usageDoc = await getUsageQuotaDoc(db)
+  usageDoc.usageQuota.apps = appCount
+  await db.put(usageDoc)
+}
 
 exports.runIfRequired = async () => {
   await migrateIfRequired(
     MIGRATION_DBS.GLOBAL_DB,
     MIGRATIONS.SYNC_APP_AND_RESET_ROWS_QUOTAS,
     async () => {
+      if (!useQuotas()) {
+        return
+      }
       const db = getGlobalDB()
-      const usageDoc = await getUsageQuotaDoc(db)
-
-      // reset the rows
-      usageDoc.usageQuota.rows = 0
-
-      // sync the apps
-      const apps = await getAllApps(CouchDB, { dev: true })
-      const appCount = apps ? apps.length : 0
-      usageDoc.usageQuota.apps = appCount
-
-      await db.put(usageDoc)
+      await syncAppsQuota(db)
+      await syncRowsQuota(db)
     }
   )
 }

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

@@ -0,0 +1,72 @@
+const getTenantId = jest.fn()
+jest.mock("@budibase/backend-core/tenancy", () => ({
+  getTenantId
+}))
+const usageQuota = require("../usageQuota")
+const env = require("../../environment")
+
+class TestConfiguration {
+  constructor() {
+    this.enableQuotas()
+  }
+
+  enableQuotas = () => {
+    env.USE_QUOTAS = 1
+  }
+
+  disableQuotas = () => {
+    env.USE_QUOTAS = null
+  }
+
+  setTenantId = (tenantId) => {
+    getTenantId.mockReturnValue(tenantId)
+  } 
+
+  setExcludedTenants = (tenants) => {
+    env.EXCLUDE_QUOTAS_TENANTS = tenants
+  }
+
+  reset = () => {
+    this.disableQuotas()
+    this.setExcludedTenants(null)
+  }
+}
+
+describe("usageQuota", () => {
+  let config
+
+  beforeEach(() => {
+    config = new TestConfiguration()
+  })
+
+  afterEach(() => {
+    config.reset()
+  })
+
+  describe("useQuotas", () => {
+    it("works when no settings have been provided", () => {
+      config.reset()
+      expect(usageQuota.useQuotas()).toBe(false)
+    })
+    it("honours USE_QUOTAS setting", () => {
+      config.disableQuotas()
+      expect(usageQuota.useQuotas()).toBe(false)
+
+      config.enableQuotas()
+      expect(usageQuota.useQuotas()).toBe(true)
+    })
+    it("honours EXCLUDE_QUOTAS_TENANTS setting", () => {
+      config.setTenantId("test")
+
+      // tenantId is in the list
+      config.setExcludedTenants("test, test2, test2")
+      expect(usageQuota.useQuotas()).toBe(false)
+      config.setExcludedTenants("test,test2,test2")
+      expect(usageQuota.useQuotas()).toBe(false)
+
+      // tenantId is not in the list
+      config.setTenantId("other")
+      expect(usageQuota.useQuotas()).toBe(true)
+    })
+  })
+})

packages/server/src/utilities/usageQuota.js

@@ -1,12 +1,31 @@
 const env = require("../environment")
-const { getGlobalDB } = require("@budibase/backend-core/tenancy")
+const { getGlobalDB, getTenantId } = require("@budibase/backend-core/tenancy")
 const {
   StaticDatabases,
   generateNewUsageQuotaDoc,
 } = require("@budibase/backend-core/db")
 
+exports.useQuotas = () => {
+  // check if quotas are enabled
+  if (env.USE_QUOTAS) {
+    // check if there are any tenants without limits
+    if (env.EXCLUDE_QUOTAS_TENANTS) {
+      const excludedTenants = env.EXCLUDE_QUOTAS_TENANTS.replace(
+        /\s/g,
+        ""
+      ).split(",")
+      const tenantId = getTenantId()
+      if (excludedTenants.includes(tenantId)) {
+        return false
+      }
+    }
+    return true
+  }
+  return false
+}
+
 exports.Properties = {
-  ROW: "rows", // mostly works - disabled - app / table deletion not yet accounted for
+  ROW: "rows", // mostly works - app / table deletion not yet accounted for
   UPLOAD: "storage", // doesn't work yet
   VIEW: "views", // doesn't work yet
   USER: "users", // doesn't work yet
@@ -37,7 +56,7 @@ exports.getUsageQuotaDoc = async db => {
  * also been reset after this call.
  */
 exports.update = async (property, usage) => {
-  if (!env.USE_QUOTAS) {
+  if (!exports.useQuotas()) {
     return
   }