kcode-k1-persons-store 0.1.0

Current SQLite storage and recovery for the K1 persons projection
Documentation
use rusqlite::ffi::ErrorCode::{
    ConstraintViolation, DatabaseCorrupt, NotADatabase, SchemaChanged, TypeMismatch, Unknown,
};
use rusqlite::{Connection, Error, TransactionBehavior, params};
use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
    time::Duration,
};

pub use kcode_k1_person_types::PersonId;
use kcode_k1_txn_ordering::K1TxnOrdering;
pub use kcode_k1_txn_ordering::TxId;

const APPLICATION_ID: i64 = 0x4b31_5050;
const SCHEMA_VERSION: i64 = 1;
const META_SCHEMA: &str = "CREATE TABLE meta(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), checkpoint BLOB CHECK(checkpoint IS NULL OR length(checkpoint) = 12)) STRICT";
const PERSONS_SCHEMA: &str = "CREATE TABLE persons(id BLOB PRIMARY KEY CHECK(length(id) = 12), root BLOB NOT NULL REFERENCES persons(id), name TEXT, CHECK((id = root AND name IS NOT NULL) OR (id <> root AND name IS NULL))) STRICT, WITHOUT ROWID";
const ROOT_INDEX_SCHEMA: &str = "CREATE INDEX persons_root ON persons(root)";

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoredPerson {
    id: PersonId,
    root: PersonId,
    name: Option<String>,
}

impl StoredPerson {
    pub fn id(&self) -> PersonId {
        self.id
    }

    pub fn root(&self) -> PersonId {
        self.root
    }

    pub fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoredSnapshot {
    checkpoint: Option<TxId>,
    persons: Vec<StoredPerson>,
}

impl StoredSnapshot {
    pub fn checkpoint(&self) -> Option<TxId> {
        self.checkpoint
    }

    pub fn persons(&self) -> &[StoredPerson] {
        &self.persons
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoreChange(Change);

#[derive(Clone, Debug, Eq, PartialEq)]
enum Change {
    Create(PersonId, String),
    Update(PersonId, String),
    Resolve(PersonId, PersonId, usize),
}

impl StoreChange {
    pub fn create(id: PersonId, name: String) -> Self {
        Self(Change::Create(id, name))
    }

    pub fn update(id: PersonId, name: String) -> Self {
        Self(Change::Update(id, name))
    }

    pub fn resolve(from: PersonId, to: PersonId, expected_class_size: usize) -> Self {
        Self(Change::Resolve(from, to, expected_class_size))
    }
}

pub struct Store {
    connection: Connection,
}

impl Store {
    pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, StoredSnapshot), String> {
        fs::create_dir_all(root).map_err(|error| format!("create store root: {error}"))?;
        let path = root.join("persons.sqlite3");
        let exists = match fs::symlink_metadata(&path) {
            Ok(metadata) if metadata.file_type().is_file() => true,
            Ok(_) => return Err("persons.sqlite3 is a symlink or non-file".to_owned()),
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => false,
            Err(error) => return Err(format!("inspect persons.sqlite3: {error}")),
        };
        if !exists {
            return empty_store(create_database(&path)?);
        }
        let connection =
            Connection::open(&path).map_err(|error| format!("open database: {error}"))?;
        if let Err(error) = configure(&connection) {
            let failure = load_error("configure database", error);
            drop(connection);
            return match failure {
                LoadFailure::Recoverable => replace_database(&path),
                LoadFailure::Fatal(message) => Err(message),
            };
        }
        match load_database(&connection, ordering) {
            Ok(snapshot) => Ok((Self { connection }, snapshot)),
            Err(LoadFailure::Fatal(message)) => Err(message),
            Err(LoadFailure::Recoverable) => {
                drop(connection);
                replace_database(&path)
            }
        }
    }

    pub fn commit(&mut self, callback: TxId, change: Option<&StoreChange>) -> Result<(), String> {
        let transaction = self
            .connection
            .transaction_with_behavior(TransactionBehavior::Immediate)
            .map_err(|error| format!("begin commit: {error}"))?;
        if let Some(StoreChange(change)) = change {
            apply_change(&transaction, change)?;
        }
        let changed = transaction
            .execute(
                "UPDATE meta SET checkpoint = ?1 WHERE singleton = 1",
                params![&callback.as_bytes()[..]],
            )
            .map_err(|error| format!("write checkpoint: {error}"))?;
        if changed != 1 {
            return Err(format!("checkpoint update affected {changed} rows"));
        }
        transaction
            .commit()
            .map_err(|error| format!("commit database: {error}"))
    }

    pub fn clear(&mut self) -> Result<(), String> {
        let transaction = self
            .connection
            .transaction_with_behavior(TransactionBehavior::Immediate)
            .map_err(|error| format!("begin clear: {error}"))?;
        transaction
            .execute("DELETE FROM persons", [])
            .map_err(|error| format!("delete persons: {error}"))?;
        let changed = transaction
            .execute("UPDATE meta SET checkpoint = NULL WHERE singleton = 1", [])
            .map_err(|error| format!("clear checkpoint: {error}"))?;
        if changed != 1 {
            return Err(format!("checkpoint clear affected {changed} rows"));
        }
        transaction
            .commit()
            .map_err(|error| format!("commit clear: {error}"))
    }
}

fn apply_change(transaction: &rusqlite::Transaction<'_>, change: &Change) -> Result<(), String> {
    let changed = match change {
        Change::Create(id, name) => {
            let id = id.as_tx_id().into_bytes();
            transaction.execute(
                "INSERT INTO persons(id, root, name) VALUES(?1, ?1, ?2)",
                params![&id[..], name],
            )
        }
        Change::Update(id, name) => {
            let id = id.as_tx_id().into_bytes();
            transaction.execute(
                "UPDATE persons SET name = ?1 WHERE id = ?2 AND root = id AND name IS NOT NULL",
                params![name, &id[..]],
            )
        }
        Change::Resolve(from, to, expected) => {
            let from = from.as_tx_id().into_bytes();
            let to = to.as_tx_id().into_bytes();
            let changed = transaction
                .execute(
                    "UPDATE persons SET root = ?1, name = CASE WHEN id = ?2 THEN NULL ELSE name END WHERE root = ?2",
                    params![&to[..], &from[..]],
                )
                .map_err(|error| format!("apply resolve: {error}"))?;
            if changed == 0 || changed != *expected {
                return Err(format!(
                    "resolve affected {changed} rows, expected {expected}"
                ));
            }
            return Ok(());
        }
    }
    .map_err(|error| format!("apply store change: {error}"))?;
    if changed != 1 {
        return Err(format!("store change affected {changed} rows"));
    }
    Ok(())
}

fn configure(connection: &Connection) -> rusqlite::Result<()> {
    connection.busy_timeout(Duration::from_secs(5))?;
    connection.pragma_update(None, "journal_mode", "WAL")?;
    connection.pragma_update(None, "synchronous", "FULL")?;
    connection.pragma_update(None, "foreign_keys", "ON")
}

fn create_database(path: &Path) -> Result<Connection, String> {
    let connection = Connection::open(path).map_err(|error| format!("create database: {error}"))?;
    configure(&connection).map_err(|error| format!("configure new database: {error}"))?;
    let setup = format!(
        "BEGIN IMMEDIATE;PRAGMA application_id={APPLICATION_ID};PRAGMA user_version={SCHEMA_VERSION};{META_SCHEMA};{PERSONS_SCHEMA};{ROOT_INDEX_SCHEMA};INSERT INTO meta(singleton, checkpoint) VALUES(1, NULL);COMMIT;"
    );
    connection
        .execute_batch(&setup)
        .map_err(|error| format!("create schema: {error}"))?;
    Ok(connection)
}

fn replace_database(path: &Path) -> Result<(Store, StoredSnapshot), String> {
    for candidate in [
        sidecar(path, "-wal"),
        sidecar(path, "-shm"),
        path.to_path_buf(),
    ] {
        match fs::remove_file(&candidate) {
            Ok(()) => {}
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
            Err(error) => return Err(format!("remove {}: {error}", candidate.display())),
        }
    }
    empty_store(create_database(path)?)
}

fn sidecar(path: &Path, suffix: &str) -> PathBuf {
    let mut value = path.as_os_str().to_os_string();
    value.push(suffix);
    PathBuf::from(value)
}

fn empty_store(connection: Connection) -> Result<(Store, StoredSnapshot), String> {
    Ok((
        Store { connection },
        StoredSnapshot {
            checkpoint: None,
            persons: Vec::new(),
        },
    ))
}

enum LoadFailure {
    Recoverable,
    Fatal(String),
}

fn load_error(context: &str, error: Error) -> LoadFailure {
    let recoverable = match &error {
        Error::SqliteFailure(failure, _) => matches!(
            failure.code,
            DatabaseCorrupt
                | NotADatabase
                | SchemaChanged
                | ConstraintViolation
                | TypeMismatch
                | Unknown
        ),
        _ => true,
    };
    if recoverable {
        LoadFailure::Recoverable
    } else {
        LoadFailure::Fatal(format!("{context}: {error}"))
    }
}

fn load_database(
    connection: &Connection,
    ordering: &K1TxnOrdering,
) -> Result<StoredSnapshot, LoadFailure> {
    let application_id: i64 = connection
        .query_row("PRAGMA application_id", [], |row| row.get(0))
        .map_err(|error| load_error("read application id", error))?;
    let version: i64 = connection
        .query_row("PRAGMA user_version", [], |row| row.get(0))
        .map_err(|error| load_error("read schema version", error))?;
    let integrity: String = connection
        .query_row("PRAGMA integrity_check", [], |row| row.get(0))
        .map_err(|error| load_error("check database integrity", error))?;
    if application_id != APPLICATION_ID || version != SCHEMA_VERSION || integrity != "ok" {
        return Err(LoadFailure::Recoverable);
    }
    let schema = schema_rows(connection).map_err(|error| load_error("read schema", error))?;
    let expected = vec![
        (
            "index".to_owned(),
            "persons_root".to_owned(),
            "persons".to_owned(),
            Some(ROOT_INDEX_SCHEMA.to_owned()),
        ),
        (
            "table".to_owned(),
            "meta".to_owned(),
            "meta".to_owned(),
            Some(META_SCHEMA.to_owned()),
        ),
        (
            "table".to_owned(),
            "persons".to_owned(),
            "persons".to_owned(),
            Some(PERSONS_SCHEMA.to_owned()),
        ),
    ];
    if schema != expected || foreign_key_error(connection)? {
        return Err(LoadFailure::Recoverable);
    }
    let meta = meta_rows(connection).map_err(|error| load_error("read metadata", error))?;
    if meta.len() != 1 || meta[0].0 != 1 {
        return Err(LoadFailure::Recoverable);
    }
    let checkpoint = match &meta[0].1 {
        Some(bytes) => Some(TxId::from_bytes(fixed_bytes(bytes)?)),
        None => None,
    };
    if checkpoint.is_some_and(|id| !ordering.contains(id)) {
        return Err(LoadFailure::Recoverable);
    }
    let rows = person_rows(connection).map_err(|error| load_error("read persons", error))?;
    if !rows.is_empty() && checkpoint.is_none() {
        return Err(LoadFailure::Recoverable);
    }
    let mut persons = Vec::with_capacity(rows.len());
    for (id, root, name) in rows {
        let id = PersonId::from_tx_id(TxId::from_bytes(fixed_bytes(&id)?));
        let root = PersonId::from_tx_id(TxId::from_bytes(fixed_bytes(&root)?));
        if !ordering.contains(id.as_tx_id())
            || (id == root && name.is_none())
            || (id != root && name.is_some())
        {
            return Err(LoadFailure::Recoverable);
        }
        persons.push(StoredPerson { id, root, name });
    }
    let by_id: HashMap<PersonId, &StoredPerson> =
        persons.iter().map(|person| (person.id, person)).collect();
    for person in &persons {
        let Some(root) = by_id.get(&person.root) else {
            return Err(LoadFailure::Recoverable);
        };
        if root.id != root.root || root.name.is_none() {
            return Err(LoadFailure::Recoverable);
        }
    }
    Ok(StoredSnapshot {
        checkpoint,
        persons,
    })
}

fn fixed_bytes(bytes: &[u8]) -> Result<[u8; 12], LoadFailure> {
    bytes.try_into().map_err(|_| LoadFailure::Recoverable)
}

type SchemaRow = (String, String, String, Option<String>);
type PersonRow = (Vec<u8>, Vec<u8>, Option<String>);

fn schema_rows(connection: &Connection) -> rusqlite::Result<Vec<SchemaRow>> {
    let mut statement = connection.prepare(
        "SELECT type, name, tbl_name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
    )?;
    statement
        .query_map([], |row| {
            Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
        })?
        .collect()
}

fn meta_rows(connection: &Connection) -> rusqlite::Result<Vec<(i64, Option<Vec<u8>>)>> {
    let mut statement = connection.prepare("SELECT singleton, checkpoint FROM meta")?;
    statement
        .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
        .collect()
}

fn person_rows(connection: &Connection) -> rusqlite::Result<Vec<PersonRow>> {
    let mut statement = connection.prepare("SELECT id, root, name FROM persons ORDER BY id")?;
    statement
        .query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))?
        .collect()
}

fn foreign_key_error(connection: &Connection) -> Result<bool, LoadFailure> {
    let mut statement = connection
        .prepare("PRAGMA foreign_key_check")
        .map_err(|error| load_error("prepare foreign key check", error))?;
    let mut rows = statement
        .query([])
        .map_err(|error| load_error("run foreign key check", error))?;
    rows.next()
        .map(|row| row.is_some())
        .map_err(|error| load_error("read foreign key check", error))
}