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
17pub trait StoredRegroleResolver {
18    fn resolve_stored_regrole(&self, name: &str) -> Result<Option<i64>, SQLError>;
19}
20impl<T: EngineHook + ?Sized> StoredRegroleResolver for T {
21    fn resolve_stored_regrole(&self, name: &str) -> Result<Option<i64>, SQLError> {
22        EngineHook::resolve_regrole(self, name)
23    }
24}
25
26fn scalar_regrole_type(ty: &ColumnType) -> bool {
27    match ty {
28        ColumnType::Regrole => true,
29        ColumnType::Domain { base, .. } => scalar_regrole_type(base),
30        _ => false,
31    }
32}
33
34fn scalar_regrole_type_name(name: &str) -> bool {
35    ColumnType::from_sql_name(name).is_ok_and(|ty| scalar_regrole_type(&ty))
36}
37
38fn regrole_constant_error() -> SQLError {
39    SQLError::Routine {
40        sqlstate: "0A000".into(),
41        message: "constant of the type regrole cannot be used here".into(),
42    }
43}
44
45fn scalar_regrole_literal(expression: &ScalarExpr) -> Option<&str> {
46    let ScalarExpr::Cast { expr, ty } = expression else {
47        return None;
48    };
49    if !scalar_regrole_type_name(ty) {
50        return None;
51    }
52    match expr.as_ref() {
53        ScalarExpr::Literal(Value::Str(input)) => Some(input),
54        _ => None,
55    }
56}
57
58#[derive(Default)]
59pub struct StoredRegroleConstants {
60    inputs: Vec<String>,
61}
62
63impl StoredRegroleConstants {
64    pub fn collect_expression(
65        &mut self,
66        expression: &Expr,
67        assignment_target: Option<&ColumnType>,
68    ) {
69        if assignment_target.is_some_and(scalar_regrole_type) {
70            if let Expr::Literal(Value::Str(input)) = expression {
71                self.inputs.push(input.clone());
72            }
73        }
74        let scalar = crate::plan::ExpressionPlan::lower(expression.clone()).scalar;
75        scalar.visit(&mut |node| {
76            if let Some(input) = scalar_regrole_literal(node) {
77                self.inputs.push(input.to_string());
78            }
79        });
80    }
81
82    pub fn collect_query_plan(&mut self, plan: &mut QueryPlan) {
83        plan.rewrite_scalar_expressions(&mut |expression| {
84            if let Some(input) = scalar_regrole_literal(expression) {
85                self.inputs.push(input.to_string());
86            }
87        });
88    }
89
90    pub fn collect_plan(&mut self, plan: &mut UnifiedPlan) {
91        plan.rewrite_scalar_expressions(&mut |expression| {
92            if let Some(input) = scalar_regrole_literal(expression) {
93                self.inputs.push(input.to_string());
94            }
95        });
96    }
97
98    pub fn validate_inputs(&self, context: &dyn EngineHook) -> Result<(), SQLError> {
99        self.validate_inputs_with(context)
100    }
101
102    pub fn validate_inputs_with<C: StoredRegroleResolver + ?Sized>(
103        &self,
104        context: &C,
105    ) -> Result<(), SQLError> {
106        for input in &self.inputs {
107            context.resolve_stored_regrole(input)?;
108        }
109        Ok(())
110    }
111
112    pub fn reject(&self, context: &dyn EngineHook) -> Result<(), SQLError> {
113        self.reject_with(context)
114    }
115
116    pub fn reject_with<C: StoredRegroleResolver + ?Sized>(
117        &self,
118        context: &C,
119    ) -> Result<(), SQLError> {
120        self.validate_inputs_with(context)?;
121        if self.inputs.is_empty() {
122            Ok(())
123        } else {
124            Err(regrole_constant_error())
125        }
126    }
127}
128
129pub fn reject_stored_regrole_constants(
130    context: &dyn EngineHook,
131    expression: &Expr,
132    assignment_target: Option<&ColumnType>,
133) -> Result<(), SQLError> {
134    reject_stored_regrole_constants_with(context, expression, assignment_target)
135}
136
137pub fn reject_stored_regrole_constants_with<C: StoredRegroleResolver + ?Sized>(
138    context: &C,
139    expression: &Expr,
140    assignment_target: Option<&ColumnType>,
141) -> Result<(), SQLError> {
142    let mut constants = StoredRegroleConstants::default();
143    constants.collect_expression(expression, assignment_target);
144    constants.reject_with(context)
145}
146
147pub fn reject_stored_query_regrole_constants(
148    context: &dyn EngineHook,
149    plan: &mut QueryPlan,
150) -> Result<(), SQLError> {
151    let mut constants = StoredRegroleConstants::default();
152    constants.collect_query_plan(plan);
153    constants.reject(context)
154}
155
156pub fn reject_stored_plan_regrole_constants(
157    context: &dyn EngineHook,
158    plan: &mut UnifiedPlan,
159) -> Result<(), SQLError> {
160    reject_stored_plan_regrole_constants_with(context, plan)
161}
162
163pub fn reject_stored_plan_regrole_constants_with<C: StoredRegroleResolver + ?Sized>(
164    context: &C,
165    plan: &mut UnifiedPlan,
166) -> Result<(), SQLError> {
167    let mut constants = StoredRegroleConstants::default();
168    constants.collect_plan(plan);
169    constants.reject_with(context)
170}