Geoffrey-rs/geoffrey_db/src/database.rs

use crate::error::{GeoffreyDBError, Result};
use crate::migration::do_migration;
use crate::query::QueryBuilder;
use geoffrey_models::models::db_metadata::DBMetadata;
use geoffrey_models::GeoffreyDatabaseModel;
use sled::IVec;
use std::convert::TryInto;
use std::path::Path;

const DB_VERSION: u64 = 4;
const DB_METADATA_ID: u64 = 1;

fn option_bytes_to_model<T: GeoffreyDatabaseModel>(bytes: Option<IVec>, id: u64) -> Result<T> {
    if let Some(bytes) = 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 struct Database {
    pub(crate) db: sled::Db,
}

impl Database {
    pub fn new(db_path: &Path, force_migration: bool) -> Result<Self> {
        let db = sled::open(db_path)?;
        let mut db = Self { db };
        let ver = db.version().unwrap_or(0);

        if force_migration && ver > 0 {
            db.set_version(ver - 1)?;
        }

        do_migration(&db, DB_VERSION)?;

        log::info!("Geoffrey Database V{}", db.version()?);

        Ok(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_id, o: &T| o_id != id && !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::<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<T>
    where
        T: GeoffreyDatabaseModel,
    {
        let tree = self.get_tree::<T>()?;
        option_bytes_to_model(tree.get(id.to_be_bytes())?, id)
    }

    pub fn clear_tree<T>(&self) -> Result<()>
    where
        T: GeoffreyDatabaseModel,
    {
        self.db.drop_tree(T::tree())?;

        Ok(())
    }

    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 = match T::try_from_bytes(&data) {
                    Ok(data) => data,
                    Err(err) => {
                        log::debug!(
                            "Invalid data: {}",
                            String::from_utf8(data.to_vec()).unwrap_or_default()
                        );
                        panic!("Unable to parse {} model from bytes: {}", T::tree(), err);
                    }
                };

                if f(id, &data) {
                    Some(data)
                } else {
                    None
                }
            } else {
                None
            }
        }))
    }

    pub fn run_query<T>(&self, query_builder: QueryBuilder<T>) -> Result<Vec<T>>
    where
        T: GeoffreyDatabaseModel,
    {
        let result: Vec<T> = self
            .filter(|id, loc: &T| {
                for query in &query_builder.queries {
                    let res = query(id, loc);

                    if !res {
                        return false;
                    }
                }

                true
            })?
            .collect();

        if result.is_empty() {
            Err(GeoffreyDBError::NotFound)
        } else {
            Ok(result)
        }
    }

    pub fn remove<T>(&self, id: u64) -> Result<T>
    where
        T: GeoffreyDatabaseModel,
    {
        let tree = self.db.open_tree(T::tree())?;
        option_bytes_to_model(tree.remove(id.to_be_bytes())?, id)
    }

    pub fn tree_iter<T>(&self) -> Result<sled::Iter>
    where
        T: GeoffreyDatabaseModel,
    {
        Ok(self.db.open_tree(T::tree()).map(|tree| tree.iter())?)
    }

    pub fn version(&self) -> Result<u64> {
        Ok(self.get::<DBMetadata>(DB_METADATA_ID)?.version)
    }

    pub(crate) fn set_version(&mut self, version: u64) -> Result<()> {
        let mut md = self.get::<DBMetadata>(DB_METADATA_ID)?;

        md.version = version;

        self.insert(md)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::test::{cleanup, DB, LOCK};
    use geoffrey_models::models::locations::shop::Shop;
    use geoffrey_models::models::locations::{LocationDataDb, LocationDb};
    use geoffrey_models::models::player::{Player, UserID};
    use geoffrey_models::models::{Dimension, Position};
    use geoffrey_models::GeoffreyDatabaseModel;
    use std::time::Instant;

    #[test]
    fn test_insert() {
        let _lock = LOCK.lock().unwrap();
        cleanup();
        let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 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 _lock = LOCK.lock().unwrap();
        cleanup();
        let player1 = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 });
        let player2 = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 });

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

        cleanup();
    }

    #[test]
    fn test_get() {
        let _lock = LOCK.lock().unwrap();
        cleanup();
        let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 });

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

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

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

    #[test]
    fn test_filter() {
        let _lock = LOCK.lock().unwrap();
        cleanup();
        let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 });
        let player = DB.insert::<Player>(player.clone()).unwrap();
        let loc = LocationDb::new(
            "Test Shop",
            Position::new(0, 0, 0, Dimension::Overworld),
            player.id.unwrap(),
            None,
            LocationDataDb::Shop(Shop {
                item_listings: Default::default(),
            }),
        );
        let loc = DB.insert::<LocationDb>(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_remove() {
        let _lock = LOCK.lock().unwrap();
        cleanup();
        let player = Player::new("CoolZero123", UserID::DiscordUUID { discord_uuid: 0u64 });

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

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

        assert!(DB.get::<Player>(p3.id.unwrap()).is_err());
        assert_eq!(p3.id.unwrap(), p2.id.unwrap());
        cleanup();
    }

    #[test]
    fn test_insert_speed() {
        let _lock = LOCK.lock().unwrap();
        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>(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()
    }
}