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 `_wire_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 = "_wire_migrations";

/** 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();
}

/** 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 TEXT PRIMARY KEY, applied_at TEXT 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);
}

/** Apply all pending migrations (each in a transaction). Returns applied names. */
export async function migrate(db: Db, dir: string): Promise<string[]> {
  const applied = new Set(await appliedMigrations(db));
  const pending = loadMigrations(dir).filter((m) => !applied.has(m.name));
  const done: string[] = [];
  for (const m of pending) {
    await db.tx(async (tx) => {
      if (m.up) await tx.exec(m.up);
      await tx.exec(`INSERT INTO ${MIGRATIONS_TABLE} (name) VALUES (?)`, [m.name]);
    });
    done.push(m.name);
  }
  return done;
}

/** Roll back the most recently applied migration. Returns its name, or null. */
export async function rollback(db: Db, dir: string): Promise<string | null> {
  const applied = await appliedMigrations(db);
  const last = applied[applied.length - 1];
  if (!last) return null;
  const migration = loadMigrations(dir).find((m) => m.name === last);
  await db.tx(async (tx) => {
    if (migration?.down) await tx.exec(migration.down);
    await tx.exec(`DELETE FROM ${MIGRATIONS_TABLE} WHERE name = ?`, [last]);
  });
  return last;
}

/** 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;
}