Skip to main content

uqa_sql/plpgsql/
binding.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Variable binding across expressions, queries, and statements.
8
9use super::{Expr, FromClause, MergeWhen, Projection, Result, SelectStmt, Statement, Value, CTE};
10use crate::ast::InternalColumnRef;
11
12/// Runtime datum value together with the concrete SQL type declared by PL/pgSQL. The type is optional for composite fields and pseudo-types whose runtime carrier already identifies their category.
13#[derive(Debug, Clone)]
14pub struct ResolvedVariable {
15    pub value: Value,
16    pub declared_type: Option<String>,
17}
18
19impl ResolvedVariable {
20    #[must_use]
21    pub fn untyped(value: Value) -> Self {
22        Self {
23            value,
24            declared_type: None,
25        }
26    }
27
28    fn into_expression(self) -> Expr {
29        match self.declared_type {
30            Some(ty) => Expr::Cast {
31                expr: Box::new(Expr::Literal(self.value)),
32                ty,
33            },
34            None => Expr::Literal(self.value),
35        }
36    }
37}
38
39/// Resolves routine variables while a compiled expression / statement
40/// is being specialized for one execution.
41pub trait VariableResolver {
42    /// Current value of an unqualified name. `Ok(None)` leaves the
43    /// column reference for the engine to resolve.
44    fn resolve_name(&mut self, name: &str) -> Result<Option<ResolvedVariable>>;
45    /// Current value of `qualifier.column` (record field access).
46    fn resolve_qualified(
47        &mut self,
48        qualifier: &str,
49        column: &str,
50    ) -> Result<Option<ResolvedVariable>>;
51    /// Value of a positional `$n` reference (function arguments).
52    fn resolve_param(&mut self, index: usize) -> Result<Option<ResolvedVariable>>;
53
54    /// Optional expression-level rewrite hook. The default preserves the variable-substitution behavior used by PL/pgSQL while allowing catalog lifecycle code to rewrite a reference without fabricating a literal value.
55    fn rewrite_name(&mut self, name: &str) -> Result<Option<Expr>> {
56        Ok(self
57            .resolve_name(name)?
58            .map(ResolvedVariable::into_expression))
59    }
60
61    /// Expression-level counterpart of [`Self::resolve_qualified`].
62    fn rewrite_qualified(&mut self, qualifier: &str, column: &str) -> Result<Option<Expr>> {
63        Ok(self
64            .resolve_qualified(qualifier, column)?
65            .map(ResolvedVariable::into_expression))
66    }
67
68    /// Expand `qualifier.*` when the containing SQL construct treats it as a
69    /// list item. Scalar binding deliberately does not call this hook because
70    /// `PostgreSQL` treats the same syntax as a composite whole-row value in
71    /// scalar contexts.
72    fn rewrite_qualified_star(&mut self, _qualifier: &str) -> Result<Option<Vec<Expr>>> {
73        Ok(None)
74    }
75
76    /// Resolve `qualifier.*` when it appears in a scalar context and therefore denotes one composite whole-row value rather than a projection list.
77    fn rewrite_qualified_whole_row(&mut self, _qualifier: &str) -> Result<Option<Expr>> {
78        Ok(None)
79    }
80
81    /// Expression-level counterpart of [`Self::resolve_param`].
82    fn rewrite_param(&mut self, index: usize) -> Result<Option<Expr>> {
83        Ok(self
84            .resolve_param(index)?
85            .map(ResolvedVariable::into_expression))
86    }
87
88    /// Observe or replace an executor-only structural column reference. SQL
89    /// variable resolvers normally leave these untouched.
90    fn rewrite_internal(&mut self, _column: InternalColumnRef) -> Result<Option<Expr>> {
91        Ok(None)
92    }
93}
94
95/// Rewrite an expression, substituting resolvable variable references
96/// with literals. References the resolver declines stay untouched.
97#[expect(
98    clippy::too_many_lines,
99    reason = "PL/pgSQL lowering preserves parser order and datum validation"
100)]
101pub fn bind_expr(expr: &Expr, r: &mut dyn VariableResolver) -> Result<Expr> {
102    Ok(match expr {
103        Expr::Column(name) => match r.rewrite_name(name)? {
104            Some(value) => value,
105            None => expr.clone(),
106        },
107        Expr::QualifiedColumn {
108            qualifier, column, ..
109        } => match r.rewrite_qualified(qualifier, column)? {
110            Some(value) => value,
111            None => expr.clone(),
112        },
113        Expr::Param(index) => match r.rewrite_param(*index)? {
114            Some(value) => value,
115            None => expr.clone(),
116        },
117        Expr::InternalColumn(column) => match r.rewrite_internal(*column)? {
118            Some(value) => value,
119            None => expr.clone(),
120        },
121        Expr::QualifiedStar(qualifier) => r
122            .rewrite_qualified_whole_row(qualifier)?
123            .unwrap_or_else(|| expr.clone()),
124        Expr::Default | Expr::Literal(_) | Expr::Star => expr.clone(),
125        Expr::Func {
126            name,
127            binding,
128            args,
129            distinct,
130            order_by,
131            filter,
132        } => Expr::Func {
133            name: name.clone(),
134            binding: binding.clone(),
135            args: bind_exprs(args, r)?,
136            distinct: *distinct,
137            order_by: bind_order_by(order_by, r)?,
138            filter: match filter {
139                Some(f) => Some(Box::new(bind_expr(f, r)?)),
140                None => None,
141            },
142        },
143        Expr::Array(items) => Expr::Array(bind_exprs(items, r)?),
144        Expr::Row(items) => Expr::Row(bind_exprs(items, r)?),
145        Expr::Binary { op, lhs, rhs } => Expr::Binary {
146            op: *op,
147            lhs: Box::new(bind_expr(lhs, r)?),
148            rhs: Box::new(bind_expr(rhs, r)?),
149        },
150        Expr::UnaryMinus(inner) => Expr::UnaryMinus(Box::new(bind_expr(inner, r)?)),
151        Expr::Not(inner) => Expr::Not(Box::new(bind_expr(inner, r)?)),
152        Expr::And(items) => Expr::And(bind_exprs(items, r)?),
153        Expr::Or(items) => Expr::Or(bind_exprs(items, r)?),
154        Expr::IsNull { expr, negated } => Expr::IsNull {
155            expr: Box::new(bind_expr(expr, r)?),
156            negated: *negated,
157        },
158        Expr::Between { expr, low, high } => Expr::Between {
159            expr: Box::new(bind_expr(expr, r)?),
160            low: Box::new(bind_expr(low, r)?),
161            high: Box::new(bind_expr(high, r)?),
162        },
163        Expr::InList {
164            expr,
165            list,
166            negated,
167        } => Expr::InList {
168            expr: Box::new(bind_expr(expr, r)?),
169            list: bind_exprs(list, r)?,
170            negated: *negated,
171        },
172        Expr::WindowCall { name, args, spec } => Expr::WindowCall {
173            name: name.clone(),
174            args: bind_exprs(args, r)?,
175            spec: crate::ast::WindowSpec {
176                reference: spec.reference.clone(),
177                partition_by: bind_exprs(&spec.partition_by, r)?,
178                order_by: bind_order_by(&spec.order_by, r)?,
179                frame: spec.frame.clone(),
180            },
181        },
182        Expr::Case {
183            base,
184            when,
185            else_branch,
186        } => Expr::Case {
187            base: match base {
188                Some(b) => Some(Box::new(bind_expr(b, r)?)),
189                None => None,
190            },
191            when: when
192                .iter()
193                .map(|(c, v)| Ok((bind_expr(c, r)?, bind_expr(v, r)?)))
194                .collect::<Result<Vec<_>>>()?,
195            else_branch: match else_branch {
196                Some(e) => Some(Box::new(bind_expr(e, r)?)),
197                None => None,
198            },
199        },
200        Expr::Cast { expr, ty } => Expr::Cast {
201            expr: Box::new(bind_expr(expr, r)?),
202            ty: ty.clone(),
203        },
204        Expr::ScalarSubquery(body) => Expr::ScalarSubquery(Box::new(bind_select(body, r)?)),
205        Expr::Exists { body, negated } => Expr::Exists {
206            body: Box::new(bind_select(body, r)?),
207            negated: *negated,
208        },
209        Expr::InSubquery {
210            expr,
211            body,
212            negated,
213        } => Expr::InSubquery {
214            expr: Box::new(bind_expr(expr, r)?),
215            body: Box::new(bind_select(body, r)?),
216            negated: *negated,
217        },
218    })
219}
220
221pub(super) fn bind_exprs(exprs: &[Expr], r: &mut dyn VariableResolver) -> Result<Vec<Expr>> {
222    exprs.iter().map(|e| bind_expr(e, r)).collect()
223}
224
225pub(super) fn bind_opt_expr(
226    expr: Option<&Expr>,
227    r: &mut dyn VariableResolver,
228) -> Result<Option<Expr>> {
229    match expr {
230        Some(e) => Ok(Some(bind_expr(e, r)?)),
231        None => Ok(None),
232    }
233}
234
235pub(super) fn bind_order_by(
236    items: &[crate::ast::OrderBy],
237    r: &mut dyn VariableResolver,
238) -> Result<Vec<crate::ast::OrderBy>> {
239    items
240        .iter()
241        .map(|o| {
242            Ok(crate::ast::OrderBy {
243                expr: bind_expr(&o.expr, r)?,
244                descending: o.descending,
245                nulls: o.nulls,
246            })
247        })
248        .collect()
249}
250
251pub(super) fn bind_projections(
252    items: &[Projection],
253    r: &mut dyn VariableResolver,
254) -> Result<Vec<Projection>> {
255    items
256        .iter()
257        .map(|p| {
258            Ok(Projection {
259                expr: bind_expr(&p.expr, r)?,
260                alias: p.alias.clone(),
261            })
262        })
263        .collect()
264}
265
266pub(super) fn bind_assignments(
267    items: &[(String, Expr)],
268    r: &mut dyn VariableResolver,
269) -> Result<Vec<(String, Expr)>> {
270    items
271        .iter()
272        .map(|(name, e)| Ok((name.clone(), bind_expr(e, r)?)))
273        .collect()
274}
275
276pub(super) fn bind_ctes(items: &[CTE], r: &mut dyn VariableResolver) -> Result<Vec<CTE>> {
277    items
278        .iter()
279        .map(|cte| {
280            Ok(CTE {
281                name: cte.name.clone(),
282                columns: cte.columns.clone(),
283                recursive: cte.recursive,
284                materialization: cte.materialization,
285                search: cte.search.clone(),
286                cycle: cte
287                    .cycle
288                    .as_ref()
289                    .map(|cycle| -> Result<crate::ast::CteCycleClause> {
290                        Ok(crate::ast::CteCycleClause {
291                            columns: cycle.columns.clone(),
292                            mark_column: cycle.mark_column.clone(),
293                            mark_value: bind_expr(&cycle.mark_value, r)?,
294                            mark_default: bind_expr(&cycle.mark_default, r)?,
295                            path_column: cycle.path_column.clone(),
296                        })
297                    })
298                    .transpose()?,
299                query: Box::new(bind_select(&cte.query, r)?),
300            })
301        })
302        .collect()
303}
304
305pub(super) fn bind_rows(
306    rows: &[Vec<Expr>],
307    r: &mut dyn VariableResolver,
308) -> Result<Vec<Vec<Expr>>> {
309    rows.iter().map(|row| bind_exprs(row, r)).collect()
310}
311
312/// Rewrite a `SELECT` body, substituting resolvable variables.
313pub fn bind_select(stmt: &SelectStmt, r: &mut dyn VariableResolver) -> Result<SelectStmt> {
314    Ok(SelectStmt {
315        projections: bind_projections(&stmt.projections, r)?,
316        values: bind_rows(&stmt.values, r)?,
317        from: match stmt.from.as_ref() {
318            Some(f) => Some(bind_from(f, r)?),
319            None => None,
320        },
321        r#where: bind_opt_expr(stmt.r#where.as_ref(), r)?,
322        group_by: bind_exprs(&stmt.group_by, r)?,
323        grouping_sets: stmt
324            .grouping_sets
325            .iter()
326            .map(|set| bind_exprs(set, r))
327            .collect::<Result<Vec<_>>>()?,
328        group_distinct: stmt.group_distinct,
329        having: bind_opt_expr(stmt.having.as_ref(), r)?,
330        order_by: bind_order_by(&stmt.order_by, r)?,
331        limit: bind_opt_expr(stmt.limit.as_ref(), r)?,
332        with_ties: stmt.with_ties,
333        offset: bind_opt_expr(stmt.offset.as_ref(), r)?,
334        with: bind_ctes(&stmt.with, r)?,
335        set_op: match stmt.set_op.as_ref() {
336            Some(op) => Some(Box::new(crate::ast::SetOp {
337                kind: op.kind,
338                all: op.all,
339                left: op
340                    .left
341                    .as_ref()
342                    .map(|left| bind_select(left, r).map(Box::new))
343                    .transpose()?,
344                right: bind_select(&op.right, r)?,
345                combined_order_by: bind_order_by(&op.combined_order_by, r)?,
346                combined_limit: bind_opt_expr(op.combined_limit.as_ref(), r)?,
347                combined_with_ties: op.combined_with_ties,
348                combined_offset: bind_opt_expr(op.combined_offset.as_ref(), r)?,
349            })),
350            None => None,
351        },
352        distinct: stmt.distinct,
353        distinct_on: bind_exprs(&stmt.distinct_on, r)?,
354        locking: stmt.locking.clone(),
355    })
356}
357
358pub(super) fn bind_from(from: &FromClause, r: &mut dyn VariableResolver) -> Result<FromClause> {
359    Ok(match from {
360        FromClause::Table { .. } => from.clone(),
361        FromClause::Join {
362            left,
363            right,
364            kind,
365            on,
366            using,
367            natural,
368            alias,
369            column_aliases,
370            lateral,
371        } => FromClause::Join {
372            left: Box::new(bind_from(left, r)?),
373            right: Box::new(bind_from(right, r)?),
374            kind: *kind,
375            on: bind_opt_expr(on.as_ref(), r)?,
376            using: using.clone(),
377            natural: *natural,
378            alias: alias.clone(),
379            column_aliases: column_aliases.clone(),
380            lateral: *lateral,
381        },
382        FromClause::Values {
383            rows,
384            alias,
385            column_aliases,
386            internal_relation,
387            internal_column_types,
388        } => FromClause::Values {
389            rows: bind_rows(rows, r)?,
390            alias: alias.clone(),
391            column_aliases: column_aliases.clone(),
392            internal_relation: *internal_relation,
393            internal_column_types: internal_column_types.clone(),
394        },
395        FromClause::Function {
396            name,
397            binding,
398            output_name,
399            relations,
400            args,
401            alias,
402            column_aliases,
403            ordinality,
404            column_types,
405        } => FromClause::Function {
406            name: name.clone(),
407            binding: binding.clone(),
408            output_name: output_name.clone(),
409            relations: relations.clone(),
410            args: bind_exprs(args, r)?,
411            alias: alias.clone(),
412            column_aliases: column_aliases.clone(),
413            ordinality: *ordinality,
414            column_types: column_types.clone(),
415        },
416        FromClause::FunctionGroup {
417            functions,
418            alias,
419            column_aliases,
420            ordinality,
421        } => FromClause::FunctionGroup {
422            functions: functions
423                .iter()
424                .map(|function| {
425                    Ok(crate::ast::TableFunction {
426                        name: function.name.clone(),
427                        binding: function.binding.clone(),
428                        output_name: function.output_name.clone(),
429                        relations: function.relations.clone(),
430                        args: bind_exprs(&function.args, r)?,
431                        column_aliases: function.column_aliases.clone(),
432                        column_types: function.column_types.clone(),
433                    })
434                })
435                .collect::<Result<Vec<_>>>()?,
436            alias: alias.clone(),
437            column_aliases: column_aliases.clone(),
438            ordinality: *ordinality,
439        },
440        FromClause::Subquery {
441            body,
442            alias,
443            column_aliases,
444        } => FromClause::Subquery {
445            body: Box::new(bind_select(body, r)?),
446            alias: alias.clone(),
447            column_aliases: column_aliases.clone(),
448        },
449    })
450}
451
452/// Rewrite a full statement, substituting resolvable variables in
453/// every expression position. Statements without expression payloads
454/// pass through unchanged.
455#[expect(
456    clippy::too_many_lines,
457    reason = "PL/pgSQL lowering preserves parser order and datum validation"
458)]
459pub fn bind_statement(stmt: &Statement, r: &mut dyn VariableResolver) -> Result<Statement> {
460    Ok(match stmt {
461        Statement::Select(body) => Statement::Select(Box::new(bind_select(body, r)?)),
462        Statement::Insert(insert) => {
463            let mut out = insert.clone();
464            out.with = bind_ctes(&insert.with, r)?;
465            out.rows = bind_rows(&insert.rows, r)?;
466            out.select_source = match insert.select_source.as_ref() {
467                Some(body) => Some(Box::new(bind_select(body, r)?)),
468                None => None,
469            };
470            out.on_conflict = match insert.on_conflict.as_ref() {
471                Some(oc) => Some(crate::ast::OnConflict {
472                    predicate: bind_opt_expr(oc.predicate.as_deref(), r)?.map(Box::new),
473                    constraint: oc.constraint.clone(),
474                    conflict_columns: oc.conflict_columns.clone(),
475                    expressions: oc
476                        .expressions
477                        .iter()
478                        .map(|expr| bind_expr(expr, r))
479                        .collect::<Result<Vec<_>>>()?,
480                    action: match &oc.action {
481                        crate::ast::OnConflictAction::Nothing => {
482                            crate::ast::OnConflictAction::Nothing
483                        }
484                        crate::ast::OnConflictAction::Update {
485                            assignments,
486                            r#where,
487                        } => crate::ast::OnConflictAction::Update {
488                            assignments: bind_assignments(assignments, r)?,
489                            r#where: bind_opt_expr(r#where.as_deref(), r)?.map(Box::new),
490                        },
491                    },
492                }),
493                None => None,
494            };
495            out.returning = bind_projections(&insert.returning, r)?;
496            Statement::Insert(out)
497        }
498        Statement::Update(update) => {
499            let mut out = update.clone();
500            out.assignments = bind_assignments(&update.assignments, r)?;
501            out.r#where = bind_opt_expr(update.r#where.as_ref(), r)?;
502            out.with = bind_ctes(&update.with, r)?;
503            out.from = match update.from.as_ref() {
504                Some(f) => Some(bind_from(f, r)?),
505                None => None,
506            };
507            out.returning = bind_projections(&update.returning, r)?;
508            Statement::Update(out)
509        }
510        Statement::Delete(delete) => {
511            let mut out = delete.clone();
512            out.r#where = bind_opt_expr(delete.r#where.as_ref(), r)?;
513            out.with = bind_ctes(&delete.with, r)?;
514            out.using = match delete.using.as_ref() {
515                Some(f) => Some(bind_from(f, r)?),
516                None => None,
517            };
518            out.returning = bind_projections(&delete.returning, r)?;
519            Statement::Delete(out)
520        }
521        Statement::Values { rows } => Statement::Values {
522            rows: bind_rows(rows, r)?,
523        },
524        Statement::CreateTableAs {
525            name,
526            if_not_exists,
527            column_names,
528            with_no_data,
529            persistence,
530            on_commit,
531            body,
532        } => Statement::CreateTableAs {
533            name: name.clone(),
534            if_not_exists: *if_not_exists,
535            column_names: column_names.clone(),
536            with_no_data: *with_no_data,
537            persistence: *persistence,
538            on_commit: *on_commit,
539            body: Box::new(bind_select(body, r)?),
540        },
541        Statement::CreateMaterializedView {
542            name,
543            column_names,
544            if_not_exists,
545            with_no_data,
546            options,
547            body,
548        } => Statement::CreateMaterializedView {
549            name: name.clone(),
550            column_names: column_names.clone(),
551            if_not_exists: *if_not_exists,
552            with_no_data: *with_no_data,
553            options: options.clone(),
554            body: Box::new(bind_select(body, r)?),
555        },
556        Statement::Explain {
557            analyze,
558            verbose,
559            format,
560            body,
561        } => Statement::Explain {
562            analyze: *analyze,
563            verbose: *verbose,
564            format: format.clone(),
565            body: Box::new(bind_statement(body, r)?),
566        },
567        Statement::DeclareCursor(cursor) => {
568            let mut out = cursor.clone();
569            out.query = Box::new(bind_select(&cursor.query, r)?);
570            Statement::DeclareCursor(out)
571        }
572        Statement::Merge(merge) => {
573            let mut out = merge.clone();
574            out.source = bind_from(&merge.source, r)?;
575            out.join_condition = bind_expr(&merge.join_condition, r)?;
576            out.when_clauses = merge
577                .when_clauses
578                .iter()
579                .map(|w| bind_merge_when(w, r))
580                .collect::<Result<Vec<_>>>()?;
581            out.returning = bind_projections(&merge.returning, r)?;
582            Statement::Merge(out)
583        }
584        Statement::Call { name, args } => Statement::Call {
585            name: name.clone(),
586            args: bind_exprs(args, r)?,
587        },
588        other => other.clone(),
589    })
590}
591
592pub(super) fn bind_merge_when(when: &MergeWhen, r: &mut dyn VariableResolver) -> Result<MergeWhen> {
593    Ok(match when {
594        MergeWhen::UpdateMatched {
595            condition,
596            assignments,
597        } => MergeWhen::UpdateMatched {
598            condition: bind_opt_expr(condition.as_ref(), r)?,
599            assignments: bind_assignments(assignments, r)?,
600        },
601        MergeWhen::DeleteMatched { condition } => MergeWhen::DeleteMatched {
602            condition: bind_opt_expr(condition.as_ref(), r)?,
603        },
604        MergeWhen::UpdateNotMatchedBySource {
605            condition,
606            assignments,
607        } => MergeWhen::UpdateNotMatchedBySource {
608            condition: bind_opt_expr(condition.as_ref(), r)?,
609            assignments: bind_assignments(assignments, r)?,
610        },
611        MergeWhen::DeleteNotMatchedBySource { condition } => MergeWhen::DeleteNotMatchedBySource {
612            condition: bind_opt_expr(condition.as_ref(), r)?,
613        },
614        MergeWhen::InsertNotMatched {
615            condition,
616            columns,
617            values,
618        } => MergeWhen::InsertNotMatched {
619            condition: bind_opt_expr(condition.as_ref(), r)?,
620            columns: columns.clone(),
621            values: bind_exprs(values, r)?,
622        },
623        MergeWhen::NothingMatched { condition } => MergeWhen::NothingMatched {
624            condition: bind_opt_expr(condition.as_ref(), r)?,
625        },
626        MergeWhen::NothingNotMatched { condition } => MergeWhen::NothingNotMatched {
627            condition: bind_opt_expr(condition.as_ref(), r)?,
628        },
629        MergeWhen::NothingNotMatchedBySource { condition } => {
630            MergeWhen::NothingNotMatchedBySource {
631                condition: bind_opt_expr(condition.as_ref(), r)?,
632            }
633        }
634    })
635}