feat(query): UserOption add workload_group (#17913)

* feat(query): UserOption add workload_group

* add more test

Author: TCeason

Cargo.lock

@@ -191,6 +191,7 @@ pub struct UserOption {
     default_role: Option<String>,
     network_policy: Option<String>,
     password_policy: Option<String>,
+    workload_group: Option<String>,
     disabled: Option<bool>,
     must_change_password: Option<bool>,
 }
@@ -202,6 +203,7 @@ impl UserOption {
             default_role: None,
             network_policy: None,
             password_policy: None,
+            workload_group: None,
             disabled: None,
             must_change_password: None,
         }
@@ -231,6 +233,11 @@ impl UserOption {
         self
     }
 
+    pub fn with_workload_group(mut self, workload_group: Option<String>) -> Self {
+        self.workload_group = workload_group;
+        self
+    }
+
     pub fn with_disabled(mut self, disabled: Option<bool>) -> Self {
         self.disabled = disabled;
         self
@@ -262,6 +269,11 @@ impl UserOption {
         self.password_policy.as_ref()
     }
 
+    // This is workload_group id
+    pub fn workload_group(&self) -> Option<&String> {
+        self.workload_group.as_ref()
+    }
+
     pub fn disabled(&self) -> Option<&bool> {
         self.disabled.as_ref()
     }
@@ -330,6 +342,8 @@ impl UserOption {
             UserOptionItem::UnsetPasswordPolicy => self.password_policy = None,
             UserOptionItem::Disabled(v) => self.disabled = Some(*v),
             UserOptionItem::MustChangePassword(v) => self.must_change_password = Some(*v),
+            UserOptionItem::SetWorkloadGroup(v) => self.workload_group = Some(v.clone()),
+            UserOptionItem::UnsetWorkloadGroup => self.workload_group = None,
         }
     }
 }

src/meta/app/src/principal/user_info.rs

@@ -103,6 +103,7 @@ impl FromToProto for mt::principal::UserOption {
             .with_default_role(p.default_role)
             .with_network_policy(p.network_policy)
             .with_password_policy(p.password_policy)
+            .with_workload_group(p.workload_group)
             .with_disabled(p.disabled)
             .with_must_change_password(p.must_change_password))
     }
@@ -115,6 +116,7 @@ impl FromToProto for mt::principal::UserOption {
             default_role: self.default_role().cloned(),
             network_policy: self.network_policy().cloned(),
             password_policy: self.password_policy().cloned(),
+            workload_group: self.workload_group().cloned(),
             disabled: self.disabled().cloned(),
             must_change_password: self.must_change_password().cloned(),
         })

src/meta/proto-conv/src/user_from_to_protobuf_impl.rs

@@ -156,6 +156,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (124, "2025-04-01: Add: add headers in udf.proto/UDFServer"),
     (125, "2025-04-16: Add: add index_type in table.proto/TableIndex"),
     (126, "2025-05-11: Add: add iceberg Storage catalog type and add address field in Glue catalog"),
+    (127, "2025-05-18: Add: UserOption::workload_group"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)

src/meta/proto-conv/src/util.rs

@@ -118,3 +118,4 @@ mod v123_parquet_format_params;
 mod v124_udf_server_headers;
 mod v125_table_index;
 mod v126_iceberg_storage_catalog_option;
+mod v127_user_option_workload_group;

src/meta/proto-conv/tests/it/main.rs

@@ -0,0 +1,46 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use fastrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+// The message bytes are built from the output of `test_pb_from_to()`
+#[test]
+fn test_decode_v127_user_iption() -> anyhow::Result<()> {
+    let user_option_v127 = vec![
+        8, 1, 18, 5, 114, 111, 108, 101, 49, 26, 8, 109, 121, 112, 111, 108, 105, 99, 121, 58, 4,
+        109, 121, 119, 103, 160, 6, 127, 168, 6, 24,
+    ];
+
+    let want = || {
+        databend_common_meta_app::principal::UserOption::default()
+            .with_set_flag(databend_common_meta_app::principal::UserOptionFlag::TenantSetting)
+            .with_default_role(Some("role1".into()))
+            .with_network_policy(Some("mypolicy".to_string()))
+            .with_workload_group(Some("mywg".to_string()))
+    };
+
+    common::test_pb_from_to(func_name!(), want())?;
+    common::test_load_old(func_name!(), user_option_v127.as_slice(), 127, want())
+}

src/meta/proto-conv/tests/it/v127_user_option_workload_group.rs

@@ -128,6 +128,7 @@ message UserOption {
   optional string password_policy = 4;
   optional bool disabled = 5;
   optional bool must_change_password = 6;
+  optional string workload_group = 7;
 }
 
 message UserInfo {

src/meta/protos/proto/user.proto

@@ -266,6 +266,8 @@ pub enum UserOptionItem {
     SetPasswordPolicy(String),
     UnsetPasswordPolicy,
     MustChangePassword(bool),
+    SetWorkloadGroup(String),
+    UnsetWorkloadGroup,
 }
 
 impl Display for UserOptionItem {
@@ -277,6 +279,8 @@ impl Display for UserOptionItem {
             UserOptionItem::SetNetworkPolicy(v) => write!(f, "SET NETWORK POLICY = '{}'", v),
             UserOptionItem::UnsetNetworkPolicy => write!(f, "UNSET NETWORK POLICY"),
             UserOptionItem::SetPasswordPolicy(v) => write!(f, "SET PASSWORD POLICY = '{}'", v),
+            UserOptionItem::SetWorkloadGroup(v) => write!(f, "SET WORKLOAD GROUP = '{}'", v),
+            UserOptionItem::UnsetWorkloadGroup => write!(f, "UNSET WORKLOAD GROUP"),
             UserOptionItem::UnsetPasswordPolicy => write!(f, "UNSET PASSWORD POLICY"),
             UserOptionItem::Disabled(v) => write!(f, "DISABLED = {}", v),
             UserOptionItem::MustChangePassword(v) => write!(f, "MUST_CHANGE_PASSWORD = {}", v),

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

@@ -4630,6 +4630,18 @@ pub fn user_option(i: Input) -> IResult<UserOptionItem> {
         },
         |(_, _, val)| UserOptionItem::MustChangePassword(val),
     );
+    let set_workload_group = map(
+        rule! {
+            SET ~ WORKLOAD ~ ^GROUP ~ ^"=" ~ ^#literal_string
+        },
+        |(_, _, _, _, wg)| UserOptionItem::SetWorkloadGroup(wg),
+    );
+    let unset_workload_group = map(
+        rule! {
+            UNSET ~ WORKLOAD ~ ^GROUP
+        },
+        |(_, _, _)| UserOptionItem::UnsetWorkloadGroup,
+    );
 
     rule!(
         #tenant_setting
@@ -4641,6 +4653,8 @@ pub fn user_option(i: Input) -> IResult<UserOptionItem> {
         | #unset_password_policy
         | #set_disabled_option
         | #must_change_password
+        | #set_workload_group
+        | #unset_workload_group
     )(i)
 }
 

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

@@ -194,6 +194,9 @@ fn test_statement() {
         r#"ALTER USER u1 WITH DEFAULT_ROLE = role1, DISABLED=true, TENANTSETTING;"#,
         r#"ALTER USER u1 WITH SET NETWORK POLICY = 'policy1';"#,
         r#"ALTER USER u1 WITH UNSET NETWORK POLICY;"#,
+        r#"CREATE USER u1 IDENTIFIED BY '123456' WITH SET WORKLOAD GROUP='W1'"#,
+        r#"ALTER USER u1 WITH SET WORKLOAD GROUP = 'W1';"#,
+        r#"ALTER USER u1 WITH UNSET WORKLOAD GROUP;"#,
         r#"CREATE USER u1 IDENTIFIED BY '123456' WITH DEFAULT_ROLE='role123', TENANTSETTING"#,
         r#"CREATE USER u1 IDENTIFIED BY '123456' WITH SET NETWORK POLICY='policy1'"#,
         r#"CREATE USER u1 IDENTIFIED BY '123456' WITH disabled=true"#,