churn-v2/crates/churn/src/cli.rs

use std::{collections::BTreeMap, net::SocketAddr};

use clap::{Parser, Subcommand};

use crate::{agent, server};

pub async fn execute() -> anyhow::Result<()> {
    let cli = Command::parse();
    match cli.command.expect("to have a subcommand") {
        Commands::Serve {
            host,
            grpc_host,
            config,
        } => {
            tracing::info!("Starting service");
            server::execute(host, grpc_host, config).await?;
        }
        Commands::Agent { commands } => match commands {
            AgentCommands::Start {} => {
                tracing::info!("starting agent");
                agent::execute().await?;
                tracing::info!("shut down agent");
            }
            AgentCommands::Setup {
                force,
                discovery,
                labels,
            } => {
                let mut setup_labels = BTreeMap::new();
                for (k, v) in labels {
                    setup_labels.insert(k, v);
                }

                if !setup_labels.contains_key("node_name") {
                    setup_labels.insert(
                        "node_name".into(),
                        petname::petname(2, "-").expect("to be able to generate a valid petname"),
                    );
                }

                agent::setup_config(discovery, force, setup_labels).await?;
                tracing::info!("wrote default agent config");
            }
        },
    }

    Ok(())
}

#[derive(Parser)]
#[command(author, version, about, long_about = None, subcommand_required = true)]
struct Command {
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    Serve {
        #[arg(env = "SERVICE_HOST", long, default_value = "127.0.0.1:3000")]
        host: SocketAddr,

        #[arg(env = "SERVICE_GRPC_HOST", long, default_value = "127.0.0.1:7900")]
        grpc_host: SocketAddr,

        #[clap(flatten)]
        config: server::config::ServerConfig,
    },
    Agent {
        #[command(subcommand)]
        commands: AgentCommands,
    },
}

#[derive(Subcommand)]
enum AgentCommands {
    Start {},
    Setup {
        #[arg(long, default_value = "false")]
        force: bool,

        #[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()))
}