Override configuration using environment variables (#541)

* Added the ability to update config settings from env vars

* Added tests

* Documented that you can override configuration with environment
variables

* Refactored the config get() methods to use toml-query

* Made the `Updateable` trait more generic

Cargo.toml

@@ -34,6 +34,7 @@ tempdir = "0.3.4"
 itertools = "0.7.4"
 tempfile = "2.2.0"
 shlex = "0.1.1"
+toml-query = "0.6"
 
 # Watch feature
 notify = { version = "4.0", optional = true }

book-example/src/format/config.md

@@ -106,3 +106,43 @@ additional-js = ["custom.js"]
 editor = "./path/to/editor"
 editable = false
 ```
+
+
+## Environment Variables
+
+All configuration values van be overridden from the command line by setting the
+corresponding environment variable. Because many operating systems restrict
+environment variables to be alphanumeric characters or `_`, the configuration
+key needs to be formatted slightly differently to the normal `foo.bar.baz` form.
+
+Variables starting with `MDBOOK_` are used for configuration. The key is
+created by removing the `MDBOOK_` prefix and turning the resulting
+string into `kebab-case`. Double underscores (`__`) separate nested
+keys, while a single underscore (`_`) is replaced with a dash (`-`).
+
+For example:
+
+- `MDBOOK_foo` -> `foo`
+- `MDBOOK_FOO` -> `foo`
+- `MDBOOK_FOO__BAR` -> `foo.bar`
+- `MDBOOK_FOO_BAR` -> `foo-bar`
+- `MDBOOK_FOO_bar__baz` -> `foo-bar.baz`
+
+So by setting the `MDBOOK_BOOK__TITLE` environment variable you can
+override the book's title without needing to touch your `book.toml`.
+
+> **Note:** To facilitate setting more complex config items, the value
+> of an environment variable is first parsed as JSON, falling back to a
+> string if the parse fails.
+>
+> This means, if you so desired, you could override all book metadata
+> when building the book with something like
+>
+> ```text
+> $ export MDBOOK_BOOK="{'title': 'My Awesome Book', authors: ['Michael-F-Bryan']}"
+> $ mdbook build
+> ```
+
+The latter case may be useful in situations where `mdbook` is invoked
+from a script or CI, where it sometimes isn't possible to update the
+`book.toml` before building.

src/book/mod.rs

@@ -58,13 +58,15 @@ impl MDBook {
             warn!("\thttps://rust-lang-nursery.github.io/mdBook/format/config.html");
         }
 
-        let config = if config_location.exists() {
+        let mut config = if config_location.exists() {
             debug!("[*] Loading config from {}", config_location.display());
             Config::from_disk(&config_location)?
         } else {
             Config::default()
         };
 
+        config.update_from_env();
+
         if log_enabled!(::log::Level::Trace) {
             for line in format!("Config: {:#?}", config).lines() {
                 trace!("{}", line);

src/config.rs

@@ -1,19 +1,25 @@
+//! Mdbook's configuration system.
+
 use std::path::{Path, PathBuf};
 use std::fs::File;
 use std::io::Read;
+use std::env;
 use toml::{self, Value};
 use toml::value::Table;
+use toml_query::read::TomlValueReadExt;
+use toml_query::insert::TomlValueInsertExt;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use serde_json;
 
 use errors::*;
 
 /// The overall configuration object for MDBook.
-#[derive(Debug, Clone, Default, PartialEq)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct Config {
     /// Metadata about the book.
     pub book: BookConfig,
     pub build: BuildConfig,
-    rest: Table,
+    rest: Value,
 }
 
 impl Config {
@@ -33,20 +39,74 @@ impl Config {
         Config::from_str(&buffer)
     }
 
+    /// Updates the `Config` from the available environment variables.
+    ///
+    /// Variables starting with `MDBOOK_` are used for configuration. The key is
+    /// created by removing the `MDBOOK_` prefix and turning the resulting
+    /// string into `kebab-case`. Double underscores (`__`) separate nested
+    /// keys, while a single underscore (`_`) is replaced with a dash (`-`).
+    ///
+    /// For example:
+    ///
+    /// - `MDBOOK_foo` -> `foo`
+    /// - `MDBOOK_FOO` -> `foo`
+    /// - `MDBOOK_FOO__BAR` -> `foo.bar`
+    /// - `MDBOOK_FOO_BAR` -> `foo-bar`
+    /// - `MDBOOK_FOO_bar__baz` -> `foo-bar.baz`
+    ///
+    /// So by setting the `MDBOOK_BOOK__TITLE` environment variable you can
+    /// override the book's title without needing to touch your `book.toml`.
+    ///
+    /// > **Note:** To facilitate setting more complex config items, the value
+    /// > of an environment variable is first parsed as JSON, falling back to a
+    /// > string if the parse fails.
+    /// >
+    /// > This means, if you so desired, you could override all book metadata
+    /// > when building the book with something like
+    /// >
+    /// > ```text
+    /// > $ export MDBOOK_BOOK="{'title': 'My Awesome Book', authors: ['Michael-F-Bryan']}"
+    /// > $ mdbook build
+    /// > ```
+    ///
+    /// The latter case may be useful in situations where `mdbook` is invoked
+    /// from a script or CI, where it sometimes isn't possible to update the
+    /// `book.toml` before building.
+    pub fn update_from_env(&mut self) {
+        debug!("Updating the config from environment variables");
+
+        let overrides = env::vars().filter_map(|(key, value)| match parse_env(&key) {
+            Some(index) => Some((index, value)),
+            None => None,
+        });
+
+        for (key, value) in overrides {
+            trace!("{} => {}", key, value);
+            let parsed_value = serde_json::from_str(&value)
+                .unwrap_or_else(|_| serde_json::Value::String(value.to_string()));
+
+            self.set(key, parsed_value).expect("unreachable");
+        }
+    }
+
     /// Fetch an arbitrary item from the `Config` as a `toml::Value`.
     ///
     /// You can use dotted indices to access nested items (e.g.
     /// `output.html.playpen` will fetch the "playpen" out of the html output
     /// table).
     pub fn get(&self, key: &str) -> Option<&Value> {
-        let pieces: Vec<_> = key.split(".").collect();
+        match self.rest.read(key) {
-        recursive_get(&pieces, &self.rest)
+            Ok(inner) => inner,
+            Err(_) => None,
+        }
     }
 
     /// Fetch a value from the `Config` so you can mutate it.
     pub fn get_mut<'a>(&'a mut self, key: &str) -> Option<&'a mut Value> {
-        let pieces: Vec<_> = key.split(".").collect();
+        match self.rest.read_mut(key) {
-        recursive_get_mut(&pieces, &mut self.rest)
+            Ok(inner) => inner,
+            Err(_) => None,
+        }
     }
 
     /// Convenience method for getting the html renderer's configuration.
@@ -80,10 +140,18 @@ impl Config {
     /// The only way this can fail is if we can't serialize `value` into a
     /// `toml::Value`.
     pub fn set<S: Serialize, I: AsRef<str>>(&mut self, index: I, value: S) -> Result<()> {
-        let pieces: Vec<_> = index.as_ref().split(".").collect();
+        let index = index.as_ref();
+
         let value =
             Value::try_from(value).chain_err(|| "Unable to represent the item as a JSON Value")?;
-        recursive_set(&pieces, &mut self.rest, value);
+
+        if index.starts_with("book.") {
+            self.book.update_value(&index[5..], value);
+        } else if index.starts_with("build.") {
+            self.build.update_value(&index[6..], value);
+        } else {
+            self.rest.insert(index, value)?;
+        }
 
         Ok(())
     }
@@ -122,73 +190,20 @@ impl Config {
             cfg.build.build_dir = dest;
         }
 
-        cfg.rest = table;
+        cfg.rest = Value::Table(table);
         cfg
     }
 }
 
-/// Recursively walk down a table and try to set some `foo.bar.baz` value.
+impl Default for Config {
-///
+    fn default() -> Config {
-/// If at any table along the way doesn't exist (or isn't itself a `Table`!) an
+        Config {
-/// empty `Table` will be inserted. e.g. if the `foo` table didn't contain a
+            book: BookConfig::default(),
-/// nested table called `bar`, we'd insert one and then keep recursing.
+            build: BuildConfig::default(),
-fn recursive_set(key: &[&str], table: &mut Table, value: Value) {
+            rest: Value::Table(Table::default()),
-    if key.is_empty() {
-        unreachable!();
-    } else if key.len() == 1 {
-        table.insert(key[0].to_string(), value);
-    } else {
-        let first = key[0];
-        let rest = &key[1..];
-
-        // if `table[first]` isn't a table, replace whatever is there with a
-        // new table.
-        if table.get(first).and_then(|t| t.as_table()).is_none() {
-            table.insert(first.to_string(), Value::Table(Table::new()));
-        }
-
-        let nested = table.get_mut(first).and_then(|t| t.as_table_mut()).unwrap();
-        recursive_set(rest, nested, value);
         }
     }
-
-/// The "getter" version of `recursive_set()`.
-fn recursive_get<'a>(key: &[&str], table: &'a Table) -> Option<&'a Value> {
-    if key.is_empty() {
-        return None;
-    } else if key.len() == 1 {
-        return table.get(key[0]);
 }
-
-    let first = key[0];
-    let rest = &key[1..];
-
-    if let Some(&Value::Table(ref nested)) = table.get(first) {
-        recursive_get(rest, nested)
-    } else {
-        None
-    }
-}
-
-/// The mutable version of `recursive_get()`.
-fn recursive_get_mut<'a>(key: &[&str], table: &'a mut Table) -> Option<&'a mut Value> {
-    // TODO: Figure out how to abstract over mutability to reduce copy-pasta
-    if key.is_empty() {
-        return None;
-    } else if key.len() == 1 {
-        return table.get_mut(key[0]);
-    }
-
-    let first = key[0];
-    let rest = &key[1..];
-
-    if let Some(&mut Value::Table(ref mut nested)) = table.get_mut(first) {
-        recursive_get_mut(rest, nested)
-    } else {
-        None
-    }
-}
-
 impl<'de> Deserialize<'de> for Config {
     fn deserialize<D: Deserializer<'de>>(de: D) -> ::std::result::Result<Self, D::Error> {
         let raw = Value::deserialize(de)?;
@@ -228,26 +243,38 @@ impl<'de> Deserialize<'de> for Config {
         Ok(Config {
             book: book,
             build: build,
-            rest: table,
+            rest: Value::Table(table),
         })
     }
 }
 
 impl Serialize for Config {
     fn serialize<S: Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
+        use serde::ser::Error;
+
         let mut table = self.rest.clone();
 
         let book_config = match Value::try_from(self.book.clone()) {
             Ok(cfg) => cfg,
             Err(_) => {
-                use serde::ser::Error;
                 return Err(S::Error::custom("Unable to serialize the BookConfig"));
             }
         };
 
-        table.insert("book".to_string(), book_config);
+        table.insert("book", book_config).expect("unreachable");
+        table.serialize(s)
+    }
+}
 
-        Value::Table(table).serialize(s)
+fn parse_env(key: &str) -> Option<String> {
+    const PREFIX: &str = "MDBOOK_";
+
+    if key.starts_with(PREFIX) {
+        let key = &key[PREFIX.len()..];
+
+        Some(key.to_lowercase().replace("__", ".").replace("_", "-"))
+    } else {
+        None
     }
 }
 
@@ -346,6 +373,35 @@ impl Default for Playpen {
     }
 }
 
+/// Allows you to "update" any arbitrary field in a struct by round-tripping via
+/// a `toml::Value`.
+///
+/// This is definitely not the most performant way to do things, which means you
+/// should probably keep it away from tight loops...
+trait Updateable<'de>: Serialize + Deserialize<'de> {
+    fn update_value<S: Serialize>(&mut self, key: &str, value: S) {
+        let mut raw = Value::try_from(&self).expect("unreachable");
+
+        {
+            if let Ok(value) = Value::try_from(value) {
+                let _ = raw.insert(key, value);
+            } else {
+                return;
+            }
+        }
+
+        if let Ok(updated) = raw.try_into() {
+            *self = updated;
+        }
+    }
+}
+
+impl<'de, T> Updateable<'de> for T
+where
+    T: Serialize + Deserialize<'de>,
+{
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -518,4 +574,77 @@ mod tests {
         let got: String = cfg.get_deserialized(key).unwrap();
         assert_eq!(got, value);
     }
+
+    #[test]
+    fn parse_env_vars() {
+        let inputs = vec![
+            ("FOO", None),
+            ("MDBOOK_foo", Some("foo")),
+            ("MDBOOK_FOO__bar__baz", Some("foo.bar.baz")),
+            ("MDBOOK_FOO_bar__baz", Some("foo-bar.baz")),
+        ];
+
+        for (src, should_be) in inputs {
+            let got = parse_env(src);
+            let should_be = should_be.map(|s| s.to_string());
+
+            assert_eq!(got, should_be);
+        }
+    }
+
+    fn encode_env_var(key: &str) -> String {
+        format!(
+            "MDBOOK_{}",
+            key.to_uppercase().replace('.', "__").replace("-", "_")
+        )
+    }
+
+    #[test]
+    fn update_config_using_env_var() {
+        let mut cfg = Config::default();
+        let key = "foo.bar";
+        let value = "baz";
+
+        assert!(cfg.get(key).is_none());
+
+        let encoded_key = encode_env_var(key);
+        env::set_var(encoded_key, value);
+
+        cfg.update_from_env();
+
+        assert_eq!(cfg.get_deserialized::<String, _>(key).unwrap(), value);
+    }
+
+    #[test]
+    fn update_config_using_env_var_and_complex_value() {
+        let mut cfg = Config::default();
+        let key = "foo-bar.baz";
+        let value = json!({"array": [1, 2, 3], "number": 3.14});
+        let value_str = serde_json::to_string(&value).unwrap();
+
+        assert!(cfg.get(key).is_none());
+
+        let encoded_key = encode_env_var(key);
+        env::set_var(encoded_key, value_str);
+
+        cfg.update_from_env();
+
+        assert_eq!(
+            cfg.get_deserialized::<serde_json::Value, _>(key).unwrap(),
+            value
+        );
+    }
+
+    #[test]
+    fn update_book_title_via_env() {
+        let mut cfg = Config::default();
+        let should_be = "Something else".to_string();
+
+        assert_ne!(cfg.book.title, Some(should_be.clone()));
+
+        env::set_var("MDBOOK_BOOK__TITLE", &should_be);
+        cfg.update_from_env();
+
+        assert_eq!(cfg.book.title, Some(should_be));
+    }
 }

src/lib.rs

@@ -113,6 +113,7 @@ extern crate shlex;
 extern crate tempdir;
 extern crate tempfile;
 extern crate toml;
+extern crate toml_query;
 
 #[cfg(test)]
 #[macro_use]
@@ -141,6 +142,10 @@ pub mod errors {
             Utf8(::std::string::FromUtf8Error);
         }
 
+        links {
+            TomlQuery(::toml_query::error::Error, ::toml_query::error::ErrorKind);
+        }
+
         errors {
             Subprocess(message: String, output: ::std::process::Output) {
                 description("A subprocess failed")