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    /// Resolved identity from a catalog-aware parse; absent in catalog-free JSON.
181    pub type_oid: Option<u32>,
182    /// Normalized type name (`integer`, `text`, ...). The engine resolves
183    /// catalog-backed references such as `%TYPE` before execution.
184    pub type_name: String,
185    /// Exact relation-column identity emitted by the PL/pgSQL parser for a table-backed `%TYPE` declaration.
186    pub type_reference: Option<RoutineColumnTypeReference>,
187    pub default: Option<Expr>,
188    pub constant: bool,
189    pub not_null: bool,
190    /// Definition of a bound cursor declared with `CURSOR (...) FOR query`.
191    pub cursor: Option<PLpgSQLCursor>,
192    /// Source line of the declaration; used to disambiguate loop
193    /// variables that shadow outer names.
194    pub lineno: Option<i64>,
195}
196
197#[derive(Debug, Clone)]
198pub struct PLpgSQLCursor {
199    pub query: Statement,
200    pub argument_row: Option<usize>,
201    /// Explicit declaration scroll mode. `None` leaves scrollability query-dependent.
202    pub scroll: Option<bool>,
203}
204
205#[derive(Debug, Clone)]
206pub struct PLpgSQLCursorArgument {
207    pub name: Option<String>,
208    pub expr: Expr,
209}
210
211/// Query source selected by one `OPEN` statement.
212#[derive(Debug, Clone)]
213pub enum PLpgSQLCursorOpen {
214    Bound {
215        arguments: Vec<PLpgSQLCursorArgument>,
216    },
217    Static {
218        query: Box<Statement>,
219        scroll: Option<bool>,
220    },
221    Dynamic {
222        query: Expr,
223        params: Vec<Expr>,
224        scroll: Option<bool>,
225    },
226}
227
228/// Constant or run-time expression controlling cursor movement.
229#[derive(Debug, Clone)]
230pub enum PLpgSQLCursorCount {
231    Constant(i64),
232    Expression(Expr),
233}
234
235/// `name -> datum` slot of a row target.
236#[derive(Debug, Clone)]
237pub struct PLpgSQLRowField {
238    pub name: String,
239    pub varno: usize,
240}
241
242/// `[DECLARE ...] BEGIN ... [EXCEPTION ...] END` block.
243#[derive(Debug, Clone)]
244pub struct PLpgSQLBlock {
245    pub label: Option<String>,
246    pub body: Vec<PLpgSQLStmt>,
247    pub exceptions: Vec<PLpgSQLExceptionArm>,
248}
249
250/// One `WHEN cond [OR cond ...] THEN stmts` arm of an exception
251/// section.
252#[derive(Debug, Clone)]
253pub struct PLpgSQLExceptionArm {
254    /// Lower-cased condition names (`others`, `division_by_zero`,
255    /// ...). Explicit `SQLSTATE 'xxxxx'` conditions arrive as the
256    /// five-character code.
257    pub conditions: Vec<String>,
258    pub body: Vec<PLpgSQLStmt>,
259}
260
261/// `RAISE` severity.
262#[derive(Debug, Clone, Copy, PartialEq, Eq)]
263pub enum RaiseLevel {
264    Debug,
265    Log,
266    Info,
267    Notice,
268    Warning,
269    Error,
270}
271
272impl RaiseLevel {
273    pub fn as_str(self) -> &'static str {
274        match self {
275            RaiseLevel::Debug => "DEBUG",
276            RaiseLevel::Log => "LOG",
277            RaiseLevel::Info => "INFO",
278            RaiseLevel::Notice => "NOTICE",
279            RaiseLevel::Warning => "WARNING",
280            RaiseLevel::Error => "ERROR",
281        }
282    }
283}
284
285/// Assignment / `INTO` target.
286#[derive(Debug, Clone)]
287pub enum IntoTarget {
288    /// A `RECORD` variable receives the whole row.
289    Rec(usize),
290    /// Positional list of scalar targets.
291    Row(Vec<PLpgSQLRowField>),
292}
293
294/// Executable `PL/pgSQL` statement.
295#[derive(Debug, Clone)]
296pub enum PLpgSQLStmt {
297    Block(PLpgSQLBlock),
298    /// `target := expr` (also `=`). `target` indexes the datum table.
299    Assign {
300        target: usize,
301        expr: Expr,
302    },
303    If {
304        cond: Expr,
305        then_body: Vec<PLpgSQLStmt>,
306        elsifs: Vec<(Expr, Vec<PLpgSQLStmt>)>,
307        else_body: Option<Vec<PLpgSQLStmt>>,
308    },
309    /// CASE statement. Simple form carries `t_expr` + the temporary
310    /// datum the compiler references from each rewritten WHEN.
311    Case {
312        t_expr: Option<Expr>,
313        t_varno: Option<usize>,
314        arms: Vec<(Expr, Vec<PLpgSQLStmt>)>,
315        else_body: Option<Vec<PLpgSQLStmt>>,
316    },
317    Loop {
318        label: Option<String>,
319        body: Vec<PLpgSQLStmt>,
320    },
321    While {
322        label: Option<String>,
323        cond: Expr,
324        body: Vec<PLpgSQLStmt>,
325    },
326    /// `FOR i IN [REVERSE] lower..upper [BY step] LOOP`.
327    ForI {
328        label: Option<String>,
329        var: usize,
330        lower: Expr,
331        upper: Expr,
332        step: Option<Expr>,
333        reverse: bool,
334        body: Vec<PLpgSQLStmt>,
335    },
336    /// `FOR target IN <query> LOOP`.
337    ForQuery {
338        label: Option<String>,
339        target: IntoTarget,
340        query: Statement,
341        body: Vec<PLpgSQLStmt>,
342    },
343    /// `FOR target IN EXECUTE query [USING params] LOOP`.
344    ForDynamic {
345        label: Option<String>,
346        target: IntoTarget,
347        query: Expr,
348        params: Vec<Expr>,
349        body: Vec<PLpgSQLStmt>,
350    },
351    /// `FOR recordvar IN bound_cursor [(arguments)] LOOP`.
352    ForCursor {
353        label: Option<String>,
354        target: usize,
355        cursor: usize,
356        arguments: Vec<PLpgSQLCursorArgument>,
357        body: Vec<PLpgSQLStmt>,
358    },
359    /// `FOREACH target [SLICE n] IN ARRAY expression LOOP`.
360    ForeachArray {
361        label: Option<String>,
362        target: usize,
363        slice: usize,
364        expr: Expr,
365        body: Vec<PLpgSQLStmt>,
366    },
367    /// `EXIT` (`is_exit`) or `CONTINUE`, optionally labelled and
368    /// conditional (`WHEN cond`).
369    Exit {
370        is_exit: bool,
371        label: Option<String>,
372        cond: Option<Expr>,
373    },
374    Return {
375        value: Option<PLpgSQLReturnValue>,
376    },
377    /// `RETURN NEXT [expr]` - bare form emits the current OUT /
378    /// TABLE column values.
379    ReturnNext {
380        value: Option<PLpgSQLReturnValue>,
381    },
382    ReturnQuery {
383        query: Statement,
384    },
385    ReturnQueryExecute {
386        query: Expr,
387        params: Vec<Expr>,
388    },
389    Raise {
390        level: RaiseLevel,
391        condition: Option<String>,
392        message: Option<String>,
393        params: Vec<Expr>,
394    },
395    /// `ASSERT condition [, message]`.
396    Assert {
397        condition: Expr,
398        message: Option<Expr>,
399    },
400    /// Embedded SQL statement, optionally `INTO [STRICT] target`.
401    ExecSQL {
402        stmt: Statement,
403        into: Option<IntoTarget>,
404        strict: bool,
405    },
406    /// `EXECUTE <string> [INTO [STRICT] target] [USING params]`.
407    DynExecute {
408        query: Expr,
409        params: Vec<Expr>,
410        into: Option<IntoTarget>,
411        strict: bool,
412    },
413    Perform {
414        query: Statement,
415    },
416    OpenCursor {
417        cursor: usize,
418        open: PLpgSQLCursorOpen,
419    },
420    FetchCursor {
421        cursor: usize,
422        target: IntoTarget,
423        direction: CursorDirection,
424        count: PLpgSQLCursorCount,
425    },
426    MoveCursor {
427        cursor: usize,
428        direction: CursorDirection,
429        count: PLpgSQLCursorCount,
430    },
431    CloseCursor {
432        cursor: usize,
433    },
434    /// Procedural `COMMIT [AND [NO] CHAIN]`.
435    Commit {
436        chain: bool,
437    },
438    /// Procedural `ROLLBACK [AND [NO] CHAIN]`.
439    Rollback {
440        chain: bool,
441    },
442    /// `GET DIAGNOSTICS var = KIND [, ...]` as `(kind, target datum)`.
443    GetDiagnostics {
444        items: Vec<(String, usize)>,
445    },
446}
447
448/// Value source for `RETURN` and `RETURN NEXT`. `PostgreSQL` 18 stores a simple
449/// datum reference in `retvarno`, distinct from a general SQL expression.
450#[derive(Debug, Clone)]
451pub enum PLpgSQLReturnValue {
452    Expr(Expr),
453    Datum(usize),
454}
455
456// ---------------------------------------------------------------------
457// Parsing: definition -> canonical text -> libpg_query JSON -> AST
458// ---------------------------------------------------------------------
459
460/// Parse the `PL/pgSQL` body of a stored definition. The definition
461/// is re-serialized into a canonical `CREATE FUNCTION` statement so
462/// restore-from-catalog and fresh DDL take the same path.
463mod binding;
464mod conditions;
465mod json_validation;
466mod lowering_expression;
467mod lowering_statement;
468mod parsing;
469
470use json_validation::{
471    ensure_single_tag, expect_tag, json_bool_or_false, json_i64_or_zero, json_kind,
472    json_optional_i64, json_optional_str, json_optional_usize, json_usize_or_zero,
473    normalize_plpgsql_type, optional_array, require, require_i64, require_nonempty_str,
474    validate_assignable_datum, validate_record_datum, validate_scalar_datum,
475};
476use lowering_expression::{lower_expr, lower_expr_list, lower_full_statement};
477use lowering_statement::{lower_block, lower_cursor_scroll_options};
478use parsing::{lower_row_fields, normalize_condition};
479
480pub use binding::{bind_expr, bind_select, bind_statement, ResolvedVariable, VariableResolver};
481pub use conditions::{condition_sqlstate, condition_sqlstates};
482pub use lowering_expression::compile_expression_text;
483pub use parsing::{
484    parse_do_block, parse_do_block_with_catalog, parse_function, parse_function_with_catalog,
485};
486pub use pg_query::{PlpgsqlCatalog, PlpgsqlType};
487
488#[cfg(test)]
489mod tests;
490
491pub mod runtime_diagnostics;
uqa_sql/plpgsql.rs

uqa_sql/
plpgsql.rs
