nixos-rebuild-ng: implement _NIXOS_REBUILD_REEXEC

pkgs/by-name/ni/nixos-rebuild-ng/README.md

@@ -91,14 +91,6 @@ ruff format .
   conflicting with the current `nixos-rebuild`. This means you can keep both in
   your system at the same time, but it also means that a few things like bash
   completion are broken right now (since it looks at `nixos-rebuild` binary)
-- `_NIXOS_REBUILD_EXEC` is **not** implemented yet, so different from
-  `nixos-rebuild`, this will use the current version of `nixos-rebuild-ng` in
-  your `PATH` to build/set profile/switch, while `nixos-rebuild` builds the new
-  version (the one that will be switched) and re-exec to it instead. This means
-  that in case of bugs in `nixos-rebuild-ng`, the only way that you will get
-  them fixed is **after** you switch to a new version
-- Ignore any performance advantages of the rewrite right now, because of the 2
-  caveats above
 - Bugs in the profile manipulation can cause corruption of your profile that
   may be difficult to fix, so right now I only recommend using
   `nixos-rebuild-ng` if you are testing in a VM or in a filesystem with
@@ -111,7 +103,7 @@ ruff format .
 - [x] Remote host/builders (via SSH)
 - [x] Improve nix arguments handling (e.g.: `nixFlags` vs `copyFlags` in the
   old `nixos-rebuild`)
-- [ ] `_NIXOS_REBUILD_EXEC`
+- [x] `_NIXOS_REBUILD_REEXEC`
 - [ ] Port `nixos-rebuild.passthru.tests`
 - [ ] Change module system to allow easier opt-in, like
   `system.switch.enableNg` for `switch-to-configuration-ng`

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

@@ -2,6 +2,7 @@ import argparse
 import atexit
 import json
 import logging
+import os
 import sys
 from pathlib import Path
 from subprocess import run
@@ -87,6 +88,7 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
     main_parser.add_argument("--ask-sudo-password", action="store_true")
     main_parser.add_argument("--use-remote-sudo", action="store_true")  # deprecated
     main_parser.add_argument("--no-ssh-tty", action="store_true")  # deprecated
+    main_parser.add_argument("--fast", action="store_true")
     main_parser.add_argument("--build-host")
     main_parser.add_argument("--target-host")
     main_parser.add_argument("action", choices=Action.values(), nargs="?")
@@ -174,10 +176,6 @@ def execute(argv: list[str]) -> None:
     target_host = Remote.from_arg(args.target_host, args.ask_sudo_password, tmpdir_path)
     build_attr = BuildAttr.from_arg(args.attr, args.file)
     flake = Flake.from_arg(args.flake, target_host)
-
-    if args.upgrade or args.upgrade_all:
-        nix.upgrade_channels(bool(args.upgrade_all))
-
     action = Action(args.action)
     # Only run shell scripts from the Nixpkgs tree if the action is
     # "switch", "boot", or "test". With other actions (such as "build"),
@@ -185,6 +183,30 @@ def execute(argv: list[str]) -> None:
     # executed, so it's safe to run nixos-rebuild against a potentially
     # untrusted tree.
     can_run = action in (Action.SWITCH, Action.BOOT, Action.TEST)
+
+    # Re-exec to a newer version of the script before building to ensure we get
+    # the latest fixes
+    if can_run and not args.fast and not os.environ.get("_NIXOS_REBUILD_REEXEC"):
+        if flake:
+            drv = nix.build_flake("pkgs.nixos-rebuild-ng", flake, **flake_build_flags)
+        else:
+            drv = nix.build("pkgs.nixos-rebuild-ng", build_attr, **build_flags)
+        new = drv / "bin/nixos-rebuild-ng"
+        current = Path(argv[0])
+        # Disable re-exec during development
+        if current.name != "__main__.py" and new != current:
+            logging.debug(
+                "detected newer version of script, re-exec'ing, current=%s, new=%s",
+                argv[0],
+                new,
+            )
+            cleanup_ssh(tmpdir_path)
+            tmpdir.cleanup()
+            os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
+
+    if args.upgrade or args.upgrade_all:
+        nix.upgrade_channels(bool(args.upgrade_all))
+
     if can_run and not flake:
         nixpkgs_path = nix.find_file("nixpkgs", **build_flags)
         rev = nix.get_nixpkgs_rev(nixpkgs_path)

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

@@ -95,7 +95,7 @@ def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None:
         CompletedProcess([], 0),
     ]
 
-    nr.execute(["nixos-rebuild", "boot", "--no-flake", "-vvv"])
+    nr.execute(["nixos-rebuild", "boot", "--no-flake", "-vvv", "--fast"])
 
     assert mock_run.call_count == 6
     mock_run.assert_has_calls(
@@ -173,6 +173,7 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
             "--install-bootloader",
             "--sudo",
             "--verbose",
+            "--fast",
         ]
     )
 
@@ -246,6 +247,7 @@ def test_execute_nix_switch_flake_target_host(
             "--use-remote-sudo",
             "--target-host",
             "user@localhost",
+            "--fast",
         ]
     )
 
@@ -339,6 +341,7 @@ def test_execute_nix_switch_flake_build_host(
             "/path/to/config#hostname",
             "--build-host",
             "user@localhost",
+            "--fast",
         ]
     )
 
@@ -415,7 +418,9 @@ def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
     nixpkgs_path = tmp_path / "nixpkgs"
     nixpkgs_path.touch()
 
-    nr.execute(["nixos-rebuild", "switch", "--rollback", "--install-bootloader"])
+    nr.execute(
+        ["nixos-rebuild", "switch", "--rollback", "--install-bootloader", "--fast"]
+    )
 
     assert mock_run.call_count >= 2
     # ignoring update_nixpkgs_rev calls
@@ -467,13 +472,7 @@ def test_execute_test_rollback(
     ]
 
     nr.execute(
-        [
-            "nixos-rebuild",
-            "test",
-            "--rollback",
-            "--profile-name",
-            "foo",
-        ]
+        ["nixos-rebuild", "test", "--rollback", "--profile-name", "foo", "--fast"]
     )
 
     assert mock_run.call_count == 2