ddb_core/output/
formatter.rs

1// Output formatting for query results
2use crate::engine::Row;
3use crate::error::Result;
4use prettytable::{Cell, Row as PrettyRow, Table};
5
6#[derive(Debug, Clone, Copy)]
7pub enum OutputFormat {
8    Table,
9    Json,
10    Yaml,
11    Csv,
12}
13
14impl OutputFormat {
15    pub fn from_str(s: &str) -> Self {
16        match s.to_lowercase().as_str() {
17            "json" => OutputFormat::Json,
18            "yaml" => OutputFormat::Yaml,
19            "csv" => OutputFormat::Csv,
20            _ => OutputFormat::Table,
21        }
22    }
23}
24
25/// Format query results according to the specified output format
26pub fn format_results(rows: &[Row], format: OutputFormat) -> Result<String> {
27    if rows.is_empty() {
28        return Ok("No results".to_string());
29    }
30
31    match format {
32        OutputFormat::Table => format_table(rows),
33        OutputFormat::Json => format_json(rows),
34        OutputFormat::Yaml => format_yaml(rows),
35        OutputFormat::Csv => format_csv(rows),
36    }
37}
38
39/// Format results as a pretty ASCII table
40fn format_table(rows: &[Row]) -> Result<String> {
41    if rows.is_empty() {
42        return Ok("No results".to_string());
43    }
44
45    let mut table = Table::new();
46
47    // Get column names from first row
48    let columns: Vec<String> = rows[0].column_names();
49
50    // Add header row
51    let header: Vec<Cell> = columns.iter().map(|c| Cell::new(c)).collect();
52    table.add_row(PrettyRow::new(header));
53
54    // Add data rows
55    for row in rows {
56        let values = row.values(&columns);
57        let cells: Vec<Cell> = values.iter().map(|v| Cell::new(&v.to_string())).collect();
58        table.add_row(PrettyRow::new(cells));
59    }
60
61    Ok(table.to_string())
62}
63
64/// Format results as JSON
65fn format_json(rows: &[Row]) -> Result<String> {
66    let json_rows: Vec<serde_json::Value> = rows
67        .iter()
68        .map(|row| {
69            let mut map = serde_json::Map::new();
70            for (key, value) in row.as_map() {
71                map.insert(key.clone(), value_to_json(value));
72            }
73            serde_json::Value::Object(map)
74        })
75        .collect();
76
77    serde_json::to_string_pretty(&json_rows)
78        .map_err(|e| crate::error::DdbError::SerializationError(e.to_string()))
79}
80
81/// Format results as YAML
82fn format_yaml(rows: &[Row]) -> Result<String> {
83    let yaml_rows: Vec<serde_yaml::Value> = rows
84        .iter()
85        .map(|row| {
86            let mut map = serde_yaml::Mapping::new();
87            for (key, value) in row.as_map() {
88                map.insert(
89                    serde_yaml::Value::String(key.clone()),
90                    value_to_yaml(value),
91                );
92            }
93            serde_yaml::Value::Mapping(map)
94        })
95        .collect();
96
97    serde_yaml::to_string(&yaml_rows)
98        .map_err(|e| crate::error::DdbError::SerializationError(e.to_string()))
99}
100
101/// Format results as CSV
102fn format_csv(rows: &[Row]) -> Result<String> {
103    if rows.is_empty() {
104        return Ok(String::new());
105    }
106
107    let mut output = String::new();
108    let columns: Vec<String> = rows[0].column_names();
109
110    // Header
111    output.push_str(&columns.join(","));
112    output.push('\n');
113
114    // Data rows
115    for row in rows {
116        let values: Vec<String> = row
117            .values(&columns)
118            .iter()
119            .map(|v| escape_csv(&v.to_string()))
120            .collect();
121        output.push_str(&values.join(","));
122        output.push('\n');
123    }
124
125    Ok(output)
126}
127
128/// Escape CSV field if needed
129fn escape_csv(s: &str) -> String {
130    if s.contains(',') || s.contains('"') || s.contains('\n') {
131        format!("\"{}\"", s.replace('"', "\"\""))
132    } else {
133        s.to_string()
134    }
135}
136
137/// Convert Value to JSON value
138fn value_to_json(value: &crate::functions::Value) -> serde_json::Value {
139    use crate::functions::Value;
140
141    match value {
142        Value::Null => serde_json::Value::Null,
143        Value::Boolean(b) => serde_json::Value::Bool(*b),
144        Value::Integer(i) => serde_json::Value::Number((*i).into()),
145        Value::Float(f) => serde_json::Number::from_f64(*f)
146            .map(serde_json::Value::Number)
147            .unwrap_or(serde_json::Value::Null),
148        Value::String(s) => serde_json::Value::String(s.clone()),
149        Value::Date(d) => serde_json::Value::String(d.format("%Y-%m-%d").to_string()),
150        Value::DateTime(dt) => serde_json::Value::String(dt.format("%Y-%m-%d %H:%M:%S").to_string()),
151        Value::Time(t) => serde_json::Value::String(t.format("%H:%M:%S").to_string()),
152    }
153}
154
155/// Convert Value to YAML value
156fn value_to_yaml(value: &crate::functions::Value) -> serde_yaml::Value {
157    use crate::functions::Value;
158
159    match value {
160        Value::Null => serde_yaml::Value::Null,
161        Value::Boolean(b) => serde_yaml::Value::Bool(*b),
162        Value::Integer(i) => serde_yaml::Value::Number((*i).into()),
163        Value::Float(f) => serde_yaml::Value::Number(serde_yaml::Number::from(*f)),
164        Value::String(s) => serde_yaml::Value::String(s.clone()),
165        Value::Date(d) => serde_yaml::Value::String(d.format("%Y-%m-%d").to_string()),
166        Value::DateTime(dt) => serde_yaml::Value::String(dt.format("%Y-%m-%d %H:%M:%S").to_string()),
167        Value::Time(t) => serde_yaml::Value::String(t.format("%H:%M:%S").to_string()),
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use crate::functions::Value;
175
176    fn create_test_rows() -> Vec<Row> {
177        let columns = vec!["id".to_string(), "name".to_string(), "age".to_string()];
178        let values1 = vec![
179            Value::Integer(1),
180            Value::String("Alice".to_string()),
181            Value::Integer(30),
182        ];
183        let values2 = vec![
184            Value::Integer(2),
185            Value::String("Bob".to_string()),
186            Value::Integer(25),
187        ];
188
189        vec![
190            Row::from_values(&columns, values1),
191            Row::from_values(&columns, values2),
192        ]
193    }
194
195    #[test]
196    fn test_format_csv() {
197        let rows = create_test_rows();
198        let output = format_csv(&rows).unwrap();
199
200        // Check that all columns are present (order may vary due to HashMap)
201        assert!(output.contains("id"));
202        assert!(output.contains("name"));
203        assert!(output.contains("age"));
204        // Check data is present
205        assert!(output.contains("Alice"));
206        assert!(output.contains("Bob"));
207        assert!(output.contains("30"));
208        assert!(output.contains("25"));
209    }
210
211    #[test]
212    fn test_format_json() {
213        let rows = create_test_rows();
214        let output = format_json(&rows).unwrap();
215
216        assert!(output.contains("\"id\""));
217        assert!(output.contains("\"name\""));
218        assert!(output.contains("Alice"));
219        assert!(output.contains("Bob"));
220    }
221
222    #[test]
223    fn test_format_yaml() {
224        let rows = create_test_rows();
225        let output = format_yaml(&rows).unwrap();
226
227        assert!(output.contains("id:"));
228        assert!(output.contains("name:"));
229        assert!(output.contains("Alice"));
230        assert!(output.contains("Bob"));
231    }
232}
ddb_core/output/formatter.rs

ddb_core/output/
formatter.rs
