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