Merge pull request #14725 from Budibase/feature-flag-helper

Test helper for setting feature flags.

packages/backend-core/src/features/features.ts

@@ -0,0 +1,300 @@
+import env from "../environment"
+import * as context from "../context"
+import { PostHog, PostHogOptions } from "posthog-node"
+import { FeatureFlag, IdentityType, UserCtx } from "@budibase/types"
+import tracer from "dd-trace"
+import { Duration } from "../utils"
+
+let posthog: PostHog | undefined
+export function init(opts?: PostHogOptions) {
+  if (
+    env.POSTHOG_TOKEN &&
+    env.POSTHOG_API_HOST &&
+    !env.SELF_HOSTED &&
+    env.POSTHOG_FEATURE_FLAGS_ENABLED
+  ) {
+    console.log("initializing posthog client...")
+    posthog = new PostHog(env.POSTHOG_TOKEN, {
+      host: env.POSTHOG_API_HOST,
+      personalApiKey: env.POSTHOG_PERSONAL_TOKEN,
+      featureFlagsPollingInterval: Duration.fromMinutes(3).toMs(),
+      ...opts,
+    })
+  } else {
+    console.log("posthog disabled")
+  }
+}
+
+export function shutdown() {
+  posthog?.shutdown()
+}
+
+export abstract class Flag<T> {
+  static boolean(defaultValue: boolean): Flag<boolean> {
+    return new BooleanFlag(defaultValue)
+  }
+
+  static string(defaultValue: string): Flag<string> {
+    return new StringFlag(defaultValue)
+  }
+
+  static number(defaultValue: number): Flag<number> {
+    return new NumberFlag(defaultValue)
+  }
+
+  protected constructor(public defaultValue: T) {}
+
+  abstract parse(value: any): T
+}
+
+type UnwrapFlag<F> = F extends Flag<infer U> ? U : never
+
+export type FlagValues<T> = {
+  [K in keyof T]: UnwrapFlag<T[K]>
+}
+
+type KeysOfType<T, U> = {
+  [K in keyof T]: T[K] extends Flag<U> ? K : never
+}[keyof T]
+
+class BooleanFlag extends Flag<boolean> {
+  parse(value: any) {
+    if (typeof value === "string") {
+      return ["true", "t", "1"].includes(value.toLowerCase())
+    }
+
+    if (typeof value === "boolean") {
+      return value
+    }
+
+    throw new Error(`could not parse value "${value}" as boolean`)
+  }
+}
+
+class StringFlag extends Flag<string> {
+  parse(value: any) {
+    if (typeof value === "string") {
+      return value
+    }
+    throw new Error(`could not parse value "${value}" as string`)
+  }
+}
+
+class NumberFlag extends Flag<number> {
+  parse(value: any) {
+    if (typeof value === "number") {
+      return value
+    }
+
+    if (typeof value === "string") {
+      const parsed = parseFloat(value)
+      if (!isNaN(parsed)) {
+        return parsed
+      }
+    }
+
+    throw new Error(`could not parse value "${value}" as number`)
+  }
+}
+
+export interface EnvFlagEntry {
+  tenantId: string
+  key: string
+  value: boolean
+}
+
+export function parseEnvFlags(flags: string): EnvFlagEntry[] {
+  const split = flags.split(",").map(x => x.split(":"))
+  const result: EnvFlagEntry[] = []
+  for (const [tenantId, ...features] of split) {
+    for (let feature of features) {
+      let value = true
+      if (feature.startsWith("!")) {
+        feature = feature.slice(1)
+        value = false
+      }
+      result.push({ tenantId, key: feature, value })
+    }
+  }
+  return result
+}
+
+export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
+  // This is used to safely cache flags sets in the current request context.
+  // Because multiple sets could theoretically exist, we don't want the cache of
+  // one to leak into another.
+  private readonly setId: string
+
+  constructor(private readonly flagSchema: T) {
+    this.setId = crypto.randomUUID()
+  }
+
+  defaults(): FlagValues<T> {
+    return Object.keys(this.flagSchema).reduce((acc, key) => {
+      const typedKey = key as keyof T
+      acc[typedKey] = this.flagSchema[key].defaultValue
+      return acc
+    }, {} as FlagValues<T>)
+  }
+
+  isFlagName(name: string | number | symbol): name is keyof T {
+    return this.flagSchema[name as keyof T] !== undefined
+  }
+
+  async get<K extends keyof T>(
+    key: K,
+    ctx?: UserCtx
+  ): Promise<FlagValues<T>[K]> {
+    const flags = await this.fetch(ctx)
+    return flags[key]
+  }
+
+  async isEnabled<K extends KeysOfType<T, boolean>>(
+    key: K,
+    ctx?: UserCtx
+  ): Promise<boolean> {
+    const flags = await this.fetch(ctx)
+    return flags[key]
+  }
+
+  async fetch(ctx?: UserCtx): Promise<FlagValues<T>> {
+    return await tracer.trace("features.fetch", async span => {
+      const cachedFlags = context.getFeatureFlags<FlagValues<T>>(this.setId)
+      if (cachedFlags) {
+        span?.addTags({ fromCache: true })
+        return cachedFlags
+      }
+
+      const tags: Record<string, any> = {}
+      const flagValues = this.defaults()
+      const currentTenantId = context.getTenantId()
+      const specificallySetFalse = new Set<string>()
+
+      for (const { tenantId, key, value } of parseEnvFlags(
+        env.TENANT_FEATURE_FLAGS || ""
+      )) {
+        if (!tenantId || (tenantId !== "*" && tenantId !== currentTenantId)) {
+          continue
+        }
+
+        tags[`readFromEnvironmentVars`] = true
+
+        if (value === false) {
+          specificallySetFalse.add(key)
+        }
+
+        // ignore unknown flags
+        if (!this.isFlagName(key)) {
+          continue
+        }
+
+        if (typeof flagValues[key] !== "boolean") {
+          throw new Error(`Feature: ${key} is not a boolean`)
+        }
+
+        // @ts-expect-error - TS does not like you writing into a generic type,
+        // but we know that it's okay in this case because it's just an object.
+        flagValues[key as keyof FlagValues] = value
+        tags[`flags.${key}.source`] = "environment"
+      }
+
+      const license = ctx?.user?.license
+      if (license) {
+        tags[`readFromLicense`] = true
+
+        for (const feature of license.features) {
+          if (!this.isFlagName(feature)) {
+            continue
+          }
+
+          if (
+            flagValues[feature] === true ||
+            specificallySetFalse.has(feature)
+          ) {
+            // If the flag is already set to through environment variables, we
+            // don't want to override it back to false here.
+            continue
+          }
+
+          // @ts-expect-error - TS does not like you writing into a generic type,
+          // but we know that it's okay in this case because it's just an object.
+          flagValues[feature] = true
+          tags[`flags.${feature}.source`] = "license"
+        }
+      }
+
+      const identity = context.getIdentity()
+      tags[`identity.type`] = identity?.type
+      tags[`identity.tenantId`] = identity?.tenantId
+      tags[`identity._id`] = identity?._id
+
+      if (posthog && identity?.type === IdentityType.USER) {
+        tags[`readFromPostHog`] = true
+
+        const personProperties: Record<string, string> = {}
+        if (identity.tenantId) {
+          personProperties.tenantId = identity.tenantId
+        }
+
+        const posthogFlags = await posthog.getAllFlagsAndPayloads(
+          identity._id,
+          {
+            personProperties,
+          }
+        )
+
+        for (const [name, value] of Object.entries(posthogFlags.featureFlags)) {
+          if (!this.isFlagName(name)) {
+            // We don't want an unexpected PostHog flag to break the app, so we
+            // just log it and continue.
+            console.warn(`Unexpected posthog flag "${name}": ${value}`)
+            continue
+          }
+
+          if (flagValues[name] === true || specificallySetFalse.has(name)) {
+            // If the flag is already set to through environment variables, we
+            // don't want to override it back to false here.
+            continue
+          }
+
+          const payload = posthogFlags.featureFlagPayloads?.[name]
+          const flag = this.flagSchema[name]
+          try {
+            // @ts-expect-error - TS does not like you writing into a generic
+            // type, but we know that it's okay in this case because it's just
+            // an object.
+            flagValues[name] = flag.parse(payload || value)
+            tags[`flags.${name}.source`] = "posthog"
+          } catch (err) {
+            // We don't want an invalid PostHog flag to break the app, so we just
+            // log it and continue.
+            console.warn(`Error parsing posthog flag "${name}": ${value}`, err)
+          }
+        }
+      }
+
+      context.setFeatureFlags(this.setId, flagValues)
+      for (const [key, value] of Object.entries(flagValues)) {
+        tags[`flags.${key}.value`] = value
+      }
+      span?.addTags(tags)
+
+      return flagValues
+    })
+  }
+}
+
+// This is the primary source of truth for feature flags. If you want to add a
+// new flag, add it here and use the `fetch` and `get` functions to access it.
+// All of the machinery in this file is to make sure that flags have their
+// default values set correctly and their types flow through the system.
+export const flags = new FlagSet({
+  DEFAULT_VALUES: Flag.boolean(env.isDev()),
+  AUTOMATION_BRANCHING: Flag.boolean(env.isDev()),
+  SQS: Flag.boolean(env.isDev()),
+  [FeatureFlag.AI_CUSTOM_CONFIGS]: Flag.boolean(env.isDev()),
+  [FeatureFlag.ENRICHED_RELATIONSHIPS]: Flag.boolean(env.isDev()),
+})
+
+type UnwrapPromise<T> = T extends Promise<infer U> ? U : T
+export type FeatureFlags = UnwrapPromise<ReturnType<typeof flags.fetch>>

packages/backend-core/src/features/index.ts

@@ -1,281 +1,2 @@
-import env from "../environment"
+export * from "./features"
-import * as context from "../context"
+export * as testutils from "./tests/utils"
-import { PostHog, PostHogOptions } from "posthog-node"
-import { FeatureFlag, IdentityType, UserCtx } from "@budibase/types"
-import tracer from "dd-trace"
-import { Duration } from "../utils"
-
-let posthog: PostHog | undefined
-export function init(opts?: PostHogOptions) {
-  if (
-    env.POSTHOG_TOKEN &&
-    env.POSTHOG_API_HOST &&
-    !env.SELF_HOSTED &&
-    env.POSTHOG_FEATURE_FLAGS_ENABLED
-  ) {
-    console.log("initializing posthog client...")
-    posthog = new PostHog(env.POSTHOG_TOKEN, {
-      host: env.POSTHOG_API_HOST,
-      personalApiKey: env.POSTHOG_PERSONAL_TOKEN,
-      featureFlagsPollingInterval: Duration.fromMinutes(3).toMs(),
-      ...opts,
-    })
-  } else {
-    console.log("posthog disabled")
-  }
-}
-
-export function shutdown() {
-  posthog?.shutdown()
-}
-
-export abstract class Flag<T> {
-  static boolean(defaultValue: boolean): Flag<boolean> {
-    return new BooleanFlag(defaultValue)
-  }
-
-  static string(defaultValue: string): Flag<string> {
-    return new StringFlag(defaultValue)
-  }
-
-  static number(defaultValue: number): Flag<number> {
-    return new NumberFlag(defaultValue)
-  }
-
-  protected constructor(public defaultValue: T) {}
-
-  abstract parse(value: any): T
-}
-
-type UnwrapFlag<F> = F extends Flag<infer U> ? U : never
-
-export type FlagValues<T> = {
-  [K in keyof T]: UnwrapFlag<T[K]>
-}
-
-type KeysOfType<T, U> = {
-  [K in keyof T]: T[K] extends Flag<U> ? K : never
-}[keyof T]
-
-class BooleanFlag extends Flag<boolean> {
-  parse(value: any) {
-    if (typeof value === "string") {
-      return ["true", "t", "1"].includes(value.toLowerCase())
-    }
-
-    if (typeof value === "boolean") {
-      return value
-    }
-
-    throw new Error(`could not parse value "${value}" as boolean`)
-  }
-}
-
-class StringFlag extends Flag<string> {
-  parse(value: any) {
-    if (typeof value === "string") {
-      return value
-    }
-    throw new Error(`could not parse value "${value}" as string`)
-  }
-}
-
-class NumberFlag extends Flag<number> {
-  parse(value: any) {
-    if (typeof value === "number") {
-      return value
-    }
-
-    if (typeof value === "string") {
-      const parsed = parseFloat(value)
-      if (!isNaN(parsed)) {
-        return parsed
-      }
-    }
-
-    throw new Error(`could not parse value "${value}" as number`)
-  }
-}
-
-export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
-  // This is used to safely cache flags sets in the current request context.
-  // Because multiple sets could theoretically exist, we don't want the cache of
-  // one to leak into another.
-  private readonly setId: string
-
-  constructor(private readonly flagSchema: T) {
-    this.setId = crypto.randomUUID()
-  }
-
-  defaults(): FlagValues<T> {
-    return Object.keys(this.flagSchema).reduce((acc, key) => {
-      const typedKey = key as keyof T
-      acc[typedKey] = this.flagSchema[key].defaultValue
-      return acc
-    }, {} as FlagValues<T>)
-  }
-
-  isFlagName(name: string | number | symbol): name is keyof T {
-    return this.flagSchema[name as keyof T] !== undefined
-  }
-
-  async get<K extends keyof T>(
-    key: K,
-    ctx?: UserCtx
-  ): Promise<FlagValues<T>[K]> {
-    const flags = await this.fetch(ctx)
-    return flags[key]
-  }
-
-  async isEnabled<K extends KeysOfType<T, boolean>>(
-    key: K,
-    ctx?: UserCtx
-  ): Promise<boolean> {
-    const flags = await this.fetch(ctx)
-    return flags[key]
-  }
-
-  async fetch(ctx?: UserCtx): Promise<FlagValues<T>> {
-    return await tracer.trace("features.fetch", async span => {
-      const cachedFlags = context.getFeatureFlags<FlagValues<T>>(this.setId)
-      if (cachedFlags) {
-        span?.addTags({ fromCache: true })
-        return cachedFlags
-      }
-
-      const tags: Record<string, any> = {}
-      const flagValues = this.defaults()
-      const currentTenantId = context.getTenantId()
-      const specificallySetFalse = new Set<string>()
-
-      const split = (env.TENANT_FEATURE_FLAGS || "")
-        .split(",")
-        .map(x => x.split(":"))
-      for (const [tenantId, ...features] of split) {
-        if (!tenantId || (tenantId !== "*" && tenantId !== currentTenantId)) {
-          continue
-        }
-
-        tags[`readFromEnvironmentVars`] = true
-
-        for (let feature of features) {
-          let value = true
-          if (feature.startsWith("!")) {
-            feature = feature.slice(1)
-            value = false
-            specificallySetFalse.add(feature)
-          }
-
-          // ignore unknown flags
-          if (!this.isFlagName(feature)) {
-            continue
-          }
-
-          if (typeof flagValues[feature] !== "boolean") {
-            throw new Error(`Feature: ${feature} is not a boolean`)
-          }
-
-          // @ts-expect-error - TS does not like you writing into a generic type,
-          // but we know that it's okay in this case because it's just an object.
-          flagValues[feature as keyof FlagValues] = value
-          tags[`flags.${feature}.source`] = "environment"
-        }
-      }
-
-      const license = ctx?.user?.license
-      if (license) {
-        tags[`readFromLicense`] = true
-
-        for (const feature of license.features) {
-          if (!this.isFlagName(feature)) {
-            continue
-          }
-
-          if (
-            flagValues[feature] === true ||
-            specificallySetFalse.has(feature)
-          ) {
-            // If the flag is already set to through environment variables, we
-            // don't want to override it back to false here.
-            continue
-          }
-
-          // @ts-expect-error - TS does not like you writing into a generic type,
-          // but we know that it's okay in this case because it's just an object.
-          flagValues[feature] = true
-          tags[`flags.${feature}.source`] = "license"
-        }
-      }
-
-      const identity = context.getIdentity()
-      tags[`identity.type`] = identity?.type
-      tags[`identity.tenantId`] = identity?.tenantId
-      tags[`identity._id`] = identity?._id
-
-      if (posthog && identity?.type === IdentityType.USER) {
-        tags[`readFromPostHog`] = true
-
-        const personProperties: Record<string, string> = {}
-        if (identity.tenantId) {
-          personProperties.tenantId = identity.tenantId
-        }
-
-        const posthogFlags = await posthog.getAllFlagsAndPayloads(
-          identity._id,
-          {
-            personProperties,
-          }
-        )
-
-        for (const [name, value] of Object.entries(posthogFlags.featureFlags)) {
-          if (!this.isFlagName(name)) {
-            // We don't want an unexpected PostHog flag to break the app, so we
-            // just log it and continue.
-            console.warn(`Unexpected posthog flag "${name}": ${value}`)
-            continue
-          }
-
-          if (flagValues[name] === true || specificallySetFalse.has(name)) {
-            // If the flag is already set to through environment variables, we
-            // don't want to override it back to false here.
-            continue
-          }
-
-          const payload = posthogFlags.featureFlagPayloads?.[name]
-          const flag = this.flagSchema[name]
-          try {
-            // @ts-expect-error - TS does not like you writing into a generic
-            // type, but we know that it's okay in this case because it's just
-            // an object.
-            flagValues[name] = flag.parse(payload || value)
-            tags[`flags.${name}.source`] = "posthog"
-          } catch (err) {
-            // We don't want an invalid PostHog flag to break the app, so we just
-            // log it and continue.
-            console.warn(`Error parsing posthog flag "${name}": ${value}`, err)
-          }
-        }
-      }
-
-      context.setFeatureFlags(this.setId, flagValues)
-      for (const [key, value] of Object.entries(flagValues)) {
-        tags[`flags.${key}.value`] = value
-      }
-      span?.addTags(tags)
-
-      return flagValues
-    })
-  }
-}
-
-// This is the primary source of truth for feature flags. If you want to add a
-// new flag, add it here and use the `fetch` and `get` functions to access it.
-// All of the machinery in this file is to make sure that flags have their
-// default values set correctly and their types flow through the system.
-export const flags = new FlagSet({
-  DEFAULT_VALUES: Flag.boolean(env.isDev()),
-  AUTOMATION_BRANCHING: Flag.boolean(env.isDev()),
-  SQS: Flag.boolean(env.isDev()),
-  [FeatureFlag.AI_CUSTOM_CONFIGS]: Flag.boolean(env.isDev()),
-  [FeatureFlag.ENRICHED_RELATIONSHIPS]: Flag.boolean(env.isDev()),
-})

packages/backend-core/src/features/tests/utils.ts

@@ -0,0 +1,64 @@
+import { FeatureFlags, parseEnvFlags } from ".."
+import { setEnv } from "../../environment"
+
+function getCurrentFlags(): Record<string, Record<string, boolean>> {
+  const result: Record<string, Record<string, boolean>> = {}
+  for (const { tenantId, key, value } of parseEnvFlags(
+    process.env.TENANT_FEATURE_FLAGS || ""
+  )) {
+    const tenantFlags = result[tenantId] || {}
+    // Don't allow overwriting specifically false flags, to match the beheaviour
+    // of FlagSet.
+    if (tenantFlags[key] === false) {
+      continue
+    }
+    tenantFlags[key] = value
+    result[tenantId] = tenantFlags
+  }
+  return result
+}
+
+function buildFlagString(
+  flags: Record<string, Record<string, boolean>>
+): string {
+  const parts: string[] = []
+  for (const [tenantId, tenantFlags] of Object.entries(flags)) {
+    for (const [key, value] of Object.entries(tenantFlags)) {
+      if (value === false) {
+        parts.push(`${tenantId}:!${key}`)
+      } else {
+        parts.push(`${tenantId}:${key}`)
+      }
+    }
+  }
+  return parts.join(",")
+}
+
+export function setFeatureFlags(
+  tenantId: string,
+  flags: Partial<FeatureFlags>
+): () => void {
+  const current = getCurrentFlags()
+  for (const [key, value] of Object.entries(flags)) {
+    const tenantFlags = current[tenantId] || {}
+    tenantFlags[key] = value
+    current[tenantId] = tenantFlags
+  }
+  const flagString = buildFlagString(current)
+  return setEnv({ TENANT_FEATURE_FLAGS: flagString })
+}
+
+export function withFeatureFlags<T>(
+  tenantId: string,
+  flags: Partial<FeatureFlags>,
+  f: () => T
+) {
+  const cleanup = setFeatureFlags(tenantId, flags)
+  const result = f()
+  if (result instanceof Promise) {
+    return result.finally(cleanup)
+  } else {
+    cleanup()
+    return result
+  }
+}

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

@@ -14,12 +14,7 @@ jest.mock("../../../utilities/redis", () => ({
 import { checkBuilderEndpoint } from "./utilities/TestFunctions"
 import * as setup from "./utilities"
 import { AppStatus } from "../../../db/utils"
-import {
+import { events, utils, context, features } from "@budibase/backend-core"
-  events,
-  utils,
-  context,
-  withEnv as withCoreEnv,
-} from "@budibase/backend-core"
 import env from "../../../environment"
 import { type App } from "@budibase/types"
 import tk from "timekeeper"
@@ -358,9 +353,13 @@ describe("/applications", () => {
         .delete(`/api/global/roles/${prodAppId}`)
         .reply(200, {})
 
-      await withCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" }, async () => {
+      await features.testutils.withFeatureFlags(
+        "*",
+        { SQS: true },
+        async () => {
           await config.api.application.delete(app.appId)
-      })
+        }
+      )
     })
   })
 

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

@@ -13,8 +13,7 @@ import {
   context,
   InternalTable,
   tenancy,
-  withEnv as withCoreEnv,
+  features,
-  setEnv as setCoreEnv,
 } from "@budibase/backend-core"
 import { quotas } from "@budibase/pro"
 import {
@@ -40,7 +39,6 @@ import {
   TableSchema,
   JsonFieldSubType,
   RowExportFormat,
-  FeatureFlag,
   RelationSchemaField,
 } from "@budibase/types"
 import { generator, mocks } from "@budibase/backend-core/tests"
@@ -98,12 +96,12 @@ describe.each([
   let envCleanup: (() => void) | undefined
 
   beforeAll(async () => {
-    await withCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" }, () => config.init())
+    await features.testutils.withFeatureFlags("*", { SQS: true }, () =>
-    if (isLucene) {
+      config.init()
-      envCleanup = setCoreEnv({ TENANT_FEATURE_FLAGS: "*:!SQS" })
+    )
-    } else if (isSqs) {
+    envCleanup = features.testutils.setFeatureFlags("*", {
-      envCleanup = setCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" })
+      SQS: isSqs,
-    }
+    })
 
     if (dsProvider) {
       const rawDatasource = await dsProvider
@@ -2517,15 +2515,9 @@ describe.each([
       let flagCleanup: (() => void) | undefined
 
       beforeAll(async () => {
-        const env = {
+        flagCleanup = features.testutils.setFeatureFlags("*", {
-          TENANT_FEATURE_FLAGS: `*:${FeatureFlag.ENRICHED_RELATIONSHIPS}`,
+          ENRICHED_RELATIONSHIPS: true,
-        }
+        })
-        if (isSqs) {
-          env.TENANT_FEATURE_FLAGS = `${env.TENANT_FEATURE_FLAGS},*:SQS`
-        } else {
-          env.TENANT_FEATURE_FLAGS = `${env.TENANT_FEATURE_FLAGS},*:!SQS`
-        }
-        flagCleanup = setCoreEnv(env)
 
         const aux2Table = await config.api.table.save(saveTableRequest())
         const aux2Data = await config.api.row.save(aux2Table._id!, {})
@@ -2752,9 +2744,10 @@ describe.each([
       it.each(testScenarios)(
         "does not enrich relationships when not enabled (via %s)",
         async (__, retrieveDelegate) => {
-          await withCoreEnv(
+          await features.testutils.withFeatureFlags(
+            "*",
             {
-              TENANT_FEATURE_FLAGS: `*:!${FeatureFlag.ENRICHED_RELATIONSHIPS}`,
+              ENRICHED_RELATIONSHIPS: false,
             },
             async () => {
               const otherRows = _.sampleSize(auxData, 5)

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

@@ -7,9 +7,9 @@ import {
 import {
   context,
   db as dbCore,
+  features,
   MAX_VALID_DATE,
   MIN_VALID_DATE,
-  setEnv as setCoreEnv,
   SQLITE_DESIGN_DOC_ID,
   utils,
   withEnv as withCoreEnv,
@@ -94,16 +94,12 @@ describe.each([
   }
 
   beforeAll(async () => {
-    await withCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" }, () => config.init())
+    await features.testutils.withFeatureFlags("*", { SQS: true }, () =>
-    if (isLucene) {
+      config.init()
-      envCleanup = setCoreEnv({
+    )
-        TENANT_FEATURE_FLAGS: "*:!SQS",
+    envCleanup = features.testutils.setFeatureFlags("*", {
+      SQS: isSqs,
     })
-    } else if (isSqs) {
-      envCleanup = setCoreEnv({
-        TENANT_FEATURE_FLAGS: "*:SQS",
-      })
-    }
 
     if (config.app?.appId) {
       config.app = await config.api.application.update(config.app?.appId, {

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

@@ -2,7 +2,7 @@ import * as setup from "./utilities"
 import path from "path"
 import nock from "nock"
 import { generator } from "@budibase/backend-core/tests"
-import { withEnv as withCoreEnv, env as coreEnv } from "@budibase/backend-core"
+import { features } from "@budibase/backend-core"
 
 interface App {
   background: string
@@ -85,11 +85,10 @@ describe("/templates", () => {
     it.each(["sqs", "lucene"])(
       `should be able to create an app from a template (%s)`,
       async source => {
-        const env: Partial<typeof coreEnv> = {
+        await features.testutils.withFeatureFlags(
-          TENANT_FEATURE_FLAGS: source === "sqs" ? "*:SQS" : "",
+          "*",
-        }
+          { SQS: source === "sqs" },
-
+          async () => {
-        await withCoreEnv(env, async () => {
             const name = generator.guid().replaceAll("-", "")
             const url = `/${name}`
 
@@ -112,14 +111,18 @@ describe("/templates", () => {
               expect(agencyProjects.name).toBe("Agency Projects")
               expect(users.name).toBe("Users")
 
-            const { rows } = await config.api.row.search(agencyProjects._id!, {
+              const { rows } = await config.api.row.search(
+                agencyProjects._id!,
+                {
                   tableId: agencyProjects._id!,
                   query: {},
-            })
+                }
+              )
 
               expect(rows).toHaveLength(3)
             })
-        })
+          }
+        )
       }
     )
   })

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

@@ -22,7 +22,6 @@ import {
   RelationshipType,
   TableSchema,
   RenameColumn,
-  FeatureFlag,
   BBReferenceFieldSubType,
   NumericCalculationFieldMetadata,
   ViewV2Schema,
@@ -32,13 +31,7 @@ import { generator, mocks } from "@budibase/backend-core/tests"
 import { DatabaseName, getDatasource } from "../../../integrations/tests/utils"
 import merge from "lodash/merge"
 import { quotas } from "@budibase/pro"
-import {
+import { db, roles, features } from "@budibase/backend-core"
-  db,
-  roles,
-  withEnv as withCoreEnv,
-  setEnv as setCoreEnv,
-  env,
-} from "@budibase/backend-core"
 
 describe.each([
   ["lucene", undefined],
@@ -103,18 +96,13 @@ describe.each([
   }
 
   beforeAll(async () => {
-    await withCoreEnv({ TENANT_FEATURE_FLAGS: isSqs ? "*:SQS" : "" }, () =>
+    await features.testutils.withFeatureFlags("*", { SQS: isSqs }, () =>
       config.init()
     )
-    if (isLucene) {
+
-      envCleanup = setCoreEnv({
+    envCleanup = features.testutils.setFeatureFlags("*", {
-        TENANT_FEATURE_FLAGS: "*:!SQS",
+      SQS: isSqs,
     })
-    } else if (isSqs) {
-      envCleanup = setCoreEnv({
-        TENANT_FEATURE_FLAGS: "*:SQS",
-      })
-    }
 
     if (dsProvider) {
       datasource = await config.createDatasource({
@@ -2666,12 +2654,8 @@ describe.each([
       describe("foreign relationship columns", () => {
         let envCleanup: () => void
         beforeAll(() => {
-          const flags = [`*:${FeatureFlag.ENRICHED_RELATIONSHIPS}`]
+          envCleanup = features.testutils.setFeatureFlags("*", {
-          if (env.TENANT_FEATURE_FLAGS) {
+            ENRICHED_RELATIONSHIPS: true,
-            flags.push(...env.TENANT_FEATURE_FLAGS.split(","))
-          }
-          envCleanup = setCoreEnv({
-            TENANT_FEATURE_FLAGS: flags.join(","),
           })
         })
 

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

@@ -2,8 +2,8 @@ import * as setup from "../../../api/routes/tests/utilities"
 import { basicTable } from "../../../tests/utilities/structures"
 import {
   db as dbCore,
+  features,
   SQLITE_DESIGN_DOC_ID,
-  withEnv as withCoreEnv,
 } from "@budibase/backend-core"
 import {
   LinkDocument,
@@ -71,11 +71,11 @@ function oldLinkDocument(): Omit<LinkDocument, "tableId"> {
 }
 
 async function sqsDisabled(cb: () => Promise<void>) {
-  await withCoreEnv({ TENANT_FEATURE_FLAGS: "*:!SQS" }, cb)
+  await features.testutils.withFeatureFlags("*", { SQS: false }, cb)
 }
 
 async function sqsEnabled(cb: () => Promise<void>) {
-  await withCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" }, cb)
+  await features.testutils.withFeatureFlags("*", { SQS: true }, cb)
 }
 
 describe("SQS migration", () => {

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

@@ -10,10 +10,7 @@ import {
 import TestConfiguration from "../../../../../tests/utilities/TestConfiguration"
 import { search } from "../../../../../sdk/app/rows/search"
 import { generator } from "@budibase/backend-core/tests"
-import {
+import { features } from "@budibase/backend-core"
-  withEnv as withCoreEnv,
-  setEnv as setCoreEnv,
-} from "@budibase/backend-core"
 import {
   DatabaseName,
   getDatasource,
@@ -41,19 +38,13 @@ describe.each([
   let table: Table
 
   beforeAll(async () => {
-    await withCoreEnv({ TENANT_FEATURE_FLAGS: isSqs ? "*:SQS" : "" }, () =>
+    await features.testutils.withFeatureFlags("*", { SQS: isSqs }, () =>
       config.init()
     )
 
-    if (isLucene) {
+    envCleanup = features.testutils.setFeatureFlags("*", {
-      envCleanup = setCoreEnv({
+      SQS: isSqs,
-        TENANT_FEATURE_FLAGS: "*:!SQS",
     })
-    } else if (isSqs) {
-      envCleanup = setCoreEnv({
-        TENANT_FEATURE_FLAGS: "*:SQS",
-      })
-    }
 
     if (dsProvider) {
       datasource = await config.createDatasource({

packages/server/src/utilities/rowProcessor/tests/outputProcessing.spec.ts

@@ -8,7 +8,7 @@ import {
 } from "@budibase/types"
 import { outputProcessing } from ".."
 import { generator, structures } from "@budibase/backend-core/tests"
-import { setEnv as setCoreEnv } from "@budibase/backend-core"
+import { features } from "@budibase/backend-core"
 import * as bbReferenceProcessor from "../bbReferenceProcessor"
 import TestConfiguration from "../../../tests/utilities/TestConfiguration"
 
@@ -21,7 +21,7 @@ jest.mock("../bbReferenceProcessor", (): typeof bbReferenceProcessor => ({
 
 describe("rowProcessor - outputProcessing", () => {
   const config = new TestConfiguration()
-  let cleanupEnv: () => void = () => {}
+  let cleanupFlags: () => void = () => {}
 
   beforeAll(async () => {
     await config.init()
@@ -33,11 +33,11 @@ describe("rowProcessor - outputProcessing", () => {
 
   beforeEach(() => {
     jest.resetAllMocks()
-    cleanupEnv = setCoreEnv({ TENANT_FEATURE_FLAGS: "*SQS" })
+    cleanupFlags = features.testutils.setFeatureFlags("*", { SQS: true })
   })
 
   afterEach(() => {
-    cleanupEnv()
+    cleanupFlags()
   })
 
   const processOutputBBReferenceMock =

packages/worker/src/api/routes/global/tests/auditLogs.spec.ts

@@ -1,5 +1,5 @@
 import { mocks, structures } from "@budibase/backend-core/tests"
-import { context, events, setEnv as setCoreEnv } from "@budibase/backend-core"
+import { context, events, features } from "@budibase/backend-core"
 import { Event, IdentityType } from "@budibase/types"
 import { TestConfiguration } from "../../../../tests"
 
@@ -17,11 +17,9 @@ describe.each(["lucene", "sql"])("/api/global/auditlogs (%s)", method => {
   let envCleanup: (() => void) | undefined
 
   beforeAll(async () => {
-    if (method === "lucene") {
+    envCleanup = features.testutils.setFeatureFlags("*", {
-      envCleanup = setCoreEnv({ TENANT_FEATURE_FLAGS: "*:!SQS" })
+      SQS: method === "sql",
-    } else if (method === "sql") {
+    })
-      envCleanup = setCoreEnv({ TENANT_FEATURE_FLAGS: "*:SQS" })
-    }
     await config.beforeAll()
   })