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::metadata::RecordSize;
70    use crate::paths::IpAddrWithMask;
71
72    use super::*;
73
74    fn seed_simple_db() -> Database {
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
81        db
82    }
83
84    #[test]
85    fn test_simple() {
86        let db = seed_simple_db();
87        let raw_db = db.to_vec().unwrap();
88
89        let reader = maxminddb::Reader::from_source(&raw_db).unwrap();
90        let expected_data_42: u32 = reader.lookup([0, 0, 0, 0].into()).unwrap();
91        let expected_data_foo: &str = reader.lookup([1, 0, 0, 0].into()).unwrap();
92
93        assert_eq!(expected_data_42, 42);
94        assert_eq!(expected_data_foo, "foo");
95    }
96
97    #[test]
98    fn test_small_record_write() {
99        let mut db = seed_simple_db();
100
101        // Force small record size to test the small node writing
102        db.metadata.record_size = RecordSize::Small;
103        let raw_db = db.to_vec().unwrap();
104
105        let reader = maxminddb::Reader::from_source(&raw_db).unwrap();
106        let expected_data: u32 = reader.lookup([0, 0, 0, 0].into()).unwrap();
107
108        assert_eq!(expected_data, 42);
109        assert!(matches!(db.metadata.record_size, RecordSize::Small));
110    }
111
112    #[test]
113    fn test_medium_record_write() {
114        let mut db = seed_simple_db();
115
116        // Force medium record size to test the medium node writing
117        db.metadata.record_size = RecordSize::Medium;
118        let raw_db = db.to_vec().unwrap();
119
120        let reader = maxminddb::Reader::from_source(&raw_db).unwrap();
121        let expected_data: u32 = reader.lookup([0, 0, 0, 0].into()).unwrap();
122
123        assert_eq!(expected_data, 42);
124        assert!(matches!(db.metadata.record_size, RecordSize::Medium));
125    }
126
127    #[test]
128    fn test_large_record_write() {
129        let mut db = seed_simple_db();
130
131        // Force large record size to test the large node writing
132        db.metadata.record_size = RecordSize::Large;
133        let raw_db = db.to_vec().unwrap();
134
135        let reader = maxminddb::Reader::from_source(&raw_db).unwrap();
136        let expected_data: u32 = reader.lookup([0, 0, 0, 0].into()).unwrap();
137
138        assert_eq!(expected_data, 42);
139        assert!(matches!(db.metadata.record_size, RecordSize::Large));
140    }
141}
maxminddb_writer/lib.rs

maxminddb_writer/
lib.rs
