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
23const WAL_AUTOCHECKPOINT_PAGES: u32 = 256;
27const 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#[derive(Copy, Clone, Debug, Eq, PartialEq)]
39pub enum RepairCheck {
40 None,
43 Quick,
46 Full,
49}
50
51pub type RepairCallback = Arc<dyn Fn(&RepairReport) + Send + Sync>;
53
54#[derive(Clone)]
56pub struct DatabaseConfig {
57 pub columns: u32,
60 pub flags: OpenFlags,
62 pub num_conns: usize,
64 pub vacuum_mode: VacuumMode,
66 pub repair_on_corrupt: bool,
69 pub repair_check: RepairCheck,
71 pub repair_check_byte_limit: Option<u64>,
76 pub on_repair: Option<RepairCallback>,
78 pub wal_autocheckpoint_pages: u32,
81 pub wal_journal_size_limit: u32,
85}
86
87impl DatabaseConfig {
88 pub fn new() -> Self {
90 Default::default()
91 }
92
93 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 pub fn with_repair_on_corrupt(self, repair_on_corrupt: bool) -> Self {
101 Self {
102 repair_on_corrupt,
103 ..self
104 }
105 }
106
107 pub fn with_repair_check(self, repair_check: RepairCheck) -> Self {
109 Self {
110 repair_check,
111 ..self
112 }
113 }
114
115 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 pub fn with_on_repair(self, on_repair: RepairCallback) -> Self {
125 Self {
126 on_repair: Some(on_repair),
127 ..self
128 }
129 }
130
131 pub fn with_wal_autocheckpoint_pages(self, pages: u32) -> Self {
133 Self {
134 wal_autocheckpoint_pages: pages,
135 ..self
136 }
137 }
138
139 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 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 pub fn with_flags(self, flags: OpenFlags) -> Self {
161 Self { flags, ..self }
162 }
163
164 pub fn with_num_conns(self, num_conns: usize) -> Self {
166 Self { num_conns, ..self }
167 }
168
169 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
197fn 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
210fn error_is_corruption(e: &Error) -> bool {
212 matches!(e, Error::Rusqlite(e) if code_is_corruption(e))
213}
214
215fn 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
230pub 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#[derive(Copy, Clone, Debug, Default)]
301pub struct CheckpointResult {
302 pub busy: bool,
304 pub log_frames: i64,
306 pub checkpointed_frames: i64,
308}
309
310impl CheckpointResult {
311 pub fn complete(&self) -> bool {
313 !self.busy && self.log_frames == self.checkpointed_frames
314 }
315}
316
317fn 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#[derive(Debug, Clone)]
332pub struct RepairReport {
333 pub detected: String,
335 pub rows_recovered: u64,
337 pub partial_tables: u32,
339 pub corrupt_path: PathBuf,
341}
342
343struct DatabaseState {
347 pool: Pool,
348 generation: u64,
349 repair_report: Option<RepairReport>,
350}
351
352#[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
368fn 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
374fn 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
389pub 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 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 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 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 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 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 conn.pragma_update(None, "busy_timeout", 2000)?;
486 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 pub fn repair_report(&self) -> Option<RepairReport> {
569 self.unlocked_inner.state.read().repair_report.clone()
570 }
571
572 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 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 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 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 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 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 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 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 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 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 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 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 if self.current().1 == seen_generation {
861 return Err(e);
862 }
863 Ok(())
864 }
865
866 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 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 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 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 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 pub fn transaction(&self) -> DBTransaction {
948 DBTransaction::new()
949 }
950
951 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 pub async fn checkpoint(&self) -> Result<CheckpointResult, Error> {
974 self.conn(wal_checkpoint_truncate).await
975 }
976
977 pub fn checkpoint_blocking(&self) -> Result<CheckpointResult, Error> {
979 self.conn_blocking(wal_checkpoint_truncate)
980 }
981
982 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 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 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 conn.execute("CREATE TABLE IF NOT EXISTS control (id INTEGER PRIMARY KEY AUTOINCREMENT, [key] TEXT UNIQUE, value TEXT)", [])?;
1142
1143 let on_disk_columns =
1145 Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
1146
1147 if columns <= on_disk_columns {
1149 return Ok(());
1150 }
1151
1152 for cn in on_disk_columns..columns {
1154 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 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 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 Ok(Some(out)) => break Ok(Some(out)),
1430 Err(e) => break Err(e),
1432 }
1433 }
1434 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 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 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 Ok(Some(out)) => break Ok(Some(out)),
1543 Err(e) => break Err(e),
1545 }
1546 }
1547 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 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 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 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 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 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 {
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 {
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 {
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 {
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}