Skip to main content

lean_ctx/core/patterns/
psql.rs

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