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) {
123 use pg_query::NodeEnum;
124 for stmt in &mut root.stmts {
125 let Some(node) = stmt.stmt.as_mut().and_then(|n| n.node.as_mut()) else {
126 continue;
127 };
128 if let NodeEnum::SelectStmt(sel) = node {
129 strip_qualifiers_in_select(sel);
130 }
131 }
132}
133
134fn strip_qualifiers_in_select(sel: &mut pg_query::protobuf::SelectStmt) {
135 if let Some(larg) = sel.larg.as_mut() {
137 strip_qualifiers_in_select(larg);
138 }
139 if let Some(rarg) = sel.rarg.as_mut() {
140 strip_qualifiers_in_select(rarg);
141 }
142
143 let from_names = collect_from_qualifiers(&sel.from_clause);
144 let Some(unique_name) = unique_from_qualifier(&from_names) else {
145 return;
146 };
147
148 for n in &mut sel.target_list {
150 strip_qualifier_in_node(n, &unique_name);
151 }
152 if let Some(w) = sel.where_clause.as_mut() {
153 strip_qualifier_in_node(w, &unique_name);
154 }
155 for n in &mut sel.group_clause {
156 strip_qualifier_in_node(n, &unique_name);
157 }
158 if let Some(h) = sel.having_clause.as_mut() {
159 strip_qualifier_in_node(h, &unique_name);
160 }
161 for n in &mut sel.sort_clause {
162 strip_qualifier_in_node(n, &unique_name);
163 }
164}
165
166fn collect_from_qualifiers(from: &[pg_query::protobuf::Node]) -> Vec<String> {
167 use pg_query::NodeEnum;
168 let mut names = Vec::new();
169 for n in from {
170 let Some(node) = n.node.as_ref() else {
171 continue;
172 };
173 if let NodeEnum::RangeVar(rv) = node {
174 let qual = rv
175 .alias
176 .as_ref()
177 .map_or_else(|| rv.relname.clone(), |a| a.aliasname.clone());
178 if !qual.is_empty() {
179 names.push(qual);
180 }
181 }
182 }
183 names
184}
185
186fn unique_from_qualifier(names: &[String]) -> Option<String> {
187 if names.len() == 1 {
188 Some(names[0].clone())
189 } else {
190 None
191 }
192}
193
194fn strip_qualifier_in_node(n: &mut pg_query::protobuf::Node, qualifier: &str) {
195 use pg_query::NodeEnum;
196 let Some(node) = n.node.as_mut() else { return };
197 match node {
198 NodeEnum::ColumnRef(cref) => {
199 if cref.fields.len() == 2
200 && let Some(first) = cref.fields.first()
201 && let Some(NodeEnum::String(s)) = first.node.as_ref()
202 && s.sval == qualifier
203 {
204 cref.fields.remove(0);
205 }
206 }
207 NodeEnum::ResTarget(rt) => {
208 if let Some(v) = rt.val.as_mut() {
209 strip_qualifier_in_node(v, qualifier);
210 }
211 }
212 NodeEnum::AExpr(e) => {
213 if let Some(l) = e.lexpr.as_mut() {
214 strip_qualifier_in_node(l, qualifier);
215 }
216 if let Some(r) = e.rexpr.as_mut() {
217 strip_qualifier_in_node(r, qualifier);
218 }
219 }
220 NodeEnum::BoolExpr(e) => {
221 for a in &mut e.args {
222 strip_qualifier_in_node(a, qualifier);
223 }
224 }
225 NodeEnum::FuncCall(fc) => {
226 for a in &mut fc.args {
227 strip_qualifier_in_node(a, qualifier);
228 }
229 if let Some(filt) = fc.agg_filter.as_mut() {
230 strip_qualifier_in_node(filt, qualifier);
231 }
232 for a in &mut fc.agg_order {
233 strip_qualifier_in_node(a, qualifier);
234 }
235 }
236 NodeEnum::CoalesceExpr(c) => {
237 for a in &mut c.args {
238 strip_qualifier_in_node(a, qualifier);
239 }
240 }
241 NodeEnum::CaseExpr(c) => {
242 if let Some(arg) = c.arg.as_mut() {
243 strip_qualifier_in_node(arg, qualifier);
244 }
245 for w in &mut c.args {
246 strip_qualifier_in_node(w, qualifier);
247 }
248 if let Some(d) = c.defresult.as_mut() {
249 strip_qualifier_in_node(d, qualifier);
250 }
251 }
252 NodeEnum::CaseWhen(w) => {
253 if let Some(e) = w.expr.as_mut() {
254 strip_qualifier_in_node(e, qualifier);
255 }
256 if let Some(r) = w.result.as_mut() {
257 strip_qualifier_in_node(r, qualifier);
258 }
259 }
260 NodeEnum::TypeCast(tc) => {
261 if let Some(arg) = tc.arg.as_mut() {
262 strip_qualifier_in_node(arg, qualifier);
263 }
264 }
265 NodeEnum::SortBy(sb) => {
266 if let Some(arg) = sb.node.as_mut() {
267 strip_qualifier_in_node(arg, qualifier);
268 }
269 }
270 NodeEnum::List(l) => {
271 for item in &mut l.items {
272 strip_qualifier_in_node(item, qualifier);
273 }
274 }
275 NodeEnum::SubLink(sl) => {
276 if let Some(t) = sl.testexpr.as_mut() {
277 strip_qualifier_in_node(t, qualifier);
278 }
279 }
280 NodeEnum::NullTest(nt) => {
281 if let Some(arg) = nt.arg.as_mut() {
282 strip_qualifier_in_node(arg, qualifier);
283 }
284 }
285 NodeEnum::BooleanTest(bt) => {
286 if let Some(arg) = bt.arg.as_mut() {
287 strip_qualifier_in_node(arg, qualifier);
288 }
289 }
290 _ => {}
291 }
292}
293
294fn collapse_whitespace(s: &str) -> String {
295 s.split_whitespace().collect::<Vec<_>>().join(" ")
296}
297
298fn hash_canonical(text: &str) -> [u8; 32] {
299 let mut h = blake3::Hasher::new();
300 h.update(b"pgevolve-normalized-body-v1\n");
301 h.update(text.as_bytes());
302 *h.finalize().as_bytes()
303}
304
305#[cfg(test)]
306mod tests {
307 use super::*;
308
309 #[test]
310 fn strip_qualifiers_equates_pg14_and_source_form() {
311 let pg14 = "SELECT products.category, count(*) AS cnt, avg(products.price) AS avg_price FROM app.products GROUP BY products.category";
312 let src = "SELECT category, count(*) AS cnt, avg(price) AS avg_price FROM app.products GROUP BY category";
313 let a = NormalizedBody::from_sql(pg14).unwrap();
314 let b = NormalizedBody::from_sql(src).unwrap();
315 assert_eq!(
316 a.canonical_text(),
317 b.canonical_text(),
318 "left: {} right: {}",
319 a.canonical_text(),
320 b.canonical_text()
321 );
322 }
323}