Skip to main content

lean_ctx_patterns/
psql.rs

1// SPDX-License-Identifier: Apache-2.0
2
3pub fn compress(cmd: &str, output: &str) -> Option<String> {
4    let trimmed = output.trim();
5    if trimmed.is_empty() {
6        return Some("ok".to_string());
7    }
8
9    if is_table_output(trimmed) {
10        return Some(compress_table(trimmed));
11    }
12
13    if cmd.contains("\\dt") || cmd.contains("\\d") {
14        return Some(compress_describe(trimmed));
15    }
16
17    if trimmed.starts_with("INSERT")
18        || trimmed.starts_with("UPDATE")
19        || trimmed.starts_with("DELETE")
20        || trimmed.starts_with("CREATE")
21        || trimmed.starts_with("ALTER")
22        || trimmed.starts_with("DROP")
23    {
24        return Some(trimmed.lines().next().unwrap_or(trimmed).to_string());
25    }
26
27    Some(compact_lines(trimmed, 20))
28}
29
30fn is_table_output(output: &str) -> bool {
31    let lines: Vec<&str> = output.lines().collect();
32    lines.len() >= 3
33        && lines
34            .iter()
35            .any(|l| l.contains("---+---") || l.contains("-+-"))
36}
37
38fn compress_table(output: &str) -> String {
39    let lines: Vec<&str> = output.lines().collect();
40    let mut separator_idx = 0;
41    let mut data_rows = 0u32;
42
43    for (i, line) in lines.iter().enumerate() {
44        if line.contains("---+---") || line.contains("-+-") {
45            separator_idx = i;
46            break;
47        }
48    }
49
50    for line in lines.iter().skip(separator_idx + 1) {
51        let trimmed = line.trim();
52        if trimmed.is_empty() || trimmed.starts_with('(') {
53            continue;
54        }
55        data_rows += 1;
56    }
57
58    let row_count_line = lines
59        .iter()
60        .rev()
61        .find(|l| l.trim().starts_with('(') && l.contains("row"));
62    let count_str =
63        row_count_line.map_or_else(|| format!("({data_rows} rows)"), |l| l.trim().to_string());
64
65    if data_rows <= 20 {
66        return output.to_string();
67    }
68
69    let preview_end = (separator_idx + 11).min(lines.len());
70    let preview: Vec<&str> = lines[..preview_end].to_vec();
71    format!("{}\n... {count_str}", preview.join("\n"))
72}
73
74fn compress_describe(output: &str) -> String {
75    let lines: Vec<&str> = output.lines().filter(|l| !l.trim().is_empty()).collect();
76    if lines.len() <= 30 {
77        return lines.join("\n");
78    }
79    format!(
80        "{}\n{}",
81        lines[..20].join("\n"),
82        super::elision_marker(lines.len() - 20)
83    )
84}
85
86fn compact_lines(text: &str, max: usize) -> String {
87    let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
88    if lines.len() <= max {
89        return lines.join("\n");
90    }
91    format!(
92        "{}\n{}",
93        lines[..max].join("\n"),
94        super::elision_marker(lines.len() - max)
95    )
96}