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

Merge pull request 'clan-cli: Add -L option to nixos-rebuild switch to still have build output, simplify logging code' (#2513) from Qubasa/clan-core:Qubasa-main into main

Browse Source 
Reviewed-on: https://git.clan.lol/clan/clan-core/pulls/2513
This commit is contained in:
Luis Hebendanz
2024-11-28 09:59:12 +00:00
parent 32332e8460 9c2c33d577
commit 781557a36a
8 changed files with 93 additions and 91 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():

39
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 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:
ansi_color_str = getattr(record, "stdout_color", None)
if ansi_color_str is not None:
ansi_color = json.loads(ansi_color_str)
else:
ansi_color = AnsiColor.BLUE.value
msg_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
msg_color = AnsiColor.RED.value
elif record.levelno == logging.WARNING:
ansi_color = AnsiColor.YELLOW.value
msg_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)
)

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

@@ -4,14 +4,14 @@ import logging
import math
import os
import shlex
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from shlex import quote
from typing import IO, Any
from clan_cli.cmd import Log, MsgColor
from clan_cli.cmd import CmdOut, 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 +38,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}"
@@ -71,7 +65,7 @@ class Host:
msg_color: MsgColor | None = None,
shell: bool = False,
timeout: float = math.inf,
) -> subprocess.CompletedProcess[str]:
) -> CmdOut:
res = local_run(
cmd,
shell=shell,
@@ -83,18 +77,12 @@ class Host:
env=env,
cwd=cwd,
log=log,
logger=cmdlog,
check=check,
error_msg=error_msg,
msg_color=msg_color,
needs_user_terminal=needs_user_terminal,
)
return subprocess.CompletedProcess(
args=res.command_list,
returncode=res.returncode,
stdout=res.stdout,
stderr=res.stderr,
)
return res
def run_local(
self,
@@ -108,7 +96,7 @@ class Host:
shell: bool = False,
needs_user_terminal: bool = False,
log: Log = Log.BOTH,
) -> subprocess.CompletedProcess[str]:
) -> CmdOut:
"""
Command to run locally for the host
"""
@@ -117,7 +105,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,
@@ -146,7 +140,7 @@ class Host:
msg_color: MsgColor | None = None,
shell: bool = False,
log: Log = Log.BOTH,
) -> subprocess.CompletedProcess[str]:
) -> CmdOut:
"""
Command to run on the host via ssh
"""
@@ -170,7 +164,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

4
pkgs/clan-cli/clan_cli/ssh/results.py
View File

@@ -1,7 +1,7 @@
import subprocess
from dataclasses import dataclass
from typing import Generic
from clan_cli.errors import CmdOut
from clan_cli.ssh import T
from clan_cli.ssh.host import Host
@@ -30,4 +30,4 @@ class HostResult(Generic[T]):
return self._result
Results = list[HostResult[subprocess.CompletedProcess[str]]]
Results = list[HostResult[CmdOut]]

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,

6
pkgs/clan-cli/tests/test_ssh_remote.py
View File

@@ -1,8 +1,6 @@
import subprocess
import pytest
from clan_cli.cmd import Log
from clan_cli.errors import ClanError
from clan_cli.errors import ClanError, CmdOut
from clan_cli.ssh.host import Host
from clan_cli.ssh.host_group import HostGroup
from clan_cli.ssh.host_key import HostKeyCheck
@@ -74,7 +72,7 @@ def test_run_exception(host_group: HostGroup) -> None:
def test_run_function_exception(host_group: HostGroup) -> None:
def some_func(h: Host) -> subprocess.CompletedProcess[str]:
def some_func(h: Host) -> CmdOut:
return h.run_local(["exit 1"], shell=True)
try: