Skip to content

Commit 20e788e

Browse files
b41shb41sh
committed
Merge branch 'main' into feat-vector-hnsw
2 parents 4bb13c7 + b9e4e50 commit 20e788e

File tree

13 files changed

+101
-18
lines changed

13 files changed

+101
-18
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -783,6 +783,9 @@ impl Display for Statement {
783783
match option {
784784
SecondaryRolesOption::None => write!(f, "NONE")?,
785785
SecondaryRolesOption::All => write!(f, "ALL")?,
786+
SecondaryRolesOption::SpecifyRole(roles) => {
787+
write_comma_separated_list(f, roles)?
788+
}
786789
}
787790
}
788791
Statement::ShowCatalogs(stmt) => write!(f, "{stmt}")?,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,7 @@ impl Display for AccountMgrLevel {
271271
pub enum SecondaryRolesOption {
272272
None,
273273
All,
274+
SpecifyRole(Vec<String>),
274275
}
275276

276277
#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,15 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
435435
},
436436
);
437437

438+
let set_secondary_specify_roles = map(
439+
rule! {
440+
SET ~ SECONDARY ~ ROLES ~ #comma_separated_list1(role_name)
441+
},
442+
|(_, _, _, roles)| Statement::SetSecondaryRoles {
443+
option: SecondaryRolesOption::SpecifyRole(roles),
444+
},
445+
);
446+
438447
let set_stmt = alt((
439448
map(
440449
rule! {
@@ -2576,6 +2585,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
25762585
| #alter_udf : "`ALTER FUNCTION <udf_name> <udf_definition> [DESC = <description>]`"
25772586
| #set_role: "`SET [DEFAULT] ROLE <role>`"
25782587
| #set_secondary_roles: "`SET SECONDARY ROLES (ALL | NONE)`"
2588+
| #set_secondary_specify_roles: "`SET SECONDARY ROLES [role_name,...]`"
25792589
| #show_user_functions : "`SHOW USER FUNCTIONS [<show_limit>]`"
25802590
),
25812591
rule!(

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,9 @@ fn test_statement() {
8989
let cases = &[
9090
r#"show databases"#,
9191
r#"show drop databases"#,
92+
r#"SET SECONDARY ROLES ALL"#,
93+
r#"SET SECONDARY ROLES NONE"#,
94+
r#"SET SECONDARY ROLES role1, role2"#,
9295
r#"show drop databases like 'db%'"#,
9396
r#"show databases format TabSeparatedWithNamesAndTypes;"#,
9497
r#"show tables"#,
@@ -964,6 +967,7 @@ fn test_statement_error() {
964967

965968
let cases = &[
966969
r#"create table a.b (c integer not null 1, b float(10))"#,
970+
r#"SET SECONDARY ROLES"#,
967971
r#"create table a (c float(10))"#,
968972
r#"create table a (c varch)"#,
969973
r#"create table a (c tuple())"#,

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,16 @@ error:
1010
| while parsing `CREATE [OR REPLACE] TABLE [IF NOT EXISTS] [<database>.]<table> [<source>] [<table_options>]`
1111

1212

13+
---------- Input ----------
14+
SET SECONDARY ROLES
15+
---------- Output ---------
16+
error:
17+
--> SQL:1:20
18+
|
19+
1 | SET SECONDARY ROLES
20+
| ^ unexpected end of input, expecting `ALL`, `NONE`, <Ident>, <LiteralString>, or `IDENTIFIER`
21+
22+
1323
---------- Input ----------
1424
create table a (c float(10))
1525
---------- Output ---------

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,41 @@ ShowDropDatabases(
2525
)
2626

2727

28+
---------- Input ----------
29+
SET SECONDARY ROLES ALL
30+
---------- Output ---------
31+
SET SECONDARY ROLES ALL
32+
---------- AST ------------
33+
SetSecondaryRoles {
34+
option: All,
35+
}
36+
37+
38+
---------- Input ----------
39+
SET SECONDARY ROLES NONE
40+
---------- Output ---------
41+
SET SECONDARY ROLES NONE
42+
---------- AST ------------
43+
SetSecondaryRoles {
44+
option: None,
45+
}
46+
47+
48+
---------- Input ----------
49+
SET SECONDARY ROLES role1, role2
50+
---------- Output ---------
51+
SET SECONDARY ROLES role1, role2
52+
---------- AST ------------
53+
SetSecondaryRoles {
54+
option: SpecifyRole(
55+
[
56+
"role1",
57+
"role2",
58+
],
59+
),
60+
}
61+
62+
2863
---------- Input ----------
2964
show drop databases like 'db%'
3065
---------- Output ---------

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,9 @@ impl Interpreter for SetSecondaryRolesInterpreter {
5252

5353
let session = self.ctx.get_current_session();
5454

55-
let secondary_roles = match self.plan {
55+
let secondary_roles = match &self.plan {
5656
SetSecondaryRolesPlan::None => Some(vec![]),
57+
SetSecondaryRolesPlan::SpecifyRole(roles) => Some(roles.clone()),
5758
SetSecondaryRolesPlan::All => None,
5859
};
5960
session.set_secondary_roles_checked(secondary_roles).await?;

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

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -171,16 +171,9 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
171171
/// `secondary_roles` will be set to None, which is default.
172172
/// 2. NONE: only the current_role has effects on validate_privilge, `secondary_roles`
173173
/// will be set to Some([]).
174-
/// We can release this restriction in the future if any user needed. If an user want to set
175-
/// secondary_roles to be a subset of the roles granted to the current user, he can pass
176-
/// `secondary_roles` as Some([role1, role2, .. etc.]).
174+
/// 3. SpecifyRoles: Some([role1, role2, .. etc.]).
177175
#[async_backtrace::framed]
178176
async fn set_secondary_roles(&self, secondary_roles: Option<Vec<String>>) -> Result<()> {
179-
if secondary_roles.is_some() && !secondary_roles.as_ref().unwrap().is_empty() {
180-
return Err(ErrorCode::InvalidArgument(
181-
"only ALL or NONE is allowed on setting secondary roles",
182-
));
183-
}
184177
self.session_ctx.set_secondary_roles(secondary_roles);
185178
Ok(())
186179
}
@@ -206,7 +199,7 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
206199
async fn get_all_effective_roles(&self) -> Result<Vec<RoleInfo>> {
207200
let secondary_roles = self.session_ctx.get_secondary_roles();
208201

209-
// if secondary_roles is not set, return all the available roles
202+
// SET SECONDARY ROLES ALL, return all the available roles
210203
if secondary_roles.is_none() {
211204
let available_roles = self.get_all_available_roles().await?;
212205
return Ok(available_roles);

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1595,16 +1595,14 @@ async fn test_session_secondary_roles() -> Result<()> {
15951595

15961596
let route = create_endpoint()?;
15971597

1598-
// failed input: only ALL or NONE is allowed
15991598
let json = serde_json::json!({"sql": "SELECT 1", "session": {"secondary_roles": vec!["role1".to_string()]}});
16001599
let (_, result) = post_json_to_endpoint(&route, &json, HeaderMap::default()).await?;
1601-
assert!(result.error.is_some());
1602-
assert!(result
1603-
.error
1604-
.unwrap()
1605-
.message
1606-
.contains("only ALL or NONE is allowed on setting secondary roles"));
1607-
assert_eq!(result.state, ExecuteStateKind::Failed);
1600+
assert!(result.error.is_none());
1601+
assert_eq!(result.state, ExecuteStateKind::Succeeded);
1602+
assert_eq!(
1603+
result.session.unwrap().secondary_roles,
1604+
Some(vec!["role1".to_string()])
1605+
);
16081606

16091607
let json = serde_json::json!({"sql": "select 1", "session": {"role": "public", "secondary_roles": Vec::<String>::new()}});
16101608
let (_, result) = post_json_to_endpoint(&route, &json, HeaderMap::default()).await?;

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
use ahash::HashSet;
1516
use databend_common_ast::ast::SecondaryRolesOption;
17+
use databend_common_exception::ErrorCode;
1618
use databend_common_exception::Result;
1719

1820
use crate::plans::Plan;
@@ -44,6 +46,20 @@ impl Binder {
4446
let plan = match option {
4547
SecondaryRolesOption::None => SetSecondaryRolesPlan::None,
4648
SecondaryRolesOption::All => SetSecondaryRolesPlan::All,
49+
SecondaryRolesOption::SpecifyRole(roles) => {
50+
let all_roles = self
51+
.ctx
52+
.get_all_available_roles()
53+
.await?
54+
.iter()
55+
.map(|r| r.name.to_string())
56+
.collect::<HashSet<String>>();
57+
if roles.iter().all(|r| all_roles.contains(r)) {
58+
SetSecondaryRolesPlan::SpecifyRole(roles.clone())
59+
} else {
60+
return Err(ErrorCode::InvalidRole("object is invalid"));
61+
}
62+
}
4763
};
4864
Ok(Plan::SetSecondaryRoles(Box::new(plan)))
4965
}

0 commit comments

Comments
 (0)