sql_splitter/duckdb/
output.rs

1//! Output formatting for query results.
2
3use super::QueryResult;
4use std::io::Write;
5
6/// Output format for query results
7#[derive(Debug, Clone, Copy, PartialEq, Default)]
8pub enum OutputFormat {
9    /// ASCII table format (default)
10    #[default]
11    Table,
12    /// JSON array format
13    Json,
14    /// JSON lines format (one object per line)
15    JsonLines,
16    /// CSV format
17    Csv,
18    /// Tab-separated values
19    Tsv,
20}
21
22impl std::str::FromStr for OutputFormat {
23    type Err = String;
24
25    fn from_str(s: &str) -> Result<Self, Self::Err> {
26        match s.to_lowercase().as_str() {
27            "table" => Ok(OutputFormat::Table),
28            "json" => Ok(OutputFormat::Json),
29            "jsonl" | "jsonlines" | "ndjson" => Ok(OutputFormat::JsonLines),
30            "csv" => Ok(OutputFormat::Csv),
31            "tsv" => Ok(OutputFormat::Tsv),
32            _ => Err(format!(
33                "Unknown format: {}. Valid: table, json, jsonl, csv, tsv",
34                s
35            )),
36        }
37    }
38}
39
40impl std::fmt::Display for OutputFormat {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            OutputFormat::Table => write!(f, "table"),
44            OutputFormat::Json => write!(f, "json"),
45            OutputFormat::JsonLines => write!(f, "jsonl"),
46            OutputFormat::Csv => write!(f, "csv"),
47            OutputFormat::Tsv => write!(f, "tsv"),
48        }
49    }
50}
51
52/// Formatter for query results
53pub struct QueryResultFormatter;
54
55impl QueryResultFormatter {
56    /// Format a query result to a string
57    pub fn format(result: &QueryResult, format: OutputFormat) -> String {
58        match format {
59            OutputFormat::Table => Self::format_table(result),
60            OutputFormat::Json => Self::format_json(result),
61            OutputFormat::JsonLines => Self::format_jsonl(result),
62            OutputFormat::Csv => Self::format_csv(result),
63            OutputFormat::Tsv => Self::format_tsv(result),
64        }
65    }
66
67    /// Write formatted result to a writer
68    pub fn write<W: Write>(
69        result: &QueryResult,
70        format: OutputFormat,
71        writer: &mut W,
72    ) -> std::io::Result<()> {
73        let output = Self::format(result, format);
74        writer.write_all(output.as_bytes())
75    }
76
77    /// Format as ASCII table
78    fn format_table(result: &QueryResult) -> String {
79        if result.columns.is_empty() {
80            return String::new();
81        }
82
83        // Calculate column widths
84        let mut widths: Vec<usize> = result.columns.iter().map(|c| c.len()).collect();
85
86        for row in &result.rows {
87            for (i, val) in row.iter().enumerate() {
88                if i < widths.len() {
89                    widths[i] = widths[i].max(val.len());
90                }
91            }
92        }
93
94        // Cap widths at 50 chars for readability
95        let max_width = 50;
96        widths.iter_mut().for_each(|w| *w = (*w).min(max_width));
97
98        let mut output = String::new();
99
100        // Top border
101        output.push('┌');
102        for (i, width) in widths.iter().enumerate() {
103            output.push_str(&"─".repeat(*width + 2));
104            if i < widths.len() - 1 {
105                output.push('┬');
106            }
107        }
108        output.push_str("┐\n");
109
110        // Header row
111        output.push('│');
112        for (i, col) in result.columns.iter().enumerate() {
113            let truncated = Self::truncate(col, widths[i]);
114            output.push_str(&format!(" {:width$} │", truncated, width = widths[i]));
115        }
116        output.push('\n');
117
118        // Header separator
119        output.push('├');
120        for (i, width) in widths.iter().enumerate() {
121            output.push_str(&"─".repeat(*width + 2));
122            if i < widths.len() - 1 {
123                output.push('┼');
124            }
125        }
126        output.push_str("┤\n");
127
128        // Data rows
129        for row in &result.rows {
130            output.push('│');
131            for (i, val) in row.iter().enumerate() {
132                if i < widths.len() {
133                    let truncated = Self::truncate(val, widths[i]);
134                    output.push_str(&format!(" {:width$} │", truncated, width = widths[i]));
135                }
136            }
137            output.push('\n');
138        }
139
140        // Bottom border
141        output.push('└');
142        for (i, width) in widths.iter().enumerate() {
143            output.push_str(&"─".repeat(*width + 2));
144            if i < widths.len() - 1 {
145                output.push('┴');
146            }
147        }
148        output.push_str("┘\n");
149
150        // Row count
151        output.push_str(&format!(
152            "{} row{}\n",
153            result.rows.len(),
154            if result.rows.len() == 1 { "" } else { "s" }
155        ));
156
157        output
158    }
159
160    /// Truncate a string to a maximum length
161    fn truncate(s: &str, max_len: usize) -> String {
162        if s.len() <= max_len {
163            s.to_string()
164        } else {
165            format!("{}…", &s[..max_len - 1])
166        }
167    }
168
169    /// Format as JSON array
170    fn format_json(result: &QueryResult) -> String {
171        let rows: Vec<serde_json::Value> = result
172            .rows
173            .iter()
174            .map(|row| {
175                let obj: serde_json::Map<String, serde_json::Value> = result
176                    .columns
177                    .iter()
178                    .zip(row.iter())
179                    .map(|(col, val)| (col.clone(), Self::json_value(val)))
180                    .collect();
181                serde_json::Value::Object(obj)
182            })
183            .collect();
184
185        serde_json::to_string_pretty(&rows).unwrap_or_else(|_| "[]".to_string())
186    }
187
188    /// Format as JSON lines (NDJSON)
189    fn format_jsonl(result: &QueryResult) -> String {
190        result
191            .rows
192            .iter()
193            .map(|row| {
194                let obj: serde_json::Map<String, serde_json::Value> = result
195                    .columns
196                    .iter()
197                    .zip(row.iter())
198                    .map(|(col, val)| (col.clone(), Self::json_value(val)))
199                    .collect();
200                serde_json::to_string(&serde_json::Value::Object(obj))
201                    .unwrap_or_else(|_| "{}".to_string())
202            })
203            .collect::<Vec<_>>()
204            .join("\n")
205    }
206
207    /// Convert a string value to appropriate JSON type
208    fn json_value(val: &str) -> serde_json::Value {
209        if val == "NULL" {
210            return serde_json::Value::Null;
211        }
212
213        // Try to parse as number
214        if let Ok(n) = val.parse::<i64>() {
215            return serde_json::Value::Number(n.into());
216        }
217        if let Ok(n) = val.parse::<f64>() {
218            if let Some(num) = serde_json::Number::from_f64(n) {
219                return serde_json::Value::Number(num);
220            }
221        }
222
223        // Try to parse as boolean
224        if val.eq_ignore_ascii_case("true") {
225            return serde_json::Value::Bool(true);
226        }
227        if val.eq_ignore_ascii_case("false") {
228            return serde_json::Value::Bool(false);
229        }
230
231        // Return as string
232        serde_json::Value::String(val.to_string())
233    }
234
235    /// Format as CSV
236    fn format_csv(result: &QueryResult) -> String {
237        let mut output = String::new();
238
239        // Header
240        output.push_str(&Self::csv_row(&result.columns));
241        output.push('\n');
242
243        // Data rows
244        for row in &result.rows {
245            output.push_str(&Self::csv_row(row));
246            output.push('\n');
247        }
248
249        output
250    }
251
252    /// Format a single CSV row
253    fn csv_row(values: &[String]) -> String {
254        values
255            .iter()
256            .map(|v| Self::csv_escape(v))
257            .collect::<Vec<_>>()
258            .join(",")
259    }
260
261    /// Escape a value for CSV
262    fn csv_escape(val: &str) -> String {
263        if val.contains(',') || val.contains('"') || val.contains('\n') || val.contains('\r') {
264            format!("\"{}\"", val.replace('"', "\"\""))
265        } else {
266            val.to_string()
267        }
268    }
269
270    /// Format as TSV
271    fn format_tsv(result: &QueryResult) -> String {
272        let mut output = String::new();
273
274        // Header
275        output.push_str(&result.columns.join("\t"));
276        output.push('\n');
277
278        // Data rows
279        for row in &result.rows {
280            let escaped: Vec<String> = row
281                .iter()
282                .map(|v| v.replace('\t', "\\t").replace('\n', "\\n"))
283                .collect();
284            output.push_str(&escaped.join("\t"));
285            output.push('\n');
286        }
287
288        output
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    fn sample_result() -> QueryResult {
297        QueryResult {
298            columns: vec!["id".to_string(), "name".to_string(), "age".to_string()],
299            column_types: vec![
300                "INTEGER".to_string(),
301                "VARCHAR".to_string(),
302                "INTEGER".to_string(),
303            ],
304            rows: vec![
305                vec!["1".to_string(), "Alice".to_string(), "30".to_string()],
306                vec!["2".to_string(), "Bob".to_string(), "25".to_string()],
307            ],
308            execution_time_secs: 0.001,
309        }
310    }
311
312    #[test]
313    fn test_format_table() {
314        let result = sample_result();
315        let output = QueryResultFormatter::format(&result, OutputFormat::Table);
316        assert!(output.contains("Alice"));
317        assert!(output.contains("Bob"));
318        assert!(output.contains("2 rows"));
319    }
320
321    #[test]
322    fn test_format_json() {
323        let result = sample_result();
324        let output = QueryResultFormatter::format(&result, OutputFormat::Json);
325        let parsed: Vec<serde_json::Value> = serde_json::from_str(&output).unwrap();
326        assert_eq!(parsed.len(), 2);
327        assert_eq!(parsed[0]["name"], "Alice");
328        assert_eq!(parsed[0]["age"], 30);
329    }
330
331    #[test]
332    fn test_format_csv() {
333        let result = sample_result();
334        let output = QueryResultFormatter::format(&result, OutputFormat::Csv);
335        assert!(output.starts_with("id,name,age\n"));
336        assert!(output.contains("1,Alice,30"));
337    }
338
339    #[test]
340    fn test_csv_escape() {
341        assert_eq!(QueryResultFormatter::csv_escape("hello"), "hello");
342        assert_eq!(
343            QueryResultFormatter::csv_escape("hello,world"),
344            "\"hello,world\""
345        );
346        assert_eq!(
347            QueryResultFormatter::csv_escape("say \"hi\""),
348            "\"say \"\"hi\"\"\""
349        );
350    }
351
352    #[test]
353    fn test_format_tsv() {
354        let result = sample_result();
355        let output = QueryResultFormatter::format(&result, OutputFormat::Tsv);
356        assert!(output.starts_with("id\tname\tage\n"));
357        assert!(output.contains("1\tAlice\t30"));
358    }
359
360    #[test]
361    fn test_json_value_conversion() {
362        assert_eq!(
363            QueryResultFormatter::json_value("NULL"),
364            serde_json::Value::Null
365        );
366        assert_eq!(
367            QueryResultFormatter::json_value("42"),
368            serde_json::json!(42)
369        );
370        assert_eq!(
371            QueryResultFormatter::json_value("3.14"),
372            serde_json::json!(3.14)
373        );
374        assert_eq!(
375            QueryResultFormatter::json_value("true"),
376            serde_json::json!(true)
377        );
378        assert_eq!(
379            QueryResultFormatter::json_value("hello"),
380            serde_json::json!("hello")
381        );
382    }
383}
sql_splitter/duckdb/output.rs

sql_splitter/duckdb/
output.rs
