NianAIGC/scripts/migrate-postgres.mjs

import { createHash } from "node:crypto";
import { existsSync, readFileSync, readdirSync } from "node:fs";
import { readFile } from "node:fs/promises";
import { basename, join } from "node:path";
import { fileURLToPath } from "node:url";
import {
  closePostgresPool,
  createPostgresPool,
  getScriptDataBackend,
  quotePostgresIdentifier
} from "./postgres-client.mjs";

loadEnvFile(".env");
loadEnvFile(".env.local");

const MIGRATION_LOCK_ID = "7308731946202608";
const migrationsDirectory = fileURLToPath(new URL("../database/migrations/", import.meta.url));
let pool;
let client;
let locked = false;

try {
  if (getScriptDataBackend() !== "postgres") {
    throw new Error("Database migrations require ZHINIAN_DATA_BACKEND=postgres");
  }
  const applicationRole = process.env.DATABASE_APP_ROLE?.trim();
  if (!applicationRole && process.env.NODE_ENV === "production") {
    throw new Error("DATABASE_APP_ROLE is required in production so application privileges can be provisioned");
  }
  if (applicationRole) quotePostgresIdentifier(applicationRole);
  pool = createPostgresPool({ applicationName: "zhinian-migrate" });
  client = await pool.connect();
  await client.query("SELECT pg_advisory_lock($1)", [MIGRATION_LOCK_ID]);
  locked = true;
  await client.query(`
    CREATE TABLE IF NOT EXISTS schema_migrations (
      version text PRIMARY KEY,
      checksum text NOT NULL,
      applied_at timestamptz NOT NULL DEFAULT now()
    )
  `);

  const migrations = await discoverMigrations(migrationsDirectory);
  const { rows: appliedRows } = await client.query("SELECT version, checksum FROM schema_migrations");
  const applied = new Map(appliedRows.map((row) => [row.version, row.checksum]));
  const discoveredVersions = new Set(migrations.map((migration) => migration.version));

  for (const version of applied.keys()) {
    if (!discoveredVersions.has(version)) {
      throw new Error(`Applied migration ${version} is missing from the migration directory; refusing to continue`);
    }
  }

  for (const migration of migrations) {
    const recordedChecksum = applied.get(migration.version);
    if (recordedChecksum && recordedChecksum !== migration.checksum) {
      throw new Error(`Applied migration ${migration.version} has changed; refusing to continue`);
    }
  }

  for (const migration of migrations) {
    if (applied.has(migration.version)) continue;
    await client.query("BEGIN");
    try {
      await client.query(migration.sql);
      await client.query(
        "INSERT INTO schema_migrations(version, checksum) VALUES ($1, $2)",
        [migration.version, migration.checksum]
      );
      await client.query("COMMIT");
      console.log(`Applied migration ${migration.version}`);
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    }
  }

  if (applicationRole) {
    await provisionApplicationRole(client, applicationRole);
    await verifyApplicationRole(client, applicationRole);
    console.log("Provisioned PostgreSQL privileges for the configured application role");
  }

  console.log(`Database migrations are current (${migrations.length} discovered)`);
} catch (error) {
  console.error(`Database migration failed: ${safeErrorMessage(error)}`);
  process.exitCode = 1;
} finally {
  if (client) {
    if (locked) {
      try {
        await client.query("SELECT pg_advisory_unlock($1)", [MIGRATION_LOCK_ID]);
      } catch {
        // Closing the session below also releases the advisory lock.
      }
    }
    client.release();
  }
  await closePostgresPool(pool);
}

async function provisionApplicationRole(client, role) {
  const quotedRole = quotePostgresIdentifier(role);
  const applicationTables = applicationRoleTablePrivileges();
  const managedTableNames = applicationTables.map(([table]) => `public.${quotePostgresIdentifier(table)}`).join(", ");
  await client.query("BEGIN");
  try {
    await client.query("REVOKE CREATE ON SCHEMA public FROM PUBLIC");
    await client.query(`GRANT USAGE ON SCHEMA public TO ${quotedRole}`);
    await client.query(`REVOKE ALL ON TABLE ${managedTableNames} FROM ${quotedRole}`);
    for (const [table, privileges] of applicationTables) {
      await client.query(
        `GRANT ${privileges} ON TABLE public.${quotePostgresIdentifier(table)} TO ${quotedRole}`
      );
    }
    await client.query("REVOKE ALL ON FUNCTION public.claim_generation_jobs(text, integer, integer) FROM PUBLIC");
    await client.query(
      "REVOKE ALL ON FUNCTION public.billing_post_wallet_entry(text, text, text, text, text, bigint, text, text, text, jsonb) FROM PUBLIC"
    );
    await client.query(
      `GRANT EXECUTE ON FUNCTION public.claim_generation_jobs(text, integer, integer) TO ${quotedRole}`
    );
    await client.query(
      `GRANT EXECUTE ON FUNCTION public.billing_post_wallet_entry(text, text, text, text, text, bigint, text, text, text, jsonb) TO ${quotedRole}`
    );
    await client.query("COMMIT");
  } catch (error) {
    await client.query("ROLLBACK");
    throw error;
  }
}

async function verifyApplicationRole(client, role) {
  const checks = applicationRoleTablePrivileges().map(([table, privileges]) => [
    `public.${table}`,
    privileges.replaceAll(" ", "")
  ]);
  const result = await client.query(
    `SELECT
       bool_and(has_table_privilege($1, table_name, privileges)) AS tables_ready,
       has_function_privilege($1, 'public.claim_generation_jobs(text,integer,integer)', 'EXECUTE') AS claim_ready,
       has_function_privilege(
         $1,
         'public.billing_post_wallet_entry(text,text,text,text,text,bigint,text,text,text,jsonb)',
         'EXECUTE'
       ) AS billing_ready
     FROM unnest($2::text[], $3::text[]) AS required(table_name, privileges)`,
    [role, checks.map(([table]) => table), checks.map(([, privileges]) => privileges)]
  );
  const status = result.rows[0];
  if (!status?.tables_ready || !status.claim_ready || !status.billing_ready) {
    throw new Error("Application role privilege verification failed");
  }
}

function applicationRoleTablePrivileges() {
  return [
    ["assets", "SELECT, INSERT, DELETE"],
    ["generation_jobs", "SELECT, INSERT, UPDATE, DELETE"],
    ["usage_events", "SELECT, INSERT, UPDATE"],
    ["projects", "SELECT, UPDATE"],
    ["image_templates", "SELECT, INSERT, UPDATE, DELETE"],
    ["platform_organizations", "SELECT, INSERT, UPDATE, DELETE"],
    ["platform_users", "SELECT, INSERT, UPDATE, DELETE"],
    ["platform_account_migrations", "SELECT, INSERT, UPDATE"],
    ["billing_price_rules", "SELECT, INSERT, UPDATE"],
    ["billing_wallets", "SELECT, INSERT, UPDATE"],
    ["billing_ledger", "SELECT, INSERT"]
  ];
}

async function discoverMigrations(directory) {
  const files = readdirSync(directory, { withFileTypes: true })
    .filter((entry) => entry.isFile() && entry.name.endsWith(".sql"))
    .map((entry) => entry.name)
    .sort((left, right) => left.localeCompare(right, "en"));
  const migrations = [];
  for (const file of files) {
    const sql = await readFile(join(directory, file), "utf8");
    migrations.push({
      version: basename(file, ".sql"),
      checksum: createHash("sha256").update(sql).digest("hex"),
      sql
    });
  }
  return migrations;
}

function loadEnvFile(path) {
  if (!existsSync(path)) return;
  for (const line of readFileSync(path, "utf8").split(/\r?\n/)) {
    const match = line.match(/^\s*([A-Z][A-Z0-9_]*)\s*=\s*(.*)\s*$/);
    if (!match || process.env[match[1]]) continue;
    process.env[match[1]] = match[2].replace(/^['"]|['"]$/g, "");
  }
}

function safeErrorMessage(error) {
  if (!(error instanceof Error)) return "Unknown error";
  let message = error.message;
  const connectionString = process.env.DATABASE_URL?.trim();
  if (connectionString) message = message.replaceAll(connectionString, "[redacted DATABASE_URL]");
  return message.replace(/postgres(?:ql)?:\/\/[^\s@]+@/gi, "postgresql://[redacted]@");
}