haematite-cli 0.21.7

//! F1: a fork inherits its parent's companion.

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

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

use super::Inherited;
use crate::serve::access::{Standalone, request};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentWrite, Fork, Metric, Mutation, Selection, Version,
};
use crate::serve::storage::embed_job::{Made, Page, field_key};
use crate::serve::storage::embed_pending::{Pending, Waiting};
use crate::serve::storage::embed_search::{Searched, search};
use crate::serve::storage::embedded::{self, Declaration, ModelName, Position, VectorKind};
use crate::serve::storage::release::{PendingDiscards, Source};
use crate::serve::storage::{Failure, Storage, branch_name, version};

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

const FIELD: &str = "body_vec";
const LATER: &str = "title_vec";
const RECORDS: usize = 2;

/// No discard is pending.
struct NoDiscards;

impl PendingDiscards for NoDiscards {
    fn pending(&self, _storage: &Storage) -> Result<Vec<(String, Version)>, Failure> {
        Ok(Vec::new())
    }
}

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(),
    )?;
    let change = Change::Collection(CollectionWrite {
        collection: "notes".to_owned(),
        schema: None,
    });
    mutate(&mut storage, "main", change)?;
    Ok((directory, storage))
}

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

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

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

/// `feature` forked from `main` at its head.
fn fork(storage: &mut Storage) -> TestResult {
    let from = head(storage, "main")?;
    let change = Change::Fork(Fork {
        name: "feature".to_owned(),
        from,
    });
    mutate(storage, "main", change)?;
    Ok(())
}

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

/// `field` declared embedded on `main`, made from `source`.
fn declare(storage: &mut Storage, field: &str, source: &str) -> TestResult {
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: source.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 pending(
    storage: &Storage,
    branch: &str,
    field: &str,
) -> Result<Option<Pending>, Box<dyn Error>> {
    let records = NonZeroUsize::new(RECORDS).ok_or("no records")?;
    let input = NonZeroUsize::new(64).ok_or("no input")?;
    Ok(storage.embed_pending(branch, "notes", field, records, input)?)
}

/// Commits `pending` as the job would, every vector `[1, 0]`.
fn commit(storage: &mut Storage, branch: &str, field: &str, pending: &Pending) -> TestResult {
    let mut made = Vec::new();
    let mut removed = Vec::new();
    for waiting in &pending.waiting {
        match waiting {
            Waiting::Text { id, made_from, .. } => made.push(Made {
                id: id.clone(),
                made_from: *made_from,
                vector: vec![1.0, 0.0],
            }),
            Waiting::Removed { id } => removed.push(id.clone()),
        }
    }
    let page = Page {
        collection: "notes".to_owned(),
        field: field.to_owned(),
        model: model(),
        made,
        removed,
        cleared: Vec::new(),
        windowed: Vec::new(),
        toward: pending.toward,
        reach: pending.reach.clone(),
    };
    let records = NonZeroUsize::new(RECORDS).ok_or("no records")?;
    storage.embed_commit(branch, &page, records)?;
    Ok(())
}

/// Runs `branch`'s job on `field` to the end of its walk; every id listed.
fn run(storage: &mut Storage, branch: &str, field: &str) -> Result<Vec<String>, Box<dyn Error>> {
    let mut listed = Vec::new();
    for _page in 0..100 {
        let Some(next) = pending(storage, branch, field)? else {
            listed.sort();
            return Ok(listed);
        };
        for waiting in &next.waiting {
            match waiting {
                Waiting::Text { id, .. } | Waiting::Removed { id } => listed.push(id.clone()),
            }
        }
        commit(storage, branch, field, &next)?;
    }
    Err("the walk did not end in 100 pages".into())
}

/// The position of `field` in `owner`'s companion.
fn position(storage: &Storage, owner: &str, field: &str) -> Result<Option<Position>, Failure> {
    let companion = embedded::companion_name(owner);
    storage.position(&companion, &embedded::position_key("notes", field))
}

fn companion_version(storage: &Storage) -> Option<Version> {
    storage
        .refs
        .get(&embedded::companion_name("feature"))
        .map(version)
}

/// Fails while the companion is not made from the parent's: after the
/// inherit, the fork's walk lists only the record written on the fork, and
/// a search on the fork once it ends compares the parent's records too.
#[test]
fn a_fork_walks_only_its_own_records_and_searches_its_parents() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    for id in ["a", "b", "c"] {
        put(&mut storage, "main", id)?;
    }
    assert_eq!(run(&mut storage, "main", FIELD)?, ["a", "b", "c"]);
    fork(&mut storage)?;
    put(&mut storage, "feature", "d")?;
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Done);
    assert_eq!(run(&mut storage, "feature", FIELD)?, ["d"]);
    let at = storage.capture(&Selection {
        branch: "feature".to_owned(),
        version: None,
    })?;
    let (declaration, companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let searched = Searched {
        collection: "notes",
        field: FIELD,
        declaration: &declaration,
        at: &at,
        companion: companion.as_ref(),
    };
    let access = request(&Standalone)?;
    let k = NonZeroUsize::new(10).ok_or("ten is not zero")?;
    let found = search(&access, &storage.reader, &searched, &[1.0, 0.0], k, 100)?;
    assert_eq!(found.compared, 4);
    assert_eq!(found.stale, 0);
    Ok(())
}

/// Fails while an open walk in the parent's companion is not waited on: the
/// answer names the field, and the fork has no companion.
#[test]
fn a_parents_open_walk_is_waited_on() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    for id in ["a", "b", "c"] {
        put(&mut storage, "main", id)?;
    }
    let first = pending(&storage, "main", FIELD)?.ok_or("nothing pending")?;
    commit(&mut storage, "main", FIELD, &first)?;
    fork(&mut storage)?;
    let waiting = Inherited::Waiting {
        collection: "notes".to_owned(),
        field: FIELD.to_owned(),
    };
    assert_eq!(storage.embed_inherit("feature")?, waiting);
    assert_eq!(companion_version(&storage), None);
    Ok(())
}

/// Fails while step two is skipped: a field declared on the parent after the
/// fork point has no position in the fork's companion, and the parent's
/// floor holds nothing for it, while the field both declare is held.
#[test]
fn a_field_declared_after_the_fork_is_not_inherited() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    put(&mut storage, "main", "a")?;
    run(&mut storage, "main", FIELD)?;
    fork(&mut storage)?;
    declare(&mut storage, LATER, "title")?;
    run(&mut storage, "main", LATER)?;
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Done);
    assert_eq!(position(&storage, "feature", LATER)?, None);
    assert!(position(&storage, "feature", FIELD)?.is_some());
    let floor = storage.relied_on("main", &NoDiscards)?;
    let names: Vec<&str> = floor
        .held_by
        .iter()
        .filter(|held| held.source == Source::EmbedOf)
        .map(|held| held.name.as_str())
        .collect();
    assert_eq!(names, ["feature:notes.body_vec"]);
    Ok(())
}

/// Fails while a second call writes: it answers `Already` and the fork's
/// companion stays at its version.
#[test]
fn a_second_inherit_writes_nothing() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    put(&mut storage, "main", "a")?;
    run(&mut storage, "main", FIELD)?;
    fork(&mut storage)?;
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Done);
    let made = companion_version(&storage);
    assert!(made.is_some());
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Already);
    assert_eq!(companion_version(&storage), made);
    Ok(())
}

/// Fails while a repeated call leaves a stray position: a companion made
/// with a position for a field the fork does not declare, as a call stopped
/// before step two leaves it, loses that position on the next call, and the
/// call after answers `Already`.
#[test]
fn a_call_after_an_unfinished_one_removes_the_stray_position() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    put(&mut storage, "main", "a")?;
    run(&mut storage, "main", FIELD)?;
    fork(&mut storage)?;
    let stray = Position {
        branch: "main".to_owned(),
        through: Some(head(&storage, "main")?),
        toward: None,
    };
    storage.position_for_tests("feature", "notes", LATER, &stray)?;
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Done);
    assert_eq!(position(&storage, "feature", LATER)?, None);
    assert_eq!(storage.embed_inherit("feature")?, Inherited::Already);
    Ok(())
}

/// Writes a position key whose field part is not UTF-8 into `owner`'s
/// companion.
fn corrupt_key(storage: &mut Storage, owner: &str) -> TestResult {
    let position = Position {
        branch: owner.to_owned(),
        through: Some(head(storage, owner)?),
        toward: None,
    };
    let bytes = embedded::encode_position(&position)?;
    let key = b"e\0notes\0\xff".to_vec();
    storage.commit_namespace(&embedded::companion_name(owner), &[(key, Some(bytes))])?;
    Ok(())
}

/// Fails while a corrupt position key is guessed at: a parent's companion
/// holding a key whose field part is not UTF-8 makes the inherit refuse by
/// name, and the fork has no companion.
#[test]
fn a_corrupt_position_key_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage, FIELD, "body")?;
    put(&mut storage, "main", "a")?;
    run(&mut storage, "main", FIELD)?;
    corrupt_key(&mut storage, "main")?;
    fork(&mut storage)?;
    let refused = storage.embed_inherit("feature").err();
    let code = refused.map(|failure| failure.code);
    assert_eq!(code.as_deref(), Some("embed_position_key_corrupt"));
    assert_eq!(companion_version(&storage), None);
    Ok(())
}

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
}