Skip to content

Commit

Permalink
chore(query): rename enable_stage_udf_priv_check to experiment_enable…
Browse files Browse the repository at this point in the history 
…_new_rbac_check
  • Loading branch information
@TCeason
TCeason committed Nov 30, 2023
1 parent 2f7e1be commit 08fa6b4
Show file tree
Hide file tree
Showing 12 changed files with 55 additions and 53 deletions.
64 changes: 34 additions & 30 deletions src/query/service/src/interpreters/access/privilege_access.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,8 @@ impl AccessChecker for PrivilegeAccess {
let user = self.ctx.get_current_user()?;
let (identity, grant_set) = (user.identity().to_string(), user.grants);
let tenant = self.ctx.get_tenant();
let enable_stage_udf_priv_check =
self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
let enable_experimental_rbac_check =
self.ctx.get_settings().get_enable_experimental_rbac_check()?;

match plan {
Plan::Query {
Expand Down Expand Up @@ -227,7 +227,7 @@ impl AccessChecker for PrivilegeAccess {
}
_ => {}
};
if enable_stage_udf_priv_check {
if enable_experimental_rbac_check {
match s_expr.get_udfs() {
Ok(udfs) => {
if !udfs.is_empty() {
Expand All @@ -249,7 +249,7 @@ impl AccessChecker for PrivilegeAccess {
let metadata = metadata.read().clone();

for table in metadata.tables() {
if table.is_source_of_stage() {
if enable_experimental_rbac_check && table.is_source_of_stage() {
match table.table().get_data_source_info() {
DataSourceInfo::StageSource(stage_info) => {
self.validate_access_stage(&stage_info.stage_info, UserPrivilegeType::Read).await?;
Expand Down Expand Up @@ -529,9 +529,11 @@ impl AccessChecker for PrivilegeAccess {
.await?;
}
Plan::ReclusterTable(plan) => {
if let Some(scalar) = &plan.push_downs {
let udf = get_udf_names(scalar)?;
self.check_udf_priv(udf).await?;
if enable_experimental_rbac_check {
if let Some(scalar) = &plan.push_downs {
let udf = get_udf_names(scalar)?;
self.check_udf_priv(udf).await?;
}
}
self.validate_access(
&GrantObject::Table(
Expand Down Expand Up @@ -628,7 +630,7 @@ impl AccessChecker for PrivilegeAccess {
.await?;
}
Plan::MergeInto(plan) => {
if enable_stage_udf_priv_check {
if enable_experimental_rbac_check {
let s_expr = &plan.input;
match s_expr.get_udfs() {
Ok(udfs) => {
Expand Down Expand Up @@ -683,7 +685,7 @@ impl AccessChecker for PrivilegeAccess {
.await?;
}
Plan::Delete(plan) => {
if enable_stage_udf_priv_check {
if enable_experimental_rbac_check {
if let Some(selection) = &plan.selection {
let udf = get_udf_names(selection)?;
self.check_udf_priv(udf).await?;
Expand Down Expand Up @@ -719,29 +721,31 @@ impl AccessChecker for PrivilegeAccess {
.await?;
}
Plan::Update(plan) => {
for scalar in plan.update_list.values() {
let udf = get_udf_names(scalar)?;
self.check_udf_priv(udf).await?;
}
if let Some(selection) = &plan.selection {
let udf = get_udf_names(selection)?;
self.check_udf_priv(udf).await?;
}
for subquery in &plan.subquery_desc {
match subquery.input_expr.get_udfs() {
Ok(udfs) => {
if !udfs.is_empty() {
for udf in udfs {
self.validate_access(
&GrantObject::UDF(udf.clone()),
vec![UserPrivilegeType::Usage],
false,
).await?
if enable_experimental_rbac_check {
for scalar in plan.update_list.values() {
let udf = get_udf_names(scalar)?;
self.check_udf_priv(udf).await?;
}
if let Some(selection) = &plan.selection {
let udf = get_udf_names(selection)?;
self.check_udf_priv(udf).await?;
}
for subquery in &plan.subquery_desc {
match subquery.input_expr.get_udfs() {
Ok(udfs) => {
if !udfs.is_empty() {
for udf in udfs {
self.validate_access(
&GrantObject::UDF(udf.clone()),
vec![UserPrivilegeType::Usage],
false,
).await?
}
}
}
}
Err(err) => {
return Err(err.add_message("get udf error on validating access"));
Err(err) => {
return Err(err.add_message("get udf error on validating access"));
}
}
}
}
Expand Down
5 changes: 2 additions & 3 deletions src/query/service/src/table_functions/infer_schema/infer_schema_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,9 +181,8 @@ impl AsyncSource for InferSchemaSource {

let (stage_info, path) =
resolve_stage_location(&self.ctx, &self.args_parsed.location).await?;
let enable_stage_udf_priv_check =
self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
if enable_stage_udf_priv_check {
let enable_experimental_rbac_check = self.ctx.get_settings().get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
if !stage_info.is_from_uri
&& !visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
Expand Down
5 changes: 2 additions & 3 deletions src/query/service/src/table_functions/inspect_parquet/inspect_parquet_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -210,9 +210,8 @@ impl AsyncSource for InspectParquetSource {
self.is_finished = true;
let uri = self.uri.strip_prefix('@').unwrap().to_string();
let (stage_info, path) = resolve_stage_location(&self.ctx, &uri).await?;
let enable_stage_udf_priv_check =
self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
if enable_stage_udf_priv_check {
let enable_experimental_rbac_check = self.ctx.get_settings().get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
if !stage_info.is_from_uri
&& !visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
Expand Down
5 changes: 2 additions & 3 deletions src/query/service/src/table_functions/list_stage/list_stage_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,9 +185,8 @@ impl AsyncSource for ListStagesSource {

let (stage_info, path) =
resolve_stage_location(&self.ctx, &self.args_parsed.location).await?;
let enable_stage_udf_priv_check =
self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
if enable_stage_udf_priv_check {
let enable_experimental_rbac_check = self.ctx.get_settings().get_enable_experimental_rbac_check()?;
if enable_experimental_rbac_check {
let visibility_checker = self.ctx.get_visibility_checker().await?;
if !stage_info.is_from_uri
&& !visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
Expand Down
2 changes: 1 addition & 1 deletion src/query/service/tests/it/storages/testdata/settings_table.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ DB.Table: 'system'.'settings', Table: settings-table_id:1, ver:0, Engine: System
| 'enable_distributed_replace_into' | '0' | '0' | 'SESSION' | 'Enable distributed execution of replace into.' | 'UInt64' |
| 'enable_dphyp' | '1' | '1' | 'SESSION' | 'Enables dphyp join order algorithm.' | 'UInt64' |
| 'enable_experimental_merge_into' | '0' | '0' | 'SESSION' | 'Enable experimental merge into.' | 'UInt64' |
| 'enable_experimental_rbac_check' | '0' | '0' | 'SESSION' | 'experiment setting disables stage and udf privilege check(disable by default).' | 'UInt64' |
| 'enable_hive_parquet_predict_pushdown' | '1' | '1' | 'SESSION' | 'Enable hive parquet predict pushdown by setting this variable to 1, default value: 1' | 'UInt64' |
| 'enable_parquet_page_index' | '1' | '1' | 'SESSION' | 'Enables parquet page index' | 'UInt64' |
| 'enable_parquet_prewhere' | '0' | '0' | 'SESSION' | 'Enables parquet prewhere' | 'UInt64' |
Expand All @@ -31,7 +32,6 @@ DB.Table: 'system'.'settings', Table: settings-table_id:1, ver:0, Engine: System
| 'enable_replace_into_partitioning' | '1' | '1' | 'SESSION' | 'Enables partitioning for replace-into statement (if table has cluster keys).' | 'UInt64' |
| 'enable_runtime_filter' | '0' | '0' | 'SESSION' | 'Enables runtime filter optimization for JOIN.' | 'UInt64' |
| 'enable_table_lock' | '1' | '1' | 'SESSION' | 'Enables table lock if necessary (enabled by default).' | 'UInt64' |
| 'experiment_enable_stage_udf_priv_check' | '0' | '0' | 'SESSION' | 'experiment setting disables stage and udf privilege check(disable by default).' | 'UInt64' |
| 'external_server_connect_timeout_secs' | '10' | '10' | 'SESSION' | 'Connection timeout to external server' | 'UInt64' |
| 'external_server_request_timeout_secs' | '180' | '180' | 'SESSION' | 'Request timeout to external server' | 'UInt64' |
| 'flight_client_timeout' | '60' | '60' | 'SESSION' | 'Sets the maximum time in seconds that a flight client request can be processed.' | 'UInt64' |
Expand Down
2 changes: 1 addition & 1 deletion src/query/settings/src/settings_default.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -519,7 +519,7 @@ impl DefaultSettings {
possible_values: Some(vec!["rounding", "truncating"]),
mode: SettingMode::Both,
}),
("experiment_enable_stage_udf_priv_check", DefaultSettingValue {
("enable_experimental_rbac_check", DefaultSettingValue {
value: UserSettingValue::UInt64(0),
desc: "experiment setting disables stage and udf privilege check(disable by default).",
possible_values: None,
Expand Down
4 changes: 2 additions & 2 deletions src/query/settings/src/settings_getter_setter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -339,8 +339,8 @@ impl Settings {
Ok(self.try_get_u64("enable_table_lock")? != 0)
}

pub fn get_enable_stage_udf_priv_check(&self) -> Result<bool> {
Ok(self.try_get_u64("experiment_enable_stage_udf_priv_check")? != 0)
pub fn get_enable_experimental_rbac_check(&self) -> Result<bool> {
Ok(self.try_get_u64("enable_experimental_rbac_check")? != 0)
}

pub fn get_table_lock_expire_secs(&self) -> Result<u64> {
Expand Down
4 changes: 2 additions & 2 deletions src/query/storages/system/src/functions_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,8 @@ impl AsyncSystemTable for FunctionsTable {
scalar_func_names.sort();
let aggregate_function_factory = AggregateFunctionFactory::instance();
let aggr_func_names = aggregate_function_factory.registered_names();
let enable_stage_udf_priv_check = ctx.get_settings().get_enable_stage_udf_priv_check()?;
let udfs = if enable_stage_udf_priv_check {
let enable_experimental_rbac_check = ctx.get_settings().get_enable_experimental_rbac_check()?;
let udfs = if enable_experimental_rbac_check {
let visibility_checker = ctx.get_visibility_checker().await?;
let udfs = FunctionsTable::get_udfs(ctx).await?;
udfs.into_iter()
Expand Down
4 changes: 2 additions & 2 deletions src/query/storages/system/src/stages_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,8 @@ impl AsyncSystemTable for StagesTable {
) -> Result<DataBlock> {
let tenant = ctx.get_tenant();
let stages = UserApiProvider::instance().get_stages(&tenant).await?;
let enable_stage_udf_priv_check = ctx.get_settings().get_enable_stage_udf_priv_check()?;
let stages = if enable_stage_udf_priv_check {
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?;
stages
.into_iter()
Expand Down
4 changes: 2 additions & 2 deletions tests/suites/0_stateless/18_rbac/18_0001_udf_priv.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ echo "=== test UDF priv"
export TEST_USER_PASSWORD="password"
export TEST_USER_CONNECT="bendsql --user=test-user --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"

echo "set global experiment_enable_stage_udf_priv_check=1" | $BENDSQL_CLIENT_CONNECT
echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT

echo "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
echo "DROP FUNCTION IF EXISTS f1;" | $BENDSQL_CLIENT_CONNECT
Expand Down Expand Up @@ -128,4 +128,4 @@ echo "DROP FUNCTION IF EXISTS f1;" | $BENDSQL_CLIENT_CONNECT
echo "DROP FUNCTION IF EXISTS f2;" | $BENDSQL_CLIENT_CONNECT
echo "drop table if exists default.t;" | $BENDSQL_CLIENT_CONNECT
echo "drop table if exists default.t2;" | $BENDSQL_CLIENT_CONNECT
echo "unset experiment_enable_stage_udf_priv_check" | $BENDSQL_CLIENT_CONNECT
echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT
4 changes: 2 additions & 2 deletions tests/suites/0_stateless/18_rbac/20_0012_privilege_access.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
export TEST_USER_PASSWORD="password"
export TEST_USER_CONNECT="bendsql --user=test-user --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"

echo "set global experiment_enable_stage_udf_priv_check=1" | $BENDSQL_CLIENT_CONNECT
echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT

echo "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
echo "drop role if exists 'test-role1'" | $BENDSQL_CLIENT_CONNECT
Expand Down Expand Up @@ -179,4 +179,4 @@ echo "drop user a" | $BENDSQL_CLIENT_CONNECT
echo "drop database if exists no_grant" | $BENDSQL_CLIENT_CONNECT
echo "drop database grant_db" | $BENDSQL_CLIENT_CONNECT

echo "unset experiment_enable_stage_udf_priv_check" | $BENDSQL_CLIENT_CONNECT
echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT
5 changes: 3 additions & 2 deletions tests/suites/1_stateful/00_stage/00_0012_stage_priv.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ export TEST_USER_NAME="u1"
export TEST_USER_PASSWORD="password"
export TEST_USER_CONNECT="bendsql --user=u1 --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"

echo "set global experiment_enable_stage_udf_priv_check=1" | $BENDSQL_CLIENT_CONNECT
echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT

echo "drop table if exists test_table;" | $BENDSQL_CLIENT_CONNECT
echo "drop user if exists u1;" | $BENDSQL_CLIENT_CONNECT
Expand Down Expand Up @@ -121,4 +121,5 @@ echo "drop stage if exists s3" | $BENDSQL_CLIENT_CONNECT
echo "drop user u1" | $BENDSQL_CLIENT_CONNECT
echo "drop table if exists t" | $BENDSQL_CLIENT_CONNECT
rm -rf /tmp/00_0012
echo "unset experiment_enable_stage_udf_priv_check" | $BENDSQL_CLIENT_CONNECT

echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT

0 comments on commit 08fa6b4

Please sign in to comment.