Merge pull request #4451 from Junnplus/grant-role

[RBAC] Support grant role and revoke role

common/management/src/role/role_api.rs

@@ -21,27 +21,41 @@ use common_meta_types::UserPrivilegeSet;
 
 #[async_trait::async_trait]
 pub trait RoleApi: Sync + Send {
-    async fn add_role(&self, role_info: &RoleInfo) -> Result<u64>;
+    async fn add_role(&self, role_info: RoleInfo) -> Result<u64>;
 
-    async fn get_role(&self, role: &RoleIdentity, seq: Option<u64>) -> Result<SeqV<RoleInfo>>;
+    async fn get_role(&self, role: RoleIdentity, seq: Option<u64>) -> Result<SeqV<RoleInfo>>;
 
     async fn get_roles(&self) -> Result<Vec<SeqV<RoleInfo>>>;
 
-    async fn grant_role_privileges(
+    async fn grant_privileges(
         &self,
-        role: &RoleIdentity,
+        role: RoleIdentity,
         object: GrantObject,
         privileges: UserPrivilegeSet,
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
-    async fn revoke_role_privileges(
+    async fn revoke_privileges(
         &self,
-        role: &RoleIdentity,
+        role: RoleIdentity,
         object: GrantObject,
         privileges: UserPrivilegeSet,
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
-    async fn drop_role(&self, role: &RoleIdentity, seq: Option<u64>) -> Result<()>;
+    async fn grant_role(
+        &self,
+        role: RoleIdentity,
+        grant_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>>;
+
+    async fn revoke_role(
+        &self,
+        role: RoleIdentity,
+        revoke_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>>;
+
+    async fn drop_role(&self, role: RoleIdentity, seq: Option<u64>) -> Result<()>;
 }

common/management/src/role/role_mgr.rs

@@ -91,7 +91,7 @@ impl RoleMgr {
 
 #[async_trait::async_trait]
 impl RoleApi for RoleMgr {
-    async fn add_role(&self, role_info: &RoleInfo) -> common_exception::Result<u64> {
+    async fn add_role(&self, role_info: RoleInfo) -> common_exception::Result<u64> {
         let match_seq = MatchSeq::Exact(0);
         let key = self.make_role_key(&role_info.identity());
         let value = serde_json::to_vec(&role_info)?;
@@ -113,23 +113,16 @@ impl RoleApi for RoleMgr {
         }
     }
 
-    async fn get_role(
-        &self,
-        role_identity: &RoleIdentity,
-        seq: Option<u64>,
-    ) -> Result<SeqV<RoleInfo>> {
-        let key = self.make_role_key(role_identity);
+    async fn get_role(&self, role: RoleIdentity, seq: Option<u64>) -> Result<SeqV<RoleInfo>> {
+        let key = self.make_role_key(&role);
         let kv_api = self.kv_api.clone();
         let res = kv_api.get_kv(&key).await?;
         let seq_value =
-            res.ok_or_else(|| ErrorCode::UnknownRole(format!("unknown role {}", role_identity)))?;
+            res.ok_or_else(|| ErrorCode::UnknownRole(format!("unknown role {}", role)))?;
 
         match MatchSeq::from(seq).match_seq(&seq_value) {
             Ok(_) => Ok(seq_value.into_seqv()?),
-            Err(_) => Err(ErrorCode::UnknownRole(format!(
-                "unknown role {}",
-                role_identity
-            ))),
+            Err(_) => Err(ErrorCode::UnknownRole(format!("unknown role {}", role))),
         }
     }
 
@@ -149,9 +142,9 @@ impl RoleApi for RoleMgr {
         Ok(r)
     }
 
-    async fn grant_role_privileges(
+    async fn grant_privileges(
         &self,
-        role: &RoleIdentity,
+        role: RoleIdentity,
         object: GrantObject,
         privileges: UserPrivilegeSet,
         seq: Option<u64>,
@@ -160,14 +153,14 @@ impl RoleApi for RoleMgr {
         let mut role_info = role_val_seq.await?.data;
         role_info
             .grants
-            .grant_privileges(&role.name, "", &object, privileges);
+            .grant_privileges(&role_info.name, "", &object, privileges);
         let seq = self.upsert_role_info(&role_info, seq).await?;
         Ok(Some(seq))
     }
 
-    async fn revoke_role_privileges(
+    async fn revoke_privileges(
         &self,
-        role: &RoleIdentity,
+        role: RoleIdentity,
         object: GrantObject,
         privileges: UserPrivilegeSet,
         seq: Option<u64>,
@@ -179,8 +172,34 @@ impl RoleApi for RoleMgr {
         Ok(Some(seq))
     }
 
-    async fn drop_role(&self, role: &RoleIdentity, seq: Option<u64>) -> Result<()> {
-        let key = self.make_role_key(role);
+    async fn grant_role(
+        &self,
+        role: RoleIdentity,
+        grant_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>> {
+        let role_val_seq = self.get_role(role, seq);
+        let mut role_info = role_val_seq.await?.data;
+        role_info.grants.grant_role(grant_role);
+        let seq = self.upsert_role_info(&role_info, seq).await?;
+        Ok(Some(seq))
+    }
+
+    async fn revoke_role(
+        &self,
+        role: RoleIdentity,
+        revoke_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>> {
+        let role_val_seq = self.get_role(role, seq);
+        let mut role_info = role_val_seq.await?.data;
+        role_info.grants.revoke_role(&revoke_role);
+        let seq = self.upsert_role_info(&role_info, seq).await?;
+        Ok(Some(seq))
+    }
+
+    async fn drop_role(&self, role: RoleIdentity, seq: Option<u64>) -> Result<()> {
+        let key = self.make_role_key(&role);
         let kv_api = self.kv_api.clone();
         let res = kv_api
             .upsert_kv(UpsertKVAction::new(

common/management/src/user/user_api.rs

@@ -15,6 +15,7 @@
 use common_exception::Result;
 use common_meta_types::AuthInfo;
 use common_meta_types::GrantObject;
+use common_meta_types::RoleIdentity;
 use common_meta_types::SeqV;
 use common_meta_types::UserInfo;
 use common_meta_types::UserPrivilegeSet;
@@ -40,7 +41,7 @@ pub trait UserApi: Sync + Send {
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
-    async fn grant_user_privileges(
+    async fn grant_privileges(
         &self,
         username: String,
         hostname: String,
@@ -49,7 +50,7 @@ pub trait UserApi: Sync + Send {
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
-    async fn revoke_user_privileges(
+    async fn revoke_privileges(
         &self,
         username: String,
         hostname: String,
@@ -58,5 +59,21 @@ pub trait UserApi: Sync + Send {
         seq: Option<u64>,
     ) -> Result<Option<u64>>;
 
+    async fn grant_role(
+        &self,
+        username: String,
+        hostname: String,
+        grant_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>>;
+
+    async fn revoke_role(
+        &self,
+        username: String,
+        hostname: String,
+        revoke_role: RoleIdentity,
+        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

@@ -25,6 +25,7 @@ use common_meta_types::MatchSeq;
 use common_meta_types::MatchSeqExt;
 use common_meta_types::OkOrExist;
 use common_meta_types::Operation;
+use common_meta_types::RoleIdentity;
 use common_meta_types::SeqV;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UserInfo;
@@ -184,7 +185,7 @@ impl UserApi for UserMgr {
         }
     }
 
-    async fn grant_user_privileges(
+    async fn grant_privileges(
         &self,
         username: String,
         hostname: String,
@@ -201,7 +202,7 @@ impl UserApi for UserMgr {
         Ok(Some(seq))
     }
 
-    async fn revoke_user_privileges(
+    async fn revoke_privileges(
         &self,
         username: String,
         hostname: String,
@@ -216,6 +217,34 @@ impl UserApi for UserMgr {
         Ok(Some(seq))
     }
 
+    async fn grant_role(
+        &self,
+        username: String,
+        hostname: String,
+        grant_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>> {
+        let user_val_seq = self.get_user(username, hostname, seq);
+        let mut user_info = user_val_seq.await?.data;
+        user_info.grants.grant_role(grant_role);
+        let seq = self.upsert_user_info(&user_info, seq).await?;
+        Ok(Some(seq))
+    }
+
+    async fn revoke_role(
+        &self,
+        username: String,
+        hostname: String,
+        revoke_role: RoleIdentity,
+        seq: Option<u64>,
+    ) -> Result<Option<u64>> {
+        let user_val_seq = self.get_user(username, hostname, seq);
+        let mut user_info = user_val_seq.await?.data;
+        user_info.grants.revoke_role(&revoke_role);
+        let seq = self.upsert_user_info(&user_info, seq).await?;
+        Ok(Some(seq))
+    }
+
     async fn drop_user(&self, username: String, hostname: String, seq: Option<u64>) -> Result<()> {
         let user_key = format_user_key(&username, &hostname);
         let key = format!("{}/{}", self.user_prefix, user_key);

common/management/tests/it/user.rs

@@ -701,7 +701,7 @@ mod set_user_privileges {
         let kv = Arc::new(kv);
         let user_mgr = UserMgr::create(kv, "tenant1")?;
 
-        let res = user_mgr.grant_user_privileges(
+        let res = user_mgr.grant_privileges(
             test_user_name.to_string(),
             test_hostname.to_string(),
             GrantObject::Global,

common/planners/src/lib.rs

@@ -52,11 +52,15 @@ mod plan_node_stage;
 mod plan_node_statistics;
 mod plan_node_visitor;
 mod plan_partition;
+mod plan_privilege_grant;
+mod plan_privilege_revoke;
 mod plan_projection;
 mod plan_read_datasource;
 mod plan_remote;
 mod plan_role_create;
 mod plan_role_drop;
+mod plan_role_grant;
+mod plan_role_revoke;
 mod plan_select;
 mod plan_setting;
 mod plan_show;
@@ -84,8 +88,6 @@ mod plan_use_database;
 mod plan_user_alter;
 mod plan_user_create;
 mod plan_user_drop;
-mod plan_user_privilege_grant;
-mod plan_user_privilege_revoke;
 mod plan_user_stage_create;
 mod plan_user_stage_describe;
 mod plan_user_stage_drop;
@@ -162,12 +164,16 @@ pub use plan_node_visitor::PlanVisitor;
 pub use plan_partition::PartInfo;
 pub use plan_partition::PartInfoPtr;
 pub use plan_partition::Partitions;
+pub use plan_privilege_grant::GrantPrivilegePlan;
+pub use plan_privilege_revoke::RevokePrivilegePlan;
 pub use plan_projection::ProjectionPlan;
 pub use plan_read_datasource::ReadDataSourcePlan;
 pub use plan_read_datasource::SourceInfo;
 pub use plan_remote::RemotePlan;
 pub use plan_role_create::CreateRolePlan;
 pub use plan_role_drop::DropRolePlan;
+pub use plan_role_grant::GrantRolePlan;
+pub use plan_role_revoke::RevokeRolePlan;
 pub use plan_select::SelectPlan;
 pub use plan_setting::SettingPlan;
 pub use plan_setting::VarValue;
@@ -201,8 +207,6 @@ pub use plan_use_database::UseDatabasePlan;
 pub use plan_user_alter::AlterUserPlan;
 pub use plan_user_create::CreateUserPlan;
 pub use plan_user_drop::DropUserPlan;
-pub use plan_user_privilege_grant::GrantPrivilegePlan;
-pub use plan_user_privilege_revoke::RevokePrivilegePlan;
 pub use plan_user_stage_create::CreateUserStagePlan;
 pub use plan_user_stage_describe::DescribeUserStagePlan;
 pub use plan_user_stage_drop::DropUserStagePlan;

common/planners/src/plan_node.rs

@@ -43,6 +43,7 @@ use crate::ExplainPlan;
 use crate::ExpressionPlan;
 use crate::FilterPlan;
 use crate::GrantPrivilegePlan;
+use crate::GrantRolePlan;
 use crate::HavingPlan;
 use crate::InsertPlan;
 use crate::KillPlan;
@@ -54,6 +55,7 @@ use crate::ReadDataSourcePlan;
 use crate::RemotePlan;
 use crate::RenameTablePlan;
 use crate::RevokePrivilegePlan;
+use crate::RevokeRolePlan;
 use crate::SelectPlan;
 use crate::SettingPlan;
 use crate::ShowCreateDatabasePlan;
@@ -123,8 +125,14 @@ pub enum PlanNode {
     CreateUser(CreateUserPlan),
     AlterUser(AlterUserPlan),
     DropUser(DropUserPlan),
+
+    // Grant.
     GrantPrivilege(GrantPrivilegePlan),
+    GrantRole(GrantRolePlan),
+
+    // Revoke.
     RevokePrivilege(RevokePrivilegePlan),
+    RevokeRole(RevokeRolePlan),
 
     // Role.
     CreateRole(CreateRolePlan),
@@ -211,8 +219,14 @@ impl PlanNode {
             PlanNode::CreateUser(v) => v.schema(),
             PlanNode::AlterUser(v) => v.schema(),
             PlanNode::DropUser(v) => v.schema(),
+
+            // Grant.
             PlanNode::GrantPrivilege(v) => v.schema(),
+            PlanNode::GrantRole(v) => v.schema(),
+
+            // Revoke.
             PlanNode::RevokePrivilege(v) => v.schema(),
+            PlanNode::RevokeRole(v) => v.schema(),
 
             // Role.
             PlanNode::CreateRole(v) => v.schema(),
@@ -298,8 +312,14 @@ impl PlanNode {
             PlanNode::CreateUser(_) => "CreateUser",
             PlanNode::AlterUser(_) => "AlterUser",
             PlanNode::DropUser(_) => "DropUser",
+
+            // Grant.
             PlanNode::GrantPrivilege(_) => "GrantPrivilegePlan",
+            PlanNode::GrantRole(_) => "GrantRolePlan",
+
+            // Revoke.
             PlanNode::RevokePrivilege(_) => "RevokePrivilegePlan",
+            PlanNode::RevokeRole(_) => "RevokeRolePlan",
 
             // Role.
             PlanNode::CreateRole(_) => "CreateRole",