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

extern crate csv;
extern crate rustc_serialize;
extern crate walkdir;

use std::path::PathBuf;
use std::fs;
use std::fs::OpenOptions;
use std::io::Write;

pub struct CSVStore {
    datadir: PathBuf,
}

impl CSVStore {
    pub fn new(datadir: PathBuf) -> CSVStore {
        CSVStore { datadir: datadir }
    }

    pub fn write<V: rustc_serialize::Encodable>(&self, key: Vec<&str>, value: V) {
        let key_path = self.get_path_from_key(key);

        fs::create_dir_all(key_path.parent()
                .expect("can not get parent for key path while trying to write data"))
            .expect("can not create folder for writing data");


        let mut file = match OpenOptions::new().append(true).open(&key_path) {
            Ok(file) => file,
            Err(_) => {
                OpenOptions::new()
                    .append(true)
                    .create(true)
                    .open(&key_path)
                    .expect("can not open key file for writing")
            }
        };

        let mut wtr = csv::Writer::from_memory();
        wtr.encode(value).expect("can not serialize value into csv");
        file.write_all(wtr.as_bytes()).expect("can not write csv data to file");
    }

    pub fn get<V: rustc_serialize::Decodable>(&self, key: Vec<&str>) -> Vec<V> {
        let key_path = self.get_path_from_key(key);
        let mut rdr =
            csv::Reader::from_file(key_path).expect("read csv data from file").has_headers(false);

        let mut out = Vec::new();
        for record in rdr.decode() {
            let entry: V = record.unwrap();
            out.push(entry);
        }

        out
    }

    pub fn destroy(&self) {
        fs::remove_dir_all(&self.datadir).expect("can not remove data store folder")
    }

    pub fn get_path_from_key(&self, key: Vec<&str>) -> PathBuf {
        let mut key_path = self.datadir.clone();
        key_path.push(format!("{}.csv", key.join("/")));

        key_path
    }
}

#[test]
fn test_write_read() {
    let store = CSVStore::new(PathBuf::from("/tmp/csvstore_test"));
    let key = vec!["test"];
    let data = "test";

    store.write(key.clone(), data);

    let read_data: Vec<String> = store.get(key);

    assert_eq!(data, read_data[0]);

    store.destroy();
}