pg_query/
truncate.rs

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

pg_query/
truncate.rs
