Quote: Return random quote on missing user arg

1 files changed, 47 insertions, 8 deletions

diff --git a/src/plugins/quote/database.rs b/src/plugins/quote/database.rs
index 49d6058..c8f550e 100644
--- a/src/plugins/quote/database.rs
+++ b/src/plugins/quote/database.rs
@@ -41,8 +41,10 @@ pub struct NewQuote<'a> {
 
 pub trait Database: Send + Sync {
     fn insert_quote(&mut self, quote: &NewQuote) -> Result<(), QuoteError>;
-    fn get_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError>;
-    fn count_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError>;
+    fn get_user_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError>;
+    fn get_channel_quote(&self, channel: &str, idx: i32) -> Result<Quote, QuoteError>;
+    fn count_user_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError>;
+    fn count_channel_quotes(&self, channel: &str) -> Result<i32, QuoteError>;
 }
 
 // HashMap
@@ -67,19 +69,34 @@ impl<S: ::std::hash::BuildHasher + Send + Sync> Database
         }
     }
 
-    fn get_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
+    fn get_user_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
         Ok(self
             .get(&(quotee.to_owned(), channel.to_owned(), idx))
             .cloned()
             .ok_or(ErrorKind::NotFound)?)
     }
 
-    fn count_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError> {
+    fn get_channel_quote(&self, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
+        Ok(self
+            .iter()
+            .filter(|&(&(_, ref c, _), _)| c == channel)
+            .nth(idx as usize - 1)
+            .ok_or(ErrorKind::NotFound)?.1.clone())
+    }
+
+    fn count_user_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError> {
         Ok(self
             .iter()
             .filter(|&(&(ref n, ref c, _), _)| n == quotee && c == channel)
             .count() as i32)
     }
+
+    fn count_channel_quotes(&self, channel: &str) -> Result<i32, QuoteError> {
+        Ok(self
+            .iter()
+            .filter(|&(&(_, ref c, _), _)| c == channel)
+            .count() as i32)
+    }
 }
 
 // Diesel automatically defines the quotes module as public.
@@ -104,7 +121,6 @@ use self::schema::quotes;
 #[cfg(feature = "mysql")]
 impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
     fn insert_quote(&mut self, quote: &NewQuote) -> Result<(), QuoteError> {
-        use diesel;
 
         let conn = &*self.get().context(ErrorKind::NoConnection)?;
         diesel::insert_into(quotes::table)
@@ -115,7 +131,7 @@ impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
         Ok(())
     }
 
-    fn get_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
+    fn get_user_quote(&self, quotee: &str, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
         let conn = &*self.get().context(ErrorKind::NoConnection)?;
         Ok(quotes::table
             .find((quotee, channel, idx))
@@ -123,8 +139,16 @@ impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
             .context(ErrorKind::MysqlError)?)
     }
 
-    fn count_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError> {
-        use diesel;
+    fn get_channel_quote(&self, channel: &str, idx: i32) -> Result<Quote, QuoteError> {
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
+        Ok(quotes::table
+            .filter(quotes::columns::channel.eq(channel))
+            .offset(idx as i64 - 1)
+            .first(conn)
+            .context(ErrorKind::MysqlError)?)
+    }
+
+    fn count_user_quotes(&self, quotee: &str, channel: &str) -> Result<i32, QuoteError> {
 
         let conn = &*self.get().context(ErrorKind::NoConnection)?;
         let count: Result<i64, _> = quotes::table
@@ -139,4 +163,19 @@ impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> {
             Err(e) => Err(e).context(ErrorKind::MysqlError)?,
         }
     }
+
+    fn count_channel_quotes(&self, channel: &str) -> Result<i32, QuoteError> {
+
+        let conn = &*self.get().context(ErrorKind::NoConnection)?;
+        let count: Result<i64, _> = quotes::table
+            .filter(quotes::columns::channel.eq(channel))
+            .count()
+            .get_result(conn);
+
+        match count {
+            Ok(c) => Ok(c as i32),
+            Err(diesel::NotFound) => Ok(0),
+            Err(e) => Err(e).context(ErrorKind::MysqlError)?,
+        }
+    }
 }