Skip to content

Commit

Permalink
feat: object warehouse support rbac
Browse files Browse the repository at this point in the history
  • Loading branch information
TCeason committed Jan 15, 2025
1 parent 00f4bd2 commit a4bbacd
Show file tree
Hide file tree
Showing 24 changed files with 485 additions and 86 deletions.
12 changes: 11 additions & 1 deletion src/meta/app/src/principal/ownership_object.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@ pub enum OwnershipObject {
UDF {
name: String,
},

Warehouse {
uid: String,
},
}

impl OwnershipObject {
Expand Down Expand Up @@ -80,6 +84,7 @@ impl fmt::Display for OwnershipObject {
}
OwnershipObject::UDF { name } => write!(f, "UDF {name}"),
OwnershipObject::Stage { name } => write!(f, "STAGE {name}"),
OwnershipObject::Warehouse { uid } => write!(f, "Warehouse {uid}"),
}
}
}
Expand Down Expand Up @@ -119,6 +124,7 @@ impl KeyCodec for OwnershipObject {
}
OwnershipObject::Stage { name } => b.push_raw("stage-by-name").push_str(name),
OwnershipObject::UDF { name } => b.push_raw("udf-by-name").push_str(name),
OwnershipObject::Warehouse { uid } => b.push_raw("warehouse-by-uid").push_str(uid),
}
}

Expand Down Expand Up @@ -165,9 +171,13 @@ impl KeyCodec for OwnershipObject {
let name = p.next_str()?;
Ok(OwnershipObject::UDF { name })
}
"warehouse-by-uid" => {
let uid = p.next_str()?;
Ok(OwnershipObject::Warehouse { uid })
}
_ => Err(kvapi::KeyError::InvalidSegment {
i: p.index(),
expect: "database-by-id|database-by-catalog-id|table-by-id|table-by-catalog-id|stage-by-name|udf-by-name"
expect: "database-by-id|database-by-catalog-id|table-by-id|table-by-catalog-id|stage-by-name|udf-by-name|warehouse-by-uid"
.to_string(),
got: q.to_string(),
}),
Expand Down
16 changes: 16 additions & 0 deletions src/meta/app/src/principal/tenant_ownership_object_ident.rs
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,22 @@ mod tests {
let parsed = TenantOwnershipObjectIdent::from_str_key(&key).unwrap();
assert_eq!(role_grantee, parsed);
}

// warehouse
{
let role_grantee = TenantOwnershipObjectIdent::new_unchecked(
Tenant::new_literal("test"),
OwnershipObject::Warehouse {
uid: "n87s".to_string(),
},
);

let key = role_grantee.to_string_key();
assert_eq!("__fd_object_owners/test/warehouse-by-uid/n87s", key);

let parsed = TenantOwnershipObjectIdent::from_str_key(&key).unwrap();
assert_eq!(role_grantee, parsed);
}
}

#[test]
Expand Down
4 changes: 3 additions & 1 deletion src/meta/app/src/principal/user_grant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,9 @@ impl GrantObject {
GrantObject::Stage(_) => {
UserPrivilegeSet::available_privileges_on_stage(available_ownership)
}
GrantObject::Warehouse(_) => UserPrivilegeSet::available_privileges_on_warehouse(),
GrantObject::Warehouse(_) => {
UserPrivilegeSet::available_privileges_on_warehouse(available_ownership)
}
}
}

Expand Down
19 changes: 15 additions & 4 deletions src/meta/app/src/principal/user_privilege.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ pub enum UserPrivilegeType {
Write = 1 << 19,
// Privilege to Create database
CreateDatabase = 1 << 20,
// Privilege to Create warehouse
CreateWarehouse = 1 << 21,
// Discard Privilege Type
Set = 1 << 4,
}
Expand All @@ -102,6 +104,7 @@ const ALL_PRIVILEGES: BitFlags<UserPrivilegeType> = make_bitflags!(
| Read
| Write
| CreateDatabase
| CreateWarehouse
}
);

Expand Down Expand Up @@ -129,6 +132,7 @@ impl Display for UserPrivilegeType {
UserPrivilegeType::Read => "Read",
UserPrivilegeType::Write => "Write",
UserPrivilegeType::CreateDatabase => "CREATE DATABASE",
UserPrivilegeType::CreateWarehouse => "CREATE WAREHOUSE",
})
}
}
Expand Down Expand Up @@ -166,6 +170,9 @@ impl From<databend_common_ast::ast::UserPrivilegeType> for UserPrivilegeType {
databend_common_ast::ast::UserPrivilegeType::CreateDatabase => {
UserPrivilegeType::CreateDatabase
}
databend_common_ast::ast::UserPrivilegeType::CreateWarehouse => {
UserPrivilegeType::CreateWarehouse
}
databend_common_ast::ast::UserPrivilegeType::Set => UserPrivilegeType::Set,
}
}
Expand Down Expand Up @@ -199,8 +206,8 @@ impl UserPrivilegeSet {
let database_privs = Self::available_privileges_on_database(false);
let stage_privs_without_ownership = Self::available_privileges_on_stage(false);
let udf_privs_without_ownership = Self::available_privileges_on_udf(false);
let wh_privs_without_ownership = Self::available_privileges_on_warehouse();
let privs = make_bitflags!(UserPrivilegeType::{ Usage | Super | CreateUser | DropUser | CreateRole | DropRole | CreateDatabase | Grant | CreateDataMask });
let wh_privs_without_ownership = Self::available_privileges_on_warehouse(false);
let privs = make_bitflags!(UserPrivilegeType::{ Usage | Super | CreateUser | DropUser | CreateRole | DropRole | CreateDatabase | Grant | CreateDataMask | CreateWarehouse });
(database_privs.privileges
| privs
| stage_privs_without_ownership.privileges
Expand Down Expand Up @@ -234,8 +241,12 @@ impl UserPrivilegeSet {
}
}

pub fn available_privileges_on_warehouse() -> Self {
make_bitflags!(UserPrivilegeType::{ Usage }).into()
pub fn available_privileges_on_warehouse(available_ownership: bool) -> Self {
if available_ownership {
make_bitflags!(UserPrivilegeType::{ Usage | Ownership }).into()
} else {
make_bitflags!(UserPrivilegeType::{ Usage }).into()
}
}

pub fn available_privileges_on_udf(available_ownership: bool) -> Self {
Expand Down
8 changes: 8 additions & 0 deletions src/meta/proto-conv/src/ownership_from_to_protobuf_impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,9 @@ impl FromToProto for mt::principal::OwnershipObject {
pb::ownership_object::Object::Stage(pb::ownership_object::OwnershipStageObject {
stage,
}) => Ok(mt::principal::OwnershipObject::Stage { name: stage }),
pb::ownership_object::Object::Warehouse(
pb::ownership_object::OwnershipWarehouseObject { uid },
) => Ok(mt::principal::OwnershipObject::Warehouse { uid }),
}
}

Expand Down Expand Up @@ -123,6 +126,11 @@ impl FromToProto for mt::principal::OwnershipObject {
stage: name.clone(),
}),
),
mt::principal::OwnershipObject::Warehouse { uid } => {
Some(pb::ownership_object::Object::Warehouse(
pb::ownership_object::OwnershipWarehouseObject { uid: uid.clone() },
))
}
};
Ok(pb::OwnershipObject {
ver: VER,
Expand Down
1 change: 1 addition & 0 deletions src/meta/proto-conv/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
(114, "2024-12-12: Add: New DataType Interval."),
(115, "2024-12-16: Add: udf.proto: add UDAFScript and UDAFServer"),
(116, "2025-01-09: Add: MarkedDeletedIndexMeta"),
(117, "2025-01-13: Add: Add new UserPrivilege CreateWarehouse and new OwnershipObject::Warehouse"),
// 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`)
Expand Down
1 change: 1 addition & 0 deletions src/meta/proto-conv/tests/it/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,3 +114,4 @@ mod v113_warehouse_grantobject;
mod v114_interval_datatype;
mod v115_add_udaf_script;
mod v116_marked_deleted_index_meta;
mod v117_warehouse_ownershipobject;
116 changes: 116 additions & 0 deletions src/meta/proto-conv/tests/it/v117_warehouse_ownershipobject.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
// 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 std::collections::HashSet;

use chrono::DateTime;
use chrono::Utc;
use databend_common_meta_app as mt;
use databend_common_meta_app::principal::OwnershipObject;
use databend_common_meta_app::principal::UserGrantSet;
use databend_common_meta_app::principal::UserPrivilegeType;
use enumflags2::make_bitflags;
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. *
// *************************************************************
//

#[test]
fn test_decode_v117_grant_object() -> anyhow::Result<()> {
let role_info_v117 = vec![
10, 2, 114, 49, 18, 172, 1, 10, 21, 10, 8, 10, 0, 160, 6, 117, 168, 6, 24, 16, 128, 128,
128, 1, 160, 6, 117, 168, 6, 24, 10, 31, 10, 21, 18, 13, 10, 7, 100, 101, 102, 97, 117,
108, 116, 18, 2, 100, 98, 160, 6, 117, 168, 6, 24, 16, 2, 160, 6, 117, 168, 6, 24, 10, 35,
10, 25, 26, 17, 10, 7, 100, 101, 102, 97, 117, 108, 116, 18, 2, 100, 98, 26, 2, 116, 98,
160, 6, 117, 168, 6, 24, 16, 2, 160, 6, 117, 168, 6, 24, 10, 22, 10, 12, 34, 4, 10, 2, 102,
49, 160, 6, 117, 168, 6, 24, 16, 1, 160, 6, 117, 168, 6, 24, 10, 24, 10, 12, 42, 4, 10, 2,
115, 49, 160, 6, 117, 168, 6, 24, 16, 128, 128, 32, 160, 6, 117, 168, 6, 24, 10, 21, 10, 8,
10, 0, 160, 6, 117, 168, 6, 24, 16, 254, 255, 191, 1, 160, 6, 117, 168, 6, 24, 160, 6, 117,
168, 6, 24, 26, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 48,
48, 32, 85, 84, 67, 34, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48,
58, 48, 48, 32, 85, 84, 67, 160, 6, 117, 168, 6, 24,
];
let want = || mt::principal::RoleInfo {
name: "r1".to_string(),
grants: UserGrantSet::new(
vec![
mt::principal::GrantEntry::new(
mt::principal::GrantObject::Global,
make_bitflags!(UserPrivilegeType::{CreateWarehouse}),
),
mt::principal::GrantEntry::new(
mt::principal::GrantObject::Database("default".to_string(), "db".to_string()),
make_bitflags!(UserPrivilegeType::{Create}),
),
mt::principal::GrantEntry::new(
mt::principal::GrantObject::Table(
"default".to_string(),
"db".to_string(),
"tb".to_string(),
),
make_bitflags!(UserPrivilegeType::{Create}),
),
mt::principal::GrantEntry::new(
mt::principal::GrantObject::UDF("f1".to_string()),
make_bitflags!(UserPrivilegeType::{Usage}),
),
mt::principal::GrantEntry::new(
mt::principal::GrantObject::Stage("s1".to_string()),
make_bitflags!(UserPrivilegeType::{Write}),
),
// test new global privilege CreateWarehouse
mt::principal::GrantEntry::new(
mt::principal::GrantObject::Global,
make_bitflags!(UserPrivilegeType::{Create | Select | Insert | Update | Delete | Drop | Alter | Super | CreateUser | DropUser | CreateRole | DropRole | Grant | CreateStage | Set | CreateDataMask | Ownership | Read | Write | CreateWarehouse }),
),
],
HashSet::new(),
),
created_on: DateTime::<Utc>::default(),
update_on: DateTime::<Utc>::default(),
};

common::test_pb_from_to(func_name!(), want())?;
common::test_load_old(func_name!(), role_info_v117.as_slice(), 117, want())?;

Ok(())
}

#[test]
fn test_decode_v117_ownership() -> anyhow::Result<()> {
let ownership_info_v117 = vec![
10, 2, 114, 49, 18, 19, 42, 11, 10, 9, 97, 117, 110, 105, 113, 117, 101, 105, 100, 160, 6,
117, 168, 6, 24, 160, 6, 117, 168, 6, 24,
];

let want = || mt::principal::OwnershipInfo {
role: "r1".to_string(),
object: OwnershipObject::Warehouse {
uid: "auniqueid".to_string(),
},
};
common::test_pb_from_to(func_name!(), want())?;
common::test_load_old(func_name!(), ownership_info_v117.as_slice(), 117, want())?;

Ok(())
}
5 changes: 5 additions & 0 deletions src/meta/protos/proto/ownership.proto
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,15 @@ message OwnershipObject {
string stage = 1;
}

message OwnershipWarehouseObject {
string uid = 1;
}

oneof object {
OwnershipDatabaseObject database = 1;
OwnershipTableObject table = 2;
OwnershipUdfObject udf = 3;
OwnershipStageObject stage = 4;
OwnershipWarehouseObject warehouse = 5;
}
}
3 changes: 3 additions & 0 deletions src/query/ast/src/ast/statements/principal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,8 @@ pub enum UserPrivilegeType {
Write,
// Privilege to Create database
CreateDatabase,
// Privilege to Create warehouse
CreateWarehouse,
// Discard Privilege Type
Set,
}
Expand Down Expand Up @@ -178,6 +180,7 @@ impl Display for UserPrivilegeType {
UserPrivilegeType::Read => "Read",
UserPrivilegeType::Write => "Write",
UserPrivilegeType::CreateDatabase => "CREATE DATABASE",
UserPrivilegeType::CreateWarehouse => "CREATE WAREHOUSE",
})
}
}
Expand Down
2 changes: 2 additions & 0 deletions src/query/ast/src/ast/statements/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ pub enum GrantObjectName {
Table(Option<String>, String),
UDF(String),
Stage(String),
Warehouse(String),
}

impl Display for GrantObjectName {
Expand All @@ -164,6 +165,7 @@ impl Display for GrantObjectName {
}
GrantObjectName::UDF(udf) => write!(f, " UDF {udf}"),
GrantObjectName::Stage(stage) => write!(f, " STAGE {stage}"),
GrantObjectName::Warehouse(w) => write!(f, " WAREHOUSE {w}"),
}
}
}
Expand Down
14 changes: 13 additions & 1 deletion src/query/ast/src/parser/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3278,6 +3278,10 @@ pub fn priv_type(i: Input) -> IResult<UserPrivilegeType> {
UserPrivilegeType::CreateDatabase,
rule! { CREATE ~ DATABASE },
),
value(
UserPrivilegeType::CreateWarehouse,
rule! { CREATE ~ WAREHOUSE },
),
value(UserPrivilegeType::DropUser, rule! { DROP ~ USER }),
value(UserPrivilegeType::CreateRole, rule! { CREATE ~ ROLE }),
value(UserPrivilegeType::DropRole, rule! { DROP ~ ROLE }),
Expand Down Expand Up @@ -3337,11 +3341,16 @@ pub fn on_object_name(i: Input) -> IResult<GrantObjectName> {
GrantObjectName::UDF(udf_name.to_string())
});

let warehouse = map(rule! { WAREHOUSE ~ #ident}, |(_, w)| {
GrantObjectName::Warehouse(w.to_string())
});

rule!(
#database : "DATABASE <database>"
| #table : "TABLE <database>.<table>"
| #stage : "STAGE <stage_name>"
| #udf : "UDF <udf_name>"
| #warehouse : "WAREHOUSE <warehouse_name>"
)(i)
}

Expand Down Expand Up @@ -3436,10 +3445,12 @@ pub fn grant_ownership_level(i: Input) -> IResult<AccountMgrLevel> {
enum Object {
Stage,
Udf,
Warehouse,
}
let object = alt((
value(Object::Udf, rule! { UDF }),
value(Object::Stage, rule! { STAGE }),
value(Object::Warehouse, rule! { STAGE }),
));

// Object object_name
Expand All @@ -3448,13 +3459,14 @@ pub fn grant_ownership_level(i: Input) -> IResult<AccountMgrLevel> {
|(object, object_name)| match object {
Object::Stage => AccountMgrLevel::Stage(object_name.to_string()),
Object::Udf => AccountMgrLevel::UDF(object_name.to_string()),
Object::Warehouse => AccountMgrLevel::Warehouse(object_name.to_string()),
},
);

rule!(
#db : "<database>.*"
| #table : "<database>.<table>"
| #object : "STAGE | UDF <object_name>"
| #object : "STAGE | UDF | WAREHOUSE <object_name>"
)(i)
}

Expand Down
Loading

0 comments on commit a4bbacd

Please sign in to comment.