feat: add explicit FLUSH PRIVILEGES to refresh role cache for query node

Author: camilesing

src/meta/app/src/tenant/tenant.rs

@@ -24,7 +24,7 @@ use crate::app_error::TenantIsEmpty;
 /// Tenant is not stored directly in meta-store.
 ///
 /// It is just a type for use on the client side.
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
 pub struct Tenant {
     // TODO: consider using NonEmptyString?
     pub tenant: String,

src/query/ast/src/ast/statements/system_action.rs

@@ -32,6 +32,7 @@ impl Display for SystemStmt {
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub enum SystemAction {
     Backtrace(bool),
+    FlushPrivileges,
 }
 
 impl Display for SystemAction {
@@ -41,6 +42,7 @@ impl Display for SystemAction {
                 true => write!(f, "ENABLE EXCEPTION_BACKTRACE"),
                 false => write!(f, "DISABLE EXCEPTION_BACKTRACE"),
             },
+            SystemAction::FlushPrivileges => write!(f, "FLUSH PRIVILEGES"),
         }
     }
 }

src/query/ast/src/parser/statement.rs

@@ -5054,15 +5054,21 @@ pub fn priority(i: Input) -> IResult<Priority> {
 }
 
 pub fn action(i: Input) -> IResult<SystemAction> {
-    let mut backtrace = parser_fn(map(
+    let backtrace = parser_fn(map(
         rule! {
              #switch ~ EXCEPTION_BACKTRACE
         },
         |(switch, _)| SystemAction::Backtrace(switch),
     ));
+    let flush_privileges = parser_fn(map(
+        rule! {
+             FLUSH ~ PRIVILEGES
+        },
+        |_| SystemAction::FlushPrivileges,
+    ));
     // add other system action type here
     rule!(
-        #backtrace
+        #backtrace | #flush_privileges
     )
     .parse(i)
 }

src/query/ast/src/parser/token.rs

@@ -683,6 +683,8 @@ pub enum TokenKind {
     FORMAT_NAME,
     #[token("FORMATS", ignore(ascii_case))]
     FORMATS,
+    #[token("FLUSH", ignore(ascii_case))]
+    FLUSH,
     #[token("FRAGMENTS", ignore(ascii_case))]
     FRAGMENTS,
     #[token("FRIDAY", ignore(ascii_case))]

src/query/service/src/interpreters/interpreter_system_action.rs

@@ -20,6 +20,7 @@ use databend_common_exception::set_backtrace;
 use databend_common_exception::Result;
 use databend_common_sql::plans::SystemAction;
 use databend_common_sql::plans::SystemPlan;
+use databend_common_users::RoleCacheManager;
 
 use crate::clusters::ClusterHelper;
 use crate::clusters::FlightParams;
@@ -90,6 +91,10 @@ impl Interpreter for SystemActionInterpreter {
             SystemAction::Backtrace(switch) => {
                 set_backtrace(switch);
             }
+            SystemAction::FlushPrivileges => {
+                let tenant = self.ctx.get_tenant();
+                RoleCacheManager::instance().force_reload(&tenant).await?;
+            }
         }
         Ok(PipelineBuildResult::create())
     }

src/query/service/src/servers/flight/v1/actions/kill_query.rs

@@ -24,7 +24,7 @@ use crate::servers::flight::v1::actions::create_session;
 pub static KILL_QUERY: &str = "/actions/kill_query";
 
 pub async fn kill_query(plan: KillPlan) -> Result<bool> {
-    let session = create_session()?;
+    let session = create_session(None)?;
     let version = GlobalConfig::version();
     let query_context = session.create_query_context(version).await?;
     let interpreter = KillInterpreter::from_flight(query_context, plan)?;

src/query/service/src/servers/flight/v1/actions/mod.rs

@@ -27,6 +27,7 @@ use std::sync::Arc;
 use databend_common_catalog::session_type::SessionType;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
+use databend_common_meta_app::tenant::Tenant;
 use databend_common_settings::Settings;
 pub use flight_actions::flight_actions;
 pub use flight_actions::FlightActions;
@@ -43,9 +44,10 @@ pub use truncate_table::TRUNCATE_TABLE;
 use crate::sessions::Session;
 use crate::sessions::SessionManager;
 
-pub(crate) fn create_session() -> Result<Arc<Session>> {
+pub(crate) fn create_session(tenant: Option<Tenant>) -> Result<Arc<Session>> {
     let config = GlobalConfig::instance();
-    let settings = Settings::create(config.query.tenant_id.clone());
+    let tenant_id = tenant.unwrap_or(config.query.tenant_id.clone());
+    let settings = Settings::create(tenant_id.clone());
     match SessionManager::instance().create_with_settings(SessionType::FlightRPC, settings, None) {
         Err(cause) => Err(cause),
         Ok(session) => Ok(Arc::new(session)),

src/query/service/src/servers/flight/v1/actions/set_priority.rs

@@ -24,7 +24,7 @@ use crate::servers::flight::v1::actions::create_session;
 pub static SET_PRIORITY: &str = "/actions/set_priority";
 
 pub async fn set_priority(plan: SetPriorityPlan) -> Result<bool> {
-    let session = create_session()?;
+    let session = create_session(None)?;
     let version = GlobalConfig::version();
     let query_context = session.create_query_context(version).await?;
     let interpreter = SetPriorityInterpreter::from_flight(query_context, plan)?;

src/query/service/src/servers/flight/v1/actions/system_action.rs

@@ -23,7 +23,7 @@ use crate::servers::flight::v1::actions::create_session;
 pub static SYSTEM_ACTION: &str = "/actions/system_action";
 
 pub async fn system_action(plan: SystemPlan) -> Result<()> {
-    let session = create_session()?;
+    let session = create_session(Some(plan.clone().tenant))?;
     let version = GlobalConfig::version();
     let query_context = session.create_query_context(version).await?;
     let interpreter = SystemActionInterpreter::from_flight(query_context, plan)?;

src/query/service/src/servers/flight/v1/actions/truncate_table.rs

@@ -23,7 +23,7 @@ use crate::servers::flight::v1::actions::create_session;
 pub static TRUNCATE_TABLE: &str = "/actions/truncate_table";
 
 pub async fn truncate_table(plan: TruncateTablePlan) -> Result<()> {
-    let session = create_session()?;
+    let session = create_session(None)?;
     let version = GlobalConfig::version();
     let query_context = session.create_query_context(version).await?;
     let interpreter = TruncateTableInterpreter::from_flight(query_context, plan)?;