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

Microshards support for PostgreSQL

import { createConnectedClient } from "./createConnectedClient";
import type { TableInfo } from "./getTablesInSchema";
import { inspectError } from "./inspectError";
import { log, progress } from "./logging";
import { subName } from "./names";
import { subtractLsn } from "./parseLsn";
import { pollDelay } from "./pollDelay";
import { ident, join, sql } from "./quote";
import { tookToHuman } from "./tookToHuman";

export type Unlock = (status?: "error") => Promise<void>;

/**
 * Locks the tables for write and waits until incremental flow to the
 * destination tables exhausts.
 *
 * If succeeded, the function returns a callback which, when called, unlocks the
 * tables on the source. This callback must be called after deactivating the
 * shard on the source. (But even if it's not called, it'll likely be okay
 * still, because the process will end, PG will disconnect, and the tables will
 * be auto-unlocked.)
 */
export async function waitUntilIncrementalCompletes(
  {
    fromDsn,
    schema,
    tables,
  }: {
    fromDsn: string;
    schema: string;
    tables: TableInfo[];
  },
  throwIfAborted: () => void,
): Promise<Unlock> {
  if (tables.length === 0) {
    return async () => {};
  }

  const waitStartedAt = performance.now();
  const fromClient = await createConnectedClient(fromDsn);
  try {
    const fromServerVersion = await fromClient.serverVersion();
    await fromClient.query("BEGIN");
    await fromClient.query("SET LOCAL statement_timeout TO 0");
    await fromClient.query(
      "SET LOCAL idle_in_transaction_session_timeout TO 0",
    );
    fromServerVersion[0] >= 17 &&
      (await fromClient.query("SET LOCAL transaction_timeout TO 0"));
    await fromClient.query(`SET LOCAL search_path TO ${schema}`);
    throwIfAborted();

    const query = sql`LOCK TABLE ONLY ${join(
      tables.map(({ name }) => sql`${ident(name)}`),
      ", ",
    )} IN EXCLUSIVE MODE`.toString();
    log.shellCmd({
      comment: "Locking source tables to pause all WRITES",
      cmd: `node-postgres ${fromDsn}#${schema}`,
      input: query,
    });
    await fromClient.query(query);
    throwIfAborted();

    const srcLsn = await fromClient.currentWalInsertLsn();
    while (true) {
      const dstLsn = await fromClient.confirmedFlushLsn(subName(schema));
      // srcLsn is constant and in the future; dstLsn is in the past, grows
      // monotonically and catches up srcLsn eventually. Once srcLsn-dstLsn
      // crosses 0, we continue.
      const gap = subtractLsn(srcLsn, dstLsn)!;
      progress(
        `Waiting for dst=${dstLsn} to be greater src=${srcLsn}, gap=${gap}`,
      );
      if (gap <= 0) {
        break;
      }

      await pollDelay();
      throwIfAborted();
    }
  } catch (e: unknown) {
    try {
      await fromClient.query("ROLLBACK");
      await fromClient.end();
    } catch (e: unknown) {
      log.error(`Failed to rollback transaction: ${inspectError(e)}`);
    }

    throw e;
  }

  progress.clear();
  log.shellCmd({
    comment: "Incremental logical replication completed",
    cmd: "",
    input: `Tables were write-locked for ${tookToHuman(waitStartedAt)}`,
  });

  return async (status) => {
    // It doesn't really matter for the locks, whether we commit or rollback
    // here. But in case of a success the "ROLLBACK" message in the logs scare
    // people. So we use "COMMIT" in case of success and "ROLLBACK" otherwise.
    const query = status === "error" ? "ROLLBACK" : "COMMIT";
    log.shellCmd({
      comment: `Unlocking source tables; write lock was held for ${tookToHuman(waitStartedAt)}`,
      cmd: `node-postgres ${fromDsn}`,
      input: query,
    });
    await fromClient.query(query);
    await fromClient.end();
  };
}