From 6f4319a2ec1e35b6d65477c1ba537b8d8276d1f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <PopDaph@users.noreply.github.com>
Date: Thu, 25 Jul 2024 13:33:08 +0200
Subject: [PATCH] Enforce one workspace group per workspace (#6505)

---
 front/lib/resources/storage/models/groups.ts  | 19 +++++++--------
 .../20240724_workspaces_groups_backfill.ts    | 23 ++++++++++++-------
 types/src/front/groups.ts                     |  3 +++
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/front/lib/resources/storage/models/groups.ts b/front/lib/resources/storage/models/groups.ts
index 585e554fbb64..0d6ed021587b 100644
--- a/front/lib/resources/storage/models/groups.ts
+++ b/front/lib/resources/storage/models/groups.ts
@@ -1,5 +1,5 @@
 import type { GroupType } from "@dust-tt/types";
-import { isSystemGroupType } from "@dust-tt/types";
+import { isSystemGroupType, isWorkspaceGroupType } from "@dust-tt/types";
 import type {
   CreationOptional,
   ForeignKey,
@@ -59,21 +59,22 @@ GroupModel.init(
 );
 
 GroupModel.addHook(
-  "beforeSave",
-  "enforce_one_system_group_per_workspace",
+  "beforeCreate",
+  "enforce_one_system_and_workspace_group_per_workspace",
   async (group: GroupModel, options: { transaction: Transaction }) => {
-    if (isSystemGroupType(group.type)) {
-      const existingSystemGroupType = await GroupModel.findOne({
+    const groupType = group.type;
+    if (isSystemGroupType(groupType) || isWorkspaceGroupType(groupType)) {
+      const existingSystemOrWorkspaceGroupType = await GroupModel.findOne({
         where: {
           workspaceId: group.workspaceId,
-          type: group.type,
+          type: groupType,
         },
         transaction: options.transaction,
       });
 
-      if (existingSystemGroupType) {
-        throw new Error("A system group exists for this workspace.", {
-          cause: "enforce_one_system_group_per_workspace",
+      if (existingSystemOrWorkspaceGroupType) {
+        throw new Error(`A ${groupType} group exists for this workspace.`, {
+          cause: `enforce_one_${groupType}_group_per_workspace`,
         });
       }
     }
diff --git a/front/migrations/20240724_workspaces_groups_backfill.ts b/front/migrations/20240724_workspaces_groups_backfill.ts
index d2d7d16b7ece..2c845ab7c4f2 100644
--- a/front/migrations/20240724_workspaces_groups_backfill.ts
+++ b/front/migrations/20240724_workspaces_groups_backfill.ts
@@ -31,14 +31,21 @@ async function backfillWorkspacesGroup(execute: boolean) {
               });
               console.log(`System group created for workspace ${w.id}`);
             } catch (error) {
-              if (
-                error instanceof Error &&
-                error.cause &&
-                error.cause === "enforce_one_system_group_per_workspace"
-              ) {
-                console.log(
-                  `System group already exists for workspace ${w.id}`
-                );
+              if (error instanceof Error && error.cause) {
+                switch (error.cause) {
+                  case "enforce_one_system_group_per_workspace":
+                    console.log(
+                      `System group already exists for workspace ${w.id}`
+                    );
+                    break;
+                  case "enforce_one_workspace_group_per_workspace":
+                    console.log(
+                      `Workspace group already exists for workspace ${w.id}`
+                    );
+                    break;
+                  default:
+                    console.error(error);
+                }
               } else {
                 console.error(error);
               }
diff --git a/types/src/front/groups.ts b/types/src/front/groups.ts
index 0c1828428f51..cc4f824b3109 100644
--- a/types/src/front/groups.ts
+++ b/types/src/front/groups.ts
@@ -7,3 +7,6 @@ export function isValidGroupType(value: unknown): value is GroupType {
 export function isSystemGroupType(value: GroupType): boolean {
   return value === "system";
 }
+export function isWorkspaceGroupType(value: GroupType): boolean {
+  return value === "workspace";
+}