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