Geoffrey-rs/geoffrey_db/src/database.rs

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<Self> {
        let db = sled::open(db_path)?;

        Ok(Self { db })
    }

    fn get_tree<T>(&self) -> Result<sled::Tree>
    where
        T: GeoffreyDatabaseModel,
    {
        Ok(self.db.open_tree::<String>(T::tree())?)
    }

    pub fn insert<T>(&self, mut model: T) -> Result<T>
    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 {
            return Err(GeoffreyDBError::NotUnique);
        }

        let tree = self.get_tree::<T>()?;
        let id_bytes = id.to_be_bytes();

        tree.insert(id_bytes, model.to_bytes()?)?;

        Ok(model)
    }

    pub fn get<T>(&self, id: u64) -> Result<Option<T>>
    where
        T: GeoffreyDatabaseModel,
    {
        let tree = self.get_tree::<T>()?;
        let id_bytes = id.to_be_bytes();

        if let Some(bytes) = tree.get(id_bytes)? {
            Ok(Some(T::try_from_bytes(&bytes)?))
        } else {
            Ok(None)
        }
    }

    pub fn clear_tree<T>(&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<impl Iterator<Item = T> + '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<T>(&self) -> Result<sled::Iter>
    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::{Location, LocationType};
    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;

    lazy_static! {
        static ref DB: Database = Database::new(Path::new("../test_database")).unwrap();
    }

    fn cleanup() {
        DB.clear_tree::<Player>().unwrap();
        DB.clear_tree::<Location>().unwrap();
        DB.db.flush().unwrap();
    }

    #[test]
    fn test_insert() {
        let player = Player::new("CoolZero123", UserID::DiscordUUID(0u64));

        let p2 = DB.insert::<Player>(player.clone()).unwrap();

        assert!(p2.id().is_some());
        assert_eq!(player.name, p2.name);

        cleanup();
    }

    #[test]
    fn test_unique_insert() {
        let player1 = Player::new("CoolZero123", UserID::DiscordUUID(0u64));
        let player2 = Player::new("CoolZero123", UserID::DiscordUUID(0u64));

        DB.insert::<Player>(player1.clone()).unwrap();
        assert_eq!(DB.insert::<Player>(player2.clone()).is_err(), true);

        cleanup();
    }

    #[test]
    fn test_get() {
        let player = Player::new("CoolZero123", UserID::DiscordUUID(0u64));

        let p2 = DB.insert::<Player>(player.clone()).unwrap();

        let p3 = DB.get::<Player>(p2.id().unwrap()).unwrap();

        assert!(p3.is_some());
        assert_eq!(p3.unwrap().name, player.name);
        cleanup();
    }

    #[test]
    fn test_filter() {
        let player = Player::new("CoolZero123", UserID::DiscordUUID(0u64));
        let player = DB.insert::<Player>(player.clone()).unwrap();
        let loc = Location::new(
            "Test Shop",
            Position::new(0, 0, Dimension::Overworld),
            player.id.unwrap(),
            None,
            LocationType::Shop(0),
        );
        let loc = DB.insert::<Location>(loc.clone()).unwrap();

        let count = DB
            .filter(|id: u64, l: &Location| {
                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(i));

            DB.insert::<Player>(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()
    }
}