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#[derive(Copy, Clone, Debug, Eq, PartialEq)]
26pub enum VacuumMode {
27 None,
28 Incremental,
29 Full,
30}
31
32#[derive(Copy, Clone, Debug, Eq, PartialEq)]
34pub enum RepairCheck {
35 None,
38 Quick,
41 Full,
44}
45
46pub type RepairCallback = Arc<dyn Fn(&RepairReport) + Send + Sync>;
48
49#[derive(Clone)]
51pub struct DatabaseConfig {
52 pub columns: u32,
55 pub flags: OpenFlags,
57 pub num_conns: usize,
59 pub vacuum_mode: VacuumMode,
61 pub repair_on_corrupt: bool,
64 pub repair_check: RepairCheck,
66 pub on_repair: Option<RepairCallback>,
68}
69
70impl DatabaseConfig {
71 pub fn new() -> Self {
73 Default::default()
74 }
75
76 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 pub fn with_repair_on_corrupt(self, repair_on_corrupt: bool) -> Self {
84 Self {
85 repair_on_corrupt,
86 ..self
87 }
88 }
89
90 pub fn with_repair_check(self, repair_check: RepairCheck) -> Self {
92 Self {
93 repair_check,
94 ..self
95 }
96 }
97
98 pub fn with_on_repair(self, on_repair: RepairCallback) -> Self {
100 Self {
101 on_repair: Some(on_repair),
102 ..self
103 }
104 }
105
106 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 pub fn with_flags(self, flags: OpenFlags) -> Self {
119 Self { flags, ..self }
120 }
121
122 pub fn with_num_conns(self, num_conns: usize) -> Self {
124 Self { num_conns, ..self }
125 }
126
127 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
152fn 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
165fn error_is_corruption(e: &Error) -> bool {
167 matches!(e, Error::Rusqlite(e) if code_is_corruption(e))
168}
169
170fn 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
185pub 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#[derive(Debug, Clone)]
256pub struct RepairReport {
257 pub detected: String,
259 pub rows_recovered: u64,
261 pub partial_tables: u32,
263 pub corrupt_path: PathBuf,
265}
266
267struct DatabaseState {
271 pool: Pool,
272 generation: u64,
273 repair_report: Option<RepairReport>,
274}
275
276#[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
292fn 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
307pub 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 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 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 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 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 conn.pragma_update(None, "busy_timeout", 2000)?;
394 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 pub fn repair_report(&self) -> Option<RepairReport> {
477 self.unlocked_inner.state.read().repair_report.clone()
478 }
479
480 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 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 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 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 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 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 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 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 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 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 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 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 if self.current().1 == seen_generation {
769 return Err(e);
770 }
771 Ok(())
772 }
773
774 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 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 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 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 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 pub fn transaction(&self) -> DBTransaction {
856 DBTransaction::new()
857 }
858
859 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 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 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 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 conn.execute("CREATE TABLE IF NOT EXISTS control (id INTEGER PRIMARY KEY AUTOINCREMENT, [key] TEXT UNIQUE, value TEXT)", [])?;
1040
1041 let on_disk_columns =
1043 Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
1044
1045 if columns <= on_disk_columns {
1047 return Ok(());
1048 }
1049
1050 for cn in on_disk_columns..columns {
1052 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 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 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 Ok(Some(out)) => break Ok(Some(out)),
1328 Err(e) => break Err(e),
1330 }
1331 }
1332 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 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 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 Ok(Some(out)) => break Ok(Some(out)),
1441 Err(e) => break Err(e),
1443 }
1444 }
1445 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 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 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 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 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 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 {
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 {
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 {
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 {
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}