haematite-cli 0.21.1

use std::collections::{BTreeMap, BTreeSet};
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::path::PathBuf;

use haematite::Hash;
use haematite::db::{DatabaseConfig, ReadOnlyDatabase};
use haematite::store::NodeCacheBudget;
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::RecordSource;
use serde_json::Value;
use serde_json::value::RawValue;

use super::{AuthorizedSource, Place, Scope};
use crate::serve::access::{self, Access, Standalone, TestSource};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentWrite, Mutation, QueryPage, QueryRequest, Selection, Version,
};
use crate::serve::query::admission::Slots;
use crate::serve::query::run::{self, Context};
use crate::serve::query::{Rows, Seal, Settings};
use crate::serve::storage::{Captured, Storage};

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

const SHARDS: usize = 4;

/// Written out of order, so an ordered answer is the source's doing.
const WRITES: [(&str, &str); 8] = [
    ("gamma", "g2"),
    ("alpha", "a3"),
    ("beta", "b1"),
    ("alpha", "a1"),
    ("gamma", "g1"),
    ("beta", "b2"),
    ("alpha", "a2"),
    ("gamma", "g3"),
];

const ORDERED: [&str; 8] = [
    "alpha/a1", "alpha/a2", "alpha/a3", "beta/b1", "beta/b2", "gamma/g1", "gamma/g2", "gamma/g3",
];

struct Moment {
    _directory: tempfile::TempDir,
    storage: Storage,
    reader: ReadOnlyDatabase,
    at: Captured,
    /// The version each document was written at, from its write's receipt.
    written: BTreeMap<String, Version>,
}

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

fn mutate(
    storage: &mut Storage,
    request_id: String,
    expected: Version,
    change: Change,
) -> Result<Version, Box<dyn std::error::Error>> {
    let mutation = Mutation {
        request_id,
        branch: "main".to_owned(),
        expected,
        change,
    };
    Ok(storage.mutate(&mutation)?.version)
}

fn put(collection: &str, id: &str) -> Result<Change, Box<dyn std::error::Error>> {
    Ok(Change::Put(DocumentWrite {
        collection: collection.to_owned(),
        document_id: id.to_owned(),
        document: RawValue::from_string(format!(r#"{{"name":"{collection}/{id}"}}"#))?,
    }))
}

/// `writes` committed on main, then captured, over `shard_count` shards.
fn moment(
    shard_count: usize,
    writes: &[(&str, &str)],
) -> Result<Moment, Box<dyn std::error::Error>> {
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    let mut version = storage.capture(&selection())?.version;
    let mut declared = BTreeSet::new();
    let mut written = BTreeMap::new();
    for (collection, id) in writes {
        if declared.insert(*collection) {
            let change = Change::Collection(CollectionWrite {
                collection: (*collection).to_owned(),
                schema: None,
            });
            version = mutate(
                &mut storage,
                format!("declare/{collection}"),
                version,
                change,
            )?;
        }
        let name = format!("{collection}/{id}");
        version = mutate(&mut storage, name.clone(), version, put(collection, id)?)?;
        written.insert(name, version);
    }
    let at = storage.capture(&selection())?;
    let reader = storage.reader()?;
    Ok(Moment {
        _directory: directory,
        storage,
        reader,
        at,
        written,
    })
}

fn refused(error: &Error) -> Box<dyn std::error::Error> {
    format!("{:?}: {}", error.refusal.code, error.refusal.message).into()
}

/// A source over `moment`'s capture.
fn open<'a>(
    access: &'a Access,
    moment: &'a Moment,
    scope: &Scope,
    start: Option<&Place>,
    scan_limit: u64,
) -> Result<AuthorizedSource<'a>, Box<dyn std::error::Error>> {
    AuthorizedSource::new(access, &moment.reader, &moment.at, scope, start, scan_limit)
        .map_err(|error| refused(&error))
}

/// Every record the source hands on, as `collection/id`, checking each one's
/// version is the one its document was written at.
fn names(
    source: &mut AuthorizedSource<'_>,
    written: &BTreeMap<String, Version>,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut names = Vec::new();
    while let Some(record) = source.next_record().map_err(|error| refused(&error))? {
        let name = format!("{}/{}", record.origin.collection, record.origin.id);
        let version = written
            .get(&name)
            .ok_or("a record that was never written")?;
        assert_eq!(
            record.origin.version.get(),
            serde_json::to_string(version)?,
            "{name}"
        );
        names.push(name);
    }
    Ok(names)
}

/// Test (a), with the moment check: a write after the capture is not seen.
#[test]
fn records_come_in_collection_and_id_order_across_shards_at_the_captured_moment() -> TestResult {
    let mut moment = moment(SHARDS, &WRITES)?;
    let shards: BTreeSet<usize> = WRITES
        .iter()
        .map(|(collection, id)| format!("d\0{collection}\0{id}"))
        .map(|key| moment.reader.shard_for(key.as_bytes()))
        .collect();
    assert!(shards.len() >= 2, "every document fell in one shard");
    let after = moment.at.version;
    mutate(
        &mut moment.storage,
        "later".to_owned(),
        after,
        put("alpha", "a0")?,
    )?;
    let access = access::request(&Standalone)?;

    let mut all = open(&access, &moment, &Scope::All, None, 99)?;
    assert_eq!(names(&mut all, &moment.written)?, ORDERED);

    let beta = Scope::Collection("beta".to_owned());
    let mut one = open(&access, &moment, &beta, None, 99)?;
    assert_eq!(names(&mut one, &moment.written)?, ["beta/b1", "beta/b2"]);
    Ok(())
}

/// Test (b).
#[test]
fn a_denied_record_is_counted_but_never_returned_or_decoded() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let access = access::request(&TestSource::denying(&[("beta", "b1")]))?;
    let mut source = open(&access, &moment, &Scope::All, None, 99)?;
    let names = names(&mut source, &moment.written)?;
    assert!(
        !names.iter().any(|name| name == "beta/b1"),
        "a denied record was returned"
    );
    assert_eq!(names.len(), ORDERED.len() - 1);
    assert_eq!(
        source.decoded.get(),
        7,
        "a denied record's body was decoded"
    );
    assert_eq!(source.scanned, 8, "the denied record was not counted");
    Ok(())
}

/// Test (c): allowance 3, four present, the second denied. The source stops at
/// the third and names it, and a source after it gives the fourth.
#[test]
fn the_source_stops_at_its_allowance_and_names_the_last_place_it_read() -> TestResult {
    let writes = [
        ("alpha", "a1"),
        ("alpha", "a2"),
        ("alpha", "a3"),
        ("alpha", "a4"),
    ];
    let moment = moment(SHARDS, &writes)?;
    let access = access::request(&TestSource::denying(&[("alpha", "a2")]))?;
    let mut source = open(&access, &moment, &Scope::All, None, 3)?;
    let first = source.next_record().map_err(|error| refused(&error))?;
    assert_eq!(first.ok_or("no first record")?.origin.id, "a1");
    let second = source.next_record().map_err(|error| refused(&error))?;
    assert_eq!(second.ok_or("no second record")?.origin.id, "a3");
    let stopped = source.next_record().map_err(|error| refused(&error))?;
    assert!(stopped.is_none(), "a fourth candidate was given");
    assert_eq!(source.scanned, 3);
    let passed = source.passed().ok_or("the source named no place")?;
    assert_eq!(passed.collection, "alpha");
    assert_eq!(passed.id, "a3");
    let after = Place {
        collection: passed.collection,
        id: passed.id,
        after: true,
    };
    let mut rest = open(&access, &moment, &Scope::All, Some(&after), 3)?;
    assert_eq!(names(&mut rest, &moment.written)?, ["alpha/a4"]);
    assert!(rest.passed().is_none(), "a source that ended named a place");
    Ok(())
}

/// A run of denied records longer than the allowance is crossed page by page:
/// every readable record is given once, none twice and none missed, and a
/// page inside the run gives nothing and names a place in it.
#[test]
fn a_denied_run_longer_than_the_allowance_is_crossed_page_by_page() -> TestResult {
    let writes = [
        ("alpha", "a1"),
        ("alpha", "a2"),
        ("alpha", "a3"),
        ("alpha", "a4"),
        ("alpha", "a5"),
        ("alpha", "a6"),
        ("alpha", "a7"),
        ("alpha", "a8"),
    ];
    let denied = [
        ("alpha", "a2"),
        ("alpha", "a3"),
        ("alpha", "a4"),
        ("alpha", "a5"),
        ("alpha", "a6"),
    ];
    let moment = moment(SHARDS, &writes)?;
    let access = access::request(&TestSource::denying(&denied))?;
    let mut given = Vec::new();
    let mut empty = 0;
    let mut start = None;
    loop {
        let mut source = open(&access, &moment, &Scope::All, start.as_ref(), 2)?;
        let page = names(&mut source, &moment.written)?;
        assert!(source.scanned <= 2, "a page read past its allowance");
        if page.is_empty() {
            empty += 1;
        }
        given.extend(page);
        let Some(passed) = source.passed() else {
            break;
        };
        start = Some(Place {
            collection: passed.collection,
            id: passed.id,
            after: true,
        });
    }
    assert_eq!(given, ["alpha/a1", "alpha/a7", "alpha/a8"]);
    assert!(empty > 0, "no page fell inside the denied run");
    Ok(())
}

/// Test (d): one candidate per shard on the first call, then one refill and one
/// revision per record; nothing is read before it is asked for.
#[test]
fn the_source_reads_only_what_each_call_needs() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let shards = u64::try_from(moment.at.shards.len())?;
    let access = access::request(&Standalone)?;
    let mut source = open(&access, &moment, &Scope::All, None, 99)?;
    assert_eq!(source.reads.get(), 0, "the source read before it was asked");
    source.next_record().map_err(|error| refused(&error))?;
    assert_eq!(source.reads.get(), shards + 1);
    source.next_record().map_err(|error| refused(&error))?;
    assert_eq!(source.reads.get(), shards + 3);
    while source
        .next_record()
        .map_err(|error| refused(&error))?
        .is_some()
    {}
    // Six more records at two reads each, then the call that finds none.
    assert_eq!(source.reads.get(), shards + 3 + 6 * 2 + 1);
    Ok(())
}

/// Test (e).
#[test]
fn a_start_place_is_inclusive_unless_after_and_must_lie_in_the_scope() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let access = access::request(&Standalone)?;
    let start = Place {
        collection: "beta".to_owned(),
        id: "b2".to_owned(),
        after: false,
    };
    let mut source = open(&access, &moment, &Scope::All, Some(&start), 99)?;
    let expected = ["beta/b2", "gamma/g1", "gamma/g2", "gamma/g3"];
    assert_eq!(names(&mut source, &moment.written)?, expected);
    let after = Place {
        collection: "beta".to_owned(),
        id: "b2".to_owned(),
        after: true,
    };
    let mut source = open(&access, &moment, &Scope::All, Some(&after), 99)?;
    let expected = ["gamma/g1", "gamma/g2", "gamma/g3"];
    assert_eq!(names(&mut source, &moment.written)?, expected);

    let alpha = Scope::Collection("alpha".to_owned());
    let outside = AuthorizedSource::new(
        &access,
        &moment.reader,
        &moment.at,
        &alpha,
        Some(&start),
        99,
    );
    let error = outside
        .err()
        .ok_or("a start outside the scope was accepted")?;
    assert_eq!(error.refusal.code, Refusal::QueryProtocol);
    Ok(())
}

/// Test (f): a capture whose roots the store does not hold.
#[test]
fn a_moment_whose_nodes_are_gone_is_refused_by_name() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let mut gone = moment.at.clone();
    for shard in &mut gone.shards {
        shard.head = Hash::of(b"a root this store never held");
    }
    let access = access::request(&Standalone)?;
    let mut source = AuthorizedSource::new(&access, &moment.reader, &gone, &Scope::All, None, 99)
        .map_err(|error| refused(&error))?;
    let error = source
        .next_record()
        .err()
        .ok_or("a record from a moment that is gone")?;
    assert_eq!(error.refusal.code, Refusal::QueryMomentGone);
    Ok(())
}

/// A source of origins reads revision entries in batches and no body: eight
/// origins on four shards take one range read for each shard and decode none.
/// Under a limit of three, a shard reads at most three, and at most one more
/// read of one entry is made for each shard drained at the limit.
#[test]
fn a_page_of_origins_reads_one_batch_per_shard_and_no_body() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let shards = u64::try_from(moment.at.shards.len())?;
    let access = access::request(&Standalone)?;
    let mut source =
        AuthorizedSource::origins(&access, &moment.reader, &moment.at, &Scope::All, None, 99)
            .map_err(|error| refused(&error))?;
    assert_eq!(names(&mut source, &moment.written)?, ORDERED);
    assert_eq!(source.reads.get(), shards);
    assert_eq!(source.decoded.get(), 0, "a body was decoded");
    let mut limited =
        AuthorizedSource::origins(&access, &moment.reader, &moment.at, &Scope::All, None, 3)
            .map_err(|error| refused(&error))?;
    let first = ["alpha/a1", "alpha/a2", "alpha/a3"];
    assert_eq!(names(&mut limited, &moment.written)?, first);
    let reads = limited.reads.get();
    assert!(reads <= 2 * shards, "{reads} reads");
    assert_eq!(limited.decoded.get(), 0, "a body was decoded");
    let passed = limited.passed().ok_or("the source named no place")?;
    assert_eq!(passed.id, "a3");
    Ok(())
}

/// Origins of `.` with no worker (its path does not exist): a denied record is
/// absent and not counted, and no row carries a value.
#[test]
fn origins_leave_out_a_denied_record_and_do_not_count_it() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let access = access::request(&TestSource::denying(&[("beta", "b1")]))?;
    let settings = Settings {
        time_ms: NonZeroU64::MIN.saturating_add(999),
        memory_bytes: NonZeroU64::MIN.saturating_add((1 << 26) - 1),
        outputs_per_record: NonZeroU64::MIN,
        scan_records: NonZeroU64::MIN.saturating_add(98),
        max_running: NonZeroUsize::MIN,
        worker: PathBuf::from("/no/such/haematite-jq"),
    };
    let catalogue = moment.storage.catalogue();
    let slots = Slots::new(NonZeroUsize::MIN);
    let seal = Seal::new()?;
    let context = Context {
        access: &access,
        reader: &moment.reader,
        catalogue: &catalogue,
        slots: &slots,
        settings: Some(&settings),
        max_page: NonZeroU32::MIN.saturating_add(99),
        seal: &seal,
    };
    let request = QueryRequest {
        selection: selection(),
        query: ".".to_owned(),
        parameters: BTreeMap::new(),
        page: QueryPage {
            rows: NonZeroU64::MIN.saturating_add(99),
            bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
        },
        continuation: None,
        rows: Rows::Origins,
    };
    let answer: Value = serde_json::from_str(run::run(&context, &request, None)?.get())?;
    let rows = answer["rows"].as_array().ok_or("no rows array")?;
    let mut given = Vec::new();
    for row in rows {
        let fields = row.as_object().ok_or("a row is not an object")?;
        assert_eq!(fields.len(), 1, "a row carries more than its origin: {row}");
        let origin = &row["origin"];
        let collection = origin["collection"].as_str().ok_or("no collection")?;
        let id = origin["id"].as_str().ok_or("no id")?;
        given.push(format!("{collection}/{id}"));
    }
    let expected: Vec<&str> = ORDERED
        .into_iter()
        .filter(|name| *name != "beta/b1")
        .collect();
    assert_eq!(given, expected);
    assert_eq!(answer["records_read"], 7);
    assert!(answer["engine"].is_null(), "{answer}");
    assert!(answer["next"].is_null(), "{answer}");
    Ok(())
}

/// The page of origins of `.` after `previous`, or the first.
fn origins_after(previous: Option<&Value>) -> Result<QueryRequest, Box<dyn std::error::Error>> {
    let mut at = selection();
    let continuation = if let Some(answer) = previous {
        let version = serde_json::from_value(answer["selection"]["version"].clone())?;
        at.version = Some(version);
        Some(answer["next"].as_str().ok_or("no next")?.to_owned())
    } else {
        None
    };
    Ok(QueryRequest {
        selection: at,
        query: ".".to_owned(),
        parameters: BTreeMap::new(),
        page: QueryPage {
            rows: NonZeroU64::MIN.saturating_add(99),
            bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
        },
        continuation,
        rows: Rows::Origins,
    })
}

/// Origins of `.` under a limit of three, the record at the place the first
/// page stops denied: it is absent and not counted, and the record after it is
/// given once.
#[test]
fn origins_cut_at_a_denied_place_give_the_record_after_it_once() -> TestResult {
    let moment = moment(SHARDS, &WRITES)?;
    let access = access::request(&TestSource::denying(&[("alpha", "a3")]))?;
    let settings = Settings {
        time_ms: NonZeroU64::MIN.saturating_add(999),
        memory_bytes: NonZeroU64::MIN.saturating_add((1 << 26) - 1),
        outputs_per_record: NonZeroU64::MIN,
        scan_records: NonZeroU64::MIN.saturating_add(2),
        max_running: NonZeroUsize::MIN,
        worker: PathBuf::from("/no/such/haematite-jq"),
    };
    let catalogue = moment.storage.catalogue();
    let slots = Slots::new(NonZeroUsize::MIN);
    let seal = Seal::new()?;
    let context = Context {
        access: &access,
        reader: &moment.reader,
        catalogue: &catalogue,
        slots: &slots,
        settings: Some(&settings),
        max_page: NonZeroU32::MIN.saturating_add(99),
        seal: &seal,
    };
    let mut given = Vec::new();
    let mut pages = Vec::new();
    let mut previous: Option<Value> = None;
    loop {
        let text = run::run(&context, &origins_after(previous.as_ref())?, None)?;
        let answer: Value = serde_json::from_str(text.get())?;
        let rows = answer["rows"].as_array().ok_or("no rows array")?;
        pages.push(rows.len());
        for row in rows {
            let collection = row["origin"]["collection"]
                .as_str()
                .ok_or("no collection")?;
            let id = row["origin"]["id"].as_str().ok_or("no id")?;
            given.push(format!("{collection}/{id}"));
        }
        if answer["next"].is_null() {
            assert_eq!(answer["records_read"], 7);
            break;
        }
        previous = Some(answer);
    }
    let expected: Vec<&str> = ORDERED
        .into_iter()
        .filter(|name| *name != "alpha/a3")
        .collect();
    assert_eq!(given, expected);
    // The first page stops at the denied a3, having given a1 and a2.
    assert_eq!(pages.first(), Some(&2), "{pages:?}");
    Ok(())
}

#[test]
fn nonzero_origin_batch_readback() -> TestResult {
    let directory = tempfile::tempdir()?;
    let storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count: 1,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    let reader = storage.reader()?;
    let at = storage.capture(&selection())?;
    let access = access::request(&Standalone)?;
    for (remaining, expected) in [(0, 1), (1, 1), (17, 17), (u64::MAX, 64)] {
        let source = AuthorizedSource::origins(
            &access,
            &reader,
            &at,
            &Scope::All,
            None,
            std::hint::black_box(remaining),
        )?;
        assert_eq!(source.batch().get(), expected);
    }
    let source = AuthorizedSource::new(&access, &reader, &at, &Scope::All, None, 17)?;
    assert_eq!(source.batch().get(), 1);
    Ok(())
}

#[test]
fn nonzero_query_response_read_stops() -> TestResult {
    let moment = moment(1, &[("records", "a"), ("records", "b")])?;
    let access = access::request(&Standalone)?;
    let settings = Settings {
        time_ms: NonZeroU64::MIN.saturating_add(999),
        memory_bytes: NonZeroU64::MIN.saturating_add((1 << 26) - 1),
        outputs_per_record: NonZeroU64::MIN,
        scan_records: NonZeroU64::MIN.saturating_add(9),
        max_running: NonZeroUsize::MIN,
        worker: PathBuf::from("unused-worker"),
    };
    let catalogue = moment.storage.catalogue();
    let slots = Slots::new(NonZeroUsize::MIN);
    let seal = Seal::new()?;
    let context = Context {
        access: &access,
        reader: &moment.reader,
        catalogue: &catalogue,
        slots: &slots,
        settings: Some(&settings),
        max_page: NonZeroU32::MIN.saturating_add(9),
        seal: &seal,
    };
    let request = || QueryRequest {
        selection: selection(),
        query: ".".to_owned(),
        parameters: BTreeMap::new(),
        page: QueryPage {
            rows: NonZeroU64::MIN.saturating_add(9),
            bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
        },
        continuation: None,
        rows: Rows::Origins,
    };
    for given in [None, Some(65_536), Some(usize::MAX)] {
        let limit = given
            .map(|value| NonZeroUsize::new(value).ok_or("zero byte limit"))
            .transpose()?;
        let text = run::run(&context, &request(), std::hint::black_box(limit))?;
        let answer: Value = serde_json::from_str(text.get())?;
        assert_eq!(answer["rows"].as_array().ok_or("rows absent")?.len(), 2);
        if let Some(limit) = limit {
            assert!(text.get().len() <= limit.get());
        }
    }
    let refused = run::run(&context, &request(), Some(NonZeroUsize::MIN))
        .err()
        .ok_or("minimum returned an empty success")?;
    assert_eq!(refused.code, "response_limit");
    let needed: usize = refused
        .message
        .split_whitespace()
        .find_map(|word| word.parse().ok())
        .ok_or("the refusal names no size")?;
    let enough = NonZeroUsize::new(needed + 4096).ok_or("zero byte limit")?;
    let text = run::run(&context, &request(), Some(enough))?;
    assert!(text.get().len() <= enough.get());
    Ok(())
}

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