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

Microshards support for PostgreSQL

import chalk from "chalk";
import chunk from "lodash/chunk";
import compact from "lodash/compact";
import first from "lodash/first";
import flatten from "lodash/flatten";
import mapValues from "lodash/mapValues";
import max from "lodash/max";
import range from "lodash/range";
import round from "lodash/round";
import sum from "lodash/sum";
import sumBy from "lodash/sumBy";
import { table } from "table";
import { weights } from "../api/weights";
import { chunkIntoN } from "../internal/chunkIntoN";
import { WEIGHT_COL_NAME_DEFAULT } from "../internal/getSchemaWeights";
import { indent } from "../internal/indent";
import { print, progress, stdout } from "../internal/logging";
import { dsnShort } from "../internal/names";
import { normalizeDsns } from "../internal/normalizeDsns";
import type { Args } from "../internal/parseArgs";
import { pluralize } from "../internal/pluralize";
import { promiseAllMap } from "../internal/promiseAllMap";

const DEBUG_RENDER_MUL = parseInt(process.env["DEBUG_RENDER_MUL"] ?? "0") || 0;

/**
 * Shows the list of microshards and their weights.
 */
export async function actionList(args: Args): Promise<boolean> {
  const weightSql = args["weight-sql"] || undefined;
  const verbose = !!args["verbose"];
  const json = !!args["json"];

  const dsns = await normalizeDsns(args["dsns"] || args["dsn"]);
  if (dsns.length === 0) {
    throw "Please provide --dsn or --dsns, DB DSNs to list the microshards on";
  }

  const { islandNosToDsn, islands } = await calcIslandWeights({
    dsns,
    weightSql,
    verbose,
  });

  if (DEBUG_RENDER_MUL > 0) {
    [...islands.values()].forEach((shards) =>
      shards.push(...flatten(range(DEBUG_RENDER_MUL - 1).map(() => shards))),
    );
  }

  const grandTotalWeightColName =
    first(flatten([...islands.values()]).map((s) => s.weightColName)) ??
    WEIGHT_COL_NAME_DEFAULT;
  const grandTotalWeight = round(
    sum(flatten([...islands.values()]).map((s) => s.weight)),
    2,
  );
  const grandTotalUnit = first(
    compact(flatten([...islands.values()]).map((s) => s.unit)),
  );
  const grandTotalShards = sum([...islands.values()].map((s) => s.length));

  const firstShard = first(
    first([...islands.values()].filter((s) => s.length > 0)),
  );
  const header = ["Microshard"].concat(
    firstShard
      ? [
          firstShard.weightColName,
          `${firstShard.weightColName} %`,
          ...Object.keys(firstShard.other),
        ]
      : [WEIGHT_COL_NAME_DEFAULT, `${WEIGHT_COL_NAME_DEFAULT} %`],
  );

  // Width in characters of the longest rendered table row.
  const tableWidth =
    2 + // indent
    header.length * 3 + // left borders and paddings around each cell
    1 + // last border
    (max(
      // content of the longest row
      flatten([...islands.values()]).map(
        ({ weight, appliedFactor, unit, otherNormalized, no }) =>
          // Microshard
          Math.max(header[0].length, no.toString().length) +
          // Weight or Size or ...
          Math.max(
            header[2].length,
            round(weight / appliedFactor, 5).toString().length +
              (unit ? unit.length + 1 : 0) +
              (appliedFactor !== 1 ? ` ✖${appliedFactor}`.length : 0),
          ) +
          // Weight or Size or ... %
          Math.max(header[2].length, "100%".length) +
          // Other columns
          sum(
            Object.entries(otherNormalized).map(([k, v]) =>
              Math.max(k.length, v.length),
            ),
          ),
      ),
    ) ?? 0);

  const maxTablesSideBySide = Math.trunc(
    ((stdout.columns || 80) - 10) / tableWidth,
  );

  const toPrint = {
    islands: [...islands].map(([islandNo, shards]) => {
      const dsn = islandNosToDsn.get(islandNo)!;
      const totalWeight = sumBy(shards, (shard) => shard.weight);
      const totalUnit = first(shards)?.unit;
      const totalOther = mapValues(first(shards)?.other, (_, k) => {
        const values = shards
          .map(({ other }) => parseFloat(other[k]))
          .filter((v) => !isNaN(v));
        return values.length > 0 ? sum(values).toString() : undefined;
      });
      return {
        no: islandNo,
        caption:
          `${dsnShort(dsn)} — ` + `${pluralize(shards.length, "microshard")}`,
        dsn: dsnShort(dsn),
        header,
        shards: shards.map(
          ({ no, weight, appliedFactor, unit, otherNormalized }) => [
            no.toString(),
            round(weight / appliedFactor, 5).toString() +
              (unit ? ` ${unit}` : "") +
              (appliedFactor !== 1 ? ` ✖${appliedFactor}` : ""),
            totalWeight !== 0
              ? `${((weight / totalWeight) * 100).toFixed(1)}%`
              : "-",
            ...Object.values(otherNormalized),
          ],
        ),
        total:
          shards.length > 0
            ? [
                "total",
                totalWeight + (totalUnit ? ` ${totalUnit}` : ""),
                "100%",
                ...Object.values(totalOther).map((v) => v ?? ""),
              ]
            : undefined,
      };
    }),
    total:
      `TOTAL ${grandTotalWeightColName.toUpperCase()}: ${grandTotalWeight}` +
      (grandTotalUnit ? ` ${grandTotalUnit}` : "") +
      `, ${pluralize(islands.size, "island")}, ${pluralize(grandTotalShards, "microshard")}`,
  };

  if (json) {
    print(JSON.stringify(toPrint, null, 2) + "\n");
    return true;
  }

  for (const island of toPrint.islands) {
    print.section(island.caption);

    if (island.shards.length === 0) {
      print(indent(table([["No microshards"]])));
      continue;
    }

    const tables =
      json ||
      islands.size === 1 ||
      maxTablesSideBySide === 0 ||
      grandTotalShards + islands.size * 8 < (stdout.rows || 25) ||
      island.shards.length <= 4
        ? [island.shards]
        : island.shards.length / maxTablesSideBySide < 2
          ? chunk(island.shards, 2)
          : chunkIntoN(island.shards, maxTablesSideBySide);
    print(
      glueHorizontally(
        tables.map((tbl, tblNo) =>
          indent(
            table(
              compact([
                island.header.map((cell) => chalk.bold(cell)),
                ...tbl.map((row) =>
                  row[0].trim() === "0"
                    ? row.map((cell) => chalk.yellow(cell))
                    : row,
                ),
                tblNo === 0 &&
                  island.total &&
                  island.total.map((cell) => chalk.bold(cell)),
              ]),
              {
                drawHorizontalLine: (i, rowCount) =>
                  i === 0 ||
                  i === 1 ||
                  (tblNo === 0 && i === rowCount - 1) ||
                  i === rowCount,
              },
            ),
          ),
        ),
      ),
    );
  }

  print(
    chalk.bgBlue(" " + chalk.whiteBright(toPrint.total) + " ") +
      (verbose && !weightSql ? "" : " (try --verbose)"),
  );

  return true;
}

export async function calcIslandWeights({
  dsns,
  weightSql,
  verbose,
}: {
  dsns: string[];
  weightSql: string | undefined;
  verbose?: boolean;
}): Promise<{
  islandNosToDsn: Map<number, string>;
  islands: Map<number, NonNullable<Awaited<ReturnType<typeof weights>>>>;
}> {
  try {
    const pendingPrefix = "Calculating microshard sizes on: ";
    const pendingDsns = new Set<string>();
    const islandNosToDsn = new Map([...dsns.entries()]);
    const islands = new Map(
      await promiseAllMap(
        [...islandNosToDsn.entries()],
        async ([islandNo, dsn]) => {
          try {
            pendingDsns.add(dsnShort(dsn));
            progress(pendingPrefix + [...pendingDsns].join(", "));
            const shards = await weights({
              dsn,
              weightSql,
              verbose,
            });
            return [islandNo, shards];
          } finally {
            pendingDsns.delete(dsnShort(dsn));
            if (pendingDsns.size > 0) {
              progress(pendingPrefix + [...pendingDsns].join(", "));
            }
          }
        },
      ),
    );
    return { islandNosToDsn, islands };
  } finally {
    progress.clear();
  }
}

function glueHorizontally(tables: string[]): string {
  const rowsByTable = tables.map((table) => table.split("\n"));
  const maxRows = max(rowsByTable.map((rows) => rows.length)) ?? 0;

  for (const rows of rowsByTable) {
    const maxWidth = rows[0].length;
    for (let i = 0; i < maxRows; i++) {
      rows[i] ||= "";
      rows[i] = rows[i].padEnd(maxWidth, " ");
    }
  }

  const result = range(maxRows).map(() => [] as string[]);
  for (const rows of rowsByTable) {
    for (const [i, col] of rows.entries()) {
      result[i].push(col);
    }
  }

  return result.map((lineArray) => lineArray.join("")).join("\n");
}