feat: add create or replace user support

src/query/ast/src/ast/statements/user.rs

@@ -21,22 +21,29 @@ use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserOption;
 use databend_common_meta_app::principal::UserOptionFlag;
 use databend_common_meta_app::principal::UserPrivilegeType;
+use databend_common_meta_app::schema::CreateOption;
 
 use crate::ast::write_comma_separated_list;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CreateUserStmt {
-    pub if_not_exists: bool,
+    pub create_option: CreateOption,
     pub user: UserIdentity,
     pub auth_option: AuthOption,
     pub user_options: Vec<UserOptionItem>,
 }
 
 impl Display for CreateUserStmt {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "CREATE USER")?;
-        if self.if_not_exists {
-            write!(f, " IF NOT EXISTS")?;
+        write!(f, "CREATE")?;
+        if let CreateOption::CreateOrReplace = self.create_option {
+            write!(f, " OR REPLACE")?;
+        }
+        write!(f, " USER")?;
+        if let CreateOption::CreateIfNotExists(if_not_exists) = self.create_option {
+            if if_not_exists {
+                write!(f, " IF NOT EXISTS")?;
+            }
         }
         write!(f, " {} IDENTIFIED", self.user)?;
         write!(f, " {}", self.auth_option)?;

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

@@ -1061,16 +1061,28 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
     );
 
     let show_users = value(Statement::ShowUsers, rule! { SHOW ~ USERS });
-    let create_user = map(
+    let create_user = map_res(
         rule! {
-            CREATE ~ USER ~ ( IF ~ ^NOT ~ ^EXISTS )?
+            CREATE ~  (OR ~ REPLACE)? ~ USER ~ ( IF ~ ^NOT ~ ^EXISTS )?
             ~ #user_identity
             ~ IDENTIFIED ~ ( WITH ~ ^#auth_type )? ~ ( BY ~ ^#literal_string )?
             ~ ( WITH ~ ^#comma_separated_list1(user_option))?
         },
-        |(_, _, opt_if_not_exists, user, _, opt_auth_type, opt_password, opt_user_option)| {
-            Statement::CreateUser(CreateUserStmt {
-                if_not_exists: opt_if_not_exists.is_some(),
+        |(
+            _,
+            opt_or_replace,
+            _,
+            opt_if_not_exists,
+            user,
+            _,
+            opt_auth_type,
+            opt_password,
+            opt_user_option,
+        )| {
+            let create_option =
+                parse_create_option(opt_or_replace.is_some(), opt_if_not_exists.is_some())?;
+            Ok(Statement::CreateUser(CreateUserStmt {
+                create_option,
                 user,
                 auth_option: AuthOption {
                     auth_type: opt_auth_type.map(|(_, auth_type)| auth_type),
@@ -1079,7 +1091,7 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
                 user_options: opt_user_option
                     .map(|(_, user_options)| user_options)
                     .unwrap_or_default(),
-            })
+            }))
         },
     );
     let alter_user = map(
@@ -1887,7 +1899,7 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
         ),
         rule!(
             #show_users : "`SHOW USERS`"
-            | #create_user : "`CREATE USER [IF NOT EXISTS] '<username>'@'hostname' IDENTIFIED [WITH <auth_type>] [BY <password>] [WITH <user_option>, ...]`"
+            | #create_user : "`CREATE [OR REPLACE] USER [IF NOT EXISTS] '<username>'@'hostname' IDENTIFIED [WITH <auth_type>] [BY <password>] [WITH <user_option>, ...]`"
             | #alter_user : "`ALTER USER ('<username>'@'hostname' | USER()) [IDENTIFIED [WITH <auth_type>] [BY <password>]] [WITH <user_option>, ...]`"
             | #drop_user : "`DROP USER [IF EXISTS] '<username>'@'hostname'`"
             | #show_roles : "`SHOW ROLES`"

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

@@ -3370,7 +3370,9 @@ CREATE USER 'u1'@'%' IDENTIFIED BY '123456' WITH DEFAULT_ROLE = 'role123' TENANT
 ---------- AST ------------
 CreateUser(
     CreateUserStmt {
-        if_not_exists: false,
+        create_option: CreateIfNotExists(
+            false,
+        ),
         user: UserIdentity {
             username: "u1",
             hostname: "%",
@@ -3400,7 +3402,9 @@ CREATE USER 'u1'@'%' IDENTIFIED BY '123456' WITH SET NETWORK POLICY = 'policy1'
 ---------- AST ------------
 CreateUser(
     CreateUserStmt {
-        if_not_exists: false,
+        create_option: CreateIfNotExists(
+            false,
+        ),
         user: UserIdentity {
             username: "u1",
             hostname: "%",
@@ -7719,7 +7723,9 @@ CREATE USER 'test-e'@'%' IDENTIFIED BY 'password'
 ---------- AST ------------
 CreateUser(
     CreateUserStmt {
-        if_not_exists: false,
+        create_option: CreateIfNotExists(
+            false,
+        ),
         user: UserIdentity {
             username: "test-e",
             hostname: "%",

src/query/management/src/user/user_api.rs

@@ -15,12 +15,13 @@
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_meta_types::MatchSeq;
 use databend_common_meta_types::SeqV;
 
 #[async_trait::async_trait]
 pub trait UserApi: Sync + Send {
-    async fn add_user(&self, user_info: UserInfo) -> Result<u64>;
+    async fn add_user(&self, user_info: UserInfo, create_option: &CreateOption) -> Result<()>;
 
     async fn get_user(&self, user: UserIdentity, seq: MatchSeq) -> Result<SeqV<UserInfo>>;
 

src/query/management/src/user/user_mgr.rs

@@ -19,6 +19,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_meta_kvapi::kvapi;
 use databend_common_meta_kvapi::kvapi::UpsertKVReq;
 use databend_common_meta_types::MatchSeq;
@@ -81,25 +82,34 @@ impl UserMgr {
 impl UserApi for UserMgr {
     #[async_backtrace::framed]
     #[minitrace::trace]
-    async fn add_user(&self, user_info: UserInfo) -> databend_common_exception::Result<u64> {
-        let match_seq = MatchSeq::Exact(0);
+    async fn add_user(
+        &self,
+        user_info: UserInfo,
+        create_option: &CreateOption,
+    ) -> databend_common_exception::Result<()> {
         let user_key = format_user_key(&user_info.name, &user_info.hostname);
         let key = format!("{}/{}", self.user_prefix, escape_for_key(&user_key)?);
         let value = serialize_struct(&user_info, ErrorCode::IllegalUserInfoFormat, || "")?;
 
-        let kv_api = self.kv_api.clone();
-        let upsert_kv = kv_api.upsert_kv(UpsertKVReq::new(
-            &key,
-            match_seq,
-            Operation::Update(value),
-            None,
-        ));
-
-        let res_seq = upsert_kv.await?.added_seq_or_else(|_v| {
-            ErrorCode::UserAlreadyExists(format!("User {} already exists.", user_key))
-        })?;
-
-        Ok(res_seq)
+        let kv_api = &self.kv_api;
+        let seq = match create_option {
+            CreateOption::CreateIfNotExists(_) => MatchSeq::Exact(0),
+            CreateOption::CreateOrReplace => MatchSeq::GE(0),
+        };
+        let res = kv_api
+            .upsert_kv(UpsertKVReq::new(&key, seq, Operation::Update(value), None))
+            .await?;
+
+        if let CreateOption::CreateIfNotExists(false) = create_option {
+            if res.prev.is_some() {
+                return Err(ErrorCode::UserAlreadyExists(format!(
+                    "User {} already exists.",
+                    user_key
+                )));
+            }
+        }
+
+        Ok(())
     }
 
     #[async_backtrace::framed]

src/query/management/tests/it/user.rs

@@ -81,6 +81,7 @@ fn default_test_auth_info() -> AuthInfo {
 
 mod add {
     use databend_common_meta_app::principal::UserInfo;
+    use databend_common_meta_app::schema::CreateOption;
     use databend_common_meta_types::Operation;
 
     use super::*;
@@ -119,7 +120,7 @@ mod add {
                 .return_once(|_u| Ok(UpsertKVReply::new(None, Some(SeqV::new(1, v)))));
             let api = Arc::new(api);
             let user_mgr = UserMgr::create(api, "tenant1")?;
-            let res = user_mgr.add_user(user_info);
+            let res = user_mgr.add_user(user_info, &CreateOption::CreateIfNotExists(false));
 
             assert!(res.await.is_ok());
         }
@@ -148,7 +149,9 @@ mod add {
 
             let user_info = UserInfo::new(test_user_name, test_hostname, default_test_auth_info());
 
-            let res = user_mgr.add_user(user_info).await;
+            let res = user_mgr
+                .add_user(user_info, &CreateOption::CreateIfNotExists(false))
+                .await;
 
             assert_eq!(
                 res.unwrap_err().code(),

src/query/service/src/auth.rs

@@ -21,6 +21,7 @@ use databend_common_exception::Result;
 use databend_common_meta_app::principal::AuthInfo;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_users::JwtAuthenticator;
 use databend_common_users::UserApiProvider;
 
@@ -116,7 +117,13 @@ impl AuthMgr {
                                 user_info.grants.grant_role(role);
                             }
                         }
-                        user_api.add_user(&tenant, user_info.clone(), true).await?;
+                        user_api
+                            .add_user(
+                                &tenant,
+                                user_info.clone(),
+                                &CreateOption::CreateIfNotExists(true),
+                            )
+                            .await?;
                         user_info
                     }
                 };

src/query/service/src/interpreters/interpreter_user_create.rs

@@ -80,7 +80,7 @@ impl Interpreter for CreateUserInterpreter {
             lockout_time: None,
         };
         user_mgr
-            .add_user(&tenant, user_info, plan.if_not_exists)
+            .add_user(&tenant, user_info, &plan.create_option)
             .await?;
 
         Ok(PipelineBuildResult::create())

src/query/service/tests/it/auth.rs

@@ -18,6 +18,7 @@ use databend_common_base::base::tokio;
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::AuthInfo;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_users::CustomClaims;
 use databend_common_users::EnsureUser;
 use databend_common_users::UserApiProvider;
@@ -152,7 +153,11 @@ async fn test_auth_mgr_with_jwt_multi_sources() -> Result<()> {
         let tenant = ctx.get_current_session().get_current_tenant();
         let user2_info = UserInfo::new(user2, "%", AuthInfo::JWT);
         UserApiProvider::instance()
-            .add_user(tenant.as_str(), user2_info.clone(), true)
+            .add_user(
+                tenant.as_str(),
+                user2_info.clone(),
+                &CreateOption::CreateIfNotExists(true),
+            )
             .await?;
         let res2 = auth_mgr
             .auth(ctx.get_current_session(), &Credential::Jwt {

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

@@ -27,6 +27,7 @@ use databend_common_meta_app::principal::UserGrantSet;
 use databend_common_meta_app::principal::UserInfo;
 use databend_common_meta_app::principal::UserOption;
 use databend_common_meta_app::principal::UserQuota;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_meta_app::storage::StorageFsConfig;
 use databend_common_meta_app::storage::StorageParams;
 use databend_common_meta_app::storage::StorageS3Config;
@@ -398,7 +399,7 @@ async fn test_users_table() -> Result<()> {
                 password_update_on: None,
                 lockout_time: None,
             },
-            false,
+            &CreateOption::CreateIfNotExists(false),
         )
         .await?;
     let auth_data = AuthInfo::new(AuthType::Sha256Password, &Some("123456789".to_string()));
@@ -418,7 +419,7 @@ async fn test_users_table() -> Result<()> {
                 password_update_on: None,
                 lockout_time: None,
             },
-            false,
+            &CreateOption::CreateIfNotExists(false),
         )
         .await?;
 

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

@@ -228,7 +228,7 @@ impl Binder {
         stmt: &CreateUserStmt,
     ) -> Result<Plan> {
         let CreateUserStmt {
-            if_not_exists,
+            create_option,
             user,
             auth_option,
             user_options,
@@ -248,7 +248,7 @@ impl Binder {
             .await?;
 
         let plan = CreateUserPlan {
-            if_not_exists: *if_not_exists,
+            create_option: create_option.clone(),
             user: user.clone(),
             auth_info: AuthInfo::create2(&auth_option.auth_type, &auth_option.password)?,
             user_option,

src/query/sql/src/planner/plans/ddl/account.rs

@@ -27,10 +27,11 @@ use databend_common_meta_app::principal::PrincipalIdentity;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserOption;
 use databend_common_meta_app::principal::UserPrivilegeSet;
+use databend_common_meta_app::schema::CreateOption;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct CreateUserPlan {
-    pub if_not_exists: bool,
+    pub create_option: CreateOption,
     pub user: UserIdentity,
     pub auth_info: AuthInfo,
     pub user_option: UserOption,

src/query/users/src/user_mgr.rs

@@ -26,6 +26,7 @@ use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
 use databend_common_meta_app::principal::UserOption;
 use databend_common_meta_app::principal::UserPrivilegeSet;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_meta_types::MatchSeq;
 
 use crate::role_mgr::BUILTIN_ROLE_ACCOUNT_ADMIN;
@@ -128,8 +129,8 @@ impl UserApiProvider {
         &self,
         tenant: &str,
         user_info: UserInfo,
-        if_not_exists: bool,
-    ) -> Result<u64> {
+        create_option: &CreateOption,
+    ) -> Result<()> {
         if let Some(name) = user_info.option.network_policy() {
             if self.get_network_policy(tenant, name).await.is_err() {
                 return Err(ErrorCode::UnknownNetworkPolicy(format!(
@@ -153,17 +154,7 @@ impl UserApiProvider {
             )));
         }
         let client = self.get_user_api_client(tenant)?;
-        let add_user = client.add_user(user_info);
-        match add_user.await {
-            Ok(res) => Ok(res),
-            Err(e) => {
-                if if_not_exists && e.code() == ErrorCode::USER_ALREADY_EXISTS {
-                    Ok(0)
-                } else {
-                    Err(e.add_message_back("(while add user)"))
-                }
-            }
-        }
+        client.add_user(user_info, create_option).await
     }
 
     #[async_backtrace::framed]

src/query/users/tests/it/network_policy.rs

@@ -23,6 +23,7 @@ use databend_common_meta_app::principal::PasswordHashMethod;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
 use databend_common_meta_app::principal::UserOption;
+use databend_common_meta_app::schema::CreateOption;
 use databend_common_users::UserApiProvider;
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
@@ -62,7 +63,9 @@ async fn test_network_policy() -> Result<()> {
     let mut option = UserOption::empty();
     option = option.with_network_policy(Some(policy_name.clone()));
     user_info.update_auth_option(None, Some(option));
-    user_mgr.add_user(tenant, user_info, false).await?;
+    user_mgr
+        .add_user(tenant, user_info, &CreateOption::CreateIfNotExists(false))
+        .await?;
 
     let user = UserIdentity::new(username, hostname);