Skip to content

Commit a22c263

Browse files
committed
fix(search): keep the purge on its locked session and refresh the dev projection function
- Re-check the session advisory lock inside every purge batch: the client reconnects silently after a dropped connection, and the new backend does not hold the old one's lock - Build the purge client without the sslmode/sslrootcert URL parameters, enforcing them as ssl verify-full; postgres.js forwards sslrootcert as a startup parameter the server rejects - Pin the purge connection's search_path to public, pg_temp - Reinstall sync_embedding_search() on a dev push that already recorded 0016, since push runs no versioned migration
1 parent 293513f commit a22c263

5 files changed

Lines changed: 166 additions & 37 deletions

File tree

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
import { readFileSync } from 'node:fs'
2+
import { SYNC_EMBEDDING_SEARCH_FUNCTION } from '@sim/db/knowledge-projection'
23
import { backfillEmbeddingSearch } from '@sim/db/script-migrations/0015_backfill_embedding_search'
34
import {
45
backfillSearchVectors,
56
buildSearchIndexes,
7+
runDevSearchBackfill,
68
} from '@sim/db/script-migrations/0016_backfill_search_vectors'
79
import { runScriptMigrations, scriptMigrations } from '@sim/db/script-migrations/index'
810
import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure'
@@ -351,4 +353,20 @@ describe('search projection upgrade in PostgreSQL', () => {
351353
await runScriptMigrations(sql)
352354
await sql`DELETE FROM embedding WHERE id LIKE 'upgrade-%'`
353355
}, 60_000)
356+
357+
it('reinstalls a changed projection function on a dev push that already recorded 0016', async () => {
358+
await sql`CREATE TABLE IF NOT EXISTS script_migrations (name text PRIMARY KEY, applied_at timestamptz NOT NULL DEFAULT now())`
359+
await sql`INSERT INTO script_migrations (name) VALUES ('0016_backfill_search_vectors')
360+
ON CONFLICT (name) DO NOTHING`
361+
await sql.unsafe(`CREATE OR REPLACE FUNCTION sync_embedding_search()
362+
RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RETURN NEW; END; $$`)
363+
await runDevSearchBackfill(sql)
364+
const [installed] = await sql<{ body: string }[]>`
365+
SELECT prosrc AS body FROM pg_proc
366+
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+
)
371+
})
354372
})

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,11 @@ export const backfillSearchVectorsMigration: ScriptMigration = {
140140
},
141141
}
142142

143-
/** Records the dev push backfill only after success, so later deploys skip the full scan. */
143+
/**
144+
* Records the dev push backfill only after success, so later deploys skip the full scan. A push
145+
* runs no versioned migration, so an already-backfilled database still reinstalls the trigger
146+
* function, which takes no table lock, to pick up a changed body.
147+
*/
144148
export async function runDevSearchBackfill(sql: Sql): Promise<void> {
145149
await sql`
146150
CREATE TABLE IF NOT EXISTS script_migrations (
@@ -152,7 +156,8 @@ export async function runDevSearchBackfill(sql: Sql): Promise<void> {
152156
SELECT name FROM script_migrations WHERE name = ${backfillSearchVectorsMigration.name}
153157
`
154158
if (applied) {
155-
logger.info('Dev search backfill already applied; skipping')
159+
await sql.unsafe(SYNC_EMBEDDING_SEARCH_FUNCTION)
160+
logger.info('Dev search backfill already applied; refreshed the projection function only')
156161
return
157162
}
158163

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,41 @@ describe('Search candidate retirement in PostgreSQL', () => {
306306
expect(await candidateCount('search')).toBe(3)
307307
})
308308

309+
it('stops when its connection is replaced, since the new session lacks the purge lock', async () => {
310+
await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 5 }, true)
311+
let closed: () => void = () => undefined
312+
const purger = postgres(fixtureUrl.toString(), {
313+
max: 1,
314+
prepare: false,
315+
onnotice: () => undefined,
316+
onclose: () => closed(),
317+
})
318+
const [{ pid }] = await purger<{ pid: number }[]>`SELECT pg_backend_pid() AS pid`
319+
try {
320+
await expect(
321+
purgeSearchProjection(purger, {
322+
batchSize: 2,
323+
pauseMs: 0,
324+
onBatch: async () => {
325+
// Wait until the client has seen the close, as it has after a real pause.
326+
const seen = new Promise<void>((resolve) => {
327+
closed = resolve
328+
})
329+
await sql`SELECT pg_terminate_backend(${pid})`
330+
await seen
331+
},
332+
})
333+
).rejects.toBeInstanceOf(SearchProjectionPurgeError)
334+
} finally {
335+
await purger.end()
336+
}
337+
expect(await candidateCount('search')).toBe(3)
338+
expect(await purgeSearchProjection(sql, { batchSize: 2, pauseMs: 0 })).toMatchObject({
339+
deleted: 3,
340+
remaining: false,
341+
})
342+
})
343+
309344
it('runs as an operator command: dry run by default, deletes only with --execute', async () => {
310345
await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 3 }, true)
311346
await seedBase({ id: 'ordinary', width: 1536, model: 'text-embedding-3-small', chunks: 2 })
@@ -332,4 +367,28 @@ describe('Search candidate retirement in PostgreSQL', () => {
332367

333368
expect(run('--execute', '--batch-size', '50000').status).toBe(1)
334369
})
370+
371+
it('runs the operator command on the public tables whatever the database search_path', async () => {
372+
await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 3 }, true)
373+
await sql`CREATE SCHEMA shadow`
374+
await sql`CREATE TABLE shadow.knowledge_base (id text, is_search_index boolean)`
375+
await sql.unsafe(`ALTER DATABASE "${database}" SET search_path = shadow, public`)
376+
try {
377+
const execute = spawnSync(
378+
'bun',
379+
['--no-env-file', './scripts/purge-search-projection.ts', '--execute', '--pause-ms', '0'],
380+
{
381+
cwd: packageRoot,
382+
env: { ...process.env, MIGRATION_DATABASE_URL: fixtureUrl.toString() },
383+
encoding: 'utf8',
384+
timeout: 30_000,
385+
}
386+
)
387+
expect(execute.status, execute.stdout + execute.stderr).toBe(0)
388+
} finally {
389+
await sql.unsafe(`ALTER DATABASE "${database}" RESET search_path`)
390+
await sql`DROP SCHEMA shadow CASCADE`
391+
}
392+
expect(await candidateCount('search')).toBe(0)
393+
})
335394
})

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import {
2+
connectOperatorDatabase,
23
parseOperatorDatabaseUrl,
34
SearchProjectionPurgeError,
45
} from '@sim/db/scripts/purge-search-projection'
@@ -7,7 +8,7 @@ import { describe, expect, it } from 'vitest'
78
const primary = 'postgresql://operator:secret@fixture.pg.psdb.cloud:5432/app'
89
const tls = 'sslmode=verify-full&sslrootcert=system'
910

10-
describe('purge connection refusal', () => {
11+
describe('purge operator connection', () => {
1112
it.each([
1213
['a missing URL', undefined],
1314
['a non-PostgreSQL URL', 'mysql://operator:secret@fixture.pg.psdb.cloud:5432/app'],
@@ -34,4 +35,15 @@ describe('purge connection refusal', () => {
3435
false
3536
)
3637
})
38+
39+
it('sends no TLS URL parameter as a startup parameter, enforcing it as ssl instead', async () => {
40+
const sql = connectOperatorDatabase(`${primary}?${tls}`)
41+
try {
42+
expect(sql.options.ssl).toBe('verify-full')
43+
expect(Object.keys(sql.options.connection)).not.toContain('sslrootcert')
44+
expect(Object.keys(sql.options.connection)).not.toContain('sslmode')
45+
} finally {
46+
await sql.end()
47+
}
48+
})
3749
})

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

Lines changed: 69 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,14 @@ const LOCK_TIMEOUT_MS = 1_000
1515
/** A probe may walk the dead index entries a purge leaves until VACUUM, so it gets longer. */
1616
const PROBE_STATEMENT_TIMEOUT_MS = 30_000
1717
const PURGE_LOCK = 'sim:search-projection-purge'
18+
/**
19+
* Whether this backend holds the session-level purge lock. A bigint advisory key appears in
20+
* `pg_locks` split into `classid` (high 32 bits) and `objid` (low 32 bits) with `objsubid` 1.
21+
*/
22+
const HOLDS_PURGE_LOCK = `EXISTS (
23+
SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND pid = pg_backend_pid() AND granted
24+
AND objsubid = 1 AND ((classid::bigint << 32) | objid::bigint) = hashtextextended('${PURGE_LOCK}', 0)
25+
)`
1826
/**
1927
* Keeps every scan on `embedding_search_kb_idx`: a sequential scan would restart from the heap's
2028
* start each batch, over the dead rows earlier batches left, and a bitmap scan would collect a
@@ -108,6 +116,35 @@ export function parseOperatorDatabaseUrl(rawUrl: string | undefined): OperatorDa
108116
return { url, hosted }
109117
}
110118

119+
/**
120+
* The purge's single-backend client for a validated operator URL. The TLS query parameters are
121+
* enforced here as `ssl` and removed from the URL: postgres.js forwards `sslrootcert` to the
122+
* server as a startup parameter, which PostgreSQL rejects as unrecognized. The search path is
123+
* pinned because the purge reads and deletes unqualified tables. Connects on the first query.
124+
*/
125+
export function connectOperatorDatabase(rawUrl: string | undefined): Sql {
126+
const { url, hosted } = parseOperatorDatabaseUrl(rawUrl)
127+
const clientUrl = new URL(url)
128+
clientUrl.searchParams.delete('sslmode')
129+
clientUrl.searchParams.delete('sslrootcert')
130+
return postgres(clientUrl.toString(), {
131+
port: Number(url.port || '5432'),
132+
ssl: hosted ? 'verify-full' : false,
133+
max: 1,
134+
prepare: false,
135+
connect_timeout: 5,
136+
max_lifetime: null,
137+
connection: {
138+
application_name: 'sim-search-projection-purge',
139+
search_path: 'public, pg_temp',
140+
statement_timeout: STATEMENT_TIMEOUT_MS,
141+
lock_timeout: LOCK_TIMEOUT_MS,
142+
idle_in_transaction_session_timeout: 5_000,
143+
},
144+
onnotice: () => undefined,
145+
})
146+
}
147+
111148
/** Search knowledge bases in id order; the table is small and read once per run. */
112149
async function searchKnowledgeBases(sql: Sql): Promise<string[]> {
113150
const rows = await sql<{ id: string }[]>`
@@ -167,10 +204,18 @@ interface SearchProjectionPurgeResult {
167204
* Deletes one batch of a Search knowledge base's candidate rows. The rows are chosen through
168205
* `embedding_search_kb_idx` and deleted by tuple id, re-checking the knowledge base so a row
169206
* moved or reassigned since the scan is left alone. Deletes nothing when the knowledge base is gone.
207+
* Refuses when this backend no longer holds the purge lock: the client reconnects silently after a
208+
* dropped connection, and the new backend does not inherit the old one's session lock.
170209
*/
171210
async function purgeBatch(sql: Sql, knowledgeBaseId: string, batchSize: number): Promise<number> {
172211
return sql.begin(async (tx) => {
173-
await tx`SELECT ${tx.unsafe(INDEX_SCANS_ONLY)}`
212+
const [session] = await tx<{ locked: boolean }[]>`
213+
SELECT ${tx.unsafe(HOLDS_PURGE_LOCK)} AS locked, ${tx.unsafe(INDEX_SCANS_ONLY)}`
214+
if (!session.locked) {
215+
throw new SearchProjectionPurgeError(
216+
'The database connection was replaced and no longer holds the purge lock; rerun to continue'
217+
)
218+
}
174219
const [base] = await tx<{ is_search_index: boolean }[]>`
175220
SELECT is_search_index FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR SHARE`
176221
if (!base) return 0
@@ -194,9 +239,10 @@ async function purgeBatch(sql: Sql, knowledgeBaseId: string, batchSize: number):
194239
}
195240

196241
/**
197-
* Deletes the candidate rows of every Search knowledge base, batch by batch, pausing after each.
198-
* Stateless: whatever a stopped run leaves, the next run finds and deletes. Canonical
199-
* `embedding` and `document` rows are never touched.
242+
* Deletes the candidate rows of every Search knowledge base, batch by batch, pausing after each,
243+
* while holding a session lock that keeps a second run out. Stateless: whatever a stopped run
244+
* leaves, the next run finds and deletes. Canonical `embedding` and `document` rows are never
245+
* touched.
200246
*/
201247
export async function purgeSearchProjection(
202248
sql: Sql,
@@ -205,20 +251,28 @@ export async function purgeSearchProjection(
205251
const batchSize = options.batchSize ?? DEFAULT_BATCH_SIZE
206252
const pauseMs = options.pauseMs ?? DEFAULT_PAUSE_MS
207253
const maxBatches = options.maxBatches ?? Number.POSITIVE_INFINITY
254+
const [{ acquired }] = await sql<{ acquired: boolean }[]>`
255+
SELECT pg_try_advisory_lock(hashtextextended(${PURGE_LOCK}, 0)) AS acquired`
256+
if (!acquired) throw new SearchProjectionPurgeError('Another purge is running')
208257
let deleted = 0
209258
let batches = 0
210-
for (const knowledgeBaseId of await searchKnowledgeBases(sql)) {
211-
for (;;) {
212-
if (batches >= maxBatches) return { deleted, batches, remaining: true }
213-
const count = await purgeBatch(sql, knowledgeBaseId, batchSize)
214-
if (count === 0) break
215-
batches += 1
216-
deleted += count
217-
await options.onBatch?.({ knowledgeBaseId, deleted: count })
218-
if (pauseMs > 0) await sleep(pauseMs)
259+
try {
260+
for (const knowledgeBaseId of await searchKnowledgeBases(sql)) {
261+
for (;;) {
262+
if (batches >= maxBatches) return { deleted, batches, remaining: true }
263+
const count = await purgeBatch(sql, knowledgeBaseId, batchSize)
264+
if (count === 0) break
265+
batches += 1
266+
deleted += count
267+
await options.onBatch?.({ knowledgeBaseId, deleted: count })
268+
if (pauseMs > 0) await sleep(pauseMs)
269+
}
219270
}
271+
return { deleted, batches, remaining: false }
272+
} finally {
273+
// On a replaced connection this unlocks nothing; the old backend's lock ended with it.
274+
await sql`SELECT pg_advisory_unlock(hashtextextended(${PURGE_LOCK}, 0))`
220275
}
221-
return { deleted, batches, remaining: false }
222276
}
223277

224278
function integerOption(
@@ -253,23 +307,7 @@ async function main(): Promise<void> {
253307
const pauseMs = integerOption(values['pause-ms'], DEFAULT_PAUSE_MS, 0, MAX_PAUSE_MS)
254308
const maxBatches =
255309
values['max-batches'] === undefined ? undefined : integerOption(values['max-batches'], 1, 1)
256-
const rawUrl = process.env.MIGRATION_DATABASE_URL
257-
const { url, hosted } = parseOperatorDatabaseUrl(rawUrl)
258-
const sql = postgres(url.toString(), {
259-
port: Number(url.port || '5432'),
260-
ssl: hosted ? 'verify-full' : false,
261-
max: 1,
262-
prepare: false,
263-
connect_timeout: 5,
264-
max_lifetime: null,
265-
connection: {
266-
application_name: 'sim-search-projection-purge',
267-
statement_timeout: STATEMENT_TIMEOUT_MS,
268-
lock_timeout: LOCK_TIMEOUT_MS,
269-
idle_in_transaction_session_timeout: 5_000,
270-
},
271-
onnotice: () => undefined,
272-
})
310+
const sql = connectOperatorDatabase(process.env.MIGRATION_DATABASE_URL)
273311
try {
274312
if (!values.execute) {
275313
const states = await inspectSearchProjection(sql)
@@ -288,9 +326,6 @@ async function main(): Promise<void> {
288326
throw new SearchProjectionPurgeError(
289327
'A shadow-copy Search retirement job is still installed; abort it with the release that shipped it first'
290328
)
291-
const [{ acquired }] = await sql<{ acquired: boolean }[]>`
292-
SELECT pg_try_advisory_lock(hashtextextended(${PURGE_LOCK}, 0)) AS acquired`
293-
if (!acquired) throw new SearchProjectionPurgeError('Another purge is running')
294329
const result = await purgeSearchProjection(sql, {
295330
batchSize,
296331
pauseMs,

0 commit comments

Comments
 (0)