pg_query/
truncate.rs

1use std::cmp::Ordering;
2use std::collections::VecDeque;
3
4use crate::*;
5
6#[derive(Debug)]
7enum TruncationAttr {
8    TargetList,
9    WhereClause,
10    ValuesLists,
11    CTEQuery,
12    Cols,
13}
14
15#[derive(Debug)]
16struct PossibleTruncation {
17    attr: TruncationAttr,
18    node: NodeMut,
19    depth: i32,
20    length: i32,
21}
22
23pub fn truncate(protobuf: &protobuf::ParseResult, max_length: usize) -> Result<String> {
24    let mut output = protobuf.deparse()?;
25    if output.len() <= max_length {
26        return Ok(output);
27    }
28
29    // SAFETY: within this scope nobody expects to have exclusive access to `protobuf`'s contents, so we can have multiple shared accesses.
30    //
31    // Raw pointer documentation:
32    //
33    // https://doc.rust-lang.org/book/ch19-01-unsafe-rust.html#dereferencing-a-raw-pointer
34    // https://doc.rust-lang.org/std/primitive.pointer.html
35    // https://manishearth.github.io/blog/2015/05/17/the-problem-with-shared-mutability
36    // https://ricardomartins.cc/2016/07/11/interior-mutability-behind-the-curtain
37    unsafe {
38        let mut protobuf = protobuf.clone();
39        let mut truncations: VecDeque<PossibleTruncation> = VecDeque::new();
40        for (node, depth, _context) in protobuf.nodes_mut().into_iter() {
41            match node {
42                NodeMut::SelectStmt(s) => {
43                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
44                    if !s.target_list.is_empty() {
45                        truncations.push_back(PossibleTruncation {
46                            attr: TruncationAttr::TargetList,
47                            node,
48                            depth,
49                            length: select_target_list_len(s.target_list.clone())?,
50                        });
51                    }
52                    if let Some(clause) = s.where_clause.as_ref() {
53                        truncations.push_back(PossibleTruncation {
54                            attr: TruncationAttr::WhereClause,
55                            node,
56                            depth,
57                            length: where_clause_len((*clause).clone())?,
58                        });
59                    }
60                    if !s.values_lists.is_empty() {
61                        truncations.push_back(PossibleTruncation {
62                            attr: TruncationAttr::ValuesLists,
63                            node,
64                            depth,
65                            length: select_values_lists_len(s.values_lists.clone())?,
66                        });
67                    }
68                }
69                NodeMut::UpdateStmt(s) => {
70                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
71                    if !s.target_list.is_empty() {
72                        truncations.push_back(PossibleTruncation {
73                            attr: TruncationAttr::TargetList,
74                            node,
75                            depth,
76                            length: update_target_list_len(s.target_list.clone())?,
77                        });
78                    }
79                    if let Some(clause) = s.where_clause.as_ref() {
80                        truncations.push_back(PossibleTruncation {
81                            attr: TruncationAttr::WhereClause,
82                            node,
83                            depth,
84                            length: where_clause_len((*clause).clone())?,
85                        });
86                    }
87                }
88                NodeMut::DeleteStmt(s) => {
89                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
90                    if let Some(clause) = s.where_clause.as_ref() {
91                        truncations.push_back(PossibleTruncation {
92                            attr: TruncationAttr::WhereClause,
93                            node,
94                            depth,
95                            length: where_clause_len((*clause).clone())?,
96                        });
97                    }
98                }
99                NodeMut::CopyStmt(s) => {
100                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
101                    if let Some(clause) = s.where_clause.as_ref() {
102                        truncations.push_back(PossibleTruncation {
103                            attr: TruncationAttr::WhereClause,
104                            node,
105                            depth,
106                            length: where_clause_len((*clause).clone())?,
107                        });
108                    }
109                }
110                NodeMut::InsertStmt(s) => {
111                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
112                    if !s.cols.is_empty() {
113                        truncations.push_back(PossibleTruncation {
114                            attr: TruncationAttr::Cols,
115                            node,
116                            depth,
117                            length: cols_len(s.cols.clone())?,
118                        });
119                    }
120                }
121                NodeMut::IndexStmt(s) => {
122                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
123                    if let Some(clause) = s.where_clause.as_ref() {
124                        truncations.push_back(PossibleTruncation {
125                            attr: TruncationAttr::WhereClause,
126                            node,
127                            depth,
128                            length: where_clause_len((*clause).clone())?,
129                        });
130                    }
131                }
132                NodeMut::RuleStmt(s) => {
133                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
134                    if let Some(clause) = s.where_clause.as_ref() {
135                        truncations.push_back(PossibleTruncation {
136                            attr: TruncationAttr::WhereClause,
137                            node,
138                            depth,
139                            length: where_clause_len((*clause).clone())?,
140                        });
141                    }
142                }
143                NodeMut::CommonTableExpr(s) => {
144                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
145                    if let Some(cte) = s.ctequery.as_ref() {
146                        truncations.push_back(PossibleTruncation {
147                            attr: TruncationAttr::CTEQuery,
148                            node,
149                            depth: depth + 1,
150                            length: cte.deparse()?.len() as i32,
151                        });
152                    }
153                }
154                NodeMut::InferClause(s) => {
155                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
156                    if let Some(clause) = s.where_clause.as_ref() {
157                        truncations.push_back(PossibleTruncation {
158                            attr: TruncationAttr::WhereClause,
159                            node,
160                            depth,
161                            length: where_clause_len((*clause).clone())?,
162                        });
163                    }
164                }
165                NodeMut::OnConflictClause(s) => {
166                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
167                    if !s.target_list.is_empty() {
168                        truncations.push_back(PossibleTruncation {
169                            attr: TruncationAttr::TargetList,
170                            node,
171                            depth,
172                            length: update_target_list_len(s.target_list.clone())?,
173                        });
174                    }
175                    if let Some(clause) = s.where_clause.as_ref() {
176                        truncations.push_back(PossibleTruncation {
177                            attr: TruncationAttr::WhereClause,
178                            node,
179                            depth,
180                            length: where_clause_len((*clause).clone())?,
181                        });
182                    }
183                }
184                _ => (),
185            }
186        }
187
188        truncations
189            .make_contiguous()
190            .sort_by(|a, b| match a.depth.cmp(&b.depth).reverse() {
191                Ordering::Equal => a.length.cmp(&b.length).reverse(),
192                other => other,
193            });
194
195        while let Some(truncation) = truncations.pop_front() {
196            match (truncation.node, truncation.attr) {
197                (NodeMut::SelectStmt(s), TruncationAttr::TargetList) => {
198                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
199                    s.target_list = vec![dummy_target()];
200                }
201                (NodeMut::SelectStmt(s), TruncationAttr::WhereClause) => {
202                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
203                    s.where_clause = Some(dummy_column());
204                }
205                (NodeMut::SelectStmt(s), TruncationAttr::ValuesLists) => {
206                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
207                    s.values_lists = vec![Node {
208                        node: Some(NodeEnum::List(protobuf::List {
209                            items: vec![*dummy_column()],
210                        })),
211                    }]
212                }
213                (NodeMut::UpdateStmt(s), TruncationAttr::TargetList) => {
214                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
215                    s.target_list = vec![dummy_target()];
216                }
217                (NodeMut::UpdateStmt(s), TruncationAttr::WhereClause) => {
218                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
219                    s.where_clause = Some(dummy_column());
220                }
221                (NodeMut::DeleteStmt(s), TruncationAttr::WhereClause) => {
222                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
223                    s.where_clause = Some(dummy_column());
224                }
225                (NodeMut::CopyStmt(s), TruncationAttr::WhereClause) => {
226                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
227                    s.where_clause = Some(dummy_column());
228                }
229                (NodeMut::InsertStmt(s), TruncationAttr::Cols) => {
230                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
231                    s.cols = vec![dummy_target()];
232                }
233                (NodeMut::IndexStmt(s), TruncationAttr::WhereClause) => {
234                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
235                    s.where_clause = Some(dummy_column());
236                }
237                (NodeMut::RuleStmt(s), TruncationAttr::WhereClause) => {
238                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
239                    s.where_clause = Some(dummy_column());
240                }
241                (NodeMut::CommonTableExpr(s), TruncationAttr::CTEQuery) => {
242                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
243                    let old = std::mem::replace(
244                        &mut s.ctequery,
245                        Some(dummy_select(vec![], Some(dummy_column()), vec![])),
246                    );
247                    if let Some(s) = old {
248                        let node = s.node.ok_or(Error::InvalidPointer)?;
249                        truncations.retain(|t| t.node.to_enum().unwrap() != node);
250                    }
251                }
252                (NodeMut::InferClause(s), TruncationAttr::WhereClause) => {
253                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
254                    s.where_clause = Some(dummy_column());
255                }
256                (NodeMut::OnConflictClause(s), TruncationAttr::TargetList) => {
257                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
258                    s.target_list = vec![dummy_target()];
259                }
260                (NodeMut::OnConflictClause(s), TruncationAttr::WhereClause) => {
261                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
262                    s.where_clause = Some(dummy_column());
263                }
264                _ => panic!("unimplemented truncation"),
265            }
266            output = protobuf.deparse()?;
267            output = output
268                .replace("SELECT WHERE \"…\"", "...")
269                .replace("\"…\"", "...");
270            // the unwanted AS doesn't happen in the Ruby version. I'm not sure where it's coming from
271            output = output.replace("SELECT ... AS ...", "SELECT ...");
272            if output.len() <= max_length {
273                return Ok(output);
274            }
275        }
276    }
277
278    // We couldn't do a proper smart truncation, so we need a hard cut-off
279    Ok(format!("{}...", truncate_str(&output, max_length - 3)))
280}
281
282// Truncates at character boundaries to prevent panics.
283fn truncate_str(string: &str, max_chars: usize) -> &str {
284    match string.char_indices().nth(max_chars) {
285        None => string,
286        Some((idx, _)) => &string[..idx],
287    }
288}
289
290fn select_target_list_len(nodes: Vec<Node>) -> Result<i32> {
291    let fragment = dummy_select(nodes, None, vec![]).deparse()?;
292    Ok(fragment.len() as i32 - 7) // "SELECT "
293}
294
295fn select_values_lists_len(nodes: Vec<Node>) -> Result<i32> {
296    let fragment = dummy_select(vec![], None, nodes).deparse()?;
297    Ok(fragment.len() as i32 - 7) // "SELECT "
298}
299
300fn update_target_list_len(nodes: Vec<Node>) -> Result<i32> {
301    let fragment = dummy_update(nodes).deparse()?;
302    Ok(fragment.len() as i32 - 13) // "UPDATE x SET "
303}
304
305fn where_clause_len(node: Box<Node>) -> Result<i32> {
306    let fragment = dummy_select(vec![], Some(node), vec![]).deparse()?;
307    Ok(fragment.len() as i32 - 13) // "SELECT WHERE "
308}
309
310fn cols_len(nodes: Vec<Node>) -> Result<i32> {
311    let fragment = dummy_insert(nodes).deparse()?;
312    Ok(fragment.len() as i32 - 31) // "INSERT INTO x () DEFAULT VALUES"
313}
314
315fn dummy_column() -> Box<Node> {
316    Box::new(Node {
317        node: Some(NodeEnum::ColumnRef(protobuf::ColumnRef {
318            location: 0,
319            fields: vec![Node {
320                node: Some(NodeEnum::String(protobuf::String {
321                    sval: "…".to_string(),
322                })),
323            }],
324        })),
325    })
326}
327
328fn dummy_target() -> Node {
329    Node {
330        node: Some(NodeEnum::ResTarget(Box::new(protobuf::ResTarget {
331            name: "…".to_string(),
332            location: 0,
333            indirection: vec![],
334            val: Some(dummy_column()),
335        }))),
336    }
337}
338
339fn dummy_select(
340    target_list: Vec<Node>,
341    where_clause: Option<Box<Node>>,
342    values_lists: Vec<Node>,
343) -> Box<Node> {
344    Box::new(Node {
345        node: Some(NodeEnum::SelectStmt(Box::new(protobuf::SelectStmt {
346            distinct_clause: vec![],
347            into_clause: None,
348            target_list,
349            from_clause: vec![],
350            where_clause,
351            group_clause: vec![],
352            having_clause: None,
353            window_clause: vec![],
354            values_lists,
355            sort_clause: vec![],
356            limit_offset: None,
357            limit_count: None,
358            limit_option: 1,
359            locking_clause: vec![],
360            with_clause: None,
361            op: 1,
362            all: false,
363            larg: None,
364            rarg: None,
365            group_distinct: false,
366        }))),
367    })
368}
369
370fn dummy_insert(cols: Vec<Node>) -> Box<Node> {
371    Box::new(Node {
372        node: Some(NodeEnum::InsertStmt(Box::new(protobuf::InsertStmt {
373            relation: Some(protobuf::RangeVar {
374                catalogname: "".to_string(),
375                schemaname: "".to_string(),
376                relname: "x".to_string(),
377                inh: true,
378                relpersistence: "p".to_string(),
379                alias: None,
380                location: 0,
381            }),
382            cols,
383            select_stmt: None,
384            on_conflict_clause: None,
385            returning_clause: None,
386            with_clause: None,
387            r#override: 1,
388        }))),
389    })
390}
391
392fn dummy_update(target_list: Vec<Node>) -> Box<Node> {
393    Box::new(Node {
394        node: Some(NodeEnum::UpdateStmt(Box::new(protobuf::UpdateStmt {
395            relation: Some(protobuf::RangeVar {
396                catalogname: "".to_string(),
397                schemaname: "".to_string(),
398                relname: "x".to_string(),
399                inh: true,
400                relpersistence: "p".to_string(),
401                alias: None,
402                location: 0,
403            }),
404            from_clause: vec![],
405            target_list,
406            where_clause: None,
407            returning_clause: None,
408            with_clause: None,
409        }))),
410    })
411}
pg_query/truncate.rs

pg_query/
truncate.rs
