use crate::ast::Expr;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum PolicyTarget {
All,
Select,
Insert,
Update,
Delete,
}
impl std::fmt::Display for PolicyTarget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PolicyTarget::All => write!(f, "ALL"),
PolicyTarget::Select => write!(f, "SELECT"),
PolicyTarget::Insert => write!(f, "INSERT"),
PolicyTarget::Update => write!(f, "UPDATE"),
PolicyTarget::Delete => write!(f, "DELETE"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum PolicyPermissiveness {
Permissive,
Restrictive,
}
impl std::fmt::Display for PolicyPermissiveness {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PolicyPermissiveness::Permissive => write!(f, "PERMISSIVE"),
PolicyPermissiveness::Restrictive => write!(f, "RESTRICTIVE"),
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct RlsPolicy {
pub name: String,
pub table: String,
pub target: PolicyTarget,
pub permissiveness: PolicyPermissiveness,
pub using: Option<Expr>,
pub with_check: Option<Expr>,
pub role: Option<String>,
}
impl RlsPolicy {
pub fn create(name: impl Into<String>, table: impl Into<String>) -> Self {
Self {
name: name.into(),
table: table.into(),
target: PolicyTarget::All,
permissiveness: PolicyPermissiveness::Permissive,
using: None,
with_check: None,
role: None,
}
}
pub fn for_all(mut self) -> Self {
self.target = PolicyTarget::All;
self
}
pub fn for_select(mut self) -> Self {
self.target = PolicyTarget::Select;
self
}
pub fn for_insert(mut self) -> Self {
self.target = PolicyTarget::Insert;
self
}
pub fn for_update(mut self) -> Self {
self.target = PolicyTarget::Update;
self
}
pub fn for_delete(mut self) -> Self {
self.target = PolicyTarget::Delete;
self
}
pub fn restrictive(mut self) -> Self {
self.permissiveness = PolicyPermissiveness::Restrictive;
self
}
pub fn using(mut self, expr: Expr) -> Self {
self.using = Some(expr);
self
}
pub fn with_check(mut self, expr: Expr) -> Self {
self.with_check = Some(expr);
self
}
pub fn to_role(mut self, role: impl Into<String>) -> Self {
self.role = Some(role.into());
self
}
}
pub fn tenant_check(
column: impl Into<String>,
session_var: impl Into<String>,
cast_type: impl Into<String>,
) -> Expr {
use crate::ast::{BinaryOp, Value};
Expr::Binary {
left: Box::new(Expr::Named(column.into())),
op: BinaryOp::Eq,
right: Box::new(Expr::Cast {
expr: Box::new(Expr::FunctionCall {
name: "current_setting".into(),
args: vec![Expr::Literal(Value::String(session_var.into()))],
alias: None,
}),
target_type: cast_type.into(),
alias: None,
}),
alias: None,
}
}
pub fn session_bool_check(session_var: impl Into<String>) -> Expr {
use crate::ast::{BinaryOp, Value};
Expr::Binary {
left: Box::new(Expr::Cast {
expr: Box::new(Expr::FunctionCall {
name: "current_setting".into(),
args: vec![Expr::Literal(Value::String(session_var.into()))],
alias: None,
}),
target_type: "boolean".into(),
alias: None,
}),
op: BinaryOp::Eq,
right: Box::new(Expr::Literal(Value::Bool(true))),
alias: None,
}
}
pub fn or(left: Expr, right: Expr) -> Expr {
use crate::ast::BinaryOp;
Expr::Binary {
left: Box::new(left),
op: BinaryOp::Or,
right: Box::new(right),
alias: None,
}
}
pub fn and(left: Expr, right: Expr) -> Expr {
use crate::ast::BinaryOp;
Expr::Binary {
left: Box::new(left),
op: BinaryOp::And,
right: Box::new(right),
alias: None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::BinaryOp;
#[test]
fn test_policy_builder() {
let policy = RlsPolicy::create("orders_isolation", "orders")
.for_all()
.using(tenant_check(
"operator_id",
"app.current_operator_id",
"uuid",
))
.with_check(tenant_check(
"operator_id",
"app.current_operator_id",
"uuid",
));
assert_eq!(policy.name, "orders_isolation");
assert_eq!(policy.table, "orders");
assert_eq!(policy.target, PolicyTarget::All);
assert!(policy.using.is_some());
assert!(policy.with_check.is_some());
}
#[test]
fn test_policy_restrictive() {
let policy = RlsPolicy::create("admin_only", "secrets")
.for_select()
.restrictive()
.to_role("app_user");
assert_eq!(policy.target, PolicyTarget::Select);
assert_eq!(policy.permissiveness, PolicyPermissiveness::Restrictive);
assert_eq!(policy.role.as_deref(), Some("app_user"));
}
#[test]
fn test_tenant_check_helper() {
let expr = tenant_check("operator_id", "app.current_operator_id", "uuid");
let Expr::Binary {
left, op, right, ..
} = &expr
else {
panic!("Expected Binary, got {expr:?}");
};
assert_eq!(*op, BinaryOp::Eq);
let Expr::Named(n) = left.as_ref() else {
panic!("Expected Named, got {left:?}");
};
assert_eq!(n, "operator_id");
let Expr::Cast {
expr: cast_expr,
target_type,
..
} = right.as_ref()
else {
panic!("Expected Cast, got {right:?}");
};
assert_eq!(target_type, "uuid");
let Expr::FunctionCall { name, args, .. } = cast_expr.as_ref() else {
panic!("Expected FunctionCall, got {cast_expr:?}");
};
assert_eq!(name, "current_setting");
assert_eq!(args.len(), 1);
}
#[test]
fn test_super_admin_bypass() {
let expr = or(
tenant_check("operator_id", "app.current_operator_id", "uuid"),
session_bool_check("app.is_super_admin"),
);
assert!(
matches!(
&expr,
Expr::Binary {
op: BinaryOp::Or,
..
}
),
"Expected Binary OR, got {expr:?}"
);
}
#[test]
fn test_and_combinator() {
let expr = and(
tenant_check("operator_id", "app.current_operator_id", "uuid"),
tenant_check("agent_id", "app.current_agent_id", "uuid"),
);
assert!(
matches!(
&expr,
Expr::Binary {
op: BinaryOp::And,
..
}
),
"Expected Binary AND, got {expr:?}"
);
}
#[test]
fn test_policy_target_display() {
assert_eq!(PolicyTarget::All.to_string(), "ALL");
assert_eq!(PolicyTarget::Select.to_string(), "SELECT");
assert_eq!(PolicyTarget::Insert.to_string(), "INSERT");
assert_eq!(PolicyTarget::Update.to_string(), "UPDATE");
assert_eq!(PolicyTarget::Delete.to_string(), "DELETE");
}
}