Skip to main content

keyvaluedb_sqlite/
lib.rs

1#![deny(clippy::all)]
2
3mod tools;
4
5pub use async_sqlite::rusqlite::OpenFlags;
6use async_sqlite::rusqlite::{params, OptionalExtension as _};
7use async_sqlite::*;
8use keyvaluedb::{
9    DBKeyRef, DBKeyValue, DBKeyValueRef, DBOp, DBTransaction, DBTransactionError, DBValue, IoStats,
10    IoStatsKind, KeyValueDB, KeyValueDBPinBoxFuture,
11};
12use parking_lot::{Mutex, RwLock};
13use std::sync::atomic::{AtomicBool, Ordering};
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16use std::{
17    io,
18    path::{Path, PathBuf},
19    str::FromStr,
20};
21use tools::*;
22
23///////////////////////////////////////////////////////////////////////////////
24
25/// Default for [DatabaseConfig::wal_autocheckpoint_pages] (1MB at 4KB pages)
26const WAL_AUTOCHECKPOINT_PAGES: u32 = 256;
27/// Default for [DatabaseConfig::wal_journal_size_limit]
28const WAL_JOURNAL_SIZE_LIMIT: u32 = 1024 * 1024;
29
30#[derive(Copy, Clone, Debug, Eq, PartialEq)]
31pub enum VacuumMode {
32    None,
33    Incremental,
34    Full,
35}
36
37/// Which integrity check gates the open-time repair
38#[derive(Copy, Clone, Debug, Eq, PartialEq)]
39pub enum RepairCheck {
40    /// No open-time check; repair still runs when the file will not open or
41    /// an operation hits corruption at runtime
42    None,
43    /// PRAGMA quick_check: page structure only, misses index-vs-table
44    /// disagreement
45    Quick,
46    /// PRAGMA integrity_check: the full verdict, catches everything sqlite
47    /// can detect
48    Full,
49}
50
51/// Called with a report every time a repair rebuilds the database
52pub type RepairCallback = Arc<dyn Fn(&RepairReport) + Send + Sync>;
53
54/// Database configuration
55#[derive(Clone)]
56pub struct DatabaseConfig {
57    /// Set number of columns.
58    /// The number of columns must not be zero.
59    pub columns: u32,
60    /// Set flags used to open the database
61    pub flags: OpenFlags,
62    /// Number of connections to open
63    pub num_conns: usize,
64    /// Vacuum mode
65    pub vacuum_mode: VacuumMode,
66    /// On corruption, at open or at runtime, salvage the readable rows into a
67    /// fresh database, keeping the damaged files beside it
68    pub repair_on_corrupt: bool,
69    /// Which integrity check gates the open-time repair
70    pub repair_check: RepairCheck,
71    /// Skip the open-time check when the database (plus WAL) is larger than
72    /// this; the check reads every page, so on a large store it can hold up
73    /// open for a long time, and runtime repair still covers what it would
74    /// have caught. None checks regardless of size.
75    pub repair_check_byte_limit: Option<u64>,
76    /// Called with a report every time a repair runs
77    pub on_repair: Option<RepairCallback>,
78    /// Checkpoint every N WAL pages: bounds the crash window and how much WAL
79    /// every open has to replay. 0 disables autocheckpointing.
80    pub wal_autocheckpoint_pages: u32,
81    /// Truncate the WAL file to this many bytes on checkpoint restart; must
82    /// sit at or below the autocheckpoint watermark or the file never shrinks
83    /// from its high-water mark
84    pub wal_journal_size_limit: u32,
85}
86
87impl DatabaseConfig {
88    /// Create new `DatabaseConfig` with default parameters
89    pub fn new() -> Self {
90        Default::default()
91    }
92
93    /// Set the number of columns. `columns` must not be zero.
94    pub fn with_columns(self, columns: u32) -> Self {
95        assert!(columns > 0, "the number of columns must not be zero");
96        Self { columns, ..self }
97    }
98
99    /// Enable corruption detection and salvage, at open and at runtime
100    pub fn with_repair_on_corrupt(self, repair_on_corrupt: bool) -> Self {
101        Self {
102            repair_on_corrupt,
103            ..self
104        }
105    }
106
107    /// Set which integrity check gates the open-time repair
108    pub fn with_repair_check(self, repair_check: RepairCheck) -> Self {
109        Self {
110            repair_check,
111            ..self
112        }
113    }
114
115    /// Skip the open-time check for databases larger than `bytes`
116    pub fn with_repair_check_byte_limit(self, bytes: u64) -> Self {
117        Self {
118            repair_check_byte_limit: Some(bytes),
119            ..self
120        }
121    }
122
123    /// Set a callback invoked with a report every time a repair runs
124    pub fn with_on_repair(self, on_repair: RepairCallback) -> Self {
125        Self {
126            on_repair: Some(on_repair),
127            ..self
128        }
129    }
130
131    /// Checkpoint every `pages` WAL pages (0 disables autocheckpointing)
132    pub fn with_wal_autocheckpoint_pages(self, pages: u32) -> Self {
133        Self {
134            wal_autocheckpoint_pages: pages,
135            ..self
136        }
137    }
138
139    /// Truncate the WAL to `bytes` on checkpoint restart; keep at or below the
140    /// autocheckpoint watermark or the file never shrinks
141    pub fn with_wal_journal_size_limit(self, bytes: u32) -> Self {
142        Self {
143            wal_journal_size_limit: bytes,
144            ..self
145        }
146    }
147
148    /// Sets the flags to 'in-memory database'
149    pub fn with_in_memory(self) -> Self {
150        Self {
151            flags: OpenFlags::SQLITE_OPEN_READ_WRITE
152                | OpenFlags::SQLITE_OPEN_CREATE
153                | OpenFlags::SQLITE_OPEN_NO_MUTEX
154                | OpenFlags::SQLITE_OPEN_MEMORY,
155            ..self
156        }
157    }
158
159    /// Replaces all the flags
160    pub fn with_flags(self, flags: OpenFlags) -> Self {
161        Self { flags, ..self }
162    }
163
164    /// Sets the number of connections for this database
165    pub fn with_num_conns(self, num_conns: usize) -> Self {
166        Self { num_conns, ..self }
167    }
168
169    /// Set the vacuum mode used by 'cleanup'
170    pub fn with_vacuum_mode(self, vacuum_mode: VacuumMode) -> Self {
171        Self {
172            vacuum_mode,
173            ..self
174        }
175    }
176}
177
178impl Default for DatabaseConfig {
179    fn default() -> DatabaseConfig {
180        DatabaseConfig {
181            columns: 1,
182            flags: OpenFlags::SQLITE_OPEN_READ_WRITE
183                | OpenFlags::SQLITE_OPEN_CREATE
184                | OpenFlags::SQLITE_OPEN_NO_MUTEX,
185            num_conns: 1,
186            vacuum_mode: VacuumMode::None,
187            repair_on_corrupt: false,
188            repair_check: RepairCheck::Full,
189            repair_check_byte_limit: None,
190            on_repair: None,
191            wal_autocheckpoint_pages: WAL_AUTOCHECKPOINT_PAGES,
192            wal_journal_size_limit: WAL_JOURNAL_SIZE_LIMIT,
193        }
194    }
195}
196
197///////////////////////////////////////////////////////////////////////////////
198
199/// Whether a rusqlite error says the file is damaged
200fn code_is_corruption(e: &rusqlite::Error) -> bool {
201    matches!(
202        e,
203        rusqlite::Error::SqliteFailure(f, _) if matches!(
204            f.code,
205            rusqlite::ErrorCode::DatabaseCorrupt | rusqlite::ErrorCode::NotADatabase
206        )
207    )
208}
209
210/// Whether a pool error says the file is damaged
211fn error_is_corruption(e: &Error) -> bool {
212    matches!(e, Error::Rusqlite(e) if code_is_corruption(e))
213}
214
215/// Whether an error chain bottoms out in sqlite saying the file is damaged
216fn io_error_is_corruption(e: &io::Error) -> bool {
217    let mut src: Option<&(dyn std::error::Error + 'static)> = e.get_ref().map(|b| b as _);
218    while let Some(cur) = src {
219        if let Some(e) = cur.downcast_ref::<rusqlite::Error>() {
220            return code_is_corruption(e);
221        }
222        if let Some(Error::Rusqlite(e)) = cur.downcast_ref::<Error>() {
223            return code_is_corruption(e);
224        }
225        src = cur.source();
226    }
227    false
228}
229
230///////////////////////////////////////////////////////////////////////////////
231
232/// An sqlite table with its statement strings
233pub struct DatabaseTable {
234    _table: String,
235    str_has_value: String,
236    str_has_value_like: String,
237    str_get_unique_value: String,
238    str_get_first_value_like: String,
239    str_set_unique_value: String,
240    str_remove_unique_value: String,
241    str_remove_and_return_unique_value: String,
242    str_remove_unique_value_like: String,
243    str_iter_with_prefix: String,
244    str_iter_no_prefix: String,
245    str_iter_keys_with_prefix: String,
246    str_iter_keys_no_prefix: String,
247}
248
249impl DatabaseTable {
250    pub fn new(table: String) -> Self {
251        let str_has_value = format!("SELECT 1 FROM {} WHERE [key] = ? LIMIT 1", table);
252        let str_has_value_like = format!(
253            "SELECT 1 FROM {} WHERE [key] LIKE ? ESCAPE '\\' LIMIT 1",
254            table
255        );
256        let str_get_unique_value = format!("SELECT value FROM {} WHERE [key] = ? LIMIT 1", table);
257        let str_get_first_value_like = format!(
258            "SELECT key, value FROM {} WHERE [key] LIKE ? ESCAPE '\\' LIMIT 1",
259            table
260        );
261        // Update in place; OR REPLACE reinserts under a new rowid and rewrites more pages
262        let str_set_unique_value = format!(
263            "INSERT INTO {} ([key], value) VALUES(?, ?) ON CONFLICT([key]) DO UPDATE SET value = excluded.value",
264            table
265        );
266        let str_remove_unique_value = format!("DELETE FROM {} WHERE [key] = ?", table);
267        let str_remove_and_return_unique_value =
268            format!("DELETE FROM {} WHERE [key] = ? RETURNING value", table);
269        let str_remove_unique_value_like =
270            format!("DELETE FROM {} WHERE [key] LIKE ? ESCAPE '\\'", table);
271        let str_iter_with_prefix = format!(
272            "SELECT key, value FROM {} WHERE [key] LIKE ? ESCAPE '\\'",
273            table
274        );
275        let str_iter_no_prefix = format!("SELECT key, value FROM {}", table);
276        let str_iter_keys_with_prefix =
277            format!("SELECT key FROM {} WHERE [key] LIKE ? ESCAPE '\\'", table);
278        let str_iter_keys_no_prefix = format!("SELECT key FROM {}", table);
279
280        Self {
281            _table: table,
282            str_has_value,
283            str_has_value_like,
284            str_get_unique_value,
285            str_get_first_value_like,
286            str_set_unique_value,
287            str_remove_unique_value,
288            str_remove_and_return_unique_value,
289            str_remove_unique_value_like,
290            str_iter_with_prefix,
291            str_iter_no_prefix,
292            str_iter_keys_with_prefix,
293            str_iter_keys_no_prefix,
294        }
295    }
296}
297
298///////////////////////////////////////////////////////////////////////////////
299
300/// What a WAL truncate checkpoint accomplished
301#[derive(Copy, Clone, Debug, Default)]
302pub struct CheckpointResult {
303    /// Another connection kept the checkpoint from completing
304    pub busy: bool,
305    /// Frames in the WAL
306    pub log_frames: i64,
307    /// Frames folded into the main database file
308    pub checkpointed_frames: i64,
309}
310
311impl CheckpointResult {
312    /// Every frame made it into the main file and the WAL was reset
313    pub fn complete(&self) -> bool {
314        !self.busy && self.log_frames == self.checkpointed_frames
315    }
316}
317
318/// Fold the WAL into the main file and truncate it, reporting the verdict;
319/// PRAGMA wal_checkpoint returns a row, so pragma_update would discard a
320/// busy (failed) checkpoint silently
321fn wal_checkpoint_truncate(conn: &rusqlite::Connection) -> rusqlite::Result<CheckpointResult> {
322    conn.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |row| {
323        Ok(CheckpointResult {
324            busy: row.get::<_, i64>(0)? != 0,
325            log_frames: row.get(1)?,
326            checkpointed_frames: row.get(2)?,
327        })
328    })
329}
330
331/// What the automatic corruption repair did, kept on the database it produced
332#[derive(Debug, Clone)]
333pub struct RepairReport {
334    /// The integrity failure that triggered the repair
335    pub detected: String,
336    /// Rows salvaged across the control table and all column tables
337    pub rows_recovered: u64,
338    /// Tables whose salvage scan ended early on a damaged page
339    pub partial_tables: u32,
340    /// Where the damaged database files were moved
341    pub corrupt_path: PathBuf,
342}
343
344/// The swappable engine state: a repair closes the pool, swaps the files and
345/// installs a fresh pool, bumping the generation so racing operations know to
346/// retry rather than fail
347struct DatabaseState {
348    pool: Pool,
349    generation: u64,
350    repair_report: Option<RepairReport>,
351}
352
353/// What a salvage scan pulled out of a damaged database
354#[derive(Default)]
355struct Salvage {
356    tables: Vec<(String, Vec<(String, rusqlite::types::Value)>)>,
357    partial_tables: u32,
358}
359
360const SIBLING_SUFFIXES: [&str; 3] = ["", "-wal", "-shm"];
361
362fn sibling(base: &Path, ext: &str, suffix: &str) -> PathBuf {
363    let mut os = base.as_os_str().to_owned();
364    os.push(ext);
365    os.push(suffix);
366    PathBuf::from(os)
367}
368
369/// The database's on-disk size: main file plus WAL
370fn db_file_bytes(path: &Path) -> u64 {
371    let size = |p: PathBuf| std::fs::metadata(p).map(|m| m.len()).unwrap_or(0);
372    size(sibling(path, "", "")) + size(sibling(path, "", "-wal"))
373}
374
375/// Move the damaged files aside as `<name>.corrupt` and the rebuilt
376/// `.repairing` files into place, returning where the damaged files went
377fn swap_repaired_files(path: &Path) -> PathBuf {
378    let corrupt_path = sibling(path, ".corrupt", "");
379    for suffix in &SIBLING_SUFFIXES {
380        let _ = std::fs::remove_file(sibling(path, ".corrupt", suffix));
381        let _ = std::fs::rename(sibling(path, "", suffix), sibling(path, ".corrupt", suffix));
382        let _ = std::fs::rename(
383            sibling(path, ".repairing", suffix),
384            sibling(path, "", suffix),
385        );
386    }
387    corrupt_path
388}
389
390/// An sqlite key-value database fulfilling the `KeyValueDB` trait
391pub struct DatabaseUnlockedInner {
392    path: PathBuf,
393    config: DatabaseConfig,
394    state: RwLock<DatabaseState>,
395    control_table: Arc<DatabaseTable>,
396    column_tables: Vec<Arc<DatabaseTable>>,
397}
398
399impl Drop for DatabaseUnlockedInner {
400    fn drop(&mut self) {
401        let _ = self.state.get_mut().pool.close_blocking();
402    }
403}
404
405pub struct DatabaseInner {
406    overall_stats: IoStats,
407    current_stats: IoStats,
408}
409
410#[derive(Clone)]
411pub struct Database {
412    unlocked_inner: Arc<DatabaseUnlockedInner>,
413    inner: Arc<Mutex<DatabaseInner>>,
414}
415
416impl Database {
417    ////////////////////////////////////////////////////////////////
418    // Initialization
419
420    pub fn open<P: AsRef<Path>>(path: P, config: DatabaseConfig) -> io::Result<Self> {
421        let path = PathBuf::from(path.as_ref());
422        let db = match Self::open_raw(&path, &config, None) {
423            Ok(db) => db,
424            Err(e) if config.repair_on_corrupt && io_error_is_corruption(&e) => {
425                Self::repair(None, &path, &config, e.to_string())?
426            }
427            Err(e) => return Err(e),
428        };
429        let db = if config.repair_on_corrupt {
430            // The open-time check reads every page; above the byte limit skip
431            // it and let runtime repair cover what it would have caught
432            let check = match config.repair_check_byte_limit {
433                Some(limit) if db_file_bytes(&path) > limit => RepairCheck::None,
434                _ => config.repair_check,
435            };
436            match db.check_blocking(check) {
437                Ok(None) => db,
438                Ok(Some(detected)) => Self::repair(Some(db), &path, &config, detected)?,
439                Err(e) if io_error_is_corruption(&e) => {
440                    Self::repair(Some(db), &path, &config, e.to_string())?
441                }
442                Err(e) => return Err(e),
443            }
444        } else {
445            db
446        };
447
448        // Fold the WAL only after the integrity verdict; a damaged WAL must be
449        // quarantined by the repair, never folded into the main file. Busy is
450        // tolerated here: another process may hold the file, and the
451        // autocheckpoint keeps the WAL bounded from then on.
452        if !db.config().flags.contains(OpenFlags::SQLITE_OPEN_MEMORY) {
453            db.checkpoint_blocking().map_err(io::Error::other)?;
454        }
455
456        if let (Some(on_repair), Some(report)) = (&db.config().on_repair, db.repair_report()) {
457            on_repair(&report);
458        }
459        Ok(db)
460    }
461
462    /// Open the connection pool and apply the session pragmas to every
463    /// connection in it
464    fn open_pool(path: &Path, config: &DatabaseConfig) -> io::Result<Pool> {
465        let in_memory = config.flags.contains(OpenFlags::SQLITE_OPEN_MEMORY);
466        let mut pool_builder = PoolBuilder::new()
467            .path(path)
468            .flags(config.flags)
469            .num_conns(config.num_conns);
470        if !in_memory {
471            pool_builder = pool_builder.journal_mode(JournalMode::Wal);
472        }
473        let pool = pool_builder.open_blocking().map_err(io::Error::other)?;
474
475        let wal_journal_size_limit = config.wal_journal_size_limit;
476        let wal_autocheckpoint_pages = config.wal_autocheckpoint_pages;
477        for res in pool.conn_for_each_blocking(move |conn| {
478            // Don't rely on STATEMENT_CACHE_DEFAULT_CAPACITY in rusqlite, set it explicitly
479            conn.set_prepared_statement_cache_capacity(256);
480
481            conn.pragma_update(None, "case_sensitive_like", "ON")?;
482            conn.pragma_update(None, "synchronous", "normal")?;
483            conn.pragma_update(None, "journal_size_limit", wal_journal_size_limit)?;
484            conn.pragma_update(None, "wal_autocheckpoint", wal_autocheckpoint_pages)?;
485            // Wait out cross-process lock contention instead of failing immediately
486            conn.pragma_update(None, "busy_timeout", 2000)?;
487            // Catch page scribbles at the operation that hits them, where the
488            // runtime repair can act, instead of letting them spread
489            conn.pragma_update(None, "cell_size_check", "ON")?;
490            Ok(())
491        }) {
492            res.map_err(io::Error::other)?;
493        }
494        Ok(pool)
495    }
496
497    fn open_raw(
498        path: &Path,
499        config: &DatabaseConfig,
500        repair_report: Option<RepairReport>,
501    ) -> io::Result<Self> {
502        let config = config.clone();
503        assert_ne!(config.columns, 0, "number of columns must be >= 1");
504
505        let path = PathBuf::from(path);
506
507        let mut column_tables = vec![];
508        for n in 0..config.columns {
509            column_tables.push(Arc::new(DatabaseTable::new(get_column_table_name(n))))
510        }
511        let control_table = Arc::new(DatabaseTable::new("control".to_string()));
512
513        let pool = Self::open_pool(&path, &config)?;
514
515        let out = Self {
516            unlocked_inner: Arc::new(DatabaseUnlockedInner {
517                path,
518                config,
519                state: RwLock::new(DatabaseState {
520                    pool,
521                    generation: 0,
522                    repair_report,
523                }),
524                control_table,
525                column_tables,
526            }),
527            inner: Arc::new(Mutex::new(DatabaseInner {
528                overall_stats: IoStats::empty(),
529                current_stats: IoStats::empty(),
530            })),
531        };
532
533        let vacuum_mode = out.config().vacuum_mode;
534
535        out.conn_blocking(move |conn| {
536            match vacuum_mode {
537                VacuumMode::None | VacuumMode::Full => {
538                    let current: u32 =
539                        conn.pragma_query_value(None, "auto_vacuum", |x| x.get(0))?;
540                    if current != 0 {
541                        conn.execute("VACUUM", [])?;
542                        conn.pragma_update(None, "auto_vacuum", 0)?;
543                    }
544                }
545                VacuumMode::Incremental => {
546                    let current: u32 =
547                        conn.pragma_query_value(None, "auto_vacuum", |x| x.get(0))?;
548                    if current != 2 {
549                        conn.execute("VACUUM", [])?;
550                        conn.pragma_update(None, "auto_vacuum", "2")?;
551                    }
552                }
553            }
554
555            Ok(())
556        })
557        .map_err(io::Error::other)?;
558
559        out.open_resize_columns()?;
560
561        Ok(out)
562    }
563
564    pub fn path(&self) -> PathBuf {
565        self.unlocked_inner.path.clone()
566    }
567
568    /// What the most recent repair did, if any has run
569    pub fn repair_report(&self) -> Option<RepairReport> {
570        self.unlocked_inner.state.read().repair_report.clone()
571    }
572
573    /// The current pool and its generation; the generation changes when a
574    /// repair swaps the pool out
575    fn current(&self) -> (Pool, u64) {
576        let state = self.unlocked_inner.state.read();
577        (state.pool.clone(), state.generation)
578    }
579
580    fn repair_enabled(&self) -> bool {
581        self.unlocked_inner.config.repair_on_corrupt
582    }
583
584    /// Integrity check on one connection. None = clean; Some = what is wrong.
585    /// A check that cannot even run to completion errs with the corruption it hit.
586    fn check_blocking(&self, check: RepairCheck) -> io::Result<Option<String>> {
587        let pragma = match check {
588            RepairCheck::None => return Ok(None),
589            RepairCheck::Quick => "PRAGMA quick_check(8)",
590            RepairCheck::Full => "PRAGMA integrity_check(8)",
591        };
592        let problems = self
593            .conn_blocking(move |conn| {
594                let mut stmt = conn.prepare(pragma)?;
595                let mut rows = stmt.query([])?;
596                let mut problems: Vec<String> = Vec::new();
597                while let Some(row) = rows.next()? {
598                    problems.push(row.get(0)?);
599                }
600                Ok(problems)
601            })
602            .map_err(io::Error::other)?;
603        if problems.len() == 1 && problems[0] == "ok" {
604            return Ok(None);
605        }
606        Ok(Some(problems.join("; ")))
607    }
608
609    /// Rebuild a damaged database from whatever rows are still readable.
610    ///
611    /// The damaged files move beside the store as `<name>.corrupt` rather than
612    /// being destroyed: corruption that sqlite's own crash-safety should have
613    /// made impossible is evidence worth keeping. `old` is the still-open
614    /// handle when the damage was found by the integrity check; None when the
615    /// file would not even open, in which case nothing is salvageable and the
616    /// result is a fresh empty database.
617    fn repair(
618        old: Option<Self>,
619        path: &Path,
620        config: &DatabaseConfig,
621        detected: String,
622    ) -> io::Result<Self> {
623        let mut rows_recovered = 0u64;
624        let mut partial_tables = 0u32;
625        if let Some(old) = &old {
626            let (old_pool, _) = old.current();
627            let salvage = Self::salvage_all(&old_pool, config)?;
628            partial_tables = salvage.partial_tables;
629            rows_recovered = Self::build_replacement(path, config, salvage)?;
630        } else {
631            Self::build_replacement(path, config, Salvage::default())?;
632        }
633
634        // Close the old handle so every file is settled before the swap
635        drop(old);
636        let corrupt_path = swap_repaired_files(path);
637
638        Self::open_raw(
639            path,
640            config,
641            Some(RepairReport {
642                detected,
643                rows_recovered,
644                partial_tables,
645                corrupt_path,
646            }),
647        )
648    }
649
650    /// Repair in place after an operation hit corruption at runtime.
651    ///
652    /// Holds the state write lock throughout: new operations wait, operations
653    /// already queued drain when the old pool closes, and the generation bump
654    /// tells racing callers that failed meanwhile to retry. Ok means the
655    /// database was repaired (or a racing caller already had); Err means the
656    /// repair could not run and the original operation error stands.
657    fn repair_live(&self, detected: String, seen_generation: u64) -> io::Result<()> {
658        let path = self.unlocked_inner.path.clone();
659        let config = self.unlocked_inner.config.clone();
660        let mut state = self.unlocked_inner.state.write();
661        if state.generation != seen_generation {
662            return Ok(());
663        }
664
665        let salvage = Self::salvage_all(&state.pool, &config)?;
666        let partial_tables = salvage.partial_tables;
667        let rows_recovered = Self::build_replacement(&path, &config, salvage)?;
668
669        if let Err(e) = state.pool.close_blocking() {
670            // Keep serving from the damaged file rather than die closed
671            if let Ok(pool) = Self::open_pool(&path, &config) {
672                state.pool = pool;
673                state.generation += 1;
674            }
675            return Err(io::Error::other(e));
676        }
677        let corrupt_path = swap_repaired_files(&path);
678
679        state.pool = Self::open_pool(&path, &config)?;
680        state.generation += 1;
681        let report = RepairReport {
682            detected,
683            rows_recovered,
684            partial_tables,
685            corrupt_path,
686        };
687        state.repair_report = Some(report.clone());
688        drop(state);
689
690        if let Some(on_repair) = &config.on_repair {
691            on_repair(&report);
692        }
693        Ok(())
694    }
695
696    /// Read every row a damaged database will still yield, all tables through
697    /// one connection under an exclusive transaction: repair must not scan, or
698    /// swap files afterward, while any other writer, in this process or
699    /// another, is mid-write. Errs busy if another writer holds the database.
700    fn salvage_all(pool: &Pool, config: &DatabaseConfig) -> io::Result<Salvage> {
701        let config_columns = config.columns;
702        pool.conn_mut_blocking(move |conn| {
703            let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Exclusive)?;
704            // Salvage every table the database says it has, or at least the
705            // configured set if its control table cannot say
706            let columns = tx
707                .query_row(
708                    "SELECT value FROM control WHERE [key] = 'columns'",
709                    [],
710                    |row| row.get::<_, String>(0),
711                )
712                .ok()
713                .and_then(|v| v.parse::<u32>().ok())
714                .unwrap_or(0)
715                .max(config_columns);
716            let mut tables = vec!["control".to_string()];
717            for cn in 0..columns {
718                tables.push(get_column_table_name(cn));
719            }
720            let mut out = Salvage::default();
721            for table in tables {
722                let (rows, partial) = Self::salvage_table_rows(&tx, &table);
723                if partial {
724                    out.partial_tables += 1;
725                }
726                out.tables.push((table, rows));
727            }
728            Ok(out)
729        })
730        .map_err(io::Error::other)
731    }
732
733    /// Build the replacement database at `<db>.repairing` from the salvaged
734    /// rows, returning how many were restored. The replacement is closed and
735    /// settled on return, ready to swap in.
736    fn build_replacement(
737        path: &Path,
738        config: &DatabaseConfig,
739        salvage: Salvage,
740    ) -> io::Result<u64> {
741        let tmp = sibling(path, ".repairing", "");
742        for suffix in &SIBLING_SUFFIXES {
743            let _ = std::fs::remove_file(sibling(path, ".repairing", suffix));
744        }
745        let fresh = Self::open_raw(&tmp, config, None)?;
746
747        let mut rows_recovered = 0u64;
748        for (table, rows) in salvage.tables {
749            if rows.is_empty() {
750                continue;
751            }
752            rows_recovered += rows.len() as u64;
753            // INSERT OR IGNORE: values the fresh open already wrote
754            // (control's own column count) win over salvaged ones
755            fresh
756                .conn_blocking(move |conn| {
757                    let tx = conn.unchecked_transaction()?;
758                    {
759                        let mut stmt = tx.prepare(&format!(
760                            "INSERT OR IGNORE INTO {} ([key], value) VALUES (?, ?)",
761                            table
762                        ))?;
763                        for (k, v) in &rows {
764                            stmt.execute(params![k, v])?;
765                        }
766                    }
767                    tx.commit()
768                })
769                .map_err(io::Error::other)?;
770        }
771        Ok(rows_recovered)
772    }
773
774    /// Read every row a damaged table will still yield, stopping at the first
775    /// page it cannot, keeping what came before it
776    fn salvage_table_rows(
777        conn: &rusqlite::Connection,
778        table: &str,
779    ) -> (Vec<(String, rusqlite::types::Value)>, bool) {
780        let mut out = Vec::new();
781        let mut partial = true;
782        if let Ok(mut stmt) = conn.prepare(&format!("SELECT [key], value FROM {}", table)) {
783            if let Ok(mut rows) = stmt.query([]) {
784                loop {
785                    match rows.next() {
786                        Ok(Some(row)) => {
787                            if let (Ok(k), Ok(v)) = (row.get(0), row.get(1)) {
788                                out.push((k, v));
789                            }
790                        }
791                        Ok(None) => {
792                            partial = false;
793                            break;
794                        }
795                        Err(_) => break,
796                    }
797                }
798            }
799        }
800        (out, partial)
801    }
802
803    pub fn config(&self) -> DatabaseConfig {
804        self.unlocked_inner.config.clone()
805    }
806
807    pub fn columns(&self) -> u32 {
808        self.unlocked_inner.config.columns
809    }
810
811    pub fn control_table(&self) -> Arc<DatabaseTable> {
812        self.unlocked_inner.control_table.clone()
813    }
814
815    pub fn column_table(&self, col: u32) -> Arc<DatabaseTable> {
816        self.unlocked_inner.column_tables[col as usize].clone()
817    }
818
819    pub fn conn_blocking<T, F>(&self, func: F) -> Result<T, Error>
820    where
821        F: FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
822        T: Send + 'static,
823    {
824        self.current().0.conn_blocking(func)
825    }
826
827    pub async fn conn<T, F>(&self, func: F) -> Result<T, Error>
828    where
829        F: FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
830        T: Send + 'static,
831    {
832        self.current().0.conn(func).await
833    }
834
835    pub async fn conn_mut<T, F>(&self, func: F) -> Result<T, Error>
836    where
837        F: FnOnce(&mut rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
838        T: Send + 'static,
839    {
840        self.current().0.conn_mut(func).await
841    }
842
843    /// Decide what a failed operation should do: repair and retry, retry
844    /// because a racing repair swapped the pool, or give up with the error.
845    ///
846    /// Reading the generation blocks while a repair holds the write lock, so a
847    /// caller that failed because of a concurrent repair waits it out here and
848    /// then retries against the healthy pool.
849    fn recover(&self, e: Error, seen_generation: u64) -> Result<(), Error> {
850        if !self.repair_enabled() {
851            return Err(e);
852        }
853        if error_is_corruption(&e) {
854            if self.repair_live(e.to_string(), seen_generation).is_err() {
855                return Err(e);
856            }
857            return Ok(());
858        }
859        // Not corruption: worth one retry only if a repair swapped the pool
860        // out from under this operation
861        if self.current().1 == seen_generation {
862            return Err(e);
863        }
864        Ok(())
865    }
866
867    /// Run `make()`'s closure on a pool connection, repairing and retrying
868    /// once if it hits corruption (or a concurrent repair)
869    async fn conn_retry<T, MK>(&self, make: MK) -> Result<T, Error>
870    where
871        MK: Fn() -> Box<dyn FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send>,
872        T: Send + 'static,
873    {
874        let (pool, generation) = self.current();
875        match pool.conn(make()).await {
876            Err(e) => {
877                self.recover(e, generation)?;
878                self.current().0.conn(make()).await
879            }
880            r => r,
881        }
882    }
883
884    /// Blocking variant of `conn_retry`
885    fn conn_retry_blocking<T, MK>(&self, make: MK) -> Result<T, Error>
886    where
887        MK: Fn() -> Box<dyn FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send>,
888        T: Send + 'static,
889    {
890        let (pool, generation) = self.current();
891        match pool.conn_blocking(make()) {
892            Err(e) => {
893                self.recover(e, generation)?;
894                self.current().0.conn_blocking(make())
895            }
896            r => r,
897        }
898    }
899
900    /// Mutable-connection variant of `conn_retry`
901    async fn conn_mut_retry<T, MK>(&self, make: MK) -> Result<T, Error>
902    where
903        MK: Fn() -> Box<dyn FnOnce(&mut rusqlite::Connection) -> Result<T, rusqlite::Error> + Send>,
904        T: Send + 'static,
905    {
906        let (pool, generation) = self.current();
907        match pool.conn_mut(make()).await {
908            Err(e) => {
909                self.recover(e, generation)?;
910                self.current().0.conn_mut(make()).await
911            }
912            r => r,
913        }
914    }
915
916    ////////////////////////////////////////////////////////////////
917    // Low level operations
918
919    /// Remove the last column family in the database. The deletion is definitive.
920    pub fn remove_last_column(&self) -> Result<(), Error> {
921        let this = self.clone();
922        self.conn_blocking(move |conn| {
923            let columns = Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
924            if columns == 0 {
925                return Err(rusqlite::Error::QueryReturnedNoRows);
926            }
927            Self::set_unique_value(conn, this.control_table(), "columns", columns - 1)?;
928
929            conn.execute(
930                &format!("DROP TABLE {}", get_column_table_name(columns - 1)),
931                [],
932            )?;
933            Ok(())
934        })
935    }
936
937    /// Add a new column family to the DB.
938    pub fn add_column(&self) -> Result<(), Error> {
939        let this = self.clone();
940
941        self.conn_blocking(move |conn| {
942            let columns = Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
943            Self::set_unique_value(conn, this.control_table(), "columns", columns + 1)?;
944            Self::create_column_table(conn, columns)
945        })
946    }
947    /// Helper to create new transaction for this database.
948    pub fn transaction(&self) -> DBTransaction {
949        DBTransaction::new()
950    }
951
952    /// Vacuum database
953    pub async fn vacuum(&self) -> Result<(), Error> {
954        let vacuum_mode = self.config().vacuum_mode;
955        self.conn_retry(|| {
956            Box::new(move |conn: &rusqlite::Connection| {
957                match vacuum_mode {
958                    VacuumMode::None => {}
959                    VacuumMode::Incremental => {
960                        conn.execute("PRAGMA incremental_vacuum", [])?;
961                    }
962                    VacuumMode::Full => {
963                        conn.execute("VACUUM", [])?;
964                    }
965                }
966                let _ = wal_checkpoint_truncate(conn)?;
967                Ok(())
968            })
969        })
970        .await
971    }
972
973    /// Fold the WAL into the main file and truncate it, reporting the verdict
974    pub async fn checkpoint(&self) -> Result<CheckpointResult, Error> {
975        self.conn(wal_checkpoint_truncate).await
976    }
977
978    /// Blocking variant of `checkpoint`
979    pub fn checkpoint_blocking(&self) -> Result<CheckpointResult, Error> {
980        self.conn_blocking(wal_checkpoint_truncate)
981    }
982
983    ////////////////////////////////////////////////////////////////
984    // Internal helpers
985
986    fn validate_column(&self, col: u32) -> rusqlite::Result<()> {
987        if col >= self.columns() {
988            return Err(rusqlite::Error::InvalidColumnIndex(col as usize));
989        }
990        Ok(())
991    }
992
993    fn create_column_table(conn: &rusqlite::Connection, column: u32) -> rusqlite::Result<()> {
994        conn.execute(&format!("CREATE TABLE IF NOT EXISTS {} (id INTEGER PRIMARY KEY AUTOINCREMENT, [key] TEXT UNIQUE, value BLOB)", get_column_table_name(column)), []).map(drop)
995    }
996
997    fn get_unique_value<V>(
998        conn: &rusqlite::Connection,
999        table: Arc<DatabaseTable>,
1000        key: &str,
1001        default: V,
1002    ) -> rusqlite::Result<V>
1003    where
1004        V: FromStr,
1005    {
1006        let mut stmt = conn.prepare_cached(&table.str_get_unique_value)?;
1007
1008        if let Ok(found) = stmt.query_row([key], |row| -> rusqlite::Result<String> { row.get(0) }) {
1009            if let Ok(v) = V::from_str(&found) {
1010                return Ok(v);
1011            }
1012        }
1013        Ok(default)
1014    }
1015
1016    fn set_unique_value<V>(
1017        conn: &rusqlite::Connection,
1018        table: Arc<DatabaseTable>,
1019        key: &str,
1020        value: V,
1021    ) -> rusqlite::Result<()>
1022    where
1023        V: ToString,
1024    {
1025        let mut stmt = conn.prepare_cached(&table.str_set_unique_value)?;
1026
1027        let changed = stmt.execute([key, value.to_string().as_str()])?;
1028
1029        // Never panic in a pool closure: a panic mid-transaction leaves the
1030        // connection holding a stale open transaction
1031        if changed > 1 {
1032            return Err(rusqlite::Error::StatementChangedRows(changed));
1033        }
1034        if changed == 0 {
1035            return Err(rusqlite::Error::QueryReturnedNoRows);
1036        }
1037
1038        Ok(())
1039    }
1040
1041    fn has_value(
1042        conn: &rusqlite::Connection,
1043        table: Arc<DatabaseTable>,
1044        key: &str,
1045    ) -> rusqlite::Result<bool> {
1046        let mut stmt = conn.prepare_cached(&table.str_has_value)?;
1047        stmt.exists([key])
1048    }
1049
1050    fn has_value_like(
1051        conn: &rusqlite::Connection,
1052        table: Arc<DatabaseTable>,
1053        key: &str,
1054    ) -> rusqlite::Result<bool> {
1055        let mut stmt = conn.prepare_cached(&table.str_has_value_like)?;
1056        stmt.exists([key])
1057    }
1058
1059    fn load_unique_value_blob(
1060        conn: &rusqlite::Connection,
1061        table: Arc<DatabaseTable>,
1062        key: &str,
1063    ) -> rusqlite::Result<Option<Vec<u8>>> {
1064        let mut stmt = conn.prepare_cached(&table.str_get_unique_value)?;
1065
1066        stmt.query_row([key], |row| -> rusqlite::Result<Vec<u8>> { row.get(0) })
1067            .optional()
1068    }
1069
1070    fn load_first_value_blob_like(
1071        conn: &rusqlite::Connection,
1072        table: Arc<DatabaseTable>,
1073        like: &str,
1074    ) -> rusqlite::Result<Option<(String, Vec<u8>)>> {
1075        let mut stmt = conn.prepare_cached(&table.str_get_first_value_like)?;
1076
1077        stmt.query_row([like], |row| -> rusqlite::Result<(String, Vec<u8>)> {
1078            Ok((row.get(0)?, row.get(1)?))
1079        })
1080        .optional()
1081    }
1082
1083    fn store_unique_value_blob(
1084        conn: &rusqlite::Connection,
1085        table: Arc<DatabaseTable>,
1086        key: &str,
1087        value: &[u8],
1088    ) -> rusqlite::Result<()> {
1089        let mut stmt = conn.prepare_cached(&table.str_set_unique_value)?;
1090
1091        let changed = stmt.execute(params![key, value])?;
1092        // Never panic in a pool closure: a panic mid-transaction leaves the
1093        // connection holding a stale open transaction
1094        if changed > 1 {
1095            return Err(rusqlite::Error::StatementChangedRows(changed));
1096        }
1097        if changed == 0 {
1098            return Err(rusqlite::Error::QueryReturnedNoRows);
1099        }
1100        Ok(())
1101    }
1102
1103    fn remove_unique_value_blob(
1104        conn: &rusqlite::Connection,
1105        table: Arc<DatabaseTable>,
1106        key: &str,
1107    ) -> rusqlite::Result<()> {
1108        let mut stmt = conn.prepare_cached(&table.str_remove_unique_value)?;
1109
1110        let _ = stmt.execute([key])?;
1111
1112        Ok(())
1113    }
1114
1115    fn remove_and_return_unique_value_blob(
1116        conn: &rusqlite::Connection,
1117        table: Arc<DatabaseTable>,
1118        key: &str,
1119    ) -> rusqlite::Result<Option<Vec<u8>>> {
1120        let mut stmt = conn.prepare_cached(&table.str_remove_and_return_unique_value)?;
1121
1122        stmt.query_row([key], |row| -> rusqlite::Result<Vec<u8>> { row.get(0) })
1123            .optional()
1124    }
1125
1126    fn remove_unique_value_blob_like(
1127        conn: &rusqlite::Connection,
1128        table: Arc<DatabaseTable>,
1129        like: &str,
1130    ) -> rusqlite::Result<usize> {
1131        let mut stmt = conn.prepare_cached(&table.str_remove_unique_value_like)?;
1132
1133        let changed = stmt.execute([like])?;
1134        Ok(changed)
1135    }
1136
1137    fn open_resize_columns(&self) -> io::Result<()> {
1138        let columns = self.columns();
1139        let this = self.clone();
1140        self.conn_blocking(move |conn| {
1141			// First see if we have a control table with the number of columns
1142			conn.execute("CREATE TABLE IF NOT EXISTS control (id INTEGER PRIMARY KEY AUTOINCREMENT, [key] TEXT UNIQUE, value TEXT)", [])?;
1143
1144            // Get column count
1145            let on_disk_columns =
1146                Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
1147
1148            // If desired column count is less than or equal to current column count, then allow it, but restrict access to columns
1149            if columns <= on_disk_columns {
1150                return Ok(());
1151            }
1152
1153            // Otherwise resize and add other columns
1154            for cn in on_disk_columns..columns {
1155                // Create the column table if we don't have it
1156                Self::create_column_table(conn, cn)?;
1157            }
1158            Self::set_unique_value(
1159                conn,
1160                this.control_table(),
1161                "columns",
1162                columns,
1163            )?;
1164            Ok(())
1165        }).map_err(io::Error::other)
1166    }
1167
1168    fn stats_read(&self, count: usize, bytes: usize) {
1169        let mut inner = self.inner.lock();
1170        inner.current_stats.reads += count as u64;
1171        inner.overall_stats.reads += count as u64;
1172        inner.current_stats.bytes_read += bytes as u64;
1173        inner.overall_stats.bytes_read += bytes as u64;
1174    }
1175
1176    fn stats_write(&self, sizes: &[usize]) {
1177        if sizes.is_empty() {
1178            return;
1179        }
1180
1181        let mut inner = self.inner.lock();
1182        for &size in sizes {
1183            inner.current_stats.record_write(size);
1184            inner.overall_stats.record_write(size);
1185        }
1186    }
1187
1188    fn stats_tx_write(&self, size: usize, duration: Duration) {
1189        let mut inner = self.inner.lock();
1190        inner
1191            .current_stats
1192            .record_tx_write(size, duration.as_micros() as f64);
1193        inner
1194            .overall_stats
1195            .record_tx_write(size, duration.as_micros() as f64);
1196    }
1197
1198    fn stats_delete(&self, count: usize) {
1199        if count == 0 {
1200            return;
1201        }
1202
1203        let mut inner = self.inner.lock();
1204        inner.current_stats.deletes += count as u64;
1205        inner.overall_stats.deletes += count as u64;
1206    }
1207
1208    fn stats_delete_prefix(&self, count: usize) {
1209        if count == 0 {
1210            return;
1211        }
1212
1213        let mut inner = self.inner.lock();
1214        inner.current_stats.prefix_deletes += count as u64;
1215        inner.overall_stats.prefix_deletes += count as u64;
1216    }
1217
1218    fn stats_transaction(&self, count: usize) {
1219        let mut inner = self.inner.lock();
1220        inner.current_stats.transactions += count as u64;
1221        inner.overall_stats.transactions += count as u64;
1222    }
1223}
1224
1225impl KeyValueDB for Database {
1226    fn get(&self, col: u32, key: &[u8]) -> KeyValueDBPinBoxFuture<'_, io::Result<Option<DBValue>>> {
1227        let key_text = key_to_text(key);
1228        let key_len = key.len();
1229
1230        Box::pin(async move {
1231            self.validate_column(col).map_err(io::Error::other)?;
1232            let someval = self
1233                .conn_retry_blocking(|| {
1234                    let that = self.clone();
1235                    let key_text = key_text.clone();
1236                    Box::new(move |conn: &rusqlite::Connection| {
1237                        Self::load_unique_value_blob(conn, that.column_table(col), &key_text)
1238                    })
1239                })
1240                .map_err(io::Error::other)?;
1241            {
1242                match &someval {
1243                    Some(val) => self.stats_read(1, key_len + val.len()),
1244                    None => self.stats_read(1, key_len),
1245                }
1246            }
1247
1248            Ok(someval)
1249        })
1250    }
1251
1252    /// Remove a value by key, returning the old value
1253    fn delete(
1254        &self,
1255        col: u32,
1256        key: &[u8],
1257    ) -> KeyValueDBPinBoxFuture<'_, io::Result<Option<DBValue>>> {
1258        let key_text = key_to_text(key);
1259        let key_len = key.len();
1260
1261        Box::pin(async move {
1262            self.validate_column(col).map_err(io::Error::other)?;
1263            self.conn_retry_blocking(|| {
1264                let that = self.clone();
1265                let key_text = key_text.clone();
1266                Box::new(move |conn: &rusqlite::Connection| {
1267                    let someval = Self::remove_and_return_unique_value_blob(
1268                        conn,
1269                        that.column_table(col),
1270                        &key_text,
1271                    )?;
1272
1273                    match &someval {
1274                        Some(val) => {
1275                            that.stats_read(1, key_len + val.len());
1276                        }
1277                        None => that.stats_read(1, key_len),
1278                    }
1279
1280                    Ok(someval)
1281                })
1282            })
1283            .map_err(io::Error::other)
1284        })
1285    }
1286
1287    fn write(
1288        &self,
1289        transaction: DBTransaction,
1290    ) -> KeyValueDBPinBoxFuture<'_, Result<(), DBTransactionError>> {
1291        let transaction = Arc::new(transaction);
1292        Box::pin(async move {
1293            self.stats_transaction(1);
1294
1295            self.conn_mut_retry(|| {
1296                let that = self.clone();
1297                let transaction_clone = transaction.clone();
1298                Box::new(move |conn: &mut rusqlite::Connection| {
1299                    let mut sizes = Vec::with_capacity(transaction_clone.ops.len());
1300                    let mut total_tx_size = 0;
1301                    let mut deletes = 0usize;
1302                    let mut prefix_deletes = 0usize;
1303                    let start = Instant::now();
1304
1305                    let tx = conn.transaction()?;
1306
1307                    for op in &transaction_clone.ops {
1308                        match op {
1309                            DBOp::Insert { col, key, value } => {
1310                                that.validate_column(*col)?;
1311                                Self::store_unique_value_blob(
1312                                    &tx,
1313                                    that.column_table(*col),
1314                                    &key_to_text(key),
1315                                    value,
1316                                )?;
1317                                sizes.push(key.len() + value.len());
1318                                total_tx_size += key.len() + value.len();
1319                            }
1320                            DBOp::Delete { col, key } => {
1321                                that.validate_column(*col)?;
1322                                Self::remove_unique_value_blob(
1323                                    &tx,
1324                                    that.column_table(*col),
1325                                    &key_to_text(key),
1326                                )?;
1327                                deletes += 1;
1328                            }
1329                            DBOp::DeletePrefix { col, prefix } => {
1330                                that.validate_column(*col)?;
1331                                Self::remove_unique_value_blob_like(
1332                                    &tx,
1333                                    that.column_table(*col),
1334                                    &(like_key_to_text(prefix) + "%"),
1335                                )?;
1336                                prefix_deletes += 1;
1337                            }
1338                        }
1339                    }
1340                    tx.commit()?;
1341
1342                    let duration = Instant::now() - start;
1343                    that.stats_write(&sizes);
1344                    that.stats_tx_write(total_tx_size, duration);
1345                    that.stats_delete(deletes);
1346                    that.stats_delete_prefix(prefix_deletes);
1347
1348                    Ok(())
1349                })
1350            })
1351            .await
1352            .map_err(io::Error::other)
1353            .map_err(|error| {
1354                let transaction = transaction.as_ref().clone();
1355                DBTransactionError { error, transaction }
1356            })
1357        })
1358    }
1359
1360    fn iter<
1361        'a,
1362        T: Send + 'static,
1363        C: Send + 'static,
1364        F: FnMut(&mut C, DBKeyValueRef) -> io::Result<Option<T>> + Send + Sync + 'static,
1365    >(
1366        &'a self,
1367        col: u32,
1368        prefix: Option<&'a [u8]>,
1369        context: C,
1370        f: F,
1371    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
1372        let opt_prefix_query = prefix.map(|p| like_key_to_text(p) + "%");
1373        Box::pin(async move {
1374            if col >= self.columns() {
1375                return Err(io::Error::from(io::ErrorKind::NotFound));
1376            }
1377
1378            let context = Arc::new(Mutex::new(Some(context)));
1379            let f = Arc::new(Mutex::new(f));
1380            let called = Arc::new(AtomicBool::new(false));
1381
1382            let make = || {
1383                let that = self.clone();
1384                let context_ref = context.clone();
1385                let f_ref = f.clone();
1386                let called = called.clone();
1387                let opt_prefix_query = opt_prefix_query.clone();
1388                Box::new(move |conn: &rusqlite::Connection| {
1389                    let mut context = context_ref.lock();
1390                    let context = context.as_mut().unwrap();
1391                    let mut f_guard = f_ref.lock();
1392                    let f = &mut *f_guard;
1393
1394                    let mut stmt;
1395                    let mut rows;
1396                    if let Some(prefix_query) = opt_prefix_query {
1397                        stmt = conn.prepare_cached(&that.column_table(col).str_iter_with_prefix)?;
1398                        rows = stmt.query([prefix_query])?;
1399                    } else {
1400                        stmt = conn.prepare_cached(&that.column_table(col).str_iter_no_prefix)?;
1401                        rows = stmt.query([])?;
1402                    }
1403
1404                    let mut sw = 0usize;
1405                    let mut sbw = 0usize;
1406
1407                    let out = loop {
1408                        match rows.next()? {
1409                            // Iterated value
1410                            Some(row) => {
1411                                let kt: String = row.get(0)?;
1412                                let v: Vec<u8> = row.get(1)?;
1413                                let k: Vec<u8> = match text_to_key(&kt) {
1414                                    Err(e) => {
1415                                        break Err(io::Error::other(format!(
1416                                            "SQLite row get column 0 text convert error: {:?}",
1417                                            e
1418                                        )));
1419                                    }
1420                                    Ok(v) => v,
1421                                };
1422
1423                                sw += 1;
1424                                sbw += k.len() + v.len();
1425
1426                                called.store(true, Ordering::Relaxed);
1427                                match f(context, (&k, &v)) {
1428                                    Ok(None) => (),
1429                                    // Callback early termination
1430                                    Ok(Some(out)) => break Ok(Some(out)),
1431                                    // Callback error termination
1432                                    Err(e) => break Err(e),
1433                                }
1434                            }
1435                            // Natural iterator termination
1436                            None => {
1437                                break Ok(None);
1438                            }
1439                        }
1440                    };
1441
1442                    that.stats_read(sw, sbw);
1443
1444                    Ok(out)
1445                })
1446            };
1447
1448            let (pool, generation) = self.current();
1449            let res = match pool.conn(make()).await {
1450                Err(e) if !called.load(Ordering::Relaxed) => match self.recover(e, generation) {
1451                    Ok(()) => self.current().0.conn(make()).await,
1452                    Err(e) => Err(e),
1453                },
1454                Err(e) => {
1455                    // Rows already reached the callback: repair for the next
1456                    // caller, but this scan cannot safely restart
1457                    if self.repair_enabled() && error_is_corruption(&e) {
1458                        let _ = self.repair_live(e.to_string(), generation);
1459                    }
1460                    Err(e)
1461                }
1462                r => r,
1463            };
1464            let res = res.map_err(io::Error::other)?;
1465
1466            let context = context.lock().take().unwrap();
1467
1468            res.map(|x| (context, x))
1469        })
1470    }
1471
1472    fn iter_keys<
1473        'a,
1474        T: Send + 'static,
1475        C: Send + 'static,
1476        F: FnMut(&mut C, DBKeyRef) -> io::Result<Option<T>> + Send + Sync + 'static,
1477    >(
1478        &'a self,
1479        col: u32,
1480        prefix: Option<&'a [u8]>,
1481        context: C,
1482        f: F,
1483    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
1484        let opt_prefix_query = prefix.map(|p| like_key_to_text(p) + "%");
1485        Box::pin(async move {
1486            if col >= self.columns() {
1487                return Err(io::Error::from(io::ErrorKind::NotFound));
1488            }
1489
1490            let context = Arc::new(Mutex::new(Some(context)));
1491            let f = Arc::new(Mutex::new(f));
1492            let called = Arc::new(AtomicBool::new(false));
1493
1494            let make = || {
1495                let that = self.clone();
1496                let context_ref = context.clone();
1497                let f_ref = f.clone();
1498                let called = called.clone();
1499                let opt_prefix_query = opt_prefix_query.clone();
1500                Box::new(move |conn: &rusqlite::Connection| {
1501                    let mut context = context_ref.lock();
1502                    let context = context.as_mut().unwrap();
1503                    let mut f_guard = f_ref.lock();
1504                    let f = &mut *f_guard;
1505
1506                    let mut stmt;
1507                    let mut rows;
1508                    if let Some(prefix_query) = opt_prefix_query {
1509                        stmt =
1510                            conn.prepare_cached(&that.column_table(col).str_iter_keys_with_prefix)?;
1511                        rows = stmt.query([prefix_query])?;
1512                    } else {
1513                        stmt =
1514                            conn.prepare_cached(&that.column_table(col).str_iter_keys_no_prefix)?;
1515                        rows = stmt.query([])?;
1516                    }
1517
1518                    let mut sw = 0usize;
1519                    let mut sbw = 0usize;
1520
1521                    let out = loop {
1522                        match rows.next()? {
1523                            // Iterated value
1524                            Some(row) => {
1525                                let kt: String = row.get(0)?;
1526                                let k: Vec<u8> = match text_to_key(&kt) {
1527                                    Err(e) => {
1528                                        break Err(io::Error::other(format!(
1529                                            "SQLite row get column 0 text convert error: {:?}",
1530                                            e
1531                                        )));
1532                                    }
1533                                    Ok(v) => v,
1534                                };
1535
1536                                sw += 1;
1537                                sbw += k.len();
1538
1539                                called.store(true, Ordering::Relaxed);
1540                                match f(context, &k) {
1541                                    Ok(None) => (),
1542                                    // Callback early termination
1543                                    Ok(Some(out)) => break Ok(Some(out)),
1544                                    // Callback error termination
1545                                    Err(e) => break Err(e),
1546                                }
1547                            }
1548                            // Natural iterator termination
1549                            None => {
1550                                break Ok(None);
1551                            }
1552                        }
1553                    };
1554
1555                    that.stats_read(sw, sbw);
1556
1557                    Ok(out)
1558                })
1559            };
1560
1561            let (pool, generation) = self.current();
1562            let res = match pool.conn(make()).await {
1563                Err(e) if !called.load(Ordering::Relaxed) => match self.recover(e, generation) {
1564                    Ok(()) => self.current().0.conn(make()).await,
1565                    Err(e) => Err(e),
1566                },
1567                Err(e) => {
1568                    // Rows already reached the callback: repair for the next
1569                    // caller, but this scan cannot safely restart
1570                    if self.repair_enabled() && error_is_corruption(&e) {
1571                        let _ = self.repair_live(e.to_string(), generation);
1572                    }
1573                    Err(e)
1574                }
1575                r => r,
1576            };
1577            let res = res.map_err(io::Error::other)?;
1578
1579            let context = context.lock().take().unwrap();
1580
1581            res.map(|x| (context, x))
1582        })
1583    }
1584
1585    fn io_stats(&self, kind: IoStatsKind) -> IoStats {
1586        fn duration_since(timestamp_microseconds: u64) -> Duration {
1587            std::time::SystemTime::now()
1588                .duration_since(std::time::UNIX_EPOCH)
1589                .map_or(Duration::from_micros(0), |time| {
1590                    let now = time.as_micros() as u64;
1591                    if now >= timestamp_microseconds {
1592                        Duration::from_micros(now - timestamp_microseconds)
1593                    } else {
1594                        Duration::from_micros(0)
1595                    }
1596                })
1597        }
1598
1599        let mut inner = self.inner.lock();
1600        match kind {
1601            IoStatsKind::Overall => {
1602                let mut stats = inner.overall_stats.clone();
1603                stats.span = duration_since(stats.started);
1604                stats
1605            }
1606            IoStatsKind::SincePrevious => {
1607                let mut stats = inner.current_stats.clone();
1608                stats.span = duration_since(stats.started);
1609                inner.current_stats = IoStats::empty();
1610                stats
1611            }
1612        }
1613    }
1614
1615    fn num_columns(&self) -> io::Result<u32> {
1616        self.conn_retry_blocking(|| {
1617            let this = self.clone();
1618            Box::new(move |conn: &rusqlite::Connection| {
1619                Self::get_unique_value(conn, this.control_table(), "columns", 0u32)
1620            })
1621        })
1622        .map_err(io::Error::other)
1623    }
1624
1625    fn num_keys(&self, col: u32) -> KeyValueDBPinBoxFuture<'_, io::Result<u64>> {
1626        Box::pin(async move {
1627            self.conn_retry(|| {
1628                Box::new(move |conn: &rusqlite::Connection| {
1629                    conn.query_row(
1630                        &format!("SELECT Count(*) FROM {}", get_column_table_name(col)),
1631                        [],
1632                        |row| -> rusqlite::Result<u64> { row.get(0) },
1633                    )
1634                })
1635            })
1636            .await
1637            .map_err(|_| io::Error::from(io::ErrorKind::NotFound))
1638        })
1639    }
1640
1641    /// Check for the existence of a value by key.
1642    fn has_key<'a>(
1643        &'a self,
1644        col: u32,
1645        key: &'a [u8],
1646    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
1647        let key_text = key_to_text(key);
1648        let key_len = key.len();
1649
1650        Box::pin(async move {
1651            self.validate_column(col).map_err(io::Error::other)?;
1652            let someval = self
1653                .conn_retry_blocking(|| {
1654                    let that = self.clone();
1655                    let key_text = key_text.clone();
1656                    Box::new(move |conn: &rusqlite::Connection| {
1657                        Self::has_value(conn, that.column_table(col), &key_text)
1658                    })
1659                })
1660                .map_err(io::Error::other)?;
1661
1662            self.stats_read(1, key_len);
1663
1664            Ok(someval)
1665        })
1666    }
1667
1668    /// Check for the existence of a value by prefix.
1669    fn has_prefix<'a>(
1670        &'a self,
1671        col: u32,
1672        prefix: &'a [u8],
1673    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
1674        let prefix_len = prefix.len();
1675        let prefix_text = like_key_to_text(prefix) + "%";
1676
1677        Box::pin(async move {
1678            self.validate_column(col).map_err(io::Error::other)?;
1679            let someval = self
1680                .conn_retry_blocking(|| {
1681                    let that = self.clone();
1682                    let prefix_text = prefix_text.clone();
1683                    Box::new(move |conn: &rusqlite::Connection| {
1684                        Self::has_value_like(conn, that.column_table(col), &prefix_text)
1685                    })
1686                })
1687                .map_err(io::Error::other)?;
1688
1689            self.stats_read(1, prefix_len);
1690
1691            Ok(someval)
1692        })
1693    }
1694
1695    /// Get the first value matching the given prefix.
1696    fn first_with_prefix<'a>(
1697        &'a self,
1698        col: u32,
1699        prefix: &'a [u8],
1700    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBKeyValue>>> {
1701        let prefix_len = prefix.len();
1702        let like = like_key_to_text(prefix) + "%";
1703
1704        Box::pin(async move {
1705            self.validate_column(col).map_err(io::Error::other)?;
1706            let someval = self
1707                .conn_retry_blocking(|| {
1708                    let that = self.clone();
1709                    let like = like.clone();
1710                    Box::new(move |conn: &rusqlite::Connection| {
1711                        Self::load_first_value_blob_like(conn, that.column_table(col), &like)
1712                    })
1713                })
1714                .map_err(io::Error::other)?;
1715
1716            self.stats_read(1, prefix_len);
1717
1718            match someval {
1719                Some((kt, val)) => match text_to_key(&kt) {
1720                    Err(e) => Err(io::Error::other(format!(
1721                        "SQLite row get column 0 text convert error: {:?}",
1722                        e
1723                    ))),
1724                    Ok(k) => Ok(Some((k, val))),
1725                },
1726                None => Ok(None),
1727            }
1728        })
1729    }
1730
1731    /// Vacuum database
1732    fn cleanup(&self) -> KeyValueDBPinBoxFuture<'_, io::Result<()>> {
1733        Box::pin(async { self.vacuum().await.map_err(io::Error::other) })
1734    }
1735}
1736
1737#[cfg(test)]
1738mod tests {
1739
1740    use super::*;
1741    use keyvaluedb_shared_tests as st;
1742    use tempfile::Builder as TempfileBuilder;
1743
1744    fn create(columns: u32) -> io::Result<Database> {
1745        let tempfile = TempfileBuilder::new()
1746            .prefix("")
1747            .tempfile()?
1748            .path()
1749            .to_path_buf();
1750        let config = DatabaseConfig::new().with_columns(columns);
1751        Database::open(tempfile, config)
1752    }
1753
1754    fn create_vacuum_mode(columns: u32, vacuum_mode: VacuumMode) -> io::Result<Database> {
1755        let tempfile = TempfileBuilder::new()
1756            .prefix("")
1757            .tempfile()?
1758            .path()
1759            .to_path_buf();
1760        let config = DatabaseConfig::new()
1761            .with_columns(columns)
1762            .with_vacuum_mode(vacuum_mode);
1763        Database::open(tempfile, config)
1764    }
1765
1766    #[tokio::test]
1767    async fn get_fails_with_non_existing_column() -> io::Result<()> {
1768        let db = create(1)?;
1769        st::test_get_fails_with_non_existing_column(db).await
1770    }
1771
1772    #[tokio::test]
1773    async fn num_keys() -> io::Result<()> {
1774        let db = create(1)?;
1775        st::test_num_keys(db).await
1776    }
1777
1778    #[tokio::test]
1779    async fn put_and_get() -> io::Result<()> {
1780        let db = create(1)?;
1781        st::test_put_and_get(db).await
1782    }
1783
1784    #[tokio::test]
1785    async fn delete_and_get() -> io::Result<()> {
1786        let db = create(1)?;
1787        st::test_delete_and_get(db).await
1788    }
1789
1790    #[tokio::test]
1791    async fn delete_and_get_single() -> io::Result<()> {
1792        let db = create(1)?;
1793        st::test_delete_and_get_single(db).await
1794    }
1795
1796    #[tokio::test]
1797    async fn delete_prefix() -> io::Result<()> {
1798        let db = create(st::DELETE_PREFIX_NUM_COLUMNS)?;
1799        st::test_delete_prefix(db).await
1800    }
1801
1802    #[tokio::test]
1803    async fn iter() -> io::Result<()> {
1804        let db = create(1)?;
1805        st::test_iter(db).await
1806    }
1807
1808    #[tokio::test]
1809    async fn iter_keys() -> io::Result<()> {
1810        let db = create(1)?;
1811        st::test_iter_keys(db).await
1812    }
1813
1814    #[tokio::test]
1815    async fn iter_with_prefix() -> io::Result<()> {
1816        let db = create(1)?;
1817        st::test_iter_with_prefix(db).await
1818    }
1819
1820    #[tokio::test]
1821    async fn complex() -> io::Result<()> {
1822        let db = create(1)?;
1823        st::test_complex(db).await
1824    }
1825
1826    #[tokio::test]
1827    async fn cleanup() -> io::Result<()> {
1828        let db = create(1)?;
1829        st::test_cleanup(db).await?;
1830
1831        let db = create_vacuum_mode(1, VacuumMode::None)?;
1832        st::test_cleanup(db).await?;
1833
1834        let db = create_vacuum_mode(1, VacuumMode::Incremental)?;
1835        st::test_cleanup(db).await?;
1836
1837        let db = create_vacuum_mode(1, VacuumMode::Full)?;
1838        st::test_cleanup(db).await?;
1839
1840        let tempfile = TempfileBuilder::new()
1841            .prefix("")
1842            .tempfile()?
1843            .path()
1844            .to_path_buf();
1845        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::None);
1846        let db = Database::open(tempfile.clone(), config)?;
1847        st::test_cleanup(db).await?;
1848
1849        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::Incremental);
1850        let db = Database::open(tempfile.clone(), config)?;
1851        st::test_cleanup(db).await?;
1852
1853        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::Full);
1854        let db = Database::open(tempfile.clone(), config)?;
1855        st::test_cleanup(db).await?;
1856
1857        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::None);
1858        let db = Database::open(tempfile, config)?;
1859        st::test_cleanup(db).await?;
1860
1861        Ok(())
1862    }
1863
1864    #[tokio::test]
1865    async fn stats() -> io::Result<()> {
1866        let db = create(st::IO_STATS_NUM_COLUMNS)?;
1867        st::test_io_stats(db).await
1868    }
1869
1870    #[tokio::test]
1871    #[should_panic]
1872    async fn db_config_with_zero_columns() {
1873        let _cfg = DatabaseConfig::new().with_columns(0);
1874    }
1875
1876    #[tokio::test]
1877    #[should_panic]
1878    async fn open_db_with_zero_columns() {
1879        let cfg = DatabaseConfig::new().with_columns(0);
1880        let _db = Database::open("", cfg);
1881    }
1882
1883    #[tokio::test]
1884    async fn add_columns() {
1885        let config_1 = DatabaseConfig::default();
1886        let config_5 = DatabaseConfig::new().with_columns(5);
1887
1888        let tempfile = TempfileBuilder::new()
1889            .prefix("")
1890            .tempfile()
1891            .unwrap()
1892            .path()
1893            .to_path_buf();
1894
1895        // open 1, add 4.
1896        {
1897            let db = Database::open(&tempfile, config_1).unwrap();
1898            assert_eq!(db.num_columns().unwrap(), 1);
1899
1900            for i in 2..=5 {
1901                db.add_column().unwrap();
1902                assert_eq!(db.num_columns().unwrap(), i);
1903            }
1904        }
1905
1906        // reopen as 5.
1907        {
1908            let db = Database::open(&tempfile, config_5).unwrap();
1909            assert_eq!(db.num_columns().unwrap(), 5);
1910        }
1911    }
1912
1913    #[tokio::test]
1914    async fn remove_columns() {
1915        let config_1 = DatabaseConfig::default();
1916        let config_5 = DatabaseConfig::new().with_columns(5);
1917
1918        let tempfile = TempfileBuilder::new()
1919            .prefix("drop_columns")
1920            .tempfile()
1921            .unwrap()
1922            .path()
1923            .to_path_buf();
1924
1925        // open 5, remove 4.
1926        {
1927            let db = Database::open(&tempfile, config_5).expect("open with 5 columns");
1928            assert_eq!(db.num_columns().unwrap(), 5);
1929
1930            for i in (1..5).rev() {
1931                db.remove_last_column().unwrap();
1932                assert_eq!(db.num_columns().unwrap(), i);
1933            }
1934        }
1935
1936        // reopen as 1.
1937        {
1938            let db = Database::open(&tempfile, config_1).unwrap();
1939            assert_eq!(db.num_columns().unwrap(), 1);
1940        }
1941    }
1942
1943    #[tokio::test]
1944    async fn test_num_keys() {
1945        let tempfile = TempfileBuilder::new()
1946            .prefix("")
1947            .tempfile()
1948            .unwrap()
1949            .path()
1950            .to_path_buf();
1951        let config = DatabaseConfig::new().with_columns(1);
1952        let db = Database::open(tempfile, config).unwrap();
1953
1954        assert_eq!(
1955            db.num_keys(0).await.unwrap(),
1956            0,
1957            "database is empty after creation"
1958        );
1959        let key1 = b"beef";
1960        let mut batch = db.transaction();
1961        batch.put(0, key1, key1);
1962        db.write(batch).await.unwrap();
1963        assert_eq!(
1964            db.num_keys(0).await.unwrap(),
1965            1,
1966            "adding a key increases the count"
1967        );
1968    }
1969}