From 4adad7f6648ae7e373c6bf4ffeac032c7088a9d4 Mon Sep 17 00:00:00 2001 From: Thiago Kenji Okada Date: Mon, 25 Nov 2024 20:04:35 +0000 Subject: [PATCH] nixos-rebuild-ng: implement `--target-host` for `--rollback` --- .../src/nixos_rebuild/__init__.py | 18 +++++- .../nixos-rebuild-ng/src/nixos_rebuild/nix.py | 34 ++++++++-- .../ni/nixos-rebuild-ng/src/tests/test_nix.py | 63 ++++++++++++++++--- 3 files changed, 100 insertions(+), 15 deletions(-) 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 081c179ad970..aea83f9a1842 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 @@ -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: 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 edc1628145da..51799f89a1fc 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 @@ -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: 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 23c2567057da..61686a038f03 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 @@ -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)