Skip to content

Commit cdae802

Browse files
author
Theodore Li
committed
Handle name getting taken before restore
1 parent 62d5e04 commit cdae802

6 files changed

Lines changed: 150 additions & 77 deletions

File tree

‎apps/sim/app/api/knowledge/[id]/restore/route.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { eq } from 'drizzle-orm'
55
import { type NextRequest, NextResponse } from 'next/server'
66
import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
77
import { generateRequestId } from '@/lib/core/utils/request'
8-
import { restoreKnowledgeBase } from '@/lib/knowledge/service'
8+
import { KnowledgeBaseConflictError, restoreKnowledgeBase } from '@/lib/knowledge/service'
99
import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
1010

1111
const logger = createLogger('RestoreKnowledgeBaseAPI')
@@ -49,6 +49,10 @@ export async function POST(request: NextRequest, { params }: { params: Promise<{
4949

5050
return NextResponse.json({ success: true })
5151
} catch (error) {
52+
if (error instanceof KnowledgeBaseConflictError) {
53+
return NextResponse.json({ error: error.message }, { status: 409 })
54+
}
55+
5256
logger.error(`[${requestId}] Error restoring knowledge base ${id}`, error)
5357
return NextResponse.json(
5458
{ error: error instanceof Error ? error.message : 'Internal server error' },

‎apps/sim/app/api/table/[tableId]/restore/route.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { createLogger } from '@sim/logger'
22
import { type NextRequest, NextResponse } from 'next/server'
33
import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
44
import { generateRequestId } from '@/lib/core/utils/request'
5-
import { getTableById, restoreTable } from '@/lib/table'
5+
import { getTableById, restoreTable, TableConflictError } from '@/lib/table'
66
import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
77

88
const logger = createLogger('RestoreTableAPI')
@@ -36,6 +36,10 @@ export async function POST(
3636

3737
return NextResponse.json({ success: true })
3838
} catch (error) {
39+
if (error instanceof TableConflictError) {
40+
return NextResponse.json({ error: error.message }, { status: 409 })
41+
}
42+
3943
logger.error(`[${requestId}] Error restoring table ${tableId}`, error)
4044
return NextResponse.json(
4145
{ error: error instanceof Error ? error.message : 'Internal server error' },

‎apps/sim/lib/core/utils/pg-error.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,13 @@ describe('getPostgresErrorCode', () => {
1616
expect(getPostgresErrorCode(err)).toBe('23505')
1717
})
1818

19+
it('walks nested Error causes', () => {
20+
const pgErr = new Error('unique_violation') as Error & { code: string }
21+
pgErr.code = '23505'
22+
const err = new Error('outer', { cause: new Error('inner', { cause: pgErr }) })
23+
expect(getPostgresErrorCode(err)).toBe('23505')
24+
})
25+
1926
it('returns undefined for non-errors', () => {
2027
expect(getPostgresErrorCode(undefined)).toBeUndefined()
2128
expect(getPostgresErrorCode('23505')).toBeUndefined()

‎apps/sim/lib/core/utils/pg-error.ts‎

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,30 @@
11
/**
22
* Returns PostgreSQL error code (e.g. `23505` for unique_violation) when present on a thrown value.
3-
* Normalizes common Drizzle / `postgres` driver shapes.
3+
* Normalizes common Drizzle / `postgres` driver shapes and walks `cause` chains.
44
*/
55
export function getPostgresErrorCode(error: unknown): string | undefined {
6-
if (!(error instanceof Error)) {
7-
return undefined
8-
}
6+
const seen = new Set<unknown>()
7+
let current: unknown = error
98

10-
const err = error as Error & { code?: string }
11-
if (typeof err.code === 'string') {
12-
return err.code
13-
}
9+
while (current !== undefined && current !== null) {
10+
if (seen.has(current)) {
11+
break
12+
}
13+
seen.add(current)
14+
15+
if (typeof current === 'object') {
16+
const code = (current as { code?: unknown }).code
17+
if (typeof code === 'string') {
18+
return code
19+
}
20+
}
21+
22+
if (current instanceof Error && current.cause !== undefined) {
23+
current = current.cause
24+
continue
25+
}
1426

15-
const cause = err.cause
16-
if (cause && typeof cause === 'object' && 'code' in cause && typeof cause.code === 'string') {
17-
return cause.code
27+
break
1828
}
1929

2030
return undefined

‎apps/sim/lib/knowledge/service.ts‎

Lines changed: 68 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -473,54 +473,77 @@ export async function restoreKnowledgeBase(
473473
}
474474
}
475475

476-
const newName = await generateRestoreName(kb.name, async (candidate) => {
477-
if (!kb.workspaceId) return false
478-
const [match] = await db
479-
.select({ id: knowledgeBase.id })
480-
.from(knowledgeBase)
481-
.where(
482-
and(
483-
eq(knowledgeBase.workspaceId, kb.workspaceId),
484-
eq(knowledgeBase.name, candidate),
485-
isNull(knowledgeBase.deletedAt)
486-
)
487-
)
488-
.limit(1)
489-
return !!match
490-
})
476+
/**
477+
* A concurrent create/rename can commit the same active name after `generateRestoreName`'s check
478+
* (MVCC) and before this transaction commits. Retries pick a new random suffix; 23505 is still
479+
* mapped to {@link KnowledgeBaseConflictError} if exhaustion occurs.
480+
*/
481+
const maxUniqueViolationRetries = 8
482+
let attemptedRestoreName = ''
483+
484+
for (let attempt = 0; attempt < maxUniqueViolationRetries; attempt++) {
485+
attemptedRestoreName = ''
486+
try {
487+
await db.transaction(async (tx) => {
488+
await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
489+
490+
attemptedRestoreName = await generateRestoreName(kb.name, async (candidate) => {
491+
if (!kb.workspaceId) return false
492+
const [match] = await tx
493+
.select({ id: knowledgeBase.id })
494+
.from(knowledgeBase)
495+
.where(
496+
and(
497+
eq(knowledgeBase.workspaceId, kb.workspaceId),
498+
eq(knowledgeBase.name, candidate),
499+
isNull(knowledgeBase.deletedAt)
500+
)
501+
)
502+
.limit(1)
503+
return !!match
504+
})
491505

492-
const now = new Date()
506+
const now = new Date()
493507

494-
await db.transaction(async (tx) => {
495-
await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
508+
await tx
509+
.update(knowledgeBase)
510+
.set({ deletedAt: null, updatedAt: now, name: attemptedRestoreName })
511+
.where(eq(knowledgeBase.id, knowledgeBaseId))
496512

497-
await tx
498-
.update(knowledgeBase)
499-
.set({ deletedAt: null, updatedAt: now, name: newName })
500-
.where(eq(knowledgeBase.id, knowledgeBaseId))
501-
502-
await tx
503-
.update(document)
504-
.set({ archivedAt: null })
505-
.where(
506-
and(
507-
eq(document.knowledgeBaseId, knowledgeBaseId),
508-
isNotNull(document.archivedAt),
509-
isNull(document.deletedAt)
510-
)
511-
)
513+
await tx
514+
.update(document)
515+
.set({ archivedAt: null })
516+
.where(
517+
and(
518+
eq(document.knowledgeBaseId, knowledgeBaseId),
519+
isNotNull(document.archivedAt),
520+
isNull(document.deletedAt)
521+
)
522+
)
512523

513-
await tx
514-
.update(knowledgeConnector)
515-
.set({ archivedAt: null, status: 'active', updatedAt: now })
516-
.where(
517-
and(
518-
eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId),
519-
isNotNull(knowledgeConnector.archivedAt),
520-
isNull(knowledgeConnector.deletedAt)
521-
)
522-
)
523-
})
524+
await tx
525+
.update(knowledgeConnector)
526+
.set({ archivedAt: null, status: 'active', updatedAt: now })
527+
.where(
528+
and(
529+
eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId),
530+
isNotNull(knowledgeConnector.archivedAt),
531+
isNull(knowledgeConnector.deletedAt)
532+
)
533+
)
534+
})
535+
break
536+
} catch (error: unknown) {
537+
if (getPostgresErrorCode(error) !== '23505') {
538+
throw error
539+
}
540+
if (attempt === maxUniqueViolationRetries - 1) {
541+
throw new KnowledgeBaseConflictError(attemptedRestoreName || kb.name)
542+
}
543+
}
544+
}
524545

525-
logger.info(`[${requestId}] Restored knowledge base: ${knowledgeBaseId} as "${newName}"`)
546+
logger.info(
547+
`[${requestId}] Restored knowledge base: ${knowledgeBaseId} as "${attemptedRestoreName}"`
548+
)
526549
}

‎apps/sim/lib/table/service.ts‎

Lines changed: 44 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -484,27 +484,52 @@ export async function restoreTable(tableId: string, requestId: string): Promise<
484484
}
485485
}
486486

487-
const newName = await generateRestoreName(table.name, async (candidate) => {
488-
const [match] = await db
489-
.select({ id: userTableDefinitions.id })
490-
.from(userTableDefinitions)
491-
.where(
492-
and(
493-
eq(userTableDefinitions.workspaceId, table.workspaceId),
494-
eq(userTableDefinitions.name, candidate),
495-
isNull(userTableDefinitions.archivedAt)
496-
)
497-
)
498-
.limit(1)
499-
return !!match
500-
})
487+
/**
488+
* A concurrent rename/create can claim the chosen name after `generateRestoreName`'s check (MVCC).
489+
* Retries pick a new random suffix; 23505 maps to {@link TableConflictError} after exhaustion.
490+
*/
491+
const maxUniqueViolationRetries = 8
492+
let attemptedRestoreName = ''
493+
494+
for (let attempt = 0; attempt < maxUniqueViolationRetries; attempt++) {
495+
attemptedRestoreName = ''
496+
try {
497+
await db.transaction(async (tx) => {
498+
await tx.execute(sql`SELECT 1 FROM user_table_definitions WHERE id = ${tableId} FOR UPDATE`)
499+
500+
attemptedRestoreName = await generateRestoreName(table.name, async (candidate) => {
501+
const [match] = await tx
502+
.select({ id: userTableDefinitions.id })
503+
.from(userTableDefinitions)
504+
.where(
505+
and(
506+
eq(userTableDefinitions.workspaceId, table.workspaceId),
507+
eq(userTableDefinitions.name, candidate),
508+
isNull(userTableDefinitions.archivedAt)
509+
)
510+
)
511+
.limit(1)
512+
return !!match
513+
})
501514

502-
await db
503-
.update(userTableDefinitions)
504-
.set({ archivedAt: null, updatedAt: new Date(), name: newName })
505-
.where(eq(userTableDefinitions.id, tableId))
515+
const now = new Date()
516+
await tx
517+
.update(userTableDefinitions)
518+
.set({ archivedAt: null, updatedAt: now, name: attemptedRestoreName })
519+
.where(eq(userTableDefinitions.id, tableId))
520+
})
521+
break
522+
} catch (error: unknown) {
523+
if (getPostgresErrorCode(error) !== '23505') {
524+
throw error
525+
}
526+
if (attempt === maxUniqueViolationRetries - 1) {
527+
throw new TableConflictError(attemptedRestoreName || table.name)
528+
}
529+
}
530+
}
506531

507-
logger.info(`[${requestId}] Restored table ${tableId} as "${newName}"`)
532+
logger.info(`[${requestId}] Restored table ${tableId} as "${attemptedRestoreName}"`)
508533
}
509534

510535
/**

0 commit comments

Comments
 (0)