1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
use std::any::Any; use std::sync::Arc; use std::error::Error as StdError; use nickel::{Request, Response, Middleware, Continue, MiddlewareResult}; use r2d2_diesel::{ConnectionManager}; use r2d2::{Pool, HandleError, PooledConnection}; use typemap::{Key}; use plugin::Extensible; use diesel::{Connection}; pub struct DieselMiddleware<T> where T: Connection + Send + Any { pub pool: Arc<Pool<ConnectionManager<T>>> } impl<T> DieselMiddleware<T> where T: Connection + Send + Any { pub fn new(connect_str: &str, num_connections: u32, error_handler: Box<HandleError<::r2d2_diesel::Error>>) -> Result<DieselMiddleware<T>, Box<StdError>> { let manager = ConnectionManager::<T>::new(connect_str); let pool = Pool::builder() .max_size(num_connections) .error_handler(error_handler) .build(manager)?; Ok(DieselMiddleware { pool: Arc::new(pool) }) } pub fn from_pool(pool: Pool<ConnectionManager<T>>) -> DieselMiddleware<T> { DieselMiddleware { pool: Arc::new(pool) } } } impl<T> Key for DieselMiddleware<T> where T: Connection + Send + Any { type Value = Arc<Pool<ConnectionManager<T>>>; } impl<T, D> Middleware<D> for DieselMiddleware<T> where T: Connection + Send + Any { fn invoke<'mw, 'conn>(&self, req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D>) -> MiddlewareResult<'mw, D> { req.extensions_mut().insert::<DieselMiddleware<T>>(self.pool.clone()); Ok(Continue(res)) } } pub trait DieselRequestExtensions<T> where T: Connection + Send + Any { fn db_conn(&self) -> PooledConnection<ConnectionManager<T>>; } impl<'a, 'b, T, D> DieselRequestExtensions<T> for Request<'a, 'b, D> where T: Connection + Send + Any { fn db_conn(&self) -> PooledConnection<ConnectionManager<T>> { self.extensions().get::<DieselMiddleware<T>>().unwrap().get().unwrap() } }