pwned/packages/import/src/bulkImport.ts

import { createHash } from "node:crypto";
import { once } from "node:events";
import { availableParallelism } from "node:os";
import { Readable } from "node:stream";
import { finished, pipeline } from "node:stream/promises";
import { parseArgs } from "node:util";
import { from as copyFrom } from "pg-copy-streams";
import type { ClientBase } from "pg";
import { pool } from "./db.ts";
import { readLines } from "./lines.ts";
import { createLimiter } from "./concurrencyLimit.ts";
import { WorkerPool } from "./workerPool.ts";
import { resolveWorkerUrl } from "./resolveWorkerUrl.ts";
import type { Entry } from "./upsertBatch.ts";

const HASH_LINE = /^([0-9A-Fa-f]{40})(?::(\d+))?$/;
const CHUNK_SIZE = 5000;
const hashWorkerUrl = resolveWorkerUrl("hashWorker", import.meta.url);

function usage(): never {
  console.error(
    `Usage: npm run bulk-import -w import -- <file> --format=<wordlist|hashes> [options]\n\n` +
      `Rebuilds the pwned_passwords table by merging its existing contents with\n` +
      `the given file. Loads via COPY into an unindexed shadow table and builds\n` +
      `the primary key once at the end (a single sorted bulk build), instead of\n` +
      `the incremental "import" script's per-row random-order index maintenance.\n` +
      `Much faster for large files, at the cost of needing ~2x disk space for the\n` +
      `duration of the run and a brief swap where pwned_passwords is unavailable.\n` +
      `The whole run is one transaction, so a failure at any point leaves\n` +
      `pwned_passwords completely untouched rather than half-migrated.\n\n` +
      `  wordlist   file contains one plaintext password per line\n` +
      `  hashes     file contains one SHA-1 hash per line, optionally "HASH:COUNT"\n` +
      `  file may be gzip-compressed (.gz)\n\n` +
      `Options:\n` +
      `  --jobs, -j   worker threads for wordlist hashing (default: all CPU cores);\n` +
      `               COPY itself only ever uses one connection, so this only\n` +
      `               helps when hashing - not network/DB throughput - is the\n` +
      `               bottleneck\n\n` +
      `Requires the pwned_passwords table to already exist (run migrate first).\n`,
  );
  process.exit(1);
}

function sha1(value: string): string {
  return createHash("sha1").update(value).digest("hex").toUpperCase();
}

function toCopyLine(prefix: string, suffix: string, count: number): string {
  // Prefix/suffix are always plain hex and count is always numeric, so none
  // of COPY's TEXT-format special characters (tab, newline, backslash) can
  // appear and no escaping is needed.
  return `${prefix}\t${suffix}\t${count}\n`;
}

function parseHashLine(line: string): Entry | null {
  const match = HASH_LINE.exec(line.trim());
  if (!match) return null;
  const [, hash, count] = match;
  const upper = hash.toUpperCase();
  return { prefix: upper.slice(0, 5), suffix: upper.slice(5), count: count ? Number(count) : 1 };
}

// Sequential path: used for the hashes format (no CPU-bound work to
// parallelize) and for wordlist with --jobs=1.
async function* toStagingLines(file: string, format: "wordlist" | "hashes") {
  let skipped = 0;
  let seen = 0;
  for await (const line of readLines(file)) {
    seen++;
    let entry: Entry | null;
    if (format === "wordlist") {
      const hash = sha1(line);
      entry = { prefix: hash.slice(0, 5), suffix: hash.slice(5), count: 1 };
    } else {
      entry = parseHashLine(line);
    }
    if (!entry) {
      skipped++;
      continue;
    }
    yield toCopyLine(entry.prefix, entry.suffix, entry.count);
  }
  console.log(`staged ${(seen - skipped).toLocaleString()} lines (${skipped} skipped)`);
}

// Parallel path: hashing is CPU-bound and single-threaded JS can't use
// multiple cores for it, so a worker pool hashes chunks of lines while the
// main thread just writes the results into the one COPY connection (COPY
// itself can't be parallelized here - the whole run is one transaction on
// one connection, for atomicity).
async function copyWordlistParallel(copyStream: NodeJS.WritableStream, file: string, jobs: number) {
  const workers = new WorkerPool<string[], Entry[]>(hashWorkerUrl, jobs);
  const limit = createLimiter(jobs);
  const inFlight: Promise<void>[] = [];

  let seen = 0;
  let chunk: string[] = [];
  const submit = (lines: string[]) =>
    limit(async () => {
      const entries = await workers.run(lines);
      const text = entries.map((e) => toCopyLine(e.prefix, e.suffix, e.count)).join("");
      if (!copyStream.write(text)) {
        await once(copyStream, "drain");
      }
    });

  for await (const line of readLines(file)) {
    chunk.push(line);
    seen++;
    if (chunk.length >= CHUNK_SIZE) {
      inFlight.push(submit(chunk));
      chunk = [];
    }
  }
  if (chunk.length > 0) inFlight.push(submit(chunk));

  await Promise.all(inFlight);
  await workers.close();
  console.log(`staged ${seen.toLocaleString()} lines`);
}

async function timePhase<T>(label: string, fn: () => Promise<T>): Promise<T> {
  const startedAt = Date.now();
  try {
    return await fn();
  } finally {
    console.log(`[${label}] ${((Date.now() - startedAt) / 1000).toFixed(2)}s`);
  }
}

async function loadStaging(client: ClientBase, file: string, format: "wordlist" | "hashes", jobs: number) {
  const copyStream = client.query(copyFrom("COPY pwned_staging (prefix, suffix, count) FROM STDIN"));

  if (format === "wordlist" && jobs > 1) {
    const done = finished(copyStream);
    await copyWordlistParallel(copyStream, file, jobs);
    copyStream.end();
    await done;
    return;
  }

  await pipeline(Readable.from(toStagingLines(file, format)), copyStream);
}

async function main() {
  const { positionals, values } = parseArgs({
    allowPositionals: true,
    options: {
      format: { type: "string" },
      jobs: { type: "string", short: "j" },
    },
  });

  const [file] = positionals;
  const format = values.format;
  if (!file || (format !== "wordlist" && format !== "hashes")) {
    usage();
  }
  const jobs = values.jobs ? Number(values.jobs) : availableParallelism();
  const startedAt = Date.now();

  const client = await pool.connect();
  try {
    // Everything runs in one transaction: pwned_passwords_new is created and
    // fully loaded before anyone outside this transaction can see it, so
    // Postgres skips WAL-logging its data entirely (if we crash, the table
    // just never existed). It also means a failure at any point rolls back
    // to the exact starting state instead of leaving a half-swapped table.
    await client.query("BEGIN");

    console.log("loading into a temporary staging table via COPY...");
    await client.query(
      "CREATE TEMP TABLE pwned_staging (prefix char(5) NOT NULL, suffix char(35) NOT NULL, count bigint NOT NULL)",
    );
    await timePhase("hash+copy", () => loadStaging(client, file, format, jobs));

    console.log("merging with existing data into a fresh, unindexed table...");
    await timePhase("merge", async () => {
      await client.query("CREATE TABLE pwned_passwords_new (LIKE pwned_passwords INCLUDING DEFAULTS)");
      await client.query(`
        INSERT INTO pwned_passwords_new (prefix, suffix, count)
        SELECT prefix, suffix, SUM(count)
        FROM (
          SELECT prefix, suffix, count FROM pwned_staging
          UNION ALL
          SELECT prefix, suffix, count FROM pwned_passwords
        ) combined
        GROUP BY prefix, suffix
      `);
    });

    console.log("building index (one sorted bulk build instead of incremental random inserts)...");
    await timePhase("reindex", async () => {
      await client.query(
        "ALTER TABLE pwned_passwords_new ADD CONSTRAINT pwned_passwords_new_pkey PRIMARY KEY (prefix, suffix)",
      );
      await client.query("CREATE INDEX pwned_passwords_new_prefix_idx ON pwned_passwords_new (prefix)");
    });

    console.log("swapping into place...");
    await timePhase("swap", async () => {
      await client.query("DROP TABLE pwned_passwords");
      await client.query("ALTER TABLE pwned_passwords_new RENAME TO pwned_passwords");
      // Renaming a constraint (rather than its backing index directly) also
      // renames the index, keeping the two in sync.
      await client.query(
        "ALTER TABLE pwned_passwords RENAME CONSTRAINT pwned_passwords_new_pkey TO pwned_passwords_pkey",
      );
      await client.query("ALTER INDEX pwned_passwords_new_prefix_idx RENAME TO pwned_passwords_prefix_idx");
    });

    const { rows } = await client.query("SELECT count(*) FROM pwned_passwords");
    await client.query("COMMIT");
    const totalSeconds = (Date.now() - startedAt) / 1000;
    console.log(`[total] ${totalSeconds.toFixed(2)}s`);
    console.log(`done: pwned_passwords now has ${Number(rows[0].count).toLocaleString()} rows`);
  } catch (err) {
    await client.query("ROLLBACK").catch(() => {});
    throw err;
  } finally {
    client.release();
    await pool.end();
  }
}

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