maxminddb_writer/
lib.rs

1use paths::IntoBitPath;
2use serde::Serialize;
3
4pub(crate) mod data;
5pub mod metadata;
6pub(crate) mod node;
7pub mod paths;
8pub(crate) mod serializer;
9
10#[derive(Debug, Default)]
11pub struct Database {
12    nodes: node::NodeTree,
13    data: data::Datastore,
14    pub metadata: metadata::Metadata,
15}
16
17impl Database {
18    fn update_size(&mut self) {
19        // make sure we have correct node count
20        let node_count = self.nodes.len();
21        self.metadata.node_count = node_count.try_into().unwrap();
22
23        // update record size if needed
24        let data_size = self.data.len();
25        let max_ptr_value = node_count + data_size + 16;
26        self.metadata.record_size = metadata::RecordSize::choose(max_ptr_value);
27    }
28
29    pub fn insert_value<T: serde::Serialize>(
30        &mut self,
31        value: T,
32    ) -> Result<data::DataRef, serializer::Error> {
33        let result = self.data.insert(value);
34        self.update_size();
35        result
36    }
37
38    pub fn insert_node(&mut self, path: impl IntoBitPath, data: data::DataRef) {
39        self.nodes.insert(path, data);
40        self.update_size();
41    }
42
43    pub fn write_to<W: std::io::Write>(&self, writer: W) -> Result<W, serializer::Error> {
44        // write node tree
45        let mut writer = self.nodes.write_to(writer, self.metadata.record_size)?;
46        // write data section separator
47        writer.write_all(&[0u8; 16])?;
48        // write data section
49        writer.write_all(self.data.serialized_data())?;
50        // write metadata marker
51        writer.write_all(metadata::METADATA_START_MARKER)?;
52        // serialize metadata
53        let mut serializer = serializer::Serializer::new(writer);
54        self.metadata.serialize(&mut serializer)?;
55        // all done
56        Ok(serializer.into_inner())
57    }
58
59    #[cfg(test)]
60    pub(crate) fn to_vec(&self) -> Result<Vec<u8>, serializer::Error> {
61        let mut result = Vec::new();
62        self.write_to(&mut result)?;
63        Ok(result)
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use crate::paths::IpAddrWithMask;
70
71    use super::*;
72
73    #[test]
74    fn test_simple() {
75        let mut db = Database::default();
76        let data_42 = db.insert_value(42u32).unwrap();
77        let data_foo = db.insert_value("foo".to_string()).unwrap();
78        db.insert_node("0.0.0.0/16".parse::<IpAddrWithMask>().unwrap(), data_42);
79        db.insert_node("1.0.0.0/16".parse::<IpAddrWithMask>().unwrap(), data_foo);
80        let raw_db = db.to_vec().unwrap();
81
82        let reader = maxminddb::Reader::from_source(&raw_db).unwrap();
83        let expected_data_42: u32 = reader.lookup([0, 0, 0, 0].into()).unwrap();
84        let expected_data_foo: &str = reader.lookup([1, 0, 0, 0].into()).unwrap();
85
86        assert_eq!(expected_data_42, 42);
87        assert_eq!(expected_data_foo, "foo");
88    }
89}
maxminddb_writer/lib.rs

maxminddb_writer/
lib.rs
