uqa_planner/statement_planning/
rule_inputs.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Prune mutation inputs discarded by rewrite rules before constant evaluation.
8use crate::{
9    CommandPlan, CtePlan, CtePlanBody, QueryPlan, RelationalPlan, SourcePlan, UnifiedPlan,
10};
11use std::collections::BTreeSet;
12use uqa_core::Value;
13use uqa_sql::{
14    ast::RuleEvent,
15    semantics::{rules::RuleCatalog, view_rewrite::context::ViewRewriteContext},
16    SQLError, ScalarExpr,
17};
18
19pub trait RuleInputColumns {
20    fn source_column_names(
21        &self,
22        table: &str,
23        relations_bound: bool,
24    ) -> Result<Option<Vec<String>>, SQLError>;
25}
26pub struct RuleInputPlanningContext<'a> {
27    pub rules: &'a dyn RuleCatalog,
28    pub views: ViewRewriteContext<'a>,
29    pub columns: &'a dyn RuleInputColumns,
30}
31pub fn rewrite_plan(
32    context: &RuleInputPlanningContext<'_>,
33    plan: &mut UnifiedPlan,
34) -> Result<(), SQLError> {
35    match plan {
36        UnifiedPlan::Query(query) => rewrite_query(context, query),
37        UnifiedPlan::Command(command) => rewrite_command(context, command),
38    }
39}
40
41fn rewrite_ctes(
42    context: &RuleInputPlanningContext<'_>,
43    ctes: &mut [CtePlan],
44) -> Result<(), SQLError> {
45    for cte in ctes {
46        match &mut cte.body {
47            CtePlanBody::Query(query) => rewrite_query(context, query)?,
48            CtePlanBody::Command(command) => rewrite_command(context, command)?,
49        }
50    }
51    Ok(())
52}
53
54fn rewrite_query(
55    context: &RuleInputPlanningContext<'_>,
56    query: &mut QueryPlan,
57) -> Result<(), SQLError> {
58    rewrite_ctes(context, &mut query.ctes)?;
59    let subqueries = match &mut query.root {
60        RelationalPlan::QueryBlock(block) => {
61            if let Some(source) = &mut block.from {
62                rewrite_source(context, source)?;
63            }
64            &mut block.subqueries
65        }
66        RelationalPlan::Values { subqueries, .. } => subqueries,
67        RelationalPlan::SetOp {
68            left,
69            right,
70            subqueries,
71            ..
72        } => {
73            rewrite_query(context, left)?;
74            rewrite_query(context, right)?;
75            subqueries
76        }
77    };
78    for query in subqueries {
79        rewrite_query(context, query)?;
80    }
81    Ok(())
82}
83
84fn rewrite_source(
85    context: &RuleInputPlanningContext<'_>,
86    source: &mut SourcePlan,
87) -> Result<(), SQLError> {
88    match source {
89        SourcePlan::Subquery { body, .. } => rewrite_query(context, body)?,
90        SourcePlan::Join { left, right, .. } => {
91            rewrite_source(context, left)?;
92            rewrite_source(context, right)?;
93        }
94        _ => {}
95    }
96    Ok(())
97}
98
99fn rewrite_command(
100    context: &RuleInputPlanningContext<'_>,
101    command: &mut CommandPlan,
102) -> Result<(), SQLError> {
103    match command {
104        CommandPlan::Explain { body, .. } => return rewrite_plan(context, body),
105        CommandPlan::DeclareCursor { query, .. } => return rewrite_query(context, query),
106        _ => {}
107    }
108    prune_rule_inputs(context, command)?;
109    if let Some(ctes) = command.ctes_mut() {
110        rewrite_ctes(context, ctes)?;
111    }
112    if let Some(source) = command.source_input_mut() {
113        rewrite_source(context, source)?;
114    }
115    for query in command.query_inputs_mut() {
116        rewrite_query(context, query)?;
117    }
118    Ok(())
119}
120
121fn prune_rule_inputs(
122    context: &RuleInputPlanningContext<'_>,
123    command: &mut CommandPlan,
124) -> Result<(), SQLError> {
125    let (table, bound, event) = match command {
126        CommandPlan::Insert(plan) => (&plan.table, plan.target_relation_bound, RuleEvent::Insert),
127        CommandPlan::Update(plan) => (&plan.table, plan.target_relation_bound, RuleEvent::Update),
128        CommandPlan::Delete(plan) => (&plan.table, plan.target_relation_bound, RuleEvent::Delete),
129        _ => return Ok(()),
130    };
131    let table = context.rules.resolve_mutation_target(table, bound)?;
132    let Some(requirements) =
133        uqa_sql::semantics::view_rewrite::rule_input_requirements(context.views, &table, event)?
134    else {
135        return Ok(());
136    };
137    uqa_sql::semantics::rules::validate_rule_returning_contract(
138        context.rules,
139        &table,
140        event,
141        command
142            .returning()
143            .is_some_and(|returning| !returning.is_empty()),
144    )?;
145    let requires_rows = requirements.requires_rows;
146    let required = requirements.columns;
147    match command {
148        CommandPlan::Insert(plan) => {
149            let columns = if plan.columns.is_empty() {
150                context
151                    .columns
152                    .source_column_names(&table, true)?
153                    .unwrap_or_default()
154            } else {
155                plan.columns.clone()
156            };
157            let positions = columns
158                .iter()
159                .enumerate()
160                .filter_map(|(position, column)| required.contains(column).then_some(position))
161                .collect::<BTreeSet<_>>();
162            for row in &mut plan.rows {
163                for (position, expression) in row.iter_mut().enumerate() {
164                    if !positions.contains(&position) {
165                        *expression = ScalarExpr::Literal(Value::Null);
166                    }
167                }
168            }
169            if !requires_rows {
170                plan.source = None;
171                plan.ctes.clear();
172                plan.subqueries.clear();
173            } else if let Some(source) = &mut plan.source {
174                crate::mutation_outputs::prune_unused_query_outputs(
175                    source,
176                    &positions,
177                    columns.len(),
178                );
179            }
180        }
181        CommandPlan::Update(plan) => {
182            for assignment in &mut plan.assignments {
183                if !required.contains(&assignment.column) {
184                    assignment.value = ScalarExpr::Literal(Value::Null);
185                }
186            }
187            if !requires_rows {
188                plan.source = None;
189                plan.predicate = None;
190                plan.ctes.clear();
191                plan.subqueries.clear();
192            }
193        }
194        CommandPlan::Delete(plan) if !requires_rows => {
195            plan.source = None;
196            plan.predicate = None;
197            plan.ctes.clear();
198            plan.subqueries.clear();
199        }
200        _ => {}
201    }
202    Ok(())
203}
uqa_planner/statement_planning/rule_inputs.rs

uqa_planner/statement_planning/
rule_inputs.rs
