Merge pull request 'vars: refactor - make shared generators carry machines list' (#5443) from dave into main

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

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

@@ -166,16 +166,16 @@ def test_generate_public_and_secret_vars(
     assert shared_value.startswith("shared")
     vars_text = stringify_all_vars(machine)
     flake_obj = Flake(str(flake.path))
-    my_generator = Generator("my_generator", machine="my_machine", _flake=flake_obj)
+    my_generator = Generator("my_generator", machines=["my_machine"], _flake=flake_obj)
     shared_generator = Generator(
         "my_shared_generator",
         share=True,
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     dependent_generator = Generator(
         "dependent_generator",
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     in_repo_store = in_repo.FactStore(flake=flake_obj)
@@ -340,12 +340,12 @@ def test_generate_secret_var_sops_with_default_group(
     flake_obj = Flake(str(flake.path))
     first_generator = Generator(
         "first_generator",
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     second_generator = Generator(
         "second_generator",
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     in_repo_store = in_repo.FactStore(flake=flake_obj)
@@ -375,13 +375,13 @@ def test_generate_secret_var_sops_with_default_group(
     first_generator_with_share = Generator(
         "first_generator",
         share=False,
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     second_generator_with_share = Generator(
         "second_generator",
         share=False,
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     assert sops_store.user_has_access("user2", first_generator_with_share, "my_secret")
@@ -512,28 +512,28 @@ def test_generate_secret_var_password_store(
         "my_generator",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     my_generator_shared = Generator(
         "my_generator",
         share=True,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     my_shared_generator = Generator(
         "my_shared_generator",
         share=True,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     my_shared_generator_not_shared = Generator(
         "my_shared_generator",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     assert store.exists(my_generator, "my_secret")
@@ -545,7 +545,7 @@ def test_generate_secret_var_password_store(
         name="my_generator",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     assert store.get(generator, "my_secret").decode() == "hello\n"
@@ -556,7 +556,7 @@ def test_generate_secret_var_password_store(
         "my_generator",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     var_name = "my_secret"
@@ -569,7 +569,7 @@ def test_generate_secret_var_password_store(
         "my_generator2",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     var_name = "my_secret2"
@@ -581,7 +581,7 @@ def test_generate_secret_var_password_store(
         "my_shared_generator",
         share=True,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     var_name = "my_shared_secret"
@@ -628,8 +628,8 @@ def test_generate_secret_for_multiple_machines(
     in_repo_store2 = in_repo.FactStore(flake=flake_obj)
 
     # Create generators for each machine
-    gen1 = Generator("my_generator", machine="machine1", _flake=flake_obj)
-    gen2 = Generator("my_generator", machine="machine2", _flake=flake_obj)
+    gen1 = Generator("my_generator", machines=["machine1"], _flake=flake_obj)
+    gen2 = Generator("my_generator", machines=["machine2"], _flake=flake_obj)
 
     assert in_repo_store1.exists(gen1, "my_value")
     assert in_repo_store2.exists(gen2, "my_value")
@@ -693,12 +693,12 @@ def test_prompt(
 
     # Set up objects for testing the results
     flake_obj = Flake(str(flake.path))
-    my_generator = Generator("my_generator", machine="my_machine", _flake=flake_obj)
+    my_generator = Generator("my_generator", machines=["my_machine"], _flake=flake_obj)
     my_generator_with_details = Generator(
         name="my_generator",
         share=False,
         files=[],
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
 
@@ -783,10 +783,10 @@ def test_shared_vars_regeneration(
     in_repo_store_2 = in_repo.FactStore(machine2.flake)
     # Create generators with machine context for testing
     child_gen_m1 = Generator(
-        "child_generator", share=False, machine="machine1", _flake=machine1.flake
+        "child_generator", share=False, machines=["machine1"], _flake=machine1.flake
     )
     child_gen_m2 = Generator(
-        "child_generator", share=False, machine="machine2", _flake=machine2.flake
+        "child_generator", share=False, machines=["machine2"], _flake=machine2.flake
     )
     # generate for machine 1
     cli.run(["vars", "generate", "--flake", str(flake.path), "machine1"])
@@ -854,13 +854,13 @@ def test_multi_machine_shared_vars(
     generator_m1 = Generator(
         "shared_generator",
         share=True,
-        machine="machine1",
+        machines=["machine1"],
         _flake=machine1.flake,
     )
     generator_m2 = Generator(
         "shared_generator",
         share=True,
-        machine="machine2",
+        machines=["machine2"],
         _flake=machine2.flake,
     )
     # generate for machine 1
@@ -916,7 +916,9 @@ def test_api_set_prompts(
     )
     machine = Machine(name="my_machine", flake=Flake(str(flake.path)))
     store = in_repo.FactStore(machine.flake)
-    my_generator = Generator("my_generator", machine="my_machine", _flake=machine.flake)
+    my_generator = Generator(
+        "my_generator", machines=["my_machine"], _flake=machine.flake
+    )
     assert store.exists(my_generator, "prompt1")
     assert store.get(my_generator, "prompt1").decode() == "input1"
     run_generators(
@@ -1060,10 +1062,10 @@ def test_migration(
     assert "Migrated var my_generator/my_value" in caplog.text
     assert "Migrated secret var my_generator/my_secret" in caplog.text
     flake_obj = Flake(str(flake.path))
-    my_generator = Generator("my_generator", machine="my_machine", _flake=flake_obj)
+    my_generator = Generator("my_generator", machines=["my_machine"], _flake=flake_obj)
     other_generator = Generator(
         "other_generator",
-        machine="my_machine",
+        machines=["my_machine"],
         _flake=flake_obj,
     )
     in_repo_store = in_repo.FactStore(flake=flake_obj)
@@ -1209,7 +1211,7 @@ def test_share_mode_switch_regenerates_secret(
     sops_store = sops.SecretStore(flake=flake_obj)
 
     generator_not_shared = Generator(
-        "my_generator", share=False, machine="my_machine", _flake=flake_obj
+        "my_generator", share=False, machines=["my_machine"], _flake=flake_obj
     )
 
     initial_public = in_repo_store.get(generator_not_shared, "my_value").decode()
@@ -1228,7 +1230,7 @@ def test_share_mode_switch_regenerates_secret(
 
     # Read the new values with shared generator
     generator_shared = Generator(
-        "my_generator", share=True, machine="my_machine", _flake=flake_obj
+        "my_generator", share=True, machines=["my_machine"], _flake=flake_obj
     )
 
     new_public = in_repo_store.get(generator_shared, "my_value").decode()

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

@@ -40,12 +40,15 @@ class StoreBase(ABC):
 
     def get_machine(self, generator: "Generator") -> str:
         """Get machine name from generator, asserting it's not None for now."""
-        if generator.machine is None:
         if generator.share:
             return "__shared"
+        if not generator.machines:
             msg = f"Generator '{generator.name}' has no machine associated"
             raise ClanError(msg)
-        return generator.machine
+        if len(generator.machines) != 1:
+            msg = f"Generator '{generator.name}' has {len(generator.machines)} machines, expected exactly 1"
+            raise ClanError(msg)
+        return generator.machines[0]
 
     # get a single fact
     @abstractmethod
@@ -147,7 +150,7 @@ class StoreBase(ABC):
         prev_generator = dataclasses.replace(
             generator,
             share=not generator.share,
-            machine=machine if generator.share else None,
+            machines=[] if not generator.share else [machine],
         )
         if self.exists(prev_generator, var.name):
             changed_files += self.delete(prev_generator, var.name)
@@ -165,12 +168,12 @@ class StoreBase(ABC):
         new_file = self._set(generator, var, value, machine)
         action_str = "Migrated" if is_migration else "Updated"
         log_info: Callable
-        if generator.machine is None:
+        if generator.share:
             log_info = log.info
         else:
             from clan_lib.machines.machines import Machine  # noqa: PLC0415
 
-            machine_obj = Machine(name=generator.machine, flake=self.flake)
+            machine_obj = Machine(name=generator.machines[0], flake=self.flake)
             log_info = machine_obj.info
         if self.is_secret_store:
             log.info(f"{action_str} secret var {generator.name}/{var.name}\n")

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

@@ -61,14 +61,22 @@ class Generator:
     migrate_fact: str | None = None
     validation_hash: str | None = None
 
-    machine: str | None = None
+    machines: list[str] = field(default_factory=list)
     _flake: "Flake | None" = None
     _public_store: "StoreBase | None" = None
     _secret_store: "StoreBase | None" = None
 
     @property
     def key(self) -> GeneratorKey:
-        return GeneratorKey(machine=self.machine, name=self.name)
+        if self.share:
+            # must be a shared generator
+            machine = None
+        elif len(self.machines) != 1:
+            msg = f"Shared generator {self.name} must have exactly one machine, but has {len(self.machines)}: {', '.join(self.machines)}"
+            raise ClanError(msg)
+        else:
+            machine = self.machines[0]
+        return GeneratorKey(machine=machine, name=self.name)
 
     def __hash__(self) -> int:
         return hash(self.key)
@@ -143,7 +151,7 @@ class Generator:
         files_selector = "config.clan.core.vars.generators.*.files.*.{secret,deploy,owner,group,mode,neededFor}"
         flake.precache(cls.get_machine_selectors(machine_names))
 
-        generators = []
+        generators: list[Generator] = []
         shared_generators_raw: dict[
             str, tuple[str, dict, dict]
         ] = {}  # name -> (machine_name, gen_data, files_data)
@@ -244,14 +252,26 @@ class Generator:
                     migrate_fact=gen_data.get("migrateFact"),
                     validation_hash=gen_data.get("validationHash"),
                     prompts=prompts,
-                    # only set machine for machine-specific generators
-                    # this is essential for the graph algorithms to work correctly
-                    machine=None if share else machine_name,
+                    # shared generators can have multiple machines, machine-specific have one
+                    machines=[machine_name],
                     _flake=flake,
                     _public_store=pub_store,
                     _secret_store=sec_store,
                 )
 
+                if share:
+                    # For shared generators, check if we already created it
+                    existing = next(
+                        (g for g in generators if g.name == gen_name and g.share), None
+                    )
+                    if existing:
+                        # Just append the machine to the existing generator
+                        existing.machines.append(machine_name)
+                    else:
+                        # Add the new shared generator
+                        generators.append(generator)
+                else:
+                    # Always add per-machine generators
                     generators.append(generator)
 
             # TODO: This should be done in a non-mutable way.

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

@@ -49,28 +49,28 @@ def test_required_generators() -> None:
     gen_1 = Generator(
         name="gen_1",
         dependencies=[],
-        machine=machine_name,
+        machines=[machine_name],
         _public_store=public_store,
         _secret_store=secret_store,
     )
     gen_2 = Generator(
         name="gen_2",
         dependencies=[gen_1.key],
-        machine=machine_name,
+        machines=[machine_name],
         _public_store=public_store,
         _secret_store=secret_store,
     )
     gen_2a = Generator(
         name="gen_2a",
         dependencies=[gen_2.key],
-        machine=machine_name,
+        machines=[machine_name],
         _public_store=public_store,
         _secret_store=secret_store,
     )
     gen_2b = Generator(
         name="gen_2b",
         dependencies=[gen_2.key],
-        machine=machine_name,
+        machines=[machine_name],
         _public_store=public_store,
         _secret_store=secret_store,
     )
@@ -118,21 +118,22 @@ def test_shared_generator_invalidates_multiple_machines_dependents() -> None:
     shared_gen = Generator(
         name="shared_gen",
         dependencies=[],
-        machine=None,  # Shared generator
+        share=True,  # Mark as shared generator
+        machines=[machine_1, machine_2],  # Shared across both machines
         _public_store=public_store,
         _secret_store=secret_store,
     )
     gen_1 = Generator(
         name="gen_1",
         dependencies=[shared_gen.key],
-        machine=machine_1,
+        machines=[machine_1],
         _public_store=public_store,
         _secret_store=secret_store,
     )
     gen_2 = Generator(
         name="gen_2",
         dependencies=[shared_gen.key],
-        machine=machine_2,
+        machines=[machine_2],
         _public_store=public_store,
         _secret_store=secret_store,
     )

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

@@ -181,10 +181,10 @@ def run_generators(
     flake = machines[0].flake
 
     def get_generator_machine(generator: Generator) -> Machine:
-        if generator.machine is None:
-            # return first machine if generator is not tied to a specific one
+        if generator.share:
+            # return first machine if generator is shared
             return machines[0]
-        return Machine(name=generator.machine, flake=flake)
+        return Machine(name=generator.machines[0], flake=flake)
 
     # preheat the select cache, to reduce repeated calls during execution
     selectors = []