@clickup/pg-microsharding/src/internal/getDump.ts

Microshards support for PostgreSQL

import compact from "lodash/compact";
import partition from "lodash/partition";
import { pgDump } from "./names";
import type { Sql } from "./quote";
import { join, shellQuote } from "./quote";
import { runShell } from "./runShell";

/**
 * Runs pg_dump, does some basic parsing and returns the result.
 *
 * By default, the tool does not validate FKs on the destination, since it's
 * implied they were already valid on the source. So the FKs on the destination
 * are created with "NOT VALID" clause. This speeds up microshards moving
 * severely. The downside is that, if we later run pg_dump over the migrated
 * destination shard, it will emit "ALTER TABLE ... ADD CONSTRAINT ... NOT
 * VALID" as well (i.e. PG remembers that the FKs were not validated). This is
 * not the end of the world, just inaccurate.
 */
export async function getDump({
  fromDsn,
  schema,
  validateFKs,
}: {
  fromDsn: string;
  schema: string;
  validateFKs: boolean;
}): Promise<{
  preData: Sql;
  primaryKeys: string[];
  postData: Sql;
}> {
  const preData = parseDump(
    await runShell(
      `${pgDump(fromDsn)} -n ${shellQuote(schema)} --section=pre-data`,
      null,
      `Reading pre-data DDL (except indexes) for ${schema} from ${fromDsn}`,
    ),
  );
  const postDataAll = parseDump(
    await runShell(
      `${pgDump(fromDsn)} -n ${shellQuote(schema)} --section=post-data`,
      null,
      `Reading post-data DDL (indexes, foreign keys) for ${schema} from ${fromDsn}`,
    ),
  );
  // For logical replication to work, each table on the subscription part must
  // have an index that allows the logical worker to find the row it receives
  // from the publication end. PostgreSQL 16+ can use regular indexes for that,
  // but the earlier versions require the subscription end to have either a
  // REPLICA IDENTITY or a PRIMARY KEY. Thus, we extract such indexes from the
  // post-data section of the dump and create them before creating the
  // subscription.
  const [primaryKeys, postData] = partition(
    postDataAll.map((entry) =>
      entry.replace(
        /^(ALTER TABLE ONLY [^\n]+\n +ADD CONSTRAINT [^\n]+ FOREIGN KEY [^\n]+ REFERENCES [^\n]+?)(?: NOT VALID)?(;)$/s,
        (_, sql, end) => sql + (validateFKs ? "" : " NOT VALID") + end,
      ),
    ),
    (entry) =>
      entry.match(
        /^ALTER TABLE .* ADD CONSTRAINT .* PRIMARY KEY \([^)]+\);$/s,
      ) ||
      entry.match(
        /^CREATE UNIQUE INDEX (.*?) ON .*\nALTER TABLE .*? REPLICA IDENTITY USING INDEX \1;$/s,
      ),
  );
  return {
    preData: join(preData, "\n"),
    primaryKeys,
    postData: join(postData, "\n"),
  };
}

/**
 * The text generated by pg_dump consists of entries divided by 3 lines of
 * "--"-comments. We split the dump into that entries, and only then analyze it
 * in the further logic.
 */
function parseDump(dump: string): string[] {
  const entries = compact(
    dump
      .split(/^--\n-- \w+[^\n]*\n--\n/m)
      .map((entry) => entry.trim())
      // Added in PG17+ and is not supported in earlier versions.
      .map((entry) =>
        entry.replace(/^SET transaction_timeout = [^\n]+/m, "-- $&"),
      ),
  );
  entries[0] = entries[0].replace(/\n\n+/gs, "\n");
  return entries;
}