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

sql/statements/
option_map.rs
