pgevolve_core/parse/builder/
shared.rs1use 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
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 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
105pub 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
134fn 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
143pub fn type_name_to_column_type(
155 type_name: &TypeName,
156 location: &SourceLocation,
157) -> Result<ColumnType, ParseError> {
158 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 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
195pub 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
219fn 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
229fn 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
253fn 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 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
291fn 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 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 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 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}