chore(query): add function current_available_roles and current_secondary_roles (#18290)

Author: TCeason

Cargo.lock

@@ -214,6 +214,7 @@ pub trait TableContext: Send + Sync {
     fn get_current_database(&self) -> String;
     fn get_current_user(&self) -> Result<UserInfo>;
     fn get_current_role(&self) -> Option<RoleInfo>;
+    fn get_secondary_roles(&self) -> Option<Vec<String>>;
     fn get_current_session_id(&self) -> String {
         unimplemented!()
     }

src/query/catalog/src/table_context.rs

@@ -881,6 +881,11 @@ impl TableContext for QueryContext {
     fn get_current_role(&self) -> Option<RoleInfo> {
         self.shared.get_current_role()
     }
+
+    fn get_secondary_roles(&self) -> Option<Vec<String>> {
+        self.shared.get_secondary_roles()
+    }
+
     async fn get_all_available_roles(&self) -> Result<Vec<RoleInfo>> {
         self.get_current_session().get_all_available_roles().await
     }

src/query/service/src/sessions/query_ctx.rs

@@ -434,6 +434,10 @@ impl QueryContextShared {
         self.session.get_current_role()
     }
 
+    pub fn get_secondary_roles(&self) -> Option<Vec<String>> {
+        self.session.get_secondary_roles()
+    }
+
     /// Get all tables that already attached in this query.
     pub fn get_tables_refs(&self) -> Vec<Arc<dyn Table>> {
         let tables = self.tables_refs.lock();

src/query/service/src/sessions/query_ctx_shared.rs

@@ -661,6 +661,9 @@ impl TableContext for CtxDelegation {
     fn get_current_role(&self) -> Option<RoleInfo> {
         todo!()
     }
+    fn get_secondary_roles(&self) -> Option<Vec<String>> {
+        todo!()
+    }
     async fn get_all_available_roles(&self) -> Result<Vec<RoleInfo>> {
         todo!()
     }

src/query/service/tests/it/sql/exec/get_table_bind_test.rs

@@ -33,7 +33,7 @@ use databend_query::sessions::QueryContext;
 use databend_query::test_kits::TestFixture;
 use parking_lot::RwLock;
 
-#[tokio::test]
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn test_equivalent_constants() -> Result<()> {
     let fixture = TestFixture::setup().await?;
     let ctx = fixture.new_query_ctx().await?;

src/query/service/tests/it/sql/planner/optimizer/optimizers/operator/filter/equivalent_constants_visitor.rs

@@ -573,6 +573,9 @@ impl TableContext for CtxDelegation {
     fn get_current_role(&self) -> Option<RoleInfo> {
         todo!()
     }
+    fn get_secondary_roles(&self) -> Option<Vec<String>> {
+        todo!()
+    }
     async fn get_all_available_roles(&self) -> Result<Vec<RoleInfo>> {
         todo!()
     }

src/query/service/tests/it/storages/fuse/operations/commit.rs

@@ -67,6 +67,7 @@ recursive = { workspace = true }
 regex = { workspace = true }
 roaring = { workspace = true }
 serde = { workspace = true }
+serde_json = { workspace = true }
 sha2 = { workspace = true }
 similar = { workspace = true }
 simsearch = { workspace = true }

src/query/sql/Cargo.toml

@@ -118,6 +118,8 @@ use itertools::Itertools;
 use jsonb::keypath::parse_key_paths;
 use jsonb::keypath::KeyPath;
 use jsonb::keypath::KeyPaths;
+use serde_json::json;
+use serde_json::to_string;
 use simsearch::SimSearch;
 use unicase::Ascii;
 
@@ -2871,9 +2873,9 @@ impl<'a> TypeChecker<'a> {
         arguments: &[&Expr],
     ) -> Result<Box<(ScalarExpr, DataType)>> {
         // Check if current function is a virtual function, e.g. `database`, `version`
-        if let Some(rewritten_func_result) =
-            self.try_rewrite_sugar_function(span, func_name, arguments)
-        {
+        if let Some(rewritten_func_result) = databend_common_base::runtime::block_on(
+            self.try_rewrite_sugar_function(span, func_name, arguments),
+        ) {
             return rewritten_func_result;
         }
 
@@ -3610,6 +3612,8 @@ impl<'a> TypeChecker<'a> {
             Ascii::new("currentuser"),
             Ascii::new("current_user"),
             Ascii::new("current_role"),
+            Ascii::new("current_secondary_roles"),
+            Ascii::new("current_available_roles"),
             Ascii::new("connection_id"),
             Ascii::new("timezone"),
             Ascii::new("nullif"),
@@ -3637,7 +3641,7 @@ impl<'a> TypeChecker<'a> {
         FUNCTIONS
     }
 
-    fn try_rewrite_sugar_function(
+    async fn try_rewrite_sugar_function(
         &mut self,
         span: Span,
         func_name: &str,
@@ -3676,6 +3680,60 @@ impl<'a> TypeChecker<'a> {
                     ),
                 }),
             ),
+            ("current_secondary_roles", &[]) => {
+                let mut res = self
+                    .ctx
+                    .get_all_effective_roles()
+                    .await
+                    .unwrap_or_default()
+                    .iter()
+                    .map(|r| r.name.clone())
+                    .collect::<Vec<String>>();
+                res.sort();
+                let roles_comma_separated_string = res.iter().join(",");
+                let res = if self.ctx.get_secondary_roles().is_none() {
+                    json!({
+                        "roles": roles_comma_separated_string,
+                        "value": "ALL"
+                    })
+                } else {
+                    json!({
+                        "roles": roles_comma_separated_string,
+                        "value": "None"
+                    })
+                };
+                match to_string(&res) {
+                    Ok(res) => Some(self.resolve(&Expr::Literal {
+                        span,
+                        value: Literal::String(res),
+                    })),
+                    Err(e) => Some(Err(ErrorCode::IllegalRole(format!(
+                        "Failed to serialize secondary roles into JSON string: {}",
+                        e
+                    )))),
+                }
+            }
+            ("current_available_roles", &[]) => {
+                let mut res = self
+                    .ctx
+                    .get_all_available_roles()
+                    .await
+                    .unwrap_or_default()
+                    .iter()
+                    .map(|r| r.name.clone())
+                    .collect::<Vec<String>>();
+                res.sort();
+                match to_string(&res) {
+                    Ok(res) => Some(self.resolve(&Expr::Literal {
+                        span,
+                        value: Literal::String(res),
+                    })),
+                    Err(e) => Some(Err(ErrorCode::IllegalRole(format!(
+                        "Failed to serialize available roles into JSON string: {}",
+                        e
+                    )))),
+                }
+            }
             ("connection_id", &[]) => Some(self.resolve(&Expr::Literal {
                 span,
                 value: Literal::String(self.ctx.get_connection_id()),

src/query/sql/src/planner/semantic/type_check.rs

@@ -14,9 +14,13 @@ Error: APIError: QueryFailed: [1063]Permission denied: privilege [Insert] is req
 0
 -- test 5: set role as testrole3, secondary roles as NONE, can access table2, can not access table1, because role3 inherited from role2
 Error: APIError: QueryFailed: [1063]Permission denied: privilege [Insert] is required on 'default'.'default'.'t20_0015_table1' for user 'testuser1'@'%' with roles [testrole3,public,testrole2]
+"{""roles"":""public"",""value"":""None""}"
+"[""public"",""testrole1"",""testrole2"",""testrole3""]"
 0
 -- test 6: set role as testrole1, secondary roles as ALL, can access both table1 and table2
 0
+"{""roles"":""public,testrole1,testrole2,testrole3"",""value"":""ALL""}"
+"[""public"",""testrole1"",""testrole2"",""testrole3""]"
 0
 -- test 7: set role as testrole1, testrole2, secondary roles defaults as ALL, can both table1 and table2
 0