Merge pull request #2651 from GrapeBaBa/issue2160_plan

[Feature]Add host_name field

Author: databend-bot

common/management/src/user/user_api.rs

@@ -23,14 +23,21 @@ use common_meta_types::SeqV;
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub struct UserInfo {
     pub name: String,
+    pub host_name: String,
     pub password: Vec<u8>,
     pub auth_type: AuthType,
 }
 
 impl UserInfo {
-    pub(crate) fn new(name: String, password: Vec<u8>, auth_type: AuthType) -> Self {
+    pub(crate) fn new(
+        name: String,
+        host_name: String,
+        password: Vec<u8>,
+        auth_type: AuthType,
+    ) -> Self {
         UserInfo {
             name,
+            host_name,
             password,
             auth_type,
         }
@@ -48,6 +55,7 @@ pub trait UserMgrApi: Sync + Send {
     async fn update_user(
         &self,
         username: String,
+        new_host_name: Option<String>,
         new_password: Option<Vec<u8>>,
         new_auth: Option<AuthType>,
         seq: Option<u64>,

common/management/src/user/user_mgr.rs

@@ -100,24 +100,31 @@ impl UserMgrApi for UserMgr {
     async fn update_user(
         &self,
         username: String,
+        new_host_name: Option<String>,
         new_password: Option<Vec<u8>>,
         new_auth: Option<AuthType>,
         seq: Option<u64>,
     ) -> Result<Option<u64>> {
-        if new_password.is_none() && new_auth.is_none() {
+        if new_password.is_none() && new_auth.is_none() && new_host_name.is_none() {
             return Ok(seq);
         }
-        let partial_update = new_auth.is_none() || new_password.is_none();
-        let user_info = if partial_update {
+        let full_update = new_auth.is_some() && new_password.is_some() && new_host_name.is_some();
+        let user_info = if full_update {
+            UserInfo::new(
+                username.clone(),
+                new_host_name.unwrap(),
+                new_password.unwrap(),
+                new_auth.unwrap(),
+            )
+        } else {
             let user_val_seq = self.get_user(username.clone(), seq);
             let user_info = user_val_seq.await?.data;
             UserInfo::new(
                 username.clone(),
+                new_host_name.unwrap_or(user_info.host_name),
                 new_password.map_or(user_info.password, |v| v.to_vec()),
-                new_auth.map_or(user_info.auth_type, |v| v),
+                new_auth.unwrap_or(user_info.auth_type),
             )
-        } else {
-            UserInfo::new(username.clone(), new_password.unwrap(), new_auth.unwrap())
         };
 
         let key = format!("{}/{}", self.user_prefix, user_info.name);

common/management/src/user/user_mgr_test.rs

@@ -73,9 +73,11 @@ mod add {
     async fn test_add_user() -> common_exception::Result<()> {
         let test_user_name = "test_user";
         let test_password = "test_password";
+        let test_hostname = "localhost";
         let auth_type = AuthType::Sha256;
         let user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_hostname.to_string(),
             Vec::from(test_password),
             auth_type.clone(),
         );
@@ -131,6 +133,7 @@ mod add {
 
             let user_info = UserInfo::new(
                 test_user_name.to_string(),
+                test_hostname.to_string(),
                 Vec::from(test_password),
                 auth_type.clone(),
             );
@@ -161,6 +164,7 @@ mod add {
             let user_mgr = UserMgr::new(kv, "tenant1");
             let user_info = UserInfo::new(
                 test_user_name.to_string(),
+                test_hostname.to_string(),
                 Vec::from(test_password),
                 auth_type,
             );
@@ -184,10 +188,12 @@ mod get {
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn test_get_user_seq_match() -> common_exception::Result<()> {
         let test_user_name = "test";
+        let test_host_name = "localhost";
         let test_key = format!("__fd_users/tenant1/{}", test_user_name);
 
         let user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_host_name.to_string(),
             Vec::from("pass"),
             AuthType::Sha256,
         );
@@ -214,10 +220,12 @@ mod get {
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn test_get_user_do_not_care_seq() -> common_exception::Result<()> {
         let test_user_name = "test";
+        let test_host_name = "localhost";
         let test_key = format!("__fd_users/tenant1/{}", test_user_name);
 
         let user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_host_name.to_string(),
             Vec::from("pass"),
             AuthType::Sha256,
         );
@@ -316,17 +324,21 @@ mod get_users {
 
     type FakeKeys = Vec<(String, SeqV<Vec<u8>>)>;
     type UserInfos = Vec<SeqV<UserInfo>>;
+
     fn prepare() -> common_exception::Result<(FakeKeys, UserInfos)> {
         let mut names = vec![];
+        let mut host_names = vec![];
         let mut keys = vec![];
         let mut res = vec![];
         let mut user_infos = vec![];
         for i in 0..9 {
             let name = format!("test_user_{}", i);
             names.push(name.clone());
+            let host_name = format!("test_host_name_{}", i);
+            host_names.push(host_name.clone());
             let key = format!("{}/{}", "tenant1", name);
             keys.push(key);
-            let user_info = UserInfo::new(name, Vec::from("pass"), AuthType::Sha256);
+            let user_info = UserInfo::new(name, host_name, Vec::from("pass"), AuthType::Sha256);
             res.push((
                 "fake_key".to_string(),
                 SeqV::new(i, serde_json::to_vec(&user_info)?),
@@ -389,7 +401,6 @@ mod get_users {
 }
 
 mod drop {
-
     use super::*;
 
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
@@ -447,6 +458,7 @@ mod update {
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn test_update_user_normal_partial_update() -> common_exception::Result<()> {
         let test_user_name = "name";
+        let test_host_name = "localhost";
         let test_key = format!("__fd_users/tenant1/{}", test_user_name);
         let test_seq = None;
 
@@ -455,6 +467,7 @@ mod update {
 
         let user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_host_name.to_string(),
             Vec::from(old_pass),
             old_auth_type,
         );
@@ -479,6 +492,7 @@ mod update {
         let new_pass = "new pass";
         let new_user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_host_name.to_string(),
             Vec::from(new_pass),
             AuthType::DoubleSha1,
         );
@@ -501,6 +515,7 @@ mod update {
 
         let res = user_mgr.update_user(
             test_user_name.to_string(),
+            Some(new_user_info.host_name),
             Some(new_user_info.password),
             None,
             test_seq,
@@ -513,6 +528,7 @@ mod update {
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn test_update_user_normal_full_update() -> common_exception::Result<()> {
         let test_user_name = "name";
+        let test_host_name = "localhost";
         let test_key = format!("__fd_users/tenant1/{}", test_user_name);
         let test_seq = None;
 
@@ -524,6 +540,7 @@ mod update {
 
         let new_user_info = UserInfo::new(
             test_user_name.to_string(),
+            test_host_name.to_string(),
             Vec::from(new_pass),
             new_auth_type.clone(),
         );
@@ -547,6 +564,7 @@ mod update {
 
         let res = user_mgr.update_user(
             test_user_name.to_string(),
+            Some(test_host_name.to_string()),
             Some(new_user_info.password),
             Some(new_auth_type),
             test_seq,
@@ -565,7 +583,7 @@ mod update {
         let user_mgr = UserMgr::new(kv, "tenant1");
 
         let new_password: Option<Vec<u8>> = None;
-        let res = user_mgr.update_user(test_name.to_string(), new_password, None, None);
+        let res = user_mgr.update_user(test_name.to_string(), None, new_password, None, None);
         assert!(res.await.is_ok());
         Ok(())
     }
@@ -589,6 +607,7 @@ mod update {
 
         let res = user_mgr.update_user(
             test_user_name.to_string(),
+            None,
             Some(Vec::from("new_pass".as_bytes())),
             None,
             test_seq,
@@ -625,6 +644,7 @@ mod update {
 
         let res = user_mgr.update_user(
             test_user_name.to_string(),
+            Some(String::from("localhost")),
             Some(Vec::from("new_pass".as_bytes())),
             Some(AuthType::Sha256),
             test_seq,

query/src/interpreters/interpreter_user_create.rs

@@ -52,6 +52,7 @@ impl Interpreter for CreatUserInterpreter {
         let user_mgr = self.ctx.get_sessions_manager().get_user_manager();
         let user_info = UserInfo {
             name: plan.name,
+            host_name: plan.host_name,
             password: plan.password,
             auth_type: plan.auth_type,
         };

query/src/users/user.rs

@@ -19,14 +19,21 @@ use common_meta_types::AuthType;
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub struct User {
     name: String,
+    host_name: String,
     password: String,
     auth_type: AuthType,
 }
 
 impl User {
-    pub fn new(name: impl Into<String>, password: impl Into<String>, auth_type: AuthType) -> Self {
+    pub fn new(
+        name: impl Into<String>,
+        host_name: impl Into<String>,
+        password: impl Into<String>,
+        auth_type: AuthType,
+    ) -> Self {
         User {
             name: name.into(),
+            host_name: host_name.into(),
             password: password.into(),
             auth_type,
         }
@@ -37,6 +44,7 @@ impl From<&User> for UserInfo {
     fn from(user: &User) -> Self {
         UserInfo {
             name: user.name.clone(),
+            host_name: user.host_name.clone(),
             password: Vec::from(user.password.clone()),
             auth_type: user.auth_type.clone(),
         }

query/src/users/user_mgr.rs

@@ -54,7 +54,7 @@ impl UserManager {
         match user {
             // TODO(BohuTANG): Mock, need removed.
             "default" | "" | "root" => {
-                let user = User::new(user, "", AuthType::None);
+                let user = User::new(user, "%", "", AuthType::None);
                 Ok(user.into())
             }
             _ => {

query/src/users/user_mgr_test.rs

@@ -27,12 +27,13 @@ async fn test_user_manager() -> Result<()> {
     config.query.tenant = "tenant1".to_string();
 
     let user = "test-user1";
+    let host_name = "localhost";
     let pwd = "test-pwd";
     let user_mgr = UserManager::create_global(config).await?;
 
     // add.
     {
-        let user_info = User::new(user, pwd, AuthType::PlainText);
+        let user_info = User::new(user, host_name, pwd, AuthType::PlainText);
         user_mgr.add_user(user_info.into()).await?;
     }