Merge pull request 'fast-vars-gen' (#3216) from fast-vars-gen into main

Reviewed-on: https://git.clan.lol/clan/clan-core/pulls/3216

pkgs/clan-cli/clan_cli/flake.py

@@ -518,6 +518,23 @@ class Flake:
         selectors: list[str],
         nix_options: list[str] | None = None,
     ) -> None:
+        """
+        Retrieves specific attributes from a Nix flake using the provided selectors.
+
+        This function interacts with the Nix build system to fetch and process
+        attributes from a flake. It uses the provided selectors to determine which
+        attributes to retrieve and optionally accepts additional Nix options for
+        customization. The results are cached for future use.
+        Used mostly as a lowlevel function for `precache` and `select` methods.
+
+        Args:
+            selectors (list[str]): A list of attribute selectors to fetch from the flake.
+            nix_options (list[str] | None): Optional additional options to pass to the Nix build command.
+
+        Raises:
+            ClanError: If the number of outputs does not match the number of selectors.
+            AssertionError: If the cache or flake cache path is not properly initialized.
+        """
         if self._cache is None:
             self.prefetch()
         assert self._cache is not None
@@ -531,7 +548,7 @@ class Flake:
               flake = builtins.getFlake("path:{self.store_path}?narHash={self.hash}");
             in
               flake.inputs.nixpkgs.legacyPackages.{config["system"]}.writeText "clan-flake-select" (
-                builtins.toJSON [ ({" ".join([f"flake.clanInternals.clanLib.select ''{attr}'' flake" for attr in selectors])}) ]
+                builtins.toJSON [ {" ".join([f"(flake.clanInternals.clanLib.select ''{attr}'' flake)" for attr in selectors])} ]
               )
         """
         if tmp_store := nix_test_store():
@@ -556,11 +573,46 @@ class Flake:
         if self.flake_cache_path:
             self._cache.save_to_file(self.flake_cache_path)
 
+    def precache(
+        self,
+        selectors: list[str],
+        nix_options: list[str] | None = None,
+    ) -> None:
+        """
+        Ensures that the specified selectors are cached locally.
+
+        This function checks if the given selectors are already cached. If not, it
+        fetches them using the Nix build system and stores them in the local cache.
+        It ensures that the cache is initialized before performing these operations.
+
+        Args:
+            selectors (list[str]): A list of attribute selectors to check and cache.
+            nix_options (list[str] | None): Optional additional options to pass to the Nix build command.
+        """
+        if self._cache is None:
+            self.prefetch()
+        assert self._cache is not None
+        assert self.flake_cache_path is not None
+        not_fetched_selectors = []
+        for selector in selectors:
+            if not self._cache.is_cached(selector):
+                not_fetched_selectors.append(selector)
+        if not_fetched_selectors:
+            self.get_from_nix(not_fetched_selectors, nix_options)
+
     def select(
         self,
         selector: str,
         nix_options: list[str] | None = None,
     ) -> Any:
+        """
+        Selects a value from the cache based on the provided selector string.
+        Fetches it via nix_build if it is not already cached.
+
+        Args:
+            selector (str): The attribute selector string to fetch the value for.
+            nix_options (list[str] | None): Optional additional options to pass to the Nix build command.
+        """
         if self._cache is None:
             self.prefetch()
         assert self._cache is not None

pkgs/clan-cli/clan_cli/machines/update.py

@@ -20,7 +20,7 @@ from clan_cli.facts.upload import upload_secrets
 from clan_cli.flake import Flake
 from clan_cli.inventory import Machine as InventoryMachine
 from clan_cli.machines.machines import Machine
-from clan_cli.nix import nix_command, nix_metadata
+from clan_cli.nix import nix_command, nix_config, nix_metadata
 from clan_cli.ssh.host import Host, HostKeyCheck
 from clan_cli.vars.generate import generate_vars
 from clan_cli.vars.upload import upload_secret_vars
@@ -265,6 +265,16 @@ def update_command(args: argparse.Namespace) -> None:
                 machine.error("Updating macOS machines is not yet supported")
                 sys.exit(1)
 
+        config = nix_config()
+        system = config["system"]
+        machine_names = [machine.name for machine in machines]
+        args.flake.precache(
+            [
+                f"clanInternals.machines.{system}.{{{','.join(machine_names)}}}.config.clan.core.vars.generators.*.validationHash",
+                f"clanInternals.machines.{system}.{{{','.join(machine_names)}}}.config.system.clan.deployment.file",
+            ]
+        )
+
         deploy_machines(machines)
     except KeyboardInterrupt:
         log.warning("Interrupted by user")

pkgs/clan-cli/clan_cli/tests/test_vars.py

@@ -691,9 +691,9 @@ def test_stdout_of_generate(
     flake_with_sops: ClanFlake,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
-    flake = flake_with_sops
+    flake_ = flake_with_sops
 
-    config = flake.machines["my_machine"]
+    config = flake_.machines["my_machine"]
     config["nixpkgs"]["hostPlatform"] = "x86_64-linux"
     my_generator = config["clan"]["core"]["vars"]["generators"]["my_generator"]
     my_generator["files"]["my_value"]["secret"] = False
@@ -703,14 +703,15 @@ def test_stdout_of_generate(
     ]
     my_secret_generator["files"]["my_secret"]["secret"] = True
     my_secret_generator["script"] = 'echo -n hello > "$out"/my_secret'
-    flake.refresh()
-    monkeypatch.chdir(flake.path)
+    flake_.refresh()
+    monkeypatch.chdir(flake_.path)
+    flake = Flake(str(flake_.path))
     from clan_cli.vars.generate import generate_vars_for_machine
 
     # with capture_output as output:
     with caplog.at_level(logging.INFO):
         generate_vars_for_machine(
-            Machine(name="my_machine", flake=Flake(str(flake.path))),
+            Machine(name="my_machine", flake=flake),
             "my_generator",
             regenerate=False,
         )
@@ -720,10 +721,10 @@ def test_stdout_of_generate(
     assert "new: hello" in caplog.text
     caplog.clear()
 
-    set_var("my_machine", "my_generator/my_value", b"world", Flake(str(flake.path)))
+    set_var("my_machine", "my_generator/my_value", b"world", flake)
     with caplog.at_level(logging.INFO):
         generate_vars_for_machine(
-            Machine(name="my_machine", flake=Flake(str(flake.path))),
+            Machine(name="my_machine", flake=flake),
             "my_generator",
             regenerate=True,
         )
@@ -734,7 +735,7 @@ def test_stdout_of_generate(
     # check the output when nothing gets regenerated
     with caplog.at_level(logging.INFO):
         generate_vars_for_machine(
-            Machine(name="my_machine", flake=Flake(str(flake.path))),
+            Machine(name="my_machine", flake=flake),
             "my_generator",
             regenerate=True,
         )
@@ -743,7 +744,7 @@ def test_stdout_of_generate(
     caplog.clear()
     with caplog.at_level(logging.INFO):
         generate_vars_for_machine(
-            Machine(name="my_machine", flake=Flake(str(flake.path))),
+            Machine(name="my_machine", flake=flake),
             "my_secret_generator",
             regenerate=False,
         )
@@ -758,7 +759,7 @@ def test_stdout_of_generate(
     )
     with caplog.at_level(logging.INFO):
         generate_vars_for_machine(
-            Machine(name="my_machine", flake=Flake(str(flake.path))),
+            Machine(name="my_machine", flake=flake),
             "my_secret_generator",
             regenerate=True,
         )

pkgs/clan-cli/clan_cli/vars/generate.py

@@ -17,7 +17,7 @@ from clan_cli.completions import (
 from clan_cli.errors import ClanError
 from clan_cli.git import commit_files
 from clan_cli.machines.inventory import get_all_machines, get_selected_machines
-from clan_cli.nix import nix_shell, nix_test_store
+from clan_cli.nix import nix_config, nix_shell, nix_test_store
 from clan_cli.vars._types import StoreBase
 
 from .check import check_vars
@@ -489,7 +489,6 @@ def generate_vars(
             was_regenerated |= generate_vars_for_machine(
                 machine, generator_name, regenerate, no_sandbox=no_sandbox
             )
-            machine.flush_caches()
         except Exception as exc:
             errors += [(machine, exc)]
         if len(errors) == 1:
@@ -515,6 +514,18 @@ def generate_command(args: argparse.Namespace) -> None:
         machines = get_all_machines(args.flake, args.option)
     else:
         machines = get_selected_machines(args.flake, args.option, args.machines)
+
+    # prefetch all vars
+    config = nix_config()
+    system = config["system"]
+    machine_names = [machine.name for machine in machines]
+    # test
+    args.flake.precache(
+        [
+            f"clanInternals.machines.{system}.{{{','.join(machine_names)}}}.config.clan.core.vars.generators.*.validationHash",
+            f"clanInternals.machines.{system}.{{{','.join(machine_names)}}}.config.system.clan.deployment.file",
+        ]
+    )
     generate_vars(machines, args.generator, args.regenerate, no_sandbox=args.no_sandbox)