diff --git a/apps/sim/background/cleanup-logs.test.ts b/apps/sim/background/cleanup-logs.test.ts index 6999171e630..9df5ec1ab77 100644 --- a/apps/sim/background/cleanup-logs.test.ts +++ b/apps/sim/background/cleanup-logs.test.ts @@ -80,12 +80,17 @@ const { } }) -vi.mock('@sim/db', () => ({ - db: { +vi.mock('@sim/db', () => { + const db = { execute: mockExecute, select: mockSelect, - }, -})) + } + return { + db, + // Cleanup-pool client shares the instance so the seeded chains still apply. + dbFor: () => db, + } +}) vi.mock('@sim/db/schema', () => ({ executionLargeValueDependencies: { @@ -255,7 +260,7 @@ describe('cleanup logs worker', () => { workspaceIds: ['workspace-1'], }) - expect(mockMarkLargeValuesDeleted).toHaveBeenCalledWith([largeValueKey]) + expect(mockMarkLargeValuesDeleted).toHaveBeenCalledWith([largeValueKey], expect.anything()) expect(mockDeleteFileMetadata).toHaveBeenCalledTimes(2) }) diff --git a/apps/sim/background/cleanup-logs.ts b/apps/sim/background/cleanup-logs.ts index bd451aa1e09..343e388ed2b 100644 --- a/apps/sim/background/cleanup-logs.ts +++ b/apps/sim/background/cleanup-logs.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { executionLargeValueDependencies, executionLargeValueReferences, @@ -30,6 +30,9 @@ import { deleteFileMetadata } from '@/lib/uploads/server/metadata' const logger = createLogger('CleanupLogs') +/** All cleanup queries run on the dedicated cleanup pool. */ +const cleanupDb = dbFor('cleanup') + interface FileDeleteStats { filesTotal: number filesDeleted: number @@ -107,7 +110,7 @@ async function deleteLargeValueKeys(keys: string[]): Promise<{ deleted: number; if (deletedKeys.length > 0) { try { - await markLargeValuesDeleted(deletedKeys) + await markLargeValuesDeleted(deletedKeys, cleanupDb) } catch (error) { logger.error('Failed to mark large execution values as deleted:', { error }) return { deleted: 0, failed: result.failed.length + deletedKeys.length } @@ -153,7 +156,7 @@ async function cleanupLargeExecutionValues( LARGE_VALUE_CLEANUP_BATCH_SIZE, LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted ) - const rows = await db + const rows = await cleanupDb .select({ key: executionLargeValues.key }) .from(executionLargeValues) .where( @@ -219,7 +222,7 @@ async function cleanupLegacyLargeExecutionValues( LARGE_VALUE_CLEANUP_BATCH_SIZE, LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted ) - const rows = await db + const rows = await cleanupDb .select({ key: workspaceFiles.key }) .from(workspaceFiles) .where( @@ -353,7 +356,11 @@ async function cleanupLargeValueMetadata(workspaceIds: string[], label: string): const tombstonesDeletedBefore = new Date( Date.now() - LARGE_VALUE_TOMBSTONE_RETENTION_HOURS * 60 * 60 * 1000 ) - const result = await pruneLargeValueMetadata({ workspaceIds, tombstonesDeletedBefore }) + const result = await pruneLargeValueMetadata({ + workspaceIds, + tombstonesDeletedBefore, + dbClient: cleanupDb, + }) logger.info( `[${label}/execution_large_value_metadata] Pruned ${result.referencesDeleted} stale references, ${result.dependenciesDeleted} dependencies, ${result.tombstonesDeleted} tombstones` ) @@ -377,8 +384,9 @@ async function cleanupWorkflowExecutionLogs( tableDef: workflowExecutionLogs, workspaceIds, tableName: `${label}/workflow_execution_logs`, + dbClient: cleanupDb, selectChunk: (chunkIds, limit) => - db + cleanupDb .select({ id: workflowExecutionLogs.id, files: workflowExecutionLogs.files, @@ -465,6 +473,7 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise workspaceIds, retentionDate, tableName: `${label}/job_execution_logs`, + dbClient: cleanupDb, }) if (runGlobalHousekeeping && plan === 'free') { diff --git a/apps/sim/background/cleanup-soft-deletes.test.ts b/apps/sim/background/cleanup-soft-deletes.test.ts index c104d9b7351..7c0e29603bf 100644 --- a/apps/sim/background/cleanup-soft-deletes.test.ts +++ b/apps/sim/background/cleanup-soft-deletes.test.ts @@ -66,13 +66,18 @@ const { } }) -vi.mock('@sim/db', () => ({ - db: { +vi.mock('@sim/db', () => { + const db = { delete: mockDelete, select: mockSelect, transaction: mockTransaction, - }, -})) + } + return { + db, + // Cleanup-pool client shares the instance so the seeded chains still apply. + dbFor: () => db, + } +}) vi.mock('@sim/db/schema', () => { const table = (cols: string[]) => diff --git a/apps/sim/background/cleanup-soft-deletes.ts b/apps/sim/background/cleanup-soft-deletes.ts index 75e9c42ae20..7db8b2cc8f0 100644 --- a/apps/sim/background/cleanup-soft-deletes.ts +++ b/apps/sim/background/cleanup-soft-deletes.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { db, dbFor } from '@sim/db' import { copilotChats, document, @@ -36,6 +36,13 @@ import { deleteFileMetadata } from '@/lib/uploads/server/metadata' 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. + */ +const cleanupDb = dbFor('cleanup') + const KB_ORPHAN_BINDING_BATCH_SIZE = 500 const KB_ORPHAN_BINDING_TOTAL_LIMIT = 5_000 /** @@ -80,7 +87,7 @@ async function selectExpiredWorkspaceFiles( ): Promise { const [legacyRows, multiContextRows] = await Promise.all([ selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workspaceFile.id, key: workspaceFile.key, @@ -97,7 +104,7 @@ async function selectExpiredWorkspaceFiles( .limit(chunkLimit) ), selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workspaceFiles.id, key: workspaceFiles.key, @@ -199,7 +206,7 @@ async function deleteExpiredLegacyWorkspaceFileRows( const result = { deleted: 0, failed: 0 } for (const batch of chunkArray(rows, DEFAULT_DELETE_CHUNK_SIZE)) { try { - const deleted = await db + const deleted = await cleanupDb .delete(workspaceFile) .where( and( @@ -239,7 +246,7 @@ async function deleteExpiredUnbilledWorkspaceFileRows( for (const [context, contextRows] of rowsByContext) { for (const batch of chunkArray(contextRows, DEFAULT_DELETE_CHUNK_SIZE)) { try { - const deleted = await db + const deleted = await cleanupDb .delete(workspaceFiles) .where( and( @@ -342,7 +349,7 @@ async function hardDeleteKnowledgeBaseDocuments( label: string ): Promise { for (let batch = 0; batch < KB_DOCUMENT_DELETE_MAX_BATCHES; batch++) { - const documentRows = await db + const documentRows = await cleanupDb .select({ id: document.id }) .from(document) .where(inArray(document.knowledgeBaseId, knowledgeBaseIds)) @@ -357,7 +364,7 @@ async function hardDeleteKnowledgeBaseDocuments( } } - const remaining = await db + const remaining = await cleanupDb .select({ id: document.id }) .from(document) .where(inArray(document.knowledgeBaseId, knowledgeBaseIds)) @@ -377,8 +384,9 @@ async function cleanupExpiredKnowledgeBases( workspaceIds, tableName: `${label}/knowledgeBase`, batchSize: KB_RETENTION_BATCH_SIZE, + dbClient: cleanupDb, selectChunk: (chunkIds, limit) => - db + cleanupDb .select({ id: knowledgeBase.id }) .from(knowledgeBase) .where( @@ -457,7 +465,7 @@ async function cleanupOrphanedKnowledgeBaseBindings( KB_ORPHAN_BINDING_BATCH_SIZE, KB_ORPHAN_BINDING_TOTAL_LIMIT - attempted ) - const rows = await db + const rows = await cleanupDb .select({ key: workspaceFiles.key }) .from(workspaceFiles) .where( @@ -535,7 +543,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise // prematurely purge data. const [doomedWorkflows, fileScope, expiredSoftDeletedChats] = await Promise.all([ selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workflow.id }) .from(workflow) .where( @@ -549,7 +557,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise ), selectExpiredWorkspaceFiles(workspaceIds, retentionDate), selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where( @@ -570,7 +578,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise const doomedChatIds = new Set(softDeletedChatIds) if (doomedWorkflowIds.length > 0) { const workflowChats = await selectRowsByIdChunks(doomedWorkflowIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where(inArray(copilotChats.workflowId, chunkIds)) @@ -595,7 +603,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise // spares their backend data and files. for (const batch of chunkArray(doomedWorkflowIds, DEFAULT_DELETE_CHUNK_SIZE)) { try { - const deleted = await db + const deleted = await cleanupDb .delete(workflow) .where( and( @@ -618,7 +626,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise // also re-checks row existence before purging external data). for (const batch of chunkArray(softDeletedChatIds, DEFAULT_DELETE_CHUNK_SIZE)) { try { - const deleted = await db + const deleted = await cleanupDb .delete(copilotChats) .where( and( @@ -667,6 +675,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise retentionDate, tableName: `${label}/${target.name}`, requireTimestampNotNull: true, + dbClient: cleanupDb, }) totalDeleted += result.deleted } diff --git a/apps/sim/background/cleanup-tasks.ts b/apps/sim/background/cleanup-tasks.ts index e987ac41078..1745e3bba14 100644 --- a/apps/sim/background/cleanup-tasks.ts +++ b/apps/sim/background/cleanup-tasks.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { copilotAsyncToolCalls, copilotChats, @@ -22,6 +22,9 @@ import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup' const logger = createLogger('CleanupTasks') +/** All cleanup queries run on the dedicated cleanup pool. */ +const cleanupDb = dbFor('cleanup') + /** * Delete copilot run checkpoints and async tool calls via join through copilotRuns. * These tables don't have a direct workspaceId — we find qualifying run IDs first. @@ -47,7 +50,7 @@ async function cleanupRunChildren( if (workspaceIds.length === 0) return [] const runIds = await selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotRuns.id }) .from(copilotRuns) .where( @@ -63,7 +66,9 @@ async function cleanupRunChildren( const ids = runIds.map((r) => r.id) return Promise.all( - RUN_CHILD_TABLES.map((t) => deleteRowsById(t.table, t.runIdCol, ids, `${label}/${t.name}`)) + RUN_CHILD_TABLES.map((t) => + deleteRowsById(t.table, t.runIdCol, ids, `${label}/${t.name}`, cleanupDb) + ) ) } @@ -82,7 +87,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise ) const doomedChats = await selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where( @@ -110,6 +115,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise workspaceIds, retentionDate, tableName: `${label}/copilotRuns`, + dbClient: cleanupDb, }) // Delete copilot chats using the exact IDs collected above so the chat @@ -122,7 +128,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise const chatsResult = { deleted: 0, failed: 0 } for (const batch of chunkArray(doomedChatIds, DEFAULT_DELETE_CHUNK_SIZE)) { try { - const deleted = await db + const deleted = await cleanupDb .delete(copilotChats) .where(and(inArray(copilotChats.id, batch), lt(copilotChats.updatedAt, retentionDate))) .returning({ id: copilotChats.id }) @@ -141,6 +147,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise workspaceIds, retentionDate, tableName: `${label}/mothershipInboxTask`, + dbClient: cleanupDb, }) const totalDeleted = diff --git a/apps/sim/lib/cleanup/batch-delete.ts b/apps/sim/lib/cleanup/batch-delete.ts index d523800cd59..7b0ede50fd9 100644 --- a/apps/sim/lib/cleanup/batch-delete.ts +++ b/apps/sim/lib/cleanup/batch-delete.ts @@ -5,6 +5,13 @@ import type { PgColumn, PgTable } from 'drizzle-orm/pg-core' const logger = createLogger('BatchDelete') +/** + * Structural client surface the delete helpers need. Satisfied by the global + * `db`, a `dbFor(...)` sub-pool client, and a transaction handle, so callers + * pick which pool the deletes run on (cleanup jobs pass `dbFor('cleanup')`). + */ +export type BatchDeleteClient = Pick + export const DEFAULT_BATCH_SIZE = 2000 /** 50 × 2000 = 100K row cap per cleanup run; drains long-tail tenants in days, not weeks. */ export const DEFAULT_MAX_BATCHES_PER_TABLE = 50 @@ -84,6 +91,8 @@ export interface ChunkedBatchDeleteOptions { */ totalRowLimit?: number workspaceChunkSize?: number + /** Client the DELETEs run on. Defaults to the global pool. */ + dbClient?: BatchDeleteClient } /** @@ -107,6 +116,7 @@ export async function chunkedBatchDelete({ maxBatches = DEFAULT_MAX_BATCHES_PER_TABLE, totalRowLimit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, workspaceChunkSize = DEFAULT_WORKSPACE_CHUNK_SIZE, + dbClient = db, }: ChunkedBatchDeleteOptions): Promise { const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 } @@ -149,7 +159,7 @@ export async function chunkedBatchDelete({ if (onBatch) await onBatch(rows) const ids = rows.map((r) => r.id) - const deleted = await db + const deleted = await dbClient .delete(tableDef) .where(inArray(sql`id`, ids)) .returning({ id: sql`id` }) @@ -189,6 +199,8 @@ export interface BatchDeleteOptions { batchSize?: number maxBatches?: number workspaceChunkSize?: number + /** Client the SELECTs and DELETEs run on. Defaults to the global pool. */ + dbClient?: BatchDeleteClient } /** @@ -204,16 +216,18 @@ export async function batchDeleteByWorkspaceAndTimestamp({ retentionDate, tableName, requireTimestampNotNull = false, + dbClient = db, ...rest }: BatchDeleteOptions): Promise { return chunkedBatchDelete({ tableDef, workspaceIds, tableName, + dbClient, selectChunk: (chunkIds, limit) => { const predicates = [inArray(workspaceIdCol, chunkIds), lt(timestampCol, retentionDate)] if (requireTimestampNotNull) predicates.push(isNotNull(timestampCol)) - return db + return dbClient .select({ id: sql`id` }) .from(tableDef) .where(and(...predicates)) @@ -232,6 +246,7 @@ export async function deleteRowsById( idCol: PgColumn, ids: string[], tableName: string, + dbClient: BatchDeleteClient = db, chunkSize: number = DEFAULT_DELETE_CHUNK_SIZE ): Promise { const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 } @@ -240,7 +255,7 @@ export async function deleteRowsById( const chunks = chunkArray(ids, chunkSize) for (const [chunkIdx, chunkIds] of chunks.entries()) { try { - const deleted = await db + const deleted = await dbClient .delete(tableDef) .where(inArray(idCol, chunkIds)) .returning({ id: idCol }) diff --git a/apps/sim/lib/cleanup/chat-cleanup.ts b/apps/sim/lib/cleanup/chat-cleanup.ts index 55f6a96832c..ba84f1f0279 100644 --- a/apps/sim/lib/cleanup/chat-cleanup.ts +++ b/apps/sim/lib/cleanup/chat-cleanup.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { copilotChats, copilotMessages, workspaceFiles } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { and, inArray, isNull } from 'drizzle-orm' @@ -10,6 +10,9 @@ import { isUsingCloudStorage, StorageService } from '@/lib/uploads' const logger = createLogger('ChatCleanup') +/** Chat cleanup only ever runs from cleanup jobs, so its reads use the cleanup pool. */ +const cleanupDb = dbFor('cleanup') + const COPILOT_CLEANUP_BATCH_SIZE = 1000 /** Bounds how many chats' `copilot_messages` rows are scanned per query. */ const CHAT_FILE_COLLECT_CHUNK_SIZE = 500 @@ -42,7 +45,7 @@ export async function collectChatFiles(chatIds: string[]): Promise { for (const chunk of chunkArray(chatIds, CHAT_FILE_COLLECT_CHUNK_SIZE)) { const [linkedFiles, messageRows] = await Promise.all([ - db + cleanupDb .select({ key: workspaceFiles.key, context: workspaceFiles.context, @@ -58,7 +61,7 @@ export async function collectChatFiles(chatIds: string[]): Promise { ), // Scan every message row for the chat (no deleted_at filter): this is a // deletion path collecting blob keys, so attachments on any row count. - db + cleanupDb .select({ content: copilotMessages.content, chatId: copilotMessages.chatId }) .from(copilotMessages) .where(inArray(copilotMessages.chatId, chunk)), @@ -205,7 +208,7 @@ export async function prepareChatCleanup( // whose rows are actually gone, so a surviving row never loses its data. const survivors = new Set() for (const chunk of chunkArray(chatIds, CHAT_FILE_COLLECT_CHUNK_SIZE)) { - const rows = await db + const rows = await cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where(inArray(copilotChats.id, chunk)) diff --git a/apps/sim/lib/core/config/env.ts b/apps/sim/lib/core/config/env.ts index 27aa7019f63..d7d51b7b518 100644 --- a/apps/sim/lib/core/config/env.ts +++ b/apps/sim/lib/core/config/env.ts @@ -25,6 +25,8 @@ export const env = createEnv({ DATABASE_URL_WEB: z.string().url().optional(), // Per-role primary URL override; @sim/db falls back to DATABASE_URL DATABASE_URL_TRIGGER: z.string().url().optional(), // Per-role primary URL override (trigger) DATABASE_URL_REALTIME: z.string().url().optional(), // Per-role primary URL override (realtime) + DATABASE_URL_CLEANUP: z.string().url().optional(), // Sub-process pool URL override (cleanup jobs, via dbFor) + DATABASE_URL_EXEC: z.string().url().optional(), // Sub-process pool URL override (inline execution writes, via dbFor) DATABASE_REPLICA_URL_WEB: z.string().url().optional(), // Per-role replica URL override; falls back to DATABASE_REPLICA_URL DATABASE_REPLICA_URL_TRIGGER: z.string().url().optional(), // Per-role replica URL override (trigger) DATABASE_REPLICA_URL_REALTIME: z.string().url().optional(), // Per-role replica URL override (realtime) diff --git a/apps/sim/lib/execution/payloads/large-value-metadata.test.ts b/apps/sim/lib/execution/payloads/large-value-metadata.test.ts index 6df41ce599b..b7b0b337f15 100644 --- a/apps/sim/lib/execution/payloads/large-value-metadata.test.ts +++ b/apps/sim/lib/execution/payloads/large-value-metadata.test.ts @@ -83,15 +83,20 @@ const { } }) -vi.mock('@sim/db', () => ({ - db: { +vi.mock('@sim/db', () => { + const db = { delete: mockDelete, execute: mockExecute, insert: mockInsert, select: mockSelect, transaction: mockTransaction, - }, -})) + } + return { + db, + // Exec-pool client shares the instance so the seeded chains still apply. + dbFor: () => db, + } +}) vi.mock('@sim/db/schema', () => ({ executionLargeValueDependencies: { diff --git a/apps/sim/lib/execution/payloads/large-value-metadata.ts b/apps/sim/lib/execution/payloads/large-value-metadata.ts index fd95938c40f..41d56fbf2e2 100644 --- a/apps/sim/lib/execution/payloads/large-value-metadata.ts +++ b/apps/sim/lib/execution/payloads/large-value-metadata.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { db, dbFor } from '@sim/db' import { executionLargeValueDependencies, executionLargeValueReferences, @@ -54,6 +54,8 @@ interface PruneLargeValueMetadataOptions { tombstonesDeletedBefore: Date batchSize?: number maxRowsPerTable?: number + /** Client the prune DELETEs run on. Defaults to the global pool; cleanup jobs pass `dbFor('cleanup')`. */ + dbClient?: LargeValueMetadataClient } function parseLargeValueStorageKey(key: string): LargeValueStorageKeyParts | null { @@ -186,7 +188,7 @@ export async function registerLargeValueOwner( return false } - await db.transaction(async (tx) => { + await dbFor('exec').transaction(async (tx) => { await tx .insert(executionLargeValues) .values({ @@ -309,7 +311,8 @@ export async function addLargeValueReference( return } - const [existingRef] = await db + const execDb = dbFor('exec') + const [existingRef] = await execDb .select({ key: executionLargeValueReferences.key }) .from(executionLargeValueReferences) .where( @@ -326,7 +329,7 @@ export async function addLargeValueReference( return } - const existingRefs = await db + const existingRefs = await execDb .select({ key: executionLargeValueReferences.key }) .from(executionLargeValueReferences) .where( @@ -344,7 +347,7 @@ export async function addLargeValueReference( ) } - await db + await execDb .insert(executionLargeValueReferences) .values({ key: boundedKey, @@ -363,24 +366,31 @@ export async function replaceLargeValueReferences( const referenceKeys = scope.workspaceId ? collectLargeValueReferenceKeys(value, scope.workspaceId) : [] - await db.transaction(async (tx) => { + await dbFor('exec').transaction(async (tx) => { await replaceLargeValueReferenceKeysWithClient(tx, scope, referenceKeys) }) } -export async function markLargeValuesDeleted(keys: string[]): Promise { +export async function markLargeValuesDeleted( + keys: string[], + dbClient: LargeValueMetadataClient = db +): Promise { if (keys.length === 0) { return } - await db + await dbClient .update(executionLargeValues) .set({ deletedAt: new Date() }) .where(inArray(executionLargeValues.key, keys)) } -async function pruneStaleReferences(workspaceIds: string[], batchSize: number): Promise { - const rows = await db.execute<{ count: number }>(sql` +async function pruneStaleReferences( + workspaceIds: string[], + batchSize: number, + dbClient: LargeValueMetadataClient +): Promise { + const rows = await dbClient.execute<{ count: number }>(sql` WITH deleted AS ( DELETE FROM ${executionLargeValueReferences} AS ref WHERE ref.ctid IN ( @@ -418,9 +428,10 @@ async function pruneStaleReferences(workspaceIds: string[], batchSize: number): async function pruneDeletedParentDependencies( workspaceIds: string[], - batchSize: number + batchSize: number, + dbClient: LargeValueMetadataClient ): Promise { - const rows = await db.execute<{ count: number }>(sql` + const rows = await dbClient.execute<{ count: number }>(sql` WITH deleted AS ( DELETE FROM ${executionLargeValueDependencies} AS dependency WHERE dependency.ctid IN ( @@ -452,9 +463,10 @@ async function pruneDeletedParentDependencies( async function pruneDeletedLargeValueTombstones( workspaceIds: string[], deletedBefore: Date, - batchSize: number + batchSize: number, + dbClient: LargeValueMetadataClient ): Promise { - const rows = await db.execute<{ count: number }>(sql` + const rows = await dbClient.execute<{ count: number }>(sql` WITH deleted AS ( DELETE FROM ${executionLargeValues} AS value WHERE value.ctid IN ( @@ -482,6 +494,7 @@ export async function pruneLargeValueMetadata({ tombstonesDeletedBefore, batchSize = LARGE_VALUE_METADATA_PRUNE_BATCH_SIZE, maxRowsPerTable = LARGE_VALUE_METADATA_PRUNE_MAX_ROWS_PER_TABLE, + dbClient = db, }: PruneLargeValueMetadataOptions): Promise { const result: LargeValueMetadataPruneResult = { referencesDeleted: 0, @@ -498,7 +511,8 @@ export async function pruneLargeValueMetadata({ if (referencesRemaining > 0) { result.referencesDeleted += await pruneStaleReferences( workspaceChunk, - Math.min(batchSize, referencesRemaining) + Math.min(batchSize, referencesRemaining), + dbClient ) } @@ -506,7 +520,8 @@ export async function pruneLargeValueMetadata({ if (dependenciesRemaining > 0) { result.dependenciesDeleted += await pruneDeletedParentDependencies( workspaceChunk, - Math.min(batchSize, dependenciesRemaining) + Math.min(batchSize, dependenciesRemaining), + dbClient ) } @@ -515,7 +530,8 @@ export async function pruneLargeValueMetadata({ result.tombstonesDeleted += await pruneDeletedLargeValueTombstones( workspaceChunk, tombstonesDeletedBefore, - Math.min(batchSize, tombstonesRemaining) + Math.min(batchSize, tombstonesRemaining), + dbClient ) } diff --git a/apps/sim/lib/logs/execution/logger.test.ts b/apps/sim/lib/logs/execution/logger.test.ts index e1c6e326af2..4b83f955541 100644 --- a/apps/sim/lib/logs/execution/logger.test.ts +++ b/apps/sim/lib/logs/execution/logger.test.ts @@ -21,14 +21,17 @@ vi.mock('@sim/db', () => { update: txUpdateMock, execute: dbExecuteMock, } + const db = { + select: dbSelectMock, + insert: vi.fn(), + update: vi.fn(), + execute: dbExecuteMock, + transaction: vi.fn(async (cb: (txArg: typeof tx) => Promise) => cb(tx)), + } return { - db: { - select: dbSelectMock, - insert: vi.fn(), - update: vi.fn(), - execute: dbExecuteMock, - transaction: vi.fn(async (cb: (txArg: typeof tx) => Promise) => cb(tx)), - }, + db, + // Exec-pool client shares the instance so call-order seeding still applies. + dbFor: () => db, } }) diff --git a/apps/sim/lib/logs/execution/logger.ts b/apps/sim/lib/logs/execution/logger.ts index eb48ec1768f..0f6c642e457 100644 --- a/apps/sim/lib/logs/execution/logger.ts +++ b/apps/sim/lib/logs/execution/logger.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { db, dbFor } from '@sim/db' import { member, organization, @@ -69,6 +69,13 @@ import { emitExecutionCompletedEvent } from '@/lib/workspace-events/emitter' import type { SerializableExecutionState } from '@/executor/execution/types' const logger = createLogger('ExecutionLogger') + +/** + * Execution-log persistence (reads and writes on `workflow_execution_logs`, + * including the completion transaction) runs on the dedicated exec pool. + * Billing/usage-ledger work stays on the global `db`. + */ +const execDb = dbFor('exec') const MAX_EXECUTION_DATA_BYTES = 3 * 1024 * 1024 const MAX_TRACE_IO_BYTES = 8 * 1024 const MAX_WORKFLOW_VALUE_BYTES = 512 * 1024 @@ -546,7 +553,7 @@ export class ExecutionLogger implements IExecutionLoggerService { execLog.debug('Starting workflow execution') // Check if execution log already exists (idempotency check) - const existingLog = await db + const existingLog = await execDb .select() .from(workflowExecutionLogs) .where(eq(workflowExecutionLogs.executionId, executionId)) @@ -583,7 +590,7 @@ export class ExecutionLogger implements IExecutionLoggerService { const startTime = new Date() - const [workflowLog] = await db + const [workflowLog] = await execDb .insert(workflowExecutionLogs) .values({ id: generateId(), @@ -742,7 +749,7 @@ export class ExecutionLogger implements IExecutionLoggerService { let execLog = logger.withMetadata({ executionId }) execLog.debug('Completing workflow execution', { isResume }) - const [existingLog] = await db + const [existingLog] = await execDb .select() .from(workflowExecutionLogs) .where(eq(workflowExecutionLogs.executionId, executionId)) @@ -933,7 +940,7 @@ export class ExecutionLogger implements IExecutionLoggerService { } const completedExecutionLargeValueKeys = collectLargeValueReferenceKeys(storedExecutionData) - const updatedLog = await db.transaction(async (tx) => { + const updatedLog = await execDb.transaction(async (tx) => { await setExecutionLogWriteTimeouts(tx) const [log] = await tx @@ -1163,7 +1170,7 @@ export class ExecutionLogger implements IExecutionLoggerService { } async getWorkflowExecution(executionId: string): Promise { - const [workflowLog] = await db + const [workflowLog] = await execDb .select() .from(workflowExecutionLogs) .where(eq(workflowExecutionLogs.executionId, executionId)) diff --git a/apps/sim/lib/logs/execution/logging-session.test.ts b/apps/sim/lib/logs/execution/logging-session.test.ts index 21947c152c2..d07fe618eb5 100644 --- a/apps/sim/lib/logs/execution/logging-session.test.ts +++ b/apps/sim/lib/logs/execution/logging-session.test.ts @@ -47,13 +47,18 @@ const { releaseExecutionSlotMock: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { +vi.mock('@sim/db', () => { + const db = { select: dbMocks.select, update: dbMocks.update, execute: dbMocks.execute, - }, -})) + } + return { + db, + // Exec-pool client shares the instance so the seeded chains still apply. + dbFor: () => db, + } +}) vi.mock('drizzle-orm', () => ({ eq: dbMocks.eq, diff --git a/apps/sim/lib/logs/execution/logging-session.ts b/apps/sim/lib/logs/execution/logging-session.ts index ed63effb669..d8ba08e5ad5 100644 --- a/apps/sim/lib/logs/execution/logging-session.ts +++ b/apps/sim/lib/logs/execution/logging-session.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { workflowExecutionLogs } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { describeError, toError } from '@sim/utils/errors' @@ -79,6 +79,9 @@ function buildCompletedMarkerPersistenceQuery(params: { ) <= ${params.marker.endedAt}` } +/** Progress-marker and status writes on `workflow_execution_logs` use the exec pool. */ +const execDb = dbFor('exec') + const logger = createLogger('LoggingSession') type CompletionAttempt = 'complete' | 'error' | 'cancelled' | 'paused' @@ -193,7 +196,7 @@ export class LoggingSession { return } try { - await db.execute( + await execDb.execute( buildStartedMarkerPersistenceQuery({ executionId: this.executionId, workflowId: this.workflowId, @@ -219,7 +222,7 @@ export class LoggingSession { return } try { - await db.execute( + await execDb.execute( buildCompletedMarkerPersistenceQuery({ executionId: this.executionId, workflowId: this.workflowId, @@ -486,7 +489,7 @@ export class LoggingSession { this.completing = true try { - const currentLog = await db + const currentLog = await execDb .select({ status: workflowExecutionLogs.status }) .from(workflowExecutionLogs) .where( @@ -617,7 +620,7 @@ export class LoggingSession { const endTime = endedAt ? new Date(endedAt) : new Date() const durationMs = typeof totalDurationMs === 'number' ? totalDurationMs : 0 - const currentLog = await db + const currentLog = await execDb .select({ status: workflowExecutionLogs.status }) .from(workflowExecutionLogs) .where( @@ -711,7 +714,7 @@ export class LoggingSession { const endTime = endedAt ? new Date(endedAt) : new Date() const durationMs = typeof totalDurationMs === 'number' ? totalDurationMs : 0 - const currentLog = await db + const currentLog = await execDb .select({ status: workflowExecutionLogs.status }) .from(workflowExecutionLogs) .where( @@ -1090,7 +1093,7 @@ export class LoggingSession { ELSE ${executionData} END` } - await db + await execDb .update(workflowExecutionLogs) .set({ level: 'error', status: 'failed', executionData }) .where( diff --git a/apps/sim/lib/logs/execution/snapshot/service.ts b/apps/sim/lib/logs/execution/snapshot/service.ts index 21b14a4ea6b..c6cc2d22ddd 100644 --- a/apps/sim/lib/logs/execution/snapshot/service.ts +++ b/apps/sim/lib/logs/execution/snapshot/service.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { workflowExecutionLogs, workflowExecutionSnapshots } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' @@ -51,7 +51,7 @@ export class SnapshotService implements ISnapshotService { * out-of-line storage is reused, so the per-execution write drops from the * full blob to a tiny heap tuple. */ - const [upsertedSnapshot] = await db + const [upsertedSnapshot] = await dbFor('exec') .insert(workflowExecutionSnapshots) .values(snapshotData) .onConflictDoUpdate({ @@ -81,7 +81,7 @@ export class SnapshotService implements ISnapshotService { } async getSnapshot(id: string): Promise { - const [snapshot] = await db + const [snapshot] = await dbFor('exec') .select() .from(workflowExecutionSnapshots) .where(eq(workflowExecutionSnapshots.id, id)) @@ -102,7 +102,9 @@ export class SnapshotService implements ISnapshotService { return sha256Hex(stateString) } + /** Only invoked from the cleanup-logs background job, so it runs on the cleanup pool. */ async cleanupOrphanedSnapshots(olderThanDays: number): Promise { + const cleanupDb = dbFor('cleanup') const cutoffDate = new Date() cutoffDate.setDate(cutoffDate.getDate() - olderThanDays) @@ -113,14 +115,14 @@ export class SnapshotService implements ISnapshotService { let stoppedEarly = false for (let batch = 0; batch < MAX_BATCHES; batch++) { - const candidates = await db + const candidates = await cleanupDb .select({ id: workflowExecutionSnapshots.id }) .from(workflowExecutionSnapshots) .where( and( lt(workflowExecutionSnapshots.createdAt, cutoffDate), notExists( - db + cleanupDb .select({ one: sql`1` }) .from(workflowExecutionLogs) .where(eq(workflowExecutionLogs.stateSnapshotId, workflowExecutionSnapshots.id)) @@ -132,13 +134,13 @@ export class SnapshotService implements ISnapshotService { if (candidates.length === 0) break const ids = candidates.map((c) => c.id) - const deleted = await db + const deleted = await cleanupDb .delete(workflowExecutionSnapshots) .where( and( inArray(workflowExecutionSnapshots.id, ids), notExists( - db + cleanupDb .select({ one: sql`1` }) .from(workflowExecutionLogs) .where(eq(workflowExecutionLogs.stateSnapshotId, workflowExecutionSnapshots.id)) diff --git a/apps/sim/lib/workflows/executor/human-in-the-loop-manager.ts b/apps/sim/lib/workflows/executor/human-in-the-loop-manager.ts index 343bed6a853..3be39c26d7f 100644 --- a/apps/sim/lib/workflows/executor/human-in-the-loop-manager.ts +++ b/apps/sim/lib/workflows/executor/human-in-the-loop-manager.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { dbFor } from '@sim/db' import { pausedExecutions, resumeQueue, workflowExecutionLogs } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' @@ -54,6 +54,13 @@ import { hasExecutionResult } from '@/executor/utils/errors' import { filterOutputForLog } from '@/executor/utils/output-filter' import type { SerializedConnection } from '@/serializer/types' +/** + * All paused-execution / resume-queue / execution-log persistence in this + * module runs on the exec pool, mirroring the completion writes in + * `lib/logs/execution/logger.ts`. + */ +const execDb = dbFor('exec') + const logger = createLogger('HumanInTheLoopManager') const RUN_BUFFER_UNAVAILABLE_ERROR = 'Run buffer temporarily unavailable' const TERMINAL_PUBLISH_ERROR = 'Run buffer terminal event publish failed' @@ -373,7 +380,7 @@ export class PauseResumeManager { ...resumeMetadata, } - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { const existing = await tx .select() .from(pausedExecutions) @@ -489,7 +496,7 @@ export class PauseResumeManager { static async enqueueOrStartResume(args: EnqueueResumeArgs): Promise { const { executionId, workflowId, contextId, resumeInput, userId, allowedPauseKinds } = args - return await db.transaction(async (tx) => { + return await execDb.transaction(async (tx) => { const pausedExecution = await tx .select() .from(pausedExecutions) @@ -815,7 +822,7 @@ export class PauseResumeManager { } = args const parentExecutionId = pausedExecution.executionId - await db + await execDb .update(workflowExecutionLogs) .set({ status: 'running' }) .where(eq(workflowExecutionLogs.executionId, parentExecutionId)) @@ -1653,7 +1660,7 @@ export class PauseResumeManager { const { resumeEntryId, pausedExecutionId, parentExecutionId, contextId } = args const now = new Date() - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { await tx .update(resumeQueue) .set({ status: 'completed', completedAt: now, failureReason: null }) @@ -1720,7 +1727,7 @@ export class PauseResumeManager { }): Promise { const now = new Date() - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { await tx .update(resumeQueue) .set({ status: 'failed', failureReason: args.failureReason, completedAt: now }) @@ -1752,7 +1759,7 @@ export class PauseResumeManager { }): Promise { const now = new Date() - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { const pausedExecution = args.preserveForRetry ? await tx .select({ @@ -1843,7 +1850,7 @@ export class PauseResumeManager { }): Promise { const { pausedExecutionId, contextId, pauseBlockId, executionState } = args - const pausedExecution = await db + const pausedExecution = await execDb .select() .from(pausedExecutions) .where(eq(pausedExecutions.id, pausedExecutionId)) @@ -1905,7 +1912,7 @@ export class PauseResumeManager { ? collectLargeValueReferenceKeys(snapshotReferenceValue, snapshotWorkspaceId) : [] - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { await tx .update(pausedExecutions) .set({ @@ -1937,7 +1944,7 @@ export class PauseResumeManager { static async beginPausedCancellation(executionId: string, workflowId: string): Promise { const now = new Date() - return await db.transaction(async (tx) => { + return await execDb.transaction(async (tx) => { const pausedExecution = await tx .select({ id: pausedExecutions.id, status: pausedExecutions.status }) .from(pausedExecutions) @@ -1992,7 +1999,7 @@ export class PauseResumeManager { ): Promise { const now = new Date() - return await db.transaction(async (tx) => { + return await execDb.transaction(async (tx) => { const pausedExecution = await tx .select({ id: pausedExecutions.id, status: pausedExecutions.status }) .from(pausedExecutions) @@ -2033,7 +2040,7 @@ export class PauseResumeManager { ): Promise { const now = new Date() - return await db.transaction(async (tx) => { + return await execDb.transaction(async (tx) => { const pausedExecution = await tx .select({ id: pausedExecutions.id }) .from(pausedExecutions) @@ -2077,7 +2084,7 @@ export class PauseResumeManager { workflowId: string ): Promise { const now = new Date() - await db + await execDb .update(pausedExecutions) .set({ status: sql`CASE WHEN resumed_count > 0 THEN 'partially_resumed' ELSE 'paused' END`, @@ -2097,7 +2104,7 @@ export class PauseResumeManager { executionId: string, workflowId: string ): Promise<'cancelling' | 'cancelled' | null> { - const activeResume = await db + const activeResume = await execDb .select({ id: resumeQueue.id }) .from(resumeQueue) .where(and(eq(resumeQueue.parentExecutionId, executionId), eq(resumeQueue.status, 'claimed'))) @@ -2108,7 +2115,7 @@ export class PauseResumeManager { return null } - const pausedExecution = await db + const pausedExecution = await execDb .select({ status: pausedExecutions.status }) .from(pausedExecutions) .where( @@ -2135,7 +2142,7 @@ export class PauseResumeManager { }): Promise { const now = new Date() - await db.transaction(async (tx) => { + await execDb.transaction(async (tx) => { const pausedExecution = await tx .select({ automaticResumeRetryCount: pausedExecutions.automaticResumeRetryCount, @@ -2228,7 +2235,7 @@ export class PauseResumeManager { pausedExecutionId: string nextResumeAt: Date | null }): Promise { - await db + await execDb .update(pausedExecutions) .set({ nextResumeAt: args.nextResumeAt }) .where( @@ -2260,7 +2267,7 @@ export class PauseResumeManager { } } - const rows = await db + const rows = await execDb .select() .from(pausedExecutions) .where(whereClause) @@ -2278,7 +2285,7 @@ export class PauseResumeManager { static async getPausedExecutionById( id: string ): Promise { - const rows = await db + const rows = await execDb .select() .from(pausedExecutions) .where(eq(pausedExecutions.id, id)) @@ -2292,7 +2299,7 @@ export class PauseResumeManager { }): Promise { const { workflowId, executionId } = options - const row = await db + const row = await execDb .select() .from(pausedExecutions) .where( @@ -2308,7 +2315,7 @@ export class PauseResumeManager { return null } - const queueEntries = await db + const queueEntries = await execDb .select() .from(resumeQueue) .where(eq(resumeQueue.parentExecutionId, executionId)) @@ -2386,7 +2393,7 @@ export class PauseResumeManager { } | null = null while (!pendingEntry) { - const selection = await db.transaction(async (tx) => { + const selection = await execDb.transaction(async (tx) => { const pausedExecution = await tx .select() .from(pausedExecutions) diff --git a/packages/db/db.ts b/packages/db/db.ts index a122948995e..df1ed7de7d3 100644 --- a/packages/db/db.ts +++ b/packages/db/db.ts @@ -1,9 +1,12 @@ +import { createLogger } from '@sim/logger' import { drizzle } from 'drizzle-orm/postgres-js' import postgres from 'postgres' import { resolveDbUrl } from './connection-url' import * as schema from './schema' import { instrumentPoolClient } from './tx-tripwire' +const logger = createLogger('Db') + /** * Per-role pool profiles. Starting numbers — validate against real per-role * process counts (PgBouncer transaction mode, max_connections=200). @@ -14,17 +17,24 @@ export const DB_POOL_PROFILES = { // overlapping logging writes); 3 risks intra-run deadlock. trigger: { primaryMax: 5, replicaMax: 2, appName: 'sim-trigger' }, realtime: { primaryMax: 5, replicaMax: 3, appName: 'sim-realtime' }, + // Sub-process pools, selected per call-site via dbFor() — never via SIM_DB_ROLE. + cleanup: { primaryMax: 5, replicaMax: 2, appName: 'sim-cleanup' }, + exec: { primaryMax: 10, replicaMax: 4, appName: 'sim-exec' }, } as const -type DbRole = keyof typeof DB_POOL_PROFILES +/** Roles a whole process runs as (via SIM_DB_ROLE). */ +const PROCESS_ROLES = ['web', 'trigger', 'realtime'] as const + +type ProcessDbRole = (typeof PROCESS_ROLES)[number] +type SubProcessDbRole = Exclude const roleEnv = process.env.SIM_DB_ROLE?.trim() -if (roleEnv && !Object.hasOwn(DB_POOL_PROFILES, roleEnv)) { +if (roleEnv && !PROCESS_ROLES.includes(roleEnv as ProcessDbRole)) { throw new Error( - `Invalid SIM_DB_ROLE '${roleEnv}' — expected one of ${Object.keys(DB_POOL_PROFILES).join(', ')} (or unset for web)` + `Invalid SIM_DB_ROLE '${roleEnv}' — expected one of ${PROCESS_ROLES.join(', ')} (or unset for web)` ) } -const role = (roleEnv as DbRole) || 'web' +const role = (roleEnv as ProcessDbRole) || 'web' const profile = DB_POOL_PROFILES[role] const connectionString = resolveDbUrl('DATABASE_URL', role) @@ -71,3 +81,58 @@ export const dbReplica: typeof db = replicaUrl } ) : db + +const subPoolClients = new Map() + +/** Which env var the process connection came from — named in dbFor fallback logs. */ +const processUrlEnvVar = process.env[`DATABASE_URL_${role.toUpperCase()}`] + ? `DATABASE_URL_${role.toUpperCase()}` + : 'DATABASE_URL' + +/** + * Per-workload drizzle client with its own pool, built lazily on first call and + * cached per role. Unlike the process-wide `db` (selected by `SIM_DB_ROLE`), + * these are selected per call-site so a workload running inside an existing + * process — cleanup jobs in the trigger worker, inline execution log writes in + * the web server — gets its own connection budget and PgBouncer pool. + * + * Resolves `DATABASE_URL_` with fallback to the URL the process itself + * resolved (`DATABASE_URL_`, then base `DATABASE_URL`), so an + * unset sub-pool URL changes nothing about where this process's traffic lands. + * Always uses the role profile's `appName` — the `DB_APP_NAME` override applies + * only to the process-wide clients. + */ +export function dbFor(role: SubProcessDbRole): typeof db { + const existing = subPoolClients.get(role) + if (existing) return existing + + const keyedEnvVar = `DATABASE_URL_${role.toUpperCase()}` + const keyedUrl = process.env[keyedEnvVar] + const url = keyedUrl ?? connectionString + if (!url) { + throw new Error('Missing DATABASE_URL environment variable') + } + + if (keyedUrl) { + logger.info(`'${role}' pool using dedicated ${keyedEnvVar}`) + } else { + logger.info( + `${keyedEnvVar} not set — '${role}' pool falling back to the process connection (${processUrlEnvVar})` + ) + } + + const subProfile = DB_POOL_PROFILES[role] + const client = drizzle( + instrumentPoolClient( + postgres(url, { + ...poolOptions, + max: subProfile.primaryMax, + connection: { application_name: subProfile.appName }, + }), + role + ), + { schema } + ) + subPoolClients.set(role, client) + return client +} diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index 3c832d79652..d701630239a 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -260,6 +260,8 @@ export const dbChainMock = { db: dbChainInstance, /** Same instance as `db` so per-test chain overrides cover both clients. */ dbReplica: dbChainInstance, + /** Sub-pool clients (`dbFor('cleanup' | 'exec')`) share the same instance too. */ + dbFor: () => dbChainInstance, runOutsideTransactionContext: (fn: () => T): T => fn(), instrumentPoolClient: (client: T): T => client, } @@ -333,6 +335,8 @@ export const databaseMock = { db: mockDbInstance, /** Same instance as `db` so per-test overrides cover both clients. */ dbReplica: mockDbInstance, + /** Sub-pool clients (`dbFor('cleanup' | 'exec')`) share the same instance too. */ + dbFor: () => mockDbInstance, sql: createMockSql(), runOutsideTransactionContext: (fn: () => T): T => fn(), instrumentPoolClient: (client: T): T => client,