ic-sqlite-vfs 2.0.0

//! Public canister API for the product-facing SQLite facade.
//!
//! Methods are intentionally synchronous. They call `Db::update` or `Db::query`,
//! so a SQLite transaction cannot cross an `await` boundary.

use crate::db::migrate::Migration;
use crate::db::value::{to_sql_ref, ToSql};
use crate::stable::memory;
use crate::stable::memory_manager::{MemoryId, MemoryManager};
use crate::stable::meta::Superblock;
use crate::stable::raw_memory::DefaultMemoryImpl;
use crate::Db;
use candid::CandidType;
use serde::Deserialize;
use std::cell::RefCell;

#[cfg(feature = "canister-api-test-failpoints")]
mod sqlite_feature_probe;

const MIGRATIONS: &[Migration] = &[
    Migration {
        version: 1,
        sql: "CREATE TABLE kv (
            key TEXT PRIMARY KEY NOT NULL,
            value TEXT NOT NULL
        );",
    },
    Migration {
        version: 2,
        sql: "ALTER TABLE kv ADD COLUMN note TEXT;",
    },
];
const MAX_KV_GET_MANY_KEYS: usize = 1_000;
const MAX_KV_KEY_BYTES: usize = 256;
const MAX_KV_VALUE_BYTES: usize = 64 * 1024;
const MAX_KV_NOTE_BYTES: usize = 4 * 1024;
const MAX_CHECKSUM_REFRESH_BYTES: u64 = 4 * 1024 * 1024;
const SQLITE_MEMORY_ID: MemoryId = MemoryId::new(120);

thread_local! {
    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
        RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
}

#[derive(CandidType, Deserialize)]
pub struct DbMeta {
    pub db_size: u64,
    pub stable_pages: u64,
    pub stable_bytes: u64,
    pub schema_version: u64,
    pub last_tx_id: u64,
    pub flags: u64,
    pub checksum: u64,
    pub checksum_stale: bool,
    pub checksum_refreshing: bool,
    pub checksum_refresh_offset: u64,
    pub layout_version: u64,
    pub page_count: u64,
    pub page_table_bytes: u64,
    pub zero_extent_count: u64,
    pub active_bytes: u64,
    pub allocated_bytes: u64,
    pub orphan_bytes_estimate: u64,
    pub orphan_ratio_basis_points: u64,
}

#[derive(CandidType, Deserialize)]
pub struct ChecksumRefresh {
    pub complete: bool,
    pub checksum: u64,
    pub scanned_bytes: u64,
    pub db_size: u64,
}

#[ic_cdk::init]
fn init() {
    init_db();
    must(Db::migrate(MIGRATIONS));
}

#[ic_cdk::post_upgrade]
fn post_upgrade() {
    init_db();
    must(Db::migrate(MIGRATIONS));
}

fn init_db() {
    MEMORY_MANAGER.with(|manager| {
        must(Db::init(manager.borrow().get(SQLITE_MEMORY_ID)));
    });
}

#[ic_cdk::update]
fn kv_put(key: String, value: String) -> Result<(), String> {
    require_controller()?;
    validate_text_len("key", &key, MAX_KV_KEY_BYTES)?;
    validate_text_len("value", &value, MAX_KV_VALUE_BYTES)?;
    Db::update(|connection| {
        connection.execute_text_text(
            "INSERT INTO kv(key, value) VALUES (?1, ?2)
             ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            &key,
            &value,
        )
    })
    .map_err(error_text)
}

#[ic_cdk::query]
fn kv_get(key: String) -> Result<Option<String>, String> {
    validate_text_len("key", &key, MAX_KV_KEY_BYTES)?;
    Db::query(|connection| {
        connection.query_optional_string_text("SELECT value FROM kv WHERE key = ?1", &key)
    })
    .map_err(error_text)
}

#[ic_cdk::query]
fn kv_get_many(keys: Vec<String>) -> Result<Vec<Option<String>>, String> {
    if keys.is_empty() {
        return Ok(Vec::new());
    }
    if keys.len() > MAX_KV_GET_MANY_KEYS {
        return Err(format!(
            "kv_get_many accepts at most {MAX_KV_GET_MANY_KEYS} keys"
        ));
    }
    for key in &keys {
        validate_text_len("key", key, MAX_KV_KEY_BYTES)?;
    }
    let sql = values_lookup_sql(keys.len(), "value")?;
    Db::query(|connection| {
        let values = keys
            .iter()
            .map(|key| to_sql_ref(key) as &dyn ToSql)
            .collect::<Vec<_>>();
        connection.query_column::<Option<String>>(&sql, &values)
    })
    .map_err(error_text)
}

#[ic_cdk::update]
fn kv_set_note(key: String, note: String) -> Result<(), String> {
    require_controller()?;
    validate_text_len("key", &key, MAX_KV_KEY_BYTES)?;
    validate_text_len("note", &note, MAX_KV_NOTE_BYTES)?;
    Db::update(|connection| {
        connection.execute_text_text("UPDATE kv SET note = ?1 WHERE key = ?2", &note, &key)
    })
    .map_err(error_text)
}

#[ic_cdk::query]
fn kv_get_note(key: String) -> Result<Option<String>, String> {
    validate_text_len("key", &key, MAX_KV_KEY_BYTES)?;
    Db::query(|connection| {
        connection.query_optional_string_text("SELECT note FROM kv WHERE key = ?1", &key)
    })
    .map_err(error_text)
}

fn values_lookup_sql(count: usize, column: &str) -> Result<String, String> {
    let mut sql = String::from("WITH lookup(ord, key) AS (VALUES ");
    for index in 0..count {
        if index > 0 {
            sql.push(',');
        }
        let ordinal = index + 1;
        sql.push_str(&format!("({index}, ?{ordinal})"));
    }
    sql.push_str(") SELECT kv.");
    sql.push_str(column);
    sql.push_str(" FROM lookup LEFT JOIN kv ON kv.key = lookup.key ORDER BY lookup.ord");
    Ok(sql)
}

fn validate_text_len(field: &str, value: &str, max_bytes: usize) -> Result<(), String> {
    let len = value.len();
    if len > max_bytes {
        Err(format!(
            "{field} accepts at most {max_bytes} bytes; got {len} bytes"
        ))
    } else {
        Ok(())
    }
}

fn validate_nonzero_u64_max(
    field: &str,
    value: u64,
    max: u64,
    endpoint: &str,
) -> Result<(), String> {
    if value == 0 {
        Err(format!("{endpoint} requires {field} greater than zero"))
    } else if value > max {
        Err(format!(
            "{endpoint} accepts {field} at most {max} bytes; got {value} bytes"
        ))
    } else {
        Ok(())
    }
}

#[ic_cdk::query]
fn kv_count() -> Result<u64, String> {
    require_controller()?;
    let count = Db::query(|connection| {
        connection.query_scalar::<i64>("SELECT COUNT(*) FROM kv", crate::params![])
    })
    .map_err(error_text)?;
    u64::try_from(count).map_err(|_| "negative row count".to_string())
}

#[ic_cdk::query]
fn db_meta() -> Result<DbMeta, String> {
    require_controller()?;
    let block = Superblock::load().map_err(|error| error.to_string())?;
    let stats =
        crate::sqlite_vfs::stable_blob::storage_stats().map_err(|error| error.to_string())?;
    let stable_pages = memory::size_pages();
    Ok(DbMeta {
        db_size: block.db_size,
        stable_pages,
        stable_bytes: stable_pages
            .checked_mul(crate::config::STABLE_PAGE_SIZE)
            .ok_or_else(|| "stable byte size overflow".to_string())?,
        schema_version: block.schema_version,
        last_tx_id: block.last_tx_id,
        flags: block.flags,
        checksum: block.checksum,
        checksum_stale: block.is_checksum_stale(),
        checksum_refreshing: block.is_checksum_refreshing(),
        checksum_refresh_offset: block.checksum_refresh_offset,
        layout_version: stats.layout_version,
        page_count: stats.page_count,
        page_table_bytes: stats.page_table_bytes,
        zero_extent_count: stats.zero_extent_count,
        active_bytes: stats.active_bytes,
        allocated_bytes: stats.allocated_bytes,
        orphan_bytes_estimate: stats.orphan_bytes_estimate,
        orphan_ratio_basis_points: stats.orphan_ratio_basis_points,
    })
}

#[ic_cdk::query]
fn db_integrity_check() -> Result<String, String> {
    require_controller()?;
    Db::integrity_check().map_err(error_text)
}

#[ic_cdk::query]
fn db_checksum() -> Result<u64, String> {
    require_controller()?;
    Db::db_checksum().map_err(error_text)
}

#[ic_cdk::update]
fn db_refresh_checksum() -> Result<u64, String> {
    require_controller()?;
    Err("use db_refresh_checksum_chunk".to_string())
}

#[ic_cdk::update]
fn db_refresh_checksum_chunk(max_bytes: u64) -> Result<ChecksumRefresh, String> {
    require_controller()?;
    validate_nonzero_u64_max(
        "max_bytes",
        max_bytes,
        MAX_CHECKSUM_REFRESH_BYTES,
        "db_refresh_checksum_chunk",
    )?;
    let report = Db::refresh_checksum_chunk(max_bytes).map_err(error_text)?;
    Ok(ChecksumRefresh {
        complete: report.complete,
        checksum: report.checksum,
        scanned_bytes: report.scanned_bytes,
        db_size: report.db_size,
    })
}

#[cfg(feature = "canister-api-test-failpoints")]
#[ic_cdk::update]
fn db_test_trap_after_stable_write(ordinal: u64) -> Result<(), String> {
    require_controller()?;
    crate::stable::memory::set_failpoint(crate::stable::memory::MemoryFailpoint::TrapAfterWrite {
        ordinal,
    });
    Ok(())
}

#[cfg(feature = "canister-api-test-failpoints")]
#[ic_cdk::update]
fn db_test_sqlite_feature_probe() -> Result<(), String> {
    require_controller()?;
    sqlite_feature_probe::run()
}

#[cfg(feature = "canister-api-test-failpoints")]
#[ic_cdk::update]
fn db_test_clear_failpoints() -> Result<(), String> {
    require_controller()?;
    crate::stable::memory::clear_failpoint();
    crate::db::statement::clear_step_failpoint();
    crate::sqlite_vfs::stable_blob::rollback_update();
    Ok(())
}

fn must(result: Result<(), crate::DbError>) {
    if let Err(error) = result {
        ic_cdk::trap(error.to_string());
    }
}

fn error_text(error: crate::DbError) -> String {
    error.to_string()
}

fn require_controller() -> Result<(), String> {
    let caller = ic_cdk::api::msg_caller();
    if ic_cdk::api::is_controller(&caller) {
        Ok(())
    } else {
        Err("caller is not a controller".to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::{
        validate_text_len, MAX_KV_GET_MANY_KEYS, MAX_KV_KEY_BYTES, MAX_KV_NOTE_BYTES,
        MAX_KV_VALUE_BYTES,
    };

    #[test]
    fn kv_input_limits_accept_boundary_values() {
        let key = "k".repeat(MAX_KV_KEY_BYTES);
        let value = "v".repeat(MAX_KV_VALUE_BYTES);
        let note = "n".repeat(MAX_KV_NOTE_BYTES);

        assert_eq!(validate_text_len("key", &key, MAX_KV_KEY_BYTES), Ok(()));
        assert_eq!(
            validate_text_len("value", &value, MAX_KV_VALUE_BYTES),
            Ok(())
        );
        assert_eq!(validate_text_len("note", &note, MAX_KV_NOTE_BYTES), Ok(()));
        assert_eq!(MAX_KV_GET_MANY_KEYS, 1_000);
    }

    #[test]
    fn kv_input_limits_reject_one_byte_over_boundary() {
        let key = "k".repeat(MAX_KV_KEY_BYTES + 1);
        let value = "v".repeat(MAX_KV_VALUE_BYTES + 1);
        let note = "n".repeat(MAX_KV_NOTE_BYTES + 1);

        assert!(validate_text_len("key", &key, MAX_KV_KEY_BYTES)
            .unwrap_err()
            .contains("at most 256 bytes"));
        assert!(validate_text_len("value", &value, MAX_KV_VALUE_BYTES)
            .unwrap_err()
            .contains("at most 65536 bytes"));
        assert!(validate_text_len("note", &note, MAX_KV_NOTE_BYTES)
            .unwrap_err()
            .contains("at most 4096 bytes"));
    }
}