common_config/
config.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use config::{Environment, File, FileFormat};
16use serde::de::DeserializeOwned;
17use serde::Serialize;
18use snafu::ResultExt;
19
20use crate::error::{LoadLayeredConfigSnafu, Result, SerdeJsonSnafu, TomlFormatSnafu};
21
22/// Separator for environment variables. For example, `DATANODE__STORAGE__MANIFEST__CHECKPOINT_MARGIN`.
23pub const ENV_VAR_SEP: &str = "__";
24
25/// Separator for list values in environment variables. For example, `localhost:3001,localhost:3002,localhost:3003`.
26pub const ENV_LIST_SEP: &str = ",";
27
28/// Configuration trait defines the common interface for configuration that can be loaded from multiple sources and serialized to TOML.
29pub trait Configurable: Serialize + DeserializeOwned + Default + Sized {
30    /// Load the configuration from multiple sources and merge them.
31    /// The precedence order is: config file > environment variables > default values.
32    /// `env_prefix` is the prefix of environment variables, e.g. "FRONTEND__xxx".
33    /// The function will use dunder(double underscore) `__` as the separator for environment variables, for example:
34    /// `DATANODE__STORAGE__MANIFEST__CHECKPOINT_MARGIN` will be mapped to `DatanodeOptions.storage.manifest.checkpoint_margin` field in the configuration.
35    /// `list_keys` is the list of keys that should be parsed as a list, for example, you can pass `Some(&["meta_client_options.metasrv_addrs"]` to parse `GREPTIMEDB_METASRV__META_CLIENT_OPTIONS__METASRV_ADDRS` as a list.
36    /// The function will use comma `,` as the separator for list values, for example: `127.0.0.1:3001,127.0.0.1:3002,127.0.0.1:3003`.
37    fn load_layered_options(config_file: Option<&str>, env_prefix: &str) -> Result<Self> {
38        let default_opts = Self::default();
39
40        let env_source = {
41            let mut env = Environment::default();
42
43            if !env_prefix.is_empty() {
44                env = env.prefix(env_prefix);
45            }
46
47            if let Some(list_keys) = Self::env_list_keys() {
48                env = env.list_separator(ENV_LIST_SEP);
49                for key in list_keys {
50                    env = env.with_list_parse_key(key);
51                }
52            }
53
54            env.try_parsing(true)
55                .separator(ENV_VAR_SEP)
56                .ignore_empty(true)
57        };
58
59        // Workaround: Replacement for `Config::try_from(&default_opts)` due to
60        // `ConfigSerializer` cannot handle the case of an empty struct contained
61        // within an iterative structure.
62        // See: https://github.com/mehcode/config-rs/issues/461
63        let json_str = serde_json::to_string(&default_opts).context(SerdeJsonSnafu)?;
64        let default_config = File::from_str(&json_str, FileFormat::Json);
65
66        // Add default values and environment variables as the sources of the configuration.
67        let mut layered_config = config::Config::builder()
68            .add_source(default_config)
69            .add_source(env_source);
70
71        // Add config file as the source of the configuration if it is specified.
72        if let Some(config_file) = config_file {
73            layered_config = layered_config.add_source(File::new(config_file, FileFormat::Toml));
74        }
75
76        let mut opts: Self = layered_config
77            .build()
78            .and_then(|x| x.try_deserialize())
79            .context(LoadLayeredConfigSnafu)?;
80
81        opts.validate_sanitize()?;
82
83        Ok(opts)
84    }
85
86    /// Validate(and possibly sanitize) the configuration.
87    fn validate_sanitize(&mut self) -> Result<()> {
88        Ok(())
89    }
90
91    /// List of toml keys that should be parsed as a list.
92    fn env_list_keys() -> Option<&'static [&'static str]> {
93        None
94    }
95
96    /// Serialize the configuration to a TOML string.
97    fn to_toml(&self) -> Result<String> {
98        toml::to_string(&self).context(TomlFormatSnafu)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use std::io::Write;
105
106    use common_telemetry::logging::LoggingOptions;
107    use common_test_util::temp_dir::create_named_temp_file;
108    use common_wal::config::DatanodeWalConfig;
109    use datanode::config::{ObjectStoreConfig, StorageConfig};
110    use meta_client::MetaClientOptions;
111    use serde::{Deserialize, Serialize};
112
113    use super::*;
114    use crate::Mode;
115
116    #[derive(Debug, Serialize, Deserialize)]
117    struct TestDatanodeConfig {
118        mode: Mode,
119        node_id: Option<u64>,
120        logging: LoggingOptions,
121        meta_client: Option<MetaClientOptions>,
122        wal: DatanodeWalConfig,
123        storage: StorageConfig,
124    }
125
126    impl Default for TestDatanodeConfig {
127        fn default() -> Self {
128            Self {
129                mode: Mode::Distributed,
130                node_id: None,
131                logging: LoggingOptions::default(),
132                meta_client: None,
133                wal: DatanodeWalConfig::default(),
134                storage: StorageConfig::default(),
135            }
136        }
137    }
138
139    impl Configurable for TestDatanodeConfig {
140        fn env_list_keys() -> Option<&'static [&'static str]> {
141            Some(&["meta_client.metasrv_addrs"])
142        }
143    }
144
145    #[test]
146    fn test_load_layered_options() {
147        let mut file = create_named_temp_file();
148        let toml_str = r#"
149            mode = "distributed"
150            enable_memory_catalog = false
151            rpc_addr = "127.0.0.1:3001"
152            rpc_hostname = "127.0.0.1"
153            rpc_runtime_size = 8
154            mysql_addr = "127.0.0.1:4406"
155            mysql_runtime_size = 2
156
157            [meta_client]
158            timeout = "3s"
159            connect_timeout = "5s"
160            tcp_nodelay = true
161
162            [wal]
163            provider = "raft_engine"
164            dir = "./greptimedb_data/wal"
165            file_size = "1GB"
166            purge_threshold = "50GB"
167            purge_interval = "10m"
168            read_batch_size = 128
169            sync_write = false
170
171            [logging]
172            level = "debug"
173            dir = "./greptimedb_data/test/logs"
174        "#;
175        write!(file, "{}", toml_str).unwrap();
176
177        let env_prefix = "DATANODE_UT";
178        temp_env::with_vars(
179            // The following environment variables will be used to override the values in the config file.
180            [
181                (
182                    // storage.type = S3
183                    [
184                        env_prefix.to_string(),
185                        "storage".to_uppercase(),
186                        "type".to_uppercase(),
187                    ]
188                    .join(ENV_VAR_SEP),
189                    Some("S3"),
190                ),
191                (
192                    // storage.bucket = mybucket
193                    [
194                        env_prefix.to_string(),
195                        "storage".to_uppercase(),
196                        "bucket".to_uppercase(),
197                    ]
198                    .join(ENV_VAR_SEP),
199                    Some("mybucket"),
200                ),
201                (
202                    // wal.dir = /other/wal/dir
203                    [
204                        env_prefix.to_string(),
205                        "wal".to_uppercase(),
206                        "dir".to_uppercase(),
207                    ]
208                    .join(ENV_VAR_SEP),
209                    Some("/other/wal/dir"),
210                ),
211                (
212                    // meta_client.metasrv_addrs = 127.0.0.1:3001,127.0.0.1:3002,127.0.0.1:3003
213                    [
214                        env_prefix.to_string(),
215                        "meta_client".to_uppercase(),
216                        "metasrv_addrs".to_uppercase(),
217                    ]
218                    .join(ENV_VAR_SEP),
219                    Some("127.0.0.1:3001,127.0.0.1:3002,127.0.0.1:3003"),
220                ),
221            ],
222            || {
223                let opts = TestDatanodeConfig::load_layered_options(
224                    Some(file.path().to_str().unwrap()),
225                    env_prefix,
226                )
227                .unwrap();
228
229                // Check the configs from environment variables.
230                match &opts.storage.store {
231                    ObjectStoreConfig::S3(s3_config) => {
232                        assert_eq!(s3_config.bucket, "mybucket".to_string());
233                    }
234                    _ => panic!("unexpected store type"),
235                }
236                assert_eq!(
237                    opts.meta_client.unwrap().metasrv_addrs,
238                    vec![
239                        "127.0.0.1:3001".to_string(),
240                        "127.0.0.1:3002".to_string(),
241                        "127.0.0.1:3003".to_string()
242                    ]
243                );
244
245                // Should be the values from config file, not environment variables.
246                let DatanodeWalConfig::RaftEngine(raft_engine_config) = opts.wal else {
247                    unreachable!()
248                };
249                assert_eq!(raft_engine_config.dir.unwrap(), "./greptimedb_data/wal");
250
251                // Should be default values.
252                assert_eq!(opts.node_id, None);
253            },
254        );
255    }
256}