1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use std::{
    cell::RefCell,
    fs::File,
    io::{BufWriter, Write},
    path::Path,
};

use crate::{core::item::ItemWriter, BatchError};

pub struct JsonItemWriter<T: Write> {
    stream: RefCell<BufWriter<T>>,
    use_pretty_formatter: bool,
}

impl<T: Write, R: serde::Serialize> ItemWriter<R> for JsonItemWriter<T> {
    fn write(&self, item: &R) -> Result<(), BatchError> {
        let json = if self.use_pretty_formatter {
            serde_json::to_string_pretty(item)
        } else {
            serde_json::to_string(item)
        };
        let result = self.stream.borrow_mut().write_all(json.unwrap().as_bytes());

        match result {
            Ok(_ser) => Ok(()),
            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
        }
    }

    fn flush(&self) -> Result<(), BatchError> {
        let result = self.stream.borrow_mut().flush();

        match result {
            Ok(()) => Ok(()),
            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
        }
    }

    fn open(&self) -> Result<(), BatchError> {
        let begin_array = if self.use_pretty_formatter {
            b"[\n".to_vec()
        } else {
            b"[".to_vec()
        };

        let result = self.stream.borrow_mut().write_all(&begin_array);

        match result {
            Ok(()) => Ok(()),
            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
        }
    }

    fn update(&self, is_first_item: bool) -> Result<(), BatchError> {
        if !is_first_item {
            let separator = if self.use_pretty_formatter {
                b",\n".to_vec()
            } else {
                b",".to_vec()
            };

            let result = self.stream.borrow_mut().write_all(&separator);

            return match result {
                Ok(()) => Ok(()),
                Err(error) => Err(BatchError::ItemWriter(error.to_string())),
            };
        }
        Ok(())
    }

    fn close(&self) -> Result<(), BatchError> {
        let end_array = if self.use_pretty_formatter {
            b"\n]\n".to_vec()
        } else {
            b"]\n".to_vec()
        };

        let result = self.stream.borrow_mut().write_all(&end_array);
        let _ = self.stream.borrow_mut().flush();

        match result {
            Ok(()) => Ok(()),
            Err(error) => Err(BatchError::ItemWriter(error.to_string())),
        }
    }
}

#[derive(Default)]
pub struct JsonItemWriterBuilder {
    indent: Box<[u8]>,
    pretty_formatter: bool,
}

impl JsonItemWriterBuilder {
    pub fn new() -> JsonItemWriterBuilder {
        JsonItemWriterBuilder {
            indent: Box::from(b"  ".to_vec()),
            pretty_formatter: false,
        }
    }

    pub fn indent(mut self, indent: &[u8]) -> JsonItemWriterBuilder {
        self.indent = Box::from(indent);
        self
    }

    pub fn pretty_formatter(mut self, yes: bool) -> JsonItemWriterBuilder {
        self.pretty_formatter = yes;
        self
    }

    pub fn from_path<R: AsRef<Path>>(self, path: R) -> JsonItemWriter<File> {
        let file = File::create(path).expect("Unable to open file");

        let buf_writer = BufWriter::new(file);

        JsonItemWriter {
            stream: RefCell::new(buf_writer),
            use_pretty_formatter: self.pretty_formatter,
        }
    }
}