feat(query): support create or replace role (#18500)

* feat(query): support create or replace role

* fix conversation

* fix ci: Handle ExistsError

* delete 306 307 add 311 rbac compat test

In ths pr add replace role, and it's based on PB.

In 306/307 role is serde_json not support PB.

Add 311 rbac read write compat test. 311 is the first version that
ownership and role serialize to pb. https://github.com/databendlabs/databend/releases/tag/v1.2.311-nightly

Author: TCeason

.github/actions/test_compat_fuse/action.yml

@@ -13,12 +13,11 @@ runs:
       run: |
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.46   --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  base
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.241  --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  revoke
-          bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.306  --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  rbac
-          bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.307  --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  rbac
+          bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.311  --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  rbac
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.318  --reader-version  current  --meta-versions  1.2.527  1.2.677  --logictest-path  rbac
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  1.2.680  --reader-version  current  --meta-versions  1.2.527  1.2.680  --logictest-path  udf
 
-          bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  current  --reader-version  1.2.307  --meta-versions  1.2.677           --logictest-path  rbac
+          bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  current  --reader-version  1.2.311  --meta-versions  1.2.677           --logictest-path  rbac
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  current  --reader-version  1.2.318  --meta-versions  1.2.677           --logictest-path  rbac
           bash ./tests/compat_fuse/test_compat_fuse.sh  --writer-version  current  --reader-version  1.2.680  --meta-versions  1.2.680           --logictest-path  udf
     - name: Upload failure

src/meta/app/src/principal/mod.rs

@@ -83,6 +83,7 @@ pub use procedure_id_ident::ProcedureId;
 pub use procedure_id_to_name::ProcedureIdToNameIdent;
 pub use procedure_identity::ProcedureIdentity;
 pub use procedure_name_ident::ProcedureNameIdent;
+pub use role_ident::Resource;
 pub use role_ident::RoleIdent;
 pub use role_ident::RoleIdentRaw;
 pub use role_info::RoleInfo;

src/query/ast/src/ast/statements/statement.rs

@@ -244,7 +244,7 @@ pub enum Statement {
         show_options: Option<ShowOptions>,
     },
     CreateRole {
-        if_not_exists: bool,
+        create_option: CreateOption,
         role_name: String,
     },
     DropRole {
@@ -875,11 +875,15 @@ impl Display for Statement {
                 write!(f, " {}", user)?;
             }
             Statement::CreateRole {
-                if_not_exists,
+                create_option,
                 role_name: role,
             } => {
-                write!(f, "CREATE ROLE")?;
-                if *if_not_exists {
+                write!(f, "CREATE")?;
+                if let CreateOption::CreateOrReplace = create_option {
+                    write!(f, " OR REPLACE")?;
+                }
+                write!(f, " ROLE")?;
+                if let CreateOption::CreateIfNotExists = create_option {
                     write!(f, " IF NOT EXISTS")?;
                 }
                 write!(f, " {}", QuotedString(role, '\''))?;

src/query/ast/src/parser/statement.rs

@@ -1643,13 +1643,17 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
         |(_, _, show_options)| Statement::ShowRoles { show_options },
     );
-    let create_role = map(
+    let create_role = map_res(
         rule! {
-            CREATE ~ ROLE ~ ( IF ~ ^NOT ~ ^EXISTS )? ~ #role_name
+            CREATE ~ ( OR ~ ^REPLACE )? ~ ROLE ~ ( IF ~ ^NOT ~ ^EXISTS )? ~ #role_name
         },
-        |(_, _, opt_if_not_exists, role_name)| Statement::CreateRole {
-            if_not_exists: opt_if_not_exists.is_some(),
-            role_name,
+        |(_, opt_or_replace, _, opt_if_not_exists, role_name)| {
+            let create_option =
+                parse_create_option(opt_or_replace.is_some(), opt_if_not_exists.is_some())?;
+            Ok(Statement::CreateRole {
+                create_option,
+                role_name,
+            })
         },
     );
     let drop_role = map(

src/query/ast/tests/it/testdata/stmt.txt

@@ -13055,7 +13055,7 @@ create role test
 CREATE ROLE 'test'
 ---------- AST ------------
 CreateRole {
-    if_not_exists: false,
+    create_option: Create,
     role_name: "test",
 }
 
@@ -13066,7 +13066,7 @@ create role 'test'
 CREATE ROLE 'test'
 ---------- AST ------------
 CreateRole {
-    if_not_exists: false,
+    create_option: Create,
     role_name: "test",
 }
 

src/query/management/src/role/role_api.rs

@@ -13,16 +13,23 @@
 // limitations under the License.
 
 use databend_common_exception::Result;
+use databend_common_meta_app::principal::role_ident;
 use databend_common_meta_app::principal::OwnershipInfo;
 use databend_common_meta_app::principal::OwnershipObject;
 use databend_common_meta_app::principal::RoleInfo;
+use databend_common_meta_app::tenant_key::errors::ExistError;
 use databend_common_meta_kvapi::kvapi::ListKVReply;
 use databend_common_meta_types::MatchSeq;
+use databend_common_meta_types::MetaError;
 use databend_common_meta_types::SeqV;
 
 #[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,
+        can_replace: bool,
+    ) -> std::result::Result<std::result::Result<(), ExistError<role_ident::Resource>>, MetaError>;
 
     #[allow(clippy::ptr_arg)]
     async fn get_role(&self, role: &String, seq: MatchSeq) -> Result<SeqV<RoleInfo>>;

src/query/management/src/role/role_mgr.rs

@@ -17,13 +17,15 @@ use std::sync::Arc;
 use databend_common_base::base::tokio::sync::Mutex;
 use databend_common_exception::ErrorCode;
 use databend_common_meta_api::kv_pb_api::KVPbApi;
+use databend_common_meta_api::kv_pb_api::UpsertPB;
 use databend_common_meta_api::reply::unpack_txn_reply;
 use databend_common_meta_api::txn_backoff::txn_backoff;
 use databend_common_meta_api::txn_cond_seq;
 use databend_common_meta_api::txn_op_del;
 use databend_common_meta_api::txn_op_put;
 use databend_common_meta_app::app_error::AppError;
 use databend_common_meta_app::app_error::TxnRetryMaxTimes;
+use databend_common_meta_app::principal::role_ident;
 use databend_common_meta_app::principal::GrantObject;
 use databend_common_meta_app::principal::OwnershipInfo;
 use databend_common_meta_app::principal::OwnershipObject;
@@ -32,6 +34,7 @@ use databend_common_meta_app::principal::RoleInfo;
 use databend_common_meta_app::principal::TenantOwnershipObjectIdent;
 use databend_common_meta_app::principal::UserPrivilegeType;
 use databend_common_meta_app::tenant::Tenant;
+use databend_common_meta_app::tenant_key::errors::ExistError;
 use databend_common_meta_app::KeyWithTenant;
 use databend_common_meta_cache::Cache;
 use databend_common_meta_client::ClientHandle;
@@ -48,6 +51,7 @@ use databend_common_meta_types::Operation;
 use databend_common_meta_types::SeqV;
 use databend_common_meta_types::TxnRequest;
 use databend_common_meta_types::UpsertKV;
+use databend_common_meta_types::With;
 use enumflags2::make_bitflags;
 use fastrace::func_name;
 use futures::TryStreamExt;
@@ -175,23 +179,28 @@ impl RoleMgr {
 impl RoleApi for RoleMgr {
     #[async_backtrace::framed]
     #[fastrace::trace]
-    async fn add_role(&self, role_info: RoleInfo) -> databend_common_exception::Result<u64> {
-        let match_seq = MatchSeq::Exact(0);
-        let key = self.role_ident(role_info.identity()).to_string_key();
-        let value = serialize_struct(&role_info, ErrorCode::IllegalUserInfoFormat, || "")?;
+    async fn add_role(
+        &self,
+        info: RoleInfo,
+        can_replace: bool,
+    ) -> Result<Result<(), ExistError<role_ident::Resource>>, MetaError> {
+        let seq = if can_replace {
+            MatchSeq::GE(0)
+        } else {
+            MatchSeq::Exact(0)
+        };
 
-        let upsert_kv = self.kv_api.upsert_kv(UpsertKV::new(
-            &key,
-            match_seq,
-            Operation::Update(value),
-            None,
-        ));
+        let key = RoleIdent::new(&self.tenant, &info.name);
+        let req = UpsertPB::insert(key, info.clone()).with(seq);
+        let res = self.kv_api.upsert_pb(&req).await?;
 
-        let res_seq = upsert_kv.await?.added_seq_or_else(|_v| {
-            ErrorCode::RoleAlreadyExists(format!("Role '{}' already exists.", role_info.name))
-        })?;
+        if !can_replace && res.prev.is_some() {
+            return Ok(Err(
+                RoleIdent::new(&self.tenant, &info.name).exist_error(func_name!())
+            ));
+        }
 
-        Ok(res_seq)
+        Ok(Ok(()))
     }
 
     #[async_backtrace::framed]

src/query/service/src/interpreters/interpreter_role_create.rs

@@ -70,7 +70,7 @@ impl Interpreter for CreateRoleInterpreter {
         let tenant = self.ctx.get_tenant();
         let user_mgr = UserApiProvider::instance();
         user_mgr
-            .add_role(&tenant, RoleInfo::new(&role_name), plan.if_not_exists)
+            .add_role(&tenant, RoleInfo::new(&role_name), &plan.create_option)
             .await?;
         RoleCacheManager::instance().force_reload(&tenant).await?;
         Ok(PipelineBuildResult::create())

src/query/service/tests/it/storages/system.rs

@@ -25,6 +25,9 @@ use databend_common_meta_app::principal::AuthType;
 use databend_common_meta_app::principal::RoleInfo;
 use databend_common_meta_app::principal::UserInfo;
 use databend_common_meta_app::schema::CreateOption;
+use databend_common_meta_app::schema::CreateOption::Create;
+use databend_common_meta_app::schema::CreateOption::CreateIfNotExists;
+use databend_common_meta_app::schema::CreateOption::CreateOrReplace;
 use databend_common_meta_app::storage::StorageFsConfig;
 use databend_common_meta_app::storage::StorageParams;
 use databend_common_meta_app::storage::StorageS3Config;
@@ -381,17 +384,26 @@ async fn test_roles_table() -> Result<()> {
     {
         let role_info = RoleInfo::new("test");
         UserApiProvider::instance()
-            .add_role(&tenant, role_info, false)
+            .add_role(&tenant, role_info, &CreateOrReplace)
             .await?;
     }
 
     {
         let mut role_info = RoleInfo::new("test1");
         role_info.grants.grant_role("test".to_string());
         UserApiProvider::instance()
-            .add_role(&tenant, role_info, false)
+            .add_role(&tenant, role_info, &Create)
             .await?;
     }
+
+    {
+        let mut role_info = RoleInfo::new("test1");
+        role_info.grants.grant_role("t2".to_string());
+        UserApiProvider::instance()
+            .add_role(&tenant, role_info, &CreateIfNotExists)
+            .await?;
+    }
+
     let table = RolesTable::create(1);
     let source_plan = table
         .read_plan(ctx.clone(), None, None, false, true)

src/query/sql/src/planner/binder/binder.rs

@@ -388,7 +388,7 @@ impl Binder {
                 self.bind_show_roles(bind_context, show_options).await?
             }
             Statement::CreateRole {
-                if_not_exists,
+                create_option,
                 role_name,
             } => {
                 if illegal_ident_name(role_name) {
@@ -397,7 +397,7 @@ impl Binder {
                     ));
                 }
                 Plan::CreateRole(Box::new(CreateRolePlan {
-                    if_not_exists: *if_not_exists,
+                    create_option: create_option.clone().into(),
                     role_name: role_name.to_string(),
                 }))
             }