1use crate::{QueryResult, ValueOut};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum RowFormat {
5 Table,
6 Json,
7 Ndjson,
8 Csv,
9 Tsv,
10 Toon,
11}
12
13impl RowFormat {
14 pub fn parse(value: &str) -> Option<Self> {
15 match value {
16 "table" => Some(Self::Table),
17 "json" => Some(Self::Json),
18 "ndjson" => Some(Self::Ndjson),
19 "csv" => Some(Self::Csv),
20 "tsv" => Some(Self::Tsv),
21 "toon" => Some(Self::Toon),
22 _ => None,
23 }
24 }
25
26 pub fn vocabulary() -> &'static str {
27 "table, json, ndjson, csv, tsv, toon"
28 }
29}
30
31pub fn format_query_result(result: &QueryResult, format: RowFormat) -> Vec<u8> {
32 match format {
33 RowFormat::Table => format_table(result).into_bytes(),
34 RowFormat::Json => format_json(result).into_bytes(),
35 RowFormat::Ndjson => format_ndjson(result).into_bytes(),
36 RowFormat::Csv => format_delimited(result, b','),
37 RowFormat::Tsv => format_delimited(result, b'\t'),
38 RowFormat::Toon => format_toon(result).into_bytes(),
39 }
40}
41
42fn columns(result: &QueryResult) -> Vec<String> {
43 if !result.columns.is_empty() {
44 return result.columns.clone();
45 }
46 result
47 .rows
48 .first()
49 .map(|row| row.iter().map(|(key, _)| key.clone()).collect())
50 .unwrap_or_default()
51}
52
53fn value_at<'a>(row: &'a [(String, ValueOut)], column: &str) -> Option<&'a ValueOut> {
54 row.iter()
55 .find(|(key, _)| key == column)
56 .map(|(_, value)| value)
57}
58
59fn format_table(result: &QueryResult) -> String {
60 let columns = columns(result);
61 if result.rows.is_empty() {
62 let mut out = "(no rows)\n".to_string();
63 if let Some(notice) = result.notice.as_deref() {
64 out.push_str("Note: ");
65 out.push_str(notice);
66 out.push('\n');
67 }
68 return out;
69 }
70
71 let mut widths: Vec<usize> = columns.iter().map(|column| column.len()).collect();
72 for row in &result.rows {
73 for (index, column) in columns.iter().enumerate() {
74 let value = value_at(row, column)
75 .map(table_value)
76 .unwrap_or_else(|| "null".to_string());
77 widths[index] = widths[index].max(value.len());
78 }
79 }
80
81 let mut out = String::new();
82 write_table_line(&mut out, &columns, &widths);
83 let separator: Vec<String> = widths.iter().map(|width| "-".repeat(*width)).collect();
84 write_table_line(&mut out, &separator, &widths);
85 for row in &result.rows {
86 let values: Vec<String> = columns
87 .iter()
88 .map(|column| {
89 value_at(row, column)
90 .map(table_value)
91 .unwrap_or_else(|| "null".to_string())
92 })
93 .collect();
94 write_table_line(&mut out, &values, &widths);
95 }
96 out
97}
98
99fn write_table_line(out: &mut String, cells: &[String], widths: &[usize]) {
100 for (index, cell) in cells.iter().enumerate() {
101 if index > 0 {
102 out.push_str(" ");
103 }
104 out.push_str(cell);
105 if index + 1 < cells.len() {
106 for _ in cell.len()..widths[index] {
107 out.push(' ');
108 }
109 }
110 }
111 out.push('\n');
112}
113
114fn format_json(result: &QueryResult) -> String {
115 let columns = columns(result);
116 let mut out = String::from("[");
117 for (row_index, row) in result.rows.iter().enumerate() {
118 if row_index > 0 {
119 out.push(',');
120 }
121 write_json_row(&mut out, row, &columns);
122 }
123 out.push_str("]\n");
124 out
125}
126
127fn format_ndjson(result: &QueryResult) -> String {
128 let columns = columns(result);
129 let mut out = String::new();
130 for row in &result.rows {
131 write_json_row(&mut out, row, &columns);
132 out.push('\n');
133 }
134 out
135}
136
137fn format_toon(result: &QueryResult) -> String {
138 let columns = columns(result);
139 let mut out = String::new();
140 out.push('[');
141 out.push_str(&result.rows.len().to_string());
142 out.push_str("]{");
143 for (index, column) in columns.iter().enumerate() {
144 if index > 0 {
145 out.push(',');
146 }
147 write_toon_key(&mut out, column);
148 }
149 out.push_str("}:\n");
150 for row in &result.rows {
151 out.push_str(" ");
152 for (index, column) in columns.iter().enumerate() {
153 if index > 0 {
154 out.push(',');
155 }
156 if let Some(value) = value_at(row, column) {
157 write_toon_value(&mut out, value);
158 } else {
159 out.push_str("null");
160 }
161 }
162 out.push('\n');
163 }
164 out
165}
166
167fn write_json_row(out: &mut String, row: &[(String, ValueOut)], columns: &[String]) {
168 out.push('{');
169 for (index, column) in columns.iter().enumerate() {
170 if index > 0 {
171 out.push(',');
172 }
173 write_json_string(out, column);
174 out.push(':');
175 if let Some(value) = value_at(row, column) {
176 write_json_value(out, value);
177 } else {
178 out.push_str("null");
179 }
180 }
181 out.push('}');
182}
183
184fn write_toon_key(out: &mut String, value: &str) {
185 if is_toon_bare_key(value) {
186 out.push_str(value);
187 } else {
188 write_toon_string(out, value);
189 }
190}
191
192fn write_toon_value(out: &mut String, value: &ValueOut) {
193 match value {
194 ValueOut::Null => out.push_str("null"),
195 ValueOut::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
196 ValueOut::Integer(value) => out.push_str(&value.to_string()),
197 ValueOut::Float(value) => out.push_str(&value.to_string()),
198 ValueOut::String(value) => {
199 if is_toon_bare_string(value) {
200 out.push_str(value);
201 } else {
202 write_toon_string(out, value);
203 }
204 }
205 }
206}
207
208fn is_toon_bare_key(value: &str) -> bool {
209 let mut chars = value.chars();
210 let Some(first) = chars.next() else {
211 return false;
212 };
213 if !(first == '_' || first.is_ascii_alphabetic()) {
214 return false;
215 }
216 chars.all(|c| c == '_' || c == '-' || c == '.' || c.is_ascii_alphanumeric())
217}
218
219fn is_toon_bare_string(value: &str) -> bool {
220 !value.is_empty()
221 && !matches!(value, "true" | "false" | "null")
222 && value.parse::<f64>().is_err()
223 && !value.starts_with('-')
224 && !value.chars().any(|c| {
225 matches!(
226 c,
227 ',' | ':' | '{' | '}' | '[' | ']' | '"' | '\\' | '\n' | '\r' | '\t'
228 )
229 })
230}
231
232fn write_toon_string(out: &mut String, value: &str) {
233 out.push('"');
234 for c in value.chars() {
235 match c {
236 '"' => out.push_str("\\\""),
237 '\\' => out.push_str("\\\\"),
238 '\n' => out.push_str("\\n"),
239 '\r' => out.push_str("\\r"),
240 '\t' => out.push_str("\\t"),
241 c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
242 c => out.push(c),
243 }
244 }
245 out.push('"');
246}
247
248fn format_delimited(result: &QueryResult, delimiter: u8) -> Vec<u8> {
249 let columns = columns(result);
250 let mut out = Vec::new();
251 write_delimited_record(&mut out, &columns, delimiter);
252 for row in &result.rows {
253 let values: Vec<String> = columns
254 .iter()
255 .map(|column| value_at(row, column).map(raw_value).unwrap_or_default())
256 .collect();
257 write_delimited_record(&mut out, &values, delimiter);
258 }
259 out
260}
261
262fn write_delimited_record(out: &mut Vec<u8>, fields: &[String], delimiter: u8) {
263 for (index, field) in fields.iter().enumerate() {
264 if index > 0 {
265 out.push(delimiter);
266 }
267 write_delimited_field(out, field, delimiter);
268 }
269 out.push(b'\n');
270}
271
272fn write_delimited_field(out: &mut Vec<u8>, field: &str, delimiter: u8) {
273 let needs_quotes = field
274 .bytes()
275 .any(|byte| byte == delimiter || byte == b'"' || byte == b'\n' || byte == b'\r');
276 if !needs_quotes {
277 out.extend_from_slice(field.as_bytes());
278 return;
279 }
280 out.push(b'"');
281 for byte in field.bytes() {
282 if byte == b'"' {
283 out.extend_from_slice(b"\"\"");
284 } else {
285 out.push(byte);
286 }
287 }
288 out.push(b'"');
289}
290
291fn table_value(value: &ValueOut) -> String {
292 raw_value(value)
293 .replace('\\', "\\\\")
294 .replace('\n', "\\n")
295 .replace('\r', "\\r")
296 .replace('\t', "\\t")
297}
298
299fn raw_value(value: &ValueOut) -> String {
300 match value {
301 ValueOut::Null => String::new(),
302 ValueOut::Bool(value) => value.to_string(),
303 ValueOut::Integer(value) => value.to_string(),
304 ValueOut::Float(value) => value.to_string(),
305 ValueOut::String(value) => value.clone(),
306 }
307}
308
309fn write_json_value(out: &mut String, value: &ValueOut) {
310 match value {
311 ValueOut::Null => out.push_str("null"),
312 ValueOut::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
313 ValueOut::Integer(value) => out.push_str(&value.to_string()),
314 ValueOut::Float(value) => out.push_str(&value.to_string()),
315 ValueOut::String(value) => write_json_string(out, value),
316 }
317}
318
319fn write_json_string(out: &mut String, value: &str) {
320 out.push('"');
321 for c in value.chars() {
322 match c {
323 '"' => out.push_str("\\\""),
324 '\\' => out.push_str("\\\\"),
325 '\n' => out.push_str("\\n"),
326 '\r' => out.push_str("\\r"),
327 '\t' => out.push_str("\\t"),
328 c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
329 c => out.push(c),
330 }
331 }
332 out.push('"');
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338
339 fn sample() -> QueryResult {
340 QueryResult {
341 statement: "select".to_string(),
342 affected: 0,
343 columns: vec!["id".to_string(), "name".to_string(), "note".to_string()],
344 rows: vec![
345 vec![
346 ("id".to_string(), ValueOut::Integer(1)),
347 ("name".to_string(), ValueOut::String("Ada".to_string())),
348 (
349 "note".to_string(),
350 ValueOut::String("quote \" comma ,".to_string()),
351 ),
352 ],
353 vec![
354 ("id".to_string(), ValueOut::Integer(2)),
355 ("name".to_string(), ValueOut::String("Linus".to_string())),
356 (
357 "note".to_string(),
358 ValueOut::String("tab\tline\nend".to_string()),
359 ),
360 ],
361 ],
362 notice: None,
363 }
364 }
365
366 #[test]
367 fn formats_table_byte_exact() {
368 assert_eq!(
369 format_query_result(&sample(), RowFormat::Table),
370 b"id name note\n-- ----- ---------------\n1 Ada quote \" comma ,\n2 Linus tab\\tline\\nend\n"
371 );
372 }
373
374 #[test]
375 fn formats_json_byte_exact() {
376 assert_eq!(
377 format_query_result(&sample(), RowFormat::Json),
378 br#"[{"id":1,"name":"Ada","note":"quote \" comma ,"},{"id":2,"name":"Linus","note":"tab\tline\nend"}]
379"#
380 );
381 }
382
383 #[test]
384 fn formats_ndjson_one_object_per_row() {
385 assert_eq!(
386 format_query_result(&sample(), RowFormat::Ndjson),
387 br#"{"id":1,"name":"Ada","note":"quote \" comma ,"}
388{"id":2,"name":"Linus","note":"tab\tline\nend"}
389"#
390 );
391 }
392
393 #[test]
394 fn formats_csv_byte_exact() {
395 assert_eq!(
396 format_query_result(&sample(), RowFormat::Csv),
397 b"id,name,note\n1,Ada,\"quote \"\" comma ,\"\n2,Linus,\"tab\tline\nend\"\n"
398 );
399 }
400
401 #[test]
402 fn formats_tsv_byte_exact() {
403 assert_eq!(
404 format_query_result(&sample(), RowFormat::Tsv),
405 b"id\tname\tnote\n1\tAda\t\"quote \"\" comma ,\"\n2\tLinus\t\"tab\tline\nend\"\n"
406 );
407 }
408
409 #[test]
410 fn formats_toon_byte_exact() {
411 assert_eq!(
412 format_query_result(&sample(), RowFormat::Toon),
413 br#"[2]{id,name,note}:
414 1,Ada,"quote \" comma ,"
415 2,Linus,"tab\tline\nend"
416"#
417 );
418 }
419
420 #[test]
421 fn formats_toon_quotes_delimiters_and_scalar_like_strings() {
422 let result = QueryResult {
423 statement: "select".to_string(),
424 affected: 0,
425 columns: vec![
426 "label".to_string(),
427 "empty".to_string(),
428 "dash".to_string(),
429 "truthy".to_string(),
430 "nullable".to_string(),
431 "count".to_string(),
432 "enabled".to_string(),
433 ],
434 rows: vec![vec![
435 ("label".to_string(), ValueOut::String("a,b".to_string())),
436 ("empty".to_string(), ValueOut::String(String::new())),
437 ("dash".to_string(), ValueOut::String("-x".to_string())),
438 ("truthy".to_string(), ValueOut::String("true".to_string())),
439 ("nullable".to_string(), ValueOut::Null),
440 ("count".to_string(), ValueOut::Integer(7)),
441 ("enabled".to_string(), ValueOut::Bool(false)),
442 ]],
443 notice: None,
444 };
445
446 assert_eq!(
447 format_query_result(&result, RowFormat::Toon),
448 br#"[1]{label,empty,dash,truthy,nullable,count,enabled}:
449 "a,b","","-x","true",null,7,false
450"#
451 );
452 }
453
454 #[test]
455 fn csv_and_tsv_round_trip_through_standard_parser() {
456 for (format, delimiter) in [(RowFormat::Csv, b','), (RowFormat::Tsv, b'\t')] {
457 let bytes = format_query_result(&sample(), format);
458 let mut reader = csv::ReaderBuilder::new()
459 .delimiter(delimiter)
460 .from_reader(bytes.as_slice());
461 let records: Vec<csv::StringRecord> =
462 reader.records().collect::<Result<_, _>>().unwrap();
463 assert_eq!(reader.headers().unwrap(), &["id", "name", "note"][..]);
464 assert_eq!(records[0].get(2), Some("quote \" comma ,"));
465 assert_eq!(records[1].get(2), Some("tab\tline\nend"));
466 }
467 }
468}