user grant privileges

common/management/src/user/user_api.rs

@@ -47,6 +47,10 @@ impl UserInfo {
             quota,
         }
     }
+
+    pub fn set_privileges(&mut self, privileges: UserPrivilege) {
+        self.privileges |= privileges;
+    }
 }
 
 #[async_trait::async_trait]
@@ -71,6 +75,14 @@ pub trait UserMgrApi: Sync + Send {
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
+    async fn set_user_privileges(
+        &self,
+        username: String,
+        hostname: String,
+        privileges: UserPrivilege,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>>;
+
     async fn drop_user(&self, username: String, hostname: String, seq: Option<u64>) -> Result<()>;
 }
 

common/management/src/user/user_mgr.rs

@@ -27,6 +27,7 @@ use common_meta_types::MatchSeqExt;
 use common_meta_types::Operation;
 use common_meta_types::SeqV;
 use common_meta_types::UpsertKVAction;
+use common_meta_types::UserPrivilege;
 
 use crate::user::user_api::UserInfo;
 use crate::user::user_api::UserMgrApi;
@@ -121,26 +122,58 @@ impl UserMgrApi for UserMgr {
         if new_password.is_none() && new_auth.is_none() {
             return Ok(seq);
         }
-        let partial_update = new_auth.is_none() || new_password.is_none();
-        let user_info = if partial_update {
-            let user_val_seq = self.get_user(username.clone(), hostname.clone(), seq);
-            let user_info = user_val_seq.await?.data;
-
-            UserInfo::new(
-                username.clone(),
-                hostname.clone(),
-                new_password.map_or(user_info.password.clone(), |v| v.to_vec()),
-                new_auth.unwrap_or(user_info.auth_type),
-            )
-        } else {
-            UserInfo::new(
-                username.clone(),
-                hostname.clone(),
-                new_password.unwrap(),
-                new_auth.unwrap(),
-            )
+        let user_val_seq = self.get_user(username.clone(), hostname.clone(), seq);
+        let user_info = user_val_seq.await?.data;
+
+        let mut new_user_info = UserInfo::new(
+            username.clone(),
+            hostname.clone(),
+            new_password.map_or(user_info.password.clone(), |v| v.to_vec()),
+            new_auth.unwrap_or(user_info.auth_type),
+        );
+        new_user_info.set_privileges(user_info.privileges);
+
+        let user_key = format_user_key(&new_user_info.name, &new_user_info.hostname);
+        let key = format!("{}/{}", self.user_prefix, user_key);
+        let value = serde_json::to_vec(&new_user_info)?;
+
+        let match_seq = match seq {
+            None => MatchSeq::GE(1),
+            Some(s) => MatchSeq::Exact(s),
         };
 
+        let kv_api = self.kv_api.clone();
+        let upsert_kv = async move {
+            kv_api
+                .upsert_kv(UpsertKVAction::new(
+                    &key,
+                    match_seq,
+                    Operation::Update(value),
+                    None,
+                ))
+                .await
+        };
+        let res = upsert_kv.await?;
+        match res.result {
+            Some(SeqV { seq: s, .. }) => Ok(Some(s)),
+            None => Err(ErrorCode::UnknownUser(format!(
+                "unknown user, or seq not match {}",
+                username
+            ))),
+        }
+    }
+
+    async fn set_user_privileges(
+        &self,
+        username: String,
+        hostname: String,
+        privileges: UserPrivilege,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>> {
+        let user_val_seq = self.get_user(username.clone(), hostname.clone(), seq);
+        let mut user_info = user_val_seq.await?.data;
+        user_info.set_privileges(privileges);
+
         let user_key = format_user_key(&user_info.name, &user_info.hostname);
         let key = format!("{}/{}", self.user_prefix, user_key);
         let value = serde_json::to_vec(&user_info)?;

common/management/tests/it/user.rs

@@ -545,7 +545,26 @@ mod update {
         );
         let test_seq = None;
 
-        // - get_kv should NOT be called
+        let old_pass = "old_key";
+        let old_auth_type = AuthType::DoubleSha1;
+
+        let user_info = UserInfo::new(
+            test_user_name.to_string(),
+            test_hostname.to_string(),
+            Vec::from(old_pass),
+            old_auth_type,
+        );
+        let prev_value = serde_json::to_vec(&user_info)?;
+
+        // - get_kv should be called
+        let mut kv = MockKV::new();
+        {
+            let test_key = test_key.clone();
+            kv.expect_get_kv()
+                .with(predicate::function(move |v| v == test_key.as_str()))
+                .times(1)
+                .return_once(move |_k| Ok(Some(SeqV::new(0, prev_value))));
+        }
         // - update_kv should be called
 
         let new_pass = "new_pass";
@@ -559,7 +578,6 @@ mod update {
         );
         let new_value = serde_json::to_vec(&new_user_info)?;
 
-        let mut kv = MockKV::new();
         kv.expect_upsert_kv()
             .with(predicate::eq(UpsertKVAction::new(
                 &test_key,
@@ -651,8 +669,26 @@ mod update {
         );
         let test_seq = None;
 
-        // get_kv should not be called
+        let old_pass = "old_key";
+        let old_auth_type = AuthType::DoubleSha1;
+
+        let user_info = UserInfo::new(
+            test_user_name.to_string(),
+            test_hostname.to_string(),
+            Vec::from(old_pass),
+            old_auth_type,
+        );
+        let prev_value = serde_json::to_vec(&user_info)?;
+
+        // - get_kv should be called
         let mut kv = MockKV::new();
+        {
+            let test_key = test_key.clone();
+            kv.expect_get_kv()
+                .with(predicate::function(move |v| v == test_key.as_str()))
+                .times(1)
+                .return_once(move |_k| Ok(Some(SeqV::new(0, prev_value))));
+        }
 
         // upsert should be called
         kv.expect_upsert_kv()

common/meta/types/src/user_privilege.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::ops;
+
 use enumflags2::bitflags;
 use enumflags2::make_bitflags;
 use enumflags2::BitFlags;
@@ -55,11 +57,28 @@ impl UserPrivilege {
         self.privileges |= privilege;
     }
 
-    pub fn has_privilege(&mut self, privilege: UserPrivilegeType) -> bool {
+    pub fn has_privilege(&self, privilege: UserPrivilegeType) -> bool {
         self.privileges.contains(privilege)
     }
 
     pub fn set_all_privileges(&mut self) {
         self.privileges |= ALL_PRIVILEGES;
     }
 }
+
+impl ops::BitOr for UserPrivilege {
+    type Output = Self;
+    #[inline(always)]
+    fn bitor(self, other: Self) -> Self {
+        Self {
+            privileges: self.privileges | other.privileges,
+        }
+    }
+}
+
+impl ops::BitOrAssign for UserPrivilege {
+    #[inline(always)]
+    fn bitor_assign(&mut self, other: Self) {
+        self.privileges |= other.privileges
+    }
+}

common/planners/src/lib.rs

@@ -65,6 +65,7 @@ mod plan_expression_validator;
 mod plan_expression_visitor;
 mod plan_extras;
 mod plan_filter;
+mod plan_grant_privilege;
 mod plan_having;
 mod plan_insert_into;
 mod plan_kill;
@@ -133,6 +134,7 @@ pub use plan_expression_visitor::ExpressionVisitor;
 pub use plan_expression_visitor::Recursion;
 pub use plan_extras::Extras;
 pub use plan_filter::FilterPlan;
+pub use plan_grant_privilege::GrantPrivilegePlan;
 pub use plan_having::HavingPlan;
 pub use plan_insert_into::InsertIntoPlan;
 pub use plan_kill::KillPlan;

common/planners/src/plan_grant_privilege.rs

@@ -0,0 +1,32 @@
+// Copyright 2021 Datafuse Labs.
+//
+// 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::sync::Arc;
+
+use common_datavalues::DataSchema;
+use common_datavalues::DataSchemaRef;
+use common_meta_types::UserPrivilege;
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq)]
+pub struct GrantPrivilegePlan {
+    pub name: String,
+    pub hostname: String,
+    pub priv_types: UserPrivilege,
+}
+
+impl GrantPrivilegePlan {
+    pub fn schema(&self) -> DataSchemaRef {
+        Arc::new(DataSchema::empty())
+    }
+}

common/planners/src/plan_node.rs

@@ -33,6 +33,7 @@ use crate::EmptyPlan;
 use crate::ExplainPlan;
 use crate::ExpressionPlan;
 use crate::FilterPlan;
+use crate::GrantPrivilegePlan;
 use crate::HavingPlan;
 use crate::InsertIntoPlan;
 use crate::KillPlan;
@@ -82,6 +83,7 @@ pub enum PlanNode {
     Kill(KillPlan),
     CreateUser(CreateUserPlan),
     AlterUser(AlterUserPlan),
+    GrantPrivilege(GrantPrivilegePlan),
 }
 
 impl PlanNode {
@@ -118,6 +120,7 @@ impl PlanNode {
             PlanNode::Kill(v) => v.schema(),
             PlanNode::CreateUser(v) => v.schema(),
             PlanNode::AlterUser(v) => v.schema(),
+            PlanNode::GrantPrivilege(v) => v.schema(),
         }
     }
 
@@ -153,6 +156,7 @@ impl PlanNode {
             PlanNode::Kill(_) => "KillQuery",
             PlanNode::CreateUser(_) => "CreateUser",
             PlanNode::AlterUser(_) => "AlterUser",
+            PlanNode::GrantPrivilege(_) => "GrantPrivilegePlan",
         }
     }
 

common/planners/src/plan_rewriter.rs

@@ -38,6 +38,7 @@ use crate::Expression;
 use crate::ExpressionPlan;
 use crate::Expressions;
 use crate::FilterPlan;
+use crate::GrantPrivilegePlan;
 use crate::HavingPlan;
 use crate::InsertIntoPlan;
 use crate::KillPlan;
@@ -108,6 +109,7 @@ pub trait PlanRewriter {
             PlanNode::Kill(plan) => self.rewrite_kill(plan),
             PlanNode::CreateUser(plan) => self.create_user(plan),
             PlanNode::AlterUser(plan) => self.alter_user(plan),
+            PlanNode::GrantPrivilege(plan) => self.grant_privilege(plan),
         }
     }
 
@@ -346,6 +348,10 @@ pub trait PlanRewriter {
     fn alter_user(&mut self, plan: &AlterUserPlan) -> Result<PlanNode> {
         Ok(PlanNode::AlterUser(plan.clone()))
     }
+
+    fn grant_privilege(&mut self, plan: &GrantPrivilegePlan) -> Result<PlanNode> {
+        Ok(PlanNode::GrantPrivilege(plan.clone()))
+    }
 }
 
 pub struct RewriteHelper {}

common/planners/src/plan_visitor.rs

@@ -30,6 +30,7 @@ use crate::ExplainPlan;
 use crate::Expression;
 use crate::ExpressionPlan;
 use crate::FilterPlan;
+use crate::GrantPrivilegePlan;
 use crate::HavingPlan;
 use crate::InsertIntoPlan;
 use crate::KillPlan;
@@ -121,6 +122,7 @@ pub trait PlanVisitor {
             PlanNode::Kill(plan) => self.visit_kill_query(plan),
             PlanNode::CreateUser(plan) => self.visit_create_user(plan),
             PlanNode::AlterUser(plan) => self.visit_alter_user(plan),
+            PlanNode::GrantPrivilege(plan) => self.visit_grant_privilege(plan),
         }
     }
 
@@ -248,6 +250,10 @@ pub trait PlanVisitor {
         Ok(())
     }
 
+    fn visit_grant_privilege(&mut self, _: &GrantPrivilegePlan) -> Result<()> {
+        Ok(())
+    }
+
     fn visit_describe_table(&mut self, _: &DescribeTablePlan) -> Result<()> {
         Ok(())
     }

query/src/interpreters/interpreter_factory.rs

@@ -27,6 +27,7 @@ use crate::interpreters::DescribeTableInterpreter;
 use crate::interpreters::DropDatabaseInterpreter;
 use crate::interpreters::DropTableInterpreter;
 use crate::interpreters::ExplainInterpreter;
+use crate::interpreters::GrantPrivilegeInterpreter;
 use crate::interpreters::InsertIntoInterpreter;
 use crate::interpreters::Interpreter;
 use crate::interpreters::SelectInterpreter;
@@ -56,6 +57,7 @@ impl InterpreterFactory {
             PlanNode::Kill(v) => KillInterpreter::try_create(ctx, v),
             PlanNode::CreateUser(v) => CreatUserInterpreter::try_create(ctx, v),
             PlanNode::AlterUser(v) => AlterUserInterpreter::try_create(ctx, v),
+            PlanNode::GrantPrivilege(v) => GrantPrivilegeInterpreter::try_create(ctx, v),
             _ => Result::Err(ErrorCode::UnknownTypeOfQuery(format!(
                 "Can't get the interpreter by plan:{}",
                 plan.name()