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 { attr: TruncationAttr::Cols, node, depth, length: cols_len(s.cols.clone())? });
114                    }
115                }
116                NodeMut::IndexStmt(s) => {
117                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
118                    if let Some(clause) = s.where_clause.as_ref() {
119                        truncations.push_back(PossibleTruncation {
120                            attr: TruncationAttr::WhereClause,
121                            node,
122                            depth,
123                            length: where_clause_len((*clause).clone())?,
124                        });
125                    }
126                }
127                NodeMut::RuleStmt(s) => {
128                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
129                    if let Some(clause) = s.where_clause.as_ref() {
130                        truncations.push_back(PossibleTruncation {
131                            attr: TruncationAttr::WhereClause,
132                            node,
133                            depth,
134                            length: where_clause_len((*clause).clone())?,
135                        });
136                    }
137                }
138                NodeMut::CommonTableExpr(s) => {
139                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
140                    if let Some(cte) = s.ctequery.as_ref() {
141                        truncations.push_back(PossibleTruncation {
142                            attr: TruncationAttr::CTEQuery,
143                            node,
144                            depth: depth + 1,
145                            length: cte.deparse()?.len() as i32,
146                        });
147                    }
148                }
149                NodeMut::InferClause(s) => {
150                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
151                    if let Some(clause) = s.where_clause.as_ref() {
152                        truncations.push_back(PossibleTruncation {
153                            attr: TruncationAttr::WhereClause,
154                            node,
155                            depth,
156                            length: where_clause_len((*clause).clone())?,
157                        });
158                    }
159                }
160                NodeMut::OnConflictClause(s) => {
161                    let s = s.as_mut().ok_or(Error::InvalidPointer)?;
162                    if !s.target_list.is_empty() {
163                        truncations.push_back(PossibleTruncation {
164                            attr: TruncationAttr::TargetList,
165                            node,
166                            depth,
167                            length: update_target_list_len(s.target_list.clone())?,
168                        });
169                    }
170                    if let Some(clause) = s.where_clause.as_ref() {
171                        truncations.push_back(PossibleTruncation {
172                            attr: TruncationAttr::WhereClause,
173                            node,
174                            depth,
175                            length: where_clause_len((*clause).clone())?,
176                        });
177                    }
178                }
179                _ => (),
180            }
181        }
182
183        truncations.make_contiguous().sort_by(|a, b| match a.depth.cmp(&b.depth).reverse() {
184            Ordering::Equal => a.length.cmp(&b.length).reverse(),
185            other => other,
186        });
187
188        while let Some(truncation) = truncations.pop_front() {
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_clause: None,
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_clause: None,
384            with_clause: None,
385        }))),
386    })
387}
pg_query/truncate.rs

pg_query/
truncate.rs
