yadunut/clan-core
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

clan-cli: Add -L option to nixos-rebuild switch to still have build output, simplify logging code

Browse Source
This commit is contained in:
Qubasa
2024-11-28 10:42:43 +01:00
parent 00064ee98b
commit 55e080a89f
6 changed files with 84 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
pkgs/clan-cli/clan_cli/cmd.py
View File

@@ -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,

13
pkgs/clan-cli/clan_cli/colors.py
View File

@@ -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():

45
pkgs/clan-cli/clan_cli/custom_logger.py
View File

@@ -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"

11
pkgs/clan-cli/clan_cli/machines/update.py
View File

@@ -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)
)

24
pkgs/clan-cli/clan_cli/ssh/host.py
View File

@@ -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

1
pkgs/clan-cli/clan_cli/vms/run.py
View File

@@ -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,