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    /// Convert the option map to a string map.
94    ///
95    /// Notes: Not all values can be converted to a string, refer to [OptionValue::expr_as_string] for more details.
96    pub fn to_str_map(&self) -> HashMap<&str, &str> {
97        let mut map = HashMap::with_capacity(self.len());
98        map.extend(
99            self.options
100                .iter()
101                .filter_map(|(k, v)| v.as_string().map(|v| (k.as_str(), v))),
102        );
103        map.extend(
104            self.secrets
105                .iter()
106                .map(|(k, v)| (k.as_str(), v.expose_secret().as_str())),
107        );
108        map
109    }
110
111    /// Convert the option map to a string map.
112    ///
113    /// Notes: Not all values can be converted to a string, refer to [OptionValue::expr_as_string] for more details.
114    pub fn into_map(self) -> HashMap<String, String> {
115        let mut map = HashMap::with_capacity(self.len());
116        map.extend(
117            self.options
118                .into_iter()
119                .filter_map(|(k, v)| v.as_string().map(|v| (k, v.to_string()))),
120        );
121        map.extend(
122            self.secrets
123                .into_iter()
124                .map(|(k, v)| (k, v.expose_secret().clone())),
125        );
126        map
127    }
128
129    pub fn kv_pairs(&self) -> Vec<String> {
130        let mut result = Vec::with_capacity(self.options.len() + self.secrets.len());
131        for (k, v) in self
132            .options
133            .iter()
134            .filter_map(|(k, v)| v.as_string().map(|v| (k, v)))
135        {
136            if k.contains(".") {
137                result.push(format!("'{k}' = '{}'", v.escape_debug()));
138            } else {
139                result.push(format!("{k} = '{}'", v.escape_debug()));
140            }
141        }
142        for (k, _) in self.secrets.iter() {
143            if k.contains(".") {
144                result.push(format!("'{k}' = '******'"));
145            } else {
146                result.push(format!("{k} = '******'"));
147            }
148        }
149        result
150    }
151
152    pub fn entries(&self) -> impl Iterator<Item = (&str, Either<&OptionValue, &str>)> {
153        let options = self
154            .options
155            .iter()
156            .map(|(k, v)| (k.as_str(), Either::Left(v)));
157        let secrets = self
158            .secrets
159            .keys()
160            .map(|k| (k.as_str(), Either::Right("******")));
161        std::iter::chain(options, secrets)
162    }
163}
164
165impl<I: IntoIterator<Item = (String, String)>> From<I> for OptionMap {
166    fn from(value: I) -> Self {
167        let mut result = OptionMap::default();
168        for (k, v) in value {
169            result.insert(k, v);
170        }
171        result
172    }
173}
174
175impl PartialEq for OptionMap {
176    fn eq(&self, other: &Self) -> bool {
177        if self.options.ne(&other.options) {
178            return false;
179        }
180
181        if self.secrets.len() != other.secrets.len() {
182            return false;
183        }
184
185        self.secrets.iter().all(|(key, value)| {
186            other
187                .secrets
188                .get(key)
189                .is_some_and(|v| value.expose_secret() == v.expose_secret())
190        })
191    }
192}
193
194impl Eq for OptionMap {}
195
196impl Visit for OptionMap {
197    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
198        for (k, v) in &self.options {
199            k.visit(visitor)?;
200            v.visit(visitor)?;
201        }
202        for (k, v) in &self.secrets {
203            k.visit(visitor)?;
204            v.expose_secret().visit(visitor)?;
205        }
206        ControlFlow::Continue(())
207    }
208}
209
210impl VisitMut for OptionMap {
211    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
212        for (_, v) in self.options.iter_mut() {
213            v.visit(visitor)?;
214        }
215        for (_, v) in self.secrets.iter_mut() {
216            v.expose_secret_mut().visit(visitor)?;
217        }
218        ControlFlow::Continue(())
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use crate::statements::OptionMap;
225
226    #[test]
227    fn test_format() {
228        let mut map = OptionMap::default();
229        map.insert("comment".to_string(), "中文comment".to_string());
230        assert_eq!("comment = '中文comment'", map.kv_pairs()[0]);
231
232        let mut map = OptionMap::default();
233        map.insert("a.b".to_string(), "中文comment".to_string());
234        assert_eq!("'a.b' = '中文comment'", map.kv_pairs()[0]);
235
236        let mut map = OptionMap::default();
237        map.insert("a.b".to_string(), "中文comment\n".to_string());
238        assert_eq!("'a.b' = '中文comment\\n'", map.kv_pairs()[0]);
239    }
240}
sql/statements/option_map.rs

sql/statements/
option_map.rs
