haematite-cli 0.21.0

//! Point facts and bounded changes between immutable origin and head roots.

use std::num::NonZeroUsize;

use haematite::db::{ObservedChange, ReadOnlyDatabase};
use haematite::tree::{DiffBudget, DiffInterval};

use crate::serve::contracts::ReadDependency;
use crate::serve::protocol::Change;
use crate::serve::storage::{
    Captured, Collection, DocumentRevision, Failure, collection_key, document_key, identifier,
    reads, revision_key,
};

pub(super) fn require_supported(change: &Change) -> Result<(), Failure> {
    match change {
        Change::Collection(_)
        | Change::Put(_)
        | Change::Delete(_)
        | Change::DeleteMany(_)
        | Change::PutMany(_) => Ok(()),
        Change::AppendPrune(_) => Err(Failure::new(
            "dependency_kind_not_admitted",
            "append-prune has no complete dependency validation",
        )),
        Change::EdgePut(_)
        | Change::EdgeDelete(_)
        | Change::EdgePutMany(_)
        | Change::IndexDeclare(_)
        | Change::VectorDeclare(_)
        | Change::VectorPut(_)
        | Change::VectorPutMany(_)
        | Change::VectorDelete(_)
        | Change::RelationPut(_)
        | Change::RelationDelete(_)
        | Change::Fork(_)
        | Change::Discard(_)
        | Change::Merge(_) => Err(Failure::new(
            "dependency_kind_not_admitted",
            "this mutation kind has no complete dependency validation yet",
        )),
    }
}

pub(super) fn validate(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    declared: &[ReadDependency],
    change: &Change,
) -> Result<(), Failure> {
    for dependency in declared {
        match dependency {
            ReadDependency::Document {
                collection,
                document_id,
            } => document(reader, origin, head, collection, document_id)?,
            ReadDependency::Collection { collection: name } => {
                collection(reader, origin, head, name)?;
            }
            ReadDependency::SourceInterval {
                collection: name,
                lower,
                upper,
            } => {
                collection(reader, origin, head, name)?;
                for tag in ["d", "m"] {
                    let prefix = format!("{tag}\0{name}\0").into_bytes();
                    let lower = bound(&prefix, lower.as_deref());
                    let upper = upper
                        .as_ref()
                        .map_or_else(|| prefix_end(&prefix), |id| Ok(bound(&prefix, Some(id))))?;
                    unchanged(
                        reader,
                        origin,
                        head,
                        &lower,
                        &upper,
                        |_| Ok(true),
                        "source interval",
                    )?;
                }
            }
        }
    }
    match change {
        Change::Collection(write) => collection(reader, origin, head, &write.collection),
        Change::Put(write) => document(reader, origin, head, &write.collection, &write.document_id),
        Change::Delete(write) => {
            document(reader, origin, head, &write.collection, &write.document_id)?;
            cascades(reader, origin, head, &write.collection, &write.document_id)
        }
        Change::DeleteMany(delete) => {
            collection(reader, origin, head, &delete.collection)?;
            for id in &delete.document_ids {
                document_rows(reader, origin, head, &delete.collection, id)?;
                cascades(reader, origin, head, &delete.collection, id)?;
            }
            Ok(())
        }
        Change::PutMany(import) => {
            collection(reader, origin, head, &import.collection)?;
            for row in &import.documents {
                document_rows(reader, origin, head, &import.collection, &row.document_id)?;
            }
            Ok(())
        }
        Change::AppendPrune(_)
        | Change::EdgePut(_)
        | Change::EdgeDelete(_)
        | Change::EdgePutMany(_)
        | Change::IndexDeclare(_)
        | Change::VectorDeclare(_)
        | Change::VectorPut(_)
        | Change::VectorPutMany(_)
        | Change::VectorDelete(_)
        | Change::RelationPut(_)
        | Change::RelationDelete(_)
        | Change::Fork(_)
        | Change::Discard(_)
        | Change::Merge(_) => require_supported(change),
    }
}

fn collection(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    name: &str,
) -> Result<(), Failure> {
    identifier(name)?;
    let key = collection_key(name);
    let before = reads::get_bytes(reader, origin, &key)?;
    let after = reads::get_bytes(reader, head, &key)?;
    if before != after {
        let revision = after
            .as_deref()
            .map(|bytes| serde_json::from_slice::<Collection>(bytes).map(|row| row.changed_at))
            .transpose()
            .map_err(Failure::storage)?;
        let mut failure = conflict(head, "collection definition", revision);
        failure.path = Some(format!("/collections/{}", pointer_part(name)));
        return Err(failure);
    }
    Ok(())
}

fn document(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    name: &str,
    id: &str,
) -> Result<(), Failure> {
    collection(reader, origin, head, name)?;
    document_rows(reader, origin, head, name, id)
}

/// The body and revision rows of one document, for a caller that has already
/// compared the collection's definition.
fn document_rows(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    name: &str,
    id: &str,
) -> Result<(), Failure> {
    identifier(name)?;
    identifier(id)?;
    let body = document_key(name, id);
    let revision = revision_key(name, id);
    let old_revision = reads::get_bytes(reader, origin, &revision)?;
    let new_revision = reads::get_bytes(reader, head, &revision)?;
    if old_revision != new_revision
        || reads::get_bytes(reader, origin, &body)? != reads::get_bytes(reader, head, &body)?
    {
        let current = new_revision
            .as_deref()
            .map(|bytes| {
                serde_json::from_slice::<DocumentRevision>(bytes).map(|row| row.written_at)
            })
            .transpose()
            .map_err(Failure::storage)?;
        let mut failure = conflict(head, "document or absence", current);
        failure.path = Some(format!(
            "/documents/{}/{}",
            pointer_part(name),
            pointer_part(id)
        ));
        return Err(failure);
    }
    Ok(())
}

fn conflict(
    head: &Captured,
    fact: &str,
    current: Option<crate::serve::protocol::Version>,
) -> Failure {
    let current = current.map_or_else(
        || "absent".to_owned(),
        |at| format!("{}:{}", at.generation, at.sequence),
    );
    Failure::new(
        "conflict",
        format!(
            "{fact} changed; current_revision={current}; head={}:{}",
            head.version.generation, head.version.sequence
        ),
    )
}

fn changed_fact(
    reader: &ReadOnlyDatabase,
    head: &Captured,
    fact: &str,
    change: &ObservedChange,
) -> Result<Failure, Failure> {
    let key = std::str::from_utf8(change.key()).map_err(Failure::storage)?;
    let parts: Vec<_> = key.split('\0').collect();
    let ["d" | "m", collection, id] = parts.as_slice() else {
        let current = if matches!(change, ObservedChange::Removed { .. }) {
            "absent"
        } else {
            "unversioned"
        };
        return Ok(Failure::new(
            "conflict",
            format!(
                "{fact} changed; current_revision={current}; head={}:{}",
                head.version.generation, head.version.sequence
            ),
        ));
    };
    identifier(collection)?;
    identifier(id)?;
    let revision = reads::get_bytes(reader, head, &revision_key(collection, id))?
        .as_deref()
        .map(|bytes| serde_json::from_slice::<DocumentRevision>(bytes).map(|row| row.written_at))
        .transpose()
        .map_err(Failure::storage)?;
    let mut failure = conflict(head, fact, revision);
    failure.path = Some(format!(
        "/documents/{}/{}",
        pointer_part(collection),
        pointer_part(id)
    ));
    Ok(failure)
}

fn pointer_part(part: &str) -> String {
    part.replace('~', "~0").replace('/', "~1")
}

fn bound(prefix: &[u8], id: Option<&str>) -> Vec<u8> {
    let mut key = prefix.to_vec();
    if let Some(id) = id {
        key.extend_from_slice(id.as_bytes());
    }
    key
}

fn prefix_end(prefix: &[u8]) -> Result<Vec<u8>, Failure> {
    // Composite prefixes end in a delimiter, so their successor is finite.
    let mut end = prefix.to_vec();
    let last = end
        .last_mut()
        .ok_or_else(|| Failure::new("dependency_bound", "composite prefix is empty"))?;
    *last = last.checked_add(1).ok_or_else(|| {
        Failure::new(
            "dependency_bound",
            "composite prefix has no finite successor",
        )
    })?;
    Ok(end)
}

fn cascades(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    collection: &str,
    id: &str,
) -> Result<(), Failure> {
    for tag in ["e", "i"] {
        let prefix = format!("{tag}\0{collection}\0{id}\0").into_bytes();
        unchanged(
            reader,
            origin,
            head,
            &prefix,
            &prefix_end(&prefix)?,
            |_| Ok(true),
            "incident edge",
        )?;
    }
    let dimensions = format!("w\0{collection}\0").into_bytes();
    unchanged(
        reader,
        origin,
        head,
        &dimensions,
        &prefix_end(&dimensions)?,
        |_| Ok(true),
        "vector field definition",
    )?;
    let vectors = format!("v\0{collection}\0").into_bytes();
    unchanged(
        reader,
        origin,
        head,
        &vectors,
        &prefix_end(&vectors)?,
        |change| {
            let key = std::str::from_utf8(change.key()).map_err(Failure::storage)?;
            let mut parts = key.split('\0');
            match (
                parts.next(),
                parts.next(),
                parts.next(),
                parts.next(),
                parts.next(),
            ) {
                (Some("v"), Some(name), Some(field), Some(document_id), None)
                    if name == collection =>
                {
                    identifier(field)?;
                    identifier(document_id)?;
                    Ok(document_id == id)
                }
                _ => Err(Failure::new(
                    "vector_key",
                    "changed vector key has invalid components",
                )),
            }
        },
        "document vector",
    )
}

fn unchanged(
    reader: &ReadOnlyDatabase,
    origin: &Captured,
    head: &Captured,
    lower: &[u8],
    upper: &[u8],
    mut relevant: impl FnMut(&ObservedChange) -> Result<bool, Failure>,
    fact: &str,
) -> Result<(), Failure> {
    let budget = DiffBudget {
        max_changes: NonZeroUsize::new(64)
            .ok_or_else(|| Failure::new("diff_budget", "change budget must be nonzero"))?,
        max_nodes: NonZeroUsize::new(4096)
            .ok_or_else(|| Failure::new("diff_budget", "node budget must be nonzero"))?,
    };
    if origin.shards.len() != head.shards.len() {
        return Err(Failure::new(
            "incomplete_roots",
            "origin and head shard sets differ",
        ));
    }
    for (from, to) in origin.shards.iter().zip(&head.shards) {
        if from.shard_id != to.shard_id {
            return Err(Failure::new(
                "incomplete_roots",
                "origin and head shard order differs",
            ));
        }
        let mut after = None;
        loop {
            let page = reader
                .changes_page_in_range(
                    from.shard_id,
                    from.head,
                    to.head,
                    DiffInterval {
                        lower: Some(lower),
                        upper: Some(upper),
                    },
                    after.as_deref(),
                    budget,
                )
                .map_err(Failure::observer)?;
            for change in &page.entries {
                if relevant(change)? {
                    return Err(changed_fact(reader, head, fact, change)?);
                }
            }
            let Some(next) = page.next_after else {
                break;
            };
            if after.as_ref().is_some_and(|previous| next <= *previous) {
                return Err(Failure::new(
                    "diff_cursor",
                    "change comparison cursor did not advance",
                ));
            }
            after = Some(next);
        }
    }
    Ok(())
}

#[cfg(test)]
mod declaration_tests {
    #[test]
    fn vector_declaration_requires_its_own_dependency_validation()
    -> Result<(), Box<dyn std::error::Error>> {
        let change =
            serde_json::from_value(serde_json::json!({"kind":"vector.declare", "value": {
                "collection":"records", "field":"embedding", "declaration": {
                    "kind":"dense", "source":"body", "model": {
                        "name":"local-model", "revision":"v1", "manifest_sha256":"a".repeat(64)
                    }, "dimensions":2, "distance":"cosine"
                }
            }}))?;
        let failure = super::require_supported(&change)
            .err()
            .ok_or("declaration dependencies admitted")?;
        assert_eq!(failure.code, "dependency_kind_not_admitted");
        Ok(())
    }
}