mithril_persistence/sqlite/
cleaner.rs

1use slog::{debug, Logger};
2
3use mithril_common::logging::LoggerExtensions;
4use mithril_common::StdResult;
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
79            .tasks
80            .contains(&SqliteCleaningTask::WalCheckpointTruncate)
81        {
82            debug!(
83                self.logger,
84                "{}",
85                SqliteCleaningTask::WalCheckpointTruncate.log_message()
86            );
87            self.connection
88                .execute("PRAGMA wal_checkpoint(TRUNCATE);")?;
89        } else {
90            self.connection.execute("PRAGMA wal_checkpoint(PASSIVE);")?;
91        }
92
93        Ok(())
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use std::ops::Range;
100    use std::path::Path;
101
102    use mithril_common::test_utils::TempDir;
103
104    use crate::sqlite::{ConnectionBuilder, ConnectionOptions, SqliteConnection};
105
106    use super::*;
107
108    fn add_test_table(connection: &SqliteConnection) {
109        connection
110            .execute("CREATE TABLE IF NOT EXISTS test (id INTEGER PRIMARY KEY, text TEXT);")
111            .unwrap();
112    }
113
114    fn fill_test_table(connection: &SqliteConnection, ids: Range<u64>) {
115        connection
116            .execute(format!(
117                "INSERT INTO test (id, text) VALUES {}",
118                ids.map(|i| format!("({}, 'some text to fill the db')", i))
119                    .collect::<Vec<String>>()
120                    .join(", ")
121            ))
122            .unwrap();
123    }
124
125    fn delete_test_rows(connection: &SqliteConnection, ids: Range<u64>) {
126        connection
127            .execute(format!(
128                "DELETE FROM test WHERE id >= {} and id < {}",
129                ids.start, ids.end
130            ))
131            .unwrap();
132    }
133
134    /// Apply migrations, disable auto_vacuum and mangle the database to create some free pages
135    /// for the vacuum to reclaim
136    fn prepare_db_for_vacuum(connection: &SqliteConnection) {
137        // Disable Auto vacuum to allow the test to check if the vacuum was run
138        connection
139            .execute("pragma auto_vacuum = none; vacuum;")
140            .unwrap();
141        add_test_table(connection);
142        fill_test_table(connection, 0..10_000);
143        // Checkpoint before deletion so entries are transferred from the WAL file to the main db
144        connection
145            .execute("PRAGMA wal_checkpoint(PASSIVE)")
146            .unwrap();
147        delete_test_rows(connection, 0..5_000);
148        // Checkpoint after deletion to create free pages in the main db
149        connection
150            .execute("PRAGMA wal_checkpoint(PASSIVE)")
151            .unwrap();
152    }
153
154    fn file_size(path: &Path) -> u64 {
155        path.metadata()
156            .unwrap_or_else(|_| panic!("Failed to read len of '{}'", path.display()))
157            .len()
158    }
159
160    #[test]
161    fn cleanup_empty_in_memory_db_should_not_crash() {
162        let connection = ConnectionBuilder::open_memory().build().unwrap();
163
164        SqliteCleaner::new(&connection)
165            .with_tasks(&[SqliteCleaningTask::Vacuum])
166            .run()
167            .expect("Vacuum should not fail");
168        SqliteCleaner::new(&connection)
169            .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
170            .run()
171            .expect("WalCheckpointTruncate should not fail");
172    }
173
174    #[test]
175    fn cleanup_empty_file_without_wal_db_should_not_crash() {
176        let db_path = TempDir::create(
177            "sqlite_cleaner",
178            "cleanup_empty_file_without_wal_db_should_not_crash",
179        )
180        .join("test.db");
181        let connection = ConnectionBuilder::open_file(&db_path).build().unwrap();
182
183        SqliteCleaner::new(&connection)
184            .with_tasks(&[SqliteCleaningTask::Vacuum])
185            .run()
186            .expect("Vacuum should not fail");
187        SqliteCleaner::new(&connection)
188            .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
189            .run()
190            .expect("WalCheckpointTruncate should not fail");
191    }
192
193    #[test]
194    fn test_vacuum() {
195        let db_dir = TempDir::create("sqlite_cleaner", "test_vacuum");
196        let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
197        let connection = ConnectionBuilder::open_file(&db_path)
198            .with_options(&[ConnectionOptions::EnableWriteAheadLog])
199            .build()
200            .unwrap();
201        prepare_db_for_vacuum(&connection);
202
203        let db_initial_size = file_size(&db_path);
204        assert!(db_initial_size > 0);
205
206        SqliteCleaner::new(&connection)
207            .with_tasks(&[SqliteCleaningTask::Vacuum])
208            .run()
209            .unwrap();
210
211        let db_after_vacuum_size = file_size(&db_path);
212
213        assert!(
214            db_initial_size > db_after_vacuum_size,
215            "db size should have decreased (vacuum enabled)"
216        );
217        assert!(
218            file_size(&db_wal_path) > 0,
219            "db wal file should not have been truncated (truncate disabled)"
220        );
221    }
222
223    #[test]
224    fn test_truncate_wal() {
225        let db_dir = TempDir::create("sqlite_cleaner", "test_truncate_wal");
226        let (db_path, db_wal_path) = (db_dir.join("test.db"), db_dir.join("test.db-wal"));
227        let connection = ConnectionBuilder::open_file(&db_path)
228            .with_options(&[ConnectionOptions::EnableWriteAheadLog])
229            .build()
230            .unwrap();
231
232        // Make "neutral" changes to the db, this will fill the WAL files with some data
233        // but won't change the db size after cleaning up.
234        add_test_table(&connection);
235        fill_test_table(&connection, 0..10_000);
236        delete_test_rows(&connection, 0..10_000);
237
238        assert!(file_size(&db_wal_path) > 0);
239
240        SqliteCleaner::new(&connection)
241            .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
242            .run()
243            .unwrap();
244
245        assert_eq!(
246            file_size(&db_wal_path),
247            0,
248            "db wal file should have been truncated"
249        );
250    }
251}
mithril_persistence/sqlite/cleaner.rs

mithril_persistence/sqlite/
cleaner.rs
