1use crate::ast::FunctionBinding;
10use crate::plan::{
11 AccessPathPlan, CommandPlan, ComputePlan, ConflictActionPlan, CtePlan, DeletePlan, InsertPlan,
12 JoinExecutionStrategy, MergePlan, ProjectionPlan, QueryBlockPlan, QueryPlan, RelationalPlan,
13 SourcePlan, UnifiedPlan, UpdatePlan,
14};
15use crate::SQLError;
16use crate::{RowSchema, ScalarExpr};
17
18use crate::{binding::context::BindingContext, routines::RoutineResolution};
19use uqa_core::Value;
20
21pub struct CatalogRoutineContext<'a, 'q> {
23 pub routines: &'a dyn RoutineResolution,
24 pub binding: &'a BindingContext<'q>,
25}
26
27pub struct BoundStatementRoutines {
28 pub query: Option<QueryPlan>,
29 pub sites: super::syntax_sites::SyntaxSites,
31 pub query_output: Option<RowSchema>,
33}
34
35#[derive(Debug, Clone)]
36pub struct BoundRoutineReference {
37 pub name: String,
38 pub binding: Option<FunctionBinding>,
39}
40
41struct CommandRoutineInputs {
42 ctes: Vec<CtePlan>,
43 source: Option<SourcePlan>,
44 expressions: Vec<ScalarExpr>,
45 subqueries: Vec<QueryPlan>,
46 outer: RowSchema,
47}
48
49pub fn bind_catalog_statement_routines(
51 context: &CatalogRoutineContext<'_, '_>,
52 plan: &UnifiedPlan,
53 params: &[crate::SQLParam],
54) -> Result<BoundStatementRoutines, SQLError> {
55 let lowered = match plan {
56 UnifiedPlan::Query(query) => {
57 let mut query = (**query).clone();
58 mark_query_relations_bound(&mut query);
59 Some((query, None))
60 }
61 UnifiedPlan::Command(command) => command_statement_query(context, command)?,
62 };
63 let Some((lowered, outer)) = lowered else {
64 return Ok(BoundStatementRoutines {
65 query: None,
66 sites: super::syntax_sites::SyntaxSites::default(),
67 query_output: None,
68 });
69 };
70 let mut query = lowered.clone();
71 let output = crate::binding::bind_syntax_query_plan_routines(
72 context.routines,
73 &mut query,
74 params,
75 context.binding,
76 outer.as_ref(),
77 )?;
78 let sites = if matches!(plan, UnifiedPlan::Command(_)) {
79 super::syntax_sites::command_query_syntax_sites(&lowered, &query)?
80 } else {
81 super::syntax_sites::query_syntax_sites(&lowered, &query)?
82 };
83 Ok(BoundStatementRoutines {
84 query: Some(query),
85 sites,
86 query_output: matches!(plan, UnifiedPlan::Query(_)).then_some(output),
87 })
88}
89
90pub fn mark_catalog_statement_relations_bound(plan: &mut UnifiedPlan) -> Result<(), SQLError> {
91 match plan {
92 UnifiedPlan::Query(query) => mark_query_relations_bound(query),
93 UnifiedPlan::Command(command) => match command.as_mut() {
94 CommandPlan::Insert(plan) => {
95 plan.relations_bound = true;
96 for cte in &mut plan.ctes {
97 mark_cte_relations_bound(&mut cte.body);
98 }
99 if let Some(source) = &mut plan.source {
100 mark_query_relations_bound(source);
101 }
102 for subquery in &mut plan.subqueries {
103 mark_query_relations_bound(subquery);
104 }
105 }
106 CommandPlan::Update(plan) => {
107 plan.relations_bound = true;
108 for cte in &mut plan.ctes {
109 mark_cte_relations_bound(&mut cte.body);
110 }
111 if let Some(source) = &mut plan.source {
112 mark_source_relations_bound(source);
113 }
114 for subquery in &mut plan.subqueries {
115 mark_query_relations_bound(subquery);
116 }
117 }
118 CommandPlan::Delete(plan) => {
119 plan.relations_bound = true;
120 for cte in &mut plan.ctes {
121 mark_cte_relations_bound(&mut cte.body);
122 }
123 if let Some(source) = &mut plan.source {
124 mark_source_relations_bound(source);
125 }
126 for subquery in &mut plan.subqueries {
127 mark_query_relations_bound(subquery);
128 }
129 }
130 CommandPlan::Notify { .. } => {}
131 CommandPlan::Merge(plan) => {
132 for cte in &mut plan.ctes {
133 mark_cte_relations_bound(&mut cte.body);
134 }
135 mark_source_relations_bound(&mut plan.source);
136 for subquery in &mut plan.subqueries {
137 mark_query_relations_bound(subquery);
138 }
139 }
140 _ => {
141 return Err(SQLError::Internal(
142 "catalog-owned statement lowered to an unsupported command".into(),
143 ));
144 }
145 },
146 }
147 Ok(())
148}
149
150fn command_statement_query(
152 context: &CatalogRoutineContext<'_, '_>,
153 command: &CommandPlan,
154) -> Result<Option<(QueryPlan, Option<RowSchema>)>, SQLError> {
155 let Some(inputs) = command_statement_routine_inputs(context, command)? else {
156 return Ok(None);
157 };
158 let projections = inputs
159 .expressions
160 .into_iter()
161 .map(|expr| ProjectionPlan {
162 expr: if matches!(expr, ScalarExpr::Default) {
164 ScalarExpr::Literal(Value::Int(1))
165 } else {
166 expr
167 },
168 alias: None,
169 })
170 .chain(std::iter::once(ProjectionPlan {
171 expr: ScalarExpr::Literal(Value::Int(1)),
172 alias: None,
173 }))
174 .collect();
175 let mut query = QueryPlan {
176 relations_bound: true,
177 ctes: inputs.ctes,
178 root: RelationalPlan::QueryBlock(Box::new(QueryBlockPlan {
179 privilege_columns: std::collections::BTreeSet::default(),
180 windows: Vec::new(),
181 projections,
182 from: inputs.source,
183 r#where: None,
184 compute: ComputePlan::Project,
185 group_by: Vec::new(),
186 grouping_sets: Vec::new(),
187 group_distinct: false,
188 having: None,
189 order_by: Vec::new(),
190 limit: None,
191 with_ties: false,
192 offset: None,
193 distinct: false,
194 distinct_on: Vec::new(),
195 subqueries: inputs.subqueries,
196 access: AccessPathPlan::Row,
197 locking: Vec::new(),
198 })),
199 };
200 mark_query_relations_bound(&mut query);
201 Ok(Some((query, Some(inputs.outer))))
202}
203
204fn command_statement_routine_inputs(
205 context: &CatalogRoutineContext<'_, '_>,
206 command: &CommandPlan,
207) -> Result<Option<CommandRoutineInputs>, SQLError> {
208 match command {
209 CommandPlan::Insert(plan) => insert_statement_routine_inputs(context, plan).map(Some),
210 CommandPlan::Update(plan) => update_statement_routine_inputs(context, plan).map(Some),
211 CommandPlan::Delete(plan) => delete_statement_routine_inputs(context, plan).map(Some),
212 CommandPlan::Merge(plan) => Ok(Some(merge_statement_routine_inputs(plan))),
213 CommandPlan::Notify { .. } => Ok(None),
214 _ => Err(SQLError::Internal(
215 "catalog-owned statement lowered to an unsupported command".into(),
216 )),
217 }
218}
219
220fn merge_statement_routine_inputs(plan: &MergePlan) -> CommandRoutineInputs {
221 let target = SourcePlan::Table {
222 bound_columns: None,
223 name: plan.target.clone(),
224 qualifier: plan.target_qualifier.clone(),
225 alias: plan.target_alias.clone(),
226 column_aliases: Vec::new(),
227 include_descendants: plan.include_descendants,
228 };
229 let source = SourcePlan::Join {
230 left: Box::new(target),
231 right: plan.source.clone(),
232 kind: crate::ast::JoinKind::Cross,
233 on: None,
234 using: None,
235 natural: false,
236 alias: None,
237 column_aliases: Vec::new(),
238 lateral: false,
239 strategy: JoinExecutionStrategy::default(),
240 };
241 let mut expressions = vec![plan.join_condition.clone()];
242 for clause in &plan.when_clauses {
243 match clause {
244 crate::plan::MergeWhenPlan::UpdateMatched {
245 condition,
246 assignments,
247 }
248 | crate::plan::MergeWhenPlan::UpdateNotMatchedBySource {
249 condition,
250 assignments,
251 } => {
252 expressions.extend(condition.iter().cloned());
253 expressions.extend(
254 assignments
255 .iter()
256 .flat_map(crate::plan::AssignmentPlan::expressions)
257 .cloned(),
258 );
259 }
260 crate::plan::MergeWhenPlan::InsertNotMatched {
261 condition,
262 columns,
263 values,
264 ..
265 } => {
266 expressions.extend(condition.iter().cloned());
267 expressions.extend(
268 columns
269 .iter()
270 .flat_map(crate::ast::AssignmentTarget::expressions)
271 .cloned(),
272 );
273 expressions.extend(values.iter().cloned());
274 }
275 crate::plan::MergeWhenPlan::DeleteMatched { condition }
276 | crate::plan::MergeWhenPlan::DeleteNotMatchedBySource { condition }
277 | crate::plan::MergeWhenPlan::NothingMatched { condition }
278 | crate::plan::MergeWhenPlan::NothingNotMatched { condition }
279 | crate::plan::MergeWhenPlan::NothingNotMatchedBySource { condition } => {
280 expressions.extend(condition.iter().cloned());
281 }
282 }
283 }
284 expressions.extend(
285 plan.returning
286 .iter()
287 .map(|projection| projection.expr.clone()),
288 );
289 CommandRoutineInputs {
290 ctes: plan.ctes.clone(),
291 source: Some(source),
292 expressions,
293 subqueries: plan.subqueries.clone(),
294 outer: RowSchema::default(),
295 }
296}
297
298fn insert_statement_routine_inputs(
299 context: &CatalogRoutineContext<'_, '_>,
300 plan: &InsertPlan,
301) -> Result<CommandRoutineInputs, SQLError> {
302 let mut expressions = plan
303 .columns
304 .iter()
305 .flat_map(crate::ast::AssignmentTarget::expressions)
306 .chain(plan.rows.iter().flatten())
307 .cloned()
308 .collect::<Vec<_>>();
309 if let Some(conflict) = &plan.on_conflict {
310 expressions.extend(conflict.expressions.iter().cloned());
311 expressions.extend(conflict.predicate.iter().map(Box::as_ref).cloned());
312 if let ConflictActionPlan::Update {
313 assignments,
314 predicate,
315 } = &conflict.action
316 {
317 expressions.extend(
318 assignments
319 .iter()
320 .flat_map(crate::plan::AssignmentPlan::expressions)
321 .cloned(),
322 );
323 expressions.extend(predicate.iter().map(Box::as_ref).cloned());
324 }
325 }
326 expressions.extend(
327 plan.returning
328 .iter()
329 .map(|projection| projection.expr.clone()),
330 );
331 let source = plan.source.as_ref().map(|source| SourcePlan::Subquery {
332 body: Box::new((**source).clone()),
333 alias: Some("__uqa_catalog_statement_source".into()),
334 column_aliases: Vec::new(),
335 });
336 let mut outer = statement_target_outer_schema(
337 context,
338 &plan.table,
339 &plan.target_qualifier,
340 &plan.returning_aliases,
341 )?;
342 if matches!(
344 plan.on_conflict.as_ref().map(|conflict| &conflict.action),
345 Some(ConflictActionPlan::Update { .. })
346 ) {
347 let target = statement_target_schema(context, &plan.table, &plan.target_qualifier)?;
348 let excluded = RowSchema::with_qualified_types(
349 "excluded",
350 target.columns().to_vec(),
351 target.column_types().to_vec(),
352 );
353 outer = RowSchema::join(&outer, &excluded, std::iter::empty::<String>());
354 }
355 Ok(CommandRoutineInputs {
356 ctes: plan.ctes.clone(),
357 source,
358 expressions,
359 subqueries: plan.subqueries.clone(),
360 outer,
361 })
362}
363
364fn update_statement_routine_inputs(
365 context: &CatalogRoutineContext<'_, '_>,
366 plan: &UpdatePlan,
367) -> Result<CommandRoutineInputs, SQLError> {
368 let mut expressions = plan
369 .assignments
370 .iter()
371 .flat_map(crate::plan::AssignmentPlan::expressions)
372 .cloned()
373 .collect::<Vec<_>>();
374 expressions.extend(plan.predicate.iter().cloned());
375 expressions.extend(
376 plan.returning
377 .iter()
378 .map(|projection| projection.expr.clone()),
379 );
380 Ok(CommandRoutineInputs {
381 ctes: plan.ctes.clone(),
382 source: plan.source.as_deref().cloned(),
383 expressions,
384 subqueries: plan.subqueries.clone(),
385 outer: statement_target_outer_schema(
386 context,
387 &plan.table,
388 &plan.target_qualifier,
389 &plan.returning_aliases,
390 )?,
391 })
392}
393
394fn delete_statement_routine_inputs(
395 context: &CatalogRoutineContext<'_, '_>,
396 plan: &DeletePlan,
397) -> Result<CommandRoutineInputs, SQLError> {
398 let mut expressions = plan.predicate.iter().cloned().collect::<Vec<_>>();
399 expressions.extend(
400 plan.returning
401 .iter()
402 .map(|projection| projection.expr.clone()),
403 );
404 Ok(CommandRoutineInputs {
405 ctes: plan.ctes.clone(),
406 source: plan.source.as_deref().cloned(),
407 expressions,
408 subqueries: plan.subqueries.clone(),
409 outer: statement_target_outer_schema(
410 context,
411 &plan.table,
412 &plan.target_qualifier,
413 &plan.returning_aliases,
414 )?,
415 })
416}
417
418fn statement_target_outer_schema(
419 context: &CatalogRoutineContext<'_, '_>,
420 table: &str,
421 target_qualifier: &str,
422 aliases: &crate::ast::ReturningAliases,
423) -> Result<RowSchema, SQLError> {
424 let target = statement_target_schema(context, table, target_qualifier)?;
425 Ok(crate::semantics::returning_expression_schema(
426 &target,
427 target_qualifier,
428 aliases,
429 None,
430 ))
431}
432
433fn statement_target_schema(
434 context: &CatalogRoutineContext<'_, '_>,
435 table: &str,
436 target_qualifier: &str,
437) -> Result<RowSchema, SQLError> {
438 let target = crate::binding::analyze_source_plan_schema(
439 context.routines,
440 &SourcePlan::Table {
441 bound_columns: None,
442 name: table.to_string(),
443 qualifier: target_qualifier.to_string(),
444 alias: None,
445 column_aliases: Vec::new(),
446 include_descendants: true,
447 },
448 &[],
449 context.binding,
450 None,
451 )?;
452 Ok(RowSchema::with_types(
453 target.columns().to_vec(),
454 target.column_types().to_vec(),
455 ))
456}
457
458pub fn collect_expression_routine_references(
460 expression: &crate::plan::ExpressionPlan,
461) -> Result<Vec<BoundRoutineReference>, SQLError> {
462 Ok(super::syntax_sites::expression_syntax_sites(expression, expression)?.routines)
463}
464
465fn mark_cte_relations_bound(body: &mut crate::plan::CtePlanBody) {
466 match body {
467 crate::plan::CtePlanBody::Query(query) => mark_query_relations_bound(query),
468 crate::plan::CtePlanBody::Command(command) => {
469 match command.as_mut() {
470 CommandPlan::Insert(plan) => {
471 plan.relations_bound = true;
472 plan.target_relation_bound = true;
473 }
474 CommandPlan::Update(plan) => {
475 plan.relations_bound = true;
476 plan.target_relation_bound = true;
477 }
478 CommandPlan::Delete(plan) => {
479 plan.relations_bound = true;
480 plan.target_relation_bound = true;
481 }
482 _ => {}
483 }
484 if let Some(ctes) = command.ctes_mut() {
485 for cte in ctes {
486 mark_cte_relations_bound(&mut cte.body);
487 }
488 }
489 for query in command.query_inputs_mut() {
490 mark_query_relations_bound(query);
491 }
492 if let Some(source) = command.source_input_mut() {
493 mark_source_relations_bound(source);
494 }
495 }
496 }
497}
498
499fn mark_query_relations_bound(query: &mut QueryPlan) {
500 query.relations_bound = true;
501 for cte in &mut query.ctes {
502 mark_cte_relations_bound(&mut cte.body);
503 }
504 match &mut query.root {
505 RelationalPlan::QueryBlock(block) => {
506 if let Some(source) = &mut block.from {
507 mark_source_relations_bound(source);
508 }
509 for subquery in &mut block.subqueries {
510 mark_query_relations_bound(subquery);
511 }
512 }
513 RelationalPlan::SetOp {
514 left,
515 right,
516 subqueries,
517 ..
518 } => {
519 mark_query_relations_bound(left);
520 mark_query_relations_bound(right);
521 for subquery in subqueries {
522 mark_query_relations_bound(subquery);
523 }
524 }
525 RelationalPlan::Values { subqueries, .. } => {
526 for subquery in subqueries {
527 mark_query_relations_bound(subquery);
528 }
529 }
530 }
531}
532
533fn mark_source_relations_bound(source: &mut SourcePlan) {
534 match source {
535 SourcePlan::Join { left, right, .. } => {
536 mark_source_relations_bound(left);
537 mark_source_relations_bound(right);
538 }
539 SourcePlan::Subquery { body, .. } => mark_query_relations_bound(body),
540 SourcePlan::Table { .. }
541 | SourcePlan::Values { .. }
542 | SourcePlan::Function { .. }
543 | SourcePlan::FunctionGroup { .. } => {}
544 }
545}
546
547pub mod analysis;