haematite-cli 0.21.3

The haematite command: serve a store, its door and Studio, answer MCP, query with jq, and host the store's assistant seat
Documentation
use std::num::{NonZeroU64, NonZeroUsize};

use haematite_query::Ready;
use haematite_query::page::Row;
use serde_json::value::RawValue;
use serde_json::{Value, json};

use super::super::Rows;
use super::super::origins::Named;
use super::super::source::Scope;
use crate::serve::storage::{Captured, Failure};

/// A page's rows as asked, as text (NB1): each value beside its origin, the
/// value the text the engine gave and never decoded again; or the origin
/// alone.
pub(super) fn rows(rows: Vec<Row>, form: Rows) -> Result<Box<RawValue>, Failure> {
    match form {
        Rows::Values => serde_json::value::to_raw_value(&rows),
        Rows::Origins => {
            let named: Vec<Named> = rows
                .into_iter()
                .map(|row| Named { origin: row.origin })
                .collect();
            serde_json::value::to_raw_value(&named)
        }
    }
    .map_err(Failure::storage)
}

/// `envelope` with `rows` as text, once it is known to fit its bytes.
pub(super) fn fitted(
    envelope: Value,
    rows: &RawValue,
    budget: NonZeroUsize,
) -> Result<Box<RawValue>, Failure> {
    let answer = crate::serve::changes::with_rows(envelope, rows)?;
    if answer.get().len() > budget.get() {
        return Err(Failure::new(
            "response_limit",
            "the answer exceeds its requested bytes",
        ));
    }
    Ok(answer)
}

/// The bytes left for the rows once the answer's `envelope` is counted.
pub(super) fn rows_bytes(budget: NonZeroUsize, envelope: &Value) -> Result<NonZeroU64, Failure> {
    let envelope = serde_json::to_vec(envelope)
        .map_err(Failure::storage)?
        .len();
    // The envelope's empty `[]` is counted again by the page as its rows.
    // It holds room for the widest continuation the page could need, so the
    // refusal names the size: a caller asks once more, with enough.
    let needed = envelope - 2;
    u64::try_from(budget.get())
        .ok()
        .and_then(|budget| budget.checked_sub(u64::try_from(needed).ok()?))
        .and_then(NonZeroU64::new)
        .ok_or_else(|| {
            Failure::new(
                "response_limit",
                format!(
                    "the answer's envelope takes {needed} bytes before any row; ask for more than that"
                ),
            )
        })
}

pub(super) fn engine(ready: &Ready) -> Value {
    json!({
        "engine": ready.engine,
        "worker": ready.worker,
        "protocol": ready.protocol,
        "jaq_core": ready.jaq_core,
        "jaq_std": ready.jaq_std,
        "jaq_json": ready.jaq_json,
    })
}

/// The collection a scope names; none for every record.
pub(super) fn named(scope: &Scope) -> Option<&str> {
    match scope {
        Scope::Collection(name) => Some(name),
        Scope::All => None,
    }
}

/// What an answer says beside its rows.
pub(super) struct Head<'a> {
    pub(super) stages: &'a [&'static str],
    pub(super) reading: &'a str,
    /// The collection a pipeline read; none for every record.
    pub(super) collection: Option<&'a str>,
    /// The records given to the program, on this page and every earlier one.
    pub(super) read: u64,
    pub(super) skipped: u64,
    /// Null without a search; otherwise that it was exact, and how many
    /// vectors this page's search compared.
    pub(super) search: &'a Value,
    pub(super) unscored: Option<u64>,
    pub(super) next: Option<&'a str>,
}

pub(super) fn answer(at: &Captured, rows: &Value, engine: &Value, head: &Head<'_>) -> Value {
    json!({
        "selection": {
            "branch": at.branch,
            "version": at.version,
            "roots": at.shards.iter().map(|shard| {
                json!({"shard": shard.shard_id, "root": shard.head.to_string()})
            }).collect::<Vec<_>>(),
        },
        "rows": rows,
        "next": head.next,
        "engine": engine,
        "reading": head.reading,
        "stages": head.stages,
        "collection": head.collection,
        "records_read": head.read,
        "skipped": head.skipped,
        "search": head.search,
        "unscored": head.unscored,
    })
}