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

Microshards support for PostgreSQL

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

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

/**
 * Sets the weight factor for some shards.
 */
export async function actionFactor(args: Args): Promise<boolean> {
  const { matchingShards } = await parseMatchingShards(args);

  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 microshard(s)';
  }

  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;
}