Add script to toggle feature flags in batch.

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<T> = {
     ? boolean
     : T[P] extends { type: "string" }
     ? string
+    : T[P] extends { type: "array" }
+    ? string[]
     : never;
 } & { execute?: boolean };
 

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
+        );
+      }
+    }
+  }
+);