use crate::helpers::config;
use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
use diesel::SqliteConnection;
use diesel_migrations::{embed_migrations, EmbeddedMigrations};
use std::sync::{OnceLock, RwLock};
#[deprecated(note = "embed and run migrations from the consuming application")]
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
struct PoolState {
pool: Pool<ConnectionManager<SqliteConnection>>,
database_url: String,
}
static POOL: OnceLock<RwLock<Option<PoolState>>> = OnceLock::new();
pub fn get_connection() -> PooledConnection<ConnectionManager<SqliteConnection>> {
let current_db_url = config::database_url();
let pool_state = POOL.get_or_init(|| RwLock::new(Some(create_pool_state(¤t_db_url))));
{
let state = pool_state.read().unwrap();
if let Some(ref ps) = *state {
if ps.database_url == current_db_url {
return ps.pool.get().expect("Failed to get connection from pool");
}
}
}
{
let mut state = pool_state.write().unwrap();
let needs_recreate = match *state {
None => true,
Some(ref ps) => ps.database_url != current_db_url,
};
if needs_recreate {
*state = Some(create_pool_state(¤t_db_url));
}
state
.as_ref()
.unwrap()
.pool
.get()
.expect("Failed to get connection from pool")
}
}
#[allow(dead_code)]
pub fn get_connection_pool() -> Pool<ConnectionManager<SqliteConnection>> {
let current_db_url = config::database_url();
let pool_state = POOL.get_or_init(|| RwLock::new(Some(create_pool_state(¤t_db_url))));
{
let state = pool_state.read().unwrap();
if let Some(ref ps) = *state {
if ps.database_url == current_db_url {
return ps.pool.clone();
}
}
}
{
let mut state = pool_state.write().unwrap();
let needs_recreate = match *state {
None => true,
Some(ref ps) => ps.database_url != current_db_url,
};
if needs_recreate {
*state = Some(create_pool_state(¤t_db_url));
}
state.as_ref().unwrap().pool.clone()
}
}
#[cfg(any(test, feature = "test-utils"))]
pub fn reset_pool() {
if let Some(pool_state) = POOL.get() {
let mut state = pool_state.write().unwrap();
*state = None;
}
}
fn create_pool_state(database_url: &str) -> PoolState {
let pool_limit: u32 = config::database_pool_limit();
let manager = ConnectionManager::<SqliteConnection>::new(database_url);
let pool = Pool::builder()
.max_size(pool_limit)
.test_on_check_out(true)
.build(manager)
.expect("Could not build connection pool");
PoolState {
pool,
database_url: database_url.to_string(),
}
}