From f161c7ca58f7e01772217078371213028d56409e Mon Sep 17 00:00:00 2001 From: DavHau Date: Fri, 6 Sep 2024 14:30:23 +0200 Subject: [PATCH 1/2] vars: improve generator execution pipeline - ensure all dependents are re-generated as well - refactor: separate out computation of generator update closure --- pkgs/clan-cli/clan_cli/machines/machines.py | 16 +++ pkgs/clan-cli/clan_cli/vars/_types.py | 1 + pkgs/clan-cli/clan_cli/vars/generate.py | 105 ++++++++++++++------ pkgs/clan-cli/clan_cli/vars/list.py | 16 ++- pkgs/clan-cli/tests/test_vars.py | 4 +- 5 files changed, 104 insertions(+), 38 deletions(-) diff --git a/pkgs/clan-cli/clan_cli/machines/machines.py b/pkgs/clan-cli/clan_cli/machines/machines.py index 1ccaeb42e..7302e2b4d 100644 --- a/pkgs/clan-cli/clan_cli/machines/machines.py +++ b/pkgs/clan-cli/clan_cli/machines/machines.py @@ -1,6 +1,10 @@ +from __future__ import annotations + +import importlib import json import logging from dataclasses import dataclass, field +from functools import cached_property from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any, Literal @@ -10,6 +14,8 @@ from clan_cli.cmd import run_no_stdout from clan_cli.errors import ClanError from clan_cli.nix import nix_build, nix_config, nix_eval, nix_metadata from clan_cli.ssh import Host, parse_deployment_address +from clan_cli.vars.public_modules import FactStoreBase +from clan_cli.vars.secret_modules import SecretStoreBase log = logging.getLogger(__name__) @@ -90,6 +96,16 @@ class Machine: def public_vars_module(self) -> str: return self.deployment["vars"]["publicModule"] + @cached_property + def secret_vars_store(self) -> SecretStoreBase: + module = importlib.import_module(self.secret_vars_module) + return module.SecretStore(machine=self) + + @cached_property + def public_vars_store(self) -> FactStoreBase: + module = importlib.import_module(self.public_vars_module) + return module.FactStore(machine=self) + @property def facts_data(self) -> dict[str, dict[str, Any]]: if self.deployment["facts"]["services"]: diff --git a/pkgs/clan-cli/clan_cli/vars/_types.py b/pkgs/clan-cli/clan_cli/vars/_types.py index 99789aa68..37ab5ee5f 100644 --- a/pkgs/clan-cli/clan_cli/vars/_types.py +++ b/pkgs/clan-cli/clan_cli/vars/_types.py @@ -18,6 +18,7 @@ class Prompt: previous_value: str | None = None +# TODO: add flag 'pending' generator needs to be executed @dataclass class Generator: name: str diff --git a/pkgs/clan-cli/clan_cli/vars/generate.py b/pkgs/clan-cli/clan_cli/vars/generate.py index 2a315a6a7..cd04f8a1e 100644 --- a/pkgs/clan-cli/clan_cli/vars/generate.py +++ b/pkgs/clan-cli/clan_cli/vars/generate.py @@ -1,5 +1,4 @@ import argparse -import importlib import logging import os import sys @@ -197,9 +196,9 @@ def execute_generator( return True -def _get_subgraph(graph: dict[str, set], vertex: str) -> dict[str, set]: +def _get_subgraph(graph: dict[str, set], vertices: list[str]) -> dict[str, set]: visited = set() - queue = [vertex] + queue = vertices while queue: vertex = queue.pop(0) if vertex not in visited: @@ -208,50 +207,90 @@ def _get_subgraph(graph: dict[str, set], vertex: str) -> dict[str, set]: return {k: v for k, v in graph.items() if k in visited} +def _dependency_graph( + machine: Machine, entry_nodes: None | list[str] = None +) -> dict[str, set]: + graph = { + gen_name: set(generator["dependencies"]) + for gen_name, generator in machine.vars_generators.items() + } + if entry_nodes: + return _get_subgraph(graph, entry_nodes) + return graph + + +def _reverse_dependency_graph( + machine: Machine, entry_nodes: None | list[str] = None +) -> dict[str, set]: + graph = _dependency_graph(machine) + reverse_graph: dict[str, set] = {gen_name: set() for gen_name in graph} + for gen_name, dependencies in graph.items(): + for dep in dependencies: + reverse_graph[dep].add(gen_name) + if entry_nodes: + return _get_subgraph(reverse_graph, entry_nodes) + return reverse_graph + + +def _required_generators( + machine: Machine, + desired_generators: list[str], +) -> list[str]: + """ + Receives list fo desired generators to update and returns list of required generators to update. + + This is needed because some generators might depend on others, so we need to update them first. + The returned list is sorted topologically. + """ + + dependency_graph = _dependency_graph(machine) + # extract sub-graph if specific generators selected + dependency_graph = _get_subgraph(dependency_graph, desired_generators) + + # check if all dependencies actually exist + for gen_name, dependencies in dependency_graph.items(): + for dep in dependencies: + if dep not in dependency_graph: + msg = f"Generator {gen_name} has a dependency on {dep}, which does not exist" + raise ClanError(msg) + + # ensure that all dependents are regenerated as well as their vars might depend on the current generator + reverse_dependency_graph = _reverse_dependency_graph(machine, desired_generators) + final_graph = _dependency_graph( + machine, entry_nodes=list(reverse_dependency_graph.keys()) + ) + + # process generators in topological order (dependencies first) + sorter = TopologicalSorter(final_graph) + return list(sorter.static_order()) + + def _generate_vars_for_machine( machine: Machine, generator_name: str | None, regenerate: bool, ) -> bool: - secret_vars_module = importlib.import_module(machine.secret_vars_module) - secret_vars_store = secret_vars_module.SecretStore(machine=machine) + return _generate_vars_for_machine_multi( + machine, [generator_name] if generator_name else [], regenerate + ) - public_vars_module = importlib.import_module(machine.public_vars_module) - public_vars_store = public_vars_module.FactStore(machine=machine) +def _generate_vars_for_machine_multi( + machine: Machine, + generator_names: list[str], + regenerate: bool, +) -> bool: machine_updated = False - if generator_name and generator_name not in machine.vars_generators: - generators = list(machine.vars_generators.keys()) - msg = f"Could not find generator with name: {generator_name}. The following generators are available: {generators}" - raise ClanError(msg) - - graph = { - gen_name: set(generator["dependencies"]) - for gen_name, generator in machine.vars_generators.items() - } - - # extract sub-graph if specific generator selected - if generator_name: - graph = _get_subgraph(graph, generator_name) - - # check if all dependencies actually exist - for gen_name, dependencies in graph.items(): - for dep in dependencies: - if dep not in graph: - msg = f"Generator {gen_name} has a dependency on {dep}, which does not exist" - raise ClanError(msg) - - # process generators in topological order - sorter = TopologicalSorter(graph) - for generator_name in sorter.static_order(): + generators_to_update = _required_generators(machine, generator_names) + for generator_name in generators_to_update: assert generator_name is not None machine_updated |= execute_generator( machine=machine, generator_name=generator_name, regenerate=regenerate, - secret_vars_store=secret_vars_store, - public_vars_store=public_vars_store, + secret_vars_store=machine.secret_vars_store, + public_vars_store=machine.public_vars_store, ) if machine_updated: # flush caches to make sure the new secrets are available in evaluation diff --git a/pkgs/clan-cli/clan_cli/vars/list.py b/pkgs/clan-cli/clan_cli/vars/list.py index 7b7bdc8f7..7e677bcea 100644 --- a/pkgs/clan-cli/clan_cli/vars/list.py +++ b/pkgs/clan-cli/clan_cli/vars/list.py @@ -30,8 +30,10 @@ def get_vars(machine: Machine) -> list[Var]: return pub_store.get_all() + sec_store.get_all() -def _get_prompt_value( - machine: Machine, generator: Generator, prompt: Prompt +def _get_previous_value( + machine: Machine, + generator: Generator, + prompt: Prompt, ) -> str | None: if not prompt.has_file: return None @@ -40,10 +42,16 @@ def _get_prompt_value( return pub_store.get( generator.name, prompt.name, shared=generator.share ).decode() + sec_store = secret_store(machine) + if sec_store.exists(generator.name, prompt.name, shared=generator.share): + return sec_store.get( + generator.name, prompt.name, shared=generator.share + ).decode() return None @API.register +# TODO: use machine_name def get_prompts(machine: Machine) -> list[Generator]: generators = [] for gen_name, generator in machine.vars_generators.items(): @@ -61,7 +69,7 @@ def get_prompts(machine: Machine) -> list[Generator]: has_file=prompt["createFile"], generator=gen_name, ) - prompt.previous_value = _get_prompt_value(machine, gen, prompt) + prompt.previous_value = _get_previous_value(machine, gen, prompt) prompts.append(prompt) generators.append(gen) @@ -69,6 +77,8 @@ def get_prompts(machine: Machine) -> list[Generator]: # TODO: Ensure generator dependencies are met (executed in correct order etc.) +# TODO: for missing prompts, default to existing values +# TODO: raise error if mandatory prompt not provided @API.register def set_prompts(machine: Machine, updates: list[GeneratorUpdate]) -> None: for update in updates: diff --git a/pkgs/clan-cli/tests/test_vars.py b/pkgs/clan-cli/tests/test_vars.py index 48d02486a..b173eabcb 100644 --- a/pkgs/clan-cli/tests/test_vars.py +++ b/pkgs/clan-cli/tests/test_vars.py @@ -27,12 +27,12 @@ def test_get_subgraph() -> None: "c": set(), "d": set(), } - assert _get_subgraph(graph, "a") == { + assert _get_subgraph(graph, ["a"]) == { "a": {"b", "c"}, "b": {"c"}, "c": set(), } - assert _get_subgraph(graph, "b") == {"b": {"c"}, "c": set()} + assert _get_subgraph(graph, ["b"]) == {"b": {"c"}, "c": set()} def test_dependencies_as_files() -> None: From 9d86c5474f83fb893c2f5eef6c6f9dc7b7e3a2f1 Mon Sep 17 00:00:00 2001 From: DavHau Date: Fri, 6 Sep 2024 15:24:50 +0200 Subject: [PATCH 2/2] clan-cli: fix cyclic import problem --- pkgs/clan-app/tests/conftest.py | 2 ++ pkgs/clan-cli/clan_cli/__init__.py | 4 ++-- pkgs/clan-cli/clan_cli/machines/machines.py | 2 -- pkgs/clan-cli/clan_cli/vars/_types.py | 5 ++--- pkgs/clan-cli/clan_cli/vars/{__init__.py => cli.py} | 0 5 files changed, 6 insertions(+), 7 deletions(-) rename pkgs/clan-cli/clan_cli/vars/{__init__.py => cli.py} (100%) diff --git a/pkgs/clan-app/tests/conftest.py b/pkgs/clan-app/tests/conftest.py index 9f107889f..530f4ad1c 100644 --- a/pkgs/clan-app/tests/conftest.py +++ b/pkgs/clan-app/tests/conftest.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import subprocess from pathlib import Path diff --git a/pkgs/clan-cli/clan_cli/__init__.py b/pkgs/clan-cli/clan_cli/__init__.py index a8e423349..80cd932dc 100644 --- a/pkgs/clan-cli/clan_cli/__init__.py +++ b/pkgs/clan-cli/clan_cli/__init__.py @@ -20,7 +20,6 @@ from . import ( history, secrets, state, - vars, vms, ) from .clan_uri import FlakeId @@ -32,6 +31,7 @@ from .hyperlink import help_hyperlink from .machines import cli as machines from .profiler import profile from .ssh import cli as ssh_cli +from .vars import cli as vars_cli log = logging.getLogger(__name__) @@ -293,7 +293,7 @@ For more detailed information, visit: {help_hyperlink("secrets", "https://docs.c ), formatter_class=argparse.RawTextHelpFormatter, ) - vars.register_parser(parser_vars) + vars_cli.register_parser(parser_vars) parser_machine = subparsers.add_parser( "machines", diff --git a/pkgs/clan-cli/clan_cli/machines/machines.py b/pkgs/clan-cli/clan_cli/machines/machines.py index 7302e2b4d..9eb2b2849 100644 --- a/pkgs/clan-cli/clan_cli/machines/machines.py +++ b/pkgs/clan-cli/clan_cli/machines/machines.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import importlib import json import logging diff --git a/pkgs/clan-cli/clan_cli/vars/_types.py b/pkgs/clan-cli/clan_cli/vars/_types.py index 37ab5ee5f..b09286945 100644 --- a/pkgs/clan-cli/clan_cli/vars/_types.py +++ b/pkgs/clan-cli/clan_cli/vars/_types.py @@ -1,11 +1,10 @@ -# !/usr/bin/env python3 import json import shutil from abc import ABC, abstractmethod from dataclasses import dataclass from pathlib import Path -from clan_cli.machines.machines import Machine +from clan_cli.machines import machines @dataclass @@ -73,7 +72,7 @@ class Var: class StoreBase(ABC): - def __init__(self, machine: Machine) -> None: + def __init__(self, machine: "machines.Machine") -> None: self.machine = machine @property diff --git a/pkgs/clan-cli/clan_cli/vars/__init__.py b/pkgs/clan-cli/clan_cli/vars/cli.py similarity index 100% rename from pkgs/clan-cli/clan_cli/vars/__init__.py rename to pkgs/clan-cli/clan_cli/vars/cli.py