haematite-cli 0.21.0

The haematite command: serve a store, its door and Studio, answer MCP, query with jq, and host the store's assistant seat
use std::error::Error;
use std::num::NonZeroU32;

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

use super::{Captured, Failure, Receipt, Storage};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentImport, ImportedDocument, Mutation, Selection, VectorWrite,
};

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

pub(super) const ONE_WRITE: SyncCounts = SyncCounts {
    record_installs: 1,
    barrier_device_flushes: 1,
    ref_file_full_flushes: 1,
    directory_full_flushes: 1,
};

pub(super) fn at(storage: &Storage) -> Result<Captured, Failure> {
    storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })
}

pub(super) fn apply(storage: &mut Storage, id: &str, change: Change) -> Result<Receipt, Failure> {
    storage.mutate_bounded(
        &Mutation {
            request_id: id.to_owned(),
            branch: "main".to_owned(),
            expected: at(storage)?.version,
            change,
        },
        NonZeroU32::MIN.saturating_add(7),
    )
}

pub(super) fn record_id(index: usize) -> String {
    format!("record-{index:03}")
}

pub(super) fn vector(index: usize) -> VectorWrite {
    VectorWrite {
        collection: "records".to_owned(),
        document_id: record_id(index),
        field: "supplied".to_owned(),
        vector: vec![1.0, 2.0, 3.0],
    }
}

pub(super) fn store(rows: usize) -> 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: 1,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::reclaim(0),
        },
        NodeCacheBudget::Unlimited,
        std::time::Duration::from_secs(1),
    )?;
    apply(
        &mut storage,
        "collection",
        Change::Collection(CollectionWrite {
            collection: "records".to_owned(),
            schema: None,
        }),
    )?;
    let mut documents = (0..rows)
        .map(|index| {
            Ok(ImportedDocument {
                document_id: record_id(index),
                document: RawValue::from_string("{}".to_owned())?,
            })
        })
        .collect::<Result<Vec<_>, serde_json::Error>>()?;
    documents.push(ImportedDocument {
        document_id: "anchor".to_owned(),
        document: RawValue::from_string("{}".to_owned())?,
    });
    apply(
        &mut storage,
        "records",
        Change::PutMany(DocumentImport {
            collection: "records".to_owned(),
            documents,
        }),
    )?;
    let mut declaration = vector(0);
    declaration.document_id = "anchor".to_owned();
    apply(&mut storage, "supplied", Change::VectorPut(declaration))?;
    Ok((directory, storage))
}