1use 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
17pub 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
42pub 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
96pub 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
108pub 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
140pub 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
149pub 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
175fn 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
189fn 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
198fn 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
212pub fn type_name_to_column_type(
224 type_name: &TypeName,
225 location: &SourceLocation,
226) -> Result<ColumnType, ParseError> {
227 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 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
264pub 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
288fn 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
298fn 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
322fn 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 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
360fn 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 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 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 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 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 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 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}