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 = row_count_line
61        .map(|l| l.trim().to_string())
62        .unwrap_or_else(|| format!("({data_rows} rows)"));
63
64    if data_rows <= 20 {
65        return output.to_string();
66    }
67
68    let preview_end = (separator_idx + 11).min(lines.len());
69    let preview: Vec<&str> = lines[..preview_end].to_vec();
70    format!("{}\n... {count_str}", preview.join("\n"))
71}
72
73fn compress_describe(output: &str) -> String {
74    let lines: Vec<&str> = output.lines().filter(|l| !l.trim().is_empty()).collect();
75    if lines.len() <= 30 {
76        return lines.join("\n");
77    }
78    format!(
79        "{}\n... ({} more lines)",
80        lines[..20].join("\n"),
81        lines.len() - 20
82    )
83}
84
85fn compact_lines(text: &str, max: usize) -> String {
86    let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
87    if lines.len() <= max {
88        return lines.join("\n");
89    }
90    format!(
91        "{}\n... ({} more lines)",
92        lines[..max].join("\n"),
93        lines.len() - max
94    )
95}