mod backup;
mod dao;
mod migration;
mod schema;
#[cfg(test)]
mod tests;
pub use dao::FailoverQueueItem;
use crate::config::get_app_config_dir;
use crate::error::AppError;
use rusqlite::Connection;
use serde::Serialize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
const DB_BACKUP_RETAIN: usize = 10;
pub(crate) const SCHEMA_VERSION: i32 = 12;
pub(crate) fn to_json_string<T: Serialize>(value: &T) -> Result<String, AppError> {
serde_json::to_string(value)
.map_err(|e| AppError::Config(format!("JSON serialization failed: {e}")))
}
macro_rules! lock_conn {
($mutex:expr) => {
$mutex
.lock()
.map_err(|e| AppError::Database(format!("Mutex lock failed: {}", e)))?
};
}
pub(crate) use lock_conn;
pub struct Database {
pub(crate) conn: Mutex<Connection>,
runtime_key: String,
}
impl Database {
pub fn init() -> Result<Self, AppError> {
let db_path = get_app_config_dir().join("cc-switch.db");
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent).map_err(|e| AppError::io(parent, e))?;
}
let conn = Connection::open(&db_path).map_err(|e| AppError::Database(e.to_string()))?;
conn.execute("PRAGMA foreign_keys = ON;", [])
.map_err(|e| AppError::Database(e.to_string()))?;
let db = Self {
conn: Mutex::new(conn),
runtime_key: format!("file:{}", db_path.display()),
};
{
let conn = lock_conn!(db.conn);
let version = Self::get_user_version(&conn)?;
drop(conn);
if version > SCHEMA_VERSION {
return Err(Self::future_schema_error(version));
}
if version > 0 && version < SCHEMA_VERSION {
log::info!(
"Creating pre-migration database backup (v{version} -> v{SCHEMA_VERSION})"
);
if let Err(err) = db.backup_database_file() {
log::warn!("Pre-migration backup failed, continuing migration: {err}");
}
}
}
db.create_tables()?;
db.apply_schema_migrations()?;
db.ensure_model_pricing_seeded()?;
Ok(db)
}
pub fn memory() -> Result<Self, AppError> {
static NEXT_MEMORY_DB_ID: AtomicU64 = AtomicU64::new(1);
let conn = Connection::open_in_memory().map_err(|e| AppError::Database(e.to_string()))?;
conn.execute("PRAGMA foreign_keys = ON;", [])
.map_err(|e| AppError::Database(e.to_string()))?;
let db = Self {
conn: Mutex::new(conn),
runtime_key: format!(
"memory:{}",
NEXT_MEMORY_DB_ID.fetch_add(1, Ordering::Relaxed)
),
};
db.create_tables()?;
db.ensure_model_pricing_seeded()?;
Ok(db)
}
pub fn is_mcp_table_empty(&self) -> Result<bool, AppError> {
let conn = lock_conn!(self.conn);
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM mcp_servers", [], |row| row.get(0))
.map_err(|e| AppError::Database(e.to_string()))?;
Ok(count == 0)
}
pub fn is_prompts_table_empty(&self) -> Result<bool, AppError> {
let conn = lock_conn!(self.conn);
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM prompts", [], |row| row.get(0))
.map_err(|e| AppError::Database(e.to_string()))?;
Ok(count == 0)
}
pub(crate) fn runtime_key(&self) -> &str {
&self.runtime_key
}
}