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

Microshards support for PostgreSQL

import compact from "lodash/compact";
import pRetry from "p-retry";
import type { TableInfo } from "./getTablesInSchema";
import { log } from "./logging";
import { libSchema, schemaNew, schemaOld, shardNo } from "./names";
import { ident, join, sql } from "./quote";
import { runSql } from "./runSql";

const ACTIVATION_RETRY_OPTIONS: pRetry.Options = {
  factor: 1.5,
  minTimeout: 100,
  maxTimeout: 3000,
  retries: 15,
};

/**
 * Cleans up after the shard migration success.
 */
export async function resultCommit({
  activateOnDestination,
  deactivateSQL,
  fromDsn,
  tables,
  toDsn,
  schema,
}: {
  activateOnDestination: boolean;
  deactivateSQL?: string;
  fromDsn: string;
  tables: TableInfo[];
  toDsn: string;
  schema: string;
}): Promise<void> {
  if (!activateOnDestination) {
    await runSql(
      toDsn,
      sql`
        DROP SCHEMA IF EXISTS ${ident(schemaNew(schema))} CASCADE;
        ALTER SCHEMA ${ident(schema)} RENAME TO ${ident(schemaNew(schema))};
      `,
      `Renaming just-migrated microshard on the destination to ${schemaNew(schema)}`,
    );
    return;
  }

  // E.g. 2023-01-09T07:09:54.253Z
  const dateSuffix = new Date()
    .toISOString()
    .replace(/\..*$/, "")
    .replace(/[-T:Z]/g, "");

  const shard = shardNo(schema);

  await runSql(
    fromDsn,
    join(
      compact([
        sql`BEGIN`,
        shard !== null &&
          sql`SELECT ${libSchema()}.microsharding_ensure_inactive(${shard})`,
        sql`ALTER schema ${ident(schema)} RENAME TO ${ident(schemaOld(schema, dateSuffix))}`,
        sql`COMMIT`,
      ]),
      "; ",
    ),
    "Renaming & deactivating microshard on the source",
  );

  if (shard !== null) {
    await pRetry(
      async () =>
        runSql(
          toDsn,
          sql`SELECT ${libSchema()}.microsharding_ensure_active(${shard})`,
          "DON'T INTERRUPT: Activating microshard on the destination while holding the write lock",
        ).catch((e: unknown) => {
          throw e instanceof Error ? e : Error(`${e}`);
        }),
      {
        ...ACTIVATION_RETRY_OPTIONS,
        onFailedAttempt: (error) =>
          log.warning(
            `Activation attempt ${error.attemptNumber} of ${error.attemptNumber + error.retriesLeft} failed: ${error.message}`,
          ),
      },
    );
  }

  if (deactivateSQL && tables.length > 0) {
    await runSql(
      fromDsn,
      join(
        tables.map(
          ({ schema, name }) =>
            deactivateSQL.replace(
              /\$1/g,
              sql`${sql`${ident(schema)}.${ident(name)}`.toString()}`.toString(),
            ) + ";",
        ),
        "\n",
      ),
      "Running custom deactivation script for microshard tables",
    );
  }
}