fix(tables): roll back partial folder restores

Author: j15z

apps/sim/lib/folders/cascade.test.ts

@@ -26,10 +26,12 @@ import { folderMutationStatus } from '@/lib/folders/status'
 
 const tableServiceMocks = vi.hoisted(() => ({
   deleteTables: vi.fn(),
+  restoreTable: vi.fn(),
 }))
 
 vi.mock('@/lib/table/service', () => ({
   deleteTables: tableServiceMocks.deleteTables,
+  restoreTable: tableServiceMocks.restoreTable,
 }))
 
 interface SelectCall {
@@ -614,6 +616,42 @@ describe('knowledge_base and table folder resources', () => {
       })
     ).rejects.toMatchObject({ code: 'locked', message: 'Table deletion is locked' })
   })
+
+  it('re-archives a successful restore prefix when a later cohort table fails', async () => {
+    resetDbChainMock()
+    queueTableRows(schemaMock.userTableDefinitions, [
+      { id: 'tbl_accounts' },
+      { id: 'tbl_contacts' },
+    ])
+    const restoreError = new Error('restore failed')
+    tableServiceMocks.restoreTable
+      .mockResolvedValueOnce(undefined)
+      .mockRejectedValueOnce(restoreError)
+    tableServiceMocks.deleteTables.mockResolvedValueOnce({
+      archived: [{ id: 'tbl_accounts', name: 'Accounts', workspaceId: 'ws-1' }],
+      failed: [],
+      notFound: [],
+    })
+
+    await expect(
+      tableConfig.restoreChildren?.({
+        workspaceId: 'ws-1',
+        folderIds: ['folder-1'],
+        timestamp: TIMESTAMP,
+      })
+    ).rejects.toBe(restoreError)
+
+    expect(tableServiceMocks.deleteTables).toHaveBeenCalledWith(
+      ['tbl_accounts'],
+      'folder-restore-rollback-folder-1',
+      {
+        expectedWorkspaceId: 'ws-1',
+        archivedAt: TIMESTAMP,
+        skipNotify: true,
+        archiveAsCohort: true,
+      }
+    )
+  })
 })
 
 describe('table folder deletion guard', () => {

apps/sim/lib/folders/config.ts

@@ -336,21 +336,46 @@ async function archiveTableChildren(context: CascadeChildrenContext): Promise<nu
 /**
  * Restores the tables this cascade archived, through the canonical table restore so a table
  * whose name was taken while it was gone is renamed instead of tripping the active-name
- * unique index.
+ * unique index. If a later member fails, the successful prefix is re-archived under the
+ * original timestamp so no active referrer can point at a still-archived cohort target.
  */
 async function restoreTableChildren(context: CascadeChildrenContext): Promise<number> {
-  const { restoreTable } = await import('@/lib/table/service')
+  const { deleteTables, restoreTable } = await import('@/lib/table/service')
   const ids = await selectChildIds(FOLDER_RESOURCES.table, context, 'archived')
   const restoringFolderIds = new Set(context.folderIds)
   const restoringTableIds = new Set(ids)
-
-  for (const id of ids) {
-    // restoreFolder fires one folder-level live-list notify for the whole subtree.
-    await restoreTable(id, `folder-cascade-${context.folderIds[0]}`, {
-      restoringFolderIds,
-      restoringTableIds,
-      skipNotify: true,
-    })
+  const restoredIds: string[] = []
+
+  try {
+    for (const id of ids) {
+      // restoreFolder fires one folder-level live-list notify for the whole subtree.
+      await restoreTable(id, `folder-cascade-${context.folderIds[0]}`, {
+        restoringFolderIds,
+        restoringTableIds,
+        skipNotify: true,
+      })
+      restoredIds.push(id)
+    }
+  } catch (error) {
+    const rollback = await deleteTables(
+      restoredIds,
+      `folder-restore-rollback-${context.folderIds[0]}`,
+      {
+        expectedWorkspaceId: context.workspaceId,
+        archivedAt: context.timestamp,
+        skipNotify: true,
+        archiveAsCohort: true,
+      }
+    )
+    if (rollback.terminalError) throw rollback.terminalError
+    if (
+      rollback.failed.length > 0 ||
+      rollback.notFound.length > 0 ||
+      rollback.archived.length !== restoredIds.length
+    ) {
+      throw new OrchestrationError('internal', 'Failed to roll back the table restore cohort')
+    }
+    throw error
   }
 
   return ids.length