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::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();
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue