forked from DGNum/colmena
Convert exec arguments to type-safe clap derive
This commit is contained in:
parent
5b8611744b
commit
367c253a47
1 changed files with 50 additions and 63 deletions
|
@ -2,62 +2,58 @@ use std::env;
|
|||
use std::path::PathBuf;
|
||||
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 tokio::sync::Semaphore;
|
||||
|
||||
use crate::error::ColmenaError;
|
||||
use crate::job::{JobMonitor, JobState, JobType};
|
||||
use crate::nix::hive::HiveArgs;
|
||||
use crate::nix::NodeFilter;
|
||||
use crate::nix::node_filter::NodeFilterOpts;
|
||||
use crate::progress::SimpleProgressOutput;
|
||||
use crate::util;
|
||||
|
||||
pub fn subcommand() -> ClapCommand {
|
||||
let command = ClapCommand::new("exec")
|
||||
.about("Run a command on remote machines")
|
||||
.arg(
|
||||
Arg::new("parallel")
|
||||
.short('p')
|
||||
.long("parallel")
|
||||
.value_name("LIMIT")
|
||||
.help("Deploy parallelism limit")
|
||||
.long_help(
|
||||
r#"Limits the maximum number of hosts to run the command in parallel.
|
||||
#[derive(Debug, Args)]
|
||||
#[command(name = "exec", about = "Run a command on remote machines")]
|
||||
struct Opts {
|
||||
#[arg(
|
||||
short,
|
||||
long,
|
||||
default_value_t = 0,
|
||||
value_name = "LIMIT",
|
||||
help = "Deploy parallelism limit",
|
||||
long_help = r#"Limits the maximum number of hosts to run the command in parallel.
|
||||
|
||||
In `colmena exec`, the parallelism limit is disabled (0) by default.
|
||||
"#,
|
||||
)
|
||||
.default_value("0")
|
||||
.num_args(1)
|
||||
.value_parser(value_parser!(usize)),
|
||||
)
|
||||
.arg(
|
||||
Arg::new("verbose")
|
||||
.short('v')
|
||||
.long("verbose")
|
||||
.help("Be verbose")
|
||||
.long_help("Deactivates the progress spinner and prints every line of output.")
|
||||
.num_args(0),
|
||||
)
|
||||
.arg(
|
||||
Arg::new("command")
|
||||
.value_name("COMMAND")
|
||||
.trailing_var_arg(true)
|
||||
.help("Command")
|
||||
.required(true)
|
||||
.num_args(1..)
|
||||
.long_help(
|
||||
r#"Command to run
|
||||
"#
|
||||
)]
|
||||
parallel: usize,
|
||||
#[arg(
|
||||
short,
|
||||
long,
|
||||
help = "Be verbose",
|
||||
long_help = "Deactivates the progress spinner and prints every line of output."
|
||||
)]
|
||||
verbose: bool,
|
||||
#[command(flatten)]
|
||||
nodes: NodeFilterOpts,
|
||||
#[arg(
|
||||
trailing_var_arg = true,
|
||||
required = true,
|
||||
value_name = "COMMAND",
|
||||
help = "Command",
|
||||
long_help = r#"Command to run
|
||||
|
||||
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
|
||||
"#,
|
||||
),
|
||||
);
|
||||
"#
|
||||
)]
|
||||
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> {
|
||||
|
@ -68,33 +64,24 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
|
|||
.unwrap();
|
||||
let ssh_config = env::var("SSH_CONFIG_FILE").ok().map(PathBuf::from);
|
||||
|
||||
// FIXME: Just get_one::<NodeFilter>
|
||||
let filter = local_args
|
||||
.get_one::<String>("on")
|
||||
.map(NodeFilter::new)
|
||||
.transpose()?;
|
||||
let Opts {
|
||||
parallel,
|
||||
verbose,
|
||||
nodes,
|
||||
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 limit = local_args.get_one::<usize>("parallel").unwrap().to_owned();
|
||||
|
||||
if limit > 0 {
|
||||
Some(Semaphore::new(limit))
|
||||
let parallel_sp = Arc::new(if parallel > 0 {
|
||||
Some(Semaphore::new(parallel))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
let command: Arc<Vec<String>> = Arc::new(
|
||||
local_args
|
||||
.get_many::<String>("command")
|
||||
.unwrap()
|
||||
.cloned()
|
||||
.collect(),
|
||||
);
|
||||
let command = Arc::new(command);
|
||||
|
||||
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());
|
||||
|
||||
|
@ -107,7 +94,7 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
|
|||
|
||||
for (name, target) in targets.drain() {
|
||||
let parallel_sp = parallel_sp.clone();
|
||||
let command = command.clone();
|
||||
let command = Arc::clone(&command);
|
||||
|
||||
let mut host = target.into_host().unwrap();
|
||||
|
||||
|
|
Loading…
Reference in a new issue