nixos-rebuild-ng: implement --target-host

This commit is contained in:
Thiago Kenji Okada 2024-11-22 19:10:37 +00:00
parent fd1cd69315
commit a2cbe67701
8 changed files with 242 additions and 50 deletions

View File

@ -2,11 +2,14 @@ import argparse
import json
import os
import sys
from pathlib import Path
from subprocess import run
from tempfile import TemporaryDirectory
from typing import assert_never
from .models import Action, Flake, NRError, Profile
from .models import Action, Flake, NRError, Profile, Ssh
from .nix import (
copy_closure,
edit,
find_file,
get_nixpkgs_rev,
@ -37,8 +40,7 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
parser.add_argument("--flake", nargs="?", const=True)
parser.add_argument("--no-flake", dest="flake", action="store_false")
parser.add_argument("--install-bootloader", action="store_true")
# TODO: add deprecated=True in Python >=3.13
parser.add_argument("--install-grub", action="store_true")
parser.add_argument("--install-grub", action="store_true") # deprecated
parser.add_argument("--profile-name", "-p", default="system")
parser.add_argument("--specialisation", "-c")
parser.add_argument("--rollback", action="store_true")
@ -46,12 +48,14 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
parser.add_argument("--upgrade-all", action="store_true")
parser.add_argument("--json", action="store_true")
parser.add_argument("--sudo", action="store_true")
# TODO: add deprecated=True in Python >=3.13
parser.add_argument("--use-remote-sudo", dest="sudo", action="store_true")
parser.add_argument("action", choices=Action.values(), nargs="?")
parser.add_argument("--use-remote-sudo", action="store_true") # deprecated
parser.add_argument("--no-ssh-tty", action="store_true")
# parser.add_argument("--build-host") # TODO
parser.add_argument("--target-host")
parser.add_argument("--verbose", "-v", action="count", default=0)
parser.add_argument("action", choices=Action.values(), nargs="?")
common_group = parser.add_argument_group("Common flags")
common_group = parser.add_argument_group("nix flags")
common_group.add_argument("--include", "-I")
common_group.add_argument("--max-jobs", "-j")
common_group.add_argument("--cores")
@ -65,10 +69,10 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
common_group.add_argument("--repair", action="store_true")
common_group.add_argument("--option", nargs=2)
nix_group = parser.add_argument_group("Classic Nix flags")
nix_group = parser.add_argument_group("nix classic flags")
nix_group.add_argument("--no-build-output", "-Q", action="store_true")
flake_group = parser.add_argument_group("Flake flags")
flake_group = parser.add_argument_group("nix flakes flags")
flake_group.add_argument("--accept-flake-config", action="store_true")
flake_group.add_argument("--refresh", action="store_true")
flake_group.add_argument("--impure", action="store_true")
@ -82,6 +86,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
flake_group.add_argument("--update-input")
flake_group.add_argument("--override-input", nargs=2)
copy_group = parser.add_argument_group("nix-copy-closure flags")
copy_group.add_argument("--use-substitutes", "-s", action="store_true")
args = parser.parse_args(argv[1:])
global VERBOSE
@ -92,12 +99,20 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
if args.action == Action.DRY_RUN.value:
args.action = Action.DRY_BUILD.value
# TODO: use deprecated=True in Python >=3.13
if args.install_grub:
info(
f"{parser.prog}: warning: --install-grub deprecated, use --install-bootloader instead"
)
args.install_bootloader = True
# TODO: use deprecated=True in Python >=3.13
if args.use_remote_sudo:
info(
f"{parser.prog}: warning: --use-remote-sudo deprecated, use --sudo instead"
)
args.sudo = True
if args.action == Action.EDIT.value and (args.file or args.attr):
parser.error("--file and --attr are not supported with 'edit'")
@ -117,23 +132,28 @@ def execute(argv: list[str]) -> None:
common_flags = flags_to_dict(
args,
[
"verbose",
"include",
"max_jobs",
"cores",
"log_format",
"quiet",
"print_build_logs",
"show_trace",
"keep_going",
"keep_failed",
"fallback",
"repair",
"verbose",
"option",
],
)
nix_flags = common_flags | flags_to_dict(args, ["no_build_output"])
flake_flags = common_flags | flags_to_dict(
common_build_flags = common_flags | flags_to_dict(
args,
[
"include",
"quiet",
"print_build_logs",
"show_trace",
],
)
build_flags = common_build_flags | flags_to_dict(args, ["no_build_output"])
flake_build_flags = common_build_flags | flags_to_dict(
args,
[
"accept_flake_config",
@ -150,9 +170,15 @@ def execute(argv: list[str]) -> None:
"override_input",
],
)
copy_flags = common_flags | flags_to_dict(args, ["use_substitutes"])
# Will be cleaned up on exit automatically.
tmpdir = TemporaryDirectory(prefix="nixos-rebuild.")
tmpdir_path = Path(tmpdir.name)
profile = Profile.from_name(args.profile_name)
flake = Flake.from_arg(args.flake)
target_host = Ssh.from_arg(args.target_host, not args.no_ssh_tty, tmpdir_path)
if args.upgrade or args.upgrade_all:
upgrade_channels(bool(args.upgrade_all))
@ -165,7 +191,7 @@ def execute(argv: list[str]) -> None:
# untrusted tree.
can_run = action in (Action.SWITCH, Action.BOOT, Action.TEST)
if can_run and not flake:
nixpkgs_path = find_file("nixpkgs", **nix_flags)
nixpkgs_path = find_file("nixpkgs", **build_flags)
rev = get_nixpkgs_rev(nixpkgs_path)
if nixpkgs_path and rev:
(nixpkgs_path / ".version-suffix").write_text(rev)
@ -175,26 +201,28 @@ def execute(argv: list[str]) -> None:
info("building the system configuration...")
if args.rollback:
path_to_config = rollback(profile)
elif flake:
path_to_config = nixos_build_flake(
"toplevel",
flake,
no_link=True,
**flake_flags,
)
set_profile(profile, path_to_config, sudo=args.sudo)
else:
path_to_config = nixos_build(
"system",
args.attr,
args.file,
no_out_link=True,
**nix_flags,
)
set_profile(profile, path_to_config, sudo=args.sudo)
if flake:
path_to_config = nixos_build_flake(
"toplevel",
flake,
no_link=True,
**flake_build_flags,
)
else:
path_to_config = nixos_build(
"system",
args.attr,
args.file,
no_out_link=True,
**build_flags,
)
copy_closure(path_to_config, target_host, **copy_flags)
set_profile(profile, path_to_config, target_host, sudo=args.sudo)
switch_to_configuration(
path_to_config,
action,
target_host,
sudo=args.sudo,
specialisation=args.specialisation,
install_bootloader=args.install_bootloader,
@ -214,7 +242,7 @@ def execute(argv: list[str]) -> None:
flake,
keep_going=True,
dry_run=dry_run,
**flake_flags,
**flake_build_flags,
)
else:
path_to_config = nixos_build(
@ -223,12 +251,13 @@ def execute(argv: list[str]) -> None:
args.file,
keep_going=True,
dry_run=dry_run,
**nix_flags,
**build_flags,
)
if action in (Action.TEST, Action.DRY_ACTIVATE):
switch_to_configuration(
path_to_config,
action,
target_host,
sudo=args.sudo,
specialisation=args.specialisation,
)
@ -240,7 +269,7 @@ def execute(argv: list[str]) -> None:
attr,
flake,
keep_going=True,
**flake_flags,
**flake_build_flags,
)
else:
path_to_config = nixos_build(
@ -248,12 +277,12 @@ def execute(argv: list[str]) -> None:
args.attr,
args.file,
keep_going=True,
**nix_flags,
**build_flags,
)
vm_path = next(path_to_config.glob("bin/run-*-vm"), "./result/bin/run-*-vm")
print(f"Done. The virtual machine can be started by running '{vm_path}'")
case Action.EDIT:
edit(flake, **flake_flags)
edit(flake, **flake_build_flags)
case Action.DRY_RUN:
assert False, "DRY_RUN should be a DRY_BUILD alias"
case Action.LIST_GENERATIONS:

View File

@ -116,7 +116,7 @@ class Profile:
@dataclass(frozen=True)
class SSH:
class Ssh:
host: str
opts: list[str]
tty: bool

View File

@ -12,6 +12,7 @@ from .models import (
GenerationJson,
NRError,
Profile,
Ssh,
)
from .process import run_wrapper
from .utils import Args, dict_to_flags, info
@ -19,6 +20,28 @@ from .utils import Args, dict_to_flags, info
FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
def copy_closure(
closure: Path,
target_host: Ssh | None,
**copy_flags: Args,
) -> None:
host = target_host
if not host:
return
run_wrapper(
[
"nix-copy-closure",
*(dict_to_flags(copy_flags)),
"--to",
host.host,
closure,
],
check=True,
extra_env={"NIX_SSHOPTS": " ".join(host.opts)},
)
def edit(flake: Flake | None, **flake_flags: Args) -> None:
"Try to find and open NixOS configuration file in editor."
if flake:
@ -277,16 +300,25 @@ def rollback_temporary_profile(profile: Profile) -> Path | None:
return None
def set_profile(profile: Profile, path_to_config: Path, sudo: bool) -> None:
def set_profile(
profile: Profile,
path_to_config: Path,
target_host: Ssh | None,
sudo: bool,
) -> None:
"Set a path as the current active Nix profile."
run_wrapper(
["nix-env", "-p", profile.path, "--set", path_to_config], check=True, sudo=sudo
["nix-env", "-p", profile.path, "--set", path_to_config],
check=True,
remote=target_host,
sudo=sudo,
)
def switch_to_configuration(
path_to_config: Path,
action: Action,
target_host: Ssh | None,
sudo: bool,
install_bootloader: bool = False,
specialisation: str | None = None,
@ -310,6 +342,7 @@ def switch_to_configuration(
[path_to_config / "bin/switch-to-configuration", str(action)],
extra_env={"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0"},
check=True,
remote=target_host,
sudo=sudo,
)

View File

@ -4,7 +4,7 @@ import os
import subprocess
from typing import Sequence, TypedDict, Unpack
from .models import SSH
from .models import Ssh
# Not exhaustive, but we can always extend it later.
@ -22,7 +22,7 @@ def run_wrapper(
check: bool, # make it explicit so we always know if the code is handling errors
env: dict[str, str] | None = None, # replaces the current environment
extra_env: dict[str, str] | None = None, # appends to the current environment
remote: SSH | None = None,
remote: Ssh | None = None,
sudo: bool = False,
**kwargs: Unpack[RunKwargs],
) -> subprocess.CompletedProcess[str]:

View File

@ -208,6 +208,109 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
)
@patch.dict(nr.process.os.environ, {}, clear=True)
@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
@patch(get_qualified_name(nr.TemporaryDirectory, nr)) # can't autospec
def test_execute_nix_switch_flake_remote(
mock_tmpdir: Any,
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),
# copy_closure
CompletedProcess([], 0),
# switch_to_configuration
CompletedProcess([], 0),
]
mock_tmpdir.return_value.name = "/tmp/test"
nr.execute(
[
"nixos-rebuild",
"switch",
"--flake",
"/path/to/config#hostname",
"--use-remote-sudo",
"--target-host",
"user@localhost",
]
)
assert mock_run.call_count == 4
mock_run.assert_has_calls(
[
call(
[
"nix",
"--extra-experimental-features",
"nix-command flakes",
"build",
"--print-out-paths",
"/path/to/config#nixosConfigurations.hostname.config.system.build.toplevel",
"--no-link",
],
check=True,
stdout=PIPE,
**DEFAULT_RUN_KWARGS,
),
call(
["nix-copy-closure", "--to", "user@localhost", config_path],
check=True,
**DEFAULT_RUN_KWARGS,
),
call(
[
"ssh",
"-t",
"-o",
"ControlMaster=auto",
"-o",
"ControlPath=/tmp/test/ssh-%n",
"-o",
"ControlPersist=60",
"user@localhost",
"--",
"sudo",
"nix-env",
"-p",
Path("/nix/var/nix/profiles/system"),
"--set",
config_path,
],
check=True,
**DEFAULT_RUN_KWARGS,
),
call(
[
"ssh",
"-t",
"-o",
"ControlMaster=auto",
"-o",
"ControlPath=/tmp/test/ssh-%n",
"-o",
"ControlPersist=60",
"user@localhost",
"--",
"sudo",
"env",
"NIXOS_INSTALL_BOOTLOADER=0",
config_path / "bin/switch-to-configuration",
"switch",
],
check=True,
**DEFAULT_RUN_KWARGS,
),
]
)
@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
nixpkgs_path = tmp_path / "nixpkgs"

View File

@ -101,7 +101,8 @@ def test_profile_from_name(mock_mkdir: Any) -> None:
def test_ssh_from_name(monkeypatch: Any, tmpdir: Path) -> None:
assert m.SSH.from_arg("user@localhost", None, tmpdir) == m.SSH(
monkeypatch.setenv("NIX_SSHOPTS", "")
assert m.Ssh.from_arg("user@localhost", None, tmpdir) == m.Ssh(
"user@localhost",
[
"-o",
@ -114,8 +115,8 @@ def test_ssh_from_name(monkeypatch: Any, tmpdir: Path) -> None:
False,
)
monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar", tmpdir)
assert m.SSH.from_arg("user@localhost", True, tmpdir) == m.SSH(
monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar")
assert m.Ssh.from_arg("user@localhost", True, tmpdir) == m.Ssh(
"user@localhost",
[
"-f",

View File

@ -12,6 +12,21 @@ import nixos_rebuild.nix as n
from .helpers import get_qualified_name
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_copy_closure(mock_run: Any) -> None:
closure = Path("/path/to/closure")
n.copy_closure(closure, None)
mock_run.assert_not_called()
target_host = m.Ssh("user@host", ["--ssh", "opt"], False)
n.copy_closure(closure, target_host)
mock_run.assert_called_with(
["nix-copy-closure", "--to", "user@host", closure],
check=True,
extra_env={"NIX_SSHOPTS": "--ssh opt"},
)
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_edit(mock_run: Any, monkeypatch: Any, tmpdir: Any) -> None:
# Flake
@ -290,11 +305,17 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
def test_set_profile(mock_run: Any) -> None:
profile_path = Path("/path/to/profile")
config_path = Path("/path/to/config")
n.set_profile(m.Profile("system", profile_path), config_path, sudo=False)
n.set_profile(
m.Profile("system", profile_path),
config_path,
target_host=None,
sudo=False,
)
mock_run.assert_called_with(
["nix-env", "-p", profile_path, "--set", config_path],
check=True,
remote=None,
sudo=False,
)
@ -311,6 +332,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
profile_path,
m.Action.SWITCH,
sudo=False,
target_host=None,
specialisation=None,
install_bootloader=False,
)
@ -319,6 +341,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"},
check=True,
sudo=False,
remote=None,
)
with pytest.raises(m.NRError) as e:
@ -326,6 +349,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
config_path,
m.Action.BOOT,
sudo=False,
target_host=None,
specialisation="special",
)
assert (
@ -342,6 +366,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
Path("/path/to/config"),
m.Action.TEST,
sudo=True,
target_host=m.Ssh("user@localhost", [], False),
install_bootloader=True,
specialisation="special",
)
@ -353,6 +378,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"},
check=True,
sudo=True,
remote=m.Ssh("user@localhost", [], False),
)

View File

@ -41,7 +41,7 @@ def test_run(mock_run: Any) -> None:
p.run_wrapper(
["test", "--with", "flags"],
check=True,
remote=m.SSH("user@localhost", ["--ssh", "opt"], False),
remote=m.Ssh("user@localhost", ["--ssh", "opt"], False),
)
mock_run.assert_called_with(
["ssh", "--ssh", "opt", "user@localhost", "--", "test", "--with", "flags"],
@ -69,7 +69,7 @@ def test_run(mock_run: Any) -> None:
check=True,
sudo=True,
extra_env={"FOO": "bar"},
remote=m.SSH("user@localhost", ["--ssh", "opt"], True),
remote=m.Ssh("user@localhost", ["--ssh", "opt"], True),
)
mock_run.assert_called_with(
[
@ -97,5 +97,5 @@ def test_run(mock_run: Any) -> None:
["test", "--with", "flags"],
check=False,
env={"foo": "bar"},
remote=m.SSH("user@localhost", [], False),
remote=m.Ssh("user@localhost", [], False),
)