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

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

A script migration can now defer: `up` throws `ScriptMigrationDeferred`, the runner
logs it, leaves the name unrecorded and runs the remaining migrations. The Tin
projection defers when the database offers `tin` but refuses to create it, so the
upgrade after the extension is permitted installs it instead of skipping it forever.

* fix(db): keep a deferred Tin projection from failing a direct run

`db:push` runs the migration file directly, where a deferral has no
migration record to leave unwritten, so a refused extension aborted the
push. The direct entry point now adopts the projection where the database
allows it and logs the refusal otherwise, while the registered runner
still sees the deferral. Also covers continuation after a deferral with a
synthetic migration list, which the registry cannot express while the only
deferring migration is its last entry.

Author: waleedlatif1

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'
+}