Fix race conditions in partner create-and-enroll upsert

apps/web/lib/api/partners/create-and-enroll-partner.ts

@@ -8,12 +8,13 @@ import { CreatePartnerProps, ProgramProps, WorkspaceProps } from "@/lib/types";
 import { sendWorkspaceWebhook } from "@/lib/webhook/publish";
 import { EnrolledPartnerSchema } from "@/lib/zod/schemas/partners";
 import { nanoid } from "@dub/utils";
-import { Prisma, ProgramEnrollmentStatus } from "@prisma/client";
+import { ProgramEnrollmentStatus } from "@prisma/client";
 import { waitUntil } from "@vercel/functions";
 import { DubApiError } from "../errors";
 import { getGroupOrThrow } from "../groups/get-group-or-throw";
+import { createOrGetProgramEnrollment } from "./create-or-get-program-enrollment";
 import { createPartnerDefaultLinks } from "./create-partner-default-links";
-import { generatePartnerUsername } from "./generate-partner-username";
+import { getOrCreatePartner } from "./get-or-create-partner";
 import { throwIfExistingTenantEnrollmentExists } from "./throw-if-existing-tenant-id-exists";
 
 interface CreateAndEnrollPartnerInput {
@@ -136,60 +137,50 @@ export const createAndEnrollPartner = async ({
     });
   }
 
-  const group = await getGroupOrThrow({
-    programId: program.id,
-    groupId: finalGroupId,
-    includeExpandedFields: true,
-  });
-
-  const payload: Pick<Prisma.PartnerUpdateInput, "programs"> = {
-    programs: {
-      create: {
-        id: createId({ prefix: "pge_" }),
-        programId: program.id,
-        tenantId: partner.tenantId,
-        status,
-        groupId: group.id,
-        clickRewardId: group.clickRewardId,
-        leadRewardId: group.leadRewardId,
-        saleRewardId: group.saleRewardId,
-        referralRewardId: group.referralRewardId,
-        discountId: group.discountId,
-        ...(enrolledAt && {
-          createdAt: enrolledAt,
-        }),
-      },
-    },
-  };
+  const [group, { partner: existingOrNewPartner }] = await Promise.all([
+    getGroupOrThrow({
+      programId: program.id,
+      groupId: finalGroupId,
+      includeExpandedFields: true,
+    }),
 
-  const upsertedPartner = await prisma.partner.upsert({
-    where: {
+    getOrCreatePartner({
       email: partner.email,
-    },
-    update: payload,
-    create: {
-      ...payload,
-      id: createId({ prefix: "pn_" }),
-      name: partner.name || partner.email,
-      email: partner.email,
-      username: await generatePartnerUsername({
+      create: {
+        id: createId({ prefix: "pn_" }),
+        name: partner.name || partner.email,
         email: partner.email,
-        name: partner.name,
-      }),
-      image: partner.image && !isStored(partner.image) ? null : partner.image,
-      country: partner.country,
-      description: partner.description,
-    },
-    include: {
-      platforms: true,
-      programs: {
-        where: {
-          programId: program.id,
-        },
+        image: partner.image && !isStored(partner.image) ? null : partner.image,
+        country: partner.country,
+        description: partner.description,
       },
-    },
+    }),
+  ]);
+
+  const { programEnrollment, created } = await createOrGetProgramEnrollment({
+    partnerId: existingOrNewPartner.id,
+    programId: program.id,
+    tenantId: partner.tenantId ?? null,
+    status,
+    groupId: group.id,
+    clickRewardId: group.clickRewardId,
+    leadRewardId: group.leadRewardId,
+    saleRewardId: group.saleRewardId,
+    referralRewardId: group.referralRewardId,
+    discountId: group.discountId,
+    enrolledAt,
   });
 
+  if (!created) {
+    return EnrolledPartnerSchema.parse({
+      ...programEnrollment.partner,
+      ...programEnrollment,
+      id: programEnrollment.partner.id,
+      links: programEnrollment.links,
+      ...polyfillSocialMediaFields(programEnrollment.partner.platforms),
+    });
+  }
+
   // Create the partner links based on group defaults
   const links = await createPartnerDefaultLinks({
     workspace: {
@@ -201,7 +192,7 @@ export const createAndEnrollPartner = async ({
       defaultFolderId: program.defaultFolderId,
     },
     partner: {
-      id: upsertedPartner.id,
+      id: existingOrNewPartner.id,
       name: partner.name,
       email: partner.email,
       username: partner.username,
@@ -216,11 +207,11 @@ export const createAndEnrollPartner = async ({
   });
 
   const enrolledPartner = EnrolledPartnerSchema.parse({
-    ...upsertedPartner,
-    ...upsertedPartner.programs[0],
-    id: upsertedPartner.id,
+    ...programEnrollment.partner,
+    ...programEnrollment,
+    id: programEnrollment.partner.id,
     links,
-    ...polyfillSocialMediaFields(upsertedPartner.platforms),
+    ...polyfillSocialMediaFields(programEnrollment.partner.platforms),
   });
 
   waitUntil(
@@ -242,13 +233,13 @@ export const createAndEnrollPartner = async ({
         !isStored(partner.image) &&
         storage
           .upload({
-            key: `partners/${upsertedPartner.id}/image_${nanoid(7)}`,
+            key: `partners/${existingOrNewPartner.id}/image_${nanoid(7)}`,
             body: partner.image,
           })
           .then(async ({ url }) => {
             await prisma.partner.update({
               where: {
-                id: upsertedPartner.id,
+                id: existingOrNewPartner.id,
               },
               data: {
                 image: url,

apps/web/lib/api/partners/create-or-get-program-enrollment.ts

@@ -0,0 +1,112 @@
+import { createId } from "@/lib/api/create-id";
+import { prisma } from "@/lib/prisma";
+import { Prisma, ProgramEnrollment } from "@prisma/client";
+import { throwIfExistingTenantEnrollmentExists } from "./throw-if-existing-tenant-id-exists";
+
+// Attempts to create a program enrollment.
+// If another request creates it concurrently, returns the existing enrollment.
+export async function createOrGetProgramEnrollment({
+  programId,
+  partnerId,
+  tenantId,
+  status,
+  groupId,
+  clickRewardId,
+  leadRewardId,
+  saleRewardId,
+  referralRewardId,
+  discountId,
+  enrolledAt,
+}: Pick<
+  ProgramEnrollment,
+  | "partnerId"
+  | "programId"
+  | "tenantId"
+  | "status"
+  | "groupId"
+  | "clickRewardId"
+  | "leadRewardId"
+  | "saleRewardId"
+  | "referralRewardId"
+  | "discountId"
+> & {
+  enrolledAt?: Date;
+}) {
+  try {
+    const programEnrollment = await prisma.programEnrollment.create({
+      data: {
+        id: createId({ prefix: "pge_" }),
+        partnerId,
+        programId,
+        tenantId,
+        status,
+        groupId,
+        clickRewardId,
+        leadRewardId,
+        saleRewardId,
+        referralRewardId,
+        discountId,
+        ...(enrolledAt && {
+          createdAt: enrolledAt,
+        }),
+      },
+      include: {
+        partner: {
+          include: {
+            platforms: true,
+          },
+        },
+      },
+    });
+
+    return {
+      programEnrollment,
+      created: true as const,
+    };
+  } catch (error) {
+    if (
+      error instanceof Prisma.PrismaClientKnownRequestError &&
+      error.code === "P2002"
+    ) {
+      console.info(
+        "[createOrGetProgramEnrollment] Unique constraint conflict (P2002), falling back to find",
+        { target: error.meta?.target },
+      );
+
+      // Same partner already enrolled (concurrent create on partnerId_programId)
+      const programEnrollment = await prisma.programEnrollment.findUnique({
+        where: {
+          partnerId_programId: {
+            partnerId,
+            programId,
+          },
+        },
+        include: {
+          links: true,
+          partner: {
+            include: {
+              platforms: true,
+            },
+          },
+        },
+      });
+
+      if (programEnrollment) {
+        return {
+          programEnrollment,
+          created: false as const,
+        };
+      }
+
+      // Conflict was on tenantId_programId — another partner owns this tenant
+      if (tenantId) {
+        await throwIfExistingTenantEnrollmentExists({
+          tenantId,
+          programId,
+        });
+      }
+    }
+
+    throw error;
+  }
+}

apps/web/lib/api/partners/get-or-create-partner.ts

@@ -0,0 +1,68 @@
+import { prisma } from "@/lib/prisma";
+import { Prisma } from "@prisma/client";
+import { generatePartnerUsername } from "./generate-partner-username";
+
+// This helper finds first, creates if missing,
+// and on unique-constraint conflict falls back to find by email.
+export async function getOrCreatePartner({
+  email,
+  create,
+}: {
+  email: string;
+  create: Omit<Prisma.PartnerUncheckedCreateInput, "username">;
+}) {
+  const partner = await prisma.partner.findUnique({
+    where: {
+      email,
+    },
+  });
+
+  if (partner) {
+    return {
+      partner,
+      created: false,
+    };
+  }
+
+  try {
+    const username = await generatePartnerUsername({
+      email,
+      name: typeof create.name === "string" ? create.name : null,
+    });
+
+    const partner = await prisma.partner.create({
+      data: {
+        ...create,
+        username,
+      },
+    });
+
+    return {
+      partner,
+      created: true,
+    };
+  } catch (error) {
+    if (
+      error instanceof Prisma.PrismaClientKnownRequestError &&
+      error.code === "P2002"
+    ) {
+      console.info(
+        "[getOrCreatePartner] Unique constraint conflict (P2002), falling back to find",
+        { target: error.meta?.target },
+      );
+
+      const partner = await prisma.partner.findUniqueOrThrow({
+        where: {
+          email,
+        },
+      });
+
+      return {
+        partner,
+        created: false,
+      };
+    }
+
+    throw error;
+  }
+}