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
use std::{
    fs::File,
    io::{BufReader, BufWriter, Read, Write},
    path::PathBuf,
};

use csv::WriterBuilder;
use tempdir::TempDir;
use zip::ZipWriter;

use crate::{csv_maker::CsvMaker, CsvCustomizer, CsvZipError};

pub struct CsvZipMaker {
    tempdir: TempDir,
    writer: ZipWriter<BufWriter<File>>,
    file_path: PathBuf,
}

impl CsvZipMaker {
    pub fn new(prefix: &str, name: &str) -> Result<Self, CsvZipError> {
        let tempdir = TempDir::new(prefix)?;
        let file_path = tempdir.path().join(&format!("{}.zip", name));
        let buf_writer = BufWriter::new(File::create(file_path.clone())?);
        let writer = ZipWriter::new(buf_writer);
        Ok(Self {
            tempdir,
            writer,
            file_path,
        })
    }

    pub fn make_csv_maker(
        &self,
        name: &str,
        customizer: Option<Box<dyn CsvCustomizer>>,
    ) -> Result<CsvMaker, CsvZipError> {
        let file_name = format!("{}.csv", name);
        let file_path = self.tempdir.path().join(&file_name);
        let mut buf_writer = BufWriter::new(File::create(file_path.clone())?);
        let mut writer_builder = WriterBuilder::new();

        if let Some(customizer) = customizer {
            customizer.customize(&mut buf_writer, &mut writer_builder)?;
        }

        Ok(CsvMaker {
            writer: writer_builder.from_writer(buf_writer),
            file_name,
            file_path,
        })
    }

    pub fn add_csv(&mut self, csv_maker: &mut CsvMaker) -> Result<(), CsvZipError> {
        csv_maker.flush()?;

        self.writer
            .start_file(&csv_maker.file_name, Default::default())?;
        let mut f = BufReader::new(File::open(csv_maker.file_path.clone())?);
        let mut buf = [0; 1024];
        loop {
            match f.read(&mut buf)? {
                0 => break,
                n => {
                    self.writer.write(&buf[0..n])?;
                }
            }
        }
        Ok(())
    }

    pub fn make_zip_file(&mut self) -> Result<&PathBuf, CsvZipError> {
        let mut res = self.writer.finish()?;
        res.flush()?;
        Ok(&self.file_path)
    }

    pub fn make_zip_binary(&mut self) -> Result<Vec<u8>, CsvZipError> {
        let mut res = self.writer.finish()?;
        res.flush()?;
        let mut file = File::open(&self.file_path)?;
        let mut buf = Vec::new();
        let _ = file.read_to_end(&mut buf)?;
        Ok(buf)
    }
}