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 serde::Serialize;
20use sqlparser::ast::{Visit, VisitMut, Visitor, VisitorMut};
21
22use crate::util::OptionValue;
23
24const REDACTED_OPTIONS: [&str; 2] = ["access_key_id", "secret_access_key"];
25
26/// Options hashmap.
27#[derive(Clone, Debug, Default, Serialize)]
28pub struct OptionMap {
29    options: BTreeMap<String, OptionValue>,
30    #[serde(skip_serializing)]
31    secrets: BTreeMap<String, SecretString>,
32}
33
34impl OptionMap {
35    pub fn new<I: IntoIterator<Item = (String, OptionValue)>>(options: I) -> Self {
36        let (secrets, options): (Vec<_>, Vec<_>) = options
37            .into_iter()
38            .partition(|(k, _)| REDACTED_OPTIONS.contains(&k.as_str()));
39        Self {
40            options: options.into_iter().collect(),
41            secrets: secrets
42                .into_iter()
43                .filter_map(|(k, v)| {
44                    v.as_string()
45                        .map(|v| (k, SecretString::new(Box::new(v.to_string()))))
46                })
47                .collect(),
48        }
49    }
50
51    pub fn insert(&mut self, k: String, v: String) {
52        if REDACTED_OPTIONS.contains(&k.as_str()) {
53            self.secrets.insert(k, SecretString::new(Box::new(v)));
54        } else {
55            self.options.insert(k, v.into());
56        }
57    }
58
59    pub fn get(&self, k: &str) -> Option<&str> {
60        if let Some(value) = self.options.get(k) {
61            value.as_string()
62        } else if let Some(value) = self.secrets.get(k) {
63            Some(value.expose_secret())
64        } else {
65            None
66        }
67    }
68
69    pub fn value(&self, k: &str) -> Option<&OptionValue> {
70        self.options.get(k)
71    }
72
73    pub fn is_empty(&self) -> bool {
74        self.options.is_empty() && self.secrets.is_empty()
75    }
76
77    pub fn len(&self) -> usize {
78        self.options.len() + self.secrets.len()
79    }
80
81    pub fn to_str_map(&self) -> HashMap<&str, &str> {
82        let mut map = HashMap::with_capacity(self.len());
83        map.extend(
84            self.options
85                .iter()
86                .filter_map(|(k, v)| v.as_string().map(|v| (k.as_str(), v))),
87        );
88        map.extend(
89            self.secrets
90                .iter()
91                .map(|(k, v)| (k.as_str(), v.expose_secret().as_str())),
92        );
93        map
94    }
95
96    pub fn into_map(self) -> HashMap<String, String> {
97        let mut map = HashMap::with_capacity(self.len());
98        map.extend(
99            self.options
100                .into_iter()
101                .filter_map(|(k, v)| v.as_string().map(|v| (k, v.to_string()))),
102        );
103        map.extend(
104            self.secrets
105                .into_iter()
106                .map(|(k, v)| (k, v.expose_secret().clone())),
107        );
108        map
109    }
110
111    pub fn kv_pairs(&self) -> Vec<String> {
112        let mut result = Vec::with_capacity(self.options.len() + self.secrets.len());
113        for (k, v) in self
114            .options
115            .iter()
116            .filter_map(|(k, v)| v.as_string().map(|v| (k, v)))
117        {
118            if k.contains(".") {
119                result.push(format!("'{k}' = '{}'", v.escape_debug()));
120            } else {
121                result.push(format!("{k} = '{}'", v.escape_debug()));
122            }
123        }
124        for (k, _) in self.secrets.iter() {
125            if k.contains(".") {
126                result.push(format!("'{k}' = '******'"));
127            } else {
128                result.push(format!("{k} = '******'"));
129            }
130        }
131        result
132    }
133}
134
135impl<I: IntoIterator<Item = (String, String)>> From<I> for OptionMap {
136    fn from(value: I) -> Self {
137        let mut result = OptionMap::default();
138        for (k, v) in value {
139            result.insert(k, v);
140        }
141        result
142    }
143}
144
145impl PartialEq for OptionMap {
146    fn eq(&self, other: &Self) -> bool {
147        if self.options.ne(&other.options) {
148            return false;
149        }
150
151        if self.secrets.len() != other.secrets.len() {
152            return false;
153        }
154
155        self.secrets.iter().all(|(key, value)| {
156            other
157                .secrets
158                .get(key)
159                .is_some_and(|v| value.expose_secret() == v.expose_secret())
160        })
161    }
162}
163
164impl Eq for OptionMap {}
165
166impl Visit for OptionMap {
167    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
168        for (k, v) in &self.options {
169            k.visit(visitor)?;
170            v.visit(visitor)?;
171        }
172        for (k, v) in &self.secrets {
173            k.visit(visitor)?;
174            v.expose_secret().visit(visitor)?;
175        }
176        ControlFlow::Continue(())
177    }
178}
179
180impl VisitMut for OptionMap {
181    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
182        for (_, v) in self.options.iter_mut() {
183            v.visit(visitor)?;
184        }
185        for (_, v) in self.secrets.iter_mut() {
186            v.expose_secret_mut().visit(visitor)?;
187        }
188        ControlFlow::Continue(())
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use crate::statements::OptionMap;
195
196    #[test]
197    fn test_format() {
198        let mut map = OptionMap::default();
199        map.insert("comment".to_string(), "中文comment".to_string());
200        assert_eq!("comment = '中文comment'", map.kv_pairs()[0]);
201
202        let mut map = OptionMap::default();
203        map.insert("a.b".to_string(), "中文comment".to_string());
204        assert_eq!("'a.b' = '中文comment'", map.kv_pairs()[0]);
205
206        let mut map = OptionMap::default();
207        map.insert("a.b".to_string(), "中文comment\n".to_string());
208        assert_eq!("'a.b' = '中文comment\\n'", map.kv_pairs()[0]);
209    }
210}
sql/statements/option_map.rs

sql/statements/
option_map.rs
