Skip to content

Commit 70814bc

Browse files
feat(db): role-keyed dbFor clients for cleanup and exec workloads (#5583)
* 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 * 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy * 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy * feat(db): log which connection each dbFor sub-pool resolved to One line per role at first use: the dedicated DATABASE_URL_<ROLE> when set, otherwise an explicit fallback message naming the process connection it shares — so a missing/typo'd env var is visible at rollout instead of silent. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy * feat(db): route pause/resume and large-value metadata persistence to the exec pool Coverage scan follow-up: the paused_executions / resume_queue / workflow_execution_logs transactions in human-in-the-loop-manager.ts and the large-value owner/reference registration writes on the execution path now use dbFor('exec'), matching the completion writes in the execution logger. All are self-contained; billing calls remain outside the moved transactions on the default client. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NNy9Fzpfc1FdHAzYyb6Ycy --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent e841e4e commit 70814bc

18 files changed

Lines changed: 291 additions & 119 deletions

apps/sim/background/cleanup-logs.test.ts

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,17 @@ const {
8080
}
8181
})
8282

83-
vi.mock('@sim/db', () => ({
84-
db: {
83+
vi.mock('@sim/db', () => {
84+
const db = {
8585
execute: mockExecute,
8686
select: mockSelect,
87-
},
88-
}))
87+
}
88+
return {
89+
db,
90+
// Cleanup-pool client shares the instance so the seeded chains still apply.
91+
dbFor: () => db,
92+
}
93+
})
8994

9095
vi.mock('@sim/db/schema', () => ({
9196
executionLargeValueDependencies: {
@@ -255,7 +260,7 @@ describe('cleanup logs worker', () => {
255260
workspaceIds: ['workspace-1'],
256261
})
257262

258-
expect(mockMarkLargeValuesDeleted).toHaveBeenCalledWith([largeValueKey])
263+
expect(mockMarkLargeValuesDeleted).toHaveBeenCalledWith([largeValueKey], expect.anything())
259264
expect(mockDeleteFileMetadata).toHaveBeenCalledTimes(2)
260265
})
261266

apps/sim/background/cleanup-logs.ts

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { db } from '@sim/db'
1+
import { dbFor } from '@sim/db'
22
import {
33
executionLargeValueDependencies,
44
executionLargeValueReferences,
@@ -30,6 +30,9 @@ import { deleteFileMetadata } from '@/lib/uploads/server/metadata'
3030

3131
const logger = createLogger('CleanupLogs')
3232

33+
/** All cleanup queries run on the dedicated cleanup pool. */
34+
const cleanupDb = dbFor('cleanup')
35+
3336
interface FileDeleteStats {
3437
filesTotal: number
3538
filesDeleted: number
@@ -107,7 +110,7 @@ async function deleteLargeValueKeys(keys: string[]): Promise<{ deleted: number;
107110

108111
if (deletedKeys.length > 0) {
109112
try {
110-
await markLargeValuesDeleted(deletedKeys)
113+
await markLargeValuesDeleted(deletedKeys, cleanupDb)
111114
} catch (error) {
112115
logger.error('Failed to mark large execution values as deleted:', { error })
113116
return { deleted: 0, failed: result.failed.length + deletedKeys.length }
@@ -153,7 +156,7 @@ async function cleanupLargeExecutionValues(
153156
LARGE_VALUE_CLEANUP_BATCH_SIZE,
154157
LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted
155158
)
156-
const rows = await db
159+
const rows = await cleanupDb
157160
.select({ key: executionLargeValues.key })
158161
.from(executionLargeValues)
159162
.where(
@@ -219,7 +222,7 @@ async function cleanupLegacyLargeExecutionValues(
219222
LARGE_VALUE_CLEANUP_BATCH_SIZE,
220223
LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted
221224
)
222-
const rows = await db
225+
const rows = await cleanupDb
223226
.select({ key: workspaceFiles.key })
224227
.from(workspaceFiles)
225228
.where(
@@ -353,7 +356,11 @@ async function cleanupLargeValueMetadata(workspaceIds: string[], label: string):
353356
const tombstonesDeletedBefore = new Date(
354357
Date.now() - LARGE_VALUE_TOMBSTONE_RETENTION_HOURS * 60 * 60 * 1000
355358
)
356-
const result = await pruneLargeValueMetadata({ workspaceIds, tombstonesDeletedBefore })
359+
const result = await pruneLargeValueMetadata({
360+
workspaceIds,
361+
tombstonesDeletedBefore,
362+
dbClient: cleanupDb,
363+
})
357364
logger.info(
358365
`[${label}/execution_large_value_metadata] Pruned ${result.referencesDeleted} stale references, ${result.dependenciesDeleted} dependencies, ${result.tombstonesDeleted} tombstones`
359366
)
@@ -377,8 +384,9 @@ async function cleanupWorkflowExecutionLogs(
377384
tableDef: workflowExecutionLogs,
378385
workspaceIds,
379386
tableName: `${label}/workflow_execution_logs`,
387+
dbClient: cleanupDb,
380388
selectChunk: (chunkIds, limit) =>
381-
db
389+
cleanupDb
382390
.select({
383391
id: workflowExecutionLogs.id,
384392
files: workflowExecutionLogs.files,
@@ -465,6 +473,7 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise<void>
465473
workspaceIds,
466474
retentionDate,
467475
tableName: `${label}/job_execution_logs`,
476+
dbClient: cleanupDb,
468477
})
469478

470479
if (runGlobalHousekeeping && plan === 'free') {

apps/sim/background/cleanup-soft-deletes.test.ts

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,18 @@ const {
6666
}
6767
})
6868

69-
vi.mock('@sim/db', () => ({
70-
db: {
69+
vi.mock('@sim/db', () => {
70+
const db = {
7171
delete: mockDelete,
7272
select: mockSelect,
7373
transaction: mockTransaction,
74-
},
75-
}))
74+
}
75+
return {
76+
db,
77+
// Cleanup-pool client shares the instance so the seeded chains still apply.
78+
dbFor: () => db,
79+
}
80+
})
7681

7782
vi.mock('@sim/db/schema', () => {
7883
const table = (cols: string[]) =>

apps/sim/background/cleanup-soft-deletes.ts

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { db } from '@sim/db'
1+
import { db, dbFor } from '@sim/db'
22
import {
33
copilotChats,
44
document,
@@ -36,6 +36,13 @@ import { deleteFileMetadata } from '@/lib/uploads/server/metadata'
3636

3737
const logger = createLogger('CleanupSoftDeletes')
3838

39+
/**
40+
* Cleanup queries run on the dedicated cleanup pool. The one exception is the
41+
* billable-file transaction below, which couples row deletion with a storage
42+
* billing decrement — billing writes stay on the default client.
43+
*/
44+
const cleanupDb = dbFor('cleanup')
45+
3946
const KB_ORPHAN_BINDING_BATCH_SIZE = 500
4047
const KB_ORPHAN_BINDING_TOTAL_LIMIT = 5_000
4148
/**
@@ -80,7 +87,7 @@ async function selectExpiredWorkspaceFiles(
8087
): Promise<WorkspaceFileScope> {
8188
const [legacyRows, multiContextRows] = await Promise.all([
8289
selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
83-
db
90+
cleanupDb
8491
.select({
8592
id: workspaceFile.id,
8693
key: workspaceFile.key,
@@ -97,7 +104,7 @@ async function selectExpiredWorkspaceFiles(
97104
.limit(chunkLimit)
98105
),
99106
selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
100-
db
107+
cleanupDb
101108
.select({
102109
id: workspaceFiles.id,
103110
key: workspaceFiles.key,
@@ -199,7 +206,7 @@ async function deleteExpiredLegacyWorkspaceFileRows(
199206
const result = { deleted: 0, failed: 0 }
200207
for (const batch of chunkArray(rows, DEFAULT_DELETE_CHUNK_SIZE)) {
201208
try {
202-
const deleted = await db
209+
const deleted = await cleanupDb
203210
.delete(workspaceFile)
204211
.where(
205212
and(
@@ -239,7 +246,7 @@ async function deleteExpiredUnbilledWorkspaceFileRows(
239246
for (const [context, contextRows] of rowsByContext) {
240247
for (const batch of chunkArray(contextRows, DEFAULT_DELETE_CHUNK_SIZE)) {
241248
try {
242-
const deleted = await db
249+
const deleted = await cleanupDb
243250
.delete(workspaceFiles)
244251
.where(
245252
and(
@@ -342,7 +349,7 @@ async function hardDeleteKnowledgeBaseDocuments(
342349
label: string
343350
): Promise<void> {
344351
for (let batch = 0; batch < KB_DOCUMENT_DELETE_MAX_BATCHES; batch++) {
345-
const documentRows = await db
352+
const documentRows = await cleanupDb
346353
.select({ id: document.id })
347354
.from(document)
348355
.where(inArray(document.knowledgeBaseId, knowledgeBaseIds))
@@ -357,7 +364,7 @@ async function hardDeleteKnowledgeBaseDocuments(
357364
}
358365
}
359366

360-
const remaining = await db
367+
const remaining = await cleanupDb
361368
.select({ id: document.id })
362369
.from(document)
363370
.where(inArray(document.knowledgeBaseId, knowledgeBaseIds))
@@ -377,8 +384,9 @@ async function cleanupExpiredKnowledgeBases(
377384
workspaceIds,
378385
tableName: `${label}/knowledgeBase`,
379386
batchSize: KB_RETENTION_BATCH_SIZE,
387+
dbClient: cleanupDb,
380388
selectChunk: (chunkIds, limit) =>
381-
db
389+
cleanupDb
382390
.select({ id: knowledgeBase.id })
383391
.from(knowledgeBase)
384392
.where(
@@ -457,7 +465,7 @@ async function cleanupOrphanedKnowledgeBaseBindings(
457465
KB_ORPHAN_BINDING_BATCH_SIZE,
458466
KB_ORPHAN_BINDING_TOTAL_LIMIT - attempted
459467
)
460-
const rows = await db
468+
const rows = await cleanupDb
461469
.select({ key: workspaceFiles.key })
462470
.from(workspaceFiles)
463471
.where(
@@ -535,7 +543,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
535543
// prematurely purge data.
536544
const [doomedWorkflows, fileScope, expiredSoftDeletedChats] = await Promise.all([
537545
selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
538-
db
546+
cleanupDb
539547
.select({ id: workflow.id })
540548
.from(workflow)
541549
.where(
@@ -549,7 +557,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
549557
),
550558
selectExpiredWorkspaceFiles(workspaceIds, retentionDate),
551559
selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
552-
db
560+
cleanupDb
553561
.select({ id: copilotChats.id })
554562
.from(copilotChats)
555563
.where(
@@ -570,7 +578,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
570578
const doomedChatIds = new Set(softDeletedChatIds)
571579
if (doomedWorkflowIds.length > 0) {
572580
const workflowChats = await selectRowsByIdChunks(doomedWorkflowIds, (chunkIds, chunkLimit) =>
573-
db
581+
cleanupDb
574582
.select({ id: copilotChats.id })
575583
.from(copilotChats)
576584
.where(inArray(copilotChats.workflowId, chunkIds))
@@ -595,7 +603,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
595603
// spares their backend data and files.
596604
for (const batch of chunkArray(doomedWorkflowIds, DEFAULT_DELETE_CHUNK_SIZE)) {
597605
try {
598-
const deleted = await db
606+
const deleted = await cleanupDb
599607
.delete(workflow)
600608
.where(
601609
and(
@@ -618,7 +626,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
618626
// also re-checks row existence before purging external data).
619627
for (const batch of chunkArray(softDeletedChatIds, DEFAULT_DELETE_CHUNK_SIZE)) {
620628
try {
621-
const deleted = await db
629+
const deleted = await cleanupDb
622630
.delete(copilotChats)
623631
.where(
624632
and(
@@ -667,6 +675,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
667675
retentionDate,
668676
tableName: `${label}/${target.name}`,
669677
requireTimestampNotNull: true,
678+
dbClient: cleanupDb,
670679
})
671680
totalDeleted += result.deleted
672681
}

apps/sim/background/cleanup-tasks.ts

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { db } from '@sim/db'
1+
import { dbFor } from '@sim/db'
22
import {
33
copilotAsyncToolCalls,
44
copilotChats,
@@ -22,6 +22,9 @@ import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup'
2222

2323
const logger = createLogger('CleanupTasks')
2424

25+
/** All cleanup queries run on the dedicated cleanup pool. */
26+
const cleanupDb = dbFor('cleanup')
27+
2528
/**
2629
* Delete copilot run checkpoints and async tool calls via join through copilotRuns.
2730
* These tables don't have a direct workspaceId — we find qualifying run IDs first.
@@ -47,7 +50,7 @@ async function cleanupRunChildren(
4750
if (workspaceIds.length === 0) return []
4851

4952
const runIds = await selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
50-
db
53+
cleanupDb
5154
.select({ id: copilotRuns.id })
5255
.from(copilotRuns)
5356
.where(
@@ -63,7 +66,9 @@ async function cleanupRunChildren(
6366
const ids = runIds.map((r) => r.id)
6467

6568
return Promise.all(
66-
RUN_CHILD_TABLES.map((t) => deleteRowsById(t.table, t.runIdCol, ids, `${label}/${t.name}`))
69+
RUN_CHILD_TABLES.map((t) =>
70+
deleteRowsById(t.table, t.runIdCol, ids, `${label}/${t.name}`, cleanupDb)
71+
)
6772
)
6873
}
6974

@@ -82,7 +87,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise<void>
8287
)
8388

8489
const doomedChats = await selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) =>
85-
db
90+
cleanupDb
8691
.select({ id: copilotChats.id })
8792
.from(copilotChats)
8893
.where(
@@ -110,6 +115,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise<void>
110115
workspaceIds,
111116
retentionDate,
112117
tableName: `${label}/copilotRuns`,
118+
dbClient: cleanupDb,
113119
})
114120

115121
// Delete copilot chats using the exact IDs collected above so the chat
@@ -122,7 +128,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise<void>
122128
const chatsResult = { deleted: 0, failed: 0 }
123129
for (const batch of chunkArray(doomedChatIds, DEFAULT_DELETE_CHUNK_SIZE)) {
124130
try {
125-
const deleted = await db
131+
const deleted = await cleanupDb
126132
.delete(copilotChats)
127133
.where(and(inArray(copilotChats.id, batch), lt(copilotChats.updatedAt, retentionDate)))
128134
.returning({ id: copilotChats.id })
@@ -141,6 +147,7 @@ export async function runCleanupTasks(payload: CleanupJobPayload): Promise<void>
141147
workspaceIds,
142148
retentionDate,
143149
tableName: `${label}/mothershipInboxTask`,
150+
dbClient: cleanupDb,
144151
})
145152

146153
const totalDeleted =

0 commit comments

Comments
 (0)