spring_batch_rs/item/csv/
csv_reader.rs

1use csv::{ReaderBuilder, StringRecordsIntoIter, Terminator, Trim};
2use serde::de::DeserializeOwned;
3use std::{cell::RefCell, fs::File, io::Read, path::Path};
4
5use crate::{
6    core::item::{ItemReader, ItemReaderResult},
7    error::BatchError,
8};
9
10/// A CSV item reader that implements the `ItemReader` trait.
11pub struct CsvItemReader<R> {
12    records: RefCell<StringRecordsIntoIter<R>>,
13}
14
15impl<R: Read, T: DeserializeOwned> ItemReader<T> for CsvItemReader<R> {
16    /// Reads the next item from the CSV file.
17    ///
18    /// Returns `Ok(Some(record))` if a record is successfully read,
19    /// `Ok(None)` if there are no more records to read, and
20    /// `Err(BatchError::ItemReader(error))` if an error occurs during reading.
21    fn read(&self) -> ItemReaderResult<T> {
22        if let Some(result) = self.records.borrow_mut().next() {
23            match result {
24                Ok(string_record) => {
25                    let result: Result<T, _> = string_record.deserialize(None);
26
27                    match result {
28                        Ok(record) => Ok(Some(record)),
29                        Err(error) => Err(BatchError::ItemReader(error.to_string())),
30                    }
31                }
32                Err(error) => Err(BatchError::ItemReader(error.to_string())),
33            }
34        } else {
35            Ok(None)
36        }
37    }
38}
39
40/// A builder for configuring CSV item reading.
41#[derive(Default)]
42pub struct CsvItemReaderBuilder {
43    delimiter: u8,
44    terminator: Terminator,
45    has_headers: bool,
46}
47
48impl CsvItemReaderBuilder {
49    /// Creates a new `CsvItemReaderBuilder` with default configuration.
50    pub fn new() -> Self {
51        Self {
52            delimiter: b',',
53            terminator: Terminator::CRLF,
54            has_headers: false,
55        }
56    }
57
58    /// Sets the delimiter character for the CSV parsing.
59    pub fn delimiter(mut self, delimiter: u8) -> Self {
60        self.delimiter = delimiter;
61        self
62    }
63
64    /// Sets the line terminator for the CSV parsing.
65    pub fn terminator(mut self, terminator: Terminator) -> Self {
66        self.terminator = terminator;
67        self
68    }
69
70    /// Sets whether the CSV file has headers.
71    pub fn has_headers(mut self, yes: bool) -> Self {
72        self.has_headers = yes;
73        self
74    }
75
76    /// Creates a `CsvItemReader` from a reader.
77    pub fn from_reader<R: Read>(self, rdr: R) -> CsvItemReader<R> {
78        let rdr = ReaderBuilder::new()
79            .trim(Trim::All)
80            .delimiter(self.delimiter)
81            .terminator(self.terminator)
82            .has_headers(self.has_headers)
83            .flexible(false)
84            .from_reader(rdr);
85
86        let records = rdr.into_records();
87
88        CsvItemReader {
89            records: RefCell::new(records),
90        }
91    }
92
93    /// Creates a `CsvItemReader` from a file path.
94    pub fn from_path<R: AsRef<Path>>(self, path: R) -> CsvItemReader<File> {
95        let rdr = ReaderBuilder::new()
96            .trim(Trim::All)
97            .delimiter(self.delimiter)
98            .terminator(self.terminator)
99            .has_headers(self.has_headers)
100            .flexible(false)
101            .from_path(path);
102
103        let records = rdr.unwrap().into_records();
104
105        CsvItemReader {
106            records: RefCell::new(records),
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use std::error::Error;
114
115    use csv::StringRecord;
116
117    use crate::item::csv::csv_reader::CsvItemReaderBuilder;
118
119    #[test]
120    fn this_test_will_pass() -> Result<(), Box<dyn Error>> {
121        let data = "city,country,pop
122        Boston,United States,4628910
123        Concord,United States,42695";
124
125        let reader = CsvItemReaderBuilder::new()
126            .has_headers(true)
127            .delimiter(b',')
128            .from_reader(data.as_bytes());
129
130        let records = reader
131            .records
132            .into_inner()
133            .collect::<Result<Vec<StringRecord>, csv::Error>>()?;
134
135        assert_eq!(
136            records,
137            vec![
138                vec!["Boston", "United States", "4628910"],
139                vec!["Concord", "United States", "42695"],
140            ]
141        );
142
143        Ok(())
144    }
145}
spring_batch_rs/item/csv/csv_reader.rs

spring_batch_rs/item/csv/
csv_reader.rs
