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    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
180/// 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.
181fn 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            // DEFAULT has a syntax position but no independently bound value in this analysis query.
193            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    // `ON CONFLICT DO UPDATE` also sees the proposed row as `excluded`.
373    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
488/// Routine identities of a bound stored expression in syntax order.
489pub 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;