Skip to main content

uqa_sql/plan/
command_children.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Direct relational children of data-modifying commands.
8
9use super::{CommandPlan, CtePlan, ProjectionPlan, QueryPlan, SourcePlan};
10
11impl CommandPlan {
12    /// WITH definitions owned by this command, in declaration order.
13    pub fn ctes(&self) -> &[CtePlan] {
14        match self {
15            Self::Insert(plan) => &plan.ctes,
16            Self::Update(plan) => &plan.ctes,
17            Self::Delete(plan) => &plan.ctes,
18            Self::Merge(plan) => &plan.ctes,
19            _ => &[],
20        }
21    }
22
23    pub fn ctes_mut(&mut self) -> Option<&mut Vec<CtePlan>> {
24        match self {
25            Self::Insert(plan) => Some(&mut plan.ctes),
26            Self::Update(plan) => Some(&mut plan.ctes),
27            Self::Delete(plan) => Some(&mut plan.ctes),
28            Self::Merge(plan) => Some(&mut plan.ctes),
29            _ => None,
30        }
31    }
32
33    /// Query children evaluated in the command's WITH scope, including every `scalar_subqueries` entry. Source-plan subqueries are owned by `source_input`; visitors should not traverse `scalar_subqueries` separately.
34    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
35        match self {
36            Self::Insert(plan) => plan
37                .source
38                .iter()
39                .map(Box::as_ref)
40                .chain(plan.subqueries.iter())
41                .collect(),
42            Self::Update(plan) => plan.subqueries.iter().collect(),
43            Self::Delete(plan) => plan.subqueries.iter().collect(),
44            Self::Merge(plan) => plan.subqueries.iter().collect(),
45            _ => Vec::new(),
46        }
47    }
48
49    pub fn query_inputs_mut(&mut self) -> Vec<&mut QueryPlan> {
50        match self {
51            Self::Insert(plan) => plan
52                .source
53                .iter_mut()
54                .map(Box::as_mut)
55                .chain(plan.subqueries.iter_mut())
56                .collect(),
57            Self::Update(plan) => plan.subqueries.iter_mut().collect(),
58            Self::Delete(plan) => plan.subqueries.iter_mut().collect(),
59            Self::Merge(plan) => plan.subqueries.iter_mut().collect(),
60            _ => Vec::new(),
61        }
62    }
63
64    pub fn source_input(&self) -> Option<&SourcePlan> {
65        match self {
66            Self::Update(plan) => plan.source.as_deref(),
67            Self::Delete(plan) => plan.source.as_deref(),
68            Self::Merge(plan) => Some(&plan.source),
69            _ => None,
70        }
71    }
72
73    pub fn source_input_mut(&mut self) -> Option<&mut SourcePlan> {
74        match self {
75            Self::Update(plan) => plan.source.as_deref_mut(),
76            Self::Delete(plan) => plan.source.as_deref_mut(),
77            Self::Merge(plan) => Some(&mut plan.source),
78            _ => None,
79        }
80    }
81
82    pub fn mutation_target(&self) -> Option<&str> {
83        match self {
84            Self::Insert(plan) => Some(&plan.table),
85            Self::Update(plan) => Some(&plan.table),
86            Self::Delete(plan) => Some(&plan.table),
87            Self::Merge(plan) => Some(&plan.target),
88            _ => None,
89        }
90    }
91
92    pub fn mutation_target_mut(&mut self) -> Option<&mut String> {
93        match self {
94            Self::Insert(plan) => Some(&mut plan.table),
95            Self::Update(plan) => Some(&mut plan.table),
96            Self::Delete(plan) => Some(&mut plan.table),
97            Self::Merge(plan) => Some(&mut plan.target),
98            _ => None,
99        }
100    }
101
102    pub fn returning(&self) -> Option<&[ProjectionPlan]> {
103        match self {
104            Self::Insert(plan) => Some(&plan.returning),
105            Self::Update(plan) => Some(&plan.returning),
106            Self::Delete(plan) => Some(&plan.returning),
107            Self::Merge(plan) => Some(&plan.returning),
108            _ => None,
109        }
110    }
111}
112
113impl CommandPlan {
114    /// Scalar expressions owned by the command, excluding its relational children.
115    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
116        use super::{ConflictActionPlan, MergeWhenPlan};
117        let mut expressions = Vec::new();
118        match self {
119            Self::Insert(plan) => {
120                expressions.extend(plan.rows.iter().flatten());
121                if let Some(conflict) = &plan.on_conflict {
122                    expressions.extend(&conflict.expressions);
123                    expressions.extend(conflict.predicate.as_deref());
124                    if let ConflictActionPlan::Update {
125                        assignments,
126                        predicate,
127                    } = &conflict.action
128                    {
129                        expressions.extend(assignments.iter().map(|assignment| &assignment.value));
130                        expressions.extend(predicate.as_deref());
131                    }
132                }
133                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
134                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
135            }
136            Self::Update(plan) => {
137                expressions.extend(plan.assignments.iter().map(|assignment| &assignment.value));
138                expressions.extend(plan.predicate.as_ref());
139                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
140                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
141            }
142            Self::Delete(plan) => {
143                expressions.extend(plan.predicate.as_ref());
144                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
145            }
146            Self::Merge(plan) => {
147                expressions.extend(plan.target_predicate.as_ref());
148                expressions.push(&plan.join_condition);
149                for clause in &plan.when_clauses {
150                    match clause {
151                        MergeWhenPlan::UpdateMatched {
152                            condition,
153                            assignments,
154                        }
155                        | MergeWhenPlan::UpdateNotMatchedBySource {
156                            condition,
157                            assignments,
158                        } => {
159                            expressions.extend(condition.as_ref());
160                            expressions
161                                .extend(assignments.iter().map(|assignment| &assignment.value));
162                        }
163                        MergeWhenPlan::InsertNotMatched {
164                            condition, values, ..
165                        } => {
166                            expressions.extend(condition.as_ref());
167                            expressions.extend(values);
168                        }
169                        MergeWhenPlan::DeleteMatched { condition }
170                        | MergeWhenPlan::DeleteNotMatchedBySource { condition }
171                        | MergeWhenPlan::NothingMatched { condition }
172                        | MergeWhenPlan::NothingNotMatched { condition }
173                        | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
174                            expressions.extend(condition.as_ref());
175                        }
176                    }
177                }
178                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
179                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
180            }
181            _ => {}
182        }
183        expressions
184    }
185
186    /// Scalar expressions owned by the command, excluding its relational children.
187    #[expect(
188        clippy::too_many_lines,
189        reason = "enumerates scalar ownership for every mutation command"
190    )]
191    pub fn expressions_mut(&mut self) -> Vec<&mut super::ScalarExpr> {
192        use super::{ConflictActionPlan, MergeWhenPlan};
193        let mut expressions = Vec::new();
194        match self {
195            Self::Insert(plan) => {
196                expressions.extend(plan.rows.iter_mut().flatten());
197                if let Some(conflict) = &mut plan.on_conflict {
198                    expressions.extend(&mut conflict.expressions);
199                    expressions.extend(conflict.predicate.as_deref_mut());
200                    if let ConflictActionPlan::Update {
201                        assignments,
202                        predicate,
203                    } = &mut conflict.action
204                    {
205                        expressions.extend(
206                            assignments
207                                .iter_mut()
208                                .map(|assignment| &mut assignment.value),
209                        );
210                        expressions.extend(predicate.as_deref_mut());
211                    }
212                }
213                expressions.extend(
214                    plan.returning
215                        .iter_mut()
216                        .map(|projection| &mut projection.expr),
217                );
218                expressions.extend(
219                    plan.view_checks
220                        .iter_mut()
221                        .map(|check| &mut check.predicate),
222                );
223            }
224            Self::Update(plan) => {
225                expressions.extend(
226                    plan.assignments
227                        .iter_mut()
228                        .map(|assignment| &mut assignment.value),
229                );
230                expressions.extend(plan.predicate.as_mut());
231                expressions.extend(
232                    plan.returning
233                        .iter_mut()
234                        .map(|projection| &mut projection.expr),
235                );
236                expressions.extend(
237                    plan.view_checks
238                        .iter_mut()
239                        .map(|check| &mut check.predicate),
240                );
241            }
242            Self::Delete(plan) => {
243                expressions.extend(plan.predicate.as_mut());
244                expressions.extend(
245                    plan.returning
246                        .iter_mut()
247                        .map(|projection| &mut projection.expr),
248                );
249            }
250            Self::Merge(plan) => {
251                expressions.extend(plan.target_predicate.as_mut());
252                expressions.push(&mut plan.join_condition);
253                for clause in &mut plan.when_clauses {
254                    match clause {
255                        MergeWhenPlan::UpdateMatched {
256                            condition,
257                            assignments,
258                        }
259                        | MergeWhenPlan::UpdateNotMatchedBySource {
260                            condition,
261                            assignments,
262                        } => {
263                            expressions.extend(condition.as_mut());
264                            expressions.extend(
265                                assignments
266                                    .iter_mut()
267                                    .map(|assignment| &mut assignment.value),
268                            );
269                        }
270                        MergeWhenPlan::InsertNotMatched {
271                            condition, values, ..
272                        } => {
273                            expressions.extend(condition.as_mut());
274                            expressions.extend(values);
275                        }
276                        MergeWhenPlan::DeleteMatched { condition }
277                        | MergeWhenPlan::DeleteNotMatchedBySource { condition }
278                        | MergeWhenPlan::NothingMatched { condition }
279                        | MergeWhenPlan::NothingNotMatched { condition }
280                        | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
281                            expressions.extend(condition.as_mut());
282                        }
283                    }
284                }
285                expressions.extend(
286                    plan.returning
287                        .iter_mut()
288                        .map(|projection| &mut projection.expr),
289                );
290                expressions.extend(
291                    plan.view_checks
292                        .iter_mut()
293                        .map(|check| &mut check.predicate),
294                );
295            }
296            _ => {}
297        }
298        expressions
299    }
300
301    pub fn scalar_subqueries(&self) -> &[QueryPlan] {
302        match self {
303            Self::Insert(plan) => &plan.subqueries,
304            Self::Update(plan) => &plan.subqueries,
305            Self::Delete(plan) => &plan.subqueries,
306            Self::Merge(plan) => &plan.subqueries,
307            _ => &[],
308        }
309    }
310
311    pub fn target_qualifier(&self) -> Option<&str> {
312        match self {
313            Self::Insert(plan) => Some(&plan.target_qualifier),
314            Self::Update(plan) => Some(&plan.target_qualifier),
315            Self::Delete(plan) => Some(&plan.target_qualifier),
316            Self::Merge(plan) => Some(&plan.target_qualifier),
317            _ => None,
318        }
319    }
320
321    pub fn returning_aliases(&self) -> Option<&crate::ast::ReturningAliases> {
322        match self {
323            Self::Insert(plan) => Some(&plan.returning_aliases),
324            Self::Update(plan) => Some(&plan.returning_aliases),
325            Self::Delete(plan) => Some(&plan.returning_aliases),
326            Self::Merge(plan) => Some(&plan.returning_aliases),
327            _ => None,
328        }
329    }
330}