fix(knowledge): retry the Tin projection after a database refuses the extension

apps/sim/lib/knowledge/search/tin-keyword.ts

@@ -33,7 +33,10 @@ const indexReadiness = new LRUCache<'index', boolean, SearchBudget | undefined>(
   },
 })
 
-/** Only organization search indexes are projected; `is_search_index` is fixed at creation. */
+/**
+ * Only organization search indexes are projected. `is_search_index` is only ever turned on, when a
+ * legacy base is adopted, so a stale answer just keeps that base on the GIN projection for one TTL.
+ */
 const searchIndexBases = new LRUCache<string, boolean>({
   max: 10_000,
   ttl: SEARCH_INDEX_TTL_MS,

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

@@ -0,0 +1,66 @@
+/**
+ * @vitest-environment node
+ */
+import type { Sql } from 'postgres'
+import { describe, expect, it, vi } from 'vitest'
+import {
+  adoptTinKeywordProjection,
+  installTinKeywordProjection,
+} from './0019_tin_keyword_projection'
+import { ScriptMigrationDeferred } from './types'
+
+/** A session that offers `tin` or not, and answers `CREATE EXTENSION` with `createError`. */
+function createSqlHarness(options: { available: boolean; createError?: { code: string } }) {
+  const statements: string[] = []
+  const run = (strings: TemplateStringsArray) => {
+    const text = strings.join('?').replace(/\s+/g, ' ').trim()
+    statements.push(text)
+    if (text.includes('pg_available_extensions')) {
+      return Promise.resolve(options.available ? [{ '?column?': 1 }] : [])
+    }
+    return Promise.resolve([])
+  }
+  const sql = run as unknown as Sql
+  sql.unsafe = vi.fn(async (text: string) => {
+    statements.push(text)
+    if (text.startsWith('CREATE EXTENSION') && options.createError) throw options.createError
+    return []
+  }) as unknown as Sql['unsafe']
+  return { sql, statements }
+}
+
+describe('installTinKeywordProjection', () => {
+  it('records a no-op where the database does not offer tin', async () => {
+    const { sql, statements } = createSqlHarness({ available: false })
+    expect(await installTinKeywordProjection(sql)).toBeUndefined()
+    expect(statements.some((text) => text.startsWith('CREATE EXTENSION'))).toBe(false)
+  })
+
+  it.each(['42501', '0A000'])(
+    'defers without installing anything when the database refuses the extension (%s)',
+    async (code) => {
+      const { sql, statements } = createSqlHarness({ available: true, createError: { code } })
+      await expect(installTinKeywordProjection(sql)).rejects.toBeInstanceOf(ScriptMigrationDeferred)
+      expect(statements.at(-1)).toBe('CREATE EXTENSION IF NOT EXISTS tin')
+    }
+  )
+
+  it('is a no-op when run directly against a database that refuses the extension', async () => {
+    const { sql, statements } = createSqlHarness({
+      available: true,
+      createError: { code: '42501' },
+    })
+    await expect(adoptTinKeywordProjection(sql)).resolves.toBeUndefined()
+    expect(statements.at(-1)).toBe('CREATE EXTENSION IF NOT EXISTS tin')
+  })
+
+  it('fails a direct run on any other extension error', async () => {
+    const { sql } = createSqlHarness({ available: true, createError: { code: '53100' } })
+    await expect(adoptTinKeywordProjection(sql)).rejects.toEqual({ code: '53100' })
+  })
+
+  it('fails the migration on any other extension error', async () => {
+    const { sql } = createSqlHarness({ available: true, createError: { code: '53100' } })
+    await expect(installTinKeywordProjection(sql)).rejects.toEqual({ code: '53100' })
+  })
+})

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

@@ -1,5 +1,5 @@
 import { EMBEDDING_KEYWORD_TIN_INDEX } from '@sim/db/schema'
-import type { ScriptMigration } from '@sim/db/script-migrations/types'
+import { type ScriptMigration, ScriptMigrationDeferred } from '@sim/db/script-migrations/types'
 import { createLogger } from '@sim/logger'
 import postgres, { type Sql } from 'postgres'
 
@@ -12,7 +12,8 @@ const EXTENSION_REFUSED_CODES = new Set(['42501', '0A000'])
 /**
  * Installs the extension, or reports that this database refuses it: listed as available is not
  * the same as creatable by the migration role. Tin is an optimization, so a refusal leaves keyword
- * search on the GIN projection instead of failing the deploy.
+ * search on the GIN projection instead of failing the deploy, and defers the migration so the
+ * upgrade after the extension is allowed installs it.
  */
 async function createTinExtension(sql: Sql): Promise<boolean> {
   try {
@@ -209,15 +210,18 @@ async function buildProjectionIndex(sql: Sql): Promise<void> {
 
 /**
  * Installs and fills the Tin keyword projection where the database offers `tin`, and records a
- * no-op elsewhere. A database that gains the extension later runs this file directly (see below)
- * to adopt it; the whole migration is idempotent.
+ * no-op elsewhere. A database that refuses the extension it offers defers instead, so a later
+ * upgrade retries it. A database that gains the extension after recording the no-op runs this file
+ * directly (see below) to adopt it; the whole migration is idempotent.
  */
 export async function installTinKeywordProjection(sql: Sql): Promise<void> {
   if (!(await tinAvailable(sql))) {
     logger.info('Tin is unavailable; keyword search keeps the GIN projection')
     return
   }
-  if (!(await createTinExtension(sql))) return
+  if (!(await createTinExtension(sql))) {
+    throw new ScriptMigrationDeferred('the database refused the tin extension')
+  }
   await installProjection(sql)
   const rows = await backfillProjection(sql)
   await buildProjectionIndex(sql)
@@ -229,12 +233,28 @@ export const tinKeywordProjectionMigration: ScriptMigration = {
   up: installTinKeywordProjection,
 }
 
+/**
+ * Installs the projection where this database allows it, treating a refused extension as a no-op:
+ * run directly — `db:push`, or adopting Tin after a cluster gains it — there is no migration
+ * record to leave unwritten, and keyword search keeps the GIN projection either way.
+ */
+export async function adoptTinKeywordProjection(sql: Sql): Promise<void> {
+  try {
+    await installTinKeywordProjection(sql)
+  } catch (error) {
+    if (!(error instanceof ScriptMigrationDeferred)) throw error
+    logger.warn('Tin projection deferred; keyword search keeps the GIN projection', {
+      reason: error.message,
+    })
+  }
+}
+
 if (import.meta.main) {
   const url = process.env.MIGRATION_DATABASE_URL ?? process.env.DATABASE_URL
   if (!url) throw new Error('DATABASE_URL is required to install the Tin keyword projection')
   const sql = postgres(url, { max: 1, max_lifetime: null, onnotice: () => undefined })
   try {
-    await installTinKeywordProjection(sql)
+    await adoptTinKeywordProjection(sql)
   } finally {
     await sql.end()
   }

packages/db/script-migrations/index.test.ts

@@ -0,0 +1,56 @@
+/**
+ * @vitest-environment node
+ */
+import type { Sql } from 'postgres'
+import { describe, expect, it, vi } from 'vitest'
+import { runScriptMigrations, scriptMigrations } from './index'
+import { type ScriptMigration, ScriptMigrationDeferred } from './types'
+
+const TIN = '0019_tin_keyword_projection'
+
+/** Every registered migration but Tin's, so Tin is the only pending one. */
+const APPLIED_BEFORE_TIN = scriptMigrations
+  .filter(({ name }) => name !== TIN)
+  .map(({ name }) => name)
+
+/** A session where `applied` is already recorded and the database refuses `tin`. */
+function createSqlHarness(applied: readonly string[]) {
+  const recorded: string[] = []
+  const run = (strings: TemplateStringsArray, ...values: unknown[]) => {
+    const text = strings.join('?').replace(/\s+/g, ' ').trim()
+    if (text.startsWith('SELECT name FROM script_migrations')) {
+      return Promise.resolve(applied.map((name) => ({ name })))
+    }
+    if (text.includes('pg_available_extensions')) return Promise.resolve([{ '?column?': 1 }])
+    if (text.startsWith('INSERT INTO script_migrations')) recorded.push(values[0] as string)
+    return Promise.resolve([])
+  }
+  const sql = run as unknown as Sql
+  sql.unsafe = vi.fn(async (text: string) => {
+    if (text.startsWith('CREATE EXTENSION')) throw { code: '42501' }
+    return []
+  }) as unknown as Sql['unsafe']
+  sql.begin = vi.fn(async (callback) => (callback as (tx: Sql) => unknown)(sql)) as Sql['begin']
+  return { sql, recorded }
+}
+
+describe('runScriptMigrations', () => {
+  it('leaves a deferred migration unrecorded without failing the upgrade', async () => {
+    const { sql, recorded } = createSqlHarness(APPLIED_BEFORE_TIN)
+    await expect(runScriptMigrations(sql)).resolves.toBeUndefined()
+    expect(recorded).toEqual([])
+  })
+
+  it('applies and records the migrations that follow a deferred one', async () => {
+    const deferring: ScriptMigration = {
+      name: 'test_deferring',
+      up: async () => {
+        throw new ScriptMigrationDeferred('the database refused the test migration')
+      },
+    }
+    const following: ScriptMigration = { name: 'test_following', up: async () => {} }
+    const { sql, recorded } = createSqlHarness([])
+    await expect(runScriptMigrations(sql, [deferring, following])).resolves.toBeUndefined()
+    expect(recorded).toEqual(['test_following'])
+  })
+})

packages/db/script-migrations/index.ts

@@ -15,7 +15,7 @@ import { repairUnknownTableRowProvenanceSecondPass } from './0006_repair_unknown
 import { repairUnknownWorkspaceFileProvenance } from './0007_repair_unknown_workspace_file_provenance'
 import { backfillCredentialGroupResourcePolicies } from './0010_backfill_credential_group_resource_policies'
 import { remapLegacyKnowledgeConnectorCredentialsMigration } from './0011_remap_legacy_knowledge_connector_credentials'
-import type { ScriptMigration } from './types'
+import { type ScriptMigration, ScriptMigrationDeferred } from './types'
 
 export type { ScriptMigration } from './types'
 
@@ -54,10 +54,18 @@ export const scriptMigrations: readonly ScriptMigration[] = [
  *
  * Fails fast: a missing required env var or a throwing `up` aborts the run
  * before the name is recorded, so the migration retries on the next upgrade.
+ * A deferred `up` is not recorded either, but lets the later migrations run.
+ *
+ * `migrations` defaults to the registry and exists so a test can apply a
+ * synthetic list: a deferral followed by a later migration is otherwise
+ * uncoverable while the only deferring migration is the last registered entry.
  */
-export async function runScriptMigrations(sql: Sql): Promise<void> {
+export async function runScriptMigrations(
+  sql: Sql,
+  migrations: readonly ScriptMigration[] = scriptMigrations
+): Promise<void> {
   const names = new Set<string>()
-  for (const migration of scriptMigrations) {
+  for (const migration of migrations) {
     if (names.has(migration.name)) {
       throw new Error(`Duplicate script migration name: ${migration.name}`)
     }
@@ -85,7 +93,7 @@ export async function runScriptMigrations(sql: Sql): Promise<void> {
   const appliedRows = await sql<{ name: string }[]>`SELECT name FROM script_migrations`
   const applied = new Set(appliedRows.map((row) => row.name))
 
-  const pending = scriptMigrations.filter((migration) => !applied.has(migration.name))
+  const pending = migrations.filter((migration) => !applied.has(migration.name))
   if (pending.length === 0) {
     console.log('No pending script migrations.')
     return
@@ -101,7 +109,13 @@ export async function runScriptMigrations(sql: Sql): Promise<void> {
     }
     console.log(`Applying script migration ${migration.name}...`)
     const startedAt = Date.now()
-    await migration.up(sql)
+    try {
+      await migration.up(sql)
+    } catch (error) {
+      if (!(error instanceof ScriptMigrationDeferred)) throw error
+      console.log(`Script migration ${migration.name} deferred: ${error.message}`)
+      continue
+    }
     await sql.begin(async (tx) => {
       for (const name of [migration.name, ...(migration.supersedes ?? [])]) {
         await tx`INSERT INTO script_migrations (name) VALUES (${name}) ON CONFLICT (name) DO NOTHING`

packages/db/script-migrations/types.ts

@@ -35,3 +35,12 @@ export interface ScriptMigration {
    */
   up(sql: Sql): Promise<void>
 }
+
+/**
+ * Thrown by `up` to leave the migration unrecorded without failing the upgrade,
+ * so the next upgrade runs it again: for work this database refuses today but
+ * may accept later, such as an extension the migration role may not yet create.
+ */
+export class ScriptMigrationDeferred extends Error {
+  override name = 'ScriptMigrationDeferred'
+}