use std::future::{Future, IntoFuture};
use std::pin::Pin;
use crate::model::{Backend, DbError, DbPool};
pub struct TestDb {
pool: DbPool,
transactional: bool,
}
enum Target {
#[cfg(feature = "model-sqlite")]
Sqlite,
#[cfg(feature = "model-postgres")]
Postgres(String),
#[cfg(feature = "model-mysql")]
MySql(String),
}
#[must_use = "a TestDbBuilder does nothing until it is awaited"]
pub struct TestDbBuilder {
target: Target,
migrations: Vec<String>,
}
impl TestDb {
#[cfg(feature = "model-sqlite")]
pub fn sqlite() -> TestDbBuilder {
TestDbBuilder { target: Target::Sqlite, migrations: Vec::new() }
}
#[cfg(feature = "model-postgres")]
pub fn postgres(url: &str) -> TestDbBuilder {
TestDbBuilder { target: Target::Postgres(url.to_string()), migrations: Vec::new() }
}
#[cfg(feature = "model-mysql")]
pub fn mysql(url: &str) -> TestDbBuilder {
TestDbBuilder { target: Target::MySql(url.to_string()), migrations: Vec::new() }
}
#[cfg(feature = "model-postgres")]
pub fn postgres_from_env() -> Option<TestDbBuilder> {
url_from_env("RWS_TEST_POSTGRES_URL", "postgres://user:password@localhost/app_test").map(|url| TestDb::postgres(&url))
}
#[cfg(feature = "model-mysql")]
pub fn mysql_from_env() -> Option<TestDbBuilder> {
url_from_env("RWS_TEST_MYSQL_URL", "mysql://user:password@localhost/app_test").map(|url| TestDb::mysql(&url))
}
pub fn pool(&self) -> DbPool {
self.pool.clone()
}
pub fn backend(&self) -> Backend {
self.pool.backend()
}
pub fn is_transactional(&self) -> bool {
self.transactional
}
pub async fn rollback(mut self) -> Result<(), DbError> {
let transactional = std::mem::replace(&mut self.transactional, false);
if !transactional {
return Ok(());
}
match &self.pool {
#[cfg(feature = "model-sqlite")]
DbPool::Sqlite(_) => Ok(()),
#[cfg(feature = "model-postgres")]
DbPool::Postgres(pool) => rollback_and_close(pool).await,
#[cfg(feature = "model-mysql")]
DbPool::MySql(pool) => rollback_and_close(pool).await,
}
}
}
impl std::ops::Deref for TestDb {
type Target = DbPool;
fn deref(&self) -> &DbPool {
&self.pool
}
}
impl Drop for TestDb {
fn drop(&mut self) {
if !self.transactional {
return;
}
match &self.pool {
#[cfg(feature = "model-sqlite")]
DbPool::Sqlite(_) => {}
#[cfg(feature = "model-postgres")]
DbPool::Postgres(pool) => close_soon(pool.clone()),
#[cfg(feature = "model-mysql")]
DbPool::MySql(pool) => close_soon(pool.clone()),
}
}
}
impl TestDbBuilder {
pub fn migrations(mut self, dir: &str) -> Self {
self.migrations.push(dir.to_string());
self
}
pub async fn build(self) -> Result<TestDb, DbError> {
match self.target {
#[cfg(feature = "model-sqlite")]
Target::Sqlite => {
let pool = sqlx::pool::PoolOptions::<sqlx::Sqlite>::new()
.max_connections(1)
.idle_timeout(None)
.max_lifetime(None)
.connect("sqlite::memory:")
.await
.map_err(|e| DbError::new(e.to_string()))?;
let pool = DbPool::Sqlite(pool);
for dir in &self.migrations {
pool.migrate(dir).await?;
}
Ok(TestDb { pool, transactional: false })
}
#[cfg(feature = "model-postgres")]
Target::Postgres(url) => {
let _ = crate::crypto_provider().install_default();
let options: sqlx::postgres::PgConnectOptions =
url.parse().map_err(|e: sqlx::Error| DbError::new(format!("RWS test database URL: {e}")))?;
migrate_once(&url, &self.migrations, || async {
let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
.max_connections(1)
.connect_with(options.clone())
.await
.map_err(|e| DbError::new(e.to_string()))?;
Ok(DbPool::Postgres(pool))
})
.await?;
let pool = transactional_pool::<sqlx::Postgres>()
.connect_with(options)
.await
.map_err(|e| DbError::new(e.to_string()))?;
Ok(TestDb { pool: DbPool::Postgres(pool), transactional: true })
}
#[cfg(feature = "model-mysql")]
Target::MySql(url) => {
let _ = crate::crypto_provider().install_default();
let options: sqlx::mysql::MySqlConnectOptions =
url.parse().map_err(|e: sqlx::Error| DbError::new(format!("RWS test database URL: {e}")))?;
migrate_once(&url, &self.migrations, || async {
let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
.max_connections(1)
.connect_with(options.clone())
.await
.map_err(|e| DbError::new(e.to_string()))?;
Ok(DbPool::MySql(pool))
})
.await?;
let pool = transactional_pool::<sqlx::MySql>()
.connect_with(options)
.await
.map_err(|e| DbError::new(e.to_string()))?;
Ok(TestDb { pool: DbPool::MySql(pool), transactional: true })
}
}
}
}
impl IntoFuture for TestDbBuilder {
type Output = Result<TestDb, DbError>;
type IntoFuture = Pin<Box<dyn Future<Output = Result<TestDb, DbError>> + Send>>;
fn into_future(self) -> Self::IntoFuture {
Box::pin(self.build())
}
}
#[cfg(any(feature = "model-postgres", feature = "model-mysql"))]
fn url_from_env(name: &str, example: &str) -> Option<String> {
match std::env::var(name).map(|v| v.trim().to_string()) {
Ok(url) if !url.is_empty() => Some(url),
_ => {
eprintln!("skipped: {name} is not set (the URL of a database for tests, e.g. {example})");
None
}
}
}
#[cfg(any(feature = "model-postgres", feature = "model-mysql"))]
fn transactional_pool<DB: sqlx::Database>() -> sqlx::pool::PoolOptions<DB> {
sqlx::pool::PoolOptions::<DB>::new()
.max_connections(1)
.min_connections(0)
.idle_timeout(None)
.max_lifetime(None)
.after_connect(|conn, _meta| <DB::TransactionManager as sqlx::TransactionManager>::begin(conn, None))
}
#[cfg(any(feature = "model-postgres", feature = "model-mysql"))]
async fn rollback_and_close<DB: sqlx::Database>(pool: &sqlx::Pool<DB>) -> Result<(), DbError> {
let result = async {
let mut conn = pool.acquire().await.map_err(|e| DbError::new(e.to_string()))?;
<DB::TransactionManager as sqlx::TransactionManager>::rollback(&mut conn).await.map_err(|e| DbError::new(e.to_string()))
}
.await;
pool.close().await;
result
}
#[cfg(any(feature = "model-postgres", feature = "model-mysql"))]
fn close_soon<DB: sqlx::Database>(pool: sqlx::Pool<DB>) {
drop(pool.close());
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(async move { pool.close().await });
}
}
#[cfg(any(feature = "model-postgres", feature = "model-mysql"))]
async fn migrate_once<F, Fut>(url: &str, dirs: &[String], connect: F) -> Result<(), DbError>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<DbPool, DbError>>,
{
use std::collections::HashSet;
use std::sync::OnceLock;
static DONE: OnceLock<tokio::sync::Mutex<HashSet<String>>> = OnceLock::new();
if dirs.is_empty() {
return Ok(());
}
let mut done = DONE.get_or_init(Default::default).lock().await;
let keys: Vec<String> = dirs.iter().map(|dir| format!("{url}\n{dir}")).collect();
if keys.iter().all(|key| done.contains(key)) {
return Ok(());
}
let pool = connect().await?;
let mut result = Ok(());
for (dir, key) in dirs.iter().zip(keys) {
if done.contains(&key) {
continue;
}
if let Err(e) = pool.migrate(dir).await {
result = Err(e);
break;
}
done.insert(key);
}
match &pool {
#[cfg(feature = "model-sqlite")]
DbPool::Sqlite(p) => p.close().await,
#[cfg(feature = "model-postgres")]
DbPool::Postgres(p) => p.close().await,
#[cfg(feature = "model-mysql")]
DbPool::MySql(p) => p.close().await,
}
result
}