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