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, CursorDirection, Expr, FromClause, FunctionBody, FunctionParamMode,
31    FunctionReturns, MergeWhen, 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 synthesized for loop-local variables. The interpreter
55    /// binds these names only while their loop runs so an outer variable with
56    /// the same name stays visible elsewhere.
57    pub fn loop_local_variable_datums(&self) -> std::collections::BTreeSet<usize> {
58        let mut out = std::collections::BTreeSet::new();
59        collect_loop_local_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_loop_local_vars_block(
84    block: &PLpgSQLBlock,
85    out: &mut std::collections::BTreeSet<usize>,
86) {
87    collect_loop_local_vars_stmts(&block.body, out);
88    for arm in &block.exceptions {
89        collect_loop_local_vars_stmts(&arm.body, out);
90    }
91}
92
93fn collect_loop_local_vars_stmts(
94    stmts: &[PLpgSQLStmt],
95    out: &mut std::collections::BTreeSet<usize>,
96) {
97    for stmt in stmts {
98        match stmt {
99            PLpgSQLStmt::Block(block) => collect_loop_local_vars_block(block, out),
100            PLpgSQLStmt::If {
101                then_body,
102                elsifs,
103                else_body,
104                ..
105            } => {
106                collect_loop_local_vars_stmts(then_body, out);
107                for (_, body) in elsifs {
108                    collect_loop_local_vars_stmts(body, out);
109                }
110                if let Some(body) = else_body {
111                    collect_loop_local_vars_stmts(body, out);
112                }
113            }
114            PLpgSQLStmt::Case {
115                arms, else_body, ..
116            } => {
117                for (_, body) in arms {
118                    collect_loop_local_vars_stmts(body, out);
119                }
120                if let Some(body) = else_body {
121                    collect_loop_local_vars_stmts(body, out);
122                }
123            }
124            PLpgSQLStmt::Loop { body, .. } | PLpgSQLStmt::While { body, .. } => {
125                collect_loop_local_vars_stmts(body, out);
126            }
127            PLpgSQLStmt::ForI { var, body, .. } => {
128                out.insert(*var);
129                collect_loop_local_vars_stmts(body, out);
130            }
131            PLpgSQLStmt::ForCursor { target, body, .. } => {
132                out.insert(*target);
133                collect_loop_local_vars_stmts(body, out);
134            }
135            PLpgSQLStmt::ForQuery { body, .. }
136            | PLpgSQLStmt::ForDynamic { body, .. }
137            | PLpgSQLStmt::ForeachArray { body, .. } => {
138                collect_loop_local_vars_stmts(body, out);
139            }
140            _ => {}
141        }
142    }
143}
144
145/// One entry in the function's flat datum table. `varno` / `dno`
146/// references inside statements index into this table.
147#[derive(Debug, Clone)]
148pub enum PLpgSQLDatum {
149    Var(Box<PLpgSQLVar>),
150    /// `RECORD` variable (also `FOR rec IN ...` loop targets).
151    Rec {
152        name: String,
153    },
154    /// `rec.field` assignment target.
155    RecField {
156        field: String,
157        parent: usize,
158    },
159    /// Multi-variable target list (`SELECT ... INTO a, b`).
160    Row {
161        fields: Vec<PLpgSQLRowField>,
162    },
163}
164
165impl PLpgSQLDatum {
166    pub fn name(&self) -> Option<&str> {
167        match self {
168            PLpgSQLDatum::Var(v) => Some(&v.name),
169            PLpgSQLDatum::Rec { name } => Some(name),
170            PLpgSQLDatum::RecField { .. } | PLpgSQLDatum::Row { .. } => None,
171        }
172    }
173}
174
175/// Scalar `PL/pgSQL` variable (declared variable, parameter, loop
176/// counter, or an internal compiler temporary).
177#[derive(Debug, Clone)]
178pub struct PLpgSQLVar {
179    pub name: String,
180    /// Normalized type name (`integer`, `text`, ...). The engine resolves
181    /// catalog-backed references such as `%TYPE` before execution.
182    pub type_name: String,
183    /// Exact relation-column identity emitted by the PL/pgSQL parser for a table-backed `%TYPE` declaration.
184    pub type_reference: Option<RoutineColumnTypeReference>,
185    pub default: Option<Expr>,
186    pub constant: bool,
187    pub not_null: bool,
188    /// Definition of a bound cursor declared with `CURSOR (...) FOR query`.
189    pub cursor: Option<PLpgSQLCursor>,
190    /// Source line of the declaration; used to disambiguate loop
191    /// variables that shadow outer names.
192    pub lineno: Option<i64>,
193}
194
195#[derive(Debug, Clone)]
196pub struct PLpgSQLCursor {
197    pub query: Statement,
198    pub argument_row: Option<usize>,
199    /// Explicit declaration scroll mode. `None` leaves scrollability query-dependent.
200    pub scroll: Option<bool>,
201}
202
203#[derive(Debug, Clone)]
204pub struct PLpgSQLCursorArgument {
205    pub name: Option<String>,
206    pub expr: Expr,
207}
208
209/// Query source selected by one `OPEN` statement.
210#[derive(Debug, Clone)]
211pub enum PLpgSQLCursorOpen {
212    Bound {
213        arguments: Vec<PLpgSQLCursorArgument>,
214    },
215    Static {
216        query: Box<Statement>,
217        scroll: Option<bool>,
218    },
219    Dynamic {
220        query: Expr,
221        params: Vec<Expr>,
222        scroll: Option<bool>,
223    },
224}
225
226/// Constant or run-time expression controlling cursor movement.
227#[derive(Debug, Clone)]
228pub enum PLpgSQLCursorCount {
229    Constant(i64),
230    Expression(Expr),
231}
232
233/// `name -> datum` slot of a row target.
234#[derive(Debug, Clone)]
235pub struct PLpgSQLRowField {
236    pub name: String,
237    pub varno: usize,
238}
239
240/// `[DECLARE ...] BEGIN ... [EXCEPTION ...] END` block.
241#[derive(Debug, Clone)]
242pub struct PLpgSQLBlock {
243    pub label: Option<String>,
244    pub body: Vec<PLpgSQLStmt>,
245    pub exceptions: Vec<PLpgSQLExceptionArm>,
246}
247
248/// One `WHEN cond [OR cond ...] THEN stmts` arm of an exception
249/// section.
250#[derive(Debug, Clone)]
251pub struct PLpgSQLExceptionArm {
252    /// Lower-cased condition names (`others`, `division_by_zero`,
253    /// ...). Explicit `SQLSTATE 'xxxxx'` conditions arrive as the
254    /// five-character code.
255    pub conditions: Vec<String>,
256    pub body: Vec<PLpgSQLStmt>,
257}
258
259/// `RAISE` severity.
260#[derive(Debug, Clone, Copy, PartialEq, Eq)]
261pub enum RaiseLevel {
262    Debug,
263    Log,
264    Info,
265    Notice,
266    Warning,
267    Error,
268}
269
270impl RaiseLevel {
271    pub fn as_str(self) -> &'static str {
272        match self {
273            RaiseLevel::Debug => "DEBUG",
274            RaiseLevel::Log => "LOG",
275            RaiseLevel::Info => "INFO",
276            RaiseLevel::Notice => "NOTICE",
277            RaiseLevel::Warning => "WARNING",
278            RaiseLevel::Error => "ERROR",
279        }
280    }
281}
282
283/// Assignment / `INTO` target.
284#[derive(Debug, Clone)]
285pub enum IntoTarget {
286    /// A `RECORD` variable receives the whole row.
287    Rec(usize),
288    /// Positional list of scalar targets.
289    Row(Vec<PLpgSQLRowField>),
290}
291
292/// Executable `PL/pgSQL` statement.
293#[derive(Debug, Clone)]
294pub enum PLpgSQLStmt {
295    Block(PLpgSQLBlock),
296    /// `target := expr` (also `=`). `target` indexes the datum table.
297    Assign {
298        target: usize,
299        expr: Expr,
300    },
301    If {
302        cond: Expr,
303        then_body: Vec<PLpgSQLStmt>,
304        elsifs: Vec<(Expr, Vec<PLpgSQLStmt>)>,
305        else_body: Option<Vec<PLpgSQLStmt>>,
306    },
307    /// CASE statement. Simple form carries `t_expr` + the temporary
308    /// datum the compiler references from each rewritten WHEN.
309    Case {
310        t_expr: Option<Expr>,
311        t_varno: Option<usize>,
312        arms: Vec<(Expr, Vec<PLpgSQLStmt>)>,
313        else_body: Option<Vec<PLpgSQLStmt>>,
314    },
315    Loop {
316        label: Option<String>,
317        body: Vec<PLpgSQLStmt>,
318    },
319    While {
320        label: Option<String>,
321        cond: Expr,
322        body: Vec<PLpgSQLStmt>,
323    },
324    /// `FOR i IN [REVERSE] lower..upper [BY step] LOOP`.
325    ForI {
326        label: Option<String>,
327        var: usize,
328        lower: Expr,
329        upper: Expr,
330        step: Option<Expr>,
331        reverse: bool,
332        body: Vec<PLpgSQLStmt>,
333    },
334    /// `FOR target IN <query> LOOP`.
335    ForQuery {
336        label: Option<String>,
337        target: IntoTarget,
338        query: Statement,
339        body: Vec<PLpgSQLStmt>,
340    },
341    /// `FOR target IN EXECUTE query [USING params] LOOP`.
342    ForDynamic {
343        label: Option<String>,
344        target: IntoTarget,
345        query: Expr,
346        params: Vec<Expr>,
347        body: Vec<PLpgSQLStmt>,
348    },
349    /// `FOR recordvar IN bound_cursor [(arguments)] LOOP`.
350    ForCursor {
351        label: Option<String>,
352        target: usize,
353        cursor: usize,
354        arguments: Vec<PLpgSQLCursorArgument>,
355        body: Vec<PLpgSQLStmt>,
356    },
357    /// `FOREACH target [SLICE n] IN ARRAY expression LOOP`.
358    ForeachArray {
359        label: Option<String>,
360        target: usize,
361        slice: usize,
362        expr: Expr,
363        body: Vec<PLpgSQLStmt>,
364    },
365    /// `EXIT` (`is_exit`) or `CONTINUE`, optionally labelled and
366    /// conditional (`WHEN cond`).
367    Exit {
368        is_exit: bool,
369        label: Option<String>,
370        cond: Option<Expr>,
371    },
372    Return {
373        value: Option<PLpgSQLReturnValue>,
374    },
375    /// `RETURN NEXT [expr]` - bare form emits the current OUT /
376    /// TABLE column values.
377    ReturnNext {
378        value: Option<PLpgSQLReturnValue>,
379    },
380    ReturnQuery {
381        query: Statement,
382    },
383    ReturnQueryExecute {
384        query: Expr,
385        params: Vec<Expr>,
386    },
387    Raise {
388        level: RaiseLevel,
389        condition: Option<String>,
390        message: Option<String>,
391        params: Vec<Expr>,
392    },
393    /// `ASSERT condition [, message]`.
394    Assert {
395        condition: Expr,
396        message: Option<Expr>,
397    },
398    /// Embedded SQL statement, optionally `INTO [STRICT] target`.
399    ExecSQL {
400        stmt: Statement,
401        into: Option<IntoTarget>,
402        strict: bool,
403    },
404    /// `EXECUTE <string> [INTO [STRICT] target] [USING params]`.
405    DynExecute {
406        query: Expr,
407        params: Vec<Expr>,
408        into: Option<IntoTarget>,
409        strict: bool,
410    },
411    Perform {
412        query: Statement,
413    },
414    OpenCursor {
415        cursor: usize,
416        open: PLpgSQLCursorOpen,
417    },
418    FetchCursor {
419        cursor: usize,
420        target: IntoTarget,
421        direction: CursorDirection,
422        count: PLpgSQLCursorCount,
423    },
424    MoveCursor {
425        cursor: usize,
426        direction: CursorDirection,
427        count: PLpgSQLCursorCount,
428    },
429    CloseCursor {
430        cursor: usize,
431    },
432    /// Procedural `COMMIT [AND [NO] CHAIN]`.
433    Commit {
434        chain: bool,
435    },
436    /// Procedural `ROLLBACK [AND [NO] CHAIN]`.
437    Rollback {
438        chain: bool,
439    },
440    /// `GET DIAGNOSTICS var = KIND [, ...]` as `(kind, target datum)`.
441    GetDiagnostics {
442        items: Vec<(String, usize)>,
443    },
444}
445
446/// Value source for `RETURN` and `RETURN NEXT`. `PostgreSQL` 18 stores a simple
447/// datum reference in `retvarno`, distinct from a general SQL expression.
448#[derive(Debug, Clone)]
449pub enum PLpgSQLReturnValue {
450    Expr(Expr),
451    Datum(usize),
452}
453
454// ---------------------------------------------------------------------
455// Parsing: definition -> canonical text -> libpg_query JSON -> AST
456// ---------------------------------------------------------------------
457
458/// Parse the `PL/pgSQL` body of a stored definition. The definition
459/// is re-serialized into a canonical `CREATE FUNCTION` statement so
460/// restore-from-catalog and fresh DDL take the same path.
461mod binding;
462mod conditions;
463mod json_validation;
464mod lowering_expression;
465mod lowering_statement;
466mod parsing;
467
468use json_validation::{
469    ensure_single_tag, expect_tag, json_bool_or_false, json_i64_or_zero, json_kind,
470    json_optional_i64, json_optional_str, json_optional_usize, json_usize_or_zero,
471    normalize_plpgsql_type, optional_array, require, require_i64, require_nonempty_str,
472    validate_assignable_datum, validate_record_datum, validate_scalar_datum,
473};
474use lowering_expression::{lower_expr, lower_expr_list, lower_full_statement};
475use lowering_statement::{lower_block, lower_cursor_scroll_options};
476use parsing::{lower_row_fields, normalize_condition};
477
478pub use binding::{bind_expr, bind_select, bind_statement, ResolvedVariable, VariableResolver};
479pub use conditions::{condition_sqlstate, condition_sqlstates};
480pub use lowering_expression::compile_expression_text;
481pub use parsing::{parse_do_block, parse_function};
482
483#[cfg(test)]
484mod tests;
uqa_sql/plpgsql.rs

uqa_sql/
plpgsql.rs
