Skip to content

Commit 9093e6e

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 9093e6e

File tree

4 files changed

+65
-137
lines changed

4 files changed

+65
-137
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: 34 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -392,39 +392,47 @@ impl SessionPrivilegeManager for SessionPrivilegeManagerImpl<'_> {
392392
vec![]
393393
} else {
394394
let user_api = UserApiProvider::instance();
395+
let role_api = user_api.role_api(&self.session_ctx.get_current_tenant());
395396
let ownerships = match object {
396-
Object::All => user_api
397-
.role_api(&self.session_ctx.get_current_tenant())
397+
Object::All => role_api
398398
.list_ownerships()
399399
.await
400400
.map_err(meta_service_error)?
401401
.into_iter()
402402
.map(|item| item.data)
403403
.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?,
404+
Object::UDF => role_api
405+
.list_ownerships_by_type(OwnershipObject::UDF {
406+
name: String::new(),
407+
})
408+
.await
409+
.map_err(meta_service_error)?,
410+
Object::Stage => role_api
411+
.list_ownerships_by_type(OwnershipObject::Stage {
412+
name: String::new(),
413+
})
414+
.await
415+
.map_err(meta_service_error)?,
416+
Object::Sequence => role_api
417+
.list_ownerships_by_type(OwnershipObject::Sequence {
418+
name: String::new(),
419+
})
420+
.await
421+
.map_err(meta_service_error)?,
422+
Object::Connection => role_api
423+
.list_ownerships_by_type(OwnershipObject::Connection {
424+
name: String::new(),
425+
})
426+
.await
427+
.map_err(meta_service_error)?,
428+
Object::Warehouse => role_api
429+
.list_ownerships_by_type(OwnershipObject::Warehouse { id: String::new() })
430+
.await
431+
.map_err(meta_service_error)?,
432+
Object::Procedure => role_api
433+
.list_ownerships_by_type(OwnershipObject::Procedure { procedure_id: 0 })
434+
.await
435+
.map_err(meta_service_error)?,
428436
};
429437
let mut ownership_infos = vec![];
430438
for ownership in ownerships {

‎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)