diff --git a/apps/sim/background/cleanup-file-versions.ts b/apps/sim/background/cleanup-file-versions.ts index 76f1c321b4c..593b3ac42d6 100644 --- a/apps/sim/background/cleanup-file-versions.ts +++ b/apps/sim/background/cleanup-file-versions.ts @@ -1,198 +1,27 @@ -import { dbFor } from '@sim/db' -import { workspaceFileVersion } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { chunkArray } from '@sim/utils/helpers' import { task } from '@trigger.dev/sdk' -import { and, count, gt, inArray, isNotNull, lt, min, or, sql } from 'drizzle-orm' import type { CleanupJobPayload } from '@/lib/billing/cleanup-dispatcher' -import { - DEFAULT_BATCH_SIZE, - DEFAULT_DELETE_CHUNK_SIZE, - DEFAULT_MAX_BATCHES_PER_TABLE, - DEFAULT_WORKSPACE_CHUNK_SIZE, -} from '@/lib/cleanup/batch-delete' +import { DEFAULT_BATCH_SIZE, DEFAULT_MAX_BATCHES_PER_TABLE } from '@/lib/cleanup/batch-delete' import { retentionCleanupQueue } from '@/lib/cleanup/queue' -import { cleanupProjectFileVersions } from '@/lib/projects/files/retention' -import { enqueueWorkspaceFileStorageCleanups } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' -import { MAX_SUPERSEDED_FILE_VERSIONS } from '@/lib/uploads/contexts/workspace/workspace-file-versions' +import { cleanupFileVersions, fileRetentionOwners } from '@/lib/file-retention' const logger = createLogger('CleanupFileVersions') - -/** All cleanup queries run on the dedicated cleanup pool. */ -const cleanupDb = dbFor('cleanup') - -/** Candidate files whose histories are ranked in one query. */ -const FILES_PER_QUERY = 500 - -/** - * Bounds one run like the other cleanup jobs: {@link DEFAULT_MAX_BATCHES_PER_TABLE} batches per - * workspace chunk and this many versions overall. The next run resumes where this one stopped. - */ const MAX_VERSIONS_PER_RUN = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE -/** - * Superseded versions a free file keeps (its newest 100 with the current one); versions beyond it - * are pruned whatever their age. Paid plans are bounded only by the inline write-time ceiling. - */ -const FREE_MAX_SUPERSEDED_VERSIONS = 99 - -/** - * Newest superseded versions retention never prunes, whatever their age, so a file always keeps its - * newest ten versions with the current one. - */ -const KEEP_SUPERSEDED = 9 - -/** - * Files in the group that can lose any version: more superseded versions than the keep-latest floor, - * and either one older than the cutoff or more than the plan allows. Ranking only these keeps the - * sort to the histories that need pruning instead of every superseded row in the group. - */ -async function selectCandidateFileIds( - workspaceIds: string[], - cutoff: Date, - maxSuperseded: number -): Promise { - const rows = await cleanupDb - .select({ fileId: workspaceFileVersion.fileId }) - .from(workspaceFileVersion) - .where( - and( - inArray(workspaceFileVersion.workspaceId, workspaceIds), - isNotNull(workspaceFileVersion.supersededAt) - ) - ) - .groupBy(workspaceFileVersion.fileId) - .having( - and( - gt(count(), KEEP_SUPERSEDED), - or( - lt( - min(workspaceFileVersion.supersededAt), - sql.param(cutoff, workspaceFileVersion.supersededAt) - ), - gt(count(), maxSuperseded) - ) - ) - ) - return rows.map((row) => row.fileId) -} - -/** - * Superseded versions of the given files past retention: older than the cutoff or beyond the plan's - * count, but never among the newest {@link KEEP_SUPERSEDED} superseded versions of a file. - */ -function selectExpiredVersions( - fileIds: string[], - cutoff: Date, - maxSuperseded: number, - batchSize: number -) { - const ranked = cleanupDb - .select({ - id: workspaceFileVersion.id, - supersededAt: workspaceFileVersion.supersededAt, - rank: sql`row_number() over (partition by ${workspaceFileVersion.fileId} order by ${workspaceFileVersion.version} desc)`.as( - 'rank' - ), - }) - .from(workspaceFileVersion) - .where( - and( - inArray(workspaceFileVersion.fileId, fileIds), - isNotNull(workspaceFileVersion.supersededAt) - ) - ) - .as('ranked') - - return cleanupDb - .select({ id: ranked.id }) - .from(ranked) - .where( - and( - gt(ranked.rank, KEEP_SUPERSEDED), - or(lt(ranked.supersededAt, cutoff), gt(ranked.rank, maxSuperseded)) - ) - ) - .limit(batchSize) -} - -/** - * Deletes the expired version rows and enqueues their stored objects on the storage-cleanup outbox - * in the same transaction, so a row never outlives its release and every released object is - * deleted durably — the outbox retries failures and treats an already-missing object as done. - */ -function deleteVersions(rows: Array<{ id: string }>): Promise { - return cleanupDb.transaction(async (tx) => { - const removed = await tx - .delete(workspaceFileVersion) - .where( - and( - inArray( - workspaceFileVersion.id, - rows.map((row) => row.id) - ), - isNotNull(workspaceFileVersion.supersededAt) - ) - ) - .returning({ key: workspaceFileVersion.key }) - await enqueueWorkspaceFileStorageCleanups( - tx, - removed.map((row) => row.key) - ) - return removed.length - }) -} - export async function runCleanupFileVersions(payload: CleanupJobPayload): Promise { const startTime = Date.now() - const { workspaceIds, retentionHours, label, plan } = payload - let projectDeleted = 0 - for (const projectId of [...new Set(payload.projectIds ?? [])].sort()) { - if (projectDeleted >= MAX_VERSIONS_PER_RUN) break - projectDeleted += await cleanupProjectFileVersions( - projectId, - MAX_VERSIONS_PER_RUN - projectDeleted - ) - } - if (workspaceIds.length === 0) { - logger.info(`[${label}] No workspaces to process`) - return - } - - const cutoff = new Date(Date.now() - retentionHours * 60 * 60 * 1000) - const maxSuperseded = - plan === 'free' ? FREE_MAX_SUPERSEDED_VERSIONS : MAX_SUPERSEDED_FILE_VERSIONS - logger.info( - `[${label}] Processing ${workspaceIds.length} workspaces, cutoff: ${cutoff.toISOString()}` + const owners = fileRetentionOwners(payload) + const deleted = await cleanupFileVersions( + owners, + { + plan: payload.plan, + cutoff: new Date(Date.now() - payload.retentionHours * 60 * 60 * 1000), + label: payload.label, + }, + MAX_VERSIONS_PER_RUN ) - - let deleted = projectDeleted - let attempted = projectDeleted - for (const group of chunkArray(workspaceIds, DEFAULT_WORKSPACE_CHUNK_SIZE)) { - if (attempted >= MAX_VERSIONS_PER_RUN) break - const candidates = await selectCandidateFileIds(group, cutoff, maxSuperseded) - let batches = 0 - for (const fileIds of chunkArray(candidates, FILES_PER_QUERY)) { - let exhausted = false - while ( - !exhausted && - batches < DEFAULT_MAX_BATCHES_PER_TABLE && - attempted < MAX_VERSIONS_PER_RUN - ) { - batches++ - const batchSize = Math.min(DEFAULT_DELETE_CHUNK_SIZE, MAX_VERSIONS_PER_RUN - attempted) - const expired = await selectExpiredVersions(fileIds, cutoff, maxSuperseded, batchSize) - attempted += expired.length - const removed = expired.length > 0 ? await deleteVersions(expired) : 0 - deleted += removed - exhausted = expired.length < batchSize || removed === 0 - } - if (!exhausted) break - } - } - const elapsed = ((Date.now() - startTime) / 1000).toFixed(2) - logger.info(`[${label}] File version cleanup: ${deleted} released in ${elapsed}s`) + logger.info(`[${payload.label}] File version cleanup: ${deleted} released in ${elapsed}s`) } export const cleanupFileVersionsTask = task({ diff --git a/apps/sim/background/cleanup-soft-deletes.test.ts b/apps/sim/background/cleanup-soft-deletes.test.ts index ef07a1a9cb7..22169f396ea 100644 --- a/apps/sim/background/cleanup-soft-deletes.test.ts +++ b/apps/sim/background/cleanup-soft-deletes.test.ts @@ -1,5 +1,4 @@ import { - dbChainMock, dbChainMockFns, hasMockCondition, queueTableRows, @@ -18,10 +17,6 @@ import { uploadsMetadataMockFns, } from '@sim/testing/mocks/uploads-metadata.mock' import { workflowsUtilsMock, workflowsUtilsMockFns } from '@sim/testing/mocks/workflows-utils.mock' -import { - workspaceFileManagerMock, - workspaceFileManagerMockFns, -} from '@sim/testing/mocks/workspace-file-manager.mock' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { @@ -77,19 +72,14 @@ vi.mock('@/lib/uploads', () => uploadsMock) vi.mock('@/lib/uploads/server/metadata', () => uploadsMetadataMock) -const { mockReleaseWorkspaceFileVersionsForPurgeInTx } = vi.hoisted(() => ({ - mockReleaseWorkspaceFileVersionsForPurgeInTx: vi.fn(), -})) -vi.mock('@/lib/uploads/contexts/workspace/workspace-file-versions', () => ({ - releaseWorkspaceFileVersionsForPurgeInTx: mockReleaseWorkspaceFileVersionsForPurgeInTx, +vi.mock('@/lib/file-retention/folders', () => ({ + cleanupArchivedWorkspaceFileFolders: vi.fn(async () => 0), })) vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock) vi.mock('@/lib/folders/naming', () => ({ deduplicateFolderName: mockDeduplicateFolderName })) -vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => workspaceFileManagerMock) - import { runCleanupSoftDeletes } from '@/background/cleanup-soft-deletes' const { mockDecrementStorageUsageForBillingContextInTx, mockResolveStorageBillingContext } = @@ -102,10 +92,8 @@ const { mockDeleteFiles } = storageServiceMockFns const { mockDeleteFileMetadata } = uploadsMetadataMockFns const { mockIsUsingCloudStorage } = uploadsMockFns const { mockDeduplicateWorkflowName } = workflowsUtilsMockFns -const { mockAllocateUniqueWorkspaceFileName } = workspaceFileManagerMockFns mockDeleteFileMetadata.mockImplementation(async () => true) mockDeduplicateWorkflowName.mockImplementation(async (name: string) => name) -mockAllocateUniqueWorkspaceFileName.mockImplementation(async (_ws: string, name: string) => name) const basePayload = { label: 'free/1', @@ -134,67 +122,22 @@ describe('cleanup soft deletes', () => { }) }) - it('releases version history inside the transaction that purges the file rows', async () => { - mockSelectRowsByIdChunks - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([ - { - id: 'file-purged', - key: 'workspace/ws-1/file-purged', - workspaceId: 'ws-1', - context: 'workspace', - sizeBytes: 5, - }, - { - id: 'file-failed', - key: 'workspace/ws-1/file-failed', - workspaceId: 'ws-1', - context: 'workspace', - sizeBytes: 4, - }, - ]) - mockDeleteFiles.mockResolvedValueOnce({ - deleted: 1, - failed: [{ key: 'workspace/ws-1/file-failed', error: 'storage unavailable' }], - }) - dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'file-purged', sizeBytes: 5 }]) - - await runCleanupSoftDeletes(basePayload) - - expect(mockReleaseWorkspaceFileVersionsForPurgeInTx).toHaveBeenCalledOnce() - expect(mockReleaseWorkspaceFileVersionsForPurgeInTx).toHaveBeenCalledWith( - dbChainMock.db, - ['file-purged'], - expect.any(Date) - ) - expect(dbChainMockFns.transaction.mock.invocationCallOrder[0]).toBeLessThan( - mockReleaseWorkspaceFileVersionsForPurgeInTx.mock.invocationCallOrder[0] - ) - const fileDelete = dbChainMockFns.delete.mock.calls.findIndex( - ([table]) => table === schemaMock.workspaceFiles - ) - expect(mockReleaseWorkspaceFileVersionsForPurgeInTx.mock.invocationCallOrder[0]).toBeLessThan( - dbChainMockFns.delete.mock.invocationCallOrder[fileDelete] - ) - }) - - it('keeps metadata rows whose object deletion failed', async () => { + it('keeps unbilled metadata rows whose object deletion failed', async () => { mockSelectRowsByIdChunks .mockResolvedValueOnce([]) .mockResolvedValueOnce([]) .mockResolvedValueOnce([ { id: 'file-failed', - key: 'workspace/ws-1/file-failed', + key: 'knowledge-base/file-failed', workspaceId: 'ws-1', - context: 'workspace', + context: 'knowledge-base', sizeBytes: 11, }, ]) mockDeleteFiles.mockResolvedValueOnce({ deleted: 0, - failed: [{ key: 'workspace/ws-1/file-failed', error: 'storage unavailable' }], + failed: [{ key: 'knowledge-base/file-failed', error: 'storage unavailable' }], }) await runCleanupSoftDeletes(basePayload) @@ -206,42 +149,6 @@ describe('cleanup soft deletes', () => { ).toBe(false) }) - it('decrements the current workspace payer only for rows conditionally deleted', async () => { - mockSelectRowsByIdChunks - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([ - { - id: 'file-deleted', - key: 'workspace/ws-1/file-deleted', - workspaceId: 'ws-1', - context: 'workspace', - sizeBytes: 7, - }, - { - id: 'file-restored', - key: 'workspace/ws-1/file-restored', - workspaceId: 'ws-1', - context: 'workspace', - sizeBytes: 13, - }, - ]) - mockDeleteFiles.mockResolvedValueOnce({ deleted: 2, failed: [] }) - dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'file-deleted', sizeBytes: 7 }]) - - await runCleanupSoftDeletes(basePayload) - - expect(mockResolveStorageBillingContext).toHaveBeenCalledOnce() - expect(mockDecrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - dbChainMock.db, - expect.objectContaining({ workspaceId: 'ws-1' }), - 7 - ) - expect(mockDeleteFiles.mock.invocationCallOrder[0]).toBeLessThan( - dbChainMockFns.transaction.mock.invocationCallOrder[0] - ) - }) - it('hard-deletes mothership metadata without touching stored-byte counters', async () => { mockSelectRowsByIdChunks .mockResolvedValueOnce([]) @@ -423,12 +330,7 @@ describe('folder cleanup target', () => { expect(target?.additionalPredicate).toBeDefined() expect(target?.additionalPredicate?.type).toBe('inArray') expect(target?.additionalPredicate?.column).toBe(schemaMock.folder.resourceType) - expect(target?.additionalPredicate?.values).toEqual([ - 'workflow', - 'file', - 'knowledge_base', - 'table', - ]) + expect(target?.additionalPredicate?.values).toEqual(['workflow', 'knowledge_base', 'table']) }) /** @@ -465,24 +367,6 @@ describe('folder cleanup target', () => { expect(dbChainMockFns.set).toHaveBeenCalledWith({ folderId: null, name: 'Report (2)' }) }) - it('falls back to an id-suffixed name when the copy-suffix range is exhausted', async () => { - // Letting the allocator throw would abort the sweep — the exact stall this guards against. - const onBatch = await getFolderOnBatch() - queueTableRows(schemaMock.folder, [{ id: 'folder-1' }]) - queueTableRows(schemaMock.workflow, []) - queueTableRows(schemaMock.workspaceFiles, [ - { id: 'f1', originalName: 'report.pdf', workspaceId: 'ws-1' }, - ]) - mockAllocateUniqueWorkspaceFileName.mockRejectedValueOnce(new Error('conflict')) - - await expect(onBatch([{ id: 'folder-1' }])).resolves.toBeUndefined() - - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - folderId: null, - originalName: 'report.pdf (f1)', - }) - }) - it('re-roots an active SUBFOLDER, which hits the same unique index', async () => { /** * `folder.parentId` is also ON DELETE SET NULL and @@ -518,32 +402,6 @@ describe('folder cleanup target', () => { expect(dbChainMockFns.set).toHaveBeenCalledWith({ parentId: null, name: 'Reports (1)' }) }) - it('recovers when the allocator RETURNS a colliding name and the update raises', async () => { - /** - * `allocateUniqueWorkspaceFileName` fails open — `fileExistsInWorkspace` swallows query - * errors and returns false — so it can hand back a name already taken at the root. Only - * the UPDATE discovers that, and an uncaught 23505 aborts the batch: the exact stall this - * hook prevents. Guarding the name lookup alone is not enough. - */ - const onBatch = await getFolderOnBatch() - queueTableRows(schemaMock.folder, [{ id: 'folder-1' }]) - queueTableRows(schemaMock.workflow, []) - queueTableRows(schemaMock.workspaceFiles, [ - { id: 'f1', originalName: 'report.pdf', workspaceId: 'ws-1' }, - ]) - mockAllocateUniqueWorkspaceFileName.mockResolvedValueOnce('taken.pdf') - dbChainMockFns.update.mockImplementationOnce(() => ({ - set: () => ({ where: () => Promise.reject(new Error('duplicate key value (23505)')) }), - })) - - await expect(onBatch([{ id: 'folder-1' }])).resolves.toBeUndefined() - - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - folderId: null, - originalName: 'report.pdf (f1)', - }) - }) - it('leaves children alone when the folder was restored between select and onBatch', async () => { /** * The DELETE re-asserts eligibility and so correctly skips a restored folder. Without the @@ -564,7 +422,6 @@ describe('folder cleanup target', () => { await onBatch([{ id: 'folder-1' }]) expect(mockDeduplicateWorkflowName).not.toHaveBeenCalled() - expect(mockAllocateUniqueWorkspaceFileName).not.toHaveBeenCalled() expect(dbChainMockFns.update).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/background/cleanup-soft-deletes.ts b/apps/sim/background/cleanup-soft-deletes.ts index cc7bc97c020..49e82024b83 100644 --- a/apps/sim/background/cleanup-soft-deletes.ts +++ b/apps/sim/background/cleanup-soft-deletes.ts @@ -1,6 +1,5 @@ -import { db, dbFor } from '@sim/db' +import { dbFor } from '@sim/db' import { - changelogRelease, copilotChats, document, folder as folderTable, @@ -10,8 +9,6 @@ import { userTableDefinitions, workflow, workflowMcpServer, - workflowTest, - workspaceFile, workspaceFiles, } from '@sim/db/schema' import { createLogger } from '@sim/logger' @@ -19,18 +16,12 @@ import { chunkArray } from '@sim/utils/helpers' import { task } from '@trigger.dev/sdk' import { and, asc, eq, inArray, isNotNull, isNull, lt, sql } from 'drizzle-orm' import { type CleanupJobPayload, runCleanupWithLimits } from '@/lib/billing/cleanup-dispatcher' -import { - decrementStorageUsageForBillingContextInTx, - lockWorkspaceStorageForMutationInTx, - resolveStorageBillingContext, -} from '@/lib/billing/storage' import { batchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete, chunkedBatchDeleteByScope, consumeRowBudget, DEFAULT_DELETE_CHUNK_SIZE, - DEFAULT_MAX_BATCHES_PER_TABLE, type RowBudget, selectRowsByIdChunks, } from '@/lib/cleanup/batch-delete' @@ -42,30 +33,18 @@ import { cleanupOwnerCondition, resolveCleanupOwnerScope, } from '@/lib/cleanup/resource-scope' +import { beginFileArchiveCleanup, fileRetentionOwners } from '@/lib/file-retention' import { deduplicateFolderName } from '@/lib/folders/naming' import { requireWorkspaceFolder } from '@/lib/folders/scope' import { settleDetachedConnectorReservations } from '@/lib/knowledge/connectors/detachment' import { hardDeleteDocuments } from '@/lib/knowledge/documents/service' -import { - cleanupArchivedProjectFileFolders, - cleanupArchivedProjectFiles, -} from '@/lib/projects/files/retention' -import { lockWorkspaceProject } from '@/lib/projects/membership' -import type { StorageContext } from '@/lib/uploads' import { isUsingCloudStorage, StorageService } from '@/lib/uploads' -import { allocateUniqueWorkspaceFileName } from '@/lib/uploads/contexts/workspace/workspace-file-manager' -import { releaseWorkspaceFileVersionsForPurgeInTx } from '@/lib/uploads/contexts/workspace/workspace-file-versions' import { deleteFileMetadata } from '@/lib/uploads/server/metadata' -import { getWorkspaceFileSize } from '@/lib/uploads/shared/types' import { deduplicateWorkflowName } from '@/lib/workflows/utils' const logger = createLogger('CleanupSoftDeletes') -/** - * Cleanup queries run on the dedicated cleanup pool. The one exception is the - * billable-file transaction below, which couples row deletion with a storage - * billing decrement — billing writes stay on the default client. - */ +/** Resource cleanup queries run on the dedicated cleanup pool. */ const cleanupDb = dbFor('cleanup') const KB_ORPHAN_BINDING_BATCH_SIZE = 500 @@ -81,324 +60,6 @@ const KB_RETENTION_BATCH_SIZE = 100 const KB_DOCUMENT_DELETE_BATCH_SIZE = 500 const KB_DOCUMENT_DELETE_MAX_BATCHES = 50 -interface WorkspaceFileScope { - /** Rows from `workspace_file` (singular, legacy workspace-context only). */ - legacyRows: Array<{ id: string; key: string; workspaceId: string }> - /** Rows from `workspace_files` (plural, multi-context). */ - multiContextRows: Array<{ - id: string - key: string - workspaceId: string | null - context: StorageContext - size: number - }> -} - -interface WorkspaceFileStorageCleanupResult { - filesDeleted: number - filesFailed: number - legacyRows: WorkspaceFileScope['legacyRows'] - multiContextRows: WorkspaceFileScope['multiContextRows'] -} - -/** - * Select every soft-deleted file row that's eligible for permanent removal. - * Returned once and reused for both S3 deletion and DB deletion so the external - * cleanup cannot drift from the row-level cleanup. - */ -async function selectExpiredWorkspaceFiles( - scope: CleanupOwnerScope, - retentionDate: Date, - budgets?: CleanupBudgets -): Promise { - const [legacyRows, multiContextRows] = await Promise.all([ - selectRowsByIdChunks( - scope.kind === 'workspace' ? scope.ids : [], - (chunkIds, chunkLimit) => - cleanupDb - .select({ - id: workspaceFile.id, - key: workspaceFile.key, - workspaceId: workspaceFile.workspaceId, - }) - .from(workspaceFile) - .where( - and( - inArray(workspaceFile.workspaceId, chunkIds), - isNotNull(workspaceFile.deletedAt), - lt(workspaceFile.deletedAt, retentionDate) - ) - ) - .limit(chunkLimit), - { budget: budgets?.legacyFiles } - ), - selectRowsByIdChunks( - scope.ids, - (chunkIds, chunkLimit) => - cleanupDb - .select({ - id: workspaceFiles.id, - key: workspaceFiles.key, - workspaceId: workspaceFiles.workspaceId, - context: workspaceFiles.context, - sizeBytes: workspaceFiles.sizeBytes, - }) - .from(workspaceFiles) - .where( - and( - cleanupOwnerCondition(workspaceFiles, scope, chunkIds), - scope.kind === 'organization' - ? eq(workspaceFiles.context, 'knowledge-base') - : undefined, - isNotNull(workspaceFiles.deletedAt), - lt(workspaceFiles.deletedAt, retentionDate) - ) - ) - .limit(chunkLimit), - { budget: budgets?.files } - ), - ]) - - return { - legacyRows, - multiContextRows: multiContextRows.map((r) => ({ - id: r.id, - key: r.key, - workspaceId: r.workspaceId, - context: r.context as StorageContext, - size: getWorkspaceFileSize(r), - })), - } -} - -async function cleanupWorkspaceFileStorage( - scope: WorkspaceFileScope -): Promise { - type Candidate = - | { source: 'legacy'; context: StorageContext; row: WorkspaceFileScope['legacyRows'][number] } - | { - source: 'multiContext' - context: StorageContext - row: WorkspaceFileScope['multiContextRows'][number] - } - - const result: WorkspaceFileStorageCleanupResult = { - filesDeleted: 0, - filesFailed: 0, - legacyRows: [], - multiContextRows: [], - } - if (!isUsingCloudStorage()) { - return { - ...result, - legacyRows: scope.legacyRows, - multiContextRows: scope.multiContextRows, - } - } - - const candidatesByContext = new Map() - const addCandidate = (candidate: Candidate) => { - const bucket = candidatesByContext.get(candidate.context) - if (bucket) bucket.push(candidate) - else candidatesByContext.set(candidate.context, [candidate]) - } - for (const row of scope.legacyRows) { - addCandidate({ source: 'legacy', context: 'workspace', row }) - } - for (const row of scope.multiContextRows) { - addCandidate({ source: 'multiContext', context: row.context, row }) - } - - for (const [context, candidates] of candidatesByContext) { - for (const batch of chunkArray(candidates, DEFAULT_DELETE_CHUNK_SIZE)) { - const deletion = await StorageService.deleteFiles( - batch.map(({ row }) => row.key), - context - ) - const failedKeys = new Set(deletion.failed.map(({ key }) => key)) - result.filesDeleted += batch.filter(({ row }) => !failedKeys.has(row.key)).length - result.filesFailed += deletion.failed.length - - for (const candidate of batch) { - if (failedKeys.has(candidate.row.key)) continue - if (candidate.source === 'legacy') result.legacyRows.push(candidate.row) - else result.multiContextRows.push(candidate.row) - } - for (const { key, error } of deletion.failed) { - logger.error(`Failed to delete storage file ${key} (context: ${context}):`, { error }) - } - } - } - - return result -} - -async function deleteExpiredLegacyWorkspaceFileRows( - rows: WorkspaceFileScope['legacyRows'], - retentionDate: Date, - label: string -): Promise<{ deleted: number; failed: number }> { - const result = { deleted: 0, failed: 0 } - for (const batch of chunkArray(rows, DEFAULT_DELETE_CHUNK_SIZE)) { - try { - const deleted = await cleanupDb - .delete(workspaceFile) - .where( - and( - inArray( - workspaceFile.id, - batch.map(({ id }) => id) - ), - isNotNull(workspaceFile.deletedAt), - lt(workspaceFile.deletedAt, retentionDate) - ) - ) - .returning({ id: workspaceFile.id }) - result.deleted += deleted.length - result.failed += batch.length - deleted.length - } catch (error) { - result.failed += batch.length - logger.error(`[${label}/workspaceFile] Exact-row delete failed`, { error }) - } - } - return result -} - -/** Contexts whose bytes count toward the workspace's billed storage: owned bodies are billed too. */ -const BILLED_FILE_CONTEXTS: readonly StorageContext[] = ['workspace', 'test', 'changelog'] - -async function deleteExpiredUnbilledWorkspaceFileRows( - rows: WorkspaceFileScope['multiContextRows'], - retentionDate: Date, - label: string -): Promise<{ deleted: number; failed: number }> { - const result = { deleted: 0, failed: 0 } - const rowsByContext = new Map() - for (const row of rows) { - if (BILLED_FILE_CONTEXTS.includes(row.context)) continue - const bucket = rowsByContext.get(row.context) - if (bucket) bucket.push(row) - else rowsByContext.set(row.context, [row]) - } - - for (const [context, contextRows] of rowsByContext) { - for (const batch of chunkArray(contextRows, DEFAULT_DELETE_CHUNK_SIZE)) { - try { - const deleted = await cleanupDb - .delete(workspaceFiles) - .where( - and( - inArray( - workspaceFiles.id, - batch.map(({ id }) => id) - ), - eq(workspaceFiles.context, context), - isNotNull(workspaceFiles.deletedAt), - lt(workspaceFiles.deletedAt, retentionDate) - ) - ) - .returning({ id: workspaceFiles.id }) - result.deleted += deleted.length - result.failed += batch.length - deleted.length - } catch (error) { - result.failed += batch.length - logger.error(`[${label}/workspaceFiles] Exact-row ${context} delete failed`, { error }) - } - } - } - return result -} - -async function deleteExpiredBillableWorkspaceFileRows( - rows: WorkspaceFileScope['multiContextRows'], - retentionDate: Date, - label: string -): Promise<{ deleted: number; failed: number }> { - const result = { deleted: 0, failed: 0 } - const rowsByWorkspace = new Map() - for (const row of rows) { - if (!BILLED_FILE_CONTEXTS.includes(row.context)) continue - if (!row.workspaceId) { - result.failed++ - logger.error(`[${label}/workspaceFiles] Billable row has no workspace attribution`, { - fileId: row.id, - }) - continue - } - const bucket = rowsByWorkspace.get(row.workspaceId) - if (bucket) bucket.push(row) - else rowsByWorkspace.set(row.workspaceId, [row]) - } - - for (const [workspaceId, workspaceRows] of rowsByWorkspace) { - for (const batch of chunkArray(workspaceRows, DEFAULT_DELETE_CHUNK_SIZE)) { - try { - const deletedCount = await db.transaction(async (tx) => { - await lockWorkspaceProject(tx, workspaceId) - await lockWorkspaceStorageForMutationInTx(tx, workspaceId) - const billingContext = await resolveStorageBillingContext(workspaceId, tx) - const fileIds = batch.map(({ id }) => id) - await tx - .delete(workflowTest) - .where( - and( - inArray(workflowTest.bodyFileId, fileIds), - isNotNull(workflowTest.deletedAt), - lt(workflowTest.deletedAt, retentionDate) - ) - ) - await tx.delete(changelogRelease).where( - inArray( - changelogRelease.bodyFileId, - tx - .select({ id: workspaceFiles.id }) - .from(workspaceFiles) - .where( - and( - inArray(workspaceFiles.id, fileIds), - isNotNull(workspaceFiles.deletedAt), - lt(workspaceFiles.deletedAt, retentionDate) - ) - ) - ) - ) - await releaseWorkspaceFileVersionsForPurgeInTx(tx, fileIds, retentionDate) - const deletedRows = await tx - .delete(workspaceFiles) - .where( - and( - inArray(workspaceFiles.id, fileIds), - eq(workspaceFiles.workspaceId, workspaceId), - inArray(workspaceFiles.context, BILLED_FILE_CONTEXTS), - isNotNull(workspaceFiles.deletedAt), - lt(workspaceFiles.deletedAt, retentionDate) - ) - ) - .returning({ - id: workspaceFiles.id, - sizeBytes: workspaceFiles.sizeBytes, - }) - const deletedBytes = deletedRows.reduce( - (total, row) => total + getWorkspaceFileSize(row), - 0 - ) - await decrementStorageUsageForBillingContextInTx(tx, billingContext, deletedBytes) - return deletedRows.length - }) - result.deleted += deletedCount - result.failed += batch.length - deletedCount - } catch (error) { - result.failed += batch.length - logger.error(`[${label}/workspaceFiles] Atomic delete and decrement failed`, { - error, - workspaceId, - }) - } - } - } - return result -} - async function hardDeleteKnowledgeBaseDocuments( knowledgeBaseIds: string[], label: string @@ -496,20 +157,7 @@ interface CleanupBatchContext { label: string } -/** - * Applies one re-root, treating the deduplicated name as a HINT rather than a guarantee. - * - * Both name allocators can hand back a name that is already taken: `fileExistsInWorkspace` - * swallows query errors and returns `false`, so `allocateUniqueWorkspaceFileName` fails OPEN, - * and `deduplicateWorkflowName`'s lookups can throw outright. Either way the UPDATE raises - * 23505, and an uncaught 23505 here aborts the batch — precisely the permanent retention stall - * this whole hook exists to prevent. So any failure retries with the row id, which is unique by - * construction and cannot collide. - * - * A failure of that retry is swallowed too: one unfixable row must not stop the other children - * from being made safe. It is logged at error level because the folder's DELETE can then still - * stall on that row via the FK's SET NULL. - */ +/** Retry a colliding name without letting one child permanently stall unrelated folder cleanup. */ async function reRootOne( preferred: () => Promise, withUniqueName: () => Promise, @@ -532,20 +180,7 @@ async function reRootOne( } } -/** - * Re-roots any still-active workflow or workspace file filed under a folder that is about to be - * hard-deleted, giving it a collision-free name first. - * - * The `folder_id` FKs are `ON DELETE SET NULL`, so Postgres already re-roots these rows on its - * own. The problem is the name: both tables carry a partial unique index keyed on - * `coalesce(folder_id, '')`, so an implicit SET NULL can land a row on a name the workspace root - * already holds and abort the whole DELETE with a 23505. `chunkedBatchDelete` counts that as a - * failed batch and stops, and the same poison row re-fails on every later run — folder retention - * would stall permanently for that workspace chunk. Renaming here leaves the SET NULL a no-op. - * - * An active child inside a soft-deleted folder is already an anomaly — the delete cascade - * archives children — so this normally selects nothing, which is why the per-row loop is fine. - */ +/** Surviving workflow children need collision-free root names before their folder FK clears. */ async function reRootActiveFolderChildren( folderIds: string[], retentionDate: Date, @@ -621,52 +256,11 @@ async function reRootActiveFolderChildrenUnguarded( ) } - const files = await cleanupDb - .select({ - id: workspaceFiles.id, - originalName: workspaceFiles.originalName, - workspaceId: workspaceFiles.workspaceId, - }) - .from(workspaceFiles) - .where( - and( - inArray(workspaceFiles.folderId, expiredIds), - isNull(workspaceFiles.deletedAt), - eq(workspaceFiles.context, 'workspace') - ) - ) - - for (const row of files) { - const workspaceId = row.workspaceId - if (!workspaceId) continue - await reRootOne( - async () => { - const originalName = await allocateUniqueWorkspaceFileName( - workspaceId, - row.originalName, - null - ) - await cleanupDb - .update(workspaceFiles) - .set({ folderId: null, originalName }) - .where(eq(workspaceFiles.id, row.id)) - }, - () => - cleanupDb - .update(workspaceFiles) - .set({ folderId: null, originalName: `${row.originalName} (${row.id})` }) - .where(eq(workspaceFiles.id, row.id)), - `workspace file ${row.id}`, - label - ) - } - /** * Subfolders are exposed to exactly the same failure. `folder.parentId` is itself * `ON DELETE SET NULL`, and `folder_workspace_resource_parent_name_active_unique` keys on * `coalesce(parent_id, '')`, so purging a parent re-roots a surviving active child into a - * namespace where its name may already be taken — the identical 23505 stall. Covering only - * workflows and files would leave the class half-closed. + * namespace where its name may already be taken — the identical 23505 stall. Subfolder names therefore need the same collision handling. */ const childFolders = await cleanupDb .select({ @@ -705,10 +299,8 @@ async function reRootActiveFolderChildrenUnguarded( ) } - if (workflows.length > 0 || files.length > 0) { - logger.warn( - `[${label}] Re-rooted ${workflows.length} workflow(s) and ${files.length} file(s) out of folders being purged` - ) + if (workflows.length > 0) { + logger.warn(`[${label}] Re-rooted ${workflows.length} workflow(s) out of folders being purged`) } } @@ -717,19 +309,7 @@ const CLEANUP_TARGETS = [ table: folderTable, softDeleteCol: folderTable.deletedAt, wsCol: folderTable.workspaceId, - /** - * `folder` is shared by all four resource types, every one of which now writes here. - * The predicate is kept (rather than dropped) so a resource type added to the enum - * before its cutover lands is not silently hard-deleted by this pass. One widened - * predicate rather than a target per type: same table, same soft-delete column, same - * workspace scoping — splitting it would only multiply the batched scans. - */ - additionalPredicate: inArray(folderTable.resourceType, [ - 'workflow', - 'file', - 'knowledge_base', - 'table', - ]), + additionalPredicate: inArray(folderTable.resourceType, ['workflow', 'knowledge_base', 'table']), onBatch: (rows: { id: string }[], ctx: CleanupBatchContext) => reRootActiveFolderChildren( rows.map(({ id }) => id), @@ -863,19 +443,8 @@ export async function runCleanupSoftDeletes( const startTime = Date.now() const { workspaceIds, retentionHours, label } = payload const scope = resolveCleanupOwnerScope(payload) - const projectFileBudget = budgets?.files ?? { - remaining: DEFAULT_DELETE_CHUNK_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, - } - const projectFolderBudget = budgets?.folders ?? { - remaining: DEFAULT_DELETE_CHUNK_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, - } - for (const projectId of [...new Set(payload.projectIds ?? [])].sort()) { - if (projectFileBudget.remaining <= 0 && projectFolderBudget.remaining <= 0) break - await cleanupArchivedProjectFiles(projectId, projectFileBudget) - await cleanupArchivedProjectFileFolders(projectId, projectFolderBudget) - } - - if (scope.ids.length === 0) { + const fileOwners = fileRetentionOwners(payload) + if (scope.ids.length === 0 && fileOwners.length === 0) { logger.info(`[${label}] No resource owners to process`) return } @@ -889,7 +458,7 @@ export async function runCleanupSoftDeletes( // external cleanup (chats + S3) AND the DB deletes below — selecting twice // could return different subsets above the LIMIT cap and orphan or // prematurely purge data. - const [doomedWorkflows, fileScope, expiredSoftDeletedChats] = await Promise.all([ + const [doomedWorkflows, fileArchiveCleanup, expiredSoftDeletedChats] = await Promise.all([ selectRowsByIdChunks( workspaceIds, (chunkIds, chunkLimit) => @@ -906,7 +475,12 @@ export async function runCleanupSoftDeletes( .limit(chunkLimit), { budget: budgets?.workflows } ), - selectExpiredWorkspaceFiles(scope, retentionDate, budgets), + beginFileArchiveCleanup(fileOwners, { + cutoff: retentionDate, + plan: payload.plan, + label, + budgets, + }), selectRowsByIdChunks( scope.ids, (chunkIds, chunkLimit) => @@ -946,8 +520,7 @@ export async function runCleanupSoftDeletes( chatCleanup = await prepareChatCleanup([...doomedChatIds], label) } - const fileCleanup = await cleanupWorkspaceFileStorage(fileScope) - if (budgets && fileCleanup.filesFailed) throw new Error('File storage cleanup failed') + const fileCleanup = await fileArchiveCleanup.cleanupStorage() let totalDeleted = 0 @@ -1000,31 +573,7 @@ export async function runCleanupSoftDeletes( } } - const legacyFileResult = await deleteExpiredLegacyWorkspaceFileRows( - fileCleanup.legacyRows, - retentionDate, - label - ) - totalDeleted += legacyFileResult.deleted - - const billableFileResult = await deleteExpiredBillableWorkspaceFileRows( - fileCleanup.multiContextRows, - retentionDate, - label - ) - totalDeleted += billableFileResult.deleted - - const unbilledFileResult = await deleteExpiredUnbilledWorkspaceFileRows( - fileCleanup.multiContextRows, - retentionDate, - label - ) - totalDeleted += unbilledFileResult.deleted - if ( - budgets && - (legacyFileResult.failed || billableFileResult.failed || unbilledFileResult.failed) - ) - throw new Error('File row cleanup failed') + totalDeleted += await fileCleanup.deleteRows() const knowledgeBaseResult = await cleanupExpiredKnowledgeBases( scope, diff --git a/apps/sim/lib/billing/storage/accounting.ts b/apps/sim/lib/billing/storage/accounting.ts new file mode 100644 index 00000000000..5e9b4846f90 --- /dev/null +++ b/apps/sim/lib/billing/storage/accounting.ts @@ -0,0 +1,59 @@ +import { project } from '@sim/db/schema' +import { eq } from 'drizzle-orm' +import { + type ProjectStorageBillingContext, + resolveProjectStorageBillingContext, + resolveStorageBillingContext, + type StorageBillingContext, +} from '@/lib/billing/storage/context' +import { prepareFileStorageMutationInTx } from '@/lib/billing/storage/tracking' +import type { DbTransaction } from '@/lib/db/types' +import { lockProject, lockWorkspaceProject } from '@/lib/projects/membership' +import type { EditableFileOwner } from '@/lib/workspace-files/ownership' + +interface PreparedFileAccounting { + billing: Billing + mutation: { applyDelta(deltaBytes: number): Promise } +} + +export function prepareFileAccountingInTx( + tx: DbTransaction, + owner: { entityType: 'workspace'; entityId: string } +): Promise> +export function prepareFileAccountingInTx( + tx: DbTransaction, + owner: { entityType: 'project'; entityId: string } +): Promise> +export function prepareFileAccountingInTx( + tx: DbTransaction, + owner: EditableFileOwner +): Promise +/** + * Prepares retained-head accounting in the caller's transaction, after authorization and before + * directory, file, history or cleanup locks. Locks lifecycle, owner row, then payer; multi-owner + * callers must acquire all lifecycle locks first. Resolve from the canonical owner, never actor + * or creator. Apply one signed delta at finalization and notify only after the caller commits. + */ +export async function prepareFileAccountingInTx( + tx: DbTransaction, + owner: EditableFileOwner +): Promise { + let billing: StorageBillingContext | ProjectStorageBillingContext + if (owner.entityType === 'workspace') { + await lockWorkspaceProject(tx, owner.entityId) + billing = await resolveStorageBillingContext(owner.entityId, tx) + } else { + await lockProject(tx, owner.entityId) + const [canonical] = await tx + .select({ ownerId: project.ownerId, organizationId: project.organizationId }) + .from(project) + .where(eq(project.id, owner.entityId)) + .limit(1) + if (!canonical) throw new Error(`Project ${owner.entityId} not found for storage accounting`) + billing = await resolveProjectStorageBillingContext( + { projectId: owner.entityId, ...canonical }, + tx + ) + } + return { billing, mutation: await prepareFileStorageMutationInTx(tx, billing) } +} diff --git a/apps/sim/lib/billing/storage/index.ts b/apps/sim/lib/billing/storage/index.ts index 9e4ad54898e..315bb0a7dd4 100644 --- a/apps/sim/lib/billing/storage/index.ts +++ b/apps/sim/lib/billing/storage/index.ts @@ -1,3 +1,4 @@ +export { prepareFileAccountingInTx } from './accounting' export { resolveStorageBillingContext, type StorageBillingContext, diff --git a/apps/sim/lib/billing/storage/payer-transfer.ts b/apps/sim/lib/billing/storage/payer-transfer.ts index 5af74809f89..f24d00935a1 100644 --- a/apps/sim/lib/billing/storage/payer-transfer.ts +++ b/apps/sim/lib/billing/storage/payer-transfer.ts @@ -49,7 +49,8 @@ export interface ChangeProjectStoragePayerParams extends ProjectStorageOwnerSnap expectedCurrentOwner: Pick } -interface ChangeProjectStoragePayerResult { +/** Payer transfer contribution: current file heads, including archived files, excluding historical versions. */ +export interface ChangeProjectStoragePayerResult { projectId: string billableBytes: number oldPayer: BillingEntity diff --git a/apps/sim/lib/billing/storage/project-accounting.integration.ts b/apps/sim/lib/billing/storage/project-accounting.integration.ts index 7da500c4b39..175c79bc419 100644 --- a/apps/sim/lib/billing/storage/project-accounting.integration.ts +++ b/apps/sim/lib/billing/storage/project-accounting.integration.ts @@ -11,6 +11,7 @@ import { sql as query } from 'drizzle-orm' import { drizzle } from 'drizzle-orm/postgres-js' import postgres, { type Sql } from 'postgres' import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { prepareFileAccountingInTx } from '@/lib/billing/storage/accounting' import { type ProjectStorageBillingContext, resolveProjectStorageBillingContext, @@ -21,10 +22,7 @@ import { changeProjectAndWorkspaceStoragePayersInTx, changeProjectStoragePayersInTx, } from '@/lib/billing/storage/payer-transfer' -import { - incrementStorageUsageForBillingContextInTx, - prepareProjectStorageMutationInTx, -} from '@/lib/billing/storage/tracking' +import { prepareFileStorageMutationInTx } from '@/lib/billing/storage/tracking' import { prepareProjectsForAccountDeletion } from '@/lib/projects/account-deletion' import { transferWorkspaceProjects } from '@/lib/projects/membership' @@ -176,7 +174,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' plan: 'team_25000', customStorageLimitGB: 2, }) - const prepared = await prepareProjectStorageMutationInTx(tx, organizationContext) + const prepared = await prepareFileStorageMutationInTx(tx, organizationContext) await prepared.applyDelta(40) await tx.execute(query`UPDATE project SET organization_id = NULL WHERE id = 'project-b'`) const personalContext = await resolveProjectStorageBillingContext( @@ -189,7 +187,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' plan: 'pro_4000', customStorageLimitGB: null, }) - const personalPrepared = await prepareProjectStorageMutationInTx(tx, personalContext) + const personalPrepared = await prepareFileStorageMutationInTx(tx, personalContext) await personalPrepared.applyDelta(20) }) const [organizationPayer] = await sql`SELECT storage_used_bytes::integer AS bytes @@ -299,7 +297,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' const releaseWrite = createDeferred() const teardownPid = createDeferred() const write = database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, { + const prepared = await prepareFileStorageMutationInTx(tx, { ...context(), organizationId: null, billingEntity: { type: 'user', id: 'user-a' }, @@ -359,23 +357,30 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' check( 'two Projects and a workspace share one locked quota without committing rejected file rows', async () => { - const workspaceContext: StorageBillingContext = { - ...context('project-a', 60), - workspaceId: 'workspace-a', - } + await sql`INSERT INTO member (id, organization_id, user_id, role) + VALUES ('payer-owner', 'organization-a', 'user-b', 'owner')` + await sql`INSERT INTO subscription (id, plan, reference_id, status, metadata) + VALUES ('payer-plan', 'team', 'organization-a', 'active', ${JSON.stringify({ customStorageLimitGB: 60 / 1024 ** 3 })})` const attempts = await Promise.allSettled([ ...['project-a', 'project-b'].map((projectId) => database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context(projectId, 60)) + const { mutation: prepared } = await prepareFileAccountingInTx(tx, { + entityType: 'project', + entityId: projectId, + }) await tx.execute(query`INSERT INTO workspace_files (id, project_id, context, size_bytes) VALUES (${projectId}, ${projectId}, 'project', 40)`) await prepared.applyDelta(40) }) ), database.transaction(async (tx) => { + const { mutation } = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: 'workspace-a', + }) await tx.execute(query`INSERT INTO workspace_files (id, workspace_id, context, size_bytes) VALUES ('workspace-file', 'workspace-a', 'workspace', 40)`) - await incrementStorageUsageForBillingContextInTx(tx, workspaceContext, 40) + await mutation.applyDelta(40) }), ]) expect(attempts.filter((result) => result.status === 'fulfilled')).toHaveLength(1) @@ -387,15 +392,90 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' await sql`SELECT storage_used_bytes::integer AS bytes FROM organization WHERE id = 'organization-a'` expect(payer.bytes).toBe(40) expect(await sql`SELECT id FROM workspace_files`).toHaveLength(1) + const [workspaceUsage] = + await sql`SELECT storage_used_bytes::integer AS bytes FROM workspace WHERE id = 'workspace-a'` + const [workspaceHeads] = + await sql`SELECT COALESCE(sum(size_bytes), 0)::integer AS bytes FROM workspace_files WHERE workspace_id = 'workspace-a'` + expect(workspaceUsage.bytes).toBe(workspaceHeads.bytes) } ) + for (const entityType of ['workspace', 'project'] as const) { + check( + `${entityType} admits exact limit, zero and shrinking over quota with consistent ledgers`, + async () => { + const entityId = entityType === 'workspace' ? 'workspace-a' : 'project-a' + await sql`INSERT INTO member (id, organization_id, user_id, role) + VALUES ('payer-owner', 'organization-a', 'user-b', 'owner')` + await sql`INSERT INTO subscription (id, plan, reference_id, status, metadata) + VALUES ('payer-plan', 'team', 'organization-a', 'active', ${JSON.stringify({ customStorageLimitGB: 60 / 1024 ** 3 })})` + await database.transaction(async (tx) => { + const { billing, mutation } = await prepareFileAccountingInTx(tx, { + entityType, + entityId, + }) + expect(billing.billingEntity).toEqual({ type: 'organization', id: 'organization-a' }) + expect(await mutation.applyDelta(60)).toBe(60) + }) + await expect( + database.transaction(async (tx) => { + const { mutation } = await prepareFileAccountingInTx(tx, { entityType, entityId }) + await tx.execute( + query`INSERT INTO workspace_files (id, context, size_bytes) VALUES ('rejected', ${entityType}, 1)` + ) + await mutation.applyDelta(1) + }) + ).rejects.toBeInstanceOf(StorageLimitExceededError) + expect(await sql`SELECT id FROM workspace_files`).toEqual([]) + await sql`UPDATE organization SET storage_used_bytes = 100 WHERE id = 'organization-a'` + for (const [delta, expected] of [ + [0, 100], + [-20, 80], + ] as const) { + await database.transaction(async (tx) => { + const { mutation } = await prepareFileAccountingInTx(tx, { entityType, entityId }) + expect(await mutation.applyDelta(delta)).toBe(expected) + }) + } + expect( + await sql`SELECT storage_used_bytes::integer AS bytes FROM organization WHERE id = 'organization-a'` + ).toEqual([{ bytes: 80 }]) + expect( + await sql`SELECT storage_used_bytes::integer AS bytes FROM workspace WHERE id = 'workspace-a'` + ).toEqual([{ bytes: entityType === 'workspace' ? 40 : 0 }]) + } + ) + } + + check('rejects a stale workspace payer snapshot without charging either ledger', async () => { + const snapshot: StorageBillingContext = { + workspaceId: 'workspace-a', + billedAccountUserId: 'user-a', + billingEntity: { type: 'organization', id: 'organization-a' }, + plan: 'team', + customStorageLimitGB: 1, + } + await sql`UPDATE workspace SET organization_id = 'organization-b' WHERE id = 'workspace-a'` + await expect( + database.transaction(async (tx) => { + const mutation = await prepareFileStorageMutationInTx(tx, snapshot) + await mutation.applyDelta(20) + }) + ).rejects.toThrow(/payer changed/) + expect( + await sql`SELECT storage_used_bytes::integer AS bytes FROM organization ORDER BY id` + ).toEqual([{ bytes: 0 }, { bytes: 0 }]) + expect(await sql`SELECT storage_used_bytes::integer AS bytes FROM workspace`).toEqual([ + { bytes: 0 }, + ]) + }) + check( 'rejects stale owner snapshots and invalid deltas before a charge can survive', async () => { await sql`UPDATE project SET owner_id = 'user-b' WHERE id = 'project-a'` await expect( - database.transaction((tx) => prepareProjectStorageMutationInTx(tx, context())) + database.transaction((tx) => prepareFileStorageMutationInTx(tx, context())) ).rejects.toThrow(/changed/) await sql`UPDATE project SET owner_id = 'user-a' WHERE id = 'project-a'` for (const delta of [ @@ -406,14 +486,14 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' ]) { await expect( database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await prepared.applyDelta(delta) }) ).rejects.toThrow(/Invalid/) } await expect( database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await prepared.applyDelta(1) await prepared.applyDelta(1) }) @@ -428,7 +508,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' 'keeps current-head deltas separate from retained history and rolls back failed finalization', async () => { await database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await tx.execute( query`INSERT INTO workspace_files (id, project_id, workspace_id, context, size_bytes, deleted_at) VALUES ('file', 'project-a', NULL, 'project', 100, NULL)` ) @@ -438,7 +518,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' await prepared.applyDelta(100) }) await database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await tx.execute( query`UPDATE workspace_files SET size_bytes = 40, deleted_at = now() WHERE id = 'file'` ) @@ -446,7 +526,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' }) await expect( database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await prepared.applyDelta(10) throw new Error('finalization failed') }) @@ -455,7 +535,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' await sql`SELECT storage_used_bytes::integer AS bytes FROM organization WHERE id = 'organization-a'` expect(payer.bytes).toBe(40) await database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await tx.execute(query`DELETE FROM workspace_files WHERE id = 'file'`) await prepared.applyDelta(-40) }) @@ -534,11 +614,65 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' } ) + check( + 'a workspace payer transfer waits for prepared finalization and moves its committed head once', + async () => { + const written = createDeferred() + const release = createDeferred() + const waiting = createDeferred() + const write = database.transaction(async (tx) => { + const { mutation } = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: 'workspace-a', + }) + await tx.execute(query`INSERT INTO workspace_files (id, workspace_id, context, size_bytes) + VALUES ('workspace-file', 'workspace-a', 'workspace', 40)`) + await mutation.applyDelta(40) + written.resolve() + await release.promise + }) + await Promise.race([written.promise, write]) + const transfer = database.transaction(async (tx) => { + const [backend] = await tx.execute<{ pid: number }>(query`SELECT pg_backend_pid() AS pid`) + waiting.resolve(backend.pid) + return changeProjectAndWorkspaceStoragePayersInTx(tx, { + projectChanges: [], + workspaceChanges: [ + { + workspaceId: 'workspace-a', + billedAccountUserId: 'user-b', + organizationId: 'organization-b', + expectedCurrentPayer: { + billedAccountUserId: 'user-a', + organizationId: 'organization-a', + }, + }, + ], + }) + }) + try { + await waitForDatabaseLock(await waiting.promise) + } finally { + release.resolve() + } + await Promise.all([write, transfer]) + expect( + await sql`SELECT id, storage_used_bytes::integer AS bytes FROM organization ORDER BY id` + ).toEqual([ + { id: 'organization-a', bytes: 0 }, + { id: 'organization-b', bytes: 40 }, + ]) + expect( + await sql`SELECT storage_used_bytes::integer AS bytes FROM workspace WHERE id = 'workspace-a'` + ).toEqual([{ bytes: 40 }]) + } + ) + check('reconciliation waits for a Project write and keeps the committed charge', async () => { const changed = createDeferred() const release = createDeferred() const write = database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await tx.execute( query`INSERT INTO workspace_files (id, project_id, workspace_id, context, size_bytes, deleted_at) VALUES ('file', 'project-a', NULL, 'project', 40, NULL)` ) @@ -575,7 +709,7 @@ describe('Project storage admission, transfer, and reconciliation in PostgreSQL' const changed = createDeferred() const release = createDeferred() const write = database.transaction(async (tx) => { - const prepared = await prepareProjectStorageMutationInTx(tx, context()) + const prepared = await prepareFileStorageMutationInTx(tx, context()) await tx.execute( query`INSERT INTO workspace_files (id, project_id, workspace_id, context, size_bytes, deleted_at) VALUES ('file', 'project-a', NULL, 'project', 40, NULL)` ) diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts index 17e9dc47708..a206de18107 100644 --- a/apps/sim/lib/billing/storage/tracking.ts +++ b/apps/sim/lib/billing/storage/tracking.ts @@ -219,36 +219,41 @@ function assertWorkspaceStorageContext( } /** - * Establishes Project then payer locks before directory/file/version mutation. The returned - * one-shot delta belongs to this transaction; it is finalization admission, not a reservation. - * Project contribution is derived from retained file heads, so only the payer counter changes. + * Locks the asserted owner and payer before directory/file/version mutation. The caller holds + * lifecycle locks and owns the transaction. The one-shot signed delta admits finalization, not + * a reservation. Workspace files update both ledgers; Projects update only their payer ledger. */ -export async function prepareProjectStorageMutationInTx( +export async function prepareFileStorageMutationInTx( tx: DbTransaction, - context: ProjectStorageBillingContext + context: ProjectStorageBillingContext | StorageBillingContext ): Promise<{ applyDelta(deltaBytes: number): Promise }> { - await lockProject(tx, context.projectId) - const [owner] = await tx - .select({ ownerId: project.ownerId, organizationId: project.organizationId }) - .from(project) - .where(eq(project.id, context.projectId)) - .for('update') - .limit(1) - if (!owner) throw new Error(`Project ${context.projectId} not found for storage accounting`) - const billingEntity: BillingEntity = owner.organizationId - ? { type: 'organization', id: owner.organizationId } - : { type: 'user', id: owner.ownerId } - if ( - owner.ownerId !== context.ownerId || - owner.organizationId !== context.organizationId || - billingEntity.type !== context.billingEntity.type || - billingEntity.id !== context.billingEntity.id || - (!owner.organizationId && context.billedAccountUserId !== owner.ownerId) - ) { - throw new Error( - `Storage payer changed for Project ${context.projectId}; resolve a fresh billing context` - ) + if ('workspaceId' in context) { + await mutateWorkspaceStorageUsage(tx, context.workspaceId, 0, 'decrement', undefined, context) + } else { + await lockProject(tx, context.projectId) + const [owner] = await tx + .select({ ownerId: project.ownerId, organizationId: project.organizationId }) + .from(project) + .where(eq(project.id, context.projectId)) + .for('update') + .limit(1) + if (!owner) throw new Error(`Project ${context.projectId} not found for storage accounting`) + const billingEntity: BillingEntity = owner.organizationId + ? { type: 'organization', id: owner.organizationId } + : { type: 'user', id: owner.ownerId } + if ( + owner.ownerId !== context.ownerId || + owner.organizationId !== context.organizationId || + billingEntity.type !== context.billingEntity.type || + billingEntity.id !== context.billingEntity.id || + (!owner.organizationId && context.billedAccountUserId !== owner.ownerId) + ) { + throw new Error( + `Storage payer changed for Project ${context.projectId}; resolve a fresh billing context` + ) + } } + const billingEntity = context.billingEntity const currentUsage = await lockStorageUsageForMutation(tx, billingEntity) if (!Number.isSafeInteger(currentUsage) || currentUsage < 0) { throw new Error(`Invalid storage usage for payer ${getPayerKey(billingEntity)}`) @@ -259,18 +264,29 @@ export async function prepareProjectStorageMutationInTx( let applied = false return { async applyDelta(deltaBytes) { - if (applied) throw new Error('Project storage delta was already applied') - if (!Number.isSafeInteger(deltaBytes)) throw new Error('Invalid Project storage delta') + if (applied) throw new Error('File storage delta was already applied') + if (!Number.isSafeInteger(deltaBytes)) throw new Error('Invalid file storage delta') applied = true const latestUsage = await lockStorageUsageForMutation(tx, billingEntity) const nextUsage = Math.max(0, latestUsage + deltaBytes) - if (!Number.isSafeInteger(nextUsage)) throw new Error('Invalid Project storage total') + if (!Number.isSafeInteger(nextUsage)) throw new Error('Invalid file storage total') if (deltaBytes > 0 && limit !== undefined && nextUsage > limit) { throw new StorageLimitExceededError( `Storage limit exceeded. Used: ${(nextUsage / 1024 ** 3).toFixed(2)}GB, Limit: ${(limit / 1024 ** 3).toFixed(0)}GB` ) } if (deltaBytes === 0) return latestUsage + if ('workspaceId' in context) { + const result = await mutateWorkspaceStorageUsage( + tx, + context.workspaceId, + Math.abs(deltaBytes), + deltaBytes > 0 ? 'increment' : 'decrement', + limit, + context + ) + return result.updatedUsage + } if (deltaBytes < 0 && latestUsage < -deltaBytes) { logger.error('Clamping Project storage payer ledger underflow', { projectId: context.projectId, @@ -531,14 +547,10 @@ export async function applyStorageUsageDeltasInTx( /** * Locks the canonical workspace row before a file mutation that changes its storage ledger. * - * Every transaction that writes a `context = 'workspace'` row of `workspace_files` and charges - * storage takes its locks in one order: the folder advisory lock (only for operations that take - * one), then this workspace row, then the file rows and their writes (whose search-index trigger - * upserts the one-per-workspace `workspace_file_search_dispatch_queue` row), then the payer row. - * Call this before the first file insert or update: an insert that reaches the queue row first - * and only then waits on the workspace row deadlocks with a content write that holds the - * workspace row and waits on the queue row. Re-locking later in the same transaction, as the - * storage increment does, never waits. + * Locks the workspace ledger before file writes and their search-dispatch trigger. Shared file + * accounting acquires lifecycle, workspace, then payer locks before directory/file/history locks. + * Callers preparing accounting must do so before their first file mutation; this low-level helper + * only locks the workspace row. Re-locking it during the final signed delta never waits. */ export async function lockWorkspaceStorageForMutationInTx( tx: DbOrTx, diff --git a/apps/sim/lib/file-retention/folders.ts b/apps/sim/lib/file-retention/folders.ts new file mode 100644 index 00000000000..ea6b9c6576f --- /dev/null +++ b/apps/sim/lib/file-retention/folders.ts @@ -0,0 +1,157 @@ +import { dbFor } from '@sim/db' +import { folder, workspaceFiles } from '@sim/db/schema' +import { and, asc, eq, inArray, isNotNull, isNull, lt, sql } from 'drizzle-orm' +import { + consumeRowBudget, + DEFAULT_MAX_BATCHES_PER_TABLE, + type RowBudget, +} from '@/lib/cleanup/batch-delete' +import { generateRestoreName } from '@/lib/core/utils/restore-name' +import type { DbTransaction } from '@/lib/db/types' +import { FILES_PER_QUERY } from '@/lib/file-retention/versions' +import { deduplicateFolderNameInScope } from '@/lib/folders/naming' +import { lockWorkspaceProject } from '@/lib/projects/membership' +import { workspaceFileNameFolderCondition } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' +import { lockFileDirectories } from '@/lib/workspace-files/locks' +import type { EditableFileOwner } from '@/lib/workspace-files/ownership' +import { fileFolderOwnerCondition, fileOwnerCondition } from '@/lib/workspace-files/ownership-query' + +const cleanupDb = dbFor('cleanup') + +/** Recheck expiry and keep child re-rooting atomic with the parent deletion. */ +export async function cleanupExpiredFileFolderInTx( + tx: DbTransaction, + owner: EditableFileOwner, + cutoff: Date, + budget?: RowBudget +): Promise { + const folderScope = fileFolderOwnerCondition(owner) + await lockFileDirectories(tx, [owner]) + const [expired] = await tx + .select({ id: folder.id }) + .from(folder) + .where(and(folderScope, isNotNull(folder.deletedAt), lt(folder.deletedAt, cutoff))) + .orderBy(asc(folder.id)) + .limit(1) + .for('update') + if (!expired) return 0 + consumeRowBudget(budget, 1) + for (;;) { + const children = await tx + .select({ id: folder.id, name: folder.name }) + .from(folder) + .where(and(folderScope, eq(folder.parentId, expired.id), isNull(folder.deletedAt))) + .orderBy(asc(folder.id)) + .limit(FILES_PER_QUERY) + .for('update') + if (!children.length) break + for (const child of children) { + const name = await deduplicateFolderNameInScope(tx, folderScope, null, child.name) + await tx + .update(folder) + .set({ parentId: null, name, updatedAt: new Date() }) + .where(and(folderScope, eq(folder.id, child.id))) + } + } + for (;;) { + const children = await tx + .select({ id: workspaceFiles.id, name: workspaceFiles.originalName }) + .from(workspaceFiles) + .where( + and( + fileOwnerCondition(owner), + eq(workspaceFiles.folderId, expired.id), + isNull(workspaceFiles.deletedAt) + ) + ) + .orderBy(asc(workspaceFiles.id)) + .limit(FILES_PER_QUERY) + .for('update') + if (!children.length) break + const candidateNames = new Set(children.map(({ name }) => name)) + const rootFiles = await tx + .select({ name: workspaceFiles.originalName }) + .from(workspaceFiles) + .where( + and( + fileOwnerCondition(owner), + workspaceFileNameFolderCondition(null), + isNull(workspaceFiles.deletedAt), + inArray(workspaceFiles.originalName, [...candidateNames]) + ) + ) + const rootNames = new Set(rootFiles.map(({ name }) => name)) + const renamed: { id: string; name: string }[] = [] + for (const child of children) { + const name = await generateRestoreName( + child.name, + async (candidate) => { + if (rootNames.has(candidate)) return true + if (candidateNames.has(candidate)) return false + const [existing] = await tx + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where( + and( + fileOwnerCondition(owner), + workspaceFileNameFolderCondition(null), + isNull(workspaceFiles.deletedAt), + eq(workspaceFiles.originalName, candidate) + ) + ) + .limit(1) + return Boolean(existing) + }, + { hasExtension: true } + ) + rootNames.add(name) + renamed.push({ id: child.id, name }) + } + await tx + .update(workspaceFiles) + .set({ + folderId: null, + originalName: sql`CASE ${workspaceFiles.id} ${sql.join( + renamed.map(({ id, name }) => sql`WHEN ${id} THEN ${name}`), + sql` ` + )} END`, + updatedAt: new Date(), + }) + .where( + and( + fileOwnerCondition(owner), + inArray( + workspaceFiles.id, + renamed.map(({ id }) => id) + ) + ) + ) + } + await tx.delete(folder).where(and(folderScope, eq(folder.id, expired.id))) + return 1 +} + +/** Workspace file folders share the directory transaction used by Project folders. */ +export async function cleanupArchivedWorkspaceFileFolders( + workspaceIds: string[], + cutoff: Date, + budget?: RowBudget +): Promise { + let total = 0 + for (const workspaceId of workspaceIds) { + for (let batch = 0; batch < DEFAULT_MAX_BATCHES_PER_TABLE && budget?.remaining !== 0; batch++) { + const removed = await cleanupDb.transaction(async (tx) => { + await lockWorkspaceProject(tx, workspaceId) + return cleanupExpiredFileFolderInTx( + tx, + { entityType: 'workspace', entityId: workspaceId }, + cutoff, + budget + ) + }) + total += removed + if (!removed) break + } + } + return total +} diff --git a/apps/sim/lib/file-retention/index.ts b/apps/sim/lib/file-retention/index.ts new file mode 100644 index 00000000000..d3bac9b0f93 --- /dev/null +++ b/apps/sim/lib/file-retention/index.ts @@ -0,0 +1 @@ +export { beginFileArchiveCleanup, cleanupFileVersions, fileRetentionOwners } from './service' diff --git a/apps/sim/lib/file-retention/legacy-archive.ts b/apps/sim/lib/file-retention/legacy-archive.ts new file mode 100644 index 00000000000..dc9dd75c426 --- /dev/null +++ b/apps/sim/lib/file-retention/legacy-archive.ts @@ -0,0 +1,385 @@ +import { db, dbFor } from '@sim/db' +import { changelogRelease, workflowTest, workspaceFile, workspaceFiles } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { chunkArray } from '@sim/utils/helpers' +import { and, eq, inArray, isNotNull, lt } from 'drizzle-orm' +import { prepareFileAccountingInTx } from '@/lib/billing/storage/accounting' +import { DEFAULT_DELETE_CHUNK_SIZE, selectRowsByIdChunks } from '@/lib/cleanup/batch-delete' +import type { CleanupBudgets } from '@/lib/cleanup/limits' +import { type CleanupOwnerScope, cleanupOwnerCondition } from '@/lib/cleanup/resource-scope' +import { cleanupArchivedWorkspaceFileFolders } from '@/lib/file-retention/folders' +import type { FileArchiveCleanup, FileRetentionOptions } from '@/lib/file-retention/types' +import { isUsingCloudStorage, type StorageContext, StorageService } from '@/lib/uploads' +import { enqueueWorkspaceFileStorageCleanups } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' +import { releaseWorkspaceFileVersionsForPurgeInTx } from '@/lib/uploads/contexts/workspace/workspace-file-versions' +import { getWorkspaceFileSize } from '@/lib/uploads/shared/types' + +const logger = createLogger('FileArchiveCleanup') +/** Billing deletion retains the default pool; selection and unbilled cleanup use the cleanup pool. */ +const cleanupDb = dbFor('cleanup') +const BILLED_FILE_CONTEXTS: readonly StorageContext[] = ['workspace', 'test', 'changelog'] + +interface WorkspaceFileScope { + /** Rows from `workspace_file` (singular, legacy workspace-context only). */ + legacyRows: Array<{ id: string; key: string; workspaceId: string }> + /** Rows from `workspace_files` (plural, multi-context). */ + multiContextRows: Array<{ + id: string + key: string + workspaceId: string | null + context: StorageContext + size: number + }> +} + +interface WorkspaceFileStorageCleanupResult { + filesDeleted: number + filesFailed: number + legacyRows: WorkspaceFileScope['legacyRows'] + multiContextRows: WorkspaceFileScope['multiContextRows'] +} + +/** + * Select every soft-deleted file row that's eligible for permanent removal. + * Workspace objects are released through the purge transaction; other legacy + * contexts retain the selected rows for their storage-first cleanup. + */ +async function selectExpiredWorkspaceFiles( + scope: CleanupOwnerScope, + retentionDate: Date, + budgets?: CleanupBudgets +): Promise { + const [legacyRows, multiContextRows] = await Promise.all([ + selectRowsByIdChunks( + scope.kind === 'workspace' ? scope.ids : [], + (chunkIds, chunkLimit) => + cleanupDb + .select({ + id: workspaceFile.id, + key: workspaceFile.key, + workspaceId: workspaceFile.workspaceId, + }) + .from(workspaceFile) + .where( + and( + inArray(workspaceFile.workspaceId, chunkIds), + isNotNull(workspaceFile.deletedAt), + lt(workspaceFile.deletedAt, retentionDate) + ) + ) + .limit(chunkLimit), + { budget: budgets?.legacyFiles } + ), + selectRowsByIdChunks( + scope.ids, + (chunkIds, chunkLimit) => + cleanupDb + .select({ + id: workspaceFiles.id, + key: workspaceFiles.key, + workspaceId: workspaceFiles.workspaceId, + context: workspaceFiles.context, + sizeBytes: workspaceFiles.sizeBytes, + }) + .from(workspaceFiles) + .where( + and( + cleanupOwnerCondition(workspaceFiles, scope, chunkIds), + scope.kind === 'organization' + ? eq(workspaceFiles.context, 'knowledge-base') + : undefined, + isNotNull(workspaceFiles.deletedAt), + lt(workspaceFiles.deletedAt, retentionDate) + ) + ) + .limit(chunkLimit), + { budget: budgets?.files } + ), + ]) + + return { + legacyRows, + multiContextRows: multiContextRows.map((r) => ({ + id: r.id, + key: r.key, + workspaceId: r.workspaceId, + context: r.context as StorageContext, + size: getWorkspaceFileSize(r), + })), + } +} + +async function cleanupWorkspaceFileStorage( + scope: WorkspaceFileScope +): Promise { + const result: WorkspaceFileStorageCleanupResult = { + filesDeleted: 0, + filesFailed: 0, + legacyRows: [], + multiContextRows: [], + } + if (!isUsingCloudStorage()) { + return { + ...result, + legacyRows: scope.legacyRows, + multiContextRows: scope.multiContextRows, + } + } + + const candidatesByContext = new Map() + result.legacyRows = scope.legacyRows + for (const row of scope.multiContextRows) { + if (BILLED_FILE_CONTEXTS.includes(row.context)) { + result.multiContextRows.push(row) + continue + } + const bucket = candidatesByContext.get(row.context) + if (bucket) bucket.push(row) + else candidatesByContext.set(row.context, [row]) + } + + for (const [context, candidates] of candidatesByContext) { + for (const batch of chunkArray(candidates, DEFAULT_DELETE_CHUNK_SIZE)) { + const deletion = await StorageService.deleteFiles( + batch.map((row) => row.key), + context + ) + const failedKeys = new Set(deletion.failed.map(({ key }) => key)) + result.filesDeleted += batch.filter((row) => !failedKeys.has(row.key)).length + result.filesFailed += deletion.failed.length + + for (const row of batch) { + if (!failedKeys.has(row.key)) result.multiContextRows.push(row) + } + for (const { key, error } of deletion.failed) { + logger.error(`Failed to delete storage file ${key} (context: ${context}):`, { error }) + } + } + } + + return result +} + +async function deleteExpiredLegacyWorkspaceFileRows( + rows: WorkspaceFileScope['legacyRows'], + retentionDate: Date, + label: string +): Promise<{ deleted: number; failed: number }> { + const result = { deleted: 0, failed: 0 } + for (const batch of chunkArray(rows, DEFAULT_DELETE_CHUNK_SIZE)) { + try { + const deleted = await cleanupDb.transaction(async (tx) => { + const rows = await tx + .delete(workspaceFile) + .where( + and( + inArray( + workspaceFile.id, + batch.map(({ id }) => id) + ), + isNotNull(workspaceFile.deletedAt), + lt(workspaceFile.deletedAt, retentionDate) + ) + ) + .returning({ id: workspaceFile.id, key: workspaceFile.key }) + await enqueueWorkspaceFileStorageCleanups( + tx, + rows.map(({ key }) => key) + ) + return rows + }) + result.deleted += deleted.length + result.failed += batch.length - deleted.length + } catch (error) { + result.failed += batch.length + logger.error(`[${label}/workspaceFile] Exact-row delete failed`, { error }) + } + } + return result +} + +async function deleteExpiredUnbilledWorkspaceFileRows( + rows: WorkspaceFileScope['multiContextRows'], + retentionDate: Date, + label: string +): Promise<{ deleted: number; failed: number }> { + const result = { deleted: 0, failed: 0 } + const rowsByContext = new Map() + for (const row of rows) { + if (BILLED_FILE_CONTEXTS.includes(row.context)) continue + const bucket = rowsByContext.get(row.context) + if (bucket) bucket.push(row) + else rowsByContext.set(row.context, [row]) + } + + for (const [context, contextRows] of rowsByContext) { + for (const batch of chunkArray(contextRows, DEFAULT_DELETE_CHUNK_SIZE)) { + try { + const deleted = await cleanupDb + .delete(workspaceFiles) + .where( + and( + inArray( + workspaceFiles.id, + batch.map(({ id }) => id) + ), + eq(workspaceFiles.context, context), + isNotNull(workspaceFiles.deletedAt), + lt(workspaceFiles.deletedAt, retentionDate) + ) + ) + .returning({ id: workspaceFiles.id }) + result.deleted += deleted.length + result.failed += batch.length - deleted.length + } catch (error) { + result.failed += batch.length + logger.error(`[${label}/workspaceFiles] Exact-row ${context} delete failed`, { error }) + } + } + } + return result +} + +async function deleteExpiredBillableWorkspaceFileRows( + rows: WorkspaceFileScope['multiContextRows'], + retentionDate: Date, + label: string +): Promise<{ deleted: number; failed: number }> { + const result = { deleted: 0, failed: 0 } + const rowsByWorkspace = new Map() + for (const row of rows) { + if (!BILLED_FILE_CONTEXTS.includes(row.context)) continue + if (!row.workspaceId) { + result.failed++ + logger.error(`[${label}/workspaceFiles] Billable row has no workspace attribution`, { + fileId: row.id, + }) + continue + } + const bucket = rowsByWorkspace.get(row.workspaceId) + if (bucket) bucket.push(row) + else rowsByWorkspace.set(row.workspaceId, [row]) + } + + for (const [workspaceId, workspaceRows] of rowsByWorkspace) { + for (const batch of chunkArray(workspaceRows, DEFAULT_DELETE_CHUNK_SIZE)) { + try { + const deletedCount = await db.transaction(async (tx) => { + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: workspaceId, + }) + const fileIds = batch.map(({ id }) => id) + await tx + .delete(workflowTest) + .where( + and( + inArray(workflowTest.bodyFileId, fileIds), + isNotNull(workflowTest.deletedAt), + lt(workflowTest.deletedAt, retentionDate) + ) + ) + await tx.delete(changelogRelease).where( + inArray( + changelogRelease.bodyFileId, + tx + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where( + and( + inArray(workspaceFiles.id, fileIds), + isNotNull(workspaceFiles.deletedAt), + lt(workspaceFiles.deletedAt, retentionDate) + ) + ) + ) + ) + await releaseWorkspaceFileVersionsForPurgeInTx( + tx, + batch.map(({ id }) => id), + retentionDate + ) + const deletedRows = await tx + .delete(workspaceFiles) + .where( + and( + inArray( + workspaceFiles.id, + batch.map(({ id }) => id) + ), + eq(workspaceFiles.workspaceId, workspaceId), + inArray(workspaceFiles.context, BILLED_FILE_CONTEXTS), + isNotNull(workspaceFiles.deletedAt), + lt(workspaceFiles.deletedAt, retentionDate) + ) + ) + .returning({ + id: workspaceFiles.id, + key: workspaceFiles.key, + sizeBytes: workspaceFiles.sizeBytes, + }) + await enqueueWorkspaceFileStorageCleanups( + tx, + deletedRows.map(({ key }) => key) + ) + const deletedBytes = deletedRows.reduce( + (total, row) => total + getWorkspaceFileSize(row), + 0 + ) + await accounting.mutation.applyDelta(-deletedBytes) + return deletedRows.length + }) + result.deleted += deletedCount + result.failed += batch.length - deletedCount + } catch (error) { + result.failed += batch.length + logger.error(`[${label}/workspaceFiles] Atomic delete and decrement failed`, { + error, + workspaceId, + }) + } + } + } + return result +} + +/** Preserve legacy workspace/multi-context cleanup while selecting each owner batch once. */ +export async function prepareLegacyFileArchiveCleanup( + scope: CleanupOwnerScope, + { cutoff, budgets, label }: FileRetentionOptions +): Promise { + const selected = await selectExpiredWorkspaceFiles(scope, cutoff, budgets) + return { + async cleanupStorage() { + const storage = await cleanupWorkspaceFileStorage(selected) + if (budgets && storage.filesFailed) throw new Error('File storage cleanup failed') + return { + filesDeleted: storage.filesDeleted, + async deleteRows() { + const legacy = await deleteExpiredLegacyWorkspaceFileRows( + storage.legacyRows, + cutoff, + label + ) + const billable = await deleteExpiredBillableWorkspaceFileRows( + storage.multiContextRows, + cutoff, + label + ) + const unbilled = await deleteExpiredUnbilledWorkspaceFileRows( + storage.multiContextRows, + cutoff, + label + ) + if (budgets && (legacy.failed || billable.failed || unbilled.failed)) + throw new Error('File row cleanup failed') + const foldersDeleted = + scope.kind === 'workspace' + ? await cleanupArchivedWorkspaceFileFolders(scope.ids, cutoff, budgets?.folders) + : 0 + return legacy.deleted + billable.deleted + unbilled.deleted + foldersDeleted + }, + } + }, + } +} diff --git a/apps/sim/lib/file-retention/service.ts b/apps/sim/lib/file-retention/service.ts new file mode 100644 index 00000000000..51225fb6d72 --- /dev/null +++ b/apps/sim/lib/file-retention/service.ts @@ -0,0 +1,147 @@ +import type { CleanupJobPayload } from '@/lib/billing/cleanup-dispatcher' +import { + DEFAULT_DELETE_CHUNK_SIZE, + DEFAULT_MAX_BATCHES_PER_TABLE, +} from '@/lib/cleanup/batch-delete' +import { prepareLegacyFileArchiveCleanup } from '@/lib/file-retention/legacy-archive' +import type { + FileArchiveCleanup, + FileArchiveDeletion, + FileRetentionOptions, + FileVersionCleanupResult, +} from '@/lib/file-retention/types' +import { cleanupWorkspaceFileVersions } from '@/lib/file-retention/workspace-versions' +import { + cleanupArchivedProjectFileFolders, + cleanupArchivedProjectFiles, + cleanupProjectFileVersions, +} from '@/lib/projects/files/retention' +import { + type FileOwnerAdapters, + requireFileOwnerAdapter, +} from '@/lib/workspace-files/owner-adapters' +import type { FileOwner } from '@/lib/workspace-files/ownership' + +interface RetentionAdapter { + versions?( + ids: string[], + options: FileRetentionOptions, + limit: number + ): Promise + archive(ids: string[], options: FileRetentionOptions): Promise +} + +const OWNER_ADAPTERS: FileOwnerAdapters = { + workspace: { + versions: cleanupWorkspaceFileVersions, + archive: (ids, options) => prepareLegacyFileArchiveCleanup({ kind: 'workspace', ids }, options), + }, + organization: { + archive: (ids, options) => + prepareLegacyFileArchiveCleanup({ kind: 'organization', ids }, options), + }, + project: { + async versions(ids, _options, limit) { + let deleted = 0 + for (const id of ids) { + if (deleted >= limit) break + deleted += await cleanupProjectFileVersions(id, limit - deleted) + } + return { deleted, attempted: deleted } + }, + async archive(ids, { budgets }) { + const files = budgets?.files ?? { + remaining: DEFAULT_DELETE_CHUNK_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, + } + const folders = budgets?.folders ?? { + remaining: DEFAULT_DELETE_CHUNK_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, + } + let deleted = 0 + for (const id of ids) { + if (files.remaining <= 0 && folders.remaining <= 0) break + deleted += await cleanupArchivedProjectFiles(id, files) + deleted += await cleanupArchivedProjectFileFolders(id, folders) + } + return { + cleanupStorage: async () => ({ filesDeleted: 0, deleteRows: async () => deleted }), + } + }, + }, +} + +/** Keep queued payloads compatible while file operations consume one owner interface. */ +export function fileRetentionOwners( + payload: Pick +): FileOwner[] { + return [ + ...(payload.projectIds ?? []).map( + (entityId): FileOwner => ({ entityType: 'project', entityId }) + ), + ...payload.workspaceIds.map((entityId): FileOwner => ({ entityType: 'workspace', entityId })), + ...(payload.organizationIds ?? []).map( + (entityId): FileOwner => ({ entityType: 'organization', entityId }) + ), + ] +} + +function ownerBatches(owners: FileOwner[], operation: 'versions' | 'archive') { + const batches = new Map< + FileOwner['entityType'], + { adapter: RetentionAdapter; ids: Set } + >() + for (const owner of owners) { + const adapter = requireFileOwnerAdapter(OWNER_ADAPTERS, owner) + if (!adapter[operation]) + throw new Error(`File ${operation} retention does not support this owner`) + const batch = batches.get(owner.entityType) + if (batch) batch.ids.add(owner.entityId) + else batches.set(owner.entityType, { adapter, ids: new Set([owner.entityId]) }) + } + return [...batches.entries()].map(([type, { adapter, ids }]) => ({ + adapter, + ids: type === 'project' ? [...ids].sort() : [...ids], + })) +} + +/** Validate every owner before destructive work and share one attempt cap across batches. */ +export async function cleanupFileVersions( + owners: FileOwner[], + options: FileRetentionOptions, + limit: number +): Promise { + const batches = ownerBatches(owners, 'versions') + let deleted = 0 + let attempted = 0 + for (const { adapter, ids } of batches) { + if (attempted >= limit) break + if (!adapter.versions) throw new Error('File version retention is unavailable') + const result = await adapter.versions(ids, options, limit - attempted) + deleted += result.deleted + attempted += result.attempted + } + return deleted +} + +/** Begin owner cleanup; Project commits immediately, while legacy storage and row phases stay ordered. */ +export async function beginFileArchiveCleanup( + owners: FileOwner[], + options: FileRetentionOptions +): Promise { + const batches = ownerBatches(owners, 'archive') + const cleanups: FileArchiveCleanup[] = [] + for (const { adapter, ids } of batches) cleanups.push(await adapter.archive(ids, options)) + return { + async cleanupStorage() { + const deletions: FileArchiveDeletion[] = [] + for (const cleanup of cleanups) deletions.push(await cleanup.cleanupStorage()) + return { + filesDeleted: deletions.reduce((sum, deletion) => sum + deletion.filesDeleted, 0), + async deleteRows() { + let deleted = 0 + for (const deletion of deletions) deleted += await deletion.deleteRows() + return deleted + }, + } + }, + } +} diff --git a/apps/sim/lib/file-retention/types.ts b/apps/sim/lib/file-retention/types.ts new file mode 100644 index 00000000000..f48c094860f --- /dev/null +++ b/apps/sim/lib/file-retention/types.ts @@ -0,0 +1,27 @@ +import type { PlanCategory } from '@/lib/billing/plan-helpers' +import type { CleanupBudgets } from '@/lib/cleanup/limits' + +/** Queued workspace policy and shared budgets; Projects resolve their current payer policy in-transaction. */ +export interface FileRetentionOptions { + plan: PlanCategory + cutoff: Date + label: string + budgets?: CleanupBudgets +} + +/** Released versions and attempted selections, used to preserve the shared run limit. */ +export interface FileVersionCleanupResult { + deleted: number + attempted: number +} + +/** Exact selected rows survive the storage phase; failed objects never enter row deletion. */ +export interface FileArchiveCleanup { + cleanupStorage(): Promise +} + +/** Completed storage-phase count and deferred row cleanup; outbox-backed deletions are not counted as physical deletes here. */ +export interface FileArchiveDeletion { + filesDeleted: number + deleteRows(): Promise +} diff --git a/apps/sim/lib/file-retention/versions.ts b/apps/sim/lib/file-retention/versions.ts new file mode 100644 index 00000000000..b9c72ace8a7 --- /dev/null +++ b/apps/sim/lib/file-retention/versions.ts @@ -0,0 +1,72 @@ +import type { db } from '@sim/db' +import { workspaceFileVersion } from '@sim/db/schema' +import { and, gt, inArray, isNotNull, lt, or, sql } from 'drizzle-orm' +import type { DbTransaction } from '@/lib/db/types' +import { enqueueWorkspaceFileStorageCleanups } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' + +/** Always retain the current head and its newest nine predecessors. */ +export const KEEP_SUPERSEDED = 9 +export const FILES_PER_QUERY = 500 + +/** Rank each file independently; expiry never removes its protected recent history. */ +export function selectExpiredFileVersions( + executor: Pick, + fileIds: string[], + cutoff: Date, + maxSuperseded: number, + limit: number +) { + const ranked = executor + .select({ + id: workspaceFileVersion.id, + supersededAt: workspaceFileVersion.supersededAt, + rank: sql`row_number() over (partition by ${workspaceFileVersion.fileId} order by ${workspaceFileVersion.version} desc)`.as( + 'rank' + ), + }) + .from(workspaceFileVersion) + .where( + and( + inArray(workspaceFileVersion.fileId, fileIds), + isNotNull(workspaceFileVersion.supersededAt) + ) + ) + .as('ranked') + return executor + .select({ id: ranked.id }) + .from(ranked) + .where( + and( + gt(ranked.rank, KEEP_SUPERSEDED), + or(lt(ranked.supersededAt, cutoff), gt(ranked.rank, maxSuperseded)) + ) + ) + .limit(limit) +} + +/** Release history and its storage-cleanup intent in the caller's existing transaction. */ +export async function releaseExpiredFileVersions( + tx: DbTransaction, + rows: Array<{ id: string }>, + context: 'workspace' | 'project' +): Promise { + if (!rows.length) return 0 + const removed = await tx + .delete(workspaceFileVersion) + .where( + and( + inArray( + workspaceFileVersion.id, + rows.map(({ id }) => id) + ), + isNotNull(workspaceFileVersion.supersededAt) + ) + ) + .returning({ key: workspaceFileVersion.key }) + await enqueueWorkspaceFileStorageCleanups( + tx, + removed.map(({ key }) => key), + context === 'workspace' ? undefined : context + ) + return removed.length +} diff --git a/apps/sim/lib/file-retention/workspace-versions.ts b/apps/sim/lib/file-retention/workspace-versions.ts new file mode 100644 index 00000000000..fa5ada96540 --- /dev/null +++ b/apps/sim/lib/file-retention/workspace-versions.ts @@ -0,0 +1,109 @@ +import { dbFor } from '@sim/db' +import { workspaceFiles, workspaceFileVersion } from '@sim/db/schema' +import { chunkArray } from '@sim/utils/helpers' +import { and, asc, count, gt, inArray, isNotNull, lt, min, or, sql } from 'drizzle-orm' +import { + DEFAULT_DELETE_CHUNK_SIZE, + DEFAULT_MAX_BATCHES_PER_TABLE, + DEFAULT_WORKSPACE_CHUNK_SIZE, +} from '@/lib/cleanup/batch-delete' +import type { FileRetentionOptions, FileVersionCleanupResult } from '@/lib/file-retention/types' +import { + FILES_PER_QUERY, + KEEP_SUPERSEDED, + releaseExpiredFileVersions, + selectExpiredFileVersions, +} from '@/lib/file-retention/versions' +import { MAX_SUPERSEDED_FILE_VERSIONS } from '@/lib/uploads/contexts/workspace/workspace-file-versions' + +const cleanupDb = dbFor('cleanup') +const FREE_MAX_SUPERSEDED_VERSIONS = 99 + +async function selectCandidateFileIds( + workspaceIds: string[], + cutoff: Date, + maxSuperseded: number, + afterId: string +): Promise { + const rows = await cleanupDb + .select({ fileId: workspaceFileVersion.fileId }) + .from(workspaceFileVersion) + .where( + and( + inArray(workspaceFileVersion.workspaceId, workspaceIds), + isNotNull(workspaceFileVersion.supersededAt), + gt(workspaceFileVersion.fileId, afterId) + ) + ) + .groupBy(workspaceFileVersion.fileId) + .having( + and( + gt(count(), KEEP_SUPERSEDED), + or( + lt( + min(workspaceFileVersion.supersededAt), + sql.param(cutoff, workspaceFileVersion.supersededAt) + ), + gt(count(), maxSuperseded) + ) + ) + ) + .orderBy(asc(workspaceFileVersion.fileId)) + .limit(FILES_PER_QUERY) + return rows.map((row) => row.fileId) +} + +/** Preserve workspace batching and the queued workspace retention policy. */ +export async function cleanupWorkspaceFileVersions( + workspaceIds: string[], + { cutoff, plan }: FileRetentionOptions, + limit: number +): Promise { + const maxSuperseded = + plan === 'free' ? FREE_MAX_SUPERSEDED_VERSIONS : MAX_SUPERSEDED_FILE_VERSIONS + let deleted = 0 + let attempted = 0 + for (const group of chunkArray(workspaceIds, DEFAULT_WORKSPACE_CHUNK_SIZE)) { + if (attempted >= limit) break + let batches = 0 + let afterId = '' + while (batches < DEFAULT_MAX_BATCHES_PER_TABLE && attempted < limit) { + const fileIds = await selectCandidateFileIds(group, cutoff, maxSuperseded, afterId) + if (fileIds.length === 0) break + afterId = fileIds[fileIds.length - 1] + let exhausted = false + while (!exhausted && batches < DEFAULT_MAX_BATCHES_PER_TABLE && attempted < limit) { + batches++ + const batchSize = Math.min(DEFAULT_DELETE_CHUNK_SIZE, limit - attempted) + const { selected, removed } = await cleanupDb.transaction(async (tx) => { + const locked = await tx + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where( + and(inArray(workspaceFiles.id, fileIds), inArray(workspaceFiles.workspaceId, group)) + ) + .orderBy(asc(workspaceFiles.id)) + .for('update') + if (!locked.length) return { selected: 0, removed: 0 } + const expired = await selectExpiredFileVersions( + tx, + locked.map(({ id }) => id), + cutoff, + maxSuperseded, + batchSize + ) + return { + selected: expired.length, + removed: await releaseExpiredFileVersions(tx, expired, 'workspace'), + } + }) + attempted += selected + deleted += removed + exhausted = selected < batchSize || removed === 0 + } + if (!exhausted) break + } + } + + return { deleted, attempted } +} diff --git a/apps/sim/lib/projects/files/__integration__/private-purge.integration.ts b/apps/sim/lib/projects/files/__integration__/private-purge.integration.ts index 3993fce5a51..dfb464cceaa 100644 --- a/apps/sim/lib/projects/files/__integration__/private-purge.integration.ts +++ b/apps/sim/lib/projects/files/__integration__/private-purge.integration.ts @@ -32,7 +32,7 @@ import { afterAll, describe, expect, it, vi } from 'vitest' vi.mock('@/lib/uploads/core/setup.server', () => uploadsSetupMock) import { resolveProjectStorageBillingContext } from '@/lib/billing/storage/context' -import { prepareProjectStorageMutationInTx } from '@/lib/billing/storage/tracking' +import { prepareFileStorageMutationInTx } from '@/lib/billing/storage/tracking' import { processOutboxEventById } from '@/lib/core/outbox/service' import { prepareProjectsForAccountDeletion } from '@/lib/projects/account-deletion' import { workspaceFileStorageCleanupOutboxHandlers } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' @@ -314,7 +314,7 @@ describe('Private Project teardown and durable object cleanup', () => { { projectId: f.projectId, ownerId: f.userId, organizationId: null }, tx ) - const accounting = await prepareProjectStorageMutationInTx(tx, context) + const accounting = await prepareFileStorageMutationInTx(tx, context) await tx .update(workspaceFiles) .set({ sizeBytes: 10 }) diff --git a/apps/sim/lib/projects/files/__integration__/retention.integration.ts b/apps/sim/lib/projects/files/__integration__/retention.integration.ts index b8dc25bd9c4..14d913fe360 100644 --- a/apps/sim/lib/projects/files/__integration__/retention.integration.ts +++ b/apps/sim/lib/projects/files/__integration__/retention.integration.ts @@ -11,6 +11,7 @@ import { user, userStats, workspace, + workspaceFile, workspaceFiles, workspaceFileVersion, } from '@sim/db/schema' @@ -18,6 +19,7 @@ import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/workspace-fixtures' import { createDeferred } from '@sim/testing/helpers/deferred' import { envFlagsMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock' +import { featureFlagsMock, featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock' import { getErrorMessage } from '@sim/utils/errors' import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' @@ -26,13 +28,20 @@ import postgres from 'postgres' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@/lib/core/config/env-flags', () => envFlagsMock) +vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock) import { createSessionPrincipal } from '@sim/testing/factories/principal.factory' import { type CleanupJobPayload, runCleanupWithLimits } from '@/lib/billing/cleanup-dispatcher' import { changeWorkspaceStoragePayerInTx } from '@/lib/billing/storage/payer-transfer' import { createCleanupBudgets } from '@/lib/cleanup/limits' +import { beginFileArchiveCleanup, cleanupFileVersions } from '@/lib/file-retention' import { authorizeProject } from '@/lib/projects/application/authorization' import { projectOperations } from '@/lib/projects/application/operations' +import { lockWorkspaceProject } from '@/lib/projects/membership' +import { processWorkspaceFileStorageCleanupsNow } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' +import { deleteWorkspaceFileVersionInTx } from '@/lib/uploads/contexts/workspace/workspace-file-versions' +import { deleteFile, downloadFile, uploadFile } from '@/lib/uploads/core/storage-service' +import type { FileOwner } from '@/lib/workspace-files/ownership' import { runCleanupFileVersions } from '@/background/cleanup-file-versions' import { runCleanupSoftDeletes } from '@/background/cleanup-soft-deletes' @@ -45,6 +54,7 @@ const fixtures: { workspaceId: string projectId: string }[] = [] +const storedKeys: string[] = [] const HOUR = 60 * 60 * 1000 const control = postgres(readTestDatabaseUrl(), { max: 2, onnotice: () => undefined }) @@ -67,6 +77,7 @@ function check(name: string, run: () => Promise) { } beforeEach(() => { + featureFlagsMockFns.mockIsFeatureEnabled.mockResolvedValue(false) setEnvFlags({ isBillingEnabled: true, isDataRetentionEnabled: true }) }) @@ -206,19 +217,167 @@ async function versions(fileId: string) { .orderBy(workspaceFileVersion.version) } -async function cleanupEvents(projectId: string) { +async function cleanupEvents(entityId: string, kind = 'project') { return db - .select({ payload: outboxEvent.payload }) + .select({ id: outboxEvent.id, payload: outboxEvent.payload }) .from(outboxEvent) .where( and( eq(outboxEvent.eventType, 'workspace-file.storage.cleanup'), - sql`${outboxEvent.payload}->>'key' LIKE ${`project/${projectId}/%`}` + sql`${outboxEvent.payload}->>'key' LIKE ${`${kind}/${entityId}/%`}` ) ) } describe('Project file retention follows the current payer in PostgreSQL', () => { + check( + 'workspace retention reranks history after a concurrent version deletion commits', + async () => { + const f = await fixture() + const seeded = await seedFile(f, 40 * 24, false, 'workspace') + const acquired = createDeferred() + const release = createDeferred() + const deletion = db.transaction(async (tx) => { + await lockWorkspaceProject(tx, f.workspaceId) + await tx + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where(eq(workspaceFiles.id, seeded.fileId)) + .for('update') + expect(await deleteWorkspaceFileVersionInTx(tx, seeded.fileId, 11)).toMatchObject({ + status: 'deleted', + }) + const result = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`) + acquired.resolve(result[0].pid) + await release.promise + }) + const pid = await Promise.race([ + acquired.promise, + deletion.then(() => { + throw new Error('Deletion ended before acquiring its lock') + }), + ]) + const cleanup = cleanupFileVersions( + [{ entityType: 'workspace', entityId: f.workspaceId }], + { plan: 'free', cutoff: new Date(), label: 'history-rerank' }, + 100 + ) + const settled = Promise.allSettled([deletion, cleanup]) + try { + await expect + .poll( + async () => { + const rows = + await control`SELECT pid FROM pg_stat_activity WHERE ${pid} = ANY(pg_blocking_pids(pid)) AND wait_event_type = 'Lock'` + return rows.length + }, + { timeout: 3000 } + ) + .toBe(1) + } finally { + release.resolve() + await settled + } + expect((await settled).filter((result) => result.status === 'rejected')).toEqual([]) + expect(await cleanup).toBe(1) + expect((await versions(seeded.fileId)).map(({ version }) => version)).toEqual([ + 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, + ]) + const events = await cleanupEvents(f.workspaceId, 'workspace') + expect(events).toHaveLength(1) + expect(events[0].payload).toMatchObject({ key: seeded.keys[0] }) + } + ) + + check( + 'unsupported file owners reject the whole batch before any history or archive deletion', + async () => { + const f = await fixture() + const expired = await seedFile(f, 40 * 24, true) + const owners: FileOwner[] = [ + { entityType: 'project', entityId: f.projectId }, + { entityType: 'user', entityId: f.ownerId }, + ] + const options = { plan: 'free' as const, cutoff: new Date(), label: 'invalid-owner-batch' } + await expect(cleanupFileVersions(owners, options, 10)).rejects.toThrow( + 'File owner is unavailable' + ) + await expect(beginFileArchiveCleanup(owners, options)).rejects.toThrow( + 'File owner is unavailable' + ) + expect(await versions(expired.fileId)).toHaveLength(12) + expect( + await db + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where(eq(workspaceFiles.id, expired.fileId)) + ).toEqual([{ id: expired.fileId }]) + expect(await cleanupEvents(f.projectId)).toEqual([]) + const [usage] = await db + .select({ bytes: userStats.storageUsedBytes }) + .from(userStats) + .where(eq(userStats.userId, f.ownerId)) + expect(usage.bytes).toBe(15) + } + ) + + check( + 'a mixed queued batch retains each owner policy and leaves unselected workspace history alone', + async () => { + const f = await fixture('pro') + const other = await fixture() + const projectFile = await seedFile(f, 60 * 24) + const workspaceFile = await seedFile(f, 60 * 24, false, 'workspace') + const unrelated = await seedFile(other, 60 * 24, false, 'workspace') + await runCleanupFileVersions({ + ...payload(f.projectId), + projectIds: [f.projectId, f.projectId], + workspaceIds: [f.workspaceId, f.workspaceId], + }) + expect(await versions(projectFile.fileId)).toHaveLength(12) + expect(await versions(workspaceFile.fileId)).toHaveLength(10) + expect(await versions(unrelated.fileId)).toHaveLength(12) + expect(await cleanupEvents(f.projectId)).toEqual([]) + const events = await db + .select({ payload: outboxEvent.payload }) + .from(outboxEvent) + .where( + and( + eq(outboxEvent.eventType, 'workspace-file.storage.cleanup'), + sql`${outboxEvent.payload}->>'key' LIKE ${`workspace/${f.workspaceId}/%`}` + ) + ) + expect(events.map(({ payload }) => payload)).toEqual( + expect.arrayContaining(workspaceFile.keys.slice(0, 2).map((key) => ({ key }))) + ) + expect(events).toHaveLength(2) + } + ) + + check( + 'one version-deletion budget bounds a mixed owner batch without changing either current head', + async () => { + const f = await fixture() + const projectFile = await seedFile(f, 60 * 24) + const workspaceFile = await seedFile(f, 60 * 24, false, 'workspace') + const deleted = await cleanupFileVersions( + [ + { entityType: 'project', entityId: f.projectId }, + { entityType: 'workspace', entityId: f.workspaceId }, + ], + { plan: 'free', cutoff: new Date(), label: 'bounded-owner-batch' }, + 3 + ) + expect(deleted).toBe(3) + expect(await versions(projectFile.fileId)).toEqual( + Array.from({ length: 10 }, (_, index) => ({ version: index + 3 })) + ) + const remaining = await versions(workspaceFile.fileId) + expect(remaining).toHaveLength(11) + expect(remaining.at(-1)).toEqual({ version: 12 }) + } + ) + check( 'omits empty Projects and batches retained Projects with the same cleanup policy', async () => { @@ -356,6 +515,268 @@ describe('Project file retention follows the current payer in PostgreSQL', () => } ) + check( + 'workspace archive release preserves a restored head and durably deletes only expired objects', + async () => { + const f = await fixture() + const expired = await seedFile(f, 40 * 24, true, 'workspace') + const restored = await seedFile(f, 40 * 24, true, 'workspace') + const legacyKey = `workspace/${f.workspaceId}/${generateId()}.md` + await db.insert(workspaceFile).values({ + id: generateId(), + workspaceId: f.workspaceId, + name: 'legacy.md', + key: legacyKey, + size: 5, + type: 'text/markdown', + uploadedBy: f.ownerId, + deletedAt: new Date(Date.now() - 40 * 24 * HOUR), + }) + await db + .update(workspace) + .set({ storageUsedBytes: 10 }) + .where(eq(workspace.id, f.workspaceId)) + for (const key of [expired.keys[11], restored.keys[11], legacyKey]) { + storedKeys.push(key) + await uploadFile({ + file: Buffer.from('draft'), + fileName: 'note.md', + contentType: 'text/markdown', + customKey: key, + preserveKey: true, + context: 'workspace', + persistMetadata: false, + }) + } + const cleanup = await beginFileArchiveCleanup( + [{ entityType: 'workspace', entityId: f.workspaceId }], + { plan: 'free', cutoff: new Date(), label: 'restore-selected-workspace-file' } + ) + await db + .update(workspaceFiles) + .set({ deletedAt: null }) + .where(eq(workspaceFiles.id, restored.fileId)) + const deletion = await cleanup.cleanupStorage() + expect(await deletion.deleteRows()).toBe(2) + const events = await cleanupEvents(f.workspaceId, 'workspace') + expect(new Set(events.map(({ payload }) => (payload as { key: string }).key))).toEqual( + new Set([...expired.keys, legacyKey]) + ) + expect( + events.find(({ payload }) => (payload as { key: string }).key === expired.keys[11])?.payload + ).toEqual({ key: expired.keys[11] }) + await processWorkspaceFileStorageCleanupsNow( + events.map(({ id }) => id), + { label: 'retention-fixture' } + ) + expect(await downloadFile({ key: restored.keys[11], context: 'workspace' })).toEqual( + Buffer.from('draft') + ) + await expect( + downloadFile({ key: expired.keys[11], context: 'workspace' }) + ).rejects.toMatchObject({ code: 'ENOENT' }) + await expect(downloadFile({ key: legacyKey, context: 'workspace' })).rejects.toMatchObject({ + code: 'ENOENT', + }) + expect(await versions(restored.fileId)).toHaveLength(12) + const [usage] = await db + .select({ bytes: userStats.storageUsedBytes }) + .from(userStats) + .where(eq(userStats.userId, f.ownerId)) + expect(usage.bytes).toBe(10) + } + ) + + check( + 'folder retention deduplicates more than one file batch before deleting its parent', + async () => { + const f = await fixture() + const folderId = generateId() + await db.insert(folder).values({ + id: folderId, + projectId: f.projectId, + userId: f.ownerId, + resourceType: 'file', + name: 'Expired parent', + deletedAt: new Date(Date.now() - 40 * 24 * HOUR), + }) + const files = Array.from({ length: 502 }, (_, index) => ({ + id: generateId(), + projectId: f.projectId, + userId: f.ownerId, + context: 'project', + key: `project/${f.projectId}/batch-${index}`, + originalName: `batch-${index}.md`, + contentType: 'text/markdown', + sizeBytes: 0, + folderId, + })) + await db.insert(workspaceFiles).values(files) + await db.insert(workspaceFiles).values({ + ...files[0], + id: generateId(), + key: `project/${f.projectId}/root`, + folderId: null, + }) + await runCleanupSoftDeletes(payload(f.projectId), createCleanupBudgets({ folders: 1 })) + expect(await db.select().from(folder).where(eq(folder.id, folderId))).toEqual([]) + const retained = await db + .select() + .from(workspaceFiles) + .where(eq(workspaceFiles.projectId, f.projectId)) + expect(retained).toHaveLength(503) + expect(retained.every((file) => file.folderId === null)).toBe(true) + expect(new Set(retained.map((file) => file.originalName)).size).toBe(503) + expect(retained.find((file) => file.id === files[0].id)?.originalName).not.toBe('batch-0.md') + } + ) + + check('workspace file folder deletion failure rolls back child locations and names', async () => { + const f = await fixture() + const parentId = generateId() + const childId = generateId() + const owner = { workspaceId: f.workspaceId, userId: f.ownerId, resourceType: 'file' as const } + await db.insert(folder).values({ + ...owner, + id: parentId, + name: 'Archived parent', + deletedAt: new Date(Date.now() - 40 * 24 * HOUR), + }) + await db.insert(folder).values({ ...owner, id: childId, name: 'Child', parentId }) + const file = await seedFile(f, 1, false, 'workspace') + await db + .update(workspaceFiles) + .set({ folderId: parentId, originalName: 'Child.md' }) + .where(eq(workspaceFiles.id, file.fileId)) + const trigger = `folder_retention_failure_${generateId().replaceAll('-', '')}` + await db.execute( + sql.raw(`CREATE FUNCTION ${trigger}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN + IF OLD.id = '${parentId}' THEN RAISE EXCEPTION 'folder retention rollback fixture'; END IF; RETURN OLD; END $$`) + ) + await db.execute( + sql.raw( + `CREATE TRIGGER ${trigger} BEFORE DELETE ON folder FOR EACH ROW EXECUTE FUNCTION ${trigger}()` + ) + ) + try { + await expect( + runCleanupSoftDeletes( + { + workspaceIds: [f.workspaceId], + plan: 'free', + retentionHours: 30 * 24, + label: 'workspace-folder-rollback', + }, + createCleanupBudgets({ folders: 1 }) + ) + ).rejects.toMatchObject({ cause: { message: 'folder retention rollback fixture' } }) + const [child] = await db + .select({ parentId: folder.parentId, name: folder.name }) + .from(folder) + .where(eq(folder.id, childId)) + expect(child).toEqual({ parentId, name: 'Child' }) + const [head] = await db + .select({ folderId: workspaceFiles.folderId, name: workspaceFiles.originalName }) + .from(workspaceFiles) + .where(eq(workspaceFiles.id, file.fileId)) + expect(head).toEqual({ folderId: parentId, name: 'Child.md' }) + } finally { + await db.execute(sql.raw(`DROP TRIGGER ${trigger} ON folder`)) + await db.execute(sql.raw(`DROP FUNCTION ${trigger}()`)) + } + }) + + check( + 'workspace file folder purge waits for a concurrent restore and leaves its children in place', + async () => { + const f = await fixture() + const parentId = generateId() + const childId = generateId() + const owner = { workspaceId: f.workspaceId, userId: f.ownerId, resourceType: 'file' as const } + await db.insert(folder).values({ + ...owner, + id: parentId, + name: 'Restoring parent', + deletedAt: new Date(Date.now() - 40 * 24 * HOUR), + }) + await db.insert(folder).values({ ...owner, id: childId, name: 'Child', parentId }) + const acquired = createDeferred() + const restore = createDeferred() + const authority = db.transaction(async (tx) => { + await authorizeProject( + tx, + createSessionPrincipal({ userId: f.ownerId }), + projectOperations.get, + { projectId: f.projectId }, + 'hold' + ) + await tx + .select({ id: workspace.id }) + .from(workspace) + .where(eq(workspace.id, f.workspaceId)) + .for('share') + const result = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`) + acquired.resolve(result[0].pid) + await restore.promise + await tx.update(folder).set({ deletedAt: null }).where(eq(folder.id, parentId)) + }) + const authPid = await Promise.race([ + acquired.promise, + authority.then(() => { + throw new Error('Restore ended before publishing its lock') + }), + ]) + let settled = false + const cleanup = runCleanupSoftDeletes( + { + workspaceIds: [f.workspaceId], + plan: 'free', + retentionHours: 30 * 24, + label: 'workspace-folder-restore', + }, + createCleanupBudgets({ folders: 1 }) + ).then( + () => { + settled = true + return null + }, + (error: unknown) => { + settled = true + return error + } + ) + let observed: string | null = null + try { + for (let attempt = 0; attempt < 100 && !settled; attempt++) { + const [waiting] = await control< + { wait_event: string }[] + >`SELECT wait_event FROM pg_stat_activity WHERE ${authPid} = ANY(pg_blocking_pids(pid)) AND wait_event_type = 'Lock'` + if (waiting) { + observed = waiting.wait_event + break + } + await sleep(10) + } + } finally { + restore.resolve() + await authority + await cleanup + } + expect(observed).toBe('advisory') + expect(await cleanup).toBeNull() + const [parent] = await db + .select({ deletedAt: folder.deletedAt }) + .from(folder) + .where(eq(folder.id, parentId)) + expect(parent).toEqual({ deletedAt: null }) + const [child] = await db + .select({ parentId: folder.parentId, name: folder.name }) + .from(folder) + .where(eq(folder.id, childId)) + expect(child).toEqual({ parentId, name: 'Child' }) + } + ) + check( 'free history is pruned with its ten-version floor while a paid Project ignores a stale free queue policy and uploader plan', async () => { @@ -485,42 +906,66 @@ describe('Project file retention follows the current payer in PostgreSQL', () => } ) - check( - 'a failed archive deletion rolls back current bytes, version removal, and the cleanup outbox', - async () => { - const f = await fixture() - const expired = await seedFile(f, 40 * 24, true) - const trigger = `retention_failure_${generateId().replaceAll('-', '')}` - await db.execute( - sql.raw(`CREATE FUNCTION ${trigger}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN + for (const kind of ['project', 'workspace'] as const) + check( + `a failed ${kind} archive deletion rolls back current bytes, version removal, and the cleanup outbox`, + async () => { + const f = await fixture() + const expired = await seedFile(f, 40 * 24, true, kind) + if (kind === 'workspace') + await db + .update(workspace) + .set({ storageUsedBytes: 5 }) + .where(eq(workspace.id, f.workspaceId)) + const trigger = `retention_failure_${generateId().replaceAll('-', '')}` + await db.execute( + sql.raw(`CREATE FUNCTION ${trigger}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN IF OLD.id = '${expired.fileId}' THEN RAISE EXCEPTION 'retention rollback fixture'; END IF; RETURN OLD; END $$`) - ) - await db.execute( - sql.raw( - `CREATE TRIGGER ${trigger} BEFORE DELETE ON workspace_files FOR EACH ROW EXECUTE FUNCTION ${trigger}()` ) - ) - try { - await expect( - runCleanupSoftDeletes(payload(f.projectId), createCleanupBudgets({ files: 1 })) - ).rejects.toMatchObject({ cause: { message: 'retention rollback fixture' } }) - expect(await versions(expired.fileId)).toHaveLength(12) - expect(await cleanupEvents(f.projectId)).toEqual([]) - const [usage] = await db - .select({ bytes: userStats.storageUsedBytes }) - .from(userStats) - .where(eq(userStats.userId, f.ownerId)) - expect(usage.bytes).toBe(15) - } finally { - await db.execute(sql.raw(`DROP TRIGGER ${trigger} ON workspace_files`)) - await db.execute(sql.raw(`DROP FUNCTION ${trigger}()`)) + await db.execute( + sql.raw( + `CREATE TRIGGER ${trigger} BEFORE DELETE ON workspace_files FOR EACH ROW EXECUTE FUNCTION ${trigger}()` + ) + ) + try { + await expect( + runCleanupSoftDeletes( + kind === 'project' + ? payload(f.projectId) + : { + workspaceIds: [f.workspaceId], + plan: 'free', + retentionHours: 30 * 24, + label: 'workspace-rollback', + }, + createCleanupBudgets({ files: 1 }) + ) + ).rejects.toMatchObject( + kind === 'project' + ? { cause: { message: 'retention rollback fixture' } } + : { message: 'File row cleanup failed' } + ) + expect(await versions(expired.fileId)).toHaveLength(12) + expect( + await cleanupEvents(kind === 'project' ? f.projectId : f.workspaceId, kind) + ).toEqual([]) + const [usage] = await db + .select({ bytes: userStats.storageUsedBytes }) + .from(userStats) + .where(eq(userStats.userId, f.ownerId)) + expect(usage.bytes).toBe(15) + } finally { + await db.execute(sql.raw(`DROP TRIGGER ${trigger} ON workspace_files`)) + await db.execute(sql.raw(`DROP FUNCTION ${trigger}()`)) + } } - } - ) + ) }) afterAll(async () => { try { + for (const key of storedKeys) + await deleteFile({ key, context: 'workspace' }).catch(() => undefined) for (const f of fixtures) { await db .delete(outboxEvent) diff --git a/apps/sim/lib/projects/files/retention.ts b/apps/sim/lib/projects/files/retention.ts index 865a4dd2fc8..fa7299ab29e 100644 --- a/apps/sim/lib/projects/files/retention.ts +++ b/apps/sim/lib/projects/files/retention.ts @@ -1,10 +1,10 @@ import { dbFor } from '@sim/db' -import { folder, organization, project, workspaceFiles, workspaceFileVersion } from '@sim/db/schema' -import { and, asc, count, eq, gt, inArray, isNotNull, isNull, lt, min, or, sql } from 'drizzle-orm' +import { organization, project, workspaceFiles, workspaceFileVersion } from '@sim/db/schema' +import { and, asc, count, eq, gt, inArray, isNotNull, lt, min, or, sql } from 'drizzle-orm' import { CLEANUP_CONFIG } from '@/lib/billing/cleanup-dispatcher' import { getPlanType } from '@/lib/billing/plan-helpers' import { resolveProjectStorageBillingContext } from '@/lib/billing/storage/context' -import { prepareProjectStorageMutationInTx } from '@/lib/billing/storage/tracking' +import { prepareFileStorageMutationInTx } from '@/lib/billing/storage/tracking' import { consumeRowBudget, DEFAULT_DELETE_CHUNK_SIZE, @@ -12,22 +12,24 @@ import { type RowBudget, } from '@/lib/cleanup/batch-delete' import { isBillingEnabled } from '@/lib/core/config/env-flags' -import { generateRestoreName } from '@/lib/core/utils/restore-name' import type { DbTransaction } from '@/lib/db/types' -import { acquireFolderMutationLock } from '@/lib/folders/locks' -import { deduplicateFolderNameInScope } from '@/lib/folders/naming' +import { cleanupExpiredFileFolderInTx } from '@/lib/file-retention/folders' +import { + FILES_PER_QUERY, + KEEP_SUPERSEDED, + releaseExpiredFileVersions, + selectExpiredFileVersions, +} from '@/lib/file-retention/versions' import { lockProject } from '@/lib/projects/membership' -import { workspaceFileNameFolderCondition } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' import { enqueueWorkspaceFileStorageCleanups } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' import { MAX_SUPERSEDED_FILE_VERSIONS, releaseWorkspaceFileVersionsForPurgeInTx, } from '@/lib/uploads/contexts/workspace/workspace-file-versions' -import { fileFolderOwnerCondition, fileOwnerCondition } from '@/lib/workspace-files/ownership-query' +import { lockFileDirectories } from '@/lib/workspace-files/locks' +import { fileOwnerCondition } from '@/lib/workspace-files/ownership-query' const cleanupDb = dbFor('cleanup') -const KEEP_SUPERSEDED = 9 -const FILES_PER_QUERY = 500 type FileRetentionJob = 'cleanup-file-versions' | 'cleanup-soft-deletes' @@ -41,7 +43,7 @@ async function prepareRetention(tx: DbTransaction, projectId: string, job: FileR tx ) if (isBillingEnabled && owner.organizationId && billing.plan === null) return null - const mutation = await prepareProjectStorageMutationInTx(tx, billing) + const mutation = await prepareFileStorageMutationInTx(tx, billing) const plan = isBillingEnabled ? getPlanType(billing.plan) : 'enterprise' const config = CLEANUP_CONFIG[job] let hours: number | null = config.defaults[plan] @@ -112,58 +114,18 @@ export async function cleanupProjectFileVersions( ) .orderBy(asc(workspaceFiles.id)) .for('update') - const ranked = tx - .select({ - id: workspaceFileVersion.id, - supersededAt: workspaceFileVersion.supersededAt, - rank: sql`row_number() over (partition by ${workspaceFileVersion.fileId} order by ${workspaceFileVersion.version} desc)`.as( - 'rank' - ), - }) - .from(workspaceFileVersion) - .where( - and( - inArray( - workspaceFileVersion.fileId, - candidates.map((row) => row.id) - ), - isNotNull(workspaceFileVersion.supersededAt) - ) - ) - .as('ranked') const versionLimit = Math.min(DEFAULT_DELETE_CHUNK_SIZE, limit - released) - const expired = await tx - .select({ id: ranked.id }) - .from(ranked) - .where( - and( - gt(ranked.rank, KEEP_SUPERSEDED), - or(lt(ranked.supersededAt, policy.cutoff), gt(ranked.rank, maximum)) - ) - ) - .limit(versionLimit) - const removed = expired.length - ? await tx - .delete(workspaceFileVersion) - .where( - and( - inArray( - workspaceFileVersion.id, - expired.map((row) => row.id) - ), - isNotNull(workspaceFileVersion.supersededAt) - ) - ) - .returning({ key: workspaceFileVersion.key }) - : [] - await enqueueWorkspaceFileStorageCleanups( + const expired = await selectExpiredFileVersions( tx, - removed.map((row) => row.key), - 'project' + candidates.map(({ id }) => id), + policy.cutoff, + maximum, + versionLimit ) + const removed = await releaseExpiredFileVersions(tx, expired, 'project') return { - lastId: removed.length === versionLimit ? afterId : candidates[candidates.length - 1].id, - removed: removed.length, + lastId: removed === versionLimit ? afterId : candidates[candidates.length - 1].id, + removed, } }) if (!batch) break @@ -183,7 +145,7 @@ export async function cleanupArchivedProjectFiles( const removed = await cleanupDb.transaction(async (tx) => { const policy = await prepareRetention(tx, projectId, 'cleanup-soft-deletes') if (!policy) return 0 - await acquireFolderMutationLock(tx, `project:${projectId}`, 'file') + await lockFileDirectories(tx, [{ entityType: 'project', entityId: projectId }]) const rows = await tx .select({ id: workspaceFiles.id, @@ -235,88 +197,22 @@ export async function cleanupArchivedProjectFiles( return total } -/** Surviving children are re-rooted under the directory lock before an expired folder is removed. */ +/** Surviving children are re-rooted atomically under current Project retention policy. */ export async function cleanupArchivedProjectFileFolders( projectId: string, budget?: RowBudget ): Promise { let total = 0 - const owner = { entityType: 'project' as const, entityId: projectId } - const folderScope = fileFolderOwnerCondition(owner) for (let batch = 0; batch < DEFAULT_MAX_BATCHES_PER_TABLE && budget?.remaining !== 0; batch++) { const removed = await cleanupDb.transaction(async (tx) => { const policy = await prepareRetention(tx, projectId, 'cleanup-soft-deletes') if (!policy) return 0 - await acquireFolderMutationLock(tx, `project:${projectId}`, 'file') - const [expired] = await tx - .select({ id: folder.id }) - .from(folder) - .where(and(folderScope, isNotNull(folder.deletedAt), lt(folder.deletedAt, policy.cutoff))) - .orderBy(asc(folder.id)) - .limit(1) - .for('update') - if (!expired) return 0 - consumeRowBudget(budget, 1) - for (;;) { - const children = await tx - .select({ id: folder.id, name: folder.name }) - .from(folder) - .where(and(folderScope, eq(folder.parentId, expired.id), isNull(folder.deletedAt))) - .orderBy(asc(folder.id)) - .limit(FILES_PER_QUERY) - .for('update') - if (!children.length) break - for (const child of children) { - const name = await deduplicateFolderNameInScope(tx, folderScope, null, child.name) - await tx - .update(folder) - .set({ parentId: null, name, updatedAt: new Date() }) - .where(and(folderScope, eq(folder.id, child.id))) - } - } - for (;;) { - const children = await tx - .select({ id: workspaceFiles.id, name: workspaceFiles.originalName }) - .from(workspaceFiles) - .where( - and( - projectFiles(projectId), - eq(workspaceFiles.folderId, expired.id), - isNull(workspaceFiles.deletedAt) - ) - ) - .orderBy(asc(workspaceFiles.id)) - .limit(FILES_PER_QUERY) - .for('update') - if (!children.length) break - for (const child of children) { - const originalName = await generateRestoreName( - child.name, - async (name) => { - const [existing] = await tx - .select({ id: workspaceFiles.id }) - .from(workspaceFiles) - .where( - and( - projectFiles(projectId), - eq(workspaceFiles.originalName, name), - workspaceFileNameFolderCondition(null), - isNull(workspaceFiles.deletedAt) - ) - ) - .limit(1) - return Boolean(existing) - }, - { hasExtension: true } - ) - await tx - .update(workspaceFiles) - .set({ folderId: null, originalName, updatedAt: new Date() }) - .where(and(projectFiles(projectId), eq(workspaceFiles.id, child.id))) - } - } - await tx.delete(folder).where(and(folderScope, eq(folder.id, expired.id))) - return 1 + return cleanupExpiredFileFolderInTx( + tx, + { entityType: 'project', entityId: projectId }, + policy.cutoff, + budget + ) }) total += removed if (!removed) break diff --git a/apps/sim/lib/uploads/contexts/workspace/__integration__/file-versions.integration.ts b/apps/sim/lib/uploads/contexts/workspace/__integration__/file-versions.integration.ts index 1cc36f7b5a5..53160ab84fd 100644 --- a/apps/sim/lib/uploads/contexts/workspace/__integration__/file-versions.integration.ts +++ b/apps/sim/lib/uploads/contexts/workspace/__integration__/file-versions.integration.ts @@ -210,7 +210,7 @@ describe('workspace file version history in PostgreSQL', () => { phase === 'storage' ? null : vi - .spyOn(storageBilling, 'incrementStorageUsageForBillingContextInTx') + .spyOn(storageBilling, 'prepareFileAccountingInTx') .mockRejectedValueOnce(operationError) try { const operation = diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-folder-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-folder-manager.ts index 9d9ee52478b..07b38663228 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-folder-manager.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-folder-manager.ts @@ -582,7 +582,7 @@ async function getFileFolder( } /** Resolves one active folder by ID or path within the authorized owner. */ -async function resolveFileFolderTarget( +export async function resolveFileFolderTarget( owner: O, target: { folderId?: string | null; folderPath?: string }, client: DbOrTx = db diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts index 1487c43f5ce..92cf3008363 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts @@ -20,6 +20,7 @@ import { getErrorMessage, getPostgresConstraintName, getPostgresErrorCode, + isPostgresCommitRejection, } from '@sim/utils/errors' import { generateShortId } from '@sim/utils/id' import { omit } from '@sim/utils/object' @@ -41,11 +42,9 @@ import { timestampKey, } from '@/lib/api/list-query' import { - decrementStorageUsageForBillingContextInTx, - incrementStorageUsageForBillingContextInTx, lockWorkspaceStorageForMutationInTx, maybeNotifyStorageLimitForBillingContext, - resolveStorageBillingContext, + prepareFileAccountingInTx, } from '@/lib/billing/storage' import { CollabDocStateConflictError, @@ -66,10 +65,7 @@ import { lockWorkspaceProject } from '@/lib/projects/membership' import { notifyWorkspaceFilesChanged } from '@/lib/realtime/notify' import { getServePathPrefix } from '@/lib/uploads' import type { WorkspaceFileFolderRecord } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' -import { - enqueueWorkspaceFileLiveDocReconciliation, - processWorkspaceFileLiveDocReconciliationNow, -} from '@/lib/uploads/contexts/workspace/workspace-file-live-doc-outbox' +import { enqueueWorkspaceFileLiveDocReconciliation } from '@/lib/uploads/contexts/workspace/workspace-file-live-doc-outbox' import { applyWorkspaceFileSecretProvenancePolicyInTx, EXACT_EMPTY_WORKSPACE_FILE_SECRET_PROVENANCE, @@ -101,6 +97,7 @@ import { headObject, uploadFile, } from '@/lib/uploads/core/storage-service' +import { finishFileContentEffects } from '@/lib/uploads/server/content-effects' import { getWorkspaceFileSize, MAX_WORKSPACE_FILE_SIZE } from '@/lib/uploads/shared/types' import { isMarkdownFile } from '@/lib/uploads/utils/file-utils' import type { ServableFile } from '@/lib/uploads/utils/file-utils.server' @@ -120,7 +117,7 @@ import { import { type EditableFileOwner, editableFileOwnerColumns, - type FileOwner, + matchesFileOwner, resolveFileOwner, } from '@/lib/workspace-files/ownership' import { fileFolderOwnerCondition, fileOwnerCondition } from '@/lib/workspace-files/ownership-query' @@ -142,6 +139,7 @@ import { listFileFolders, listWorkspaceFileFolders, normalizeWorkspaceFileItemName, + resolveFileFolderTarget, resolveWorkspaceFileFolderTarget, workspaceFileNameFolderCondition, } from './workspace-file-folder-manager' @@ -489,6 +487,20 @@ interface StagedFileContent { } const stagedFileContents = new WeakSet() + +interface PlannedFileIdentity { + readonly id: string +} + +const plannedFileOwners = new WeakMap>() + +/** Reserves an in-process destination identity so a compound copy can rewrite selected references. */ +function planFileIdentity(owner: EditableFileOwner): PlannedFileIdentity { + const identity = Object.freeze({ id: `wf_${generateShortId()}` }) + plannedFileOwners.set(identity, Object.freeze({ ...owner })) + return identity +} + /** Writes private bytes before the owner-authorized metadata transaction begins. */ async function stageFileContent(args: { owner: EditableFileOwner @@ -497,7 +509,9 @@ async function stageFileContent(args: { contentType: string content: Buffer folderId?: string | null + signal?: AbortSignal }): Promise { + args.signal?.throwIfAborted() const name = normalizeWorkspaceFileItemName(args.name, 'File') const owner = Object.freeze({ ...args.owner }) const key = @@ -532,6 +546,7 @@ async function stageFileContent(args: { ...(args.folderId ? { folderId: args.folderId } : {}), }, persistMetadata: false, + signal: args.signal, }) if (uploaded.key !== key) throw new Error('Storage returned an unexpected file key') } catch (error) { @@ -581,30 +596,6 @@ async function discardStagedFileContent(staged: StagedFileContent): Promise( staged: StagedFileContent, @@ -634,6 +625,81 @@ async function finalizeStagedFileContent( } } +/** Commits new canonical metadata and provenance under the owner's directory lock. */ +async function commitFileCreateInTx( + tx: DbTransaction, + args: { + owner: EditableFileOwner + staged: StagedFileContent + userId: string + folderId?: string | null + folderPath?: string + exactName?: boolean + identity?: PlannedFileIdentity + secretProvenance?: WorkspaceFileSecretProvenance + } +): Promise { + assertStagedFileOwner(args.owner, args.staged) + if (args.identity) { + const plannedOwner = plannedFileOwners.get(args.identity) + if ( + plannedOwner?.entityType !== args.owner.entityType || + plannedOwner.entityId !== args.owner.entityId + ) + throw new Error('Planned file identity does not belong to this owner') + } + if (args.owner.entityType === 'workspace') await lockWorkspaceProject(tx, args.owner.entityId) + await lockFileDirectories(tx, [args.owner]) + const target = await resolveFileFolderTarget(args.owner, args, tx) + const folderId = target?.id ?? null + const exists = async (name: string) => { + const [row] = await tx + .select({ id: workspaceFiles.id }) + .from(workspaceFiles) + .where( + and( + fileOwnerCondition(args.owner), + eq(workspaceFiles.originalName, name), + workspaceFileNameFolderCondition(folderId), + isNull(workspaceFiles.deletedAt) + ) + ) + .limit(1) + return Boolean(row) + } + let name = args.staged.name + if (await exists(name)) { + if (args.exactName) throw new FileConflictError(name) + let suffix = 1 + do { + name = withCopySuffix( + args.staged.name, + suffix <= MAX_NUMBERED_COPY_SUFFIX ? suffix : generateShortId(8) + ) + suffix += 1 + } while (suffix <= MAX_NUMBERED_COPY_SUFFIX + 2 && (await exists(name))) + } + const inserted = await insertFileMetadataInTx(tx, args.owner, { + id: args.identity?.id ?? `wf_${generateShortId()}`, + key: args.staged.key, + userId: args.userId, + folderId, + originalName: name, + contentType: args.staged.contentType, + size: args.staged.size, + }) + if (!inserted) throw new FileConflictError(name) + if (args.secretProvenance) { + await replaceWorkspaceFileSecretProvenanceInTx( + tx, + inserted.id, + inserted.contentUpdatedAt, + args.secretProvenance + ) + } + return inserted +} + /** * Upload a file to workspace-scoped storage */ @@ -679,7 +745,6 @@ export async function uploadWorkspaceFile( const effectiveName = pageRestore?.name ?? normalizedFileName const effectiveContentType = pageRestore ? SIM_PAGE_CONTENT_TYPE : contentType const exactName = options?.exactName ?? false - const storageBillingContext = await resolveStorageBillingContext(workspaceId) let lastError: unknown const maxAttempts = exactName ? 1 : MAX_UPLOAD_UNIQUE_RETRIES @@ -690,7 +755,7 @@ export async function uploadWorkspaceFile( if (exactName && (await fileExistsInWorkspace(workspaceId, uniqueName, folderId))) { throw new FileConflictError(uniqueName) } - const fileId = `wf_${generateShortId()}` + const identity = planFileIdentity({ entityType: 'workspace', entityId: workspaceId }) try { const uploadResult = await stageFileContent({ @@ -703,49 +768,34 @@ export async function uploadWorkspaceFile( }) const finalized = await finalizeStagedFileContent(uploadResult, async (tx) => { - await lockWorkspaceProject(tx, workspaceId) - await lockFileDirectories(tx, [{ entityType: 'workspace', entityId: workspaceId }]) - let activeFolderId: string | null + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: workspaceId, + }) + let activeFolderId = folderId if (options?.folderPath !== undefined) { + await lockFileDirectories(tx, [{ entityType: 'workspace', entityId: workspaceId }]) const folderIndex = await loadActiveFolderPathIndex(workspaceId, 'file', tx) const resolvedFolderId = resolveFolderPathFromIndex(folderIndex, options.folderPath) if (resolvedFolderId === undefined) { throw new OrchestrationError('not_found', 'Target folder not found') } activeFolderId = resolvedFolderId - } else { - activeFolderId = await assertWorkspaceFileFolderTarget(workspaceId, folderId, tx) } - const inserted = await insertWorkspaceFileMetadataInTx(tx, { - id: fileId, - key: uploadResult.key, + const inserted = await commitFileCreateInTx(tx, { + owner: { entityType: 'workspace', entityId: workspaceId }, + staged: uploadResult, + identity, userId, - workspaceId, folderId: activeFolderId, - originalName: uniqueName, - contentType: effectiveContentType, - size: effectiveBuffer.length, + exactName: true, + secretProvenance: options?.secretProvenance, }) - if (!inserted) { - throw new FileConflictError(uniqueName) - } - if (options?.secretProvenance) { - await replaceWorkspaceFileSecretProvenanceInTx( - tx, - inserted.id, - inserted.contentUpdatedAt, - options.secretProvenance - ) - } - const usage = await incrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - effectiveBuffer.length - ) - return { inserted, updatedUsage: usage } + const usage = await accounting.mutation.applyDelta(effectiveBuffer.length) + return { inserted, updatedUsage: usage, billing: accounting.billing } }) - void maybeNotifyStorageLimitForBillingContext(storageBillingContext, finalized.updatedUsage) + void maybeNotifyStorageLimitForBillingContext(finalized.billing, finalized.updatedUsage) logger.info( `Successfully uploaded workspace file: ${uniqueName} with key: ${uploadResult.key}` @@ -812,7 +862,6 @@ export async function createOwnedWorkspaceFile(params: { const fileId = `wf_${generateShortId()}` const name = params.fileName(fileId) const buffer = Buffer.from(params.content, 'utf-8') - const storageBillingContext = await resolveStorageBillingContext(params.workspaceId) const staged = await stageFileContent({ owner: { entityType: 'workspace', entityId: params.workspaceId }, userId: params.userId, @@ -820,7 +869,11 @@ export async function createOwnedWorkspaceFile(params: { contentType: params.contentType, content: buffer, }) - const { owner, updatedUsage } = await finalizeStagedFileContent(staged, async (tx) => { + const { owner, updatedUsage, billing } = await finalizeStagedFileContent(staged, async (tx) => { + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: params.workspaceId, + }) const inserted = await insertWorkspaceFileMetadataInTx(tx, { id: fileId, key: staged.key, @@ -833,14 +886,14 @@ export async function createOwnedWorkspaceFile(params: { context: params.context, }) if (!inserted) throw new Error(`Owned ${params.context} file ${fileId} was not inserted`) - const updatedUsage = await incrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - buffer.length - ) - return { owner: await params.insertOwner(tx, fileId), updatedUsage } + const updatedUsage = await accounting.mutation.applyDelta(buffer.length) + return { + owner: await params.insertOwner(tx, fileId), + updatedUsage, + billing: accounting.billing, + } }) - void maybeNotifyStorageLimitForBillingContext(storageBillingContext, updatedUsage) + void maybeNotifyStorageLimitForBillingContext(billing, updatedUsage) return { fileId, owner } } @@ -1012,13 +1065,15 @@ export async function registerUploadedWorkspaceFile(params: { const folderId = params.folderId ?? null - const storageBillingContext = await resolveStorageBillingContext(workspaceId) for (let attempt = 0; attempt < MAX_UPLOAD_UNIQUE_RETRIES; attempt++) { const fileId = `wf_${generateShortId()}` const displayName = await allocateUniqueWorkspaceFileName(workspaceId, effectiveName, folderId) const finalized = await db.transaction(async (tx) => { - await lockWorkspaceProject(tx, workspaceId) + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: workspaceId, + }) await lockFileDirectories(tx, [{ entityType: 'workspace', entityId: workspaceId }]) const activeFolderId = await assertWorkspaceFileFolderTarget(workspaceId, folderId, tx) const inserted = await insertWorkspaceFileMetadataInTx(tx, { @@ -1055,11 +1110,7 @@ export async function registerUploadedWorkspaceFile(params: { return { kind: 'existing', file: raceWinner } as const } - const updatedUsage = await incrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - effectiveSize - ) + const updatedUsage = await accounting.mutation.applyDelta(effectiveSize) await replaceWorkspaceFileSecretProvenanceInTx( tx, inserted.id, @@ -1067,7 +1118,7 @@ export async function registerUploadedWorkspaceFile(params: { secretProvenance ) await markUploadSessionFileRegistered(tx, params.uploadSessionId, workspaceId, inserted.id) - return { kind: 'created', file: inserted, updatedUsage } as const + return { kind: 'created', file: inserted, updatedUsage, billing: accounting.billing } as const }) if (finalized.kind === 'name-conflict') { @@ -1078,7 +1129,7 @@ export async function registerUploadedWorkspaceFile(params: { } if (finalized.kind === 'created') { - void maybeNotifyStorageLimitForBillingContext(storageBillingContext, finalized.updatedUsage) + void maybeNotifyStorageLimitForBillingContext(finalized.billing, finalized.updatedUsage) } await commitPageRestoreRewrite() @@ -1694,6 +1745,7 @@ const WORKSPACE_FILE_SORTS = { } satisfies Record[]> export interface QueryWorkspaceFilesOptions { + fileIds?: readonly string[] scope?: WorkspaceFileScope /** Restrict to one file folder. */ /** `undefined` lists every folder, `null` lists only root files. */ @@ -1801,6 +1853,7 @@ async function queryFileRecords( const conditions = [ fileOwnerCondition(owner), + options.fileIds ? inArray(workspaceFiles.id, [...options.fileIds]) : undefined, scope === 'all' ? undefined : scope === 'archived' @@ -2073,15 +2126,13 @@ async function listWorkspaceFileReferenceCandidates( return hydrateWorkspaceFilePaths(rows, workspaceId) } -/** - * Load the canonical authorization context for an active workspace file by resource ID. - * Database failures propagate so callers never confuse unavailable state with a missing file. - * Chat uploads are admitted only on explicit opt-in (see {@link WorkspaceFileLookupOptions}). - */ -export async function loadActiveWorkspaceFileContext( +async function loadWorkspaceFileContext( fileId: string, - options?: WorkspaceFileLookupOptions & { includeDeleted?: boolean } -): Promise { + options: WorkspaceFileLookupOptions & { + includeDeleted?: boolean + includeArchivedWorkspace?: boolean + } +): Promise { const [context] = await db .select({ fileId: workspaceFiles.id, @@ -2090,6 +2141,7 @@ export async function loadActiveWorkspaceFileContext( workspaceOrganizationId: workspace.organizationId, allowPersonalApiKeys: workspace.allowPersonalApiKeys, billedAccountUserId: workspace.billedAccountUserId, + deletedAt: workspaceFiles.deletedAt, ownership: { projectId: workspaceFiles.projectId, context: workspaceFiles.context, @@ -2105,63 +2157,42 @@ export async function loadActiveWorkspaceFileContext( and( eq(workspaceFiles.id, fileId), workspaceFileContextCondition(options), - ...(options?.includeDeleted ? [] : [isNull(workspaceFiles.deletedAt)]), - isNull(workspace.archivedAt) + options.includeDeleted ? undefined : isNull(workspaceFiles.deletedAt), + options.includeArchivedWorkspace ? undefined : isNull(workspace.archivedAt) ) ) .limit(1) if ( !context || - !matchesWorkspaceFileOwner(resolveFileOwner(context.ownership), context.workspaceId) + !matchesFileOwner(resolveFileOwner(context.ownership), { + entityType: 'workspace', + entityId: context.workspaceId, + }) ) { return null } return omit(context, ['ownership']) } -function matchesWorkspaceFileOwner(owner: FileOwner | null, workspaceId: string): boolean { - return owner?.entityType === 'workspace' && owner.entityId === workspaceId +/** Loads canonical active-workspace identity; chat uploads and deleted files require explicit opt-in. */ +export async function loadActiveWorkspaceFileContext( + fileId: string, + options?: WorkspaceFileLookupOptions & { includeDeleted?: boolean } +): Promise { + const context = await loadWorkspaceFileContext(fileId, { + includeDeleted: options?.includeDeleted, + includeChatUploads: options?.includeChatUploads, + includeOwnedFiles: options?.includeOwnedFiles, + }) + return context ? omit(context, ['deletedAt']) : null } -/** - * Load a workspace file for a lifecycle transition, including archived files. - * The workspace archive state is returned by the canonical workspace record and is enforced by - * the operation's manager primitive where the transition requires an active workspace. - */ +/** Loads workspace files for lifecycle transitions; the manager enforces the operation's archive policy. */ export async function loadWorkspaceFileLifecycleContext( fileId: string ): Promise { - const [context] = await db - .select({ - fileId: workspaceFiles.id, - fileContext: workspaceFiles.context, - workspaceId: workspace.id, - workspaceOrganizationId: workspace.organizationId, - allowPersonalApiKeys: workspace.allowPersonalApiKeys, - billedAccountUserId: workspace.billedAccountUserId, - deletedAt: workspaceFiles.deletedAt, - ownership: { - projectId: workspaceFiles.projectId, - context: workspaceFiles.context, - workspaceId: workspaceFiles.workspaceId, - organizationId: workspaceFiles.organizationId, - userId: workspaceFiles.userId, - chatId: workspaceFiles.chatId, - }, - }) - .from(workspaceFiles) - .innerJoin(workspace, eq(workspaceFiles.workspaceId, workspace.id)) - .where(and(eq(workspaceFiles.id, fileId), eq(workspaceFiles.context, 'workspace'))) - .limit(1) - - if ( - !context || - !matchesWorkspaceFileOwner(resolveFileOwner(context.ownership), context.workspaceId) - ) { - return null - } - return omit(context, ['ownership']) + return loadWorkspaceFileContext(fileId, { includeDeleted: true, includeArchivedWorkspace: true }) } /** @@ -2379,13 +2410,10 @@ interface CommitFileContentOptions { secretProvenancePolicy?: WorkspaceFileSecretProvenancePolicy } -/** Commits bytes, history and provenance under canonical scope and storage locks. */ +/** Commits bytes, history and provenance after the caller has locked its canonical billing payer. */ async function commitFileContentInTx(tx: DbTransaction, options: CommitFileContentOptions) { const { owner, fileId, staged } = options - if (owner.entityType === 'workspace') { - await lockWorkspaceProject(tx, owner.entityId) - await lockWorkspaceStorageForMutationInTx(tx, owner.entityId) - } + if (owner.entityType === 'workspace') await lockWorkspaceProject(tx, owner.entityId) if (staged.contentHash === null) throw new Error('Content replacement requires staged bytes with a verified content hash') assertStagedFileOwner(owner, staged) @@ -2557,7 +2585,6 @@ export async function updateWorkspaceFileContent( throw new OrchestrationError('not_found', 'File not found') } - const storageBillingContext = await resolveStorageBillingContext(workspaceId) const nextContentType = contentType || fileRecord.type try { const staged = await stageFileContent({ @@ -2570,6 +2597,10 @@ export async function updateWorkspaceFileContent( }) const finalized = await finalizeStagedFileContent(staged, async (tx) => { + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: workspaceId, + }) const committed = await commitFileContentInTx(tx, { owner: { entityType: 'workspace', entityId: workspaceId }, fileId, @@ -2582,20 +2613,7 @@ export async function updateWorkspaceFileContent( sizeDiff, storageCleanupEventIds, } = committed - let updatedUsage: number | undefined - if (sizeDiff > 0) { - updatedUsage = await incrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - sizeDiff - ) - } else if (sizeDiff < 0) { - await decrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - Math.abs(sizeDiff) - ) - } + const updatedUsage = await accounting.mutation.applyDelta(sizeDiff) const liveDocEventId = options.syncLiveDoc !== false && @@ -2612,6 +2630,7 @@ export async function updateWorkspaceFileContent( file: updatedFile, sizeDiff, updatedUsage, + billing: accounting.billing, liveDocEventId, storageCleanupEventIds, currentVersion: committed.currentVersion, @@ -2620,37 +2639,15 @@ export async function updateWorkspaceFileContent( if (finalized.sizeDiff !== 0) { void maybeNotifyStorageLimitForBillingContext( - storageBillingContext, + finalized.billing, finalized.updatedUsage, finalized.sizeDiff < 0 ) } - await processWorkspaceFileStorageCleanupsNow(finalized.storageCleanupEventIds, { - workspaceId, - fileId, - reason: 'released version', - }) - - if (finalized.liveDocEventId) { - try { - const result = await processWorkspaceFileLiveDocReconciliationNow(finalized.liveDocEventId) - if (result !== 'completed') { - logger.warn('Live document reconciliation deferred to outbox retry', { - workspaceId, - fileId, - eventId: finalized.liveDocEventId, - result, - }) - } - } catch (error) { - logger.warn('Live document reconciliation deferred after inline processing error', { - workspaceId, - fileId, - eventId: finalized.liveDocEventId, - error: getErrorMessage(error), - }) - } - } + await finishFileContentEffects( + { cleanupIds: finalized.storageCleanupEventIds, liveDocEventId: finalized.liveDocEventId }, + { workspaceId, fileId, reason: 'released version' } + ) const currentFolderPath = finalized.file.folderId === fileRecord.folderId ? fileRecord.folderPath : null @@ -2851,7 +2848,6 @@ export async function purgeCreatedWorkspaceFile(params: { expectedFolderId: string | null expectedUpdatedAt: Date }): Promise { - const storageBillingContext = await resolveStorageBillingContext(params.workspaceId) const expectedFolder = params.expectedFolderId === null ? isNull(workspaceFiles.folderId) @@ -2868,8 +2864,10 @@ export async function purgeCreatedWorkspaceFile(params: { isNull(workspaceFiles.deletedAt) ) const cleanupEventIds = await db.transaction(async (tx) => { - await lockWorkspaceProject(tx, params.workspaceId) - await lockWorkspaceStorageForMutationInTx(tx, params.workspaceId) + const accounting = await prepareFileAccountingInTx(tx, { + entityType: 'workspace', + entityId: params.workspaceId, + }) const [lockedFile] = await tx .select({ id: workspaceFiles.id, @@ -2890,11 +2888,7 @@ export async function purgeCreatedWorkspaceFile(params: { .returning({ id: workspaceFiles.id }) if (!deleted) throw new Error('Locked archive-created file could not be deleted') - await decrementStorageUsageForBillingContextInTx( - tx, - storageBillingContext, - getWorkspaceFileSize(lockedFile) - ) + await accounting.mutation.applyDelta(-getWorkspaceFileSize(lockedFile)) const keys = new Set([lockedFile.key, ...versionKeys]) return enqueueWorkspaceFileStorageCleanups(tx, [...keys]) }) diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-accounting.test.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-accounting.test.ts index 570e7ae5788..38d8389cc6e 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-accounting.test.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-accounting.test.ts @@ -16,6 +16,7 @@ import { workspaceFileSecretProvenanceMock, workspaceFileSecretProvenanceMockFns, } from '@sim/testing/mocks/workspace-file-secret-provenance.mock' +import { toRecord } from '@sim/utils/object' import { eq } from 'drizzle-orm' import { beforeEach, describe, expect, it, vi } from 'vitest' @@ -116,13 +117,10 @@ workspaceFileFoldersMockFns.mockNormalizeWorkspaceFileItemName.mockImplementatio ) const mockResolveRestoredFolderId = folderQueriesMockFns.mockResolveRestoredFolderId -const mockDecrementStorageUsageForBillingContextInTx = - billingStorageMockFns.mockDecrementStorageUsageForBillingContextInTx -const mockIncrementStorageUsageForBillingContextInTx = - billingStorageMockFns.mockIncrementStorageUsageForBillingContextInTx +const mockApplyFileStorageDelta = billingStorageMockFns.mockApplyFileStorageDelta const mockMaybeNotifyStorageLimitForBillingContext = billingStorageMockFns.mockMaybeNotifyStorageLimitForBillingContext -const mockResolveStorageBillingContext = billingStorageMockFns.mockResolveStorageBillingContext +const mockPrepareFileAccountingInTx = billingStorageMockFns.mockPrepareFileAccountingInTx const mockInitializeWorkspaceFileSecretProvenanceInTx = workspaceFileSecretProvenanceMockFns.mockInitializeWorkspaceFileSecretProvenanceInTx @@ -179,7 +177,10 @@ describe('workspace file metadata and storage accounting', () => { let stagedKey = '' beforeEach(() => { resetDbChainMock() - mockResolveStorageBillingContext.mockResolvedValue(STORAGE_CONTEXT) + mockPrepareFileAccountingInTx.mockResolvedValue({ + billing: STORAGE_CONTEXT, + mutation: { applyDelta: mockApplyFileStorageDelta }, + }) mockResolveWorkspaceFileFolderTarget.mockResolvedValue(null) mockAssertWorkspaceFileFolderTarget.mockResolvedValue(null) mockHasCloudStorage.mockReturnValue(false) @@ -197,9 +198,8 @@ describe('workspace file metadata and storage accounting', () => { mockGetWorkspaceWithOwner.mockResolvedValue({ archivedAt: null }) mockFileNameExistsInWorkspaceFolder.mockResolvedValue(false) mockResolveRestoredFolderId.mockResolvedValue(null) - mockIncrementStorageUsageForBillingContextInTx.mockReset().mockResolvedValue(10) + mockApplyFileStorageDelta.mockReset().mockResolvedValue(10) mockInitializeWorkspaceFileSecretProvenanceInTx.mockResolvedValue(undefined) - mockDecrementStorageUsageForBillingContextInTx.mockResolvedValue(undefined) mockMaybeNotifyStorageLimitForBillingContext.mockResolvedValue(undefined) mockDeleteFile.mockResolvedValue(undefined) mockEnqueueWorkspaceFileStorageCleanups.mockImplementation(async (_executor, keys: string[]) => @@ -228,6 +228,12 @@ describe('workspace file metadata and storage accounting', () => { mockResolveFolderPathFromIndex .mockReturnValueOnce('folder-initial') .mockReturnValueOnce('folder-final') + workspaceFileFoldersMockFns.mockResolveFileFolderTarget.mockImplementationOnce( + async (_owner, target) => { + const folderId = toRecord(target).folderId + return typeof folderId === 'string' ? { id: folderId } : null + } + ) dbChainMockFns.returning.mockResolvedValueOnce([inserted]) await uploadWorkspaceFile( @@ -260,9 +266,7 @@ describe('workspace file metadata and storage accounting', () => { it('cleans up a newly uploaded object when atomic metadata finalization rolls back', async () => { dbChainMockFns.returning.mockResolvedValueOnce([FILE_ROW]) - mockIncrementStorageUsageForBillingContextInTx.mockRejectedValueOnce( - new Error('payer update failed') - ) + mockApplyFileStorageDelta.mockRejectedValueOnce(new Error('payer update failed')) await expect( uploadWorkspaceFile( @@ -306,18 +310,14 @@ describe('workspace file metadata and storage accounting', () => { expect(eq).toHaveBeenCalledWith(workspaceFiles.originalName, FILE_ROW.originalName) expect(eq).toHaveBeenCalledWith(workspaceFiles.folderId, extractedRow.folderId) expect(eq).toHaveBeenCalledWith(workspaceFiles.updatedAt, FILE_ROW.updatedAt) - expect(mockDecrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - expect.any(Object), - STORAGE_CONTEXT, - FILE_ROW.size - ) + expect(mockApplyFileStorageDelta).toHaveBeenCalledWith(-FILE_ROW.size) expect(mockEnqueueWorkspaceFileStorageCleanups).toHaveBeenCalledWith(expect.any(Object), [ FILE_ROW.key, ]) expect(dbChainMockFns.delete.mock.invocationCallOrder[0]).toBeLessThan( - mockDecrementStorageUsageForBillingContextInTx.mock.invocationCallOrder[0] + mockApplyFileStorageDelta.mock.invocationCallOrder[0] ) - expect(mockDecrementStorageUsageForBillingContextInTx.mock.invocationCallOrder[0]).toBeLessThan( + expect(mockApplyFileStorageDelta.mock.invocationCallOrder[0]).toBeLessThan( mockEnqueueWorkspaceFileStorageCleanups.mock.invocationCallOrder[0] ) expect(mockProcessWorkspaceFileStorageCleanupsNow).toHaveBeenCalledWith( @@ -341,7 +341,7 @@ describe('workspace file metadata and storage accounting', () => { }) ).resolves.toBe(false) - expect(mockDecrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() expect(mockEnqueueWorkspaceFileStorageCleanups).not.toHaveBeenCalled() expect(mockProcessWorkspaceFileStorageCleanupsNow).not.toHaveBeenCalled() expect(mockDeleteFile).not.toHaveBeenCalled() @@ -373,7 +373,7 @@ describe('workspace file metadata and storage accounting', () => { ]) expect([first.created, second.created].sort()).toEqual([false, true]) - expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledTimes(1) + expect(mockApplyFileStorageDelta).toHaveBeenCalledTimes(1) expect(mockReplaceWorkspaceFileSecretProvenanceInTx).toHaveBeenCalledTimes(1) expect(mockReplaceWorkspaceFileSecretProvenanceInTx).toHaveBeenCalledWith( expect.any(Object), @@ -411,7 +411,7 @@ describe('workspace file metadata and storage accounting', () => { { status: 'exact', entries: [] } ) expect(mockReplaceWorkspaceFileSecretProvenanceInTx).not.toHaveBeenCalled() - expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() }) it('does not delete an object when a registration race finds a different operation', async () => { @@ -431,7 +431,7 @@ describe('workspace file metadata and storage accounting', () => { }) ).rejects.toThrow('already registered to a different workspace file operation') - expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() expect(mockDeleteFile).not.toHaveBeenCalled() }) @@ -453,7 +453,7 @@ describe('workspace file metadata and storage accounting', () => { }) ).rejects.toMatchObject({ code: 'conflict' }) expect(dbChainMockFns.returning).not.toHaveBeenCalled() - expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() }) @@ -468,7 +468,7 @@ describe('workspace file metadata and storage accounting', () => { await deleteWorkspaceFile(FILE_ROW.workspaceId, FILE_ROW.id) await deleteWorkspaceFile(FILE_ROW.workspaceId, FILE_ROW.id) - expect(mockDecrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() }) it('re-roots a restored file whose folder has since been archived', async () => { @@ -516,11 +516,7 @@ describe('workspace file metadata and storage accounting', () => { persistMetadata: false, }) ) - expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - expect.any(Object), - STORAGE_CONTEXT, - 3 - ) + expect(mockApplyFileStorageDelta).toHaveBeenCalledWith(3) expect(mockApplyWorkspaceFileSecretProvenancePolicyInTx).toHaveBeenCalledWith( expect.any(Object), FILE_ROW.id, @@ -599,9 +595,7 @@ describe('workspace file metadata and storage accounting', () => { dbChainMockFns.limit.mockResolvedValueOnce([FILE_ROW]).mockResolvedValueOnce([FILE_ROW]) dbChainMockFns.returning.mockResolvedValueOnce([updatedFile]) - mockIncrementStorageUsageForBillingContextInTx.mockRejectedValueOnce( - new Error('Storage limit exceeded') - ) + mockApplyFileStorageDelta.mockRejectedValueOnce(new Error('Storage limit exceeded')) await expect( updateWorkspaceFileContent( @@ -642,11 +636,7 @@ describe('workspace file metadata and storage accounting', () => { MD_ROW.id, PREPARED_COLLAB_STATE ) - expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - transaction, - STORAGE_CONTEXT, - content.length - MD_ROW.sizeBytes - ) + expect(mockApplyFileStorageDelta).toHaveBeenCalledWith(content.length - MD_ROW.sizeBytes) expect(mockEnqueueWorkspaceFileLiveDocReconciliation).toHaveBeenCalledWith( transaction, expect.objectContaining({ fileId: MD_ROW.id }) @@ -722,8 +712,7 @@ describe('workspace file metadata and storage accounting', () => { expect(mockSaveCollabDocStateInTx).toHaveBeenCalledWith(transaction, MD_ROW.id, preparedState) expect(dbChainMockFns.update).not.toHaveBeenCalled() expect(mockApplyWorkspaceFileSecretProvenancePolicyInTx).not.toHaveBeenCalled() - expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() - expect(mockDecrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockApplyFileStorageDelta).not.toHaveBeenCalled() expect(mockEnqueueWorkspaceFileLiveDocReconciliation).not.toHaveBeenCalled() expect(mockProcessWorkspaceFileLiveDocReconciliationNow).not.toHaveBeenCalled() expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() @@ -776,9 +765,7 @@ describe('workspace file metadata and storage accounting', () => { { ...MD_ROW, key: replacementKey, sizeBytes: 13 }, ]) - mockIncrementStorageUsageForBillingContextInTx.mockRejectedValueOnce( - new Error('accounting unavailable') - ) + mockApplyFileStorageDelta.mockRejectedValueOnce(new Error('accounting unavailable')) await expect( updateWorkspaceFileContent( @@ -801,11 +788,7 @@ describe('workspace file metadata and storage accounting', () => { MD_ROW.id, PREPARED_COLLAB_STATE ) - expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - transaction, - STORAGE_CONTEXT, - 8 - ) + expect(mockApplyFileStorageDelta).toHaveBeenCalledWith(8) expect(mockEnqueueWorkspaceFileLiveDocReconciliation).not.toHaveBeenCalled() expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() expect(mockEnqueueWorkspaceFileStorageCleanups).toHaveBeenCalledWith( diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-billing.test.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-billing.test.ts deleted file mode 100644 index 93a76b0b06e..00000000000 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-storage-billing.test.ts +++ /dev/null @@ -1,108 +0,0 @@ -import { dbChainMockFns, resetDbChainMock } from '@sim/testing' -import { billingStorageMock, billingStorageMockFns } from '@sim/testing/mocks/billing-storage.mock' -import { projectMembershipMock } from '@sim/testing/mocks/project-membership.mock' -import { storageServiceMock, storageServiceMockFns } from '@sim/testing/mocks/storage-service.mock' -import { uploadsMock, uploadsMockFns } from '@sim/testing/mocks/uploads.mock' -import { - workspaceFileFoldersMock, - workspaceFileFoldersMockFns, -} from '@sim/testing/mocks/workspace-file-folders.mock' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -vi.mock('@/lib/projects/membership', () => projectMembershipMock) -vi.mock('@/lib/billing/storage', () => billingStorageMock) - -vi.mock('@/lib/uploads', () => uploadsMock) - -vi.mock('@/lib/uploads/core/storage-service', () => storageServiceMock) - -vi.mock( - '@/lib/uploads/contexts/workspace/workspace-file-folder-manager', - () => workspaceFileFoldersMock -) - -import { uploadWorkspaceFile } from '@/lib/uploads/contexts/workspace/workspace-file-manager' - -const mockUploadFile = storageServiceMockFns.mockUploadFile -const mockIncrementStorageUsageForBillingContextInTx = - billingStorageMockFns.mockIncrementStorageUsageForBillingContextInTx -const mockMaybeNotifyStorageLimitForBillingContext = - billingStorageMockFns.mockMaybeNotifyStorageLimitForBillingContext -const mockResolveStorageBillingContext = billingStorageMockFns.mockResolveStorageBillingContext -const mockResolveWorkspaceFileFolderTarget = - workspaceFileFoldersMockFns.mockResolveWorkspaceFileFolderTarget - -workspaceFileFoldersMockFns.mockBuildWorkspaceFileFolderPathMap.mockImplementation(() => new Map()) -workspaceFileFoldersMockFns.mockNormalizeWorkspaceFileItemName.mockImplementation( - (name: string) => name -) - -uploadsMockFns.mockGetServePathPrefix.mockImplementation(() => '/api/files/serve/s3/') - -const STORAGE_CONTEXT = { - workspaceId: 'workspace-1', - billedAccountUserId: 'workspace-owner', - billingEntity: { type: 'organization' as const, id: 'workspace-org' }, - plan: 'team_25000', - customStorageLimitGB: null, -} - -describe('workspace file storage attribution', () => { - beforeEach(() => { - resetDbChainMock() - mockResolveStorageBillingContext.mockResolvedValue(STORAGE_CONTEXT) - mockResolveWorkspaceFileFolderTarget.mockResolvedValue(null) - mockIncrementStorageUsageForBillingContextInTx.mockResolvedValue(5) - mockMaybeNotifyStorageLimitForBillingContext.mockResolvedValue(undefined) - mockUploadFile.mockImplementation(async ({ customKey }) => ({ key: customKey })) - }) - - it.each(['external-collaborator', 'personal-api-key-user'])( - 'charges the workspace payer while retaining %s as uploader metadata', - async (actorUserId) => { - dbChainMockFns.returning - .mockImplementationOnce(async () => [ - { - id: 'file-1', - key: mockUploadFile.mock.calls.at(-1)?.[0].customKey, - userId: actorUserId, - workspaceId: 'workspace-1', - folderId: null, - context: 'workspace', - chatId: null, - originalName: 'note.txt', - displayName: 'note.txt', - contentType: 'text/plain', - size: 5, - sizeBytes: 5, - deletedAt: null, - uploadedAt: new Date(), - updatedAt: new Date(), - contentUpdatedAt: new Date(), - }, - ]) - .mockResolvedValueOnce([{ id: 'file-1' }]) - - await uploadWorkspaceFile( - 'workspace-1', - actorUserId, - Buffer.from('hello'), - 'note.txt', - 'text/plain' - ) - - expect(mockResolveStorageBillingContext).toHaveBeenCalledWith('workspace-1') - expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith( - expect.any(Object), - STORAGE_CONTEXT, - 5 - ) - expect(dbChainMockFns.values).toHaveBeenCalledWith( - expect.objectContaining({ - userId: actorUserId, - workspaceId: 'workspace-1', - }) - ) - } - ) -}) diff --git a/apps/sim/lib/uploads/server/content-effects.test.ts b/apps/sim/lib/uploads/server/content-effects.test.ts new file mode 100644 index 00000000000..7345b98375f --- /dev/null +++ b/apps/sim/lib/uploads/server/content-effects.test.ts @@ -0,0 +1,67 @@ +import { getMockLogger } from '@sim/testing/mocks/logger.mock' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + cleanup: vi.fn(), + reconcile: vi.fn(), +})) +vi.mock('@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox', () => ({ + processWorkspaceFileStorageCleanupsNow: mocks.cleanup, +})) +vi.mock('@/lib/uploads/contexts/workspace/workspace-file-live-doc-outbox', () => ({ + processWorkspaceFileLiveDocReconciliationNow: mocks.reconcile, +})) + +const logger = getMockLogger('FileContentEffects') + +import { finishFileContentEffects } from '@/lib/uploads/server/content-effects' + +const effects = { cleanupIds: ['cleanup-1'], liveDocEventId: 'event-1' } +describe('committed file reconciliation policy', () => { + beforeEach(() => { + mocks.cleanup.mockResolvedValue(undefined) + }) + it.each(['completed', 'pending', 'processing', 'lease_lost'])( + 'allows a %s result to finish or retry durably', + async (status) => { + mocks.reconcile.mockResolvedValue(status) + await expect(finishFileContentEffects(effects, { fileId: 'file-1' })).resolves.toBeUndefined() + expect(mocks.cleanup).toHaveBeenCalledExactlyOnceWith(['cleanup-1'], { fileId: 'file-1' }) + expect(mocks.reconcile).toHaveBeenCalledExactlyOnceWith('event-1') + expect(logger.error).not.toHaveBeenCalled() + if (status === 'completed') expect(logger.warn).not.toHaveBeenCalled() + else + expect(logger.warn).toHaveBeenCalledExactlyOnceWith( + 'Live document reconciliation remains pending', + { fileId: 'file-1', eventId: 'event-1', result: status } + ) + } + ) + it.each(['dead_letter', 'not_found'])( + 'reports a terminal %s result without failing the committed write', + async (status) => { + mocks.reconcile.mockResolvedValue(status) + await expect(finishFileContentEffects(effects, { fileId: 'file-1' })).resolves.toBeUndefined() + expect(mocks.cleanup).toHaveBeenCalledExactlyOnceWith(['cleanup-1'], { fileId: 'file-1' }) + expect(mocks.reconcile).toHaveBeenCalledExactlyOnceWith('event-1') + expect(logger.error).toHaveBeenCalledExactlyOnceWith( + 'Committed file live-document reconciliation requires intervention', + { fileId: 'file-1', eventId: 'event-1', result: status } + ) + } + ) + it('reports an inline processing exception without failing the committed write', async () => { + mocks.reconcile.mockRejectedValue(new Error('database unavailable')) + await expect(finishFileContentEffects(effects, { fileId: 'file-1' })).resolves.toBeUndefined() + expect(mocks.cleanup).toHaveBeenCalledExactlyOnceWith(['cleanup-1'], { fileId: 'file-1' }) + expect(mocks.reconcile).toHaveBeenCalledExactlyOnceWith('event-1') + expect(logger.error).toHaveBeenCalledExactlyOnceWith( + 'Committed file live-document reconciliation failed inline', + expect.objectContaining({ + fileId: 'file-1', + eventId: 'event-1', + error: expect.objectContaining({ message: 'database unavailable' }), + }) + ) + }) +}) diff --git a/apps/sim/lib/uploads/server/content-effects.ts b/apps/sim/lib/uploads/server/content-effects.ts new file mode 100644 index 00000000000..4ab1822747f --- /dev/null +++ b/apps/sim/lib/uploads/server/content-effects.ts @@ -0,0 +1,32 @@ +import { createLogger } from '@sim/logger' +import { describeError } from '@sim/utils/errors' +import { processWorkspaceFileLiveDocReconciliationNow } from '@/lib/uploads/contexts/workspace/workspace-file-live-doc-outbox' +import { processWorkspaceFileStorageCleanupsNow } from '@/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox' + +const logger = createLogger('FileContentEffects') + +/** Reports durable effect failures without changing the outcome of an already committed write. */ +export async function finishFileContentEffects( + effects: { cleanupIds: readonly string[]; liveDocEventId?: string }, + logContext: Record +): Promise { + await processWorkspaceFileStorageCleanupsNow(effects.cleanupIds, logContext) + if (!effects.liveDocEventId) return + const context = { ...logContext, eventId: effects.liveDocEventId } + try { + const result = await processWorkspaceFileLiveDocReconciliationNow(effects.liveDocEventId) + if (result === 'dead_letter' || result === 'not_found') { + logger.error('Committed file live-document reconciliation requires intervention', { + ...context, + result, + }) + } else if (result !== 'completed') { + logger.warn('Live document reconciliation remains pending', { ...context, result }) + } + } catch (error) { + logger.error('Committed file live-document reconciliation failed inline', { + ...context, + error: describeError(error), + }) + } +} diff --git a/apps/sim/lib/workspace-files/ownership.ts b/apps/sim/lib/workspace-files/ownership.ts index 058467c6e24..c1618254b3c 100644 --- a/apps/sim/lib/workspace-files/ownership.ts +++ b/apps/sim/lib/workspace-files/ownership.ts @@ -62,3 +62,8 @@ export function resolveFileOwner(file: PersistedFileOwnership): FileOwner | null } return null } + +/** Compares canonical owner identity without treating attribution as authority. */ +export function matchesFileOwner(owner: FileOwner | null, assertedOwner: FileOwner): boolean { + return owner?.entityType === assertedOwner.entityType && owner.entityId === assertedOwner.entityId +} diff --git a/packages/testing/src/mocks/billing-storage.mock.ts b/packages/testing/src/mocks/billing-storage.mock.ts index 334280ec593..0c0291d6837 100644 --- a/packages/testing/src/mocks/billing-storage.mock.ts +++ b/packages/testing/src/mocks/billing-storage.mock.ts @@ -1,5 +1,7 @@ import { vi } from 'vitest' +const mockApplyFileStorageDelta = vi.fn<(deltaBytes: number) => Promise>(async () => 0) + /** * Stand-in for `StorageLimitExceededError` with the real `name`, constructor, and * `code: 'payload_too_large'`. It is NOT a subclass of the real `OrchestrationError`; a test @@ -17,9 +19,8 @@ export class MockStorageLimitExceededError extends Error { /** * Controllable mock functions for `@/lib/billing/storage`. * - * Every export is a bare `vi.fn()` (the async ones resolve `undefined`, which covers the - * fire-and-forget `maybeNotifyStorageLimitForBillingContext` and the `*InTx` counters); set the - * billing context, quota result, or limit a test needs per case. + * The file preparer returns a controllable signed mutation; other exports are bare `vi.fn()` + * functions. Set the billing context, quota result, or limit a test needs per case. * * @example * ```ts @@ -30,6 +31,12 @@ export class MockStorageLimitExceededError extends Error { * ``` */ export const billingStorageMockFns = { + mockLockWorkspaceStorageForMutationInTx: vi.fn(async () => undefined), + mockApplyFileStorageDelta, + mockPrepareFileAccountingInTx: vi.fn(async () => ({ + billing: {}, + mutation: { applyDelta: mockApplyFileStorageDelta }, + })), mockResolveStorageBillingContext: vi.fn(), mockCheckStorageQuota: vi.fn(), mockCheckStorageQuotaForBillingContext: vi.fn(), @@ -38,7 +45,6 @@ export const billingStorageMockFns = { mockGetUserStorageLimit: vi.fn(), mockGetUserStorageUsage: vi.fn(), mockApplyStorageUsageDeltasInTx: vi.fn(), - mockLockWorkspaceStorageForMutationInTx: vi.fn(), mockCheckAndIncrementStorageUsageInTx: vi.fn(), mockDecrementStorageUsageForBillingContextInTx: vi.fn(), mockIncrementAdmittedStorageUsageForBillingContextInTx: vi.fn(), @@ -56,6 +62,9 @@ export const billingStorageMockFns = { * ``` */ export const billingStorageMock = { + lockWorkspaceStorageForMutationInTx: + billingStorageMockFns.mockLockWorkspaceStorageForMutationInTx, + prepareFileAccountingInTx: billingStorageMockFns.mockPrepareFileAccountingInTx, StorageLimitExceededError: MockStorageLimitExceededError, resolveStorageBillingContext: billingStorageMockFns.mockResolveStorageBillingContext, checkStorageQuota: billingStorageMockFns.mockCheckStorageQuota, @@ -65,8 +74,6 @@ export const billingStorageMock = { getUserStorageLimit: billingStorageMockFns.mockGetUserStorageLimit, getUserStorageUsage: billingStorageMockFns.mockGetUserStorageUsage, applyStorageUsageDeltasInTx: billingStorageMockFns.mockApplyStorageUsageDeltasInTx, - lockWorkspaceStorageForMutationInTx: - billingStorageMockFns.mockLockWorkspaceStorageForMutationInTx, checkAndIncrementStorageUsageInTx: billingStorageMockFns.mockCheckAndIncrementStorageUsageInTx, decrementStorageUsageForBillingContextInTx: billingStorageMockFns.mockDecrementStorageUsageForBillingContextInTx, diff --git a/packages/testing/src/mocks/workspace-file-folders.mock.ts b/packages/testing/src/mocks/workspace-file-folders.mock.ts index a056dac1549..791b4109d52 100644 --- a/packages/testing/src/mocks/workspace-file-folders.mock.ts +++ b/packages/testing/src/mocks/workspace-file-folders.mock.ts @@ -74,6 +74,9 @@ export const workspaceFileFoldersMockFns = { mockListWorkspaceFileFolders: vi.fn(async (..._args: unknown[]): Promise => []), mockListFileFolders: vi.fn(async (..._args: unknown[]): Promise => []), mockGetWorkspaceFileFolder: vi.fn(), + mockResolveFileFolderTarget: vi.fn( + async (..._args: unknown[]): Promise<{ id: string } | null> => null + ), mockResolveWorkspaceFileFolderTarget: vi.fn( async (..._args: unknown[]): Promise => null ), @@ -123,6 +126,7 @@ export const workspaceFileFoldersMock = { listWorkspaceFileFolders: fns.mockListWorkspaceFileFolders, listFileFolders: fns.mockListFileFolders, getWorkspaceFileFolder: fns.mockGetWorkspaceFileFolder, + resolveFileFolderTarget: fns.mockResolveFileFolderTarget, resolveWorkspaceFileFolderTarget: fns.mockResolveWorkspaceFileFolderTarget, assertWorkspaceFileFolderTarget: fns.mockAssertWorkspaceFileFolderTarget, createWorkspaceFileFolder: fns.mockCreateWorkspaceFileFolder, diff --git a/packages/utils/src/errors.ts b/packages/utils/src/errors.ts index 14248fb4e25..863772624ac 100644 --- a/packages/utils/src/errors.ts +++ b/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'], diff --git a/scripts/check-file-names.baseline.json b/scripts/check-file-names.baseline.json index ee3a89c5661..9294e57c23b 100644 --- a/scripts/check-file-names.baseline.json +++ b/scripts/check-file-names.baseline.json @@ -108,7 +108,6 @@ "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-secret-provenance.test.ts", "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-secret-provenance.ts", "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-storage-accounting.test.ts", - "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-storage-billing.test.ts", "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox.test.ts", "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-storage-cleanup-outbox.ts", "stutter\tapps/sim/lib/uploads/contexts/workspace/workspace-file-versions.ts",