feat(query): Enhance SET SECONDARY ROLES to Specify Role Lists (#18337)

We propose extending the SET SECONDARY ROLES command to allow users to specify a comma-separated list of desired secondary roles.

Syntax Example:

```sql
SET SECONDARY ROLES role_name_1, role_name_2;
```

Author: TCeason

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

@@ -783,6 +783,9 @@ impl Display for Statement {
                 match option {
                     SecondaryRolesOption::None => write!(f, "NONE")?,
                     SecondaryRolesOption::All => write!(f, "ALL")?,
+                    SecondaryRolesOption::SpecifyRole(roles) => {
+                        write_comma_separated_list(f, roles)?
+                    }
                 }
             }
             Statement::ShowCatalogs(stmt) => write!(f, "{stmt}")?,

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

@@ -271,6 +271,7 @@ impl Display for AccountMgrLevel {
 pub enum SecondaryRolesOption {
     None,
     All,
+    SpecifyRole(Vec<String>),
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]

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

@@ -435,6 +435,15 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
     );
 
+    let set_secondary_specify_roles = map(
+        rule! {
+            SET ~ SECONDARY ~ ROLES ~ #comma_separated_list1(role_name)
+        },
+        |(_, _, _, roles)| Statement::SetSecondaryRoles {
+            option: SecondaryRolesOption::SpecifyRole(roles),
+        },
+    );
+
     let set_stmt = alt((
         map(
             rule! {
@@ -2576,6 +2585,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #alter_udf : "`ALTER FUNCTION <udf_name> <udf_definition> [DESC = <description>]`"
             | #set_role: "`SET [DEFAULT] ROLE <role>`"
             | #set_secondary_roles: "`SET SECONDARY ROLES (ALL | NONE)`"
+            | #set_secondary_specify_roles: "`SET SECONDARY ROLES [role_name,...]`"
             | #show_user_functions : "`SHOW USER FUNCTIONS [<show_limit>]`"
         ),
         rule!(

src/query/ast/tests/it/parser.rs

@@ -89,6 +89,9 @@ fn test_statement() {
     let cases = &[
         r#"show databases"#,
         r#"show drop databases"#,
+        r#"SET SECONDARY ROLES ALL"#,
+        r#"SET SECONDARY ROLES NONE"#,
+        r#"SET SECONDARY ROLES role1, role2"#,
         r#"show drop databases like 'db%'"#,
         r#"show databases format TabSeparatedWithNamesAndTypes;"#,
         r#"show tables"#,
@@ -964,6 +967,7 @@ fn test_statement_error() {
 
     let cases = &[
         r#"create table a.b (c integer not null 1, b float(10))"#,
+        r#"SET SECONDARY ROLES"#,
         r#"create table a (c float(10))"#,
         r#"create table a (c varch)"#,
         r#"create table a (c tuple())"#,

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

@@ -10,6 +10,16 @@ error:
   | while parsing `CREATE [OR REPLACE] TABLE [IF NOT EXISTS] [<database>.]<table> [<source>] [<table_options>]`
 
 
+---------- Input ----------
+SET SECONDARY ROLES
+---------- Output ---------
+error: 
+  --> SQL:1:20
+  |
+1 | SET SECONDARY ROLES
+  |                    ^ unexpected end of input, expecting `ALL`, `NONE`, <Ident>, <LiteralString>, or `IDENTIFIER`
+
+
 ---------- Input ----------
 create table a (c float(10))
 ---------- Output ---------

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

@@ -25,6 +25,41 @@ ShowDropDatabases(
 )
 
 
+---------- Input ----------
+SET SECONDARY ROLES ALL
+---------- Output ---------
+SET SECONDARY ROLES ALL
+---------- AST ------------
+SetSecondaryRoles {
+    option: All,
+}
+
+
+---------- Input ----------
+SET SECONDARY ROLES NONE
+---------- Output ---------
+SET SECONDARY ROLES NONE
+---------- AST ------------
+SetSecondaryRoles {
+    option: None,
+}
+
+
+---------- Input ----------
+SET SECONDARY ROLES role1, role2
+---------- Output ---------
+SET SECONDARY ROLES role1, role2
+---------- AST ------------
+SetSecondaryRoles {
+    option: SpecifyRole(
+        [
+            "role1",
+            "role2",
+        ],
+    ),
+}
+
+
 ---------- Input ----------
 show drop databases like 'db%'
 ---------- Output ---------

src/query/service/src/interpreters/interpreter_role_set_secondary.rs

@@ -52,8 +52,9 @@ impl Interpreter for SetSecondaryRolesInterpreter {
 
         let session = self.ctx.get_current_session();
 
-        let secondary_roles = match self.plan {
+        let secondary_roles = match &self.plan {
             SetSecondaryRolesPlan::None => Some(vec![]),
+            SetSecondaryRolesPlan::SpecifyRole(roles) => Some(roles.clone()),
             SetSecondaryRolesPlan::All => None,
         };
         session.set_secondary_roles_checked(secondary_roles).await?;

src/query/service/src/sessions/session_privilege_mgr.rs

@@ -171,16 +171,9 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
     ///    `secondary_roles` will be set to None, which is default.
     /// 2. NONE: only the current_role has effects on validate_privilge, `secondary_roles`
     ///    will be set to Some([]).
-    /// We can release this restriction in the future if any user needed. If an user want to set
-    /// secondary_roles to be a subset of the roles granted to the current user, he can pass
-    /// `secondary_roles` as Some([role1, role2, .. etc.]).
+    /// 3. SpecifyRoles: Some([role1, role2, .. etc.]).
     #[async_backtrace::framed]
     async fn set_secondary_roles(&self, secondary_roles: Option<Vec<String>>) -> Result<()> {
-        if secondary_roles.is_some() && !secondary_roles.as_ref().unwrap().is_empty() {
-            return Err(ErrorCode::InvalidArgument(
-                "only ALL or NONE is allowed on setting secondary roles",
-            ));
-        }
         self.session_ctx.set_secondary_roles(secondary_roles);
         Ok(())
     }
@@ -206,7 +199,7 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
     async fn get_all_effective_roles(&self) -> Result<Vec<RoleInfo>> {
         let secondary_roles = self.session_ctx.get_secondary_roles();
 
-        // if secondary_roles is not set, return all the available roles
+        // SET SECONDARY ROLES ALL, return all the available roles
         if secondary_roles.is_none() {
             let available_roles = self.get_all_available_roles().await?;
             return Ok(available_roles);

src/query/service/tests/it/servers/http/http_query_handlers.rs

@@ -1595,16 +1595,14 @@ async fn test_session_secondary_roles() -> Result<()> {
 
     let route = create_endpoint()?;
 
-    // failed input: only ALL or NONE is allowed
     let json = serde_json::json!({"sql":  "SELECT 1", "session": {"secondary_roles": vec!["role1".to_string()]}});
     let (_, result) = post_json_to_endpoint(&route, &json, HeaderMap::default()).await?;
-    assert!(result.error.is_some());
-    assert!(result
-        .error
-        .unwrap()
-        .message
-        .contains("only ALL or NONE is allowed on setting secondary roles"));
-    assert_eq!(result.state, ExecuteStateKind::Failed);
+    assert!(result.error.is_none());
+    assert_eq!(result.state, ExecuteStateKind::Succeeded);
+    assert_eq!(
+        result.session.unwrap().secondary_roles,
+        Some(vec!["role1".to_string()])
+    );
 
     let json = serde_json::json!({"sql":  "select 1", "session": {"role": "public", "secondary_roles": Vec::<String>::new()}});
     let (_, result) = post_json_to_endpoint(&route, &json, HeaderMap::default()).await?;

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

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use ahash::HashSet;
 use databend_common_ast::ast::SecondaryRolesOption;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 
 use crate::plans::Plan;
@@ -44,6 +46,20 @@ impl Binder {
         let plan = match option {
             SecondaryRolesOption::None => SetSecondaryRolesPlan::None,
             SecondaryRolesOption::All => SetSecondaryRolesPlan::All,
+            SecondaryRolesOption::SpecifyRole(roles) => {
+                let all_roles = self
+                    .ctx
+                    .get_all_available_roles()
+                    .await?
+                    .iter()
+                    .map(|r| r.name.to_string())
+                    .collect::<HashSet<String>>();
+                if roles.iter().all(|r| all_roles.contains(r)) {
+                    SetSecondaryRolesPlan::SpecifyRole(roles.clone())
+                } else {
+                    return Err(ErrorCode::InvalidRole("object is invalid"));
+                }
+            }
         };
         Ok(Plan::SetSecondaryRoles(Box::new(plan)))
     }