spring_batch_rs/item/json/
json_writer.rs

1use std::{
2    cell::{Cell, RefCell},
3    fs::File,
4    io::{BufWriter, Write},
5    path::Path,
6};
7
8use crate::{
9    core::item::{ItemWriter, ItemWriterResult},
10    BatchError,
11};
12
13pub struct JsonItemWriter<T: Write> {
14    stream: RefCell<BufWriter<T>>,
15    use_pretty_formatter: bool,
16    is_first_element: Cell<bool>,
17}
18
19impl<T: Write, R: serde::Serialize> ItemWriter<R> for JsonItemWriter<T> {
20    fn write(&self, items: &[R]) -> ItemWriterResult {
21        let mut json_chunk = String::new();
22
23        for item in items.iter() {
24            if !self.is_first_element.get() {
25                json_chunk.push(',');
26            } else {
27                self.is_first_element.set(false);
28            }
29
30            let result = if self.use_pretty_formatter {
31                serde_json::to_string_pretty(item)
32            } else {
33                serde_json::to_string(item)
34            };
35
36            json_chunk.push_str(&result.unwrap());
37
38            if self.use_pretty_formatter {
39                json_chunk.push('\n');
40            }
41        }
42
43        let result = self.stream.borrow_mut().write_all(json_chunk.as_bytes());
44
45        match result {
46            Ok(_ser) => Ok(()),
47            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
48        }
49    }
50
51    fn flush(&self) -> ItemWriterResult {
52        let result = self.stream.borrow_mut().flush();
53
54        match result {
55            Ok(()) => Ok(()),
56            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
57        }
58    }
59
60    fn open(&self) -> ItemWriterResult {
61        let begin_array = if self.use_pretty_formatter {
62            b"[\n".to_vec()
63        } else {
64            b"[".to_vec()
65        };
66
67        let result = self.stream.borrow_mut().write_all(&begin_array);
68
69        match result {
70            Ok(()) => Ok(()),
71            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
72        }
73    }
74
75    fn close(&self) -> ItemWriterResult {
76        let end_array = if self.use_pretty_formatter {
77            b"\n]\n".to_vec()
78        } else {
79            b"]\n".to_vec()
80        };
81
82        let result = self.stream.borrow_mut().write_all(&end_array);
83        let _ = self.stream.borrow_mut().flush();
84
85        match result {
86            Ok(()) => Ok(()),
87            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
88        }
89    }
90}
91
92#[derive(Default)]
93pub struct JsonItemWriterBuilder {
94    indent: Box<[u8]>,
95    pretty_formatter: bool,
96}
97
98impl JsonItemWriterBuilder {
99    pub fn new() -> Self {
100        Self {
101            indent: Box::from(b"  ".to_vec()),
102            pretty_formatter: false,
103        }
104    }
105
106    pub fn indent(mut self, indent: &[u8]) -> Self {
107        self.indent = Box::from(indent);
108        self
109    }
110
111    pub fn pretty_formatter(mut self, yes: bool) -> Self {
112        self.pretty_formatter = yes;
113        self
114    }
115
116    pub fn from_path<R: AsRef<Path>>(self, path: R) -> JsonItemWriter<File> {
117        let file = File::create(path).expect("Unable to open file");
118
119        let buf_writer = BufWriter::new(file);
120
121        JsonItemWriter {
122            stream: RefCell::new(buf_writer),
123            use_pretty_formatter: self.pretty_formatter,
124            is_first_element: Cell::new(true),
125        }
126    }
127
128    pub fn from_writer<W: Write>(self, wtr: W) -> JsonItemWriter<W> {
129        let buf_writer = BufWriter::new(wtr);
130
131        JsonItemWriter {
132            stream: RefCell::new(buf_writer),
133            use_pretty_formatter: self.pretty_formatter,
134            is_first_element: Cell::new(true),
135        }
136    }
137}