Skip to content

Commit 341bb95

Browse files
TCeasonTCeason
committed
refactor(query): clean up role management API surface
 - Merge 6 per-type list_*_ownerships methods into a single list_ownerships_by_type(OwnershipObject) to eliminate duplicated logic across RoleApi trait and RoleMgr implementation. - Simplify exists_role to call management layer get_role() directly, avoiding the unnecessary ErrorCode roundtrip through UnknownRole → or_unknown_role → Option. - Change return type of grant_privileges_to_role, revoke_privileges_from_role, grant_role_to_role, revoke_role_from_role from Result<Option<u64>> to Result<()>. The Option was always Some on success — no caller ever inspected the u64 or checked for None.
1 parent 7096c77 commit 341bb95

File tree

4 files changed

+91
-156
lines changed

4 files changed

+91
-156
lines changed

‎src/query/management/src/role/role_api.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,13 @@ pub trait RoleApi: Sync + Send {
4040

4141
async fn list_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>, MetaError>;
4242

43-
async fn list_udf_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
44-
async fn list_stage_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
45-
async fn list_seq_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
46-
async fn list_procedure_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
47-
async fn list_connection_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
48-
async fn list_warehouse_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError>;
43+
/// List ownerships filtered by type.
44+
///
45+
/// The `obj` parameter is used as a dummy to determine the ownership type prefix.
46+
async fn list_ownerships_by_type(
47+
&self,
48+
obj: OwnershipObject,
49+
) -> Result<Vec<OwnershipInfo>, MetaError>;
4950

5051
/// General role update.
5152
///

‎src/query/management/src/role/role_mgr.rs‎

Lines changed: 6 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -301,102 +301,18 @@ impl RoleApi for RoleMgr {
301301

302302
#[async_backtrace::framed]
303303
#[fastrace::trace]
304-
async fn list_udf_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
305-
let obj = OwnershipObject::UDF {
306-
name: "foo".to_string(),
307-
};
308-
309-
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
310-
let dir_name = DirName::new(ident);
311-
let values = self
312-
.kv_api
313-
.list_pb_values(ListOptions::unlimited(&dir_name))
314-
.await?;
315-
let udfs = values.try_collect().await?;
316-
Ok(udfs)
317-
}
318-
319-
#[async_backtrace::framed]
320-
#[fastrace::trace]
321-
async fn list_stage_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
322-
let obj = OwnershipObject::Stage {
323-
name: "s1".to_string(),
324-
};
325-
326-
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
327-
let dir_name = DirName::new(ident);
328-
let values = self
329-
.kv_api
330-
.list_pb_values(ListOptions::unlimited(&dir_name))
331-
.await?;
332-
let stages = values.try_collect().await?;
333-
Ok(stages)
334-
}
335-
336-
#[async_backtrace::framed]
337-
#[fastrace::trace]
338-
async fn list_seq_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
339-
let obj = OwnershipObject::Sequence {
340-
name: "seq1".to_string(),
341-
};
342-
343-
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
344-
let dir_name = DirName::new(ident);
345-
let values = self
346-
.kv_api
347-
.list_pb_values(ListOptions::unlimited(&dir_name))
348-
.await?;
349-
let seqs = values.try_collect().await?;
350-
Ok(seqs)
351-
}
352-
353-
#[async_backtrace::framed]
354-
#[fastrace::trace]
355-
async fn list_procedure_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
356-
let obj = OwnershipObject::Procedure { procedure_id: 1 };
357-
358-
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
359-
let dir_name = DirName::new(ident);
360-
let values = self
361-
.kv_api
362-
.list_pb_values(ListOptions::unlimited(&dir_name))
363-
.await?;
364-
let seqs = values.try_collect().await?;
365-
Ok(seqs)
366-
}
367-
368-
#[async_backtrace::framed]
369-
#[fastrace::trace]
370-
async fn list_connection_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
371-
let obj = OwnershipObject::Connection {
372-
name: "con".to_string(),
373-
};
374-
375-
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
376-
let dir_name = DirName::new(ident);
377-
let values = self
378-
.kv_api
379-
.list_pb_values(ListOptions::unlimited(&dir_name))
380-
.await?;
381-
let conns = values.try_collect().await?;
382-
Ok(conns)
383-
}
384-
385-
#[async_backtrace::framed]
386-
#[fastrace::trace]
387-
async fn list_warehouse_ownerships(&self) -> Result<Vec<OwnershipInfo>, MetaError> {
388-
let obj = OwnershipObject::Warehouse {
389-
id: "w".to_string(),
390-
};
391-
304+
async fn list_ownerships_by_type(
305+
&self,
306+
obj: OwnershipObject,
307+
) -> Result<Vec<OwnershipInfo>, MetaError> {
392308
let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
393309
let dir_name = DirName::new(ident);
394310
let values = self
395311
.kv_api
396312
.list_pb_values(ListOptions::unlimited(&dir_name))
397313
.await?;
398-
let ws = values.try_collect().await?;
399-
Ok(ws)
314+
let items = values.try_collect().await?;
315+
Ok(items)
400316
}
401317

402318
/// General role update.

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

Lines changed: 60 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -387,53 +387,68 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
387387
let roles = self.get_all_effective_roles().await?;
388388
let roles_name: Vec<String> = roles.iter().map(|role| role.name.to_string()).collect();
389389

390-
let ownership_infos =
391-
if roles_name.contains(&BUILTIN_ROLE_ACCOUNT_ADMIN.to_string()) || ignore_ownership {
392-
vec![]
393-
} else {
394-
let user_api = UserApiProvider::instance();
395-
let ownerships = match object {
396-
Object::All => user_api
397-
.role_api(&self.session_ctx.get_current_tenant())
398-
.list_ownerships()
399-
.await
400-
.map_err(meta_service_error)?
401-
.into_iter()
402-
.map(|item| item.data)
403-
.collect(),
404-
Object::UDF => user_api
405-
.role_api(&self.session_ctx.get_current_tenant())
406-
.list_udf_ownerships()
407-
.await?,
408-
Object::Stage => user_api
409-
.role_api(&self.session_ctx.get_current_tenant())
410-
.list_stage_ownerships()
411-
.await?,
412-
Object::Sequence => user_api
413-
.role_api(&self.session_ctx.get_current_tenant())
414-
.list_seq_ownerships()
415-
.await?,
416-
Object::Connection => user_api
417-
.role_api(&self.session_ctx.get_current_tenant())
418-
.list_connection_ownerships()
419-
.await?,
420-
Object::Warehouse => user_api
421-
.role_api(&self.session_ctx.get_current_tenant())
422-
.list_warehouse_ownerships()
423-
.await?,
424-
Object::Procedure => user_api
425-
.role_api(&self.session_ctx.get_current_tenant())
426-
.list_procedure_ownerships()
427-
.await?,
428-
};
429-
let mut ownership_infos = vec![];
430-
for ownership in ownerships {
431-
if roles_name.contains(&ownership.role) {
432-
ownership_infos.push(ownership);
433-
}
390+
let ownership_infos = if roles_name.contains(&BUILTIN_ROLE_ACCOUNT_ADMIN.to_string())
391+
|| ignore_ownership
392+
{
393+
vec![]
394+
} else {
395+
let user_api = UserApiProvider::instance();
396+
let role_api = user_api.role_api(&self.session_ctx.get_current_tenant());
397+
let ownerships = match object {
398+
Object::All => role_api
399+
.list_ownerships()
400+
.await
401+
.map_err(meta_service_error)?
402+
.into_iter()
403+
.map(|item| item.data)
404+
.collect(),
405+
Object::UDF => {
406+
role_api
407+
.list_ownerships_by_type(OwnershipObject::UDF {
408+
name: String::new(),
409+
})
410+
.await?
411+
}
412+
Object::Stage => {
413+
role_api
414+
.list_ownerships_by_type(OwnershipObject::Stage {
415+
name: String::new(),
416+
})
417+
.await?
418+
}
419+
Object::Sequence => {
420+
role_api
421+
.list_ownerships_by_type(OwnershipObject::Sequence {
422+
name: String::new(),
423+
})
424+
.await?
425+
}
426+
Object::Connection => {
427+
role_api
428+
.list_ownerships_by_type(OwnershipObject::Connection {
429+
name: String::new(),
430+
})
431+
.await?
432+
}
433+
Object::Warehouse => {
434+
role_api
435+
.list_ownerships_by_type(OwnershipObject::Warehouse { id: String::new() })
436+
.await?
437+
}
438+
Object::Procedure => {
439+
role_api
440+
.list_ownerships_by_type(OwnershipObject::Procedure { procedure_id: 0 })
441+
.await?
434442
}
435-
ownership_infos
436443
};
444+
let mut ownership_infos = vec![];
445+
for ownership in ownerships {
446+
if roles_name.contains(&ownership.role) {
447+
ownership_infos.push(ownership);
448+
}
449+
}
450+
ownership_infos
451+
};
437452

438453
Ok(GrantObjectVisibilityChecker::new(
439454
&self.get_current_user()?,

‎src/query/users/src/role_mgr.rs‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
use std::collections::HashMap;
1616

1717
use databend_common_exception::ErrorCode;
18-
use databend_common_exception::ErrorCodeResultExt;
1918
use databend_common_exception::Result;
2019
use databend_common_management::RoleApi;
2120
use databend_common_meta_app::principal::BUILTIN_ROLE_ACCOUNT_ADMIN;
@@ -108,10 +107,14 @@ impl UserApiProvider {
108107

109108
#[async_backtrace::framed]
110109
pub async fn exists_role(&self, tenant: &Tenant, role: String) -> Result<bool> {
110+
if self.builtin_roles().contains_key(&role) {
111+
return Ok(true);
112+
}
111113
Ok(self
112-
.get_role(tenant, role)
114+
.role_api(tenant)
115+
.get_role(&role)
113116
.await
114-
.or_unknown_role()?
117+
.map_err(meta_service_error)?
115118
.is_some())
116119
}
117120

@@ -230,17 +233,17 @@ impl UserApiProvider {
230233
role: &str,
231234
object: GrantObject,
232235
privileges: UserPrivilegeSet,
233-
) -> Result<Option<u64>> {
236+
) -> Result<()> {
234237
let client = self.role_api(tenant);
235-
let res = client
238+
client
236239
.update_role_with(role, |ri: &mut RoleInfo| {
237240
ri.update_role_time();
238241
ri.grants.grant_privileges(&object, privileges)
239242
})
240243
.await
241244
.map_err(meta_service_error)?
242245
.map_err(|e| ErrorCode::from(e).add_message_back("(while set role privileges)"))?;
243-
Ok(Some(res))
246+
Ok(())
244247
}
245248

246249
#[async_backtrace::framed]
@@ -250,17 +253,17 @@ impl UserApiProvider {
250253
role: &str,
251254
object: GrantObject,
252255
privileges: UserPrivilegeSet,
253-
) -> Result<Option<u64>> {
256+
) -> Result<()> {
254257
let client = self.role_api(tenant);
255-
let res = client
258+
client
256259
.update_role_with(role, |ri: &mut RoleInfo| {
257260
ri.update_role_time();
258261
ri.grants.revoke_privileges(&object, privileges)
259262
})
260263
.await
261264
.map_err(meta_service_error)?
262265
.map_err(|e| ErrorCode::from(e).add_message_back("(while revoke role privileges)"))?;
263-
Ok(Some(res))
266+
Ok(())
264267
}
265268

266269
// the grant_role can not have cycle with target_role.
@@ -270,7 +273,7 @@ impl UserApiProvider {
270273
tenant: &Tenant,
271274
target_role: &String,
272275
grant_role: String,
273-
) -> Result<Option<u64>> {
276+
) -> Result<()> {
274277
let related_roles = self
275278
.find_related_roles(tenant, &[grant_role.clone()])
276279
.await?;
@@ -283,15 +286,15 @@ impl UserApiProvider {
283286
}
284287

285288
let client = self.role_api(tenant);
286-
let res = client
289+
client
287290
.update_role_with(target_role, |ri: &mut RoleInfo| {
288291
ri.update_role_time();
289292
ri.grants.grant_role(grant_role)
290293
})
291294
.await
292295
.map_err(meta_service_error)?
293296
.map_err(|e| ErrorCode::from(e).add_message_back("(while grant role to role)"))?;
294-
Ok(Some(res))
297+
Ok(())
295298
}
296299

297300
#[async_backtrace::framed]
@@ -300,17 +303,17 @@ impl UserApiProvider {
300303
tenant: &Tenant,
301304
role: &str,
302305
revoke_role: &String,
303-
) -> Result<Option<u64>> {
306+
) -> Result<()> {
304307
let client = self.role_api(tenant);
305-
let res = client
308+
client
306309
.update_role_with(role, |ri: &mut RoleInfo| {
307310
ri.update_role_time();
308311
ri.grants.revoke_role(revoke_role)
309312
})
310313
.await
311314
.map_err(meta_service_error)?
312315
.map_err(|e| ErrorCode::from(e).add_message_back("(while revoke role from role)"))?;
313-
Ok(Some(res))
316+
Ok(())
314317
}
315318

316319
// Drop a role by name

0 commit comments

Comments
 (0)