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

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

@@ -118,6 +118,16 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     if args.action == Action.EDIT.value and (args.file or args.attr):
         parser.error("--file and --attr are not supported with 'edit'")
 
+    if args.target_host and args.action not in (
+        Action.SWITCH.value,
+        Action.BOOT.value,
+        Action.TEST.value,
+        Action.BUILD.value,
+        Action.DRY_BUILD.value,
+        Action.DRY_ACTIVATE.value,
+    ):
+        parser.error(f"--target-host is not supported with '{args.action}'")
+
     if args.flake and (args.file or args.attr):
         parser.error("--flake cannot be used with --file or --attr")
 
@@ -203,7 +213,7 @@ def execute(argv: list[str]) -> None:
         case Action.SWITCH | Action.BOOT:
             info("building the system configuration...")
             if args.rollback:
-                path_to_config = rollback(profile)
+                path_to_config = rollback(profile, target_host, sudo=args.sudo)
             else:
                 if flake:
                     path_to_config = nixos_build_flake(
@@ -236,7 +246,11 @@ def execute(argv: list[str]) -> None:
             if args.rollback:
                 if action not in (Action.TEST, Action.BUILD):
                     raise NRError(f"--rollback is incompatible with '{action}'")
-                maybe_path_to_config = rollback_temporary_profile(profile)
+                maybe_path_to_config = 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:

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

@@ -149,7 +149,12 @@ def _parse_generation_from_nix_env(line: str) -> Generation:
     )
 
 
-def get_generations(profile: Profile, lock_profile: bool = False) -> list[Generation]:
+def get_generations(
+    profile: Profile,
+    target_host: Remote | None = None,
+    using_nix_env: bool = False,
+    sudo: bool = False,
+) -> list[Generation]:
     """Get all NixOS generations from profile.
 
     Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created)
@@ -161,7 +166,7 @@ def get_generations(profile: Profile, lock_profile: bool = False) -> list[Genera
         raise NRError(f"no profile '{profile.name}' found")
 
     result = []
-    if lock_profile:
+    if using_nix_env:
         # Using `nix-env --list-generations` needs root to lock the profile
         # TODO: do we actually need to lock profile for e.g.: rollback?
         # https://github.com/NixOS/nix/issues/5144
@@ -169,10 +174,13 @@ def get_generations(profile: Profile, lock_profile: bool = False) -> list[Genera
             ["nix-env", "-p", profile.path, "--list-generations"],
             stdout=PIPE,
             check=True,
+            remote=target_host,
+            sudo=sudo,
         )
         for line in r.stdout.splitlines():
             result.append(_parse_generation_from_nix_env(line))
     else:
+        assert not target_host, "target_host is not supported when using_nix_env=False"
         for p in profile.path.parent.glob("system-*-link"):
             result.append(_parse_generation_from_nix_store(p, profile))
     return sorted(result, key=lambda d: d.id)
@@ -279,16 +287,30 @@ def nixos_build_flake(
     return Path(r.stdout.strip())
 
 
-def rollback(profile: Profile) -> Path:
+def rollback(profile: Profile, target_host: Remote | None, sudo: bool) -> Path:
     "Rollback Nix profile, like one created by `nixos-rebuild switch`."
-    run_wrapper(["nix-env", "--rollback", "-p", profile.path], check=True)
+    run_wrapper(
+        ["nix-env", "--rollback", "-p", profile.path],
+        check=True,
+        remote=target_host,
+        sudo=sudo,
+    )
     # Rollback config PATH is the own profile
     return profile.path
 
 
-def rollback_temporary_profile(profile: Profile) -> Path | None:
+def rollback_temporary_profile(
+    profile: Profile,
+    target_host: Remote | None,
+    sudo: bool,
+) -> Path | None:
     "Rollback a temporary Nix profile, like one created by `nixos-rebuild test`."
-    generations = get_generations(profile, lock_profile=True)
+    generations = get_generations(
+        profile,
+        target_host=target_host,
+        using_nix_env=True,
+        sudo=sudo,
+    )
     previous_gen_id = None
     for generation in generations:
         if not generation.current:

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

@@ -122,7 +122,7 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
 
     assert n.get_generations(
         m.Profile("system", tmp_path / "system"),
-        lock_profile=False,
+        using_nix_env=False,
     ) == [
         m.Generation(id=1, current=False, timestamp=ANY),
         m.Generation(id=2, current=True, timestamp=ANY),
@@ -147,7 +147,7 @@ def test_get_generations_from_nix_env(mock_run: Any, tmp_path: Path) -> None:
     path = tmp_path / "test"
     path.touch()
 
-    assert n.get_generations(m.Profile("system", path), lock_profile=True) == [
+    assert n.get_generations(m.Profile("system", path), using_nix_env=True) == [
         m.Generation(id=2082, current=False, timestamp="2024-11-07 22:58:56"),
         m.Generation(id=2083, current=False, timestamp="2024-11-07 22:59:41"),
         m.Generation(id=2084, current=True, timestamp="2024-11-07 23:54:17"),
@@ -268,8 +268,28 @@ def test_rollback(mock_run: Any, tmp_path: Path) -> None:
 
     profile = m.Profile("system", path)
 
-    assert n.rollback(profile) == profile.path
-    mock_run.assert_called_with(["nix-env", "--rollback", "-p", path], check=True)
+    assert n.rollback(profile, None, False) == profile.path
+    mock_run.assert_called_with(
+        ["nix-env", "--rollback", "-p", path],
+        check=True,
+        remote=None,
+        sudo=False,
+    )
+
+    assert (
+        n.rollback(
+            profile,
+            m.Remote("user@localhost", [], False),
+            True,
+        )
+        == profile.path
+    )
+    mock_run.assert_called_with(
+        ["nix-env", "--rollback", "-p", path],
+        check=True,
+        remote=m.Remote("user@localhost", [], False),
+        sudo=True,
+    )
 
 
 def test_rollback_temporary_profile(tmp_path: Path) -> None:
@@ -288,17 +308,46 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
                 """),
         )
         assert (
-            n.rollback_temporary_profile(m.Profile("system", path))
+            n.rollback_temporary_profile(m.Profile("system", path), None, False)
             == path.parent / "system-2083-link"
         )
+        mock_run.assert_called_with(
+            [
+                "nix-env",
+                "-p",
+                path,
+                "--list-generations",
+            ],
+            stdout=PIPE,
+            check=True,
+            remote=None,
+            sudo=False,
+        )
+
         assert (
-            n.rollback_temporary_profile(m.Profile("foo", path))
+            n.rollback_temporary_profile(
+                m.Profile("foo", path),
+                m.Remote("user@localhost", [], False),
+                True,
+            )
             == path.parent / "foo-2083-link"
         )
+        mock_run.assert_called_with(
+            [
+                "nix-env",
+                "-p",
+                path,
+                "--list-generations",
+            ],
+            stdout=PIPE,
+            check=True,
+            remote=m.Remote("user@localhost", [], False),
+            sudo=True,
+        )
 
     with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
         mock_run.return_value = CompletedProcess([], 0, stdout="")
-        assert n.rollback_temporary_profile(profile) is None
+        assert n.rollback_temporary_profile(profile, None, False) is None
 
 
 @patch(get_qualified_name(n.run_wrapper, n), autospec=True)