1 files changed, 38 insertions, 50 deletions
diff --git a/src/plugins/factoids/database.rs b/src/plugins/factoids/database.rs
index 88aa0fd..b1fe8dd 100644
--- a/src/plugins/factoids/database.rs
+++ b/src/plugins/factoids/database.rs
@@ -13,13 +13,12 @@ use diesel::mysql::MysqlConnection;
 use r2d2::Pool;
 #[cfg(feature = "mysql")]
 use r2d2_diesel::ConnectionManager;
+#[cfg(feature = "mysql")]
+use failure::ResultExt;
 
 use chrono::NaiveDateTime;
 
-pub enum DbResponse {
-    Success,
-    Failed(&'static str),
-}
+use super::error::*;
 
 #[cfg_attr(feature = "mysql", derive(Queryable))]
 #[derive(Clone, Debug)]
@@ -42,15 +41,15 @@ pub struct NewFactoid<'a> {
 }
 
 pub trait Database: Send {
-    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse;
-    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid>;
-    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse;
-    fn count_factoids(&self, name: &str) -> Result<i32, &'static str>;
+    fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError>;
+    fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError>;
+    fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError>;
+    fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError>;
 }
 
 // HashMap
 impl Database for HashMap<(String, i32), Factoid> {
-    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
+    fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError> {
         let factoid = Factoid {
             name: String::from(factoid.name),
             idx: factoid.idx,
@@ -61,23 +60,25 @@ impl Database for HashMap<(String, i32), Factoid> {
 
         let name = factoid.name.clone();
         match self.insert((name, factoid.idx), factoid) {
-            None => DbResponse::Success,
-            Some(_) => DbResponse::Failed("Factoid was overwritten"),
+            None => Ok(()),
+            Some(_) => Err(ErrorKind::Duplicate)?,
         }
     }
 
-    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
-        self.get(&(String::from(name), idx)).cloned()
+    fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError> {
+        Ok(self.get(&(String::from(name), idx))
+            .cloned()
+            .ok_or(ErrorKind::NotFound)?)
     }
 
-    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
+    fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError> {
         match self.remove(&(String::from(name), idx)) {
-            Some(_) => DbResponse::Success,
-            None => DbResponse::Failed("Factoid not found"),
+            Some(_) => Ok(()),
+            None => Err(ErrorKind::NotFound)?,
         }
     }
 
-    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
+    fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError> {
         Ok(self.iter().filter(|&(&(ref n, _), _)| n == name).count() as i32)
     }
 }
@@ -102,38 +103,31 @@ use self::schema::factoids;
 
 #[cfg(feature = "mysql")]
 impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
-    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
+    fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError> {
         use diesel;
 
-        let conn = &*self.get().expect("Failed to get connection");
-        match diesel::insert_into(factoids::table)
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
+        diesel::insert_into(factoids::table)
             .values(factoid)
             .execute(conn)
-        {
-            Ok(_) => DbResponse::Success,
-            Err(e) => {
-                error!("DB Insertion Error: {}", e);
-                DbResponse::Failed("Failed to add factoid")
-            }
-        }
+            .context(ErrorKind::MysqlError)?;
+
+        Ok(())
     }
 
-    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
-        let conn = &*self.get().expect("Failed to get connection");
-        match factoids::table.find((name, idx)).first(conn) {
-            Ok(f) => Some(f),
-            Err(e) => {
-                error!("DB Count Error: {}", e);
-                None
-            }
-        }
+    fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError> {
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
+        Ok(factoids::table
+            .find((name, idx))
+            .first(conn)
+            .context(ErrorKind::MysqlError)?)
     }
 
-    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
+    fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError> {
         use diesel;
         use self::factoids::columns;
 
-        let conn = &*self.get().expect("Failed to get connection");
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
         match diesel::delete(
             factoids::table
                 .filter(columns::name.eq(name))
@@ -142,22 +136,19 @@ impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
         {
             Ok(v) => {
                 if v > 0 {
-                    DbResponse::Success
+                    Ok(())
                 } else {
-                    DbResponse::Failed("Could not find any factoid with that name")
+                    Err(ErrorKind::NotFound)?
                 }
             }
-            Err(e) => {
-                error!("DB Deletion Error: {}", e);
-                DbResponse::Failed("Failed to delete factoid")
-            }
+            Err(e) => Err(e).context(ErrorKind::MysqlError)?,
         }
     }
 
-    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
+    fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError> {
         use diesel;
 
-        let conn = &*self.get().expect("Failed to get connection");
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
         let count: Result<i64, _> = factoids::table
             .filter(factoids::columns::name.eq(name))
             .count()
@@ -166,10 +157,7 @@ impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
         match count {
             Ok(c) => Ok(c as i32),
             Err(diesel::NotFound) => Ok(0),
-            Err(e) => {
-                error!("DB Count Error: {}", e);
-                Err("Database Error")
-            }
+            Err(e) => Err(e).context(ErrorKind::MysqlError)?,
         }
     }
 }