cli: match against node class for client commands

For `deploy` and `ssh` commands, a new `--class` argument is added
that allows selecting only clients in a given Nomad node_class.

Author: nrdxp

nix/cli/packages/cli/src/cli.rs

@@ -29,6 +29,7 @@ pub(crate) async fn ssh(sub: &ArgMatches, cluster: ClusterHandle) -> Result<()>
     let mut args: Vec<String> = sub.get_many("args").unwrap_or_default().cloned().collect();
     let job: Vec<String> = sub.get_many("job").unwrap_or_default().cloned().collect();
     let delay = Duration::from_secs(*sub.get_one::<u64>("delay").unwrap_or(&0));
+    let node_class = sub.get_one::<String>("class").cloned();
 
     let namespace = sub
         .get_one::<String>("namespace")
@@ -41,7 +42,7 @@ pub(crate) async fn ssh(sub: &ArgMatches, cluster: ClusterHandle) -> Result<()>
 
     if sub.is_present("all") {
         let nodes = if sub.is_present("clients") {
-            cluster.nodes.find_clients()
+            cluster.nodes.find_clients(node_class)
         } else {
             cluster.nodes
         };
@@ -58,7 +59,7 @@ pub(crate) async fn ssh(sub: &ArgMatches, cluster: ClusterHandle) -> Result<()>
         return Ok(());
     } else if sub.is_present("parallel") {
         let nodes = if sub.is_present("clients") {
-            cluster.nodes.find_clients()
+            cluster.nodes.find_clients(node_class)
         } else {
             cluster.nodes
         };
@@ -143,11 +144,12 @@ async fn init_ssh(ip: IpAddr, args: Vec<String>, cluster: String) -> Result<()>
 pub(crate) async fn deploy(sub: &ArgMatches, cluster: ClusterHandle) -> Result<()> {
     let opts = <subs::Deploy as FromArgMatches>::from_arg_matches(sub).unwrap_or_default();
     let cluster = cluster.await??;
+    let node_class = sub.get_one::<String>("class").cloned();
 
     info!("node needles: {:?}", opts.nodes);
 
     let instances = if opts.clients {
-        cluster.nodes.find_clients()
+        cluster.nodes.find_clients(node_class)
     } else {
         cluster
             .nodes
@@ -219,7 +221,6 @@ async fn info_print(cluster: ClusterHandle, json: bool) -> Result<()> {
     if json {
         let stdout = io::stdout();
         let handle = stdout.lock();
-        env::set_var("BITTE_INFO_NO_ALLOCS", "");
         serde_json::to_writer_pretty(handle, &cluster)?;
     } else {
         let mut core_nodes_table = Table::new();
@@ -234,23 +235,14 @@ async fn info_print(cluster: ClusterHandle, json: bool) -> Result<()> {
 
         for node in nodes.into_iter() {
             match node.asg {
-                Some(asg) => {
-                    let name: String = asg.to_string();
-                    // TODO extract true client group
+                Some(_) => {
                     let group: String = {
-                        let suffix = name.split('-').last().unwrap_or_default().to_owned();
-                        let i_type = node
-                            .node_type
-                            .clone()
-                            .unwrap_or_default()
-                            .split('.')
-                            .last()
-                            .unwrap_or_default()
-                            .to_owned();
-                        if suffix == i_type {
-                            "".to_string()
-                        } else {
-                            format!(" ({})", suffix)
+                        match node.nomad_client {
+                            Some(client) => match client.node_class {
+                                Some(class) => format!(" ({})", class),
+                                None => "".to_string(),
+                            },
+                            None => "".to_string(),
                         }
                     };
 

nix/cli/packages/cli/src/cli/subs.rs

@@ -21,6 +21,8 @@ pub struct Info {
     #[clap(short, long)]
     /// output as JSON
     json: bool,
+    #[clap(flatten)]
+    nomad: Nomad,
 }
 
 #[derive(Parser, Default)]
@@ -39,6 +41,13 @@ pub struct Deploy {
     /// nodes to deploy; takes one or more needles to match against:
     /// private & public ip, node name and aws client id
     pub nodes: Vec<String>,
+
+    #[clap(flatten)]
+    nomad: Nomad,
+
+    #[clap(long, short = 'o',  requires_all = &["nomad", "clients"])]
+    /// the Nomad node class to filter clients against
+    class: Option<String>,
 }
 #[derive(Parser)]
 /// Generate completions for the given shell
@@ -79,6 +88,19 @@ struct Globals {
     aws_asg_regions: Option<Vec<String>>,
 }
 
+#[derive(Parser, Default)]
+struct Nomad {
+    #[clap(
+        long,
+        value_name = "TOKEN",
+        env = "NOMAD_TOKEN",
+        value_parser = clap::value_parser!(Uuid),
+        setting = ArgSettings::HideEnvValues
+    )]
+    /// The Nomad token used to query node information
+    nomad: Option<Uuid>,
+}
+
 #[derive(Parser)]
 /// SSH to instances
 pub struct Ssh {
@@ -94,15 +116,11 @@ pub struct Ssh {
     /// specify client by: job, group, alloc_index;
     /// this will also 'cd' to the alloc dir if <ARGS> is empty
     job: Option<String>,
-    #[clap(
-        long,
-        value_name = "TOKEN",
-        env = "NOMAD_TOKEN",
-        value_parser = clap::value_parser!(Uuid),
-        setting = ArgSettings::HideEnvValues
-    )]
-    /// for '-j': The Nomad token used to query node information
-    nomad: Option<Uuid>,
+    #[clap(long, short, env = "NOMAD_NAMESPACE")]
+    /// Nomad namespace to search for jobs with `-j`
+    namespace: Option<String>,
+    #[clap(flatten)]
+    nomad: Nomad,
     #[clap(
         long,
         short,
@@ -121,12 +139,12 @@ pub struct Ssh {
     )]
     /// run <ARGS> on nodes in parallel
     parallel: bool,
-    #[clap(long, short, env = "NOMAD_NAMESPACE")]
-    /// for '-j': specify nomad namespace to search for <JOB>
-    namespace: Option<String>,
     #[clap(long, short = 'l', requires = "multi")]
     /// for '-a' or '-p': execute commands only on Nomad clients
     clients: bool,
+    #[clap(long, short = 'o',  requires_all = &["nomad", "clients"])]
+    /// the Nomad node class to filter clients against
+    class: Option<String>,
     #[clap(long, short, requires = "all")]
     /// for '-a': seconds to delay between commands
     delay: Option<usize>,

nix/cli/packages/cli/src/main.rs

@@ -17,24 +17,18 @@ async fn main() -> Result<()> {
 
     let matches = app.clone().get_matches();
 
-    let run = |sub: &ArgMatches, init_log: bool, token| {
+    let run = |sub: &ArgMatches, init_log: bool| {
         if init_log {
             cli::init_log(matches.occurrences_of("verbose"))
         };
+        let token = sub.get_one::<Uuid>("nomad").copied();
         BitteCluster::init(sub.clone(), token)
     };
 
     match matches.subcommand() {
-        Some(("deploy", sub)) => cli::deploy(sub, run(sub, false, None)).await?,
-        Some(("info", sub)) => cli::info(sub, run(sub, true, None)).await?,
-        Some(("ssh", sub)) => {
-            let token: Option<Uuid> = if sub.is_present("job") {
-                sub.get_one::<Uuid>("nomad").copied()
-            } else {
-                None
-            };
-            cli::ssh(sub, run(sub, true, token)).await?
-        }
+        Some(("deploy", sub)) => cli::deploy(sub, run(sub, false)).await?,
+        Some(("info", sub)) => cli::info(sub, run(sub, true)).await?,
+        Some(("ssh", sub)) => cli::ssh(sub, run(sub, true)).await?,
         Some(("completions", sub)) => {
             if let Some(shell) = sub.get_one::<Shell>("shell").copied() {
                 cli::completions(shell, app).await;

nix/cli/packages/cli/src/types.rs

@@ -8,7 +8,6 @@ use clap::{ArgEnum, ArgMatches};
 use serde::{de::Deserializer, Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::collections::hash_set::HashSet;
-use std::env;
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
 use std::sync::Arc;
@@ -98,6 +97,8 @@ pub struct NomadClient {
     pub allocs: Option<NomadAllocs>,
     #[serde(rename = "Address")]
     pub address: Option<IpAddr>,
+    #[serde(rename = "NodeClass")]
+    pub node_class: Option<String>,
 }
 
 impl NomadClient {
@@ -122,7 +123,7 @@ pub struct BitteNode {
     pub priv_ip: IpAddr,
     pub pub_ip: IpAddr,
     pub nixos: String,
-    #[serde(skip_serializing_if = "skip_info")]
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub nomad_client: Option<NomadClient>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub node_type: Option<String>,
@@ -132,17 +133,13 @@ pub struct BitteNode {
     pub asg: Option<String>,
 }
 
-fn skip_info<T>(_: &Option<T>) -> bool {
-    env::var("BITTE_INFO_NO_ALLOCS").is_ok()
-}
-
 pub trait BitteFind
 where
     Self: IntoIterator,
 {
     fn find_needle(self, needle: &str) -> Result<Self::Item>;
     fn find_needles(self, needles: Vec<&str>) -> Self;
-    fn find_clients(self) -> Self;
+    fn find_clients(self, node_class: Option<String>) -> Self;
     fn find_with_job(
         self,
         name: &str,
@@ -248,8 +245,17 @@ impl BitteFind for BitteNodes {
             .with_context(|| format!("{} does not match any nodes", needle))
     }
 
-    fn find_clients(self) -> Self {
-        self.into_iter().filter(|node| node.asg.is_some()).collect()
+    fn find_clients(self, node_class: Option<String>) -> Self {
+        match node_class {
+            Some(class) => self
+                .into_iter()
+                .filter(|node| match &node.nomad_client {
+                    Some(client) => client.node_class.clone().unwrap_or_default() == class,
+                    None => false,
+                })
+                .collect(),
+            None => self.into_iter().filter(|node| node.asg.is_some()).collect(),
+        }
     }
 
     fn find_needles(self, needles: Vec<&str>) -> Self {