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

Microshards support for PostgreSQL

import type { ChildProcessWithoutNullStreams } from "child_process";
import { spawn } from "child_process";
import chalk from "chalk";
import compact from "lodash/compact";
import { lastLines } from "./lastLines";
import { log, progress, stderr, stdlog, stdout } from "./logging";
import { pluralize } from "./pluralize";

const MAX_INPUT_LINES = 2;
const MAX_PROGRESS_LINES = 3;

/**
 * Runs a shell command passing it an optional input as stdin. Throws on errors.
 * Returns the lines of stdout. If onOutArg is absent, then stderr of the
 * process to piped to logging's stderr. Otherwise, onOutArg() is called with
 * the full output (combined stdout + stderr) on every data chunk and also every
 * ~1 second if there is no activity (that allows the caller to e.g. prepend
 * ticking timestamps to the output lines to emphasize that the process is still
 * alive). If onOutArg="progress", then last MAX_PROGRESS_LINES lines of the
 * output are progress-printed (with ticking timestamps prepended) and then
 * erased in the end.
 */
export async function runShell(
  cmd: string,
  input: string | null,
  comment?: string,
  showProgress?: "progress.unlogged" | "progress",
  showEnv?: string[],
): Promise<string> {
  let inputStr = null;
  if (input) {
    const inputLines = compact(input.split("\n"));
    const inputLinesLess = inputLines.slice(0, MAX_INPUT_LINES);
    const inputLinesMore = inputLines.length - inputLinesLess.length;
    inputStr = compact([
      ...inputLinesLess,
      inputLinesMore > 0 && `...and ${pluralize(inputLinesMore, "more line")}`,
    ]).join("\n");
  }

  if (comment) {
    log.shellCmd({
      comment,
      cmd,
      input: inputStr,
      env: Object.fromEntries(
        compact(showEnv?.map((k) => process.env[k] && [k, process.env[k]!])),
      ),
    });
  }

  const onOut = showProgress
    ? (out: string) =>
        (showProgress === "progress" ? progress : progress.unlogged)(
          chalk
            .gray(
              lastLines(out, MAX_PROGRESS_LINES, (stdout.columns || 80) - 32),
            )
            .replace(/^/gm, "> "),
        )
    : undefined;

  let timeout: NodeJS.Timeout | undefined;
  let proc: ChildProcessWithoutNullStreams | undefined;

  const sigIntHandler = (): unknown =>
    proc?.pid && process.kill(proc.pid, "SIGINT");
  process.on("SIGINT", sigIntHandler);
  stdout.setMaxListeners(stdout.getMaxListeners() + 1);
  stderr.setMaxListeners(stderr.getMaxListeners() + 1);
  process.setMaxListeners(process.getMaxListeners() + 1);

  try {
    return await new Promise((resolve, reject) => {
      proc = spawn("/bin/bash", ["-o", "pipefail", "-c", cmd], {
        shell: false,
        env: {
          ...process.env,
          PGOPTIONS: compact([
            "--client-min-messages=warning",
            process.env["PGOPTIONS"],
          ]).join(" "),
        },
      });

      if (input) {
        proc.stdin.write(input);
        proc.stdin.end();
      }

      if (!onOut) {
        proc.stderr.pipe(stderr, { end: false });
      }

      let stdoutStr = "";
      let stderrStr = "";
      let outStr = "";

      const onOutWithTimer = (): void => {
        clearTimeout(timeout);
        onOut?.(outStr);
        timeout = setTimeout(onOutWithTimer, 1000);
      };

      proc.stdout.on("data", (data) => {
        const str = data.toString();
        stdoutStr += str;
        outStr += str;
        onOutWithTimer();
      });

      proc.stderr.on("data", (data) => {
        const str = data.toString();
        stdlog.write(str);
        stderrStr += str;
        outStr += str;
        onOutWithTimer();
      });

      proc.on("error", (err) => {
        onOutWithTimer();
        reject(err);
      });

      proc.on("close", (code, signal) => {
        onOutWithTimer();
        if (code !== 0) {
          reject(
            compact([
              `Process exited with ${signal ? `signal ${signal}` : `code ${code}`}`,
              onOut && stderrStr,
            ]).join("\n"),
          );
        } else {
          resolve(stdoutStr);
        }
      });
    });
  } catch (e: unknown) {
    if (!comment) {
      log.shellCmd({
        comment: "Error running the command",
        cmd,
        input: inputStr,
        isError: true,
      });
    }

    throw e;
  } finally {
    clearTimeout(timeout);
    proc?.stderr.unpipe(stderr);
    process.setMaxListeners(process.getMaxListeners() - 1);
    stderr.setMaxListeners(stderr.getMaxListeners() - 1);
    stdout.setMaxListeners(stdout.getMaxListeners() - 1);
    process.removeListener("SIGINT", sigIntHandler);
    onOut && progress.clear();
  }
}