Skip to main content

uqa_sql/binding/
stored_routines.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Exact routine binding for catalog-owned statements.
8
9use 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
21/// Routine metadata and immutable namespace inputs for a stored statement.
22pub 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    /// What binding recorded for the statement's syntax, in syntax order.
30    pub sites: super::syntax_sites::SyntaxSites,
31    /// The columns a query statement returns; a command's binding query is not its result.
32    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
49/// Bind a copy of a catalog-owned statement lowered from its stored syntax; `params` types the positional parameters its syntax references.
50pub 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
150/// A command's syntax as one query whose select list holds the command's expressions in syntax order, so the command binds and reads back like a query. `DEFAULT` markers have no scalar syntax to bind.
151fn 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            // DEFAULT has a syntax position but no independently bound value in this analysis query.
163            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    // `ON CONFLICT DO UPDATE` also sees the proposed row as `excluded`.
343    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
458/// Routine identities of a bound stored expression in syntax order.
459pub 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;