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        self.sql_pairs(false).map(|pair| pair.to_string()).collect()
143    }
144
145    /// Streams the same pairs as `kv_pairs`, or all actual values for size validation.
146    pub(crate) fn write_sql(
147        &self,
148        f: &mut dyn std::fmt::Write,
149        separator: &str,
150        prefix: &str,
151        complete: bool,
152    ) -> std::fmt::Result {
153        for (index, pair) in self.sql_pairs(complete).enumerate() {
154            if index > 0 {
155                f.write_str(separator)?;
156            }
157            write!(f, "{prefix}{pair}")?;
158        }
159        Ok(())
160    }
161
162    fn sql_pairs(&self, complete: bool) -> impl Iterator<Item = SqlPair<'_>> {
163        let options = self.options.iter().filter_map(move |(key, value)| {
164            let value = match value.as_string() {
165                Some(value) => Either::Right(value),
166                None if complete => Either::Left(value),
167                None => return None,
168            };
169            Some(SqlPair { key, value })
170        });
171        let secrets = self.secrets.iter().map(move |(key, value)| SqlPair {
172            key,
173            value: Either::Right(if complete {
174                value.expose_secret()
175            } else {
176                "******"
177            }),
178        });
179        options.chain(secrets)
180    }
181
182    pub fn entries(&self) -> impl Iterator<Item = (&str, Either<&OptionValue, &str>)> {
183        let options = self
184            .options
185            .iter()
186            .map(|(k, v)| (k.as_str(), Either::Left(v)));
187        let secrets = self
188            .secrets
189            .keys()
190            .map(|k| (k.as_str(), Either::Right("******")));
191        std::iter::chain(options, secrets)
192    }
193}
194
195struct SqlPair<'a> {
196    key: &'a str,
197    value: Either<&'a OptionValue, &'a str>,
198}
199
200impl std::fmt::Display for SqlPair<'_> {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        if self.key.contains('.') {
203            write!(f, "'{}'", self.key)?;
204        } else {
205            f.write_str(self.key)?;
206        }
207        match self.value {
208            Either::Left(value) => write!(f, " = {value}"),
209            Either::Right(value) => write!(f, " = '{}'", value.escape_debug()),
210        }
211    }
212}
213
214impl<I: IntoIterator<Item = (String, String)>> From<I> for OptionMap {
215    fn from(value: I) -> Self {
216        let mut result = OptionMap::default();
217        for (k, v) in value {
218            result.insert(k, v);
219        }
220        result
221    }
222}
223
224impl PartialEq for OptionMap {
225    fn eq(&self, other: &Self) -> bool {
226        if self.options.ne(&other.options) {
227            return false;
228        }
229
230        if self.secrets.len() != other.secrets.len() {
231            return false;
232        }
233
234        self.secrets.iter().all(|(key, value)| {
235            other
236                .secrets
237                .get(key)
238                .is_some_and(|v| value.expose_secret() == v.expose_secret())
239        })
240    }
241}
242
243impl Eq for OptionMap {}
244
245impl Visit for OptionMap {
246    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
247        for (k, v) in &self.options {
248            k.visit(visitor)?;
249            v.visit(visitor)?;
250        }
251        for (k, v) in &self.secrets {
252            k.visit(visitor)?;
253            v.expose_secret().visit(visitor)?;
254        }
255        ControlFlow::Continue(())
256    }
257}
258
259impl VisitMut for OptionMap {
260    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
261        for v in self.options.values_mut() {
262            v.visit(visitor)?;
263        }
264        for v in self.secrets.values_mut() {
265            v.expose_secret_mut().visit(visitor)?;
266        }
267        ControlFlow::Continue(())
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use std::collections::HashMap;
274
275    use crate::statements::OptionMap;
276
277    #[test]
278    fn test_format() {
279        let mut map = OptionMap::default();
280        map.insert("comment".to_string(), "中文comment".to_string());
281        assert_eq!("comment = '中文comment'", map.kv_pairs()[0]);
282
283        let mut map = OptionMap::default();
284        map.insert("a.b".to_string(), "中文comment".to_string());
285        assert_eq!("'a.b' = '中文comment'", map.kv_pairs()[0]);
286
287        let mut map = OptionMap::default();
288        map.insert("a.b".to_string(), "中文comment\n".to_string());
289        assert_eq!("'a.b' = '中文comment\\n'", map.kv_pairs()[0]);
290    }
291
292    #[test]
293    fn test_from_filtered_string_map() {
294        let map = OptionMap::from_filtered_string_map(
295            &HashMap::from([
296                ("visible".to_string(), "1".to_string()),
297                ("hidden".to_string(), "2".to_string()),
298            ]),
299            &["hidden"],
300        );
301
302        assert_eq!(map.get("visible"), Some("1"));
303        assert_eq!(map.get("hidden"), None);
304    }
305}
sql/statements/option_map.rs

sql/statements/
option_map.rs
