From a263c6debdb34c5f0974ff1a16b71f2558a3e22b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 2 Oct 2026 00:36:59 -0700 Subject: [PATCH 1/3] fix(knowledge): count knowledge-base list totals with a lateral aggregate The counted knowledge-base list left-joined documents under the reader's ACL predicate and grouped by base. When the planner overestimated how many active bases a workspace has (soft-deleted bases skew the estimate), it switched to a hash join fed by a bitmap over the document ACL GIN index. Every ordinary document carries the shared `ws` token, so that bitmap covered every tenant's documents just to total one base's. Totals are now a LATERAL aggregate correlated on the base, so each base is counted through its own knowledge_base_id index and the page limit stops counting past the page. Results are unchanged: same filters, same order, and an aggregate always yields one row, so empty bases stay at zero. The shared drizzle mock gains innerJoinLateral/leftJoinLateral and a count() that aliases like drizzle's SQL expression. --- apps/sim/lib/knowledge/service.test.ts | 2 +- apps/sim/lib/knowledge/service.ts | 33 ++++++++++++++------- packages/testing/src/mocks/database.mock.ts | 21 ++++++++++++- 3 files changed, 43 insertions(+), 13 deletions(-) diff --git a/apps/sim/lib/knowledge/service.test.ts b/apps/sim/lib/knowledge/service.test.ts index db333fd2ba2..6823ef63660 100644 --- a/apps/sim/lib/knowledge/service.test.ts +++ b/apps/sim/lib/knowledge/service.test.ts @@ -304,7 +304,7 @@ describe('knowledge base counts with live source permissions', () => { expect(dbChainMockFns.select).not.toHaveBeenCalledWith({ connectorId: schemaMock.knowledgeConnector.id, }) - expect(dbChainMockFns.groupBy).toHaveBeenCalledOnce() + expect(dbChainMockFns.groupBy).not.toHaveBeenCalled() }) it('does not retain stale totals when a live source no longer authorizes its documents', async () => { diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index f6e43d4bc4f..b34a3304b4a 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -200,25 +200,36 @@ async function readCountedKnowledgeBaseRows( Array> > { const scope = 'get' in access ? await access.get() : access - const query = db + /** + * A lateral aggregate, so each base is counted through its own `knowledge_base_id` index and + * never through a bitmap of every tenant's documents sharing the `ws` token. An aggregate always + * yields one row, so an empty base stays at zero, and the limit stops counting past the page. + */ + const totals = db .select({ - ...ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS, - tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number), - docCount: count(document.knowledgeBaseId), + tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)` + .mapWith(Number) + .as('readable_token_count'), + docCount: count().as('readable_doc_count'), }) - .from(knowledgeBase) - .leftJoin( - document, + .from(document) + .where( and( eq(document.knowledgeBaseId, knowledgeBase.id), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), + ...ACTIVE_DOCUMENT_CONDITIONS, knowledgeAccessCondition(scope) ) ) + .as('totals') + const query = db + .select({ + ...ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS, + tokenCount: totals.tokenCount, + docCount: totals.docCount, + }) + .from(knowledgeBase) + .innerJoinLateral(totals, sql`true`) .where(where) - .groupBy(knowledgeBase.id) .orderBy(...orderBy) const rows = limit === undefined ? await query : await query.limit(limit) diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index 14cd6f2fc37..2895ee2ba44 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -91,7 +91,16 @@ function createMockSqlOperators() { gte: vi.fn((a, b) => ({ type: 'gte', left: a, right: b })), lt: vi.fn((a, b) => ({ type: 'lt', left: a, right: b })), lte: vi.fn((a, b) => ({ type: 'lte', left: a, right: b })), - count: vi.fn((column) => ({ type: 'count', column })), + /** Drizzle's `count()` is an `SQL` expression, so it aliases and decodes like one. */ + count: vi.fn((column) => { + const expression = { + type: 'count', + column, + as: (alias: string) => ({ ...expression, alias }), + mapWith: (decoder: unknown) => ({ ...expression, decoder }), + } + return expression + }), avg: vi.fn((column) => ({ type: 'avg', column })), sum: vi.fn((column) => ({ type: 'sum', column })), min: vi.fn((column) => ({ type: 'min', column })), @@ -235,6 +244,8 @@ const asAlias = chainSpy() const forClause = chainSpy() const innerJoin = chainSpy() const leftJoin = chainSpy() +const innerJoinLateral = chainSpy() +const leftJoinLateral = chainSpy() const insert = chainSpy() const update = chainSpy() const set = chainSpy() @@ -334,6 +345,12 @@ const joinBuilder = (tables: unknown[], fields: SelectedFields = {}): any => { builder.leftJoin = spyOrDefault(leftJoin, (table: unknown) => joinBuilder([...tables, table], fields) ) + builder.innerJoinLateral = spyOrDefault(innerJoinLateral, (subquery: unknown) => + joinBuilder([...tables, subquery], fields) + ) + builder.leftJoinLateral = spyOrDefault(leftJoinLateral, (subquery: unknown) => + joinBuilder([...tables, subquery], fields) + ) return builder } @@ -359,6 +376,8 @@ export const dbChainMockFns = { returning, innerJoin, leftJoin, + innerJoinLateral, + leftJoinLateral, groupBy, having, as: asAlias, From 3fcf7e26ed5872a620a1e300a43cdda037e368bb Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 2 Oct 2026 00:59:16 -0700 Subject: [PATCH 2/3] test(knowledge): pin per-base document reads in the counted list plan EXPLAIN the exact counted-list SQL with nested loops disabled and require every document scan to be correlated on its knowledge base, so a set-based join over the ACL index fails the build. Drop the comment's page-limit claim, which is the planner's cost choice rather than a guarantee, and a redundant mock-call assertion. --- .../list-totals-plan.integration.ts | 131 ++++++++++++++++++ apps/sim/lib/knowledge/service.test.ts | 1 - apps/sim/lib/knowledge/service.ts | 2 +- 3 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts diff --git a/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts b/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts new file mode 100644 index 00000000000..10dcfec96a7 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts @@ -0,0 +1,131 @@ +/** + * The counted knowledge-base list must total each base through that base's own documents, never + * through a set-based join over the whole document table. The reader's ACL tokens are shared by + * every tenant, so a join the planner drives from the ACL index instead of `knowledge_base_id` + * reads every tenant's documents to total one base. Nested loops are disabled while planning so + * the planner reaches for exactly that join wherever the query shape still allows it. + */ +import type { Principal } from '@sim/auth/principal' +import * as schema from '@sim/db/schema' +import { document, workspace } from '@sim/db/schema' +import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' +import { withUtcTimestamps } from '@sim/db/timestamps' +import { generateId } from '@sim/utils/id' +import { eq } from 'drizzle-orm' +import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js' +import postgres from 'postgres' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' + +const database = vi.hoisted(() => ({ + current: undefined as PostgresJsDatabase | undefined, + /** Every statement the app issues, so the test can EXPLAIN the exact SQL it ran. */ + statements: [] as { sql: string; params: unknown[] }[], +})) +vi.mock('@sim/db', async (importOriginal) => ({ + ...(await importOriginal()), + get db() { + if (!database.current) throw new Error('Test database not initialized') + return database.current + }, +})) + +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { listKnowledgeBases } from '@/lib/knowledge/application/knowledge-bases' + +interface PlanNode { + 'Node Type': string + 'Relation Name'?: string + 'Index Cond'?: string + 'Recheck Cond'?: string + Filter?: string + Plans?: PlanNode[] +} + +function documentScans(node: PlanNode): PlanNode[] { + const own = node['Relation Name'] === 'document' ? [node] : [] + return [...own, ...(node.Plans ?? []).flatMap(documentScans)] +} + +const ids = createKnowledgeAclFixtureIds() +const reader: Principal = { kind: 'session', userId: ids.bobId, sessionId: 'fixture-reader' } +const connection = postgres( + readTestDatabaseUrl(), + withUtcTimestamps({ + max: 2, + prepare: false, + fetch_types: false, + connection: {}, + onnotice: () => {}, + }) +) + +describe('knowledge-base list totals plan', () => { + beforeAll(async () => { + vi.stubGlobal('fetch', async () => { + throw new Error('Unexpected provider request in knowledge-base totals plan test') + }) + database.current = drizzle(connection, { + schema, + logger: { + logQuery(sql, params) { + database.statements.push({ sql, params }) + }, + }, + }) + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + await database.current.insert(document).values({ + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + filename: 'Workspace handbook', + fileUrl: 'https://fixture.test/shared', + fileSize: 10, + mimeType: 'text/plain', + tokenCount: 10, + processingStatus: 'completed', + }) + }) + + afterAll(async () => { + await database.current?.delete(workspace).where(eq(workspace.id, ids.workspaceId)) + await connection.end() + vi.unstubAllGlobals() + }) + + it('totals each base through its own documents even when the planner prefers a hash join', async () => { + database.statements.length = 0 + const { knowledgeBases } = await listKnowledgeBases.execute({ + principal: reader, + input: { workspaceId: ids.workspaceId }, + }) + expect(knowledgeBases.map(({ knowledgeBase }) => knowledgeBase)).toEqual([ + expect.objectContaining({ id: ids.knowledgeBaseId, docCount: 1, tokenCount: 10 }), + ]) + + const counted = database.statements.filter( + ({ sql }) => + /^select .* from "knowledge_base" /.test(sql) && + sql.includes('"document"."knowledge_base_id" = "knowledge_base"."id"') + ) + expect(counted).toHaveLength(1) + const plan = await connection.begin(async (tx) => { + await tx`SET LOCAL enable_nestloop = off` + const [row] = await tx.unsafe( + `EXPLAIN (FORMAT JSON) ${counted[0].sql}`, + counted[0].params as Parameters[1] + ) + return (row['QUERY PLAN'] as { Plan: PlanNode }[])[0].Plan + }) + + const scans = documentScans(plan) + expect(scans.length).toBeGreaterThan(0) + for (const scan of scans) { + const conditions = [scan['Index Cond'], scan['Recheck Cond'], scan.Filter].join(' ') + expect(conditions, `${scan['Node Type']} on document`).toContain( + '(knowledge_base_id = knowledge_base.id)' + ) + } + }) +}) diff --git a/apps/sim/lib/knowledge/service.test.ts b/apps/sim/lib/knowledge/service.test.ts index 6823ef63660..6d76f82f89b 100644 --- a/apps/sim/lib/knowledge/service.test.ts +++ b/apps/sim/lib/knowledge/service.test.ts @@ -304,7 +304,6 @@ describe('knowledge base counts with live source permissions', () => { expect(dbChainMockFns.select).not.toHaveBeenCalledWith({ connectorId: schemaMock.knowledgeConnector.id, }) - expect(dbChainMockFns.groupBy).not.toHaveBeenCalled() }) it('does not retain stale totals when a live source no longer authorizes its documents', async () => { diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index b34a3304b4a..32bbb366615 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -203,7 +203,7 @@ async function readCountedKnowledgeBaseRows( /** * A lateral aggregate, so each base is counted through its own `knowledge_base_id` index and * never through a bitmap of every tenant's documents sharing the `ws` token. An aggregate always - * yields one row, so an empty base stays at zero, and the limit stops counting past the page. + * yields one row, so an empty base stays at zero. */ const totals = db .select({ From 4f72a21024cba313d99cec01a1176640b075f413 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 2 Oct 2026 01:09:07 -0700 Subject: [PATCH 3/3] test(knowledge): tear down every fixture row in the list totals plan test Delete the fixture organization and both users alongside the workspace, matching the v1 knowledge route test that seeds the same fixture; permissions cascade from the users. --- .../__integration__/list-totals-plan.integration.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts b/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts index 10dcfec96a7..b56acdda302 100644 --- a/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts @@ -7,11 +7,11 @@ */ import type { Principal } from '@sim/auth/principal' import * as schema from '@sim/db/schema' -import { document, workspace } from '@sim/db/schema' +import { document, organization, user, workspace } from '@sim/db/schema' import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' import { withUtcTimestamps } from '@sim/db/timestamps' import { generateId } from '@sim/utils/id' -import { eq } from 'drizzle-orm' +import { eq, inArray } from 'drizzle-orm' import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js' import postgres from 'postgres' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' @@ -90,6 +90,8 @@ describe('knowledge-base list totals plan', () => { afterAll(async () => { await database.current?.delete(workspace).where(eq(workspace.id, ids.workspaceId)) + await database.current?.delete(organization).where(eq(organization.id, ids.organizationId)) + await database.current?.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId])) await connection.end() vi.unstubAllGlobals() })