fix(tables): preserve foreign key integrity

Author: j15z

apps/sim/ee/workspace-forking/lib/copy/copy-resources.test.ts

@@ -55,6 +55,8 @@ import {
   rebindKnowledgeDocumentSecretProvenance,
   replaceKnowledgeDocumentSecretProvenanceInTx,
 } from '@/lib/knowledge/secret-provenance'
+import { getColumnId } from '@/lib/table/column-keys'
+import { collectColumnReferencedTableIds } from '@/lib/table/column-types/registry.server'
 import { DEFAULT_TABLE_VIEW_NAME } from '@/lib/table/constants'
 import { nKeysBetween } from '@/lib/table/order-key'
 import {
@@ -90,7 +92,10 @@ import {
   type ForkReferenceResolver,
   rewriteEnvRefsInText,
 } from '@/ee/workspace-forking/lib/remap/remap-references'
-import { remapForkTableWorkflowGroups } from '@/ee/workspace-forking/lib/remap/remap-table-groups'
+import {
+  remapForkTableReferences,
+  remapForkTableWorkflowGroups,
+} from '@/ee/workspace-forking/lib/remap/remap-table-groups'
 
 const logger = createLogger('WorkspaceForkCopyResources')
 
@@ -99,6 +104,8 @@ const CONTENT_PAGE = 500
 const PROVENANCE_CONTENT_PAGE = 8
 const MAX_FORK_PROVENANCE_ENTRIES = 10_000
 const MAX_FORK_PROVENANCE_BYTES = 8 * 1024 * 1024
+/** Matches the fork contract's per-resource selection ceiling after dependencies are expanded. */
+export const MAX_FORK_TABLES_WITH_DEPENDENCIES = 2_000
 
 function isForkProvenancePageWithinBudget(sidecars: readonly { entries: unknown }[]): boolean {
   let entries = 0
@@ -216,6 +223,11 @@ export interface CopyResourcesParams {
    * plan resolver); omitted by fork-create, which preserves env names verbatim (no rewrite).
    */
   resolveEnvName?: (key: string) => string | null | undefined
+  /**
+   * Detect whether a referenced source table already maps to a target during promote. Row-level
+   * mappings do not exist yet, so the copy fails instead of inventing target row identities.
+   */
+  resolveMappedTableReference?: (sourceTableId: string) => string | null | undefined
   /**
    * Resolve a source block id to its target block id for copied tables' workflow-group
    * `outputs[].blockId`. Promote passes the SAME persisted-pair resolver its workflow writes
@@ -237,6 +249,13 @@ export interface ForkContentPlanEntry {
   childId: string
 }
 
+export interface ForkContentTableEntry extends ForkContentPlanEntry {
+  /** Copied tables this table's reference columns require to remain available. */
+  dependsOnChildIds?: string[]
+  /** Stable column id to copied target-table id, used to derive copied referenced-row ids. */
+  referenceColumnTargetTableIds?: Record<string, string>
+}
+
 /**
  * A KB to copy post-commit, plus the source-document -> child-document id map for the
  * documents that were pre-created as placeholders in the transaction (referenced by copied
@@ -288,7 +307,7 @@ export interface ForkContentPlan {
   childWorkspaceId: string
   /** Initiating user, recorded as the owner of copied KB-document blob bindings in the child. */
   userId: string
-  tables: ForkContentPlanEntry[]
+  tables: ForkContentTableEntry[]
   knowledgeBases: ForkContentKbEntry[]
   skills: ForkContentSkillEntry[]
   /** Documents copied into an already-existing target KB (sync-only; empty at fork create). */
@@ -359,6 +378,102 @@ function setId(idMap: Map<ForkResourceType, Map<string, string>>, type: ForkReso
  */
 type SkillSkeletonInsert = Omit<typeof skill.$inferInsert, 'content'> & { content: SQL }
 
+/** Derives the copied row identity without retaining an unbounded source-row map in memory. */
+function deriveCopiedTableRowId(childTableId: string, sourceRowId: string): string {
+  return `row_${sha256Hex(`table-row:${childTableId}:${sourceRowId}`).slice(0, 32)}`
+}
+
+/** Rewrites reference cells through the same deterministic identity used by copied target rows. */
+function remapCopiedReferenceCells(
+  data: unknown,
+  referenceColumnTargetTableIds: Readonly<Record<string, string>> | undefined
+): unknown {
+  if (!referenceColumnTargetTableIds || !isRecordLike(data)) return data
+  let remapped: Record<string, unknown> | undefined
+  for (const [columnId, childTableId] of Object.entries(referenceColumnTargetTableIds)) {
+    const sourceRowId = data[columnId]
+    if (typeof sourceRowId !== 'string' || sourceRowId.length === 0) continue
+    remapped ??= { ...data }
+    remapped[columnId] = deriveCopiedTableRowId(childTableId, sourceRowId)
+  }
+  return remapped ?? data
+}
+
+/**
+ * Loads the selected tables plus the transitive closure of tables named by their reference
+ * columns. Each layer is workspace-scoped and active-only; an unavailable dependency fails the
+ * copy instead of persisting a source-workspace table id into the child schema.
+ */
+async function loadTableDefinitionsWithDependencies(
+  tx: DbOrTx,
+  sourceWorkspaceId: string,
+  selectedTableIds: readonly string[],
+  resolveMappedTableReference?: (sourceTableId: string) => string | null | undefined
+): Promise<Array<typeof userTableDefinitions.$inferSelect>> {
+  const orderedIds = [...new Set(selectedTableIds)]
+  if (orderedIds.length > MAX_FORK_TABLES_WITH_DEPENDENCIES) {
+    throw new Error(
+      `Cannot copy more than ${MAX_FORK_TABLES_WITH_DEPENDENCIES} tables including referenced dependencies`
+    )
+  }
+  const scheduledIds = new Set(orderedIds)
+  const dependencyIds = new Set<string>()
+  const definitionsById = new Map<string, typeof userTableDefinitions.$inferSelect>()
+  let pendingIds = [...orderedIds]
+
+  while (pendingIds.length > 0) {
+    const batchIds = pendingIds
+    const batchIdSet = new Set(batchIds)
+    pendingIds = []
+    const rows = await tx
+      .select()
+      .from(userTableDefinitions)
+      .where(
+        and(
+          inArray(userTableDefinitions.id, batchIds),
+          eq(userTableDefinitions.workspaceId, sourceWorkspaceId),
+          isNull(userTableDefinitions.archivedAt)
+        )
+      )
+    const batchRows = rows.filter((row) => batchIdSet.has(row.id))
+
+    for (const row of batchRows) {
+      definitionsById.set(row.id, row)
+      const referencedIds = collectColumnReferencedTableIds((row.schema as TableSchema).columns)
+      for (const referencedId of referencedIds) {
+        dependencyIds.add(referencedId)
+        if (scheduledIds.has(referencedId)) continue
+        const mappedTableId = resolveMappedTableReference?.(referencedId)
+        if (mappedTableId) {
+          throw new Error(
+            `Referenced table ${referencedId} is mapped to ${mappedTableId}, but referenced row mappings are unavailable`
+          )
+        }
+        if (scheduledIds.size >= MAX_FORK_TABLES_WITH_DEPENDENCIES) {
+          throw new Error(
+            `Cannot copy more than ${MAX_FORK_TABLES_WITH_DEPENDENCIES} tables including referenced dependencies`
+          )
+        }
+        scheduledIds.add(referencedId)
+        orderedIds.push(referencedId)
+        pendingIds.push(referencedId)
+      }
+    }
+
+    const missingDependencyId = batchIds.find(
+      (id) => dependencyIds.has(id) && !definitionsById.has(id)
+    )
+    if (missingDependencyId) {
+      throw new Error(`Referenced table ${missingDependencyId} is unavailable for copy`)
+    }
+  }
+
+  return orderedIds.flatMap((id) => {
+    const definition = definitionsById.get(id)
+    return definition ? [definition] : []
+  })
+}
+
 /**
  * Copy the selected resources' **container rows** into the child workspace inside
  * the fork transaction: custom tools, skills, and MCP server configs (each a
@@ -627,16 +742,12 @@ export async function copyForkResourceContainers(
   }
 
   if (selection.tables.length > 0) {
-    const definitions = await tx
-      .select()
-      .from(userTableDefinitions)
-      .where(
-        and(
-          inArray(userTableDefinitions.id, selection.tables),
-          eq(userTableDefinitions.workspaceId, sourceWorkspaceId),
-          isNull(userTableDefinitions.archivedAt)
-        )
-      )
+    const definitions = await loadTableDefinitionsWithDependencies(
+      tx,
+      sourceWorkspaceId,
+      selection.tables,
+      params.resolveMappedTableReference
+    )
     const sourceViews =
       definitions.length > 0
         ? await tx
@@ -671,12 +782,22 @@ export async function copyForkResourceContainers(
 
     const inserts: (typeof userTableDefinitions.$inferInsert)[] = []
     const viewInserts: (typeof tableViews.$inferInsert)[] = []
+    const tableIdMap = new Map(
+      definitions.map((definition) => [definition.id, generateId()] as const)
+    )
+    for (const [sourceTableId, childTableId] of tableIdMap) {
+      record('table', sourceTableId, childTableId)
+    }
     for (const definition of definitions) {
-      const childTableId = generateId()
-      const remappedSchema = remapForkTableWorkflowGroups(
-        definition.schema as TableSchema,
-        workflowIdMap,
-        params.resolveBlockId
+      const childTableId = tableIdMap.get(definition.id)
+      if (!childTableId) throw new Error(`Missing copied table identity for ${definition.id}`)
+      const remappedSchema = remapForkTableReferences(
+        remapForkTableWorkflowGroups(
+          definition.schema as TableSchema,
+          workflowIdMap,
+          params.resolveBlockId
+        ),
+        tableIdMap
       )
       inserts.push({
         ...definition,
@@ -733,8 +854,27 @@ export async function copyForkResourceContainers(
           updatedAt: now,
         })
       }
-      record('table', definition.id, childTableId)
-      contentPlan.tables.push({ sourceId: definition.id, childId: childTableId })
+      const dependsOnChildIds = collectColumnReferencedTableIds(
+        (definition.schema as TableSchema).columns
+      ).flatMap((sourceId) => {
+        const dependencyId = tableIdMap.get(sourceId)
+        return dependencyId && dependencyId !== childTableId ? [dependencyId] : []
+      })
+      const referenceColumnTargetTableIds = Object.fromEntries(
+        (definition.schema as TableSchema).columns.flatMap((column) => {
+          const [sourceTargetId] = collectColumnReferencedTableIds([column])
+          const childTargetId = sourceTargetId ? tableIdMap.get(sourceTargetId) : undefined
+          return childTargetId ? [[getColumnId(column), childTargetId]] : []
+        })
+      )
+      contentPlan.tables.push({
+        sourceId: definition.id,
+        childId: childTableId,
+        ...(dependsOnChildIds.length > 0 ? { dependsOnChildIds } : {}),
+        ...(Object.keys(referenceColumnTargetTableIds).length > 0
+          ? { referenceColumnTargetTableIds }
+          : {}),
+      })
       names.tables.push(definition.name)
     }
     if (inserts.length > 0) await tx.insert(userTableDefinitions).values(inserts)
@@ -1245,14 +1385,17 @@ export async function copyForkResourceContent(params: {
           return {
             row: {
               ...row,
-              id: generateId(),
+              id: deriveCopiedTableRowId(table.childId, row.id),
               tableId: table.childId,
               workspaceId: childWorkspaceId,
               orderKey: row.orderKey ?? mintedKeys[mintedIdx++] ?? null,
               secretProvenanceVersion:
                 classification.mode === 'legacy' ? null : TABLE_ROW_SECRET_PROVENANCE_VERSION,
               // Repoint resource-chip URLs in cell data at the child copies (no-op when no maps).
-              data: contentRefMaps ? remapTableRowResourceUrls(row.data, contentRefMaps) : row.data,
+              data: remapCopiedReferenceCells(
+                contentRefMaps ? remapTableRowResourceUrls(row.data, contentRefMaps) : row.data,
+                table.referenceColumnTargetTableIds
+              ),
             },
             provenance: classification.mode === 'tracked' ? classification : undefined,
           }
@@ -1295,6 +1438,29 @@ export async function copyForkResourceContent(params: {
     }
   }
 
+  const failedTableIds = new Set(
+    failures.flatMap((failure) => (failure.kind === 'table' ? [failure.childId] : []))
+  )
+  let foundFailedDependent = true
+  while (foundFailedDependent) {
+    foundFailedDependent = false
+    for (const table of contentPlan.tables) {
+      if (failedTableIds.has(table.childId)) continue
+      if (!table.dependsOnChildIds?.some((dependencyId) => failedTableIds.has(dependencyId))) {
+        continue
+      }
+      failedTableIds.add(table.childId)
+      failures.push({ kind: 'table', childId: table.childId })
+      copiedResources -= 1
+      failedResources += 1
+      foundFailedDependent = true
+      logger.warn(`[${requestId}] Failed copied table because a referenced table copy failed`, {
+        sourceTableId: table.sourceId,
+        childTableId: table.childId,
+      })
+    }
+  }
+
   for (const kb of contentPlan.knowledgeBases) {
     try {
       await logSkippedConnectorDocuments(kb)

apps/sim/ee/workspace-forking/lib/copy/copy-resources.ts

@@ -272,6 +272,9 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
   })
 
   it('threads push orientation through the shared container and mapping boundaries', async () => {
+    const resolver = vi.fn((kind: ForkRemapKind, sourceId: string) =>
+      kind === 'table' && sourceId === 'mapped-table' ? 'target-table' : null
+    )
     await copyPromoteUnmappedResources({
       tx,
       edge,
@@ -290,7 +293,7 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
       },
       workflowIdMap: new Map(),
       folderIdMap: new Map(),
-      resolver: () => null,
+      resolver,
       resolveBlockId,
       referencedDocumentIds: [],
     })
@@ -303,6 +306,9 @@ describe('copyPromoteUnmappedResources - files + folder content-refs', () => {
         },
       })
     )
+    const containerParams = mockCopyForkResourceContainers.mock.calls.at(-1)?.[0]
+    expect(containerParams?.resolveMappedTableReference('mapped-table')).toBe('target-table')
+    expect(resolver).toHaveBeenCalledWith('table', 'mapped-table')
     expect(mockPersistCopiedResourceMappings).toHaveBeenCalledWith(
       expect.objectContaining({
         edgeChildWorkspaceId: 'edge-child',

apps/sim/ee/workspace-forking/lib/promote/copy-unmapped.test.ts

@@ -232,6 +232,7 @@ export async function copyPromoteUnmappedResources(params: {
     // A sync can rename env vars, so a copied custom tool's `code` must have its `{{ENV}}` refs
     // rewritten through the same plan resolver that remaps subblock-value env refs.
     resolveEnvName: (key) => resolver('env-var', key),
+    resolveMappedTableReference: (sourceTableId) => resolver('table', sourceTableId),
     resolveBlockId,
     documentMappingContext: {
       edgeChildWorkspaceId: edge.childWorkspaceId,

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

@@ -1,3 +1,4 @@
+import { remapColumnReferencedTableIds } from '@/lib/table/column-types/registry.server'
 import type { TableSchema } from '@/lib/table/types'
 import {
   deriveForkBlockId,
@@ -60,3 +61,14 @@ export function remapForkTableWorkflowGroups(
 
   return { ...schema, columns, workflowGroups: remappedGroups }
 }
+
+/** Rewrites copied reference columns to the copied target table identities. */
+export function remapForkTableReferences(
+  schema: TableSchema,
+  tableIdMap: ReadonlyMap<string, string>
+): TableSchema {
+  const columns = remapColumnReferencedTableIds(schema.columns, tableIdMap)
+  return columns.some((column, index) => column !== schema.columns[index])
+    ? { ...schema, columns }
+    : schema
+}