feat(db): auto-apply tracked script data migrations in db:migrate (#5497)

* feat(db): auto-apply tracked script data migrations in db:migrate

* fix(db): reset session lock_timeout before script migrations, guard journal insert

* improvement(db): re-verify advisory-lock session before script migrations

Author: TheodoreSpeaks

apps/sim/lib/table/order-key.ts

@@ -11,10 +11,7 @@
  * repeated same-spot inserts.
  */
 
-import {
-  generateKeyBetween,
-  generateNKeysBetween,
-} from '@/lib/fractional-indexing/fractional-indexing'
+import { generateKeyBetween, generateNKeysBetween } from '@sim/utils/fractional-indexing'
 
 /**
  * Returns a key that sorts strictly between `a` and `b`. Pass `null` for an open

apps/sim/scripts/backfill-table-order-keys.ts

@@ -20,7 +20,8 @@
  * the original backfill also used — minting a fresh, evenly-spaced, distinct run
  * with `nKeysBetween`.
  *
- * Distinct from `backfill-table-order-keys.ts`, which keys tables with NULL keys;
+ * Distinct from the `0001_backfill_table_order_keys` script migration
+ * (`packages/db/script-migrations/`), which keys tables with NULL keys;
  * this one repairs tables that are fully keyed but bytewise-disordered. Run it
  * AFTER migration 0228 so the re-key writes and sorts under `COLLATE "C"`.
  *
@@ -41,7 +42,7 @@ import { drizzle } from 'drizzle-orm/postgres-js'
 import postgres from 'postgres'
 import { nKeysBetween } from '@/lib/table/order-key'
 
-/** See backfill-table-order-keys.ts — keeps each VALUES list well under the param/stack ceilings. */
+/** Keeps each VALUES list well under Postgres's 65535-bound-param and JS call-stack ceilings. */
 const WRITE_CHUNK_SIZE = 5000
 
 export async function runRepair(): Promise<void> {

apps/sim/scripts/repair-table-order-key-collation.ts

@@ -0,0 +1,85 @@
+import { getErrorMessage } from '@sim/utils/errors'
+import { generateNKeysBetween } from '@sim/utils/fractional-indexing'
+import type { ScriptMigration } from './types'
+
+/**
+ * Rows written per `UPDATE … FROM unnest(...)` statement. The two unnest
+ * arrays are single bind parameters, so chunking only bounds per-statement
+ * work and wire-message size, not parameter count. Chunks share the one
+ * per-table transaction, so the table is still keyed atomically.
+ */
+const WRITE_CHUNK_SIZE = 5000
+
+/**
+ * Backfills the `order_key` column on `user_table_rows`.
+ *
+ * Row ordering moved from the contiguous integer `position` to a fractional
+ * string `order_key` (O(1) insert/delete — no reshift/recompact). Each existing
+ * row gets a key derived from its current `position` order, so the fractional
+ * ordering matches today's once `TABLES_FRACTIONAL_ORDERING` is on.
+ *
+ * Per-table-atomic: each table is keyed inside one transaction holding the same
+ * per-table advisory lock the app uses for inserts, so a concurrent insert
+ * can't interleave. Idempotent: tables already fully keyed are skipped; a table
+ * with any NULL key is fully re-keyed from `position` order (deterministic, so
+ * a re-run after a partial failure is safe). Per-table failures are isolated —
+ * remaining tables still run, then the migration throws so it stays pending and
+ * retries (only still-unkeyed tables) on the next upgrade.
+ */
+export const backfillTableOrderKeys: ScriptMigration = {
+  name: '0001_backfill_table_order_keys',
+  async up(sql) {
+    const pending = await sql<{ table_id: string }[]>`
+      SELECT DISTINCT table_id FROM user_table_rows WHERE order_key IS NULL
+    `
+    console.log(`order_key backfill — ${pending.length} table(s) with NULL order_key`)
+
+    const stats = { tablesKeyed: 0, rowsKeyed: 0, failed: 0 }
+
+    for (const { table_id: tableId } of pending) {
+      try {
+        const keyed = await sql.begin(async (tx) => {
+          await tx`
+            SELECT pg_advisory_xact_lock(hashtextextended(${`user_table_rows_pos:${tableId}`}, 0))
+          `
+          const rows = await tx<{ id: string }[]>`
+            SELECT id FROM user_table_rows
+            WHERE table_id = ${tableId}
+            ORDER BY position ASC, id ASC
+          `
+          if (rows.length === 0) return 0
+          const keys = generateNKeysBetween(null, null, rows.length)
+
+          for (let start = 0; start < rows.length; start += WRITE_CHUNK_SIZE) {
+            const chunk = rows.slice(start, start + WRITE_CHUNK_SIZE)
+            const ids = chunk.map((r) => r.id)
+            const chunkKeys = keys.slice(start, start + chunk.length)
+            await tx`
+              UPDATE user_table_rows AS t
+              SET order_key = v.order_key
+              FROM (
+                SELECT unnest(${ids}::text[]) AS id, unnest(${chunkKeys}::text[]) AS order_key
+              ) AS v
+              WHERE t.id = v.id AND t.table_id = ${tableId}
+            `
+          }
+          return rows.length
+        })
+        stats.tablesKeyed += 1
+        stats.rowsKeyed += keyed
+        console.log(`  ${tableId}: keyed ${keyed} rows`)
+      } catch (error) {
+        stats.failed += 1
+        console.error(`  ${tableId}: FAILED — ${getErrorMessage(error)}`)
+      }
+    }
+
+    console.log(
+      `order_key backfill done — tables keyed: ${stats.tablesKeyed}, ` +
+        `rows keyed: ${stats.rowsKeyed}, failed: ${stats.failed}`
+    )
+    if (stats.failed > 0) {
+      throw new Error(`order_key backfill failed for ${stats.failed} table(s); will retry next run`)
+    }
+  },
+}

packages/db/script-migrations/0001_backfill_table_order_keys.ts

@@ -0,0 +1,77 @@
+import type { Sql } from 'postgres'
+import { backfillTableOrderKeys } from './0001_backfill_table_order_keys'
+import type { ScriptMigration } from './types'
+
+export type { ScriptMigration } from './types'
+
+/**
+ * Ordered, append-only registry of script migrations. An entry may be deleted
+ * once a later SQL migration supersedes it (accepting that deployments which
+ * never ran it skip the backfill) — never renamed or reordered.
+ */
+export const scriptMigrations: readonly ScriptMigration[] = [backfillTableOrderKeys]
+
+/**
+ * Applies pending script migrations in registry order, recording each by name.
+ * Runs inside the migration advisory-lock session, immediately after drizzle's
+ * SQL migrations succeed — so a script always sees the fully-migrated schema of
+ * the release that ships it. The tracking table is runner-managed (like
+ * drizzle's own `__drizzle_migrations`), not part of the app schema.
+ *
+ * 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.
+ */
+export async function runScriptMigrations(sql: Sql): Promise<void> {
+  const names = new Set<string>()
+  for (const migration of scriptMigrations) {
+    if (names.has(migration.name)) {
+      throw new Error(`Duplicate script migration name: ${migration.name}`)
+    }
+    names.add(migration.name)
+  }
+
+  /**
+   * The SQL-migration phase leaves `lock_timeout = '5s'` on this session (DDL
+   * must fail fast rather than queue a table-wide stall). Script migrations
+   * want the opposite: they block on app-held row/advisory locks (e.g. the
+   * per-table insert lock in the order_key backfill) and must wait them out —
+   * exactly like the standalone scripts they replace, which ran on fresh
+   * connections with the Postgres defaults. `SET` cannot be parameterized,
+   * hence `unsafe` with constants.
+   */
+  await sql.unsafe('SET statement_timeout = 0')
+  await sql.unsafe('SET lock_timeout = 0')
+
+  await sql`
+    CREATE TABLE IF NOT EXISTS script_migrations (
+      name text PRIMARY KEY,
+      applied_at timestamptz NOT NULL DEFAULT now()
+    )
+  `
+  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))
+  if (pending.length === 0) {
+    console.log('No pending script migrations.')
+    return
+  }
+
+  for (const migration of pending) {
+    for (const key of migration.requiredEnv ?? []) {
+      if (!process.env[key]) {
+        throw new Error(
+          `Script migration ${migration.name} requires env var ${key}; set it on the migrations container.`
+        )
+      }
+    }
+    console.log(`Applying script migration ${migration.name}...`)
+    const startedAt = Date.now()
+    await migration.up(sql)
+    await sql`
+      INSERT INTO script_migrations (name) VALUES (${migration.name})
+      ON CONFLICT (name) DO NOTHING
+    `
+    console.log(`Script migration ${migration.name} applied in ${Date.now() - startedAt}ms.`)
+  }
+}

packages/db/script-migrations/index.ts

@@ -0,0 +1,35 @@
+import type { Sql } from 'postgres'
+
+/**
+ * A run-once TypeScript data migration, applied by the migration runner
+ * (`scripts/migrate.ts`) after all pending SQL migrations succeed and recorded
+ * by name in the `script_migrations` table — the code-migration analogue of
+ * drizzle's `__drizzle_migrations` journal.
+ *
+ * Authoring rules:
+ * - **Idempotent and resumable.** The name is recorded only after `up`
+ *   resolves; a crash mid-run means the whole migration re-runs from the top
+ *   on the next upgrade. Guard with cheap preconditions (`WHERE x IS NULL`,
+ *   `ON CONFLICT DO NOTHING`) so re-runs are near-free.
+ * - **db-level imports only** (`postgres`, `drizzle-orm`, `@sim/utils`). The
+ *   migrations docker image ships `packages/db` + `utils` + `logger` and
+ *   nothing from `apps/*`.
+ * - **Runs at whatever schema HEAD the release ships** — scripts are not
+ *   interleaved with SQL migrations. A SQL migration must never drop or
+ *   repurpose a column that a registered script still reads; delete the
+ *   script from the registry in the same PR instead.
+ * - **Owns its transactions.** The runner deliberately does not wrap `up` in
+ *   one: backfills commit per batch and cannot be rolled back wholesale.
+ */
+export interface ScriptMigration {
+  /** Unique stable identifier recorded in `script_migrations`; never rename after release. */
+  name: string
+  /** Env vars the migration needs; the runner throws before `up` if any is unset. */
+  requiredEnv?: readonly string[]
+  /**
+   * Applies the migration using the runner's session. The runner resets
+   * `statement_timeout` and `lock_timeout` to 0 first, so long backfills and
+   * blocking lock acquisitions wait instead of aborting with `55P03`.
+   */
+  up(sql: Sql): Promise<void>
+}