Skip to content

Commit

Permalink
optimize: optimize query system.tables when query single table
Browse files Browse the repository at this point in the history
  • Loading branch information
TCeason committed Nov 19, 2024
1 parent aac6e09 commit 58caec1
Show file tree
Hide file tree
Showing 23 changed files with 341 additions and 221 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions src/meta/app/src/principal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ pub use user_defined_function::UserDefinedFunction;
pub use user_grant::GrantEntry;
pub use user_grant::GrantObject;
pub use user_grant::UserGrantSet;
pub use user_grant::SYSTEM_TABLES_ALLOW_LIST;
pub use user_identity::UserIdentity;
pub use user_info::UserInfo;
pub use user_info::UserOption;
Expand Down
27 changes: 27 additions & 0 deletions src/meta/app/src/principal/user_grant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,33 @@ use enumflags2::BitFlags;
use crate::principal::UserPrivilegeSet;
use crate::principal::UserPrivilegeType;

// some statements like `SELECT 1`, `SHOW USERS`, `SHOW ROLES`, `SHOW TABLES` will be
// rewritten to the queries on the system tables, we need to skip the privilege check on
// these tables.
pub const SYSTEM_TABLES_ALLOW_LIST: [&str; 21] = [
"catalogs",
"columns",
"databases",
"databases_with_history",
"dictionaries",
"tables",
"views",
"tables_with_history",
"views_with_history",
"password_policies",
"streams",
"streams_terse",
"virtual_columns",
"users",
"roles",
"stages",
"one",
"processes",
"user_functions",
"functions",
"indexes",
];

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Hash)]
pub enum GrantObject {
Global,
Expand Down
5 changes: 4 additions & 1 deletion src/query/catalog/src/table_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,10 @@ pub trait TableContext: Send + Sync {
check_current_role_only: bool,
) -> Result<()>;
async fn get_available_roles(&self) -> Result<Vec<RoleInfo>>;
async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker>;
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker>;
fn get_fuse_version(&self) -> String;
fn get_format_settings(&self) -> Result<FormatSettings>;
fn get_tenant(&self) -> Tenant;
Expand Down
1 change: 1 addition & 0 deletions src/query/management/src/role/role_mgr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ impl RoleApi for RoleMgr {
#[async_backtrace::framed]
#[fastrace::trace]
async fn get_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>, ErrorCode> {
println!("will call get_ownerships");
let object_owner_prefix = self.ownership_object_prefix();
let values = self
.kv_api
Expand Down
28 changes: 1 addition & 27 deletions src/query/service/src/interpreters/access/privilege_access.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use databend_common_meta_app::principal::StageType;
use databend_common_meta_app::principal::UserGrantSet;
use databend_common_meta_app::principal::UserPrivilegeSet;
use databend_common_meta_app::principal::UserPrivilegeType;
use databend_common_meta_app::principal::SYSTEM_TABLES_ALLOW_LIST;
use databend_common_meta_app::tenant::Tenant;
use databend_common_meta_types::seq_value::SeqV;
use databend_common_sql::binder::MutationType;
Expand Down Expand Up @@ -58,33 +59,6 @@ enum ObjectId {
Table(u64, u64),
}

// some statements like `SELECT 1`, `SHOW USERS`, `SHOW ROLES`, `SHOW TABLES` will be
// rewritten to the queries on the system tables, we need to skip the privilege check on
// these tables.
const SYSTEM_TABLES_ALLOW_LIST: [&str; 21] = [
"catalogs",
"columns",
"databases",
"databases_with_history",
"dictionaries",
"tables",
"views",
"tables_with_history",
"views_with_history",
"password_policies",
"streams",
"streams_terse",
"virtual_columns",
"users",
"roles",
"stages",
"one",
"processes",
"user_functions",
"functions",
"indexes",
];

// table functions that need `Super` privilege
const SYSTEM_TABLE_FUNCTIONS: [&str; 1] = ["fuse_amend"];

Expand Down
10 changes: 8 additions & 2 deletions src/query/service/src/sessions/query_ctx.rs
Original file line number Diff line number Diff line change
Expand Up @@ -692,8 +692,14 @@ impl TableContext for QueryContext {
self.get_current_session().get_id()
}

async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker> {
self.shared.session.get_visibility_checker().await
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker> {
self.shared
.session
.get_visibility_checker(ignore_ownership)
.await
}

fn get_fuse_version(&self) -> String {
Expand Down
9 changes: 7 additions & 2 deletions src/query/service/src/sessions/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -310,8 +310,13 @@ impl Session {
}

#[async_backtrace::framed]
pub async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker> {
self.privilege_mgr().get_visibility_checker().await
pub async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker> {
self.privilege_mgr()
.get_visibility_checker(ignore_ownership)
.await
}

pub fn get_settings(&self) -> Arc<Settings> {
Expand Down
41 changes: 24 additions & 17 deletions src/query/service/src/sessions/session_privilege_mgr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,10 @@ pub trait SessionPrivilegeManager {

async fn validate_available_role(&self, role_name: &str) -> Result<RoleInfo>;

async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker>;
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker>;

// fn show_grants(&self);
}
Expand Down Expand Up @@ -336,27 +339,31 @@ impl<'a> SessionPrivilegeManager for SessionPrivilegeManagerImpl<'a> {
}

#[async_backtrace::framed]
async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker> {
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker> {
// TODO(liyz): is it check the visibility according onwerships?
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&self.session_ctx.get_current_tenant())
.get_ownerships()
.await?;
let roles = self.get_all_effective_roles().await?;
let roles_name: Vec<String> = roles.iter().map(|role| role.name.to_string()).collect();

let ownership_objects = if roles_name.contains(&"account_admin".to_string()) {
vec![]
} else {
let mut ownership_objects = vec![];
for ownership in ownerships {
if roles_name.contains(&ownership.data.role) {
ownership_objects.push(ownership.data.object);
let ownership_objects =
if roles_name.contains(&"account_admin".to_string()) || ignore_ownership {
vec![]
} else {
let user_api = UserApiProvider::instance();
let ownerships = user_api
.role_api(&self.session_ctx.get_current_tenant())
.get_ownerships()
.await?;
let mut ownership_objects = vec![];
for ownership in ownerships {
if roles_name.contains(&ownership.data.role) {
ownership_objects.push(ownership.data.object);
}
}
}
ownership_objects
};
ownership_objects
};

Ok(GrantObjectVisibilityChecker::new(
&self.get_current_user()?,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ impl AsyncSource for ParquetInferSchemaSource {
.get_settings()
.get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
let visibility_checker = self.ctx.get_visibility_checker(false).await?;
if !(stage_info.is_temporary
|| visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
|| stage_info.stage_type == StageType::User
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ impl AsyncSource for InspectParquetSource {
.get_settings()
.get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
let visibility_checker = self.ctx.get_visibility_checker(false).await?;
if !(stage_info.is_temporary
|| visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
|| stage_info.stage_type == StageType::User
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ impl ListStagesSource {
.get_settings()
.get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
let visibility_checker = self.ctx.get_visibility_checker(false).await?;
if !(stage_info.is_temporary
|| visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
|| stage_info.stage_type == StageType::User
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -566,7 +566,7 @@ async fn show_object_grant(
let tenant = ctx.get_tenant();
let user_api = UserApiProvider::instance();
let roles = user_api.get_roles(&tenant).await?;
let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;
let current_user = ctx.get_current_user()?.identity().username;
let (object, owner_object, object_id, object_name) = match grant_type {
"table" => {
Expand Down
5 changes: 4 additions & 1 deletion src/query/service/tests/it/sql/exec/get_table_bind_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -690,7 +690,10 @@ impl TableContext for CtxDelegation {
todo!()
}

async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker> {
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker> {
todo!()
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -594,7 +594,10 @@ impl TableContext for CtxDelegation {
todo!()
}

async fn get_visibility_checker(&self) -> Result<GrantObjectVisibilityChecker> {
async fn get_visibility_checker(
&self,
ignore_ownership: bool,
) -> Result<GrantObjectVisibilityChecker> {
todo!()
}

Expand Down
1 change: 1 addition & 0 deletions src/query/storages/system/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ databend-common-storages-result-cache = { workspace = true }
databend-common-storages-stream = { workspace = true }
databend-common-storages-view = { workspace = true }
databend-common-users = { workspace = true }
databend-common-management = { workspace = true }
databend-storages-common-cache = { workspace = true }
futures = { workspace = true }
itertools = { workspace = true }
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/columns_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,7 @@ pub(crate) async fn dump_tables(
}
}

let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;

let mut final_dbs: Vec<(String, u64)> = Vec::new();

Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/databases_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ where DatabasesTable<WITH_HISTORY>: HistoryAware
let mut owners: Vec<Option<String>> = vec![];
let mut dropped_on: Vec<Option<i64>> = vec![];

let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;
let catalog_dbs = visibility_checker.get_visibility_database();
// None means has global level privileges
if let Some(catalog_dbs) = catalog_dbs {
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/indexes_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ impl IndexesTable {
ctx: Arc<dyn TableContext>,
) -> Result<Vec<TableInfo>> {
let tenant = ctx.get_tenant();
let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;
let catalog = ctx.get_catalog(CATALOG_DEFAULT).await?;

let ctl_name = catalog.name();
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/system/src/stages_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ impl AsyncSystemTable for StagesTable {
let enable_experimental_rbac_check =
ctx.get_settings().get_enable_experimental_rbac_check()?;
let stages = if enable_experimental_rbac_check {
let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;
stages
.into_iter()
.filter(|stage| {
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 @@ -80,7 +80,7 @@ impl<const T: bool> AsyncSystemTable for StreamsTable<T> {
.iter()
.map(|e| (e.name(), e.clone()))
.collect::<Vec<_>>();
let visibility_checker = ctx.get_visibility_checker().await?;
let visibility_checker = ctx.get_visibility_checker(false).await?;
let user_api = UserApiProvider::instance();

let mut catalogs = vec![];
Expand Down
Loading

0 comments on commit 58caec1

Please sign in to comment.