diff --git a/packages/server/src/middleware/usageQuota.js b/packages/server/src/middleware/usageQuota.js index fb2e0722a7..2d55aac61d 100644 --- a/packages/server/src/middleware/usageQuota.js +++ b/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 = {} 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 61e6385aa6..5445919f26 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 @@ -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 diff --git a/packages/server/src/utilities/tests/usageQuota.spec.js b/packages/server/src/utilities/tests/usageQuota/usageQuota.spec.js similarity index 97% rename from packages/server/src/utilities/tests/usageQuota.spec.js rename to packages/server/src/utilities/tests/usageQuota/usageQuota.spec.js index 6023b08fbe..764d509b5e 100644 --- a/packages/server/src/utilities/tests/usageQuota.spec.js +++ b/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() { diff --git a/packages/server/src/utilities/usageQuota.js b/packages/server/src/utilities/usageQuota/index.js similarity index 81% rename from packages/server/src/utilities/usageQuota.js rename to packages/server/src/utilities/usageQuota/index.js index 5965334205..110f8ef600 100644 --- a/packages/server/src/utilities/usageQuota.js +++ b/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} 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 diff --git a/packages/server/src/utilities/usageQuota/rows.js b/packages/server/src/utilities/usageQuota/rows.js new file mode 100644 index 0000000000..3a5602958b --- /dev/null +++ b/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) +} diff --git a/packages/server/src/utilities/usageQuoteReset.js b/packages/server/src/utilities/usageQuota/usageQuoteReset.js similarity index 100% rename from packages/server/src/utilities/usageQuoteReset.js rename to packages/server/src/utilities/usageQuota/usageQuoteReset.js