Blame for backend/scripts/migrationLib.mjs

import crypto from "crypto";
import fs from "fs";
import path from "path";
import { fileURLToPath } from "url";

const __dirname = path.dirname(fileURLToPath(import.meta.url));

export const ADVISORY_LOCK_KEY = 0x6d696b65;
export const BOOTSTRAP_BASELINE_FILENAME = "000_one_shot_schema.sql";
export const MOCK_AUTH_FILENAME = "999_mock_auth.sql";

/** Incremental migrations already merged into 000_one_shot_schema.sql */
export const BASELINE_INCORPORATED_PREFIXES = ["002", "003", "004"];

/** Migrations expected on an existing production-like database before baselining */
export const BASELINE_MIGRATION_PREFIXES = [
  "001",
  "002",
  "003",
  "004",
  "005",
];

export const SCHEMA_MIGRATIONS_DDL = `
CREATE TABLE IF NOT EXISTS schema_migrations (
  filename text PRIMARY KEY,
  checksum text NOT NULL,
  applied_at timestamptz NOT NULL DEFAULT now()
);
`;

export function getDefaultMigrationsDir() {
  return path.join(__dirname, "../migrations");
}

export function isExcludedFromIncrementalDiscovery(filename) {
  return (
    filename === BOOTSTRAP_BASELINE_FILENAME || filename === MOCK_AUTH_FILENAME
  );
}

export function parseMigrationPrefix(filename) {
  const match = /^(\d{3})_.+\.sql$/u.exec(filename);
  return match ? match[1] : null;
}

export function computeChecksum(content) {
  return crypto.createHash("sha256").update(content, "utf8").digest("hex");
}

export function computeFileChecksum(filePath) {
  return computeChecksum(fs.readFileSync(filePath, "utf8"));
}

export function discoverIncrementalMigrations(migrationsDir) {
  const filenames = fs
    .readdirSync(migrationsDir)
    .filter((name) => name.endsWith(".sql"))
    .filter((name) => !isExcludedFromIncrementalDiscovery(name))
    .filter((name) => parseMigrationPrefix(name) !== null)
    .sort((left, right) =>
      parseMigrationPrefix(left).localeCompare(parseMigrationPrefix(right), undefined, {
        numeric: true,
      }),
    );

  const migrations = filenames.map((filename) => {
    const filePath = path.join(migrationsDir, filename);
    return {
      filename,
      prefix: parseMigrationPrefix(filename),
      path: filePath,
      checksum: computeFileChecksum(filePath),
    };
  });

  const filenameByPrefix = new Map();
  for (const migration of migrations) {
    const existingFilename = filenameByPrefix.get(migration.prefix);
    if (existingFilename) {
      throw new Error(
        `Duplicate migration prefix ${migration.prefix}: ${existingFilename}, ${migration.filename}`,
      );
    }
    filenameByPrefix.set(migration.prefix, migration.filename);
  }

  return migrations;
}

export async function ensureSchemaMigrationsTable(client) {
  await client.query(SCHEMA_MIGRATIONS_DDL);
}

export async function loadAppliedMigrations(client) {
  const result = await client.query(
    "SELECT filename, checksum FROM schema_migrations",
  );
  return new Map(result.rows.map((row) => [row.filename, row.checksum]));
}

async function columnExists(client, tableName, columnName) {
  const result = await client.query(
    `
      SELECT EXISTS (
        SELECT 1
        FROM information_schema.columns
        WHERE table_schema = 'public'
          AND table_name = $1
          AND column_name = $2
      ) AS exists
    `,
    [tableName, columnName],
  );
  return result.rows[0].exists;
}

async function tableExists(client, tableName) {
  const result = await client.query(
    `
      SELECT EXISTS (
        SELECT 1
        FROM information_schema.tables
        WHERE table_schema = 'public'
          AND table_name = $1
      ) AS exists
    `,
    [tableName],
  );
  return result.rows[0].exists;
}

async function indexExists(client, indexName) {
  const result = await client.query(
    `
      SELECT EXISTS (
        SELECT 1
        FROM pg_indexes
        WHERE schemaname = 'public'
          AND indexname = $1
      ) AS exists
    `,
    [indexName],
  );
  return result.rows[0].exists;
}

async function triggerExists(client, tableName, triggerName) {
  const result = await client.query(
    `
      SELECT EXISTS (
        SELECT 1
        FROM information_schema.triggers
        WHERE event_object_schema = 'public'
          AND event_object_table = $1
          AND trigger_name = $2
      ) AS exists
    `,
    [tableName, triggerName],
  );
  return result.rows[0].exists;
}

async function requireColumns(client, tableName, columnNames) {
  for (const columnName of columnNames) {
    if (!(await columnExists(client, tableName, columnName))) {
      throw new Error(
        `Expected column public.${tableName}.${columnName} for migration verification`,
      );
    }
  }
}

async function requireIndexes(client, indexNames) {
  for (const indexName of indexNames) {
    if (!(await indexExists(client, indexName))) {
      throw new Error(
        `Expected index public.${indexName} for migration verification`,
      );
    }
  }
}

async function requireTrigger(client, tableName, triggerName) {
  if (!(await triggerExists(client, tableName, triggerName))) {
    throw new Error(
      `Expected trigger public.${triggerName} on public.${tableName} for migration verification`,
    );
  }
}

export async function verifyMigrationApplied(client, prefix) {
  switch (prefix) {
    case "001":
      await requireColumns(client, "document_edits", [
        "operation",
        "target_block_id",
        "target_start",
        "target_end",
        "base_surface_id",
        "target_text_hash",
      ]);
      return;
    case "002":
      if (!(await tableExists(client, "document_comment_threads"))) {
        throw new Error(
          "Expected table public.document_comment_threads for migration verification",
        );
      }
      if (!(await tableExists(client, "document_comment_replies"))) {
        throw new Error(
          "Expected table public.document_comment_replies for migration verification",
        );
      }
      await requireColumns(client, "document_comment_threads", [
        "document_id",
        "document_version_id",
        "word_comment_id",
        "source_type",
        "body",
        "status",
        "updated_at",
      ]);
      await requireColumns(client, "document_comment_replies", [
        "thread_id",
        "document_id",
        "document_version_id",
        "word_comment_id",
        "source_type",
        "body",
        "updated_at",
      ]);
      await requireIndexes(client, [
        "document_comment_threads_document_idx",
        "document_comment_threads_version_idx",
        "document_comment_threads_version_word_comment_uidx",
        "document_comment_replies_thread_idx",
        "document_comment_replies_document_idx",
        "document_comment_replies_version_idx",
        "document_comment_replies_version_word_comment_uidx",
      ]);
      await requireTrigger(
        client,
        "document_comment_threads",
        "document_comment_threads_set_updated_at",
      );
      await requireTrigger(
        client,
        "document_comment_replies",
        "document_comment_replies_set_updated_at",
      );
      return;
    case "003":
      await requireColumns(client, "document_comment_threads", ["commented_at"]);
      await requireColumns(client, "document_comment_replies", ["commented_at"]);
      await requireIndexes(client, [
        "document_comment_threads_commented_at_idx",
        "document_comment_replies_commented_at_idx",
      ]);
      return;
    case "004":
      await requireColumns(client, "document_comment_replies", ["parent_reply_id"]);
      await requireIndexes(client, [
        "document_comment_replies_parent_reply_idx",
        "document_comment_replies_thread_parent_reply_idx",
      ]);
      return;
    case "005":
      if (!(await tableExists(client, "tabular_review_runs"))) {
        throw new Error(
          "Expected table public.tabular_review_runs for migration verification",
        );
      }
      await requireColumns(client, "tabular_cells", [
        "updated_at",
        "error_reason",
        "run_id",
      ]);
      if (!(await indexExists(client, "uq_tabular_cells_key"))) {
        throw new Error(
          "Expected index uq_tabular_cells_key for migration verification",
        );
      }
      return;
    case "006":
      await requireColumns(client, "user_profiles", ["model"]);
      return;
    default:
      throw new Error(`No verifier configured for migration prefix ${prefix}`);
  }
}

export async function isDatabaseEmpty(client) {
  const result = await client.query(`
    SELECT EXISTS (
      SELECT 1
      FROM information_schema.tables
      WHERE table_schema = 'public'
        AND table_name = 'documents'
    ) AS exists
  `);
  return !result.rows[0].exists;
}

async function executeSqlFile(client, filePath) {
  const sql = fs.readFileSync(filePath, "utf8");
  await client.query(sql);
}

export async function recordMigration(client, filename, checksum) {
  await client.query(
    "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)",
    [filename, checksum],
  );
}

export async function bootstrapFreshDatabase(
  client,
  migrationsDir,
  { allowMockAuthBootstrap = false } = {},
) {
  const mockAuthPath = path.join(migrationsDir, MOCK_AUTH_FILENAME);
  const baselinePath = path.join(migrationsDir, BOOTSTRAP_BASELINE_FILENAME);

  if (allowMockAuthBootstrap && !fs.existsSync(mockAuthPath)) {
    throw new Error(`Missing bootstrap migration ${MOCK_AUTH_FILENAME}`);
  }
  if (!fs.existsSync(baselinePath)) {
    throw new Error(`Missing bootstrap migration ${BOOTSTRAP_BASELINE_FILENAME}`);
  }

  await client.query("BEGIN");
  try {
    if (allowMockAuthBootstrap) {
      await executeSqlFile(client, mockAuthPath);
    }
    await executeSqlFile(client, baselinePath);

    const incrementalMigrations = discoverIncrementalMigrations(migrationsDir);
    for (const migration of incrementalMigrations) {
      if (!BASELINE_INCORPORATED_PREFIXES.includes(migration.prefix)) {
        continue;
      }
      await recordMigration(client, migration.filename, migration.checksum);
    }

    await client.query("COMMIT");
  } catch (error) {
    await client.query("ROLLBACK");
    throw error;
  }
}

export async function applyPendingMigrations(client, migrations, appliedMap) {
  const applied = [];

  for (const migration of migrations) {
    const appliedChecksum = appliedMap.get(migration.filename);
    if (appliedChecksum) {
      if (appliedChecksum !== migration.checksum) {
        throw new Error(
          `Migration ${migration.filename} was modified after it was applied`,
        );
      }
      continue;
    }

    await client.query("BEGIN");
    try {
      await executeSqlFile(client, migration.path);
      await recordMigration(client, migration.filename, migration.checksum);
      await client.query("COMMIT");
      applied.push(migration.filename);
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    }
  }

  return applied;
}

export async function acquireAdvisoryLock(client) {
  await client.query("SELECT pg_advisory_lock($1)", [ADVISORY_LOCK_KEY]);
}

export async function releaseAdvisoryLock(client) {
  await client.query("SELECT pg_advisory_unlock($1)", [ADVISORY_LOCK_KEY]);
}

export async function runMigrations(pool, options = {}) {
  const migrationsDir = options.migrationsDir ?? getDefaultMigrationsDir();
  const allowMockAuthBootstrap = options.allowMockAuthBootstrap ?? false;
  const client = await pool.connect();

  try {
    await acquireAdvisoryLock(client);
    await ensureSchemaMigrationsTable(client);

    let appliedMap = await loadAppliedMigrations(client);
    const migrations = discoverIncrementalMigrations(migrationsDir);

    for (const migration of migrations) {
      const appliedChecksum = appliedMap.get(migration.filename);
      if (appliedChecksum && appliedChecksum !== migration.checksum) {
        throw new Error(
          `Migration ${migration.filename} was modified after it was applied`,
        );
      }
    }

    if (appliedMap.size === 0) {
      if (await isDatabaseEmpty(client)) {
        await bootstrapFreshDatabase(client, migrationsDir, {
          allowMockAuthBootstrap,
        });
        appliedMap = await loadAppliedMigrations(client);
      } else {
        throw new Error(
          "Database has schema but no migration baseline. Run: npm run migrate:baseline",
        );
      }
    }

    const applied = await applyPendingMigrations(client, migrations, appliedMap);
    return { applied };
  } finally {
    await releaseAdvisoryLock(client);
    client.release();
  }
}

export async function baselineDatabase(pool, options = {}) {
  const migrationsDir = options.migrationsDir ?? getDefaultMigrationsDir();
  const client = await pool.connect();

  try {
    await acquireAdvisoryLock(client);
    await ensureSchemaMigrationsTable(client);

    const appliedMap = await loadAppliedMigrations(client);
    if (appliedMap.size > 0) {
      throw new Error(
        "schema_migrations already contains rows; baseline is only for empty tracking tables",
      );
    }

    const migrations = discoverIncrementalMigrations(migrationsDir);
    const baselineMigrations = migrations.filter((migration) =>
      BASELINE_MIGRATION_PREFIXES.includes(migration.prefix),
    );

    if (baselineMigrations.length !== BASELINE_MIGRATION_PREFIXES.length) {
      throw new Error(
        "Expected baseline migrations 001 through 005 to exist on disk",
      );
    }

    for (const migration of baselineMigrations) {
      await verifyMigrationApplied(client, migration.prefix);
    }

    await client.query("BEGIN");
    try {
      for (const migration of baselineMigrations) {
        await recordMigration(client, migration.filename, migration.checksum);
      }
      await client.query("COMMIT");
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    }

    return {
      recorded: baselineMigrations.map((migration) => migration.filename),
    };
  } finally {
    await releaseAdvisoryLock(client);
    client.release();
  }
}