fix(files): clean staged content after rejected commits

Author: mzxchandra

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

@@ -20,6 +20,7 @@ import {
 import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
 import { sha256Hex } from '@sim/security/hash'
 import { createDeferred } from '@sim/testing/helpers/deferred'
+import { getPostgresErrorCode } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { and, asc, eq, inArray, sql } from 'drizzle-orm'
 import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
@@ -272,21 +273,119 @@ describe('workspace file version history in PostgreSQL', () => {
     }
   )
 
+  it.each(['upload', 'content'] as const)(
+    'cleans staged %s bytes when PostgreSQL rejects COMMIT after the callback completes',
+    async (operation) => {
+      const fixture = await seedFile('original')
+      const [before] = await db
+        .select()
+        .from(workspace)
+        .where(eq(workspace.id, fixture.workspaceId))
+      const triggerName = sql.identifier(`reject_commit_${generateId().replaceAll('-', '')}`)
+      await db.execute(sql`CREATE FUNCTION ${triggerName}() RETURNS trigger LANGUAGE plpgsql AS $$
+        BEGIN
+          IF NEW.workspace_id = TG_ARGV[0] THEN
+            RAISE EXCEPTION 'Deferred file constraint rejected COMMIT' USING ERRCODE = '23514';
+          END IF;
+          RETURN NEW;
+        END;
+      $$`)
+      await db.execute(sql`CREATE CONSTRAINT TRIGGER ${triggerName}
+        AFTER INSERT OR UPDATE ON ${workspaceFiles} DEFERRABLE INITIALLY DEFERRED
+        FOR EACH ROW EXECUTE FUNCTION ${triggerName}(${sql.raw(`'${fixture.workspaceId}'`)})`)
+      const transaction = db.transaction.bind(db)
+      let callbackCompleted = false
+      const observeCallback = vi
+        .spyOn(db, 'transaction')
+        .mockImplementationOnce((callback, config) =>
+          transaction(async (tx) => {
+            const result = await callback(tx)
+            callbackCompleted = true
+            return result
+          }, config)
+        )
+      const upload = storageService.uploadFile
+      let stagedKey = ''
+      const capture = vi.spyOn(storageService, 'uploadFile').mockImplementation(async (args) => {
+        const result = await upload(args)
+        stagedKey = result.key
+        return result
+      })
+      try {
+        const result =
+          operation === 'content'
+            ? updateWorkspaceFileContent(
+                fixture.workspaceId,
+                fixture.fileId,
+                fixture.aliceId,
+                Buffer.from('rejected replacement content'),
+                undefined,
+                { version: { source: 'api', authorUserId: fixture.aliceId } }
+              )
+            : uploadWorkspaceFile(
+                fixture.workspaceId,
+                fixture.aliceId,
+                Buffer.from('rejected upload'),
+                'rejected.txt',
+                'text/plain',
+                { notifyWorkspaceChange: false }
+              )
+        const rejection = await result.catch((error: unknown) => error)
+        expect(getPostgresErrorCode(rejection)).toBe('23514')
+        expect(callbackCompleted).toBe(true)
+      } finally {
+        capture.mockRestore()
+        observeCallback.mockRestore()
+        await db.execute(sql`DROP TRIGGER ${triggerName} ON ${workspaceFiles}`)
+        await db.execute(sql`DROP FUNCTION ${triggerName}()`)
+      }
+      expect(stagedKey).not.toBe('')
+      expect(
+        await db
+          .select({ id: workspaceFiles.id })
+          .from(workspaceFiles)
+          .where(eq(workspaceFiles.key, stagedKey))
+      ).toEqual([])
+      const [after] = await db.select().from(workspace).where(eq(workspace.id, fixture.workspaceId))
+      expect(after.storageUsedBytes).toBe(before.storageUsedBytes)
+      expect(await versionRows(fixture.fileId)).toEqual([])
+      const retained = await getWorkspaceFile(fixture.workspaceId, fixture.fileId)
+      if (!retained) throw new Error('original file missing')
+      expect((await fetchWorkspaceFileBuffer(retained, { maxBytes: 1024 })).toString()).toBe(
+        'original'
+      )
+      expect(await objectExists(stagedKey)).toBe(false)
+      expect(
+        await db
+          .select({ id: outboxEvent.id })
+          .from(outboxEvent)
+          .where(
+            and(
+              eq(outboxEvent.eventType, WORKSPACE_FILE_STORAGE_CLEANUP_OUTBOX_EVENT),
+              sql`${outboxEvent.payload}->>'key' = ${stagedKey}`
+            )
+          )
+      ).toHaveLength(1)
+    }
+  )
+
   it.each([
-    { operation: 'upload', enqueueAvailable: true },
-    { operation: 'content', enqueueAvailable: true },
-    { operation: 'upload', enqueueAvailable: false },
-    { operation: 'content', enqueueAvailable: false },
+    { operation: 'upload', enqueueAvailable: true, code: undefined },
+    { operation: 'content', enqueueAvailable: true, code: undefined },
+    { operation: 'upload', enqueueAvailable: false, code: undefined },
+    { operation: 'content', enqueueAvailable: false, code: undefined },
+    { operation: 'content', enqueueAvailable: true, code: '40003' },
+    { operation: 'upload', enqueueAvailable: false, code: 'CONNECTION_CLOSED' },
   ] as const)(
-    'retains committed $operation bytes after acknowledgement loss (cleanup database available=$enqueueAvailable)',
-    async ({ operation, enqueueAvailable }) => {
+    'retains committed $operation bytes after acknowledgement loss (cleanup database available=$enqueueAvailable, code=$code)',
+    async ({ operation, enqueueAvailable, code }) => {
       const fixture = await seedFile('original')
       const transaction = db.transaction.bind(db)
       const lostAcknowledgement = vi
         .spyOn(db, 'transaction')
         .mockImplementationOnce(async (callback, config) => {
           await transaction(callback, config)
-          throw new Error('Commit acknowledgement lost')
+          throw Object.assign(new Error('Commit acknowledgement lost'), { code })
         })
       const enqueue = storageCleanup.enqueueWorkspaceFileStorageCleanups
       const enqueueFailure = vi

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

@@ -19,6 +19,7 @@ import {
   getErrorMessage,
   getPostgresConstraintName,
   getPostgresErrorCode,
+  isPostgresCommitRejection,
 } from '@sim/utils/errors'
 import { generateShortId } from '@sim/utils/id'
 import { omit } from '@sim/utils/object'
@@ -561,7 +562,7 @@ async function discardStagedFileContent(staged: StagedFileContent): Promise<void
   })
 }
 
-/** Keeps possibly committed bytes when the transaction fails after its callback has completed. */
+/** Discards known rollbacks while retaining bytes whose COMMIT outcome is uncertain. */
 async function finalizeStagedFileContent<T>(
   staged: StagedFileContent,
   prepare: (tx: DbTransaction) => Promise<T>
@@ -574,7 +575,7 @@ async function finalizeStagedFileContent<T>(
       return result
     })
   } catch (error) {
-    if (preparedForCommit) {
+    if (preparedForCommit && !isPostgresCommitRejection(error)) {
       logger.error(
         'File commit outcome is uncertain; retaining staged content for reconciliation',
         {

packages/utils/src/errors.test.ts

@@ -4,6 +4,7 @@ import {
   getPostgresCancellationReason,
   getPostgresErrorCode,
   getTransientDatabaseFailure,
+  isPostgresCommitRejection,
 } from '@sim/utils/errors'
 import { describe, expect, it } from 'vitest'
 
@@ -236,3 +237,28 @@ describe('describeError', () => {
     expect(described?.causeChain?.length).toBeLessThanOrEqual(10)
   })
 })
+
+describe('isPostgresCommitRejection', () => {
+  it('recognizes a deferred constraint rejection through a transaction wrapper', () => {
+    const rejection = Object.assign(new Error('deferred constraint failed'), { code: '23514' })
+    expect(isPostgresCommitRejection(new Error('transaction failed', { cause: rejection }))).toBe(
+      true
+    )
+  })
+
+  it.each(['40003', '08007', 'CONNECTION_CLOSED', 'ECONNRESET', undefined])(
+    'does not treat uncertain outcome %s as permission to discard staged data',
+    (code) => {
+      const error = Object.assign(new Error('commit outcome unavailable'), { code })
+      expect(isPostgresCommitRejection(error)).toBe(false)
+    }
+  )
+
+  it('does not replace an outer connection failure with a nested rejection code', () => {
+    const rejection = Object.assign(new Error('earlier constraint failure'), { code: '23514' })
+    const connection = Object.assign(new Error('connection lost', { cause: rejection }), {
+      code: 'CONNECTION_CLOSED',
+    })
+    expect(isPostgresCommitRejection(connection)).toBe(false)
+  })
+})

packages/utils/src/errors.ts

@@ -30,6 +30,30 @@ export function getPostgresErrorCode(error: unknown): string | undefined {
   return readPgErrorField(error, 'code')
 }
 
+const POSTGRES_COMMIT_REJECTION_CODES = new Set([
+  '23000',
+  '23001',
+  '23502',
+  '23503',
+  '23505',
+  '23514',
+  '23P01',
+  '40000',
+  '40001',
+  '40002',
+  '40P01',
+])
+
+/**
+ * Recognizes explicit PostgreSQL constraint/rollback rejections at COMMIT. Use only on the
+ * transaction's rejection after its callback completed, never on errors from post-commit work.
+ * Unknown outcomes, including 40003 and connection failures, are not safe cleanup authority.
+ */
+export function isPostgresCommitRejection(error: unknown): boolean {
+  const code = getPostgresErrorCode(error)
+  return code !== undefined && POSTGRES_COMMIT_REJECTION_CODES.has(code)
+}
+
 const POSTGRES_CANCELLATION_REASONS = [
   ['57014', 'canceling statement due to statement timeout', 'statement_timeout'],
   ['57014', 'canceling statement due to user request', 'user_cancel'],