haematite-cli 0.21.8

use std::error::Error;
use std::num::NonZeroU32;

use haematite::db::{DatabaseConfig, ReadOnlyDatabase};
use haematite::store::NodeCacheBudget;
use haematite_query::page::RecordSource;
use serde_json::value::RawValue;

use super::{Reach, Read, Search, hits};
use crate::serve::access::{self, Access, Standalone, TestSource};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentWrite, Metric, Mutation, Selection, VectorWrite,
};
use crate::serve::query::parser;
use crate::serve::query::source::{AuthorizedSource, Scope};
use crate::serve::storage::{Captured, Failure, Storage};

type TestResult = Result<(), Box<dyn Error>>;
/// Each hit's id and score in the order given, or the refusal of the search.
type Searched = Result<Vec<(String, Option<f64>)>, Failure>;
/// How many vectors a search compared, or the refusal of the search.
type Compared = Result<u64, Failure>;

/// A page that may compare every vector here, from the first hit.
const WIDE: Reach = Reach {
    bound: 100,
    skip: 0,
};

/// Each record's vector in `embedding`. By dot with q's own [1, 0] they score
/// a 3, b 2, c 2, d 1, q 1: written out of that order, so the order of an
/// answer is the search's doing.
const VECTORS: [(&str, [f32; 2]); 5] = [
    ("c", [2.0, -1.0]),
    ("q", [1.0, 0.0]),
    ("a", [3.0, 0.0]),
    ("d", [1.0, 4.0]),
    ("b", [2.0, 5.0]),
];

struct Moment {
    reader: ReadOnlyDatabase,
    at: Captured,
    _storage: Storage,
    _directory: tempfile::TempDir,
}

fn main_branch() -> Selection {
    Selection {
        branch: "main".to_owned(),
        version: None,
    }
}

/// Every record of `VECTORS`, each `{"name":"<id>"}` with its vector, captured.
fn moment() -> Result<Moment, 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::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    let mut changes = vec![Change::Collection(CollectionWrite {
        collection: "records".to_owned(),
        schema: None,
    })];
    for (id, vector) in VECTORS {
        changes.push(Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: id.to_owned(),
            document: RawValue::from_string(format!(r#"{{"name":"{id}"}}"#))?,
        }));
        changes.push(Change::VectorPut(VectorWrite {
            collection: "records".to_owned(),
            document_id: id.to_owned(),
            field: "embedding".to_owned(),
            vector: vector.to_vec(),
        }));
    }
    for (number, change) in changes.into_iter().enumerate() {
        let expected = storage.capture(&main_branch())?.version;
        storage.mutate(&Mutation {
            request_id: format!("write-{number}"),
            branch: "main".to_owned(),
            expected,
            change,
        })?;
    }
    Ok(Moment {
        reader: storage.reader()?,
        at: storage.capture(&main_branch())?,
        _storage: storage,
        _directory: directory,
    })
}

/// A search by `record`'s vector in `embedding`, by dot, for `limit` hits.
fn search(record: &str, limit: u32) -> Result<Search, Box<dyn Error>> {
    Ok(parser::Search {
        field: "embedding".to_owned(),
        record: record.to_owned(),
        measure: Metric::Dot,
        limit: NonZeroU32::new(limit).ok_or("zero limit")?,
    }
    .into())
}

/// Each hit's id and score in the order given, read with `access`, or the
/// refusal of the search. Each hit's value is its own record.
fn searched(
    moment: &Moment,
    access: &Access,
    ask: &Search,
    reach: &Reach,
) -> Result<Searched, Box<dyn Error>> {
    let scope = Scope::Collection("records".to_owned());
    let source = AuthorizedSource::new(access, &moment.reader, &moment.at, &scope, None, 100)?;
    let found = hits(
        source,
        &Read {
            access,
            reader: &moment.reader,
            service: None,
        },
        &moment.at,
        "records",
        ask,
        reach,
    );
    let mut listed = match found {
        Ok(listed) => listed,
        Err(failure) => return Ok(Err(failure)),
    };
    let mut rows = Vec::new();
    while let Some(record) = listed.next_record()? {
        let expected = format!(r#"{{"name":"{}"}}"#, record.origin.id);
        assert_eq!(record.value.get(), expected);
        assert_eq!(record.origin.collection, "records");
        rows.push((record.origin.id, record.origin.score));
    }
    Ok(Ok(rows))
}

/// How many vectors a search by q compared within `reach`, or its refusal.
fn compared(moment: &Moment, access: &Access, reach: &Reach) -> Result<Compared, Box<dyn Error>> {
    let scope = Scope::Collection("records".to_owned());
    let source = AuthorizedSource::new(access, &moment.reader, &moment.at, &scope, None, 100)?;
    let ask = search("q", 2)?;
    let found = hits(
        source,
        &Read {
            access,
            reader: &moment.reader,
            service: None,
        },
        &moment.at,
        "records",
        &ask,
        reach,
    );
    Ok(found.map(|found| found.compared))
}

fn scored<const N: usize>(expected: [(&str, f64); N]) -> [(String, Option<f64>); N] {
    expected.map(|(id, score)| (id.to_owned(), Some(score)))
}

#[test]
fn hits_come_best_first_with_ties_by_id_each_with_its_score() -> TestResult {
    let moment = moment()?;
    let access = access::request(&Standalone)?;
    let rows = searched(&moment, &access, &search("q", 4)?, &WIDE)??;
    let expected = scored([("a", 3.0), ("b", 2.0), ("c", 2.0), ("d", 1.0)]);
    assert_eq!(rows, expected);
    Ok(())
}

#[test]
fn a_question_record_the_request_may_not_read_is_refused_as_an_absent_one() -> TestResult {
    let moment = moment()?;
    let access = access::request(&TestSource::denying(&[("records", "q")]))?;
    let denied = searched(&moment, &access, &search("q", 4)?, &WIDE)?
        .err()
        .ok_or("a denied question was searched")?;
    let absent = searched(&moment, &access, &search("absent", 4)?, &WIDE)?
        .err()
        .ok_or("an absent question was searched")?;
    assert_eq!(denied.code, "query_search_record_missing");
    assert_eq!(serde_json::to_vec(&denied)?, serde_json::to_vec(&absent)?);
    Ok(())
}

#[test]
fn a_hit_the_request_may_not_read_is_replaced_by_the_next() -> TestResult {
    let moment = moment()?;
    let access = access::request(&TestSource::denying(&[("records", "b")]))?;
    let rows = searched(&moment, &access, &search("q", 3)?, &WIDE)??;
    let expected = scored([("a", 3.0), ("c", 2.0), ("d", 1.0)]);
    assert_eq!(rows, expected);
    Ok(())
}

#[test]
fn a_search_counts_what_it_compared_and_is_refused_past_its_bound() -> TestResult {
    let moment = moment()?;
    let access = access::request(&Standalone)?;
    let at_bound = Reach { bound: 5, skip: 0 };
    assert_eq!(compared(&moment, &access, &at_bound)??, 5);
    let under = Reach { bound: 4, skip: 0 };
    let refused = compared(&moment, &access, &under)?
        .err()
        .ok_or("a search past its bound was answered")?;
    assert_eq!(refused.code, "query_scan_limit");
    Ok(())
}

#[test]
fn a_vector_the_request_may_not_read_is_not_compared() -> TestResult {
    let moment = moment()?;
    let access = access::request(&TestSource::denying(&[("records", "b")]))?;
    assert_eq!(compared(&moment, &access, &WIDE)??, 4);
    Ok(())
}

#[test]
fn a_later_page_skips_the_hits_earlier_pages_gave() -> TestResult {
    let moment = moment()?;
    let access = access::request(&Standalone)?;
    let later = Reach {
        bound: 100,
        skip: 2,
    };
    let rows = searched(&moment, &access, &search("q", 4)?, &later)??;
    let expected = scored([("c", 2.0), ("d", 1.0)]);
    assert_eq!(rows, expected);
    Ok(())
}

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