@clickup/pg-microsharding/src/api/move.ts

Microshards support for PostgreSQL

import chalk from "chalk";
import delay from "delay";
import compact from "lodash/compact";
import { advanceSequences } from "../internal/advanceSequences";
import { analyzeTables } from "../internal/analyzeTables";
import { changeTableOwners } from "../internal/changeTableOwners";
import { cleanUpPubSub } from "../internal/cleanUpPubSub";
import { disableSubscription } from "../internal/disableSubscription";
import { enableSubscription } from "../internal/enableSubscription";
import { getDump } from "../internal/getDump";
import { getTablesInSchema } from "../internal/getTablesInSchema";
import { inspectError } from "../internal/inspectError";
import { log, progress, stdout } from "../internal/logging";
import { join } from "../internal/quote";
import { resultAbort } from "../internal/resultAbort";
import { resultCommit } from "../internal/resultCommit";
import {
  runSqlProgressNonTransactional,
  runSqlProgressTransactional,
} from "../internal/runSqlProgress";
import { startCopyingTables } from "../internal/startCopyingTables";
import { validatePrimaryKeysFromDump } from "../internal/validatePrimaryKeysFromDump";
import { waitUntilBackfillCompletes } from "../internal/waitUntilBackfillCompletes";
import { waitUntilEverythingIsNotLaggingMuch } from "../internal/waitUntilEverythingIsNotLaggingMuch";
import type { Unlock } from "../internal/waitUntilIncrementalCompletes";
import { waitUntilIncrementalCompletes } from "../internal/waitUntilIncrementalCompletes";
import { wrapSigInt } from "../internal/wrapSigInt";

/**
 * Moves or copies a schema from one master DB to another.
 */
export async function move({
  schema,
  fromDsn,
  toDsn,
  commitAction,
  deactivateSQL,
  maxReplicationLagSec = 20,
  wait = false,
}: {
  schema: string;
  fromDsn: string;
  toDsn: string;
  commitAction: "deactivate-activate" | "rename-to-new" | null;
  deactivateSQL?: string;
  maxReplicationLagSec?: number;
  wait?: boolean;
}): Promise<void> {
  let unlock: Unlock = async (): Promise<void> => {};
  try {
    const { tables, partitionedTables } = await getTablesInSchema({
      fromDsn,
      schema,
    });
    const { preData, primaryKeys, postData } = await getDump({
      fromDsn,
      schema,
    });
    validatePrimaryKeysFromDump({
      primaryKeys,
      allTables: [...tables, ...partitionedTables],
    });

    await waitUntilEverythingIsNotLaggingMuch({
      when: "before starting the work",
      fromDsn: undefined, // only destination physical replicas are checked
      tables,
      toDsn,
      schema,
      maxReplicationLagSec,
      waitForUserToContinue: false,
    });

    try {
      await cleanUpPubSub({ fromDsn, toDsn, schema });
      await runSqlProgressTransactional(
        toDsn,
        preData,
        `Applying initial DDL (except indexes) for ${schema} to ${toDsn}`,
      );
    } catch (e: unknown) {
      // No cleanup if failed here since it's one transaction.
      log.error(inspectError(e));
      throw (
        "Exited abnormally with no-op" +
        (typeof e === "string" ? ` (${e})` : "")
      );
    }

    try {
      await wrapSigInt(async (throwIfAborted) => {
        if (primaryKeys.length > 0) {
          await runSqlProgressTransactional(
            toDsn,
            join(primaryKeys, "\n"),
            "Creating minimal indexes to let the logical replication replay the data",
          );
        }

        // During the tables copying, change their owner to the toDsn user,
        // assuming that toDsn's user may have custom values for e.g. work_mem
        // or other parameters. PostgreSQL runs tablesync replication worker
        // with the privileges and settings of the table's owner, so it allows
        // us to have "per-tablesync" limits customization. E.g. we may want a
        // way higher value for work_mem (gigabytes) to prevent excessive
        // creation of temp files, and it's safe, since there are just a few
        // tablesync workers running concurrently in the cluster (depends on
        // max_sync_workers_per_subscription PostgreSQL setting).
        const newOwner = new URL(toDsn).username;
        const tablesWithOtherOwner = tables.filter(
          ({ owner }) =>
            owner !== newOwner || process.env["DEBUG_CHANGE_TABLE_OWNERS"],
        );
        await changeTableOwners({
          toDsn,
          tables: tablesWithOtherOwner.map(({ schema, name }) => ({
            schema,
            name,
            owner: newOwner,
          })),
        });

        throwIfAborted();
        await startCopyingTables({ fromDsn, tables, toDsn, schema });

        throwIfAborted();
        await waitUntilBackfillCompletes(
          { fromDsn, tables, toDsn, schema },
          throwIfAborted,
        );

        throwIfAborted();
        await changeTableOwners({
          toDsn,
          tables: tablesWithOtherOwner,
          comment:
            "Restoring table(s) ownership back, since tablesync workers have completed",
        });

        throwIfAborted();
        await waitUntilEverythingIsNotLaggingMuch(
          {
            when: "before creating heavy indexes",
            fromDsn,
            tables,
            toDsn,
            schema,
            maxReplicationLagSec,
            waitForUserToContinue: wait,
          },
          throwIfAborted,
        );

        throwIfAborted();
        await disableSubscription({ toDsn, tables, schema });

        throwIfAborted();
        await runSqlProgressNonTransactional(
          toDsn,
          postData,
          "Creating heavy indexes, foreign keys etc.",
        );

        throwIfAborted();
        await enableSubscription({ toDsn, tables, schema });

        throwIfAborted();
        await analyzeTables({ toDsn, tables });

        throwIfAborted();
        await waitUntilEverythingIsNotLaggingMuch(
          {
            when: commitAction
              ? "before activating the destination microshard"
              : "before removing the logical subscription",
            fromDsn,
            tables,
            toDsn,
            schema,
            maxReplicationLagSec,
            waitForUserToContinue: wait,
          },
          throwIfAborted,
        );

        throwIfAborted();
        unlock = await waitUntilIncrementalCompletes(
          { fromDsn, schema, tables },
          throwIfAborted,
        );

        // This runs while the fromDsn tables are still locked.
        await advanceSequences({ fromDsn, toDsn });
      });
    } catch (e: unknown) {
      log.error(inspectError(e));
      log("Cleaning up the half-migrated destination.");

      try {
        await unlock("error");
      } catch (e: unknown) {
        log.error(`Unlock failed: ${inspectError(e)}`);
      } finally {
        unlock = async () => {};
      }

      await delay(1000);
      await resultAbort({ fromDsn, toDsn, schema });
      throw "Exited abnormally" + (typeof e === "string" ? ` (${e})` : "");
    }

    try {
      await cleanUpPubSub({ fromDsn, toDsn, schema });
      if (commitAction) {
        await resultCommit({
          activateOnDestination: commitAction === "deactivate-activate",
          deactivateSQL,
          fromDsn,
          tables,
          toDsn,
          schema,
        });
      }
    } catch (e: unknown) {
      log.error(`Commit failed: ${inspectError(e)}`);

      try {
        await unlock("error");
      } catch (e: unknown) {
        log.error(`Unlock failed: ${inspectError(e)}`);
      } finally {
        unlock = async () => {};
      }

      const dangerBand = chalk.bgRed(
        chalk.white(
          "💀 DANGER! "
            .repeat(300)
            .trimEnd()
            .substring(0, (stdout.columns || 80) * 3),
        ),
      );
      throw compact([
        dangerBand,
        "Exited abnormally while committing the result!",
        `Schema ${schema} is in half-working state.`,
        commitAction &&
          `E.g. it may already be inactive on ${fromDsn}, but not yet activated on ${toDsn}.`,
        "You will likely need a manual intervention.",
        dangerBand,
      ]).join("\n");
    }
  } finally {
    progress.clear();
    await unlock();
  }
}