@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";

// The downside of turning this on is that, if we run pg_dump over the migrated
// shard, then it will emit "ALTER TABLE ... ADD CONSTRAINT ... NOT VALID" (i.e.
// PG remembers that the FKs were not validated). This is not the end of the
// world, but it's not worth the ~5% overall speedup this option gives.
const SKIP_FK_VALIDATION = false;

/**
 * Runs pg_dump, does some basic parsing and returns the result.
 */
export async function getDump({
  fromDsn,
  schema,
}: {
  fromDsn: string;
  schema: string;
}): 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,
    (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, "-- $&"),
      )
      // If enabled, makes FKs creation instant: we don't need to check
      // integrity at creation time, since the data on the source is already
      // validated. (It will still check at tuples write time though.)
      .map((entry) =>
        SKIP_FK_VALIDATION
          ? entry.replace(
              /^(ALTER TABLE ONLY [^\n]+\n +ADD CONSTRAINT [^\n]+ FOREIGN KEY [^\n]+ REFERENCES [^\n]+)(;)$/s,
              "$1 NOT VALID$2",
            )
          : entry,
      ),
  );
  entries[0] = entries[0].replace(/\n\n+/gs, "\n");
  return entries;
}