feat(risedev): support providing env var in yaml config (#18396)

Signed-off-by: xxchan <[email protected]>

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

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)
     };
 

src/risedevtool/src/bin/risedev-dev.rs

@@ -66,6 +66,7 @@ impl ProgressManager {
 fn task_main(
     manager: &mut ProgressManager,
     services: &Vec<ServiceConfig>,
+    env: Vec<String>,
 ) -> 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(_) => {

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<Path>, profile: &str) -> Result<(Option<String>, Yaml)> {
+    /// `(config_path, env, steps)`
+    pub fn expand(
+        root: impl AsRef<Path>,
+        profile: &str,
+    ) -> Result<(Option<String>, Vec<String>, Yaml)> {
         Self::expand_with_extra_info(root, profile, HashMap::new())
     }
 
@@ -72,7 +93,7 @@ impl ConfigExpander {
         root: impl AsRef<Path>,
         profile: &str,
         extra_info: HashMap<String, String>,
-    ) -> Result<(Option<String>, Yaml)> {
+    ) -> Result<(Option<String>, Vec<String>, 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.

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<String>,
+}
 
 pub const RISEDEV_NAME: &str = "risedev";
 
@@ -33,8 +35,8 @@ pub fn new_tmux_command() -> Command {
 }
 
 impl ConfigureTmuxTask {
-    pub fn new() -> Result<Self> {
-        Ok(Self)
+    pub fn new(env: Vec<String>) -> Result<Self> {
+        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"));