Merge pull request #2647 from trycompai/fix/onboarding-transaction-perf

tofikwest · web-flow · commit 10b51fee7d47 · 2026-04-23T11:41:14.000-04:00
perf(onboarding): replace sequential update loops with bulk SQL
diff --git a/apps/app/src/actions/organization/lib/initialize-organization.ts b/apps/app/src/actions/organization/lib/initialize-organization.ts
@@ -216,61 +216,61 @@ export const _upsertOrgFrameworkStructureCore = async ({
   );
 
   if (policyTemplatesForCreation.length > 0) {
+    // Pre-generate Policy and PolicyVersion IDs in a single round-trip so we can
+    // skip the post-insert findMany lookups and the per-policy update loop.
+    // Policy.currentVersionId -> PolicyVersion.id and PolicyVersion.policyId ->
+    // Policy.id form an FK cycle, so we insert Policy first (currentVersionId null),
+    // insert PolicyVersion, then set currentVersionId in one bulk UPDATE.
+    const idPairs = await tx.$queryRaw<
+      Array<{ policy_id: string; version_id: string }>
+    >`
+      SELECT
+        generate_prefixed_cuid('pol'::text) AS policy_id,
+        generate_prefixed_cuid('pv'::text) AS version_id
+      FROM generate_series(1, ${policyTemplatesForCreation.length}::int)
+    `;
+    const preparedPolicies = policyTemplatesForCreation.map((template, i) => ({
+      template,
+      policyId: idPairs[i].policy_id,
+      versionId: idPairs[i].version_id,
+      contentArray: extractTipTapContentArray(template.content),
+    }));
+
     await tx.policy.createMany({
-      data: policyTemplatesForCreation.map((policyTemplate) => {
-        const templateContent = policyTemplate.content;
-        const contentArray = extractTipTapContentArray(templateContent);
-        return {
-          name: policyTemplate.name,
-          description: policyTemplate.description,
-          department: policyTemplate.department,
-          frequency: policyTemplate.frequency,
-          content: { set: contentArray },
-          organizationId: organizationId,
-          policyTemplateId: policyTemplate.id,
-        };
-      }),
+      data: preparedPolicies.map(({ template, policyId, contentArray }) => ({
+        id: policyId,
+        name: template.name,
+        description: template.description,
+        department: template.department,
+        frequency: template.frequency,
+        content: { set: contentArray },
+        organizationId: organizationId,
+        policyTemplateId: template.id,
+      })),
     });
 
-    // Fetch newly created policies to create versions for them
-    const newlyCreatedPolicies = await tx.policy.findMany({
-      where: {
-        organizationId: organizationId,
-        policyTemplateId: {
-          in: policyTemplatesForCreation.map((t) => t.id),
-        },
-      },
-      select: { id: true, policyTemplateId: true, content: true },
+    await tx.policyVersion.createMany({
+      data: preparedPolicies.map(({ policyId, versionId, contentArray }) => ({
+        id: versionId,
+        policyId,
+        version: 1,
+        content: { set: contentArray },
+        changelog: 'Initial version from template',
+      })),
     });
 
-    // Create version 1 for each newly created policy
-    if (newlyCreatedPolicies.length > 0) {
-      await tx.policyVersion.createMany({
-        data: newlyCreatedPolicies.map((policy) => ({
-          policyId: policy.id,
-          version: 1,
-          content: { set: policy.content as Prisma.InputJsonValue[] },
-          changelog: 'Initial version from template',
-        })),
-      });
-
-      // Fetch the created versions to update policies with currentVersionId
-      const createdVersions = await tx.policyVersion.findMany({
-        where: {
-          policyId: { in: newlyCreatedPolicies.map((p) => p.id) },
-          version: 1,
-        },
-        select: { id: true, policyId: true },
-      });
-
-      // Update each policy with its currentVersionId
-      for (const version of createdVersions) {
-        await tx.policy.update({
-          where: { id: version.policyId },
-          data: { currentVersionId: version.id },
-        });
-      }
-    }
+    const currentVersionValues = Prisma.join(
+      preparedPolicies.map(
+        ({ policyId, versionId }) =>
+          Prisma.sql`(${policyId}::text, ${versionId}::text)`,
+      ),
+    );
+    await tx.$executeRaw`
+      UPDATE "Policy"
+      SET "currentVersionId" = v.version_id
+      FROM (VALUES ${currentVersionValues}) AS v(policy_id, version_id)
+      WHERE "Policy".id = v.policy_id
+    `;
   }
 
   /**
@@ -350,6 +350,8 @@ export const _upsertOrgFrameworkStructureCore = async ({
   );
 
   const requirementMapEntriesToCreate: Prisma.RequirementMapCreateManyInput[] = [];
+  const controlToPolicyPairs: Array<{ controlId: string; policyId: string }> = [];
+  const controlToTaskPairs: Array<{ controlId: string; taskId: string }> = [];
 
   for (const controlTemplateRelation of groupedControlTemplateRelations) {
     const newControlId = controlTemplateIdToInstanceIdMap.get(
@@ -363,81 +365,86 @@ export const _upsertOrgFrameworkStructureCore = async ({
       continue;
     }
 
-    const updateData: Prisma.ControlUpdateInput = {};
-    let needsUpdate = false;
-
     // --- Process Requirements for RequirementMap ---
-    if (controlTemplateRelation.requirementTemplateIds.length > 0) {
-      for (const reqTemplateId of controlTemplateRelation.requirementTemplateIds) {
-        let frameworkEditorFrameworkIdForReq: string | undefined;
-        for (const fw of frameworkEditorFrameworks) {
-          if (fw.requirements.some((r) => r.id === reqTemplateId)) {
-            frameworkEditorFrameworkIdForReq = fw.id;
-            break;
-          }
-        }
-        const frameworkInstanceId = frameworkEditorFrameworkIdForReq
-          ? editorFrameworkIdToInstanceIdMap.get(frameworkEditorFrameworkIdForReq)
-          : undefined;
-
-        if (frameworkInstanceId) {
-          requirementMapEntriesToCreate.push({
-            controlId: newControlId,
-            requirementId: reqTemplateId,
-            frameworkInstanceId: frameworkInstanceId,
-          });
-        } else {
-          console.warn(
-            `UpsertOrgFrameworkStructureCore: Could not find FrameworkInstanceId for editor requirement ID ${reqTemplateId}. Cannot create RequirementMap for Control ${newControlId}.`,
-          );
+    for (const reqTemplateId of controlTemplateRelation.requirementTemplateIds) {
+      let frameworkEditorFrameworkIdForReq: string | undefined;
+      for (const fw of frameworkEditorFrameworks) {
+        if (fw.requirements.some((r) => r.id === reqTemplateId)) {
+          frameworkEditorFrameworkIdForReq = fw.id;
+          break;
         }
       }
+      const frameworkInstanceId = frameworkEditorFrameworkIdForReq
+        ? editorFrameworkIdToInstanceIdMap.get(frameworkEditorFrameworkIdForReq)
+        : undefined;
+
+      if (frameworkInstanceId) {
+        requirementMapEntriesToCreate.push({
+          controlId: newControlId,
+          requirementId: reqTemplateId,
+          frameworkInstanceId: frameworkInstanceId,
+        });
+      } else {
+        console.warn(
+          `UpsertOrgFrameworkStructureCore: Could not find FrameworkInstanceId for editor requirement ID ${reqTemplateId}. Cannot create RequirementMap for Control ${newControlId}.`,
+        );
+      }
     }
 
-    // --- Connect Policies ---
-    if (controlTemplateRelation.policyTemplateIds.length > 0) {
-      const policiesToConnect = [];
-      for (const policyTemplateId of controlTemplateRelation.policyTemplateIds) {
-        const newPolicyId = policyTemplateIdToInstanceIdMap.get(policyTemplateId);
-        if (newPolicyId) {
-          policiesToConnect.push({ id: newPolicyId });
-        } else {
-          console.warn(
-            `UpsertOrgFrameworkStructureCore: Policy instance not found for template ID ${policyTemplateId}. Cannot connect to Control ${newControlId}.`,
-          );
-        }
-      }
-      if (policiesToConnect.length > 0) {
-        updateData.policies = { connect: policiesToConnect };
-        needsUpdate = true;
+    // --- Collect Control <-> Policy pairs ---
+    for (const policyTemplateId of controlTemplateRelation.policyTemplateIds) {
+      const newPolicyId = policyTemplateIdToInstanceIdMap.get(policyTemplateId);
+      if (newPolicyId) {
+        controlToPolicyPairs.push({ controlId: newControlId, policyId: newPolicyId });
+      } else {
+        console.warn(
+          `UpsertOrgFrameworkStructureCore: Policy instance not found for template ID ${policyTemplateId}. Cannot connect to Control ${newControlId}.`,
+        );
       }
     }
 
-    // --- Connect Tasks ---
-    if (controlTemplateRelation.taskTemplateIds.length > 0) {
-      const tasksToConnect = [];
-      for (const taskTemplateId of controlTemplateRelation.taskTemplateIds) {
-        const newTaskId = taskTemplateIdToInstanceIdMap.get(taskTemplateId);
-        if (newTaskId) {
-          tasksToConnect.push({ id: newTaskId });
-        } else {
-          console.warn(
-            `UpsertOrgFrameworkStructureCore: Task instance not found for template ID ${taskTemplateId}. Cannot connect to Control ${newControlId}.`,
-          );
-        }
-      }
-      if (tasksToConnect.length > 0) {
-        updateData.tasks = { connect: tasksToConnect };
-        needsUpdate = true;
+    // --- Collect Control <-> Task pairs ---
+    for (const taskTemplateId of controlTemplateRelation.taskTemplateIds) {
+      const newTaskId = taskTemplateIdToInstanceIdMap.get(taskTemplateId);
+      if (newTaskId) {
+        controlToTaskPairs.push({ controlId: newControlId, taskId: newTaskId });
+      } else {
+        console.warn(
+          `UpsertOrgFrameworkStructureCore: Task instance not found for template ID ${taskTemplateId}. Cannot connect to Control ${newControlId}.`,
+        );
       }
     }
+  }
 
-    if (needsUpdate) {
-      await tx.control.update({
-        where: { id: newControlId },
-        data: updateData,
-      });
-    }
+  // Bulk-insert into the implicit M2M join tables instead of N `control.update({ connect })`
+  // calls. ON CONFLICT DO NOTHING preserves the idempotency the connect loop provided for
+  // re-runs where some links already exist (e.g., adding a framework to an existing org).
+  if (controlToPolicyPairs.length > 0) {
+    const rows = Prisma.join(
+      controlToPolicyPairs.map(
+        ({ controlId, policyId }) =>
+          Prisma.sql`(${controlId}::text, ${policyId}::text)`,
+      ),
+    );
+    await tx.$executeRaw`
+      INSERT INTO "_ControlToPolicy" ("A", "B")
+      VALUES ${rows}
+      ON CONFLICT ("A", "B") DO NOTHING
+    `;
+  }
+
+  if (controlToTaskPairs.length > 0) {
+    const rows = Prisma.join(
+      controlToTaskPairs.map(
+        ({ controlId, taskId }) =>
+          Prisma.sql`(${controlId}::text, ${taskId}::text)`,
+      ),
+    );
+    await tx.$executeRaw`
+      INSERT INTO "_ControlToTask" ("A", "B")
+      VALUES ${rows}
+      ON CONFLICT ("A", "B") DO NOTHING
+    `;
   }
 
   // --- Create RequirementMap entries ---
