feat: add labels to config

Signed-off-by: kjuulh <contact@kjuulh.io>

crates/churn/src/agent/actions.rs

@@ -1,6 +1,6 @@
-use anyhow::Context;
+use apt::AptTask;
 
-use super::task::{IntoTask, Task};
+use super::task::IntoTask;
 
 pub struct Plan {}
 impl Plan {
@@ -13,48 +13,4 @@ impl Plan {
     }
 }
 
-pub struct AptTask {}
-
-impl AptTask {
-    pub fn new() -> Self {
-        Self {}
-    }
-}
-
-#[async_trait::async_trait]
-impl Task for AptTask {
-    fn id(&self) -> String {
-        "apt".into()
-    }
-
-    async fn execute(&self) -> anyhow::Result<()> {
-        let mut cmd = tokio::process::Command::new("apt-get");
-        cmd.args(["update", "-q"]);
-        let output = cmd.output().await.context("failed to run apt update")?;
-        match output.status.success() {
-            true => tracing::info!("successfully ran apt update"),
-            false => {
-                anyhow::bail!(
-                    "failed to run apt update: {}",
-                    std::str::from_utf8(&output.stderr)?
-                );
-            }
-        }
-
-        let mut cmd = tokio::process::Command::new("apt-get");
-        cmd.env("DEBIAN_FRONTEND", "noninteractive")
-            .args(["upgrade", "-y"]);
-        let output = cmd.output().await.context("failed to run apt upgrade")?;
-        match output.status.success() {
-            true => tracing::info!("successfully ran apt upgrade"),
-            false => {
-                anyhow::bail!(
-                    "failed to run apt upgrade: {}",
-                    std::str::from_utf8(&output.stderr)?
-                );
-            }
-        }
-
-        Ok(())
-    }
-}
+pub mod apt;

crates/churn/src/agent/actions/apt.rs

@@ -0,0 +1,49 @@
+use anyhow::Context;
+
+use crate::agent::task::Task;
+
+pub struct AptTask {}
+
+impl AptTask {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+#[async_trait::async_trait]
+impl Task for AptTask {
+    fn id(&self) -> String {
+        "apt".into()
+    }
+
+    async fn execute(&self) -> anyhow::Result<()> {
+        let mut cmd = tokio::process::Command::new("apt-get");
+        cmd.args(["update", "-q"]);
+        let output = cmd.output().await.context("failed to run apt update")?;
+        match output.status.success() {
+            true => tracing::info!("successfully ran apt update"),
+            false => {
+                anyhow::bail!(
+                    "failed to run apt update: {}",
+                    std::str::from_utf8(&output.stderr)?
+                );
+            }
+        }
+
+        let mut cmd = tokio::process::Command::new("apt-get");
+        cmd.env("DEBIAN_FRONTEND", "noninteractive")
+            .args(["upgrade", "-y"]);
+        let output = cmd.output().await.context("failed to run apt upgrade")?;
+        match output.status.success() {
+            true => tracing::info!("successfully ran apt upgrade"),
+            false => {
+                anyhow::bail!(
+                    "failed to run apt upgrade: {}",
+                    std::str::from_utf8(&output.stderr)?
+                );
+            }
+        }
+
+        Ok(())
+    }
+}

crates/churn/src/agent/config.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use anyhow::Context;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
@@ -23,6 +25,8 @@ impl AgentConfig {
 struct ConfigFile {
     agent_id: String,
     discovery: String,
+
+    labels: BTreeMap<String, String>,
 }
 
 impl ConfigFile {
@@ -43,10 +47,15 @@ impl ConfigFile {
         toml::from_str(&contents).context("failed to parse the contents of the churn agent config")
     }
 
-    pub async fn write_default(discovery: impl Into<String>, force: bool) -> anyhow::Result<Self> {
+    pub async fn write_default(
+        discovery: impl Into<String>,
+        force: bool,
+        labels: impl Into<BTreeMap<String, String>>,
+    ) -> anyhow::Result<Self> {
         let s = Self {
             agent_id: Uuid::new_v4().to_string(),
             discovery: discovery.into(),
+            labels: labels.into(),
         };
 
         let directory = dirs::data_dir()
@@ -73,8 +82,12 @@ impl ConfigFile {
     }
 }
 
-pub async fn setup_config(discovery: impl Into<String>, force: bool) -> anyhow::Result<()> {
-    ConfigFile::write_default(discovery, force).await?;
+pub async fn setup_config(
+    discovery: impl Into<String>,
+    force: bool,
+    labels: impl Into<BTreeMap<String, String>>,
+) -> anyhow::Result<()> {
+    ConfigFile::write_default(discovery, force, labels).await?;
 
     Ok(())
 }

crates/churn/src/cli.rs

@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::{collections::BTreeMap, net::SocketAddr};
 
 use clap::{Parser, Subcommand};
 
@@ -21,8 +21,17 @@ pub async fn execute() -> anyhow::Result<()> {
                 agent::execute().await?;
                 tracing::info!("shut down agent");
             }
-            AgentCommands::Setup { force, discovery } => {
-                agent::setup_config(discovery, force).await?;
+            AgentCommands::Setup {
+                force,
+                discovery,
+                labels,
+            } => {
+                let mut setup_labels = BTreeMap::new();
+                for (k, v) in labels {
+                    setup_labels.insert(k, v);
+                }
+
+                agent::setup_config(discovery, force, setup_labels).await?;
                 tracing::info!("wrote default agent config");
             }
         },
@@ -65,5 +74,15 @@ enum AgentCommands {
 
         #[arg(env = "DISCOVERY_HOST", long = "discovery")]
         discovery: String,
+
+        #[arg(long = "label", short = 'l', value_parser = parse_key_val, action = clap::ArgAction::Append)]
+        labels: Vec<(String, String)>,
     },
 }
+
+fn parse_key_val(s: &str) -> Result<(String, String), String> {
+    let (key, value) = s
+        .split_once("=")
+        .ok_or_else(|| format!("invalid key=value: no `=` found in `{s}`"))?;
+    Ok((key.to_string(), value.to_string()))
+}