nickel_postgres/
middleware.rs

1use std::error::Error;
2use std::result::Result;
3use nickel::{Request, Response, Middleware, Continue, MiddlewareResult};
4use nickel::status::StatusCode;
5use r2d2_postgres::{PostgresConnectionManager, SslMode};
6use r2d2::{Config, Pool, PooledConnection, GetTimeout};
7use typemap::Key;
8use plugin::Extensible;
9
10pub struct PostgresMiddleware {
11    pub pool: Pool<PostgresConnectionManager>,
12}
13
14impl PostgresMiddleware {
15    /// Create middleware using defaults
16    ///
17    /// The middleware will be setup with no ssl and the r2d2 defaults.
18    pub fn new(db_url: &str) -> Result<PostgresMiddleware, Box<Error>> {
19        let manager = try!(PostgresConnectionManager::new(db_url, SslMode::None));
20        let pool = try!(Pool::new(Config::default(), manager));
21
22        Ok(PostgresMiddleware { pool: pool })
23    }
24
25    /// Create middleware using pre-built `r2d2::Pool`
26    ///
27    /// This allows the caller to create and configure the pool with specific settings.
28    pub fn with_pool(pool: Pool<PostgresConnectionManager>) -> PostgresMiddleware {
29        PostgresMiddleware { pool: pool }
30    }
31}
32
33impl Key for PostgresMiddleware { type Value = Pool<PostgresConnectionManager>; }
34
35impl<D> Middleware<D> for PostgresMiddleware {
36    fn invoke<'mw, 'conn>(&self, req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D>) -> MiddlewareResult<'mw, D> {
37        req.extensions_mut().insert::<PostgresMiddleware>(self.pool.clone());
38
39        Ok(Continue(res))
40    }
41}
42
43/// Add `pg_conn()` helper method to `nickel::Request`
44///
45/// This trait must only be used in conjunction with `PostgresMiddleware`.
46///
47/// On error, the method returns a tuple per Nickel convention. This allows the route to use the
48/// `try_with!` macro.
49///
50/// Example:
51///
52/// ```ignore
53/// app.get("/my_counter", middleware! { |request, response|
54/// 	let db = try_with!(response, request.pg_conn());
55/// });
56/// ```
57pub trait PostgresRequestExtensions {
58    fn pg_conn(&self) -> Result<PooledConnection<PostgresConnectionManager>, (StatusCode, GetTimeout)>;
59}
60
61impl<'a, 'b, D> PostgresRequestExtensions for Request<'a, 'b, D> {
62    fn pg_conn(&self) -> Result<PooledConnection<PostgresConnectionManager>, (StatusCode, GetTimeout)> {
63        self.extensions()
64            .get::<PostgresMiddleware>()
65            .expect("PostgresMiddleware must be registered before using PostgresRequestExtensions::pg_conn()")
66            .get()
67            .or_else(|err| Err((StatusCode::InternalServerError, err)))
68    }
69}