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