From 55e080a89f57b079be854dd5a1574e1426c2c299 Mon Sep 17 00:00:00 2001 From: Qubasa Date: Thu, 28 Nov 2024 10:42:43 +0100 Subject: [PATCH] clan-cli: Add -L option to nixos-rebuild switch to still have build output, simplify logging code --- pkgs/clan-cli/clan_cli/cmd.py | 64 +++++++++++++---------- pkgs/clan-cli/clan_cli/colors.py | 13 ----- pkgs/clan-cli/clan_cli/custom_logger.py | 45 ++++++++-------- pkgs/clan-cli/clan_cli/machines/update.py | 11 +++- pkgs/clan-cli/clan_cli/ssh/host.py | 24 +++++---- pkgs/clan-cli/clan_cli/vms/run.py | 1 - 6 files changed, 84 insertions(+), 74 deletions(-) diff --git a/pkgs/clan-cli/clan_cli/cmd.py b/pkgs/clan-cli/clan_cli/cmd.py index 81ee998a5..acca58268 100644 --- a/pkgs/clan-cli/clan_cli/cmd.py +++ b/pkgs/clan-cli/clan_cli/cmd.py @@ -1,5 +1,4 @@ import contextlib -import json import logging import math import os @@ -55,7 +54,6 @@ class MsgColor: def handle_io( process: subprocess.Popen, log: Log, - cmdlog: logging.Logger, *, prefix: str | None, input_bytes: bytes | None, @@ -64,12 +62,34 @@ def handle_io( timeout: float = math.inf, msg_color: MsgColor | None = None, ) -> tuple[str, str]: - rlist = [process.stdout, process.stderr] - wlist = [process.stdin] if input_bytes is not None else [] + rlist = [ + process.stdout, + process.stderr, + ] # rlist is a list of file descriptors to be monitored for read events + wlist = ( + [process.stdin] if input_bytes is not None else [] + ) # wlist is a list of file descriptors to be monitored for write events stdout_buf = b"" stderr_buf = b"" start = time.time() + # Function to handle file descriptors + def handle_fd(fd: IO[Any] | None, readlist: list[IO[Any]]) -> bytes: + if fd and fd in readlist: + read = os.read(fd.fileno(), 4096) + if len(read) != 0: + return read + rlist.remove(fd) + return b"" + + # Extra information passed to the logger + stdout_extra = {"command_prefix": prefix} + stderr_extra = {"command_prefix": prefix} + if msg_color and msg_color.stderr: + stdout_extra["color"] = msg_color.stderr.value + if msg_color and msg_color.stdout: + stderr_extra["color"] = msg_color.stdout.value + # Loop until no more data is available while len(rlist) != 0 or len(wlist) != 0: # Check if the command has timed out @@ -86,43 +106,35 @@ def handle_io( # Process has exited break - # Function to handle file descriptors - def handle_fd(fd: IO[Any] | None, readlist: list[IO[Any]]) -> bytes: - if fd and fd in readlist: - read = os.read(fd.fileno(), 4096) - if len(read) != 0: - return read - rlist.remove(fd) - return b"" - - extra = {"command_prefix": prefix} - if msg_color and msg_color.stderr: - extra["stderr_color"] = json.dumps(msg_color.stderr.value) - if msg_color and msg_color.stdout: - extra["stdout_color"] = json.dumps(msg_color.stdout.value) - # # Process stdout # ret = handle_fd(process.stdout, readlist) + + # If Log.STDOUT is set, log the stdout output if ret and log in [Log.STDOUT, Log.BOTH]: lines = ret.decode("utf-8", "replace").rstrip("\n").split("\n") for line in lines: - cmdlog.info(line, extra=extra) + cmdlog.info(line, extra=stdout_extra) + + # If stdout file is set, stream the stdout output if ret and stdout: stdout.write(ret) stdout.flush() + stdout_buf += ret # # Process stderr # - stdout_buf += ret ret = handle_fd(process.stderr, readlist) + # If Log.STDERR is set, log the stderr output if ret and log in [Log.STDERR, Log.BOTH]: lines = ret.decode("utf-8", "replace").rstrip("\n").split("\n") for line in lines: - cmdlog.error(line, extra=extra) + cmdlog.info(line, extra=stderr_extra) + + # If stderr file is set, stream the stderr output if ret and stderr: stderr.write(ret) stderr.flush() @@ -238,7 +250,6 @@ def run( env: dict[str, str] | None = None, cwd: Path | None = None, log: Log = Log.STDERR, - logger: logging.Logger = cmdlog, prefix: str | None = None, msg_color: MsgColor | None = None, check: bool = True, @@ -259,10 +270,10 @@ def run( else: filtered_input = input.decode("ascii", "replace") print_trace( - f"$: echo '{filtered_input}' | {indent_command(cmd)}", logger, prefix + f"$: echo '{filtered_input}' | {indent_command(cmd)}", cmdlog, prefix ) - elif logger.isEnabledFor(logging.DEBUG): - print_trace(f"$: {indent_command(cmd)}", logger, prefix) + elif cmdlog.isEnabledFor(logging.DEBUG): + print_trace(f"$: {indent_command(cmd)}", cmdlog, prefix) start = timeit.default_timer() with ExitStack() as stack: @@ -291,7 +302,6 @@ def run( log, prefix=prefix, msg_color=msg_color, - cmdlog=logger, timeout=timeout, input_bytes=input, stdout=stdout, diff --git a/pkgs/clan-cli/clan_cli/colors.py b/pkgs/clan-cli/clan_cli/colors.py index 171379c75..9d45cde6a 100644 --- a/pkgs/clan-cli/clan_cli/colors.py +++ b/pkgs/clan-cli/clan_cli/colors.py @@ -1,4 +1,3 @@ -import re from enum import Enum from typing import Any @@ -150,18 +149,6 @@ def color( return color_by_tuple(message, fg.value, bg.value) -def strip_color(message: str) -> str: - """ - Remove ANSI color/style sequences from a string. The set of all possible - ANSI sequences is large, so does not try to strip every possible one. But - does strip some outliers seen not just in text generated by this module, but - by other ANSI colorizers in the wild. Those include `\x1b[K` (aka EL or - erase to end of line) and `\x1b[m`, a terse version of the more common - `\x1b[0m`. - """ - return re.sub("\x1b\\[(K|.*?m)", "", message) - - if __name__ == "__main__": print("====ANSI Colors====") for _, value in AnsiColor.__members__.items(): diff --git a/pkgs/clan-cli/clan_cli/custom_logger.py b/pkgs/clan-cli/clan_cli/custom_logger.py index e79870092..633c5680f 100644 --- a/pkgs/clan-cli/clan_cli/custom_logger.py +++ b/pkgs/clan-cli/clan_cli/custom_logger.py @@ -1,5 +1,4 @@ import inspect -import json import logging import os import sys @@ -29,47 +28,49 @@ class PrefixFormatter(logging.Formatter): def __init__( self, trace_prints: bool = False, default_prefix: str | None = None ) -> None: + super().__init__() + self.default_prefix = default_prefix self.trace_prints = trace_prints - - super().__init__() self.hostnames: list[str] = [] - self.hostname_color_offset = 1 # first host shouldn't get aggressive red + self.hostname_color_offset = 0 def format(self, record: logging.LogRecord) -> str: filepath = _get_filepath(record) - if record.levelno == logging.DEBUG: - ansi_color_str = getattr(record, "stdout_color", None) - if ansi_color_str is not None: - ansi_color = json.loads(ansi_color_str) + # If extra["color"] is set, use that color for the message. + msg_color = getattr(record, "color", None) + if not msg_color: + if record.levelno == logging.DEBUG: + msg_color = AnsiColor.BLUE.value + elif record.levelno == logging.ERROR: + msg_color = AnsiColor.RED.value + elif record.levelno == logging.WARNING: + msg_color = AnsiColor.YELLOW.value else: - ansi_color = AnsiColor.BLUE.value - elif record.levelno == logging.ERROR: - ansi_color_str = getattr(record, "stderr_color", None) - if ansi_color_str is not None: - ansi_color = json.loads(ansi_color_str) - else: - ansi_color = AnsiColor.RED.value - elif record.levelno == logging.WARNING: - ansi_color = AnsiColor.YELLOW.value - else: - ansi_color = AnsiColor.DEFAULT.value + msg_color = AnsiColor.DEFAULT.value + # If extra["command_prefix"] is set, use that as the logging prefix. command_prefix = getattr(record, "command_prefix", self.default_prefix) + # If color is disabled, don't use color. if DISABLE_COLOR: if command_prefix: format_str = f"[{command_prefix}] %(message)s" else: format_str = "%(message)s" + + # If command_prefix is set, color the prefix with a unique color. elif command_prefix: prefix_color = self.hostname_colorcode(command_prefix) format_str = color_by_tuple(f"[{command_prefix}]", fg=prefix_color) - format_str += color_by_tuple(" %(message)s", fg=ansi_color) - else: - format_str = color_by_tuple("%(message)s", fg=ansi_color) + format_str += color_by_tuple(" %(message)s", fg=msg_color) + # If command_prefix is not set, color the message with the default level color. + else: + format_str = color_by_tuple("%(message)s", fg=msg_color) + + # Add the source file and line number if trace_prints is enabled. if self.trace_prints: format_str += f"\nSource: {filepath}:%(lineno)d::%(funcName)s\n" diff --git a/pkgs/clan-cli/clan_cli/machines/update.py b/pkgs/clan-cli/clan_cli/machines/update.py index 8e226e3c5..3d69cc586 100644 --- a/pkgs/clan-cli/clan_cli/machines/update.py +++ b/pkgs/clan-cli/clan_cli/machines/update.py @@ -133,6 +133,7 @@ def deploy_machine(machines: MachineGroup) -> None: "--option", "accept-flake-config", "true", + "-L", "--build-host", "", *machine.nix_options, @@ -154,11 +155,17 @@ def deploy_machine(machines: MachineGroup) -> None: check=False, msg_color=MsgColor(stderr=AnsiColor.DEFAULT), ) + ret = host.run( + switch_cmd, + extra_env=env, + check=False, + msg_color=MsgColor(stderr=AnsiColor.DEFAULT), + ) - # if the machine is mobile, we retry to deploy with the quirk method + # if the machine is mobile, we retry to deploy with the mobile workaround method is_mobile = machine.deployment.get("nixosMobileWorkaround", False) if is_mobile and ret.returncode != 0: - log.info("Mobile machine detected, applying quirk deployment method") + log.info("Mobile machine detected, applying workaround deployment method") ret = host.run( test_cmd, extra_env=env, msg_color=MsgColor(stderr=AnsiColor.DEFAULT) ) diff --git a/pkgs/clan-cli/clan_cli/ssh/host.py b/pkgs/clan-cli/clan_cli/ssh/host.py index 6ad4c8778..e6c42a38f 100644 --- a/pkgs/clan-cli/clan_cli/ssh/host.py +++ b/pkgs/clan-cli/clan_cli/ssh/host.py @@ -12,6 +12,7 @@ from typing import IO, Any from clan_cli.cmd import Log, MsgColor from clan_cli.cmd import run as local_run +from clan_cli.colors import AnsiColor from clan_cli.ssh.host_key import HostKeyCheck cmdlog = logging.getLogger(__name__) @@ -38,12 +39,6 @@ class Host: if not self.command_prefix: self.command_prefix = self.host - def __repr__(self) -> str: - return str(self) - - def __str__(self) -> str: - return f"{self.user}@{self.host}" + str(self.port if self.port else "") - @property def target(self) -> str: return f"{self.user or 'root'}@{self.host}" @@ -83,7 +78,6 @@ class Host: env=env, cwd=cwd, log=log, - logger=cmdlog, check=check, error_msg=error_msg, msg_color=msg_color, @@ -117,7 +111,13 @@ class Host: env.update(extra_env) displayed_cmd = " ".join(cmd) - cmdlog.info(f"$ {displayed_cmd}", extra={"command_prefix": self.command_prefix}) + cmdlog.info( + f"$ {displayed_cmd}", + extra={ + "command_prefix": self.command_prefix, + "color": AnsiColor.GREEN.value, + }, + ) return self._run( cmd, shell=shell, @@ -170,7 +170,13 @@ class Host: export_cmd = f"export {' '.join(env_vars)}; " displayed_cmd += export_cmd displayed_cmd += " ".join(cmd) - cmdlog.info(f"$ {displayed_cmd}", extra={"command_prefix": self.command_prefix}) + cmdlog.info( + f"$ {displayed_cmd}", + extra={ + "command_prefix": self.command_prefix, + "color": AnsiColor.GREEN.value, + }, + ) # Build the ssh command bash_cmd = export_cmd diff --git a/pkgs/clan-cli/clan_cli/vms/run.py b/pkgs/clan-cli/clan_cli/vms/run.py index 852b27232..cc48c68f0 100644 --- a/pkgs/clan-cli/clan_cli/vms/run.py +++ b/pkgs/clan-cli/clan_cli/vms/run.py @@ -343,7 +343,6 @@ def run_vm( future = executor.submit( handle_io, vm.process, - cmdlog=log, prefix=f"[{vm_config.machine_name}] ", stdout=sys.stdout.buffer, stderr=sys.stderr.buffer,