refactor and fix bulk attribute handling (#18231)

packages/prisma/index.ts

@@ -81,6 +81,15 @@ if (process.env.NODE_ENV !== "production") {
 
 type PrismaClientWithExtensions = typeof prismaWithClientExtensions;
 export type PrismaClient = PrismaClientWithExtensions;
+
+type OmitPrismaClient = Omit<
+  PrismaClient,
+  "$connect" | "$disconnect" | "$on" | "$transaction" | "$use" | "$extends"
+>;
+
+// we cant pass tx to functions as types miss match since we have a custom prisma client https://github.com/prisma/prisma/discussions/20924#discussioncomment-10077649
+export type PrismaTransaction = OmitPrismaClient;
+
 export default prisma;
 
 export * from "./selects";

packages/trpc/server/routers/viewer/attributes/attributeUtils.test.ts

@@ -0,0 +1,259 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+
+import {
+  handleSimpleAttribute,
+  handleSelectAttribute,
+  processUserAttributes,
+  removeAttribute,
+} from "./attributeUtils";
+
+describe("Attribute Utils", () => {
+  const mockTx = {
+    membership: {
+      findFirst: vi.fn(),
+    },
+    attributeToUser: {
+      findFirst: vi.fn(),
+      upsert: vi.fn(),
+      deleteMany: vi.fn(),
+    },
+    attributeOption: {
+      update: vi.fn(),
+      create: vi.fn(),
+    },
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe("processUserAttributes", () => {
+    it("should return error if user is not a member", async () => {
+      mockTx.membership.findFirst.mockResolvedValue(null);
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, []);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: false,
+        message: "User is not part of your organization",
+      });
+    });
+
+    it("should skip attributes without type", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [{ id: "attr-1", value: "test" }]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      expect(mockTx.attributeToUser.findFirst).not.toHaveBeenCalled();
+    });
+
+    it("should process TEXT attribute", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [
+        { id: "attr-1", value: "test", type: "TEXT" },
+      ]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      expect(mockTx.attributeToUser.findFirst).toHaveBeenCalled();
+    });
+
+    it("should process NUMBER attribute", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [
+        { id: "attr-1", value: "123", type: "NUMBER" },
+      ]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      expect(mockTx.attributeToUser.findFirst).toHaveBeenCalled();
+    });
+
+    it("should process SINGLE_SELECT attribute", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [
+        { id: "attr-1", options: [{ value: "opt-1" }], type: "SINGLE_SELECT" },
+      ]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      // Should first delete existing options
+      expect(mockTx.attributeToUser.deleteMany).toHaveBeenCalledWith({
+        where: {
+          memberId: 1,
+          attributeOption: {
+            attribute: {
+              id: "attr-1",
+            },
+          },
+        },
+      });
+      // Then add the new option
+      expect(mockTx.attributeToUser.upsert).toHaveBeenCalledWith({
+        where: {
+          memberId_attributeOptionId: {
+            memberId: 1,
+            attributeOptionId: "opt-1",
+          },
+        },
+        create: {
+          memberId: 1,
+          attributeOptionId: "opt-1",
+        },
+        update: {},
+      });
+    });
+
+    it("should process MULTI_SELECT attribute", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [
+        {
+          id: "attr-1",
+          options: [{ value: "opt-1" }, { value: "opt-2" }],
+          type: "MULTI_SELECT",
+        },
+      ]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      // Should not delete existing options for multi-select
+      expect(mockTx.attributeToUser.deleteMany).not.toHaveBeenCalled();
+      expect(mockTx.attributeToUser.upsert).toHaveBeenCalledTimes(2);
+    });
+
+    it("should handle attribute removal when no value or options provided", async () => {
+      mockTx.membership.findFirst.mockResolvedValue({ id: 1 });
+
+      const result = await processUserAttributes(mockTx as any, 1, 1, [{ id: "attr-1", type: "TEXT" }]);
+
+      expect(result).toEqual({
+        userId: 1,
+        success: true,
+      });
+      expect(mockTx.attributeToUser.deleteMany).toHaveBeenCalled();
+    });
+  });
+
+  describe("handleSimpleAttribute", () => {
+    it("should update existing attribute option", async () => {
+      mockTx.attributeToUser.findFirst.mockResolvedValue({
+        id: 1,
+        attributeOption: { id: 2 },
+      });
+
+      await handleSimpleAttribute(mockTx as any, 1, {
+        id: "attr-1",
+        value: "test value",
+      });
+
+      expect(mockTx.attributeOption.update).toHaveBeenCalledWith({
+        where: { id: 2 },
+        data: {
+          value: "test value",
+          slug: "test-value",
+        },
+      });
+    });
+
+    it("should create new attribute option if none exists", async () => {
+      mockTx.attributeToUser.findFirst.mockResolvedValue(null);
+
+      await handleSimpleAttribute(mockTx as any, 1, {
+        id: "attr-1",
+        value: "new value",
+      });
+
+      expect(mockTx.attributeOption.create).toHaveBeenCalledWith({
+        data: {
+          value: "new value",
+          slug: "new-value",
+          attribute: {
+            connect: { id: "attr-1" },
+          },
+          assignedUsers: {
+            create: { memberId: 1 },
+          },
+        },
+      });
+    });
+  });
+
+  describe("handleSelectAttribute", () => {
+    it("should remove existing options for SINGLE_SELECT before adding new one", async () => {
+      await handleSelectAttribute(mockTx as any, 1, {
+        id: "attr-1",
+        options: [{ value: "opt-1" }],
+        type: "SINGLE_SELECT",
+      });
+
+      expect(mockTx.attributeToUser.deleteMany).toHaveBeenCalledWith({
+        where: {
+          memberId: 1,
+          attributeOption: {
+            attribute: {
+              id: "attr-1",
+            },
+          },
+        },
+      });
+      expect(mockTx.attributeToUser.upsert).toHaveBeenCalledWith({
+        where: {
+          memberId_attributeOptionId: {
+            memberId: 1,
+            attributeOptionId: "opt-1",
+          },
+        },
+        create: {
+          memberId: 1,
+          attributeOptionId: "opt-1",
+        },
+        update: {},
+      });
+    });
+
+    it("should not remove existing options for MULTI_SELECT", async () => {
+      await handleSelectAttribute(mockTx as any, 1, {
+        id: "attr-1",
+        options: [{ value: "opt-1" }, { value: "opt-2" }],
+        type: "MULTI_SELECT",
+      });
+
+      expect(mockTx.attributeToUser.deleteMany).not.toHaveBeenCalled();
+      expect(mockTx.attributeToUser.upsert).toHaveBeenCalledTimes(2);
+    });
+  });
+
+  describe("removeAttribute", () => {
+    it("should delete all attribute options for user", async () => {
+      await removeAttribute(mockTx as any, 1, "attr-1");
+
+      expect(mockTx.attributeToUser.deleteMany).toHaveBeenCalledWith({
+        where: {
+          memberId: 1,
+          attributeOption: {
+            attribute: {
+              id: "attr-1",
+            },
+          },
+        },
+      });
+    });
+  });
+});

packages/trpc/server/routers/viewer/attributes/attributeUtils.ts

@@ -0,0 +1,183 @@
+import { safeStringify } from "@calcom/lib/safeStringify";
+import slugify from "@calcom/lib/slugify";
+import type { PrismaTransaction } from "@calcom/prisma";
+import type { AttributeType } from "@calcom/prisma/enums";
+
+type SimpleAttributeInput = {
+  id: string;
+  value: string;
+};
+
+type SelectAttributeInput = {
+  id: string;
+  options: { value: string }[];
+  type: AttributeType;
+};
+
+type AttributeInput = {
+  id: string;
+  value?: string;
+  options?: { value: string }[];
+  type?: AttributeType;
+};
+
+type ProcessAttributesResult = {
+  userId: number;
+  success: boolean;
+  message?: string;
+};
+
+const isSimpleAttribute = (type: AttributeType) => {
+  return type === "TEXT" || type === "NUMBER";
+};
+
+const isSelectAttribute = (type: AttributeType) => {
+  return type === "SINGLE_SELECT" || type === "MULTI_SELECT";
+};
+
+export const processUserAttributes = async (
+  tx: PrismaTransaction,
+  userId: number,
+  teamId: number,
+  attributes: AttributeInput[]
+): Promise<ProcessAttributesResult> => {
+  const membership = await tx.membership.findFirst({
+    where: {
+      userId,
+      teamId,
+    },
+  });
+
+  if (!membership) {
+    return {
+      userId,
+      success: false,
+      message: "User is not part of your organization",
+    };
+  }
+
+  for (const attribute of attributes) {
+    if (!attribute.type) {
+      console.log("Skipping attribute without type", safeStringify(attribute));
+      continue;
+    }
+
+    if (isSimpleAttribute(attribute.type) && attribute.value) {
+      await handleSimpleAttribute(tx, membership.id, {
+        id: attribute.id,
+        value: attribute.value,
+      });
+    } else if (isSelectAttribute(attribute.type) && attribute.options?.length) {
+      await handleSelectAttribute(tx, membership.id, {
+        id: attribute.id,
+        options: attribute.options,
+        type: attribute.type,
+      });
+    } else if (attribute.type && !attribute.value && !attribute.options?.length) {
+      // Handle attribute removal
+      await removeAttribute(tx, membership.id, attribute.id);
+    }
+  }
+
+  return {
+    userId,
+    success: true,
+  };
+};
+
+export const handleSimpleAttribute = async (
+  tx: PrismaTransaction,
+  memberId: number,
+  attribute: SimpleAttributeInput
+) => {
+  const valueAsString = String(attribute.value);
+
+  const existingAttributeOption = await tx.attributeToUser.findFirst({
+    where: {
+      memberId,
+      attributeOption: {
+        attribute: {
+          id: attribute.id,
+        },
+      },
+    },
+    select: {
+      id: true,
+      attributeOption: {
+        select: {
+          id: true,
+        },
+      },
+    },
+  });
+
+  if (existingAttributeOption) {
+    // Update the value if it already exists
+    await tx.attributeOption.update({
+      where: {
+        id: existingAttributeOption.attributeOption.id,
+      },
+      data: {
+        value: valueAsString,
+        slug: slugify(valueAsString),
+      },
+    });
+  } else {
+    await tx.attributeOption.create({
+      data: {
+        value: valueAsString,
+        slug: slugify(valueAsString),
+        attribute: {
+          connect: {
+            id: attribute.id,
+          },
+        },
+        assignedUsers: {
+          create: {
+            memberId,
+          },
+        },
+      },
+    });
+  }
+};
+
+export const handleSelectAttribute = async (
+  tx: PrismaTransaction,
+  memberId: number,
+  attribute: SelectAttributeInput
+) => {
+  // For single select, first remove any existing options
+  if (attribute.type === "SINGLE_SELECT") {
+    await removeAttribute(tx, memberId, attribute.id);
+  }
+
+  for (const option of attribute.options) {
+    await tx.attributeToUser.upsert({
+      where: {
+        memberId_attributeOptionId: {
+          memberId,
+          attributeOptionId: option.value,
+        },
+      },
+      create: {
+        memberId,
+        attributeOptionId: option.value,
+      },
+      update: {}, // No update needed if it already exists
+    });
+  }
+};
+
+export const removeAttribute = async (tx: PrismaTransaction, memberId: number, attributeId: string) => {
+  await tx.attributeToUser.deleteMany({
+    where: {
+      memberId,
+      attributeOption: {
+        attribute: {
+          id: attributeId,
+        },
+      },
+    },
+  });
+};

packages/trpc/server/routers/viewer/attributes/bulkAssignAttributes.handler.ts

@@ -1,9 +1,9 @@
-import slugify from "@calcom/lib/slugify";
 import prisma from "@calcom/prisma";
 
 import { TRPCError } from "@trpc/server";
 
 import type { TrpcSessionUser } from "../../../trpc";
+import { processUserAttributes } from "./attributeUtils";
 import type { ZBulkAssignAttributes } from "./bulkAssignAttributes.schema";
 
 type GetOptions = {
@@ -23,6 +23,9 @@ const bulkAssignAttributesHandler = async ({ input, ctx }: GetOptions) => {
     });
   }
 
+  // Create a map of attribute types for quick lookup
+  const attributeTypes = new Map<string, string>();
+
   // Ensure this organization can access these attributes and attribute options
   const attributes = await prisma.attribute.findMany({
     where: {
@@ -45,149 +48,22 @@ const bulkAssignAttributesHandler = async ({ input, ctx }: GetOptions) => {
     });
   }
 
-  const arrayOfAttributeOptionIds = attributes.flatMap(
-    (attribute) => attribute.options?.map((option) => option.id) || []
-  );
-
-  const attributeOptionIds = Array.from(new Set(arrayOfAttributeOptionIds));
-
-  const attributeOptions = await prisma.attributeOption.findMany({
-    where: {
-      id: {
-        in: attributeOptionIds,
-      },
-      attribute: {
-        teamId: org.id,
-      },
-    },
-    select: {
-      id: true,
-      value: true,
-      slug: true,
-    },
+  // Store attribute types in the map
+  attributes.forEach((attr) => {
+    attributeTypes.set(attr.id, attr.type);
   });
 
-  if (attributeOptions.length !== attributeOptionIds.length) {
-    throw new TRPCError({
-      code: "UNAUTHORIZED",
-      message: "You do not have access to these attribute options",
-    });
-  }
-
   const results = await Promise.all(
     input.userIds.map(async (userId) => {
       return prisma.$transaction(async (tx) => {
-        const membership = await tx.membership.findFirst({
-          where: {
-            userId: userId,
-            // @ts-expect-error we check this higher in the logic
-            teamId: org?.id,
-          },
-        });
+        // Add type information to the input attributes
+        const attributesWithType = input.attributes.map((attr) => ({
+          ...attr,
+          type: attributeTypes.get(attr.id),
+        }));
 
-        if (!membership) {
-          return {
-            userId,
-            success: false,
-            message: "User is not part of your organization",
-          };
-        }
-
-        for (const attribute of input.attributes) {
-          // TEXT, NUMBER
-          if (attribute.value && !attribute.options) {
-            const valueAsString = String(attribute.value);
-
-            // Check if it is already the value
-            const existingAttributeOption = await tx.attributeToUser.findFirst({
-              where: {
-                memberId: membership.id,
-                attributeOption: {
-                  attribute: {
-                    id: attribute.id,
-                  },
-                },
-              },
-              select: {
-                id: true,
-                attributeOption: {
-                  select: {
-                    id: true,
-                  },
-                },
-              },
-            });
-
-            if (existingAttributeOption) {
-              // Update the value if it already exists
-              await tx.attributeOption.update({
-                where: {
-                  id: existingAttributeOption.attributeOption.id,
-                },
-                data: {
-                  value: valueAsString,
-                  slug: slugify(valueAsString),
-                },
-              });
-            } else {
-              await tx.attributeOption.create({
-                data: {
-                  value: valueAsString,
-                  slug: slugify(valueAsString),
-                  attribute: {
-                    connect: {
-                      id: attribute.id,
-                    },
-                  },
-                  assignedUsers: {
-                    create: {
-                      memberId: membership.id,
-                    },
-                  },
-                },
-              });
-            }
-          } else if (!attribute.value && attribute.options && attribute.options.length > 0) {
-            const options = attribute.options;
-
-            for (const option of options) {
-              // Assign the attribute option to the user
-              await tx.attributeToUser.upsert({
-                where: {
-                  memberId_attributeOptionId: {
-                    memberId: membership.id,
-                    attributeOptionId: option.value,
-                  },
-                },
-                create: {
-                  memberId: membership.id,
-                  attributeOptionId: option.value,
-                },
-                update: {}, // No update needed if it already exists
-              });
-            }
-          }
-
-          // Delete the attribute from the user
-          if (!attribute.value && !attribute.options) {
-            await tx.attributeToUser.deleteMany({
-              where: {
-                memberId: membership.id,
-                attributeOption: {
-                  attribute: {
-                    id: attribute.id,
-                  },
-                },
-              },
-            });
-          }
-        }
-
-        return {
-          userId,
-          success: true,
-          message: "Attributes assigned successfully",
-        };
+        // @ts-expect-error - org.id is being checked above for nullish
+        return processUserAttributes(tx, userId, org.id, attributesWithType);
       });
     })
   );