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::{
10    CommandPlan, CtePlan, DeletePlan, InsertPlan, MergePlan, ProjectionPlan, QueryPlan, SourcePlan,
11    UpdatePlan,
12};
13
14impl CommandPlan {
15    /// WITH definitions owned by this command, in declaration order.
16    pub fn ctes(&self) -> &[CtePlan] {
17        match self {
18            Self::Insert(plan) => &plan.ctes,
19            Self::Update(plan) => &plan.ctes,
20            Self::Delete(plan) => &plan.ctes,
21            Self::Merge(plan) => &plan.ctes,
22            _ => &[],
23        }
24    }
25
26    pub fn ctes_mut(&mut self) -> Option<&mut Vec<CtePlan>> {
27        match self {
28            Self::Insert(plan) => Some(&mut plan.ctes),
29            Self::Update(plan) => Some(&mut plan.ctes),
30            Self::Delete(plan) => Some(&mut plan.ctes),
31            Self::Merge(plan) => Some(&mut plan.ctes),
32            _ => None,
33        }
34    }
35
36    /// 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.
37    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
38        match self {
39            Self::Insert(plan) => plan.query_inputs(),
40            Self::Update(plan) => plan.query_inputs(),
41            Self::Delete(plan) => plan.query_inputs(),
42            Self::Merge(plan) => plan.query_inputs(),
43            _ => Vec::new(),
44        }
45    }
46
47    pub fn query_inputs_mut(&mut self) -> Vec<&mut QueryPlan> {
48        match self {
49            Self::Insert(plan) => plan
50                .source
51                .iter_mut()
52                .map(Box::as_mut)
53                .chain(plan.subqueries.iter_mut())
54                .collect(),
55            Self::Update(plan) => plan.subqueries.iter_mut().collect(),
56            Self::Delete(plan) => plan.subqueries.iter_mut().collect(),
57            Self::Merge(plan) => plan.subqueries.iter_mut().collect(),
58            _ => Vec::new(),
59        }
60    }
61
62    pub fn source_input(&self) -> Option<&SourcePlan> {
63        match self {
64            Self::Update(plan) => plan.source_input(),
65            Self::Delete(plan) => plan.source_input(),
66            Self::Merge(plan) => plan.source_input(),
67            _ => None,
68        }
69    }
70
71    pub fn source_input_mut(&mut self) -> Option<&mut SourcePlan> {
72        match self {
73            Self::Update(plan) => plan.source.as_deref_mut(),
74            Self::Delete(plan) => plan.source.as_deref_mut(),
75            Self::Merge(plan) => Some(&mut plan.source),
76            _ => None,
77        }
78    }
79
80    pub fn mutation_target(&self) -> Option<&str> {
81        match self {
82            Self::Insert(plan) => Some(&plan.table),
83            Self::Update(plan) => Some(&plan.table),
84            Self::Delete(plan) => Some(&plan.table),
85            Self::Merge(plan) => Some(&plan.target),
86            _ => None,
87        }
88    }
89
90    pub fn mutation_target_mut(&mut self) -> Option<&mut String> {
91        match self {
92            Self::Insert(plan) => Some(&mut plan.table),
93            Self::Update(plan) => Some(&mut plan.table),
94            Self::Delete(plan) => Some(&mut plan.table),
95            Self::Merge(plan) => Some(&mut plan.target),
96            _ => None,
97        }
98    }
99
100    pub fn returning(&self) -> Option<&[ProjectionPlan]> {
101        match self {
102            Self::Insert(plan) => Some(&plan.returning),
103            Self::Update(plan) => Some(&plan.returning),
104            Self::Delete(plan) => Some(&plan.returning),
105            Self::Merge(plan) => Some(&plan.returning),
106            _ => None,
107        }
108    }
109
110    pub fn returning_mut(&mut self) -> Option<&mut [ProjectionPlan]> {
111        match self {
112            Self::Insert(plan) => Some(&mut plan.returning),
113            Self::Update(plan) => Some(&mut plan.returning),
114            Self::Delete(plan) => Some(&mut plan.returning),
115            Self::Merge(plan) => Some(&mut plan.returning),
116            _ => None,
117        }
118    }
119}
120
121impl CommandPlan {
122    /// Scalar expressions owned by the command, excluding its relational children.
123    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
124        match self {
125            Self::Insert(plan) => plan.expressions(),
126            Self::Update(plan) => plan.expressions(),
127            Self::Delete(plan) => plan.expressions(),
128            Self::Merge(plan) => plan.expressions(),
129            _ => Vec::new(),
130        }
131    }
132
133    /// Scalar expressions owned by the command, excluding its relational children.
134    #[expect(
135        clippy::too_many_lines,
136        reason = "enumerates scalar ownership for every mutation command"
137    )]
138    pub fn expressions_mut(&mut self) -> Vec<&mut super::ScalarExpr> {
139        use super::{ConflictActionPlan, MergeWhenPlan};
140        let mut expressions = Vec::new();
141        match self {
142            Self::Insert(plan) => {
143                expressions.extend(
144                    plan.columns
145                        .iter_mut()
146                        .flat_map(crate::ast::AssignmentTarget::expressions_mut),
147                );
148                expressions.extend(plan.rows.iter_mut().flatten());
149                if let Some(conflict) = &mut plan.on_conflict {
150                    expressions.extend(&mut conflict.expressions);
151                    expressions.extend(conflict.predicate.as_deref_mut());
152                    if let ConflictActionPlan::Update {
153                        assignments,
154                        predicate,
155                    } = &mut conflict.action
156                    {
157                        expressions.extend(
158                            assignments
159                                .iter_mut()
160                                .flat_map(super::AssignmentPlan::expressions_mut),
161                        );
162                        expressions.extend(predicate.as_deref_mut());
163                    }
164                }
165                expressions.extend(
166                    plan.returning
167                        .iter_mut()
168                        .map(|projection| &mut projection.expr),
169                );
170                expressions.extend(
171                    plan.view_checks
172                        .iter_mut()
173                        .map(|check| &mut check.predicate),
174                );
175            }
176            Self::Update(plan) => {
177                expressions.extend(
178                    plan.assignments
179                        .iter_mut()
180                        .flat_map(super::AssignmentPlan::expressions_mut),
181                );
182                expressions.extend(plan.predicate.as_mut());
183                expressions.extend(
184                    plan.returning
185                        .iter_mut()
186                        .map(|projection| &mut projection.expr),
187                );
188                expressions.extend(
189                    plan.view_checks
190                        .iter_mut()
191                        .map(|check| &mut check.predicate),
192                );
193            }
194            Self::Delete(plan) => {
195                expressions.extend(plan.predicate.as_mut());
196                expressions.extend(
197                    plan.returning
198                        .iter_mut()
199                        .map(|projection| &mut projection.expr),
200                );
201            }
202            Self::Merge(plan) => {
203                expressions.extend(plan.target_predicate.as_mut());
204                expressions.push(&mut plan.join_condition);
205                for clause in &mut plan.when_clauses {
206                    match clause {
207                        MergeWhenPlan::UpdateMatched {
208                            condition,
209                            assignments,
210                        }
211                        | MergeWhenPlan::UpdateNotMatchedBySource {
212                            condition,
213                            assignments,
214                        } => {
215                            expressions.extend(condition.as_mut());
216                            expressions.extend(
217                                assignments
218                                    .iter_mut()
219                                    .flat_map(super::AssignmentPlan::expressions_mut),
220                            );
221                        }
222                        MergeWhenPlan::InsertNotMatched {
223                            condition,
224                            columns,
225                            values,
226                            ..
227                        } => {
228                            expressions.extend(condition.as_mut());
229                            expressions.extend(
230                                columns
231                                    .iter_mut()
232                                    .flat_map(crate::ast::AssignmentTarget::expressions_mut),
233                            );
234                            expressions.extend(values);
235                        }
236                        MergeWhenPlan::DeleteMatched { condition }
237                        | MergeWhenPlan::DeleteNotMatchedBySource { condition }
238                        | MergeWhenPlan::NothingMatched { condition }
239                        | MergeWhenPlan::NothingNotMatched { condition }
240                        | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
241                            expressions.extend(condition.as_mut());
242                        }
243                    }
244                }
245                expressions.extend(
246                    plan.returning
247                        .iter_mut()
248                        .map(|projection| &mut projection.expr),
249                );
250                expressions.extend(
251                    plan.view_checks
252                        .iter_mut()
253                        .map(|check| &mut check.predicate),
254                );
255            }
256            _ => {}
257        }
258        expressions
259    }
260
261    pub fn scalar_subqueries(&self) -> &[QueryPlan] {
262        match self {
263            Self::Insert(plan) => &plan.subqueries,
264            Self::Update(plan) => &plan.subqueries,
265            Self::Delete(plan) => &plan.subqueries,
266            Self::Merge(plan) => &plan.subqueries,
267            _ => &[],
268        }
269    }
270
271    pub fn target_qualifier(&self) -> Option<&str> {
272        match self {
273            Self::Insert(plan) => Some(&plan.target_qualifier),
274            Self::Update(plan) => Some(&plan.target_qualifier),
275            Self::Delete(plan) => Some(&plan.target_qualifier),
276            Self::Merge(plan) => Some(&plan.target_qualifier),
277            _ => None,
278        }
279    }
280
281    pub fn returning_aliases(&self) -> Option<&crate::ast::ReturningAliases> {
282        match self {
283            Self::Insert(plan) => Some(&plan.returning_aliases),
284            Self::Update(plan) => Some(&plan.returning_aliases),
285            Self::Delete(plan) => Some(&plan.returning_aliases),
286            Self::Merge(plan) => Some(&plan.returning_aliases),
287            _ => None,
288        }
289    }
290}
291
292impl InsertPlan {
293    /// Scalar expressions owned by this command, excluding its relational children.
294    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
295        let mut expressions = Vec::new();
296        expressions.extend(
297            self.columns
298                .iter()
299                .flat_map(crate::ast::AssignmentTarget::expressions),
300        );
301        expressions.extend(self.rows.iter().flatten());
302        if let Some(conflict) = &self.on_conflict {
303            expressions.extend(&conflict.expressions);
304            expressions.extend(conflict.predicate.as_deref());
305            expressions.extend(self.conflict_update_expressions());
306        }
307        expressions.extend(self.returning.iter().map(|projection| &projection.expr));
308        expressions.extend(self.view_checks.iter().map(|check| &check.predicate));
309
310        expressions
311    }
312
313    /// Expressions that see the additional EXCLUDED relation in ON CONFLICT DO UPDATE.
314    pub fn conflict_update_expressions(&self) -> Vec<&super::ScalarExpr> {
315        match self.on_conflict.as_ref().map(|conflict| &conflict.action) {
316            Some(super::ConflictActionPlan::Update {
317                assignments,
318                predicate,
319            }) => assignments
320                .iter()
321                .flat_map(super::AssignmentPlan::expressions)
322                .chain(predicate.as_deref())
323                .collect(),
324            _ => Vec::new(),
325        }
326    }
327
328    /// Query children evaluated in the command's WITH scope: the source query and every `subqueries` entry.
329    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
330        self.source
331            .iter()
332            .map(Box::as_ref)
333            .chain(self.subqueries.iter())
334            .collect()
335    }
336}
337
338impl UpdatePlan {
339    /// Scalar expressions owned by this command, excluding its relational children.
340    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
341        let mut expressions = Vec::new();
342        expressions.extend(
343            self.assignments
344                .iter()
345                .flat_map(super::AssignmentPlan::expressions),
346        );
347        expressions.extend(self.predicate.as_ref());
348        expressions.extend(self.returning.iter().map(|projection| &projection.expr));
349        expressions.extend(self.view_checks.iter().map(|check| &check.predicate));
350
351        expressions
352    }
353
354    /// Query children evaluated in the command's WITH scope, which are its `subqueries` entries.
355    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
356        self.subqueries.iter().collect()
357    }
358
359    /// The `FROM` relation the command joins to its target.
360    pub fn source_input(&self) -> Option<&SourcePlan> {
361        self.source.as_deref()
362    }
363}
364
365impl DeletePlan {
366    /// Scalar expressions owned by this command, excluding its relational children.
367    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
368        let mut expressions = Vec::new();
369        expressions.extend(self.predicate.as_ref());
370        expressions.extend(self.returning.iter().map(|projection| &projection.expr));
371
372        expressions
373    }
374
375    /// Query children evaluated in the command's WITH scope, which are its `subqueries` entries.
376    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
377        self.subqueries.iter().collect()
378    }
379
380    /// The `USING` relation the command joins to its target.
381    pub fn source_input(&self) -> Option<&SourcePlan> {
382        self.source.as_deref()
383    }
384}
385
386impl MergePlan {
387    /// Scalar expressions owned by this command, excluding its relational children.
388    pub fn expressions(&self) -> Vec<&super::ScalarExpr> {
389        use super::MergeWhenPlan;
390        let mut expressions = Vec::new();
391        expressions.extend(self.target_predicate.as_ref());
392        expressions.push(&self.join_condition);
393        for clause in &self.when_clauses {
394            match clause {
395                MergeWhenPlan::UpdateMatched {
396                    condition,
397                    assignments,
398                }
399                | MergeWhenPlan::UpdateNotMatchedBySource {
400                    condition,
401                    assignments,
402                } => {
403                    expressions.extend(condition.as_ref());
404                    expressions.extend(
405                        assignments
406                            .iter()
407                            .flat_map(super::AssignmentPlan::expressions),
408                    );
409                }
410                MergeWhenPlan::InsertNotMatched {
411                    condition,
412                    columns,
413                    values,
414                    ..
415                } => {
416                    expressions.extend(condition.as_ref());
417                    expressions.extend(
418                        columns
419                            .iter()
420                            .flat_map(crate::ast::AssignmentTarget::expressions),
421                    );
422                    expressions.extend(values);
423                }
424                MergeWhenPlan::DeleteMatched { condition }
425                | MergeWhenPlan::DeleteNotMatchedBySource { condition }
426                | MergeWhenPlan::NothingMatched { condition }
427                | MergeWhenPlan::NothingNotMatched { condition }
428                | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
429                    expressions.extend(condition.as_ref());
430                }
431            }
432        }
433        expressions.extend(self.returning.iter().map(|projection| &projection.expr));
434        expressions.extend(self.view_checks.iter().map(|check| &check.predicate));
435
436        expressions
437    }
438
439    /// Query children evaluated in the command's WITH scope, which are its `subqueries` entries.
440    pub fn query_inputs(&self) -> Vec<&QueryPlan> {
441        self.subqueries.iter().collect()
442    }
443
444    /// The `USING` relation the command merges into its target.
445    pub fn source_input(&self) -> Option<&SourcePlan> {
446        Some(&self.source)
447    }
448}