use slog::{debug, Logger};
use mithril_common::logging::LoggerExtensions;
use mithril_common::StdResult;
use crate::sqlite::SqliteConnection;
#[derive(Eq, PartialEq, Copy, Clone)]
pub enum SqliteCleaningTask {
Vacuum,
WalCheckpointTruncate,
}
impl SqliteCleaningTask {
pub fn log_message(self: SqliteCleaningTask) -> &'static str {
match self {
SqliteCleaningTask::Vacuum => "Running `vacuum` on the SQLite database",
SqliteCleaningTask::WalCheckpointTruncate => {
"Running `wal_checkpoint(TRUNCATE)` on the SQLite database"
}
}
}
}
pub struct SqliteCleaner<'a> {
connection: &'a SqliteConnection,
logger: Logger,
tasks: Vec<SqliteCleaningTask>,
}
impl<'a> SqliteCleaner<'a> {
pub fn new(connection: &'a SqliteConnection) -> Self {
Self {
connection,
logger: Logger::root(slog::Discard, slog::o!()),
tasks: vec![],
}
}
pub fn with_logger(mut self, logger: Logger) -> Self {
self.logger = logger.new_with_component_name::<Self>();
self
}
pub fn with_tasks(mut self, tasks: &[SqliteCleaningTask]) -> Self {
for option in tasks {
self.tasks.push(*option);
}
self
}
pub fn run(self) -> StdResult<()> {
if self.tasks.contains(&SqliteCleaningTask::Vacuum) {
debug!(self.logger, "{}", SqliteCleaningTask::Vacuum.log_message());
self.connection.execute("vacuum")?;
}
if self
.tasks
.contains(&SqliteCleaningTask::WalCheckpointTruncate)
{
debug!(
self.logger,
"{}",
SqliteCleaningTask::WalCheckpointTruncate.log_message()
);
self.connection
.execute("PRAGMA wal_checkpoint(TRUNCATE);")?;
} else {
self.connection.execute("PRAGMA wal_checkpoint(PASSIVE);")?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::ops::Range;
use std::path::Path;
use mithril_common::test_utils::TempDir;
use crate::sqlite::{ConnectionBuilder, ConnectionOptions, SqliteConnection};
use super::*;
fn add_test_table(connection: &SqliteConnection) {
connection
.execute("CREATE TABLE IF NOT EXISTS test (id INTEGER PRIMARY KEY, text TEXT);")
.unwrap();
}
fn fill_test_table(connection: &SqliteConnection, ids: Range<u64>) {
connection
.execute(format!(
"INSERT INTO test (id, text) VALUES {}",
ids.map(|i| format!("({}, 'some text to fill the db')", i))
.collect::<Vec<String>>()
.join(", ")
))
.unwrap();
}
fn delete_test_rows(connection: &SqliteConnection, ids: Range<u64>) {
connection
.execute(format!(
"DELETE FROM test WHERE id >= {} and id < {}",
ids.start, ids.end
))
.unwrap();
}
fn prepare_db_for_vacuum(connection: &SqliteConnection) {
connection
.execute("pragma auto_vacuum = none; vacuum;")
.unwrap();
add_test_table(connection);
fill_test_table(connection, 0..10_000);
connection
.execute("PRAGMA wal_checkpoint(PASSIVE)")
.unwrap();
delete_test_rows(connection, 0..5_000);
connection
.execute("PRAGMA wal_checkpoint(PASSIVE)")
.unwrap();
}
fn file_size(path: &Path) -> u64 {
path.metadata()
.unwrap_or_else(|_| panic!("Failed to read len of '{}'", path.display()))
.len()
}
#[test]
fn cleanup_empty_in_memory_db_should_not_crash() {
let connection = ConnectionBuilder::open_memory().build().unwrap();
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::Vacuum])
.run()
.expect("Vacuum should not fail");
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
.run()
.expect("WalCheckpointTruncate should not fail");
}
#[test]
fn cleanup_empty_file_without_wal_db_should_not_crash() {
let db_path = TempDir::create(
"sqlite_cleaner",
"cleanup_empty_file_without_wal_db_should_not_crash",
)
.join("test.db");
let connection = ConnectionBuilder::open_file(&db_path).build().unwrap();
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::Vacuum])
.run()
.expect("Vacuum should not fail");
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
.run()
.expect("WalCheckpointTruncate should not fail");
}
#[test]
fn test_vacuum() {
let db_dir = TempDir::create("sqlite_cleaner", "test_vacuum");
let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
let connection = ConnectionBuilder::open_file(&db_path)
.with_options(&[ConnectionOptions::EnableWriteAheadLog])
.build()
.unwrap();
prepare_db_for_vacuum(&connection);
let db_initial_size = file_size(&db_path);
assert!(db_initial_size > 0);
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::Vacuum])
.run()
.unwrap();
let db_after_vacuum_size = file_size(&db_path);
assert!(
db_initial_size > db_after_vacuum_size,
"db size should have decreased (vacuum enabled)"
);
assert!(
file_size(&db_wal_path) > 0,
"db wal file should not have been truncated (truncate disabled)"
);
}
#[test]
fn test_truncate_wal() {
let db_dir = TempDir::create("sqlite_cleaner", "test_truncate_wal");
let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
let connection = ConnectionBuilder::open_file(&db_path)
.with_options(&[ConnectionOptions::EnableWriteAheadLog])
.build()
.unwrap();
add_test_table(&connection);
fill_test_table(&connection, 0..10_000);
delete_test_rows(&connection, 0..10_000);
assert!(file_size(&db_wal_path) > 0);
SqliteCleaner::new(&connection)
.with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
.run()
.unwrap();
assert_eq!(
file_size(&db_wal_path),
0,
"db wal file should have been truncated"
);
}
}