haematite-cli 0.21.8

//! `query.run`: text in, one page of rows out, each beside its origin, read
//! through the request's access at one captured moment.
//!
//! Pages retain their captured moment and charge every readable input.
//!
//! Asked for `"rows": "origins"`, each row carries its origin alone. The text
//! Identity sources and limit-only pipelines need no worker or record body.
//! Other queries omit values after executing their stages.

mod answer;
mod entry;
mod failure;
mod matching;
mod plan;
mod rerank;
mod rerank_pipeline;
mod rerank_vectors;

#[cfg(all(test, unix))]
mod similar_after_tests;

#[cfg(all(test, unix))]
mod similar_text_tests;

#[cfg(all(test, unix))]
mod rerank_test_support;
mod rg;
mod shape;
pub(in crate::serve) mod window;

pub use entry::run;

#[cfg(test)]
#[path = "stages_tests.rs"]
mod stages_tests;

use answer::{Head, answer, fitted, named, rows_bytes};
use failure::{refused, start_failure};
use plan::present;
use rerank_pipeline::worked;

use std::collections::BTreeMap;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use haematite::db::ReadOnlyDatabase;
use haematite_query::Ready;
use haematite_query::page::{Allowance, BINDING_NAMES, Origin, Resume};
use serde_json::value::RawValue;
use serde_json::{Value, json};

use super::admission::Slots;
use super::parser::Refused;
use super::source::{AuthorizedSource, Place};
use super::{Binding, Continuation, Engine, Rows, Seal, Settings, origins};
use crate::serve::access::Access;
use crate::serve::contracts::Caller;
use crate::serve::protocol::{QueryRequest, Selection};
use crate::serve::service::Service;
use crate::serve::storage::{Captured, Catalogue, Failure};

/// The most bytes a record id may have (`contracts::identifier`).
const ID_BYTES: usize = 512;

/// What one call of `query.run` reads with, all owned by the service.
pub struct Context<'a> {
    pub access: &'a Access,
    pub reader: &'a ReadOnlyDatabase,
    pub catalogue: &'a Catalogue,
    /// The service's one admission limit, never made per request.
    pub slots: &'a Slots,
    pub settings: Option<&'a Settings>,
    pub max_page: NonZeroU32,
    /// The service's key for continuations, held only in its memory.
    pub seal: &'a Seal,
}

/// One page of `request`'s rows.
///
/// # Errors
///
/// A named `Failure` for every refusal; never a shortened answer.
pub(in crate::serve) fn served(
    context: &Context<'_>,
    request: &QueryRequest,
    response_byte_limit: Option<NonZeroUsize>,
    service: Option<&Service>,
) -> Result<Box<RawValue>, Failure> {
    let token = continuation(request, context.seal)?;
    let settings = settings(context, request.page.rows)?;
    let budget = budget(request.page.bytes, response_byte_limit)?;
    let now = token
        .as_ref()
        .map_or_else(clock, |token| Ok(token.now.clone()))?;
    let parameters = parameters(&request.parameters, &now)?;
    let parsed = shape::parse(&request.query)?;
    if let Err(refused @ Refused::Window(..)) = &parsed {
        return Err(Failure::new(refused.code(), refused.to_string()));
    }
    entry::ready(service, parsed.as_ref().ok())?;
    let at = context
        .catalogue
        .capture_held(&moment(&request.selection, token.as_ref()))?;
    let asked = Asked {
        request,
        token: token.as_ref(),
        now: &now,
        at: &at,
        budget,
        scan: settings.scan_records.get(),
    };
    if let Ok(plan) = &parsed {
        plan::check(context.reader, &at, plan)?;
    }
    if request.rows == Rows::Origins
        && let Some(executed) = plan::unworked(
            &request.query,
            parsed.as_ref().ok(),
            context,
            &at,
            &now,
            settings.scan_records.get(),
        )?
    {
        return origins_only(context, &asked, &executed);
    }
    worked(context, &asked, settings, parsed, parameters, service)
}

/// A request as one page reads it: its continuation, clock, moment and limits.
struct Asked<'a> {
    request: &'a QueryRequest,
    token: Option<&'a Continuation>,
    now: &'a RawValue,
    at: &'a Captured,
    budget: NonZeroUsize,
    scan: u64,
}

/// One page of origins the service reads itself from revision entries: no
/// worker, no body. The same bounds, cut and sealed `next` as a page of values;
/// its `engine` is null.
fn origins_only(
    context: &Context<'_>,
    asked: &Asked<'_>,
    executed: &plan::Executed,
) -> Result<Box<RawValue>, Failure> {
    let scope = &executed.scope;
    let limits = executed.limits();
    let binding = bind(asked.request, &unworked_engine(), asked.scan)?;
    if let Some(token) = asked.token {
        token.check_binding(&binding)?;
    }
    present(context.reader, asked.at, scope)?;
    let resume = asked.token.map(|token| &token.resume);
    let read_before = resume.map_or(0, |resume| resume.records_read);
    let start = resume.map(|resume| Place {
        collection: resume.pending.collection.clone(),
        id: resume.pending.id.clone(),
        after: resume.passed,
    });
    let mut source = AuthorizedSource::origins(
        context.access,
        context.reader,
        asked.at,
        scope,
        start.as_ref(),
        asked.scan,
    )
    .map_err(|error| refused(&error))?;
    let widest = widest_next(
        asked.at,
        named(scope),
        asked.now,
        &binding,
        context.seal,
        limits.len(),
    )?;
    let mut head = Head {
        reading: executed.reading,
        stages: &executed.names,
        collection: named(scope),
        read: u64::MAX,
        skipped: u64::MAX,
        search: &Value::Null,
        unscored: None,
        next: Some(&widest),
    };
    let envelope = answer(asked.at, &json!([]), &Value::Null, &head);
    let allowance = Allowance {
        rows: asked.request.page.rows,
        bytes: rows_bytes(asked.budget, &envelope)?,
        scan_records: read_before.saturating_add(asked.scan),
    };
    let page =
        origins::page(&mut source, &allowance, resume, &limits).map_err(|error| refused(&error))?;
    let next = page
        .next
        .map(|resume| {
            let at = asked.at;
            let now = asked.now.to_owned();
            Continuation::new(at.branch.clone(), at.version, resume, now, binding)?
                .encode(context.seal)
        })
        .transpose()?;
    let rows = serde_json::value::to_raw_value(&page.rows).map_err(Failure::storage)?;
    head.read = page.records_read;
    head.skipped = 0;
    head.next = next.as_deref();
    fitted(
        answer(asked.at, &json!([]), &Value::Null, &head),
        &rows,
        asked.budget,
    )
}

/// The query settings, once the page asked for is within the service's own.
fn settings<'a>(context: &Context<'a>, rows: NonZeroU64) -> Result<&'a Settings, Failure> {
    if rows.get() > u64::from(context.max_page.get()) {
        return Err(Failure::new(
            "page_limit",
            format!("maximum page size is {}", context.max_page),
        ));
    }
    context.settings.ok_or_else(|| {
        Failure::new(
            "query_engine_unavailable",
            "query settings are not configured",
        )
    })
}

/// The continuation the request presents, checked against the moment it names.
fn continuation(request: &QueryRequest, seal: &Seal) -> Result<Option<Continuation>, Failure> {
    let Some(encoded) = request.continuation.as_deref() else {
        return Ok(None);
    };
    let token = Continuation::decode(encoded, seal)?;
    token.check_moment(&request.selection.branch, request.selection.version)?;
    Ok(Some(token))
}

/// The moment a page is read at: a later page's is its continuation's.
fn moment(selection: &Selection, token: Option<&Continuation>) -> Selection {
    Selection {
        branch: selection.branch.clone(),
        version: token.map_or(selection.version, |token| Some(token.version)),
    }
}

/// The one caller a continuation is bound to until a request carries its own.
/// Every page is read through the presenting request's access, so a
/// continuation grants nothing.
fn caller() -> Result<Caller, Failure> {
    Caller::local("standalone")
}

/// The worker that answered, as a continuation is bound to it.
fn worker_engine(ready: &Ready) -> Engine {
    Engine {
        protocol: ready.protocol,
        worker: ready.worker.clone(),
        engine: ready.engine.clone(),
        jaq_core: ready.jaq_core.clone(),
        jaq_std: ready.jaq_std.clone(),
        jaq_json: ready.jaq_json.clone(),
    }
}

/// The service itself, as a page of origins with no worker is bound to it.
fn unworked_engine() -> Engine {
    Engine {
        protocol: haematite_query::PROTOCOL,
        worker: "haem".to_owned(),
        engine: "origins".to_owned(),
        jaq_core: String::new(),
        jaq_std: String::new(),
        jaq_json: String::new(),
    }
}

/// What a continuation is bound to: the text, the caller's own parameters,
/// the engine, the caller, and the page's limits and row form as asked.
fn bind(request: &QueryRequest, engine: &Engine, scan: u64) -> Result<Binding, Failure> {
    let limits = Allowance {
        rows: request.page.rows,
        bytes: request.page.bytes,
        scan_records: scan,
    };
    Binding::new(
        &request.query,
        &request.parameters,
        engine,
        &caller()?,
        &limits,
        request.rows,
    )
}

/// The longest `next` this page could give: a pending record with the longest
/// collection (when the scope names none) and id, every byte of them escaped,
/// the longest score and the largest counts.
/// Room for it is kept, so a page cut for its bytes still carries its `next`.
fn widest_next(
    at: &Captured,
    collection: Option<&str>,
    now: &RawValue,
    binding: &Binding,
    seal: &Seal,
    limit_counts: usize,
) -> Result<String, Failure> {
    let version = serde_json::value::to_raw_value(&at.version).map_err(Failure::storage)?;
    let pending = Origin {
        collection: collection.map_or_else(|| "\u{1}".repeat(ID_BYTES), str::to_owned),
        id: "\u{1}".repeat(ID_BYTES),
        version,
        score: Some(-f64::MIN_POSITIVE),
    };
    // `false` is written longer than `true`.
    let resume = Resume {
        records_read: u64::MAX,
        pending,
        discard_outputs: u64::MAX,
        passed: false,
        passed_counts: vec![u64::MAX; limit_counts],
    };
    let widest = Continuation::new(
        at.branch.clone(),
        at.version,
        resume,
        now.to_owned(),
        binding.clone(),
    )?;
    widest.encode(seal)
}

/// The smaller of the request's bytes and the transport's limit.
fn budget(
    requested: NonZeroU64,
    response_byte_limit: Option<NonZeroUsize>,
) -> Result<NonZeroUsize, Failure> {
    let requested = NonZeroUsize::try_from(requested)
        .map_err(|_| Failure::new("page_limit", "the requested bytes do not fit this machine"))?;
    let bytes = response_byte_limit
        .map_or_else(|| requested.get(), |limit| limit.get().min(requested.get()));
    NonZeroUsize::new(bytes)
        .ok_or_else(|| Failure::new("page_limit", "selected byte limit is zero"))
}

/// No one value can be larger than the whole answer, so the worker's frames
/// are bounded by it, and never below the smallest frame a refusal needs.
fn frame_bytes(budget: NonZeroUsize) -> Result<NonZeroU32, Failure> {
    let bytes = u32::try_from(budget.get())
        .unwrap_or(u32::MAX)
        .max(haematite_query::MINIMUM_FRAME_BYTES);
    NonZeroU32::new(bytes).ok_or_else(|| Failure::new("page_limit", "a frame of no bytes"))
}

/// The caller's parameters and `$now`, which a first page reads from the clock
/// and a later page keeps. The per-record names and `now` are reserved.
fn parameters(
    given: &BTreeMap<String, Box<RawValue>>,
    now: &RawValue,
) -> Result<BTreeMap<String, Box<RawValue>>, Failure> {
    for name in given.keys() {
        if name == "now" || BINDING_NAMES.contains(&name.as_str()) {
            return Err(Failure::new(
                "query_parameter_reserved",
                format!("${name} is reserved and cannot be supplied"),
            ));
        }
    }
    let mut parameters = given.clone();
    parameters.insert("now".to_owned(), now.to_owned());
    Ok(parameters)
}

/// `$now`, read from the clock once, when a query is first asked.
fn clock() -> Result<Box<RawValue>, Failure> {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|error| Failure::new("query_clock", error.to_string()))?;
    clock_json(now)
}

fn clock_json(now: Duration) -> Result<Box<RawValue>, Failure> {
    // A fixed fraction keeps continuation envelopes independent of clock precision.
    RawValue::from_string(format!("{}.{:09}", now.as_secs(), now.subsec_nanos()))
        .map_err(Failure::storage)
}

#[cfg(test)]
mod nonzero_tests;