From 9103f00756a11aca0c67c99f5657d647441c1079 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Fri, 10 Jul 2026 17:23:30 -0700 Subject: [PATCH 1/3] feat(db): role-keyed dbFor clients for cleanup and exec workloads Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy --- apps/sim/background/cleanup-logs.test.ts | 13 +++-- apps/sim/background/cleanup-logs.ts | 13 +++-- .../background/cleanup-soft-deletes.test.ts | 6 ++- apps/sim/background/cleanup-soft-deletes.ts | 25 +++++---- apps/sim/background/cleanup-tasks.ts | 21 +++++--- apps/sim/lib/cleanup/batch-delete.ts | 21 ++++++-- apps/sim/lib/core/config/env.ts | 2 + apps/sim/lib/logs/execution/logger.test.ts | 17 +++--- apps/sim/lib/logs/execution/logger.ts | 19 ++++--- .../logs/execution/logging-session.test.ts | 13 +++-- .../sim/lib/logs/execution/logging-session.ts | 17 +++--- .../lib/logs/execution/snapshot/service.ts | 14 ++--- packages/db/db.ts | 54 +++++++++++++++++-- packages/testing/src/mocks/database.mock.ts | 4 ++ 14 files changed, 178 insertions(+), 61 deletions(-) diff --git a/apps/sim/background/cleanup-logs.test.ts b/apps/sim/background/cleanup-logs.test.ts index 6999171e630..7a75ee76952 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: { diff --git a/apps/sim/background/cleanup-logs.ts b/apps/sim/background/cleanup-logs.ts index bd451aa1e09..d8a2cc99126 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 @@ -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( @@ -377,8 +380,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 +469,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 a6ede324481..bfc68018da1 100644 --- a/apps/sim/background/cleanup-soft-deletes.test.ts +++ b/apps/sim/background/cleanup-soft-deletes.test.ts @@ -43,7 +43,11 @@ const { } }) -vi.mock('@sim/db', () => ({ db: { select: mockSelect } })) +vi.mock('@sim/db', () => { + const db = { select: mockSelect } + // Cleanup-pool client shares the instance so the seeded chains still apply. + return { db, 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 ee01e297898..dd229cbb5c2 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 { dbFor } from '@sim/db' import { copilotChats, document, @@ -29,6 +29,9 @@ import { deleteFileMetadata } from '@/lib/uploads/server/metadata' const logger = createLogger('CleanupSoftDeletes') +/** All cleanup queries run on the dedicated cleanup pool. */ +const cleanupDb = dbFor('cleanup') + const KB_ORPHAN_BINDING_BATCH_SIZE = 500 const KB_ORPHAN_BINDING_TOTAL_LIMIT = 5_000 /** @@ -57,7 +60,7 @@ async function selectExpiredWorkspaceFiles( ): Promise { const [legacyRows, multiContextRows] = await Promise.all([ selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workspaceFile.id, key: workspaceFile.key }) .from(workspaceFile) .where( @@ -70,7 +73,7 @@ async function selectExpiredWorkspaceFiles( .limit(chunkLimit) ), selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workspaceFiles.id, key: workspaceFiles.key, @@ -194,7 +197,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( @@ -268,7 +271,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise // different subsets above the LIMIT cap and orphan or prematurely purge data. const [doomedWorkflows, fileScope] = await Promise.all([ selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: workflow.id }) .from(workflow) .where( @@ -288,7 +291,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise if (doomedWorkflowIds.length > 0) { const doomedChats = await selectRowsByIdChunks(doomedWorkflowIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where(inArray(copilotChats.workflowId, chunkIds)) @@ -310,7 +313,8 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise workflow, workflow.id, doomedWorkflowIds, - `${label}/workflow` + `${label}/workflow`, + cleanupDb ) totalDeleted += workflowResult.deleted @@ -318,7 +322,8 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise workspaceFile, workspaceFile.id, fileScope.legacyRows.map((r) => r.id), - `${label}/workspaceFile` + `${label}/workspaceFile`, + cleanupDb ) totalDeleted += legacyFileResult.deleted @@ -326,7 +331,8 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise workspaceFiles, workspaceFiles.id, fileScope.multiContextRows.map((r) => r.id), - `${label}/workspaceFiles` + `${label}/workspaceFiles`, + cleanupDb ) totalDeleted += multiContextFileResult.deleted @@ -339,6 +345,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 341f5132eaa..08c7ad73410 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, @@ -21,6 +21,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. @@ -46,7 +49,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( @@ -62,7 +65,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) + ) ) } @@ -81,7 +86,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise ) const doomedChats = await selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - db + cleanupDb .select({ id: copilotChats.id }) .from(copilotChats) .where( @@ -106,7 +111,8 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise copilotFeedback, copilotFeedback.chatId, doomedChatIds, - `${label}/copilotFeedback` + `${label}/copilotFeedback`, + cleanupDb ) // Delete copilot runs (has workspaceId directly, cascades checkpoints) @@ -117,6 +123,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 @@ -125,7 +132,8 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise copilotChats, copilotChats.id, doomedChatIds, - `${label}/copilotChats` + `${label}/copilotChats`, + cleanupDb ) // Delete mothership inbox tasks (has workspaceId directly) @@ -136,6 +144,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/core/config/env.ts b/apps/sim/lib/core/config/env.ts index 717a6f5eb2b..1ce20ae023d 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/logs/execution/logger.test.ts b/apps/sim/lib/logs/execution/logger.test.ts index 923542fde12..d6cde880d97 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 000d4a8cafb..a2b7ae487c0 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, @@ -65,6 +65,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 @@ -528,7 +535,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)) @@ -565,7 +572,7 @@ export class ExecutionLogger implements IExecutionLoggerService { const startTime = new Date() - const [workflowLog] = await db + const [workflowLog] = await execDb .insert(workflowExecutionLogs) .values({ id: generateId(), @@ -717,7 +724,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)) @@ -900,7 +907,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 @@ -1129,7 +1136,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 57eebb8e260..047521beb71 100644 --- a/apps/sim/lib/logs/execution/logging-session.test.ts +++ b/apps/sim/lib/logs/execution/logging-session.test.ts @@ -45,13 +45,18 @@ const { loadWorkflowStateForExecutionMock: 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 1e68ba0f060..c26d3abf773 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' @@ -78,6 +78,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' @@ -183,7 +186,7 @@ export class LoggingSession { return } try { - await db.execute( + await execDb.execute( buildStartedMarkerPersistenceQuery({ executionId: this.executionId, workflowId: this.workflowId, @@ -209,7 +212,7 @@ export class LoggingSession { return } try { - await db.execute( + await execDb.execute( buildCompletedMarkerPersistenceQuery({ executionId: this.executionId, workflowId: this.workflowId, @@ -465,7 +468,7 @@ export class LoggingSession { this.completing = true try { - const currentLog = await db + const currentLog = await execDb .select({ status: workflowExecutionLogs.status }) .from(workflowExecutionLogs) .where( @@ -596,7 +599,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( @@ -690,7 +693,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( @@ -1058,7 +1061,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..aaadd7af9bf 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 { db, 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({ @@ -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/packages/db/db.ts b/packages/db/db.ts index a122948995e..ca112cf1d51 100644 --- a/packages/db/db.ts +++ b/packages/db/db.ts @@ -14,17 +14,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 +78,42 @@ export const dbReplica: typeof db = replicaUrl } ) : db + +const subPoolClients = new Map() + +/** + * 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 base `DATABASE_URL`, so + * behavior is identical to the shared client until the keyed URL is configured. + * 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 url = resolveDbUrl('DATABASE_URL', role) + if (!url) { + throw new Error('Missing DATABASE_URL environment variable') + } + + 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 9a879e4cad8..caf7164e6df 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, From c7ac96bf7e79b6e2e8bb08aa994d0d238c12592e Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Fri, 10 Jul 2026 18:54:40 -0700 Subject: [PATCH 2/3] fix(db): keep cleanup-invoked helpers and snapshot reads on their role pools Route markLargeValuesDeleted / pruneLargeValueMetadata (optional dbClient) and chat-cleanup's file collection through the cleanup pool, and getSnapshot through the exec pool, so the cleanup and inline-execution workloads stop borrowing the process-wide pool for these queries. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy --- apps/sim/background/cleanup-logs.test.ts | 2 +- apps/sim/background/cleanup-logs.ts | 8 +++- apps/sim/lib/cleanup/chat-cleanup.ts | 9 +++-- .../payloads/large-value-metadata.ts | 37 +++++++++++++------ .../lib/logs/execution/snapshot/service.ts | 4 +- 5 files changed, 41 insertions(+), 19 deletions(-) diff --git a/apps/sim/background/cleanup-logs.test.ts b/apps/sim/background/cleanup-logs.test.ts index 7a75ee76952..9df5ec1ab77 100644 --- a/apps/sim/background/cleanup-logs.test.ts +++ b/apps/sim/background/cleanup-logs.test.ts @@ -260,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 d8a2cc99126..343e388ed2b 100644 --- a/apps/sim/background/cleanup-logs.ts +++ b/apps/sim/background/cleanup-logs.ts @@ -110,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 } @@ -356,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` ) diff --git a/apps/sim/lib/cleanup/chat-cleanup.ts b/apps/sim/lib/cleanup/chat-cleanup.ts index c5dafdf9c27..083c249863b 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 { 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 @@ -41,7 +44,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 }) .from(workspaceFiles) .where( @@ -53,7 +56,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 }) .from(copilotMessages) .where(inArray(copilotMessages.chatId, chunk)), diff --git a/apps/sim/lib/execution/payloads/large-value-metadata.ts b/apps/sim/lib/execution/payloads/large-value-metadata.ts index fd95938c40f..c5ffba5bdbc 100644 --- a/apps/sim/lib/execution/payloads/large-value-metadata.ts +++ b/apps/sim/lib/execution/payloads/large-value-metadata.ts @@ -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 { @@ -368,19 +370,26 @@ export async function replaceLargeValueReferences( }) } -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 +427,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 +462,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 +493,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 +510,8 @@ export async function pruneLargeValueMetadata({ if (referencesRemaining > 0) { result.referencesDeleted += await pruneStaleReferences( workspaceChunk, - Math.min(batchSize, referencesRemaining) + Math.min(batchSize, referencesRemaining), + dbClient ) } @@ -506,7 +519,8 @@ export async function pruneLargeValueMetadata({ if (dependenciesRemaining > 0) { result.dependenciesDeleted += await pruneDeletedParentDependencies( workspaceChunk, - Math.min(batchSize, dependenciesRemaining) + Math.min(batchSize, dependenciesRemaining), + dbClient ) } @@ -515,7 +529,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/snapshot/service.ts b/apps/sim/lib/logs/execution/snapshot/service.ts index aaadd7af9bf..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, dbFor } 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' @@ -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)) From 143b38251226352151ae6166a9d1288fbced1932 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Fri, 10 Jul 2026 19:03:10 -0700 Subject: [PATCH 3/3] fix(db): dbFor falls back to the process-role URL, not the base URL With DATABASE_URL_WEB/TRIGGER set (as in prod) and the sub-pool URLs unset, falling back to the base URL would silently shift execution-log and cleanup traffic to a different PgBouncer endpoint on deploy. Chain the fallback through the URL the process itself resolved so the rollout stays inert. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy --- packages/db/db.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/db/db.ts b/packages/db/db.ts index ca112cf1d51..26e67810c73 100644 --- a/packages/db/db.ts +++ b/packages/db/db.ts @@ -88,8 +88,9 @@ const subPoolClients = new Map() * 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 base `DATABASE_URL`, so - * behavior is identical to the shared client until the keyed URL is configured. + * 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. */ @@ -97,7 +98,7 @@ export function dbFor(role: SubProcessDbRole): typeof db { const existing = subPoolClients.get(role) if (existing) return existing - const url = resolveDbUrl('DATABASE_URL', role) + const url = process.env[`DATABASE_URL_${role.toUpperCase()}`] ?? connectionString if (!url) { throw new Error('Missing DATABASE_URL environment variable') }