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..b56acdda302 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/list-totals-plan.integration.ts @@ -0,0 +1,133 @@ +/** + * 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, 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, 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' + +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 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() + }) + + 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 db333fd2ba2..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).toHaveBeenCalledOnce() }) 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..32bbb366615 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. + */ + 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,