Skip to main content

uqa_sql/semantics/rules/action_binding/
references.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! OLD/NEW reference analysis for rewrite-rule expressions and stored plans.
8
9use std::collections::BTreeSet;
10
11use crate::ast::{Expr, SelectStmt, Statement};
12use crate::plan::ExpressionPlan;
13use crate::plpgsql::{ResolvedVariable, VariableResolver};
14use crate::SQLError;
15use crate::ScalarExpr;
16
17use super::{bind_rule_action, bind_rule_expr_scoped, bind_rule_select_scoped};
18use crate::catalog::events::RuleConditionBinding;
19
20#[derive(Default)]
21struct RuleRowReferenceDetector {
22    qualifier: Option<String>,
23    whole_row: bool,
24}
25
26#[derive(Default)]
27struct RuleRowColumnCollector {
28    columns: BTreeSet<String>,
29    qualifier: Option<&'static str>,
30    whole_row: bool,
31}
32
33impl RuleRowColumnCollector {
34    fn accepts(&self, qualifier: &str) -> bool {
35        self.qualifier.map_or_else(
36            || qualifier.eq_ignore_ascii_case("old") || qualifier.eq_ignore_ascii_case("new"),
37            |selected| selected.eq_ignore_ascii_case(qualifier),
38        )
39    }
40}
41
42impl VariableResolver for RuleRowColumnCollector {
43    fn resolve_name(&mut self, name: &str) -> Result<Option<ResolvedVariable>, SQLError> {
44        self.whole_row |= self.accepts(name);
45        Ok(None)
46    }
47
48    fn resolve_qualified(
49        &mut self,
50        qualifier: &str,
51        column: &str,
52    ) -> Result<Option<ResolvedVariable>, SQLError> {
53        if self.accepts(qualifier) {
54            self.columns.insert(column.to_string());
55        }
56        Ok(None)
57    }
58
59    fn resolve_param(&mut self, _index: usize) -> Result<Option<ResolvedVariable>, SQLError> {
60        Ok(None)
61    }
62
63    fn rewrite_qualified_whole_row(&mut self, qualifier: &str) -> Result<Option<Expr>, SQLError> {
64        self.whole_row |= self.accepts(qualifier);
65        Ok(None)
66    }
67}
68
69impl VariableResolver for RuleRowReferenceDetector {
70    fn resolve_name(&mut self, name: &str) -> Result<Option<ResolvedVariable>, SQLError> {
71        if name.eq_ignore_ascii_case("old") || name.eq_ignore_ascii_case("new") {
72            self.observe(name, true);
73        }
74        Ok(None)
75    }
76
77    fn resolve_qualified(
78        &mut self,
79        qualifier: &str,
80        _column: &str,
81    ) -> Result<Option<ResolvedVariable>, SQLError> {
82        self.observe(qualifier, false);
83        Ok(None)
84    }
85
86    fn resolve_param(&mut self, _index: usize) -> Result<Option<ResolvedVariable>, SQLError> {
87        Ok(None)
88    }
89
90    fn rewrite_qualified_whole_row(&mut self, qualifier: &str) -> Result<Option<Expr>, SQLError> {
91        self.observe(qualifier, true);
92        Ok(None)
93    }
94}
95
96impl RuleRowReferenceDetector {
97    fn observe(&mut self, qualifier: &str, whole_row: bool) {
98        if qualifier.eq_ignore_ascii_case("old") || qualifier.eq_ignore_ascii_case("new") {
99            if self.qualifier.is_none() {
100                self.qualifier = Some(qualifier.to_ascii_lowercase());
101            }
102            self.whole_row |= whole_row;
103        }
104    }
105}
106
107pub fn rule_action_has_set_operation(action: &Statement) -> bool {
108    match action {
109        Statement::Select(select) => select.set_op.is_some(),
110        Statement::Insert(insert) => insert
111            .select_source
112            .as_ref()
113            .is_some_and(|select| select.set_op.is_some()),
114        _ => false,
115    }
116}
117
118pub fn first_rule_row_reference_in_expr(
119    expr: &Expr,
120    shadowed: &BTreeSet<String>,
121) -> Option<String> {
122    let mut detector = RuleRowReferenceDetector::default();
123    let _ = bind_rule_expr_scoped(expr, &mut detector, shadowed);
124    detector.qualifier
125}
126
127pub fn first_rule_row_reference_in_select(
128    catalog: &dyn super::RuleSourceCatalog,
129    select: &SelectStmt,
130) -> Option<String> {
131    let mut detector = RuleRowReferenceDetector::default();
132    let _ = bind_rule_select_scoped(catalog, select, &mut detector);
133    detector.qualifier
134}
135
136pub fn rule_expr_references_row(expr: &Expr) -> bool {
137    first_rule_row_reference_in_expr(expr, &BTreeSet::new()).is_some()
138}
139
140pub fn rule_expr_references_whole_row(expr: &Expr) -> bool {
141    let mut detector = RuleRowReferenceDetector::default();
142    let _ = bind_rule_expr_scoped(expr, &mut detector, &BTreeSet::new());
143    detector.whole_row
144}
145
146pub fn first_rule_row_reference_in_statement(
147    catalog: &dyn super::RuleSourceCatalog,
148    statement: &Statement,
149) -> Result<Option<String>, SQLError> {
150    let target = match statement {
151        Statement::Insert(plan) => Some(plan.table.as_str()),
152        Statement::Update(plan) => Some(plan.table.as_str()),
153        Statement::Delete(plan) => Some(plan.table.as_str()),
154        Statement::Merge(plan) => Some(plan.target.as_str()),
155        _ => None,
156    };
157    let columns = target
158        .map(|table| catalog.query_source_columns(table, false))
159        .transpose()?
160        .flatten()
161        .unwrap_or_default()
162        .into_iter()
163        .collect();
164    let mut detector = RuleRowReferenceDetector::default();
165    let _ = bind_rule_action(catalog, statement, &columns, &mut detector)?;
166    Ok(detector.qualifier)
167}
168
169pub fn rule_statement_references_row(
170    catalog: &dyn super::RuleSourceCatalog,
171    statement: &Statement,
172    action_columns: &BTreeSet<String>,
173) -> Result<bool, SQLError> {
174    let mut detector = RuleRowReferenceDetector::default();
175    let _ = bind_rule_action(catalog, statement, action_columns, &mut detector)?;
176    Ok(detector.qualifier.is_some())
177}
178
179pub fn rule_statement_references_whole_row(
180    catalog: &dyn super::RuleSourceCatalog,
181    statement: &Statement,
182    action_columns: &BTreeSet<String>,
183) -> Result<bool, SQLError> {
184    let mut detector = RuleRowReferenceDetector::default();
185    let _ = bind_rule_action(catalog, statement, action_columns, &mut detector)?;
186    Ok(detector.whole_row)
187}
188
189pub fn rule_expr_row_columns(expr: &Expr) -> BTreeSet<String> {
190    let mut collector = RuleRowColumnCollector::default();
191    let _ = bind_rule_expr_scoped(expr, &mut collector, &BTreeSet::new());
192    collector.columns
193}
194
195fn collect_condition_plan_row_column(
196    expression: &ScalarExpr,
197    binding: &RuleConditionBinding,
198    columns: &mut BTreeSet<String>,
199) {
200    if let ScalarExpr::InternalColumn(column) = expression {
201        if let Some(name) = binding.column_name(*column) {
202            columns.insert(name.to_string());
203        }
204    }
205}
206
207fn condition_plan_expression_references_whole_row(expression: &ScalarExpr) -> bool {
208    matches!(
209        expression,
210        ScalarExpr::Column(qualifier) | ScalarExpr::QualifiedStar(qualifier)
211            if qualifier.eq_ignore_ascii_case("old")
212                || qualifier.eq_ignore_ascii_case("new")
213    )
214}
215
216pub fn rule_condition_plan_references_whole_row(plan: &ExpressionPlan) -> bool {
217    let mut referenced = false;
218    plan.scalar.visit(&mut |expression| {
219        referenced |= condition_plan_expression_references_whole_row(expression);
220    });
221    for subquery in &plan.subqueries {
222        let mut subquery = subquery.clone();
223        subquery.rewrite_scalar_expressions(&mut |expression| {
224            referenced |= condition_plan_expression_references_whole_row(expression);
225        });
226    }
227    referenced
228}
229
230pub fn rule_condition_plan_row_columns(
231    plan: &ExpressionPlan,
232    binding: &RuleConditionBinding,
233) -> BTreeSet<String> {
234    let mut columns = BTreeSet::new();
235    plan.scalar.visit(&mut |expression| {
236        collect_condition_plan_row_column(expression, binding, &mut columns);
237    });
238    for subquery in &plan.subqueries {
239        let mut subquery = subquery.clone();
240        subquery.rewrite_scalar_expressions(&mut |expression| {
241            collect_condition_plan_row_column(expression, binding, &mut columns);
242        });
243    }
244    columns
245}
246
247pub fn rule_statement_row_columns(
248    catalog: &dyn super::RuleSourceCatalog,
249    statement: &Statement,
250    action_columns: &BTreeSet<String>,
251) -> Result<BTreeSet<String>, SQLError> {
252    let mut collector = RuleRowColumnCollector::default();
253    let _ = bind_rule_action(catalog, statement, action_columns, &mut collector)?;
254    Ok(collector.columns)
255}
256
257/// Input assignments contribute only NEW fields; OLD references require source rows but never the corresponding new assignment expression.
258pub fn rule_new_row_columns(
259    catalog: &dyn super::RuleSourceCatalog,
260    rule: &crate::catalog::events::StoredRule,
261) -> Result<Option<BTreeSet<String>>, SQLError> {
262    let mut collector = RuleRowColumnCollector {
263        qualifier: Some("new"),
264        ..RuleRowColumnCollector::default()
265    };
266    if let Some(condition) = &rule.definition.condition {
267        bind_rule_expr_scoped(condition, &mut collector, &BTreeSet::new())?;
268    }
269    for action in &rule.definition.actions {
270        let columns = super::rule_action_target_columns(catalog, action)?;
271        bind_rule_action(catalog, action, &columns, &mut collector)?;
272    }
273    Ok((!collector.whole_row).then_some(collector.columns))
274}