sql/statements/
option_map.rs1use 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#[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}