From 9923c3abb8f87a6f643d5f753585e2b31f4de7e0 Mon Sep 17 00:00:00 2001 From: xxchan Date: Wed, 4 Sep 2024 15:52:57 +0800 Subject: [PATCH] feat(risedev): support providing env var in yaml config (#18396) Signed-off-by: xxchan --- risedev.yml | 6 ++- src/risedevtool/src/bin/risedev-compose.rs | 4 +- src/risedevtool/src/bin/risedev-dev.rs | 7 +-- src/risedevtool/src/config.rs | 45 +++++++++++++++++-- .../src/task/configure_tmux_service.rs | 15 ++++--- 5 files changed, 62 insertions(+), 15 deletions(-) diff --git a/risedev.yml b/risedev.yml index 3c7f8e0e09be4..22c4569adb610 100644 --- a/risedev.yml +++ b/risedev.yml @@ -16,8 +16,12 @@ profile: # The default configuration will start 1 compute node, 1 meta node and 1 frontend. default: - # Specify a configuration file to override the default settings + # # Specify a configuration file to override the default settings # config-path: src/config/example.toml + # # Specify custom environment variables + # env: + # RUST_LOG: "info,risingwave_storage::hummock=off" + # RW_ENABLE_PRETTY_LOG: "true" steps: # If you want to use the local s3 storage, enable the following line # - use: minio diff --git a/src/risedevtool/src/bin/risedev-compose.rs b/src/risedevtool/src/bin/risedev-compose.rs index 0547fda6b7008..e51961831056a 100644 --- a/src/risedevtool/src/bin/risedev-compose.rs +++ b/src/risedevtool/src/bin/risedev-compose.rs @@ -82,11 +82,11 @@ fn main() -> Result<()> { ) .collect(); - let (config_path, expanded_config) = + let (config_path, _env, expanded_config) = ConfigExpander::expand_with_extra_info(".", &opts.profile, extra_info)?; (expanded_config, Some(compose_deploy_config), config_path) } else { - let (config_path, expanded_config) = ConfigExpander::expand(".", &opts.profile)?; + let (config_path, _env, expanded_config) = ConfigExpander::expand(".", &opts.profile)?; (expanded_config, None, config_path) }; diff --git a/src/risedevtool/src/bin/risedev-dev.rs b/src/risedevtool/src/bin/risedev-dev.rs index c53453b3f903d..80415e321d805 100644 --- a/src/risedevtool/src/bin/risedev-dev.rs +++ b/src/risedevtool/src/bin/risedev-dev.rs @@ -66,6 +66,7 @@ impl ProgressManager { fn task_main( manager: &mut ProgressManager, services: &Vec, + env: Vec, ) -> Result<(Vec<(String, Duration)>, String)> { let log_path = env::var("PREFIX_LOG")?; @@ -82,7 +83,7 @@ fn task_main( // Start Tmux and kill previous services { let mut ctx = ExecuteContext::new(&mut logger, manager.new_progress(), status_dir.clone()); - let mut service = ConfigureTmuxTask::new()?; + let mut service = ConfigureTmuxTask::new(env)?; service.execute(&mut ctx)?; writeln!( @@ -392,7 +393,7 @@ fn main() -> Result<()> { .nth(1) .unwrap_or_else(|| "default".to_string()); - let (config_path, risedev_config) = ConfigExpander::expand(".", &task_name)?; + let (config_path, env, risedev_config) = ConfigExpander::expand(".", &task_name)?; if let Some(config_path) = &config_path { let target = Path::new(&env::var("PREFIX_CONFIG")?).join("risingwave.toml"); @@ -420,7 +421,7 @@ fn main() -> Result<()> { services.len(), task_name )); - let task_result = task_main(&mut manager, &services); + let task_result = task_main(&mut manager, &services, env); match task_result { Ok(_) => { diff --git a/src/risedevtool/src/config.rs b/src/risedevtool/src/config.rs index e4ba9acdf4e19..1abab635b88c1 100644 --- a/src/risedevtool/src/config.rs +++ b/src/risedevtool/src/config.rs @@ -50,6 +50,24 @@ impl ConfigExpander { /// Transforms `risedev.yml` and `risedev-profiles.user.yml` to a fully expanded yaml file. /// + /// Format: + /// + /// ```yaml + /// my-profile: + /// config-path: src/config/ci-recovery.toml + /// env: + /// RUST_LOG: "info,risingwave_storage::hummock=off" + /// RW_ENABLE_PRETTY_LOG: "true" + /// steps: + /// - use: minio + /// - use: sqlite + /// - use: meta-node + /// meta-backend: sqlite + /// - use: compute-node + /// parallelism: 1 + /// - use: frontend + /// ``` + /// /// # Arguments /// /// * `root` is the root directory of these YAML files. @@ -58,8 +76,11 @@ impl ConfigExpander { /// /// # Returns /// - /// A pair of `config_path` and expanded steps (items in `{profile}.steps` section in YAML) - pub fn expand(root: impl AsRef, profile: &str) -> Result<(Option, Yaml)> { + /// `(config_path, env, steps)` + pub fn expand( + root: impl AsRef, + profile: &str, + ) -> Result<(Option, Vec, Yaml)> { Self::expand_with_extra_info(root, profile, HashMap::new()) } @@ -72,7 +93,7 @@ impl ConfigExpander { root: impl AsRef, profile: &str, extra_info: HashMap, - ) -> Result<(Option, Yaml)> { + ) -> Result<(Option, Vec, Yaml)> { let global_path = root.as_ref().join(RISEDEV_CONFIG_FILE); let global_yaml = Self::load_yaml(global_path)?; let global_config = global_yaml @@ -120,6 +141,22 @@ impl ConfigExpander { .get(&Yaml::String("config-path".to_string())) .and_then(|s| s.as_str()) .map(|s| s.to_string()); + let mut env = vec![]; + if let Some(env_section) = profile_section.get(&Yaml::String("env".to_string())) { + let env_section = env_section + .as_hash() + .ok_or_else(|| anyhow!("expect `env` section to be a hashmap"))?; + + for (k, v) in env_section { + let key = k + .as_str() + .ok_or_else(|| anyhow!("expect env key to be a string"))?; + let value = v + .as_str() + .ok_or_else(|| anyhow!("expect env value to be a string"))?; + env.push(format!("{}={}", key, value)); + } + } let steps = profile_section .get(&Yaml::String("steps".to_string())) @@ -131,7 +168,7 @@ impl ConfigExpander { let steps = IdExpander::new(&steps)?.visit(steps)?; let steps = ProvideExpander::new(&steps)?.visit(steps)?; - Ok((config_path, steps)) + Ok((config_path, env, steps)) } /// Parses the expanded yaml into [`ServiceConfig`]s. diff --git a/src/risedevtool/src/task/configure_tmux_service.rs b/src/risedevtool/src/task/configure_tmux_service.rs index af4581456611b..925cb2de38444 100644 --- a/src/risedevtool/src/task/configure_tmux_service.rs +++ b/src/risedevtool/src/task/configure_tmux_service.rs @@ -22,7 +22,9 @@ use console::style; use crate::util::{risedev_cmd, stylized_risedev_subcmd}; use crate::{ExecuteContext, Task}; -pub struct ConfigureTmuxTask; +pub struct ConfigureTmuxTask { + env: Vec, +} pub const RISEDEV_NAME: &str = "risedev"; @@ -33,8 +35,8 @@ pub fn new_tmux_command() -> Command { } impl ConfigureTmuxTask { - pub fn new() -> Result { - Ok(Self) + pub fn new(env: Vec) -> Result { + Ok(Self { env }) } } @@ -78,8 +80,11 @@ impl Task for ConfigureTmuxTask { cmd.arg("new-session") // this will automatically create the `risedev` tmux server .arg("-d") .arg("-s") - .arg(RISEDEV_NAME) - .arg("-c") + .arg(RISEDEV_NAME); + for e in &self.env { + cmd.arg("-e").arg(e); + } + cmd.arg("-c") .arg(Path::new(&prefix_path)) .arg(Path::new(&prefix_bin).join("welcome.sh"));