sql/statements/
option_map.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 std::collections::{BTreeMap, HashMap};
16use std::ops::ControlFlow;
17
18use common_base::secrets::{ExposeSecret, ExposeSecretMut, SecretString};
19use either::Either;
20use serde::Serialize;
21use sqlparser::ast::{Visit, VisitMut, Visitor, VisitorMut};
22
23use crate::util::OptionValue;
24
25const REDACTED_OPTIONS: [&str; 2] = ["access_key_id", "secret_access_key"];
26
27/// Options hashmap.
28#[derive(Clone, Debug, Default, Serialize)]
29pub struct OptionMap {
30    options: BTreeMap<String, OptionValue>,
31    #[serde(skip_serializing)]
32    secrets: BTreeMap<String, SecretString>,
33}
34
35impl OptionMap {
36    pub fn new<I: IntoIterator<Item = (String, OptionValue)>>(options: I) -> Self {
37        let (secrets, options): (Vec<_>, Vec<_>) = options
38            .into_iter()
39            .partition(|(k, _)| REDACTED_OPTIONS.contains(&k.as_str()));
40        Self {
41            options: options.into_iter().collect(),
42            secrets: secrets
43                .into_iter()
44                .filter_map(|(k, v)| {
45                    v.as_string()
46                        .map(|v| (k, SecretString::new(Box::new(v.to_string()))))
47                })
48                .collect(),
49        }
50    }
51
52    pub fn from_filtered_string_map(
53        options: &HashMap<String, String>,
54        hidden_keys: &[&str],
55    ) -> Self {
56        Self::from(
57            options
58                .iter()
59                .filter(|(key, _)| !hidden_keys.contains(&key.as_str()))
60                .map(|(key, value)| (key.clone(), value.clone())),
61        )
62    }
63
64    pub fn insert(&mut self, k: String, v: String) {
65        if REDACTED_OPTIONS.contains(&k.as_str()) {
66            self.secrets.insert(k, SecretString::new(Box::new(v)));
67        } else {
68            self.options.insert(k, v.into());
69        }
70    }
71
72    pub fn insert_options(&mut self, key: &str, value: OptionValue) {
73        if REDACTED_OPTIONS.contains(&key) {
74            self.secrets.insert(
75                key.to_string(),
76                SecretString::new(Box::new(value.to_string())),
77            );
78        } else {
79            self.options.insert(key.to_string(), value);
80        }
81    }
82
83    pub fn get(&self, k: &str) -> Option<&str> {
84        if let Some(value) = self.options.get(k) {
85            value.as_string()
86        } else if let Some(value) = self.secrets.get(k) {
87            Some(value.expose_secret())
88        } else {
89            None
90        }
91    }
92
93    pub fn value(&self, k: &str) -> Option<&OptionValue> {
94        self.options.get(k)
95    }
96
97    pub fn is_empty(&self) -> bool {
98        self.options.is_empty() && self.secrets.is_empty()
99    }
100
101    pub fn len(&self) -> usize {
102        self.options.len() + self.secrets.len()
103    }
104
105    /// Convert the option map to a string map.
106    ///
107    /// Notes: Not all values can be converted to a string, refer to [OptionValue::expr_as_string] for more details.
108    pub fn to_str_map(&self) -> HashMap<&str, &str> {
109        let mut map = HashMap::with_capacity(self.len());
110        map.extend(
111            self.options
112                .iter()
113                .filter_map(|(k, v)| v.as_string().map(|v| (k.as_str(), v))),
114        );
115        map.extend(
116            self.secrets
117                .iter()
118                .map(|(k, v)| (k.as_str(), v.expose_secret().as_str())),
119        );
120        map
121    }
122
123    /// Convert the option map to a string map.
124    ///
125    /// Notes: Not all values can be converted to a string, refer to [OptionValue::expr_as_string] for more details.
126    pub fn into_map(self) -> HashMap<String, String> {
127        let mut map = HashMap::with_capacity(self.len());
128        map.extend(
129            self.options
130                .into_iter()
131                .filter_map(|(k, v)| v.as_string().map(|v| (k, v.to_string()))),
132        );
133        map.extend(
134            self.secrets
135                .into_iter()
136                .map(|(k, v)| (k, v.expose_secret().clone())),
137        );
138        map
139    }
140
141    pub fn kv_pairs(&self) -> Vec<String> {
142        let mut result = Vec::with_capacity(self.options.len() + self.secrets.len());
143        for (k, v) in self
144            .options
145            .iter()
146            .filter_map(|(k, v)| v.as_string().map(|v| (k, v)))
147        {
148            if k.contains(".") {
149                result.push(format!("'{k}' = '{}'", v.escape_debug()));
150            } else {
151                result.push(format!("{k} = '{}'", v.escape_debug()));
152            }
153        }
154        for (k, _) in self.secrets.iter() {
155            if k.contains(".") {
156                result.push(format!("'{k}' = '******'"));
157            } else {
158                result.push(format!("{k} = '******'"));
159            }
160        }
161        result
162    }
163
164    pub fn entries(&self) -> impl Iterator<Item = (&str, Either<&OptionValue, &str>)> {
165        let options = self
166            .options
167            .iter()
168            .map(|(k, v)| (k.as_str(), Either::Left(v)));
169        let secrets = self
170            .secrets
171            .keys()
172            .map(|k| (k.as_str(), Either::Right("******")));
173        std::iter::chain(options, secrets)
174    }
175}
176
177impl<I: IntoIterator<Item = (String, String)>> From<I> for OptionMap {
178    fn from(value: I) -> Self {
179        let mut result = OptionMap::default();
180        for (k, v) in value {
181            result.insert(k, v);
182        }
183        result
184    }
185}
186
187impl PartialEq for OptionMap {
188    fn eq(&self, other: &Self) -> bool {
189        if self.options.ne(&other.options) {
190            return false;
191        }
192
193        if self.secrets.len() != other.secrets.len() {
194            return false;
195        }
196
197        self.secrets.iter().all(|(key, value)| {
198            other
199                .secrets
200                .get(key)
201                .is_some_and(|v| value.expose_secret() == v.expose_secret())
202        })
203    }
204}
205
206impl Eq for OptionMap {}
207
208impl Visit for OptionMap {
209    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
210        for (k, v) in &self.options {
211            k.visit(visitor)?;
212            v.visit(visitor)?;
213        }
214        for (k, v) in &self.secrets {
215            k.visit(visitor)?;
216            v.expose_secret().visit(visitor)?;
217        }
218        ControlFlow::Continue(())
219    }
220}
221
222impl VisitMut for OptionMap {
223    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
224        for (_, v) in self.options.iter_mut() {
225            v.visit(visitor)?;
226        }
227        for (_, v) in self.secrets.iter_mut() {
228            v.expose_secret_mut().visit(visitor)?;
229        }
230        ControlFlow::Continue(())
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use std::collections::HashMap;
237
238    use crate::statements::OptionMap;
239
240    #[test]
241    fn test_format() {
242        let mut map = OptionMap::default();
243        map.insert("comment".to_string(), "中文comment".to_string());
244        assert_eq!("comment = '中文comment'", map.kv_pairs()[0]);
245
246        let mut map = OptionMap::default();
247        map.insert("a.b".to_string(), "中文comment".to_string());
248        assert_eq!("'a.b' = '中文comment'", map.kv_pairs()[0]);
249
250        let mut map = OptionMap::default();
251        map.insert("a.b".to_string(), "中文comment\n".to_string());
252        assert_eq!("'a.b' = '中文comment\\n'", map.kv_pairs()[0]);
253    }
254
255    #[test]
256    fn test_from_filtered_string_map() {
257        let map = OptionMap::from_filtered_string_map(
258            &HashMap::from([
259                ("visible".to_string(), "1".to_string()),
260                ("hidden".to_string(), "2".to_string()),
261            ]),
262            &["hidden"],
263        );
264
265        assert_eq!(map.get("visible"), Some("1"));
266        assert_eq!(map.get("hidden"), None);
267    }
268}
sql/statements/option_map.rs

sql/statements/
option_map.rs
