use std::{
cell::{Cell, RefCell},
fs::File,
io::{BufWriter, Write},
marker::PhantomData,
path::Path,
};
use crate::{
BatchError,
core::item::{ItemWriter, ItemWriterResult},
};
pub struct JsonItemWriter<O, W: Write> {
stream: RefCell<BufWriter<W>>,
use_pretty_formatter: bool,
indent: Box<[u8]>,
is_first_element: Cell<bool>,
_phantom: PhantomData<O>,
}
impl<O: serde::Serialize, W: Write> ItemWriter<O> for JsonItemWriter<O, W> {
fn write(&self, items: &[O]) -> ItemWriterResult {
let mut json_chunk = String::new();
for item in items.iter() {
if !self.is_first_element.get() {
json_chunk.push(',');
} else {
self.is_first_element.set(false);
}
let result = if self.use_pretty_formatter {
let mut buf = Vec::new();
let formatter = serde_json::ser::PrettyFormatter::with_indent(&self.indent);
let mut ser = serde_json::Serializer::with_formatter(&mut buf, formatter);
match item.serialize(&mut ser) {
Ok(_) => match String::from_utf8(buf) {
Ok(s) => Ok(s),
Err(e) => Err(BatchError::ItemWriter(e.to_string())),
},
Err(e) => Err(BatchError::ItemWriter(e.to_string())),
}
} else {
serde_json::to_string(item).map_err(|e| BatchError::ItemWriter(e.to_string()))
};
match result {
Ok(json_str) => json_chunk.push_str(&json_str),
Err(e) => return Err(e),
}
if self.use_pretty_formatter {
json_chunk.push('\n');
}
}
let result = self.stream.borrow_mut().write_all(json_chunk.as_bytes());
match result {
Ok(_ser) => Ok(()),
Err(error) => Err(BatchError::ItemWriter(error.to_string())),
}
}
fn flush(&self) -> ItemWriterResult {
let result = self.stream.borrow_mut().flush();
match result {
Ok(()) => Ok(()),
Err(error) => Err(BatchError::ItemWriter(error.to_string())),
}
}
fn open(&self) -> ItemWriterResult {
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 close(&self) -> ItemWriterResult {
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())),
}
}
}
pub struct JsonItemWriterBuilder<O> {
indent: Box<[u8]>,
pretty_formatter: bool,
_pd: PhantomData<O>,
}
impl<O> Default for JsonItemWriterBuilder<O> {
fn default() -> Self {
Self::new()
}
}
impl<O> JsonItemWriterBuilder<O> {
pub fn new() -> Self {
Self {
indent: Box::from(b" ".to_vec()),
pretty_formatter: false,
_pd: PhantomData,
}
}
pub fn indent(mut self, indent: &[u8]) -> Self {
self.indent = Box::from(indent);
self
}
pub fn pretty_formatter(mut self, yes: bool) -> Self {
self.pretty_formatter = yes;
self
}
pub fn from_path<W: AsRef<Path>>(self, path: W) -> JsonItemWriter<O, 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,
indent: self.indent.clone(),
is_first_element: Cell::new(true),
_phantom: PhantomData,
}
}
pub fn from_writer<W: Write>(self, wtr: W) -> JsonItemWriter<O, W> {
let buf_writer = BufWriter::new(wtr);
JsonItemWriter {
stream: RefCell::new(buf_writer),
use_pretty_formatter: self.pretty_formatter,
indent: self.indent,
is_first_element: Cell::new(true),
_phantom: PhantomData,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::item::ItemWriter;
use serde::Serialize;
use std::fs;
use tempfile::tempdir;
#[derive(Serialize, Debug, PartialEq)]
struct TestItem {
id: u32,
name: String,
value: f64,
}
#[test]
fn json_writer_builder_should_create_with_defaults() {
let builder = JsonItemWriterBuilder::<TestItem>::new();
assert!(!builder.pretty_formatter);
assert_eq!(builder.indent, b" ".to_vec().into_boxed_slice());
}
#[test]
fn json_writer_builder_should_set_pretty_formatter() {
let builder = JsonItemWriterBuilder::<TestItem>::new().pretty_formatter(true);
assert!(builder.pretty_formatter);
}
#[test]
fn json_writer_builder_should_set_custom_indent() {
let custom_indent = b" ";
let builder = JsonItemWriterBuilder::<TestItem>::new().indent(custom_indent);
assert_eq!(builder.indent, custom_indent.to_vec().into_boxed_slice());
}
#[test]
fn json_writer_builder_should_implement_default() {
let builder1 = JsonItemWriterBuilder::<TestItem>::new();
let builder2 = JsonItemWriterBuilder::<TestItem>::default();
assert_eq!(builder1.pretty_formatter, builder2.pretty_formatter);
assert_eq!(builder1.indent, builder2.indent);
}
#[test]
fn json_writer_builder_should_support_generic_type() {
let _builder = JsonItemWriterBuilder::<TestItem>::new();
let _builder_default = JsonItemWriterBuilder::<TestItem>::default();
}
#[test]
fn json_writer_from_path_should_create_file_writer() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("test_output.json");
let writer: JsonItemWriter<TestItem, File> =
JsonItemWriterBuilder::new().from_path(&file_path);
let item = TestItem {
id: 1,
name: "test".to_string(),
value: 42.5,
};
writer.open().unwrap();
writer.write(&[item]).unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert!(content.contains(r#"{"id":1,"name":"test","value":42.5}"#));
}
#[test]
fn json_writer_should_handle_custom_indent() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("indent_test.json");
let writer = JsonItemWriterBuilder::new()
.pretty_formatter(true)
.indent(b"\t")
.from_path(&file_path);
let item = TestItem {
id: 1,
name: "test".to_string(),
value: 42.5,
};
writer.open().unwrap();
writer.write(&[item]).unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert!(content.contains('\t'));
}
#[test]
fn json_writer_should_handle_pretty_formatting() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("pretty_test.json");
let writer = JsonItemWriterBuilder::new()
.pretty_formatter(true)
.from_path(&file_path);
let item = TestItem {
id: 1,
name: "test".to_string(),
value: 42.5,
};
writer.open().unwrap();
writer.write(&[item]).unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert!(content.contains("[\n"));
assert!(content.contains("\n]\n"));
assert!(content.contains(" \"id\": 1"));
}
#[test]
fn json_writer_should_handle_empty_items() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("empty_test.json");
let writer = JsonItemWriterBuilder::new().from_path(&file_path);
let empty_items: Vec<TestItem> = vec![];
writer.open().unwrap();
writer.write(&empty_items).unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert_eq!(content, "[]\n");
}
#[test]
fn json_writer_should_handle_multiple_writes() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("multi_test.json");
let writer = JsonItemWriterBuilder::new().from_path(&file_path);
let item1 = TestItem {
id: 1,
name: "first".to_string(),
value: 10.0,
};
let item2 = TestItem {
id: 2,
name: "second".to_string(),
value: 20.0,
};
writer.open().unwrap();
writer.write(&[item1]).unwrap();
writer.write(&[item2]).unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert!(content.contains(r#"{"id":1,"name":"first","value":10.0}"#));
assert!(content.contains(r#"{"id":2,"name":"second","value":20.0}"#));
assert!(content.contains(','));
}
#[test]
fn json_writer_should_write_to_in_memory_buffer() {
use std::io::Cursor;
let buf = Cursor::new(Vec::new());
let writer = JsonItemWriterBuilder::<TestItem>::new().from_writer(buf);
let item = TestItem {
id: 7,
name: "cursor".to_string(),
value: 0.5,
};
writer.open().unwrap();
writer.write(&[item]).unwrap();
writer.close().unwrap();
}
#[test]
fn json_writer_should_flush_without_error() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("flush_test.json");
let writer = JsonItemWriterBuilder::<TestItem>::new().from_path(&file_path);
writer.open().unwrap();
writer.flush().unwrap();
writer.close().unwrap();
}
#[test]
fn json_writer_compact_open_writes_bracket_without_newline() {
let temp_dir = tempdir().unwrap();
let file_path = temp_dir.path().join("compact_open.json");
let writer = JsonItemWriterBuilder::<TestItem>::new()
.pretty_formatter(false)
.from_path(&file_path);
writer.open().unwrap();
writer.close().unwrap();
let content = fs::read_to_string(&file_path).unwrap();
assert_eq!(
content, "[]\n",
"compact format should produce []\\n, got: {content:?}"
);
}
struct FailWriter;
impl std::io::Write for FailWriter {
fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
Err(std::io::Error::new(
std::io::ErrorKind::Other,
"write failed",
))
}
fn flush(&mut self) -> std::io::Result<()> {
Err(std::io::Error::new(
std::io::ErrorKind::Other,
"flush failed",
))
}
}
fn fail_json_writer<O: Serialize>() -> JsonItemWriter<O, FailWriter> {
JsonItemWriter {
stream: RefCell::new(BufWriter::with_capacity(0, FailWriter)),
use_pretty_formatter: false,
indent: Box::from(b" ".as_slice()),
is_first_element: Cell::new(true),
_phantom: PhantomData,
}
}
#[test]
fn should_return_error_when_open_fails_on_io() {
let writer = fail_json_writer::<String>();
let result = ((&writer) as &dyn ItemWriter<String>).open();
assert!(result.is_err(), "open should fail when writer fails");
}
#[test]
fn should_return_error_when_close_fails_on_io() {
let writer = fail_json_writer::<String>();
let result = ((&writer) as &dyn ItemWriter<String>).close();
assert!(result.is_err(), "close should fail when writer fails");
}
#[test]
fn should_return_error_when_flush_fails_on_io() {
let writer = fail_json_writer::<String>();
let result = ((&writer) as &dyn ItemWriter<String>).flush();
assert!(result.is_err(), "flush should fail when writer fails");
}
#[test]
fn should_return_error_when_write_fails_on_io() {
let writer = fail_json_writer::<String>();
let result = ((&writer) as &dyn ItemWriter<String>).write(&["hello".to_string()]);
assert!(
result.is_err(),
"write should fail when underlying IO fails"
);
}
#[test]
fn should_return_error_when_serialization_fails_with_pretty_formatter() {
use crate::BatchError;
struct NonSerializable;
impl Serialize for NonSerializable {
fn serialize<S: serde::Serializer>(&self, _s: S) -> Result<S::Ok, S::Error> {
Err(serde::ser::Error::custom(
"intentional serialization failure",
))
}
}
let buf = std::io::Cursor::new(Vec::new());
let writer = JsonItemWriterBuilder::<NonSerializable>::new()
.pretty_formatter(true)
.from_writer(buf);
let result = ((&writer) as &dyn ItemWriter<NonSerializable>).write(&[NonSerializable]);
match result.err().unwrap() {
BatchError::ItemWriter(msg) => assert!(
msg.contains("intentional"),
"error should contain serialization message, got: {msg}"
),
e => panic!("expected ItemWriter error, got {e:?}"),
}
}
}