WRNexusJS/packages/db/src/migrate.ts

/**
 * Migration runner. Migrations are `.sql` files in `app/db/migrations`, each
 * split into `-- +up` and `-- +down` sections. Applied migrations are recorded
 * in a `_wrn_migrations` table so they run exactly once, newest-last.
 *
 * `scaffoldMigration(..., models)` writes an initial migration straight from the
 * TS models — the source of truth — so you don't hand-write the first schema.
 */

import { existsSync, mkdirSync, readFileSync, readdirSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import type { Db } from "./driver.ts";
import { createTableSql, type Dialect } from "./sql.ts";
import type { Model } from "./schema.ts";

export interface Migration {
  name: string;
  up: string;
  down: string;
}

const MIGRATIONS_TABLE = "_wrn_migrations";
const MIGRATION_LOCKS_TABLE = "_wrn_migration_locks";

export interface MigrationRunOptions {
  /** Return pending migration names without executing their SQL. */
  dryRun?: boolean;
  /** Stop safely between migrations. Active database statements cannot be interrupted portably. */
  signal?: AbortSignal;
  /** Coordinate migration runners through the database. Default true. */
  lock?: boolean;
  /** Allow recovery of a lock left by a crashed process. Default 5 minutes. */
  lockTimeoutMs?: number;
}

/** Split a migration file into its `up` and `down` SQL sections. */
export function parseMigration(name: string, content: string): Migration {
  return { name, up: section(content, "up"), down: section(content, "down") };
}

function section(content: string, which: "up" | "down"): string {
  const marker = new RegExp(`^--\\s*\\+${which}\\b.*$`, "mi");
  const match = marker.exec(content);
  if (!match) {
    // A file with no markers at all is treated entirely as `up`.
    return which === "up" && !/^--\s*\+(up|down)\b/im.test(content) ? content.trim() : "";
  }
  const from = content.indexOf("\n", match.index);
  const rest = content.slice(from === -1 ? content.length : from + 1);
  const next = /^--\s*\+(up|down)\b/im.exec(rest);
  return (next ? rest.slice(0, next.index) : rest).trim();
}

function hasExecutableSql(sql: string): boolean {
  let lineComment = false;
  let blockComment = false;
  for (let index = 0; index < sql.length; index++) {
    const char = sql[index]!;
    const next = sql[index + 1] ?? "";
    if (lineComment) {
      if (char === "\n" || char === "\r") lineComment = false;
      continue;
    }
    if (blockComment) {
      if (char === "*" && next === "/") {
        blockComment = false;
        index++;
      }
      continue;
    }
    if (char === "-" && next === "-") {
      lineComment = true;
      index++;
      continue;
    }
    if (char === "/" && next === "*") {
      blockComment = true;
      index++;
      continue;
    }
    // A quote opens a literal, which is executable content in itself, so the
    // scan can stop here without tracking the literal's contents.
    if (char === "'" || char === '"' || char === "`") return true;
    if (!/\s|;/.test(char)) return true;
  }
  return false;
}

function additiveColumnTarget(sql: string): { table: string; column: string } | undefined {
  const executable = sql
    .replace(/\/\*[\s\S]*?\*\//g, " ")
    .replace(/--[^\r\n]*/g, " ")
    .trim();
  const match =
    /^ALTER\s+TABLE\s+([A-Za-z_][A-Za-z0-9_]*)\s+ADD\s+COLUMN\s+([A-Za-z_][A-Za-z0-9_]*)\b[\s\S]*;?\s*$/i.exec(
      executable,
    );
  return match ? { table: match[1]!, column: match[2]! } : undefined;
}

async function additiveColumnAlreadyExists(db: Db, sql: string): Promise<boolean> {
  const target = additiveColumnTarget(sql);
  if (!target) return false;
  if (db.driver.dialect === "sqlite") {
    const columns = await db.all<{ name: string }>(`PRAGMA table_info(${target.table})`);
    return columns.some(({ name }) => name.toLowerCase() === target.column.toLowerCase());
  }
  if (db.driver.dialect === "postgres") {
    return Boolean(
      await db.one(
        "SELECT 1 AS present FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?",
        [target.table, target.column],
      ),
    );
  }
  return Boolean(
    await db.one(
      "SELECT 1 AS present FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?",
      [target.table, target.column],
    ),
  );
}

/** Load and parse all migration files in a directory, sorted by filename. */
export function loadMigrations(dir: string): Migration[] {
  if (!existsSync(dir)) return [];
  return readdirSync(dir)
    .filter((f) => f.endsWith(".sql"))
    .sort()
    .map((f) => parseMigration(f.replace(/\.sql$/, ""), readFileSync(join(dir, f), "utf8")));
}

async function ensureTable(db: Db): Promise<void> {
  await db.exec(
    `CREATE TABLE IF NOT EXISTS ${MIGRATIONS_TABLE} (name VARCHAR(255) PRIMARY KEY, applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP)`,
  );
}

/** Names of already-applied migrations, oldest first. */
export async function appliedMigrations(db: Db): Promise<string[]> {
  await ensureTable(db);
  const rows = await db.all<{ name: string }>(
    `SELECT name FROM ${MIGRATIONS_TABLE} ORDER BY applied_at, name`,
  );
  return rows.map((r) => r.name);
}

function throwIfAborted(signal?: AbortSignal): void {
  if (signal?.aborted) throw signal.reason ?? new DOMException("Migration aborted", "AbortError");
}

async function acquireMigrationLock(db: Db, timeoutMs: number): Promise<() => Promise<void>> {
  if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
    throw new RangeError("migration lockTimeoutMs must be a positive number");
  }
  await db.exec(
    `CREATE TABLE IF NOT EXISTS ${MIGRATION_LOCKS_TABLE} (name VARCHAR(255) PRIMARY KEY, owner VARCHAR(255) NOT NULL, expires_at VARCHAR(40) NOT NULL)`,
  );
  const owner = crypto.randomUUID();
  const now = new Date();
  await db.exec(`DELETE FROM ${MIGRATION_LOCKS_TABLE} WHERE name = ? AND expires_at <= ?`, [
    "global",
    now.toISOString(),
  ]);
  try {
    await db.exec(
      `INSERT INTO ${MIGRATION_LOCKS_TABLE} (name, owner, expires_at) VALUES (?, ?, ?)`,
      ["global", owner, new Date(now.getTime() + timeoutMs).toISOString()],
    );
  } catch (error) {
    const held = await db.all(`SELECT owner FROM ${MIGRATION_LOCKS_TABLE} WHERE name = ?`, [
      "global",
    ]);
    if (held.length === 0) throw error;
    throw new Error("WRN-DB-MIGRATION-LOCKED: another process is running database migrations", {
      cause: error,
    });
  }
  return async () => {
    await db.exec(`DELETE FROM ${MIGRATION_LOCKS_TABLE} WHERE name = ? AND owner = ?`, [
      "global",
      owner,
    ]);
  };
}

/** Apply an ordered migration list (each in a transaction). Returns applied names. */
export async function applyMigrations(
  db: Db,
  migrations: readonly Migration[],
  options: MigrationRunOptions = {},
): Promise<string[]> {
  if (migrations.length === 0) return [];
  throwIfAborted(options.signal);
  const applied = new Set(await appliedMigrations(db));
  const pending = migrations.filter((migration) => !applied.has(migration.name));
  if (options.dryRun || pending.length === 0) return pending.map(({ name }) => name);
  const release =
    options.lock === false
      ? undefined
      : await acquireMigrationLock(db, options.lockTimeoutMs ?? 300_000);
  const done: string[] = [];
  try {
    // Re-read after locking because another runner may have completed while we waited.
    const current = new Set(await appliedMigrations(db));
    for (const migration of pending.filter(({ name }) => !current.has(name))) {
      throwIfAborted(options.signal);
      await db.tx(async (tx) => {
        if (
          hasExecutableSql(migration.up) &&
          !(await additiveColumnAlreadyExists(tx, migration.up))
        ) {
          await tx.exec(migration.up);
        }
        await tx.exec(`INSERT INTO ${MIGRATIONS_TABLE} (name) VALUES (?)`, [migration.name]);
      });
      done.push(migration.name);
    }
    return done;
  } finally {
    await release?.();
  }
}

/** Apply all pending migrations from a directory. */
export async function migrate(
  db: Db,
  dir: string,
  options: MigrationRunOptions = {},
): Promise<string[]> {
  return applyMigrations(db, loadMigrations(dir), options);
}

/** Roll back the most recently applied migration. Returns its name, or null. */
export async function rollback(
  db: Db,
  dir: string,
  options: Omit<MigrationRunOptions, "dryRun"> & { dryRun?: boolean } = {},
): Promise<string | null> {
  throwIfAborted(options.signal);
  const applied = await appliedMigrations(db);
  const last = applied[applied.length - 1];
  if (!last) return null;
  if (options.dryRun) return last;
  const release =
    options.lock === false
      ? undefined
      : await acquireMigrationLock(db, options.lockTimeoutMs ?? 300_000);
  const migration = loadMigrations(dir).find((m) => m.name === last);
  try {
    throwIfAborted(options.signal);
    await db.tx(async (tx) => {
      if (migration && hasExecutableSql(migration.down)) await tx.exec(migration.down);
      await tx.exec(`DELETE FROM ${MIGRATIONS_TABLE} WHERE name = ?`, [last]);
    });
    return last;
  } finally {
    await release?.();
  }
}

/** Full status: every migration file with whether it has been applied. */
export async function status(db: Db, dir: string): Promise<{ name: string; applied: boolean }[]> {
  const applied = new Set(await appliedMigrations(db));
  return loadMigrations(dir).map((m) => ({ name: m.name, applied: applied.has(m.name) }));
}

/** Order models so a referenced table is created before the table referencing it. */
function topoSort(models: Model[]): Model[] {
  const byName = new Map(models.map((m) => [m.name, m]));
  const sorted: Model[] = [];
  const visited = new Set<string>();
  const visit = (m: Model): void => {
    if (visited.has(m.name)) return;
    visited.add(m.name);
    for (const column of Object.values(m.columns)) {
      const ref = column.def.references;
      if (ref && ref.table !== m.name && byName.has(ref.table)) visit(byName.get(ref.table)!);
    }
    sorted.push(m);
  };
  for (const m of models) visit(m);
  return sorted;
}

function nextNumber(dir: string): number {
  if (!existsSync(dir)) return 1;
  let max = 0;
  for (const f of readdirSync(dir)) {
    const m = /^(\d+)/.exec(f);
    if (m) max = Math.max(max, Number(m[1]));
  }
  return max + 1;
}

/**
 * Write a new migration file. With `models`, the `up`/`down` are generated from
 * the TS models (create/drop every table); otherwise empty stubs are written.
 * Returns the created file path.
 */
export function scaffoldMigration(
  dir: string,
  name: string,
  dialect: Dialect,
  models: Model[] = [],
): string {
  mkdirSync(dir, { recursive: true });
  const num = String(nextNumber(dir)).padStart(4, "0");
  const slug =
    name
      .trim()
      .toLowerCase()
      .replace(/[^a-z0-9]+/g, "_")
      .replace(/^_+|_+$/g, "") || "migration";
  const file = join(dir, `${num}_${slug}.sql`);

  let up = "";
  let down = "";
  if (models.length > 0) {
    const ordered = topoSort(models); // referenced tables first
    up = ordered.map((m) => createTableSql(m, dialect)).join("\n\n");
    const quote = dialect === "mysql" ? (s: string) => `\`${s}\`` : (s: string) => `"${s}"`;
    down = ordered
      .slice()
      .reverse()
      .map((m) => `DROP TABLE IF EXISTS ${quote(m.name)};`)
      .join("\n");
  }

  writeFileSync(file, `-- +up\n${up}\n\n-- +down\n${down}\n`, "utf8");
  return file;
}