@clickup/pg-microsharding/src/actions/actionFactor.ts

Microshards support for PostgreSQL

import difference from "lodash/difference";
import flatten from "lodash/flatten";
import groupBy from "lodash/groupBy";
import round from "lodash/round";
import sortBy from "lodash/sortBy";
import prompts from "prompts";
import { table } from "table";
import { factor } from "../api/factor";
import { discoverShards } from "../internal/discoverShards";
import { getSchemaWeights } from "../internal/getSchemaWeights";
import { print } from "../internal/logging";
import { dsnShort } from "../internal/names";
import { normalizeDsns } from "../internal/normalizeDsns";
import type { Args } from "../internal/parseArgs";
import { promiseAllMap } from "../internal/promiseAllMap";
import { actionList } from "./actionList";

interface FactorSpec {
  op: "" | "+" | "-" | "*";
  num: number;
}

/**
 * Sets the weight factor for some shards.
 */
export async function actionFactor(args: Args): Promise<boolean> {
  const dsns = await normalizeDsns(args["dsns"] || args["dsn"]);
  if (dsns.length === 0) {
    throw "Please provide --dsn or --dsns, DB DSN to adjust microshard(s) factor at";
  }

  let shardNos: number[];
  if ((args["shard"] || args["shards"])?.match(/^(\d+(?:,\d+)*)$/)) {
    shardNos = RegExp.$1.split(",").map(Number);
  } else if (args["shards"]) {
    const matchingDsns = dsns.filter(
      (dsn) =>
        dsn.startsWith(args["shards"]!) ||
        dsnShort(dsn).startsWith(args["shards"]!),
    );
    if (matchingDsns.length === 0) {
      await actionList(args).catch(() => {});
      throw "If you provide --shards, it must be in format: N,M,... or DSN-PREFIX";
    }

    shardNos = (await discoverShards({ dsns: matchingDsns })).map(
      ({ shard }) => shard,
    );
  } else {
    await actionList(args).catch(() => {});
    throw "Please provide --shard=N or --shards=N,M..., microshard numbers to set the factor for";
  }

  let factorSpec: FactorSpec;
  if (args["factor"]?.match(/^([-+*])?(\d+(\.\d+)?)$/)) {
    factorSpec = {
      op: RegExp.$1 as any,
      num: parseFloat(RegExp.$2),
    };
  } else {
    await actionList(args).catch(() => {});
    throw 'Please provide --factor=P|+P.Q|-P.Q|"*P.Q", factor to set for the microshards';
  }

  const allShards = await discoverShards({ dsns });
  const matchingShards = sortBy(
    allShards.filter(({ shard }) => shardNos.includes(shard)),
    ({ shard }) => shard,
  );
  if (matchingShards.length === 0) {
    throw "No microshards found to change the weight factors for.";
  }

  const missingShards = difference(
    shardNos,
    allShards.map(({ shard }) => shard),
  );
  if (missingShards.length > 0) {
    throw `Microshards ${missingShards.join(", ")} not found in any of the DSNs`;
  }

  const plan = flatten(
    await promiseAllMap(
      Object.entries(groupBy(matchingShards, "dsn")),
      async ([dsn, shards]) => {
        const weights = await getSchemaWeights({
          dsn,
          schemas: shards.map(({ schema }) => schema),
          weightSql: undefined,
        });
        return shards.map(({ schema }) => {
          const info = weights.get(schema);
          const oldFactor = info?.appliedFactor || 1;
          return {
            dsn,
            schema,
            oldFactor,
            newFactor: round(
              factorSpec.op === ""
                ? factorSpec.num
                : factorSpec.op === "+"
                  ? oldFactor + factorSpec.num
                  : factorSpec.op === "-"
                    ? Math.max(oldFactor - factorSpec.num, 1)
                    : factorSpec.op === "*"
                      ? oldFactor * factorSpec.num
                      : oldFactor,
              3,
            ),
          };
        });
      },
    ),
  );

  print(
    table(
      [
        ["DSN", "Microshard", "Current Factor", "New Factor"],
        ...plan.map(({ dsn, schema, oldFactor, newFactor }) => [
          dsnShort(dsn),
          schema,
          oldFactor,
          newFactor,
        ]),
      ],
      {
        drawHorizontalLine: (i, rowCount) =>
          i === 0 || i === 1 || i === rowCount,
      },
    ).trim(),
  );
  const validResponse = "apply";
  const response = await prompts({
    type: "text",
    name: "value",
    message: `Type "${validResponse}" to apply the changes in weight factors:`,
    validate: (value: string) =>
      value !== validResponse
        ? `Enter "${validResponse}" to confirm, ^C to skip`
        : true,
  });
  if (response.value !== validResponse) {
    return false;
  }

  for (const { dsn, schema, newFactor } of plan) {
    print(`- ${dsnShort(dsn)}, ${schema}: setting factor=${newFactor}`);
    await factor({ dsn, schema, factor: newFactor });
  }

  return true;
}