@clickup/pg-microsharding/dist/internal/getSchemaWeights.js

Microshards support for PostgreSQL

"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.WEIGHT_COL_NAME_DEFAULT = void 0;
exports.getSchemaWeights = getSchemaWeights;
const mapValues_1 = __importDefault(require("lodash/mapValues"));
const omit_1 = __importDefault(require("lodash/omit"));
const names_1 = require("./names");
const quote_1 = require("./quote");
const runSql_1 = require("./runSql");
const unindent_1 = require("./unindent");
const FUNC_NAME_WEIGHT = "microsharding_schema_weights_weight_";
const FUNC_NAME_STATS = "microsharding_schema_weights_stats_";
exports.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 = (0, unindent_1.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.
 */
async function getSchemaWeights({ dsn, schemas, weightSql, verbose, }) {
    weightSql || (weightSql = WEIGHT_SQL_DEFAULT);
    const statsSql = (0, unindent_1.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 = (0, quote_1.join)([
        (0, unindent_1.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;
        $$
      `),
        (0, unindent_1.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;
        $$
      `),
        (0, quote_1.sql) `
        SELECT
          schema_name,
          obj_description(schema_name::regnamespace, 'pg_namespace'),
          pg_temp.${(0, quote_1.ident)(FUNC_NAME_WEIGHT)}(schema_name),
          pg_temp.${(0, quote_1.ident)(FUNC_NAME_STATS)}(schema_name)
        FROM json_array_elements_text(${JSON.stringify(schemas)}) AS t(schema_name)
      `,
    ], ";\n");
    const resRows = await (0, runSql_1.runSql)(dsn, query);
    return new Map(resRows.map(([schema, schemaComment, weightJson, statsJson]) => {
        var _a;
        const weightRow = weightJson
            ? JSON.parse(weightJson)
            : {};
        const statsRow = 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 = (_a = (0, names_1.shardFactorFromComment)(schemaComment)) !== null && _a !== void 0 ? _a : 1;
        return [
            schema,
            {
                weightColName: !weightColName || weightColName.startsWith("?")
                    ? exports.WEIGHT_COL_NAME_DEFAULT
                    : weightColName,
                weight: weight * appliedFactor,
                appliedFactor,
                unit,
                other: (0, mapValues_1.default)(weightColName ? (0, omit_1.default)(row, weightColName) : row, String),
            },
        ];
    }));
}
//# sourceMappingURL=getSchemaWeights.js.map