stack/packages/db/src/migrate.ts

import { readFileSync } from 'node:fs';
import { dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';
import { sql } from 'drizzle-orm';
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');
}

/* ------------------------------------------------------------------ */
/*  Shared hash-ledger migration core (#1392 / #1402)                  */
/* ------------------------------------------------------------------ */
//
// Both tiers migrate through this single core, which applies migrations in
// JOURNAL ORDER, one statement at a time, and skips by HASH — never by
// folderMillis timestamp. The previous postgres path delegated to drizzle's
// postgres-js migrator, which:
//
//   * applies only migrations with folderMillis > last-applied, silently
//     skipping journal entries whose `when` is older than the ledger's newest
//     stamp — 0009/0010 carry April-2025 timestamps below 0008's June-2026
//     one, so any database last migrated in the 0008 era silently loses
//     0009/0010 forever (#1402 D1); and
//   * wraps each migration in ONE transaction, which breaks migrations that
//     do `ALTER TYPE ADD VALUE` and then reference the new value in the same
//     migration (0009) — Postgres' check_safe_enum_use rejects it (#1402 D2).
//
// Per-statement execution (each statement autocommits) and skip-by-hash fix
// both. The PGlite path has run this way since it was written; this is the
// TODO it left behind, now shared instead of duplicated.

/** One migration as loaded from the shipped drizzle/ folder. */
export interface MigrationPlanEntry {
  hash: string;
  folderMillis: number;
  statements: string[];
}

/** The persistence operations the hash-ledger core needs, per tier. */
export interface HashLedgerDeps {
  /** Create the drizzle schema + ledger table if absent (idempotent). */
  ensureLedger(): Promise<void>;
  /** Hashes already recorded in the ledger. */
  appliedHashes(): Promise<string[]>;
  /** Record one fully-applied migration in the ledger. */
  recordApplied(hash: string, folderMillis: number): Promise<void>;
  /** Execute one SQL statement, autocommitting (never inside a wider tx). */
  runStatement(statement: string): Promise<void>;
}

function loadPlan(): MigrationPlanEntry[] {
  return readMigrationFiles({ migrationsFolder: migrationsFolder() }).map((m) => ({
    hash: m.hash,
    folderMillis: m.folderMillis,
    statements: m.sql.map((s) => s.trim()).filter((s) => s.length > 0),
  }));
}

/**
 * Apply every unapplied migration in journal order, skipping by hash.
 *
 * Failure model: each statement autocommits, and the ledger row is written
 * only after all statements of a migration succeed. A crash mid-migration
 * leaves the prefix applied with no ledger entry, so the next boot replays
 * those statements and fails loudly on "already exists". Recovery: drop the
 * partially-applied objects, or insert the migration's hash into
 * `drizzle.__drizzle_migrations` manually. The thrown error identifies the
 * statement and migration that failed.
 */
export async function applyMigrationsByHash(
  deps: HashLedgerDeps,
  plan: MigrationPlanEntry[] = loadPlan(),
): Promise<{ applied: number; skipped: number }> {
  await deps.ensureLedger();
  const alreadyApplied = new Set(await deps.appliedHashes());

  let applied = 0;
  let skipped = 0;
  for (const migration of plan) {
    if (alreadyApplied.has(migration.hash)) {
      skipped += 1;
      continue;
    }
    for (const [stmtIdx, stmt] of migration.statements.entries()) {
      try {
        await deps.runStatement(stmt);
      } catch (err) {
        const cause = err instanceof Error ? err.message : String(err);
        throw new Error(
          `migration hash=${migration.hash} statement #${stmtIdx} failed: ${cause}\n` +
            `Statement: ${stmt.slice(0, 200)}${stmt.length > 200 ? '…' : ''}`,
          { cause: err },
        );
      }
    }
    await deps.recordApplied(migration.hash, migration.folderMillis);
    applied += 1;
  }
  return { applied, skipped };
}

const LEDGER_DDL = [
  'CREATE SCHEMA IF NOT EXISTS drizzle',
  `CREATE TABLE IF NOT EXISTS drizzle.__drizzle_migrations (
      id SERIAL PRIMARY KEY,
      hash text NOT NULL,
      created_at bigint
    )`,
];

function connectionString(url?: string): string {
  return url ?? process.env['DATABASE_URL'] ?? DEFAULT_DATABASE_URL;
}

/**
 * Apply Drizzle migrations against a postgres database, hash-ledger style.
 * Idempotent: re-running against a fully-migrated database applies nothing.
 */
export async function runMigrations(url?: string): Promise<void> {
  const sqlClient = postgres(connectionString(url), { max: 1 });
  try {
    await applyMigrationsByHash({
      ensureLedger: async () => {
        for (const ddl of LEDGER_DDL) await sqlClient.unsafe(ddl);
      },
      appliedHashes: async () => {
        const rows = (await sqlClient.unsafe(
          'SELECT hash FROM drizzle.__drizzle_migrations',
        )) as Array<{ hash: string }>;
        return rows.map((r) => String(r.hash));
      },
      recordApplied: async (hash, folderMillis) => {
        await sqlClient.unsafe(
          'INSERT INTO drizzle.__drizzle_migrations (hash, created_at) VALUES ($1, $2)',
          [hash, folderMillis],
        );
      },
      runStatement: async (stmt) => {
        await sqlClient.unsafe(stmt);
      },
    });
  } finally {
    await sqlClient.end();
  }
}

/**
 * Apply Drizzle migrations against an embedded PGlite database.
 *
 * We don't reuse drizzle's pglite migrator for the same reasons as the
 * postgres path (single-transaction wrap; folderMillis skip). PGlite's
 * `exec()` runs each statement under the Simple Query protocol,
 * autocommitting between them — exactly the semantics the shared core needs.
 *
 * The ledger rows this writes are interoperable with the postgres path (same
 * schema, same hashes), because both consume the same shipped migrations.
 */
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 applyMigrationsByHash({
    ensureLedger: async () => {
      for (const ddl of LEDGER_DDL) await client.exec(ddl);
    },
    appliedHashes: async () => {
      const rows = (await handle.db.execute(
        sql`SELECT hash FROM drizzle.__drizzle_migrations`,
      )) as unknown as ExecuteRows<{ hash: string }>;
      return rows.rows.map((r) => String(r.hash));
    },
    recordApplied: async (hash, folderMillis) => {
      await handle.db.execute(
        sql`INSERT INTO drizzle.__drizzle_migrations (hash, created_at) VALUES (${hash}, ${folderMillis})`,
      );
    },
    runStatement: async (stmt) => {
      await client.exec(stmt);
    },
  });
}

/* ------------------------------------------------------------------ */
/*  Migration status (#1392: the installer must VERIFY, not assume)     */
/* ------------------------------------------------------------------ */

/** Read the journal tags (migration folder names) in journal order. */
export function readJournalTags(folder: string = migrationsFolder()): string[] {
  const journal = JSON.parse(readFileSync(resolve(folder, 'meta', '_journal.json'), 'utf8')) as {
    entries?: Array<{ tag?: string }>;
  };
  return (journal.entries ?? []).map((e) => e.tag ?? '').filter((t) => t.length > 0);
}

export interface MigrationStatus {
  /** Hashes recorded in the database's ledger (0 if no ledger exists). */
  appliedCount: number;
  /** Migrations shipped in this package's drizzle/ folder. */
  expectedCount: number;
  /** The tag (folder name) of the last journal entry, for error messages. */
  expectedLastTag: string;
  /** True iff every shipped migration's hash is in the ledger. */
  complete: boolean;
}

/**
 * Report whether a postgres database carries the full shipped schema.
 *
 * Read-only apart from `ensureLedger` semantics: it never creates the ledger
 * (unlike the migrators), so a database with NO ledger reports
 * appliedCount=0 / complete=false — the exact #1392/#1389 signature (an
 * install whose dependency set shipped no migrations at all).
 */
export async function getMigrationStatus(url?: string): Promise<MigrationStatus> {
  const sqlClient = postgres(connectionString(url), { max: 1 });
  try {
    const plan = loadPlan();
    const tags = readJournalTags();
    const expectedHashes = new Set(plan.map((m) => m.hash));

    let appliedHashes: string[] = [];
    const regRows = (await sqlClient.unsafe(
      "SELECT to_regclass('drizzle.__drizzle_migrations') AS reg",
    )) as Array<{ reg: string | null }>;
    if (regRows[0]?.reg) {
      const rows = (await sqlClient.unsafe(
        'SELECT hash FROM drizzle.__drizzle_migrations',
      )) as Array<{ hash: string }>;
      appliedHashes = rows.map((r) => String(r.hash));
    }

    const appliedSet = new Set(appliedHashes);
    return {
      appliedCount: appliedHashes.length,
      expectedCount: plan.length,
      expectedLastTag: tags[tags.length - 1] ?? '',
      complete:
        plan.length > 0 &&
        plan.every((m) => appliedSet.has(m.hash)) &&
        // A ledger with entries OUTSIDE the shipped plan means the database
        // came from a different (e.g. newer) build — not "complete" either.
        appliedHashes.every((h) => expectedHashes.has(h)),
    };
  } finally {
    await sqlClient.end();
  }
}