@clickup/pg-microsharding/src/api/weights.ts

Microshards support for PostgreSQL

import compact from "lodash/compact";
import partition from "lodash/partition";
import { getSchemaWeights } from "../internal/getSchemaWeights";
import { listActiveSchemas } from "../internal/listActiveSchemas";
import { shardNo } from "../internal/names";
import { normalizeCellNumbers } from "../internal/normalizeCellNumbers";

/**
 * Similar to listActiveSchemas(), but also returns the weight of each
 * microshard and its number.
 */
export async function weights({
  dsn,
  weightSql,
  verbose,
}: {
  dsn: string;
  weightSql: string | undefined;
  verbose?: boolean;
}): Promise<
  Array<{
    weightColName: string;
    weight: number;
    appliedFactor: number;
    unit: string | undefined;
    other: Record<string, string>;
    otherNormalized: Record<string, string>;
    schema: string;
    no: number;
  }>
> {
  const schemas = await listActiveSchemas({ dsn });

  const weights = await getSchemaWeights({ dsn, schemas, weightSql, verbose });
  const res = compact(
    [...weights.entries()].map(([schema, shard]) => {
      const no = shardNo(schema);
      return no !== null
        ? {
            ...shard,
            otherNormalized: { ...shard.other },
            schema,
            no,
          }
        : null;
    }),
  );
  for (const group of partition(res, ({ no }) => no === 0)) {
    normalizeCellNumbers(group.map(({ otherNormalized }) => otherNormalized));
  }

  return res;
}