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

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

@@ -1,130 +0,0 @@
-#!/usr/bin/env bun
-
-/**
- * Backfills the `order_key` column on `user_table_rows`.
- *
- * Row ordering is moving from the contiguous integer `position` to a fractional
- * string `order_key` (O(1) insert/delete — no reshift/recompact). This script
- * assigns each existing row a key derived from its current `position` order, so
- * the new ordering matches today's once the `TABLES_FRACTIONAL_ORDERING` flag is
- * flipped 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).
- *
- * Usage:
- *   DATABASE_URL=... bun run apps/sim/scripts/backfill-table-order-keys.ts
- *   DATABASE_URL=... bun run apps/sim/scripts/backfill-table-order-keys.ts --dry-run
- */
-
-import { userTableRows } from '@sim/db/schema'
-import { getErrorMessage } from '@sim/utils/errors'
-import { asc, eq, isNull, sql } from 'drizzle-orm'
-import { drizzle } from 'drizzle-orm/postgres-js'
-import postgres from 'postgres'
-import { nKeysBetween } from '@/lib/table/order-key'
-
-/**
- * Rows written per `UPDATE … FROM (VALUES …)`. One statement for a whole large table builds an
- * enormous VALUES list that overflows the JS call stack while drizzle assembles it (and would
- * exceed Postgres's 65535-bound-param limit at ~32k rows, 2 params/row). 5000 keeps ~10k params
- * — well under both ceilings — while minimizing round-trips. Chunks share the one per-table
- * transaction, so the table is still keyed atomically.
- */
-const WRITE_CHUNK_SIZE = 5000
-
-export async function runBackfill(): Promise<void> {
-  const dryRun = process.argv.includes('--dry-run')
-  const connectionString = process.env.DATABASE_URL ?? process.env.POSTGRES_URL
-  if (!connectionString) {
-    console.error('Missing DATABASE_URL or POSTGRES_URL')
-    process.exit(1)
-  }
-
-  const client = postgres(connectionString, {
-    prepare: false,
-    idle_timeout: 20,
-    connect_timeout: 30,
-    max: 5,
-    onnotice: () => {},
-  })
-  const db = drizzle(client)
-
-  const stats = { tables: 0, tablesKeyed: 0, rowsKeyed: 0, failed: 0 }
-
-  try {
-    // Tables that still have at least one un-keyed row.
-    const pending = await db
-      .selectDistinct({ tableId: userTableRows.tableId })
-      .from(userTableRows)
-      .where(isNull(userTableRows.orderKey))
-
-    console.log(
-      `Backfill starting — ${pending.length} table(s) with NULL order_key${dryRun ? ' [DRY RUN]' : ''}`
-    )
-
-    for (const { tableId } of pending) {
-      stats.tables += 1
-      try {
-        const keyed = await db.transaction(async (trx) => {
-          // Serialize with concurrent inserts on this table (same lock the app uses).
-          await trx.execute(
-            sql`SELECT pg_advisory_xact_lock(hashtextextended(${`user_table_rows_pos:${tableId}`}, 0))`
-          )
-          const rows = await trx
-            .select({ id: userTableRows.id })
-            .from(userTableRows)
-            .where(eq(userTableRows.tableId, tableId))
-            .orderBy(asc(userTableRows.position), asc(userTableRows.id))
-
-          if (rows.length === 0) return 0
-          const keys = nKeysBetween(null, null, rows.length)
-          if (dryRun) return rows.length
-
-          // Chunked UPDATE … FROM (VALUES …) mapping id → key (see WRITE_CHUNK_SIZE).
-          for (let start = 0; start < rows.length; start += WRITE_CHUNK_SIZE) {
-            const chunk = rows.slice(start, start + WRITE_CHUNK_SIZE)
-            const values = sql.join(
-              chunk.map((r, i) => sql`(${r.id}, ${keys[start + i]})`),
-              sql`, `
-            )
-            await trx.execute(sql`
-              UPDATE user_table_rows AS t
-              SET order_key = v.order_key
-              FROM (VALUES ${values}) AS v(id, order_key)
-              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('Backfill complete.')
-    console.log(`  tables scanned: ${stats.tables}`)
-    console.log(`  tables keyed:   ${stats.tablesKeyed}`)
-    console.log(`  rows keyed:     ${stats.rowsKeyed}`)
-    console.log(`  failed:         ${stats.failed}`)
-    if (stats.failed > 0) process.exitCode = 1
-  } finally {
-    await client.end({ timeout: 5 }).catch(() => {})
-  }
-}
-
-if ((import.meta as { main?: boolean }).main) {
-  try {
-    await runBackfill()
-  } catch (error) {
-    console.error('Backfill aborted:', getErrorMessage(error))
-    process.exitCode = 1
-  }
-}

apps/sim/scripts/repair-table-order-key-collation.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> {

packages/db/script-migrations/0001_backfill_table_order_keys.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/index.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/types.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>
+}

packages/db/scripts/migrate.ts

@@ -4,6 +4,7 @@ import { backoffWithJitter } from '@sim/utils/retry'
 import { drizzle } from 'drizzle-orm/postgres-js'
 import { migrate } from 'drizzle-orm/postgres-js/migrator'
 import postgres from 'postgres'
+import { runScriptMigrations } from '../script-migrations/index'
 
 /**
  * Concurrent-index convention: plain `CREATE INDEX` write-blocks large/hot
@@ -111,6 +112,8 @@ try {
   try {
     await runMigrationsWithRetry()
     console.log('Migrations applied successfully.')
+    await assertLockSessionHeld()
+    await runScriptMigrations(client)
     await warnOnInvalidIndexes()
   } finally {
     await releaseMigrationLock()
@@ -179,17 +182,25 @@ async function acquireMigrationLock(): Promise<void> {
  * Replays are safe: drizzle rolls the batch back on failure, and post-COMMIT
  * CONCURRENTLY statements are idempotent by convention.
  */
+/**
+ * Verify the session still holds the migration advisory lock: a changed
+ * backend pid means the connection was recycled and the lock silently dropped.
+ * Only sound on a direct connection — see `hasDirectMigrationUrl`.
+ */
+async function assertLockSessionHeld(): Promise<void> {
+  if (!hasDirectMigrationUrl) return
+  const [{ pid }] = await client`SELECT pg_backend_pid() AS pid`
+  if (pid !== lockSessionPid) {
+    throw new Error(
+      `Database session changed mid-run (backend pid ${lockSessionPid} -> ${pid}); ` +
+        'the migration advisory lock was lost. Aborting so a fresh runner can retry safely.'
+    )
+  }
+}
+
 async function runMigrationsWithRetry(): Promise<void> {
   for (let attempt = 1; ; attempt++) {
-    if (hasDirectMigrationUrl) {
-      const [{ pid }] = await client`SELECT pg_backend_pid() AS pid`
-      if (pid !== lockSessionPid) {
-        throw new Error(
-          `Database session changed mid-run (backend pid ${lockSessionPid} -> ${pid}); ` +
-            'the migration advisory lock was lost. Aborting so a fresh runner can retry safely.'
-        )
-      }
-    }
+    await assertLockSessionHeld()
     await client.unsafe('SET statement_timeout = 0')
     await client.unsafe(`SET lock_timeout = '${DDL_LOCK_TIMEOUT}'`)
     try {

packages/utils/package.json

@@ -34,6 +34,10 @@
       "types": "./src/formatting.ts",
       "default": "./src/formatting.ts"
     },
+    "./fractional-indexing": {
+      "types": "./src/fractional-indexing.ts",
+      "default": "./src/fractional-indexing.ts"
+    },
     "./object": {
       "types": "./src/object.ts",
       "default": "./src/object.ts"

packages/utils/src/fractional-indexing.test.ts

@@ -9,11 +9,7 @@
  * assertion and rows would display out of order. These tests would catch that.
  */
 import { describe, expect, it } from 'vitest'
-import {
-  BASE_62_DIGITS,
-  generateKeyBetween,
-  generateNKeysBetween,
-} from '@/lib/fractional-indexing/fractional-indexing'
+import { BASE_62_DIGITS, generateKeyBetween, generateNKeysBetween } from './fractional-indexing'
 
 /** Bytewise (ASCII / UTF-16 code-unit) comparator — what `COLLATE "C"` and the library use. */
 const byteCompare = (a: string, b: string): number => (a < b ? -1 : a > b ? 1 : 0)