Skip to content

Commit 71fb0a5

Browse files
committed
fix(db): avoid repeated dev index builds and search backfills
1 parent 4fecfc9 commit 71fb0a5

5 files changed

Lines changed: 179 additions & 42 deletions

File tree

‎.github/workflows/migrations.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ jobs:
8080
# terminal" but still exits 0, so the job goes green without applying the
8181
# change. tee keeps the output live in the log; we then fail on drizzle's
8282
# own TTY error. pipefail also preserves a non-zero db:push exit through tee.
83-
bun run db:push --force < /dev/null 2>&1 | tee /tmp/db-push.log
83+
SIM_DEV_DB_PUSH=1 bun run db:push --force < /dev/null 2>&1 | tee /tmp/db-push.log
8484
if grep -q "Interactive prompts require a TTY terminal" /tmp/db-push.log; then
8585
echo "ERROR: db:push needs an interactive rename decision; land it as a versioned migration instead of relying on push." >&2
8686
exit 1

‎packages/db/dev-push.test.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import * as schema from '@sim/db/schema'
2+
import { generateDrizzleJson, generateMigration } from 'drizzle-kit/api'
3+
import { afterEach, describe, expect, it, vi } from 'vitest'
4+
5+
interface SnapshotTable {
6+
indexes: Record<string, { method: string; with?: Record<string, unknown> }>
7+
}
8+
9+
afterEach(() => vi.unstubAllEnvs())
10+
11+
describe('dev push index options', () => {
12+
it('matches introspected options without changing ordinary migration snapshots', async () => {
13+
vi.stubEnv('SIM_DEV_DB_PUSH', undefined)
14+
const ordinary = generateDrizzleJson(schema)
15+
const introspected = structuredClone(ordinary)
16+
const hnswIndexes = Object.values<SnapshotTable>(introspected.tables).flatMap((table) =>
17+
Object.values(table.indexes).filter((index) => index.method === 'hnsw')
18+
)
19+
expect(hnswIndexes).toHaveLength(22)
20+
for (const index of hnswIndexes) {
21+
expect(index.with).toEqual({ m: 16, ef_construction: 64 })
22+
index.with = { m: '16', ef_construction: '64' }
23+
}
24+
25+
vi.stubEnv('SIM_DEV_DB_PUSH', '1')
26+
const dev = generateDrizzleJson(schema)
27+
expect(await generateMigration(introspected, dev)).toEqual([])
28+
expect({ ...dev, id: ordinary.id }).toEqual(introspected)
29+
30+
vi.stubEnv('SIM_DEV_DB_PUSH', undefined)
31+
const restored = generateDrizzleJson(schema)
32+
expect({ ...restored, id: ordinary.id }).toEqual(ordinary)
33+
})
34+
35+
it('still detects a real index-option change in dev', async () => {
36+
vi.stubEnv('SIM_DEV_DB_PUSH', '1')
37+
const current = generateDrizzleJson(schema)
38+
const previous = structuredClone(current)
39+
previous.tables['public.embedding'].indexes.embedding_vector_hnsw_idx.with = {
40+
m: '8',
41+
ef_construction: '64',
42+
}
43+
const statements = await generateMigration(previous, current)
44+
expect(statements).toHaveLength(2)
45+
expect(statements[0]).toContain('DROP INDEX "embedding_vector_hnsw_idx"')
46+
expect(statements[1]).toContain('WITH (m=16,ef_construction=64)')
47+
})
48+
})

‎packages/db/schema.ts‎

Lines changed: 33 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,17 @@ import {
2727
} from 'drizzle-orm/pg-core'
2828
import { DEFAULT_FREE_CREDITS, TAG_SLOTS } from './constants'
2929

30+
/**
31+
* Drizzle push compares index options as JSON, while Postgres introspection
32+
* returns strings. Normalize only dev pushes so unchanged HNSW indexes survive
33+
* the diff; versioned migration snapshots retain their original numeric values.
34+
*/
35+
function hnswIndexOptions() {
36+
return process.env.SIM_DEV_DB_PUSH === '1'
37+
? { m: '16', ef_construction: '64' }
38+
: { m: 16, ef_construction: 64 }
39+
}
40+
3041
// Custom tsvector type for full-text search
3142
export const tsvector = customType<{
3243
data: string
@@ -3194,28 +3205,16 @@ export const embedding = pgTable(
31943205
// Vector similarity search indexes (HNSW), one per stored width
31953206
embeddingVectorHnswIdx: index('embedding_vector_hnsw_idx')
31963207
.using('hnsw', table.embedding.op('vector_cosine_ops'))
3197-
.with({
3198-
m: 16,
3199-
ef_construction: 64,
3200-
}),
3208+
.with(hnswIndexOptions()),
32013209
embedding384VectorHnswIdx: index('embedding_384_vector_hnsw_idx')
32023210
.using('hnsw', table.embedding384.op('vector_cosine_ops'))
3203-
.with({
3204-
m: 16,
3205-
ef_construction: 64,
3206-
}),
3211+
.with(hnswIndexOptions()),
32073212
embedding768VectorHnswIdx: index('embedding_768_vector_hnsw_idx')
32083213
.using('hnsw', table.embedding768.op('vector_cosine_ops'))
3209-
.with({
3210-
m: 16,
3211-
ef_construction: 64,
3212-
}),
3214+
.with(hnswIndexOptions()),
32133215
embedding1024VectorHnswIdx: index('embedding_1024_vector_hnsw_idx')
32143216
.using('hnsw', table.embedding1024.op('vector_cosine_ops'))
3215-
.with({
3216-
m: 16,
3217-
ef_construction: 64,
3218-
}),
3217+
.with(hnswIndexOptions()),
32193218
/**
32203219
* pgvector indexes `vector` only up to 2,000 dimensions and `halfvec` up to
32213220
* 4,000, so the 3,072 column is indexed through a `halfvec` cast — the
@@ -3236,27 +3235,24 @@ export const embedding = pgTable(
32363235
*/
32373236
embedding3072VectorHnswIdx: index('embedding_3072_vector_hnsw_idx')
32383237
.using('hnsw', sql`(${table.embedding3072}::halfvec(3072)) halfvec_cosine_ops`)
3239-
.with({
3240-
m: 16,
3241-
ef_construction: 64,
3242-
}),
3238+
.with(hnswIndexOptions()),
32433239

32443240
/** contract-pending(after stored candidate retrieval is fully deployed): remove these expression indexes; old app versions still query them during rollout. */
32453241
embeddingBinaryHnswIdx: index('embedding_binary_hnsw_idx')
32463242
.using('hnsw', sql`(binary_quantize(${table.embedding})::bit(1536)) bit_hamming_ops`)
3247-
.with({ m: 16, ef_construction: 64 }),
3243+
.with(hnswIndexOptions()),
32483244
embedding384BinaryHnswIdx: index('embedding_384_binary_hnsw_idx')
32493245
.using('hnsw', sql`(binary_quantize(${table.embedding384})::bit(384)) bit_hamming_ops`)
3250-
.with({ m: 16, ef_construction: 64 }),
3246+
.with(hnswIndexOptions()),
32513247
embedding768BinaryHnswIdx: index('embedding_768_binary_hnsw_idx')
32523248
.using('hnsw', sql`(binary_quantize(${table.embedding768})::bit(768)) bit_hamming_ops`)
3253-
.with({ m: 16, ef_construction: 64 }),
3249+
.with(hnswIndexOptions()),
32543250
embedding1024BinaryHnswIdx: index('embedding_1024_binary_hnsw_idx')
32553251
.using('hnsw', sql`(binary_quantize(${table.embedding1024})::bit(1024)) bit_hamming_ops`)
3256-
.with({ m: 16, ef_construction: 64 }),
3252+
.with(hnswIndexOptions()),
32573253
embedding3072BinaryHnswIdx: index('embedding_3072_binary_hnsw_idx')
32583254
.using('hnsw', sql`(binary_quantize(${table.embedding3072})::bit(3072)) bit_hamming_ops`)
3259-
.with({ m: 16, ef_construction: 64 }),
3255+
.with(hnswIndexOptions()),
32603256

32613257
// Text tag indexes
32623258
tag1Idx: index('emb_kb_tag1_lower_idx').on(table.knowledgeBaseId, sql`lower(${table.tag1})`),
@@ -3349,37 +3345,37 @@ export const embeddingSearch = pgTable(
33493345
.where(sql`${table.enabled}`),
33503346
binaryIdx: index('embedding_search_binary_hnsw_idx')
33513347
.using('hnsw', table.binary.op('bit_hamming_ops'))
3352-
.with({ m: 16, ef_construction: 64 }),
3348+
.with(hnswIndexOptions()),
33533349
binary384Idx: index('embedding_search_384_binary_hnsw_idx')
33543350
.using('hnsw', table.binary384.op('bit_hamming_ops'))
3355-
.with({ m: 16, ef_construction: 64 }),
3351+
.with(hnswIndexOptions()),
33563352
binary768Idx: index('embedding_search_768_binary_hnsw_idx')
33573353
.using('hnsw', table.binary768.op('bit_hamming_ops'))
3358-
.with({ m: 16, ef_construction: 64 }),
3354+
.with(hnswIndexOptions()),
33593355
binary1024Idx: index('embedding_search_1024_binary_hnsw_idx')
33603356
.using('hnsw', table.binary1024.op('bit_hamming_ops'))
3361-
.with({ m: 16, ef_construction: 64 }),
3357+
.with(hnswIndexOptions()),
33623358
binary3072Idx: index('embedding_search_3072_binary_hnsw_idx')
33633359
.using('hnsw', table.binary3072.op('bit_hamming_ops'))
3364-
.with({ m: 16, ef_construction: 64 }),
3360+
.with(hnswIndexOptions()),
33653361
vectorIdx: index('embedding_search_cosine_hnsw_idx')
33663362
.using('hnsw', table.vector.op('halfvec_cosine_ops'))
3367-
.with({ m: 16, ef_construction: 64 }),
3363+
.with(hnswIndexOptions()),
33683364
vector512Idx: index('embedding_search_512_cosine_hnsw_idx')
33693365
.using('hnsw', table.vector512.op('halfvec_cosine_ops'))
3370-
.with({ m: 16, ef_construction: 64 }),
3366+
.with(hnswIndexOptions()),
33713367
vector384Idx: index('embedding_search_384_cosine_hnsw_idx')
33723368
.using('hnsw', table.vector384.op('halfvec_cosine_ops'))
3373-
.with({ m: 16, ef_construction: 64 }),
3369+
.with(hnswIndexOptions()),
33743370
vector768Idx: index('embedding_search_768_cosine_hnsw_idx')
33753371
.using('hnsw', table.vector768.op('halfvec_cosine_ops'))
3376-
.with({ m: 16, ef_construction: 64 }),
3372+
.with(hnswIndexOptions()),
33773373
vector1024Idx: index('embedding_search_1024_cosine_hnsw_idx')
33783374
.using('hnsw', table.vector1024.op('halfvec_cosine_ops'))
3379-
.with({ m: 16, ef_construction: 64 }),
3375+
.with(hnswIndexOptions()),
33803376
vector3072Idx: index('embedding_search_3072_cosine_hnsw_idx')
33813377
.using('hnsw', table.vector3072.op('halfvec_cosine_ops'))
3382-
.with({ m: 16, ef_construction: 64 }),
3378+
.with(hnswIndexOptions()),
33833379
widthCheck: check(
33843380
'embedding_search_width_check',
33853381
sql`num_nonnulls("binary", "binary_384", "binary_768", "binary_1024", "binary_3072") = 1`
@@ -3456,10 +3452,7 @@ export const docsEmbeddings = pgTable(
34563452
// Vector similarity search indexes (HNSW) - optimized for documentation embeddings
34573453
embeddingVectorHnswIdx: index('docs_embedding_vector_hnsw_idx')
34583454
.using('hnsw', table.embedding.op('vector_cosine_ops'))
3459-
.with({
3460-
m: 16,
3461-
ef_construction: 64,
3462-
}),
3455+
.with(hnswIndexOptions()),
34633456

34643457
// GIN index for JSONB metadata queries
34653458
metadataGinIdx: index('docs_emb_metadata_gin_idx').using('gin', table.metadata),
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import {
2+
backfillSearchVectorsMigration,
3+
runDevSearchBackfill,
4+
} from '@sim/db/script-migrations/0016_backfill_search_vectors'
5+
import type { Sql } from 'postgres'
6+
import { afterEach, describe, expect, it, vi } from 'vitest'
7+
8+
/** Models the ledger while keeping the expensive projection work mocked. */
9+
function createLedger(appliedNames: string[] = []) {
10+
const names = new Set(appliedNames)
11+
const query = vi.fn(async (parts: TemplateStringsArray, name?: string) => {
12+
const statement = parts.join('?').trim()
13+
if (statement.startsWith('SELECT name')) {
14+
return name && names.has(name) ? [{ name }] : []
15+
}
16+
if (statement.startsWith('INSERT INTO script_migrations') && name) names.add(name)
17+
return []
18+
})
19+
const begin = vi.fn(async (callback: (tx: typeof query) => Promise<void>) => callback(query))
20+
return { names, begin, sql: Object.assign(query, { begin }) as unknown as Sql }
21+
}
22+
23+
afterEach(() => vi.restoreAllMocks())
24+
25+
describe('dev search backfill tracking', () => {
26+
it('records success and skips projection work on the next deploy', async () => {
27+
const up = vi.spyOn(backfillSearchVectorsMigration, 'up').mockResolvedValue(undefined)
28+
const { sql, names } = createLedger()
29+
await runDevSearchBackfill(sql)
30+
expect(names).toEqual(
31+
new Set(['0015_backfill_embedding_search', '0016_backfill_search_vectors'])
32+
)
33+
await runDevSearchBackfill(sql)
34+
expect(up).toHaveBeenCalledTimes(1)
35+
})
36+
37+
it('retries failed work without recording completion', async () => {
38+
const up = vi
39+
.spyOn(backfillSearchVectorsMigration, 'up')
40+
.mockRejectedValueOnce(new Error('backfill interrupted'))
41+
.mockResolvedValue(undefined)
42+
const { sql, names, begin } = createLedger()
43+
await expect(runDevSearchBackfill(sql)).rejects.toThrow('backfill interrupted')
44+
expect(names.size).toBe(0)
45+
expect(begin).not.toHaveBeenCalled()
46+
await runDevSearchBackfill(sql)
47+
expect(names.has(backfillSearchVectorsMigration.name)).toBe(true)
48+
expect(up).toHaveBeenCalledTimes(2)
49+
})
50+
51+
it('runs the vector upgrade even when the older binary backfill is recorded', async () => {
52+
const up = vi.spyOn(backfillSearchVectorsMigration, 'up').mockResolvedValue(undefined)
53+
const { sql } = createLedger(['0015_backfill_embedding_search'])
54+
await runDevSearchBackfill(sql)
55+
expect(up).toHaveBeenCalledOnce()
56+
})
57+
58+
it('skips a backfill already recorded by the tracked migration runner', async () => {
59+
const up = vi.spyOn(backfillSearchVectorsMigration, 'up').mockResolvedValue(undefined)
60+
const { sql, begin } = createLedger([backfillSearchVectorsMigration.name])
61+
await runDevSearchBackfill(sql)
62+
expect(up).not.toHaveBeenCalled()
63+
expect(begin).not.toHaveBeenCalled()
64+
})
65+
})

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -246,12 +246,43 @@ export const backfillSearchVectorsMigration: ScriptMigration = {
246246
},
247247
}
248248

249+
/** Records the dev push backfill only after success, so later deploys skip both full scans. */
250+
export async function runDevSearchBackfill(sql: Sql): Promise<void> {
251+
await sql`
252+
CREATE TABLE IF NOT EXISTS script_migrations (
253+
name text PRIMARY KEY,
254+
applied_at timestamptz NOT NULL DEFAULT now()
255+
)
256+
`
257+
const [applied] = await sql`
258+
SELECT name FROM script_migrations WHERE name = ${backfillSearchVectorsMigration.name}
259+
`
260+
if (applied) {
261+
logger.info('Dev search backfill already applied; skipping')
262+
return
263+
}
264+
265+
await backfillSearchVectorsMigration.up(sql)
266+
await sql.begin(async (tx) => {
267+
for (const name of [
268+
backfillSearchVectorsMigration.name,
269+
...(backfillSearchVectorsMigration.supersedes ?? []),
270+
]) {
271+
await tx`INSERT INTO script_migrations (name) VALUES (${name}) ON CONFLICT (name) DO NOTHING`
272+
}
273+
})
274+
}
275+
249276
if (import.meta.main) {
250277
const url = resolveMigrationDatabaseUrl()
251278
if (!url) throw new Error('DATABASE_URL is required to initialize search vectors')
252279
const sql = postgres(url, { max: 1, max_lifetime: null, onnotice: () => undefined })
253280
try {
254-
await backfillSearchVectorsMigration.up(sql)
281+
if (process.env.SIM_DEV_DB_PUSH === '1') {
282+
await runDevSearchBackfill(sql)
283+
} else {
284+
await backfillSearchVectorsMigration.up(sql)
285+
}
255286
} finally {
256287
await sql.end()
257288
}

0 commit comments

Comments
 (0)