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 insert(&mut self, k: String, v: String) {
53        if REDACTED_OPTIONS.contains(&k.as_str()) {
54            self.secrets.insert(k, SecretString::new(Box::new(v)));
55        } else {
56            self.options.insert(k, v.into());
57        }
58    }
59
60    pub fn insert_options(&mut self, key: &str, value: OptionValue) {
61        if REDACTED_OPTIONS.contains(&key) {
62            self.secrets.insert(
63                key.to_string(),
64                SecretString::new(Box::new(value.to_string())),
65            );
66        } else {
67            self.options.insert(key.to_string(), value);
68        }
69    }
70
71    pub fn get(&self, k: &str) -> Option<&str> {
72        if let Some(value) = self.options.get(k) {
73            value.as_string()
74        } else if let Some(value) = self.secrets.get(k) {
75            Some(value.expose_secret())
76        } else {
77            None
78        }
79    }
80
81    pub fn value(&self, k: &str) -> Option<&OptionValue> {
82        self.options.get(k)
83    }
84
85    pub fn is_empty(&self) -> bool {
86        self.options.is_empty() && self.secrets.is_empty()
87    }
88
89    pub fn len(&self) -> usize {
90        self.options.len() + self.secrets.len()
91    }
92
93    pub fn to_str_map(&self) -> HashMap<&str, &str> {
94        let mut map = HashMap::with_capacity(self.len());
95        map.extend(
96            self.options
97                .iter()
98                .filter_map(|(k, v)| v.as_string().map(|v| (k.as_str(), v))),
99        );
100        map.extend(
101            self.secrets
102                .iter()
103                .map(|(k, v)| (k.as_str(), v.expose_secret().as_str())),
104        );
105        map
106    }
107
108    pub fn into_map(self) -> HashMap<String, String> {
109        let mut map = HashMap::with_capacity(self.len());
110        map.extend(
111            self.options
112                .into_iter()
113                .filter_map(|(k, v)| v.as_string().map(|v| (k, v.to_string()))),
114        );
115        map.extend(
116            self.secrets
117                .into_iter()
118                .map(|(k, v)| (k, v.expose_secret().clone())),
119        );
120        map
121    }
122
123    pub fn kv_pairs(&self) -> Vec<String> {
124        let mut result = Vec::with_capacity(self.options.len() + self.secrets.len());
125        for (k, v) in self
126            .options
127            .iter()
128            .filter_map(|(k, v)| v.as_string().map(|v| (k, v)))
129        {
130            if k.contains(".") {
131                result.push(format!("'{k}' = '{}'", v.escape_debug()));
132            } else {
133                result.push(format!("{k} = '{}'", v.escape_debug()));
134            }
135        }
136        for (k, _) in self.secrets.iter() {
137            if k.contains(".") {
138                result.push(format!("'{k}' = '******'"));
139            } else {
140                result.push(format!("{k} = '******'"));
141            }
142        }
143        result
144    }
145
146    pub fn entries(&self) -> impl Iterator<Item = (&str, Either<&OptionValue, &str>)> {
147        let options = self
148            .options
149            .iter()
150            .map(|(k, v)| (k.as_str(), Either::Left(v)));
151        let secrets = self
152            .secrets
153            .keys()
154            .map(|k| (k.as_str(), Either::Right("******")));
155        std::iter::chain(options, secrets)
156    }
157}
158
159impl<I: IntoIterator<Item = (String, String)>> From<I> for OptionMap {
160    fn from(value: I) -> Self {
161        let mut result = OptionMap::default();
162        for (k, v) in value {
163            result.insert(k, v);
164        }
165        result
166    }
167}
168
169impl PartialEq for OptionMap {
170    fn eq(&self, other: &Self) -> bool {
171        if self.options.ne(&other.options) {
172            return false;
173        }
174
175        if self.secrets.len() != other.secrets.len() {
176            return false;
177        }
178
179        self.secrets.iter().all(|(key, value)| {
180            other
181                .secrets
182                .get(key)
183                .is_some_and(|v| value.expose_secret() == v.expose_secret())
184        })
185    }
186}
187
188impl Eq for OptionMap {}
189
190impl Visit for OptionMap {
191    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
192        for (k, v) in &self.options {
193            k.visit(visitor)?;
194            v.visit(visitor)?;
195        }
196        for (k, v) in &self.secrets {
197            k.visit(visitor)?;
198            v.expose_secret().visit(visitor)?;
199        }
200        ControlFlow::Continue(())
201    }
202}
203
204impl VisitMut for OptionMap {
205    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
206        for (_, v) in self.options.iter_mut() {
207            v.visit(visitor)?;
208        }
209        for (_, v) in self.secrets.iter_mut() {
210            v.expose_secret_mut().visit(visitor)?;
211        }
212        ControlFlow::Continue(())
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use crate::statements::OptionMap;
219
220    #[test]
221    fn test_format() {
222        let mut map = OptionMap::default();
223        map.insert("comment".to_string(), "中文comment".to_string());
224        assert_eq!("comment = '中文comment'", map.kv_pairs()[0]);
225
226        let mut map = OptionMap::default();
227        map.insert("a.b".to_string(), "中文comment".to_string());
228        assert_eq!("'a.b' = '中文comment'", map.kv_pairs()[0]);
229
230        let mut map = OptionMap::default();
231        map.insert("a.b".to_string(), "中文comment\n".to_string());
232        assert_eq!("'a.b' = '中文comment\\n'", map.kv_pairs()[0]);
233    }
234}
sql/statements/option_map.rs

sql/statements/
option_map.rs
