diff --git a/front/scripts/helpers.ts b/front/scripts/helpers.ts index 2112541befd0..018655fe55c3 100644 --- a/front/scripts/helpers.ts +++ b/front/scripts/helpers.ts @@ -4,7 +4,7 @@ import { hideBin } from "yargs/helpers"; // Define a type for the argument specification object. export type ArgumentSpecs = { - [key: string]: Options & { type?: "string" | "boolean" | "number" }; + [key: string]: Options & { type?: "array" | "string" | "boolean" | "number" }; }; // Define a type for the worker function. @@ -18,6 +18,8 @@ type InferArgs = { ? boolean : T[P] extends { type: "string" } ? string + : T[P] extends { type: "array" } + ? string[] : never; } & { execute?: boolean }; diff --git a/front/scripts/toggle_feature_flags.ts b/front/scripts/toggle_feature_flags.ts new file mode 100644 index 000000000000..e04bd4d7993c --- /dev/null +++ b/front/scripts/toggle_feature_flags.ts @@ -0,0 +1,112 @@ +import type { WhitelistableFeature } from "@dust-tt/types"; +import { WHITELISTABLE_FEATURES } from "@dust-tt/types"; + +import { FeatureFlag, Workspace } from "@app/lib/models"; +import { makeScript } from "@app/scripts/helpers"; + +async function enableFeatureFlag( + workspace: Workspace, + featureFlag: WhitelistableFeature, + execute: boolean +) { + const { id: workspaceId, name } = workspace; + + if (execute) { + try { + await FeatureFlag.create({ + workspaceId, + name: featureFlag as WhitelistableFeature, + }); + } catch (err) { + console.log( + `Workspace ${name}(${workspaceId}) already has ${featureFlag} enabled -- Skipping.` + ); + + return; + } + } + + console.log( + `${ + execute ? "" : "[DRYRUN]:" + } Feature flag ${featureFlag} enabled for workspace: ${name}(${workspaceId}).` + ); +} + +async function disableFeatureFlag( + workspace: Workspace, + featureFlag: WhitelistableFeature, + execute: boolean +) { + const { id: workspaceId, name } = workspace; + + if (execute) { + const existingFlag = await FeatureFlag.findOne({ + where: { + workspaceId, + name: featureFlag, + }, + }); + if (!existingFlag) { + console.log( + `Workspace ${name}(${workspaceId}) does not have ${featureFlag} enabled -- Skipping.` + ); + return; + } + + await existingFlag.destroy(); + } + + console.log( + `${ + execute ? "" : "[DRYRUN]:" + } Feature flag ${featureFlag} disabled for workspace: ${name}(${workspaceId}).` + ); +} + +makeScript( + { + enable: { + type: "boolean", + default: false, + }, + featureFlag: { + type: "string", + choices: WHITELISTABLE_FEATURES, + demandOption: true, + }, + workspaceIds: { + type: "array", + demandOption: true, + description: + "List of workspace identifiers, separated by a space, for which the feature flag should be toggled.", + }, + }, + async ({ enable, featureFlag, workspaceIds, execute }) => { + for (const wId of workspaceIds) { + const workspace = await Workspace.findOne({ + where: { + sId: wId, + }, + }); + if (!workspace) { + console.log(`Workspace ${wId} not found -- Skipping.`); + continue; + } + + if (enable) { + await enableFeatureFlag( + workspace, + featureFlag as WhitelistableFeature, + execute + ); + } else { + await disableFeatureFlag( + workspace, + featureFlag as WhitelistableFeature, + execute + ); + } + } + } +);