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

Microshards support for PostgreSQL

import { createWriteStream, mkdirSync } from "fs";
import { PassThrough, Transform } from "stream";
import { stripVTControlCharacters } from "util";
import chalk from "chalk";
import logUpdate from "log-update";
import { dateToStr } from "./dateToStr";
import { prependDate } from "./prependDate";
import { shellQuote } from "./quote";
import { stripPassword } from "./stripPassword";

/** The directory to write the data and log files to. */
export const TMP_DIR = `/tmp/pg-microsharding.${process.getuid!()}`;
mkdirSync(TMP_DIR, { recursive: true, mode: 0o700 });

/**
 * A helper class to easier mock output in unit tests.
 */
class StdWritable extends PassThrough {
  constructor(private std: NodeJS.WriteStream) {
    super();
    this.pipe(this.std);

    if (!this.columns && parseInt(process.env["COLUMNS"] ?? "")) {
      this.columns = parseInt(process.env["COLUMNS"]!);
    }

    if (!this.rows && parseInt(process.env["ROWS"] ?? "")) {
      this.rows = parseInt(process.env["ROWS"]!);
    }
  }

  get rows(): number {
    return this.std.rows;
  }

  get columns(): number {
    return this.std.columns;
  }

  set rows(rows: number) {
    this.std.rows = rows;
  }

  set columns(columns: number) {
    this.std.columns = columns;
  }
}

/**
 * A helper class to write messages to a log file.
 */
class FileLogWritable extends Transform {
  constructor() {
    super({
      transform: (chunk, _encoding, callback) =>
        callback(null, stripVTControlCharacters(String(chunk))),
    });
    // Swallow the data, until pipeToNewLogFile() is called.
    this.pipe(
      new FileLogWritable.Writable({ write: (_, __, callback) => callback() }),
    );
  }

  pipeToNewLogFile(key: string): () => Promise<void> {
    const fileName = `${TMP_DIR}/${dateToStr(new Date())}.${key}.log`;
    const file = createWriteStream(fileName, { flags: "a" });
    this.pipe(file);
    return async () =>
      new Promise((resolve) => {
        this.unpipe(file);
        file.end(resolve);
      });
  }
}

/** Use instead of process.stdout. */
export const stdout = new StdWritable(process.stdout);

/** Use instead of process.stderr. */
export const stderr = new StdWritable(process.stderr);

/** The messages are also logged to this stream. By default, they are swallowed.
 * The caller may pipe it to a file if needed.*/
export const stdlog = new FileLogWritable();

// A log-update engine. We use stderr, to not pollute the output when --json
// flag is used for instance.
const progressObj = logUpdate.create(stderr, { showCursor: true });

// Last unique lines written with progress().
const progressLoggedLines: Set<string> = new Set();

/**
 * Logs erasable (or persistent) progress message prepended with the timestamp.
 * The unique lines of the message are also written to the file log.
 */
export function progress(...lines: string[]): void {
  progress.unlogged(...lines);

  const linesToLog = lines.filter(
    (line) => !progressLoggedLines.has(line) && line.trim(),
  );
  if (linesToLog.length > 0) {
    stdlog.write(stripPassword(prependDate(linesToLog)) + "\n");
  }

  progressLoggedLines.clear();
  for (const line of lines) {
    progressLoggedLines.add(line);
  }
}

/**
 * Same as progress(), but does not write to the file log.
 */
progress.unlogged = (...lines: string[]): void => {
  progressObj(stripPassword(prependDate(lines)));
};

/**
 * Clears the progress.
 */
progress.clear = () => {
  progressObj.clear();
  progressLoggedLines.clear();
};

/**
 * Persists the progress on screen.
 */
progress.done = () => {
  progressObj.done();
  progressLoggedLines.clear();
};

/**
 * Clears the progress and prints the raw message to console and to log file.
 */
export function print(msg: string): void {
  progress.done();
  const str = stripPassword(msg) + "\n";
  stdout.write(str);
  stdlog.write(str);
}

/**
 * Same as print(), but uses bold font.
 */
print.section = (msg: string): void => {
  print(chalk.bold(msg));
};

/**
 * Same as print(), but prepends the current date and time. The message is also
 * written to the log file.
 */
export function log(...lines: string[]): void {
  print(prependDate(lines));
}

/**
 * Same as log() (i.e. with date/time prepended), but in red.
 */
log.warning = (...lines: string[]): void => {
  print(chalk.yellow(prependDate(lines)));
};

/**
 * Same as log() (i.e. with date/time prepended), but in red.
 */
log.error = (...lines: string[]): void => {
  print(chalk.red(prependDate(lines)));
};

/**
 * Logs details about the upcoming runShell() execution.
 */
log.shellCmd = ({
  comment,
  cmd,
  input,
  isError = false,
  env = {},
}: {
  comment: string;
  cmd: string;
  input: string | null;
  isError?: boolean;
  env?: Record<string, string>;
}): void => {
  if (!comment.match(/[!?.:]$/)) {
    comment += "...";
  }

  log(chalk[isError ? "red" : "greenBright"](chalk.bold(comment)));

  for (const [k, v] of Object.entries(env)) {
    log(chalk.gray(`$ export ${k}=${shellQuote(v)}`));
  }

  if (cmd) {
    log(chalk.gray(`$ ${cmd}`));
  }

  input = input?.trim() ?? null;
  if (input) {
    log(chalk.gray(input.trim()));
  }
};