uqa_sql/
plpgsql.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `PL/pgSQL` function bodies: typed AST, parser, and the variable
8//! binding rewriter.
9//!
10//! Bodies are parsed with `libpg_query`'s `PL/pgSQL` parser
11//! (`pg_query::parse_plpgsql`), which returns the same JSON dump
12//! `PostgreSQL` itself produces. This module lowers that JSON into a
13//! typed AST whose embedded SQL fragments are pre-compiled into
14//! [`Expr`] / [`Statement`] values, ready for execution against the
15//! engine.
16//!
17//! Variable references inside embedded SQL are plain column
18//! references after compilation. At execution time the interpreter
19//! rewrites them into literals through [`VariableResolver`] /
20//! [`bind_expr`] / [`bind_statement`] before handing the statement to
21//! the engine. This matches `plpgsql.variable_conflict =
22//! use_variable` resolution: a name that is both a `PL/pgSQL`
23//! variable and a column of a queried table resolves to the variable
24//! (stock `PostgreSQL` raises an ambiguity error instead).
25
26use serde_json::Value as JSONValue;
27use uqa_core::Value;
28
29use crate::ast::{
30    CreateFunction, Expr, FromClause, FunctionBody, FunctionParamMode, FunctionReturns, MergeWhen,
31    Projection, RoutineColumnTypeReference, SelectStmt, Statement, CTE,
32};
33use crate::error::{Result, SQLError};
34
35// ---------------------------------------------------------------------
36// Typed AST
37// ---------------------------------------------------------------------
38
39/// A parsed `PL/pgSQL` function body: the flat datum table plus the
40/// outermost block.
41#[derive(Debug, Clone)]
42pub struct PLpgSQLFunction {
43    pub datums: Vec<PLpgSQLDatum>,
44    pub action: PLpgSQLBlock,
45    /// Datum holding the implicit `NEW` record for a trigger function.
46    pub new_datum: Option<usize>,
47    /// Datum holding the implicit `OLD` record for a trigger function.
48    pub old_datum: Option<usize>,
49    /// Index of the implicit `FOUND` variable in [`Self::datums`].
50    pub found_datum: Option<usize>,
51}
52
53impl PLpgSQLFunction {
54    /// Datum indices used as `FOR i IN a..b` loop counters. The
55    /// interpreter binds these names only while their loop runs so an
56    /// outer variable with the same name stays visible elsewhere.
57    pub fn fori_variable_datums(&self) -> std::collections::BTreeSet<usize> {
58        let mut out = std::collections::BTreeSet::new();
59        collect_fori_vars_block(&self.action, &mut out);
60        out
61    }
62
63    /// Datum indices used as bound-cursor arguments. They are visible only
64    /// while the cursor query is bound, not throughout the routine body.
65    pub fn cursor_argument_datums(&self) -> std::collections::BTreeSet<usize> {
66        let mut out = std::collections::BTreeSet::new();
67        for datum in &self.datums {
68            let PLpgSQLDatum::Var(var) = datum else {
69                continue;
70            };
71            let Some(argument_row) = var.cursor.as_ref().and_then(|cursor| cursor.argument_row)
72            else {
73                continue;
74            };
75            if let Some(PLpgSQLDatum::Row { fields }) = self.datums.get(argument_row) {
76                out.extend(fields.iter().map(|field| field.varno));
77            }
78        }
79        out
80    }
81}
82
83fn collect_fori_vars_block(block: &PLpgSQLBlock, out: &mut std::collections::BTreeSet<usize>) {
84    collect_fori_vars_stmts(&block.body, out);
85    for arm in &block.exceptions {
86        collect_fori_vars_stmts(&arm.body, out);
87    }
88}
89
90fn collect_fori_vars_stmts(stmts: &[PLpgSQLStmt], out: &mut std::collections::BTreeSet<usize>) {
91    for stmt in stmts {
92        match stmt {
93            PLpgSQLStmt::Block(block) => collect_fori_vars_block(block, out),
94            PLpgSQLStmt::If {
95                then_body,
96                elsifs,
97                else_body,
98                ..
99            } => {
100                collect_fori_vars_stmts(then_body, out);
101                for (_, body) in elsifs {
102                    collect_fori_vars_stmts(body, out);
103                }
104                if let Some(body) = else_body {
105                    collect_fori_vars_stmts(body, out);
106                }
107            }
108            PLpgSQLStmt::Case {
109                arms, else_body, ..
110            } => {
111                for (_, body) in arms {
112                    collect_fori_vars_stmts(body, out);
113                }
114                if let Some(body) = else_body {
115                    collect_fori_vars_stmts(body, out);
116                }
117            }
118            PLpgSQLStmt::Loop { body, .. } | PLpgSQLStmt::While { body, .. } => {
119                collect_fori_vars_stmts(body, out);
120            }
121            PLpgSQLStmt::ForI { var, body, .. } => {
122                out.insert(*var);
123                collect_fori_vars_stmts(body, out);
124            }
125            PLpgSQLStmt::ForQuery { body, .. } => collect_fori_vars_stmts(body, out),
126            _ => {}
127        }
128    }
129}
130
131/// One entry in the function's flat datum table. `varno` / `dno`
132/// references inside statements index into this table.
133#[derive(Debug, Clone)]
134pub enum PLpgSQLDatum {
135    Var(Box<PLpgSQLVar>),
136    /// `RECORD` variable (also `FOR rec IN ...` loop targets).
137    Rec {
138        name: String,
139    },
140    /// `rec.field` assignment target.
141    RecField {
142        field: String,
143        parent: usize,
144    },
145    /// Multi-variable target list (`SELECT ... INTO a, b`).
146    Row {
147        fields: Vec<PLpgSQLRowField>,
148    },
149}
150
151impl PLpgSQLDatum {
152    pub fn name(&self) -> Option<&str> {
153        match self {
154            PLpgSQLDatum::Var(v) => Some(&v.name),
155            PLpgSQLDatum::Rec { name } => Some(name),
156            PLpgSQLDatum::RecField { .. } | PLpgSQLDatum::Row { .. } => None,
157        }
158    }
159}
160
161/// Scalar `PL/pgSQL` variable (declared variable, parameter, loop
162/// counter, or an internal compiler temporary).
163#[derive(Debug, Clone)]
164pub struct PLpgSQLVar {
165    pub name: String,
166    /// Normalized type name (`integer`, `text`, ...). The engine resolves
167    /// catalog-backed references such as `%TYPE` before execution.
168    pub type_name: String,
169    /// Exact relation-column identity emitted by the PL/pgSQL parser for a table-backed `%TYPE` declaration.
170    pub type_reference: Option<RoutineColumnTypeReference>,
171    pub default: Option<Expr>,
172    pub constant: bool,
173    pub not_null: bool,
174    /// Definition of a bound cursor declared with `CURSOR (...) FOR query`.
175    pub cursor: Option<PLpgSQLCursor>,
176    /// Source line of the declaration; used to disambiguate loop
177    /// variables that shadow outer names.
178    pub lineno: Option<i64>,
179}
180
181#[derive(Debug, Clone)]
182pub struct PLpgSQLCursor {
183    pub query: Statement,
184    pub argument_row: Option<usize>,
185}
186
187#[derive(Debug, Clone)]
188pub struct PLpgSQLCursorArgument {
189    pub name: Option<String>,
190    pub expr: Expr,
191}
192
193/// `name -> datum` slot of a row target.
194#[derive(Debug, Clone)]
195pub struct PLpgSQLRowField {
196    pub name: String,
197    pub varno: usize,
198}
199
200/// `[DECLARE ...] BEGIN ... [EXCEPTION ...] END` block.
201#[derive(Debug, Clone)]
202pub struct PLpgSQLBlock {
203    pub label: Option<String>,
204    pub body: Vec<PLpgSQLStmt>,
205    pub exceptions: Vec<PLpgSQLExceptionArm>,
206}
207
208/// One `WHEN cond [OR cond ...] THEN stmts` arm of an exception
209/// section.
210#[derive(Debug, Clone)]
211pub struct PLpgSQLExceptionArm {
212    /// Lower-cased condition names (`others`, `division_by_zero`,
213    /// ...). Explicit `SQLSTATE 'xxxxx'` conditions arrive as the
214    /// five-character code.
215    pub conditions: Vec<String>,
216    pub body: Vec<PLpgSQLStmt>,
217}
218
219/// `RAISE` severity.
220#[derive(Debug, Clone, Copy, PartialEq, Eq)]
221pub enum RaiseLevel {
222    Debug,
223    Log,
224    Info,
225    Notice,
226    Warning,
227    Error,
228}
229
230impl RaiseLevel {
231    pub fn as_str(self) -> &'static str {
232        match self {
233            RaiseLevel::Debug => "DEBUG",
234            RaiseLevel::Log => "LOG",
235            RaiseLevel::Info => "INFO",
236            RaiseLevel::Notice => "NOTICE",
237            RaiseLevel::Warning => "WARNING",
238            RaiseLevel::Error => "ERROR",
239        }
240    }
241}
242
243/// Assignment / `INTO` target.
244#[derive(Debug, Clone)]
245pub enum IntoTarget {
246    /// A `RECORD` variable receives the whole row.
247    Rec(usize),
248    /// Positional list of scalar targets.
249    Row(Vec<PLpgSQLRowField>),
250}
251
252/// Executable `PL/pgSQL` statement.
253#[derive(Debug, Clone)]
254pub enum PLpgSQLStmt {
255    Block(PLpgSQLBlock),
256    /// `target := expr` (also `=`). `target` indexes the datum table.
257    Assign {
258        target: usize,
259        expr: Expr,
260    },
261    If {
262        cond: Expr,
263        then_body: Vec<PLpgSQLStmt>,
264        elsifs: Vec<(Expr, Vec<PLpgSQLStmt>)>,
265        else_body: Option<Vec<PLpgSQLStmt>>,
266    },
267    /// CASE statement. Simple form carries `t_expr` + the temporary
268    /// datum the compiler references from each rewritten WHEN.
269    Case {
270        t_expr: Option<Expr>,
271        t_varno: Option<usize>,
272        arms: Vec<(Expr, Vec<PLpgSQLStmt>)>,
273        else_body: Option<Vec<PLpgSQLStmt>>,
274    },
275    Loop {
276        label: Option<String>,
277        body: Vec<PLpgSQLStmt>,
278    },
279    While {
280        label: Option<String>,
281        cond: Expr,
282        body: Vec<PLpgSQLStmt>,
283    },
284    /// `FOR i IN [REVERSE] lower..upper [BY step] LOOP`.
285    ForI {
286        label: Option<String>,
287        var: usize,
288        lower: Expr,
289        upper: Expr,
290        step: Option<Expr>,
291        reverse: bool,
292        body: Vec<PLpgSQLStmt>,
293    },
294    /// `FOR target IN <query> LOOP`.
295    ForQuery {
296        label: Option<String>,
297        target: IntoTarget,
298        query: Statement,
299        body: Vec<PLpgSQLStmt>,
300    },
301    /// `EXIT` (`is_exit`) or `CONTINUE`, optionally labelled and
302    /// conditional (`WHEN cond`).
303    Exit {
304        is_exit: bool,
305        label: Option<String>,
306        cond: Option<Expr>,
307    },
308    Return {
309        value: Option<PLpgSQLReturnValue>,
310    },
311    /// `RETURN NEXT [expr]` - bare form emits the current OUT /
312    /// TABLE column values.
313    ReturnNext {
314        value: Option<PLpgSQLReturnValue>,
315    },
316    ReturnQuery {
317        query: Statement,
318    },
319    ReturnQueryExecute {
320        query: Expr,
321        params: Vec<Expr>,
322    },
323    Raise {
324        level: RaiseLevel,
325        condition: Option<String>,
326        message: Option<String>,
327        params: Vec<Expr>,
328    },
329    /// Embedded SQL statement, optionally `INTO [STRICT] target`.
330    ExecSQL {
331        stmt: Statement,
332        into: Option<IntoTarget>,
333        strict: bool,
334    },
335    /// `EXECUTE <string> [INTO [STRICT] target] [USING params]`.
336    DynExecute {
337        query: Expr,
338        params: Vec<Expr>,
339        into: Option<IntoTarget>,
340        strict: bool,
341    },
342    Perform {
343        query: Statement,
344    },
345    OpenCursor {
346        cursor: usize,
347        arguments: Vec<PLpgSQLCursorArgument>,
348    },
349    FetchCursor {
350        cursor: usize,
351        target: IntoTarget,
352        direction: i64,
353        count: i64,
354    },
355    CloseCursor {
356        cursor: usize,
357    },
358    /// `GET DIAGNOSTICS var = KIND [, ...]` as `(kind, target datum)`.
359    GetDiagnostics {
360        items: Vec<(String, usize)>,
361    },
362}
363
364/// Value source for `RETURN` and `RETURN NEXT`. `PostgreSQL` 18 stores a simple
365/// datum reference in `retvarno`, distinct from a general SQL expression.
366#[derive(Debug, Clone)]
367pub enum PLpgSQLReturnValue {
368    Expr(Expr),
369    Datum(usize),
370}
371
372// ---------------------------------------------------------------------
373// Parsing: definition -> canonical text -> libpg_query JSON -> AST
374// ---------------------------------------------------------------------
375
376/// Parse the `PL/pgSQL` body of a stored definition. The definition
377/// is re-serialized into a canonical `CREATE FUNCTION` statement so
378/// restore-from-catalog and fresh DDL take the same path.
379mod binding;
380mod conditions;
381mod json_validation;
382mod lowering_expression;
383mod lowering_statement;
384mod parsing;
385
386use json_validation::{
387    ensure_single_tag, expect_tag, json_bool_or_false, json_i64_or_zero, json_kind,
388    json_optional_i64, json_optional_str, json_optional_usize, json_usize_or_zero,
389    normalize_plpgsql_type, optional_array, require, require_i64, require_nonempty_str,
390    validate_assignable_datum, validate_record_datum, validate_scalar_datum,
391};
392use lowering_expression::{lower_expr, lower_expr_list, lower_full_statement};
393use lowering_statement::lower_block;
394use parsing::{lower_row_fields, normalize_condition};
395
396pub use binding::{bind_expr, bind_select, bind_statement, ResolvedVariable, VariableResolver};
397pub use conditions::{condition_sqlstate, condition_sqlstates};
398pub use lowering_expression::compile_expression_text;
399pub use parsing::{parse_do_block, parse_function};
400
401#[cfg(test)]
402mod tests;
uqa_sql/plpgsql.rs

uqa_sql/
plpgsql.rs
