Skip to content

Commit 2d873a9

Browse files
committed
fix(search): refuse the purge until migration 0449 is live
- --execute checks that public.sync_embedding_search has the retiring body and its embedding_search_sync trigger is enabled; before that, chunk writes keep projecting the rows a purge deletes. The dry run reports the same check - One whitespace-insensitive matcher for the live function source, shared by the preflight and the integration parity tests; drop the unit test that compared the migration file text - --help adds the rebuild precautions: wait out an anti-wraparound autovacuum, bound maintenance_work_mem, clean up a failed concurrent reindex
1 parent 30f50b4 commit 2d873a9

5 files changed

Lines changed: 102 additions & 62 deletions

File tree

‎packages/db/knowledge-projection.test.ts‎

Lines changed: 0 additions & 18 deletions
This file was deleted.

‎packages/db/knowledge-projection.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,18 @@ export const SYNC_EMBEDDING_SEARCH_FUNCTION = `CREATE OR REPLACE FUNCTION sync_e
132132
END;
133133
$$`
134134

135+
const comparableSource = (source: string) => source.replace(/\s+/g, ' ').trim()
136+
137+
/**
138+
* Whether a live `pg_proc.prosrc` is {@link SYNC_EMBEDDING_SEARCH_FUNCTION}'s body, ignoring
139+
* whitespace: PostgreSQL stores the body between the `$$` quotes verbatim.
140+
*/
141+
export function isCurrentSyncEmbeddingSearchSource(prosrc: string): boolean {
142+
return (
143+
comparableSource(prosrc) === comparableSource(SYNC_EMBEDDING_SEARCH_FUNCTION.split('$$')[1])
144+
)
145+
}
146+
135147
export type KnowledgeProjection = 'embedding_search'
136148

137149
/** Active vector projections that mirror document access for legacy synchronous writers. */

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { readFileSync } from 'node:fs'
2-
import { SYNC_EMBEDDING_SEARCH_FUNCTION } from '@sim/db/knowledge-projection'
2+
import { isCurrentSyncEmbeddingSearchSource } from '@sim/db/knowledge-projection'
33
import { backfillEmbeddingSearch } from '@sim/db/script-migrations/0015_backfill_embedding_search'
44
import {
55
backfillSearchVectors,
@@ -364,9 +364,6 @@ describe('search projection upgrade in PostgreSQL', () => {
364364
const [installed] = await sql<{ body: string }[]>`
365365
SELECT prosrc AS body FROM pg_proc
366366
WHERE oid = ${`"${schemaName}".sync_embedding_search`}::regproc`
367-
const normalized = (text: string) => text.replace(/\s+/g, ' ').trim()
368-
expect(normalized(installed.body)).toBe(
369-
normalized(SYNC_EMBEDDING_SEARCH_FUNCTION.split('$$')[1])
370-
)
367+
expect(isCurrentSyncEmbeddingSearchSource(installed.body)).toBe(true)
371368
})
372369
})

‎packages/db/scripts/purge-search-projection.integration.ts‎

Lines changed: 56 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { spawnSync } from 'node:child_process'
22
import { readFile } from 'node:fs/promises'
33
import { fileURLToPath } from 'node:url'
4-
import { SYNC_EMBEDDING_SEARCH_FUNCTION } from '@sim/db/knowledge-projection'
4+
import { isCurrentSyncEmbeddingSearchSource } from '@sim/db/knowledge-projection'
55
import {
66
inspectSearchProjection,
77
purgeSearchProjection,
@@ -47,6 +47,19 @@ describe('Search candidate retirement in PostgreSQL', () => {
4747
FROM pg_proc p WHERE p.oid = 'public.sync_embedding_search'::regproc`
4848
}
4949

50+
/** Applies migration 0449 as the migrator does, in one transaction. */
51+
async function applyMigration() {
52+
const migration = await readFile(
53+
new URL('../migrations/0449_retire_search_projection_writes.sql', import.meta.url),
54+
'utf8'
55+
)
56+
await sql.begin(async (tx) => {
57+
for (const statement of migration.split('--> statement-breakpoint')) {
58+
if (statement.trim()) await tx.unsafe(statement)
59+
}
60+
})
61+
}
62+
5063
async function control<T>(run: (admin: Sql) => Promise<T>): Promise<T> {
5164
const url = new URL(databaseUrl)
5265
url.pathname = '/postgres'
@@ -91,17 +104,7 @@ describe('Search candidate retirement in PostgreSQL', () => {
91104
expect(provision.status, provision.stdout + provision.stderr).toBe(0)
92105
;[provisioned] = await definitions()
93106
// Every behaviour below runs on the function the versioned migration installs, replayed twice.
94-
const migration = await readFile(
95-
new URL('../migrations/0449_retire_search_projection_writes.sql', import.meta.url),
96-
'utf8'
97-
)
98-
for (let replay = 0; replay < 2; replay++) {
99-
await sql.begin(async (tx) => {
100-
for (const statement of migration.split('--> statement-breakpoint')) {
101-
if (statement.trim()) await tx.unsafe(statement)
102-
}
103-
})
104-
}
107+
for (let replay = 0; replay < 2; replay++) await applyMigration()
105108
await sql`INSERT INTO "user" (id, name, email, email_verified, created_at, updated_at)
106109
VALUES ('reader', 'Synthetic reader', 'reader@example.test', true, now(), now())`
107110
await sql`INSERT INTO project (id, name, owner_id) VALUES ('project', 'Synthetic project', 'reader')`
@@ -159,14 +162,10 @@ describe('Search candidate retirement in PostgreSQL', () => {
159162
return row.n
160163
}
161164

162-
const normalized = (text: string) => text.replace(/\s+/g, ' ').trim()
163-
164165
it('installs the provisioning function and leaves both triggers as they were', async () => {
165166
const [migrated] = await definitions()
166-
expect(normalized(migrated.body)).toBe(normalized(provisioned.body))
167-
expect(normalized(migrated.body)).toBe(
168-
normalized(SYNC_EMBEDDING_SEARCH_FUNCTION.split('$$')[1])
169-
)
167+
expect(isCurrentSyncEmbeddingSearchSource(provisioned.body)).toBe(true)
168+
expect(isCurrentSyncEmbeddingSearchSource(migrated.body)).toBe(true)
170169
expect(migrated.sync).toBe(provisioned.sync)
171170
expect(migrated.sync).toContain("current_setting('sim.projection_mode'")
172171
expect(migrated.acl).toBe(provisioned.acl)
@@ -368,17 +367,47 @@ describe('Search candidate retirement in PostgreSQL', () => {
368367
})
369368
})
370369

370+
function runPurge(...args: string[]) {
371+
return spawnSync('bun', ['--no-env-file', './scripts/purge-search-projection.ts', ...args], {
372+
cwd: packageRoot,
373+
env: { ...process.env, MIGRATION_DATABASE_URL: fixtureUrl.toString() },
374+
encoding: 'utf8',
375+
timeout: 30_000,
376+
})
377+
}
378+
379+
it('refuses --execute until migration 0449 is live, and the dry run reports it', async () => {
380+
await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 3 }, true)
381+
await sql.unsafe(`CREATE OR REPLACE FUNCTION public.sync_embedding_search()
382+
RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RETURN NEW; END; $$`)
383+
try {
384+
const dryRun = runPurge()
385+
expect(dryRun.status, dryRun.stdout + dryRun.stderr).toBe(0)
386+
expect(dryRun.stdout).toContain('"projectionRetired":false')
387+
const stale = runPurge('--execute', '--pause-ms', '0')
388+
expect(stale.status).toBe(1)
389+
expect(stale.stdout + stale.stderr).toContain('Migration 0449 is not live')
390+
expect(await candidateCount('search')).toBe(3)
391+
392+
await applyMigration()
393+
await sql`ALTER TABLE embedding DISABLE TRIGGER embedding_search_sync`
394+
expect(runPurge('--execute', '--pause-ms', '0').status).toBe(1)
395+
expect(await candidateCount('search')).toBe(3)
396+
} finally {
397+
await sql`ALTER TABLE embedding ENABLE TRIGGER embedding_search_sync`
398+
await applyMigration()
399+
}
400+
const dryRun = runPurge()
401+
expect(dryRun.stdout).toContain('"projectionRetired":true')
402+
const execute = runPurge('--execute', '--pause-ms', '0')
403+
expect(execute.status, execute.stdout + execute.stderr).toBe(0)
404+
expect(await candidateCount('search')).toBe(0)
405+
})
406+
371407
it('runs as an operator command: dry run by default, deletes only with --execute', async () => {
372408
await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 3 }, true)
373409
await seedBase({ id: 'ordinary', width: 1536, model: 'text-embedding-3-small', chunks: 2 })
374-
const run = (...args: string[]) => {
375-
return spawnSync('bun', ['--no-env-file', './scripts/purge-search-projection.ts', ...args], {
376-
cwd: packageRoot,
377-
env: { ...process.env, MIGRATION_DATABASE_URL: fixtureUrl.toString() },
378-
encoding: 'utf8',
379-
timeout: 30_000,
380-
})
381-
}
410+
const run = runPurge
382411
const help = run('--help')
383412
expect(help.status).toBe(0)
384413
expect(help.stdout).toContain('REINDEX INDEX CONCURRENTLY embedding_search_512_cosine_hnsw_idx')
@@ -401,16 +430,7 @@ describe('Search candidate retirement in PostgreSQL', () => {
401430
await sql`CREATE TABLE shadow.knowledge_base (id text, is_search_index boolean)`
402431
await sql.unsafe(`ALTER DATABASE "${database}" SET search_path = shadow, public`)
403432
try {
404-
const execute = spawnSync(
405-
'bun',
406-
['--no-env-file', './scripts/purge-search-projection.ts', '--execute', '--pause-ms', '0'],
407-
{
408-
cwd: packageRoot,
409-
env: { ...process.env, MIGRATION_DATABASE_URL: fixtureUrl.toString() },
410-
encoding: 'utf8',
411-
timeout: 30_000,
412-
}
413-
)
433+
const execute = runPurge('--execute', '--pause-ms', '0')
414434
expect(execute.status, execute.stdout + execute.stderr).toBe(0)
415435
} finally {
416436
await sql.unsafe(`ALTER DATABASE "${database}" RESET search_path`)

‎packages/db/scripts/purge-search-projection.ts‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { parseArgs } from 'node:util'
2+
import { isCurrentSyncEmbeddingSearchSource } from '@sim/db/knowledge-projection'
23
import { createLogger, LogLevel } from '@sim/logger'
34
import { getPostgresErrorCode } from '@sim/utils/errors'
45
import { sleep } from '@sim/utils/helpers'
@@ -35,8 +36,9 @@ bases (knowledge_base.is_search_index). Canonical embedding and document rows ar
3536
3637
MIGRATION_DATABASE_URL=... bun --no-env-file packages/db/scripts/purge-search-projection.ts [--execute] [options]
3738
38-
Without --execute, prints each Search knowledge base's planner row estimate (≈) and whether any
39-
candidate row remains, and changes nothing.
39+
Without --execute, prints whether migration 0449 is live (projectionRetired), each Search knowledge
40+
base's planner row estimate (≈) and whether any candidate row remains, and changes nothing.
41+
--execute refuses until migration 0449 is live: wait for the deploy's migrate job to finish.
4042
--execute Delete in batches, each its own short transaction; any error stops the run.
4143
--batch-size N Rows per batch, 1–${MAX_BATCH_SIZE} (default ${DEFAULT_BATCH_SIZE}).
4244
--pause-ms N Pause after each batch, 0–${MAX_PAUSE_MS} (default ${DEFAULT_PAUSE_MS}).
@@ -45,9 +47,13 @@ candidate row remains, and changes nothing.
4547
MIGRATION_DATABASE_URL must be a direct PlanetScale primary (port 5432, sslmode=verify-full,
4648
sslrootcert=system); poolers are refused. Rerun until the dry run reports none remaining, then,
4749
from a direct primary session outside peak traffic (this script never runs them):
50+
-- wait while pg_stat_activity shows an anti-wraparound autovacuum of embedding_search
51+
SET maintenance_work_mem = '2GB'; SET statement_timeout = 0;
4852
REINDEX INDEX CONCURRENTLY embedding_search_512_cosine_hnsw_idx;
53+
-- if it fails: DROP INDEX CONCURRENTLY embedding_search_512_cosine_hnsw_idx_ccnew; then retry
4954
VACUUM (ANALYZE) embedding_search;
50-
The full runbook is in the pull request that added this script.`
55+
Rebuild each HNSW index a purged Search knowledge base used: text-embedding-3-small and -large
56+
bases use only the 512 one. The full runbook is in the pull request that added this script.`
5157

5258
/** Deliberate operator messages never include database values, credentials or statement text. */
5359
export class SearchProjectionPurgeError extends Error {
@@ -145,6 +151,24 @@ export function connectOperatorDatabase(rawUrl: string | undefined): Sql {
145151
})
146152
}
147153

154+
/**
155+
* Whether migration 0449 is live: `public.sync_embedding_search` has the body that deletes a
156+
* Search chunk's candidate row, and its `embedding_search_sync` trigger is enabled. Before that,
157+
* chunk writes keep projecting the rows a purge deletes.
158+
*/
159+
async function isProjectionRetired(sql: Sql): Promise<boolean> {
160+
const [live] = await sql<{ source: string | null; trigger_enabled: boolean }[]>`
161+
SELECT (SELECT prosrc FROM pg_proc WHERE oid = to_regproc('public.sync_embedding_search')) AS source,
162+
EXISTS (
163+
SELECT 1 FROM pg_trigger
164+
WHERE tgrelid = to_regclass('public.embedding') AND tgname = 'embedding_search_sync'
165+
AND tgenabled <> 'D'
166+
) AS trigger_enabled`
167+
return (
168+
live.trigger_enabled && live.source !== null && isCurrentSyncEmbeddingSearchSource(live.source)
169+
)
170+
}
171+
148172
/** Search knowledge bases in id order; the table is small and read once per run. */
149173
async function searchKnowledgeBases(sql: Sql): Promise<string[]> {
150174
const rows = await sql<{ id: string }[]>`
@@ -317,6 +341,7 @@ async function main(): Promise<void> {
317341
if (!values.execute) {
318342
const states = await inspectSearchProjection(sql)
319343
logger.info('Search candidate rows (dry run; nothing deleted; estimatedRows is ≈)', {
344+
projectionRetired: await isProjectionRetired(sql),
320345
knowledgeBases: states,
321346
remaining: states.some((state) => state.remaining),
322347
})
@@ -327,6 +352,10 @@ async function main(): Promise<void> {
327352
to_regclass('public.search_retirement_state') IS NOT NULL AS retirement_installed`
328353
if (!session.writable)
329354
throw new SearchProjectionPurgeError('The purge requires a writable primary connection')
355+
if (!(await isProjectionRetired(sql)))
356+
throw new SearchProjectionPurgeError(
357+
'Migration 0449 is not live (sync_embedding_search() is not the retiring body, or its trigger is disabled); wait for the migrate job to finish'
358+
)
330359
if (session.retirement_installed)
331360
throw new SearchProjectionPurgeError(
332361
'A shadow-copy Search retirement job is still installed; abort it with the release that shipped it first'

0 commit comments

Comments
 (0)