uqa_sql/plpgsql/
parameters.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Bind procedural variables as invocation parameters, preserving catalog inputs.
8
9use super::{
10    binding::{bind_statement, ResolvedVariable, VariableResolver},
11    variable_conflicts::{statement_variable_names, VariableSiteResolver},
12    VariableConflict,
13};
14use crate::{
15    ast::{Expr, Statement},
16    SQLError, SQLParam,
17};
18
19pub fn parameter_type_mismatch(
20    number: usize,
21    actual: &crate::ColumnType,
22    expected: &crate::ColumnType,
23) -> crate::ast::DeferredSQLError {
24    crate::ast::DeferredSQLError {
25        sqlstate: "42804".into(),
26        message: format!(
27            "type of parameter {number} ({}) does not match that when preparing the plan ({})",
28            actual.display_name(),
29            expected.display_name()
30        ),
31    }
32}
33
34#[derive(Debug, Clone)]
35pub enum PLpgSQLVariableReference {
36    Name(String),
37    Qualified { qualifier: String, column: String },
38    Parameter(usize),
39}
40impl PLpgSQLVariableReference {
41    pub fn read(&self, resolver: &mut dyn VariableResolver) -> Result<SQLParam, SQLError> {
42        let value = match self {
43            Self::Name(name) => resolver.resolve_name(name),
44            Self::Qualified { qualifier, column } => resolver.resolve_qualified(qualifier, column),
45            Self::Parameter(index) => resolver.resolve_param(*index),
46        }?
47        .ok_or_else(|| {
48            SQLError::Internal("prepared PL/pgSQL variable no longer resolves".into())
49        })?;
50        Ok(parameter(value, resolver))
51    }
52}
53fn parameter(value: ResolvedVariable, resolver: &dyn VariableResolver) -> SQLParam {
54    match value
55        .declared_type
56        .as_deref()
57        .and_then(|name| resolver.parameter_type(name))
58    {
59        Some(ty) => SQLParam::typed_scalar(value.value, ty),
60        None => SQLParam::scalar(value.value),
61    }
62}
63
64pub struct PLpgSQLVariableBindings {
65    pub plan: crate::plan::UnifiedPlan,
66    pub references: Vec<PLpgSQLVariableReference>,
67    pub parameters: Vec<SQLParam>,
68}
69
70/// Reuse the ordinary PL name/column conflict resolver, then number only names
71/// it selected as variables. Repeated reads receive fresh invocation values.
72pub fn parameterize_statement_variables(
73    statement: &Statement,
74    resolver: &mut dyn VariableResolver,
75    conflict: VariableConflict,
76    resolve_sites: VariableSiteResolver<'_>,
77    aggregates: &dyn crate::plan::AggregateClassifier,
78) -> Result<PLpgSQLVariableBindings, SQLError> {
79    let keep = statement_variable_names(statement, resolver, conflict, resolve_sites)?;
80    let mut binding = VariableBindings {
81        inner: resolver,
82        keep: &keep,
83        next_name: 0,
84        references: Vec::new(),
85        parameters: Vec::new(),
86        collations: Vec::new(),
87    };
88    let statement = bind_statement(statement, &mut binding)?;
89    let mut plan = crate::plan::UnifiedPlan::lower_with(statement, aggregates);
90    // A procedural variable is an implicitly collated input. Retain that
91    // identity on its parameter node before ordinary prepared analysis.
92    plan.rewrite_scalar_expressions(&mut |expression| {
93        let crate::ScalarExpr::Param(number) = expression else {
94            return;
95        };
96        let Some(index) = number.checked_sub(1) else {
97            return;
98        };
99        let Some(oid) = binding.collations.get(index).copied().flatten() else {
100            return;
101        };
102        let Some(ty) = binding.parameters[index].declared_scalar_type() else {
103            return;
104        };
105        if crate::schema::composites::type_is_collatable(ty) {
106            *expression = crate::ScalarExpr::Collated {
107                expr: Box::new(crate::ScalarExpr::Param(*number)),
108                ty: ty.clone(),
109                collation: crate::ir::ResultCollation::new(
110                    crate::catalog::collations::derivation::Derivation::implicit(i64::from(oid)),
111                    oid,
112                ),
113            };
114        }
115    });
116    Ok(PLpgSQLVariableBindings {
117        plan,
118        references: binding.references,
119        parameters: binding.parameters,
120    })
121}
122struct VariableBindings<'a> {
123    inner: &'a mut dyn VariableResolver,
124    keep: &'a [bool],
125    next_name: usize,
126    references: Vec<PLpgSQLVariableReference>,
127    parameters: Vec<SQLParam>,
128    collations: Vec<Option<u32>>,
129}
130impl VariableBindings<'_> {
131    fn mark(
132        &mut self,
133        reference: PLpgSQLVariableReference,
134        value: Option<ResolvedVariable>,
135        named: bool,
136    ) -> Option<Expr> {
137        let value = value?;
138        if named {
139            let site = self.next_name;
140            self.next_name += 1;
141            if self.keep.get(site).copied().unwrap_or(false) {
142                return None;
143            }
144        }
145        self.references.push(reference);
146        self.collations.push(value.collation_oid);
147        self.parameters.push(parameter(value, self.inner));
148        Some(Expr::Param(self.references.len()))
149    }
150}
151impl VariableResolver for VariableBindings<'_> {
152    fn resolve_name(&mut self, name: &str) -> Result<Option<ResolvedVariable>, SQLError> {
153        self.inner.resolve_name(name)
154    }
155    fn resolve_qualified(
156        &mut self,
157        qualifier: &str,
158        column: &str,
159    ) -> Result<Option<ResolvedVariable>, SQLError> {
160        self.inner.resolve_qualified(qualifier, column)
161    }
162    fn resolve_param(&mut self, index: usize) -> Result<Option<ResolvedVariable>, SQLError> {
163        self.inner.resolve_param(index)
164    }
165    fn rewrite_name(&mut self, name: &str) -> Result<Option<Expr>, SQLError> {
166        let value = self.inner.resolve_name(name)?;
167        Ok(self.mark(
168            PLpgSQLVariableReference::Name(name.to_string()),
169            value,
170            true,
171        ))
172    }
173    fn rewrite_qualified(
174        &mut self,
175        qualifier: &str,
176        column: &str,
177    ) -> Result<Option<Expr>, SQLError> {
178        let value = self.inner.resolve_qualified(qualifier, column)?;
179        Ok(self.mark(
180            PLpgSQLVariableReference::Qualified {
181                qualifier: qualifier.to_string(),
182                column: column.to_string(),
183            },
184            value,
185            true,
186        ))
187    }
188    fn rewrite_param(&mut self, index: usize) -> Result<Option<Expr>, SQLError> {
189        let value = self.inner.resolve_param(index)?;
190        Ok(self.mark(PLpgSQLVariableReference::Parameter(index), value, false))
191    }
192}
193
194#[cfg(test)]
195mod tests;
uqa_sql/plpgsql/parameters.rs

uqa_sql/plpgsql/
parameters.rs
