Quote: Return random quote on missing user arg

2 files changed, 82 insertions, 25 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)?,
+        }
+    }
 }
diff --git a/src/plugins/quote/mod.rs b/src/plugins/quote/mod.rs
index edeed40..c6e33fe 100644
--- a/src/plugins/quote/mod.rs
+++ b/src/plugins/quote/mod.rs
@@ -44,7 +44,7 @@ impl<T: Database, C: Client> Quote<T, C> {
         content: &str,
         author: &str,
     ) -> Result<&str, QuoteError> {
-        let count = self.quotes.read().count_quotes(quotee, channel)?;
+        let count = self.quotes.read().count_user_quotes(quotee, channel)?;
         let tm = time::now().to_timespec();
 
         let quote = database::NewQuote {
@@ -80,13 +80,13 @@ impl<T: Database, C: Client> Quote<T, C> {
     }
 
     fn get(&self, command: &PluginCommand) -> Result<String, QuoteError> {
-        if command.tokens.is_empty() {
-            Err(ErrorKind::InvalidCommand)?;
-        }
-
-        let quotee = &command.tokens[0];
+        let quotee = &command.tokens.get(0);
         let channel = &command.target;
-        let count = self.quotes.read().count_quotes(quotee, channel)?;
+        let count = if let Some(quotee) = quotee {
+            self.quotes.read().count_user_quotes(quotee, channel)?
+        } else {
+            self.quotes.read().count_channel_quotes(channel)?
+        };
 
         if count < 1 {
             Err(ErrorKind::NotFound)?;
@@ -109,11 +109,19 @@ impl<T: Database, C: Client> Quote<T, C> {
             }
         };
 
-        let quote = self
-            .quotes
-            .read()
-            .get_quote(quotee, channel, idx)
-            .context(ErrorKind::NotFound)?;
+        let quote = if let Some(quotee) = quotee {
+            self
+                .quotes
+                .read()
+                .get_user_quote(quotee, channel, idx)
+                .context(ErrorKind::NotFound)?
+        } else {
+            self
+                .quotes
+                .read()
+                .get_channel_quote(channel, idx)
+                .context(ErrorKind::NotFound)?
+        };
 
         Ok(format!(
             "\"{}\" - {}[{}/{}]",
@@ -122,12 +130,22 @@ impl<T: Database, C: Client> Quote<T, C> {
     }
 
     fn info(&self, command: &PluginCommand) -> Result<String, QuoteError> {
-        match command.tokens.len() {
-            0 => Err(ErrorKind::InvalidCommand)?,
+        let tokens = command.tokens.iter().filter(|t| !t.is_empty()).collect::<Vec<_>>();
+        match tokens.len() {
+            0 => {
+                let channel = &command.target;
+                let count = self.quotes.read().count_channel_quotes(channel)?;
+
+                Ok(match count {
+                    0 => Err(ErrorKind::NotFound)?,
+                    1 => format!("1 quote was saved in {}", channel),
+                    _ => format!("{} quotes were saved in {}", channel, count),
+                })
+            }
             1 => {
                 let quotee = &command.tokens[0];
                 let channel = &command.target;
-                let count = self.quotes.read().count_quotes(quotee, channel)?;
+                let count = self.quotes.read().count_user_quotes(quotee, channel)?;
 
                 Ok(match count {
                     0 => Err(ErrorKind::NotFound)?,
@@ -141,7 +159,7 @@ impl<T: Database, C: Client> Quote<T, C> {
                 let idx = i32::from_str(&command.tokens[1]).context(ErrorKind::InvalidIndex)?;
 
                 let idx = if idx < 0 {
-                    self.quotes.read().count_quotes(quotee, channel)? + idx + 1
+                    self.quotes.read().count_user_quotes(quotee, channel)? + idx + 1
                 } else {
                     idx
                 };
@@ -149,7 +167,7 @@ impl<T: Database, C: Client> Quote<T, C> {
                 let quote = self
                     .quotes
                     .read()
-                    .get_quote(quotee, channel, idx)
+                    .get_user_quote(quotee, channel, idx)
                     .context(ErrorKind::NotFound)?;
 
                 Ok(format!(