Preserve the comments of the "cli.toml" on changes

Cargo.toml

@@ -240,6 +240,7 @@ tokio-postgres = { version = "0.7.8", features = ["with-chrono-0_4"] }
 tokio-tungstenite = { version = "0.21", features = ["native-tls"] }
 tokio-util = { version = "0.7.4", features = ["time"] }
 toml = "0.8"
+toml_edit = "0.22.22"
 tower-http = { version = "0.5", features = ["cors"] }
 tracing = { version = "0.1.37", features = ["release_max_level_off"] }
 tracing-appender = "0.2.2"

crates/cli/Cargo.toml

@@ -63,6 +63,7 @@ thiserror.workspace = true
 tokio.workspace = true
 tokio-tungstenite.workspace = true
 toml.workspace = true
+toml_edit.workspace = true
 wasmbin.workspace = true
 wasmtime.workspace = true
 webbrowser.workspace = true

crates/cli/src/config.rs

@@ -1,11 +1,22 @@
 use crate::util::{contains_protocol, host_or_url_to_host_and_protocol};
 use anyhow::Context;
 use jsonwebtoken::DecodingKey;
-use serde::{Deserialize, Serialize};
+use spacetimedb::config::{set_opt_value, set_table_opt_value};
 use spacetimedb_fs_utils::atomic_write;
 use spacetimedb_paths::cli::CliTomlPath;
-
-#[derive(Serialize, Deserialize, Clone, Debug)]
+use std::collections::HashMap;
+use toml_edit::{ArrayOfTables, DocumentMut, Item, Table};
+
+const DEFAULT_SERVER_KEY: &str = "default_server";
+const WEB_SESSION_TOKEN_KEY: &str = "web_session_token";
+const SPACETIMEDB_TOKEN_KEY: &str = "spacetimedb_token";
+const SERVER_CONFIGS_KEY: &str = "server_configs";
+const NICKNAME_KEY: &str = "nickname";
+const HOST_KEY: &str = "host";
+const PROTOCOL_KEY: &str = "protocol";
+const ECDSA_PUBLIC_KEY: &str = "ecdsa_public_key";
+
+#[derive(Clone, Debug)]
 pub struct ServerConfig {
     pub nickname: Option<String>,
     pub host: String,
@@ -14,6 +25,21 @@ pub struct ServerConfig {
 }
 
 impl ServerConfig {
+    /// Generate a new [Table] representing this [ServerConfig].
+    pub fn as_table(&self) -> Table {
+        let mut table = Table::new();
+        Self::update_table(&mut table, self);
+        table
+    }
+
+    /// Update an existing [Table] with the values of a [ServerConfig].
+    pub fn update_table(edit: &mut Table, from: &ServerConfig) {
+        set_table_opt_value(edit, NICKNAME_KEY, from.nickname.as_deref());
+        edit[HOST_KEY] = from.host.as_str().into();
+        edit[PROTOCOL_KEY] = from.protocol.as_str().into();
+        set_table_opt_value(edit, ECDSA_PUBLIC_KEY, from.ecdsa_public_key.as_deref());
+    }
+
     fn nick_or_host(&self) -> &str {
         if let Some(nick) = &self.nickname {
             nick
@@ -33,10 +59,29 @@ impl ServerConfig {
     }
 }
 
-#[derive(Default, Deserialize, Serialize, Debug, Clone)]
+impl From<&Table> for ServerConfig {
+    fn from(table: &Table) -> Self {
+        let nickname = table.get(NICKNAME_KEY).and_then(Item::as_str).map(String::from);
+        let host = table.get(HOST_KEY).and_then(Item::as_str).map(String::from).unwrap();
+        let protocol = table
+            .get(PROTOCOL_KEY)
+            .and_then(Item::as_str)
+            .map(String::from)
+            .unwrap();
+        let ecdsa_public_key = table.get(ECDSA_PUBLIC_KEY).and_then(Item::as_str).map(String::from);
+        ServerConfig {
+            nickname,
+            host,
+            protocol,
+            ecdsa_public_key,
+        }
+    }
+}
+
+// Any change here must be coordinated with Config::doc
+#[derive(Default, Debug, Clone)]
 pub struct RawConfig {
     default_server: Option<String>,
-    #[serde(default, skip_serializing_if = "Vec::is_empty")]
     server_configs: Vec<ServerConfig>,
     // TODO: Consider how these tokens should look to be backwards-compatible with the future changes (e.g. we may want to allow users to `login` to switch between multiple accounts - what will we cache and where?)
     // TODO: Move these IDs/tokens out of config so we're no longer storing sensitive tokens in a human-edited file.
@@ -48,6 +93,9 @@ pub struct RawConfig {
 pub struct Config {
     home: RawConfig,
     home_path: CliTomlPath,
+    // The TOML document that was parsed to create `home`.
+    // We need to keep it to preserve comments and formatting when saving the config.
+    doc: DocumentMut,
 }
 
 const NO_DEFAULT_SERVER_ERROR_MESSAGE: &str = "No default server configuration.
@@ -367,6 +415,32 @@ Fetch the server's fingerprint with:
         self.spacetimedb_token = None;
     }
 }
+impl From<&DocumentMut> for RawConfig {
+    fn from(value: &DocumentMut) -> Self {
+        let default_server = value.get(DEFAULT_SERVER_KEY).and_then(Item::as_str).map(String::from);
+        let web_session_token = value
+            .get(WEB_SESSION_TOKEN_KEY)
+            .and_then(Item::as_str)
+            .map(String::from);
+        let spacetimedb_token = value
+            .get(SPACETIMEDB_TOKEN_KEY)
+            .and_then(Item::as_str)
+            .map(String::from);
+
+        let server_configs = value
+            .get(SERVER_CONFIGS_KEY)
+            .and_then(Item::as_array_of_tables)
+            .map(|arr| arr.iter().map(ServerConfig::from).collect())
+            .unwrap_or_default();
+
+        RawConfig {
+            default_server,
+            server_configs,
+            web_session_token,
+            spacetimedb_token,
+        }
+    }
+}
 
 impl Config {
     pub fn default_server_name(&self) -> Option<&str> {
@@ -490,14 +564,15 @@ impl Config {
     }
 
     pub fn load(home_path: CliTomlPath) -> anyhow::Result<Self> {
-        let home = spacetimedb::config::parse_config::<RawConfig>(home_path.as_ref())
+        let home = spacetimedb::config::parse_preserving_config::<RawConfig>(home_path.as_ref())
             .with_context(|| format!("config file {} is invalid", home_path.display()))?;
         Ok(match home {
-            Some(home) => Self { home, home_path },
+            Some((doc, home)) => Self { home, home_path, doc },
             None => {
                 let config = Self {
                     home: RawConfig::new_with_localhost(),
                     home_path,
+                    doc: Default::default(),
                 };
                 config.save();
                 config
@@ -511,15 +586,76 @@ impl Config {
         Self {
             home: RawConfig::new_with_localhost(),
             home_path,
+            doc: Default::default(),
         }
     }
 
+    /// Returns a preserving copy of [Config].
+    fn doc(&self) -> DocumentMut {
+        let mut doc = self.doc.clone();
+
+        let mut set_value = |key: &str, value: Option<&str>| {
+            set_opt_value(&mut doc, key, value);
+        };
+        // Intentionally use a destructuring assignment in case the fields change...
+        let RawConfig {
+            default_server,
+            server_configs,
+            web_session_token,
+            spacetimedb_token,
+        } = &self.home;
+
+        set_value(DEFAULT_SERVER_KEY, default_server.as_deref());
+        set_value(WEB_SESSION_TOKEN_KEY, web_session_token.as_deref());
+        set_value(SPACETIMEDB_TOKEN_KEY, spacetimedb_token.as_deref());
+
+        // Short-circuit if there are no servers.
+        if server_configs.is_empty() {
+            doc.remove(SERVER_CONFIGS_KEY);
+            return doc;
+        }
+        // ... or if there are no server_configs to edit.
+        let server_configs = if let Some(cfg) = doc.get_mut(SERVER_CONFIGS_KEY).and_then(Item::as_array_of_tables_mut) {
+            cfg
+        } else {
+            doc[SERVER_CONFIGS_KEY] = Item::ArrayOfTables(server_configs.iter().map(ServerConfig::as_table).collect());
+            return doc;
+        };
+
+        let mut new_configs = self
+            .home
+            .server_configs
+            .iter()
+            .map(|cfg| (cfg.nick_or_host(), cfg))
+            .collect::<HashMap<_, _>>();
+
+        // Update the existing servers.
+        let mut new_vec = Vec::with_capacity(server_configs.len());
+        for old_config in server_configs.iter_mut() {
+            let nick_or_host = old_config
+                .get(NICKNAME_KEY)
+                .or_else(|| old_config.get(HOST_KEY))
+                .and_then(|v| v.as_str())
+                .unwrap();
+            if let Some(new_config) = new_configs.remove(nick_or_host) {
+                ServerConfig::update_table(old_config, new_config);
+                new_vec.push(old_config.clone());
+            }
+        }
+
+        // Add the new servers.
+        new_vec.extend(new_configs.values().cloned().map(ServerConfig::as_table));
+        *server_configs = ArrayOfTables::from_iter(new_vec);
+
+        doc
+    }
+
     pub fn save(&self) {
         let home_path = &self.home_path;
         // If the `home_path` is in a directory, ensure it exists.
         home_path.create_parent().unwrap();
 
-        let config = toml::to_string_pretty(&self.home).unwrap();
+        let config = self.doc().to_string();
 
         eprintln!("Saving config to {}.", home_path.display());
         // TODO: We currently have a race condition if multiple processes are modifying the config.
@@ -632,3 +768,132 @@ Update the server's fingerprint with:
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use spacetimedb_lib::error::ResultTest;
+    use spacetimedb_paths::cli::CliTomlPath;
+    use spacetimedb_paths::FromPathUnchecked;
+    use std::fs;
+
+    const CONFIG_FULL: &str = r#"default_server = "local"
+web_session_token = "web_session"
+spacetimedb_token = "26ac38857c2bd6c5b60ec557ecd4f9add918fef577dc92c01ca96ff08af5b84d"
+
+# comment on table
+[[server_configs]]
+nickname = "local"
+host = "127.0.0.1:3000"
+protocol = "http"
+
+[[server_configs]]
+nickname = "testnet" # Comment nickname
+host = "testnet.spacetimedb.com"
+protocol = "https"
+
+# Comment end
+"#;
+    const CONFIG_FULL_NO_COMMENT: &str = r#"default_server = "local"
+web_session_token = "web_session"
+spacetimedb_token = "26ac38857c2bd6c5b60ec557ecd4f9add918fef577dc92c01ca96ff08af5b84d"
+
+[[server_configs]]
+nickname = "local"
+host = "127.0.0.1:3000"
+protocol = "http"
+
+[[server_configs]]
+nickname = "testnet"
+host = "testnet.spacetimedb.com"
+protocol = "https"
+
+# Comment end
+"#;
+    const CONFIG_CHANGE_SERVER: &str = r#"default_server = "local"
+web_session_token = "web_session"
+spacetimedb_token = "26ac38857c2bd6c5b60ec557ecd4f9add918fef577dc92c01ca96ff08af5b84d"
+
+[[server_configs]]
+nickname = "testnet" # Comment nickname
+host = "prod.spacetimedb.com"
+protocol = "https"
+
+# Comment end
+"#;
+    const CONFIG_EMPTY: &str = r#"
+# Comment end
+"#;
+    fn check_config<F>(input: &str, output: &str, f: F) -> ResultTest<()>
+    where
+        F: FnOnce(&mut Config) -> ResultTest<()>,
+    {
+        let tmp = tempfile::tempdir()?;
+        let config_path = CliTomlPath::from_path_unchecked(tmp.path().join("config.toml"));
+
+        fs::write(&config_path, input)?;
+
+        let mut config = Config::load(config_path.clone()).unwrap();
+        f(&mut config)?;
+        config.save();
+
+        let contents = fs::read_to_string(&config_path)?;
+
+        assert_eq!(contents, output);
+
+        Ok(())
+    }
+
+    // Test editing the config file.
+    #[test]
+    fn test_config_edits() -> ResultTest<()> {
+        check_config(CONFIG_FULL, CONFIG_EMPTY, |config| {
+            config.home.default_server = None;
+            config.home.server_configs.clear();
+            config.home.spacetimedb_token = None;
+            config.home.web_session_token = None;
+
+            Ok(())
+        })?;
+
+        check_config(CONFIG_FULL, CONFIG_CHANGE_SERVER, |config| {
+            config.home.server_configs.remove(0);
+            config.home.server_configs[0].host = "prod.spacetimedb.com".to_string();
+            Ok(())
+        })?;
+
+        Ok(())
+    }
+
+    // Test adding to the config file.
+    #[test]
+    fn test_config_adds() -> ResultTest<()> {
+        check_config(CONFIG_FULL, CONFIG_FULL, |_| Ok(()))?;
+        check_config(CONFIG_EMPTY, CONFIG_EMPTY, |_| Ok(()))?;
+
+        check_config(CONFIG_EMPTY, CONFIG_FULL_NO_COMMENT, |config| {
+            config.home.default_server = Some("local".to_string());
+            config.home.server_configs = vec![
+                ServerConfig {
+                    nickname: Some("local".to_string()),
+                    host: "127.0.0.1:3000".to_string(),
+                    protocol: "http".to_string(),
+                    ecdsa_public_key: None,
+                },
+                ServerConfig {
+                    nickname: Some("testnet".to_string()),
+                    host: "testnet.spacetimedb.com".to_string(),
+                    protocol: "https".to_string(),
+                    ecdsa_public_key: None,
+                },
+            ];
+            config.home.spacetimedb_token =
+                Some("26ac38857c2bd6c5b60ec557ecd4f9add918fef577dc92c01ca96ff08af5b84d".to_string());
+            config.home.web_session_token = Some("web_session".to_string());
+
+            Ok(())
+        })?;
+
+        Ok(())
+    }
+}

crates/core/Cargo.toml

@@ -92,6 +92,7 @@ tokio-util.workspace = true
 tokio.workspace = true
 tokio-stream = "0.1"
 toml.workspace = true
+toml_edit.workspace = true
 tracing-appender.workspace = true
 tracing-core.workspace = true
 tracing-flame.workspace = true