Improve delete-group feature to be more robust

apps/web/app/(ee)/api/cron/groups/remap-default-links/route.ts

@@ -240,6 +240,7 @@ export async function POST(req: Request) {
         programId,
         partnerIds,
         groupId,
+        isGroupDeleted,
       },
     });
 

apps/web/app/(ee)/api/cron/groups/remap-discount-codes/route.ts

@@ -13,11 +13,12 @@ const inputSchema = z.object({
   programId: z.string(),
   groupId: z.string(),
   partnerIds: z.array(z.string()),
+  isGroupDeleted: z.boolean().optional(),
 });
 
 // POST /api/cron/groups/remap-discount-codes
 export const POST = withCron(async ({ rawBody }) => {
-  const { programId, partnerIds, groupId } = inputSchema.parse(
+  const { programId, partnerIds, groupId, isGroupDeleted } = inputSchema.parse(
     JSON.parse(rawBody),
   );
 
@@ -41,6 +42,8 @@ export const POST = withCron(async ({ rawBody }) => {
     },
   });
 
+  const oldDiscount = programEnrollments[0]?.discountCodes[0]?.discount;
+
   if (programEnrollments.length === 0) {
     return logAndRespond("No program enrollments found.");
   }
@@ -160,5 +163,24 @@ export const POST = withCron(async ({ rawBody }) => {
     }
   }
 
+  // if the group is deleted, need to check if there are any remaining discount codes, if not, delete the discount
+  if (isGroupDeleted && oldDiscount) {
+    const remainingDiscountCodes = await prisma.discountCode.count({
+      where: {
+        discountId: oldDiscount.id,
+      },
+    });
+    if (remainingDiscountCodes === 0) {
+      await prisma.discount.deleteMany({
+        where: {
+          id: oldDiscount.id,
+        },
+      });
+      console.log(
+        `Deleted discount ${oldDiscount.id} because it has no remaining discount codes.`,
+      );
+    }
+  }
+
   return logAndRespond("Finished remapping discount codes for the group.");
 });

apps/web/app/(ee)/api/groups/[groupIdOrSlug]/route.ts

@@ -1,25 +1,20 @@
 import { recordAuditLog } from "@/lib/api/audit-logs/record-audit-log";
-import { deleteDiscountCodes } from "@/lib/api/discounts/delete-discount-code";
-import { isDiscountEquivalent } from "@/lib/api/discounts/is-discount-equivalent";
 import { DubApiError } from "@/lib/api/errors";
 import { getGroupOrThrow } from "@/lib/api/groups/get-group-or-throw";
+import { movePartnersToGroup } from "@/lib/api/groups/move-partners-to-group";
 import { upsertGroupMoveRules } from "@/lib/api/groups/upsert-group-move-rules";
-import { includeProgramEnrollment } from "@/lib/api/links/include-program-enrollment";
-import { includeTags } from "@/lib/api/links/include-tags";
 import { getDefaultProgramIdOrThrow } from "@/lib/api/programs/get-default-program-id-or-throw";
 import { parseRequestBody } from "@/lib/api/utils";
 import { extractUtmParams } from "@/lib/api/utm/extract-utm-params";
 import { withWorkspace } from "@/lib/auth";
 import { qstash } from "@/lib/cron";
-import { recordLink } from "@/lib/tinybird";
 import { GroupWithProgramSchema } from "@/lib/zod/schemas/group-with-program";
 import {
   DEFAULT_PARTNER_GROUP,
   GroupSchema,
   updateGroupSchema,
 } from "@/lib/zod/schemas/groups";
 import { prisma } from "@dub/prisma";
-import { DiscountCode } from "@dub/prisma/client";
 import { APP_DOMAIN_WITH_NGROK, constructURLFromUTMParams } from "@dub/utils";
 import { waitUntil } from "@vercel/functions";
 import { NextResponse } from "next/server";
@@ -292,10 +287,6 @@ export const DELETE = withWorkspace(
                 },
               }),
         },
-        include: {
-          partners: true,
-          discount: true,
-        },
       }),
 
       prisma.partnerGroup.findUniqueOrThrow({
@@ -305,9 +296,6 @@ export const DELETE = withWorkspace(
             slug: DEFAULT_PARTNER_GROUP.slug,
           },
         },
-        include: {
-          discount: true,
-        },
       }),
     ]);
 
@@ -318,37 +306,29 @@ export const DELETE = withWorkspace(
       });
     }
 
-    const keepDiscountCodes = isDiscountEquivalent(
-      group.discount,
-      defaultGroup.discount,
-    );
-
-    // Cache discount codes to delete them later
-    let discountCodesToDelete: DiscountCode[] = [];
-    if (group.discountId && !keepDiscountCodes) {
-      discountCodesToDelete = await prisma.discountCode.findMany({
+    while (true) {
+      const programEnrollments = await prisma.programEnrollment.findMany({
         where: {
-          discountId: group.discountId,
+          groupId: group.id,
         },
+        take: 100,
+      });
+      if (programEnrollments.length === 0) {
+        break;
+      }
+      const count = await movePartnersToGroup({
+        workspaceId: workspace.id,
+        programId,
+        partnerIds: programEnrollments.map(({ partnerId }) => partnerId),
+        userId: session.user.id,
+        group: defaultGroup,
+        isGroupDeleted: true,
       });
+      console.log(`Moved ${count} partners to the default group`);
     }
 
     const deletedGroup = await prisma.$transaction(async (tx) => {
-      // 1. Update all partners in the group to the default group
-      await tx.programEnrollment.updateMany({
-        where: {
-          groupId: group.id,
-        },
-        data: {
-          groupId: defaultGroup.id,
-          clickRewardId: defaultGroup.clickRewardId,
-          leadRewardId: defaultGroup.leadRewardId,
-          saleRewardId: defaultGroup.saleRewardId,
-          discountId: defaultGroup.discountId,
-        },
-      });
-
-      // 2. Delete the group's rewards
+      // 1. Delete the group's rewards
       if (group.clickRewardId || group.leadRewardId || group.saleRewardId) {
         await tx.reward.deleteMany({
           where: {
@@ -363,26 +343,11 @@ export const DELETE = withWorkspace(
         });
       }
 
-      if (group.discountId) {
-        // 3. Update the discount codes
-        await tx.discountCode.updateMany({
-          where: {
-            discountId: group.discountId,
-          },
-          data: {
-            discountId: keepDiscountCodes ? defaultGroup.discountId : null,
-          },
-        });
-
-        // 4. Delete the group's discount
-        await tx.discount.delete({
-          where: {
-            id: group.discountId,
-          },
-        });
-      }
+      // Note: we can't delete this group's discount yet because it is needed
+      // for `remap-discount-codes` that runs in movePartnersToGroup
+      // but we will delete the Discount in `remap-discount-codes` once there are no remaining discount codes.
 
-      // 5. Delete the group move workflow
+      // 2. Delete the group move workflow
       if (group.workflowId) {
         await tx.workflow.delete({
           where: {
@@ -391,7 +356,7 @@ export const DELETE = withWorkspace(
         });
       }
 
-      // 6. Delete the group
+      // 3. Delete the group
       await tx.partnerGroup.delete({
         where: {
           id: group.id,
@@ -403,60 +368,20 @@ export const DELETE = withWorkspace(
 
     if (deletedGroup) {
       waitUntil(
-        (async () => {
-          const partnerIds = group.partners.map(({ partnerId }) => partnerId);
-
-          // TODO:
-          // This won't work for larger groups
-          // We should split this into multiple batches
-          const partnerLinks = await prisma.link.findMany({
-            where: {
-              programId,
-              partnerId: {
-                in: partnerIds,
-              },
+        recordAuditLog({
+          workspaceId: workspace.id,
+          programId,
+          action: "group.deleted",
+          description: `Group ${group.name} (${group.id}) deleted`,
+          actor: session.user,
+          targets: [
+            {
+              type: "group",
+              id: group.id,
+              metadata: group,
             },
-            include: {
-              ...includeTags,
-              ...includeProgramEnrollment,
-            },
-          });
-
-          await Promise.allSettled([
-            partnerIds.length > 0 &&
-              qstash.publishJSON({
-                url: `${APP_DOMAIN_WITH_NGROK}/api/cron/groups/remap-default-links`,
-                body: {
-                  programId,
-                  groupId: defaultGroup.id,
-                  partnerIds,
-                  userId: session.user.id,
-                  isGroupDeleted: true,
-                },
-              }),
-
-            ...discountCodesToDelete.map((discountCode) =>
-              deleteDiscountCodes(discountCode),
-            ),
-
-            recordAuditLog({
-              workspaceId: workspace.id,
-              programId,
-              action: "group.deleted",
-              description: `Group ${group.name} (${group.id}) deleted`,
-              actor: session.user,
-              targets: [
-                {
-                  type: "group",
-                  id: group.id,
-                  metadata: group,
-                },
-              ],
-            }),
-
-            recordLink(partnerLinks),
-          ]);
-        })(),
+          ],
+        }),
       );
     }
 

apps/web/lib/api/groups/move-partners-to-group.ts

@@ -29,6 +29,7 @@ interface MovePartnersToGroupParams {
     | "saleRewardId"
     | "discountId"
   >;
+  isGroupDeleted?: boolean;
 }
 
 export async function movePartnersToGroup({
@@ -37,6 +38,7 @@ export async function movePartnersToGroup({
   partnerIds,
   userId,
   group,
+  isGroupDeleted = false,
 }: MovePartnersToGroupParams): Promise<number> {
   if (partnerIds.length === 0) {
     return 0;
@@ -167,6 +169,7 @@ export async function movePartnersToGroup({
             groupId: group.id,
             partnerIds,
             userId: workspaceUserId,
+            isGroupDeleted,
           },
         }),