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    Xml,
13}
14
15impl OutputFormat {
16    pub fn from_str(s: &str) -> Self {
17        match s.to_lowercase().as_str() {
18            "json" => OutputFormat::Json,
19            "yaml" => OutputFormat::Yaml,
20            "csv" => OutputFormat::Csv,
21            "xml" => OutputFormat::Xml,
22            _ => OutputFormat::Table,
23        }
24    }
25}
26
27/// Format query results according to the specified output format
28pub fn format_results(rows: &[Row], format: OutputFormat) -> Result<String> {
29    if rows.is_empty() {
30        return Ok("No results".to_string());
31    }
32
33    match format {
34        OutputFormat::Table => format_table(rows),
35        OutputFormat::Json => format_json(rows),
36        OutputFormat::Yaml => format_yaml(rows),
37        OutputFormat::Csv => format_csv(rows),
38        OutputFormat::Xml => format_xml(rows),
39    }
40}
41
42/// Format results as a pretty ASCII table
43fn format_table(rows: &[Row]) -> Result<String> {
44    if rows.is_empty() {
45        return Ok("No results".to_string());
46    }
47
48    let mut table = Table::new();
49
50    // Get column names from first row
51    let columns: Vec<String> = rows[0].column_names();
52
53    // Add header row
54    let header: Vec<Cell> = columns.iter().map(|c| Cell::new(c)).collect();
55    table.add_row(PrettyRow::new(header));
56
57    // Add data rows
58    for row in rows {
59        let values = row.values(&columns);
60        let cells: Vec<Cell> = values.iter().map(|v| Cell::new(&v.to_string())).collect();
61        table.add_row(PrettyRow::new(cells));
62    }
63
64    Ok(table.to_string())
65}
66
67/// Format results as JSON
68fn format_json(rows: &[Row]) -> Result<String> {
69    let json_rows: Vec<serde_json::Value> = rows
70        .iter()
71        .map(|row| {
72            let mut map = serde_json::Map::new();
73            for (key, value) in row.as_map() {
74                map.insert(key.clone(), value_to_json(value));
75            }
76            serde_json::Value::Object(map)
77        })
78        .collect();
79
80    serde_json::to_string_pretty(&json_rows)
81        .map_err(|e| crate::error::DdbError::SerializationError(e.to_string()))
82}
83
84/// Format results as YAML
85fn format_yaml(rows: &[Row]) -> Result<String> {
86    let yaml_rows: Vec<serde_yaml::Value> = rows
87        .iter()
88        .map(|row| {
89            let mut map = serde_yaml::Mapping::new();
90            for (key, value) in row.as_map() {
91                map.insert(
92                    serde_yaml::Value::String(key.clone()),
93                    value_to_yaml(value),
94                );
95            }
96            serde_yaml::Value::Mapping(map)
97        })
98        .collect();
99
100    serde_yaml::to_string(&yaml_rows)
101        .map_err(|e| crate::error::DdbError::SerializationError(e.to_string()))
102}
103
104/// Format results as CSV
105fn format_csv(rows: &[Row]) -> Result<String> {
106    if rows.is_empty() {
107        return Ok(String::new());
108    }
109
110    let mut output = String::new();
111    let columns: Vec<String> = rows[0].column_names();
112
113    // Header
114    output.push_str(&columns.join(","));
115    output.push('\n');
116
117    // Data rows
118    for row in rows {
119        let values: Vec<String> = row
120            .values(&columns)
121            .iter()
122            .map(|v| escape_csv(&v.to_string()))
123            .collect();
124        output.push_str(&values.join(","));
125        output.push('\n');
126    }
127
128    Ok(output)
129}
130
131/// Format results as XML
132fn format_xml(rows: &[Row]) -> Result<String> {
133    if rows.is_empty() {
134        return Ok("<results></results>".to_string());
135    }
136
137    let mut output = String::from("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<results>\n");
138
139    for row in rows {
140        output.push_str("  <row>\n");
141        for (key, value) in row.as_map() {
142            // Escape XML special characters in tag name
143            let safe_key = escape_xml_tag(key);
144            let value_str = escape_xml(&value_to_string(value));
145            output.push_str(&format!("    <{}>{}</{}>\n", safe_key, value_str, safe_key));
146        }
147        output.push_str("  </row>\n");
148    }
149
150    output.push_str("</results>");
151    Ok(output)
152}
153
154/// Escape XML tag name (replace invalid characters with underscores)
155fn escape_xml_tag(s: &str) -> String {
156    s.chars()
157        .map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' })
158        .collect()
159}
160
161/// Escape XML special characters in content
162fn escape_xml(s: &str) -> String {
163    s.replace('&', "&amp;")
164        .replace('<', "&lt;")
165        .replace('>', "&gt;")
166        .replace('"', "&quot;")
167        .replace('\'', "&apos;")
168}
169
170/// Convert Value to string for XML
171fn value_to_string(value: &crate::functions::Value) -> String {
172    use crate::functions::Value;
173
174    match value {
175        Value::Null => String::new(),
176        Value::Boolean(b) => b.to_string(),
177        Value::Integer(i) => i.to_string(),
178        Value::Float(f) => f.to_string(),
179        Value::String(s) => s.clone(),
180        Value::Date(d) => d.format("%Y-%m-%d").to_string(),
181        Value::DateTime(dt) => dt.format("%Y-%m-%d %H:%M:%S").to_string(),
182        Value::Time(t) => t.format("%H:%M:%S").to_string(),
183    }
184}
185
186/// Escape CSV field if needed
187fn escape_csv(s: &str) -> String {
188    if s.contains(',') || s.contains('"') || s.contains('\n') {
189        format!("\"{}\"", s.replace('"', "\"\""))
190    } else {
191        s.to_string()
192    }
193}
194
195/// Convert Value to JSON value
196fn value_to_json(value: &crate::functions::Value) -> serde_json::Value {
197    use crate::functions::Value;
198
199    match value {
200        Value::Null => serde_json::Value::Null,
201        Value::Boolean(b) => serde_json::Value::Bool(*b),
202        Value::Integer(i) => serde_json::Value::Number((*i).into()),
203        Value::Float(f) => serde_json::Number::from_f64(*f)
204            .map(serde_json::Value::Number)
205            .unwrap_or(serde_json::Value::Null),
206        Value::String(s) => serde_json::Value::String(s.clone()),
207        Value::Date(d) => serde_json::Value::String(d.format("%Y-%m-%d").to_string()),
208        Value::DateTime(dt) => serde_json::Value::String(dt.format("%Y-%m-%d %H:%M:%S").to_string()),
209        Value::Time(t) => serde_json::Value::String(t.format("%H:%M:%S").to_string()),
210    }
211}
212
213/// Convert Value to YAML value
214fn value_to_yaml(value: &crate::functions::Value) -> serde_yaml::Value {
215    use crate::functions::Value;
216
217    match value {
218        Value::Null => serde_yaml::Value::Null,
219        Value::Boolean(b) => serde_yaml::Value::Bool(*b),
220        Value::Integer(i) => serde_yaml::Value::Number((*i).into()),
221        Value::Float(f) => serde_yaml::Value::Number(serde_yaml::Number::from(*f)),
222        Value::String(s) => serde_yaml::Value::String(s.clone()),
223        Value::Date(d) => serde_yaml::Value::String(d.format("%Y-%m-%d").to_string()),
224        Value::DateTime(dt) => serde_yaml::Value::String(dt.format("%Y-%m-%d %H:%M:%S").to_string()),
225        Value::Time(t) => serde_yaml::Value::String(t.format("%H:%M:%S").to_string()),
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::functions::Value;
233
234    fn create_test_rows() -> Vec<Row> {
235        let columns = vec!["id".to_string(), "name".to_string(), "age".to_string()];
236        let values1 = vec![
237            Value::Integer(1),
238            Value::String("Alice".to_string()),
239            Value::Integer(30),
240        ];
241        let values2 = vec![
242            Value::Integer(2),
243            Value::String("Bob".to_string()),
244            Value::Integer(25),
245        ];
246
247        vec![
248            Row::from_values(&columns, values1),
249            Row::from_values(&columns, values2),
250        ]
251    }
252
253    #[test]
254    fn test_format_csv() {
255        let rows = create_test_rows();
256        let output = format_csv(&rows).unwrap();
257
258        // Check that all columns are present (order may vary due to HashMap)
259        assert!(output.contains("id"));
260        assert!(output.contains("name"));
261        assert!(output.contains("age"));
262        // Check data is present
263        assert!(output.contains("Alice"));
264        assert!(output.contains("Bob"));
265        assert!(output.contains("30"));
266        assert!(output.contains("25"));
267    }
268
269    #[test]
270    fn test_format_json() {
271        let rows = create_test_rows();
272        let output = format_json(&rows).unwrap();
273
274        assert!(output.contains("\"id\""));
275        assert!(output.contains("\"name\""));
276        assert!(output.contains("Alice"));
277        assert!(output.contains("Bob"));
278    }
279
280    #[test]
281    fn test_format_yaml() {
282        let rows = create_test_rows();
283        let output = format_yaml(&rows).unwrap();
284
285        assert!(output.contains("id:"));
286        assert!(output.contains("name:"));
287        assert!(output.contains("Alice"));
288        assert!(output.contains("Bob"));
289    }
290
291    #[test]
292    fn test_format_xml() {
293        let rows = create_test_rows();
294        let output = format_xml(&rows).unwrap();
295
296        assert!(output.contains("<?xml version=\"1.0\" encoding=\"UTF-8\"?>"));
297        assert!(output.contains("<results>"));
298        assert!(output.contains("<row>"));
299        assert!(output.contains("<id>"));
300        assert!(output.contains("<name>"));
301        assert!(output.contains("Alice"));
302        assert!(output.contains("Bob"));
303        assert!(output.contains("</results>"));
304    }
305}
ddb_core/output/formatter.rs

ddb_core/output/
formatter.rs
