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, ColumnRef, DefElem, Node, 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/// Extract the string value of a bare `String` [`Node`].
97///
98/// Returns `Some(sval)` when `node` is a [`NodeEnum::String`], otherwise `None`.
99/// Used where `pg_query` encodes a list element (publication name, column name)
100/// directly as a `String` node rather than wrapping it in a `DefElem`.
101pub fn node_string_value(node: &Node) -> Option<String> {
102    match node.node.as_ref()? {
103        NodeEnum::String(s) => Some(s.sval.clone()),
104        _ => None,
105    }
106}
107
108/// Extract a string value from a `DefElem.arg`.
109///
110/// Handles the encodings `pg_query` uses for scalar option values:
111/// - [`NodeEnum::String`] — bare identifier or quoted string.
112/// - [`NodeEnum::Integer`] / [`NodeEnum::Float`] — numeric option values,
113///   stringified to their textual form.
114/// - [`NodeEnum::List`] — dollar-quoted bodies are wrapped in a `List` whose
115///   first item is the body `String` (the optional second item is the
116///   dollar-quote tag); this returns that first `String`'s value.
117///
118/// Returns `None` for an absent arg or any other node kind.
119pub fn def_elem_string(de: &DefElem) -> Option<String> {
120    let arg = de.arg.as_ref()?;
121    match arg.node.as_ref()? {
122        NodeEnum::String(s) => Some(s.sval.clone()),
123        NodeEnum::Integer(i) => Some(i.ival.to_string()),
124        NodeEnum::Float(f) => Some(f.fval.clone()),
125        NodeEnum::List(list) => list
126            .items
127            .first()
128            .and_then(|n| n.node.as_ref())
129            .and_then(|n| {
130                if let NodeEnum::String(s) = n {
131                    Some(s.sval.clone())
132                } else {
133                    None
134                }
135            }),
136        _ => None,
137    }
138}
139
140/// Build an unquoted [`Identifier`], wrapping [`crate::ir::IrError`] into a
141/// source-located [`ParseError::Ir`].
142pub fn ident(s: &str, location: &SourceLocation) -> Result<Identifier, ParseError> {
143    Identifier::from_unquoted(s).map_err(|e| ParseError::Ir {
144        location: location.clone(),
145        source: crate::ir::IrError::InvalidIdentifier(e.to_string()),
146    })
147}
148
149/// Render a `pg_query::TypeName` into a string `ColumnType::parse_from_pg_type_string`
150/// can consume.
151///
152/// Strategy: take the *last* segment of `names` (Postgres prefixes types with
153/// `pg_catalog.` internally; the parser is alias-aware), append a parenthesized
154/// list of typmod arguments, append `[]` for each array dimension.
155pub fn render_type_name_to_string(type_name: &TypeName) -> Option<String> {
156    let last = type_name.names.last()?.node.as_ref()?;
157    let NodeEnum::String(s) = last else {
158        return None;
159    };
160    let bare = s.sval.clone();
161    let mut out = bare;
162    if !type_name.typmods.is_empty() {
163        let mut args: Vec<String> = Vec::with_capacity(type_name.typmods.len());
164        for n in &type_name.typmods {
165            args.push(typmod_arg_to_string(n.node.as_ref()?)?);
166        }
167        out = format!("{out}({})", args.join(","));
168    }
169    for _ in 0..type_name.array_bounds.len() {
170        out.push_str("[]");
171    }
172    Some(out)
173}
174
175/// Stringify a single type-modifier argument node.
176///
177/// Type modifiers are an `expr_list`, so an argument is either a literal
178/// ([`AConst`]) or a bareword. The bareword form is what `PostGIS` uses for the
179/// subtype in `geometry(Point,4326)` — Postgres parses it as a single-field
180/// [`ColumnRef`], not a constant.
181fn typmod_arg_to_string(node: &NodeEnum) -> Option<String> {
182    match node {
183        NodeEnum::AConst(c) => literal_arg_to_string(c),
184        NodeEnum::ColumnRef(cref) => columnref_ident(cref),
185        _ => None,
186    }
187}
188
189/// Convert an [`AConst`] used as a typmod argument to its canonical string form.
190fn literal_arg_to_string(c: &AConst) -> Option<String> {
191    match c.val.as_ref()? {
192        a_const::Val::Ival(i) => Some(i.ival.to_string()),
193        a_const::Val::Sval(s) => Some(s.sval.clone()),
194        _ => None,
195    }
196}
197
198/// Extract the identifier text of a single-field bareword [`ColumnRef`].
199///
200/// A multi-field ref (`a.b`) is not a valid type-modifier shape, so it is
201/// rejected (returns `None`).
202fn columnref_ident(cref: &ColumnRef) -> Option<String> {
203    let [field] = cref.fields.as_slice() else {
204        return None;
205    };
206    match field.node.as_ref()? {
207        NodeEnum::String(s) => Some(s.sval.clone()),
208        _ => None,
209    }
210}
211
212/// Convert a `TypeName` into a [`ColumnType`], propagating parser errors.
213///
214/// When `TypeName.names` contains two String nodes and the first is not
215/// `pg_catalog`, the type is schema-qualified by the user (e.g. `app.order_status`).
216/// In that case we emit `ColumnType::UserDefined(QualifiedName)` so that the
217/// AST resolution pass can validate that the referenced type is declared in source.
218///
219/// Unqualified single-segment names are handled by the usual string path; if they
220/// don't match any built-in they fall through to `ColumnType::Other`, which is
221/// intentional — unqualified user types must be resolved via `default_schema` at
222/// domain/composite parse time rather than here.
223pub fn type_name_to_column_type(
224    type_name: &TypeName,
225    location: &SourceLocation,
226) -> Result<ColumnType, ParseError> {
227    // Collect String nodes from names; fail fast on unexpected kinds.
228    let name_strings: Vec<&str> = type_name
229        .names
230        .iter()
231        .map(|n| match n.node.as_ref() {
232            Some(NodeEnum::String(s)) => Ok(s.sval.as_str()),
233            other => Err(ParseError::Structural {
234                location: location.clone(),
235                message: format!(
236                    "expected String node in type-name list, got {:?}",
237                    other.map(std::mem::discriminant),
238                ),
239            }),
240        })
241        .collect::<Result<Vec<_>, _>>()?;
242
243    // Two-segment, non-pg_catalog prefix → user-defined type reference.
244    if let [schema, name] = name_strings.as_slice()
245        && *schema != "pg_catalog"
246    {
247        let schema_id = ident(schema, location)?;
248        let name_id = ident(name, location)?;
249        return Ok(ColumnType::UserDefined(QualifiedName::new(
250            schema_id, name_id,
251        )));
252    }
253
254    let s = render_type_name_to_string(type_name).ok_or_else(|| ParseError::Structural {
255        location: location.clone(),
256        message: "could not stringify type name".into(),
257    })?;
258    ColumnType::parse_from_pg_type_string(&s).map_err(|e| ParseError::Structural {
259        location: location.clone(),
260        message: format!("invalid column type {s:?}: {e}"),
261    })
262}
263
264/// Convert a column-default expression node into a [`DefaultExpr`].
265///
266/// Recognizes:
267/// - Bare [`AConst`] integer/float/string/bool/null literals → [`DefaultExpr::Literal`].
268/// - `nextval('seq')` and `nextval('schema.seq')` → [`DefaultExpr::Sequence`].
269/// - Anything else → [`DefaultExpr::Expr`] via the normalizer.
270///
271/// `target_type` enables redundant-cast stripping in the `Expr` arm.
272pub fn build_default_expr(
273    node: &NodeEnum,
274    target_type: Option<&ColumnType>,
275    default_schema: Option<&Identifier>,
276    location: &SourceLocation,
277) -> Result<DefaultExpr, ParseError> {
278    if let Some(lit) = literal_from_node(node) {
279        return Ok(DefaultExpr::Literal(lit));
280    }
281    if let Some(seq) = nextval_target(node, default_schema, location)? {
282        return Ok(DefaultExpr::Sequence(seq));
283    }
284    let normalized = normalize_expr::from_pg_node(node, target_type, location)?;
285    Ok(DefaultExpr::Expr(normalized))
286}
287
288/// If `node` is a literal `AConst` (possibly wrapped in a `TypeCast` whose target
289/// matches the column type), return the equivalent [`LiteralValue`].
290fn literal_from_node(node: &NodeEnum) -> Option<LiteralValue> {
291    let inner = unwrap_typecast(node);
292    match inner {
293        NodeEnum::AConst(c) => aconst_to_literal(c),
294        _ => None,
295    }
296}
297
298/// Strip a single layer of `TypeCast` wrapping, if any.
299fn unwrap_typecast(node: &NodeEnum) -> &NodeEnum {
300    if let NodeEnum::TypeCast(cast) = node
301        && let Some(arg) = cast.arg.as_ref()
302        && let Some(inner) = arg.node.as_ref()
303    {
304        return inner;
305    }
306    node
307}
308
309fn aconst_to_literal(c: &AConst) -> Option<LiteralValue> {
310    if c.isnull {
311        return Some(LiteralValue::Null);
312    }
313    match c.val.as_ref()? {
314        a_const::Val::Ival(i) => Some(LiteralValue::Integer(i64::from(i.ival))),
315        a_const::Val::Fval(f) => f.fval.parse::<f64>().ok().map(LiteralValue::Float),
316        a_const::Val::Boolval(b) => Some(LiteralValue::Bool(b.boolval)),
317        a_const::Val::Sval(s) => Some(LiteralValue::Text(s.sval.clone())),
318        a_const::Val::Bsval(_) => None,
319    }
320}
321
322/// If `node` is `nextval('seq')` (with optional `::regclass` cast on the arg),
323/// return the referenced sequence's [`QualifiedName`].
324fn nextval_target(
325    node: &NodeEnum,
326    default_schema: Option<&Identifier>,
327    location: &SourceLocation,
328) -> Result<Option<QualifiedName>, ParseError> {
329    let inner = unwrap_typecast(node);
330    let NodeEnum::FuncCall(fc) = inner else {
331        return Ok(None);
332    };
333    let func = match fc.funcname.last().and_then(|n| n.node.as_ref()) {
334        Some(NodeEnum::String(s)) => s.sval.as_str(),
335        _ => return Ok(None),
336    };
337    if !func.eq_ignore_ascii_case("nextval") {
338        return Ok(None);
339    }
340    let arg = fc.args.first().and_then(|n| n.node.as_ref());
341    let Some(arg_node) = arg else {
342        return Ok(None);
343    };
344    // The argument to nextval is normally a string literal cast to regclass:
345    // `nextval('app.seq1'::regclass)`. Strip the cast and look for an AConst Sval.
346    let inner_arg = unwrap_typecast(arg_node);
347    let NodeEnum::AConst(c) = inner_arg else {
348        return Ok(None);
349    };
350    let Some(a_const::Val::Sval(s)) = c.val.as_ref() else {
351        return Ok(None);
352    };
353    Ok(Some(parse_qualified_seq_name(
354        &s.sval,
355        default_schema,
356        location,
357    )?))
358}
359
360/// Parse a possibly-qualified sequence reference like `schema.seq` or `seq`.
361fn parse_qualified_seq_name(
362    s: &str,
363    default_schema: Option<&Identifier>,
364    location: &SourceLocation,
365) -> Result<QualifiedName, ParseError> {
366    let parts: Vec<&str> = s.split('.').collect();
367    match parts.len() {
368        1 => {
369            let name = ident(parts[0], location)?;
370            let schema = default_schema
371                .cloned()
372                .ok_or_else(|| ParseError::UnqualifiedName {
373                    location: location.clone(),
374                })?;
375            Ok(QualifiedName::new(schema, name))
376        }
377        2 => {
378            let schema = ident(parts[0], location)?;
379            let name = ident(parts[1], location)?;
380            Ok(QualifiedName::new(schema, name))
381        }
382        _ => Err(ParseError::Structural {
383            location: location.clone(),
384            message: format!("unsupported qualified sequence reference {s:?}"),
385        }),
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392    use std::path::PathBuf;
393
394    fn loc() -> SourceLocation {
395        SourceLocation::new(PathBuf::from("test.sql"), 1, 1)
396    }
397
398    fn parse_first(sql: &str) -> NodeEnum {
399        let parsed = pg_query::parse(sql).expect("parses");
400        parsed
401            .protobuf
402            .stmts
403            .into_iter()
404            .next()
405            .and_then(|raw| raw.stmt)
406            .and_then(|n| n.node)
407            .expect("at least one statement")
408    }
409
410    fn parse_select_expr(expr: &str) -> NodeEnum {
411        let n = parse_first(&format!("SELECT {expr};"));
412        let NodeEnum::SelectStmt(s) = n else { panic!() };
413        let target = s.target_list.into_iter().next().unwrap().node.unwrap();
414        let NodeEnum::ResTarget(rt) = target else {
415            panic!()
416        };
417        rt.val.unwrap().node.unwrap()
418    }
419
420    #[test]
421    fn literal_integer_default() {
422        let n = parse_select_expr("42");
423        let d = build_default_expr(&n, Some(&ColumnType::Integer), None, &loc()).unwrap();
424        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Integer(42))));
425    }
426
427    #[test]
428    fn literal_text_default() {
429        let n = parse_select_expr("'hello'");
430        let d = build_default_expr(&n, Some(&ColumnType::Text), None, &loc()).unwrap();
431        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Text(s)) if s == "hello"));
432    }
433
434    #[test]
435    fn null_default() {
436        let n = parse_select_expr("NULL");
437        let d = build_default_expr(&n, None, None, &loc()).unwrap();
438        assert!(matches!(d, DefaultExpr::Literal(LiteralValue::Null)));
439    }
440
441    #[test]
442    fn bool_default() {
443        let n = parse_select_expr("true");
444        let d = build_default_expr(&n, Some(&ColumnType::Boolean), None, &loc()).unwrap();
445        // Booleans in pg_query may parse as a function call `'t'::boolean` form
446        // — accept either Bool literal or Expr containing "true".
447        match d {
448            DefaultExpr::Literal(LiteralValue::Bool(true)) => {}
449            DefaultExpr::Expr(e) => assert!(e.canonical_text.contains("true")),
450            other => panic!("unexpected default: {other:?}"),
451        }
452    }
453
454    #[test]
455    fn nextval_qualified_default() {
456        let n = parse_select_expr("nextval('app.seq1'::regclass)");
457        let d = build_default_expr(&n, None, None, &loc()).unwrap();
458        match d {
459            DefaultExpr::Sequence(q) => assert_eq!(q.to_string(), "app.seq1"),
460            other => panic!("expected Sequence, got {other:?}"),
461        }
462    }
463
464    #[test]
465    fn nextval_unqualified_uses_default_schema() {
466        let n = parse_select_expr("nextval('seq1'::regclass)");
467        let app = Identifier::from_unquoted("app").unwrap();
468        let d = build_default_expr(&n, None, Some(&app), &loc()).unwrap();
469        match d {
470            DefaultExpr::Sequence(q) => assert_eq!(q.to_string(), "app.seq1"),
471            other => panic!("expected Sequence, got {other:?}"),
472        }
473    }
474
475    #[test]
476    fn func_call_other_than_nextval_is_expr() {
477        let n = parse_select_expr("now()");
478        let d = build_default_expr(&n, None, None, &loc()).unwrap();
479        assert!(matches!(d, DefaultExpr::Expr(_)));
480    }
481
482    #[test]
483    fn cast_to_target_strips_in_expr_arm() {
484        // `'a' || 'b'` is an expression — cast stripping has nothing to strip
485        // here. This test just checks the Expr arm runs.
486        let n = parse_select_expr("'a' || 'b'");
487        let d = build_default_expr(&n, Some(&ColumnType::Text), None, &loc()).unwrap();
488        assert!(matches!(d, DefaultExpr::Expr(_)));
489    }
490
491    #[test]
492    fn type_name_renders_with_typmod() {
493        // Parse a CREATE TABLE to get a real TypeName.
494        let stmt = parse_first("CREATE TABLE t (c varchar(50));");
495        let NodeEnum::CreateStmt(create) = stmt else {
496            panic!()
497        };
498        let elt = create.table_elts.into_iter().next().unwrap();
499        let NodeEnum::ColumnDef(col) = elt.node.unwrap() else {
500            panic!()
501        };
502        let tn = col.type_name.unwrap();
503        let s = render_type_name_to_string(&tn).unwrap();
504        let ct = ColumnType::parse_from_pg_type_string(&s).unwrap();
505        assert_eq!(ct, ColumnType::Varchar { len: Some(50) });
506    }
507
508    /// Extract the first column's `TypeName` from a one-column `CREATE TABLE`.
509    fn first_column_type_name(sql: &str) -> TypeName {
510        let NodeEnum::CreateStmt(create) = parse_first(sql) else {
511            panic!("expected CREATE TABLE")
512        };
513        let elt = create.table_elts.into_iter().next().unwrap();
514        let NodeEnum::ColumnDef(col) = elt.node.unwrap() else {
515            panic!("expected ColumnDef")
516        };
517        col.type_name.unwrap()
518    }
519
520    #[test]
521    fn parameterized_postgis_types_parse() {
522        // PostGIS parameterized types put the subtype (Point / MultiPolygon / …)
523        // in the typmod list as a *bareword*, which pg_query parses as a ColumnRef
524        // (type modifiers are an expr_list), not an AConst. The renderer must
525        // stringify it; pg_query lowercases the bareword, so the canonical raw is
526        // lowercase. (issue #40)
527        let cases = [
528            ("geometry(Point,4326)", "geometry(point,4326)"),
529            (
530                "geography(MultiPolygon,4326)",
531                "geography(multipolygon,4326)",
532            ),
533            ("geometry(geometry,4326)", "geometry(geometry,4326)"),
534        ];
535        for (decl, expected_raw) in cases {
536            let tn = first_column_type_name(&format!("CREATE TABLE t (c {decl});"));
537            let ct = type_name_to_column_type(&tn, &loc())
538                .unwrap_or_else(|e| panic!("{decl} should parse, got {e:?}"));
539            assert!(
540                matches!(&ct, ColumnType::Other { raw } if raw == expected_raw),
541                "{decl} -> {ct:?}"
542            );
543        }
544    }
545
546    #[test]
547    fn heterogeneous_typmod_args_render() {
548        // A type modifier list may legally mix string/int literals and barewords.
549        // The renderer dispatches per node kind rather than assuming AConst.
550        let tn = first_column_type_name("CREATE TABLE t (c mytype('foo',1,bar));");
551        let ct = type_name_to_column_type(&tn, &loc()).expect("mixed typmods should parse");
552        assert!(
553            matches!(&ct, ColumnType::Other { raw } if raw == "mytype(foo,1,bar)"),
554            "got {ct:?}"
555        );
556    }
557}