1use super::{CommandPlan, CtePlan, ProjectionPlan, QueryPlan, SourcePlan};
10
11impl CommandPlan {
12 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 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 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 #[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}