use crate::error::{GeoffreyDBError, Result}; use geoffrey_models::GeoffreyDatabaseModel; use std::convert::TryInto; use std::path::Path; pub struct Database { db: sled::Db, } impl Database { pub fn new(db_path: &Path) -> Result { let db = sled::open(db_path)?; Ok(Self { db }) } fn get_tree(&self) -> Result where T: GeoffreyDatabaseModel, { Ok(self.db.open_tree::(T::tree())?) } pub fn insert(&self, mut model: T) -> Result where T: GeoffreyDatabaseModel, { let id = match model.id() { Some(id) => id, None => { let id = self.db.generate_id()?; model.set_id(id); id } }; let match_count = self.filter(|_, o: &T| !o.check_unique(&model))?.count(); if match_count > 0 { log::debug!("{} is not unique: {:?}", T::tree(), model); return Err(GeoffreyDBError::NotUnique); } let tree = self.get_tree::()?; let id_bytes = id.to_be_bytes(); tree.insert(id_bytes, model.to_bytes()?)?; Ok(model) } pub fn get(&self, id: u64) -> Result where T: GeoffreyDatabaseModel, { let tree = self.get_tree::()?; let id_bytes = id.to_be_bytes(); if let Some(bytes) = tree.get(id_bytes)? { Ok(T::try_from_bytes(&bytes)?) } else { log::debug!("{} of id {} was not found in the database", T::tree(), id); Err(GeoffreyDBError::NotFound) } } pub fn clear_tree(&self) -> Result<()> where T: GeoffreyDatabaseModel, { Ok(self.db.open_tree(T::tree())?.clear()?) } pub fn filter<'a, T>( &self, f: impl Fn(u64, &T) -> bool + 'a, ) -> Result + 'a> where T: GeoffreyDatabaseModel, { let tree = self.db.open_tree(T::tree())?; Ok(tree.iter().filter_map(move |e| { if let Ok((id, data)) = e { let id = u64::from_be_bytes(id.to_vec().try_into().unwrap()); let data = T::try_from_bytes(&data).unwrap(); if f(id, &data) { Some(data) } else { None } } else { None } })) } pub fn tree_iter(&self) -> Result where T: GeoffreyDatabaseModel, { Ok(self.db.open_tree(T::tree()).map(|tree| tree.iter())?) } } #[cfg(test)] mod tests { use crate::database::Database; use geoffrey_models::models::locations::{LocationDb, LocationDataDb}; use geoffrey_models::models::player::{Player, UserID}; use geoffrey_models::models::{Dimension, Position}; use geoffrey_models::GeoffreyDatabaseModel; use lazy_static::lazy_static; use std::path::Path; use std::time::Instant; use geoffrey_models::models::locations::shop::Shop; lazy_static! { static ref DB: Database = Database::new(Path::new("../test_database")).unwrap(); } fn cleanup() { DB.clear_tree::().unwrap(); DB.clear_tree::().unwrap(); DB.db.clear().unwrap(); DB.db.flush().unwrap(); } #[test] fn test_insert() { cleanup(); let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 }); let p2 = DB.insert::(player.clone()).unwrap(); assert!(p2.id().is_some()); assert_eq!(player.name, p2.name); cleanup(); } #[test] fn test_unique_insert() { cleanup(); let player1 = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 }); let player2 = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 }); DB.insert::(player1.clone()).unwrap(); assert_eq!(DB.insert::(player2.clone()).is_err(), true); cleanup(); } #[test] fn test_get() { cleanup(); let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 }); let p2 = DB.insert::(player.clone()).unwrap(); let p3 = DB.get::(p2.id().unwrap()).unwrap(); assert_eq!(p3.name, player.name); cleanup(); } #[test] fn test_filter() { cleanup(); let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 }); let player = DB.insert::(player.clone()).unwrap(); let loc = LocationDb::new( "Test Shop", Position::new(0, 0, Dimension::Overworld), player.id.unwrap(), None, LocationDataDb::Shop(Shop { item_listings: Default::default() }), ); let loc = DB.insert::(loc.clone()).unwrap(); let count = DB .filter(|id: u64, l: &LocationDb| { assert_eq!(id, l.id().unwrap()); loc.id().unwrap() == id }) .unwrap() .count(); assert_eq!(count, 1); DB.db.flush().unwrap(); cleanup(); } #[test] fn test_insert_speed() { cleanup(); let insert_count = 1000; let timer = Instant::now(); for i in 0..insert_count { let player = Player::new("test", UserID::DiscordUUID { discord_uuid: i }); DB.insert::(player).unwrap(); } DB.db.flush().unwrap(); let sec_elapsed = timer.elapsed().as_secs_f32(); println!( "Completed in {}s. {} inserts per second", sec_elapsed, insert_count as f32 / sec_elapsed ); cleanup() } }