From 4faab0a20f933618542177905cde84519ec76002 Mon Sep 17 00:00:00 2001
From: a-kenji <aks.kenji@protonmail.com>
Date: Sun, 26 May 2024 15:21:54 +0200
Subject: [PATCH] clan: `facts generate` specific service

Add `--service` flag to the `clan` cli which allows specifying a certain
service to be generated.

Example:

```
clan facts generate [MACHINE] --service [SERVICE]
```

Fixes #1395
---
 pkgs/clan-cli/clan_cli/facts/generate.py   | 35 ++++++++++++++++++----
 pkgs/clan-cli/clan_cli/machines/install.py |  2 +-
 pkgs/clan-cli/clan_cli/machines/update.py  |  2 +-
 pkgs/clan-cli/clan_cli/vms/run.py          |  2 +-
 4 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/pkgs/clan-cli/clan_cli/facts/generate.py b/pkgs/clan-cli/clan_cli/facts/generate.py
index c1da81f3..2dee7282 100644
--- a/pkgs/clan-cli/clan_cli/facts/generate.py
+++ b/pkgs/clan-cli/clan_cli/facts/generate.py
@@ -134,7 +134,10 @@ def prompt_func(text: str) -> str:
 
 
 def _generate_facts_for_machine(
-    machine: Machine, tmpdir: Path, prompt: Callable[[str], str] = prompt_func
+    machine: Machine,
+    service: str | None,
+    tmpdir: Path,
+    prompt: Callable[[str], str] = prompt_func,
 ) -> bool:
     local_temp = tmpdir / machine.name
     local_temp.mkdir()
@@ -145,7 +148,19 @@ def _generate_facts_for_machine(
     public_facts_store = public_facts_module.FactStore(machine=machine)
 
     machine_updated = False
-    for service in machine.facts_data:
+
+    if service and service not in machine.facts_data:
+        services = list(machine.facts_data.keys())
+        raise ClanError(
+            f"Could not find service with name: {service}. The following services are available: {services}"
+        )
+
+    if service:
+        machine_service_facts = {service: machine.facts_data[service]}
+    else:
+        machine_service_facts = machine.facts_data
+
+    for service in machine_service_facts:
         machine_updated |= generate_service_facts(
             machine=machine,
             service=service,
@@ -161,7 +176,9 @@ def _generate_facts_for_machine(
 
 
 def generate_facts(
-    machines: list[Machine], prompt: Callable[[str], str] = prompt_func
+    machines: list[Machine],
+    service: str | None,
+    prompt: Callable[[str], str] = prompt_func,
 ) -> bool:
     was_regenerated = False
     with TemporaryDirectory() as tmp:
@@ -170,7 +187,9 @@ def generate_facts(
         for machine in machines:
             errors = 0
             try:
-                was_regenerated |= _generate_facts_for_machine(machine, tmpdir, prompt)
+                was_regenerated |= _generate_facts_for_machine(
+                    machine, service, tmpdir, prompt
+                )
             except Exception as exc:
                 log.error(f"Failed to generate facts for {machine.name}: {exc}")
                 errors += 1
@@ -189,7 +208,7 @@ def generate_command(args: argparse.Namespace) -> None:
         machines = get_all_machines(args.flake)
     else:
         machines = get_selected_machines(args.flake, args.machines)
-    generate_facts(machines)
+    generate_facts(machines, args.service)
 
 
 def register_generate_parser(parser: argparse.ArgumentParser) -> None:
@@ -200,4 +219,10 @@ def register_generate_parser(parser: argparse.ArgumentParser) -> None:
         nargs="*",
         default=[],
     )
+    parser.add_argument(
+        "--service",
+        type=str,
+        help="service to generate facts for, if empty, generate facts for every service",
+        default=None,
+    )
     parser.set_defaults(func=generate_command)
diff --git a/pkgs/clan-cli/clan_cli/machines/install.py b/pkgs/clan-cli/clan_cli/machines/install.py
index fe9c962a..cdea55c8 100644
--- a/pkgs/clan-cli/clan_cli/machines/install.py
+++ b/pkgs/clan-cli/clan_cli/machines/install.py
@@ -35,7 +35,7 @@ def install_nixos(
     target_host = f"{h.user or 'root'}@{h.host}"
     log.info(f"target host: {target_host}")
 
-    generate_facts([machine])
+    generate_facts([machine], None)
 
     with TemporaryDirectory() as tmpdir_:
         tmpdir = Path(tmpdir_)
diff --git a/pkgs/clan-cli/clan_cli/machines/update.py b/pkgs/clan-cli/clan_cli/machines/update.py
index 34c31a25..177c6a2d 100644
--- a/pkgs/clan-cli/clan_cli/machines/update.py
+++ b/pkgs/clan-cli/clan_cli/machines/update.py
@@ -98,7 +98,7 @@ def deploy_nixos(machines: MachineGroup) -> None:
         env = os.environ.copy()
         env["NIX_SSHOPTS"] = ssh_arg
 
-        generate_facts([machine])
+        generate_facts([machine], None)
         upload_secrets(machine)
 
         path = upload_sources(".", target)
diff --git a/pkgs/clan-cli/clan_cli/vms/run.py b/pkgs/clan-cli/clan_cli/vms/run.py
index f1ff7c8a..69520763 100644
--- a/pkgs/clan-cli/clan_cli/vms/run.py
+++ b/pkgs/clan-cli/clan_cli/vms/run.py
@@ -69,7 +69,7 @@ def get_secrets(
     secret_facts_module = importlib.import_module(machine.secret_facts_module)
     secret_facts_store = secret_facts_module.SecretStore(machine=machine)
 
-    generate_facts([machine])
+    generate_facts([machine], None)
 
     secret_facts_store.upload(secrets_dir)
     return secrets_dir