use std::fs;
use std::path::{Path, PathBuf};
use crate::directory::get_code_graph_dir;
use crate::types::SchemaVersion;
use super::migrations::{
CURRENT_SCHEMA_VERSION, current_time_millis, get_current_version, run_migrations,
};
use super::sqlite_adapter::{
PragmaOptions, SqliteBackend, SqliteDatabase, SqliteError, SqliteParams, SqliteResult,
SqliteValue, create_database,
};
pub use super::sqlite_adapter::{SqliteBackend as DatabaseBackend, SqliteDatabase as Database};
fn configure_connection(db: &mut dyn SqliteDatabase) -> SqliteResult<()> {
db.pragma("busy_timeout = 5000", None)?;
db.pragma("foreign_keys = ON", None)?;
db.pragma("journal_mode = WAL", None)?;
db.pragma("synchronous = NORMAL", None)?;
db.pragma("cache_size = -64000", None)?;
db.pragma("temp_store = MEMORY", None)?;
db.pragma("mmap_size = 268435456", None)?;
Ok(())
}
pub struct DatabaseConnection {
db: Box<dyn SqliteDatabase>,
db_path: PathBuf,
backend: SqliteBackend,
}
impl DatabaseConnection {
fn new(db: Box<dyn SqliteDatabase>, db_path: PathBuf, backend: SqliteBackend) -> Self {
Self {
db,
db_path,
backend,
}
}
pub fn initialize(db_path: impl AsRef<Path>) -> SqliteResult<Self> {
let db_path = db_path.as_ref().to_path_buf();
if let Some(dir) = db_path.parent() {
fs::create_dir_all(dir).map_err(|error| {
SqliteError::new(format!("failed to create database dir: {error}"))
})?;
}
let created = create_database(&db_path)?;
let mut db = created.db;
configure_connection(db.as_mut())?;
let schema_path = Path::new(file!()).with_file_name("schema.sql");
let schema = fs::read_to_string(&schema_path).map_err(|error| {
SqliteError::new(format!(
"failed to read schema at {}: {error}",
schema_path.display()
))
})?;
db.exec(&schema)?;
let current_version = get_current_version(db.as_mut())?;
if current_version < CURRENT_SCHEMA_VERSION {
let mut stmt = db.prepare(
"INSERT OR IGNORE INTO schema_versions (version, applied_at, description) VALUES (?, ?, ?)",
)?;
stmt.run(SqliteParams::positional(vec![
SqliteValue::from(CURRENT_SCHEMA_VERSION),
SqliteValue::from(current_time_millis()),
SqliteValue::from("Initial schema includes all migrations"),
]))?;
}
Ok(Self::new(db, db_path, created.backend))
}
pub fn open(db_path: impl AsRef<Path>) -> SqliteResult<Self> {
let db_path = db_path.as_ref().to_path_buf();
if !db_path.exists() {
return Err(SqliteError::new(format!(
"Database not found: {}",
db_path.display()
)));
}
let created = create_database(&db_path)?;
let mut db = created.db;
configure_connection(db.as_mut())?;
let current_version = get_current_version(db.as_mut())?;
if current_version < CURRENT_SCHEMA_VERSION {
run_migrations(db.as_mut(), current_version)?;
}
Ok(Self::new(db, db_path, created.backend))
}
pub fn get_db(&mut self) -> &mut dyn SqliteDatabase {
self.db.as_mut()
}
pub fn get_backend(&self) -> SqliteBackend {
self.backend
}
pub fn get_path(&self) -> &Path {
&self.db_path
}
pub fn get_journal_mode(&mut self) -> SqliteResult<String> {
let raw = self
.db
.pragma("journal_mode", Some(PragmaOptions { simple: false }))?;
let mode = match raw {
Some(SqliteValue::Text(value)) => value,
Some(value) => value.into_string_lossy().unwrap_or_default(),
None => String::new(),
};
Ok(mode.to_lowercase())
}
pub fn get_schema_version(&mut self) -> SqliteResult<Option<SchemaVersion>> {
let mut stmt = self.db.prepare(
"SELECT version, applied_at, description FROM schema_versions ORDER BY version DESC LIMIT 1",
)?;
let Some(mut row) = stmt.get(SqliteParams::none())? else {
return Ok(None);
};
let version = row
.remove("version")
.and_then(|value| value.as_i64())
.ok_or_else(|| SqliteError::new("schema version row missing version"))?;
let applied_at = row
.remove("applied_at")
.and_then(|value| value.as_i64())
.ok_or_else(|| SqliteError::new("schema version row missing applied_at"))?;
let description = row
.remove("description")
.and_then(SqliteValue::into_string_lossy);
Ok(Some(SchemaVersion {
version: version as u64,
applied_at,
description,
}))
}
pub fn transaction<T>(
&mut self,
mut f: impl FnMut(&mut dyn SqliteDatabase) -> SqliteResult<T>,
) -> SqliteResult<T> {
let mut result = None;
self.db.transaction(&mut |db| {
result = Some(f(db)?);
Ok(())
})?;
result.ok_or_else(|| SqliteError::new("transaction callback did not run"))
}
pub fn get_size(&self) -> SqliteResult<u64> {
fs::metadata(&self.db_path)
.map(|metadata| metadata.len())
.map_err(|error| SqliteError::new(format!("failed to stat database: {error}")))
}
pub fn optimize(&mut self) -> SqliteResult<()> {
self.db.exec("VACUUM")?;
self.db.exec("ANALYZE")?;
Ok(())
}
pub fn run_maintenance(&mut self) {
let _ = self.db.exec("PRAGMA optimize");
let _ = self.db.exec("PRAGMA wal_checkpoint(PASSIVE)");
}
pub fn close(&mut self) -> SqliteResult<()> {
self.db.close()
}
pub fn is_open(&self) -> bool {
self.db.is_open()
}
}
pub const DATABASE_FILENAME: &str = "rustcodegraph.db";
pub fn get_database_path(project_root: impl AsRef<Path>) -> PathBuf {
get_code_graph_dir(project_root).join(DATABASE_FILENAME)
}