diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py index 50c4d9ba787f..2faa0771f870 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py @@ -1,5 +1,4 @@ import argparse -import atexit import json import logging import os @@ -8,7 +7,7 @@ from pathlib import Path from subprocess import CalledProcessError, run from typing import assert_never -from . import nix +from . import nix, tmpdir from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES from .models import Action, BuildAttr, Flake, NRError, Profile from .process import Remote, cleanup_ssh @@ -20,7 +19,14 @@ logger.setLevel(logging.INFO) def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentParser]]: common_flags = argparse.ArgumentParser(add_help=False) - common_flags.add_argument("--verbose", "-v", action="count", dest="v", default=0) + common_flags.add_argument( + "--verbose", + "-v", + action="count", + dest="v", + default=0, + help="Enable verbose logging (includes nix)", + ) common_flags.add_argument("--max-jobs", "-j") common_flags.add_argument("--cores") common_flags.add_argument("--log-format") @@ -262,14 +268,15 @@ def reexec( drv = None attr = "config.system.build.nixos-rebuild" try: - # Need to set target_host=None, to avoid connecting to remote - if flake := Flake.from_arg(args.flake, None): + # 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(attr, flake, **flake_build_flags, no_link=True) else: build_attr = BuildAttr.from_arg(args.attr, args.file) drv = nix.build(attr, build_attr, **build_flags, no_out_link=True) except CalledProcessError: - logger.warning("could not find a newer version of nixos-rebuild") + logger.warning("could not build a newer version of nixos-rebuild") if drv: new = drv / f"bin/{EXECUTABLE}" @@ -280,7 +287,10 @@ def reexec( argv[0], new, ) + # Manually call clean-up functions since os.execve() will replace + # the process immediately cleanup_ssh() + tmpdir.TMPDIR.cleanup() os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"}) @@ -290,8 +300,6 @@ def execute(argv: list[str]) -> None: if not WITH_NIX_2_18: logger.warning("you're using Nix <2.18, some features will not work correctly") - atexit.register(cleanup_ssh) - common_flags = vars(args_groups["common_flags"]) common_build_flags = common_flags | vars(args_groups["common_build_flags"]) build_flags = common_build_flags | vars(args_groups["classic_build_flags"]) @@ -478,6 +486,15 @@ def main() -> None: try: execute(sys.argv) + except CalledProcessError as ex: + if logger.level == logging.DEBUG: + import traceback + + traceback.print_exc() + else: + print(str(ex), file=sys.stderr) + # Exit with the error code of the process that failed + sys.exit(ex.returncode) except (Exception, KeyboardInterrupt) as ex: if logger.level == logging.DEBUG: raise diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py index c1f2e5d5c711..c26a771c26f8 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py @@ -7,7 +7,9 @@ from pathlib import Path from string import Template from subprocess import PIPE, CalledProcessError from typing import Final +from uuid import uuid4 +from . import tmpdir from .constants import WITH_NIX_2_18 from .models import ( Action, @@ -76,24 +78,50 @@ def remote_build( instantiate_flags: dict[str, Args] | None = None, copy_flags: dict[str, Args] | None = None, ) -> Path: + # We need to use `--add-root` otherwise Nix will print this warning: + # > warning: you did not specify '--add-root'; the result might be removed + # > by the garbage collector r = run_wrapper( [ "nix-instantiate", build_attr.path, "--attr", build_attr.to_attr(attr), + "--add-root", + tmpdir.TMPDIR_PATH / uuid4().hex, *dict_to_flags(instantiate_flags or {}), ], stdout=PIPE, ) - drv = Path(r.stdout.strip()) + drv = Path(r.stdout.strip()).resolve() copy_closure(drv, to_host=build_host, from_host=None, **(copy_flags or {})) + + # Need a temporary directory in remote to use in `nix-store --add-root` r = run_wrapper( - ["nix-store", "--realise", drv, *dict_to_flags(build_flags or {})], - remote=build_host, - stdout=PIPE, + ["mktemp", "-d", "-t", "nixos-rebuild.XXXXX"], remote=build_host, stdout=PIPE ) - return Path(r.stdout.strip()) + remote_tmpdir = Path(r.stdout.strip()) + try: + r = run_wrapper( + [ + "nix-store", + "--realise", + drv, + "--add-root", + remote_tmpdir / uuid4().hex, + *dict_to_flags(build_flags or {}), + ], + remote=build_host, + stdout=PIPE, + ) + # When you use `--add-root`, `nix-store` returns the root and not the + # path inside Nix store + r = run_wrapper( + ["readlink", "-f", r.stdout.strip()], remote=build_host, stdout=PIPE + ) + return Path(r.stdout.strip()) + finally: + run_wrapper(["rm", "-rf", remote_tmpdir], remote=build_host, check=False) def remote_build_flake( diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py index 10666b47d657..54c8b71036e9 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py @@ -1,22 +1,21 @@ +import atexit import logging import os import shlex import subprocess from dataclasses import dataclass from getpass import getpass -from pathlib import Path -from tempfile import TemporaryDirectory -from typing import Self, Sequence, TypedDict, Unpack +from typing import Final, Self, Sequence, TypedDict, Unpack + +from . import tmpdir logger = logging.getLogger(__name__) -TMPDIR = TemporaryDirectory(prefix="nixos-rebuild.") -TMPDIR_PATH = Path(TMPDIR.name) -SSH_DEFAULT_OPTS = [ +SSH_DEFAULT_OPTS: Final = [ "-o", "ControlMaster=auto", "-o", - f"ControlPath={TMPDIR_PATH / "ssh-%n"}", + f"ControlPath={tmpdir.TMPDIR_PATH / "ssh-%n"}", "-o", "ControlPersist=60", ] @@ -70,13 +69,15 @@ class RunKwargs(TypedDict, total=False): def cleanup_ssh() -> None: "Close SSH ControlMaster connection." - for ctrl in TMPDIR_PATH.glob("ssh-*"): + for ctrl in tmpdir.TMPDIR_PATH.glob("ssh-*"): run_wrapper( ["ssh", "-o", f"ControlPath={ctrl}", "-O", "exit", "dummyhost"], check=False, capture_output=True, ) - TMPDIR.cleanup() + + +atexit.register(cleanup_ssh) def run_wrapper( diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py new file mode 100644 index 000000000000..521af84ce464 --- /dev/null +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py @@ -0,0 +1,6 @@ +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Final + +TMPDIR: Final = TemporaryDirectory(prefix="nixos-rebuild.") +TMPDIR_PATH: Final = Path(TMPDIR.name) diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py index 39a50ecc405a..23fd55283fce 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py @@ -82,18 +82,18 @@ def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None: nixpkgs_path.mkdir() config_path = tmp_path / "test" config_path.touch() - mock_run.side_effect = [ - # update_nixpkgs_rev - CompletedProcess([], 0, str(nixpkgs_path)), - CompletedProcess([], 0, "nixpkgs-rev"), - CompletedProcess([], 0), - # nixos_build - CompletedProcess([], 0, str(config_path)), - # set_profile - CompletedProcess([], 0), - # switch_to_configuration - CompletedProcess([], 0), - ] + + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix-instantiate": + return CompletedProcess([], 0, str(nixpkgs_path)) + elif args[0] == "git" and "rev-parse" in args: + return CompletedProcess([], 0, "nixpkgs-rev") + elif args[0] == "nix-build": + return CompletedProcess([], 0, str(config_path)) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute(["nixos-rebuild", "boot", "--no-flake", "-vvv", "--fast"]) @@ -155,14 +155,14 @@ def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None: def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None: config_path = tmp_path / "test" config_path.touch() - mock_run.side_effect = [ - # nixos_build_flake - CompletedProcess([], 0, str(config_path)), - # set_profile - CompletedProcess([], 0), - # switch_to_configuration - CompletedProcess([], 0), - ] + + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix": + return CompletedProcess([], 0, str(config_path)) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute( [ @@ -226,16 +226,14 @@ def test_execute_nix_switch_flake_target_host( ) -> None: config_path = tmp_path / "test" config_path.touch() - mock_run.side_effect = [ - # nixos_build_flake - CompletedProcess([], 0, str(config_path)), - # set_profile - CompletedProcess([], 0), - # copy_closure - CompletedProcess([], 0), - # switch_to_configuration - CompletedProcess([], 0), - ] + + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix": + return CompletedProcess([], 0, str(config_path)) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute( [ @@ -317,18 +315,16 @@ def test_execute_nix_switch_flake_build_host( ) -> None: config_path = tmp_path / "test" config_path.touch() - mock_run.side_effect = [ - # nixos_build_flake - CompletedProcess([], 0, str(config_path)), - CompletedProcess([], 0), - CompletedProcess([], 0, str(config_path)), - # set_profile - CompletedProcess([], 0), - # copy_closure - CompletedProcess([], 0), - # switch_to_configuration - CompletedProcess([], 0), - ] + + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix" and "eval" in args: + return CompletedProcess([], 0, str(config_path)) + if args[0] == "ssh" and "nix" in args: + return CompletedProcess([], 0, str(config_path)) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute( [ @@ -478,12 +474,14 @@ def test_execute_build(mock_run: Any, tmp_path: Path) -> None: def test_execute_test_flake(mock_run: Any, tmp_path: Path) -> None: config_path = tmp_path / "test" config_path.touch() - mock_run.side_effect = [ - # nixos_build_flake - CompletedProcess([], 0, str(config_path)), - # switch_to_configuration - CompletedProcess([], 0), - ] + + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix": + return CompletedProcess([], 0, str(config_path)) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute( ["nixos-rebuild", "test", "--flake", "github:user/repo#hostname", "--fast"] @@ -522,20 +520,21 @@ def test_execute_test_rollback( mock_path_exists: Any, mock_run: Any, ) -> None: - mock_run.side_effect = [ - # rollback_temporary_profile - CompletedProcess( - [], - 0, - stdout=textwrap.dedent("""\ - 2082 2024-11-07 22:58:56 - 2083 2024-11-07 22:59:41 - 2084 2024-11-07 23:54:17 (current) - """), - ), - # switch_to_configuration - CompletedProcess([], 0), - ] + def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]: + if args[0] == "nix-env": + return CompletedProcess( + [], + 0, + stdout=textwrap.dedent("""\ + 2082 2024-11-07 22:58:56 + 2083 2024-11-07 22:59:41 + 2084 2024-11-07 23:54:17 (current) + """), + ) + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_side_effect nr.execute( ["nixos-rebuild", "test", "--rollback", "--profile-name", "foo", "--fast"] diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py index 34f47029c103..c2bef1fffe97 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py @@ -1,4 +1,5 @@ import textwrap +import uuid from pathlib import Path from subprocess import PIPE, CompletedProcess from typing import Any @@ -73,14 +74,29 @@ def test_build_flake(mock_run: Any) -> None: ) -@patch( - get_qualified_name(n.run_wrapper, n), - autospec=True, - return_value=CompletedProcess([], 0, stdout=" \n/path/to/file\n "), -) -def test_remote_build(mock_run: Any, monkeypatch: Any) -> None: +@patch(get_qualified_name(n.run_wrapper, n), autospec=True) +@patch(get_qualified_name(n.uuid4, n), autospec=True) +def test_remote_build(mock_uuid4: Any, mock_run: Any, monkeypatch: Any) -> None: build_host = m.Remote("user@host", [], None) monkeypatch.setenv("NIX_SSHOPTS", "--ssh opts") + + def run_wrapper_side_effect( + args: list[str], **kwargs: Any + ) -> CompletedProcess[str]: + if args[0] == "nix-instantiate": + return CompletedProcess([], 0, stdout=" \n/path/to/file\n ") + elif args[0] == "mktemp": + return CompletedProcess([], 0, stdout=" \n/tmp/tmpdir\n ") + elif args[0] == "nix-store": + return CompletedProcess([], 0, stdout=" \n/tmp/tmpdir/config\n ") + elif args[0] == "readlink": + return CompletedProcess([], 0, stdout=" \n/path/to/config\n ") + else: + return CompletedProcess([], 0) + + mock_run.side_effect = run_wrapper_side_effect + mock_uuid4.side_effect = [uuid.UUID(int=1), uuid.UUID(int=2)] + assert n.remote_build( "config.system.build.toplevel", m.BuildAttr("", "preAttr"), @@ -88,7 +104,8 @@ def test_remote_build(mock_run: Any, monkeypatch: Any) -> None: build_flags={"build": True}, instantiate_flags={"inst": True}, copy_flags={"copy": True}, - ) == Path("/path/to/file") + ) == Path("/path/to/config") + mock_run.assert_has_calls( [ call( @@ -97,6 +114,8 @@ def test_remote_build(mock_run: Any, monkeypatch: Any) -> None: "", "--attr", "preAttr.config.system.build.toplevel", + "--add-root", + n.tmpdir.TMPDIR_PATH / "00000000000000000000000000000001", "--inst", ], stdout=PIPE, @@ -114,10 +133,28 @@ def test_remote_build(mock_run: Any, monkeypatch: Any) -> None: }, ), call( - ["nix-store", "--realise", Path("/path/to/file"), "--build"], + ["mktemp", "-d", "-t", "nixos-rebuild.XXXXX"], remote=build_host, stdout=PIPE, ), + call( + [ + "nix-store", + "--realise", + Path("/path/to/file"), + "--add-root", + Path("/tmp/tmpdir/00000000000000000000000000000002"), + "--build", + ], + remote=build_host, + stdout=PIPE, + ), + call( + ["readlink", "-f", "/tmp/tmpdir/config"], + remote=build_host, + stdout=PIPE, + ), + call(["rm", "-rf", Path("/tmp/tmpdir")], remote=build_host, check=False), ] )