Skip to content

Commit 0a4036b

Browse files
committed
improvement(knowledge): guard keyword upserts, scope the keyword backfill, and keep adopted Tin scoped
1 parent 1c4ffa5 commit 0a4036b

7 files changed

Lines changed: 148 additions & 66 deletions

‎apps/sim/lib/knowledge/__integration__/embedding-insert-batches.integration.ts‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@ describe('bounded embedding insert transactions', () => {
4949
beforeAll(async () => {
5050
fixtures.root = mkdtempSync(path.join(tmpdir(), 'sim-embedding-batches-'))
5151
await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
52+
/** A search index, the only kind of base whose chunks the keyword projection holds. */
53+
await db
54+
.update(knowledgeBase)
55+
.set({ isSearchIndex: true })
56+
.where(eq(knowledgeBase.id, ids.knowledgeBaseId))
5257
vi.spyOn(embeddingClient, 'assertKnowledgeEmbeddingCapacity').mockResolvedValue(undefined)
5358
})
5459

@@ -151,9 +156,11 @@ describe('bounded embedding insert transactions', () => {
151156
.from(embedding)
152157
.where(eq(embedding.documentId, file.documentId))
153158
).toEqual([{ id: previousId }])
154-
expect(
155-
await db.$client`SELECT id FROM embedding_search WHERE document_id = ${file.documentId}`
156-
).toEqual([{ id: previousId }])
159+
for (const table of ['embedding_search', 'embedding_keyword_search']) {
160+
expect(
161+
await db.$client.unsafe(`SELECT id FROM ${table} WHERE document_id = $1`, [file.documentId])
162+
).toEqual([{ id: previousId }])
163+
}
157164
expect(
158165
await db
159166
.select()
@@ -173,18 +180,13 @@ describe('bounded embedding insert transactions', () => {
173180
expect(await db.select().from(document).where(eq(document.id, file.documentId))).toMatchObject([
174181
{ processingStatus: 'completed', chunkCount: 205, processingError: null },
175182
])
176-
/** A workspace knowledge base keeps no keyword projection rows: its keywords rank on `embedding`. */
177-
for (const [table, count] of [
178-
['embedding', 205],
179-
['embedding_search', 205],
180-
['embedding_keyword_search', 0],
181-
] as const) {
183+
for (const table of ['embedding', 'embedding_search', 'embedding_keyword_search']) {
182184
expect(
183185
await db.$client.unsafe(
184186
`SELECT count(*)::int AS count FROM ${table} WHERE document_id = $1`,
185187
[file.documentId]
186188
)
187-
).toEqual([{ count }])
189+
).toEqual([{ count: 205 }])
188190
}
189191
expect(
190192
await db.$client`SELECT count(*)::int AS count FROM embedding_secret_provenance p

‎apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,11 @@ describe('document processing commit lock scope', () => {
5454
beforeAll(async () => {
5555
fixtures.root = mkdtempSync(path.join(tmpdir(), 'sim-processing-lock-scope-'))
5656
await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
57+
/** A search index, the only kind of base whose chunks the keyword projection holds. */
58+
await db
59+
.update(knowledgeBase)
60+
.set({ isSearchIndex: true })
61+
.where(eq(knowledgeBase.id, ids.knowledgeBaseId))
5762
vi.spyOn(embeddingClient, 'assertKnowledgeEmbeddingCapacity').mockResolvedValue(undefined)
5863
fixtures.process.mockResolvedValue({
5964
chunks,
@@ -152,6 +157,14 @@ describe('document processing commit lock scope', () => {
152157
expect(await db.select().from(document).where(eq(document.id, file.documentId))).toMatchObject([
153158
{ processingStatus: 'completed', chunkCount: 3 },
154159
])
160+
for (const projection of ['embedding_search', 'embedding_keyword_search']) {
161+
expect(
162+
await db.$client.unsafe(
163+
`SELECT count(*)::int AS count FROM ${projection} WHERE document_id = $1`,
164+
[file.documentId]
165+
)
166+
).toEqual([{ count: 3 }])
167+
}
155168
})
156169

157170
it.each([

‎packages/db/script-migrations/0016_backfill_search_vectors.integration.ts‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,11 @@ describe('search projection upgrade in PostgreSQL', () => {
4444
await sql.unsafe(
4545
`ALTER TABLE embedding_search ${fields.map((field) => `DROP COLUMN ${field}`).join(', ')}`
4646
)
47-
await sql`INSERT INTO knowledge_base (id, user_id, name, workspace_id, embedding_model) VALUES
48-
('prefix', 'reader', 'Prefix fixture', 'workspace', 'text-embedding-3-small'),
49-
('full', 'reader', 'Full fixture', 'workspace', 'gemini-embedding-001')`
47+
/** A workspace holds one search index, so each fixture index has its own. */
48+
await sql`INSERT INTO knowledge_base (id, user_id, name, workspace_id, embedding_model, is_search_index) VALUES
49+
('prefix', 'reader', 'Prefix fixture', 'workspace', 'text-embedding-3-small', true),
50+
('full', 'reader', 'Full fixture', 'full-workspace', 'gemini-embedding-001', true),
51+
('legacy', 'reader', 'Legacy fixture', 'workspace', 'text-embedding-3-small', false)`
5052
await backfillEmbeddingSearch(sql)
5153
await sql.unsafe(`INSERT INTO embedding
5254
(id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, content_length, token_count, start_offset, end_offset, embedding)
@@ -153,6 +155,22 @@ describe('search projection upgrade in PostgreSQL', () => {
153155
expect(await backfillSearchKeywords(sql)).toBe(0)
154156
})
155157

158+
it('leaves the chunks of a base that is not a search index unfilled', async () => {
159+
await sql.unsafe(`INSERT INTO embedding
160+
(id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, content_length, token_count, start_offset, end_offset, embedding)
161+
VALUES ('legacy-chunk', 'legacy', 'legacy-document', 0, 'legacy-hash', 'Synthetic fixture', 17, 4, 0, 17,
162+
array_fill(0.01::real, ARRAY[1536])::vector(1536))`)
163+
await sql`DELETE FROM embedding_keyword_search WHERE id = 'legacy-chunk'`
164+
try {
165+
expect(await backfillSearchKeywords(sql)).toBe(0)
166+
expect(
167+
await sql`SELECT id FROM embedding_keyword_search WHERE id = 'legacy-chunk'`
168+
).toHaveLength(0)
169+
} finally {
170+
await sql`DELETE FROM embedding WHERE id = 'legacy-chunk'`
171+
}
172+
})
173+
156174
it('repairs an interrupted index build and preserves valid indexes on replay', async () => {
157175
await expect(
158176
sql.unsafe(

‎packages/db/script-migrations/0016_backfill_search_vectors.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,11 @@ export async function backfillSearchVectors(sql: Sql): Promise<number> {
115115
return count
116116
}
117117

118-
/** Keeps keyword scoring independent of the full-vector and chunk-content storage working sets. */
118+
/**
119+
* Keeps keyword scoring independent of the full-vector and chunk-content storage working sets. Only
120+
* search-index chunks are filled: they are the only rows the projection holds once
121+
* `0025_scope_keyword_projections` scopes its triggers, so a rerun cannot refill the rest.
122+
*/
119123
export async function backfillSearchKeywords(sql: Sql): Promise<number> {
120124
await sql.begin(async (tx) => {
121125
await tx.unsafe("SET LOCAL lock_timeout = '5s'")
@@ -159,6 +163,7 @@ export async function backfillSearchKeywords(sql: Sql): Promise<number> {
159163
), batch AS MATERIALIZED (
160164
SELECT e.id, e.knowledge_base_id, e.document_id, e.enabled, e.content_tsv
161165
FROM missing m INNER JOIN embedding e ON e.id = m.id
166+
INNER JOIN knowledge_base k ON k.id = e.knowledge_base_id AND k.is_search_index
162167
ORDER BY e.id FOR KEY SHARE OF e
163168
), inserted AS (
164169
INSERT INTO embedding_keyword_search (id, knowledge_base_id, document_id, enabled, content_tsv)

‎packages/db/script-migrations/0019_tin_keyword_projection.ts‎

Lines changed: 35 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { SYNCHRONOUS_PROJECTION_WHEN } from '@sim/db/knowledge-projection'
12
import { EMBEDDING_KEYWORD_TIN_INDEX } from '@sim/db/schema'
23
import { resolveMigrationDatabaseUrl } from '@sim/db/script-migrations/database-url'
34
import { type ScriptMigration, ScriptMigrationDeferred } from '@sim/db/script-migrations/types'
@@ -50,6 +51,35 @@ export async function installMembershipKey(tx: Sql | TransactionSql): Promise<vo
5051
$$`)
5152
}
5253

54+
/**
55+
* The chunk trigger's body. A chunk outside a search index writes nothing; an update removes the
56+
* row a chunk left behind by moving out of one. An insert has nothing to remove, since a new chunk
57+
* id has no row and a base adopted meanwhile waits on the membership lock the insert holds.
58+
*/
59+
export async function installTinChunkSync(tx: Sql | TransactionSql): Promise<void> {
60+
await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_embedding_keyword_tin()
61+
RETURNS trigger LANGUAGE plpgsql AS $$
62+
BEGIN
63+
PERFORM pg_advisory_xact_lock_shared(knowledge_tin_membership_key(NEW.knowledge_base_id));
64+
IF NOT EXISTS (
65+
SELECT 1 FROM knowledge_base WHERE id = NEW.knowledge_base_id AND is_search_index
66+
) THEN
67+
IF TG_OP = 'UPDATE' THEN
68+
DELETE FROM embedding_keyword_tin WHERE id = NEW.id;
69+
END IF;
70+
RETURN NEW;
71+
END IF;
72+
INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content)
73+
VALUES (NEW.id, NEW.knowledge_base_id, NEW.document_id, NEW.enabled,
74+
knowledge_tin_base_token(NEW.knowledge_base_id) || ' ' || knowledge_tin_stream(NEW.content_tsv))
75+
ON CONFLICT (id) DO UPDATE SET
76+
knowledge_base_id = EXCLUDED.knowledge_base_id, document_id = EXCLUDED.document_id,
77+
enabled = EXCLUDED.enabled, content = EXCLUDED.content;
78+
RETURN NEW;
79+
END;
80+
$$`)
81+
}
82+
5383
/**
5484
* Installs the stream and base-token functions shared by the triggers and the query path, and the
5585
* embedding and knowledge base triggers, atomically with respect to embedding writers.
@@ -71,7 +101,9 @@ export async function installMembershipKey(tx: Sql | TransactionSql): Promise<vo
71101
* other's rows. A row lock cannot do this: a key-share lock on the marker row is compatible with
72102
* the uncommitted marker update, so a writer would read the old marker without waiting. A whole
73103
* base is reached through `embedding`'s knowledge base index, since a projected row always carries
74-
* its chunk's base and the projection keeps no index but Tin's.
104+
* its chunk's base and the projection keeps no index but Tin's. The embedding trigger carries the
105+
* guard `0024_knowledge_projection_async` added, so a database adopting Tin later ends where a full
106+
* migration run does.
75107
*/
76108
export async function installProjection(sql: Sql): Promise<void> {
77109
await sql.begin(async (tx) => {
@@ -87,28 +119,10 @@ export async function installProjection(sql: Sql): Promise<void> {
87119
SELECT 'zkb' || md5(knowledge_base_id)
88120
$$`)
89121
await installMembershipKey(tx)
90-
await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_embedding_keyword_tin()
91-
RETURNS trigger LANGUAGE plpgsql AS $$
92-
BEGIN
93-
PERFORM pg_advisory_xact_lock_shared(knowledge_tin_membership_key(NEW.knowledge_base_id));
94-
IF NOT EXISTS (
95-
SELECT 1 FROM knowledge_base WHERE id = NEW.knowledge_base_id AND is_search_index
96-
) THEN
97-
DELETE FROM embedding_keyword_tin WHERE id = NEW.id;
98-
RETURN NEW;
99-
END IF;
100-
INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content)
101-
VALUES (NEW.id, NEW.knowledge_base_id, NEW.document_id, NEW.enabled,
102-
knowledge_tin_base_token(NEW.knowledge_base_id) || ' ' || knowledge_tin_stream(NEW.content_tsv))
103-
ON CONFLICT (id) DO UPDATE SET
104-
knowledge_base_id = EXCLUDED.knowledge_base_id, document_id = EXCLUDED.document_id,
105-
enabled = EXCLUDED.enabled, content = EXCLUDED.content;
106-
RETURN NEW;
107-
END;
108-
$$`)
122+
await installTinChunkSync(tx)
109123
await tx.unsafe(`CREATE OR REPLACE TRIGGER embedding_keyword_tin_sync
110124
AFTER INSERT OR UPDATE OF knowledge_base_id, document_id, enabled, content ON embedding
111-
FOR EACH ROW EXECUTE FUNCTION sync_embedding_keyword_tin()`)
125+
FOR EACH ROW WHEN (${SYNCHRONOUS_PROJECTION_WHEN}) EXECUTE FUNCTION sync_embedding_keyword_tin()`)
112126
await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_knowledge_base_keyword_tin()
113127
RETURNS trigger LANGUAGE plpgsql AS $$
114128
BEGIN

‎packages/db/script-migrations/0025_scope_keyword_projections.integration.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,21 @@ describe('keyword projections scoped to search indexes in PostgreSQL', () => {
141141
expect(await tinIds()).toEqual([])
142142
})
143143

144+
it('keeps the scoped and guarded Tin trigger when Tin is adopted after this migration', async () => {
145+
await installProjection(sql)
146+
const [trigger] = await sql<{ definition: string }[]>`
147+
SELECT pg_get_triggerdef(oid) AS definition FROM pg_trigger
148+
WHERE tgname = 'embedding_keyword_tin_sync' AND tgrelid = 'embedding'::regclass`
149+
expect(trigger?.definition).toContain('sim.projection_mode')
150+
const scans = await sql.begin(async (tx) => {
151+
const before = await scansOf(tx, ['embedding_keyword_tin'])
152+
await tx`INSERT INTO embedding (id, knowledge_base_id, document_id, content_tsv)
153+
VALUES ('in-legacy', 'legacy', 'doc', to_tsvector('english', 'Release notes'))`
154+
return (await scansOf(tx, ['embedding_keyword_tin'])) - before
155+
})
156+
expect(scans).toBe(0)
157+
})
158+
144159
it('projects a base adopted as a search index, and removes it when it is not', async () => {
145160
await sql`INSERT INTO embedding (id, knowledge_base_id, document_id, enabled, content_tsv) VALUES
146161
('legacy-1', 'legacy', 'doc', true, to_tsvector('english', 'Quarterly planning')),
@@ -154,6 +169,27 @@ describe('keyword projections scoped to search indexes in PostgreSQL', () => {
154169
expect(await keywordIds()).toEqual([])
155170
})
156171

172+
it('leaves a current keyword row unwritten when its base is adopted', async () => {
173+
await sql`INSERT INTO embedding (id, knowledge_base_id, document_id, content_tsv)
174+
VALUES ('kept', 'legacy', 'doc', to_tsvector('english', 'Kept chunk'))`
175+
/** A row written before the projection was scoped, already current. */
176+
await sql`INSERT INTO embedding_keyword_search (id, knowledge_base_id, document_id, enabled, content_tsv)
177+
SELECT id, knowledge_base_id, document_id, enabled, content_tsv FROM embedding WHERE id = 'kept'`
178+
const updated = await sql.begin(async (tx) => {
179+
const read = async () => {
180+
const [row] = await tx<Array<{ updated: number }>>`
181+
SELECT n_tup_upd::int AS updated FROM pg_stat_xact_user_tables
182+
WHERE relid = 'embedding_keyword_search'::regclass`
183+
return row.updated
184+
}
185+
const before = await read()
186+
await tx`UPDATE knowledge_base SET is_search_index = true WHERE id = 'legacy'`
187+
return (await read()) - before
188+
})
189+
expect(updated).toBe(0)
190+
expect(await keywordIds()).toEqual(['kept'])
191+
})
192+
157193
it('projects a chunk whose insert commits while its base is being adopted', async () => {
158194
let inserted!: () => void
159195
const insertedSignal = new Promise<void>((resolve) => {

‎packages/db/script-migrations/0025_scope_keyword_projections.ts‎

Lines changed: 25 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
import { installMembershipKey } from '@sim/db/script-migrations/0019_tin_keyword_projection'
1+
import {
2+
installMembershipKey,
3+
installTinChunkSync,
4+
} from '@sim/db/script-migrations/0019_tin_keyword_projection'
25
import { resolveMigrationDatabaseUrl } from '@sim/db/script-migrations/database-url'
36
import type { ScriptMigration } from '@sim/db/script-migrations/types'
47
import { createLogger } from '@sim/logger'
@@ -12,10 +15,23 @@ const TRIGGER_LOCK_TIMEOUT = '5s'
1215
const TRIGGER_LOCK_RETRY_BUDGET_MS = 20 * 60_000
1316
const TRIGGER_LOCK_RETRY_BACKOFF = { baseMs: 2_000, maxMs: 30_000 } as const
1417

15-
/** How both GIN keyword writers bring an existing row up to its chunk, as `0016_backfill_search_vectors` does. */
18+
/** The columns a GIN keyword row copies from its chunk. */
19+
const KEYWORD_SEARCH_COLUMNS = [
20+
'knowledge_base_id',
21+
'document_id',
22+
'enabled',
23+
'content_tsv',
24+
] as const
25+
26+
/**
27+
* How both GIN keyword writers bring an existing row (aliased `s`) up to its chunk. A row already
28+
* current is left unwritten, as the projector's upserts leave it, so adopting a base whose rows
29+
* survived rewrites only what changed while it holds the membership lock.
30+
*/
1631
const KEYWORD_SEARCH_UPSERT = `ON CONFLICT (id) DO UPDATE SET
17-
knowledge_base_id = EXCLUDED.knowledge_base_id, document_id = EXCLUDED.document_id,
18-
enabled = EXCLUDED.enabled, content_tsv = EXCLUDED.content_tsv`
32+
${KEYWORD_SEARCH_COLUMNS.map((column) => `${column} = EXCLUDED.${column}`).join(', ')}
33+
WHERE (${KEYWORD_SEARCH_COLUMNS.map((column) => `s.${column}`).join(', ')})
34+
IS DISTINCT FROM (${KEYWORD_SEARCH_COLUMNS.map((column) => `EXCLUDED.${column}`).join(', ')})`
1935

2036
/**
2137
* The GIN keyword projection's chunk trigger, scoped to search indexes as the Tin projection is:
@@ -36,7 +52,7 @@ async function installKeywordSearchSync(tx: TransactionSql): Promise<void> {
3652
END IF;
3753
RETURN NEW;
3854
END IF;
39-
INSERT INTO embedding_keyword_search (id, knowledge_base_id, document_id, enabled, content_tsv)
55+
INSERT INTO embedding_keyword_search AS s (id, knowledge_base_id, document_id, enabled, content_tsv)
4056
VALUES (NEW.id, NEW.knowledge_base_id, NEW.document_id, NEW.enabled, NEW.content_tsv)
4157
${KEYWORD_SEARCH_UPSERT};
4258
RETURN NEW;
@@ -58,7 +74,7 @@ async function installKeywordSearchMembership(tx: TransactionSql): Promise<void>
5874
DELETE FROM embedding_keyword_search WHERE knowledge_base_id = NEW.id;
5975
RETURN NEW;
6076
END IF;
61-
INSERT INTO embedding_keyword_search (id, knowledge_base_id, document_id, enabled, content_tsv)
77+
INSERT INTO embedding_keyword_search AS s (id, knowledge_base_id, document_id, enabled, content_tsv)
6278
SELECT id, knowledge_base_id, document_id, enabled, content_tsv
6379
FROM embedding WHERE knowledge_base_id = NEW.id
6480
${KEYWORD_SEARCH_UPSERT};
@@ -72,35 +88,13 @@ async function installKeywordSearchMembership(tx: TransactionSql): Promise<void>
7288
}
7389

7490
/**
75-
* The Tin chunk trigger's body from `0019_tin_keyword_projection`, where that migration installed
76-
* it, without the delete an insert outside a search index ran: a new chunk id has no Tin row, and a
77-
* base adopted meanwhile waits on the membership lock the insert holds.
91+
* The Tin chunk trigger's body as `0019_tin_keyword_projection` now installs it, where that
92+
* migration installed it: without the delete an insert outside a search index ran.
7893
*/
7994
async function installTinSync(tx: TransactionSql): Promise<void> {
8095
const [row] = await tx<Array<{ installed: boolean }>>`
8196
SELECT to_regprocedure('sync_embedding_keyword_tin()') IS NOT NULL AS installed`
82-
if (!row?.installed) return
83-
await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_embedding_keyword_tin()
84-
RETURNS trigger LANGUAGE plpgsql AS $$
85-
BEGIN
86-
PERFORM pg_advisory_xact_lock_shared(knowledge_tin_membership_key(NEW.knowledge_base_id));
87-
IF NOT EXISTS (
88-
SELECT 1 FROM knowledge_base WHERE id = NEW.knowledge_base_id AND is_search_index
89-
) THEN
90-
IF TG_OP = 'UPDATE' THEN
91-
DELETE FROM embedding_keyword_tin WHERE id = NEW.id;
92-
END IF;
93-
RETURN NEW;
94-
END IF;
95-
INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content)
96-
VALUES (NEW.id, NEW.knowledge_base_id, NEW.document_id, NEW.enabled,
97-
knowledge_tin_base_token(NEW.knowledge_base_id) || ' ' || knowledge_tin_stream(NEW.content_tsv))
98-
ON CONFLICT (id) DO UPDATE SET
99-
knowledge_base_id = EXCLUDED.knowledge_base_id, document_id = EXCLUDED.document_id,
100-
enabled = EXCLUDED.enabled, content = EXCLUDED.content;
101-
RETURN NEW;
102-
END;
103-
$$`)
97+
if (row?.installed) await installTinChunkSync(tx)
10498
}
10599

106100
/**

0 commit comments

Comments
 (0)