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>
This commit is contained in:
Theodore Li
2026-07-22 14:35:42 -04:00
committed by GitHub
co-authored by Claude Fable 5
parent e841e4e731
commit 70814bc4a2
18 changed files with 291 additions and 119 deletions
+10 -5
View File
@@ -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)
})
+15 -6
View File
@@ -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<void>
workspaceIds,
retentionDate,
tableName: `${label}/job_execution_logs`,
dbClient: cleanupDb,
})
if (runGlobalHousekeeping && plan === 'free') {
@@ -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[]) =>
+23 -14
View File
@@ -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<WorkspaceFileScope> {
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<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))
@@ -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
}
+12 -5
View File
@@ -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<void>
)
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<void>
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<void>
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<void>
workspaceIds,
retentionDate,
tableName: `${label}/mothershipInboxTask`,
dbClient: cleanupDb,
})
const totalDeleted =
+18 -3
View File
@@ -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 })
+7 -4
View File
@@ -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<FileRef[]> {
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<FileRef[]> {
),
// 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<string>()
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))
+2
View File
@@ -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)
@@ -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: {
@@ -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<void> {
export async function markLargeValuesDeleted(
keys: string[],
dbClient: LargeValueMetadataClient = db
): Promise<void> {
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<number> {
const rows = await db.execute<{ count: number }>(sql`
async function pruneStaleReferences(
workspaceIds: string[],
batchSize: number,
dbClient: LargeValueMetadataClient
): Promise<number> {
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<number> {
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<number> {
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<LargeValueMetadataPruneResult> {
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
)
}
+10 -7
View File
@@ -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<unknown>) => cb(tx)),
}
return {
db: {
select: dbSelectMock,
insert: vi.fn(),
update: vi.fn(),
execute: dbExecuteMock,
transaction: vi.fn(async (cb: (txArg: typeof tx) => Promise<unknown>) => cb(tx)),
},
db,
// Exec-pool client shares the instance so call-order seeding still applies.
dbFor: () => db,
}
})
+13 -6
View File
@@ -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<WorkflowExecutionLog | null> {
const [workflowLog] = await db
const [workflowLog] = await execDb
.select()
.from(workflowExecutionLogs)
.where(eq(workflowExecutionLogs.executionId, executionId))
@@ -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,
+10 -7
View File
@@ -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(
@@ -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<WorkflowExecutionSnapshot | null> {
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<number> {
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))
@@ -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<EnqueueResumeResult> {
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<void> {
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<void> {
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<void> {
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<boolean> {
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<boolean> {
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<boolean> {
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<void> {
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<void> {
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<void> {
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<typeof pausedExecutions.$inferSelect | null> {
const rows = await db
const rows = await execDb
.select()
.from(pausedExecutions)
.where(eq(pausedExecutions.id, id))
@@ -2292,7 +2299,7 @@ export class PauseResumeManager {
}): Promise<PausedExecutionDetail | null> {
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)
+69 -4
View File
@@ -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<keyof typeof DB_POOL_PROFILES, ProcessDbRole>
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<SubProcessDbRole, typeof db>()
/** 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_<ROLE>` with fallback to the URL the process itself
* resolved (`DATABASE_URL_<PROCESSROLE>`, 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
}
@@ -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: <T>(fn: () => T): T => fn(),
instrumentPoolClient: <T>(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: <T>(fn: () => T): T => fn(),
instrumentPoolClient: <T>(client: T): T => client,