improvement(db): add session statement/lock timeouts; simplify KB doc tx

TheodoreSpeaks · TheodoreSpeaks · commit 0ad209ffb73d · 2026-05-13T19:46:18.000-07:00
diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts
@@ -766,99 +766,99 @@ export async function createDocumentRecords(
   knowledgeBaseId: string,
   requestId: string
 ): Promise<DocumentData[]> {
-  return await db.transaction(async (tx) => {
-    await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
-
-    const kb = await tx
-      .select({ id: knowledgeBase.id })
-      .from(knowledgeBase)
-      .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
-      .limit(1)
+  // No tx wrapper: the bulk `db.insert(...).values([...])` is a single statement
+  // and atomic by Postgres. The KB FK constraint fails loud if the KB is
+  // concurrently deleted, so an explicit FOR UPDATE lock is unnecessary and
+  // doubles per-call pool checkouts.
+  const kb = await db
+    .select({ id: knowledgeBase.id })
+    .from(knowledgeBase)
+    .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
+    .limit(1)
 
-    if (kb.length === 0) {
-      throw new Error('Knowledge base not found')
-    }
+  if (kb.length === 0) {
+    throw new Error('Knowledge base not found')
+  }
 
-    const now = new Date()
-    const documentRecords = []
-    const returnData: DocumentData[] = []
+  const now = new Date()
+  const documentRecords = []
+  const returnData: DocumentData[] = []
 
-    for (const docData of documents) {
-      const documentId = generateId()
+  for (const docData of documents) {
+    const documentId = generateId()
 
-      let processedTags: Partial<ProcessedDocumentTags> = {}
+    let processedTags: Partial<ProcessedDocumentTags> = {}
 
-      if (docData.documentTagsData) {
-        try {
-          const tagData = JSON.parse(docData.documentTagsData)
-          if (Array.isArray(tagData)) {
-            processedTags = await processDocumentTags(knowledgeBaseId, tagData, requestId)
-          }
-        } catch (error) {
-          if (error instanceof SyntaxError) {
-            logger.warn(`[${requestId}] Failed to parse documentTagsData for bulk document:`, error)
-          } else {
-            throw error
-          }
+    if (docData.documentTagsData) {
+      try {
+        const tagData = JSON.parse(docData.documentTagsData)
+        if (Array.isArray(tagData)) {
+          processedTags = await processDocumentTags(knowledgeBaseId, tagData, requestId)
+        }
+      } catch (error) {
+        if (error instanceof SyntaxError) {
+          logger.warn(`[${requestId}] Failed to parse documentTagsData for bulk document:`, error)
+        } else {
+          throw error
         }
       }
+    }
 
-      const newDocument = {
-        id: documentId,
-        knowledgeBaseId,
-        filename: docData.filename,
-        fileUrl: docData.fileUrl,
-        fileSize: docData.fileSize,
-        mimeType: docData.mimeType,
-        chunkCount: 0,
-        tokenCount: 0,
-        characterCount: 0,
-        processingStatus: 'pending' as const,
-        enabled: true,
-        uploadedAt: now,
-        tag1: processedTags.tag1 ?? docData.tag1 ?? null,
-        tag2: processedTags.tag2 ?? docData.tag2 ?? null,
-        tag3: processedTags.tag3 ?? docData.tag3 ?? null,
-        tag4: processedTags.tag4 ?? docData.tag4 ?? null,
-        tag5: processedTags.tag5 ?? docData.tag5 ?? null,
-        tag6: processedTags.tag6 ?? docData.tag6 ?? null,
-        tag7: processedTags.tag7 ?? docData.tag7 ?? null,
-        number1: processedTags.number1 ?? null,
-        number2: processedTags.number2 ?? null,
-        number3: processedTags.number3 ?? null,
-        number4: processedTags.number4 ?? null,
-        number5: processedTags.number5 ?? null,
-        date1: processedTags.date1 ?? null,
-        date2: processedTags.date2 ?? null,
-        boolean1: processedTags.boolean1 ?? null,
-        boolean2: processedTags.boolean2 ?? null,
-        boolean3: processedTags.boolean3 ?? null,
-      }
-
-      documentRecords.push(newDocument)
-      returnData.push({
-        documentId,
-        filename: docData.filename,
-        fileUrl: docData.fileUrl,
-        fileSize: docData.fileSize,
-        mimeType: docData.mimeType,
-      })
+    const newDocument = {
+      id: documentId,
+      knowledgeBaseId,
+      filename: docData.filename,
+      fileUrl: docData.fileUrl,
+      fileSize: docData.fileSize,
+      mimeType: docData.mimeType,
+      chunkCount: 0,
+      tokenCount: 0,
+      characterCount: 0,
+      processingStatus: 'pending' as const,
+      enabled: true,
+      uploadedAt: now,
+      tag1: processedTags.tag1 ?? docData.tag1 ?? null,
+      tag2: processedTags.tag2 ?? docData.tag2 ?? null,
+      tag3: processedTags.tag3 ?? docData.tag3 ?? null,
+      tag4: processedTags.tag4 ?? docData.tag4 ?? null,
+      tag5: processedTags.tag5 ?? docData.tag5 ?? null,
+      tag6: processedTags.tag6 ?? docData.tag6 ?? null,
+      tag7: processedTags.tag7 ?? docData.tag7 ?? null,
+      number1: processedTags.number1 ?? null,
+      number2: processedTags.number2 ?? null,
+      number3: processedTags.number3 ?? null,
+      number4: processedTags.number4 ?? null,
+      number5: processedTags.number5 ?? null,
+      date1: processedTags.date1 ?? null,
+      date2: processedTags.date2 ?? null,
+      boolean1: processedTags.boolean1 ?? null,
+      boolean2: processedTags.boolean2 ?? null,
+      boolean3: processedTags.boolean3 ?? null,
     }
 
-    if (documentRecords.length > 0) {
-      await tx.insert(document).values(documentRecords)
-      logger.info(
-        `[${requestId}] Bulk created ${documentRecords.length} document records in knowledge base ${knowledgeBaseId}`
-      )
+    documentRecords.push(newDocument)
+    returnData.push({
+      documentId,
+      filename: docData.filename,
+      fileUrl: docData.fileUrl,
+      fileSize: docData.fileSize,
+      mimeType: docData.mimeType,
+    })
+  }
 
-      await tx
-        .update(knowledgeBase)
-        .set({ updatedAt: now })
-        .where(eq(knowledgeBase.id, knowledgeBaseId))
-    }
+  if (documentRecords.length > 0) {
+    await db.insert(document).values(documentRecords)
+    logger.info(
+      `[${requestId}] Bulk created ${documentRecords.length} document records in knowledge base ${knowledgeBaseId}`
+    )
 
-    return returnData
-  })
+    await db
+      .update(knowledgeBase)
+      .set({ updatedAt: now })
+      .where(eq(knowledgeBase.id, knowledgeBaseId))
+  }
+
+  return returnData
 }
 
 export interface TagFilterCondition {
@@ -1312,26 +1312,23 @@ export async function createSingleDocument(
     ...processedTags,
   }
 
-  await db.transaction(async (tx) => {
-    await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
-
-    const kb = await tx
-      .select({ id: knowledgeBase.id })
-      .from(knowledgeBase)
-      .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
-      .limit(1)
+  // No tx wrapper: single insert is atomic; KB FK fails loud on concurrent delete.
+  const kb = await db
+    .select({ id: knowledgeBase.id })
+    .from(knowledgeBase)
+    .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
+    .limit(1)
 
-    if (kb.length === 0) {
-      throw new Error('Knowledge base not found')
-    }
+  if (kb.length === 0) {
+    throw new Error('Knowledge base not found')
+  }
 
-    await tx.insert(document).values(newDocument)
+  await db.insert(document).values(newDocument)
 
-    await tx
-      .update(knowledgeBase)
-      .set({ updatedAt: now })
-      .where(eq(knowledgeBase.id, knowledgeBaseId))
-  })
+  await db
+    .update(knowledgeBase)
+    .set({ updatedAt: now })
+    .where(eq(knowledgeBase.id, knowledgeBaseId))
   logger.info(`[${requestId}] Document created: ${documentId} in knowledge base ${knowledgeBaseId}`)
 
   return newDocument as {
diff --git a/apps/sim/lib/workspaces/lifecycle.test.ts b/apps/sim/lib/workspaces/lifecycle.test.ts
@@ -55,6 +55,7 @@ describe('workspace lifecycle', () => {
     })
 
     const tx = {
+      execute: vi.fn().mockResolvedValue([]),
       select: vi.fn().mockReturnValue({
         from: vi.fn().mockReturnValue({
           where: vi.fn().mockResolvedValue([{ id: 'kb-1' }]),
diff --git a/apps/sim/lib/workspaces/lifecycle.ts b/apps/sim/lib/workspaces/lifecycle.ts
@@ -49,6 +49,13 @@ export async function archiveWorkspace(
     .where(eq(workflowMcpServer.workspaceId, workspaceId))
 
   await db.transaction(async (tx) => {
+    // Workspace archival is a rare admin/cleanup operation that touches every
+    // child table; on large workspaces it can exceed the 30s session default.
+    // Override per-tx with a generous ceiling — if it ever runs longer than
+    // this something is genuinely wrong.
+    await tx.execute(sql`SET LOCAL statement_timeout = '5min'`)
+    await tx.execute(sql`SET LOCAL lock_timeout = '30s'`)
+
     await tx
       .update(knowledgeBase)
       .set({
diff --git a/packages/db/db.ts b/packages/db/db.ts
@@ -13,6 +13,14 @@ const postgresClient = postgres(connectionString, {
   connect_timeout: 30,
   max: 30,
   onnotice: () => {},
+  // Server-side guards. lock_timeout cancels a query waiting on a row lock for
+  // >5s (e.g. another tx holding `SELECT ... FOR UPDATE`). statement_timeout
+  // cancels any query running >30s. Heavy paths that legitimately need longer
+  // (table service bulk JSONB rewrites) override per-tx with `SET LOCAL`.
+  connection: {
+    lock_timeout: 5_000,
+    statement_timeout: 30_000,
+  },
 })
 
 export const db = drizzle(postgresClient, { schema })
