1use 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
28pub 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
38pub(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("e_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("e_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("e_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
112pub(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
139pub(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 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 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}