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

uqa_sql/
plpgsql.rs
