fix(knowledge): avoid count-query fan-out for unpaged lists

Author: icecrasher321

apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts

@@ -18,6 +18,7 @@ import { getWorkspaceKnowledgeBases } from '@/lib/knowledge/service'
  */
 const ids = createKnowledgeAclFixtureIds()
 const foreign = createKnowledgeAclFixtureIds()
+const manyBases = createKnowledgeAclFixtureIds()
 const offPageId = generateId()
 const emptyId = generateId()
 const archivedId = generateId()
@@ -106,6 +107,16 @@ beforeAll(async () => {
       SELECT ${baseId} || '-' || n, ${baseId}, 'bulk.txt', 'https://fixture.invalid/bulk',
         1, 'text/plain', ARRAY['ws'], 1 FROM generate_series(1, 10000) AS n`)
   }
+  await seedKnowledgeAclFixture(manyBases, { connectorType: 'google_drive' })
+  await db.execute(sql`INSERT INTO knowledge_base (id, workspace_id, user_id, name)
+    SELECT ${manyBases.knowledgeBaseId} || '-' || n, ${manyBases.workspaceId},
+      ${manyBases.aliceId}, 'Scale fixture ' || n FROM generate_series(1, 10000) AS n`)
+  await db.execute(sql`INSERT INTO document
+    (id, knowledge_base_id, filename, file_url, file_size, mime_type, acl, token_count)
+    VALUES (${generateId()}, ${manyBases.knowledgeBaseId}, 'first.txt',
+      'https://fixture.invalid/first', 1, 'text/plain', ARRAY['ws'], 13),
+      (${generateId()}, ${`${manyBases.knowledgeBaseId}-10000`}, 'last.txt',
+      'https://fixture.invalid/last', 1, 'text/plain', ARRAY['ws'], 17)`)
   await db.execute(sql`ANALYZE knowledge_base`)
   await db.execute(sql`ANALYZE document`)
 }, 60_000)
@@ -114,7 +125,7 @@ afterAll(async () => {
   const reportPath = process.env.KNOWLEDGE_BASE_LIST_REPORT_PATH
   if (reportPath) writeFileSync(reportPath, JSON.stringify(reports, null, 2))
   try {
-    for (const fixture of [ids, foreign]) {
+    for (const fixture of [ids, foreign, manyBases]) {
       await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId))
       await db.delete(organization).where(eq(organization.id, fixture.organizationId))
       await db.delete(user).where(inArray(user.id, [fixture.aliceId, fixture.bobId]))
@@ -191,4 +202,34 @@ describe('knowledge base list counts on real Postgres', () => {
     })
     expect(archived.data.map(({ id }) => id)).toEqual([archivedId])
   })
+
+  it('counts a large unpaged workspace within a fixed database round-trip budget', async () => {
+    let documentQueries = 0
+    const previousDebug = db.$client.options.debug
+    db.$client.options.debug = (_connection, query) => {
+      if (query.includes('"document"')) documentQueries++
+    }
+    try {
+      const all = await getWorkspaceKnowledgeBases(manyBases.workspaceId, 'active', {
+        countsFor: access,
+      })
+      expect(all.data).toHaveLength(10001)
+      expect(all.nextCursorKeys).toBeNull()
+      expect(all.data.find((kb) => kb.id === manyBases.knowledgeBaseId)).toMatchObject({
+        docCount: 1,
+        tokenCount: 13,
+      })
+      expect(all.data.find((kb) => kb.id === `${manyBases.knowledgeBaseId}-10000`)).toMatchObject({
+        docCount: 1,
+        tokenCount: 17,
+      })
+      expect(all.data.reduce((total, kb) => total + kb.docCount, 0)).toBe(2)
+      expect(all.data.reduce((total, kb) => total + kb.tokenCount, 0)).toBe(30)
+      reports.push({ unpagedBases: all.data.length, documentQueries })
+      expect(documentQueries).toBeGreaterThan(0)
+      expect(documentQueries).toBeLessThan(10)
+    } finally {
+      db.$client.options.debug = previousDebug
+    }
+  })
 })

apps/sim/lib/knowledge/service.ts

@@ -2,7 +2,6 @@ import { db } from '@sim/db'
 import { document, knowledgeBase, knowledgeConnector, workspaceFiles } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getPostgresConstraintName, getPostgresErrorCode } from '@sim/utils/errors'
-import { chunkArray } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { filterUndefined } from '@sim/utils/object'
 import type { SQL } from 'drizzle-orm'
@@ -34,7 +33,7 @@ import { resourceScopeCondition } from '@/lib/core/resource-scope.server'
 import { generateRestoreName } from '@/lib/core/utils/restore-name'
 import { findActiveFolder, resolveRestoredFolderId } from '@/lib/folders/queries'
 import { isKnowledgeMemberAccessAvailable } from '@/lib/knowledge/access/availability'
-import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate'
+import { knowledgeAccessCondition, textArrayLiteral } from '@/lib/knowledge/access/predicate'
 import type { KnowledgeAccessProvider } from '@/lib/knowledge/access/types'
 import { mirrorsSourceAcls } from '@/lib/knowledge/connectors/access-modes'
 import {
@@ -53,7 +52,6 @@ import type {
 import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 
 const logger = createLogger('KnowledgeBaseService')
-const KNOWLEDGE_BASE_COUNT_BATCH_SIZE = 100
 
 /**
  * Every caller-fixable knowledge-base failure is an {@link OrchestrationError},
@@ -203,17 +201,14 @@ async function readCountedKnowledgeBaseRows(
 > {
   const scope = 'get' in access ? await access.get() : access
   const rows = await readKnowledgeBaseRows(where, orderBy, limit)
-  const counts = new Map<string, { docCount: number; tokenCount: number }>()
-  for (const batch of chunkArray(rows, KNOWLEDGE_BASE_COUNT_BATCH_SIZE)) {
-    const totals = await countDocumentsByKnowledgeBase(
-      inArray(
-        document.knowledgeBaseId,
-        batch.map((kb) => kb.id)
-      ),
-      knowledgeAccessCondition(scope)
-    )
-    for (const total of totals) counts.set(total.knowledgeBaseId, total)
-  }
+  const totals =
+    rows.length > 0
+      ? await countDocumentsByKnowledgeBase(
+          sql`${document.knowledgeBaseId} = ANY(${textArrayLiteral(rows.map((kb) => kb.id))})`,
+          knowledgeAccessCondition(scope)
+        )
+      : []
+  const counts = new Map(totals.map((total) => [total.knowledgeBaseId, total]))
 
   /**
    * The counts above already include everything the reader's stored ACL admits. Only a

apps/sim/lib/selectors/server/providers/powerbi.test.ts

@@ -136,7 +136,8 @@ describe('Power BI selector provider boundary', () => {
         manualWorkspaceId: 'stale-manual-workspace',
         [field]: workspace,
       }
-      const picker = PowerBIBlock.subBlocks.find((subBlock) => subBlock.id === 'datasetSelector')!
+      const picker = PowerBIBlock.subBlocks.find((subBlock) => subBlock.id === 'datasetSelector')
+      if (!picker) throw new Error('Power BI semantic-model selector is missing')
       const context = buildSelectorContextFromValues({
         selectorKey: 'powerbi.datasets',
         contextConfigs: getSelectorContextSubBlocks(PowerBIBlock.subBlocks, values),

apps/sim/tools/powerbi/__fixtures__/provider-fixture.ts

@@ -31,7 +31,7 @@ export type PowerBIFixtureScenario =
   | 'missing-identity'
   | 'missing-query'
 
-export interface PowerBIFixtureRequest {
+interface PowerBIFixtureRequest {
   method: string
   path: string
   body: unknown

apps/sim/tools/powerbi/execute-query.test.ts

@@ -2,6 +2,9 @@ import { jsonResponse } from '@sim/testing/helpers/http'
 import { describe, expect, it } from 'vitest'
 import { powerbiExecuteQueryTool } from '@/tools/powerbi/execute-query'
 
+const transformResponse = powerbiExecuteQueryTool.transformResponse
+if (!transformResponse) throw new Error('Power BI query tool must transform its response')
+
 const params = {
   accessToken: 'provider-token',
   groupId: 'workspace-id',
@@ -32,7 +35,7 @@ describe('Power BI DAX result handling', () => {
       ],
     })
 
-    const result = await powerbiExecuteQueryTool.transformResponse!(response, params)
+    const result = await transformResponse(response, params)
 
     expect(result).toMatchObject({
       success: false,
@@ -52,7 +55,7 @@ describe('Power BI DAX result handling', () => {
   })
 
   it('preserves provider column names, blank values and falsy values in complete results', async () => {
-    const result = await powerbiExecuteQueryTool.transformResponse!(
+    const result = await transformResponse(
       jsonResponse({
         results: [
           {
@@ -79,7 +82,7 @@ describe('Power BI DAX result handling', () => {
   })
 
   it('does not fabricate a query error message when Microsoft supplies only a code', async () => {
-    const result = await powerbiExecuteQueryTool.transformResponse!(
+    const result = await transformResponse(
       jsonResponse({ error: { code: 'DatasetExecuteQueriesError' } }),
       params
     )
@@ -98,7 +101,7 @@ describe('Power BI DAX result handling', () => {
   })
 
   it.each(['omitted', 'empty'])('accepts one result table with %s rows', async (kind) => {
-    const result = await powerbiExecuteQueryTool.transformResponse!(
+    const result = await transformResponse(
       jsonResponse({ results: [{ tables: [kind === 'empty' ? { rows: [] } : {}] }] }),
       params
     )
@@ -117,7 +120,7 @@ describe('Power BI DAX result handling', () => {
             { code: 'AnalysisServicesErrorCode', detail: { type: 1, value: '3238920194' } },
             { code: 'DetailsMessage', detail: { type: 1, value: 'The DAX query is invalid.' } },
           ]
-    const result = await powerbiExecuteQueryTool.transformResponse!(
+    const result = await transformResponse(
       jsonResponse({
         error: {
           code: 'DatasetExecuteQueriesError',
@@ -174,9 +177,7 @@ describe('Power BI DAX result handling', () => {
           : undefined
       )
 
-      await expect(powerbiExecuteQueryTool.transformResponse!(response, params)).rejects.toThrow(
-        /maximum size.*20971520/
-      )
+      await expect(transformResponse(response, params)).rejects.toThrow(/maximum size.*20971520/)
       expect(canceled).toBe(true)
     }
   )
@@ -191,7 +192,7 @@ describe('Power BI DAX result handling', () => {
         },
       })
     )
-    const transformed = powerbiExecuteQueryTool.transformResponse!(response, params, {
+    const transformed = transformResponse(response, params, {
       signal: abort.signal,
     })
     abort.abort(new Error('Execution canceled'))

apps/sim/tools/powerbi/types.ts

@@ -1,6 +1,6 @@
 import type { OutputProperty, ToolResponse } from '@/tools/types'
 
-export interface PowerBIAuthParams {
+interface PowerBIAuthParams {
   accessToken: string
 }
 
@@ -27,7 +27,7 @@ export interface PowerBIExecuteQueryParams extends PowerBIDatasetParams {
   includeNulls?: boolean
 }
 
-export type PowerBINotifyOption = 'NoNotification' | 'MailOnFailure' | 'MailOnCompletion'
+type PowerBINotifyOption = 'NoNotification' | 'MailOnFailure' | 'MailOnCompletion'
 
 export interface PowerBIRefreshDatasetParams extends PowerBIDatasetParams {
   notifyOption?: PowerBINotifyOption
@@ -71,7 +71,7 @@ export interface PowerBIDataset {
   webUrl: string | null
 }
 
-export interface PowerBIRefreshAttempt {
+interface PowerBIRefreshAttempt {
   attemptId: number | null
   type: string | null
   startTime: string | null
@@ -96,7 +96,7 @@ export interface PowerBIQueryError {
   details: unknown | null
 }
 
-export interface PowerBIInformationProtectionLabel {
+interface PowerBIInformationProtectionLabel {
   id: string | null
   name: string | null
 }
@@ -206,21 +206,6 @@ export const POWERBI_DATASET_OUTPUT_PROPERTIES = {
   webUrl: { ...nullableString, description: 'Semantic model URL in Power BI, when available' },
 } satisfies Record<string, OutputProperty>
 
-export const POWERBI_REFRESH_ATTEMPT_OUTPUT_PROPERTIES = {
-  attemptId: {
-    type: 'number',
-    nullable: true,
-    description: 'Refresh attempt index',
-  },
-  type: { ...nullableString, description: 'Provider refresh attempt type' },
-  startTime: { ...nullableString, description: 'Attempt start timestamp' },
-  endTime: { ...nullableString, description: 'Attempt end timestamp, when available' },
-  serviceExceptionJson: {
-    ...nullableString,
-    description: 'Serialized provider failure details, when available',
-  },
-} satisfies Record<string, OutputProperty>
-
 export const POWERBI_REFRESH_OUTPUT_PROPERTIES = {
   requestId: { ...nullableString, description: 'Provider refresh request ID' },
   refreshType: { ...nullableString, description: 'Provider refresh trigger type' },
@@ -237,7 +222,23 @@ export const POWERBI_REFRESH_OUTPUT_PROPERTIES = {
   refreshAttempts: {
     type: 'array',
     description: 'Refresh attempts supplied by the provider',
-    items: { type: 'object', properties: POWERBI_REFRESH_ATTEMPT_OUTPUT_PROPERTIES },
+    items: {
+      type: 'object',
+      properties: {
+        attemptId: {
+          type: 'number',
+          nullable: true,
+          description: 'Refresh attempt index',
+        },
+        type: { ...nullableString, description: 'Provider refresh attempt type' },
+        startTime: { ...nullableString, description: 'Attempt start timestamp' },
+        endTime: { ...nullableString, description: 'Attempt end timestamp, when available' },
+        serviceExceptionJson: {
+          ...nullableString,
+          description: 'Serialized provider failure details, when available',
+        },
+      },
+    },
   },
 } satisfies Record<string, OutputProperty>
 

apps/sim/tools/powerbi/utils.ts

@@ -10,7 +10,7 @@ import type {
 } from '@/tools/powerbi/types'
 import { safeUrlPathSegment } from '@/tools/url-path'
 
-export const MAX_POWERBI_RESPONSE_BYTES = 20 * 1024 * 1024
+const MAX_POWERBI_RESPONSE_BYTES = 20 * 1024 * 1024
 
 export const POWERBI_ACCESS_TOKEN_PARAM = {
   type: 'string',