Skip to main content

uqa_sql/catalog/
regrole_dependencies.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `PostgreSQL` dependency restrictions for stored `regrole` constants.
8
9use crate::ast::{ColumnType, Expr};
10use crate::plan::{QueryPlan, UnifiedPlan};
11use crate::SQLError;
12use crate::ScalarExpr;
13use uqa_core::Value;
14
15use crate::expr::EngineHook;
16
17#[cfg(test)]
18mod tests;
19
20pub trait StoredRegroleResolver {
21    fn resolve_stored_regrole(&self, name: &str) -> Result<Option<i64>, SQLError>;
22}
23impl<T: EngineHook + ?Sized> StoredRegroleResolver for T {
24    fn resolve_stored_regrole(&self, name: &str) -> Result<Option<i64>, SQLError> {
25        EngineHook::resolve_regrole(self, name)
26    }
27}
28
29fn scalar_regrole_type(ty: &ColumnType) -> bool {
30    match ty {
31        ColumnType::Regrole => true,
32        ColumnType::Domain { base, .. } => scalar_regrole_type(base),
33        _ => false,
34    }
35}
36
37fn scalar_regrole_type_name(name: &str) -> bool {
38    ColumnType::from_sql_name(name).is_ok_and(|ty| scalar_regrole_type(&ty))
39}
40
41fn regrole_constant_error() -> SQLError {
42    SQLError::Routine {
43        sqlstate: "0A000".into(),
44        message: "constant of the type regrole cannot be used here".into(),
45    }
46}
47
48fn scalar_regrole_literal(expression: &ScalarExpr) -> Option<&str> {
49    let ScalarExpr::Cast { expr, ty, .. } = expression else {
50        return None;
51    };
52    if !scalar_regrole_type_name(ty) {
53        return None;
54    }
55    match expr.as_ref() {
56        ScalarExpr::Literal(Value::Str(input)) => Some(input),
57        _ => None,
58    }
59}
60
61#[derive(Default)]
62pub struct StoredRegroleConstants {
63    inputs: Vec<String>,
64    has_bound_constant: bool,
65}
66
67impl StoredRegroleConstants {
68    pub fn collect_expression(
69        &mut self,
70        expression: &Expr,
71        assignment_target: Option<&ColumnType>,
72    ) {
73        if assignment_target.is_some_and(scalar_regrole_type) {
74            if let Expr::Literal(Value::Str(input)) = expression {
75                self.inputs.push(input.clone());
76            }
77        }
78        let scalar = crate::plan::ExpressionPlan::lower(expression.clone()).scalar;
79        self.collect_scalar_expression(&scalar);
80    }
81
82    /// Keep dependency identity before a caller replaces an expression by its optimized form.
83    pub fn collect_scalar_expression(&mut self, expression: &ScalarExpr) {
84        expression.visit(&mut |node| self.collect_scalar(node));
85    }
86
87    fn collect_scalar(&mut self, expression: &ScalarExpr) {
88        if let Some(input) = scalar_regrole_literal(expression) {
89            self.inputs.push(input.to_string());
90        }
91        if let ScalarExpr::TypedLiteral {
92            value,
93            ty,
94            bound_type,
95            ..
96        } = expression
97        {
98            self.has_bound_constant |= !matches!(value, Value::Null)
99                && (bound_type.as_ref().is_some_and(scalar_regrole_type)
100                    || scalar_regrole_type_name(ty));
101        }
102    }
103
104    pub fn collect_query_plan(&mut self, plan: &mut QueryPlan) {
105        plan.rewrite_scalar_expressions(&mut |expression| self.collect_scalar(expression));
106    }
107
108    pub fn collect_plan(&mut self, plan: &mut UnifiedPlan) {
109        plan.rewrite_scalar_expressions(&mut |expression| self.collect_scalar(expression));
110    }
111
112    pub fn validate_inputs(&self, context: &dyn EngineHook) -> Result<(), SQLError> {
113        self.validate_inputs_with(context)
114    }
115
116    pub fn validate_inputs_with<C: StoredRegroleResolver + ?Sized>(
117        &self,
118        context: &C,
119    ) -> Result<(), SQLError> {
120        for input in &self.inputs {
121            context.resolve_stored_regrole(input)?;
122        }
123        Ok(())
124    }
125
126    pub fn reject(&self, context: &dyn EngineHook) -> Result<(), SQLError> {
127        self.reject_with(context)
128    }
129
130    pub fn reject_with<C: StoredRegroleResolver + ?Sized>(
131        &self,
132        context: &C,
133    ) -> Result<(), SQLError> {
134        self.validate_inputs_with(context)?;
135        if self.inputs.is_empty() && !self.has_bound_constant {
136            Ok(())
137        } else {
138            Err(regrole_constant_error())
139        }
140    }
141}
142
143pub fn reject_stored_regrole_constants(
144    context: &dyn EngineHook,
145    expression: &Expr,
146    assignment_target: Option<&ColumnType>,
147) -> Result<(), SQLError> {
148    reject_stored_regrole_constants_with(context, expression, assignment_target)
149}
150
151pub fn reject_stored_regrole_constants_with<C: StoredRegroleResolver + ?Sized>(
152    context: &C,
153    expression: &Expr,
154    assignment_target: Option<&ColumnType>,
155) -> Result<(), SQLError> {
156    let mut constants = StoredRegroleConstants::default();
157    constants.collect_expression(expression, assignment_target);
158    constants.reject_with(context)
159}
160
161pub fn reject_stored_query_regrole_constants(
162    context: &dyn EngineHook,
163    plan: &mut QueryPlan,
164) -> Result<(), SQLError> {
165    let mut constants = StoredRegroleConstants::default();
166    constants.collect_query_plan(plan);
167    constants.reject(context)
168}
169
170pub fn reject_stored_plan_regrole_constants(
171    context: &dyn EngineHook,
172    plan: &mut UnifiedPlan,
173) -> Result<(), SQLError> {
174    reject_stored_plan_regrole_constants_with(context, plan)
175}
176
177pub fn reject_stored_plan_regrole_constants_with<C: StoredRegroleResolver + ?Sized>(
178    context: &C,
179    plan: &mut UnifiedPlan,
180) -> Result<(), SQLError> {
181    let mut constants = StoredRegroleConstants::default();
182    constants.collect_plan(plan);
183    constants.reject_with(context)
184}