nixos-rebuild-ng: add --add-root in nix.remote_build (#365592)

This commit is contained in:
Thiago Kenji Okada
2024-12-18 10:08:40 +00:00
committed by GitHub
6 changed files with 180 additions and 92 deletions
@@ -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
@@ -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(
@@ -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(
@@ -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)
@@ -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"]
@@ -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("<nixpkgs/nixos>", "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:
"<nixpkgs/nixos>",
"--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),
]
)