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.

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,49 @@
+/**
+ * @vitest-environment node
+ */
+import type { Sql } from 'postgres'
+import { describe, expect, it, vi } from 'vitest'
+import { 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('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)

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

@@ -0,0 +1,39 @@
+/**
+ * @vitest-environment node
+ */
+import type { Sql } from 'postgres'
+import { describe, expect, it, vi } from 'vitest'
+import { runScriptMigrations, scriptMigrations } from './index'
+
+const TIN = '0019_tin_keyword_projection'
+
+/** A session where every migration but Tin's is recorded and the database refuses `tin`. */
+function createSqlHarness() {
+  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(
+        scriptMigrations.filter(({ name }) => name !== TIN).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()
+    await expect(runScriptMigrations(sql)).resolves.toBeUndefined()
+    expect(recorded).toEqual([])
+  })
+})

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,6 +54,7 @@ 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.
  */
 export async function runScriptMigrations(sql: Sql): Promise<void> {
   const names = new Set<string>()
@@ -101,7 +102,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'
+}