rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-06-03 13:50:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: use secrecy SerectString to hold secrets option (#3804)

Browse Source 
* build: centralize secrecy dependency

Signed-off-by: tison <wander4096@gmail.com>

* add secrecy to sql crate

Signed-off-by: tison <wander4096@gmail.com>

* try impl

Signed-off-by: tison <wander4096@gmail.com>

* update test

Signed-off-by: tison <wander4096@gmail.com>

* make linters happy

Signed-off-by: tison <wander4096@gmail.com>

* bundle secrecy

Signed-off-by: tison <wander4096@gmail.com>

* bundle secrecy

Signed-off-by: tison <wander4096@gmail.com>

* replace secrecy

Signed-off-by: tison <wander4096@gmail.com>

* tidy clones

Signed-off-by: tison <wander4096@gmail.com>

* fixup

Signed-off-by: tison <wander4096@gmail.com>

* fixup

Signed-off-by: tison <wander4096@gmail.com>

* updated

Signed-off-by: tison <wander4096@gmail.com>

* Apply suggestions from code review

Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com>

* use BTreeMap

Signed-off-by: tison <wander4096@gmail.com>

* tidy

Signed-off-by: tison <wander4096@gmail.com>

---------

Signed-off-by: tison <wander4096@gmail.com>
Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com>
This commit is contained in:
tison
2024-04-29 10:18:18 +08:00
committed by GitHub
parent 7ef18c0915
commit c387687262
33 changed files with 419 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/sql/src/parsers/copy_parser.rs
View File

@@ -378,20 +378,17 @@ mod tests {
stmt.database_name
);
assert_eq!(
[("format".to_string(), "parquet".to_string())]
[("format", "parquet")]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.with.map
stmt.with.to_str_map()
);
assert_eq!(
[
("foo".to_string(), "Bar".to_string()),
("one".to_string(), "two".to_string())
]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.connection.map
[("foo", "Bar"), ("one", "two")]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.connection.to_str_map()
);
}
@@ -417,20 +414,17 @@ mod tests {
stmt.database_name
);
assert_eq!(
[("format".to_string(), "parquet".to_string())]
[("format", "parquet")]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.with.map
stmt.with.to_str_map()
);
assert_eq!(
[
("foo".to_string(), "Bar".to_string()),
("one".to_string(), "two".to_string())
]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.connection.map
[("foo", "Bar"), ("one", "two")]
.into_iter()
.collect::<HashMap<_, _>>(),
stmt.connection.to_str_map()
);
}
}

9
src/sql/src/parsers/create_parser.rs
View File

@@ -1638,10 +1638,11 @@ ENGINE=mito";
..
}
);
let options = &c.options;
assert_eq!(1, options.map.len());
let (k, v) = options.map.iter().next().unwrap();
assert_eq!(("ttl", "10s"), (k.as_str(), v.as_str()));
assert_eq!(1, c.options.len());
assert_eq!(
[("ttl", "10s")].into_iter().collect::<HashMap<_, _>>(),
c.options.to_str_map()
);
}
_ => unreachable!(),
}

24
src/sql/src/statements.rs
View File

@@ -43,7 +43,6 @@ use datatypes::schema::constraint::{CURRENT_TIMESTAMP, CURRENT_TIMESTAMP_FN};
use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, COMMENT_KEY};
use datatypes::types::{cast, TimestampType};
use datatypes::value::{OrderedF32, OrderedF64, Value};
use itertools::Itertools;
pub use option_map::OptionMap;
use snafu::{ensure, OptionExt, ResultExt};
use sqlparser::ast::{ExactNumberInfo, UnaryOperator};
@@ -59,29 +58,6 @@ use crate::error::{
SerializeColumnDefaultConstraintSnafu, TimestampOverflowSnafu, UnsupportedDefaultValueSnafu,
};
const REDACTED_OPTIONS: [&str; 2] = ["access_key_id", "secret_access_key"];
/// Convert the options into redacted and sorted key-value string. Options with key in
/// [REDACTED_OPTIONS] will be converted into `<key> = '******'`.
fn redact_and_sort_options(options: &OptionMap) -> Vec<String> {
let options = options.as_ref();
let mut result = Vec::with_capacity(options.len());
let keys = options.keys().sorted();
for key in keys {
if let Some(val) = options.get(key) {
let redacted = REDACTED_OPTIONS
.iter()
.any(|opt| opt.eq_ignore_ascii_case(key));
if redacted {
result.push(format!("{key} = '******'"));
} else {
result.push(format!("{key} = '{}'", val.escape_default()));
}
}
}
result
}
fn parse_string_to_value(
column_name: &str,
s: String,

18
src/sql/src/statements/copy.rs
View File

@@ -17,7 +17,7 @@ use std::fmt::Display;
use sqlparser::ast::ObjectName;
use sqlparser_derive::{Visit, VisitMut};
use crate::statements::{redact_and_sort_options, OptionMap};
use crate::statements::OptionMap;
#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
pub enum Copy {
@@ -53,12 +53,12 @@ impl Display for CopyTable {
(&args.with, &args.connection)
}
};
if !with.map.is_empty() {
let options = redact_and_sort_options(with);
if !with.is_empty() {
let options = with.kv_pairs();
write!(f, " WITH ({})", options.join(", "))?;
}
if !connection.map.is_empty() {
let options = redact_and_sort_options(connection);
if !connection.is_empty() {
let options = connection.kv_pairs();
write!(f, " CONNECTION ({})", options.join(", "))?;
}
Ok(())
@@ -84,12 +84,12 @@ impl Display for CopyDatabase {
(&args.with, &args.connection)
}
};
if !with.map.is_empty() {
let options = redact_and_sort_options(with);
if !with.is_empty() {
let options = with.kv_pairs();
write!(f, " WITH ({})", options.join(", "))?;
}
if !connection.map.is_empty() {
let options = redact_and_sort_options(connection);
if !connection.is_empty() {
let options = connection.kv_pairs();
write!(f, " CONNECTION ({})", options.join(", "))?;
}
Ok(())

12
src/sql/src/statements/create.rs
View File

@@ -20,7 +20,7 @@ use sqlparser::ast::{Expr, Query};
use sqlparser_derive::{Visit, VisitMut};
use crate::ast::{ColumnDef, Ident, ObjectName, TableConstraint, Value as SqlValue};
use crate::statements::{redact_and_sort_options, OptionMap};
use crate::statements::OptionMap;
const LINE_SEP: &str = ",\n";
const COMMA_SEP: &str = ", ";
@@ -155,8 +155,8 @@ impl Display for CreateTable {
writeln!(f, "{partitions}")?;
}
writeln!(f, "ENGINE={}", &self.engine)?;
if !self.options.map.is_empty() {
let options = redact_and_sort_options(&self.options);
if !self.options.is_empty() {
let options = self.options.kv_pairs();
write!(f, "WITH(\n{}\n)", format_list_indent!(options))?;
}
Ok(())
@@ -213,8 +213,8 @@ impl Display for CreateExternalTable {
writeln!(f, "{}", format_table_constraint(&self.constraints))?;
writeln!(f, ")")?;
writeln!(f, "ENGINE={}", &self.engine)?;
if !self.options.map.is_empty() {
let options = redact_and_sort_options(&self.options);
if !self.options.is_empty() {
let options = self.options.kv_pairs();
write!(f, "WITH(\n{}\n)", format_list_indent!(options))?;
}
Ok(())
@@ -410,7 +410,7 @@ ENGINE=mito
.unwrap();
match &result[0] {
Statement::CreateTable(c) => {
assert_eq!(2, c.options.map.len());
assert_eq!(2, c.options.len());
}
_ => unreachable!(),
}

145
src/sql/src/statements/option_map.rs
View File

@@ -12,52 +12,135 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod visit;
mod visit_mut;
use std::collections::{BTreeMap, HashMap};
use std::ops::ControlFlow;
use std::borrow::Borrow;
use std::collections::HashMap;
use std::iter::FromIterator;
use common_base::secrets::{ExposeSecret, ExposeSecretMut, SecretString};
use sqlparser::ast::{Visit, VisitMut, Visitor, VisitorMut};
const REDACTED_OPTIONS: [&str; 2] = ["access_key_id", "secret_access_key"];
/// Options hashmap.
/// Because the trait `Visit` and `VisitMut` is not implemented for `HashMap<String, String>`, we have to wrap it and implement them by ourself.
#[derive(Clone, Eq, PartialEq, Debug)]
#[derive(Clone, Debug, Default)]
pub struct OptionMap {
pub map: HashMap<String, String>,
options: BTreeMap<String, String>,
secrets: BTreeMap<String, SecretString>,
}
impl OptionMap {
pub fn insert(&mut self, k: String, v: String) {
self.map.insert(k, v);
if REDACTED_OPTIONS.contains(&k.as_str()) {
self.secrets.insert(k, SecretString::new(Box::new(v)));
} else {
self.options.insert(k, v);
}
}
pub fn get(&self, k: &str) -> Option<&String> {
self.map.get(k)
if let Some(value) = self.options.get(k) {
Some(value)
} else if let Some(value) = self.secrets.get(k) {
Some(value.expose_secret())
} else {
None
}
}
pub fn is_empty(&self) -> bool {
self.options.is_empty() && self.secrets.is_empty()
}
pub fn len(&self) -> usize {
self.options.len() + self.secrets.len()
}
pub fn to_str_map(&self) -> HashMap<&str, &str> {
let mut map = HashMap::with_capacity(self.len());
map.extend(self.options.iter().map(|(k, v)| (k.as_str(), v.as_str())));
map.extend(
self.secrets
.iter()
.map(|(k, v)| (k.as_str(), v.expose_secret().as_str())),
);
map
}
pub fn into_map(self) -> HashMap<String, String> {
let mut map = HashMap::with_capacity(self.len());
map.extend(self.options);
map.extend(
self.secrets
.into_iter()
.map(|(k, v)| (k, v.expose_secret().to_string())),
);
map
}
pub fn kv_pairs(&self) -> Vec<String> {
let mut result = Vec::with_capacity(self.options.len() + self.secrets.len());
for (k, v) in self.options.iter() {
result.push(format!("{k} = '{}'", v.escape_default()));
}
for (k, _) in self.secrets.iter() {
result.push(format!("{k} = '******'"));
}
result
}
}
impl From<HashMap<String, String>> for OptionMap {
fn from(map: HashMap<String, String>) -> Self {
Self { map }
}
}
impl AsRef<HashMap<String, String>> for OptionMap {
fn as_ref(&self) -> &HashMap<String, String> {
&self.map
}
}
impl Borrow<HashMap<String, String>> for OptionMap {
fn borrow(&self) -> &HashMap<String, String> {
&self.map
}
}
impl FromIterator<(String, String)> for OptionMap {
fn from_iter<I: IntoIterator<Item = (String, String)>>(iter: I) -> Self {
Self {
map: iter.into_iter().collect(),
fn from(value: HashMap<String, String>) -> Self {
let mut result = OptionMap::default();
for (k, v) in value.into_iter() {
result.insert(k, v);
}
result
}
}
impl PartialEq for OptionMap {
fn eq(&self, other: &Self) -> bool {
if self.options.ne(&other.options) {
return false;
}
if self.secrets.len() != other.secrets.len() {
return false;
}
self.secrets.iter().all(|(key, value)| {
other
.secrets
.get(key)
.map_or(false, |v| value.expose_secret() == v.expose_secret())
})
}
}
impl Eq for OptionMap {}
impl Visit for OptionMap {
fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
for (k, v) in &self.options {
k.visit(visitor)?;
v.visit(visitor)?;
}
for (k, v) in &self.secrets {
k.visit(visitor)?;
v.expose_secret().visit(visitor)?;
}
ControlFlow::Continue(())
}
}
impl VisitMut for OptionMap {
fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
for (_, v) in self.options.iter_mut() {
v.visit(visitor)?;
}
for (_, v) in self.secrets.iter_mut() {
v.expose_secret_mut().visit(visitor)?;
}
ControlFlow::Continue(())
}
}

29
src/sql/src/statements/option_map/visit.rs
View File

@@ -1,29 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::ops::ControlFlow;
use sqlparser::ast::{Visit, Visitor};
use crate::statements::OptionMap;
impl Visit for OptionMap {
fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
for (k, v) in &self.map {
k.visit(visitor)?;
v.visit(visitor)?;
}
ControlFlow::Continue(())
}
}

28
src/sql/src/statements/option_map/visit_mut.rs
View File

@@ -1,28 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::ops::ControlFlow;
use sqlparser::ast::{VisitMut, VisitorMut};
use crate::statements::OptionMap;
impl VisitMut for OptionMap {
fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
for (_, v) in self.map.iter_mut() {
v.visit(visitor)?;
}
ControlFlow::Continue(())
}
}