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