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 { Client } = pg;

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) {
  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}"`).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 inserts = extractInsertStatements(fs.readFileSync(filePath, 'utf8'));

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

  const prompt = prod
    ? `RESTORE ${inserts.length} rows into PRODUCTION "${dbName}" from:\n  ${filePath}\nThis TRUNCATES all public tables first.`
    : `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}"...`);
  await truncatePublicTables(client);

  // 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}".`);
}

main().catch((error) => {
  console.error(error);
  process.exit(1);
});