mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-19 09:30:49 +08:00
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:
co-authored by
Claude Fable 5
parent
e841e4e731
commit
70814bc4a2
@@ -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)
|
||||
})
|
||||
|
||||
|
||||
@@ -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[]) =>
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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 })
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user