chore(query): add settings to control enable stage/udf privilege check (

#13852)
databendlabs · Nov 29, 2023 · 271ea35 · 271ea35
1 parent 3f94663
commit 271ea35
Show file tree

Hide file tree

Showing 12 changed files with 249 additions and 187 deletions.
diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs
diff --git a/src/query/service/src/table_functions/infer_schema/infer_schema_table.rs b/src/query/service/src/table_functions/infer_schema/infer_schema_table.rs
@@ -181,15 +181,19 @@ impl AsyncSource for InferSchemaSource {
 
         let (stage_info, path) =
             resolve_stage_location(&self.ctx, &self.args_parsed.location).await?;
-        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)
-        {
-            return Err(ErrorCode::PermissionDenied(format!(
-                "Permission denied, privilege READ is required on stage {} for user {}",
-                stage_info.stage_name.clone(),
-                &self.ctx.get_current_user()?.identity(),
-            )));
+        let enable_stage_udf_priv_check =
+            self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
+        if enable_stage_udf_priv_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)
+            {
+                return Err(ErrorCode::PermissionDenied(format!(
+                    "Permission denied, privilege READ is required on stage {} for user {}",
+                    stage_info.stage_name.clone(),
+                    &self.ctx.get_current_user()?.identity(),
+                )));
+            }
         }
         let files_info = StageFilesInfo {
             path: path.clone(),

diff --git a/src/query/service/src/table_functions/inspect_parquet/inspect_parquet_table.rs b/src/query/service/src/table_functions/inspect_parquet/inspect_parquet_table.rs
@@ -210,15 +210,19 @@ 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 visibility_checker = self.ctx.get_visibility_checker().await?;
-        if !stage_info.is_from_uri
-            && !visibility_checker.check_stage_read_visibility(&stage_info.stage_name)
-        {
-            return Err(ErrorCode::PermissionDenied(format!(
-                "Permission denied, privilege READ is required on stage {} for user {}",
-                stage_info.stage_name.clone(),
-                &self.ctx.get_current_user()?.identity(),
-            )));
+        let enable_stage_udf_priv_check =
+            self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
+        if enable_stage_udf_priv_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)
+            {
+                return Err(ErrorCode::PermissionDenied(format!(
+                    "Permission denied, privilege READ is required on stage {} for user {}",
+                    stage_info.stage_name.clone(),
+                    &self.ctx.get_current_user()?.identity(),
+                )));
+            }
         }
 
         let operator = init_stage_operator(&stage_info)?;

diff --git a/src/query/service/src/table_functions/list_stage/list_stage_table.rs b/src/query/service/src/table_functions/list_stage/list_stage_table.rs
@@ -185,15 +185,19 @@ impl AsyncSource for ListStagesSource {
 
         let (stage_info, path) =
             resolve_stage_location(&self.ctx, &self.args_parsed.location).await?;
-        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)
-        {
-            return Err(ErrorCode::PermissionDenied(format!(
-                "Permission denied, privilege READ is required on stage {} for user {}",
-                stage_info.stage_name.clone(),
-                &self.ctx.get_current_user()?.identity(),
-            )));
+        let enable_stage_udf_priv_check =
+            self.ctx.get_settings().get_enable_stage_udf_priv_check()?;
+        if enable_stage_udf_priv_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)
+            {
+                return Err(ErrorCode::PermissionDenied(format!(
+                    "Permission denied, privilege READ is required on stage {} for user {}",
+                    stage_info.stage_name.clone(),
+                    &self.ctx.get_current_user()?.identity(),
+                )));
+            }
         }
         let op = StageTable::get_op(&stage_info)?;
 

diff --git a/src/query/service/tests/it/storages/testdata/settings_table.txt b/src/query/service/tests/it/storages/testdata/settings_table.txt
@@ -30,6 +30,7 @@ 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' |

diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs
@@ -508,6 +508,12 @@ impl DefaultSettings {
                     possible_values: Some(vec!["rounding", "truncating"]),
                     display_in_show_settings: true,
                 }),
+                ("experiment_enable_stage_udf_priv_check", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(0),
+                    desc: "experiment setting disables stage and udf privilege check(disable by default).",
+                    possible_values: None,
+                    display_in_show_settings: true,
+                }),
             ]);
 
             Ok(Arc::new(DefaultSettings {

diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs
@@ -297,6 +297,10 @@ 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_table_lock_expire_secs(&self) -> Result<u64> {
         self.try_get_u64("table_lock_expire_secs")
     }

diff --git a/src/query/storages/system/src/functions_table.rs b/src/query/storages/system/src/functions_table.rs
@@ -69,12 +69,16 @@ impl AsyncSystemTable for FunctionsTable {
         scalar_func_names.sort();
         let aggregate_function_factory = AggregateFunctionFactory::instance();
         let aggr_func_names = aggregate_function_factory.registered_names();
-        let visibility_checker = ctx.get_visibility_checker().await?;
-        let udfs = FunctionsTable::get_udfs(ctx).await?;
-        let udfs = udfs
-            .into_iter()
-            .filter(|udf| visibility_checker.check_udf_visibility(&udf.name))
-            .collect::<Vec<_>>();
+        let enable_stage_udf_priv_check = ctx.get_settings().get_enable_stage_udf_priv_check()?;
+        let udfs = if enable_stage_udf_priv_check {
+            let visibility_checker = ctx.get_visibility_checker().await?;
+            let udfs = FunctionsTable::get_udfs(ctx).await?;
+            udfs.into_iter()
+                .filter(|udf| visibility_checker.check_udf_visibility(&udf.name))
+                .collect::<Vec<_>>()
+        } else {
+            FunctionsTable::get_udfs(ctx).await?
+        };
 
         let names: Vec<&str> = scalar_func_names
             .iter()

diff --git a/src/query/storages/system/src/stages_table.rs b/src/query/storages/system/src/stages_table.rs
@@ -56,13 +56,20 @@ impl AsyncSystemTable for StagesTable {
     ) -> Result<DataBlock> {
         let tenant = ctx.get_tenant();
         let stages = UserApiProvider::instance().get_stages(&tenant).await?;
-        let visibility_checker = ctx.get_visibility_checker().await?;
-        let stages = stages
-            .into_iter()
-            .filter(|stage| {
-                !stage.is_from_uri && visibility_checker.check_stage_visibility(&stage.stage_name)
-            })
-            .collect::<Vec<_>>();
+        let enable_stage_udf_priv_check = ctx.get_settings().get_enable_stage_udf_priv_check()?;
+        let stages = if enable_stage_udf_priv_check {
+            let visibility_checker = ctx.get_visibility_checker().await?;
+            stages
+                .into_iter()
+                .filter(|stage| {
+                    !stage.is_from_uri
+                        && visibility_checker.check_stage_visibility(&stage.stage_name)
+                })
+                .collect::<Vec<_>>()
+        } else {
+            stages
+        };
+
         let mut name: Vec<Vec<u8>> = Vec::with_capacity(stages.len());
         let mut stage_type: Vec<Vec<u8>> = Vec::with_capacity(stages.len());
         let mut stage_params: Vec<Vec<u8>> = Vec::with_capacity(stages.len());

diff --git a/tests/suites/1_stateful/11_rbac/00_0012_stage_priv.sh b/tests/suites/1_stateful/11_rbac/00_0012_stage_priv.sh
@@ -7,9 +7,12 @@ 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 "drop table if exists test_table;" | $BENDSQL_CLIENT_CONNECT
 echo "drop user if exists u1;" | $BENDSQL_CLIENT_CONNECT
 echo "drop STAGE if exists s2;" | $BENDSQL_CLIENT_CONNECT
+echo "drop STAGE if exists s1;" | $BENDSQL_CLIENT_CONNECT
 echo "CREATE STAGE s2;" | $BENDSQL_CLIENT_CONNECT
 
 echo "CREATE TABLE test_table (
@@ -115,4 +118,5 @@ 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
 
diff --git a/tests/suites/1_stateful/11_rbac/18_0001_udf_priv.sh b/tests/suites/1_stateful/11_rbac/18_0001_udf_priv.sh
@@ -7,6 +7,8 @@ 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 "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f1;" |  $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f2;" |  $BENDSQL_CLIENT_CONNECT
@@ -126,3 +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
diff --git a/tests/suites/1_stateful/11_rbac/20_0012_privilege_access.sh b/tests/suites/1_stateful/11_rbac/20_0012_privilege_access.sh
@@ -6,6 +6,8 @@ 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 "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
 echo "drop role if exists 'test-role1'" | $BENDSQL_CLIENT_CONNECT
 echo "drop role if exists 'test-role2'" | $BENDSQL_CLIENT_CONNECT
@@ -176,3 +178,5 @@ echo "select count(1) from information_schema.tables where table_schema not in (
 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