feat: add BE configuration update tool (#23)

Author: QuakeWang

src/tools/be/be_config_update.rs

@@ -0,0 +1,112 @@
+use super::be_http_client;
+use crate::config::Config;
+use crate::error::{CliError, Result};
+use crate::tools::ExecutionResult;
+use crate::tools::Tool;
+use crate::ui;
+use dialoguer::{Confirm, Input, theme::ColorfulTheme};
+use serde::Deserialize;
+use std::path::PathBuf;
+
+#[derive(Deserialize)]
+struct ConfigUpdateResult {
+    config_name: String,
+    status: String,
+    msg: String,
+}
+
+pub struct BeUpdateConfigTool;
+
+impl Tool for BeUpdateConfigTool {
+    fn name(&self) -> &str {
+        "set-be-config"
+    }
+
+    fn description(&self) -> &str {
+        "Update BE configuration variables"
+    }
+
+    fn execute(&self, _config: &Config, _pid: u32) -> Result<ExecutionResult> {
+        let key = prompt_input("Enter BE config key to update")?;
+        let value = prompt_input(&format!("Enter value for '{key}'"))?;
+        let persist = Confirm::with_theme(&ColorfulTheme::default())
+            .with_prompt("Persist this configuration?")
+            .default(false)
+            .interact()
+            .map_err(|e| CliError::InvalidInput(format!("Input failed: {e}")))?;
+
+        ui::print_info(&format!(
+            "Updating BE config: {key}={value} (persist: {persist})"
+        ));
+
+        let endpoint = format!("/api/update_config?{key}={value}&persist={persist}");
+        handle_update_result(be_http_client::post_be_endpoint(&endpoint), &key)
+    }
+
+    fn requires_pid(&self) -> bool {
+        false
+    }
+}
+
+fn prompt_input(prompt: &str) -> Result<String> {
+    let input: String = Input::with_theme(&ColorfulTheme::default())
+        .with_prompt(prompt)
+        .interact_text()
+        .map_err(|e| CliError::InvalidInput(format!("Input failed: {e}")))?;
+
+    let trimmed = input.trim();
+    if trimmed.is_empty() {
+        ui::print_warning("Input cannot be empty!");
+        Err(CliError::GracefulExit)
+    } else {
+        Ok(trimmed.to_string())
+    }
+}
+
+fn get_current_value(key: &str) -> Option<String> {
+    be_http_client::request_be_webserver_port("/varz", Some(key))
+        .ok()?
+        .lines()
+        .next()?
+        .split('=')
+        .nth(1)
+        .map(|v| v.trim().to_string())
+}
+
+fn handle_update_result(result: Result<String>, key: &str) -> Result<ExecutionResult> {
+    let json_response = result.map_err(|e| {
+        ui::print_error(&format!("Failed to update BE config: {e}."));
+        ui::print_info("Tips: Ensure the BE service is running and accessible.");
+        e
+    })?;
+
+    let results: Vec<ConfigUpdateResult> = serde_json::from_str(&json_response)
+        .map_err(|e| CliError::ToolExecutionFailed(format!("Failed to parse response: {e}")))?;
+
+    println!();
+    ui::print_info("Results:");
+
+    let all_ok = results.iter().all(|item| {
+        if item.status == "OK" {
+            match get_current_value(&item.config_name) {
+                Some(value) => println!("  ✓ {} = {}", item.config_name, value),
+                None => println!("  ✓ {}: OK", item.config_name),
+            }
+            true
+        } else {
+            println!("  ✗ {}: FAILED - {}", item.config_name, item.msg);
+            false
+        }
+    });
+
+    if all_ok {
+        Ok(ExecutionResult {
+            output_path: PathBuf::from("console_output"),
+            message: format!("Config '{key}' updated successfully"),
+        })
+    } else {
+        Err(CliError::ToolExecutionFailed(
+            "Some configurations failed to update".to_string(),
+        ))
+    }
+}

src/tools/be/be_http_client.rs

@@ -8,31 +8,30 @@ use std::process::Command;
 
 const BE_DEFAULT_IP: &str = "127.0.0.1";
 
-/// Send an HTTP GET request to a BE API endpoint
-pub fn request_be_webserver_port(endpoint: &str, filter_pattern: Option<&str>) -> Result<String> {
-    let mut be_targets: BTreeSet<(String, u16)> = BTreeSet::new();
-
+fn get_be_targets() -> Result<BTreeSet<(String, u16)>> {
     let ports = get_be_http_ports()?;
-
     let selected_host = be::list::get_selected_be_host();
-
     let cluster_hosts = get_be_ip().unwrap_or_default();
 
     let mut all_hosts = BTreeSet::new();
-    if let Some(host) = &selected_host {
-        all_hosts.insert(host.clone());
-    }
-    for host in cluster_hosts {
+    if let Some(host) = selected_host {
         all_hosts.insert(host);
     }
+    all_hosts.extend(cluster_hosts);
 
     if all_hosts.is_empty() {
         all_hosts.insert(BE_DEFAULT_IP.to_string());
     }
 
-    for host in all_hosts {
-        be_targets.extend(ports.iter().map(|p| (host.clone(), *p)));
-    }
+    Ok(all_hosts
+        .into_iter()
+        .flat_map(|host| ports.iter().map(move |p| (host.clone(), *p)))
+        .collect())
+}
+
+/// Send an HTTP GET request to a BE API endpoint
+pub fn request_be_webserver_port(endpoint: &str, filter_pattern: Option<&str>) -> Result<String> {
+    let be_targets = get_be_targets()?;
 
     for (host, port) in &be_targets {
         let url = format!("http://{host}:{port}{endpoint}");
@@ -104,3 +103,31 @@ pub fn get_be_ip() -> Result<Vec<String>> {
 
     Ok(vec![BE_DEFAULT_IP.to_string()])
 }
+
+/// Send an HTTP POST request to a BE API endpoint
+pub fn post_be_endpoint(endpoint: &str) -> Result<String> {
+    let be_targets = get_be_targets()?;
+
+    for (host, port) in &be_targets {
+        let url = format!("http://{host}:{port}{endpoint}");
+        let mut curl_cmd = Command::new("curl");
+        curl_cmd.args(["-sS", "-X", "POST", &url]);
+
+        if let Ok(output) = executor::execute_command(&mut curl_cmd, "curl") {
+            return Ok(String::from_utf8_lossy(&output.stdout).to_string());
+        }
+    }
+
+    let ports_str = be_targets
+        .iter()
+        .map(|(h, p)| format!("{h}:{p}"))
+        .collect::<Vec<_>>()
+        .join(", ");
+
+    ui::print_warning(
+        "Could not connect to any BE http endpoint. You can select a host via 'be-list'.",
+    );
+    Err(CliError::ToolExecutionFailed(format!(
+        "Could not connect to any BE http port ({ports_str}). Check if BE is running."
+    )))
+}

src/tools/be/be_vars.rs

@@ -31,7 +31,6 @@ impl Tool for BeVarsTool {
         let result = be_http_client::request_be_webserver_port("/varz", Some(&variable_name));
 
         let handler = BeResponseHandler {
-            success_message: "Query completed!",
             empty_warning: "No variables found matching '{}'.",
             error_context: "Failed to query BE",
             tips: "Ensure the BE service is running and accessible.",

src/tools/be/mod.rs

@@ -1,3 +1,4 @@
+mod be_config_update;
 mod be_http_client;
 mod be_vars;
 mod jmap;
@@ -7,6 +8,7 @@ mod pipeline_tasks;
 mod pstack;
 mod response_handler;
 
+pub use be_config_update::BeUpdateConfigTool;
 pub use be_vars::BeVarsTool;
 pub use jmap::{JmapDumpTool, JmapHistoTool};
 pub use list::BeListTool;

src/tools/be/pipeline_tasks.rs

@@ -23,7 +23,6 @@ impl Tool for PipelineTasksTool {
         let result = be_http_client::request_be_webserver_port("/api/running_pipeline_tasks", None);
 
         let handler = BeResponseHandler {
-            success_message: "Pipeline tasks fetched successfully!",
             empty_warning: "No running pipeline tasks found.",
             error_context: "Failed to fetch pipeline tasks",
             tips: "Ensure the BE service is running and accessible.",

src/tools/be/response_handler.rs

@@ -8,42 +8,39 @@ use std::path::PathBuf;
 
 /// Configuration for handling BE API responses
 pub struct BeResponseHandler<'a> {
-    pub success_message: &'a str,
     pub empty_warning: &'a str,
     pub error_context: &'a str,
     pub tips: &'a str,
 }
 
 impl<'a> BeResponseHandler<'a> {
+    fn handle_error(&self, e: crate::error::CliError) -> crate::error::CliError {
+        ui::print_error(&format!("{}: {e}.", self.error_context));
+        ui::print_info(&format!("Tips: {}", self.tips));
+        e
+    }
+
     /// Handle response for console-only output (like be_vars)
     pub fn handle_console_result(
         &self,
         result: Result<String>,
         context: &str,
     ) -> Result<ExecutionResult> {
-        match result {
-            Ok(output) => {
-                ui::print_success(self.success_message);
-                println!();
-                ui::print_info("Results:");
+        let output = result.map_err(|e| self.handle_error(e))?;
 
-                if output.is_empty() {
-                    ui::print_warning(&self.empty_warning.replace("{}", context));
-                } else {
-                    println!("{output}");
-                }
+        println!();
+        ui::print_info("Results:");
 
-                Ok(ExecutionResult {
-                    output_path: PathBuf::from("console_output"),
-                    message: format!("Query completed for: {context}"),
-                })
-            }
-            Err(e) => {
-                ui::print_error(&format!("{}: {e}.", self.error_context));
-                ui::print_info(&format!("Tips: {}", self.tips));
-                Err(e)
-            }
+        if output.is_empty() {
+            ui::print_warning(&self.empty_warning.replace("{}", context));
+        } else {
+            println!("{output}");
         }
+
+        Ok(ExecutionResult {
+            output_path: PathBuf::from("console_output"),
+            message: format!("Query completed for: {context}"),
+        })
     }
 
     /// Handle response with file output (like pipeline_tasks)
@@ -57,47 +54,36 @@ impl<'a> BeResponseHandler<'a> {
     where
         F: Fn(&str) -> String,
     {
-        match result {
-            Ok(output) => {
-                ui::print_success(self.success_message);
-                println!();
-                ui::print_info("Results:");
+        let output = result.map_err(|e| self.handle_error(e))?;
 
-                if output.trim().is_empty() {
-                    ui::print_warning(self.empty_warning);
+        println!();
+        ui::print_info("Results:");
 
-                    Ok(ExecutionResult {
-                        output_path: PathBuf::from("console_output"),
-                        message: "No data found".to_string(),
-                    })
-                } else {
-                    let timestamp = Utc::now().format("%Y%m%d_%H%M%S");
-
-                    let filename = format!("{file_prefix}_{timestamp}.txt");
-                    let output_path = config.output_dir.join(filename);
+        if output.trim().is_empty() {
+            ui::print_warning(self.empty_warning);
+            return Ok(ExecutionResult {
+                output_path: PathBuf::from("console_output"),
+                message: "No data found".to_string(),
+            });
+        }
 
-                    fs::write(&output_path, &output)?;
+        let timestamp = Utc::now().format("%Y%m%d_%H%M%S");
+        let filename = format!("{file_prefix}_{timestamp}.txt");
+        let output_path = config.output_dir.join(filename);
 
-                    println!("{}", summary_fn(&output));
+        fs::write(&output_path, &output)?;
+        println!("{}", summary_fn(&output));
 
-                    let message = format!(
-                        "{} saved to {}",
-                        file_prefix.replace('_', " ").to_title_case(),
-                        output_path.display()
-                    );
+        let message = format!(
+            "{} saved to {}",
+            file_prefix.replace('_', " ").to_title_case(),
+            output_path.display()
+        );
 
-                    Ok(ExecutionResult {
-                        output_path,
-                        message,
-                    })
-                }
-            }
-            Err(e) => {
-                ui::print_error(&format!("{}: {e}.", self.error_context));
-                ui::print_info(&format!("Tips: {}", self.tips));
-                Err(e)
-            }
-        }
+        Ok(ExecutionResult {
+            output_path,
+            message,
+        })
     }
 }
 

src/tools/mod.rs

@@ -44,7 +44,8 @@ impl ToolRegistry {
     /// Creates a new tool registry with all available tools
     pub fn new() -> Self {
         use crate::tools::be::{
-            BeListTool, BeVarsTool, MemzGlobalTool, MemzTool, PipelineTasksTool, PstackTool,
+            BeListTool, BeUpdateConfigTool, BeVarsTool, MemzGlobalTool, MemzTool,
+            PipelineTasksTool, PstackTool,
         };
         use crate::tools::be::{JmapDumpTool as BeJmapDumpTool, JmapHistoTool as BeJmapHistoTool};
         use crate::tools::fe::routine_load::get_routine_load_tools;
@@ -71,6 +72,7 @@ impl ToolRegistry {
         registry.be_tools.push(Box::new(BeListTool));
         registry.be_tools.push(Box::new(PstackTool));
         registry.be_tools.push(Box::new(BeVarsTool));
+        registry.be_tools.push(Box::new(BeUpdateConfigTool));
         registry.be_tools.push(Box::new(BeJmapDumpTool));
         registry.be_tools.push(Box::new(BeJmapHistoTool));
         registry.be_tools.push(Box::new(PipelineTasksTool));