stack/packages/db/src/migrate.ts

import { dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';
import { sql } from 'drizzle-orm';
import { drizzle as drizzlePostgres } from 'drizzle-orm/postgres-js';
import { migrate as migratePostgres } from 'drizzle-orm/postgres-js/migrator';
import { readMigrationFiles } from 'drizzle-orm/migrator';
import postgres from 'postgres';
import { DEFAULT_DATABASE_URL } from './defaults.js';
import type { DbHandle } from './client.js';

interface PgliteExecutor {
  exec(query: string): Promise<unknown>;
}

interface ExecuteRows<T> {
  rows: T[];
}

function migrationsFolder(): string {
  const here = dirname(fileURLToPath(import.meta.url));
  return resolve(here, '../drizzle');
}

export async function runMigrations(url?: string): Promise<void> {
  const connectionString = url ?? process.env['DATABASE_URL'] ?? DEFAULT_DATABASE_URL;
  const sqlClient = postgres(connectionString, { max: 1 });
  const db = drizzlePostgres(sqlClient);
  try {
    // TODO: postgres-tier first-install also fails because (a) Drizzle wraps every
    // migration in one transaction (breaks 0009's ALTER TYPE ADD VALUE → SET DEFAULT
    // sequence) and (b) drizzle/meta/_journal.json has 0009 ordered before 0008,
    // which the postgres-js migrator skips by `created_at < folderMillis`. The
    // PGlite path below sidesteps both. A follow-up should either share the
    // per-statement loop (see runPgliteMigrations) or fix the journal ordering.
    await migratePostgres(db, { migrationsFolder: migrationsFolder() });
  } finally {
    await sqlClient.end();
  }
}

// Apply Drizzle migrations against an embedded PGlite database.
//
// We don't reuse drizzle's pglite migrator because it wraps ALL migrations in
// one outer transaction, which breaks Postgres' `check_safe_enum_use` rule —
// e.g. migration 0009 does `ALTER TYPE ADD VALUE 'pending'` then references
// `'pending'` as a default in the same tx. PGlite's `exec()` runs each
// statement under the Simple Query protocol, autocommitting between them.
//
// We still write to the standard `drizzle.__drizzle_migrations` ledger so the
// result is interoperable with `runMigrations()` on a postgres-backed deploy
// (modulo the journal-ordering bug noted above).
//
// We skip-by-hash rather than skip-by-folderMillis (which is what Drizzle's
// postgres-js migrator does). That's deliberate — out-of-order timestamps in
// `_journal.json` won't silently drop migrations.
//
// Failure model: each statement autocommits, and the ledger row is written
// only after all statements in a migration succeed. A crash mid-migration
// leaves the prefix applied with no ledger entry, so the next boot will
// replay those statements and fail loudly on "already exists". Recovery:
// drop the partially-applied objects, or insert the migration's hash into
// `drizzle.__drizzle_migrations` manually. The error log identifies which
// statement of which migration was the culprit.
export async function runPgliteMigrations(handle: DbHandle): Promise<void> {
  const client = (handle.db as unknown as { $client?: PgliteExecutor }).$client;
  if (!client || typeof client.exec !== 'function') {
    throw new Error('runPgliteMigrations: handle.db is not backed by a PGlite client');
  }

  await client.exec('CREATE SCHEMA IF NOT EXISTS drizzle');
  await client.exec(`
    CREATE TABLE IF NOT EXISTS drizzle.__drizzle_migrations (
      id SERIAL PRIMARY KEY,
      hash text NOT NULL,
      created_at bigint
    )
  `);

  const appliedRows = (await handle.db.execute(
    sql`SELECT hash FROM drizzle.__drizzle_migrations`,
  )) as unknown as ExecuteRows<{ hash: string }>;
  const applied = new Set(appliedRows.rows.map((r) => r.hash));

  const migrations = readMigrationFiles({ migrationsFolder: migrationsFolder() });
  for (const migration of migrations) {
    if (applied.has(migration.hash)) continue;

    // Run each statement-breakpoint chunk in its own exec() call so PGlite
    // commits between statements — this is what lets `ALTER TYPE ADD VALUE`
    // become visible before a subsequent statement references the new value.
    for (const [stmtIdx, stmt] of migration.sql.entries()) {
      const trimmed = stmt.trim();
      if (!trimmed) continue;
      try {
        await client.exec(trimmed);
      } catch (err) {
        const cause = err instanceof Error ? err.message : String(err);
        throw new Error(
          `runPgliteMigrations: migration hash=${migration.hash} statement #${stmtIdx} failed: ${cause}\n` +
            `Statement: ${trimmed.slice(0, 200)}${trimmed.length > 200 ? '…' : ''}`,
          { cause: err },
        );
      }
    }
    await handle.db.execute(
      sql`INSERT INTO drizzle.__drizzle_migrations (hash, created_at) VALUES (${migration.hash}, ${migration.folderMillis})`,
    );
  }
}