diff --git a/apps/app/src/actions/organization/lib/initialize-organization.test.ts b/apps/app/src/actions/organization/lib/initialize-organization.test.ts new file mode 100644 index 0000000000..ed084b6cc1 --- /dev/null +++ b/apps/app/src/actions/organization/lib/initialize-organization.test.ts @@ -0,0 +1,247 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +const mockTx = { + frameworkInstance: { + findMany: vi.fn(), + createMany: vi.fn(), + }, + control: { + findMany: vi.fn(), + createMany: vi.fn(), + }, + policy: { + findMany: vi.fn(), + createMany: vi.fn(), + update: vi.fn(), + }, + policyVersion: { + findMany: vi.fn(), + createMany: vi.fn(), + }, + task: { + findMany: vi.fn(), + createMany: vi.fn(), + }, + requirementMap: { createMany: vi.fn() }, + controlDocumentType: { createMany: vi.fn() }, + frameworkControlPolicyLink: { createMany: vi.fn() }, + frameworkControlTaskLink: { createMany: vi.fn() }, + frameworkControlDocumentTypeLink: { createMany: vi.fn() }, + $executeRaw: vi.fn(), + $queryRaw: vi.fn(), +}; + +vi.mock('./load-framework-sources', () => ({ + loadFrameworkSources: vi.fn(), +})); + +import { _upsertOrgFrameworkStructureCore } from './initialize-organization'; +import { loadFrameworkSources } from './load-framework-sources'; +import type { LoadedFrameworkSources } from './load-framework-sources'; + +const mockedLoadSources = vi.mocked(loadFrameworkSources); + +function buildSources({ + frameworkId = 'fw_1', + controlTemplateId = 'ct_1', + policyTemplateId = 'pt_1', + taskTemplateId = 'tt_1', +}: { + frameworkId?: string; + controlTemplateId?: string; + policyTemplateId?: string; + taskTemplateId?: string; +} = {}) { + return { + controlTemplates: [ + { id: controlTemplateId, name: 'C1', description: 'd', documentTypes: ['access-request'] }, + ], + policyTemplates: [ + { id: policyTemplateId, name: 'P1', description: 'd', content: [], frequency: 'yearly', department: 'none' }, + ], + taskTemplates: [ + { id: taskTemplateId, name: 'T1', description: 'd', frequency: null, department: null, automationStatus: 'AUTOMATED' }, + ], + groupedRelations: [ + { + frameworkId, + controlTemplateId, + requirementTemplateIds: ['req_1'], + policyTemplateIds: [policyTemplateId], + taskTemplateIds: [taskTemplateId], + documentTypes: ['access-request'], + }, + ], + latestVersionByFrameworkId: new Map([[frameworkId, 'fv_1']]), + frameworksWithoutVersion: [], + requirementToFrameworkId: new Map([['req_1', frameworkId]]), + } satisfies Record as unknown as LoadedFrameworkSources; +} + +describe('_upsertOrgFrameworkStructureCore', () => { + const ORG_ID = 'org_test'; + const FRAMEWORK_EDITOR_ID = 'fw_1'; + const FRAMEWORK_INSTANCE_ID = 'fi_1'; + const CONTROL_ID = 'ctrl_1'; + const POLICY_ID = 'pol_1'; + const TASK_ID = 'tsk_1'; + + beforeEach(() => { + vi.clearAllMocks(); + + mockedLoadSources.mockResolvedValue(buildSources()); + + // No existing framework instances + mockTx.frameworkInstance.findMany + .mockResolvedValueOnce([]) // existing check + .mockResolvedValueOnce([{ id: FRAMEWORK_INSTANCE_ID, frameworkId: FRAMEWORK_EDITOR_ID }]); // all instances + + mockTx.frameworkInstance.createMany.mockResolvedValue({ count: 1 }); + + // No existing controls/policies/tasks + mockTx.control.findMany + .mockResolvedValueOnce([]) // existing check + .mockResolvedValueOnce([{ id: CONTROL_ID, controlTemplateId: 'ct_1' }]); // all controls + mockTx.control.createMany.mockResolvedValue({ count: 1 }); + + mockTx.policy.findMany + .mockResolvedValueOnce([]) // existing check + .mockResolvedValueOnce([{ id: POLICY_ID, policyTemplateId: 'pt_1' }]); // all policies + mockTx.policy.createMany.mockResolvedValue({ count: 1 }); + + mockTx.$queryRaw.mockResolvedValue([{ policy_id: POLICY_ID, version_id: 'pv_1' }]); + mockTx.policyVersion.createMany.mockResolvedValue({ count: 1 }); + mockTx.$executeRaw.mockResolvedValue(1); + + mockTx.task.findMany + .mockResolvedValueOnce([]) // existing check + .mockResolvedValueOnce([{ id: TASK_ID, taskTemplateId: 'tt_1' }]); // all tasks + mockTx.task.createMany.mockResolvedValue({ count: 1 }); + + mockTx.requirementMap.createMany.mockResolvedValue({ count: 1 }); + mockTx.controlDocumentType.createMany.mockResolvedValue({ count: 1 }); + mockTx.frameworkControlPolicyLink.createMany.mockResolvedValue({ count: 1 }); + mockTx.frameworkControlTaskLink.createMany.mockResolvedValue({ count: 1 }); + mockTx.frameworkControlDocumentTypeLink.createMany.mockResolvedValue({ count: 1 }); + }); + + const callUpsert = () => + _upsertOrgFrameworkStructureCore({ + organizationId: ORG_ID, + targetFrameworkEditorIds: [FRAMEWORK_EDITOR_ID], + frameworkEditorFrameworks: [ + { id: FRAMEWORK_EDITOR_ID, name: 'SOC 2', requirements: [] } as any, + ], + tx: mockTx as any, + }); + + it('creates FrameworkControlPolicyLink entries alongside _ControlToPolicy', async () => { + await callUpsert(); + + // Old table: _ControlToPolicy via $executeRaw (called at least once for policies) + expect(mockTx.$executeRaw).toHaveBeenCalled(); + + // New table: FrameworkControlPolicyLink + expect(mockTx.frameworkControlPolicyLink.createMany).toHaveBeenCalledWith({ + data: [ + { + frameworkInstanceId: FRAMEWORK_INSTANCE_ID, + controlId: CONTROL_ID, + policyId: POLICY_ID, + }, + ], + skipDuplicates: true, + }); + }); + + it('creates FrameworkControlTaskLink entries alongside _ControlToTask', async () => { + await callUpsert(); + + // Old table: _ControlToTask via $executeRaw + expect(mockTx.$executeRaw).toHaveBeenCalled(); + + expect(mockTx.frameworkControlTaskLink.createMany).toHaveBeenCalledWith({ + data: [ + { + frameworkInstanceId: FRAMEWORK_INSTANCE_ID, + controlId: CONTROL_ID, + taskId: TASK_ID, + }, + ], + skipDuplicates: true, + }); + }); + + it('creates FrameworkControlDocumentTypeLink entries alongside ControlDocumentType', async () => { + await callUpsert(); + + expect(mockTx.controlDocumentType.createMany).toHaveBeenCalled(); + + expect(mockTx.frameworkControlDocumentTypeLink.createMany).toHaveBeenCalledWith({ + data: [ + { + frameworkInstanceId: FRAMEWORK_INSTANCE_ID, + controlId: CONTROL_ID, + formType: 'access-request', + }, + ], + skipDuplicates: true, + }); + }); + + it('skips framework-scoped links when frameworkInstanceId cannot be resolved', async () => { + mockedLoadSources.mockResolvedValue( + buildSources({ frameworkId: 'fw_unknown' }), + ); + + await callUpsert(); + + expect(mockTx.frameworkControlPolicyLink.createMany).not.toHaveBeenCalled(); + expect(mockTx.frameworkControlTaskLink.createMany).not.toHaveBeenCalled(); + expect(mockTx.frameworkControlDocumentTypeLink.createMany).not.toHaveBeenCalled(); + }); + + it('handles control appearing in multiple frameworks', async () => { + const FI_2 = 'fi_2'; + mockedLoadSources.mockResolvedValue({ + ...buildSources(), + groupedRelations: [ + { + frameworkId: FRAMEWORK_EDITOR_ID, + controlTemplateId: 'ct_1', + requirementTemplateIds: ['req_1'], + policyTemplateIds: ['pt_1'], + taskTemplateIds: [], + documentTypes: [], + }, + { + frameworkId: 'fw_2', + controlTemplateId: 'ct_1', + requirementTemplateIds: ['req_2'], + policyTemplateIds: ['pt_1'], + taskTemplateIds: [], + documentTypes: [], + }, + ], + } as LoadedFrameworkSources); + + // Return both framework instances + mockTx.frameworkInstance.findMany + .mockReset() + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([ + { id: FRAMEWORK_INSTANCE_ID, frameworkId: FRAMEWORK_EDITOR_ID }, + { id: FI_2, frameworkId: 'fw_2' }, + ]); + + await callUpsert(); + + expect(mockTx.frameworkControlPolicyLink.createMany).toHaveBeenCalledWith({ + data: expect.arrayContaining([ + { frameworkInstanceId: FRAMEWORK_INSTANCE_ID, controlId: CONTROL_ID, policyId: POLICY_ID }, + { frameworkInstanceId: FI_2, controlId: CONTROL_ID, policyId: POLICY_ID }, + ]), + skipDuplicates: true, + }); + }); +}); diff --git a/apps/app/src/actions/organization/lib/initialize-organization.ts b/apps/app/src/actions/organization/lib/initialize-organization.ts index 846a5b0ba2..3db4eed756 100644 --- a/apps/app/src/actions/organization/lib/initialize-organization.ts +++ b/apps/app/src/actions/organization/lib/initialize-organization.ts @@ -312,6 +312,9 @@ export const _upsertOrgFrameworkStructureCore = async ({ const controlToPolicyPairs: Array<{ controlId: string; policyId: string }> = []; const controlToTaskPairs: Array<{ controlId: string; taskId: string }> = []; const controlDocumentTypeEntries: Prisma.ControlDocumentTypeCreateManyInput[] = []; + const frameworkControlPolicyEntries: Prisma.FrameworkControlPolicyLinkCreateManyInput[] = []; + const frameworkControlTaskEntries: Prisma.FrameworkControlTaskLinkCreateManyInput[] = []; + const frameworkControlDocumentTypeEntries: Prisma.FrameworkControlDocumentTypeLinkCreateManyInput[] = []; const controlTemplateById = new Map(controlTemplates.map((c) => [c.id, c])); for (const controlTemplateRelation of groupedControlTemplateRelations) { @@ -326,33 +329,28 @@ export const _upsertOrgFrameworkStructureCore = async ({ continue; } - for (const reqTemplateId of controlTemplateRelation.requirementTemplateIds) { - const frameworkEditorFrameworkIdForReq = requirementToFrameworkId.get(reqTemplateId); - const frameworkInstanceId = frameworkEditorFrameworkIdForReq - ? editorFrameworkIdToInstanceIdMap.get(frameworkEditorFrameworkIdForReq) - : undefined; + const frameworkInstanceId = editorFrameworkIdToInstanceIdMap.get( + controlTemplateRelation.frameworkId, + ); + if (!frameworkInstanceId) continue; - 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) { + requirementMapEntriesToCreate.push({ + controlId: newControlId, + requirementId: reqTemplateId, + frameworkInstanceId, + }); } 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}.`, - ); + frameworkControlPolicyEntries.push({ + frameworkInstanceId, + controlId: newControlId, + policyId: newPolicyId, + }); } } @@ -360,20 +358,24 @@ export const _upsertOrgFrameworkStructureCore = async ({ 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}.`, - ); + frameworkControlTaskEntries.push({ + frameworkInstanceId, + controlId: newControlId, + taskId: newTaskId, + }); } } - // ControlDocumentType: explicit junction rows driven from manifest/live - // documentTypes so the org starts with the same evidence form types the - // published version specified. Deduped against existing rows via the - // unique constraint. - const ct = controlTemplateById.get(controlTemplateRelation.controlTemplateId); - for (const formType of ct?.documentTypes ?? []) { + const documentTypes = controlTemplateRelation.documentTypes.length > 0 + ? controlTemplateRelation.documentTypes + : (controlTemplateById.get(controlTemplateRelation.controlTemplateId)?.documentTypes ?? []); + for (const formType of documentTypes) { controlDocumentTypeEntries.push({ controlId: newControlId, formType }); + frameworkControlDocumentTypeEntries.push({ + frameworkInstanceId, + controlId: newControlId, + formType, + }); } } @@ -422,6 +424,27 @@ export const _upsertOrgFrameworkStructureCore = async ({ }); } + if (frameworkControlPolicyEntries.length > 0) { + await tx.frameworkControlPolicyLink.createMany({ + data: frameworkControlPolicyEntries, + skipDuplicates: true, + }); + } + + if (frameworkControlTaskEntries.length > 0) { + await tx.frameworkControlTaskLink.createMany({ + data: frameworkControlTaskEntries, + skipDuplicates: true, + }); + } + + if (frameworkControlDocumentTypeEntries.length > 0) { + await tx.frameworkControlDocumentTypeLink.createMany({ + data: frameworkControlDocumentTypeEntries, + skipDuplicates: true, + }); + } + return { processedFrameworks: frameworkEditorFrameworks, controlTemplates, diff --git a/apps/app/src/actions/organization/lib/load-framework-sources.ts b/apps/app/src/actions/organization/lib/load-framework-sources.ts index b0f0d11fe8..aba9c000a1 100644 --- a/apps/app/src/actions/organization/lib/load-framework-sources.ts +++ b/apps/app/src/actions/organization/lib/load-framework-sources.ts @@ -80,10 +80,12 @@ export interface LoadedFrameworkSources { automationStatus: TaskAutomationStatus; }>; groupedRelations: Array<{ + frameworkId: string; controlTemplateId: string; requirementTemplateIds: string[]; policyTemplateIds: string[]; taskTemplateIds: string[]; + documentTypes: EvidenceFormType[]; }>; latestVersionByFrameworkId: Map; frameworksWithoutVersion: string[]; @@ -128,22 +130,27 @@ export async function loadFrameworkSources({ const relationsByControl = new Map< string, { + frameworkId: string; controlTemplateId: string; requirementTemplateIds: Set; policyTemplateIds: Set; taskTemplateIds: Set; + documentTypes: Set; } >(); - const getOrCreateRelation = (controlTemplateId: string) => { - let rel = relationsByControl.get(controlTemplateId); + const getOrCreateRelation = (frameworkId: string, controlTemplateId: string) => { + const key = `${frameworkId}::${controlTemplateId}`; + let rel = relationsByControl.get(key); if (!rel) { rel = { + frameworkId, controlTemplateId, requirementTemplateIds: new Set(), policyTemplateIds: new Set(), taskTemplateIds: new Set(), + documentTypes: new Set(), }; - relationsByControl.set(controlTemplateId, rel); + relationsByControl.set(key, rel); } return rel; }; @@ -161,10 +168,11 @@ export async function loadFrameworkSources({ documentTypes: (c.documentTypes ?? []) as EvidenceFormType[], }); } - const rel = getOrCreateRelation(c.id); + const rel = getOrCreateRelation(frameworkId, c.id); for (const rid of c.requirementIds) rel.requirementTemplateIds.add(rid); for (const pid of c.policyIds) rel.policyTemplateIds.add(pid); for (const tid of c.taskIds) rel.taskTemplateIds.add(tid); + for (const dt of c.documentTypes ?? []) rel.documentTypes.add(dt as EvidenceFormType); } for (const p of manifest.policies) { if (!policiesMap.has(p.id)) { @@ -247,10 +255,23 @@ export async function loadFrameworkSources({ }, }); for (const cr of controlRelationsLive) { - const rel = getOrCreateRelation(cr.id); - for (const r of cr.requirements) rel.requirementTemplateIds.add(r.id); - for (const p of cr.policyTemplates) rel.policyTemplateIds.add(p.id); - for (const t of cr.taskTemplates) rel.taskTemplateIds.add(t.id); + const frameworkIds = new Set( + cr.requirements + .map((r) => requirementToFrameworkId.get(r.id)) + .filter((id): id is string => Boolean(id)), + ); + for (const fwId of frameworkIds) { + const rel = getOrCreateRelation(fwId, cr.id); + for (const r of cr.requirements) { + if (requirementToFrameworkId.get(r.id) === fwId) { + rel.requirementTemplateIds.add(r.id); + } + } + for (const p of cr.policyTemplates) rel.policyTemplateIds.add(p.id); + for (const t of cr.taskTemplates) rel.taskTemplateIds.add(t.id); + const controlEntry = controlsMap.get(cr.id); + for (const dt of controlEntry?.documentTypes ?? []) rel.documentTypes.add(dt); + } } const fallbackPolicyIds = controlRelationsLive.flatMap((cr) => @@ -297,10 +318,12 @@ export async function loadFrameworkSources({ } const groupedRelations = Array.from(relationsByControl.values()).map((rel) => ({ + frameworkId: rel.frameworkId, controlTemplateId: rel.controlTemplateId, requirementTemplateIds: Array.from(rel.requirementTemplateIds), policyTemplateIds: Array.from(rel.policyTemplateIds), taskTemplateIds: Array.from(rel.taskTemplateIds), + documentTypes: Array.from(rel.documentTypes), })); return { diff --git a/packages/db/prisma/migrations/20260514200000_backfill_missing_framework_control_links/migration.sql b/packages/db/prisma/migrations/20260514200000_backfill_missing_framework_control_links/migration.sql new file mode 100644 index 0000000000..7a5885beab --- /dev/null +++ b/packages/db/prisma/migrations/20260514200000_backfill_missing_framework_control_links/migration.sql @@ -0,0 +1,63 @@ +-- Backfill FrameworkControlPolicyLink for orgs that were onboarded after the +-- initial backfill migration ran but before the Next.js upsert code was updated. +-- Uses RequirementMap to scope _ControlToPolicy entries to specific framework instances. + +INSERT INTO "FrameworkControlPolicyLink" ( + "frameworkInstanceId", + "controlId", + "policyId" +) +SELECT DISTINCT + rm."frameworkInstanceId", + cp."A", + cp."B" +FROM "_ControlToPolicy" cp +JOIN "RequirementMap" rm ON rm."controlId" = cp."A" +WHERE rm."archivedAt" IS NULL + AND NOT EXISTS ( + SELECT 1 FROM "FrameworkControlPolicyLink" fpl + WHERE fpl."frameworkInstanceId" = rm."frameworkInstanceId" + AND fpl."controlId" = cp."A" + AND fpl."policyId" = cp."B" + ) +ON CONFLICT ("frameworkInstanceId", "controlId", "policyId") DO NOTHING; + +INSERT INTO "FrameworkControlTaskLink" ( + "frameworkInstanceId", + "controlId", + "taskId" +) +SELECT DISTINCT + rm."frameworkInstanceId", + ct."A", + ct."B" +FROM "_ControlToTask" ct +JOIN "RequirementMap" rm ON rm."controlId" = ct."A" +WHERE rm."archivedAt" IS NULL + AND NOT EXISTS ( + SELECT 1 FROM "FrameworkControlTaskLink" ftl + WHERE ftl."frameworkInstanceId" = rm."frameworkInstanceId" + AND ftl."controlId" = ct."A" + AND ftl."taskId" = ct."B" + ) +ON CONFLICT ("frameworkInstanceId", "controlId", "taskId") DO NOTHING; + +INSERT INTO "FrameworkControlDocumentTypeLink" ( + "frameworkInstanceId", + "controlId", + "formType" +) +SELECT DISTINCT + rm."frameworkInstanceId", + cdt."controlId", + cdt."formType" +FROM "ControlDocumentType" cdt +JOIN "RequirementMap" rm ON rm."controlId" = cdt."controlId" +WHERE rm."archivedAt" IS NULL + AND NOT EXISTS ( + SELECT 1 FROM "FrameworkControlDocumentTypeLink" fdl + WHERE fdl."frameworkInstanceId" = rm."frameworkInstanceId" + AND fdl."controlId" = cdt."controlId" + AND fdl."formType" = cdt."formType" + ) +ON CONFLICT ("frameworkInstanceId", "controlId", "formType") DO NOTHING;