use crate::persist::global::config::DatabaseConfig;
use crate::prelude::*;
use r2d2::{CustomizeConnection, Pool};
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::OpenFlags;
pub type ConnectionPool = Pool<SqliteConnectionManager>;
#[derive(Debug)]
struct TransactionCleanupCustomizer;
impl CustomizeConnection<rusqlite::Connection, rusqlite::Error> for TransactionCleanupCustomizer {
fn on_acquire(
&self,
conn: &mut rusqlite::Connection,
) -> std::result::Result<(), rusqlite::Error> {
let _ = conn.execute("ROLLBACK", []);
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct DatabaseService {
pool: ConnectionPool,
config: DatabaseConfig,
}
impl DatabaseService {
pub fn new(config: &DatabaseConfig) -> anyhow::Result<Self> {
if config.db_path == ":memory:" {
info!("Database: Using in-memory database");
return Self::new_in_memory_with_config(Some(config.db_max_connections));
}
if let Some(parent) = Path::new(&config.db_path).parent() {
std::fs::create_dir_all(parent)
.map_err(|e| anyhow::anyhow!("Failed to create directory: {}", e))?;
}
let config_clone = config.clone();
let manager = SqliteConnectionManager::file(&config.db_path)
.with_flags(
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_FULL_MUTEX,
)
.with_init(move |conn| {
if config_clone.db_enable_wal {
conn.pragma_update(None, "journal_mode", "WAL")?;
}
if config_clone.db_enable_foreign_keys {
conn.pragma_update(None, "foreign_keys", "ON")?;
}
conn.pragma_update(None, "busy_timeout", config_clone.db_busy_timeout_ms)?;
conn.pragma_update(None, "synchronous", "NORMAL")?;
conn.pragma_update(None, "cache_size", 10000)?;
conn.pragma_update(None, "temp_store", "memory")?;
conn.pragma_update(None, "mmap_size", 134217728)?;
conn.pragma_update(None, "wal_autocheckpoint", 1000)?;
Ok(())
});
let pool = Pool::builder()
.max_size(config.db_max_connections)
.connection_timeout(std::time::Duration::from_secs(
config.db_connection_timeout_seconds,
))
.connection_customizer(Box::new(TransactionCleanupCustomizer))
.build(manager)
.map_err(|e| anyhow::anyhow!("Failed to create database pool: {}", e))?;
let service = Self {
pool,
config: config.clone(),
};
service.run_migrations()?;
Ok(service)
}
pub fn new_in_memory() -> anyhow::Result<Self> {
Self::new_in_memory_with_config(None)
}
pub fn new_in_memory_with_config(max_connections: Option<u32>) -> anyhow::Result<Self> {
let pool_size = max_connections.unwrap_or(1);
let manager = SqliteConnectionManager::file("file:memdb?mode=memory&cache=shared")
.with_flags(
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_FULL_MUTEX,
)
.with_init(|conn| {
conn.pragma_update(None, "foreign_keys", "ON")?;
conn.pragma_update(None, "busy_timeout", 30000)?;
conn.pragma_update(None, "synchronous", "NORMAL")?;
conn.pragma_update(None, "cache_size", 10000)?;
conn.pragma_update(None, "temp_store", "memory")?;
Ok(())
});
let pool = Pool::builder()
.max_size(pool_size)
.connection_timeout(std::time::Duration::from_secs(5))
.connection_customizer(Box::new(TransactionCleanupCustomizer))
.build(manager)
.map_err(|e| anyhow::anyhow!("Failed to create in-memory database pool: {}", e))?;
let service = Self {
pool,
config: DatabaseConfig {
db_path: ":memory:".to_string(),
db_max_connections: pool_size,
db_connection_timeout_seconds: 5,
db_enable_wal: false,
db_enable_foreign_keys: true,
db_busy_timeout_ms: 5000,
},
};
{
let conn = service.get_connection()?;
crate::persist::global::migrations::run_migrations(conn)?;
}
Ok(service)
}
pub fn get_connection(
&self,
) -> anyhow::Result<r2d2::PooledConnection<SqliteConnectionManager>> {
self.pool.get().map_err(Into::into)
}
pub fn run_migrations(&self) -> anyhow::Result<()> {
let conn = self.get_connection()?;
crate::persist::global::migrations::run_migrations(conn)?;
tracing::debug!("Database migrations completed successfully");
Ok(())
}
pub fn health_check(&self) -> anyhow::Result<()> {
let conn = self.get_connection()?;
let _: i32 = conn.query_row("SELECT 1", [], |row| row.get(0))?;
Ok(())
}
}