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
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105

extern crate iron;
extern crate r2d2;
extern crate r2d2_sqlite;
extern crate rusqlite;

use std::error::Error;
use std::sync::Arc;

use iron::prelude::*;
use iron::{typemap, BeforeMiddleware};
use std::path::{Path};

/// Pool of `SqliteConnectionManager` kept by the Iron middleware
pub type RusqlitePool = Arc<r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>>;

/// Pooled conenction to the SQLite database
pub type SqliteConnection = r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>;

/// Iron rusqlite middleware
pub struct RusqliteMiddleware {
    /// Pool of connections to SQLite through the rusqlite library
    pub pool: RusqlitePool
}

pub struct Value(RusqlitePool);

impl typemap::Key for RusqliteMiddleware { type Value = Value; }

impl RusqliteMiddleware {

    /// Creates a new pooled connection to the SQLite database using the default options `rusqlite`.
    /// The `path` should be the path to the SQLite database file on your system.
    ///
    /// See `rusqlite::Connection::open` for mode details.
    pub fn new<P: AsRef<Path>>(path: P) -> Result<RusqliteMiddleware, Box<Error>> {
        let config = r2d2::Config::default();
        let manager = r2d2_sqlite::SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager)?;

        Ok(RusqliteMiddleware{ pool: Arc::new(pool) })
    }

    /// Creates a new pooled connection to the SQLite database using the given rusqlite flags
    /// (i.e. `rusqlite::OpenFlags`). The `path` should be the path to the SQLite database file
    /// on your system.
    ///
    /// See `rusqlite::Connection::open_with_flags` for mode details.
    pub fn new_with_flags<P: AsRef<Path>>(path: P, flags: rusqlite::OpenFlags) -> Result<RusqliteMiddleware, Box<Error>> {
        let config = r2d2::Config::default();
        let manager = r2d2_sqlite::SqliteConnectionManager::new_with_flags(path, flags);
        let pool = r2d2::Pool::new(config, manager)?;

        Ok(RusqliteMiddleware{ pool: Arc::new(pool) })
    }

    /// Get a handle to a pooled connection for the SQLite database. This can be used to execute
    /// some SQL commands prior to launching your Iron webserver. An example would be creating
    /// tables if they do not currently exist in the database.
    pub fn get_connection(&self) -> SqliteConnection {
        let poll = self.pool.clone();
        poll.get().unwrap()
    }
}

/// Implementation of the `iron::BeforeMiddleware` trait to make this actually Iron middleware.
impl BeforeMiddleware for RusqliteMiddleware {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        // Insert the the value into the request extensions
        req.extensions.insert::<RusqliteMiddleware>(Value(self.pool.clone()));
        Ok(())
    }
}

/// Trait which adds a method the `iron::Request` to enable the retrieval of a database connection.
///
/// ### Example
///
/// ```ignore
/// use iron_rusqlite_middleware::RusqliteMiddlewareExtension;
///
/// fn handler(req: &mut Request) -> IronResult<Response> {
///     let conn = req.database_connection();
///
///     // Do stuff with the rusqlite::Connection object
///
///     Ok(Response::with((status::Ok, "Done.")))
/// }
/// ```
pub trait RusqliteRequestExtension {
    /// Returns a pooled connection to the SQLite database. The connection is automatically returned
    /// to the connection pool when the pool connection is dropped.
    ///
    /// **Panics** if the `RusqliteMiddleware` has not been added to the Iron app, or if retrieving
    /// the database connection times out.
    fn database_connection(&self) -> r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>;
}

/// Implementation of the `RusqliteRequestExention` for the `iron::Request`.
impl<'a, 'b> RusqliteRequestExtension for Request<'a, 'b> {
    fn database_connection(&self) -> SqliteConnection {
        let pv = self.extensions.get::<RusqliteMiddleware>().unwrap();  // Is this safe?
        let &Value(ref poll) = pv;
        poll.get().unwrap()
    }
}