Use r2d2 as a ConnectionPool for diesel

4 files changed, 55 insertions, 34 deletions

diff --git a/src/lib.rs b/src/lib.rs
index bd2e302..aeead16 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -32,6 +32,10 @@
 #[cfg(feature = "mysql")]
 #[macro_use]
 extern crate diesel;
+#[cfg(feature = "mysql")]
+extern crate r2d2;
+#[cfg(feature = "mysql")]
+extern crate r2d2_diesel;
 
 #[macro_use]
 extern crate frippy_derive;
diff --git a/src/main.rs b/src/main.rs
index cb4e384..6605693 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,10 @@ extern crate time;
 extern crate diesel_migrations;
 #[cfg(feature = "mysql")]
 extern crate diesel;
+#[cfg(feature = "mysql")]
+extern crate r2d2;
+#[cfg(feature = "mysql")]
+extern crate r2d2_diesel;
 
 #[macro_use]
 extern crate log;
@@ -119,13 +123,16 @@ fn main() {
         #[cfg(feature = "mysql")]
         {
             if let Some(url) = mysql_url {
-                use diesel;
-                use diesel::Connection;
-                match diesel::mysql::MysqlConnection::establish(url) {
-                    Ok(conn) => {
-                        match embedded_migrations::run(&conn) {
+                use diesel::MysqlConnection;
+                use r2d2;
+                use r2d2_diesel::ConnectionManager;
+
+                let manager = ConnectionManager::<MysqlConnection>::new(url.clone());
+                match r2d2::Pool::builder().build(manager) {
+                    Ok(pool) => {
+                        match embedded_migrations::run(&*pool.get().expect("Failed to get connection")) {
                             Ok(_) => {
-                                bot.add_plugin(plugins::Factoids::new(conn));
+                                bot.add_plugin(plugins::Factoids::new(pool));
                                 info!("Connected to MySQL server")
                             }
                             Err(e) => {
diff --git a/src/plugins/factoids/database.rs b/src/plugins/factoids/database.rs
index cb5ef29..f003419 100644
--- a/src/plugins/factoids/database.rs
+++ b/src/plugins/factoids/database.rs
@@ -5,9 +5,12 @@ use std::collections::HashMap;
 
 #[cfg(feature = "mysql")]
 use diesel::prelude::*;
-
 #[cfg(feature = "mysql")]
 use diesel::mysql::MysqlConnection;
+#[cfg(feature = "mysql")]
+use r2d2::Pool;
+#[cfg(feature = "mysql")]
+use r2d2_diesel::ConnectionManager;
 
 use chrono::NaiveDateTime;
 
@@ -40,15 +43,15 @@ pub struct NewFactoid<'a> {
 
 
 pub trait Database: Send {
-    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse;
-    fn get(&self, name: &str, idx: i32) -> Option<Factoid>;
-    fn delete(&mut self, name: &str, idx: i32) -> DbResponse;
-    fn count(&self, name: &str) -> Result<i32, &'static str>;
+    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>;
 }
 
 // HashMap
 impl Database for HashMap<(String, i32), Factoid> {
-    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
+    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
         let factoid = Factoid {
             name: String::from(factoid.name),
             idx: factoid.idx,
@@ -64,18 +67,18 @@ impl Database for HashMap<(String, i32), Factoid> {
         }
     }
 
-    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
+    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
         self.get(&(String::from(name), idx)).map(|f| f.clone())
     }
 
-    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
+    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
         match self.remove(&(String::from(name), idx)) {
             Some(_) => DbResponse::Success,
             None => DbResponse::Failed("Factoid not found"),
         }
     }
 
-    fn count(&self, name: &str) -> Result<i32, &'static str> {
+    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
         Ok(self.iter()
                .filter(|&(&(ref n, _), _)| n == name)
                .count() as i32)
@@ -98,12 +101,14 @@ mod mysql {
 }
 
 #[cfg(feature = "mysql")]
-impl Database for MysqlConnection {
-    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
+impl Database for Pool<ConnectionManager<MysqlConnection>> {
+    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
         use diesel;
+
+        let conn = &*self.get().expect("Failed to get connection");
         match diesel::insert_into(factoids::table)
                   .values(factoid)
-                  .execute(self) {
+                  .execute(conn) {
             Ok(_) => DbResponse::Success,
             Err(e) => {
                 error!("DB Insertion Error: {}", e);
@@ -112,8 +117,9 @@ impl Database for MysqlConnection {
         }
     }
 
-    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
-        match factoids::table.find((name, idx)).first(self) {
+    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);
@@ -122,13 +128,15 @@ impl Database for MysqlConnection {
         }
     }
 
-    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
+    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
         use diesel;
         use self::factoids::columns;
+
+        let conn = &*self.get().expect("Failed to get connection");
         match diesel::delete(factoids::table
                                  .filter(columns::name.eq(name))
                                  .filter(columns::idx.eq(idx)))
-                      .execute(self) {
+                      .execute(conn) {
             Ok(v) => {
                 if v > 0 {
                     DbResponse::Success
@@ -143,11 +151,13 @@ impl Database for MysqlConnection {
         }
     }
 
-    fn count(&self, name: &str) -> Result<i32, &'static str> {
+    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
+
+        let conn = &*self.get().expect("Failed to get connection");
         let count: Result<i64, _> = factoids::table
             .filter(factoids::columns::name.eq(name))
             .count()
-            .first(self);
+            .first(conn);
 
         match count {
             Ok(c) => Ok(c as i32),
diff --git a/src/plugins/factoids/mod.rs b/src/plugins/factoids/mod.rs
index 49ace10..b662d22 100644
--- a/src/plugins/factoids/mod.rs
+++ b/src/plugins/factoids/mod.rs
@@ -39,7 +39,7 @@ impl<T: Database> Factoids<T> {
     }
 
     fn create_factoid(&self, name: &str, content: &str, author: &str) -> Result<&str, &str> {
-            let count = try_lock!(self.factoids).count(&name)?;
+            let count = try_lock!(self.factoids).count_factoids(&name)?;
             let tm = time::now().to_timespec();
 
             let factoid = database::NewFactoid {
@@ -50,7 +50,7 @@ impl<T: Database> Factoids<T> {
                 created: NaiveDateTime::from_timestamp(tm.sec, tm.nsec as u32),
             };
 
-            match try_lock!(self.factoids).insert(&factoid) {
+            match try_lock!(self.factoids).insert_factoid(&factoid) {
                 DbResponse::Success => Ok("Successfully added"),
                 DbResponse::Failed(e) => Err(e),
             }
@@ -93,12 +93,12 @@ impl<T: Database> Factoids<T> {
         }
 
         let name = command.tokens.remove(0);
-        let count = match try_lock!(self.factoids).count(&name) {
+        let count = match try_lock!(self.factoids).count_factoids(&name) {
             Ok(c) => c,
             Err(e) => return client.send_notice(&command.source, e),
         };
 
-        match try_lock!(self.factoids).delete(&name, count - 1) {
+        match try_lock!(self.factoids).delete_factoid(&name, count - 1) {
             DbResponse::Success => client.send_notice(&command.source, "Successfully removed"),
             DbResponse::Failed(e) => client.send_notice(&command.source, &e),
         }
@@ -110,7 +110,7 @@ impl<T: Database> Factoids<T> {
             0 => return self.invalid_command(client, command),
             1 => {
                 let name = &command.tokens[0];
-                let count = match try_lock!(self.factoids).count(name) {
+                let count = match try_lock!(self.factoids).count_factoids(name) {
                     Ok(c) => c,
                     Err(e) => return client.send_notice(&command.source, e),
                 };
@@ -132,7 +132,7 @@ impl<T: Database> Factoids<T> {
             }
         };
 
-        let factoid = match try_lock!(self.factoids).get(name, idx) {
+        let factoid = match try_lock!(self.factoids).get_factoid(name, idx) {
             Some(v) => v,
             None => {
                 return client.send_notice(&command.source,
@@ -152,7 +152,7 @@ impl<T: Database> Factoids<T> {
             0 => self.invalid_command(client, command),
             1 => {
                 let name = &command.tokens[0];
-                let count = match try_lock!(self.factoids).count(name) {
+                let count = match try_lock!(self.factoids).count_factoids(name) {
                     Ok(c) => c,
                     Err(e) => return client.send_notice(&command.source, e),
                 };
@@ -176,7 +176,7 @@ impl<T: Database> Factoids<T> {
                     Err(_) => return client.send_notice(&command.source, "Invalid index"),
                 };
 
-                let factoid = match try_lock!(self.factoids).get(name, idx) {
+                let factoid = match try_lock!(self.factoids).get_factoid(name, idx) {
                     Some(v) => v,
                     None => {
                         return client.send_notice(&command.source,
@@ -204,12 +204,12 @@ impl<T: Database> Factoids<T> {
 
         } else {
             let name = command.tokens.remove(0);
-            let count = match try_lock!(self.factoids).count(&name) {
+            let count = match try_lock!(self.factoids).count_factoids(&name) {
                 Ok(c) => c,
                 Err(e) => return client.send_notice(&command.source, e),
             };
 
-            let factoid = match try_lock!(self.factoids).get(&name, count - 1) {
+            let factoid = match try_lock!(self.factoids).get_factoid(&name, count - 1) {
                 Some(v) => v.content,
                 None if error => return self.invalid_command(client, &command),
                 None => return Ok(()),