nixos-rebuild-ng: refactor part of __init__.py code in services.py (#420352)

This commit is contained in:
Thiago Kenji Okada 2025-06-27 10:55:40 +01:00 committed by GitHub
commit 1aad15088d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 362 additions and 229 deletions

View File

@ -3,17 +3,15 @@ import json
import logging
import os
import sys
from pathlib import Path
from subprocess import CalledProcessError, run
from typing import Final, assert_never
from . import nix, tmpdir
from . import nix
from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
from .models import Action, BuildAttr, Flake, ImageVariants, NixOSRebuildError, Profile
from .process import Remote, cleanup_ssh
from .utils import Args, LogFormatter, tabulate
NIXOS_REBUILD_ATTR: Final = "config.system.build.nixos-rebuild"
from .models import Action, BuildAttr, Flake, Profile
from .process import Remote
from .services import build_and_activate_system, reexec
from .utils import LogFormatter, tabulate
logger: Final = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
@ -271,72 +269,6 @@ def parse_args(
return args, args_groups
def reexec(
argv: list[str],
args: argparse.Namespace,
build_flags: Args,
flake_build_flags: Args,
) -> None:
drv = None
try:
# Parsing the args here but ignore ask_sudo_password since it is not
# needed and we would end up asking sudo password twice
if flake := Flake.from_arg(args.flake, Remote.from_arg(args.target_host, None)):
drv = nix.build_flake(
NIXOS_REBUILD_ATTR,
flake,
flake_build_flags | {"no_link": True},
)
else:
build_attr = BuildAttr.from_arg(args.attr, args.file)
drv = nix.build(
NIXOS_REBUILD_ATTR,
build_attr,
build_flags | {"no_out_link": True},
)
except CalledProcessError:
logger.warning(
"could not build a newer version of nixos-rebuild, using current version",
exc_info=logger.isEnabledFor(logging.DEBUG),
)
if drv:
new = drv / f"bin/{EXECUTABLE}"
current = Path(argv[0])
if new != current:
logger.debug(
"detected newer version of script, re-exec'ing, current=%s, new=%s",
current,
new,
)
# Manually call clean-up functions since os.execve() will replace
# the process immediately
cleanup_ssh()
tmpdir.TMPDIR.cleanup()
try:
os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
except Exception:
# Possible errors that we can have here:
# - Missing the binary
# - Exec format error (e.g.: another OS/CPU arch)
logger.warning(
"could not re-exec in a newer version of nixos-rebuild, "
"using current version",
exc_info=logger.isEnabledFor(logging.DEBUG),
)
# We already run clean-up, let's re-exec in the current version
# to avoid issues
os.execve(current, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
def validate_image_variant(image_variant: str, variants: ImageVariants) -> None:
if image_variant not in variants:
raise NixOSRebuildError(
"please specify one of the following supported image variants via "
"--image-variant:\n" + "\n".join(f"- {v}" for v in variants)
)
def execute(argv: list[str]) -> None:
args, args_groups = parse_args(argv)
@ -395,147 +327,20 @@ def execute(argv: list[str]) -> None:
| Action.BUILD_VM
| Action.BUILD_VM_WITH_BOOTLOADER
):
logger.info("building the system configuration...")
dry_run = action == Action.DRY_BUILD
no_link = action in (Action.SWITCH, Action.BOOT)
rollback = bool(args.rollback)
match action:
case Action.BUILD_IMAGE if flake:
variants = nix.get_build_image_variants_flake(
flake,
eval_flags=flake_common_flags,
)
validate_image_variant(args.image_variant, variants)
attr = f"config.system.build.images.{args.image_variant}"
case Action.BUILD_IMAGE:
variants = nix.get_build_image_variants(
build_attr,
instantiate_flags=common_flags,
)
validate_image_variant(args.image_variant, variants)
attr = f"config.system.build.images.{args.image_variant}"
case Action.BUILD_VM:
attr = "config.system.build.vm"
case Action.BUILD_VM_WITH_BOOTLOADER:
attr = "config.system.build.vmWithBootLoader"
case _:
attr = "config.system.build.toplevel"
match (action, rollback, build_host, flake):
case (Action.SWITCH | Action.BOOT, True, _, _):
path_to_config = nix.rollback(profile, target_host, sudo=args.sudo)
case (Action.TEST | Action.BUILD, True, _, _):
maybe_path_to_config = nix.rollback_temporary_profile(
profile,
target_host,
sudo=args.sudo,
)
if maybe_path_to_config: # kinda silly but this makes mypy happy
path_to_config = maybe_path_to_config
else:
raise NixOSRebuildError("could not find previous generation")
case (_, True, _, _):
raise NixOSRebuildError(
f"--rollback is incompatible with '{action}'"
)
case (_, False, Remote(_), Flake(_)):
path_to_config = nix.build_remote_flake(
attr,
flake,
build_host,
eval_flags=flake_common_flags,
flake_build_flags=flake_build_flags
| {"no_link": no_link, "dry_run": dry_run},
copy_flags=copy_flags,
)
case (_, False, None, Flake(_)):
path_to_config = nix.build_flake(
attr,
flake,
flake_build_flags=flake_build_flags
| {"no_link": no_link, "dry_run": dry_run},
)
case (_, False, Remote(_), None):
path_to_config = nix.build_remote(
attr,
build_attr,
build_host,
realise_flags=common_flags,
instantiate_flags=build_flags,
copy_flags=copy_flags,
)
case (_, False, None, None):
path_to_config = nix.build(
attr,
build_attr,
build_flags=build_flags
| {"no_out_link": no_link, "dry_run": dry_run},
)
case never:
# should never happen, but mypy is not smart enough to
# handle this with assert_never
# https://github.com/python/mypy/issues/16650
# https://github.com/python/mypy/issues/16722
raise AssertionError(
f"expected code to be unreachable, but got: {never}"
)
if not rollback:
nix.copy_closure(
path_to_config,
to_host=target_host,
from_host=build_host,
copy_flags=copy_flags,
)
if action in (Action.SWITCH, Action.BOOT):
nix.set_profile(
profile,
path_to_config,
target_host=target_host,
sudo=args.sudo,
)
# Print only the result to stdout to make it easier to script
def print_result(msg: str, result: str | Path) -> None:
print(msg, end=" ", file=sys.stderr, flush=True)
print(result, flush=True)
match action:
case Action.SWITCH | Action.BOOT | Action.TEST | Action.DRY_ACTIVATE:
nix.switch_to_configuration(
path_to_config,
action,
target_host=target_host,
sudo=args.sudo,
specialisation=args.specialisation,
install_bootloader=args.install_bootloader,
)
print_result("Done. The new configuration is", path_to_config)
case Action.BUILD:
print_result("Done. The new configuration is", path_to_config)
case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
# If you get `not-found`, please open an issue
vm_path = next(path_to_config.glob("bin/run-*-vm"), "not-found")
print_result(
"Done. The virtual machine can be started by running", vm_path
)
case Action.BUILD_IMAGE:
if flake:
image_name = nix.get_build_image_name_flake(
flake,
args.image_variant,
eval_flags=flake_common_flags,
)
else:
image_name = nix.get_build_image_name(
build_attr,
args.image_variant,
instantiate_flags=flake_common_flags,
)
disk_path = path_to_config / image_name
print_result("Done. The disk image can be found in", disk_path)
build_and_activate_system(
action=action,
args=args,
build_host=build_host,
target_host=target_host,
profile=profile,
flake=flake,
build_attr=build_attr,
build_flags=build_flags,
common_flags=common_flags,
copy_flags=copy_flags,
flake_build_flags=flake_build_flags,
flake_common_flags=flake_common_flags,
)
case Action.EDIT:
nix.edit(flake, flake_build_flags)
@ -561,9 +366,9 @@ def execute(argv: list[str]) -> None:
case Action.REPL:
if flake:
nix.repl_flake("toplevel", flake, flake_build_flags)
nix.repl_flake(flake, flake_build_flags)
else:
nix.repl("system", build_attr, build_flags)
nix.repl(build_attr, build_flags)
case _:
assert_never(action)

View File

@ -154,7 +154,7 @@ class Flake:
return cls(path, nixos_attr)
@classmethod
def from_arg(cls, flake_arg: Any, target_host: Remote | None) -> Self | None:
def from_arg(cls, flake_arg: Any, target_host: Remote | None) -> Self | None: # noqa: ANN401
match flake_arg:
case str(s):
return cls.parse(s, target_host)

View File

@ -545,14 +545,14 @@ def list_generations(profile: Profile) -> list[GenerationJson]:
)
def repl(attr: str, build_attr: BuildAttr, nix_flags: Args | None = None) -> None:
def repl(build_attr: BuildAttr, nix_flags: Args | None = None) -> None:
run_args = ["nix", "repl", "--file", build_attr.path]
if build_attr.attr:
run_args.append(build_attr.attr)
run_wrapper([*run_args, *dict_to_flags(nix_flags)])
def repl_flake(attr: str, flake: Flake, flake_flags: Args | None = None) -> None:
def repl_flake(flake: Flake, flake_flags: Args | None = None) -> None:
expr = Template(
files(__package__).joinpath(FLAKE_REPL_TEMPLATE).read_text()
).substitute(

View File

@ -0,0 +1,317 @@
import argparse
import logging
import os
import sys
from pathlib import Path
from subprocess import CalledProcessError
from typing import Final
from . import nix, tmpdir
from .constants import EXECUTABLE
from .models import Action, BuildAttr, Flake, ImageVariants, NixOSRebuildError, Profile
from .process import Remote, cleanup_ssh
from .utils import Args
NIXOS_REBUILD_ATTR: Final = "config.system.build.nixos-rebuild"
logger: Final = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
def reexec(
argv: list[str],
args: argparse.Namespace,
build_flags: Args,
flake_build_flags: Args,
) -> None:
drv = None
try:
# Parsing the args here but ignore ask_sudo_password since it is not
# needed and we would end up asking sudo password twice
if flake := Flake.from_arg(args.flake, Remote.from_arg(args.target_host, None)):
drv = nix.build_flake(
NIXOS_REBUILD_ATTR,
flake,
flake_build_flags | {"no_link": True},
)
else:
build_attr = BuildAttr.from_arg(args.attr, args.file)
drv = nix.build(
NIXOS_REBUILD_ATTR,
build_attr,
build_flags | {"no_out_link": True},
)
except CalledProcessError:
logger.warning(
"could not build a newer version of nixos-rebuild, using current version",
exc_info=logger.isEnabledFor(logging.DEBUG),
)
if drv:
new = drv / f"bin/{EXECUTABLE}"
current = Path(argv[0])
if new != current:
logger.debug(
"detected newer version of script, re-exec'ing, current=%s, new=%s",
current,
new,
)
# Manually call clean-up functions since os.execve() will replace
# the process immediately
cleanup_ssh()
tmpdir.TMPDIR.cleanup()
try:
os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
except Exception:
# Possible errors that we can have here:
# - Missing the binary
# - Exec format error (e.g.: another OS/CPU arch)
logger.warning(
"could not re-exec in a newer version of nixos-rebuild, "
"using current version",
exc_info=logger.isEnabledFor(logging.DEBUG),
)
# We already run clean-up, let's re-exec in the current version
# to avoid issues
os.execve(current, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
def _validate_image_variant(image_variant: str, variants: ImageVariants) -> None:
if image_variant not in variants:
raise NixOSRebuildError(
"please specify one of the following supported image variants via "
"--image-variant:\n" + "\n".join(f"- {v}" for v in variants)
)
def _get_system_attr(
action: Action,
args: argparse.Namespace,
flake: Flake | None,
build_attr: BuildAttr,
common_flags: Args,
flake_common_flags: Args,
) -> str:
match action:
case Action.BUILD_IMAGE if flake:
variants = nix.get_build_image_variants_flake(
flake,
eval_flags=flake_common_flags,
)
_validate_image_variant(args.image_variant, variants)
attr = f"config.system.build.images.{args.image_variant}"
case Action.BUILD_IMAGE:
variants = nix.get_build_image_variants(
build_attr,
instantiate_flags=common_flags,
)
_validate_image_variant(args.image_variant, variants)
attr = f"config.system.build.images.{args.image_variant}"
case Action.BUILD_VM:
attr = "config.system.build.vm"
case Action.BUILD_VM_WITH_BOOTLOADER:
attr = "config.system.build.vmWithBootLoader"
case _:
attr = "config.system.build.toplevel"
return attr
def _build_system(
attr: str,
action: Action,
args: argparse.Namespace,
build_host: Remote | None,
target_host: Remote | None,
profile: Profile,
flake: Flake | None,
build_attr: BuildAttr,
build_flags: Args,
common_flags: Args,
copy_flags: Args,
flake_build_flags: Args,
flake_common_flags: Args,
) -> Path:
dry_run = action == Action.DRY_BUILD
no_link = action in (Action.SWITCH, Action.BOOT)
match (action, args.rollback, build_host, flake):
case (Action.SWITCH | Action.BOOT, True, _, _):
path_to_config = nix.rollback(profile, target_host, sudo=args.sudo)
case (Action.TEST | Action.BUILD, True, _, _):
maybe_path_to_config = nix.rollback_temporary_profile(
profile,
target_host,
sudo=args.sudo,
)
if maybe_path_to_config: # kinda silly but this makes mypy happy
path_to_config = maybe_path_to_config
else:
raise NixOSRebuildError("could not find previous generation")
case (_, True, _, _):
raise NixOSRebuildError(f"--rollback is incompatible with '{action}'")
case (_, False, Remote(_), Flake(_)):
path_to_config = nix.build_remote_flake(
attr,
flake,
build_host,
eval_flags=flake_common_flags,
flake_build_flags=flake_build_flags
| {"no_link": no_link, "dry_run": dry_run},
copy_flags=copy_flags,
)
case (_, False, None, Flake(_)):
path_to_config = nix.build_flake(
attr,
flake,
flake_build_flags=flake_build_flags
| {"no_link": no_link, "dry_run": dry_run},
)
case (_, False, Remote(_), None):
path_to_config = nix.build_remote(
attr,
build_attr,
build_host,
realise_flags=common_flags,
instantiate_flags=build_flags,
copy_flags=copy_flags,
)
case (_, False, None, None):
path_to_config = nix.build(
attr,
build_attr,
build_flags=build_flags | {"no_out_link": no_link, "dry_run": dry_run},
)
case never:
# should never happen, but mypy is not smart enough to
# handle this with assert_never
# https://github.com/python/mypy/issues/16650
# https://github.com/python/mypy/issues/16722
raise AssertionError(f"expected code to be unreachable, but got: {never}")
if not args.rollback:
nix.copy_closure(
path_to_config,
to_host=target_host,
from_host=build_host,
copy_flags=copy_flags,
)
return path_to_config
def _activate_system(
path_to_config: Path,
action: Action,
args: argparse.Namespace,
target_host: Remote | None,
profile: Profile,
flake: Flake | None,
build_attr: BuildAttr,
flake_common_flags: Args,
common_flags: Args,
) -> None:
# Print only the result to stdout to make it easier to script
def print_result(msg: str, result: str | Path) -> None:
print(msg, end=" ", file=sys.stderr, flush=True)
print(result, flush=True)
match action:
case Action.SWITCH | Action.BOOT if not args.rollback:
nix.set_profile(
profile,
path_to_config,
target_host=target_host,
sudo=args.sudo,
)
nix.switch_to_configuration(
path_to_config,
action,
target_host=target_host,
sudo=args.sudo,
specialisation=args.specialisation,
install_bootloader=args.install_bootloader,
)
print_result("Done. The new configuration is", path_to_config)
case Action.SWITCH | Action.BOOT | Action.TEST | Action.DRY_ACTIVATE:
nix.switch_to_configuration(
path_to_config,
action,
target_host=target_host,
sudo=args.sudo,
specialisation=args.specialisation,
install_bootloader=args.install_bootloader,
)
print_result("Done. The new configuration is", path_to_config)
case Action.BUILD:
print_result("Done. The new configuration is", path_to_config)
case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
# If you get `not-found`, please open an issue
vm_path = next(path_to_config.glob("bin/run-*-vm"), "not-found")
print_result("Done. The virtual machine can be started by running", vm_path)
case Action.BUILD_IMAGE:
if flake:
image_name = nix.get_build_image_name_flake(
flake,
args.image_variant,
eval_flags=flake_common_flags,
)
else:
image_name = nix.get_build_image_name(
build_attr,
args.image_variant,
instantiate_flags=common_flags,
)
disk_path = path_to_config / image_name
print_result("Done. The disk image can be found in", disk_path)
def build_and_activate_system(
action: Action,
args: argparse.Namespace,
build_host: Remote | None,
target_host: Remote | None,
profile: Profile,
flake: Flake | None,
build_attr: BuildAttr,
build_flags: Args,
common_flags: Args,
copy_flags: Args,
flake_build_flags: Args,
flake_common_flags: Args,
) -> None:
logger.info("building the system configuration...")
attr = _get_system_attr(
action=action,
args=args,
flake=flake,
build_attr=build_attr,
common_flags=common_flags,
flake_common_flags=flake_common_flags,
)
path_to_config = _build_system(
attr,
action=action,
args=args,
build_host=build_host,
target_host=target_host,
profile=profile,
flake=flake,
build_attr=build_attr,
build_flags=build_flags,
common_flags=common_flags,
copy_flags=copy_flags,
flake_build_flags=flake_build_flags,
flake_common_flags=flake_common_flags,
)
_activate_system(
path_to_config,
action=action,
args=args,
target_host=target_host,
profile=profile,
flake=flake,
build_attr=build_attr,
common_flags=common_flags,
flake_common_flags=flake_common_flags,
)

View File

@ -65,10 +65,21 @@ extend-select = [
"ISC001",
"ISC002",
"ISC003",
# unused arguments
"ARG001",
"ARG002",
"ARG003",
"ARG004",
"ARG005",
]
ignore = [
[tool.ruff.lint.per-file-ignores]
"tests/*" = [
# allow unused arguments in tests (e.g., mocks)
"ARG001",
"ARG005",
# allow Any type
"ANN401"
"ANN401"
]
[tool.pytest.ini_options]

View File

@ -141,7 +141,7 @@ def test_reexec(mock_build: Mock, mock_execve: Mock, monkeypatch: MonkeyPatch) -
mock_build.assert_has_calls(
[
call(
nr.NIXOS_REBUILD_ATTR,
nr.services.NIXOS_REBUILD_ATTR,
nr.models.BuildAttr(ANY, ANY),
{"build": True, "no_out_link": True},
)
@ -185,7 +185,7 @@ def test_reexec_flake(
nr.reexec(argv, args, {"build": True}, {"flake": True})
mock_build.assert_called_once_with(
nr.NIXOS_REBUILD_ATTR,
nr.services.NIXOS_REBUILD_ATTR,
nr.models.Flake(ANY, ANY),
{"flake": True, "no_link": True},
)
@ -536,7 +536,7 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None:
)
@patch("subprocess.run", autospec=True)
@patch("uuid.uuid4", autospec=True)
@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
@pytest.mark.skipif(
not WITH_NIX_2_18,
reason="Tests internal logic based on the assumption that Nix >= 2.18",
@ -755,7 +755,7 @@ def test_execute_nix_switch_build_target_host(
clear=True,
)
@patch("subprocess.run", autospec=True)
@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
def test_execute_nix_switch_flake_target_host(
mock_cleanup_ssh: Mock,
mock_run: Mock,
@ -862,7 +862,7 @@ def test_execute_nix_switch_flake_target_host(
clear=True,
)
@patch("subprocess.run", autospec=True)
@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
def test_execute_nix_switch_flake_build_host(
mock_cleanup_ssh: Mock,
mock_run: Mock,

View File

@ -579,18 +579,18 @@ def test_list_generations(mock_get_generations: Mock, tmp_path: Path) -> None:
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_repl(mock_run: Mock) -> None:
n.repl("attr", m.BuildAttr("<nixpkgs/nixos>", None), {"nix_flag": True})
n.repl(m.BuildAttr("<nixpkgs/nixos>", None), {"nix_flag": True})
mock_run.assert_called_with(
["nix", "repl", "--file", "<nixpkgs/nixos>", "--nix-flag"]
)
n.repl("attr", m.BuildAttr(Path("file.nix"), "myAttr"))
n.repl(m.BuildAttr(Path("file.nix"), "myAttr"))
mock_run.assert_called_with(["nix", "repl", "--file", Path("file.nix"), "myAttr"])
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_repl_flake(mock_run: Mock) -> None:
n.repl_flake("attr", m.Flake(Path("flake.nix"), "myAttr"), {"nix_flag": True})
n.repl_flake(m.Flake(Path("flake.nix"), "myAttr"), {"nix_flag": True})
# See nixos-rebuild-ng.tests.repl for a better test,
# this is mostly for sanity check
assert mock_run.call_count == 1