improvement(knowledge): fold processDocumentAsync prefetches into one JOIN

processDocumentAsync was issuing three separate SELECTs per processed
document: knowledge_base (config), workspace (billing settings), and
document (tag values). For a typical Trigger.dev fleet processing
thousands of docs, that's thousands of redundant pool checkouts.

Collapsed into a single JOIN at the top of processDocumentAsync that
fetches kb config + billed account user + document tag values in one
roundtrip. The post-embedding tag SELECT (which previously held tags
through the full embedding-generation wait) is gone; tags from the
initial prefetch are reused.

Behavior:
- Missing/archived/deleted document or KB → same 'failed' status outcome
  as before, single consolidated error message.
- Missing billed account → preserves existing error.
- All 208 KB tests pass (test mock extended for innerJoin/leftJoin).

Author: waleedlatif1

apps/sim/app/api/knowledge/utils.test.ts

@@ -116,6 +116,31 @@ vi.mock('@sim/db', async () => {
             },
           }
         },
+        innerJoin() {
+          // document × knowledge_base context JOIN — return the first kb and
+          // doc row merged (covers processDocumentAsync's prefetch).
+          return {
+            leftJoin: () => ({
+              where: () => ({
+                limit: (n: number) =>
+                  Promise.resolve(
+                    kbRows.length > 0 && docRows.length > 0
+                      ? [{ ...kbRows[0], ...docRows[0], billedAccountUserId: 'billing-user-1' }]
+                          .slice(0, n)
+                      : []
+                  ),
+              }),
+            }),
+            where: () => ({
+              limit: (n: number) =>
+                Promise.resolve(
+                  kbRows.length > 0 && docRows.length > 0
+                    ? [{ ...kbRows[0], ...docRows[0] }].slice(0, n)
+                    : []
+                ),
+            }),
+          }
+        },
       }
     },
   }

apps/sim/lib/knowledge/documents/service.ts

@@ -5,6 +5,7 @@ import {
   knowledgeBase,
   knowledgeBaseTagDefinitions,
   knowledgeConnector,
+  workspace as workspaceTable,
 } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { sha256Hex } from '@sim/security/hash'
@@ -47,7 +48,6 @@ import type { ProcessedDocumentTags } from '@/lib/knowledge/types'
 import { estimateTokenCount } from '@/lib/tokenization/estimators'
 import { deleteFile } from '@/lib/uploads/core/storage-service'
 import { extractStorageKey } from '@/lib/uploads/utils/file-utils'
-import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
 import type {
   DocumentProcessingPayload,
   processDocument as processDocumentTask,
@@ -426,34 +426,68 @@ export async function processDocumentAsync(
   try {
     logger.info(`[${documentId}] Starting document processing: ${docData.filename}`)
 
-    const kb = await db
+    // Single JOIN fetches everything we'll need: KB processing config, workspace
+    // billing context, and the document's tag values (folded onto the embedding
+    // rows later). Replaces three separate SELECTs.
+    const contextRows = await db
       .select({
         userId: knowledgeBase.userId,
         workspaceId: knowledgeBase.workspaceId,
         chunkingConfig: knowledgeBase.chunkingConfig,
         embeddingModel: knowledgeBase.embeddingModel,
+        billedAccountUserId: workspaceTable.billedAccountUserId,
+        tag1: document.tag1,
+        tag2: document.tag2,
+        tag3: document.tag3,
+        tag4: document.tag4,
+        tag5: document.tag5,
+        tag6: document.tag6,
+        tag7: document.tag7,
+        number1: document.number1,
+        number2: document.number2,
+        number3: document.number3,
+        number4: document.number4,
+        number5: document.number5,
+        date1: document.date1,
+        date2: document.date2,
+        boolean1: document.boolean1,
+        boolean2: document.boolean2,
+        boolean3: document.boolean3,
       })
-      .from(knowledgeBase)
-      .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
+      .from(document)
+      .innerJoin(knowledgeBase, eq(knowledgeBase.id, document.knowledgeBaseId))
+      .leftJoin(
+        workspaceTable,
+        and(eq(workspaceTable.id, knowledgeBase.workspaceId), isNull(workspaceTable.archivedAt))
+      )
+      .where(
+        and(
+          eq(document.id, documentId),
+          eq(knowledgeBase.id, knowledgeBaseId),
+          isNull(document.archivedAt),
+          isNull(document.deletedAt),
+          isNull(knowledgeBase.deletedAt)
+        )
+      )
       .limit(1)
 
-    if (kb.length === 0) {
+    if (contextRows.length === 0) {
       logger.warn(
-        `[${documentId}] Skipping document processing: knowledge base ${knowledgeBaseId} is deleted`
+        `[${documentId}] Skipping document processing: document or knowledge base ${knowledgeBaseId} no longer exists`
       )
       await db
         .update(document)
         .set({
           processingStatus: 'failed',
-          processingError: 'Knowledge base deleted',
+          processingError: 'Document or knowledge base no longer exists',
           processingCompletedAt: new Date(),
         })
-        .where(
-          and(eq(document.id, documentId), isNull(document.archivedAt), isNull(document.deletedAt))
-        )
+        .where(eq(document.id, documentId))
       return
     }
 
+    const ctx = contextRows[0]
+
     await db
       .update(document)
       .set({
@@ -468,7 +502,7 @@ export async function processDocumentAsync(
 
     logger.info(`[${documentId}] Status updated to 'processing', starting document processor`)
 
-    const rawConfig = kb[0].chunkingConfig as {
+    const rawConfig = ctx.chunkingConfig as {
       maxSize?: number
       minSize?: number
       overlap?: number
@@ -481,13 +515,13 @@ export async function processDocumentAsync(
       overlap: rawConfig?.overlap ?? 200,
     }
 
-    const kbEmbeddingModel = kb[0].embeddingModel
-    if (!kb[0].workspaceId) {
+    const kbEmbeddingModel = ctx.embeddingModel
+    if (!ctx.workspaceId) {
       throw new Error(`Knowledge base ${knowledgeBaseId} is missing workspace billing context`)
     }
-    const billingUserId = await getWorkspaceBilledAccountUserId(kb[0].workspaceId)
+    const billingUserId = ctx.billedAccountUserId
     if (!billingUserId) {
-      throw new Error(`Workspace ${kb[0].workspaceId} is missing billed account`)
+      throw new Error(`Workspace ${ctx.workspaceId} is missing billed account`)
     }
     let totalEmbeddingTokens = 0
     let embeddingIsBYOK = false
@@ -503,8 +537,8 @@ export async function processDocumentAsync(
           kbConfig.maxSize,
           kbConfig.overlap,
           kbConfig.minSize,
-          kb[0].userId,
-          kb[0].workspaceId,
+          ctx.userId,
+          ctx.workspaceId,
           rawConfig?.strategy,
           rawConfig?.strategyOptions
         )
@@ -542,7 +576,7 @@ export async function processDocumentAsync(
               isBYOK,
               modelName,
               pricingId,
-            } = await generateEmbeddings(batch, kbEmbeddingModel, kb[0].workspaceId)
+            } = await generateEmbeddings(batch, kbEmbeddingModel, ctx.workspaceId)
             for (const emb of batchEmbeddings) {
               embeddings.push(emb)
             }
@@ -555,41 +589,11 @@ export async function processDocumentAsync(
           }
         }
 
-        logger.info(`[${documentId}] Embeddings generated, fetching document tags`)
-
-        const documentRecord = await db
-          .select({
-            tag1: document.tag1,
-            tag2: document.tag2,
-            tag3: document.tag3,
-            tag4: document.tag4,
-            tag5: document.tag5,
-            tag6: document.tag6,
-            tag7: document.tag7,
-            number1: document.number1,
-            number2: document.number2,
-            number3: document.number3,
-            number4: document.number4,
-            number5: document.number5,
-            date1: document.date1,
-            date2: document.date2,
-            boolean1: document.boolean1,
-            boolean2: document.boolean2,
-            boolean3: document.boolean3,
-          })
-          .from(document)
-          .where(
-            and(
-              eq(document.id, documentId),
-              isNull(document.archivedAt),
-              isNull(document.deletedAt)
-            )
-          )
-          .limit(1)
-
-        const documentTags = documentRecord[0] || {}
+        // Document tag values are stable from upload time and were prefetched in
+        // the JOIN above — reuse them instead of issuing another SELECT.
+        const documentTags = ctx
 
-        logger.info(`[${documentId}] Creating embedding records with tags`)
+        logger.info(`[${documentId}] Embeddings generated, creating embedding records with tags`)
 
         const tokenizerProvider = getEmbeddingModelInfo(kbEmbeddingModel).tokenizerProvider
 
@@ -694,7 +698,7 @@ export async function processDocumentAsync(
         if (cost > 0) {
           await recordUsage({
             userId: billingUserId,
-            workspaceId: kb[0].workspaceId ?? undefined,
+            workspaceId: ctx.workspaceId ?? undefined,
             entries: [
               {
                 category: 'model',