diff --git a/rs/cli/src/commands/subnet/create.rs b/rs/cli/src/commands/subnet/create.rs index 45139d8bc..92254acb5 100644 --- a/rs/cli/src/commands/subnet/create.rs +++ b/rs/cli/src/commands/subnet/create.rs @@ -26,7 +26,7 @@ pub struct Create { #[clap(long, num_args(1..))] pub only: Vec, - #[clap(long = "add-nodes", num_args(1..), help = r#"Add the provided nodes to the subnet. Fails if a node is unavailable/unhealthy."#, visible_aliases = &["add", "add-node", "add-node-id", "add-node-ids"])] + #[clap(long, num_args(1..), help = r#"Add the provided nodes to the subnet. Fails if a node is unavailable/unhealthy."#, visible_aliases = &["add", "add-node", "add-node-id", "add-node-ids"])] pub add_nodes: Vec, /// Motivation for replacing custom nodes diff --git a/rs/cli/src/commands/subnet/replace.rs b/rs/cli/src/commands/subnet/replace.rs index acf83db9a..0ce19f422 100644 --- a/rs/cli/src/commands/subnet/replace.rs +++ b/rs/cli/src/commands/subnet/replace.rs @@ -1,5 +1,6 @@ use clap::{error::ErrorKind, Args}; +use decentralization::network::SubnetQueryBy; use ic_types::PrincipalId; use itertools::Itertools; @@ -11,8 +12,8 @@ use crate::{auth::AuthRequirement, exe::ExecutableCommand, subnet_manager::Subne #[derive(Args, Debug)] pub struct Replace { /// Specific node IDs to remove from the subnet - #[clap(long = "remove-nodes", short, num_args(1..), visible_aliases = &["nodes", "node", "node-id", "node-ids", "remove", "remove-node", "remove-nodes", "remove-node-id", "remove-node-ids"])] - pub nodes: Vec, + #[clap(long, short, num_args(1..), visible_aliases = &["nodes", "node", "node-id", "node-ids", "remove", "remove-node", "remove-nodes", "remove-node-id", "remove-node-ids"])] + pub remove_nodes: Vec, /// Do not replace unhealthy nodes #[clap(long)] @@ -35,12 +36,12 @@ pub struct Replace { pub only: Vec, /// Add specific nodes to the subnet. Fails if a node is unavailable/unhealthy. - #[clap(long = "add-nodes", num_args(1..), visible_aliases = &["add", "add-node", "add-node-id", "add-node-ids"])] + #[clap(long, num_args(1..), visible_aliases = &["add", "add-node", "add-node-id", "add-node-ids"])] pub add_nodes: Vec, - /// The ID of the subnet. - #[clap(long, short, alias = "subnet-id")] - pub id: Option, + /// The ID of the subnet. Must either match the subnet of the provided nodes, or be omitted. + #[clap(long, short, visible_aliases = &["subnet", "id"])] + pub subnet_id: Option, #[clap(flatten)] pub submission_parameters: SubmissionParameters, @@ -52,9 +53,11 @@ impl ExecutableCommand for Replace { } async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> { - let subnet_target = match &self.id { + self.validate_subnet_id_and_nodes(&ctx).await?; + + let subnet_target = match &self.subnet_id { Some(id) => SubnetTarget::FromId(*id), - _ => SubnetTarget::FromNodesIds(self.nodes.clone()), + _ => SubnetTarget::FromNodesIds(self.remove_nodes.clone()), }; let all_nodes = ctx.registry().await.nodes().await?.values().cloned().collect_vec(); @@ -98,21 +101,34 @@ impl ExecutableCommand for Replace { } fn validate(&self, _args: &GlobalArgs, cmd: &mut clap::Command) { - if !self.nodes.is_empty() && self.id.is_some() { - cmd.error( - ErrorKind::ArgumentConflict, - "Both subnet ID and a list of nodes to replace are provided. Only one of the two is allowed.", - ) - .exit() - } else if self.nodes.is_empty() && self.id.is_none() { + if self.remove_nodes.is_empty() && self.subnet_id.is_none() { cmd.error( ErrorKind::MissingRequiredArgument, "Specify either a subnet ID or a list of nodes to replace", ) .exit() - } else if !self.nodes.is_empty() && self.motivation.is_none() { + } else if !self.remove_nodes.is_empty() && self.motivation.is_none() { cmd.error(ErrorKind::MissingRequiredArgument, "Required argument motivation not found") .exit() } } } + +impl Replace { + /// If both a subnet id and nodes to remove are provided, ensure they match + async fn validate_subnet_id_and_nodes(&self, ctx: &crate::ctx::DreContext) -> anyhow::Result<()> { + let nodes_add_or_remove = [self.remove_nodes.clone(), self.add_nodes.clone()].concat(); + if let (Some(expected_subnet_id), false) = (self.subnet_id, nodes_add_or_remove.is_empty()) { + let registry = ctx.registry().await; + let nodes = registry.get_nodes_from_ids(&nodes_add_or_remove).await?; + let subnet = registry + .subnet(SubnetQueryBy::NodeList(nodes)) + .await + .map_err(|e| anyhow::anyhow!(e.to_string()))?; + if subnet.id != expected_subnet_id { + anyhow::bail!("Provided --id does not match the subnet of --remove-nodes"); + } + } + Ok(()) + } +} diff --git a/rs/cli/src/commands/subnet/resize.rs b/rs/cli/src/commands/subnet/resize.rs index 12ca34740..7f299134a 100644 --- a/rs/cli/src/commands/subnet/resize.rs +++ b/rs/cli/src/commands/subnet/resize.rs @@ -15,11 +15,11 @@ use crate::{ pub struct Resize { /// Number of nodes to be added #[clap(long, default_value = "0")] - pub add: usize, + pub add_count: usize, /// Number of nodes to be removed #[clap(long, default_value = "0")] - pub remove: usize, + pub remove_count: usize, /// Features or Node IDs to exclude from the available nodes pool #[clap(long, num_args(1..))] @@ -59,8 +59,8 @@ impl ExecutableCommand for Resize { .subnet_resize( SubnetResizeRequest { subnet: self.id, - add: self.add, - remove: self.remove, + add: self.add_count, + remove: self.remove_count, exclude: self.exclude.clone().into(), only: self.only.clone().into(), add_nodes: self.add_nodes.clone().into(), diff --git a/rs/cli/src/unit_tests/args_parse.rs b/rs/cli/src/unit_tests/args_parse.rs index bf0a8f54a..9b5ac2617 100644 --- a/rs/cli/src/unit_tests/args_parse.rs +++ b/rs/cli/src/unit_tests/args_parse.rs @@ -72,7 +72,7 @@ fn parse_replace_remove_and_count() { let Subcommands::Replace(r) = s.subcommands else { panic!("expected replace") }; - assert_eq!(r.nodes.len(), 2); + assert_eq!(r.remove_nodes.len(), 2); let args = [ "--id", @@ -103,3 +103,36 @@ fn parse_replace_remove_and_count() { }; assert_eq!(r.optimize, Some(2)); } + +#[test] +fn parse_replace_id_and_add_nodes_aliases() { + let args = [ + "--id", + &PrincipalId::new_subnet_test_id(1).to_string(), + "--motivation", + "m", + "--add-nodes", + &id(1).to_string(), + &id(2).to_string(), + ]; + let s = Subnet::parse_from(["subnet"].into_iter().chain(["replace"]).chain(args)); + let Subcommands::Replace(r) = s.subcommands else { + panic!("expected replace") + }; + assert_eq!(r.add_nodes.len(), 2); + + let args = [ + "--id", + &PrincipalId::new_subnet_test_id(1).to_string(), + "--motivation", + "m", + "--add", + &id(1).to_string(), + &id(2).to_string(), + ]; + let s = Subnet::parse_from(["subnet"].into_iter().chain(["replace"]).chain(args)); + let Subcommands::Replace(r) = s.subcommands else { + panic!("expected replace") + }; + assert_eq!(r.add_nodes.len(), 2); +}