Skip to main content

pgevolve_core/parse/builder/
shared.rs

1//! Helpers shared by all AST → IR builders.
2//!
3//! - Schema-qualified name resolution (with `-- @pgevolve schema=` defaulting).
4//! - Type-name stringification (so [`crate::ir::column_type::ColumnType::parse_from_pg_type_string`]
5//!   can decide the canonical form).
6//! - Default-expression classification (literal vs. `nextval` vs. arbitrary expr).
7
8use pg_query::NodeEnum;
9use pg_query::protobuf::{AConst, RangeVar, TypeName, a_const};
10
11use crate::identifier::{Identifier, QualifiedName};
12use crate::ir::column_type::ColumnType;
13use crate::ir::default_expr::{DefaultExpr, LiteralValue};
14use crate::parse::error::{ParseError, SourceLocation};
15use crate::parse::normalize_expr;
16
17/// Resolve a `RangeVar` into a [`QualifiedName`], filling in `default_schema`
18/// when the source did not qualify the object.
19///
20/// Returns [`ParseError::UnqualifiedName`] if neither the source nor the
21/// directive supply a schema.
22pub fn resolve_qname(
23    range_var: &RangeVar,
24    default_schema: Option<&Identifier>,
25    location: &SourceLocation,
26) -> Result<QualifiedName, ParseError> {
27    let name = ident(&range_var.relname, location)?;
28    if !range_var.schemaname.is_empty() {
29        let schema = ident(&range_var.schemaname, location)?;
30        return Ok(QualifiedName::new(schema, name));
31    }
32    default_schema.map_or_else(
33        || {
34            Err(ParseError::UnqualifiedName {
35                location: location.clone(),
36            })
37        },
38        |s| Ok(QualifiedName::new(s.clone(), name)),
39    )
40}
41
42/// Resolve a `repeated Node` type-name list (as produced by `CreateEnumStmt`,
43/// `CreateDomainStmt`, `CompositeTypeStmt`, etc.) into a [`QualifiedName`].
44///
45/// PG encodes the name as a list of [`pg_query::NodeEnum::String`] nodes:
46/// - one element → unqualified name (requires `default_schema`).
47/// - two elements → `[schema, name]`.
48///
49/// Returns [`ParseError::UnqualifiedName`] if neither the source nor the
50/// directive supply a schema.
51///
52/// This helper is shared by the enum, domain, and composite builders (T2–T4).
53pub fn qname_from_string_list(
54    nodes: &[pg_query::protobuf::Node],
55    default_schema: Option<&Identifier>,
56    location: &SourceLocation,
57) -> Result<QualifiedName, ParseError> {
58    let strings: Vec<&str> = nodes
59        .iter()
60        .map(|n| match n.node.as_ref() {
61            Some(NodeEnum::String(s)) => Ok(s.sval.as_str()),
62            other => Err(ParseError::Structural {
63                location: location.clone(),
64                message: format!(
65                    "expected String node in type-name list, got {:?}",
66                    other.map(std::mem::discriminant),
67                ),
68            }),
69        })
70        .collect::<Result<Vec<_>, _>>()?;
71    match strings.as_slice() {
72        [name] => {
73            let name_id = ident(name, location)?;
74            let schema = default_schema
75                .cloned()
76                .ok_or_else(|| ParseError::UnqualifiedName {
77                    location: location.clone(),
78                })?;
79            Ok(QualifiedName::new(schema, name_id))
80        }
81        [schema, name] => {
82            let schema_id = ident(schema, location)?;
83            let name_id = ident(name, location)?;
84            Ok(QualifiedName::new(schema_id, name_id))
85        }
86        _ => Err(ParseError::Structural {
87            location: location.clone(),
88            message: format!(
89                "unexpected type-name list length {} (expected 1 or 2 String nodes)",
90                nodes.len()
91            ),
92        }),
93    }
94}
95
96/// Build an unquoted [`Identifier`], wrapping [`crate::ir::IrError`] into a
97/// source-located [`ParseError::Ir`].
98pub fn ident(s: &str, location: &SourceLocation) -> Result<Identifier, ParseError> {
99    Identifier::from_unquoted(s).map_err(|e| ParseError::Ir {
100        location: location.clone(),
101        source: crate::ir::IrError::InvalidIdentifier(e.to_string()),
102    })
103}
104
105/// Render a `pg_query::TypeName` into a string `ColumnType::parse_from_pg_type_string`
106/// can consume.
107///
108/// Strategy: take the *last* segment of `names` (Postgres prefixes types with
109/// `pg_catalog.` internally; the parser is alias-aware), append a parenthesized
110/// list of typmod arguments, append `[]` for each array dimension.
111pub fn render_type_name_to_string(type_name: &TypeName) -> Option<String> {
112    let last = type_name.names.last()?.node.as_ref()?;
113    let NodeEnum::String(s) = last else {
114        return None;
115    };
116    let bare = s.sval.clone();
117    let mut out = bare;
118    if !type_name.typmods.is_empty() {
119        let mut args: Vec<String> = Vec::with_capacity(type_name.typmods.len());
120        for n in &type_name.typmods {
121            let Some(NodeEnum::AConst(c)) = &n.node else {
122                return None;
123            };
124            args.push(literal_arg_to_string(c)?);
125        }
126        out = format!("{out}({})", args.join(","));
127    }
128    for _ in 0..type_name.array_bounds.len() {
129        out.push_str("[]");
130    }
131    Some(out)
132}
133
134/// Convert an [`AConst`] used as a typmod argument to its canonical string form.
135fn literal_arg_to_string(c: &AConst) -> Option<String> {
136    match c.val.as_ref()? {
137        a_const::Val::Ival(i) => Some(i.ival.to_string()),
138        a_const::Val::Sval(s) => Some(s.sval.clone()),
139        _ => None,
140    }
141}
142
143/// Convert a `TypeName` into a [`ColumnType`], propagating parser errors.
144///
145/// When `TypeName.names` contains two String nodes and the first is not
146/// `pg_catalog`, the type is schema-qualified by the user (e.g. `app.order_status`).
147/// In that case we emit `ColumnType::UserDefined(QualifiedName)` so that the
148/// AST resolution pass can validate that the referenced type is declared in source.
149///
150/// Unqualified single-segment names are handled by the usual string path; if they
151/// don't match any built-in they fall through to `ColumnType::Other`, which is
152/// intentional — unqualified user types must be resolved via `default_schema` at
153/// domain/composite parse time rather than here.
154pub fn type_name_to_column_type(
155    type_name: &TypeName,
156    location: &SourceLocation,
157) -> Result<ColumnType, ParseError> {
158    // Collect String nodes from names; fail fast on unexpected kinds.
159    let name_strings: Vec<&str> = type_name
160        .names
161        .iter()
162        .map(|n| match n.node.as_ref() {
163            Some(NodeEnum::String(s)) => Ok(s.sval.as_str()),
164            other => Err(ParseError::Structural {
165                location: location.clone(),
166                message: format!(
167                    "expected String node in type-name list, got {:?}",
168                    other.map(std::mem::discriminant),
169                ),
170            }),
171        })
172        .collect::<Result<Vec<_>, _>>()?;
173
174    // Two-segment, non-pg_catalog prefix → user-defined type reference.
175    if let [schema, name] = name_strings.as_slice()
176        && *schema != "pg_catalog"
177    {
178        let schema_id = ident(schema, location)?;
179        let name_id = ident(name, location)?;
180        return Ok(ColumnType::UserDefined(QualifiedName::new(
181            schema_id, name_id,
182        )));
183    }
184
185    let s = render_type_name_to_string(type_name).ok_or_else(|| ParseError::Structural {
186        location: location.clone(),
187        message: "could not stringify type name".into(),
188    })?;
189    ColumnType::parse_from_pg_type_string(&s).map_err(|e| ParseError::Structural {
190        location: location.clone(),
191        message: format!("invalid column type {s:?}: {e}"),
192    })
193}
194
195/// Convert a column-default expression node into a [`DefaultExpr`].
196///
197/// Recognizes:
198/// - Bare [`AConst`] integer/float/string/bool/null literals → [`DefaultExpr::Literal`].
199/// - `nextval('seq')` and `nextval('schema.seq')` → [`DefaultExpr::Sequence`].
200/// - Anything else → [`DefaultExpr::Expr`] via the normalizer.
201///
202/// `target_type` enables redundant-cast stripping in the `Expr` arm.
203pub fn build_default_expr(
204    node: &NodeEnum,
205    target_type: Option<&ColumnType>,
206    default_schema: Option<&Identifier>,
207    location: &SourceLocation,
208) -> Result<DefaultExpr, ParseError> {
209    if let Some(lit) = literal_from_node(node) {
210        return Ok(DefaultExpr::Literal(lit));
211    }
212    if let Some(seq) = nextval_target(node, default_schema, location)? {
213        return Ok(DefaultExpr::Sequence(seq));
214    }
215    let normalized = normalize_expr::from_pg_node(node, target_type, location)?;
216    Ok(DefaultExpr::Expr(normalized))
217}
218
219/// If `node` is a literal `AConst` (possibly wrapped in a `TypeCast` whose target
220/// matches the column type), return the equivalent [`LiteralValue`].
221fn literal_from_node(node: &NodeEnum) -> Option<LiteralValue> {
222    let inner = unwrap_typecast(node);
223    match inner {
224        NodeEnum::AConst(c) => aconst_to_literal(c),
225        _ => None,
226    }
227}
228
229/// Strip a single layer of `TypeCast` wrapping, if any.
230fn unwrap_typecast(node: &NodeEnum) -> &NodeEnum {
231    if let NodeEnum::TypeCast(cast) = node
232        && let Some(arg) = cast.arg.as_ref()
233        && let Some(inner) = arg.node.as_ref()
234    {
235        return inner;
236    }
237    node
238}
239
240fn aconst_to_literal(c: &AConst) -> Option<LiteralValue> {
241    if c.isnull {
242        return Some(LiteralValue::Null);
243    }
244    match c.val.as_ref()? {
245        a_const::Val::Ival(i) => Some(LiteralValue::Integer(i64::from(i.ival))),
246        a_const::Val::Fval(f) => f.fval.parse::<f64>().ok().map(LiteralValue::Float),
247        a_const::Val::Boolval(b) => Some(LiteralValue::Bool(b.boolval)),
248        a_const::Val::Sval(s) => Some(LiteralValue::Text(s.sval.clone())),
249        a_const::Val::Bsval(_) => None,
250    }
251}
252
253/// If `node` is `nextval('seq')` (with optional `::regclass` cast on the arg),
254/// return the referenced sequence's [`QualifiedName`].
255fn nextval_target(
256    node: &NodeEnum,
257    default_schema: Option<&Identifier>,
258    location: &SourceLocation,
259) -> Result<Option<QualifiedName>, ParseError> {
260    let inner = unwrap_typecast(node);
261    let NodeEnum::FuncCall(fc) = inner else {
262        return Ok(None);
263    };
264    let func = match fc.funcname.last().and_then(|n| n.node.as_ref()) {
265        Some(NodeEnum::String(s)) => s.sval.as_str(),
266        _ => return Ok(None),
267    };
268    if !func.eq_ignore_ascii_case("nextval") {
269        return Ok(None);
270    }
271    let arg = fc.args.first().and_then(|n| n.node.as_ref());
272    let Some(arg_node) = arg else {
273        return Ok(None);
274    };
275    // The argument to nextval is normally a string literal cast to regclass:
276    // `nextval('app.seq1'::regclass)`. Strip the cast and look for an AConst Sval.
277    let inner_arg = unwrap_typecast(arg_node);
278    let NodeEnum::AConst(c) = inner_arg else {
279        return Ok(None);
280    };
281    let Some(a_const::Val::Sval(s)) = c.val.as_ref() else {
282        return Ok(None);
283    };
284    Ok(Some(parse_qualified_seq_name(
285        &s.sval,
286        default_schema,
287        location,
288    )?))
289}
290
291/// Parse a possibly-qualified sequence reference like `schema.seq` or `seq`.
292fn parse_qualified_seq_name(
293    s: &str,
294    default_schema: Option<&Identifier>,
295    location: &SourceLocation,
296) -> Result<QualifiedName, ParseError> {
297    let parts: Vec<&str> = s.split('.').collect();
298    match parts.len() {
299        1 => {
300            let name = ident(parts[0], location)?;
301            let schema = default_schema
302                .cloned()
303                .ok_or_else(|| ParseError::UnqualifiedName {
304                    location: location.clone(),
305                })?;
306            Ok(QualifiedName::new(schema, name))
307        }
308        2 => {
309            let schema = ident(parts[0], location)?;
310            let name = ident(parts[1], location)?;
311            Ok(QualifiedName::new(schema, name))
312        }
313        _ => Err(ParseError::Structural {
314            location: location.clone(),
315            message: format!("unsupported qualified sequence reference {s:?}"),
316        }),
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use std::path::PathBuf;
324
325    fn loc() -> SourceLocation {
326        SourceLocation::new(PathBuf::from("test.sql"), 1, 1)
327    }
328
329    fn parse_first(sql: &str) -> NodeEnum {
330        let parsed = pg_query::parse(sql).expect("parses");
331        parsed
332            .protobuf
333            .stmts
334            .into_iter()
335            .next()
336            .and_then(|raw| raw.stmt)
337            .and_then(|n| n.node)
338            .expect("at least one statement")
339    }
340
341    fn parse_select_expr(expr: &str) -> NodeEnum {
342        let n = parse_first(&format!("SELECT {expr};"));
343        let NodeEnum::SelectStmt(s) = n else { panic!() };
344        let target = s.target_list.into_iter().next().unwrap().node.unwrap();
345        let NodeEnum::ResTarget(rt) = target else {
346            panic!()
347        };
348        rt.val.unwrap().node.unwrap()
349    }
350
351    #[test]
352    fn literal_integer_default() {
353        let n = parse_select_expr("42");
354        let d = build_default_expr(&n, Some(&ColumnType::Integer), None, &loc()).unwrap();
355        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Integer(42))));
356    }
357
358    #[test]
359    fn literal_text_default() {
360        let n = parse_select_expr("'hello'");
361        let d = build_default_expr(&n, Some(&ColumnType::Text), None, &loc()).unwrap();
362        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Text(s)) if s == "hello"));
363    }
364
365    #[test]
366    fn null_default() {
367        let n = parse_select_expr("NULL");
368        let d = build_default_expr(&n, None, None, &loc()).unwrap();
369        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Null)));
370    }
371
372    #[test]
373    fn bool_default() {
374        let n = parse_select_expr("true");
375        let d = build_default_expr(&n, Some(&ColumnType::Boolean), None, &loc()).unwrap();
376        // Booleans in pg_query may parse as a function call `'t'::boolean` form
377        // — accept either Bool literal or Expr containing "true".
378        match d {
379            DefaultExpr::Literal(LiteralValue::Bool(true)) => {}
380            DefaultExpr::Expr(e) => assert!(e.canonical_text.contains("true")),
381            other => panic!("unexpected default: {other:?}"),
382        }
383    }
384
385    #[test]
386    fn nextval_qualified_default() {
387        let n = parse_select_expr("nextval('app.seq1'::regclass)");
388        let d = build_default_expr(&n, None, None, &loc()).unwrap();
389        match d {
390            DefaultExpr::Sequence(q) => assert_eq!(q.to_string(), "app.seq1"),
391            other => panic!("expected Sequence, got {other:?}"),
392        }
393    }
394
395    #[test]
396    fn nextval_unqualified_uses_default_schema() {
397        let n = parse_select_expr("nextval('seq1'::regclass)");
398        let app = Identifier::from_unquoted("app").unwrap();
399        let d = build_default_expr(&n, None, Some(&app), &loc()).unwrap();
400        match d {
401            DefaultExpr::Sequence(q) => assert_eq!(q.to_string(), "app.seq1"),
402            other => panic!("expected Sequence, got {other:?}"),
403        }
404    }
405
406    #[test]
407    fn func_call_other_than_nextval_is_expr() {
408        let n = parse_select_expr("now()");
409        let d = build_default_expr(&n, None, None, &loc()).unwrap();
410        assert!(matches!(d, DefaultExpr::Expr(_)));
411    }
412
413    #[test]
414    fn cast_to_target_strips_in_expr_arm() {
415        // `'a' || 'b'` is an expression — cast stripping has nothing to strip
416        // here. This test just checks the Expr arm runs.
417        let n = parse_select_expr("'a' || 'b'");
418        let d = build_default_expr(&n, Some(&ColumnType::Text), None, &loc()).unwrap();
419        assert!(matches!(d, DefaultExpr::Expr(_)));
420    }
421
422    #[test]
423    fn type_name_renders_with_typmod() {
424        // Parse a CREATE TABLE to get a real TypeName.
425        let stmt = parse_first("CREATE TABLE t (c varchar(50));");
426        let NodeEnum::CreateStmt(create) = stmt else {
427            panic!()
428        };
429        let elt = create.table_elts.into_iter().next().unwrap();
430        let NodeEnum::ColumnDef(col) = elt.node.unwrap() else {
431            panic!()
432        };
433        let tn = col.type_name.unwrap();
434        let s = render_type_name_to_string(&tn).unwrap();
435        let ct = ColumnType::parse_from_pg_type_string(&s).unwrap();
436        assert_eq!(ct, ColumnType::Varchar { len: Some(50) });
437    }
438}