Skip to content

chore: minor refine on RoleMgr #17683

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 7, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/query/management/src/role/role_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ pub trait RoleApi: Sync + Send {

async fn get_raw_meta_roles(&self) -> Result<ListKVReply>;

async fn get_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>>;
async fn list_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>>;

/// General role update.
///
Expand Down
11 changes: 6 additions & 5 deletions src/query/management/src/role/role_mgr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ static TXN_MAX_RETRY_TIMES: u32 = 60;

static BUILTIN_ROLE_ACCOUNT_ADMIN: &str = "account_admin";

#[derive(Clone)]
pub struct RoleMgr {
kv_api: Arc<dyn kvapi::KVApi<Error = MetaError> + Send + Sync>,
tenant: Tenant,
Expand Down Expand Up @@ -210,7 +211,7 @@ impl RoleApi for RoleMgr {

#[async_backtrace::framed]
#[fastrace::trace]
async fn get_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>, ErrorCode> {
async fn list_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>, ErrorCode> {
let object_owner_prefix = self.ownership_object_prefix();
let values = self
.kv_api
Expand All @@ -228,7 +229,7 @@ impl RoleApi for RoleMgr {
// After rollback the old version, deserialize will return Err Ownership can not be none.
// But get ownerships should try to ensure success because in this version.
Err(err) => error!(
"deserialize key {} Got err {} while (get_ownerships)",
"deserialize key {} Got err {} while (list_ownerships)",
&key, err
),
}
Expand Down Expand Up @@ -269,11 +270,11 @@ impl RoleApi for RoleMgr {

/// Only drop role will call transfer.
///
/// If a role is dropped, but the owner object is exists,
/// If a role is dropped, but the owner object is existing,
///
/// The owner role need to update to account_admin.
///
/// get_ownerships use prefix_list_kv that will generate once meta call
/// list_ownerships use prefix_list_kv that will generate once meta call
///
/// According to Txn reduce meta call. If role own n objects, will generate once meta call.
#[async_backtrace::framed]
Expand All @@ -287,7 +288,7 @@ impl RoleApi for RoleMgr {
trials.next().unwrap().map_err(AppError::from)?.await;
let mut if_then = vec![];
let mut condition = vec![];
let seq_owns = self.get_ownerships().await.map_err(|e| {
let seq_owns = self.list_ownerships().await.map_err(|e| {
e.add_message_back("(while in transfer_ownership_to_admin get ownerships).")
})?;
let mut need_transfer = false;
Expand Down
6 changes: 3 additions & 3 deletions src/query/service/src/interpreters/access/privilege_access.rs
Original file line number Diff line number Diff line change
Expand Up @@ -764,7 +764,7 @@ impl AccessChecker for PrivilegeAccess {
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&tenant)
.get_ownerships()
.list_ownerships()
.await?;
let roles = self.ctx.get_all_effective_roles().await?;
let roles_name: Vec<String> = roles.iter().map(|role| role.name.to_string()).collect();
Expand All @@ -782,7 +782,7 @@ impl AccessChecker for PrivilegeAccess {
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&tenant)
.get_ownerships()
.list_ownerships()
.await?;
let roles = self.ctx.get_all_effective_roles().await?;
let roles_name: Vec<String> = roles.iter().map(|role| role.name.to_string()).collect();
Expand Down Expand Up @@ -937,7 +937,7 @@ impl AccessChecker for PrivilegeAccess {
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&tenant)
.get_ownerships()
.list_ownerships()
.await?;
let roles = self.ctx.get_all_effective_roles().await?;
let roles_name: Vec<String> = roles.iter().map(|role| role.name.to_string()).collect();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,8 @@ impl Interpreter for CreateWarehouseInterpreter {
.await?;

if let WarehouseInfo::SystemManaged(sw) = warehouse {
let role_api = UserApiProvider::instance().role_api(&tenant);
if let Some(current_role) = self.ctx.get_current_role() {
let role_api = UserApiProvider::instance().role_api(&tenant);
role_api
.grant_ownership(
&OwnershipObject::Warehouse { id: sw.role_id },
Expand Down Expand Up @@ -114,8 +114,8 @@ impl Interpreter for CreateWarehouseInterpreter {
.await?;

if let WarehouseInfo::SystemManaged(sw) = warehouse {
let role_api = UserApiProvider::instance().role_api(&tenant);
if let Some(current_role) = self.ctx.get_current_role() {
let role_api = UserApiProvider::instance().role_api(&tenant);
role_api
.grant_ownership(
&OwnershipObject::Warehouse { id: sw.role_id },
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,8 @@ impl Interpreter for CreateDatabaseInterpreter {

// Grant ownership as the current role. The above create_db_req.meta.owner could be removed in
// the future.
let role_api = UserApiProvider::instance().role_api(&tenant);
if let Some(current_role) = self.ctx.get_current_role() {
let role_api = UserApiProvider::instance().role_api(&tenant);
role_api
.grant_ownership(
&OwnershipObject::Database {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,8 +72,8 @@ impl Interpreter for DropWarehouseInterpreter {
);

if let WarehouseInfo::SystemManaged(sw) = warehouse {
let role_api = UserApiProvider::instance().role_api(&tenant);
if let Some(current_role) = self.ctx.get_current_role() {
let role_api = UserApiProvider::instance().role_api(&tenant);
role_api
.grant_ownership(
&OwnershipObject::Warehouse { id: sw.role_id },
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,8 @@ impl Interpreter for CreateUserStageInterpreter {

// Grant ownership as the current role
let tenant = self.ctx.get_tenant();
let role_api = UserApiProvider::instance().role_api(&tenant);
if let Some(current_role) = self.ctx.get_current_role() {
let role_api = UserApiProvider::instance().role_api(&tenant);
role_api
.grant_ownership(
&OwnershipObject::Stage {
Expand Down
2 changes: 1 addition & 1 deletion src/query/service/src/sessions/session_privilege_mgr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -397,7 +397,7 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&self.session_ctx.get_current_tenant())
.get_ownerships()
.list_ownerships()
.await?;
let mut ownership_objects = vec![];
for ownership in ownerships {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -495,7 +495,7 @@ async fn show_account_grants(
// 2. not expand roles
// So no need to get ownerships.
if !roles.is_empty() {
let ownerships = user_api.role_api(&tenant).get_ownerships().await?;
let ownerships = user_api.role_api(&tenant).list_ownerships().await?;
for ownership in ownerships {
if roles.contains(&ownership.data.role) {
match ownership.data.object {
Expand Down Expand Up @@ -807,7 +807,7 @@ async fn show_object_grant(
}
}

let ownerships = user_api.role_api(&tenant).get_ownerships().await?;
let ownerships = user_api.role_api(&tenant).list_ownerships().await?;
for ownership in ownerships {
if ownership.data.object == owner_object {
privileges.push("OWNERSHIP".to_string());
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/streams_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ impl<const T: bool> AsyncSystemTable for StreamsTable<T> {
.collect::<Vec<_>>();

let ownership = if T {
user_api.get_ownerships(&tenant).await.unwrap_or_default()
user_api.list_ownerships(&tenant).await.unwrap_or_default()
} else {
HashMap::new()
};
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/tables_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -589,7 +589,7 @@ where TablesTable<WITH_HISTORY, WITHOUT_VIEW>: HistoryAware
dbs.clear();

let ownership = if get_ownership && default_catalog {
user_api.get_ownerships(&tenant).await.unwrap_or_default()
user_api.list_ownerships(&tenant).await.unwrap_or_default()
} else {
HashMap::new()
};
Expand Down
6 changes: 3 additions & 3 deletions src/query/users/src/role_mgr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,13 +84,13 @@ impl UserApiProvider {
}

#[async_backtrace::framed]
pub async fn get_ownerships(
pub async fn list_ownerships(
&self,
tenant: &Tenant,
) -> Result<HashMap<OwnershipObject, String>> {
let seq_owns = self
.role_api(tenant)
.get_ownerships()
.list_ownerships()
.await
.map_err(|e| e.add_message_back("(while get ownerships)."))?;

Expand Down Expand Up @@ -172,7 +172,7 @@ impl UserApiProvider {
if let Some(owner) = ownership {
// if object has ownership, but the owner role is not exists, set owner role to ACCOUNT_ADMIN,
// only account_admin can access this object.
// Note: get_ownerships no need to do this check.
// Note: list_ownerships no need to do this check.
// Because this can cause system.table queries to slow down
// The intention is that the account admin can grant ownership.
// So system.tables will display dropped role. It's by design.
Expand Down
8 changes: 5 additions & 3 deletions src/query/users/src/user_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ use databend_common_management::PasswordPolicyMgr;
use databend_common_management::ProcedureMgr;
use databend_common_management::QuotaApi;
use databend_common_management::QuotaMgr;
use databend_common_management::RoleApi;
use databend_common_management::RoleMgr;
use databend_common_management::SettingMgr;
use databend_common_management::StageApi;
Expand All @@ -45,6 +44,7 @@ use databend_common_meta_store::MetaStore;
use databend_common_meta_store::MetaStoreProvider;
use databend_common_meta_types::MatchSeq;
use databend_common_meta_types::MetaError;
use log::debug;

use crate::builtin::BuiltIn;
use crate::BUILTIN_ROLE_PUBLIC;
Expand All @@ -65,6 +65,7 @@ impl UserApiProvider {
) -> Result<()> {
GlobalInstance::set(Self::try_create(conf, builtin, tenant).await?);
let user_mgr = UserApiProvider::instance();

if let Some(q) = quota {
let i = user_mgr.tenant_quota_api(tenant);
let res = i.get_quota(MatchSeq::GE(0)).await?;
Expand Down Expand Up @@ -125,13 +126,14 @@ impl UserApiProvider {
Arc::new(user_mgr)
}

pub fn role_api(&self, tenant: &Tenant) -> Arc<impl RoleApi> {
pub fn role_api(&self, tenant: &Tenant) -> RoleMgr {
let role_mgr = RoleMgr::create(
self.client.clone(),
tenant,
GlobalConfig::instance().query.upgrade_to_pb,
);
Arc::new(role_mgr)
debug!("RoleMgr created");
role_mgr
}

pub fn stage_api(&self, tenant: &Tenant) -> Arc<dyn StageApi> {
Expand Down