feat(query): add is_role_in_session funciton (#19867)

Author: TCeason

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

@@ -54,6 +54,7 @@ impl<'a> TypeChecker<'a> {
             Ascii::new("current_role"),
             Ascii::new("current_secondary_roles"),
             Ascii::new("current_available_roles"),
+            Ascii::new("is_role_in_session"),
             Ascii::new("connection_id"),
             Ascii::new("client_session_id"),
             Ascii::new("timezone"),
@@ -181,6 +182,49 @@ impl<'a> TypeChecker<'a> {
                     )))),
                 }
             }
+            ("is_role_in_session", &[role]) => {
+                let effective_roles = match self.ctx.get_all_effective_roles().await {
+                    Ok(roles) => roles,
+                    Err(err) => return Some(Err(err)),
+                };
+                let (role_expr, _) = match self.resolve(role) {
+                    Ok(res) => *res,
+                    Err(err) => return Some(Err(err)),
+                };
+
+                let mut predicate_levels =
+                    Vec::with_capacity(effective_roles.len().max(1).ilog2() as usize + 1);
+                for effective_role in effective_roles {
+                    let role_literal = ScalarExpr::ConstantExpr(ConstantExpr {
+                        span,
+                        value: Scalar::String(effective_role.name),
+                    });
+                    let predicate =
+                        match self.resolve_scalar_function_call(span, "eq", vec![], vec![
+                            role_expr.clone(),
+                            role_literal,
+                        ]) {
+                            Ok(res) => {
+                                let (predicate, _) = *res;
+                                predicate
+                            }
+                            Err(err) => return Some(Err(err)),
+                        };
+                    if let Err(err) = self.merge_or_level(span, &mut predicate_levels, predicate) {
+                        return Some(Err(err));
+                    }
+                }
+
+                let predicate = match self.fold_or_levels(span, predicate_levels) {
+                    Ok(Some(predicate)) => predicate,
+                    Ok(None) => ScalarExpr::ConstantExpr(ConstantExpr {
+                        span,
+                        value: Scalar::Boolean(false),
+                    }),
+                    Err(err) => return Some(Err(err)),
+                };
+                Some(self.resolve_scalar_function_call(span, "is_true", vec![], vec![predicate]))
+            }
             ("connection_id", &[]) => Some(self.resolve(&Expr::Literal {
                 span,
                 value: Literal::String(self.ctx.get_connection_id()),

tests/sqllogictests/suites/ee/05_ee_ddl/05_0004_ddl_security_policy.test

@@ -1063,5 +1063,107 @@ DROP TABLE employees;
 statement ok
 DROP MASKING POLICY mask_ssn_conditional;
 
+## row access policy with active role hierarchy
+statement ok
+SET ROLE account_admin;
+
+statement ok
+SET SECONDARY ROLES ALL;
+
+statement ok
+DROP TABLE IF EXISTS rap_time_range_test;
+
+statement ok
+DROP ROW ACCESS POLICY IF EXISTS rap_time_range;
+
+statement ok
+DROP ROLE IF EXISTS rap_role_7_day;
+
+statement ok
+DROP ROLE IF EXISTS rap_role_1_day;
+
+statement ok
+CREATE ROLE rap_role_7_day;
+
+statement ok
+CREATE ROLE rap_role_1_day;
+
+statement ok
+CREATE ROW ACCESS POLICY rap_time_range AS (start_time TIMESTAMP) RETURNS boolean ->
+  CASE
+    WHEN IS_ROLE_IN_SESSION('rap_role_7_day') THEN
+      start_time BETWEEN to_timestamp('2024-01-08 00:00:00') AND to_timestamp('2024-01-15 00:00:00')
+    WHEN IS_ROLE_IN_SESSION('rap_role_1_day') THEN
+      start_time BETWEEN to_timestamp('2024-01-14 00:00:00') AND to_timestamp('2024-01-15 00:00:00')
+    ELSE false
+  END;
+
+statement ok
+CREATE TABLE rap_time_range_test(id INT, start_time TIMESTAMP);
+
+statement ok
+INSERT INTO rap_time_range_test VALUES
+  (1, to_timestamp('2024-01-07 00:00:00')),
+  (2, to_timestamp('2024-01-09 00:00:00')),
+  (3, to_timestamp('2024-01-14 12:00:00')),
+  (4, to_timestamp('2024-01-15 00:00:00')),
+  (5, to_timestamp('2024-01-16 00:00:00'));
+
+statement ok
+ALTER TABLE rap_time_range_test ADD ROW ACCESS POLICY rap_time_range ON (start_time);
+
+statement ok
+SET ROLE rap_role_1_day;
+
+statement ok
+SET SECONDARY ROLES NONE;
+
+query I
+SELECT id FROM rap_time_range_test ORDER BY id;
+----
+3
+4
+
+statement ok
+SET SECONDARY ROLES rap_role_7_day;
+
+query I
+SELECT id FROM rap_time_range_test ORDER BY id;
+----
+2
+3
+4
+
+statement ok
+SET ROLE rap_role_7_day;
+
+statement ok
+SET SECONDARY ROLES NONE;
+
+query I
+SELECT id FROM rap_time_range_test ORDER BY id;
+----
+2
+3
+4
+
+statement ok
+SET ROLE account_admin;
+
+statement ok
+SET SECONDARY ROLES ALL;
+
+statement ok
+DROP TABLE rap_time_range_test;
+
+statement ok
+DROP ROW ACCESS POLICY rap_time_range;
+
+statement ok
+DROP ROLE rap_role_7_day;
+
+statement ok
+DROP ROLE rap_role_1_day;
+
 statement ok
 unset global enable_planner_cache;