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

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py

@@ -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")  # deprecated
-    parser.add_argument("--install-grub", action="store_true")
     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", action="store_true")  # deprecated
-    parser.add_argument("--use-remote-sudo", dest="sudo", action="store_true")
+    parser.add_argument("--no-ssh-tty", action="store_true")
-    parser.add_argument("action", choices=Action.values(), nargs="?")
+    # 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"])
+    common_build_flags = common_flags | flags_to_dict(
-    flake_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:
+            else:
+                if flake:
                     path_to_config = nixos_build_flake(
                         "toplevel",
                         flake,
                         no_link=True,
-                    **flake_flags,
+                        **flake_build_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,
+                        **build_flags,
                     )
-                set_profile(profile, path_to_config, sudo=args.sudo)
+                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:

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py

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

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py

@@ -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,
     )
 

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py

@@ -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]:

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py

@@ -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"

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py

@@ -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)
+    monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar")
-    assert m.SSH.from_arg("user@localhost", True, tmpdir) == m.SSH(
+    assert m.Ssh.from_arg("user@localhost", True, tmpdir) == m.Ssh(
         "user@localhost",
         [
             "-f",

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py

@@ -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),
     )
 
 

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py

@@ -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),
         )