From fd5d7934a09da1de18025edf2c89ea4bf72ea85d Mon Sep 17 00:00:00 2001 From: DavHau Date: Thu, 9 Oct 2025 15:40:51 +0700 Subject: [PATCH] vars: refactor - make shared generators carry machines list This should make it simpler to improve the implementation of granting a new machine access to a shared secret. The current approach using the health_check is pretty hacky --- pkgs/clan-cli/clan_cli/tests/test_vars.py | 58 ++++++++++++----------- pkgs/clan-cli/clan_cli/vars/_types.py | 17 ++++--- pkgs/clan-cli/clan_cli/vars/generator.py | 34 ++++++++++--- pkgs/clan-cli/clan_cli/vars/graph_test.py | 15 +++--- pkgs/clan-cli/clan_lib/vars/generate.py | 6 +-- 5 files changed, 78 insertions(+), 52 deletions(-) diff --git a/pkgs/clan-cli/clan_cli/tests/test_vars.py b/pkgs/clan-cli/clan_cli/tests/test_vars.py index 3127f8723..50c20cd6d 100644 --- a/pkgs/clan-cli/clan_cli/tests/test_vars.py +++ b/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() diff --git a/pkgs/clan-cli/clan_cli/vars/_types.py b/pkgs/clan-cli/clan_cli/vars/_types.py index ff50102c5..b6fe75e35 100644 --- a/pkgs/clan-cli/clan_cli/vars/_types.py +++ b/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 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") diff --git a/pkgs/clan-cli/clan_cli/vars/generator.py b/pkgs/clan-cli/clan_cli/vars/generator.py index 454584b22..3daf33392 100644 --- a/pkgs/clan-cli/clan_cli/vars/generator.py +++ b/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,15 +252,27 @@ 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, ) - generators.append(generator) + 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. if include_previous_values: diff --git a/pkgs/clan-cli/clan_cli/vars/graph_test.py b/pkgs/clan-cli/clan_cli/vars/graph_test.py index 8860a21ab..ef3af73fc 100644 --- a/pkgs/clan-cli/clan_cli/vars/graph_test.py +++ b/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, ) diff --git a/pkgs/clan-cli/clan_lib/vars/generate.py b/pkgs/clan-cli/clan_lib/vars/generate.py index 5e0c78b27..e4ccf281e 100644 --- a/pkgs/clan-cli/clan_lib/vars/generate.py +++ b/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 = []