From 756fc26bcfd08e2be44ad2108a812fb8ee226482 Mon Sep 17 00:00:00 2001
From: Chris O'Neil <chriso83@protonmail.com>
Date: Sun, 3 Nov 2024 21:33:15 +0000
Subject: [PATCH] feat: `start` command supports custom inventory

---
 src/ansible/provisioning.rs | 29 +++++++++++++++++++++++++++--
 src/lib.rs                  |  9 +++++++--
 src/main.rs                 | 15 ++++++++++++++-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/ansible/provisioning.rs b/src/ansible/provisioning.rs
index e044d1e0..5c42f181 100644
--- a/src/ansible/provisioning.rs
+++ b/src/ansible/provisioning.rs
@@ -22,7 +22,11 @@ use evmlib::common::U256;
 use log::{debug, error, trace};
 use semver::Version;
 use sn_service_management::NodeRegistry;
-use std::{net::SocketAddr, path::PathBuf, time::{Instant, Duration}};
+use std::{
+    net::SocketAddr,
+    path::PathBuf,
+    time::{Duration, Instant},
+};
 use walkdir::WalkDir;
 
 use crate::ansible::extra_vars;
@@ -490,9 +494,30 @@ impl AnsibleProvisioner {
         Ok(())
     }
 
-    pub fn start_nodes(&self, interval: Duration) -> Result<()> {
+    pub fn start_nodes(
+        &self,
+        environment_name: &str,
+        interval: Duration,
+        custom_inventory: Option<Vec<VirtualMachine>>,
+    ) -> Result<()> {
         let mut extra_vars = ExtraVarsDocBuilder::default();
         extra_vars.add_variable("interval", &interval.as_millis().to_string());
+
+        if let Some(custom_inventory) = custom_inventory {
+            println!("Running the start telegraf playbook with a custom inventory");
+            generate_custom_environment_inventory(
+                &custom_inventory,
+                environment_name,
+                &self.ansible_runner.working_directory_path.join("inventory"),
+            )?;
+            self.ansible_runner.run_playbook(
+                AnsiblePlaybook::StartNodes,
+                AnsibleInventoryType::Custom,
+                Some(extra_vars.build()),
+            )?;
+            return Ok(());
+        }
+
         for node_inv_type in AnsibleInventoryType::iter_node_type() {
             self.ansible_runner.run_playbook(
                 AnsiblePlaybook::StartNodes,
diff --git a/src/lib.rs b/src/lib.rs
index bcc9df77..397ed1f1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -632,8 +632,13 @@ impl TestnetDeployer {
         Ok(())
     }
 
-    pub fn start(&self, interval: Duration) -> Result<()> {
-        self.ansible_provisioner.start_nodes(interval)?;
+    pub fn start(
+        &self,
+        interval: Duration,
+        custom_inventory: Option<Vec<VirtualMachine>>,
+    ) -> Result<()> {
+        self.ansible_provisioner
+            .start_nodes(&self.environment_name, interval, custom_inventory)?;
         Ok(())
     }
 
diff --git a/src/main.rs b/src/main.rs
index cabd2c70..db6248ce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -560,6 +560,11 @@ enum Commands {
     /// This can be useful if all nodes did not upgrade successfully.
     #[clap(name = "start")]
     Start {
+        /// Provide a list of VM names to use as a custom inventory.
+        ///
+        /// This will start nodes on a particular subset of VMs.
+        #[clap(name = "custom-inventory", long, use_value_delimiter = true)]
+        custom_inventory: Option<Vec<String>>,
         /// Maximum number of forks Ansible will use to execute tasks on target hosts.
         #[clap(long, default_value_t = 50)]
         forks: usize,
@@ -1978,6 +1983,7 @@ async fn main() -> Result<()> {
             Ok(())
         }
         Commands::Start {
+            custom_inventory,
             forks,
             interval,
             name,
@@ -2000,7 +2006,14 @@ async fn main() -> Result<()> {
                 return Err(eyre!("The {name} environment does not exist"));
             }
 
-            testnet_deployer.start(interval)?;
+            let custom_inventory = if let Some(custom_inventory) = custom_inventory {
+                let custom_vms = get_custom_inventory(&inventory, &custom_inventory)?;
+                Some(custom_vms)
+            } else {
+                None
+            };
+
+            testnet_deployer.start(interval, custom_inventory)?;
 
             Ok(())
         }