refactor(fork): page the re-home lookups and reuse the plan's identity rows

Self-review of the folder-transit change surfaced two scaling problems in
the re-home pass, both of which grow with the size of the fork edge rather
than the size of the sync:

- The resource lookups built `IN (...)` lists straight from the edge's
  mapping rows, so a large fork could hand Postgres a list approaching the
  bind-parameter ceiling and a pathological query plan. Page them at 500,
  matching the paging the rest of the fork copy already uses.
- The pass re-read the whole edge mapping via `getEdgeMappingRows`, which
  the promote plan had already loaded in the same transaction — a second
  full load of identical rows. Expose them on `ForkPromotePlan` and pass
  them in, which also drops a mock from the re-home tests.

Also tally moved rows from `returning()` rather than the planned batch
size, so the log line reports what the `folder_id IS NULL` guard actually
wrote instead of what was attempted.

Author: waleedlatif1

apps/sim/ee/workspace-forking/lib/promote/promote-plan.ts

@@ -8,6 +8,7 @@ import type { ForkEdge } from '@/ee/workspace-forking/lib/lineage/lineage'
 import { detectForkCascadeReferences } from '@/ee/workspace-forking/lib/mapping/cascade'
 import {
   buildForkResolver,
+  type ForkMappingRow,
   getEdgeMappingRows,
   resourceTypeToForkKind,
 } from '@/ee/workspace-forking/lib/mapping/mapping-store'
@@ -89,6 +90,12 @@ export interface ForkPromotePlan {
   willUpdate: number
   willCreate: number
   willArchive: number
+  /**
+   * The edge's persisted identity rows, already read to build {@link ForkPromotePlan.resolver}.
+   * Exposed so later stages of the same transaction reuse them instead of re-reading the whole
+   * edge mapping, which for a large fork is a second full load of the same rows.
+   */
+  mappingRows: ForkMappingRow[]
 }
 
 /**
@@ -551,5 +558,6 @@ export async function computeForkPromotePlan(params: {
     willUpdate,
     willCreate,
     willArchive: archivedTargetIds.length,
+    mappingRows,
   }
 }

apps/sim/ee/workspace-forking/lib/promote/promote.ts

@@ -567,7 +567,7 @@ export async function promoteFork(params: PromoteForkParams): Promise<PromoteFor
     // no-op and never overrides a placement chosen in the target.
     const rehomeResult = await rehomeFlattenedForkResources({
       tx,
-      edge,
+      mappingRows: plan.mappingRows,
       sourceWorkspaceId,
       targetWorkspaceId,
       direction,

apps/sim/ee/workspace-forking/lib/promote/rehome-mapped.test.ts

@@ -4,15 +4,7 @@
 import { folder as folderTable, knowledgeBase, workspaceFiles } from '@sim/db/schema'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import type { DbOrTx } from '@/lib/db/types'
-
-const { mockGetEdgeMappingRows } = vi.hoisted(() => ({
-  mockGetEdgeMappingRows: vi.fn(),
-}))
-
-vi.mock('@/ee/workspace-forking/lib/mapping/mapping-store', () => ({
-  getEdgeMappingRows: mockGetEdgeMappingRows,
-}))
-
+import type { ForkMappingRow } from '@/ee/workspace-forking/lib/mapping/mapping-store'
 import { rehomeFlattenedForkResources } from '@/ee/workspace-forking/lib/promote/rehome-mapped'
 
 interface UpdateCall {
@@ -55,21 +47,38 @@ function makeTx(rows: {
       set: (values: Record<string, unknown>) => ({
         where: () => {
           updates.push({ table, values })
-          return Promise.resolve()
+          // The real update reports the rows it actually moved; echo one id back so the
+          // caller's tally reflects a genuine write rather than the planned batch size.
+          return { returning: () => Promise.resolve([{ id: 'moved' }]) }
         },
       }),
     }),
   }
   return { tx: tx as unknown as DbOrTx, updates, insertedFolders }
 }
 
-const edge = {
-  childWorkspaceId: 'child-ws',
-  parentWorkspaceId: 'parent-ws',
-} as Parameters<typeof rehomeFlattenedForkResources>[0]['edge']
+const fileMapping: ForkMappingRow[] = [
+  {
+    id: 'map-1',
+    childWorkspaceId: 'child-ws',
+    resourceType: 'file',
+    parentResourceId: 'workspace/parent-ws/a.png',
+    childResourceId: 'workspace/child-ws/a.png',
+  },
+]
+
+const kbMapping: ForkMappingRow[] = [
+  {
+    id: 'map-2',
+    childWorkspaceId: 'child-ws',
+    resourceType: 'knowledge_base',
+    parentResourceId: 'kb-target',
+    childResourceId: 'kb-source',
+  },
+]
 
 const baseParams = {
-  edge,
+  mappingRows: [] as ForkMappingRow[],
   sourceWorkspaceId: 'child-ws',
   targetWorkspaceId: 'parent-ws',
   direction: 'push' as const,
@@ -80,20 +89,10 @@ const baseParams = {
 describe('rehomeFlattenedForkResources', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockGetEdgeMappingRows.mockResolvedValue([])
   })
 
   it('mirrors the source folder and moves a root-flattened mapped file into it', async () => {
     // Push: the child is the source, so the mapping row's child side is the source key.
-    mockGetEdgeMappingRows.mockResolvedValue([
-      {
-        id: 'map-1',
-        childWorkspaceId: 'child-ws',
-        resourceType: 'file',
-        parentResourceId: 'workspace/parent-ws/a.png',
-        childResourceId: 'workspace/child-ws/a.png',
-      },
-    ])
     const { tx, updates, insertedFolders } = makeTx({
       // Both the target lookup (flattened row) and the source lookup read this table; the
       // rows carry the fields each phase needs.
@@ -113,7 +112,11 @@ describe('rehomeFlattenedForkResources', () => {
       ],
     })
 
-    const result = await rehomeFlattenedForkResources({ ...baseParams, tx })
+    const result = await rehomeFlattenedForkResources({
+      ...baseParams,
+      mappingRows: fileMapping,
+      tx,
+    })
 
     expect(insertedFolders).toHaveLength(1)
     expect(insertedFolders[0]).toMatchObject({
@@ -130,39 +133,25 @@ describe('rehomeFlattenedForkResources', () => {
   })
 
   it('leaves a resource alone when the source itself sits at the root', async () => {
-    mockGetEdgeMappingRows.mockResolvedValue([
-      {
-        id: 'map-1',
-        childWorkspaceId: 'child-ws',
-        resourceType: 'file',
-        parentResourceId: 'workspace/parent-ws/a.png',
-        childResourceId: 'workspace/child-ws/a.png',
-      },
-    ])
     const { tx, updates, insertedFolders } = makeTx({
       files: [
         { id: 'file-target', key: 'workspace/parent-ws/a.png', folderId: null },
         { id: 'file-source', key: 'workspace/child-ws/a.png', folderId: null },
       ],
     })
 
-    const result = await rehomeFlattenedForkResources({ ...baseParams, tx })
+    const result = await rehomeFlattenedForkResources({
+      ...baseParams,
+      mappingRows: fileMapping,
+      tx,
+    })
 
     expect(insertedFolders).toHaveLength(0)
     expect(updates).toHaveLength(0)
     expect(result.rehomed.file).toBe(0)
   })
 
   it('never touches a target already placed in a folder, so a deliberate move survives a re-sync', async () => {
-    mockGetEdgeMappingRows.mockResolvedValue([
-      {
-        id: 'map-1',
-        childWorkspaceId: 'child-ws',
-        resourceType: 'knowledge_base',
-        parentResourceId: 'kb-target',
-        childResourceId: 'kb-source',
-      },
-    ])
     // The target read filters on `folderId IS NULL`, so an already-placed row is simply absent.
     const { tx, updates } = makeTx({
       knowledgeBases: [],
@@ -178,7 +167,11 @@ describe('rehomeFlattenedForkResources', () => {
       ],
     })
 
-    const result = await rehomeFlattenedForkResources({ ...baseParams, tx })
+    const result = await rehomeFlattenedForkResources({
+      ...baseParams,
+      mappingRows: kbMapping,
+      tx,
+    })
 
     expect(updates).toHaveLength(0)
     expect(result.rehomed.knowledge_base).toBe(0)
@@ -195,15 +188,6 @@ describe('rehomeFlattenedForkResources', () => {
   })
 
   it('orients pull the other way: the parent side is the source', async () => {
-    mockGetEdgeMappingRows.mockResolvedValue([
-      {
-        id: 'map-1',
-        childWorkspaceId: 'child-ws',
-        resourceType: 'file',
-        parentResourceId: 'workspace/parent-ws/a.png',
-        childResourceId: 'workspace/child-ws/a.png',
-      },
-    ])
     const { tx, updates } = makeTx({
       files: [
         // On a pull the CHILD is the target, so its key is the one that must still be flattened.
@@ -224,6 +208,7 @@ describe('rehomeFlattenedForkResources', () => {
 
     const result = await rehomeFlattenedForkResources({
       ...baseParams,
+      mappingRows: fileMapping,
       tx,
       direction: 'pull',
       sourceWorkspaceId: 'parent-ws',