mithril_persistence/sqlite/
cleaner.rs1use slog::{Logger, debug};
2
3use mithril_common::StdResult;
4use mithril_common::logging::LoggerExtensions;
5
6use crate::sqlite::SqliteConnection;
7
8#[derive(Eq, PartialEq, Copy, Clone)]
10pub enum SqliteCleaningTask {
11 Vacuum,
17 WalCheckpointTruncate,
22}
23
24impl SqliteCleaningTask {
25 pub fn log_message(self: SqliteCleaningTask) -> &'static str {
27 match self {
28 SqliteCleaningTask::Vacuum => "Running `vacuum` on the SQLite database",
29 SqliteCleaningTask::WalCheckpointTruncate => {
30 "Running `wal_checkpoint(TRUNCATE)` on the SQLite database"
31 }
32 }
33 }
34}
35
36pub struct SqliteCleaner<'a> {
39 connection: &'a SqliteConnection,
40 logger: Logger,
41 tasks: Vec<SqliteCleaningTask>,
42}
43
44impl<'a> SqliteCleaner<'a> {
45 pub fn new(connection: &'a SqliteConnection) -> Self {
47 Self {
48 connection,
49 logger: Logger::root(slog::Discard, slog::o!()),
50 tasks: vec![],
51 }
52 }
53
54 pub fn with_logger(mut self, logger: Logger) -> Self {
56 self.logger = logger.new_with_component_name::<Self>();
57 self
58 }
59
60 pub fn with_tasks(mut self, tasks: &[SqliteCleaningTask]) -> Self {
62 for option in tasks {
63 self.tasks.push(*option);
64 }
65 self
66 }
67
68 pub fn run(self) -> StdResult<()> {
70 if self.tasks.contains(&SqliteCleaningTask::Vacuum) {
71 debug!(self.logger, "{}", SqliteCleaningTask::Vacuum.log_message());
72 self.connection.execute("vacuum")?;
73 }
74
75 if self.tasks.contains(&SqliteCleaningTask::WalCheckpointTruncate) {
79 debug!(
80 self.logger,
81 "{}",
82 SqliteCleaningTask::WalCheckpointTruncate.log_message()
83 );
84 self.connection.execute("PRAGMA wal_checkpoint(TRUNCATE);")?;
85 } else {
86 self.connection.execute("PRAGMA wal_checkpoint(PASSIVE);")?;
87 }
88
89 Ok(())
90 }
91}
92
93#[cfg(test)]
94mod tests {
95 use std::ops::Range;
96 use std::path::Path;
97
98 use mithril_common::test_utils::TempDir;
99
100 use crate::sqlite::{ConnectionBuilder, ConnectionOptions, SqliteConnection};
101
102 use super::*;
103
104 fn add_test_table(connection: &SqliteConnection) {
105 connection
106 .execute("CREATE TABLE IF NOT EXISTS test (id INTEGER PRIMARY KEY, text TEXT);")
107 .unwrap();
108 }
109
110 fn fill_test_table(connection: &SqliteConnection, ids: Range<u64>) {
111 connection
112 .execute(format!(
113 "INSERT INTO test (id, text) VALUES {}",
114 ids.map(|i| format!("({i}, 'some text to fill the db')"))
115 .collect::<Vec<String>>()
116 .join(", ")
117 ))
118 .unwrap();
119 }
120
121 fn delete_test_rows(connection: &SqliteConnection, ids: Range<u64>) {
122 connection
123 .execute(format!(
124 "DELETE FROM test WHERE id >= {} and id < {}",
125 ids.start, ids.end
126 ))
127 .unwrap();
128 }
129
130 fn prepare_db_for_vacuum(connection: &SqliteConnection) {
133 connection.execute("pragma auto_vacuum = none; vacuum;").unwrap();
135 add_test_table(connection);
136 fill_test_table(connection, 0..10_000);
137 connection.execute("PRAGMA wal_checkpoint(PASSIVE)").unwrap();
139 delete_test_rows(connection, 0..5_000);
140 connection.execute("PRAGMA wal_checkpoint(PASSIVE)").unwrap();
142 }
143
144 fn file_size(path: &Path) -> u64 {
145 path.metadata()
146 .unwrap_or_else(|_| panic!("Failed to read len of '{}'", path.display()))
147 .len()
148 }
149
150 #[test]
151 fn cleanup_empty_in_memory_db_should_not_crash() {
152 let connection = ConnectionBuilder::open_memory().build().unwrap();
153
154 SqliteCleaner::new(&connection)
155 .with_tasks(&[SqliteCleaningTask::Vacuum])
156 .run()
157 .expect("Vacuum should not fail");
158 SqliteCleaner::new(&connection)
159 .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
160 .run()
161 .expect("WalCheckpointTruncate should not fail");
162 }
163
164 #[test]
165 fn cleanup_empty_file_without_wal_db_should_not_crash() {
166 let db_path = TempDir::create(
167 "sqlite_cleaner",
168 "cleanup_empty_file_without_wal_db_should_not_crash",
169 )
170 .join("test.db");
171 let connection = ConnectionBuilder::open_file(&db_path).build().unwrap();
172
173 SqliteCleaner::new(&connection)
174 .with_tasks(&[SqliteCleaningTask::Vacuum])
175 .run()
176 .expect("Vacuum should not fail");
177 SqliteCleaner::new(&connection)
178 .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
179 .run()
180 .expect("WalCheckpointTruncate should not fail");
181 }
182
183 #[test]
184 fn test_vacuum() {
185 let db_dir = TempDir::create("sqlite_cleaner", "test_vacuum");
186 let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
187 let connection = ConnectionBuilder::open_file(&db_path)
188 .with_options(&[ConnectionOptions::EnableWriteAheadLog])
189 .build()
190 .unwrap();
191 prepare_db_for_vacuum(&connection);
192
193 let db_initial_size = file_size(&db_path);
194 assert!(db_initial_size > 0);
195
196 SqliteCleaner::new(&connection)
197 .with_tasks(&[SqliteCleaningTask::Vacuum])
198 .run()
199 .unwrap();
200
201 let db_after_vacuum_size = file_size(&db_path);
202
203 assert!(
204 db_initial_size > db_after_vacuum_size,
205 "db size should have decreased (vacuum enabled)"
206 );
207 assert!(
208 file_size(&db_wal_path) > 0,
209 "db wal file should not have been truncated (truncate disabled)"
210 );
211 }
212
213 #[test]
214 fn test_truncate_wal() {
215 let db_dir = TempDir::create("sqlite_cleaner", "test_truncate_wal");
216 let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
217 let connection = ConnectionBuilder::open_file(&db_path)
218 .with_options(&[ConnectionOptions::EnableWriteAheadLog])
219 .build()
220 .unwrap();
221
222 add_test_table(&connection);
225 fill_test_table(&connection, 0..10_000);
226 delete_test_rows(&connection, 0..10_000);
227
228 assert!(file_size(&db_wal_path) > 0);
229
230 SqliteCleaner::new(&connection)
231 .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
232 .run()
233 .unwrap();
234
235 assert_eq!(
236 file_size(&db_wal_path),
237 0,
238 "db wal file should have been truncated"
239 );
240 }
241}