keyvaluedb_sqlite/

lib.rs

#![deny(clippy::all)]

mod tools;

pub use async_sqlite::rusqlite::OpenFlags;
use async_sqlite::rusqlite::{params, OptionalExtension as _};
use async_sqlite::*;
use keyvaluedb::{
    DBKeyRef, DBKeyValue, DBKeyValueRef, DBOp, DBTransaction, DBTransactionError, DBValue, IoStats,
    IoStatsKind, KeyValueDB, KeyValueDBPinBoxFuture,
};
use parking_lot::Mutex;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{
    io,
    path::{Path, PathBuf},
    str::FromStr,
};
use tools::*;

///////////////////////////////////////////////////////////////////////////////

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum VacuumMode {
    None,
    Incremental,
    Full,
}

/// Database configuration
#[derive(Clone)]
pub struct DatabaseConfig {
    /// Set number of columns.
    /// The number of columns must not be zero.
    pub columns: u32,
    /// Set flags used to open the database
    pub flags: OpenFlags,
    /// Number of connections to open
    pub num_conns: usize,
    /// Vacuum mode
    pub vacuum_mode: VacuumMode,
    /// Run an integrity check at open and, on corruption, salvage the readable
    /// rows into a fresh database, keeping the damaged files beside it
    pub repair_on_corrupt: bool,
}

impl DatabaseConfig {
    /// Create new `DatabaseConfig` with default parameters
    pub fn new() -> Self {
        Default::default()
    }

    /// Set the number of columns. `columns` must not be zero.
    pub fn with_columns(self, columns: u32) -> Self {
        assert!(columns > 0, "the number of columns must not be zero");
        Self { columns, ..self }
    }

    /// Enable integrity check and salvage at open
    pub fn with_repair_on_corrupt(self, repair_on_corrupt: bool) -> Self {
        Self {
            repair_on_corrupt,
            ..self
        }
    }

    /// Sets the flags to 'in-memory database'
    pub fn with_in_memory(self) -> Self {
        Self {
            flags: OpenFlags::SQLITE_OPEN_READ_WRITE
                | OpenFlags::SQLITE_OPEN_CREATE
                | OpenFlags::SQLITE_OPEN_NO_MUTEX
                | OpenFlags::SQLITE_OPEN_MEMORY,
            ..self
        }
    }

    /// Replaces all the flags
    pub fn with_flags(self, flags: OpenFlags) -> Self {
        Self { flags, ..self }
    }

    /// Sets the number of connections for this database
    pub fn with_num_conns(self, num_conns: usize) -> Self {
        Self { num_conns, ..self }
    }

    /// Set the vacuum mode used by 'cleanup'
    pub fn with_vacuum_mode(self, vacuum_mode: VacuumMode) -> Self {
        Self {
            vacuum_mode,
            ..self
        }
    }
}

impl Default for DatabaseConfig {
    fn default() -> DatabaseConfig {
        DatabaseConfig {
            columns: 1,
            flags: OpenFlags::SQLITE_OPEN_READ_WRITE
                | OpenFlags::SQLITE_OPEN_CREATE
                | OpenFlags::SQLITE_OPEN_NO_MUTEX,
            num_conns: 1,
            vacuum_mode: VacuumMode::None,
            repair_on_corrupt: false,
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

/// Whether an error chain bottoms out in sqlite saying the file is damaged
fn io_error_is_corruption(e: &io::Error) -> bool {
    fn code_is_corruption(e: &rusqlite::Error) -> bool {
        matches!(
            e,
            rusqlite::Error::SqliteFailure(f, _) if matches!(
                f.code,
                rusqlite::ErrorCode::DatabaseCorrupt | rusqlite::ErrorCode::NotADatabase
            )
        )
    }
    let mut src: Option<&(dyn std::error::Error + 'static)> = e.get_ref().map(|b| b as _);
    while let Some(cur) = src {
        if let Some(e) = cur.downcast_ref::<rusqlite::Error>() {
            return code_is_corruption(e);
        }
        if let Some(Error::Rusqlite(e)) = cur.downcast_ref::<Error>() {
            return code_is_corruption(e);
        }
        src = cur.source();
    }
    false
}

///////////////////////////////////////////////////////////////////////////////

/// An sqlite table with its statement strings
pub struct DatabaseTable {
    _table: String,
    str_has_value: String,
    str_has_value_like: String,
    str_get_unique_value: String,
    str_get_first_value_like: String,
    str_set_unique_value: String,
    str_remove_unique_value: String,
    str_remove_and_return_unique_value: String,
    str_remove_unique_value_like: String,
    str_iter_with_prefix: String,
    str_iter_no_prefix: String,
    str_iter_keys_with_prefix: String,
    str_iter_keys_no_prefix: String,
}

impl DatabaseTable {
    pub fn new(table: String) -> Self {
        let str_has_value = format!("SELECT 1 FROM {} WHERE [key] = ? LIMIT 1", table);
        let str_has_value_like = format!(
            "SELECT 1 FROM {} WHERE [key] LIKE ? ESCAPE '\\' LIMIT 1",
            table
        );
        let str_get_unique_value = format!("SELECT value FROM {} WHERE [key] = ? LIMIT 1", table);
        let str_get_first_value_like = format!(
            "SELECT key, value FROM {} WHERE [key] LIKE ? ESCAPE '\\' LIMIT 1",
            table
        );
        let str_set_unique_value = format!(
            "INSERT OR REPLACE INTO {} ([key], value) VALUES(?, ?)",
            table
        );
        let str_remove_unique_value = format!("DELETE FROM {} WHERE [key] = ?", table);
        let str_remove_and_return_unique_value =
            format!("DELETE FROM {} WHERE [key] = ? RETURNING value", table);
        let str_remove_unique_value_like =
            format!("DELETE FROM {} WHERE [key] LIKE ? ESCAPE '\\'", table);
        let str_iter_with_prefix = format!(
            "SELECT key, value FROM {} WHERE [key] LIKE ? ESCAPE '\\'",
            table
        );
        let str_iter_no_prefix = format!("SELECT key, value FROM {}", table);
        let str_iter_keys_with_prefix =
            format!("SELECT key FROM {} WHERE [key] LIKE ? ESCAPE '\\'", table);
        let str_iter_keys_no_prefix = format!("SELECT key FROM {}", table);

        Self {
            _table: table,
            str_has_value,
            str_has_value_like,
            str_get_unique_value,
            str_get_first_value_like,
            str_set_unique_value,
            str_remove_unique_value,
            str_remove_and_return_unique_value,
            str_remove_unique_value_like,
            str_iter_with_prefix,
            str_iter_no_prefix,
            str_iter_keys_with_prefix,
            str_iter_keys_no_prefix,
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

/// What the automatic corruption repair did, kept on the database it produced
#[derive(Debug, Clone)]
pub struct RepairReport {
    /// The integrity failure that triggered the repair
    pub detected: String,
    /// Rows salvaged across the control table and all column tables
    pub rows_recovered: u64,
    /// Tables whose salvage scan ended early on a damaged page
    pub partial_tables: u32,
    /// Where the damaged database files were moved
    pub corrupt_path: PathBuf,
}

/// An sqlite key-value database fulfilling the `KeyValueDB` trait
pub struct DatabaseUnlockedInner {
    path: PathBuf,
    config: DatabaseConfig,
    pool: Pool,
    control_table: Arc<DatabaseTable>,
    column_tables: Vec<Arc<DatabaseTable>>,
    repair_report: Option<RepairReport>,
}

impl Drop for DatabaseUnlockedInner {
    fn drop(&mut self) {
        let _ = self.pool.close_blocking();
    }
}

pub struct DatabaseInner {
    overall_stats: IoStats,
    current_stats: IoStats,
}

#[derive(Clone)]
pub struct Database {
    unlocked_inner: Arc<DatabaseUnlockedInner>,
    inner: Arc<Mutex<DatabaseInner>>,
}

impl Database {
    ////////////////////////////////////////////////////////////////
    // Initialization

    pub fn open<P: AsRef<Path>>(path: P, config: DatabaseConfig) -> io::Result<Self> {
        let path = PathBuf::from(path.as_ref());
        let db = match Self::open_raw(&path, &config, None) {
            Ok(db) => db,
            Err(e) if config.repair_on_corrupt && io_error_is_corruption(&e) => {
                return Self::repair(None, &path, &config, e.to_string());
            }
            Err(e) => return Err(e),
        };
        if !config.repair_on_corrupt {
            return Ok(db);
        }
        match db.quick_check_blocking() {
            Ok(None) => Ok(db),
            Ok(Some(detected)) => Self::repair(Some(db), &path, &config, detected),
            Err(e) if io_error_is_corruption(&e) => {
                Self::repair(Some(db), &path, &config, e.to_string())
            }
            Err(e) => Err(e),
        }
    }

    fn open_raw(
        path: &Path,
        config: &DatabaseConfig,
        repair_report: Option<RepairReport>,
    ) -> io::Result<Self> {
        let config = config.clone();
        assert_ne!(config.columns, 0, "number of columns must be >= 1");

        let path = PathBuf::from(path);
        let flags = config.flags;

        let mut column_tables = vec![];
        for n in 0..config.columns {
            column_tables.push(Arc::new(DatabaseTable::new(get_column_table_name(n))))
        }
        let control_table = Arc::new(DatabaseTable::new("control".to_string()));

        let pool_builder = PoolBuilder::new()
            .path(&path)
            .flags(flags)
            .num_conns(config.num_conns);

        let pool = pool_builder.open_blocking().map_err(io::Error::other)?;

        let out = Self {
            unlocked_inner: Arc::new(DatabaseUnlockedInner {
                path,
                config,
                pool,
                control_table,
                column_tables,
                repair_report,
            }),
            inner: Arc::new(Mutex::new(DatabaseInner {
                overall_stats: IoStats::empty(),
                current_stats: IoStats::empty(),
            })),
        };

        let vacuum_mode = out.config().vacuum_mode;

        out.conn_blocking(move |conn| {
            // Don't rely on STATEMENT_CACHE_DEFAULT_CAPACITY in rusqlite, set it explicitly
            conn.set_prepared_statement_cache_capacity(256);

            conn.pragma_update(None, "case_sensitive_like", "ON")?;
            conn.pragma_update(None, "journal_mode", "WAL")?;
            conn.pragma_update(None, "synchronous", "normal")?;
            conn.pragma_update(None, "journal_size_limit", 6144000)?;
            conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")?;

            match vacuum_mode {
                VacuumMode::None | VacuumMode::Full => {
                    let current: u32 =
                        conn.pragma_query_value(None, "auto_vacuum", |x| x.get(0))?;
                    if current != 0 {
                        conn.execute("VACUUM", [])?;
                        conn.pragma_update(None, "auto_vacuum", 0)?;
                    }
                }
                VacuumMode::Incremental => {
                    let current: u32 =
                        conn.pragma_query_value(None, "auto_vacuum", |x| x.get(0))?;
                    if current != 2 {
                        conn.execute("VACUUM", [])?;
                        conn.pragma_update(None, "auto_vacuum", "2")?;
                    }
                }
            }

            Ok(())
        })
        .map_err(io::Error::other)?;

        out.open_resize_columns()?;

        Ok(out)
    }

    pub fn path(&self) -> PathBuf {
        self.unlocked_inner.path.clone()
    }

    /// What the open-time repair did, if this database came from one
    pub fn repair_report(&self) -> Option<&RepairReport> {
        self.unlocked_inner.repair_report.as_ref()
    }

    /// Integrity check on one connection. None = clean; Some = what is wrong.
    /// A check that cannot even run to completion errs with the corruption it hit.
    fn quick_check_blocking(&self) -> io::Result<Option<String>> {
        let problems = self
            .conn_blocking(|conn| {
                let mut stmt = conn.prepare("PRAGMA quick_check(8)")?;
                let mut rows = stmt.query([])?;
                let mut problems: Vec<String> = Vec::new();
                while let Some(row) = rows.next()? {
                    problems.push(row.get(0)?);
                }
                Ok(problems)
            })
            .map_err(io::Error::other)?;
        if problems.len() == 1 && problems[0] == "ok" {
            return Ok(None);
        }
        Ok(Some(problems.join("; ")))
    }

    /// Rebuild a damaged database from whatever rows are still readable.
    ///
    /// The damaged files move beside the store as `<name>.corrupt` rather than
    /// being destroyed: corruption that sqlite's own crash-safety should have
    /// made impossible is evidence worth keeping. `old` is the still-open
    /// handle when the damage was found by the integrity check; None when the
    /// file would not even open, in which case nothing is salvageable and the
    /// result is a fresh empty database.
    fn repair(
        old: Option<Self>,
        path: &Path,
        config: &DatabaseConfig,
        detected: String,
    ) -> io::Result<Self> {
        let suffixes = ["", "-wal", "-shm"];
        let sibling = |base: &Path, ext: &str, suffix: &str| {
            let mut os = base.as_os_str().to_owned();
            os.push(ext);
            os.push(suffix);
            PathBuf::from(os)
        };

        // Build the replacement in a sibling file, never in place
        let tmp = sibling(path, ".repairing", "");
        for suffix in &suffixes {
            let _ = std::fs::remove_file(sibling(path, ".repairing", suffix));
        }
        let fresh = Self::open_raw(&tmp, config, None)?;

        let mut rows_recovered = 0u64;
        let mut partial_tables = 0u32;
        if let Some(old) = &old {
            // Salvage every table the old database says it has, or at least
            // the configured set if its control table cannot say
            let old_columns = old
                .conn_blocking({
                    let this = old.clone();
                    move |conn| Self::get_unique_value(conn, this.control_table(), "columns", 0u32)
                })
                .unwrap_or(0)
                .max(config.columns);
            let mut tables = vec!["control".to_string()];
            for cn in 0..old_columns {
                tables.push(get_column_table_name(cn));
            }
            for table in tables {
                let (rows, partial) = old
                    .conn_blocking({
                        let table = table.clone();
                        move |conn| Ok(Self::salvage_table_rows(conn, &table))
                    })
                    .map_err(io::Error::other)?;
                if partial {
                    partial_tables += 1;
                }
                if rows.is_empty() {
                    continue;
                }
                rows_recovered += rows.len() as u64;
                // INSERT OR IGNORE: values the fresh open already wrote
                // (control's own column count) win over salvaged ones
                fresh
                    .conn_blocking({
                        let table = table.clone();
                        move |conn| {
                            let tx = conn.unchecked_transaction()?;
                            {
                                let mut stmt = tx.prepare(&format!(
                                    "INSERT OR IGNORE INTO {} ([key], value) VALUES (?, ?)",
                                    table
                                ))?;
                                for (k, v) in &rows {
                                    stmt.execute(params![k, v])?;
                                }
                            }
                            tx.commit()
                        }
                    })
                    .map_err(io::Error::other)?;
            }
        }

        // Close both so every file is settled before the swap
        drop(fresh);
        drop(old);

        let corrupt_path = sibling(path, ".corrupt", "");
        for suffix in &suffixes {
            let _ = std::fs::remove_file(sibling(path, ".corrupt", suffix));
            let _ = std::fs::rename(sibling(path, "", suffix), sibling(path, ".corrupt", suffix));
            let _ = std::fs::rename(
                sibling(path, ".repairing", suffix),
                sibling(path, "", suffix),
            );
        }

        Self::open_raw(
            path,
            config,
            Some(RepairReport {
                detected,
                rows_recovered,
                partial_tables,
                corrupt_path,
            }),
        )
    }

    /// Read every row a damaged table will still yield, stopping at the first
    /// page it cannot, keeping what came before it
    fn salvage_table_rows(
        conn: &rusqlite::Connection,
        table: &str,
    ) -> (Vec<(String, rusqlite::types::Value)>, bool) {
        let mut out = Vec::new();
        let mut partial = true;
        if let Ok(mut stmt) = conn.prepare(&format!("SELECT [key], value FROM {}", table)) {
            if let Ok(mut rows) = stmt.query([]) {
                loop {
                    match rows.next() {
                        Ok(Some(row)) => {
                            if let (Ok(k), Ok(v)) = (row.get(0), row.get(1)) {
                                out.push((k, v));
                            }
                        }
                        Ok(None) => {
                            partial = false;
                            break;
                        }
                        Err(_) => break,
                    }
                }
            }
        }
        (out, partial)
    }

    pub fn config(&self) -> DatabaseConfig {
        self.unlocked_inner.config.clone()
    }

    pub fn columns(&self) -> u32 {
        self.unlocked_inner.config.columns
    }

    pub fn control_table(&self) -> Arc<DatabaseTable> {
        self.unlocked_inner.control_table.clone()
    }

    pub fn column_table(&self, col: u32) -> Arc<DatabaseTable> {
        self.unlocked_inner.column_tables[col as usize].clone()
    }

    pub fn conn_blocking<T, F>(&self, func: F) -> Result<T, Error>
    where
        F: FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
        T: Send + 'static,
    {
        self.unlocked_inner.pool.conn_blocking(func)
    }

    pub async fn conn<T, F>(&self, func: F) -> Result<T, Error>
    where
        F: FnOnce(&rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
        T: Send + 'static,
    {
        self.unlocked_inner.pool.conn(func).await
    }

    pub async fn conn_mut<T, F>(&self, func: F) -> Result<T, Error>
    where
        F: FnOnce(&mut rusqlite::Connection) -> Result<T, rusqlite::Error> + Send + 'static,
        T: Send + 'static,
    {
        self.unlocked_inner.pool.conn_mut(func).await
    }

    ////////////////////////////////////////////////////////////////
    // Low level operations

    /// Remove the last column family in the database. The deletion is definitive.
    pub fn remove_last_column(&self) -> Result<(), Error> {
        let this = self.clone();
        self.conn_blocking(move |conn| {
            let columns = Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
            if columns == 0 {
                return Err(rusqlite::Error::QueryReturnedNoRows);
            }
            Self::set_unique_value(conn, this.control_table(), "columns", columns - 1)?;

            conn.execute(
                &format!("DROP TABLE {}", get_column_table_name(columns - 1)),
                [],
            )?;
            Ok(())
        })
    }

    /// Add a new column family to the DB.
    pub fn add_column(&self) -> Result<(), Error> {
        let this = self.clone();

        self.conn_blocking(move |conn| {
            let columns = Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;
            Self::set_unique_value(conn, this.control_table(), "columns", columns + 1)?;
            Self::create_column_table(conn, columns)
        })
    }
    /// Helper to create new transaction for this database.
    pub fn transaction(&self) -> DBTransaction {
        DBTransaction::new()
    }

    /// Vacuum database
    pub async fn vacuum(&self) -> Result<(), Error> {
        match self.config().vacuum_mode {
            VacuumMode::None => {
                self.conn(move |conn| {
                    conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")?;
                    Ok(())
                })
                .await
            }
            VacuumMode::Incremental => {
                self.conn(move |conn| {
                    conn.execute("PRAGMA incremental_vacuum", [])?;
                    conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")?;
                    Ok(())
                })
                .await
            }
            VacuumMode::Full => {
                self.conn(move |conn| {
                    conn.execute("VACUUM", [])?;
                    conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")?;
                    Ok(())
                })
                .await
            }
        }
    }

    ////////////////////////////////////////////////////////////////
    // Internal helpers

    fn validate_column(&self, col: u32) -> rusqlite::Result<()> {
        if col >= self.columns() {
            return Err(rusqlite::Error::InvalidColumnIndex(col as usize));
        }
        Ok(())
    }

    fn create_column_table(conn: &rusqlite::Connection, column: u32) -> rusqlite::Result<()> {
        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)
    }

    fn get_unique_value<V>(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
        default: V,
    ) -> rusqlite::Result<V>
    where
        V: FromStr,
    {
        let mut stmt = conn.prepare_cached(&table.str_get_unique_value)?;

        if let Ok(found) = stmt.query_row([key], |row| -> rusqlite::Result<String> { row.get(0) }) {
            if let Ok(v) = V::from_str(&found) {
                return Ok(v);
            }
        }
        Ok(default)
    }

    fn set_unique_value<V>(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
        value: V,
    ) -> rusqlite::Result<()>
    where
        V: ToString,
    {
        let mut stmt = conn.prepare_cached(&table.str_set_unique_value)?;

        let changed = stmt.execute([key, value.to_string().as_str()])?;

        assert!(
            changed <= 1,
            "multiple changes to unique key should not occur"
        );
        if changed == 0 {
            return Err(rusqlite::Error::QueryReturnedNoRows);
        }

        Ok(())
    }

    fn has_value(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
    ) -> rusqlite::Result<bool> {
        let mut stmt = conn.prepare_cached(&table.str_has_value)?;
        stmt.exists([key])
    }

    fn has_value_like(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
    ) -> rusqlite::Result<bool> {
        let mut stmt = conn.prepare_cached(&table.str_has_value_like)?;
        stmt.exists([key])
    }

    fn load_unique_value_blob(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
    ) -> rusqlite::Result<Option<Vec<u8>>> {
        let mut stmt = conn.prepare_cached(&table.str_get_unique_value)?;

        stmt.query_row([key], |row| -> rusqlite::Result<Vec<u8>> { row.get(0) })
            .optional()
    }

    fn load_first_value_blob_like(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        like: &str,
    ) -> rusqlite::Result<Option<(String, Vec<u8>)>> {
        let mut stmt = conn.prepare_cached(&table.str_get_first_value_like)?;

        stmt.query_row([like], |row| -> rusqlite::Result<(String, Vec<u8>)> {
            Ok((row.get(0)?, row.get(1)?))
        })
        .optional()
    }

    fn store_unique_value_blob(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
        value: &[u8],
    ) -> rusqlite::Result<()> {
        let mut stmt = conn.prepare_cached(&table.str_set_unique_value)?;

        let changed = stmt.execute(params![key, value])?;
        assert!(
            changed <= 1,
            "multiple changes to unique key should not occur"
        );
        if changed == 0 {
            return Err(rusqlite::Error::QueryReturnedNoRows);
        }
        Ok(())
    }

    fn remove_unique_value_blob(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
    ) -> rusqlite::Result<()> {
        let mut stmt = conn.prepare_cached(&table.str_remove_unique_value)?;

        let _ = stmt.execute([key])?;

        Ok(())
    }

    fn remove_and_return_unique_value_blob(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        key: &str,
    ) -> rusqlite::Result<Option<Vec<u8>>> {
        let mut stmt = conn.prepare_cached(&table.str_remove_and_return_unique_value)?;

        stmt.query_row([key], |row| -> rusqlite::Result<Vec<u8>> { row.get(0) })
            .optional()
    }

    fn remove_unique_value_blob_like(
        conn: &rusqlite::Connection,
        table: Arc<DatabaseTable>,
        like: &str,
    ) -> rusqlite::Result<usize> {
        let mut stmt = conn.prepare_cached(&table.str_remove_unique_value_like)?;

        let changed = stmt.execute([like])?;
        Ok(changed)
    }

    fn open_resize_columns(&self) -> io::Result<()> {
        let columns = self.columns();
        let this = self.clone();
        self.conn_blocking(move |conn| {
			// First see if we have a control table with the number of columns
			conn.execute("CREATE TABLE IF NOT EXISTS control (id INTEGER PRIMARY KEY AUTOINCREMENT, [key] TEXT UNIQUE, value TEXT)", [])?;

            // Get column count
            let on_disk_columns =
                Self::get_unique_value(conn, this.control_table(), "columns", 0u32)?;

            // If desired column count is less than or equal to current column count, then allow it, but restrict access to columns
            if columns <= on_disk_columns {
                return Ok(());
            }

            // Otherwise resize and add other columns
            for cn in on_disk_columns..columns {
                // Create the column table if we don't have it
                Self::create_column_table(conn, cn)?;
            }
            Self::set_unique_value(
                conn,
                this.control_table(),
                "columns",
                columns,
            )?;
            Ok(())
        }).map_err(io::Error::other)
    }

    fn stats_read(&self, count: usize, bytes: usize) {
        let mut inner = self.inner.lock();
        inner.current_stats.reads += count as u64;
        inner.overall_stats.reads += count as u64;
        inner.current_stats.bytes_read += bytes as u64;
        inner.overall_stats.bytes_read += bytes as u64;
    }

    fn stats_write(&self, sizes: &[usize]) {
        if sizes.is_empty() {
            return;
        }

        let mut inner = self.inner.lock();
        for &size in sizes {
            inner.current_stats.record_write(size);
            inner.overall_stats.record_write(size);
        }
    }

    fn stats_tx_write(&self, size: usize, duration: Duration) {
        let mut inner = self.inner.lock();
        inner
            .current_stats
            .record_tx_write(size, duration.as_micros() as f64);
        inner
            .overall_stats
            .record_tx_write(size, duration.as_micros() as f64);
    }

    fn stats_delete(&self, count: usize) {
        if count == 0 {
            return;
        }

        let mut inner = self.inner.lock();
        inner.current_stats.deletes += count as u64;
        inner.overall_stats.deletes += count as u64;
    }

    fn stats_delete_prefix(&self, count: usize) {
        if count == 0 {
            return;
        }

        let mut inner = self.inner.lock();
        inner.current_stats.prefix_deletes += count as u64;
        inner.overall_stats.prefix_deletes += count as u64;
    }

    fn stats_transaction(&self, count: usize) {
        let mut inner = self.inner.lock();
        inner.current_stats.transactions += count as u64;
        inner.overall_stats.transactions += count as u64;
    }
}

impl KeyValueDB for Database {
    fn get(&self, col: u32, key: &[u8]) -> KeyValueDBPinBoxFuture<'_, io::Result<Option<DBValue>>> {
        let key_text = key_to_text(key);
        let key_len = key.len();

        Box::pin(async move {
            let that = self.clone();
            that.validate_column(col).map_err(io::Error::other)?;
            let someval = self
                .conn_blocking(move |conn| {
                    Self::load_unique_value_blob(conn, that.column_table(col), &key_text)
                })
                .map_err(io::Error::other)?;
            {
                match &someval {
                    Some(val) => self.stats_read(1, key_len + val.len()),
                    None => self.stats_read(1, key_len),
                }
            }

            Ok(someval)
        })
    }

    /// Remove a value by key, returning the old value
    fn delete(
        &self,
        col: u32,
        key: &[u8],
    ) -> KeyValueDBPinBoxFuture<'_, io::Result<Option<DBValue>>> {
        let key_text = key_to_text(key);
        let key_len = key.len();

        Box::pin(async move {
            let that = self.clone();
            that.validate_column(col).map_err(io::Error::other)?;
            self.conn_blocking(move |conn| {
                let someval = Self::remove_and_return_unique_value_blob(
                    conn,
                    that.column_table(col),
                    &key_text,
                )?;

                match &someval {
                    Some(val) => {
                        that.stats_read(1, key_len + val.len());
                    }
                    None => that.stats_read(1, key_len),
                }

                Ok(someval)
            })
            .map_err(io::Error::other)
        })
    }

    fn write(
        &self,
        transaction: DBTransaction,
    ) -> KeyValueDBPinBoxFuture<'_, Result<(), DBTransactionError>> {
        let transaction = Arc::new(transaction);
        Box::pin(async move {
            self.stats_transaction(1);

            let that = self.clone();
            let transaction_clone = transaction.clone();
            self.conn_mut(move |conn| {
                let mut sizes = Vec::with_capacity(transaction_clone.ops.len());
                let mut total_tx_size = 0;
                let mut deletes = 0usize;
                let mut prefix_deletes = 0usize;
                let start = Instant::now();

                let tx = conn.transaction()?;

                for op in &transaction_clone.ops {
                    match op {
                        DBOp::Insert { col, key, value } => {
                            that.validate_column(*col)?;
                            Self::store_unique_value_blob(
                                &tx,
                                that.column_table(*col),
                                &key_to_text(key),
                                value,
                            )?;
                            sizes.push(key.len() + value.len());
                            total_tx_size += key.len() + value.len();
                        }
                        DBOp::Delete { col, key } => {
                            that.validate_column(*col)?;
                            Self::remove_unique_value_blob(
                                &tx,
                                that.column_table(*col),
                                &key_to_text(key),
                            )?;
                            deletes += 1;
                        }
                        DBOp::DeletePrefix { col, prefix } => {
                            that.validate_column(*col)?;
                            Self::remove_unique_value_blob_like(
                                &tx,
                                that.column_table(*col),
                                &(like_key_to_text(prefix) + "%"),
                            )?;
                            prefix_deletes += 1;
                        }
                    }
                }
                tx.commit()?;

                let duration = Instant::now() - start;
                that.stats_write(&sizes);
                that.stats_tx_write(total_tx_size, duration);
                that.stats_delete(deletes);
                that.stats_delete_prefix(prefix_deletes);

                Ok(())
            })
            .await
            .map_err(io::Error::other)
            .map_err(|error| {
                let transaction = transaction.as_ref().clone();
                DBTransactionError { error, transaction }
            })
        })
    }

    fn iter<
        'a,
        T: Send + 'static,
        C: Send + 'static,
        F: FnMut(&mut C, DBKeyValueRef) -> io::Result<Option<T>> + Send + Sync + 'static,
    >(
        &'a self,
        col: u32,
        prefix: Option<&'a [u8]>,
        context: C,
        mut f: F,
    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
        let opt_prefix_query = prefix.map(|p| like_key_to_text(p) + "%");
        Box::pin(async move {
            if col >= self.columns() {
                return Err(io::Error::from(io::ErrorKind::NotFound));
            }

            let that = self.clone();
            let context = Arc::new(Mutex::new(Some(context)));
            let context_ref = context.clone();

            let res = self
                .conn(move |conn| {
                    let mut context = context_ref.lock();
                    let context = context.as_mut().unwrap();

                    let mut stmt;
                    let mut rows;
                    if let Some(prefix_query) = opt_prefix_query {
                        stmt = match conn
                            .prepare_cached(&that.column_table(col).str_iter_with_prefix)
                        {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                        rows = match stmt.query([prefix_query]) {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                    } else {
                        stmt = match conn.prepare_cached(&that.column_table(col).str_iter_no_prefix)
                        {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                        rows = match stmt.query([]) {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                    }

                    let mut sw = 0usize;
                    let mut sbw = 0usize;

                    let out = loop {
                        match rows.next() {
                            // Iterated value
                            Ok(Some(row)) => {
                                let kt: String = match row.get(0) {
                                    Err(e) => {
                                        break Err(io::Error::other(e));
                                    }
                                    Ok(v) => v,
                                };
                                let v: Vec<u8> = match row.get(1) {
                                    Err(e) => {
                                        break Err(io::Error::other(e));
                                    }
                                    Ok(v) => v,
                                };
                                let k: Vec<u8> = match text_to_key(&kt) {
                                    Err(e) => {
                                        break Err(io::Error::other(format!(
                                            "SQLite row get column 0 text convert error: {:?}",
                                            e
                                        )));
                                    }
                                    Ok(v) => v,
                                };

                                sw += 1;
                                sbw += k.len() + v.len();

                                match f(context, (&k, &v)) {
                                    Ok(None) => (),
                                    Ok(Some(out)) => {
                                        // Callback early termination
                                        that.stats_read(sw, sbw);
                                        break Ok(Some(out));
                                    }
                                    Err(e) => {
                                        // Callback error termination
                                        that.stats_read(sw, sbw);
                                        break Err(e);
                                    }
                                }
                            }
                            // Natural iterator termination
                            Ok(None) => {
                                break Ok(None);
                            }
                            // Error iterator termination
                            Err(_) => {
                                break Ok(None);
                            }
                        }
                    };

                    that.stats_read(sw, sbw);

                    Ok(out)
                })
                .await
                .map_err(io::Error::other)?;

            let context = context.lock().take().unwrap();

            res.map(|x| (context, x))
        })
    }

    fn iter_keys<
        'a,
        T: Send + 'static,
        C: Send + 'static,
        F: FnMut(&mut C, DBKeyRef) -> io::Result<Option<T>> + Send + Sync + 'static,
    >(
        &'a self,
        col: u32,
        prefix: Option<&'a [u8]>,
        context: C,
        mut f: F,
    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
        let opt_prefix_query = prefix.map(|p| like_key_to_text(p) + "%");
        Box::pin(async move {
            if col >= self.columns() {
                return Err(io::Error::from(io::ErrorKind::NotFound));
            }

            let that = self.clone();
            let context = Arc::new(Mutex::new(Some(context)));
            let context_ref = context.clone();

            let res = self
                .conn(move |conn| {
                    let mut context = context_ref.lock();
                    let context = context.as_mut().unwrap();

                    let mut stmt;
                    let mut rows;
                    if let Some(prefix_query) = opt_prefix_query {
                        stmt = match conn
                            .prepare_cached(&that.column_table(col).str_iter_keys_with_prefix)
                        {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                        rows = match stmt.query([prefix_query]) {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                    } else {
                        stmt = match conn
                            .prepare_cached(&that.column_table(col).str_iter_keys_no_prefix)
                        {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                        rows = match stmt.query([]) {
                            Ok(v) => v,
                            Err(e) => {
                                return Ok(Err(io::Error::other(e)));
                            }
                        };
                    }

                    let mut sw = 0usize;
                    let mut sbw = 0usize;

                    let out = loop {
                        match rows.next() {
                            // Iterated value
                            Ok(Some(row)) => {
                                let kt: String = match row.get(0) {
                                    Err(e) => {
                                        break Err(io::Error::other(e));
                                    }
                                    Ok(v) => v,
                                };
                                let k: Vec<u8> = match text_to_key(&kt) {
                                    Err(e) => {
                                        break Err(io::Error::other(format!(
                                            "SQLite row get column 0 text convert error: {:?}",
                                            e
                                        )));
                                    }
                                    Ok(v) => v,
                                };

                                sw += 1;
                                sbw += k.len();

                                match f(context, &k) {
                                    Ok(None) => (),
                                    Ok(Some(out)) => {
                                        // Callback early termination
                                        that.stats_read(sw, sbw);
                                        break Ok(Some(out));
                                    }
                                    Err(e) => {
                                        // Callback error termination
                                        that.stats_read(sw, sbw);
                                        break Err(e);
                                    }
                                }
                            }
                            // Natural iterator termination
                            Ok(None) => {
                                break Ok(None);
                            }
                            // Error iterator termination
                            Err(_) => {
                                break Ok(None);
                            }
                        }
                    };

                    that.stats_read(sw, sbw);

                    Ok(out)
                })
                .await
                .map_err(io::Error::other)?;

            let context = context.lock().take().unwrap();

            res.map(|x| (context, x))
        })
    }

    fn io_stats(&self, kind: IoStatsKind) -> IoStats {
        fn duration_since(timestamp_microseconds: u64) -> Duration {
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_or(Duration::from_micros(0), |time| {
                    let now = time.as_micros() as u64;
                    if now >= timestamp_microseconds {
                        Duration::from_micros(now - timestamp_microseconds)
                    } else {
                        Duration::from_micros(0)
                    }
                })
        }

        let mut inner = self.inner.lock();
        match kind {
            IoStatsKind::Overall => {
                let mut stats = inner.overall_stats.clone();
                stats.span = duration_since(stats.started);
                stats
            }
            IoStatsKind::SincePrevious => {
                let mut stats = inner.current_stats.clone();
                stats.span = duration_since(stats.started);
                inner.current_stats = IoStats::empty();
                stats
            }
        }
    }

    fn num_columns(&self) -> io::Result<u32> {
        let this = self.clone();
        self.conn_blocking(move |conn| {
            Self::get_unique_value(conn, this.control_table(), "columns", 0u32)
        })
        .map_err(io::Error::other)
    }

    fn num_keys(&self, col: u32) -> KeyValueDBPinBoxFuture<'_, io::Result<u64>> {
        let this = self.clone();
        Box::pin(async move {
            this.conn(move |conn| {
                conn.query_row(
                    &format!("SELECT Count(*) FROM {}", get_column_table_name(col)),
                    [],
                    |row| -> rusqlite::Result<u64> { row.get(0) },
                )
            })
            .await
            .map_err(|_| io::Error::from(io::ErrorKind::NotFound))
        })
    }

    /// Check for the existence of a value by key.
    fn has_key<'a>(
        &'a self,
        col: u32,
        key: &'a [u8],
    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
        let key_text = key_to_text(key);
        let key_len = key.len();

        Box::pin(async move {
            let that = self.clone();
            that.validate_column(col).map_err(io::Error::other)?;
            let someval = self
                .conn_blocking(move |conn| Self::has_value(conn, that.column_table(col), &key_text))
                .map_err(io::Error::other)?;

            self.stats_read(1, key_len);

            Ok(someval)
        })
    }

    /// Check for the existence of a value by prefix.
    fn has_prefix<'a>(
        &'a self,
        col: u32,
        prefix: &'a [u8],
    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
        let prefix_len = prefix.len();
        let prefix_text = like_key_to_text(prefix) + "%";

        Box::pin(async move {
            let that = self.clone();
            that.validate_column(col).map_err(io::Error::other)?;
            let someval = self
                .conn_blocking(move |conn| {
                    Self::has_value_like(conn, that.column_table(col), &prefix_text)
                })
                .map_err(io::Error::other)?;

            self.stats_read(1, prefix_len);

            Ok(someval)
        })
    }

    /// Get the first value matching the given prefix.
    fn first_with_prefix<'a>(
        &'a self,
        col: u32,
        prefix: &'a [u8],
    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBKeyValue>>> {
        let prefix_len = prefix.len();
        let like = like_key_to_text(prefix) + "%";

        Box::pin(async move {
            let that = self.clone();
            that.validate_column(col).map_err(io::Error::other)?;
            let someval = self
                .conn_blocking(move |conn| {
                    Self::load_first_value_blob_like(conn, that.column_table(col), &like)
                })
                .map_err(io::Error::other)?;

            self.stats_read(1, prefix_len);

            match someval {
                Some((kt, val)) => match text_to_key(&kt) {
                    Err(e) => Err(io::Error::other(format!(
                        "SQLite row get column 0 text convert error: {:?}",
                        e
                    ))),
                    Ok(k) => Ok(Some((k, val))),
                },
                None => Ok(None),
            }
        })
    }

    /// Vacuum database
    fn cleanup(&self) -> KeyValueDBPinBoxFuture<'_, io::Result<()>> {
        Box::pin(async { self.vacuum().await.map_err(io::Error::other) })
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use keyvaluedb_shared_tests as st;
    use tempfile::Builder as TempfileBuilder;

    fn create(columns: u32) -> io::Result<Database> {
        let tempfile = TempfileBuilder::new()
            .prefix("")
            .tempfile()?
            .path()
            .to_path_buf();
        let config = DatabaseConfig::new().with_columns(columns);
        Database::open(tempfile, config)
    }

    fn create_vacuum_mode(columns: u32, vacuum_mode: VacuumMode) -> io::Result<Database> {
        let tempfile = TempfileBuilder::new()
            .prefix("")
            .tempfile()?
            .path()
            .to_path_buf();
        let config = DatabaseConfig::new()
            .with_columns(columns)
            .with_vacuum_mode(vacuum_mode);
        Database::open(tempfile, config)
    }

    #[tokio::test]
    async fn get_fails_with_non_existing_column() -> io::Result<()> {
        let db = create(1)?;
        st::test_get_fails_with_non_existing_column(db).await
    }

    #[tokio::test]
    async fn num_keys() -> io::Result<()> {
        let db = create(1)?;
        st::test_num_keys(db).await
    }

    #[tokio::test]
    async fn put_and_get() -> io::Result<()> {
        let db = create(1)?;
        st::test_put_and_get(db).await
    }

    #[tokio::test]
    async fn delete_and_get() -> io::Result<()> {
        let db = create(1)?;
        st::test_delete_and_get(db).await
    }

    #[tokio::test]
    async fn delete_and_get_single() -> io::Result<()> {
        let db = create(1)?;
        st::test_delete_and_get_single(db).await
    }

    #[tokio::test]
    async fn delete_prefix() -> io::Result<()> {
        let db = create(st::DELETE_PREFIX_NUM_COLUMNS)?;
        st::test_delete_prefix(db).await
    }

    #[tokio::test]
    async fn iter() -> io::Result<()> {
        let db = create(1)?;
        st::test_iter(db).await
    }

    #[tokio::test]
    async fn iter_keys() -> io::Result<()> {
        let db = create(1)?;
        st::test_iter_keys(db).await
    }

    #[tokio::test]
    async fn iter_with_prefix() -> io::Result<()> {
        let db = create(1)?;
        st::test_iter_with_prefix(db).await
    }

    #[tokio::test]
    async fn complex() -> io::Result<()> {
        let db = create(1)?;
        st::test_complex(db).await
    }

    #[tokio::test]
    async fn cleanup() -> io::Result<()> {
        let db = create(1)?;
        st::test_cleanup(db).await?;

        let db = create_vacuum_mode(1, VacuumMode::None)?;
        st::test_cleanup(db).await?;

        let db = create_vacuum_mode(1, VacuumMode::Incremental)?;
        st::test_cleanup(db).await?;

        let db = create_vacuum_mode(1, VacuumMode::Full)?;
        st::test_cleanup(db).await?;

        let tempfile = TempfileBuilder::new()
            .prefix("")
            .tempfile()?
            .path()
            .to_path_buf();
        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::None);
        let db = Database::open(tempfile.clone(), config)?;
        st::test_cleanup(db).await?;

        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::Incremental);
        let db = Database::open(tempfile.clone(), config)?;
        st::test_cleanup(db).await?;

        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::Full);
        let db = Database::open(tempfile.clone(), config)?;
        st::test_cleanup(db).await?;

        let config = DatabaseConfig::new().with_vacuum_mode(VacuumMode::None);
        let db = Database::open(tempfile, config)?;
        st::test_cleanup(db).await?;

        Ok(())
    }

    #[tokio::test]
    async fn stats() -> io::Result<()> {
        let db = create(st::IO_STATS_NUM_COLUMNS)?;
        st::test_io_stats(db).await
    }

    #[tokio::test]
    #[should_panic]
    async fn db_config_with_zero_columns() {
        let _cfg = DatabaseConfig::new().with_columns(0);
    }

    #[tokio::test]
    #[should_panic]
    async fn open_db_with_zero_columns() {
        let cfg = DatabaseConfig::new().with_columns(0);
        let _db = Database::open("", cfg);
    }

    #[tokio::test]
    async fn add_columns() {
        let config_1 = DatabaseConfig::default();
        let config_5 = DatabaseConfig::new().with_columns(5);

        let tempfile = TempfileBuilder::new()
            .prefix("")
            .tempfile()
            .unwrap()
            .path()
            .to_path_buf();

        // open 1, add 4.
        {
            let db = Database::open(&tempfile, config_1).unwrap();
            assert_eq!(db.num_columns().unwrap(), 1);

            for i in 2..=5 {
                db.add_column().unwrap();
                assert_eq!(db.num_columns().unwrap(), i);
            }
        }

        // reopen as 5.
        {
            let db = Database::open(&tempfile, config_5).unwrap();
            assert_eq!(db.num_columns().unwrap(), 5);
        }
    }

    #[tokio::test]
    async fn remove_columns() {
        let config_1 = DatabaseConfig::default();
        let config_5 = DatabaseConfig::new().with_columns(5);

        let tempfile = TempfileBuilder::new()
            .prefix("drop_columns")
            .tempfile()
            .unwrap()
            .path()
            .to_path_buf();

        // open 5, remove 4.
        {
            let db = Database::open(&tempfile, config_5).expect("open with 5 columns");
            assert_eq!(db.num_columns().unwrap(), 5);

            for i in (1..5).rev() {
                db.remove_last_column().unwrap();
                assert_eq!(db.num_columns().unwrap(), i);
            }
        }

        // reopen as 1.
        {
            let db = Database::open(&tempfile, config_1).unwrap();
            assert_eq!(db.num_columns().unwrap(), 1);
        }
    }

    #[tokio::test]
    async fn test_num_keys() {
        let tempfile = TempfileBuilder::new()
            .prefix("")
            .tempfile()
            .unwrap()
            .path()
            .to_path_buf();
        let config = DatabaseConfig::new().with_columns(1);
        let db = Database::open(tempfile, config).unwrap();

        assert_eq!(
            db.num_keys(0).await.unwrap(),
            0,
            "database is empty after creation"
        );
        let key1 = b"beef";
        let mut batch = db.transaction();
        batch.put(0, key1, key1);
        db.write(batch).await.unwrap();
        assert_eq!(
            db.num_keys(0).await.unwrap(),
            1,
            "adding a key increases the count"
        );
    }
}