Skip to content

Commit

Permalink
optimize(query): first check privilege in SystemEngine get_full_data
Browse files Browse the repository at this point in the history
  • Loading branch information
@TCeason
TCeason committed Sep 9, 2024
1 parent 6253e83 commit c95b1c3
Show file tree
Hide file tree
Showing 7 changed files with 279 additions and 43 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

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

63 changes: 52 additions & 11 deletions src/query/storages/system/src/columns_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,16 +277,7 @@ pub(crate) async fn dump_tables(

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

if databases.is_empty() {
let all_databases = catalog.list_databases(&tenant).await?;
for db in all_databases {
let db_id = db.get_db_info().database_id.db_id;
let db_name = db.name();
if visibility_checker.check_database_visibility(CATALOG_DEFAULT, db_name, db_id) {
final_dbs.push((db_name.to_string(), db_id));
}
}
} else {
if !databases.is_empty() {
for db in databases {
let db_id = catalog
.get_database(&tenant, &db)
Expand All @@ -298,12 +289,62 @@ pub(crate) async fn dump_tables(
final_dbs.push((db.to_string(), db_id));
}
}
} else {
let catalog_dbs = visibility_checker.get_visibility_database();
if let Some(catalog_dbs) = catalog_dbs {
for (catalog_name, dbs) in catalog_dbs {
if catalog_name == CATALOG_DEFAULT {
let mut catalog_db_ids = vec![];
for (db_name, db_id) in dbs {
if let Some(db_name) = db_name {
let db_id = catalog
.get_database(&tenant, db_name)
.await?
.get_db_info()
.database_id
.db_id;
final_dbs.push((db_name.to_string(), db_id));
}
if let Some(db_id) = db_id {
catalog_db_ids.push(*db_id);
}
}
match catalog
.mget_database_names_by_ids(&tenant, &catalog_db_ids)
.await
{
Ok(databases) => {
for (i, db) in databases.into_iter().flatten().enumerate() {
final_dbs.push((db.to_string(), catalog_db_ids[i]));
}
}
Err(err) => {
let msg =
format!("Failed to get database: {}, {}", catalog.name(), err);
warn!("{}", msg);
}
}
}
}
} else {
let all_databases = catalog.list_databases(&tenant).await?;
for db in all_databases {
let db_id = db.get_db_info().database_id.db_id;
let db_name = db.name();
if visibility_checker.check_database_visibility(CATALOG_DEFAULT, db_name, db_id) {
final_dbs.push((db_name.to_string(), db_id));
}
}
}
}

let mut final_tables: Vec<(String, Vec<Arc<dyn Table>>)> = Vec::with_capacity(final_dbs.len());
for (database, db_id) in final_dbs {
let tables = if tables.is_empty() {
(catalog.list_tables(&tenant, &database).await).unwrap_or_default()
catalog
.list_tables(&tenant, &database)
.await
.unwrap_or_default()
} else {
let mut res = Vec::new();
for table in &tables {
Expand Down
125 changes: 93 additions & 32 deletions src/query/storages/system/src/databases_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ use databend_common_meta_app::schema::TableIdent;
use databend_common_meta_app::schema::TableInfo;
use databend_common_meta_app::schema::TableMeta;
use databend_common_users::UserApiProvider;
use log::warn;

use crate::table::AsyncOneBlockSystemTable;
use crate::table::AsyncSystemTable;
Expand Down Expand Up @@ -68,47 +69,107 @@ impl AsyncSystemTable for DatabasesTable {
let user_api = UserApiProvider::instance();
let mut catalog_names = vec![];
let mut db_names = vec![];
let mut db_id = vec![];
let mut db_ids = vec![];
let mut owners: Vec<Option<String>> = vec![];

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

for (ctl_name, catalog) in catalogs.into_iter() {
let databases = catalog.list_databases(&tenant).await?;
let final_dbs = databases
.into_iter()
.filter(|db| {
visibility_checker.check_database_visibility(
&ctl_name,
db.name(),
db.get_db_info().database_id.db_id,
)
})
.collect::<Vec<_>>();

for db in final_dbs {
catalog_names.push(ctl_name.clone());
let db_name = db.name().to_string();
db_names.push(db_name);
let id = db.get_db_info().database_id.db_id;
db_id.push(id);
owners.push(
user_api
.get_ownership(&tenant, &OwnershipObject::Database {
catalog_name: ctl_name.to_string(),
db_id: id,
})
.await
.ok()
.and_then(|ownership| ownership.map(|o| o.role.clone())),
);
let catalog_dbs = visibility_checker.get_visibility_database();
if let Some(catalog_dbs) = catalog_dbs {
for (catalog, dbs) in catalog_dbs {
let mut catalog_db_ids = vec![];
let ctl = ctx.get_catalog(catalog).await?;
for (db_name, db_id) in dbs {
if let Some(db_name) = db_name {
let db_id = ctl
.get_database(&tenant, db_name)
.await?
.get_db_info()
.database_id
.db_id;

catalog_names.push(catalog.clone());
db_names.push(db_name.to_string());
db_ids.push(db_id);
owners.push(
user_api
.get_ownership(&tenant, &OwnershipObject::Database {
catalog_name: catalog.clone(),
db_id,
})
.await
.ok()
.and_then(|ownership| ownership.map(|o| o.role.clone())),
);
}
if let Some(db_id) = db_id {
catalog_db_ids.push(*db_id);
}
}
match ctl
.mget_database_names_by_ids(&tenant, &catalog_db_ids)
.await
{
Ok(databases) => {
for (i, db) in databases.into_iter().flatten().enumerate() {
catalog_names.push(catalog.clone());
db_names.push(db);
db_ids.push(catalog_db_ids[i]);
owners.push(
user_api
.get_ownership(&tenant, &OwnershipObject::Database {
catalog_name: catalog.clone(),
db_id: catalog_db_ids[i],
})
.await
.ok()
.and_then(|ownership| ownership.map(|o| o.role.clone())),
);
}
}
Err(err) => {
let msg = format!("Failed to get database: {}, {}", ctl.name(), err);
warn!("{}", msg);
}
}
}
} else {
for (ctl_name, catalog) in catalogs.into_iter() {
let databases = catalog.list_databases(&tenant).await?;
let final_dbs = databases
.into_iter()
.filter(|db| {
visibility_checker.check_database_visibility(
&ctl_name,
db.name(),
db.get_db_info().database_id.db_id,
)
})
.collect::<Vec<_>>();

for db in final_dbs {
catalog_names.push(ctl_name.clone());
let db_name = db.name().to_string();
db_names.push(db_name);
let id = db.get_db_info().database_id.db_id;
db_ids.push(id);
owners.push(
user_api
.get_ownership(&tenant, &OwnershipObject::Database {
catalog_name: ctl_name.to_string(),
db_id: id,
})
.await
.ok()
.and_then(|ownership| ownership.map(|o| o.role.clone())),
);
}
}
}

Ok(DataBlock::new_from_columns(vec![
StringType::from_data(catalog_names),
StringType::from_data(db_names),
UInt64Type::from_data(db_id),
UInt64Type::from_data(db_ids),
StringType::from_opt_data(owners),
]))
}
Expand Down
1 change: 1 addition & 0 deletions src/query/users/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ passwords = { version = "3.1.16", features = ["common-password"] }
reqwest = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
itertools = "0.13.0"

[dev-dependencies]
databend-common-expression = { workspace = true }
Expand Down
57 changes: 57 additions & 0 deletions src/query/users/src/visibility_checker.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::collections::HashSet;
use std::string::ToString;

use databend_common_meta_app::principal::GrantObject;
use databend_common_meta_app::principal::OwnershipObject;
Expand All @@ -22,6 +24,7 @@ use databend_common_meta_app::principal::UserInfo;
use databend_common_meta_app::principal::UserPrivilegeSet;
use databend_common_meta_app::principal::UserPrivilegeType;
use enumflags2::BitFlags;
use itertools::Itertools;

/// GrantObjectVisibilityChecker is used to check whether a user has the privilege to access a
/// database or table.
Expand All @@ -36,6 +39,7 @@ pub struct GrantObjectVisibilityChecker {
granted_tables: HashSet<(String, String, String)>,
granted_tables_id: HashSet<(String, u64, u64)>,
extra_databases: HashSet<(String, String)>,
sys_databases: HashSet<(String, String)>,
extra_databases_id: HashSet<(String, u64)>,
granted_udfs: HashSet<String>,
granted_write_stages: HashSet<String>,
Expand Down Expand Up @@ -197,6 +201,10 @@ impl GrantObjectVisibilityChecker {
granted_udfs,
granted_write_stages,
granted_read_stages,
sys_databases: HashSet::from([
("default".to_string(), "information_schema".to_string()),
("default".to_string(), "system".to_string()),
]),
}
}

Expand Down Expand Up @@ -324,4 +332,53 @@ impl GrantObjectVisibilityChecker {

false
}

#[allow(clippy::type_complexity)]
pub fn get_visibility_database(
&self,
) -> Option<HashMap<&String, HashSet<(Option<&String>, Option<&u64>)>>> {
if self.granted_global_db_table {
return None;
}

let capacity = self.granted_databases.len()
+ self.granted_databases_id.len()
+ self.extra_databases.len()
+ self.extra_databases_id.len()
+ self.sys_databases.len();

let dbs = self
.granted_databases
.iter()
.map(|(catalog, db)| (catalog, (Some(db), None)))
.chain(
self.granted_databases_id
.iter()
.map(|(catalog, db_id)| (catalog, (None, Some(db_id)))),
)
.chain(
self.extra_databases
.iter()
.map(|(catalog, db)| (catalog, (Some(db), None))),
)
.chain(
self.sys_databases
.iter()
.map(|(catalog, db)| (catalog, (Some(db), None))),
)
.chain(
self.extra_databases_id
.iter()
.map(|(catalog, db_id)| (catalog, (None, Some(db_id)))),
)
.into_grouping_map()
.fold(
HashSet::with_capacity(capacity / 4),
|mut set, _key, value| {
set.insert(value);
set
},
);
Some(dbs)
}
}
27 changes: 27 additions & 0 deletions tests/suites/0_stateless/18_rbac/18_0013_column_privilege.result
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
=== show grants for a ===
SELECT default USER a GRANT SELECT ON 'default'.'default'.* TO 'a'@'%'
SELECT default.grant_db.t USER a GRANT SELECT ON 'default'.'grant_db'.'t' TO 'a'@'%'
=== show databases ===
default
grant_db
information_schema
system
=== show tables ===
Error: APIError: ResponseError with 1063: Permission denied: User 'a'@'%' does not have the required privileges for database 'system'
t
=== use db ===
Error: APIError: ResponseError with 1063: Permission denied: User 'a'@'%' does not have the required privileges for database 'system'
=== show columns ===
dummy TINYINT UNSIGNED NO NULL NULL
c1 INT NO NULL NULL
created_on TIMESTAMP NO NULL NULL
inherited_roles BIGINT UNSIGNED NO NULL NULL
inherited_roles_name VARCHAR NO NULL NULL
name VARCHAR NO NULL NULL
update_on TIMESTAMP NO NULL NULL
keywords VARCHAR NO NULL NULL
reserved TINYINT UNSIGNED NO NULL NULL
Error: APIError: ResponseError with 1063: Permission denied: User 'a'@'%' does not have the required privileges for database 'nogrant'
Error: APIError: ResponseError with 1063: Permission denied: User 'a'@'%' does not have the required privileges for table 'nogrant.t'
1
0
Loading

0 comments on commit c95b1c3

Please sign in to comment.