feat: Retain context

src/cluster.rs

@@ -221,8 +221,8 @@ impl Cluster {
     }
 
     pub async fn talosconfig_merge(&self) {
-        let cluster_path = get_configs_path().join(&self.name).join("talosconfig");
         let path = get_configs_path().join("talosconfig");
+        let cluster_path = get_configs_path().join(&self.name).join("talosconfig");
 
         let mut command = Command::new("talosctl");
         command.args([

src/lib.rs

@@ -10,6 +10,7 @@ pub mod node;
 pub mod patch;
 pub mod schematic;
 pub mod secret;
+pub mod talosconfig;
 pub mod tftp;
 
 pub(crate) static REPO_PATH: OnceLock<PathBuf> = OnceLock::new();

src/main.rs

@@ -6,6 +6,7 @@ use clap::{CommandFactory, Parser};
 use clap_complete::{Shell, generate as generate_complete};
 use crete::cluster::Cluster;
 use crete::environment::PathEnvironment;
+use crete::talosconfig::{get_current_context, set_context};
 use crete::{get_configs_path, get_repo_path, set_repo_path};
 use futures::StreamExt;
 use futures::future::join_all;
@@ -53,6 +54,8 @@ async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
             .unwrap();
     }
 
+    let context = get_current_context().await;
+
     // Remove existing config files and create output directory
     let path = get_configs_path();
     tokio::fs::remove_dir_all(&path).await.ok();
@@ -80,6 +83,10 @@ async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
         cluster.talosconfig_merge().await;
     }
 
+    if let Some(context) = context {
+        set_context(context).await;
+    }
+
     Ok(())
 }
 

src/talosconfig.rs

@@ -0,0 +1,34 @@
+use serde::Deserialize;
+use tokio::process::Command;
+
+use crate::{get_configs_path, run_command};
+
+#[derive(Debug, Deserialize)]
+pub struct TalosConfig {
+    pub context: String,
+}
+
+pub async fn get_current_context() -> Option<String> {
+    let path = get_configs_path().join("talosconfig");
+
+    tokio::fs::read(&path)
+        .await
+        .ok()
+        .and_then(|config| serde_yaml::from_slice(&config).ok())
+        .map(|config: TalosConfig| config.context)
+}
+
+pub async fn set_context(context: String) {
+    let path = get_configs_path().join("talosconfig");
+
+    let mut command = Command::new("talosctl");
+    command.args([
+        "config",
+        "--talosconfig",
+        path.to_str().expect("Path should be valid utf-8"),
+        "context",
+        &context,
+    ]);
+
+    run_command(command).await;
+}