Skip to main content

uqa_sql/plpgsql/
parsing.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! PL/pgSQL parser invocation, datum lowering, and condition normalization.
8
9use super::{
10    condition_sqlstate, ensure_single_tag, expect_tag, json_bool_or_false, json_kind,
11    json_optional_i64, json_usize_or_zero, lower_block, lower_cursor_scroll_options, lower_expr,
12    lower_full_statement, normalize_plpgsql_type, optional_array, require, require_nonempty_str,
13    validate_assignable_datum, CreateFunction, FunctionBody, FunctionParamMode, FunctionReturns,
14    JSONValue, PLpgSQLCursor, PLpgSQLDatum, PLpgSQLFunction, PLpgSQLRowField, PLpgSQLVar, Result,
15    RoutineColumnTypeReference, SQLError,
16};
17
18pub fn parse_function(def: &CreateFunction) -> Result<PLpgSQLFunction> {
19    let FunctionBody::Source(body) = &def.body else {
20        return Err(SQLError::Internal(
21            "PL/pgSQL parser invoked on a SQL-standard body".into(),
22        ));
23    };
24    let text = synthesize_create_text(def, body);
25    parse_plpgsql_text(&text)
26}
27
28/// Parse a `DO $$ ... $$` body through `PostgreSQL`'s native inline-code path.
29pub fn parse_do_block(body: &str) -> Result<PLpgSQLFunction> {
30    let tag = fresh_dollar_tag(body);
31    let text = format!("DO {tag}{body}{tag} LANGUAGE plpgsql;");
32    parse_plpgsql_text(&text)
33}
34
35/// Canonical `CREATE FUNCTION` / `CREATE PROCEDURE` text used solely
36/// to feed the `PL/pgSQL` parser (parameter DEFAULTs are resolved at
37/// call time and intentionally omitted).
38pub(super) fn synthesize_create_text(def: &CreateFunction, body: &str) -> String {
39    let mut sql = String::new();
40    sql.push_str(if def.is_procedure {
41        "CREATE PROCEDURE "
42    } else {
43        "CREATE FUNCTION "
44    });
45    sql.push_str(&quote_ident(&def.name));
46    sql.push('(');
47    let mut first = true;
48    for p in &def.params {
49        if matches!(p.mode, FunctionParamMode::Table) {
50            continue;
51        }
52        if !first {
53            sql.push_str(", ");
54        }
55        first = false;
56        match p.mode {
57            FunctionParamMode::Out => sql.push_str("OUT "),
58            FunctionParamMode::InOut => sql.push_str("INOUT "),
59            FunctionParamMode::Variadic => sql.push_str("VARIADIC "),
60            FunctionParamMode::In | FunctionParamMode::Table => {}
61        }
62        if !p.name.is_empty() {
63            sql.push_str(&quote_ident(&p.name));
64            sql.push(' ');
65        }
66        sql.push_str(&p.type_name);
67    }
68    sql.push(')');
69    match &def.returns {
70        FunctionReturns::None => {}
71        FunctionReturns::Scalar { type_name } => {
72            sql.push_str(" RETURNS ");
73            sql.push_str(type_name);
74        }
75        FunctionReturns::SetOf { type_name } => {
76            sql.push_str(" RETURNS SETOF ");
77            sql.push_str(type_name);
78        }
79        FunctionReturns::Table => {
80            sql.push_str(" RETURNS TABLE(");
81            let mut first_col = true;
82            for p in &def.params {
83                if !matches!(p.mode, FunctionParamMode::Table) {
84                    continue;
85                }
86                if !first_col {
87                    sql.push_str(", ");
88                }
89                first_col = false;
90                sql.push_str(&quote_ident(&p.name));
91                sql.push(' ');
92                sql.push_str(&p.type_name);
93            }
94            sql.push(')');
95        }
96    }
97    let tag = fresh_dollar_tag(body);
98    sql.push_str(" AS ");
99    sql.push_str(&tag);
100    sql.push_str(body);
101    sql.push_str(&tag);
102    sql.push_str(" LANGUAGE plpgsql;");
103    sql
104}
105
106pub(super) fn quote_ident(name: &str) -> String {
107    format!("\"{}\"", name.replace('"', "\"\""))
108}
109
110/// Dollar-quote tag guaranteed not to collide with the body text.
111pub(super) fn fresh_dollar_tag(body: &str) -> String {
112    let mut n = 0usize;
113    loop {
114        let tag = if n == 0 {
115            "$$".to_string()
116        } else {
117            format!("$plpgsql{n}$")
118        };
119        if !body.contains(&tag) {
120            return tag;
121        }
122        n += 1;
123    }
124}
125
126pub(super) fn parse_plpgsql_text(text: &str) -> Result<PLpgSQLFunction> {
127    let json = pg_query::parse_plpgsql(text)?;
128    let functions = json
129        .as_array()
130        .ok_or_else(|| SQLError::Internal("PL/pgSQL parse returned no function list".into()))?;
131    if functions.len() != 1 {
132        return Err(SQLError::Internal(format!(
133            "PL/pgSQL parse returned {} functions; expected exactly one",
134            functions.len()
135        )));
136    }
137    let function = expect_tag(&functions[0], "PLpgSQL_function", "parsed function")?;
138    lower_function(function)
139}
140
141// ---------------------------------------------------------------------
142// JSON lowering
143// ---------------------------------------------------------------------
144
145/// Divergence from `PostgreSQL`: the JSON dump does not carry each
146/// block's `initvarnos`, so declared-variable defaults (including
147/// those of nested `DECLARE` sections) are evaluated once at routine
148/// entry rather than on every block entry, and a nested declaration
149/// shadows its outer namesake for the whole body.
150pub(super) fn lower_function(function: &JSONValue) -> Result<PLpgSQLFunction> {
151    let raw_datums = function
152        .get("datums")
153        .and_then(JSONValue::as_array)
154        .ok_or_else(|| SQLError::Internal("PL/pgSQL function without datums".into()))?;
155    let mut datums = Vec::with_capacity(raw_datums.len());
156    for raw in raw_datums {
157        datums.push(lower_datum(raw)?);
158    }
159    validate_datums(&datums)?;
160    let trigger_datum = |field: &str, name: &str| -> Result<Option<usize>> {
161        let explicit = match json_optional_i64(function, field)? {
162            Some(index) if index >= 0 => {
163                let index = usize::try_from(index).map_err(|_| {
164                    SQLError::Internal(format!(
165                        "PL/pgSQL {field} {index} does not fit this platform"
166                    ))
167                })?;
168                if index >= datums.len() {
169                    return Err(SQLError::Internal(format!(
170                        "PL/pgSQL {field} has out-of-range datum index {index}"
171                    )));
172                }
173                Some(index)
174            }
175            Some(index) => {
176                return Err(SQLError::Internal(format!(
177                    "PL/pgSQL {field} has invalid datum index {index}"
178                )))
179            }
180            None => None,
181        };
182        Ok(explicit.or_else(|| {
183            datums.iter().position(|datum| {
184                datum
185                    .name()
186                    .is_some_and(|datum_name| datum_name.eq_ignore_ascii_case(name))
187            })
188        }))
189    };
190    let new_datum = trigger_datum("new_varno", "new")?;
191    let old_datum = trigger_datum("old_varno", "old")?;
192    let found_datum = datums
193        .iter()
194        .position(|d| matches!(d, PLpgSQLDatum::Var(v) if v.name.eq_ignore_ascii_case("found")));
195    let raw_action = require(function, "action")?;
196    let action = expect_tag(raw_action, "PLpgSQL_stmt_block", "function body")?;
197    let action = lower_block(action, &datums)?;
198    Ok(PLpgSQLFunction {
199        datums,
200        action,
201        new_datum,
202        old_datum,
203        found_datum,
204    })
205}
206
207fn has_percent_type_suffix(type_name: &str) -> bool {
208    type_name
209        .get(type_name.len().saturating_sub("%type".len())..)
210        .is_some_and(|suffix| suffix.eq_ignore_ascii_case("%type"))
211}
212
213fn lower_percent_type_reference(
214    datatype: &JSONValue,
215    variable_name: &str,
216) -> Result<RoutineColumnTypeReference> {
217    let identifiers = require(datatype, "typname_identifiers")?
218        .as_array()
219        .ok_or_else(|| {
220            SQLError::Internal(format!(
221                "PL/pgSQL variable `{variable_name}` type metadata `typname_identifiers` must be an array"
222            ))
223        })?;
224    let identifiers = identifiers
225        .iter()
226        .enumerate()
227        .map(|(index, identifier)| match identifier.as_str() {
228            Some(identifier) if !identifier.is_empty() => Ok(identifier.to_string()),
229            _ => Err(SQLError::Internal(format!(
230                "PL/pgSQL variable `{variable_name}` type metadata identifier {index} must be a non-empty string"
231            ))),
232        })
233        .collect::<Result<Vec<_>>>()?;
234    match identifiers.as_slice() {
235        [relation, column] => Ok(RoutineColumnTypeReference::new(
236            None,
237            relation.clone(),
238            column.clone(),
239        )),
240        [schema, relation, column] => Ok(RoutineColumnTypeReference::new(
241            Some(schema.clone()),
242            relation.clone(),
243            column.clone(),
244        )),
245        _ => Err(SQLError::TypeMismatch(format!(
246            "PL/pgSQL variable `{variable_name}` %TYPE must identify a relation column"
247        ))),
248    }
249}
250
251pub(super) fn lower_datum(raw: &JSONValue) -> Result<PLpgSQLDatum> {
252    ensure_single_tag(raw, "datum")?;
253    if let Some(var) = raw.get("PLpgSQL_var") {
254        let name = require_nonempty_str(var, "refname", "variable datum")?;
255        let datatype = require(var, "datatype")?;
256        let datatype = expect_tag(datatype, "PLpgSQL_type", "variable datatype")?;
257        let type_name = normalize_plpgsql_type(&require_nonempty_str(
258            datatype,
259            "typname",
260            "variable datatype",
261        )?);
262        if type_name.is_empty() {
263            return Err(SQLError::Internal(format!(
264                "PL/pgSQL variable `{name}` has an empty normalized type"
265            )));
266        }
267        let type_reference = has_percent_type_suffix(&type_name)
268            .then(|| lower_percent_type_reference(datatype, &name))
269            .transpose()?;
270        let default = match var.get("default_val") {
271            Some(node) => Some(lower_expr(node)?),
272            None => None,
273        };
274        let cursor = if let Some(query) = var.get("cursor_explicit_expr") {
275            Some(PLpgSQLCursor {
276                query: lower_full_statement(query)?,
277                argument_row: match json_optional_i64(var, "cursor_explicit_argrow")? {
278                    None | Some(-1) => None,
279                    Some(index) if index >= 0 => Some(usize::try_from(index).map_err(|_| {
280                        SQLError::Internal(format!(
281                            "PL/pgSQL cursor `{name}` argument row {index} does not fit this platform"
282                        ))
283                    })?),
284                    Some(index) => {
285                        return Err(SQLError::Internal(format!(
286                            "PL/pgSQL cursor `{name}` has invalid argument row {index}"
287                        )));
288                    }
289                },
290                scroll: lower_cursor_scroll_options(var, "cursor declaration")?,
291            })
292        } else {
293            if var.get("cursor_explicit_argrow").is_some() {
294                return Err(SQLError::Internal(format!(
295                    "PL/pgSQL cursor variable `{name}` has arguments but no query"
296                )));
297            }
298            None
299        };
300        return Ok(PLpgSQLDatum::Var(Box::new(PLpgSQLVar {
301            name,
302            type_name,
303            type_reference,
304            default,
305            constant: json_bool_or_false(var, "isconst")?,
306            not_null: json_bool_or_false(var, "notnull")?,
307            cursor,
308            lineno: json_optional_i64(var, "lineno")?,
309        })));
310    }
311    if let Some(rec) = raw.get("PLpgSQL_rec") {
312        return Ok(PLpgSQLDatum::Rec {
313            name: require_nonempty_str(rec, "refname", "record datum")?,
314        });
315    }
316    if let Some(field) = raw.get("PLpgSQL_recfield") {
317        return Ok(PLpgSQLDatum::RecField {
318            field: require_nonempty_str(field, "fieldname", "record-field datum")?,
319            // libpg_query omits a zero-valued recparentno.
320            parent: json_usize_or_zero(field, "recparentno")?,
321        });
322    }
323    if let Some(row) = raw.get("PLpgSQL_row") {
324        return Ok(PLpgSQLDatum::Row {
325            fields: lower_row_fields(row)?,
326        });
327    }
328    Err(SQLError::Unsupported(format!(
329        "PL/pgSQL datum {}",
330        json_kind(raw)
331    )))
332}
333
334pub(super) fn lower_row_fields(row: &JSONValue) -> Result<Vec<PLpgSQLRowField>> {
335    let mut out = Vec::new();
336    if let Some(fields) = optional_array(row, "fields")? {
337        for f in fields {
338            // libpg_query's JSON dump omits zero-valued fields, so a
339            // missing varno means datum 0.
340            out.push(PLpgSQLRowField {
341                name: require_nonempty_str(f, "name", "row target field")?,
342                varno: json_usize_or_zero(f, "varno")?,
343            });
344        }
345    }
346    Ok(out)
347}
348
349pub(super) fn validate_datums(datums: &[PLpgSQLDatum]) -> Result<()> {
350    for (idx, datum) in datums.iter().enumerate() {
351        match datum {
352            PLpgSQLDatum::RecField { parent, .. } => {
353                let Some(parent_datum) = datums.get(*parent) else {
354                    return Err(SQLError::Internal(format!(
355                        "PL/pgSQL record-field datum {idx} references missing parent datum {parent}"
356                    )));
357                };
358                if !matches!(parent_datum, PLpgSQLDatum::Rec { .. }) {
359                    return Err(SQLError::Internal(format!(
360                        "PL/pgSQL record-field datum {idx} parent {parent} is not a record"
361                    )));
362                }
363            }
364            PLpgSQLDatum::Row { fields } => {
365                if fields.is_empty() {
366                    return Err(SQLError::Internal(format!(
367                        "PL/pgSQL row datum {idx} has no fields"
368                    )));
369                }
370                for field in fields {
371                    validate_assignable_datum(datums, field.varno, "row target field")?;
372                }
373            }
374            PLpgSQLDatum::Var(var) => {
375                if let Some(cursor) = &var.cursor {
376                    if var.type_name != "refcursor" {
377                        return Err(SQLError::Internal(format!(
378                            "PL/pgSQL bound cursor `{}` is not a refcursor datum",
379                            var.name
380                        )));
381                    }
382                    if let Some(argument_row) = cursor.argument_row {
383                        if !matches!(datums.get(argument_row), Some(PLpgSQLDatum::Row { .. })) {
384                            return Err(SQLError::Internal(format!(
385                                "PL/pgSQL cursor `{}` references invalid argument row {argument_row}",
386                                var.name
387                            )));
388                        }
389                    }
390                }
391            }
392            PLpgSQLDatum::Rec { .. } => {}
393        }
394    }
395    Ok(())
396}
397
398pub(super) fn normalize_condition(value: String, allow_others: bool) -> Result<String> {
399    let lower = value.to_ascii_lowercase();
400    if allow_others && lower == "others" {
401        return Ok(lower);
402    }
403    if condition_sqlstate(&lower).is_some() {
404        return Ok(lower);
405    }
406    let upper = value.to_ascii_uppercase();
407    if upper.len() == 5
408        && upper
409            .bytes()
410            .all(|byte| byte.is_ascii_uppercase() || byte.is_ascii_digit())
411    {
412        return Ok(upper);
413    }
414    Err(SQLError::Internal(format!(
415        "unrecognized PL/pgSQL exception condition `{value}`"
416    )))
417}