@clickup/pg-microsharding/src/actions/internal/parseMatchingShards.ts

Microshards support for PostgreSQL

import difference from "lodash/difference";
import sortBy from "lodash/sortBy";
import type { DiscoveredShard } from "../../internal/discoverShards";
import { discoverShards } from "../../internal/discoverShards";
import { dsnShort } from "../../internal/names";
import { normalizeDsns } from "../../internal/normalizeDsns";
import type { Args } from "../../internal/parseArgs";
import { actionList } from "../actionList";

/**
 * Extracts DSNs and the list of microshard numbers from the args.
 */
export async function parseMatchingShards(args: Args): Promise<{
  dsns: string[];
  matchingShards: DiscoveredShard[];
}> {
  const dsns = await normalizeDsns(args["dsns"] || args["dsn"]);
  if (dsns.length === 0) {
    throw "Please provide --dsn or --dsns, DB DSN(s)";
  }

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

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

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

  return {
    dsns,
    matchingShards,
  };
}