Skip to content

Commit e0b359e

Browse files
committed
fix(files): retry restore names claimed by legacy writers
1 parent 4e0223e commit e0b359e

2 files changed

Lines changed: 129 additions & 3 deletions

File tree

‎apps/sim/lib/uploads/contexts/workspace/__integration__/file-names.integration.ts‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import path from 'node:path'
99
import { db, dbFor } from '@sim/db'
1010
import { copilotChats, organization, user, workspace, workspaceFiles } from '@sim/db/schema'
1111
import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
12+
import { createDeferred } from '@sim/testing/helpers/deferred'
1213
import { generateId } from '@sim/utils/id'
1314
import { and, eq, inArray, isNull, sql } from 'drizzle-orm'
1415
import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
@@ -33,6 +34,8 @@ import {
3334
workspaceFileNameFolderCondition,
3435
} from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager'
3536
import {
37+
deleteWorkspaceFile,
38+
fetchWorkspaceFileBuffer,
3639
generateWorkspaceFileKey,
3740
getWorkspaceFileByName,
3841
resolveWorkspaceFileReference,
@@ -41,6 +44,7 @@ import {
4144
workspaceFileVfsPath,
4245
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
4346
import { createWorkspaceFileDelegatedPrincipal } from '@/lib/workspace-files/application/delegated-principal'
47+
import { restoreWorkspaceFileOperation } from '@/lib/workspace-files/application/restore-workspace-file'
4448

4549
describe('workspace file names in PostgreSQL', () => {
4650
const fixtures: ReturnType<typeof createKnowledgeAclFixtureIds>[] = []
@@ -197,6 +201,105 @@ describe('workspace file names in PostgreSQL', () => {
197201
}
198202
})
199203

204+
it.each([false, true])(
205+
'restores with a new name when a deployed writer claims the candidate (nested=%s)',
206+
async (nested) => {
207+
const fixture = await seedWorkspace()
208+
const folderId = nested
209+
? (
210+
await createWorkspaceFileFolder({
211+
workspaceId: fixture.workspaceId,
212+
userId: fixture.aliceId,
213+
name: 'Restore target',
214+
})
215+
).id
216+
: null
217+
const archived = await upload(fixture.workspaceId, fixture.aliceId, 'restore.txt', folderId)
218+
await deleteWorkspaceFile(fixture.workspaceId, archived.id)
219+
const competitor = await upload(fixture.workspaceId, fixture.aliceId, 'other.txt', folderId)
220+
const triggerName = sql.identifier(`restore_race_${generateId().replaceAll('-', '')}`)
221+
await db.execute(sql`
222+
CREATE FUNCTION ${triggerName}() RETURNS trigger LANGUAGE plpgsql AS $$
223+
BEGIN
224+
IF NEW.id = TG_ARGV[0] THEN
225+
PERFORM pg_advisory_xact_lock(hashtext('restore-race:' || NEW.id));
226+
END IF;
227+
RETURN NEW;
228+
END $$
229+
`)
230+
await db.execute(sql`
231+
CREATE TRIGGER ${triggerName} BEFORE UPDATE OF deleted_at ON ${workspaceFiles}
232+
FOR EACH ROW WHEN (OLD.deleted_at IS NOT NULL AND NEW.deleted_at IS NULL)
233+
EXECUTE FUNCTION ${triggerName}(${sql.raw(`'${archived.id}'`)})
234+
`)
235+
const ready = createDeferred<number>()
236+
const release = createDeferred<void>()
237+
const blocker = db.transaction(async (tx) => {
238+
await tx.execute(
239+
sql`SELECT pg_advisory_xact_lock(hashtext(${`restore-race:${archived.id}`}))`
240+
)
241+
const [connection] = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`)
242+
ready.resolve(connection.pid)
243+
await release.promise
244+
})
245+
const blockerPid = await ready.promise
246+
const restoration = Promise.allSettled([
247+
restoreWorkspaceFileOperation.execute({
248+
principal: { kind: 'session', userId: fixture.aliceId, sessionId: generateId() },
249+
input: { fileId: archived.id, assertedWorkspaceId: fixture.workspaceId },
250+
}),
251+
])
252+
try {
253+
await expect
254+
.poll(
255+
async () =>
256+
(
257+
await db.execute(sql`
258+
SELECT pid FROM pg_stat_activity WHERE ${blockerPid} = ANY(pg_blocking_pids(pid))
259+
`)
260+
).length,
261+
{ timeout: 5000 }
262+
)
263+
.toBe(1)
264+
/** The deployed rename updates only the file row; it does not take the new directory mutex. */
265+
await db
266+
.update(workspaceFiles)
267+
.set({ originalName: archived.name, updatedAt: new Date() })
268+
.where(
269+
and(
270+
eq(workspaceFiles.id, competitor.id),
271+
eq(workspaceFiles.workspaceId, fixture.workspaceId),
272+
eq(workspaceFiles.context, 'workspace')
273+
)
274+
)
275+
} finally {
276+
release.resolve()
277+
await blocker
278+
await restoration
279+
await db.execute(sql`DROP TRIGGER ${triggerName} ON ${workspaceFiles}`)
280+
await db.execute(sql`DROP FUNCTION ${triggerName}()`)
281+
}
282+
const [result] = await restoration
283+
if (result.status === 'rejected') throw result.reason
284+
expect(result.value.file).toMatchObject({
285+
id: archived.id,
286+
name: 'restore_restored.txt',
287+
folderId,
288+
key: archived.key,
289+
})
290+
expect(
291+
(await fetchWorkspaceFileBuffer(result.value.file, { maxBytes: 1024 })).toString()
292+
).toBe('restore.txt')
293+
expect(
294+
(await getWorkspaceFileByName(fixture.workspaceId, 'restore.txt', { folderId }))?.id
295+
).toBe(competitor.id)
296+
expect(
297+
(await getWorkspaceFileByName(fixture.workspaceId, 'restore_restored.txt', { folderId }))
298+
?.id
299+
).toBe(archived.id)
300+
}
301+
)
302+
200303
it('falls back to a short-id suffix after 20 numbered copies, including under concurrency', async () => {
201304
const fixture = await seedWorkspace()
202305
await upload(fixture.workspaceId, fixture.aliceId, 'page.html')

‎apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2694,9 +2694,19 @@ export async function restoreWorkspaceFile(workspaceId: string, fileId: string):
26942694
if (!ws || ws.archivedAt) {
26952695
throw new OrchestrationError('validation', 'Cannot restore file into an archived workspace')
26962696
}
2697-
await db.transaction((tx) =>
2698-
restoreFileInTx(tx, { entityType: 'workspace', entityId: workspaceId }, fileId)
2699-
)
2697+
const maxUniqueViolationRetries = 8
2698+
for (let attempt = 0; attempt < maxUniqueViolationRetries; attempt++) {
2699+
try {
2700+
await db.transaction((tx) =>
2701+
restoreFileInTx(tx, { entityType: 'workspace', entityId: workspaceId }, fileId)
2702+
)
2703+
return
2704+
} catch (error) {
2705+
if (!(error instanceof FileConflictError) || attempt === maxUniqueViolationRetries - 1) {
2706+
throw error
2707+
}
2708+
}
2709+
}
27002710
}
27012711

27022712
/** Restores a retained head, re-rooting and deduplicating within its canonical owner. */
@@ -2756,6 +2766,19 @@ async function restoreFileInTx(
27562766
)
27572767
)
27582768
.returning()
2769+
.catch((error: unknown) => {
2770+
const nameConstraint =
2771+
owner.entityType === 'project'
2772+
? 'workspace_files_project_folder_name_active_unique'
2773+
: 'workspace_files_workspace_folder_name_active_unique'
2774+
if (
2775+
getPostgresErrorCode(error) === '23505' &&
2776+
getPostgresConstraintName(error) === nameConstraint
2777+
) {
2778+
throw new FileConflictError(originalName)
2779+
}
2780+
throw error
2781+
})
27592782
if (!restored) throw new OrchestrationError('conflict', 'File changed during restoration')
27602783
return restored
27612784
}

0 commit comments

Comments
 (0)