pgevolve_core/parse/
normalize_body.rs1use serde::{Deserialize, Serialize};
21
22#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
24pub struct NormalizedBody {
25 canonical_text: String,
26 canonical_hash: [u8; 32],
27}
28
29#[derive(Debug, thiserror::Error)]
31pub enum BodyError {
32 #[error("pg_query rejected body: {0}")]
34 Parse(String),
35}
36
37impl NormalizedBody {
38 pub const fn empty() -> Self {
43 Self {
44 canonical_text: String::new(),
45 canonical_hash: [0u8; 32],
46 }
47 }
48
49 pub fn from_raw_canonical(canonical_text: String) -> Self {
58 let canonical_hash = hash_canonical(&canonical_text);
59 Self {
60 canonical_text,
61 canonical_hash,
62 }
63 }
64
65 pub fn from_sql(sql: &str) -> Result<Self, BodyError> {
72 let parsed = pg_query::parse(sql).map_err(|e| BodyError::Parse(e.to_string()))?;
73 let mut protobuf = parsed.protobuf;
79 strip_redundant_qualifiers(&mut protobuf);
80 let deparsed = pg_query::deparse(&protobuf).unwrap_or_default();
81 let source = if deparsed.is_empty() { sql } else { &deparsed };
82 let canonical_text = collapse_whitespace(source);
83 let canonical_hash = hash_canonical(&canonical_text);
84 Ok(Self {
85 canonical_text,
86 canonical_hash,
87 })
88 }
89
90 pub fn canonical_text(&self) -> &str {
93 &self.canonical_text
94 }
95
96 #[allow(clippy::missing_const_for_fn)]
103 pub fn canonical_hash(&self) -> &[u8; 32] {
104 &self.canonical_hash
105 }
106}
107
108fn strip_redundant_qualifiers(root: &mut pg_query::protobuf::ParseResult) {
130 use pg_query::NodeEnum;
131 for stmt in &mut root.stmts {
132 let Some(node) = stmt.stmt.as_mut().and_then(|n| n.node.as_mut()) else {
133 continue;
134 };
135 if let NodeEnum::SelectStmt(sel) = node {
136 strip_qualifiers_in_select(sel);
137 }
138 }
139}
140
141fn strip_qualifiers_in_select(sel: &mut pg_query::protobuf::SelectStmt) {
142 if let Some(larg) = sel.larg.as_mut() {
149 strip_qualifiers_in_select(larg);
150 }
151 if let Some(rarg) = sel.rarg.as_mut() {
152 strip_qualifiers_in_select(rarg);
153 }
154
155 if let Some(with) = sel.with_clause.as_mut() {
157 for cte in &mut with.ctes {
158 recurse_cte(cte);
159 }
160 }
161
162 for n in &mut sel.from_clause {
164 recurse_from_node(n);
165 }
166
167 for n in &mut sel.target_list {
169 recurse_subselects_in_node(n);
170 }
171 if let Some(w) = sel.where_clause.as_mut() {
172 recurse_subselects_in_node(w);
173 }
174 if let Some(h) = sel.having_clause.as_mut() {
175 recurse_subselects_in_node(h);
176 }
177
178 let from_names = collect_from_qualifiers(&sel.from_clause);
179 let Some(unique_name) = unique_from_qualifier(&from_names) else {
180 return;
181 };
182
183 for n in &mut sel.target_list {
185 strip_qualifier_in_node(n, &unique_name);
186 }
187 if let Some(w) = sel.where_clause.as_mut() {
188 strip_qualifier_in_node(w, &unique_name);
189 }
190 for n in &mut sel.group_clause {
191 strip_qualifier_in_node(n, &unique_name);
192 }
193 if let Some(h) = sel.having_clause.as_mut() {
194 strip_qualifier_in_node(h, &unique_name);
195 }
196 for n in &mut sel.sort_clause {
197 strip_qualifier_in_node(n, &unique_name);
198 }
199}
200
201fn recurse_select_node(node: &mut pg_query::protobuf::Node) {
203 use pg_query::NodeEnum;
204 if let Some(NodeEnum::SelectStmt(sel)) = node.node.as_mut() {
205 strip_qualifiers_in_select(sel);
206 }
207}
208
209fn recurse_cte(cte: &mut pg_query::protobuf::Node) {
211 use pg_query::NodeEnum;
212 if let Some(NodeEnum::CommonTableExpr(c)) = cte.node.as_mut()
213 && let Some(q) = c.ctequery.as_mut()
214 {
215 recurse_select_node(q);
216 }
217}
218
219fn recurse_from_node(n: &mut pg_query::protobuf::Node) {
223 use pg_query::NodeEnum;
224 match n.node.as_mut() {
225 Some(NodeEnum::RangeSubselect(rs)) => {
226 if let Some(q) = rs.subquery.as_mut() {
227 recurse_select_node(q);
228 }
229 }
230 Some(NodeEnum::JoinExpr(je)) => {
231 if let Some(l) = je.larg.as_mut() {
232 recurse_from_node(l);
233 }
234 if let Some(r) = je.rarg.as_mut() {
235 recurse_from_node(r);
236 }
237 }
238 _ => {}
239 }
240}
241
242fn recurse_subselects_in_node(n: &mut pg_query::protobuf::Node) {
247 use pg_query::NodeEnum;
248 let Some(node) = n.node.as_mut() else { return };
249 match node {
250 NodeEnum::SubLink(sl) => {
251 if let Some(t) = sl.testexpr.as_mut() {
252 recurse_subselects_in_node(t);
253 }
254 if let Some(s) = sl.subselect.as_mut() {
255 recurse_select_node(s);
256 }
257 }
258 NodeEnum::ResTarget(rt) => {
259 if let Some(v) = rt.val.as_mut() {
260 recurse_subselects_in_node(v);
261 }
262 }
263 NodeEnum::AExpr(e) => {
264 if let Some(l) = e.lexpr.as_mut() {
265 recurse_subselects_in_node(l);
266 }
267 if let Some(r) = e.rexpr.as_mut() {
268 recurse_subselects_in_node(r);
269 }
270 }
271 NodeEnum::BoolExpr(e) => {
272 for a in &mut e.args {
273 recurse_subselects_in_node(a);
274 }
275 }
276 NodeEnum::FuncCall(fc) => {
277 for a in &mut fc.args {
278 recurse_subselects_in_node(a);
279 }
280 if let Some(filt) = fc.agg_filter.as_mut() {
281 recurse_subselects_in_node(filt);
282 }
283 for a in &mut fc.agg_order {
284 recurse_subselects_in_node(a);
285 }
286 }
287 NodeEnum::CoalesceExpr(c) => {
288 for a in &mut c.args {
289 recurse_subselects_in_node(a);
290 }
291 }
292 NodeEnum::CaseExpr(c) => {
293 if let Some(arg) = c.arg.as_mut() {
294 recurse_subselects_in_node(arg);
295 }
296 for w in &mut c.args {
297 recurse_subselects_in_node(w);
298 }
299 if let Some(d) = c.defresult.as_mut() {
300 recurse_subselects_in_node(d);
301 }
302 }
303 NodeEnum::CaseWhen(w) => {
304 if let Some(e) = w.expr.as_mut() {
305 recurse_subselects_in_node(e);
306 }
307 if let Some(r) = w.result.as_mut() {
308 recurse_subselects_in_node(r);
309 }
310 }
311 NodeEnum::TypeCast(tc) => {
312 if let Some(arg) = tc.arg.as_mut() {
313 recurse_subselects_in_node(arg);
314 }
315 }
316 NodeEnum::SortBy(sb) => {
317 if let Some(arg) = sb.node.as_mut() {
318 recurse_subselects_in_node(arg);
319 }
320 }
321 NodeEnum::List(l) => {
322 for item in &mut l.items {
323 recurse_subselects_in_node(item);
324 }
325 }
326 NodeEnum::NullTest(nt) => {
327 if let Some(arg) = nt.arg.as_mut() {
328 recurse_subselects_in_node(arg);
329 }
330 }
331 NodeEnum::BooleanTest(bt) => {
332 if let Some(arg) = bt.arg.as_mut() {
333 recurse_subselects_in_node(arg);
334 }
335 }
336 _ => {}
337 }
338}
339
340fn collect_from_qualifiers(from: &[pg_query::protobuf::Node]) -> Vec<String> {
341 use pg_query::NodeEnum;
342 let mut names = Vec::new();
343 for n in from {
344 let Some(node) = n.node.as_ref() else {
345 continue;
346 };
347 if let NodeEnum::RangeVar(rv) = node {
348 let qual = rv
349 .alias
350 .as_ref()
351 .map_or_else(|| rv.relname.clone(), |a| a.aliasname.clone());
352 if !qual.is_empty() {
353 names.push(qual);
354 }
355 }
356 }
357 names
358}
359
360fn unique_from_qualifier(names: &[String]) -> Option<String> {
361 if names.len() == 1 {
362 Some(names[0].clone())
363 } else {
364 None
365 }
366}
367
368fn strip_qualifier_in_node(n: &mut pg_query::protobuf::Node, qualifier: &str) {
369 use pg_query::NodeEnum;
370 let Some(node) = n.node.as_mut() else { return };
371 match node {
372 NodeEnum::ColumnRef(cref) => {
373 if cref.fields.len() == 2
374 && let Some(first) = cref.fields.first()
375 && let Some(NodeEnum::String(s)) = first.node.as_ref()
376 && s.sval == qualifier
377 {
378 cref.fields.remove(0);
379 }
380 }
381 NodeEnum::ResTarget(rt) => {
382 if let Some(v) = rt.val.as_mut() {
383 strip_qualifier_in_node(v, qualifier);
384 }
385 }
386 NodeEnum::AExpr(e) => {
387 if let Some(l) = e.lexpr.as_mut() {
388 strip_qualifier_in_node(l, qualifier);
389 }
390 if let Some(r) = e.rexpr.as_mut() {
391 strip_qualifier_in_node(r, qualifier);
392 }
393 }
394 NodeEnum::BoolExpr(e) => {
395 for a in &mut e.args {
396 strip_qualifier_in_node(a, qualifier);
397 }
398 }
399 NodeEnum::FuncCall(fc) => {
400 for a in &mut fc.args {
401 strip_qualifier_in_node(a, qualifier);
402 }
403 if let Some(filt) = fc.agg_filter.as_mut() {
404 strip_qualifier_in_node(filt, qualifier);
405 }
406 for a in &mut fc.agg_order {
407 strip_qualifier_in_node(a, qualifier);
408 }
409 }
410 NodeEnum::CoalesceExpr(c) => {
411 for a in &mut c.args {
412 strip_qualifier_in_node(a, qualifier);
413 }
414 }
415 NodeEnum::CaseExpr(c) => {
416 if let Some(arg) = c.arg.as_mut() {
417 strip_qualifier_in_node(arg, qualifier);
418 }
419 for w in &mut c.args {
420 strip_qualifier_in_node(w, qualifier);
421 }
422 if let Some(d) = c.defresult.as_mut() {
423 strip_qualifier_in_node(d, qualifier);
424 }
425 }
426 NodeEnum::CaseWhen(w) => {
427 if let Some(e) = w.expr.as_mut() {
428 strip_qualifier_in_node(e, qualifier);
429 }
430 if let Some(r) = w.result.as_mut() {
431 strip_qualifier_in_node(r, qualifier);
432 }
433 }
434 NodeEnum::TypeCast(tc) => {
435 if let Some(arg) = tc.arg.as_mut() {
436 strip_qualifier_in_node(arg, qualifier);
437 }
438 }
439 NodeEnum::SortBy(sb) => {
440 if let Some(arg) = sb.node.as_mut() {
441 strip_qualifier_in_node(arg, qualifier);
442 }
443 }
444 NodeEnum::List(l) => {
445 for item in &mut l.items {
446 strip_qualifier_in_node(item, qualifier);
447 }
448 }
449 NodeEnum::SubLink(sl) => {
450 if let Some(t) = sl.testexpr.as_mut() {
451 strip_qualifier_in_node(t, qualifier);
452 }
453 }
454 NodeEnum::NullTest(nt) => {
455 if let Some(arg) = nt.arg.as_mut() {
456 strip_qualifier_in_node(arg, qualifier);
457 }
458 }
459 NodeEnum::BooleanTest(bt) => {
460 if let Some(arg) = bt.arg.as_mut() {
461 strip_qualifier_in_node(arg, qualifier);
462 }
463 }
464 _ => {}
465 }
466}
467
468fn collapse_whitespace(s: &str) -> String {
469 s.split_whitespace().collect::<Vec<_>>().join(" ")
470}
471
472fn hash_canonical(text: &str) -> [u8; 32] {
473 let mut h = blake3::Hasher::new();
474 h.update(b"pgevolve-normalized-body-v1\n");
475 h.update(text.as_bytes());
476 *h.finalize().as_bytes()
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482
483 #[test]
484 fn strip_qualifiers_equates_pg14_and_source_form() {
485 let pg14 = "SELECT products.category, count(*) AS cnt, avg(products.price) AS avg_price FROM app.products GROUP BY products.category";
486 let src = "SELECT category, count(*) AS cnt, avg(price) AS avg_price FROM app.products GROUP BY category";
487 let a = NormalizedBody::from_sql(pg14).unwrap();
488 let b = NormalizedBody::from_sql(src).unwrap();
489 assert_eq!(
490 a.canonical_text(),
491 b.canonical_text(),
492 "left: {} right: {}",
493 a.canonical_text(),
494 b.canonical_text()
495 );
496 }
497
498 #[test]
501 fn subquery_inner_scope_strip() {
502 let qualified =
503 NormalizedBody::from_sql("SELECT s.id FROM (SELECT t.id FROM app.t) s").unwrap();
504 let unqualified =
505 NormalizedBody::from_sql("SELECT s.id FROM (SELECT id FROM app.t) s").unwrap();
506 assert_eq!(
507 qualified.canonical_text(),
508 unqualified.canonical_text(),
509 "inner t.id should strip to id: {} vs {}",
510 qualified.canonical_text(),
511 unqualified.canonical_text()
512 );
513 }
514
515 #[test]
517 fn cte_inner_scope_strip() {
518 let qualified =
519 NormalizedBody::from_sql("WITH c AS (SELECT u.x FROM app.u) SELECT * FROM c").unwrap();
520 let unqualified =
521 NormalizedBody::from_sql("WITH c AS (SELECT x FROM app.u) SELECT * FROM c").unwrap();
522 assert_eq!(
523 qualified.canonical_text(),
524 unqualified.canonical_text(),
525 "inner u.x should strip to x: {} vs {}",
526 qualified.canonical_text(),
527 unqualified.canonical_text()
528 );
529 }
530
531 #[test]
536 fn correlated_subquery_outer_qualifier_preserved() {
537 let body = NormalizedBody::from_sql(
538 "SELECT a FROM app.t WHERE EXISTS (SELECT 1 FROM app.u WHERE u.fk = t.pk)",
539 )
540 .unwrap();
541 assert!(
543 body.canonical_text().contains("t.pk"),
544 "correlated outer ref t.pk must be preserved, got: {}",
545 body.canonical_text()
546 );
547 assert!(
549 !body.canonical_text().contains("u.fk"),
550 "inner u.fk should have stripped to fk, got: {}",
551 body.canonical_text()
552 );
553
554 let variant = NormalizedBody::from_sql(
557 "SELECT a FROM app.t WHERE EXISTS (SELECT 1 FROM app.u WHERE u.fk = t.other)",
558 )
559 .unwrap();
560 assert_ne!(
561 body.canonical_hash(),
562 variant.canonical_hash(),
563 "distinct correlation predicates must not collapse: {} vs {}",
564 body.canonical_text(),
565 variant.canonical_text()
566 );
567 }
568
569 #[test]
571 fn distinct_bodies_differ() {
572 let a = NormalizedBody::from_sql("SELECT a FROM app.t").unwrap();
573 let b = NormalizedBody::from_sql("SELECT b FROM app.t").unwrap();
574 assert_ne!(a.canonical_hash(), b.canonical_hash());
575 }
576
577 #[test]
580 fn multi_table_top_level_from_unchanged() {
581 let body =
582 NormalizedBody::from_sql("SELECT t.a, u.b FROM app.t JOIN app.u ON t.k = u.k").unwrap();
583 assert!(
584 body.canonical_text().contains("t.a"),
585 "join scope must keep t.a, got: {}",
586 body.canonical_text()
587 );
588 assert!(
589 body.canonical_text().contains("u.b"),
590 "join scope must keep u.b, got: {}",
591 body.canonical_text()
592 );
593 }
594}