Convert exec arguments to type-safe clap derive

This commit is contained in:
i1i1 2023-08-10 01:55:11 +03:00 committed by Zhaofeng Li
parent 5b8611744b
commit 367c253a47

View file

@ -2,62 +2,58 @@ use std::env;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
use clap::{value_parser, Arg, ArgMatches, Command as ClapCommand, FromArgMatches}; use clap::{ArgMatches, Args, Command as ClapCommand, FromArgMatches};
use futures::future::join_all; use futures::future::join_all;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use crate::error::ColmenaError; use crate::error::ColmenaError;
use crate::job::{JobMonitor, JobState, JobType}; use crate::job::{JobMonitor, JobState, JobType};
use crate::nix::hive::HiveArgs; use crate::nix::hive::HiveArgs;
use crate::nix::NodeFilter; use crate::nix::node_filter::NodeFilterOpts;
use crate::progress::SimpleProgressOutput; use crate::progress::SimpleProgressOutput;
use crate::util; use crate::util;
pub fn subcommand() -> ClapCommand { #[derive(Debug, Args)]
let command = ClapCommand::new("exec") #[command(name = "exec", about = "Run a command on remote machines")]
.about("Run a command on remote machines") struct Opts {
.arg( #[arg(
Arg::new("parallel") short,
.short('p') long,
.long("parallel") default_value_t = 0,
.value_name("LIMIT") value_name = "LIMIT",
.help("Deploy parallelism limit") help = "Deploy parallelism limit",
.long_help( long_help = r#"Limits the maximum number of hosts to run the command in parallel.
r#"Limits the maximum number of hosts to run the command in parallel.
In `colmena exec`, the parallelism limit is disabled (0) by default. In `colmena exec`, the parallelism limit is disabled (0) by default.
"#, "#
) )]
.default_value("0") parallel: usize,
.num_args(1) #[arg(
.value_parser(value_parser!(usize)), short,
) long,
.arg( help = "Be verbose",
Arg::new("verbose") long_help = "Deactivates the progress spinner and prints every line of output."
.short('v') )]
.long("verbose") verbose: bool,
.help("Be verbose") #[command(flatten)]
.long_help("Deactivates the progress spinner and prints every line of output.") nodes: NodeFilterOpts,
.num_args(0), #[arg(
) trailing_var_arg = true,
.arg( required = true,
Arg::new("command") value_name = "COMMAND",
.value_name("COMMAND") help = "Command",
.trailing_var_arg(true) long_help = r#"Command to run
.help("Command")
.required(true)
.num_args(1..)
.long_help(
r#"Command to run
It's recommended to use -- to separate Colmena options from the command to run. For example: It's recommended to use -- to separate Colmena options from the command to run. For example:
colmena exec --on @routers -- tcpdump -vni any ip[9] == 89 colmena exec --on @routers -- tcpdump -vni any ip[9] == 89
"#, "#
), )]
); command: Vec<String>,
}
util::register_selector_args(command) pub fn subcommand() -> ClapCommand {
Opts::augment_args(ClapCommand::new("exec"))
} }
pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(), ColmenaError> { pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(), ColmenaError> {
@ -68,33 +64,24 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
.unwrap(); .unwrap();
let ssh_config = env::var("SSH_CONFIG_FILE").ok().map(PathBuf::from); let ssh_config = env::var("SSH_CONFIG_FILE").ok().map(PathBuf::from);
// FIXME: Just get_one::<NodeFilter> let Opts {
let filter = local_args parallel,
.get_one::<String>("on") verbose,
.map(NodeFilter::new) nodes,
.transpose()?; command,
} = Opts::from_arg_matches(local_args).unwrap();
let mut targets = hive.select_nodes(filter, ssh_config, true).await?; let mut targets = hive.select_nodes(nodes.on, ssh_config, true).await?;
let parallel_sp = Arc::new({ let parallel_sp = Arc::new(if parallel > 0 {
let limit = local_args.get_one::<usize>("parallel").unwrap().to_owned(); Some(Semaphore::new(parallel))
} else {
if limit > 0 { None
Some(Semaphore::new(limit))
} else {
None
}
}); });
let command: Arc<Vec<String>> = Arc::new( let command = Arc::new(command);
local_args
.get_many::<String>("command")
.unwrap()
.cloned()
.collect(),
);
let mut output = SimpleProgressOutput::new(local_args.get_flag("verbose")); let mut output = SimpleProgressOutput::new(verbose);
let (mut monitor, meta) = JobMonitor::new(output.get_sender()); let (mut monitor, meta) = JobMonitor::new(output.get_sender());
@ -107,7 +94,7 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
for (name, target) in targets.drain() { for (name, target) in targets.drain() {
let parallel_sp = parallel_sp.clone(); let parallel_sp = parallel_sp.clone();
let command = command.clone(); let command = Arc::clone(&command);
let mut host = target.into_host().unwrap(); let mut host = target.into_host().unwrap();