@clickup/pg-microsharding/src/internal/getSchemaWeights.ts

Microshards support for PostgreSQL

import mapValues from "lodash/mapValues";
import omit from "lodash/omit";
import { shardFactorFromComment } from "./names";
import { ident, join, sql } from "./quote";
import { runSql } from "./runSql";

const FUNC_NAME_WEIGHT = "microsharding_schema_weights_weight_";
const FUNC_NAME_STATS = "microsharding_schema_weights_stats_";

export const WEIGHT_COL_NAME_DEFAULT = "Weight";

const WEIGHT_SQL_DEFAULT = `
  SELECT round(sum(pg_total_relation_size(quote_ident(table_schema)||'.'||quote_ident(table_name))) / 1024 / 1024) || ' MB' AS "Size"
  FROM information_schema.tables
  LEFT JOIN pg_stat_user_tables ON schemaname = table_schema AND relname = table_name
  WHERE table_type = 'BASE TABLE' AND table_schema = current_schema()
  GROUP BY table_schema
`;

/**
 * Returns weights for each schema in the list. SQL query in weightSql should
 * return a numeric value with optional units after it; the units should be the
 * same in all responses and are ignored while sorting.
 */
export async function getSchemaWeights({
  dsn,
  schemas,
  weightSql,
  verbose,
}: {
  dsn: string;
  schemas: string[];
  weightSql: string | undefined;
  verbose?: boolean;
}): Promise<
  Map<
    string,
    {
      weightColName: string;
      weight: number;
      appliedFactor: number;
      unit: string | undefined;
      other: Record<string, string>;
    }
  >
> {
  const statsSql = `
    SELECT
      count(1) AS "Tables"
      ,sum(seq_scan)+sum(idx_scan) AS "Reads"
      ${
        verbose
          ? `,sum(seq_scan) AS seq_scan
             ,sum(idx_scan) AS idx_scan
             ,sum(vacuum_count) AS vacuum_count`
          : ""
      }
      ,sum(n_tup_ins)+sum(n_tup_upd)+sum(n_tup_del) AS "Writes"
      ${
        verbose
          ? `,sum(n_tup_ins) AS n_tup_ins
             ,sum(n_tup_upd) AS n_tup_upd
             ,sum(n_tup_del) AS n_tup_del`
          : ""
      }
    FROM information_schema.tables
    LEFT JOIN pg_stat_user_tables ON schemaname = table_schema AND relname = table_name
    WHERE table_type = 'BASE TABLE' AND table_schema = current_schema()
    GROUP BY table_schema
  `;
  const resRows = await runSql(
    dsn,
    join(
      [
        `CREATE FUNCTION pg_temp.${FUNC_NAME_WEIGHT}(schema_name text) RETURNS text LANGUAGE plpgsql
        AS $$
        BEGIN
          EXECUTE format('SET search_path TO %I', schema_name);
          RETURN (SELECT row_to_json(r) FROM (${(weightSql || WEIGHT_SQL_DEFAULT).trim()}) r);
        EXCEPTION WHEN undefined_table THEN
          RETURN NULL;
        END;
        $$`,
        `CREATE FUNCTION pg_temp.${FUNC_NAME_STATS}(schema_name text) RETURNS text LANGUAGE plpgsql
        AS $$
        BEGIN
          EXECUTE format('SET search_path TO %I', schema_name);
          RETURN (SELECT row_to_json(r) FROM (${statsSql.trim()}) r);
        END;
        $$`,
        join(
          schemas.map(
            (schema) => sql`
              SELECT
                ${schema},
                obj_description(${schema}::regnamespace, 'pg_namespace'),
                pg_temp.${ident(FUNC_NAME_WEIGHT)}(${schema}),
                pg_temp.${ident(FUNC_NAME_STATS)}(${schema})
            `,
          ),
          "\nUNION ALL\n",
        ),
      ],
      ";\n",
    ),
  );
  return new Map(
    resRows.map(([schema, schemaComment, weightJson, statsJson]) => {
      const weightRow: Record<string, unknown> = weightJson
        ? JSON.parse(weightJson)
        : {};
      const statsRow: Record<string, unknown> = statsJson
        ? JSON.parse(statsJson)
        : {};
      const row = { ...weightRow, ...statsRow };
      const weightColName = Object.keys(row)[0] || undefined;
      const weightWithUnit = weightColName
        ? String(row[weightColName] || "")
        : "";
      const [weight, unit] = weightWithUnit.match(/^(\d+)(.*)$/s)
        ? [parseFloat(RegExp.$1), RegExp.$2.trim()]
        : [0, undefined];
      const appliedFactor = shardFactorFromComment(schemaComment) ?? 1;
      return [
        schema,
        {
          weightColName:
            !weightColName || weightColName.startsWith("?")
              ? WEIGHT_COL_NAME_DEFAULT
              : weightColName,
          weight: weight * appliedFactor,
          appliedFactor,
          unit,
          other: mapValues(
            weightColName ? omit(row, weightColName) : row,
            String,
          ),
        },
      ];
    }),
  );
}