Merge pull request #14116 from Budibase/sqs-per-tenant

Per-tenant enablement of SQS

packages/backend-core/src/db/utils.ts

@@ -206,3 +206,21 @@ export function pagination<T>(
     nextPage,
   }
 }
+
+export function isSqsEnabledForTenant(): boolean {
+  const tenantId = getTenantId()
+  if (!env.SQS_SEARCH_ENABLE) {
+    return false
+  }
+
+  // This is to guard against the situation in tests where tests pass because
+  // we're not actually using SQS, we're using Lucene and the tests pass due to
+  // parity.
+  if (env.isTest() && env.SQS_SEARCH_ENABLE_TENANTS.length === 0) {
+    throw new Error(
+      "to enable SQS you must specify a list of tenants in the SQS_SEARCH_ENABLE_TENANTS env var"
+    )
+  }
+
+  return env.SQS_SEARCH_ENABLE_TENANTS.includes(tenantId)
+}

packages/backend-core/src/environment.ts

@@ -116,6 +116,9 @@ const environment = {
   COUCH_DB_URL: process.env.COUCH_DB_URL || "http://localhost:4005",
   COUCH_DB_SQL_URL: process.env.COUCH_DB_SQL_URL || "http://localhost:4006",
   SQS_SEARCH_ENABLE: process.env.SQS_SEARCH_ENABLE,
+  SQS_SEARCH_ENABLE_TENANTS:
+    process.env.SQS_SEARCH_ENABLE_TENANTS?.split(",") || [],
+  SQS_MIGRATION_ENABLE: process.env.SQS_MIGRATION_ENABLE,
   COUCH_DB_USERNAME: process.env.COUCH_DB_USER,
   COUCH_DB_PASSWORD: process.env.COUCH_DB_PASSWORD,
   GOOGLE_CLIENT_ID: process.env.GOOGLE_CLIENT_ID,

packages/server/src/api/controllers/table/utils.ts

@@ -15,7 +15,7 @@ import { getViews, saveView } from "../view/utils"
 import viewTemplate from "../view/viewBuilder"
 import { cloneDeep } from "lodash/fp"
 import { quotas } from "@budibase/pro"
-import { events, context } from "@budibase/backend-core"
+import { events, context, db as dbCore } from "@budibase/backend-core"
 import {
   AutoFieldSubType,
   ContextUser,
@@ -324,7 +324,7 @@ class TableSaveFunctions {
       importRows: this.importRows,
       user: this.user,
     })
-    if (env.SQS_SEARCH_ENABLE) {
+    if (dbCore.isSqsEnabledForTenant()) {
       await sdk.tables.sqs.addTable(table)
     }
     return table
@@ -518,7 +518,7 @@ export async function internalTableCleanup(table: Table, rows?: Row[]) {
   if (rows) {
     await AttachmentCleanup.tableDelete(table, rows)
   }
-  if (env.SQS_SEARCH_ENABLE) {
+  if (dbCore.isSqsEnabledForTenant()) {
     await sdk.tables.sqs.removeTable(table)
   }
 }

packages/server/src/api/routes/tests/search.spec.ts

@@ -54,10 +54,13 @@ describe.each([
   let rows: Row[]
 
   beforeAll(async () => {
+    await config.withCoreEnv({ SQS_SEARCH_ENABLE: "true" }, () => config.init())
     if (isSqs) {
-      envCleanup = config.setEnv({ SQS_SEARCH_ENABLE: "true" })
+      envCleanup = config.setCoreEnv({
+        SQS_SEARCH_ENABLE: "true",
+        SQS_SEARCH_ENABLE_TENANTS: [config.getTenantId()],
+      })
     }
-    await config.init()
 
     if (config.app?.appId) {
       config.app = await config.api.application.update(config.app?.appId, {

packages/server/src/api/routes/tests/templates.spec.ts

@@ -86,9 +86,10 @@ describe("/templates", () => {
       async source => {
         const env = {
           SQS_SEARCH_ENABLE: source === "sqs" ? "true" : "false",
+          SQS_SEARCH_ENABLE_TENANTS: [config.getTenantId()],
         }
 
-        await config.withEnv(env, async () => {
+        await config.withCoreEnv(env, async () => {
           const name = generator.guid().replaceAll("-", "")
           const url = `/${name}`
 

packages/server/src/api/routes/tests/viewV2.spec.ts

@@ -88,10 +88,16 @@ describe.each([
   }
 
   beforeAll(async () => {
+    await config.withCoreEnv(
+      { SQS_SEARCH_ENABLE: isSqs ? "true" : "false" },
+      () => config.init()
+    )
     if (isSqs) {
-      envCleanup = config.setEnv({ SQS_SEARCH_ENABLE: "true" })
+      envCleanup = config.setCoreEnv({
+        SQS_SEARCH_ENABLE: "true",
+        SQS_SEARCH_ENABLE_TENANTS: [config.getTenantId()],
+      })
     }
-    await config.init()
 
     if (dsProvider) {
       datasource = await config.createDatasource({

packages/server/src/appMigrations/migrations.ts

@@ -1,6 +1,6 @@
 // This file should never be manually modified, use `yarn add-app-migration` in order to add a new one
 
-import env from "../environment"
+import { env } from "@budibase/backend-core"
 import { AppMigration } from "."
 
 import m20240604153647_initial_sqs from "./migrations/20240604153647_initial_sqs"

packages/server/src/appMigrations/migrations/20240604153647_initial_sqs.ts

@@ -1,8 +1,7 @@
-import { context } from "@budibase/backend-core"
+import { context, env } from "@budibase/backend-core"
 import { allLinkDocs } from "../../db/utils"
 import LinkDocumentImpl from "../../db/linkedRows/LinkDocument"
 import sdk from "../../sdk"
-import env from "../../environment"
 
 const migration = async () => {
   const linkDocs = await allLinkDocs()

packages/server/src/appMigrations/migrations/tests/20240604153647_initial_sqs.spec.ts

@@ -69,11 +69,14 @@ function oldLinkDocument(): Omit<LinkDocument, "tableId"> {
 type SQSEnvVar = "SQS_MIGRATION_ENABLE" | "SQS_SEARCH_ENABLE"
 
 async function sqsDisabled(envVar: SQSEnvVar, cb: () => Promise<void>) {
-  await config.withEnv({ [envVar]: "" }, cb)
+  await config.withCoreEnv({ [envVar]: "", SQS_SEARCH_ENABLE_TENANTS: [] }, cb)
 }
 
 async function sqsEnabled(envVar: SQSEnvVar, cb: () => Promise<void>) {
-  await config.withEnv({ [envVar]: "1" }, cb)
+  await config.withCoreEnv(
+    { [envVar]: "1", SQS_SEARCH_ENABLE_TENANTS: [config.getTenantId()] },
+    cb
+  )
 }
 
 describe.each(["SQS_MIGRATION_ENABLE", "SQS_SEARCH_ENABLE"] as SQSEnvVar[])(

packages/server/src/environment.ts

@@ -87,8 +87,6 @@ const environment = {
   SQL_MAX_ROWS: process.env.SQL_MAX_ROWS,
   SQL_LOGGING_ENABLE: process.env.SQL_LOGGING_ENABLE,
   SQL_ALIASING_DISABLE: process.env.SQL_ALIASING_DISABLE,
-  SQS_SEARCH_ENABLE: process.env.SQS_SEARCH_ENABLE,
-  SQS_MIGRATION_ENABLE: process.env.SQS_MIGRATION_ENABLE,
   // flags
   ALLOW_DEV_AUTOMATIONS: process.env.ALLOW_DEV_AUTOMATIONS,
   DISABLE_THREADING: process.env.DISABLE_THREADING,

packages/server/src/sdk/app/rows/search.ts

@@ -11,11 +11,11 @@ import * as internal from "./search/internal"
 import * as external from "./search/external"
 import { NoEmptyFilterStrings } from "../../../constants"
 import * as sqs from "./search/sqs"
-import env from "../../../environment"
 import { ExportRowsParams, ExportRowsResult } from "./search/types"
 import { dataFilters } from "@budibase/shared-core"
 import sdk from "../../index"
 import { searchInputMapping } from "./search/utils"
+import { db as dbCore } from "@budibase/backend-core"
 
 export { isValidFilter } from "../../../integrations/utils"
 
@@ -89,7 +89,7 @@ export async function search(
 
   if (isExternalTable) {
     return external.search(options, table)
-  } else if (env.SQS_SEARCH_ENABLE) {
+  } else if (dbCore.isSqsEnabledForTenant()) {
     return sqs.search(options, table)
   } else {
     return internal.search(options, table)

packages/server/src/sdk/app/rows/search/tests/search.spec.ts

@@ -31,10 +31,17 @@ describe.each([
   let rows: Row[]
 
   beforeAll(async () => {
+    await config.withCoreEnv(
+      { SQS_SEARCH_ENABLE: isSqs ? "true" : "false" },
+      () => config.init()
+    )
+
     if (isSqs) {
-      envCleanup = config.setEnv({ SQS_SEARCH_ENABLE: "true" })
+      envCleanup = config.setCoreEnv({
+        SQS_SEARCH_ENABLE: "true",
+        SQS_SEARCH_ENABLE_TENANTS: [config.getTenantId()],
+      })
     }
-    await config.init()
 
     if (dsProvider) {
       datasource = await config.createDatasource({

packages/server/src/sdk/app/tables/getters.ts

@@ -1,4 +1,4 @@
-import { context } from "@budibase/backend-core"
+import { context, db as dbCore, env } from "@budibase/backend-core"
 import { getTableParams } from "../../../db/utils"
 import {
   breakExternalTableId,
@@ -15,7 +15,6 @@ import {
 } from "@budibase/types"
 import datasources from "../datasources"
 import sdk from "../../../sdk"
-import env from "../../../environment"
 
 export function processTable(table: Table): Table {
   if (!table) {
@@ -34,7 +33,7 @@ export function processTable(table: Table): Table {
       sourceId: table.sourceId || INTERNAL_TABLE_SOURCE_ID,
       sourceType: TableSourceType.INTERNAL,
     }
-    if (env.SQS_SEARCH_ENABLE) {
+    if (dbCore.isSqsEnabledForTenant()) {
       processed.sql = !!env.SQS_SEARCH_ENABLE
     }
     return processed

packages/server/src/sdk/tests/attachments.spec.ts

@@ -39,7 +39,9 @@ describe("should be able to re-write attachment URLs", () => {
     }
 
     const db = dbCore.getDB(config.getAppId())
-    await sdk.backups.updateAttachmentColumns(db.name, db)
+    await config.doInContext(config.getAppId(), () =>
+      sdk.backups.updateAttachmentColumns(db.name, db)
+    )
 
     return {
       db,

packages/server/src/tests/utilities/TestConfiguration.ts

@@ -245,10 +245,10 @@ export default class TestConfiguration {
     }
   }
 
-  async withEnv(newEnvVars: Partial<typeof env>, f: () => Promise<void>) {
+  async withEnv<T>(newEnvVars: Partial<typeof env>, f: () => Promise<T>) {
     let cleanup = this.setEnv(newEnvVars)
     try {
-      await f()
+      return await f()
     } finally {
       cleanup()
     }
@@ -273,13 +273,13 @@ export default class TestConfiguration {
     }
   }
 
-  async withCoreEnv(
+  async withCoreEnv<T>(
     newEnvVars: Partial<typeof coreEnv>,
-    f: () => Promise<void>
+    f: () => Promise<T>
   ) {
     let cleanup = this.setCoreEnv(newEnvVars)
     try {
-      await f()
+      return await f()
     } finally {
       cleanup()
     }

packages/worker/src/api/controllers/system/environment.ts

@@ -24,7 +24,7 @@ async function isSqsAvailable() {
 }
 
 async function isSqsMissing() {
-  return env.SQS_SEARCH_ENABLE && !(await isSqsAvailable())
+  return coreEnv.SQS_SEARCH_ENABLE && !(await isSqsAvailable())
 }
 
 export const fetch = async (ctx: Ctx) => {

packages/worker/src/api/index.ts

@@ -5,8 +5,7 @@ const compress = require("koa-compress")
 import zlib from "zlib"
 import { routes } from "./routes"
 import { middleware as pro, sdk } from "@budibase/pro"
-import { auth, middleware } from "@budibase/backend-core"
+import { auth, middleware, env } from "@budibase/backend-core"
-import env from "../environment"
 
 if (env.SQS_SEARCH_ENABLE) {
   sdk.auditLogs.useSQLSearch()

packages/worker/src/environment.ts

@@ -46,7 +46,6 @@ const environment = {
   DISABLE_ACCOUNT_PORTAL: process.env.DISABLE_ACCOUNT_PORTAL,
   SMTP_FALLBACK_ENABLED: process.env.SMTP_FALLBACK_ENABLED,
   DISABLE_DEVELOPER_LICENSE: process.env.DISABLE_DEVELOPER_LICENSE,
-  SQS_SEARCH_ENABLE: process.env.SQS_SEARCH_ENABLE,
   BUDIBASE_ENVIRONMENT: process.env.BUDIBASE_ENVIRONMENT,
   // smtp
   SMTP_USER: process.env.SMTP_USER,