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

use my_public_ip_lib::{PublicIp, Reader, Writer};

use crate::error::Result;

#[derive(Clone)]
pub struct Store {
    db: sled::Db,
    writer_tree: sled::Tree,
    reader_tree: sled::Tree,
}

impl Store {
    pub fn new(db: sled::Db) -> Result<Store> {
        let writer_tree = db.open_tree("writer")?;
        let reader_tree = db.open_tree("reader")?;
        Ok(Store {
            db,
            writer_tree,
            reader_tree,
        })
    }

    pub fn open(db_dir: &str) -> Result<Store> {
        let db = sled::open(&db_dir)?;
        Store::new(db)
    }

    pub fn get_writer(&self, name: &str) -> Result<Option<Writer>> {
        let val = self.writer_tree.get(name)?;
        val.map(|val| serde_json::from_slice(val.as_ref()))
            .transpose()
            .map_err(Into::into)
    }

    pub fn set_writer(&self, name: &str, writer: &Writer) -> Result<Option<Writer>> {
        let writer_str = serde_json::to_string(writer)?;

        let last_writer = self.writer_tree.insert(name, writer_str.as_bytes())?;

        last_writer
            .map(|val| serde_json::from_slice(val.as_ref()))
            .transpose()
            .map_err(Into::into)
    }

    pub fn get_reader(&self, name: &str) -> Result<Option<Reader>> {
        let val = self.reader_tree.get(name)?;
        val.map(|val| serde_json::from_slice(val.as_ref()))
            .transpose()
            .map_err(Into::into)
    }

    pub fn set_reader(&self, name: &str, reader: &Reader) -> Result<Option<Reader>> {
        let reader_str = serde_json::to_string(reader)?;

        let last_reader = self.reader_tree.insert(name, reader_str.as_bytes())?;

        last_reader
            .map(|val| serde_json::from_slice(val.as_ref()))
            .transpose()
            .map_err(Into::into)
    }

    pub fn list_writer(&self) -> Result<Vec<PublicIp>> {
        let iter = self.writer_tree.iter();

        let (lower, upper) = iter.size_hint();
        let mut public_ips = Vec::with_capacity(upper.unwrap_or(lower));

        for row in iter {
            let (k, v) = row?;
            let writer: Writer = serde_json::from_slice(v.as_ref())?;
            public_ips.push(PublicIp {
                name: String::from_utf8_lossy(k.as_ref()).into_owned(),
                ip: writer.ip,
                updated_at: writer.updated_at,
            });
        }

        Ok(public_ips)
    }

    pub fn flush(&self) -> Result<()> {
        self.db.flush().map(|_| ()).map_err(Into::into)
    }
}