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

uqa_sql/
plpgsql.rs
