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