spring_batch_rs/item/csv/
csv_writer.rs

1use std::{cell::RefCell, fs::File, io::Write, path::Path};
2
3use csv::{Writer, WriterBuilder};
4use serde::Serialize;
5
6use crate::{
7    core::item::{ItemWriter, ItemWriterResult},
8    BatchError,
9};
10
11pub struct CsvItemWriter<T: Write> {
12    writer: RefCell<Writer<T>>,
13}
14
15impl<T: Write, R: Serialize> ItemWriter<R> for CsvItemWriter<T> {
16    fn write(&self, items: &[R]) -> ItemWriterResult {
17        for item in items.iter() {
18            let result = self.writer.borrow_mut().serialize(item);
19
20            if result.is_err() {
21                let error = result.err().unwrap();
22                return Err(BatchError::ItemWriter(error.to_string()));
23            }
24        }
25        Ok(())
26    }
27
28    /// Flush the contents of the internal buffer to the underlying writer.
29    ///
30    /// If there was a problem writing to the underlying writer, then an error
31    /// is returned.
32    ///
33    /// Note that this also flushes the underlying writer.
34    fn flush(&self) -> ItemWriterResult {
35        let result = self.writer.borrow_mut().flush();
36        match result {
37            Ok(()) => Ok(()),
38            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
39        }
40    }
41}
42
43#[derive(Default)]
44pub struct CsvItemWriterBuilder {
45    delimiter: u8,
46    has_headers: bool,
47}
48
49impl CsvItemWriterBuilder {
50    pub fn new() -> Self {
51        Self {
52            delimiter: b',',
53            has_headers: false,
54        }
55    }
56
57    pub fn delimiter(mut self, delimiter: u8) -> Self {
58        self.delimiter = delimiter;
59        self
60    }
61
62    pub fn has_headers(mut self, yes: bool) -> Self {
63        self.has_headers = yes;
64        self
65    }
66
67    pub fn from_path<R: AsRef<Path>>(self, path: R) -> CsvItemWriter<File> {
68        let writer = WriterBuilder::new()
69            .flexible(false)
70            .has_headers(self.has_headers)
71            .from_path(path);
72
73        CsvItemWriter {
74            writer: RefCell::new(writer.unwrap()),
75        }
76    }
77
78    /// Serialize a single record using Serde.
79    ///
80    /// # Example
81    ///
82    /// This shows how to serialize normal Rust structs as CSV records. The
83    /// fields of the struct are used to write a header row automatically.
84    /// (Writing the header row automatically can be disabled by building the
85    /// CSV writer with a [`WriterBuilder`](struct.WriterBuilder.html) and
86    /// calling the `has_headers` method.)
87    ///
88    /// ```
89    /// # use std::error::Error;
90    /// # use csv::Writer;
91    /// # use spring_batch_rs::{item::csv::csv_writer::CsvItemWriterBuilder, core::item::ItemWriter};
92    /// #[derive(serde::Serialize)]
93    /// struct Row<'a> {
94    ///     city: &'a str,
95    ///     country: &'a str,
96    ///     #[serde(rename = "popcount")]
97    ///     population: u64,
98    /// }
99    ///
100    /// # fn main() { example().unwrap(); }
101    /// fn example() -> Result<(), Box<dyn Error>> {
102    ///     let wtr = CsvItemWriterBuilder::new()
103    ///         .has_headers(true)
104    ///         .from_writer(vec![]);
105    ///
106    ///     let rows = &[
107    ///         Row {
108    ///             city: "Boston",
109    ///             country: "United States",
110    ///             population: 4628910,
111    ///         },
112    ///         Row {
113    ///             city: "Concord",
114    ///             country: "United States",
115    ///             population: 42695,
116    ///         }
117    ///     ];
118    ///     wtr.write(rows);
119    ///
120    ///     Ok(())
121    /// }
122    /// ```
123    pub fn from_writer<W: Write>(self, wtr: W) -> CsvItemWriter<W> {
124        let wtr = WriterBuilder::new()
125            .flexible(false)
126            .has_headers(self.has_headers)
127            .from_writer(wtr);
128
129        CsvItemWriter {
130            writer: RefCell::new(wtr),
131        }
132    }
133}
spring_batch_rs/item/csv/csv_writer.rs

spring_batch_rs/item/csv/
csv_writer.rs
