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

uqa_sql/
plpgsql.rs
