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