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 logging
import os import os
import sys import sys
from pathlib import Path
from subprocess import CalledProcessError, run from subprocess import CalledProcessError, run
from typing import Final, assert_never 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 .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
from .models import Action, BuildAttr, Flake, ImageVariants, NixOSRebuildError, Profile from .models import Action, BuildAttr, Flake, Profile
from .process import Remote, cleanup_ssh from .process import Remote
from .utils import Args, LogFormatter, tabulate from .services import build_and_activate_system, reexec
from .utils import LogFormatter, tabulate
NIXOS_REBUILD_ATTR: Final = "config.system.build.nixos-rebuild"
logger: Final = logging.getLogger(__name__) logger: Final = logging.getLogger(__name__)
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
@ -271,72 +269,6 @@ def parse_args(
return args, args_groups 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: def execute(argv: list[str]) -> None:
args, args_groups = parse_args(argv) args, args_groups = parse_args(argv)
@ -395,147 +327,20 @@ def execute(argv: list[str]) -> None:
| Action.BUILD_VM | Action.BUILD_VM
| Action.BUILD_VM_WITH_BOOTLOADER | Action.BUILD_VM_WITH_BOOTLOADER
): ):
logger.info("building the system configuration...") build_and_activate_system(
action=action,
dry_run = action == Action.DRY_BUILD args=args,
no_link = action in (Action.SWITCH, Action.BOOT) build_host=build_host,
rollback = bool(args.rollback) target_host=target_host,
profile=profile,
match action: flake=flake,
case Action.BUILD_IMAGE if flake: build_attr=build_attr,
variants = nix.get_build_image_variants_flake( build_flags=build_flags,
flake, common_flags=common_flags,
eval_flags=flake_common_flags, copy_flags=copy_flags,
) flake_build_flags=flake_build_flags,
validate_image_variant(args.image_variant, variants) flake_common_flags=flake_common_flags,
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)
case Action.EDIT: case Action.EDIT:
nix.edit(flake, flake_build_flags) nix.edit(flake, flake_build_flags)
@ -561,9 +366,9 @@ def execute(argv: list[str]) -> None:
case Action.REPL: case Action.REPL:
if flake: if flake:
nix.repl_flake("toplevel", flake, flake_build_flags) nix.repl_flake(flake, flake_build_flags)
else: else:
nix.repl("system", build_attr, build_flags) nix.repl(build_attr, build_flags)
case _: case _:
assert_never(action) assert_never(action)

View File

@ -154,7 +154,7 @@ class Flake:
return cls(path, nixos_attr) return cls(path, nixos_attr)
@classmethod @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: match flake_arg:
case str(s): case str(s):
return cls.parse(s, target_host) 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] run_args = ["nix", "repl", "--file", build_attr.path]
if build_attr.attr: if build_attr.attr:
run_args.append(build_attr.attr) run_args.append(build_attr.attr)
run_wrapper([*run_args, *dict_to_flags(nix_flags)]) 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( expr = Template(
files(__package__).joinpath(FLAKE_REPL_TEMPLATE).read_text() files(__package__).joinpath(FLAKE_REPL_TEMPLATE).read_text()
).substitute( ).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", "ISC001",
"ISC002", "ISC002",
"ISC003", "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 # allow Any type
"ANN401" "ANN401"
] ]
[tool.pytest.ini_options] [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( mock_build.assert_has_calls(
[ [
call( call(
nr.NIXOS_REBUILD_ATTR, nr.services.NIXOS_REBUILD_ATTR,
nr.models.BuildAttr(ANY, ANY), nr.models.BuildAttr(ANY, ANY),
{"build": True, "no_out_link": True}, {"build": True, "no_out_link": True},
) )
@ -185,7 +185,7 @@ def test_reexec_flake(
nr.reexec(argv, args, {"build": True}, {"flake": True}) nr.reexec(argv, args, {"build": True}, {"flake": True})
mock_build.assert_called_once_with( mock_build.assert_called_once_with(
nr.NIXOS_REBUILD_ATTR, nr.services.NIXOS_REBUILD_ATTR,
nr.models.Flake(ANY, ANY), nr.models.Flake(ANY, ANY),
{"flake": True, "no_link": True}, {"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("subprocess.run", autospec=True)
@patch("uuid.uuid4", 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( @pytest.mark.skipif(
not WITH_NIX_2_18, not WITH_NIX_2_18,
reason="Tests internal logic based on the assumption that 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, clear=True,
) )
@patch("subprocess.run", autospec=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( def test_execute_nix_switch_flake_target_host(
mock_cleanup_ssh: Mock, mock_cleanup_ssh: Mock,
mock_run: Mock, mock_run: Mock,
@ -862,7 +862,7 @@ def test_execute_nix_switch_flake_target_host(
clear=True, clear=True,
) )
@patch("subprocess.run", autospec=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( def test_execute_nix_switch_flake_build_host(
mock_cleanup_ssh: Mock, mock_cleanup_ssh: Mock,
mock_run: 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) @patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_repl(mock_run: Mock) -> None: 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( mock_run.assert_called_with(
["nix", "repl", "--file", "<nixpkgs/nixos>", "--nix-flag"] ["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"]) mock_run.assert_called_with(["nix", "repl", "--file", Path("file.nix"), "myAttr"])
@patch(get_qualified_name(n.run_wrapper, n), autospec=True) @patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_repl_flake(mock_run: Mock) -> None: 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, # See nixos-rebuild-ng.tests.repl for a better test,
# this is mostly for sanity check # this is mostly for sanity check
assert mock_run.call_count == 1 assert mock_run.call_count == 1