feat(query): support unset quota for workload group (#17924)

* feat(query): support unset quota for workload group

* feat(query): support unset quota for workload group

Author: zhang2014

src/common/exception/src/exception_code.rs

@@ -409,6 +409,7 @@ build_exceptions! {
     AlreadyExistsWorkload(3141),
     UnknownWorkload(3142),
     WorkloadOperateConflict(3143),
+    UnknownWorkloadQuotas(3144),
 }
 
 // Storage errors [3001, 4000].

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

@@ -30,10 +30,10 @@ use crate::ast::statements::pipe::CreatePipeStmt;
 use crate::ast::statements::settings::Settings;
 use crate::ast::statements::task::CreateTaskStmt;
 use crate::ast::statements::warehouse::ShowWarehousesStmt;
-use crate::ast::statements::workload::AlterWorkloadGroupStmt;
 use crate::ast::statements::workload::CreateWorkloadGroupStmt;
 use crate::ast::statements::workload::DropWorkloadGroupStmt;
 use crate::ast::statements::workload::RenameWorkloadGroupStmt;
+use crate::ast::statements::workload::SetWorkloadGroupQuotasStmt;
 use crate::ast::statements::workload::ShowWorkloadGroupsStmt;
 use crate::ast::write_comma_separated_list;
 use crate::ast::CreateOption;
@@ -156,7 +156,8 @@ pub enum Statement {
     CreateWorkloadGroup(CreateWorkloadGroupStmt),
     DropWorkloadGroup(DropWorkloadGroupStmt),
     RenameWorkloadGroup(RenameWorkloadGroupStmt),
-    AlterWorkloadGroup(AlterWorkloadGroupStmt),
+    SetWorkloadQuotasGroup(SetWorkloadGroupQuotasStmt),
+    UnsetWorkloadQuotasGroup(UnsetWorkloadGroupQuotasStmt),
 
     // Databases
     ShowDatabases(ShowDatabasesStmt),
@@ -589,7 +590,8 @@ impl Statement {
             | Statement::CreateWorkloadGroup(..)
             | Statement::DropWorkloadGroup(..)
             | Statement::RenameWorkloadGroup(..)
-            | Statement::AlterWorkloadGroup(..) => false,
+            | Statement::SetWorkloadQuotasGroup(..)
+            | Statement::UnsetWorkloadQuotasGroup(..) => false,
             Statement::StatementWithSettings { stmt, settings: _ } => {
                 stmt.allowed_in_multi_statement()
             }
@@ -1034,7 +1036,8 @@ impl Display for Statement {
             Statement::CreateWorkloadGroup(stmt) => write!(f, "{stmt}")?,
             Statement::DropWorkloadGroup(stmt) => write!(f, "{stmt}")?,
             Statement::RenameWorkloadGroup(stmt) => write!(f, "{stmt}")?,
-            Statement::AlterWorkloadGroup(stmt) => write!(f, "{stmt}")?,
+            Statement::SetWorkloadQuotasGroup(stmt) => write!(f, "{stmt}")?,
+            Statement::UnsetWorkloadQuotasGroup(stmt) => write!(f, "{stmt}")?,
         }
         Ok(())
     }

src/query/ast/src/ast/statements/workload.rs

@@ -218,18 +218,18 @@ impl Display for RenameWorkloadGroupStmt {
 }
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
-pub struct AlterWorkloadGroupStmt {
+pub struct SetWorkloadGroupQuotasStmt {
     pub name: Identifier,
     #[drive(skip)]
     pub quotas: BTreeMap<String, QuotaValueStmt>,
 }
 
-impl Display for AlterWorkloadGroupStmt {
+impl Display for SetWorkloadGroupQuotasStmt {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "ALTER WORKLOAD GROUP {}", self.name)?;
 
         if !self.quotas.is_empty() {
-            write!(f, " SET ")?;
+            write!(f, " SET")?;
 
             for (idx, (key, value)) in self.quotas.iter().enumerate() {
                 if idx != 0 {
@@ -244,6 +244,33 @@ impl Display for AlterWorkloadGroupStmt {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct UnsetWorkloadGroupQuotasStmt {
+    pub name: Identifier,
+    #[drive(skip)]
+    pub quotas: Vec<Identifier>,
+}
+
+impl Display for UnsetWorkloadGroupQuotasStmt {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "ALTER WORKLOAD GROUP {}", self.name)?;
+
+        if !self.quotas.is_empty() {
+            write!(f, " UNSET")?;
+
+            for (idx, name) in self.quotas.iter().enumerate() {
+                if idx != 0 {
+                    write!(f, ",")?;
+                }
+
+                write!(f, " {}", name)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::time::Duration;

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

@@ -712,12 +712,21 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
     );
 
-    let alter_workload_group = map(
+    let set_workload_group_quotas = map(
         rule! {
             ALTER ~ WORKLOAD ~ GROUP ~ #ident ~ SET ~ #workload_quotas
         },
         |(_, _, _, name, _, quotas)| {
-            Statement::AlterWorkloadGroup(AlterWorkloadGroupStmt { name, quotas })
+            Statement::SetWorkloadQuotasGroup(SetWorkloadGroupQuotasStmt { name, quotas })
+        },
+    );
+
+    let unset_workload_group_quotas = map(
+        rule! {
+            ALTER ~ WORKLOAD ~ GROUP ~ #ident ~ UNSET ~ #unset_source
+        },
+        |(_, _, _, name, _, quotas)| {
+            Statement::UnsetWorkloadQuotasGroup(UnsetWorkloadGroupQuotasStmt { name, quotas })
         },
     );
 
@@ -2556,7 +2565,8 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #create_workload_group: "`CREATE WORKLOAD GROUP [IF NOT EXISTS] <name> WITH [<workload_group_quotas>]`"
             | #drop_workload_group: "`DROP WORKLOAD GROUP [IF EXISTS] <name>`"
             | #rename_workload_group: "`RENAME WORKLOAD GROUP <old_name> TO <new_name>`"
-            | #alter_workload_group: "`ALTER WORKLOAD GROUP <name> set [<workload_group_quotas>]`"
+            | #set_workload_group_quotas: "`ALTER WORKLOAD GROUP <name> SET [<workload_group_quotas>]`"
+            | #unset_workload_group_quotas: "`ALTER WORKLOAD GROUP <name> UNSET {<name> | (<name>, ...)}`"
         ),
         // database
         rule!(

src/query/management/src/workload/workload_api.rs

@@ -65,7 +65,9 @@ pub trait WorkloadApi: Sync + Send {
 
     async fn rename(&self, old_name: String, new_name: String) -> Result<()>;
 
-    async fn alter_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()>;
+    async fn set_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()>;
+
+    async fn unset_quotas(&self, name: String, quotas: Vec<String>) -> Result<()>;
 
     async fn get_all(&self) -> Result<Vec<WorkloadGroup>>;
 

src/query/management/src/workload/workload_mgr.rs

@@ -201,7 +201,7 @@ impl WorkloadApi for WorkloadMgr {
         }
     }
 
-    async fn alter_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()> {
+    async fn set_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()> {
         for _index in 0..5 {
             let workload = self.get_seq_by_name(&name).await?;
             let seq = workload.seq;
@@ -230,6 +230,40 @@ impl WorkloadApi for WorkloadMgr {
         ))
     }
 
+    async fn unset_quotas(&self, name: String, quotas: Vec<String>) -> Result<()> {
+        for _index in 0..5 {
+            let workload = self.get_seq_by_name(&name).await?;
+            let seq = workload.seq;
+            let mut workload = workload.into_value().unwrap();
+
+            for quota_name in &quotas {
+                if workload.quotas.remove(quota_name).is_none() {
+                    return Err(ErrorCode::UnknownWorkloadQuotas(format!(
+                        "Unknown workload group quota name {} in {}",
+                        quota_name, name
+                    )));
+                }
+            }
+
+            let workload_key = format!("{}/{}", self.workload_key_prefix, workload.id);
+            let mut alter_workload = TxnRequest::default();
+            alter_workload
+                .condition
+                .push(TxnCondition::eq_seq(workload_key.clone(), seq));
+            alter_workload
+                .if_then
+                .push(TxnOp::put(workload_key, serde_json::to_vec(&workload)?));
+
+            if self.metastore.transaction(alter_workload).await?.success {
+                return Ok(());
+            }
+        }
+
+        Err(ErrorCode::WorkloadOperateConflict(
+            "Workload operate conflict(tried 5 times).",
+        ))
+    }
+
     async fn get_all(&self) -> Result<Vec<WorkloadGroup>> {
         let list_reply = self
             .metastore

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

@@ -189,8 +189,7 @@ async fn test_alter_quotas() -> Result<()> {
         QuotaValue::Duration(Duration::from_secs(10)),
     );
 
-    mgr.alter_quotas("test_group".to_string(), new_quotas)
-        .await?;
+    mgr.set_quotas("test_group".to_string(), new_quotas).await?;
 
     let updated = mgr.get_by_name("test_group").await?;
     assert_eq!(updated.quotas.len(), 1);
@@ -206,7 +205,7 @@ async fn test_alter_quotas() -> Result<()> {
         QuotaValue::Duration(Duration::from_secs(20)),
     );
 
-    mgr.alter_quotas("test_group".to_string(), update_quotas)
+    mgr.set_quotas("test_group".to_string(), update_quotas)
         .await?;
 
     let updated = mgr.get_by_name("test_group").await?;

src/query/service/src/interpreters/access/privilege_access.rs

@@ -1373,7 +1373,8 @@ impl AccessChecker for PrivilegeAccess {
             Plan::CreateWorkloadGroup(_) => {}
             Plan::DropWorkloadGroup(_) => {}
             Plan::RenameWorkloadGroup(_) => {}
-            Plan::AlterWorkloadGroup(_) => {}
+            Plan::SetWorkloadGroupQuotas(_) => {}
+            Plan::UnsetWorkloadGroupQuotas(_) => {}
         }
 
         Ok(())

src/query/service/src/interpreters/interpreter_factory.rs

@@ -39,7 +39,6 @@ use super::interpreter_user_stage_drop::DropUserStageInterpreter;
 use super::*;
 use crate::interpreters::access::Accessor;
 use crate::interpreters::interpreter_add_warehouse_cluster::AddWarehouseClusterInterpreter;
-use crate::interpreters::interpreter_alter_workload_group::AlterWorkloadGroupInterpreter;
 use crate::interpreters::interpreter_assign_warehouse_nodes::AssignWarehouseNodesInterpreter;
 use crate::interpreters::interpreter_catalog_drop::DropCatalogInterpreter;
 use crate::interpreters::interpreter_connection_create::CreateConnectionInterpreter;
@@ -72,6 +71,7 @@ use crate::interpreters::interpreter_rename_warehouse_cluster::RenameWarehouseCl
 use crate::interpreters::interpreter_rename_workload_group::RenameWorkloadGroupInterpreter;
 use crate::interpreters::interpreter_resume_warehouse::ResumeWarehouseInterpreter;
 use crate::interpreters::interpreter_set_priority::SetPriorityInterpreter;
+use crate::interpreters::interpreter_set_workload_group_quotas::SetWorkloadGroupQuotasInterpreter;
 use crate::interpreters::interpreter_show_online_nodes::ShowOnlineNodesInterpreter;
 use crate::interpreters::interpreter_show_warehouses::ShowWarehousesInterpreter;
 use crate::interpreters::interpreter_show_workload_groups::ShowWorkloadGroupsInterpreter;
@@ -90,6 +90,7 @@ use crate::interpreters::interpreter_txn_abort::AbortInterpreter;
 use crate::interpreters::interpreter_txn_begin::BeginInterpreter;
 use crate::interpreters::interpreter_txn_commit::CommitInterpreter;
 use crate::interpreters::interpreter_unassign_warehouse_nodes::UnassignWarehouseNodesInterpreter;
+use crate::interpreters::interpreter_unset_workload_group_quotas::UnsetWorkloadGroupQuotasInterpreter;
 use crate::interpreters::interpreter_use_warehouse::UseWarehouseInterpreter;
 use crate::interpreters::interpreter_view_describe::DescribeViewInterpreter;
 use crate::interpreters::AlterUserInterpreter;
@@ -734,10 +735,12 @@ impl InterpreterFactory {
             Plan::RenameWorkloadGroup(p) => Ok(Arc::new(
                 RenameWorkloadGroupInterpreter::try_create(ctx, *p.clone())?,
             )),
-            Plan::AlterWorkloadGroup(p) => Ok(Arc::new(AlterWorkloadGroupInterpreter::try_create(
-                ctx,
-                *p.clone(),
-            )?)),
+            Plan::SetWorkloadGroupQuotas(p) => Ok(Arc::new(
+                SetWorkloadGroupQuotasInterpreter::try_create(ctx, *p.clone())?,
+            )),
+            Plan::UnsetWorkloadGroupQuotas(p) => Ok(Arc::new(
+                UnsetWorkloadGroupQuotasInterpreter::try_create(ctx, *p.clone())?,
+            )),
             // Plan::ShowCreateProcedure(_) => {}
             //
             // Plan::RenameProcedure(p) => Ok(Arc::new(RenameProcedureInterpreter::try_create(

…ters/interpreter_alter_workload_group.rs …interpreter_set_workload_group_quotas.rssrc/query/service/src/interpreters/interpreter_alter_workload_group.rs renamed to src/query/service/src/interpreters/interpreter_set_workload_group_quotas.rs src/query/service/src/interpreters/interpreter_alter_workload_group.rs renamed to src/query/service/src/interpreters/interpreter_set_workload_group_quotas.rs

@@ -22,29 +22,29 @@ use databend_common_license::license::Feature;
 use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_management::WorkloadApi;
 use databend_common_management::WorkloadMgr;
-use databend_common_sql::plans::AlterWorkloadGroupPlan;
+use databend_common_sql::plans::SetWorkloadGroupQuotasPlan;
 
 use crate::interpreters::interpreter_create_workload_group::to_quota_value;
 use crate::interpreters::util::AuditElement;
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
 use crate::sessions::QueryContext;
 
-pub struct AlterWorkloadGroupInterpreter {
+pub struct SetWorkloadGroupQuotasInterpreter {
     ctx: Arc<QueryContext>,
-    plan: AlterWorkloadGroupPlan,
+    plan: SetWorkloadGroupQuotasPlan,
 }
 
-impl AlterWorkloadGroupInterpreter {
-    pub fn try_create(ctx: Arc<QueryContext>, plan: AlterWorkloadGroupPlan) -> Result<Self> {
-        Ok(AlterWorkloadGroupInterpreter { ctx, plan })
+impl SetWorkloadGroupQuotasInterpreter {
+    pub fn try_create(ctx: Arc<QueryContext>, plan: SetWorkloadGroupQuotasPlan) -> Result<Self> {
+        Ok(SetWorkloadGroupQuotasInterpreter { ctx, plan })
     }
 }
 
 #[async_trait::async_trait]
-impl Interpreter for AlterWorkloadGroupInterpreter {
+impl Interpreter for SetWorkloadGroupQuotasInterpreter {
     fn name(&self) -> &str {
-        "AlterWorkloadGroupInterpreter"
+        "SetWorkloadGroupQuotasInterpreter"
     }
 
     fn is_ddl(&self) -> bool {
@@ -64,14 +64,14 @@ impl Interpreter for AlterWorkloadGroupInterpreter {
 
         let workload_manager = GlobalInstance::get::<Arc<WorkloadMgr>>();
         workload_manager
-            .alter_quotas(self.plan.name.clone(), workload_quotas)
+            .set_quotas(self.plan.name.clone(), workload_quotas)
             .await?;
 
         let user_info = self.ctx.get_current_user()?;
         log::info!(
             target: "databend::log::audit",
             "{}",
-            serde_json::to_string(&AuditElement::create(&user_info, "alter_workload", &self.plan))?
+            serde_json::to_string(&AuditElement::create(&user_info, "set_workload_quotas", &self.plan))?
         );
         Ok(PipelineBuildResult::create())
     }