haematite-cli 0.21.8

//! V2: one page of the embedding job, committed into the companion.

use std::error::Error;
use std::num::NonZeroUsize;

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;

use super::{DEFAULT_PAGE_RECORDS, Made, Page, Reach, field_key};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentWrite, Fork, Metric, Mutation, Version,
};
use crate::serve::storage::embed_pending::{Pending, Waiting};
use crate::serve::storage::embedded::{
    self, Declaration, ModelName, Position, Provenance, Toward, VectorKind,
};
use crate::serve::storage::{
    Captured, Failure, Storage, branch_name, document_key, reads, revision_key, version,
};

type TestResult = Result<(), Box<dyn Error>>;

const FIELD: &str = "body_vec";

fn create() -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count: 2,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::reclaim(0),
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    mutate(
        &mut storage,
        Change::Collection(CollectionWrite {
            collection: "notes".to_owned(),
            schema: None,
        }),
    )?;
    Ok((directory, storage))
}

fn head(storage: &Storage) -> Result<Version, Failure> {
    storage
        .refs
        .get(&branch_name("main"))
        .map(version)
        .ok_or_else(|| Failure::new("branch_missing", "main"))
}

fn mutate(storage: &mut Storage, change: Change) -> Result<Version, Failure> {
    let at = head(storage)?;
    let receipt = storage.mutate(&Mutation {
        request_id: format!("embed-{}-{}", at.generation, at.sequence),
        branch: "main".to_owned(),
        expected: at,
        change,
    })?;
    Ok(receipt.version)
}

fn put(storage: &mut Storage, id: &str, body: &str) -> Result<Version, Box<dyn Error>> {
    let document = RawValue::from_string(format!(r#"{{"body":"{body}"}}"#))?;
    let change = Change::Put(DocumentWrite {
        collection: "notes".to_owned(),
        document_id: id.to_owned(),
        document,
    });
    Ok(mutate(storage, change)?)
}

fn model() -> ModelName {
    ModelName {
        name: "stand-in".to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: "abc".to_owned(),
    }
}

/// `notes.body_vec` declared embedded from `body`, two members, cosine.
fn declare(storage: &mut Storage) -> Result<(), Box<dyn Error>> {
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: model(),
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    let bytes = embedded::encode_declaration(&declaration)?;
    let key = field_key("notes", FIELD);
    storage.commit_namespace(&branch_name("main"), &[(key, Some(bytes))])?;
    Ok(())
}

fn revision(storage: &Storage, id: &str) -> Result<Version, Box<dyn Error>> {
    let record = storage.refs.get(&branch_name("main")).ok_or("no main")?;
    let at = Captured {
        branch: "main".to_owned(),
        version: version(record),
        shards: record.shards.clone(),
    };
    Ok(storage.revision(&at, "notes", id)?.ok_or("no revision")?)
}

fn made(id: &str, made_from: Version, vector: &[f32]) -> Made {
    Made {
        id: id.to_owned(),
        made_from,
        vector: vector.to_vec(),
    }
}

/// A page that ends its walk toward `toward`.
fn page(made: Vec<Made>, removed: &[&str], toward: Version) -> Page {
    Page {
        collection: "notes".to_owned(),
        field: FIELD.to_owned(),
        model: model(),
        made,
        removed: removed.iter().map(|id| (*id).to_owned()).collect(),
        cleared: Vec::new(),
        windowed: Vec::new(),
        toward,
        reach: Reach::End,
    }
}

/// A page in the middle of its walk toward `toward`, up to `after`.
fn upto(made: Vec<Made>, toward: Version, after: &str) -> Page {
    Page {
        reach: Reach::Upto(after.to_owned()),
        ..page(made, &[], toward)
    }
}

fn records() -> Result<NonZeroUsize, Box<dyn Error>> {
    Ok(NonZeroUsize::new(DEFAULT_PAGE_RECORDS).ok_or("no default")?)
}

fn companion(storage: &Storage) -> Option<Captured> {
    let name = embedded::companion_name("main");
    storage.refs.get(&name).map(|record| Captured {
        branch: name.clone(),
        version: version(record),
        shards: record.shards.clone(),
    })
}

fn stored(storage: &Storage, id: &str) -> Result<Option<Vec<u8>>, Box<dyn Error>> {
    let at = companion(storage).ok_or("no companion")?;
    let key = embedded::vector_key("notes", FIELD, id);
    Ok(reads::get_bytes(&storage.reader, &at, &key)?)
}

fn position(storage: &Storage) -> Result<Option<Position>, Failure> {
    let name = embedded::companion_name("main");
    storage.position(&name, &embedded::position_key("notes", FIELD))
}

fn code<T>(outcome: Result<T, Failure>) -> Result<String, Box<dyn Error>> {
    Ok(outcome
        .err()
        .ok_or("a page was committed that should be refused")?
        .code)
}

/// Fails while the vectors and the position are not written together: one
/// page leaves both in the companion, each vector headed by the revision it
/// was made from and the model's tag; a page is one commit of the
/// companion; and the branch never moves.
#[test]
fn a_page_writes_its_vectors_and_position_in_one_commit() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "a")?;
    put(&mut storage, "n-2", "b")?;
    declare(&mut storage)?;
    let through = head(&storage)?;
    let first = revision(&storage, "n-1")?;
    let second = revision(&storage, "n-2")?;
    let made = vec![
        made("n-1", first, &[1.0, 0.0]),
        made("n-2", second, &[0.0, 1.0]),
    ];
    let committed = storage.embed_commit("main", &page(made, &[], through), records()?)?;
    assert_eq!((committed.written, committed.skipped.len()), (2, 0));
    let bytes = stored(&storage, "n-1")?.ok_or("no vector")?;
    let (provenance, vector) = embedded::decode_embedded(&bytes, 2)?;
    let made_by = Provenance {
        made_from: first,
        tag: embedded::model_tag(&model())?,
    };
    assert_eq!((provenance, vector), (made_by, vec![1.0, 0.0]));
    let expected = Position {
        branch: "main".to_owned(),
        through: Some(through),
        toward: None,
    };
    assert_eq!(committed.position, expected);
    assert_eq!(position(&storage)?, Some(expected));
    assert_eq!(head(&storage)?, through, "the job moved the branch");
    let next = put(&mut storage, "n-3", "c")?;
    let before = companion(&storage).ok_or("no companion")?.version;
    storage.embed_commit("main", &page(Vec::new(), &[], next), records()?)?;
    let after = companion(&storage).ok_or("no companion")?.version;
    assert_eq!(after.generation, before.generation);
    assert_eq!(after.sequence, before.sequence + 1, "a page is one commit");
    assert_eq!(head(&storage)?, next, "the job moved the branch");
    Ok(())
}

/// Fails while a record rewritten after the page's version is given the
/// vector of its old text: it is skipped, has no vector, and the position
/// still advances, since the next page lists it again.
#[test]
fn a_record_rewritten_after_the_page_is_skipped() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "a")?;
    declare(&mut storage)?;
    let through = head(&storage)?;
    let old = revision(&storage, "n-1")?;
    put(&mut storage, "n-1", "rewritten")?;
    let made = vec![made("n-1", old, &[1.0, 0.0])];
    let committed = storage.embed_commit("main", &page(made, &[], through), records()?)?;
    assert_eq!(committed.skipped, ["n-1"]);
    assert_eq!(committed.written, 0);
    assert_eq!(stored(&storage, "n-1")?, None);
    assert_eq!(position(&storage)?.and_then(|at| at.through), Some(through));
    Ok(())
}

/// Fails while a removal is not written: a record removed since its vector
/// was made loses the vector in the next page. The record's keys are removed
/// directly: until V6 teaches vectors.rs to read a declaration in `w\0`, the
/// delete verb refuses a collection with a declared field.
#[test]
fn a_removed_record_loses_its_vector() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "a")?;
    declare(&mut storage)?;
    let through = head(&storage)?;
    let made = vec![made("n-1", revision(&storage, "n-1")?, &[1.0, 0.0])];
    storage.embed_commit("main", &page(made, &[], through), records()?)?;
    assert!(stored(&storage, "n-1")?.is_some());
    let removed = [
        (document_key("notes", "n-1"), None),
        (revision_key("notes", "n-1"), None),
    ];
    storage.commit_namespace(&branch_name("main"), &removed)?;
    let gone = head(&storage)?;
    let committed = storage.embed_commit("main", &page(Vec::new(), &["n-1"], gone), records()?)?;
    assert_eq!(committed.removed, 1);
    assert_eq!(stored(&storage, "n-1")?, None);
    Ok(())
}

/// Fails while a page whose text was not read at its version is taken: a
/// revision at or before `toward` that is not the one given refuses the
/// page, and nothing is written.
#[test]
fn text_not_read_at_the_pages_version_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    let before = put(&mut storage, "n-0", "earlier")?;
    put(&mut storage, "n-1", "a")?;
    declare(&mut storage)?;
    let through = head(&storage)?;
    let made = vec![made("n-1", before, &[1.0, 0.0])];
    let refused = storage.embed_commit("main", &page(made, &[], through), records()?);
    assert_eq!(code(refused)?, "embed_page_stale_text");
    assert!(companion(&storage).is_none(), "a refused page wrote");
    Ok(())
}

/// Fails while any of these is written: a page past the page setting, a
/// vector of another length than declared, another model's vectors, or a
/// field not declared embedded. Each is refused by name and writes nothing.
#[test]
fn a_page_past_its_bounds_is_refused_and_writes_nothing() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "a")?;
    put(&mut storage, "n-2", "b")?;
    declare(&mut storage)?;
    let through = head(&storage)?;
    let first = revision(&storage, "n-1")?;
    let second = revision(&storage, "n-2")?;
    let two = vec![
        made("n-1", first, &[1.0, 0.0]),
        made("n-2", second, &[0.0, 1.0]),
    ];
    let one = NonZeroUsize::MIN;
    let refused = storage.embed_commit("main", &page(two, &[], through), one);
    assert_eq!(code(refused)?, "embed_page_too_large");
    let long = vec![made("n-1", first, &[1.0, 0.0, 0.0])];
    let refused = storage.embed_commit("main", &page(long, &[], through), records()?);
    assert_eq!(code(refused)?, "embed_vector_length");
    let mut other = page(vec![made("n-1", first, &[1.0, 0.0])], &[], through);
    other.model.manifest_sha256 = "abd".to_owned();
    let refused = storage.embed_commit("main", &other, records()?);
    assert_eq!(code(refused)?, "embed_model_mismatch");
    let mut undeclared = page(vec![made("n-1", first, &[1.0, 0.0])], &[], through);
    undeclared.field = "title_vec".to_owned();
    let refused = storage.embed_commit("main", &undeclared, records()?);
    assert_eq!(code(refused)?, "embed_field_not_declared");
    assert!(companion(&storage).is_none(), "a refused page wrote");
    Ok(())
}

/// Fails while the position can go back: a page completing a version before
/// the job's position is refused by name.
#[test]
fn the_position_never_goes_back() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    let earlier = head(&storage)?;
    put(&mut storage, "n-1", "a")?;
    let later = head(&storage)?;
    storage.embed_commit("main", &page(Vec::new(), &[], later), records()?)?;
    let refused = storage.embed_commit("main", &page(Vec::new(), &[], earlier), records()?);
    assert_eq!(code(refused)?, "embed_position_backwards");
    assert_eq!(position(&storage)?.and_then(|at| at.through), Some(later));
    Ok(())
}

/// Five records, `n-1` to `n-5`, under a declared field: their revisions,
/// and the version they were all read at.
fn five(storage: &mut Storage) -> Result<(Vec<Made>, Version), Box<dyn Error>> {
    for index in 1..=5 {
        put(storage, &format!("n-{index}"), "text")?;
    }
    declare(storage)?;
    let toward = head(storage)?;
    let mut all = Vec::new();
    for index in 1..=5 {
        let id = format!("n-{index}");
        all.push(made(&id, revision(storage, &id)?, &[1.0, 0.0]));
    }
    Ok((all, toward))
}

/// Fails while a walk cannot be resumed where it stopped: a walk of three
/// pages stopped after its first resumes after that page's last record,
/// refuses that record again, embeds the rest, and ends with `toward` as
/// the new `through`.
#[test]
fn a_walk_of_three_pages_resumes_after_an_interruption() -> TestResult {
    let (_directory, mut storage) = create()?;
    let (all, toward) = five(&mut storage)?;
    let first = all.get(..2).ok_or("two")?.to_vec();
    storage.embed_commit("main", &upto(first, toward, "n-2"), records()?)?;
    // Stopped here, as a crash would: the position says where.
    let resumed = position(&storage)?.ok_or("no position")?;
    let open = Toward {
        version: toward,
        after: "n-2".to_owned(),
    };
    assert_eq!((resumed.through, resumed.toward), (None, Some(open)));
    let again = all.get(1..4).ok_or("three")?.to_vec();
    let refused = storage.embed_commit("main", &upto(again, toward, "n-4"), records()?);
    assert_eq!(code(refused)?, "embed_page_outside_walk");
    let second = all.get(2..4).ok_or("two")?.to_vec();
    storage.embed_commit("main", &upto(second, toward, "n-4"), records()?)?;
    let last = all.get(4..).ok_or("one")?.to_vec();
    storage.embed_commit("main", &page(last, &[], toward), records()?)?;
    for index in 1..=5 {
        assert!(
            stored(&storage, &format!("n-{index}"))?.is_some(),
            "n-{index}"
        );
    }
    let ended = position(&storage)?.ok_or("no position")?;
    assert_eq!((ended.through, ended.toward), (Some(toward), None));
    Ok(())
}

/// Fails while `after` can go back within a walk, or a page can reach past
/// its own `after`: each is refused by name, and the position stays.
#[test]
fn a_walk_never_goes_back_and_a_page_stays_inside_it() -> TestResult {
    let (_directory, mut storage) = create()?;
    let (all, toward) = five(&mut storage)?;
    let first = all.get(..3).ok_or("three")?.to_vec();
    storage.embed_commit("main", &upto(first, toward, "n-3"), records()?)?;
    let refused = storage.embed_commit("main", &upto(Vec::new(), toward, "n-2"), records()?);
    assert_eq!(code(refused)?, "embed_walk_backwards");
    let past = all.get(4..).ok_or("one")?.to_vec();
    let refused = storage.embed_commit("main", &upto(past, toward, "n-4"), records()?);
    assert_eq!(code(refused)?, "embed_page_outside_walk");
    let kept = position(&storage)?
        .and_then(|at| at.toward)
        .map(|open| open.after);
    assert_eq!(kept.as_deref(), Some("n-3"));
    Ok(())
}

/// Fails while a new walk can begin over an open one: the open walk's
/// target stays, and a page toward another version is refused by name.
#[test]
fn a_new_walk_while_one_is_open_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    let (all, toward) = five(&mut storage)?;
    let first = all.get(..1).ok_or("one")?.to_vec();
    storage.embed_commit("main", &upto(first, toward, "n-1"), records()?)?;
    let later = put(&mut storage, "n-6", "text")?;
    let refused = storage.embed_commit("main", &page(Vec::new(), &[], later), records()?);
    assert_eq!(code(refused)?, "embed_walk_open");
    let open = position(&storage)?
        .and_then(|at| at.toward)
        .map(|open| open.version);
    assert_eq!(open, Some(toward));
    Ok(())
}

/// Writes record `id` on `branch` at its head.
fn put_on(storage: &mut Storage, branch: &str, id: &str) -> Result<Version, Box<dyn Error>> {
    let at = storage
        .refs
        .get(&branch_name(branch))
        .map(version)
        .ok_or("no branch")?;
    let receipt = storage.mutate(&Mutation {
        request_id: format!("{branch}-{id}"),
        branch: branch.to_owned(),
        expected: at,
        change: Change::Put(DocumentWrite {
            collection: "notes".to_owned(),
            document_id: id.to_owned(),
            document: RawValue::from_string(r#"{"body":"text"}"#.to_owned())?,
        }),
    })?;
    Ok(receipt.version)
}

/// The next page of `feature`'s walk, two records a page, committed as the
/// job would; the position it leaves.
fn feature_page(storage: &mut Storage) -> Result<Position, Box<dyn Error>> {
    let two = NonZeroUsize::new(2).ok_or("two is not zero")?;
    let input = NonZeroUsize::new(64).ok_or("no input")?;
    let pending: Pending = storage
        .embed_pending("feature", "notes", FIELD, two, input)?
        .ok_or("nothing pending")?;
    let mut made_now = Vec::new();
    for waiting in &pending.waiting {
        if let Waiting::Text { id, made_from, .. } = waiting {
            made_now.push(made(id, *made_from, &[1.0, 0.0]));
        }
    }
    let page = Page {
        reach: pending.reach.clone(),
        ..page(made_now, &[], pending.toward)
    };
    storage.embed_commit("feature", &page, two)?;
    let name = embedded::companion_name("feature");
    let position = storage.position(&name, &embedded::position_key("notes", FIELD))?;
    Ok(position.ok_or("no position")?)
}

/// Fails while a walk from an inherited position relabels the parent's
/// version as the walked branch's: over three pages the position keeps
/// main's name and version, and only the page that ends the walk writes
/// feature's own; `embed_pending` answers before every page.
#[test]
fn an_inherited_position_keeps_its_branch_until_the_walk_ends() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "text")?;
    declare(&mut storage)?;
    let forked_at = head(&storage)?;
    mutate(
        &mut storage,
        Change::Fork(Fork {
            name: "feature".to_owned(),
            from: forked_at,
        }),
    )?;
    let inherited = Position {
        branch: "main".to_owned(),
        through: Some(forked_at),
        toward: None,
    };
    storage.position_for_tests("feature", "notes", FIELD, &inherited)?;
    let mut written = Vec::new();
    for index in 1..=5 {
        written.push(put_on(&mut storage, "feature", &format!("f-{index}"))?);
    }
    let target = *written.last().ok_or("nothing written")?;
    for page in 1..=2 {
        let middle = feature_page(&mut storage)?;
        assert_eq!(middle.branch, "main", "page {page}");
        assert_eq!(middle.through, Some(forked_at), "page {page}");
        assert_eq!(
            middle.toward.map(|open| open.version),
            Some(target),
            "page {page}"
        );
    }
    let ended = feature_page(&mut storage)?;
    let expected = Position {
        branch: "feature".to_owned(),
        through: Some(target),
        toward: None,
    };
    assert_eq!(ended, expected);
    Ok(())
}

#[path = "embed_windows_tests.rs"]
mod windows;

#[path = "embed_rewalk_tests.rs"]
mod rewalk;

const fn fixture_long_wait_interval() -> std::time::Duration {
    const LONG_WAIT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
    LONG_WAIT_INTERVAL
}