Art-gallery/scripts/restore-db-data.js

/**
 * Restore a data-only backup (INSERT dumps from npm run dev:db:backup) into a database.
 *
 * Dev:  node scripts/restore-db-data.js --file db/DataBackup/gallery_dev_data_....txt
 * Prod: npm run devtoprod:db:restore -- --file db/DataBackup/gallery_dev_data_....txt
 */
const fs = require('fs');
const path = require('path');
const pg = require('pg');
const {
  assertProdDatabase,
  assertDevDatabase,
  confirmProdAction,
  loadDevPgConfig,
  loadProdPgConfig,
} = require('./db-env');
const { printCliResult } = require('./lib/cli-result');

const { Client } = pg;

// Prod restore keeps these tables untouched (no TRUNCATE, no INSERT from dev backup).
const PROD_RESTORE_SKIP_TABLES = new Set(['curator_audit_log']);

function getInsertTableName(statement) {
  const match = statement.match(/^INSERT INTO "([^"]+)"/i)
    || statement.match(/^INSERT INTO ([a-zA-Z_][a-zA-Z0-9_]*)/i);
  return match ? match[1] : null;
}

/**
 * Older backups serialized TEXT[] (e.g. users.permissions) via JSON.stringify,
 * producing '["a","b"]' which Postgres rejects as an array literal. Rewrite
 * JSON string-array literals to ARRAY['a','b']::text[].
 */
function coerceJsonStringArrayLiterals(statement) {
  return statement.replace(/'(\[(?:"(?:\\.|[^"\\])*"(?:\s*,\s*"(?:\\.|[^"\\])*")*)?\])'/g, (full, jsonBody) => {
    let arr;
    try {
      arr = JSON.parse(jsonBody);
    } catch {
      return full;
    }
    if (!Array.isArray(arr) || !arr.every((item) => typeof item === 'string')) {
      return full;
    }
    if (arr.length === 0) return `ARRAY[]::text[]`;
    const items = arr.map((s) => `'${String(s).replace(/'/g, "''")}'`).join(', ');
    return `ARRAY[${items}]::text[]`;
  });
}

function filterInsertsForRestore(inserts, prod) {
  if (!prod) return inserts;
  return inserts.filter((statement) => {
    const table = getInsertTableName(statement);
    return !table || !PROD_RESTORE_SKIP_TABLES.has(table);
  });
}

function parseArgs(argv) {
  const fileIdx = argv.indexOf('--file');
  if (fileIdx === -1 || !argv[fileIdx + 1]) {
    throw new Error('--file <path> is required');
  }
  return {
    filePath: path.resolve(argv[fileIdx + 1]),
    prod: argv.includes('--prod'),
  };
}

// A single-quoted SQL string is closed only when the running count of quote
// characters is even (doubled '' escapes count as two, so they stay even).
function quotesBalanced(text) {
  let count = 0;
  for (let i = 0; i < text.length; i += 1) {
    if (text[i] === "'") count += 1;
  }
  return count % 2 === 0;
}

// Values (e.g. artist bios) can contain embedded newlines, so a statement may
// span multiple physical lines. Accumulate lines until the quotes are balanced
// and the statement ends with ';'.
function extractInsertStatements(content) {
  const lines = content.split(/\r?\n/);
  const statements = [];
  let current = null;

  for (const line of lines) {
    if (current === null) {
      if (line.startsWith('INSERT INTO ')) {
        current = line;
      } else {
        continue;
      }
    } else {
      current += `\n${line}`;
    }

    if (quotesBalanced(current) && current.trimEnd().endsWith(';')) {
      statements.push(current);
      current = null;
    }
  }

  if (current !== null) {
    statements.push(current);
  }

  return statements;
}

async function truncatePublicTables(client, { excludeTables = new Set() } = {}) {
  const tablesRes = await client.query(
    `SELECT table_name
     FROM information_schema.tables
     WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
     ORDER BY table_name`,
  );
  const names = tablesRes.rows
    .map((r) => r.table_name)
    .filter((name) => !excludeTables.has(name))
    .map((name) => `"${name}"`)
    .join(', ');
  if (!names) return;
  await client.query(`TRUNCATE TABLE ${names} RESTART IDENTITY CASCADE`);
}

async function main() {
  const { filePath, prod } = parseArgs(process.argv.slice(2));

  if (!fs.existsSync(filePath)) {
    throw new Error(`Backup file not found: ${filePath}`);
  }

  const config = prod ? loadProdPgConfig() : loadDevPgConfig();
  const dbName = config.database;
  const allInserts = extractInsertStatements(fs.readFileSync(filePath, 'utf8')).map(
    coerceJsonStringArrayLiterals
  );
  const inserts = filterInsertsForRestore(allInserts, prod);
  const skippedInserts = prod ? allInserts.length - inserts.length : 0;

  if (inserts.length === 0) {
    throw new Error('No INSERT statements found in backup file');
  }

  const skipNote = prod && PROD_RESTORE_SKIP_TABLES.size > 0
    ? `\nExcluded from prod restore: ${[...PROD_RESTORE_SKIP_TABLES].join(', ')}`
    : '';
  const prompt = prod
    ? `RESTORE ${inserts.length} rows into PRODUCTION "${dbName}" from:\n  ${filePath}\nThis TRUNCATES public tables first (except excluded tables).${skipNote}`
    : `Restore ${inserts.length} rows into "${dbName}" from:\n  ${filePath}\nThis TRUNCATES all public tables first.`;

  if (prod) {
    assertProdDatabase(dbName);
  } else {
    assertDevDatabase(dbName);
  }
  await confirmProdAction(prompt);

  const client = new Client(config);
  await client.connect();

  console.log(`Truncating public tables in "${dbName}"...`);
  if (prod && PROD_RESTORE_SKIP_TABLES.size > 0) {
    console.log(`  skipping: ${[...PROD_RESTORE_SKIP_TABLES].join(', ')}`);
  }
  await truncatePublicTables(client, {
    excludeTables: prod ? PROD_RESTORE_SKIP_TABLES : new Set(),
  });

  // Fast path: disable FK/trigger checks for a single-pass load. Requires
  // permission to set session_replication_role (superuser, or on PG 15+ a
  // "GRANT SET ON PARAMETER session_replication_role TO <role>"). If the role
  // lacks that privilege, fall back to a multi-pass insert that retries rows
  // whose foreign keys are not yet satisfied, so no superuser is needed.
  let replicaMode = false;
  try {
    await client.query('SET session_replication_role = replica');
    replicaMode = true;
  } catch (err) {
    if (err.code === '42501') {
      console.warn(
        '  note: cannot set session_replication_role (role is not superuser); using multi-pass insert',
      );
    } else {
      throw err;
    }
  }

  console.log(`Restoring ${inserts.length} INSERT statements...`);
  let restored = 0;
  try {
    if (replicaMode) {
      for (const statement of inserts) {
        await client.query(statement);
        restored += 1;
        if (restored % 500 === 0) {
          console.log(`  ${restored}/${inserts.length}`);
        }
      }
    } else {
      // Each statement auto-commits on its own, so a foreign-key failure only
      // rejects that one row; retry it on a later pass once its parent exists.
      let remaining = inserts;
      let pass = 0;
      while (remaining.length > 0) {
        pass += 1;
        const retry = [];
        let progressed = 0;
        for (const statement of remaining) {
          try {
            await client.query(statement);
            restored += 1;
            progressed += 1;
            if (restored % 500 === 0) {
              console.log(`  ${restored}/${inserts.length}`);
            }
          } catch (err) {
            if (err.code === '23503') {
              retry.push(statement);
            } else {
              throw err;
            }
          }
        }
        if (progressed === 0) {
          throw new Error(
            `Restore stalled on pass ${pass}: ${retry.length} rows have unresolved foreign keys`,
          );
        }
        remaining = retry;
      }
    }
  } finally {
    if (replicaMode) {
      await client.query('SET session_replication_role = DEFAULT');
    }
  }

  await client.end();
  console.log(`Restore complete: ${restored} statements into "${dbName}".`);
  if (skippedInserts > 0) {
    console.log(`Skipped ${skippedInserts} INSERT statement(s) for excluded prod tables.`);
  }

  printCliResult({
    script: 'restore-db-data',
    ok: true,
    summary: `Restore complete into "${dbName}".`,
    details: [
      `Statements restored: ${restored}`,
      ...(skippedInserts > 0
        ? [`Skipped (excluded tables): ${skippedInserts}`]
        : []),
    ],
  });
}

main().catch((error) => {
  console.error(error);
  printCliResult({
    script: 'restore-db-data',
    ok: false,
    summary: error.message || String(error),
  });
});