Skip to content

Commit edb7799

Browse files
authored
fix(knowledge,files): count knowledge base documents per base and index the modified-order file list (#10026)
1 parent 377e81b commit edb7799

12 files changed

Lines changed: 33369 additions & 52 deletions

File tree

‎apps/sim/lib/billing/credits/weekly-refresh.test.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import {
1414
const FROZEN_NOW = new Date('2026-08-15T00:00:00.000Z')
1515

1616
beforeEach(() => {
17-
drizzleOrmMock.sum.mockImplementation(() => ({ as: () => 'sum' }))
1817
vi.useFakeTimers()
1918
vi.setSystemTime(FROZEN_NOW)
2019
})

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

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { writeFileSync } from 'node:fs'
22
import { db } from '@sim/db'
33
import { document, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
4+
import { explainCaptured } from '@sim/db/testing/explain'
45
import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
56
import { generateId } from '@sim/utils/id'
67
import { eq, inArray, sql } from 'drizzle-orm'
@@ -14,12 +15,14 @@ import { getWorkspaceKnowledgeBases } from '@/lib/knowledge/service'
1415

1516
/**
1617
* A small KB page must not count the rest of its workspace or other tenants before applying
17-
* its limit. Real query plans catch this even when warm caches hide it from a timing test.
18+
* its limit, and many bases must not read the documents other tenants share a token with. Real query plans catch this even when warm caches hide it from a timing test.
1819
* The same read must retain ACL/lifecycle filtering, empty bases, and keyset continuity.
1920
*/
2021
const ids = createKnowledgeAclFixtureIds()
2122
const foreign = createKnowledgeAclFixtureIds()
2223
const manyBases = createKnowledgeAclFixtureIds()
24+
const archivedBases = createKnowledgeAclFixtureIds()
25+
const ARCHIVED_BASE_COUNT = 600
2326
const offPageId = generateId()
2427
const emptyId = generateId()
2528
const archivedId = generateId()
@@ -118,6 +121,16 @@ beforeAll(async () => {
118121
'https://fixture.invalid/first', 1, 'text/plain', ARRAY['ws'], 13),
119122
(${generateId()}, ${`${manyBases.knowledgeBaseId}-10000`}, 'last.txt',
120123
'https://fixture.invalid/last', 1, 'text/plain', ARRAY['ws'], 17)`)
124+
await seedKnowledgeAclFixture(archivedBases, { connectorType: 'google_drive' })
125+
await db.execute(sql`INSERT INTO knowledge_base (id, workspace_id, user_id, name, deleted_at)
126+
SELECT ${archivedBases.knowledgeBaseId} || '-' || n, ${archivedBases.workspaceId},
127+
${archivedBases.aliceId}, 'Archived fixture ' || n, now()
128+
FROM generate_series(1, ${ARCHIVED_BASE_COUNT}) AS n`)
129+
await db.execute(sql`INSERT INTO document
130+
(id, knowledge_base_id, filename, file_url, file_size, mime_type, acl, token_count, archived_at)
131+
SELECT ${archivedBases.knowledgeBaseId} || '-doc-' || n, ${archivedBases.knowledgeBaseId} || '-' || n,
132+
'archived.txt', 'https://fixture.invalid/archived', 1, 'text/plain', ARRAY['ws'], 1, now()
133+
FROM generate_series(1, ${ARCHIVED_BASE_COUNT}) AS n`)
121134
await db.execute(sql`ANALYZE knowledge_base`)
122135
await db.execute(sql`ANALYZE document`)
123136
}, 60_000)
@@ -126,7 +139,7 @@ afterAll(async () => {
126139
const reportPath = process.env.KNOWLEDGE_BASE_LIST_REPORT_PATH
127140
if (reportPath) writeFileSync(reportPath, JSON.stringify(reports, null, 2))
128141
try {
129-
for (const fixture of [ids, foreign, manyBases]) {
142+
for (const fixture of [ids, foreign, manyBases, archivedBases]) {
130143
await deleteWorkspaceFixture(db, eq(workspace.id, fixture.workspaceId))
131144
await db.delete(organization).where(eq(organization.id, fixture.organizationId))
132145
await db.delete(user).where(inArray(user.id, [fixture.aliceId, fixture.bobId]))
@@ -172,6 +185,34 @@ describe('knowledge base list counts on real Postgres', () => {
172185
}
173186
)
174187

188+
/**
189+
* Bases whose documents are all archived, beside thousands of other tenants' live `ws`
190+
* documents. No small fixture reproduces the production estimates that once turned this count
191+
* into a scan of every tenant's documents (see `countDocumentsByKnowledgeBase`), so this pins
192+
* the shape that rules that plan out: each index read of `document` seeks one base.
193+
*/
194+
it('counts many archived bases one base at a time', async () => {
195+
const list = await getWorkspaceKnowledgeBases(archivedBases.workspaceId, 'archived', {
196+
countsFor: access,
197+
})
198+
expect(list.data).toHaveLength(ARCHIVED_BASE_COUNT)
199+
expect(list.data.every((kb) => kb.docCount === 0 && kb.tokenCount === 0)).toBe(true)
200+
const nodes = await explainCaptured(
201+
db,
202+
() =>
203+
getWorkspaceKnowledgeBases(archivedBases.workspaceId, 'archived', { countsFor: access }),
204+
{ match: /"document"/ }
205+
)
206+
expect(
207+
nodes.some((node) => node['Node Type'] === 'Seq Scan' && node['Relation Name'] === 'document')
208+
).toBe(false)
209+
const documentIndexReads = nodes.filter((node) => node['Index Name']?.startsWith('doc_'))
210+
expect(documentIndexReads.length).toBeGreaterThan(0)
211+
expect(documentIndexReads.map((node) => node['Index Cond'])).toEqual(
212+
documentIndexReads.map(() => '(knowledge_base_id = knowledge_base.id)')
213+
)
214+
})
215+
175216
it('keeps empty KBs and count visibility through pagination and unpaged reads', async () => {
176217
const first = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', {
177218
countsFor: access,

‎apps/sim/lib/knowledge/service.test.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,9 @@ describe('knowledge base counts with live source permissions', () => {
296296
createdAt: new Date('2026-01-01'),
297297
},
298298
])
299-
queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 }])
299+
queueTableRows(schemaMock.knowledgeBase, [
300+
{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 },
301+
])
300302
const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { countsFor: access })
301303
expect(result.data[0]).toMatchObject({ docCount: 2, tokenCount: 10 })
302304
expect(getForConnectors).not.toHaveBeenCalled()
@@ -307,9 +309,9 @@ describe('knowledge base counts with live source permissions', () => {
307309

308310
it('does not retain stale totals when a live source no longer authorizes its documents', async () => {
309311
const { access, getForConnectors } = reader()
310-
queueTableRows(schemaMock.document, [])
312+
queueTableRows(schemaMock.knowledgeBase, [])
311313
queueTableRows(schemaMock.knowledgeConnector, [{ connectorId: 'confluence-source' }])
312-
queueTableRows(schemaMock.document, [])
314+
queueTableRows(schemaMock.knowledgeBase, [])
313315
const base = { id: 'kb-1', docCount: 5, tokenCount: 50 } as KnowledgeBaseWithCounts
314316
await expect(attachKnowledgeBaseConnectors(base, access)).resolves.toMatchObject({
315317
docCount: 0,

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

Lines changed: 49 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -189,8 +189,8 @@ async function readKnowledgeBaseRows(
189189
}
190190

191191
/**
192-
* Pages bases before counting the documents `access` admits. Explicit document base IDs keep
193-
* the count selective instead of scanning a shared ACL token across tenants before the join.
192+
* Pages bases before counting the documents `access` admits, so a small page never counts the
193+
* rest of its workspace or another tenant's documents.
194194
*/
195195
async function readCountedKnowledgeBaseRows(
196196
where: SQL | undefined,
@@ -202,34 +202,24 @@ async function readCountedKnowledgeBaseRows(
202202
> {
203203
const scope = 'get' in access ? await access.get() : access
204204
const rows = await readKnowledgeBaseRows(where, orderBy, limit)
205-
const totals =
206-
rows.length > 0
207-
? await countDocumentsByKnowledgeBase(
208-
sql`${document.knowledgeBaseId} = ANY(${textArrayLiteral(rows.map((kb) => kb.id))})`,
209-
knowledgeAccessCondition(scope)
210-
)
211-
: []
205+
if (rows.length === 0) return []
206+
/**
207+
* The bases to count: an unpaged list is bounded by its own filter, a page by its row IDs.
208+
* Every count below runs once for the whole list, so a workspace with tens of thousands of
209+
* bases never turns into hundreds of per-batch round trips.
210+
*/
211+
const subject =
212+
limit === undefined && where
213+
? where
214+
: sql`${knowledgeBase.id} = ANY(${textArrayLiteral(rows.map((kb) => kb.id))})`
215+
const totals = await countDocumentsByKnowledgeBase(subject, knowledgeAccessCondition(scope))
212216
const counts = new Map(totals.map((total) => [total.knowledgeBaseId, total]))
213-
214217
/**
215218
* The counts above already include everything the reader's stored ACL admits. Only a
216-
* provider can add documents a live source (GitHub, Confluence) authorizes beyond that,
217-
* and that supplement is resolved once for the whole list: an unpaged list is bounded by
218-
* its own filter, a page by its row IDs, so a workspace with tens of thousands of bases
219-
* never turns into hundreds of per-batch round trips.
219+
* provider can add documents a live source (GitHub, Confluence) authorizes beyond that.
220220
*/
221221
const liveCounts =
222-
'get' in access && rows.length > 0
223-
? await readLiveSourceDocumentCounts(
224-
limit === undefined && where
225-
? where
226-
: inArray(
227-
knowledgeBase.id,
228-
rows.map((kb) => kb.id)
229-
),
230-
access
231-
)
232-
: undefined
222+
'get' in access ? await readLiveSourceDocumentCounts(subject, access) : undefined
233223
return rows.map((kb) => ({
234224
...kb,
235225
docCount: (counts.get(kb.id)?.docCount ?? 0) + (liveCounts?.get(kb.id)?.docCount ?? 0),
@@ -244,23 +234,43 @@ const ACTIVE_DOCUMENT_CONDITIONS = [
244234
] as const
245235

246236
/**
247-
* Document totals per knowledge base for one access predicate, restricted to the bases
248-
* `subject` selects. `subject` may reference `knowledge_base` columns.
237+
* Document totals per knowledge base for one access predicate, one row per base `subject`
238+
* selects from `knowledge_base` (zero when it has no readable documents).
239+
*
240+
* Each base is counted on its own. One grouped scan over many bases lets the planner
241+
* over-estimate their documents and AND the base index with bitmaps of every tenant's documents
242+
* (the shared-token ACL index, a partial index of unarchived documents) before it reads a row;
243+
* correlated on one base ID, the count is a plain scan of that base's active documents.
249244
*/
250245
async function countDocumentsByKnowledgeBase(
251246
subject: SQL,
252247
accessCondition: SQL
253248
): Promise<Array<{ knowledgeBaseId: string; docCount: number; tokenCount: number }>> {
254-
return db
249+
const totals = db
255250
.select({
256-
knowledgeBaseId: document.knowledgeBaseId,
257-
docCount: count(),
258-
tokenCount: sql<number>`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number),
251+
docCount: count().as('readable_doc_count'),
252+
tokenCount: sql<number>`COALESCE(SUM(${document.tokenCount}), 0)`
253+
.mapWith(Number)
254+
.as('readable_token_count'),
259255
})
260256
.from(document)
261-
.innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id))
262-
.where(and(subject, ...ACTIVE_DOCUMENT_CONDITIONS, accessCondition))
263-
.groupBy(document.knowledgeBaseId)
257+
.where(
258+
and(
259+
eq(document.knowledgeBaseId, knowledgeBase.id),
260+
...ACTIVE_DOCUMENT_CONDITIONS,
261+
accessCondition
262+
)
263+
)
264+
.as('totals')
265+
return db
266+
.select({
267+
knowledgeBaseId: knowledgeBase.id,
268+
docCount: totals.docCount,
269+
tokenCount: totals.tokenCount,
270+
})
271+
.from(knowledgeBase)
272+
.crossJoinLateral(totals)
273+
.where(subject)
264274
}
265275

266276
/**
@@ -952,18 +962,18 @@ export async function getActiveKnowledgeBaseReferences(
952962
* operation that only resolves its context pays for neither the count nor the connector read.
953963
*/
954964
export async function attachKnowledgeBaseConnectors(
955-
knowledgeBase: ActiveKnowledgeBaseReference,
965+
base: ActiveKnowledgeBaseReference,
956966
access: KnowledgeReadAccess
957967
): Promise<KnowledgeBaseWithCounts> {
958-
const subject = eq(document.knowledgeBaseId, knowledgeBase.id)
968+
const subject = eq(knowledgeBase.id, base.id)
959969
const scope = 'get' in access ? await access.get() : access
960970
const [ordinary] = await countDocumentsByKnowledgeBase(subject, knowledgeAccessCondition(scope))
961971
const live = 'get' in access ? await readLiveSourceDocumentCounts(subject, access) : undefined
962972
const [withConnectors] = await attachConnectorTypes([
963973
{
964-
...knowledgeBase,
965-
docCount: Number(ordinary?.docCount ?? 0) + (live?.get(knowledgeBase.id)?.docCount ?? 0),
966-
tokenCount: (ordinary?.tokenCount ?? 0) + (live?.get(knowledgeBase.id)?.tokenCount ?? 0),
974+
...base,
975+
docCount: Number(ordinary?.docCount ?? 0) + (live?.get(base.id)?.docCount ?? 0),
976+
tokenCount: (ordinary?.tokenCount ?? 0) + (live?.get(base.id)?.tokenCount ?? 0),
967977
},
968978
])
969979
return withConnectors

‎apps/sim/lib/workspace-files/__integration__/file-listing.integration.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -679,6 +679,27 @@ describe('bounded workspace file reads in PostgreSQL', () => {
679679
)
680680
}
681681
})
682+
683+
/**
684+
* The file browser's default: every live file, most recently modified first. With sorting
685+
* ruled out, an ordered plan exists only through an index keyed on that order.
686+
*/
687+
it('reads the whole workspace in modified order from its workspace index', async () => {
688+
const nodes = await explainCaptured(
689+
db,
690+
() =>
691+
queryWorkspaceFiles(sortWorkspaceId, {
692+
sortBy: 'updatedAt',
693+
sortOrder: 'desc',
694+
limit: 50,
695+
}),
696+
{ match: /order by/, gucs: { enable_sort: 'off' } }
697+
)
698+
expect(nodes.map((node) => node['Node Type'])).not.toContain('Sort')
699+
expect(nodes.map((node) => node['Index Name'])).toContain(
700+
'workspace_files_workspace_active_updated_idx'
701+
)
702+
})
682703
})
683704

684705
describe('SQL type labels and kinds', () => {
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
-- A workspace's live file list, most recently modified first: the file browser's default sort.
2+
-- 0442 indexed the upload order, but no index served this one, so each list read every live file
3+
-- of the workspace and sorted them, which for a large workspace runs from disk and reaches the
4+
-- statement timeout.
5+
--
6+
-- The predicate and the heap-only owner checks are 0442's, for the same reasons; the key is the
7+
-- (updated_at, id) keyset the list orders and resumes on, kept as narrow.
8+
--
9+
-- Replay-safety: as 0442. The build is CONCURRENTLY below an embedded COMMIT, and a replay moves
10+
-- only an INVALID index left by a failed build aside to be dropped, keeping a valid one.
11+
DO $$
12+
BEGIN
13+
IF EXISTS (
14+
SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_files_workspace_active_updated_idx_failed')
15+
AND (indisvalid OR indrelid <> to_regclass('workspace_files'))
16+
) THEN
17+
RAISE EXCEPTION 'Refusing to drop unexpected recovery index workspace_files_workspace_active_updated_idx_failed';
18+
END IF;
19+
IF EXISTS (
20+
SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_files_workspace_active_updated_idx')
21+
AND indrelid = to_regclass('workspace_files') AND NOT indisvalid
22+
) THEN
23+
ALTER INDEX "workspace_files_workspace_active_updated_idx" RENAME TO "workspace_files_workspace_active_updated_idx_failed";
24+
END IF;
25+
END $$;--> statement-breakpoint
26+
COMMIT;--> statement-breakpoint
27+
SET lock_timeout = 0;--> statement-breakpoint
28+
DROP INDEX CONCURRENTLY IF EXISTS "workspace_files_workspace_active_updated_idx_failed";--> statement-breakpoint
29+
CREATE INDEX CONCURRENTLY IF NOT EXISTS "workspace_files_workspace_active_updated_idx" ON "workspace_files" USING btree ("workspace_id","updated_at","id") WHERE "workspace_files"."deleted_at" IS NULL AND "workspace_files"."context" = 'workspace' AND "workspace_files"."workspace_id" IS NOT NULL;--> statement-breakpoint
30+
SET lock_timeout = '5s';

0 commit comments

Comments
 (0)