feat(query): support global network policy (#17050)

* feat: add scope for default settings

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

src/meta/app/src/principal/user_info.rs

@@ -102,6 +102,10 @@ impl UserInfo {
         }
     }
 
+    pub fn is_account_admin(&self) -> bool {
+        self.grants.roles().contains(&"account_admin".to_string())
+    }
+
     pub fn has_option_flag(&self, flag: UserOptionFlag) -> bool {
         self.option.has_option_flag(flag)
     }

src/query/service/src/auth.rs

@@ -97,6 +97,10 @@ impl AuthMgr {
         need_user_info: bool,
     ) -> Result<(String, Option<String>)> {
         let user_api = UserApiProvider::instance();
+        let global_network_policy = session
+            .get_settings()
+            .get_network_policy()
+            .unwrap_or_default();
         match credential {
             Credential::NoNeed => Ok(("".to_string(), None)),
             Credential::DatabendToken { token } => {
@@ -169,6 +173,17 @@ impl AuthMgr {
                     }
                 };
 
+                // check global network policy if user is not account admin
+                if !user.is_account_admin() && !global_network_policy.is_empty() {
+                    user_api
+                        .enforce_network_policy(
+                            &tenant,
+                            &global_network_policy,
+                            client_ip.as_deref(),
+                        )
+                        .await?;
+                }
+
                 session.set_authed_user(user, jwt.custom.role).await?;
                 Ok((user_name, None))
             }
@@ -182,6 +197,18 @@ impl AuthMgr {
                 let mut user = user_api
                     .get_user_with_client_ip(&tenant, identity.clone(), client_ip.as_deref())
                     .await?;
+
+                // check global network policy if user is not account admin
+                if !user.is_account_admin() && !global_network_policy.is_empty() {
+                    user_api
+                        .enforce_network_policy(
+                            &tenant,
+                            &global_network_policy,
+                            client_ip.as_deref(),
+                        )
+                        .await?;
+                }
+
                 // Check password policy for login
                 let need_change = UserApiProvider::instance()
                     .check_login_password(&tenant, identity.clone(), &user)

src/query/service/src/interpreters/interpreter_network_policy_drop.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_sql::plans::DropNetworkPolicyPlan;
 use databend_common_users::UserApiProvider;
@@ -54,6 +55,18 @@ impl Interpreter for DropNetworkPolicyInterpreter {
         let plan = self.plan.clone();
         let tenant = self.ctx.get_tenant();
 
+        let global_network_policy = self
+            .ctx
+            .get_settings()
+            .get_network_policy()
+            .unwrap_or_default();
+        if global_network_policy == plan.name {
+            return Err(ErrorCode::NetworkPolicyIsUsedByUser(format!(
+                "network policy `{}` is global network policy, can't be dropped",
+                global_network_policy,
+            )));
+        }
+
         let user_mgr = UserApiProvider::instance();
         user_mgr
             .drop_network_policy(&tenant, plan.name.as_str(), plan.if_exists)

src/query/service/src/interpreters/interpreter_set.rs

@@ -84,6 +84,16 @@ impl SetInterpreter {
                         .await?;
                     true
                 }
+                "network_policy" => {
+                    // check if the network policy exists
+                    let tenant = self.ctx.get_tenant();
+                    let _ = UserApiProvider::instance()
+                        .get_network_policy(&tenant, scalar)
+                        .await?;
+                    self.set_settings(var.to_string(), scalar.clone(), is_global)
+                        .await?;
+                    true
+                }
                 // TODO: if account_admin is built-in meta in future, we need process set sandbox_tenant in there.
                 // Like: https://github.com/datafuselabs/databend/pull/14451/files#diff-a26c9dfc9c0a37f5efa19e2b16006732b9023f42ee47cbe37fe461fb46b9dfc0R82-R85
                 "sandbox_tenant" => {

src/query/service/src/servers/flight_sql/flight_sql_service/session.rs

@@ -89,6 +89,7 @@ impl FlightSqlServiceImpl {
         password: String,
         client_ip: Option<&str>,
     ) -> Result<Arc<Session>, Status> {
+        let user_api = UserApiProvider::instance();
         let session_manager = SessionManager::instance();
         let session = session_manager
             .create_session(SessionType::FlightSQL)
@@ -100,12 +101,26 @@ impl FlightSqlServiceImpl {
         let tenant = session.get_current_tenant();
 
         let identity = UserIdentity::new(&user, "%");
-        let mut user = UserApiProvider::instance()
+        let mut user = user_api
             .get_user_with_client_ip(&tenant, identity.clone(), client_ip)
             .await
             .map_err(|e| status!("get_user fail {}", e))?;
+
+        // check global network policy if user is not account admin
+        if !user.is_account_admin() {
+            let global_network_policy = session
+                .get_settings()
+                .get_network_policy()
+                .unwrap_or_default();
+            if !global_network_policy.is_empty() {
+                user_api
+                    .enforce_network_policy(&tenant, &global_network_policy, client_ip)
+                    .await?;
+            }
+        }
+
         // Check password policy for login
-        let need_change = UserApiProvider::instance()
+        let need_change = user_api
             .check_login_password(&tenant, identity.clone(), &user)
             .await
             .map_err(|e| status!("not compliant with password policy {}", e))?;
@@ -136,7 +151,7 @@ impl FlightSqlServiceImpl {
             _ => Err(Status::unauthenticated("wrong auth type")),
         };
 
-        UserApiProvider::instance()
+        user_api
             .update_user_login_result(tenant, identity, authed.is_ok(), &user)
             .await?;
         authed?;

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -269,24 +269,36 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for InteractiveWorke
 impl InteractiveWorkerBase {
     #[async_backtrace::framed]
     async fn authenticate(&self, salt: &[u8], info: CertifiedInfo) -> Result<bool> {
+        let user_api = UserApiProvider::instance();
         let ctx = self.session.create_query_context().await?;
+        let tenant = ctx.get_tenant();
         let identity = UserIdentity::new(&info.user_name, "%");
         let client_ip = info.user_client_address.split(':').collect::<Vec<_>>()[0];
-        let mut user = UserApiProvider::instance()
-            .get_user_with_client_ip(&ctx.get_tenant(), identity.clone(), Some(client_ip))
+        let mut user = user_api
+            .get_user_with_client_ip(&tenant, identity.clone(), Some(client_ip))
             .await?;
 
+        // check global network policy if user is not account admin
+        if !user.is_account_admin() {
+            let global_network_policy = ctx.get_settings().get_network_policy().unwrap_or_default();
+            if !global_network_policy.is_empty() {
+                user_api
+                    .enforce_network_policy(&tenant, &global_network_policy, Some(client_ip))
+                    .await?;
+            }
+        }
+
         // Check password policy for login
-        let need_change = UserApiProvider::instance()
-            .check_login_password(&ctx.get_tenant(), identity.clone(), &user)
+        let need_change = user_api
+            .check_login_password(&tenant, identity.clone(), &user)
             .await?;
         if need_change {
             user.update_auth_need_change_password();
         }
 
         let authed = user.auth_info.auth_mysql(&info.user_password, salt)?;
-        UserApiProvider::instance()
-            .update_user_login_result(ctx.get_tenant(), identity, authed, &user)
+        user_api
+            .update_user_login_result(tenant, identity, authed, &user)
             .await?;
         if authed {
             self.session.set_authed_user(user, None).await?;

src/query/service/src/sessions/session_mgr.rs

@@ -116,24 +116,6 @@ impl SessionManager {
         Ok(())
     }
 
-    pub fn load_config_changes(&self, settings: &Arc<Settings>) -> Result<()> {
-        let query_config = &GlobalConfig::instance().query;
-        if let Some(parquet_fast_read_bytes) = query_config.parquet_fast_read_bytes {
-            settings.set_parquet_fast_read_bytes(parquet_fast_read_bytes)?;
-        }
-
-        if let Some(max_storage_io_requests) = query_config.max_storage_io_requests {
-            settings.set_max_storage_io_requests(max_storage_io_requests)?;
-        }
-
-        if let Some(enterprise_license_key) = query_config.databend_enterprise_license.clone() {
-            unsafe {
-                settings.set_enterprise_license(enterprise_license_key)?;
-            }
-        }
-        Ok(())
-    }
-
     pub fn create_with_settings(
         &self,
         typ: SessionType,

src/query/settings/src/lib.rs

@@ -23,4 +23,5 @@ pub use settings::Settings;
 pub use settings_default::ReplaceIntoShuffleStrategy;
 pub use settings_default::SettingMode;
 pub use settings_default::SettingRange;
+pub use settings_default::SettingScope;
 pub use settings_getter_setter::FlightCompression;