llm_cost_ops/export/
formats.rs

1// Export formats (CSV, JSON, Excel)
2
3use bytes::Bytes;
4use serde::{Deserialize, Serialize};
5use std::io::Write;
6
7use super::{ExportError, ExportResult};
8
9/// Supported export formats
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum ExportFormat {
13    Csv,
14    Json,
15    Excel,
16    JsonLines,
17}
18
19impl std::fmt::Display for ExportFormat {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        match self {
22            ExportFormat::Csv => write!(f, "csv"),
23            ExportFormat::Json => write!(f, "json"),
24            ExportFormat::Excel => write!(f, "xlsx"),
25            ExportFormat::JsonLines => write!(f, "jsonl"),
26        }
27    }
28}
29
30impl ExportFormat {
31    pub fn mime_type(&self) -> &'static str {
32        match self {
33            ExportFormat::Csv => "text/csv",
34            ExportFormat::Json => "application/json",
35            ExportFormat::Excel => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
36            ExportFormat::JsonLines => "application/x-ndjson",
37        }
38    }
39
40    pub fn file_extension(&self) -> &'static str {
41        match self {
42            ExportFormat::Csv => "csv",
43            ExportFormat::Json => "json",
44            ExportFormat::Excel => "xlsx",
45            ExportFormat::JsonLines => "jsonl",
46        }
47    }
48}
49
50/// Data to be exported
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct ExportData {
53    pub headers: Vec<String>,
54    pub rows: Vec<Vec<serde_json::Value>>,
55    pub metadata: std::collections::HashMap<String, String>,
56}
57
58impl ExportData {
59    pub fn new(headers: Vec<String>) -> Self {
60        Self {
61            headers,
62            rows: Vec::new(),
63            metadata: std::collections::HashMap::new(),
64        }
65    }
66
67    pub fn add_row(&mut self, row: Vec<serde_json::Value>) {
68        self.rows.push(row);
69    }
70
71    pub fn add_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
72        self.metadata.insert(key.into(), value.into());
73    }
74
75    pub fn row_count(&self) -> usize {
76        self.rows.len()
77    }
78}
79
80/// Export trait
81pub trait Exporter: Send + Sync {
82    fn export(&self, data: &ExportData) -> ExportResult<Bytes>;
83}
84
85/// CSV Exporter
86pub struct CsvExporter;
87
88impl Exporter for CsvExporter {
89    fn export(&self, data: &ExportData) -> ExportResult<Bytes> {
90        let mut writer = csv::Writer::from_writer(Vec::new());
91
92        // Write headers
93        writer
94            .write_record(&data.headers)
95            .map_err(|e| ExportError::FormatError(format!("CSV header error: {}", e)))?;
96
97        // Write rows
98        for row in &data.rows {
99            let string_row: Vec<String> = row
100                .iter()
101                .map(|v| match v {
102                    serde_json::Value::String(s) => s.clone(),
103                    serde_json::Value::Number(n) => n.to_string(),
104                    serde_json::Value::Bool(b) => b.to_string(),
105                    serde_json::Value::Null => String::new(),
106                    _ => v.to_string(),
107                })
108                .collect();
109
110            writer
111                .write_record(&string_row)
112                .map_err(|e| ExportError::FormatError(format!("CSV row error: {}", e)))?;
113        }
114
115        writer
116            .flush()
117            .map_err(|e| ExportError::FormatError(format!("CSV flush error: {}", e)))?;
118
119        let bytes = writer
120            .into_inner()
121            .map_err(|e| ExportError::FormatError(format!("CSV finalize error: {}", e)))?;
122
123        Ok(Bytes::from(bytes))
124    }
125}
126
127/// JSON Exporter
128pub struct JsonExporter;
129
130impl Exporter for JsonExporter {
131    fn export(&self, data: &ExportData) -> ExportResult<Bytes> {
132        // Create JSON array of objects
133        let mut records = Vec::new();
134
135        for row in &data.rows {
136            let mut record = serde_json::Map::new();
137            for (i, header) in data.headers.iter().enumerate() {
138                if let Some(value) = row.get(i) {
139                    record.insert(header.clone(), value.clone());
140                }
141            }
142            records.push(serde_json::Value::Object(record));
143        }
144
145        let output = serde_json::json!({
146            "data": records,
147            "metadata": data.metadata,
148            "count": data.rows.len(),
149        });
150
151        let json_bytes = serde_json::to_vec_pretty(&output)?;
152        Ok(Bytes::from(json_bytes))
153    }
154}
155
156/// JSON Lines Exporter (NDJSON)
157pub struct JsonLinesExporter;
158
159impl Exporter for JsonLinesExporter {
160    fn export(&self, data: &ExportData) -> ExportResult<Bytes> {
161        let mut output = Vec::new();
162
163        for row in &data.rows {
164            let mut record = serde_json::Map::new();
165            for (i, header) in data.headers.iter().enumerate() {
166                if let Some(value) = row.get(i) {
167                    record.insert(header.clone(), value.clone());
168                }
169            }
170
171            let line = serde_json::to_string(&serde_json::Value::Object(record))?;
172            writeln!(output, "{}", line)?;
173        }
174
175        Ok(Bytes::from(output))
176    }
177}
178
179/// Excel Exporter
180pub struct ExcelExporter {
181    sheet_name: String,
182}
183
184impl ExcelExporter {
185    pub fn new(sheet_name: impl Into<String>) -> Self {
186        Self {
187            sheet_name: sheet_name.into(),
188        }
189    }
190}
191
192impl Default for ExcelExporter {
193    fn default() -> Self {
194        Self::new("Data")
195    }
196}
197
198impl Exporter for ExcelExporter {
199    fn export(&self, data: &ExportData) -> ExportResult<Bytes> {
200        use rust_xlsxwriter::*;
201
202        let mut workbook = Workbook::new();
203        let worksheet = workbook
204            .add_worksheet()
205            .set_name(&self.sheet_name)
206            .map_err(|e| ExportError::FormatError(format!("Excel sheet error: {}", e)))?;
207
208        // Header formatting
209        let header_format = Format::new()
210            .set_bold()
211            .set_background_color(Color::RGB(0x4472C4))
212            .set_font_color(Color::White);
213
214        // Write headers
215        for (col, header) in data.headers.iter().enumerate() {
216            worksheet
217                .write_string_with_format(0, col as u16, header, &header_format)
218                .map_err(|e| ExportError::FormatError(format!("Excel header error: {}", e)))?;
219        }
220
221        // Write data rows
222        for (row_idx, row) in data.rows.iter().enumerate() {
223            for (col_idx, value) in row.iter().enumerate() {
224                let excel_row = (row_idx + 1) as u32;
225                let excel_col = col_idx as u16;
226
227                match value {
228                    serde_json::Value::String(s) => {
229                        worksheet
230                            .write_string(excel_row, excel_col, s)
231                            .map_err(|e| {
232                                ExportError::FormatError(format!("Excel write error: {}", e))
233                            })?;
234                    }
235                    serde_json::Value::Number(n) => {
236                        if let Some(f) = n.as_f64() {
237                            worksheet
238                                .write_number(excel_row, excel_col, f)
239                                .map_err(|e| {
240                                    ExportError::FormatError(format!("Excel write error: {}", e))
241                                })?;
242                        } else {
243                            worksheet
244                                .write_string(excel_row, excel_col, n.to_string())
245                                .map_err(|e| {
246                                    ExportError::FormatError(format!("Excel write error: {}", e))
247                                })?;
248                        }
249                    }
250                    serde_json::Value::Bool(b) => {
251                        worksheet
252                            .write_boolean(excel_row, excel_col, *b)
253                            .map_err(|e| {
254                                ExportError::FormatError(format!("Excel write error: {}", e))
255                            })?;
256                    }
257                    serde_json::Value::Null => {
258                        worksheet
259                            .write_blank(excel_row, excel_col, &Format::new())
260                            .map_err(|e| {
261                                ExportError::FormatError(format!("Excel write error: {}", e))
262                            })?;
263                    }
264                    _ => {
265                        worksheet
266                            .write_string(excel_row, excel_col, value.to_string())
267                            .map_err(|e| {
268                                ExportError::FormatError(format!("Excel write error: {}", e))
269                            })?;
270                    }
271                }
272            }
273        }
274
275        // Auto-fit columns
276        worksheet.autofit();
277
278        // Save to bytes
279        let bytes = workbook
280            .save_to_buffer()
281            .map_err(|e| ExportError::FormatError(format!("Excel save error: {}", e)))?;
282
283        Ok(Bytes::from(bytes))
284    }
285}
286
287/// Factory for creating exporters
288pub fn create_exporter(format: ExportFormat) -> Box<dyn Exporter> {
289    match format {
290        ExportFormat::Csv => Box::new(CsvExporter),
291        ExportFormat::Json => Box::new(JsonExporter),
292        ExportFormat::Excel => Box::new(ExcelExporter::default()),
293        ExportFormat::JsonLines => Box::new(JsonLinesExporter),
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    fn create_test_data() -> ExportData {
302        let mut data = ExportData::new(vec![
303            "Name".to_string(),
304            "Age".to_string(),
305            "Active".to_string(),
306        ]);
307
308        data.add_row(vec![
309            serde_json::json!("Alice"),
310            serde_json::json!(30),
311            serde_json::json!(true),
312        ]);
313
314        data.add_row(vec![
315            serde_json::json!("Bob"),
316            serde_json::json!(25),
317            serde_json::json!(false),
318        ]);
319
320        data.add_metadata("exported_by", "test");
321        data
322    }
323
324    #[test]
325    fn test_csv_export() {
326        let data = create_test_data();
327        let exporter = CsvExporter;
328        let result = exporter.export(&data);
329
330        assert!(result.is_ok());
331        let bytes = result.unwrap();
332        let content = String::from_utf8(bytes.to_vec()).unwrap();
333
334        assert!(content.contains("Name,Age,Active"));
335        assert!(content.contains("Alice,30,true"));
336    }
337
338    #[test]
339    fn test_json_export() {
340        let data = create_test_data();
341        let exporter = JsonExporter;
342        let result = exporter.export(&data);
343
344        assert!(result.is_ok());
345        let bytes = result.unwrap();
346        let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
347
348        assert_eq!(json["count"], 2);
349        assert_eq!(json["data"][0]["Name"], "Alice");
350    }
351
352    #[test]
353    fn test_jsonlines_export() {
354        let data = create_test_data();
355        let exporter = JsonLinesExporter;
356        let result = exporter.export(&data);
357
358        assert!(result.is_ok());
359        let bytes = result.unwrap();
360        let content = String::from_utf8(bytes.to_vec()).unwrap();
361
362        let lines: Vec<&str> = content.lines().collect();
363        assert_eq!(lines.len(), 2);
364    }
365
366    #[test]
367    fn test_excel_export() {
368        let data = create_test_data();
369        let exporter = ExcelExporter::default();
370        let result = exporter.export(&data);
371
372        assert!(result.is_ok());
373        let bytes = result.unwrap();
374        assert!(bytes.len() > 0);
375    }
376
377    #[test]
378    fn test_export_format_display() {
379        assert_eq!(ExportFormat::Csv.to_string(), "csv");
380        assert_eq!(ExportFormat::Json.to_string(), "json");
381        assert_eq!(ExportFormat::Excel.to_string(), "xlsx");
382    }
383
384    #[test]
385    fn test_export_format_mime_types() {
386        assert_eq!(ExportFormat::Csv.mime_type(), "text/csv");
387        assert_eq!(ExportFormat::Json.mime_type(), "application/json");
388        assert_eq!(
389            ExportFormat::Excel.mime_type(),
390            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
391        );
392    }
393}
llm_cost_ops/export/formats.rs

llm_cost_ops/export/
formats.rs
