@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";
import { unindent } from "./unindent";

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

export const WEIGHT_COL_NAME_DEFAULT = "Weight";

// Perf notes:
// - pg_relation_size() is fastest, but doesn't include TOASTs
// - pg_table_size() is 50% slower (still okay), it includes TOASTs
// - pg_total_relation_size() includes indexes and is super-slow
const WEIGHT_SQL_DEFAULT = unindent(`
  SELECT round(sum(pg_table_size(pg_class.oid)) / 1024 / 1024) || ' MB' AS "Size"
  FROM pg_class
  WHERE pg_class.relkind = 'r' AND pg_class.relnamespace = current_schema()::regnamespace
`);

/**
 * 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>;
    }
  >
> {
  weightSql ||= WEIGHT_SQL_DEFAULT;
  const statsSql = unindent(`
    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",
            ]
          : []),
      ].join(",\n  ")}
    FROM pg_stat_user_tables
    WHERE schemaname = current_schema()
  `);
  const query = join(
    [
      unindent(`
        CREATE OR REPLACE 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.trim()}) r);
        EXCEPTION WHEN undefined_table THEN
          RETURN NULL;
        END;
        $$
      `),
      unindent(`
        CREATE OR REPLACE 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;
        $$
      `),
      sql`
        SELECT
          schema_name,
          obj_description(schema_name::regnamespace, 'pg_namespace'),
          pg_temp.${ident(FUNC_NAME_WEIGHT)}(schema_name),
          pg_temp.${ident(FUNC_NAME_STATS)}(schema_name)
        FROM json_array_elements_text(${JSON.stringify(schemas)}) AS t(schema_name)
      `,
    ],
    ";\n",
  );
  const resRows = await runSql(dsn, query);
  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,
          ),
        },
      ];
    }),
  );
}