r2d2_sqlite/
lib.rs

1#![deny(warnings)]
2//! # Sqlite support for the `r2d2` connection pool.
3//!
4//! Library crate: [r2d2-sqlite](https://crates.io/crates/r2d2-sqlite/)
5//!
6//! Integrated with: [r2d2](https://crates.io/crates/r2d2)
7//! and [rusqlite](https://crates.io/crates/rusqlite)
8//!
9//! ## Example
10//!
11//! ```rust,no_run
12//! extern crate r2d2;
13//! extern crate r2d2_sqlite;
14//! extern crate rusqlite;
15//!
16//! use std::thread;
17//! use r2d2_sqlite::SqliteConnectionManager;
18//! use rusqlite::params;
19//!
20//! fn main() {
21//!     let manager = SqliteConnectionManager::file("file.db");
22//!     let pool = r2d2::Pool::new(manager).unwrap();
23//!     pool.get()
24//!         .unwrap()
25//!         .execute("CREATE TABLE IF NOT EXISTS foo (bar INTEGER)", params![])
26//!         .unwrap();
27//!
28//!     (0..10)
29//!         .map(|i| {
30//!             let pool = pool.clone();
31//!             thread::spawn(move || {
32//!                 let conn = pool.get().unwrap();
33//!                 conn.execute("INSERT INTO foo (bar) VALUES (?)", &[&i])
34//!                     .unwrap();
35//!             })
36//!         })
37//!         .collect::<Vec<_>>()
38//!         .into_iter()
39//!         .map(thread::JoinHandle::join)
40//!         .collect::<Result<_, _>>()
41//!         .unwrap()
42//! }
43//! ```
44pub use rusqlite;
45use rusqlite::{Connection, Error, OpenFlags};
46use std::fmt;
47use std::path::{Path, PathBuf};
48use std::sync::Mutex;
49use uuid::Uuid;
50
51#[derive(Debug)]
52enum Source {
53    File(PathBuf),
54    Memory(String),
55}
56
57type InitFn = dyn Fn(&mut Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static;
58
59/// An `r2d2::ManageConnection` for `rusqlite::Connection`s.
60pub struct SqliteConnectionManager {
61    source: Source,
62    flags: OpenFlags,
63    init: Option<Box<InitFn>>,
64    _persist: Mutex<Option<Connection>>,
65}
66
67impl fmt::Debug for SqliteConnectionManager {
68    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
69        let mut builder = f.debug_struct("SqliteConnectionManager");
70        let _ = builder.field("source", &self.source);
71        let _ = builder.field("flags", &self.source);
72        let _ = builder.field("init", &self.init.as_ref().map(|_| "InitFn"));
73        builder.finish()
74    }
75}
76
77impl SqliteConnectionManager {
78    /// Creates a new `SqliteConnectionManager` from file.
79    ///
80    /// See `rusqlite::Connection::open`
81    pub fn file<P: AsRef<Path>>(path: P) -> Self {
82        Self {
83            source: Source::File(path.as_ref().to_path_buf()),
84            flags: OpenFlags::default(),
85            init: None,
86            _persist: Mutex::new(None),
87        }
88    }
89
90    /// Creates a new `SqliteConnectionManager` from memory.
91    pub fn memory() -> Self {
92        Self {
93            source: Source::Memory(Uuid::new_v4().to_string()),
94            flags: OpenFlags::default(),
95            init: None,
96            _persist: Mutex::new(None),
97        }
98    }
99
100    /// Converts `SqliteConnectionManager` into one that sets OpenFlags upon
101    /// connection creation.
102    ///
103    /// See `rustqlite::OpenFlags` for a list of available flags.
104    pub fn with_flags(self, flags: OpenFlags) -> Self {
105        Self { flags, ..self }
106    }
107
108    /// Converts `SqliteConnectionManager` into one that calls an initialization
109    /// function upon connection creation. Could be used to set PRAGMAs, for
110    /// example.
111    ///
112    /// ### Example
113    ///
114    /// Make a `SqliteConnectionManager` that sets the `foreign_keys` pragma to
115    /// true for every connection.
116    ///
117    /// ```rust,no_run
118    /// # use r2d2_sqlite::{SqliteConnectionManager};
119    /// let manager = SqliteConnectionManager::file("app.db")
120    ///     .with_init(|c| c.execute_batch("PRAGMA foreign_keys=1;"));
121    /// ```
122    pub fn with_init<F>(self, init: F) -> Self
123    where
124        F: Fn(&mut Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static,
125    {
126        let init: Option<Box<InitFn>> = Some(Box::new(init));
127        Self { init, ..self }
128    }
129}
130
131impl r2d2::ManageConnection for SqliteConnectionManager {
132    type Connection = Connection;
133    type Error = rusqlite::Error;
134
135    fn connect(&self) -> Result<Connection, Error> {
136        match self.source {
137            Source::File(ref path) => Connection::open_with_flags(path, self.flags),
138            Source::Memory(ref id) => {
139                let connection = || {
140                    Connection::open_with_flags(
141                        format!("file:{}?mode=memory&cache=shared", id),
142                        self.flags,
143                    )
144                };
145
146                {
147                    let mut persist = self._persist.lock().unwrap();
148                    if persist.is_none() {
149                        *persist = Some(connection()?);
150                    }
151                }
152
153                connection()
154            }
155        }
156        .map_err(Into::into)
157        .and_then(|mut c| match self.init {
158            None => Ok(c),
159            Some(ref init) => init(&mut c).map(|_| c),
160        })
161    }
162
163    fn is_valid(&self, conn: &mut Connection) -> Result<(), Error> {
164        conn.execute_batch("").map_err(Into::into)
165    }
166
167    fn has_broken(&self, _: &mut Connection) -> bool {
168        false
169    }
170}