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

Microshards support for PostgreSQL

import { execSync } from "child_process";
import { writeFileSync } from "fs";
import compact from "lodash/compact";
import { fileSync, setGracefulCleanup } from "tmp";
import { TMP_DIR, print } from "./logging";
import { shellQuote } from "./quote";

const TMUX_SESSION = "pg-microsharding-tmux";
const ACTIVE_BG = "#000033";

// We do not want to pass the following psql envs through tmux session, because
// the data in them is already included in the DSNs in the command line.
const EXCLUDE_PG_ENV = [
  "PGUSER",
  "PGPASSWORD",
  "PGHOST",
  "PGPORT",
  "PGDATABASE",
  "PGDSNS",
];

/**
 * Returns true if this script runs inside a TMUX session.
 */
export function isInTmuxSession(): boolean {
  return !!process.env["TMUX"];
}

/**
 * If there is already an active TMUX session, reattaches to it. We need to keep
 * the strong notion of singleton for all long-running microsharding actions, so
 * that e.g. if we run `pg-microsharding rebalance` in one terminal, and then
 * run it again in another terminal, the second run will attach to the first
 * one.
 */
export async function tryReattachToTmuxSession(): Promise<void> {
  if (isInTmuxSession()) {
    return;
  }

  try {
    execSync("which tmux", { stdio: "ignore" });
  } catch {
    throw "TMUX is required to run this action; please install it";
  }

  // Try to reattach to an existing session instead of starting a new one (in
  // case we e.g. got disconnected).
  try {
    execSync(`tmux attach-session -d -t ${TMUX_SESSION} 2>/dev/null`, {
      stdio: "inherit",
    });
    process.exit(0);
  } catch {
    // No such session yet.
  }
}

/**
 * Opens M TMUX panes, in each run N commands sequentially (panes[M][N]). If a
 * single command fails, then fails the whole pane.
 *
 * In case we already run it from a TMUX session, just runs all the commands
 * sequentially.
 */
export async function runInTmux(
  panes: Array<Array<{ key: string; title: string; command: string[] }>>,
): Promise<void> {
  if (isInTmuxSession()) {
    for (const { title, command } of panes.flat()) {
      print.section(`\n##\n## ${title}\n##\n`);
      execSync(command.map((v) => shellQuote(v)).join(" "), {
        stdio: "inherit",
      });
    }

    return;
  }

  setGracefulCleanup();

  const tmuxCommands: string[] = ["set -e"];
  for (const [i, sequence] of panes.entries()) {
    const commands: string[] = [
      "set -e",
      String.raw`trap 'test $? -eq 0 && printf "\033[1;32mDONE. Press ^C to close this pane..." || printf "\033[0;31mFAILURE. Press ^C to close this pane..."; sleep 100000' EXIT`,
    ];

    for (const [i, { title, command }] of sequence.entries()) {
      const progress = `${i + 1} of ${sequence.length}`;
      commands.push(
        `printf '\\033]2;%s\\033\\' ' ${progress} | ${title} '`,
        // Pass env vars like PGOPTIONS, PGSSLCERT etc. through tmux.
        ...compact(
          Object.entries(process.env).map(
            ([k, v]) =>
              v &&
              k.match(/^PG[A-Z][A-Z_]*$/) &&
              !EXCLUDE_PG_ENV.includes(k) &&
              `export ${k}=${shellQuote(v)}`,
          ),
        ),
        command.map(shellQuote).join(" "),
        "echo",
      );
    }

    // TMUX has a limit on the command total length:
    // https://github.com/tmux/tmux/issues/254. So instead of passing the big
    // script, we put this script in a temporary file and then run it using the
    // `bash this-temp-file` plain command via TMUX. Notice that the temporary
    // files are auto-removed, but if the script is terminated cruelly with e.g.
    // -9 signal, they may linger; it's not a big deal since they are cleaned up
    // by linux tmpwatch anyway, we just don't have a better option.
    const tmpFile = fileSync({
      tmpdir: TMP_DIR,
      prefix:
        "tmux-" +
        sequence
          .map(({ key }) => key)
          .slice(0, 10)
          .join("-"),
    }).name;
    writeFileSync(tmpFile, commands.join("\n"));

    if (i === 0) {
      tmuxCommands.push(
        `tmux new-session -s ${TMUX_SESSION} -d bash ${shellQuote(tmpFile)}`,
      );
    } else {
      tmuxCommands.push(
        `tmux split-window -f bash ${shellQuote(tmpFile)}`,
        "tmux select-layout tiled",
      );
    }
  }

  tmuxCommands.push(
    'tmux bind -n "C-c" confirm-before -p "Confirm SIGINT? (y/n)" "send-keys C-c"',
    `tmux set -g window-active-style "bg=${ACTIVE_BG}"`,
    `tmux set -g pane-active-border-style "bg=${ACTIVE_BG}"`,
    "tmux set mouse on",
    "tmux set pane-border-status top",
    'tmux set pane-border-format "#T"',
    "tmux select-layout tiled",
    `tmux attach-session -d -t ${TMUX_SESSION}`,
  );

  execSync(tmuxCommands.join("\n"), { stdio: "inherit" });
}