From 4c1d1cd6d0f9713dd4b86af3e8e6325ade2d1538 Mon Sep 17 00:00:00 2001
From: Joey Hines <joey@ahines.net>
Date: Fri, 17 Dec 2021 17:24:51 -0700
Subject: [PATCH] Added a remove method to GeoffreyDB

+ clippy + fmt
---
 geoffrey_db/src/database.rs | 42 ++++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/geoffrey_db/src/database.rs b/geoffrey_db/src/database.rs
index b608e46..6648e86 100644
--- a/geoffrey_db/src/database.rs
+++ b/geoffrey_db/src/database.rs
@@ -2,12 +2,22 @@ use crate::error::{GeoffreyDBError, Result};
 use crate::migration::do_migration;
 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 = 3;
 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,
 }
@@ -66,14 +76,7 @@ impl Database {
         T: GeoffreyDatabaseModel,
     {
         let tree = self.get_tree::<T>()?;
-        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)
-        }
+        option_bytes_to_model(tree.get(id.to_be_bytes())?, id)
     }
 
     pub fn clear_tree<T>(&self) -> Result<()>
@@ -110,6 +113,14 @@ impl Database {
         }))
     }
 
+    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,
@@ -228,6 +239,21 @@ mod tests {
         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();