enhancement(generator): Adds optional file output

Signed-off-by: Ian Henry <[email protected]>

Author: eeyun

src/generate.rs

@@ -5,7 +5,12 @@ use crate::config::{
 use colored::*;
 use indexmap::IndexMap;
 use serde::Serialize;
-use std::collections::BTreeMap;
+use std::{
+    collections::BTreeMap,
+    fs::{create_dir_all, File},
+    io::Write,
+    path::PathBuf,
+};
 use structopt::StructOpt;
 use toml::Value;
 
@@ -46,6 +51,10 @@ pub struct Opts {
     /// is then up to you to restructure the `inputs` of each component to build
     /// the topology you need.
     expression: String,
+
+    /// Generate config as a file
+    #[structopt(long, parse(from_os_str))]
+    file: Option<PathBuf>,
 }
 
 #[derive(Serialize)]
@@ -71,7 +80,11 @@ pub struct Config {
     pub sinks: Option<IndexMap<String, SinkOuter>>,
 }
 
-fn generate_example(include_globals: bool, expression: &str) -> Result<String, Vec<String>> {
+fn generate_example(
+    include_globals: bool,
+    expression: &str,
+    file: &Option<PathBuf>,
+) -> Result<String, Vec<String>> {
     let components: Vec<Vec<_>> = expression
         .split(|c| c == '|' || c == '/')
         .map(|s| {
@@ -309,6 +322,16 @@ fn generate_example(include_globals: bool, expression: &str) -> Result<String, V
         }
     }
 
+    if file.is_some() {
+        match write_config(file.as_ref().unwrap(), &builder) {
+            Ok(_) => println!(
+                "Config file written to {:?}",
+                &file.as_ref().unwrap().join("\n")
+            ),
+            Err(e) => errs.push(format!("failed to write to file: {}", e)),
+        };
+    };
+
     if !errs.is_empty() {
         Err(errs)
     } else {
@@ -317,7 +340,7 @@ fn generate_example(include_globals: bool, expression: &str) -> Result<String, V
 }
 
 pub fn cmd(opts: &Opts) -> exitcode::ExitCode {
-    match generate_example(!opts.fragment, &opts.expression) {
+    match generate_example(!opts.fragment, &opts.expression, &opts.file) {
         Ok(s) => {
             println!("{}", s);
             exitcode::OK
@@ -329,6 +352,19 @@ pub fn cmd(opts: &Opts) -> exitcode::ExitCode {
     }
 }
 
+fn write_config(filepath: &PathBuf, body: &str) -> Result<(), crate::Error> {
+    if filepath.exists() {
+        // If the file exists, we don't want to overwrite, that's just rude.
+        return Err(format!("{:?} already exists", &filepath).into());
+    } else {
+        if let Some(directory) = filepath.parent() {
+            create_dir_all(directory)?;
+        }
+        File::create(filepath).and_then(|mut file| file.write(body.as_bytes()))?;
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -339,23 +375,23 @@ mod tests {
 
         for name in SourceDescription::types() {
             let param = format!("{}//", name);
-            let cfg = generate_example(true, &param).unwrap();
+            let cfg = generate_example(true, &param, &None).unwrap();
             if let Err(error) = toml::from_str::<crate::config::ConfigBuilder>(&cfg) {
                 errors.push((param, error));
             }
         }
 
         for name in TransformDescription::types() {
             let param = format!("/{}/", name);
-            let cfg = generate_example(true, &param).unwrap();
+            let cfg = generate_example(true, &param, &None).unwrap();
             if let Err(error) = toml::from_str::<crate::config::ConfigBuilder>(&cfg) {
                 errors.push((param, error));
             }
         }
 
         for name in SinkDescription::types() {
             let param = format!("//{}", name);
-            let cfg = generate_example(true, &param).unwrap();
+            let cfg = generate_example(true, &param, &None).unwrap();
             if let Err(error) = toml::from_str::<crate::config::ConfigBuilder>(&cfg) {
                 errors.push((param, error));
             }
@@ -371,7 +407,7 @@ mod tests {
     #[test]
     fn generate_basic() {
         assert_eq!(
-            generate_example(true, "stdin/json_parser/console"),
+            generate_example(true, "stdin/json_parser/console", &None),
             Ok(r#"data_dir = "/var/lib/vector/"
 
 [sources.source0]
@@ -402,7 +438,7 @@ when_full = "block"
         );
 
         assert_eq!(
-            generate_example(true, "stdin|json_parser|console"),
+            generate_example(true, "stdin|json_parser|console", &None),
             Ok(r#"data_dir = "/var/lib/vector/"
 
 [sources.source0]
@@ -433,7 +469,7 @@ when_full = "block"
         );
 
         assert_eq!(
-            generate_example(true, "stdin//console"),
+            generate_example(true, "stdin//console", &None),
             Ok(r#"data_dir = "/var/lib/vector/"
 
 [sources.source0]
@@ -458,7 +494,7 @@ when_full = "block"
         );
 
         assert_eq!(
-            generate_example(true, "//console"),
+            generate_example(true, "//console", &None),
             Ok(r#"data_dir = "/var/lib/vector/"
 
 [sinks.sink0]
@@ -479,7 +515,7 @@ when_full = "block"
         );
 
         assert_eq!(
-            generate_example(true, "/add_fields,json_parser,remove_fields"),
+            generate_example(true, "/add_fields,json_parser,remove_fields", &None),
             Ok(r#"data_dir = "/var/lib/vector/"
 
 [transforms.transform0]
@@ -504,7 +540,7 @@ type = "remove_fields"
         );
 
         assert_eq!(
-            generate_example(false, "/add_fields,json_parser,remove_fields"),
+            generate_example(false, "/add_fields,json_parser,remove_fields", &None),
             Ok(r#"
 [transforms.transform0]
 inputs = []