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    /// Expression-level counterpart of [`Self::resolve_param`].
69    fn rewrite_param(&mut self, index: usize) -> Result<Option<Expr>> {
70        Ok(self
71            .resolve_param(index)?
72            .map(ResolvedVariable::into_expression))
73    }
74
75    /// Observe or replace an executor-only structural column reference. SQL
76    /// variable resolvers normally leave these untouched.
77    fn rewrite_internal(&mut self, _column: InternalColumnRef) -> Result<Option<Expr>> {
78        Ok(None)
79    }
80}
81
82/// Rewrite an expression, substituting resolvable variable references
83/// with literals. References the resolver declines stay untouched.
84pub fn bind_expr(expr: &Expr, r: &mut dyn VariableResolver) -> Result<Expr> {
85    Ok(match expr {
86        Expr::Column(name) => match r.rewrite_name(name)? {
87            Some(value) => value,
88            None => expr.clone(),
89        },
90        Expr::QualifiedColumn {
91            qualifier, column, ..
92        } => match r.rewrite_qualified(qualifier, column)? {
93            Some(value) => value,
94            None => expr.clone(),
95        },
96        Expr::Param(index) => match r.rewrite_param(*index)? {
97            Some(value) => value,
98            None => expr.clone(),
99        },
100        Expr::InternalColumn(column) => match r.rewrite_internal(*column)? {
101            Some(value) => value,
102            None => expr.clone(),
103        },
104        Expr::Default | Expr::Literal(_) | Expr::Star | Expr::QualifiedStar(_) => expr.clone(),
105        Expr::Func {
106            name,
107            binding,
108            args,
109            distinct,
110            order_by,
111            filter,
112        } => Expr::Func {
113            name: name.clone(),
114            binding: binding.clone(),
115            args: bind_exprs(args, r)?,
116            distinct: *distinct,
117            order_by: bind_order_by(order_by, r)?,
118            filter: match filter {
119                Some(f) => Some(Box::new(bind_expr(f, r)?)),
120                None => None,
121            },
122        },
123        Expr::Array(items) => Expr::Array(bind_exprs(items, r)?),
124        Expr::Row(items) => Expr::Row(bind_exprs(items, r)?),
125        Expr::Binary { op, lhs, rhs } => Expr::Binary {
126            op: *op,
127            lhs: Box::new(bind_expr(lhs, r)?),
128            rhs: Box::new(bind_expr(rhs, r)?),
129        },
130        Expr::UnaryMinus(inner) => Expr::UnaryMinus(Box::new(bind_expr(inner, r)?)),
131        Expr::Not(inner) => Expr::Not(Box::new(bind_expr(inner, r)?)),
132        Expr::And(items) => Expr::And(bind_exprs(items, r)?),
133        Expr::Or(items) => Expr::Or(bind_exprs(items, r)?),
134        Expr::IsNull { expr, negated } => Expr::IsNull {
135            expr: Box::new(bind_expr(expr, r)?),
136            negated: *negated,
137        },
138        Expr::Between { expr, low, high } => Expr::Between {
139            expr: Box::new(bind_expr(expr, r)?),
140            low: Box::new(bind_expr(low, r)?),
141            high: Box::new(bind_expr(high, r)?),
142        },
143        Expr::InList {
144            expr,
145            list,
146            negated,
147        } => Expr::InList {
148            expr: Box::new(bind_expr(expr, r)?),
149            list: bind_exprs(list, r)?,
150            negated: *negated,
151        },
152        Expr::WindowCall { name, args, spec } => Expr::WindowCall {
153            name: name.clone(),
154            args: bind_exprs(args, r)?,
155            spec: crate::ast::WindowSpec {
156                reference: spec.reference.clone(),
157                partition_by: bind_exprs(&spec.partition_by, r)?,
158                order_by: bind_order_by(&spec.order_by, r)?,
159                frame: spec.frame.clone(),
160            },
161        },
162        Expr::Case {
163            base,
164            when,
165            else_branch,
166        } => Expr::Case {
167            base: match base {
168                Some(b) => Some(Box::new(bind_expr(b, r)?)),
169                None => None,
170            },
171            when: when
172                .iter()
173                .map(|(c, v)| Ok((bind_expr(c, r)?, bind_expr(v, r)?)))
174                .collect::<Result<Vec<_>>>()?,
175            else_branch: match else_branch {
176                Some(e) => Some(Box::new(bind_expr(e, r)?)),
177                None => None,
178            },
179        },
180        Expr::Cast { expr, ty } => Expr::Cast {
181            expr: Box::new(bind_expr(expr, r)?),
182            ty: ty.clone(),
183        },
184        Expr::ScalarSubquery(body) => Expr::ScalarSubquery(Box::new(bind_select(body, r)?)),
185        Expr::Exists { body, negated } => Expr::Exists {
186            body: Box::new(bind_select(body, r)?),
187            negated: *negated,
188        },
189        Expr::InSubquery {
190            expr,
191            body,
192            negated,
193        } => Expr::InSubquery {
194            expr: Box::new(bind_expr(expr, r)?),
195            body: Box::new(bind_select(body, r)?),
196            negated: *negated,
197        },
198    })
199}
200
201pub(super) fn bind_exprs(exprs: &[Expr], r: &mut dyn VariableResolver) -> Result<Vec<Expr>> {
202    exprs.iter().map(|e| bind_expr(e, r)).collect()
203}
204
205pub(super) fn bind_opt_expr(
206    expr: Option<&Expr>,
207    r: &mut dyn VariableResolver,
208) -> Result<Option<Expr>> {
209    match expr {
210        Some(e) => Ok(Some(bind_expr(e, r)?)),
211        None => Ok(None),
212    }
213}
214
215pub(super) fn bind_order_by(
216    items: &[crate::ast::OrderBy],
217    r: &mut dyn VariableResolver,
218) -> Result<Vec<crate::ast::OrderBy>> {
219    items
220        .iter()
221        .map(|o| {
222            Ok(crate::ast::OrderBy {
223                expr: bind_expr(&o.expr, r)?,
224                descending: o.descending,
225                nulls: o.nulls,
226            })
227        })
228        .collect()
229}
230
231pub(super) fn bind_projections(
232    items: &[Projection],
233    r: &mut dyn VariableResolver,
234) -> Result<Vec<Projection>> {
235    items
236        .iter()
237        .map(|p| {
238            Ok(Projection {
239                expr: bind_expr(&p.expr, r)?,
240                alias: p.alias.clone(),
241            })
242        })
243        .collect()
244}
245
246pub(super) fn bind_assignments(
247    items: &[(String, Expr)],
248    r: &mut dyn VariableResolver,
249) -> Result<Vec<(String, Expr)>> {
250    items
251        .iter()
252        .map(|(name, e)| Ok((name.clone(), bind_expr(e, r)?)))
253        .collect()
254}
255
256pub(super) fn bind_ctes(items: &[CTE], r: &mut dyn VariableResolver) -> Result<Vec<CTE>> {
257    items
258        .iter()
259        .map(|cte| {
260            Ok(CTE {
261                name: cte.name.clone(),
262                columns: cte.columns.clone(),
263                recursive: cte.recursive,
264                materialization: cte.materialization,
265                search: cte.search.clone(),
266                cycle: cte
267                    .cycle
268                    .as_ref()
269                    .map(|cycle| -> Result<crate::ast::CteCycleClause> {
270                        Ok(crate::ast::CteCycleClause {
271                            columns: cycle.columns.clone(),
272                            mark_column: cycle.mark_column.clone(),
273                            mark_value: bind_expr(&cycle.mark_value, r)?,
274                            mark_default: bind_expr(&cycle.mark_default, r)?,
275                            path_column: cycle.path_column.clone(),
276                        })
277                    })
278                    .transpose()?,
279                query: Box::new(bind_select(&cte.query, r)?),
280            })
281        })
282        .collect()
283}
284
285pub(super) fn bind_rows(
286    rows: &[Vec<Expr>],
287    r: &mut dyn VariableResolver,
288) -> Result<Vec<Vec<Expr>>> {
289    rows.iter().map(|row| bind_exprs(row, r)).collect()
290}
291
292/// Rewrite a `SELECT` body, substituting resolvable variables.
293pub fn bind_select(stmt: &SelectStmt, r: &mut dyn VariableResolver) -> Result<SelectStmt> {
294    Ok(SelectStmt {
295        projections: bind_projections(&stmt.projections, r)?,
296        values: bind_rows(&stmt.values, r)?,
297        from: match stmt.from.as_ref() {
298            Some(f) => Some(bind_from(f, r)?),
299            None => None,
300        },
301        r#where: bind_opt_expr(stmt.r#where.as_ref(), r)?,
302        group_by: bind_exprs(&stmt.group_by, r)?,
303        grouping_sets: stmt
304            .grouping_sets
305            .iter()
306            .map(|set| bind_exprs(set, r))
307            .collect::<Result<Vec<_>>>()?,
308        group_distinct: stmt.group_distinct,
309        having: bind_opt_expr(stmt.having.as_ref(), r)?,
310        order_by: bind_order_by(&stmt.order_by, r)?,
311        limit: bind_opt_expr(stmt.limit.as_ref(), r)?,
312        with_ties: stmt.with_ties,
313        offset: bind_opt_expr(stmt.offset.as_ref(), r)?,
314        with: bind_ctes(&stmt.with, r)?,
315        set_op: match stmt.set_op.as_ref() {
316            Some(op) => Some(Box::new(crate::ast::SetOp {
317                kind: op.kind,
318                all: op.all,
319                left: op
320                    .left
321                    .as_ref()
322                    .map(|left| bind_select(left, r).map(Box::new))
323                    .transpose()?,
324                right: bind_select(&op.right, r)?,
325                combined_order_by: bind_order_by(&op.combined_order_by, r)?,
326                combined_limit: bind_opt_expr(op.combined_limit.as_ref(), r)?,
327                combined_with_ties: op.combined_with_ties,
328                combined_offset: bind_opt_expr(op.combined_offset.as_ref(), r)?,
329            })),
330            None => None,
331        },
332        distinct: stmt.distinct,
333        distinct_on: bind_exprs(&stmt.distinct_on, r)?,
334        locking: stmt.locking.clone(),
335    })
336}
337
338pub(super) fn bind_from(from: &FromClause, r: &mut dyn VariableResolver) -> Result<FromClause> {
339    Ok(match from {
340        FromClause::Table { .. } => from.clone(),
341        FromClause::Join {
342            left,
343            right,
344            kind,
345            on,
346            using,
347            natural,
348            alias,
349            column_aliases,
350            lateral,
351        } => FromClause::Join {
352            left: Box::new(bind_from(left, r)?),
353            right: Box::new(bind_from(right, r)?),
354            kind: *kind,
355            on: bind_opt_expr(on.as_ref(), r)?,
356            using: using.clone(),
357            natural: *natural,
358            alias: alias.clone(),
359            column_aliases: column_aliases.clone(),
360            lateral: *lateral,
361        },
362        FromClause::Values {
363            rows,
364            alias,
365            column_aliases,
366            internal_relation,
367            internal_column_types,
368        } => FromClause::Values {
369            rows: bind_rows(rows, r)?,
370            alias: alias.clone(),
371            column_aliases: column_aliases.clone(),
372            internal_relation: *internal_relation,
373            internal_column_types: internal_column_types.clone(),
374        },
375        FromClause::Function {
376            name,
377            output_name,
378            relation,
379            args,
380            alias,
381            column_aliases,
382            ordinality,
383            column_types,
384        } => FromClause::Function {
385            name: name.clone(),
386            output_name: output_name.clone(),
387            relation: relation.clone(),
388            args: bind_exprs(args, r)?,
389            alias: alias.clone(),
390            column_aliases: column_aliases.clone(),
391            ordinality: *ordinality,
392            column_types: column_types.clone(),
393        },
394        FromClause::FunctionGroup {
395            functions,
396            alias,
397            column_aliases,
398            ordinality,
399        } => FromClause::FunctionGroup {
400            functions: functions
401                .iter()
402                .map(|function| {
403                    Ok(crate::ast::TableFunction {
404                        name: function.name.clone(),
405                        output_name: function.output_name.clone(),
406                        relation: function.relation.clone(),
407                        args: bind_exprs(&function.args, r)?,
408                        column_aliases: function.column_aliases.clone(),
409                        column_types: function.column_types.clone(),
410                    })
411                })
412                .collect::<Result<Vec<_>>>()?,
413            alias: alias.clone(),
414            column_aliases: column_aliases.clone(),
415            ordinality: *ordinality,
416        },
417        FromClause::Subquery {
418            body,
419            alias,
420            column_aliases,
421        } => FromClause::Subquery {
422            body: Box::new(bind_select(body, r)?),
423            alias: alias.clone(),
424            column_aliases: column_aliases.clone(),
425        },
426    })
427}
428
429/// Rewrite a full statement, substituting resolvable variables in
430/// every expression position. Statements without expression payloads
431/// pass through unchanged.
432pub fn bind_statement(stmt: &Statement, r: &mut dyn VariableResolver) -> Result<Statement> {
433    Ok(match stmt {
434        Statement::Select(body) => Statement::Select(Box::new(bind_select(body, r)?)),
435        Statement::Insert(insert) => {
436            let mut out = insert.clone();
437            out.with = bind_ctes(&insert.with, r)?;
438            out.rows = bind_rows(&insert.rows, r)?;
439            out.select_source = match insert.select_source.as_ref() {
440                Some(body) => Some(Box::new(bind_select(body, r)?)),
441                None => None,
442            };
443            out.on_conflict = match insert.on_conflict.as_ref() {
444                Some(oc) => Some(crate::ast::OnConflict {
445                    conflict_columns: oc.conflict_columns.clone(),
446                    action: match &oc.action {
447                        crate::ast::OnConflictAction::Nothing => {
448                            crate::ast::OnConflictAction::Nothing
449                        }
450                        crate::ast::OnConflictAction::Update {
451                            assignments,
452                            r#where,
453                        } => crate::ast::OnConflictAction::Update {
454                            assignments: bind_assignments(assignments, r)?,
455                            r#where: bind_opt_expr(r#where.as_ref(), r)?,
456                        },
457                    },
458                }),
459                None => None,
460            };
461            out.returning = bind_projections(&insert.returning, r)?;
462            Statement::Insert(out)
463        }
464        Statement::Update(update) => {
465            let mut out = update.clone();
466            out.assignments = bind_assignments(&update.assignments, r)?;
467            out.r#where = bind_opt_expr(update.r#where.as_ref(), r)?;
468            out.with = bind_ctes(&update.with, r)?;
469            out.from = match update.from.as_ref() {
470                Some(f) => Some(bind_from(f, r)?),
471                None => None,
472            };
473            out.returning = bind_projections(&update.returning, r)?;
474            Statement::Update(out)
475        }
476        Statement::Delete(delete) => {
477            let mut out = delete.clone();
478            out.r#where = bind_opt_expr(delete.r#where.as_ref(), r)?;
479            out.with = bind_ctes(&delete.with, r)?;
480            out.using = match delete.using.as_ref() {
481                Some(f) => Some(bind_from(f, r)?),
482                None => None,
483            };
484            out.returning = bind_projections(&delete.returning, r)?;
485            Statement::Delete(out)
486        }
487        Statement::Values { rows } => Statement::Values {
488            rows: bind_rows(rows, r)?,
489        },
490        Statement::CreateTableAs {
491            name,
492            if_not_exists,
493            column_names,
494            with_no_data,
495            persistence,
496            on_commit,
497            body,
498        } => Statement::CreateTableAs {
499            name: name.clone(),
500            if_not_exists: *if_not_exists,
501            column_names: column_names.clone(),
502            with_no_data: *with_no_data,
503            persistence: *persistence,
504            on_commit: *on_commit,
505            body: Box::new(bind_select(body, r)?),
506        },
507        Statement::CreateMaterializedView {
508            name,
509            column_names,
510            if_not_exists,
511            with_no_data,
512            options,
513            body,
514        } => Statement::CreateMaterializedView {
515            name: name.clone(),
516            column_names: column_names.clone(),
517            if_not_exists: *if_not_exists,
518            with_no_data: *with_no_data,
519            options: options.clone(),
520            body: Box::new(bind_select(body, r)?),
521        },
522        Statement::Explain {
523            analyze,
524            verbose,
525            format,
526            body,
527        } => Statement::Explain {
528            analyze: *analyze,
529            verbose: *verbose,
530            format: format.clone(),
531            body: Box::new(bind_statement(body, r)?),
532        },
533        Statement::DeclareCursor(cursor) => {
534            let mut out = cursor.clone();
535            out.query = Box::new(bind_select(&cursor.query, r)?);
536            Statement::DeclareCursor(out)
537        }
538        Statement::Merge(merge) => {
539            let mut out = merge.clone();
540            out.source = bind_from(&merge.source, r)?;
541            out.join_condition = bind_expr(&merge.join_condition, r)?;
542            out.when_clauses = merge
543                .when_clauses
544                .iter()
545                .map(|w| bind_merge_when(w, r))
546                .collect::<Result<Vec<_>>>()?;
547            out.returning = bind_projections(&merge.returning, r)?;
548            Statement::Merge(out)
549        }
550        Statement::Call { name, args } => Statement::Call {
551            name: name.clone(),
552            args: bind_exprs(args, r)?,
553        },
554        other => other.clone(),
555    })
556}
557
558pub(super) fn bind_merge_when(when: &MergeWhen, r: &mut dyn VariableResolver) -> Result<MergeWhen> {
559    Ok(match when {
560        MergeWhen::UpdateMatched {
561            condition,
562            assignments,
563        } => MergeWhen::UpdateMatched {
564            condition: bind_opt_expr(condition.as_ref(), r)?,
565            assignments: bind_assignments(assignments, r)?,
566        },
567        MergeWhen::DeleteMatched { condition } => MergeWhen::DeleteMatched {
568            condition: bind_opt_expr(condition.as_ref(), r)?,
569        },
570        MergeWhen::UpdateNotMatchedBySource {
571            condition,
572            assignments,
573        } => MergeWhen::UpdateNotMatchedBySource {
574            condition: bind_opt_expr(condition.as_ref(), r)?,
575            assignments: bind_assignments(assignments, r)?,
576        },
577        MergeWhen::DeleteNotMatchedBySource { condition } => MergeWhen::DeleteNotMatchedBySource {
578            condition: bind_opt_expr(condition.as_ref(), r)?,
579        },
580        MergeWhen::InsertNotMatched {
581            condition,
582            columns,
583            values,
584        } => MergeWhen::InsertNotMatched {
585            condition: bind_opt_expr(condition.as_ref(), r)?,
586            columns: columns.clone(),
587            values: bind_exprs(values, r)?,
588        },
589        MergeWhen::NothingMatched { condition } => MergeWhen::NothingMatched {
590            condition: bind_opt_expr(condition.as_ref(), r)?,
591        },
592        MergeWhen::NothingNotMatched { condition } => MergeWhen::NothingNotMatched {
593            condition: bind_opt_expr(condition.as_ref(), r)?,
594        },
595        MergeWhen::NothingNotMatchedBySource { condition } => {
596            MergeWhen::NothingNotMatchedBySource {
597                condition: bind_opt_expr(condition.as_ref(), r)?,
598            }
599        }
600    })
601}