feat(query): support manage query workload through sql (#17881)

* feat(query): support manage query workload through sql

* feat(query): support manage query workload through sql

* feat(query): support manage query workload through sql

Author: zhang2014

src/common/license/src/license.rs

@@ -71,12 +71,14 @@ pub enum Feature {
     StorageQuota(StorageQuota),
     #[serde(alias = "amend_table", alias = "AMEND_TABLE")]
     AmendTable,
-    #[serde(alias = "system_management", alias = "SYSTEM_MANAGEMENT")]
-    SystemManagement,
     #[serde(alias = "hilbert_clustering", alias = "HILBERT_CLUSTERING")]
     HilbertClustering,
+    #[serde(alias = "system_management", alias = "SYSTEM_MANAGEMENT")]
+    SystemManagement,
     #[serde(alias = "ngram_index", alias = "NGRAM_INDEX")]
     NgramIndex,
+    #[serde(alias = "workload_group", alias = "WORKLOAD_GROUP")]
+    WorkloadGroup,
     #[serde(other)]
     Unknown,
 }
@@ -125,6 +127,7 @@ impl fmt::Display for Feature {
             Feature::SystemManagement => write!(f, "system_management"),
             Feature::HilbertClustering => write!(f, "hilbert_clustering"),
             Feature::NgramIndex => write!(f, "ngram_index"),
+            Feature::WorkloadGroup => write!(f, "workload_group"),
             Feature::Unknown => write!(f, "unknown"),
         }
     }

src/query/ast/src/ast/statements/mod.rs

@@ -56,6 +56,7 @@ mod user;
 mod view;
 mod virtual_column;
 mod warehouse;
+mod workload;
 
 pub use call::*;
 pub use catalog::*;
@@ -101,3 +102,4 @@ pub use user::*;
 pub use view::*;
 pub use virtual_column::*;
 pub use warehouse::*;
+pub use workload::*;

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

@@ -30,6 +30,11 @@ 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::ShowWorkloadGroupsStmt;
 use crate::ast::write_comma_separated_list;
 use crate::ast::CreateOption;
 use crate::ast::Identifier;
@@ -146,6 +151,13 @@ pub enum Statement {
     AssignWarehouseNodes(AssignWarehouseNodesStmt),
     UnassignWarehouseNodes(UnassignWarehouseNodesStmt),
 
+    // Workloads
+    ShowWorkloadGroups(ShowWorkloadGroupsStmt),
+    CreateWorkloadGroup(CreateWorkloadGroupStmt),
+    DropWorkloadGroup(DropWorkloadGroupStmt),
+    RenameWorkloadGroup(RenameWorkloadGroupStmt),
+    AlterWorkloadGroup(AlterWorkloadGroupStmt),
+
     // Databases
     ShowDatabases(ShowDatabasesStmt),
     ShowDropDatabases(ShowDropDatabasesStmt),
@@ -578,7 +590,12 @@ impl Statement {
             | Statement::AssignWarehouseNodes(..)
             | Statement::UnassignWarehouseNodes(..)
             | Statement::ResumeWarehouse(..)
-            | Statement::SuspendWarehouse(..) => false,
+            | Statement::SuspendWarehouse(..)
+            | Statement::ShowWorkloadGroups(..)
+            | Statement::CreateWorkloadGroup(..)
+            | Statement::DropWorkloadGroup(..)
+            | Statement::RenameWorkloadGroup(..)
+            | Statement::AlterWorkloadGroup(..) => false,
             Statement::StatementWithSettings { stmt, settings: _ } => {
                 stmt.allowed_in_multi_statement()
             }
@@ -1022,6 +1039,11 @@ impl Display for Statement {
             Statement::RenameWarehouseCluster(stmt) => write!(f, "{stmt}")?,
             Statement::AssignWarehouseNodes(stmt) => write!(f, "{stmt}")?,
             Statement::UnassignWarehouseNodes(stmt) => write!(f, "{stmt}")?,
+            Statement::ShowWorkloadGroups(stmt) => write!(f, "{stmt}")?,
+            Statement::CreateWorkloadGroup(stmt) => write!(f, "{stmt}")?,
+            Statement::DropWorkloadGroup(stmt) => write!(f, "{stmt}")?,
+            Statement::RenameWorkloadGroup(stmt) => write!(f, "{stmt}")?,
+            Statement::AlterWorkloadGroup(stmt) => write!(f, "{stmt}")?,
         }
         Ok(())
     }