feat(db): role-keyed dbFor clients for cleanup and exec workloads

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy

Author: TheodoreSpeaks

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: {

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<void>
     workspaceIds,
     retentionDate,
     tableName: `${label}/job_execution_logs`,
+    dbClient: cleanupDb,
   })
 
   if (runGlobalHousekeeping && plan === 'free') {

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[]) =>

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,
@@ -37,6 +37,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
 /**
@@ -81,7 +88,7 @@ async function selectExpiredWorkspaceFiles(
 ): Promise<WorkspaceFileScope> {
   const [legacyRows, multiContextRows] = await Promise.all([
     selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
-      db
+      cleanupDb
         .select({
           id: workspaceFile.id,
           key: workspaceFile.key,
@@ -98,7 +105,7 @@ async function selectExpiredWorkspaceFiles(
         .limit(chunkLimit)
     ),
     selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
-      db
+      cleanupDb
         .select({
           id: workspaceFiles.id,
           key: workspaceFiles.key,
@@ -200,7 +207,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(
@@ -240,7 +247,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(
@@ -343,7 +350,7 @@ async function hardDeleteKnowledgeBaseDocuments(
   label: string
 ): Promise<void> {
   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))
@@ -358,7 +365,7 @@ async function hardDeleteKnowledgeBaseDocuments(
     }
   }
 
-  const remaining = await db
+  const remaining = await cleanupDb
     .select({ id: document.id })
     .from(document)
     .where(inArray(document.knowledgeBaseId, knowledgeBaseIds))
@@ -378,8 +385,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(
@@ -458,7 +466,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
   // 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(
@@ -555,7 +563,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))
@@ -577,7 +585,8 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
     workflow,
     workflow.id,
     doomedWorkflowIds,
-    `${label}/workflow`
+    `${label}/workflow`,
+    cleanupDb
   )
   totalDeleted += workflowResult.deleted
 
@@ -614,6 +623,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
       retentionDate,
       tableName: `${label}/${target.name}`,
       requireTimestampNotNull: true,
+      dbClient: cleanupDb,
     })
     totalDeleted += result.deleted
   }

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<void>
   )
 
   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<void>
     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<void>
     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<void>
     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<void>
     workspaceIds,
     retentionDate,
     tableName: `${label}/mothershipInboxTask`,
+    dbClient: cleanupDb,
   })
 
   const totalDeleted =

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<typeof db, 'select' | 'delete'>
+
 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<TRow extends { id: string }> {
    */
   totalRowLimit?: number
   workspaceChunkSize?: number
+  /** Client the DELETEs run on. Defaults to the global pool. */
+  dbClient?: BatchDeleteClient
 }
 
 /**
@@ -107,6 +116,7 @@ export async function chunkedBatchDelete<TRow extends { id: string }>({
   maxBatches = DEFAULT_MAX_BATCHES_PER_TABLE,
   totalRowLimit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE,
   workspaceChunkSize = DEFAULT_WORKSPACE_CHUNK_SIZE,
+  dbClient = db,
 }: ChunkedBatchDeleteOptions<TRow>): Promise<TableCleanupResult> {
   const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 }
 
@@ -149,7 +159,7 @@ export async function chunkedBatchDelete<TRow extends { id: string }>({
         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<TableCleanupResult> {
   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<string>`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<TableCleanupResult> {
   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 })