mithril_persistence/sqlite/
cleaner.rs

1use slog::{Logger, debug};
2
3use mithril_common::StdResult;
4use mithril_common::logging::LoggerExtensions;
5
6use crate::sqlite::SqliteConnection;
7
8/// Tasks that can be performed by the SqliteCleaner
9#[derive(Eq, PartialEq, Copy, Clone)]
10pub enum SqliteCleaningTask {
11    /// Reconstruct the database file, repacking it into a minimal amount of disk space.
12    ///
13    /// see: <https://www.sqlite.org/lang_vacuum.html>
14    ///
15    /// ⚠ This operation can be very slow on large databases ⚠
16    Vacuum,
17    /// Run a checkpoint to transfer the data from the WAL file to the main db file and truncate
18    /// it afterward.
19    ///
20    /// see: <https://www.sqlite.org/pragma.html#pragma_wal_checkpoint>
21    WalCheckpointTruncate,
22}
23
24impl SqliteCleaningTask {
25    /// Get the log message for the task.
26    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
36/// The SqliteCleaner is responsible for cleaning up databases by performing tasks defined
37/// in [SqliteCleaningTask].
38pub struct SqliteCleaner<'a> {
39    connection: &'a SqliteConnection,
40    logger: Logger,
41    tasks: Vec<SqliteCleaningTask>,
42}
43
44impl<'a> SqliteCleaner<'a> {
45    /// Create a new instance of the `SqliteCleaner`.
46    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    /// Set the logger to be used by the cleaner.
55    pub fn with_logger(mut self, logger: Logger) -> Self {
56        self.logger = logger.new_with_component_name::<Self>();
57        self
58    }
59
60    /// Set the [SqliteCleaningTask] to be performed by the cleaner.
61    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    /// Cleanup the database by performing the defined tasks.
69    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        // Important: If WAL is enabled Vacuuming the database will not shrink until a
76        // checkpoint is run, so it must be done after vacuuming.
77        // Note: running a checkpoint when the WAL is disabled is harmless.
78        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    /// Apply migrations, disable auto_vacuum and mangle the database to create some free pages
131    /// for the vacuum to reclaim
132    fn prepare_db_for_vacuum(connection: &SqliteConnection) {
133        // Disable Auto vacuum to allow the test to check if the vacuum was run
134        connection.execute("pragma auto_vacuum = none; vacuum;").unwrap();
135        add_test_table(connection);
136        fill_test_table(connection, 0..10_000);
137        // Checkpoint before deletion so entries are transferred from the WAL file to the main db
138        connection.execute("PRAGMA wal_checkpoint(PASSIVE)").unwrap();
139        delete_test_rows(connection, 0..5_000);
140        // Checkpoint after deletion to create free pages in the main db
141        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        // Make "neutral" changes to the db, this will fill the WAL files with some data
223        // but won't change the db size after cleaning up.
224        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}
mithril_persistence/sqlite/cleaner.rs

mithril_persistence/sqlite/
cleaner.rs
