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
115 #[derive(Debug, Serialize, Deserialize, Default)]
116 struct TestDatanodeConfig {
117 node_id: Option<u64>,
118 logging: LoggingOptions,
119 meta_client: Option<MetaClientOptions>,
120 wal: DatanodeWalConfig,
121 storage: StorageConfig,
122 }
123
124 impl Configurable for TestDatanodeConfig {
125 fn env_list_keys() -> Option<&'static [&'static str]> {
126 Some(&["meta_client.metasrv_addrs"])
127 }
128 }
129
130 #[test]
131 fn test_load_layered_options() {
132 let mut file = create_named_temp_file();
133 let toml_str = r#"
134 enable_memory_catalog = false
135 rpc_addr = "127.0.0.1:3001"
136 rpc_hostname = "127.0.0.1"
137 rpc_runtime_size = 8
138 mysql_addr = "127.0.0.1:4406"
139 mysql_runtime_size = 2
140
141 [meta_client]
142 timeout = "3s"
143 connect_timeout = "5s"
144 tcp_nodelay = true
145
146 [wal]
147 provider = "raft_engine"
148 dir = "./greptimedb_data/wal"
149 file_size = "1GB"
150 purge_threshold = "50GB"
151 purge_interval = "10m"
152 read_batch_size = 128
153 sync_write = false
154
155 [logging]
156 level = "debug"
157 dir = "./greptimedb_data/test/logs"
158 "#;
159 write!(file, "{}", toml_str).unwrap();
160
161 let env_prefix = "DATANODE_UT";
162 temp_env::with_vars(
163 // The following environment variables will be used to override the values in the config file.
164 [
165 (
166 // storage.type = S3
167 [
168 env_prefix.to_string(),
169 "storage".to_uppercase(),
170 "type".to_uppercase(),
171 ]
172 .join(ENV_VAR_SEP),
173 Some("S3"),
174 ),
175 (
176 // storage.bucket = mybucket
177 [
178 env_prefix.to_string(),
179 "storage".to_uppercase(),
180 "bucket".to_uppercase(),
181 ]
182 .join(ENV_VAR_SEP),
183 Some("mybucket"),
184 ),
185 (
186 // wal.dir = /other/wal/dir
187 [
188 env_prefix.to_string(),
189 "wal".to_uppercase(),
190 "dir".to_uppercase(),
191 ]
192 .join(ENV_VAR_SEP),
193 Some("/other/wal/dir"),
194 ),
195 (
196 // meta_client.metasrv_addrs = 127.0.0.1:3001,127.0.0.1:3002,127.0.0.1:3003
197 [
198 env_prefix.to_string(),
199 "meta_client".to_uppercase(),
200 "metasrv_addrs".to_uppercase(),
201 ]
202 .join(ENV_VAR_SEP),
203 Some("127.0.0.1:3001,127.0.0.1:3002,127.0.0.1:3003"),
204 ),
205 ],
206 || {
207 let opts = TestDatanodeConfig::load_layered_options(
208 Some(file.path().to_str().unwrap()),
209 env_prefix,
210 )
211 .unwrap();
212
213 // Check the configs from environment variables.
214 match &opts.storage.store {
215 ObjectStoreConfig::S3(s3_config) => {
216 assert_eq!(s3_config.bucket, "mybucket".to_string());
217 }
218 _ => panic!("unexpected store type"),
219 }
220 assert_eq!(
221 opts.meta_client.unwrap().metasrv_addrs,
222 vec![
223 "127.0.0.1:3001".to_string(),
224 "127.0.0.1:3002".to_string(),
225 "127.0.0.1:3003".to_string()
226 ]
227 );
228
229 // Should be the values from config file, not environment variables.
230 let DatanodeWalConfig::RaftEngine(raft_engine_config) = opts.wal else {
231 unreachable!()
232 };
233 assert_eq!(raft_engine_config.dir.unwrap(), "./greptimedb_data/wal");
234
235 // Should be default values.
236 assert_eq!(opts.node_id, None);
237 },
238 );
239 }
240}