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}