@clickup/pg-microsharding/src/cli.ts

Microshards support for PostgreSQL

#!/usr/bin/env node
import chalk from "chalk";
import mapValues from "lodash/mapValues";
import pickBy from "lodash/pickBy";
import { actionAllocate } from "./actions/actionAllocate";
import { actionCleanup } from "./actions/actionCleanup";
import { actionCopy } from "./actions/actionCopy";
import { actionFactor } from "./actions/actionFactor";
import { actionInstall } from "./actions/actionInstall";
import { actionList } from "./actions/actionList";
import { actionMove } from "./actions/actionMove";
import { actionRebalance } from "./actions/actionRebalance";
import { allocate } from "./api/allocate";
import { cleanup } from "./api/cleanup";
import { factor } from "./api/factor";
import { install } from "./api/install";
import { move } from "./api/move";
import { rebalance } from "./api/rebalance";
import { weights } from "./api/weights";
import { inspectError } from "./internal/inspectError";
import { print } from "./internal/logging";
import { parseArgs } from "./internal/parseArgs";
import { shellQuote } from "./internal/quote";
import { readConfigs } from "./internal/readConfigs";
import { unindent } from "./internal/unindent";

export {
  allocate,
  cleanup,
  factor,
  install,
  move,
  rebalance,
  shellQuote,
  weights,
};

const HELP = unindent(`
  USAGE

    pg-microsharding install
      [--schema-name-fmt=SCHEMA_NAME_FMT]
      [--dsn=DSN | --dsns=DNS1,DSN2,...]

    pg-microsharding list | ls
      [--weight-sql='SELECT returning weight with optional unit']
      [--verbose]
      [--dsn=DSN | --dsns=DNS1,DSN2,...]
      [--json]

    pg-microsharding allocate
      --shard=N | --shards=N-M
      --migrate-cmd='shell command to run migrations'
      --activate={yes | no}
      [--dsn=DSN | --dsns=DNS1,DSN2,...]

    pg-microsharding factor
      --shard=N | --shards=N,M,... | --shards=DSN-PREFIX
      --factor=P|+P.Q|-P.Q|"*P.Q"
      [--dsn=DSN | --dsns=DNS1,DSN2,...]

    pg-microsharding move
      --shard=N
      --from=DSN-OR-DSN-PREFIX-OR
      --to=DSN-OR-DSN-PREFIX
      --activate-on-destination={yes | no}
      [--wait]
      [--max-replication-lag-sec=N]
      [--dsns=DNS1,DSN2,...]
      [--deactivate-sql='SQL $1 SQL']

    pg-microsharding rebalance
      --activate-on-destination={yes | no}
      [--wait]
      [--deactivate-sql='SQL $1 SQL']
      [--weight-sql='SELECT returning weight with optional unit']
      [--decommission=DSN1,DSN2,...]
      [--max-replication-lag-sec=N]
      [--parallelism=N]
      [--dsn=DSN | --dsns=DNS1,DSN2,...]

    pg-microsharding cleanup
      [--dsn=DSN | --dsns=DNS1,DSN2,...]

    pg-microsharding copy
      --schema=SCHEMA-NAME
      --from=DSN-OR-DSN-PREFIX-OR
      --to=DSN-OR-DSN-PREFIX
      [--wait]
      [--max-replication-lag-sec=N]
      [--dsns=DNS1,DSN2,...]

  ENVIRONMENT VARIABLES

    The tool receives parameters from command line option, but allows to set
    defaults for most of them using environment variables.

    Some variables are standard for psql command:

    - PGUSER: default database user
    - PGPASSWORD: default database password
    - PGHOST: default database host
    - PGPORT: default database port
    - PGDATABASE: default database name
    - PGSSLMODE: default SSL mode (e.g. "prefer")

    Custom variables:

    - DSNS or PGDSNS: default value for --dsns option, comma-separated
      list of DSNs (see below)
    - MIGRATE_CMD: default value for --migrate-cmd option
    - WEIGHT_SQL: default value for --weight-sql option
    - DEACTIVATE_SQL: default value for --deactivate-sql option
    - PARALLELISM: default value for --parallelism option
    - MAX_REPLICATION_LAG_SEC: default value for --max-replication-lag-sec
    - SCHEMA_NAME_FMT: default value for --schema-name-fmt (applies only
      to install action; if not passed, "sh%04d" is used)

  CONFIG FILE

    Unless --skip-config flag is passed, the tool also tries to find
    pg-microsharding.config.js (or .ts) file in the current and parent
    directories. If found, it treats its exports as environment variables
    (or as a function that returns environment variables), merging them
    into process.env.

  DSN AND ADDRESSING DATABASES

    Option --dsns, if required, should be a comma separated list of DSNs.
    Also, you may pass duplicated DSNs and even DSNs of replicas: the tool
    will filter them out and remain only master DSNs in the list.

    DSN format examples (parts defaults are from environment variables):

    - postgresql://user:pass@hostname/db?options (all parts are optional)
    - hostname:port/db (all parts except the hostname are optional)

    The tool also tries to find pg-microsharding.config.js (or .ts) file in
    the current and parent directories and, if found, treats its exports as
    environment variables, merging them into process.env.

  REPLICATION LAG PREVENTION

    The tool tries hard to not affect the replication lag of the destination
    node when moving or rebalancing shards. It waits until the lag drops below
    --max-replication-lag-sec seconds before running heavy operations (or
    until the user presses Shift+S to force-continue).

    Also, if you want the tool to pause explicitly and wait until the user
    presses Shift+S before activating the shard on the destination node,
    you can use the --wait option.
`);

const ACTIONS = {
  allocate: actionAllocate,
  cleanup: actionCleanup,
  copy: actionCopy,
  factor: actionFactor,
  install: actionInstall,
  list: actionList,
  ls: actionList,
  move: actionMove,
  rebalance: actionRebalance,
};

/**
 * Tool main function.
 */
export async function main(argsIn: string[]): Promise<boolean> {
  const args = parseArgs(argsIn);

  if (!args["skip-config"]) {
    for (const config of await readConfigs("pg-microsharding.config")) {
      Object.assign(
        process.env,
        mapValues(
          pickBy(
            config,
            (v) =>
              typeof v === "string" ||
              typeof v === "number" ||
              typeof v === "boolean",
          ),
          String,
        ),
      );
    }
  }

  const DSNS = process.env["PGDSNS"] || process.env["DSNS"];
  if (!args["dsns"] && !args["dsn"] && DSNS) {
    args["dsns"] = DSNS;
  }

  const MIGRATE_CMD = process.env["MIGRATE_CMD"];
  if (!args["migrate-cmd"] && MIGRATE_CMD) {
    args["migrate-cmd"] = MIGRATE_CMD;
  }

  const WEIGHT_SQL = process.env["WEIGHT_SQL"];
  if (!args["weight-sql"] && WEIGHT_SQL) {
    args["weight-sql"] = WEIGHT_SQL;
  }

  const DEACTIVATE_SQL = process.env["DEACTIVATE_SQL"];
  if (!args["deactivate-sql"] && DEACTIVATE_SQL) {
    args["deactivate-sql"] = DEACTIVATE_SQL;
  }

  const PARALLELISM = process.env["PARALLELISM"];
  if (!args["parallelism"] && PARALLELISM) {
    args["parallelism"] = PARALLELISM;
  }

  const MAX_REPLICATION_LAG_SEC = process.env["MAX_REPLICATION_LAG_SEC"];
  if (!args["max-replication-lag-sec"] && MAX_REPLICATION_LAG_SEC) {
    args["max-replication-lag-sec"] = MAX_REPLICATION_LAG_SEC;
  }

  const SCHEMA_NAME_FMT = process.env["SCHEMA_NAME_FMT"];
  if (!args["schema-name-fmt"] && SCHEMA_NAME_FMT) {
    args["schema-name-fmt"] = SCHEMA_NAME_FMT;
  }

  const action = args._[0] as keyof typeof ACTIONS;
  if (!action) {
    print(HELP);
    return true;
  }

  if (action in ACTIONS) {
    return ACTIONS[action](args);
  }

  throw `Unknown action: ${action}`;
}

/**
 * A wrapper around main() to call it from a bin script.
 */
export function cli(): void {
  main(process.argv.slice(2))
    .then((success) => process.exit(success ? 0 : 1))
    .catch((e) => {
      print(chalk.red(inspectError(e)));
      process.exit(1);
    });
}

if (require.main === module) {
  cli();
}