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(
121                    plan.columns
122                        .iter()
123                        .flat_map(crate::ast::AssignmentTarget::expressions),
124                );
125                expressions.extend(plan.rows.iter().flatten());
126                if let Some(conflict) = &plan.on_conflict {
127                    expressions.extend(&conflict.expressions);
128                    expressions.extend(conflict.predicate.as_deref());
129                    if let ConflictActionPlan::Update {
130                        assignments,
131                        predicate,
132                    } = &conflict.action
133                    {
134                        expressions.extend(
135                            assignments
136                                .iter()
137                                .flat_map(super::AssignmentPlan::expressions),
138                        );
139                        expressions.extend(predicate.as_deref());
140                    }
141                }
142                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
143                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
144            }
145            Self::Update(plan) => {
146                expressions.extend(
147                    plan.assignments
148                        .iter()
149                        .flat_map(super::AssignmentPlan::expressions),
150                );
151                expressions.extend(plan.predicate.as_ref());
152                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
153                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
154            }
155            Self::Delete(plan) => {
156                expressions.extend(plan.predicate.as_ref());
157                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
158            }
159            Self::Merge(plan) => {
160                expressions.extend(plan.target_predicate.as_ref());
161                expressions.push(&plan.join_condition);
162                for clause in &plan.when_clauses {
163                    match clause {
164                        MergeWhenPlan::UpdateMatched {
165                            condition,
166                            assignments,
167                        }
168                        | MergeWhenPlan::UpdateNotMatchedBySource {
169                            condition,
170                            assignments,
171                        } => {
172                            expressions.extend(condition.as_ref());
173                            expressions.extend(
174                                assignments
175                                    .iter()
176                                    .flat_map(super::AssignmentPlan::expressions),
177                            );
178                        }
179                        MergeWhenPlan::InsertNotMatched {
180                            condition,
181                            columns,
182                            values,
183                        } => {
184                            expressions.extend(condition.as_ref());
185                            expressions.extend(
186                                columns
187                                    .iter()
188                                    .flat_map(crate::ast::AssignmentTarget::expressions),
189                            );
190                            expressions.extend(values);
191                        }
192                        MergeWhenPlan::DeleteMatched { condition }
193                        | MergeWhenPlan::DeleteNotMatchedBySource { condition }
194                        | MergeWhenPlan::NothingMatched { condition }
195                        | MergeWhenPlan::NothingNotMatched { condition }
196                        | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
197                            expressions.extend(condition.as_ref());
198                        }
199                    }
200                }
201                expressions.extend(plan.returning.iter().map(|projection| &projection.expr));
202                expressions.extend(plan.view_checks.iter().map(|check| &check.predicate));
203            }
204            _ => {}
205        }
206        expressions
207    }
208
209    /// Scalar expressions owned by the command, excluding its relational children.
210    #[expect(
211        clippy::too_many_lines,
212        reason = "enumerates scalar ownership for every mutation command"
213    )]
214    pub fn expressions_mut(&mut self) -> Vec<&mut super::ScalarExpr> {
215        use super::{ConflictActionPlan, MergeWhenPlan};
216        let mut expressions = Vec::new();
217        match self {
218            Self::Insert(plan) => {
219                expressions.extend(
220                    plan.columns
221                        .iter_mut()
222                        .flat_map(crate::ast::AssignmentTarget::expressions_mut),
223                );
224                expressions.extend(plan.rows.iter_mut().flatten());
225                if let Some(conflict) = &mut plan.on_conflict {
226                    expressions.extend(&mut conflict.expressions);
227                    expressions.extend(conflict.predicate.as_deref_mut());
228                    if let ConflictActionPlan::Update {
229                        assignments,
230                        predicate,
231                    } = &mut conflict.action
232                    {
233                        expressions.extend(
234                            assignments
235                                .iter_mut()
236                                .flat_map(super::AssignmentPlan::expressions_mut),
237                        );
238                        expressions.extend(predicate.as_deref_mut());
239                    }
240                }
241                expressions.extend(
242                    plan.returning
243                        .iter_mut()
244                        .map(|projection| &mut projection.expr),
245                );
246                expressions.extend(
247                    plan.view_checks
248                        .iter_mut()
249                        .map(|check| &mut check.predicate),
250                );
251            }
252            Self::Update(plan) => {
253                expressions.extend(
254                    plan.assignments
255                        .iter_mut()
256                        .flat_map(super::AssignmentPlan::expressions_mut),
257                );
258                expressions.extend(plan.predicate.as_mut());
259                expressions.extend(
260                    plan.returning
261                        .iter_mut()
262                        .map(|projection| &mut projection.expr),
263                );
264                expressions.extend(
265                    plan.view_checks
266                        .iter_mut()
267                        .map(|check| &mut check.predicate),
268                );
269            }
270            Self::Delete(plan) => {
271                expressions.extend(plan.predicate.as_mut());
272                expressions.extend(
273                    plan.returning
274                        .iter_mut()
275                        .map(|projection| &mut projection.expr),
276                );
277            }
278            Self::Merge(plan) => {
279                expressions.extend(plan.target_predicate.as_mut());
280                expressions.push(&mut plan.join_condition);
281                for clause in &mut plan.when_clauses {
282                    match clause {
283                        MergeWhenPlan::UpdateMatched {
284                            condition,
285                            assignments,
286                        }
287                        | MergeWhenPlan::UpdateNotMatchedBySource {
288                            condition,
289                            assignments,
290                        } => {
291                            expressions.extend(condition.as_mut());
292                            expressions.extend(
293                                assignments
294                                    .iter_mut()
295                                    .flat_map(super::AssignmentPlan::expressions_mut),
296                            );
297                        }
298                        MergeWhenPlan::InsertNotMatched {
299                            condition,
300                            columns,
301                            values,
302                        } => {
303                            expressions.extend(condition.as_mut());
304                            expressions.extend(
305                                columns
306                                    .iter_mut()
307                                    .flat_map(crate::ast::AssignmentTarget::expressions_mut),
308                            );
309                            expressions.extend(values);
310                        }
311                        MergeWhenPlan::DeleteMatched { condition }
312                        | MergeWhenPlan::DeleteNotMatchedBySource { condition }
313                        | MergeWhenPlan::NothingMatched { condition }
314                        | MergeWhenPlan::NothingNotMatched { condition }
315                        | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
316                            expressions.extend(condition.as_mut());
317                        }
318                    }
319                }
320                expressions.extend(
321                    plan.returning
322                        .iter_mut()
323                        .map(|projection| &mut projection.expr),
324                );
325                expressions.extend(
326                    plan.view_checks
327                        .iter_mut()
328                        .map(|check| &mut check.predicate),
329                );
330            }
331            _ => {}
332        }
333        expressions
334    }
335
336    pub fn scalar_subqueries(&self) -> &[QueryPlan] {
337        match self {
338            Self::Insert(plan) => &plan.subqueries,
339            Self::Update(plan) => &plan.subqueries,
340            Self::Delete(plan) => &plan.subqueries,
341            Self::Merge(plan) => &plan.subqueries,
342            _ => &[],
343        }
344    }
345
346    pub fn target_qualifier(&self) -> Option<&str> {
347        match self {
348            Self::Insert(plan) => Some(&plan.target_qualifier),
349            Self::Update(plan) => Some(&plan.target_qualifier),
350            Self::Delete(plan) => Some(&plan.target_qualifier),
351            Self::Merge(plan) => Some(&plan.target_qualifier),
352            _ => None,
353        }
354    }
355
356    pub fn returning_aliases(&self) -> Option<&crate::ast::ReturningAliases> {
357        match self {
358            Self::Insert(plan) => Some(&plan.returning_aliases),
359            Self::Update(plan) => Some(&plan.returning_aliases),
360            Self::Delete(plan) => Some(&plan.returning_aliases),
361            Self::Merge(plan) => Some(&plan.returning_aliases),
362            _ => None,
363        }
364    }
365}
uqa_sql/plan/command_children.rs

uqa_sql/plan/
command_children.rs
