Skip to content

Commit

Permalink
ignore the invalid built-in user&udf
Browse files Browse the repository at this point in the history
  • Loading branch information
@BohuTANG
BohuTANG committed Jul 2, 2024
1 parent bd7c199 commit 7537952
Show file tree
Hide file tree
Showing 7 changed files with 269 additions and 267 deletions.
53 changes: 20 additions & 33 deletions src/query/service/src/builtin/builtin_udfs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use databend_common_exception::Result;
use databend_common_expression::types::DataType;
use databend_common_meta_app::principal::UserDefinedFunction;
use databend_common_sql::resolve_type_name;
use log::warn;
use log::error;

pub struct BuiltinUDFs {
udf_configs: Vec<UDFConfig>,
Expand All @@ -37,11 +37,7 @@ impl BuiltinUDFs {
}

// Parse the UDF definition and return the UserDefinedFunction.
async fn parse_udf_definition(
&self,
name: &str,
definition: &str,
) -> Result<UserDefinedFunction> {
fn parse_udf_definition(&self, name: &str, definition: &str) -> Result<UserDefinedFunction> {
let tokens = tokenize_sql(definition)?;
let (stmt, _) = parse_sql(&tokens, Dialect::PostgreSQL)?;

Expand Down Expand Up @@ -83,27 +79,26 @@ impl BuiltinUDFs {
}
}

pub async fn to_meta_udfs(&self) -> Result<HashMap<String, UserDefinedFunction>> {
/// Convert to UDFs.
/// Skip invalid UDF configs.
pub fn to_udfs(&self) -> HashMap<String, UserDefinedFunction> {
let mut udf_map = HashMap::new();

for udf_config in self.udf_configs.iter() {
match self
.parse_udf_definition(&udf_config.name, &udf_config.definition)
.await
{
match self.parse_udf_definition(&udf_config.name, &udf_config.definition) {
Ok(user_defined_function) => {
udf_map.insert(user_defined_function.name.clone(), user_defined_function);
}
Err(e) => {
warn!(
"Failed to parse UDF definition for '{}': {}",
error!(
"Failed to parse built-in UDF definition for '{}': {}",
udf_config.name, e
);
}
}
}

Ok(udf_map)
udf_map
}
}

Expand Down Expand Up @@ -143,27 +138,19 @@ ADDRESS = 'https://databend.com'"

let builtin_udfs = BuiltinUDFs::create(udf_configs);

let result = builtin_udfs.to_meta_udfs().await;
let udf_map = builtin_udfs.to_udfs();

// Verify the result
match result {
Ok(udf_map) => {
// Test first UDF
assert!(udf_map.contains_key("test_udf1"));
let udf1 = udf_map.get("test_udf1").unwrap();
assert_eq!(udf1.name, "test_udf1");
// Test first UDF
assert!(udf_map.contains_key("test_udf1"));
let udf1 = udf_map.get("test_udf1").unwrap();
assert_eq!(udf1.name, "test_udf1");

// Test second UDF
assert!(udf_map.contains_key("test_udf2"));
let udf2 = udf_map.get("test_udf2").unwrap();
assert_eq!(udf2.name, "test_udf2");
// Test second UDF
assert!(udf_map.contains_key("test_udf2"));
let udf2 = udf_map.get("test_udf2").unwrap();
assert_eq!(udf2.name, "test_udf2");

// Test invalid UDF is not included
assert!(!udf_map.contains_key("invalid_udf"));
}
Err(e) => {
panic!("Unexpected error: {}", e);
}
}
// Test invalid UDF is not included
assert!(!udf_map.contains_key("invalid_udf"));
}
}
89 changes: 51 additions & 38 deletions src/query/service/src/builtin/builtin_users.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use databend_common_exception::ErrorCode;
use databend_common_exception::Result;
use databend_common_meta_app::principal::AuthInfo;
use databend_common_meta_app::principal::AuthType;
use log::error;

pub struct BuiltinUsers {
user_configs: Vec<UserConfig>,
Expand All @@ -40,40 +41,52 @@ impl BuiltinUsers {
}
}

fn parse_user_auth_config(user_config: UserConfig) -> Result<AuthInfo> {
let auth_config = user_config.auth.clone();
let auth_type = AuthType::from_str(&auth_config.auth_type)?;
match auth_type {
AuthType::NoPassword => {
Self::check_no_auth_string(auth_config.auth_string.clone(), AuthInfo::None)
}
AuthType::JWT => {
Self::check_no_auth_string(auth_config.auth_string.clone(), AuthInfo::JWT)
}
AuthType::Sha256Password | AuthType::DoubleSha1Password => {
let password_type = auth_type.get_password_type().expect("must success");
match &auth_config.auth_string {
None => Err(ErrorCode::InvalidConfig("must set auth_string")),
Some(s) => {
let p = hex::decode(s).map_err(|e| {
ErrorCode::InvalidConfig(format!("password is not hex: {e:?}"))
})?;
Ok(AuthInfo::Password {
hash_value: p,
hash_method: password_type,
})
}
}
}
}
}

/// Convert to auth infos.
pub fn to_meta_auth_infos(&self) -> Result<HashMap<String, AuthInfo>> {
/// Skip invalid user auth configs.
pub fn to_auth_infos(&self) -> HashMap<String, AuthInfo> {
let mut auth_infos = HashMap::new();
for user_config in self.user_configs.iter() {
let name = user_config.name.clone();
let auth_config = user_config.auth.clone();
let auth_type = AuthType::from_str(&auth_config.auth_type)?;
let auth_info = match auth_type {
AuthType::NoPassword => {
Self::check_no_auth_string(auth_config.auth_string.clone(), AuthInfo::None)
}
AuthType::JWT => {
Self::check_no_auth_string(auth_config.auth_string.clone(), AuthInfo::JWT)
match Self::parse_user_auth_config(user_config.clone()) {
Ok(auth_info) => {
auth_infos.insert(user_config.name.clone(), auth_info);
}
AuthType::Sha256Password | AuthType::DoubleSha1Password => {
let password_type = auth_type.get_password_type().expect("must success");
match &auth_config.auth_string {
None => Err(ErrorCode::InvalidConfig("must set auth_string")),
Some(s) => {
let p = hex::decode(s).map_err(|e| {
ErrorCode::InvalidConfig(format!("password is not hex: {e:?}"))
})?;
Ok(AuthInfo::Password {
hash_value: p,
hash_method: password_type,
})
}
}
Err(e) => {
error!(
"Failed to parse built-in user auth for '{}': {}",
user_config.name, e
);
}
}?;

auth_infos.insert(name, auth_info);
}
}
Ok(auth_infos)
auth_infos
}
}

Expand All @@ -100,7 +113,7 @@ mod tests {
let user_configs = vec![create_user_config("user1", "no_password", None)];
let builtin_users = BuiltinUsers::create(user_configs);

let auth_infos = builtin_users.to_meta_auth_infos().unwrap();
let auth_infos = builtin_users.to_auth_infos();
let auth_info = auth_infos.get("user1").unwrap();

assert_eq!(auth_info.get_type(), AuthType::NoPassword);
Expand All @@ -111,7 +124,7 @@ mod tests {
let user_configs = vec![create_user_config("user2", "jwt", None)];
let builtin_users = BuiltinUsers::create(user_configs);

let auth_infos = builtin_users.to_meta_auth_infos().unwrap();
let auth_infos = builtin_users.to_auth_infos();
let auth_info = auth_infos.get("user2").unwrap();

assert_eq!(auth_info.get_type(), AuthType::JWT);
Expand All @@ -126,7 +139,7 @@ mod tests {
)];
let builtin_users = BuiltinUsers::create(user_configs);

let auth_infos = builtin_users.to_meta_auth_infos().unwrap();
let auth_infos = builtin_users.to_auth_infos();
let auth_info = auth_infos.get("user3").unwrap();

match auth_info {
Expand All @@ -153,7 +166,7 @@ mod tests {
)];
let builtin_users = BuiltinUsers::create(user_configs);

let auth_infos = builtin_users.to_meta_auth_infos().unwrap();
let auth_infos = builtin_users.to_auth_infos();
let auth_info = auth_infos.get("user4").unwrap();

match auth_info {
Expand All @@ -180,25 +193,25 @@ mod tests {
)];
let builtin_users = BuiltinUsers::create(user_configs);

let result = builtin_users.to_meta_auth_infos();
assert!(result.is_err());
let auth_infos = builtin_users.to_auth_infos();
assert!(auth_infos.get("user5").is_none());
}

#[test]
fn test_missing_auth_string_for_password() {
let user_configs = vec![create_user_config("user6", "sha256_password", None)];
let builtin_users = BuiltinUsers::create(user_configs);

let result = builtin_users.to_meta_auth_infos();
assert!(result.is_err());
let auth_infos = builtin_users.to_auth_infos();
assert!(auth_infos.get("user6").is_none());
}

#[test]
fn test_invalid_auth_type() {
let user_configs = vec![create_user_config("user7", "InvalidAuthType", None)];
let builtin_users = BuiltinUsers::create(user_configs);

let result = builtin_users.to_meta_auth_infos();
assert!(result.is_err());
let auth_infos = builtin_users.to_auth_infos();
assert!(auth_infos.get("user7").is_none());
}
}
6 changes: 4 additions & 2 deletions src/query/service/src/global_services.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,11 @@ impl GlobalServices {
{
let built_in_users = BuiltinUsers::create(config.query.builtin.users.clone());
let built_in_udfs = BuiltinUDFs::create(config.query.builtin.udfs.clone());

// We will ignore the error here, and just log a error.
let builtin = BuiltIn {
users: built_in_users.to_meta_auth_infos()?,
udfs: built_in_udfs.to_meta_udfs().await?,
users: built_in_users.to_auth_infos(),
udfs: built_in_udfs.to_udfs(),
};
UserApiProvider::init(
config.meta.to_meta_grpc_client_conf(),
Expand Down
Loading

0 comments on commit 7537952

Please sign in to comment.