haematite-cli 0.21.9

#[path = "plan_dispatch.rs"]
mod dispatch;

use dispatch::mapped;

use std::num::NonZeroU64;

use haematite::db::ReadOnlyDatabase;
use haematite_query::handle::Error;
use haematite_query::page::{Origin, Record, RecordSource, Resume};
use haematite_query::stage::Stage;
use haematite_query::{Reading, Ready};
use serde_json::Value;
use serde_json::value::RawValue;

use super::super::parser::{self, Plan, Root, Source};
use super::super::search::{Reach, Read, Records, Search, hits};
use super::super::shaped_source::ShapedSource;
use super::super::source::{AuthorizedSource, Place, Scope};
use super::{Context, refused};
use crate::serve::schema;
use crate::serve::service::Service;
use crate::serve::storage::{Captured, Failure, reads};

pub(super) struct Executed {
    pub scope: Scope,
    pub roots: Vec<Root>,
    pub search: Option<Search>,
    pub matching: Option<super::matching::Request>,
    pub reranks: Vec<super::rerank::Request>,
    pub reading: &'static str,
    pub stages: Vec<Stage>,
    pub names: Vec<&'static str>,
}

impl Executed {
    pub fn limits(&self) -> Vec<NonZeroU64> {
        self.stages
            .iter()
            .filter_map(|stage| match stage {
                Stage::Limit(limit) => Some(*limit),
                _ => None,
            })
            .collect()
    }
}

pub(super) fn jq(plan: &Plan) -> &str {
    plan.stages
        .iter()
        .find_map(|stage| match &stage.value {
            parser::Stage::Jq(text) => Some(text.as_str()),
            _ => None,
        })
        .unwrap_or(".")
}

/// Check argument names against the captured shape before starting a worker.
pub(super) fn check(reader: &ReadOnlyDatabase, at: &Captured, plan: &Plan) -> Result<(), Failure> {
    super::matching::check(reader, at, plan)?;
    match &plan.source {
        Source::Shaped(roots) => {
            for root in roots {
                let collection = reads::collection(reader, at, &root.name)?
                    .ok_or_else(|| missing(&root.name))?;
                check_arguments(&root.name, &root.arguments, collection.schema.as_ref())?;
            }
            Ok(())
        }
        Source::Collection { name, arguments } => {
            if arguments.is_empty() && !super::shape::shaped(plan) {
                return Ok(());
            }
            let collection = reads::collection(reader, at, name)?.ok_or_else(|| missing(name))?;
            check_arguments(name, arguments, collection.schema.as_ref())
        }
        Source::All => Ok(()),
    }
}

fn check_arguments(
    name: &str,
    arguments: &[parser::Argument],
    shape: Option<&Value>,
) -> Result<(), Failure> {
    if arguments.is_empty() {
        return Ok(());
    }
    if let Some(shape) = shape {
        schema::Collection::compile(shape)
            .map_err(|error| Failure::new(error.code(), error.to_string()))?;
        let properties = shape.get("properties").and_then(Value::as_object);
        for argument in arguments {
            if !properties.is_some_and(|fields| fields.contains_key(&argument.field)) {
                return Err(Failure::new(
                    "query_field_unknown",
                    format!(
                        "collection {name} has no top-level property {}",
                        argument.field
                    ),
                ));
            }
        }
    }
    Ok(())
}

pub(super) fn present(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    scope: &Scope,
) -> Result<(), Failure> {
    if let Scope::Collection(name) = scope
        && reads::collection(reader, at, name)?.is_none()
    {
        return Err(missing(name));
    }
    Ok(())
}

fn missing(name: &str) -> Failure {
    Failure::new(
        "collection_missing",
        format!("collection {name} does not exist at the captured version"),
    )
}

/// A plain jq reading retains its own source and one worker stage.
pub(super) fn scope(
    ready: &Ready,
    plan: Option<Plan>,
    memory: NonZeroU64,
    context: &Context<'_>,
    at: &Captured,
    now: &RawValue,
    scan: u64,
) -> Result<Executed, Failure> {
    match (ready.reading, plan) {
        (Some(reading @ (Reading::Candidate | Reading::Alternate)), Some(plan))
            if reading == Reading::Alternate || super::shape::shaped(&plan) =>
        {
            mapped(
                &plan,
                Some(memory),
                Some(&WindowContext {
                    context,
                    at,
                    now,
                    scan,
                }),
            )
        }
        (Some(Reading::Candidate), _) => Ok(Executed {
            scope: Scope::All,
            roots: Vec::new(),
            search: None,
            matching: None,
            reranks: Vec::new(),
            reading: "jq",
            stages: vec![Stage::Jq],
            names: vec!["jq"],
        }),
        _ => Err(Failure::new(
            "query_protocol",
            "the worker named a reading this Start did not carry",
        )),
    }
}

struct WindowContext<'a, 'b> {
    context: &'a Context<'b>,
    at: &'a Captured,
    now: &'a RawValue,
    scan: u64,
}

/// Identity origins and limits need revision entries, without record bodies.
pub(super) fn unworked(
    text: &str,
    plan: Option<&Plan>,
    context: &Context<'_>,
    at: &Captured,
    now: &RawValue,
    scan: u64,
) -> Result<Option<Executed>, Failure> {
    if text.trim() == "." {
        return Ok(Some(Executed {
            scope: Scope::All,
            roots: Vec::new(),
            search: None,
            matching: None,
            reranks: Vec::new(),
            reading: "jq",
            stages: Vec::new(),
            names: Vec::new(),
        }));
    }
    let Some(plan) = plan else {
        return Ok(None);
    };
    let Source::Collection { arguments, .. } = &plan.source else {
        return Ok(None);
    };
    let limits_only = plan
        .stages
        .iter()
        .all(|stage| matches!(stage.value, parser::Stage::Limit(_)));
    if arguments.is_empty()
        && limits_only
        && (!plan.stages.is_empty() || text.trim_start().starts_with("collection("))
    {
        return mapped(
            plan,
            None,
            Some(&WindowContext {
                context,
                at,
                now,
                scan,
            }),
        )
        .map(Some);
    }
    Ok(None)
}

/// Keep an aggregate's origin even when its authorized source is empty.
pub(super) struct Stored<'a> {
    matching: Option<super::matching::Prepared<'a>>,
    rerank: Option<super::rerank::Prepared<'a>>,
    records: Records<'a>,
    origin: Origin,
}

impl<'a> Stored<'a> {
    pub fn prepare_match(
        &mut self,
        context: &Context<'a>,
        at: &'a Captured,
        request: Option<&super::matching::Request>,
    ) {
        self.matching = request.map(|request| super::matching::Prepared::new(context, at, request));
    }

    pub fn match_failure(&mut self) -> Option<Failure> {
        self.matching
            .as_mut()
            .and_then(super::matching::Prepared::take_failure)
    }

    /// Prepare all rerank questions before consuming the source records.
    ///
    /// # Errors
    /// Returns any refusal from validating or reading a rerank question.
    pub fn prepare(
        &mut self,
        context: &Context<'a>,
        at: &'a Captured,
        requests: &[super::rerank::Request],
        service: Option<&Service>,
        scan: u64,
    ) -> Result<(), Failure> {
        if !requests.is_empty() {
            self.rerank = Some(super::rerank::Prepared::new(
                context, at, requests, service, scan,
            )?);
        }
        Ok(())
    }

    pub fn unscored(&self) -> Option<u64> {
        self.rerank.as_ref().and_then(|rerank| rerank.unscored)
    }

    pub fn failure(&mut self, error: &Error) -> Failure {
        self.rerank
            .as_mut()
            .map_or_else(|| refused(error), |rerank| rerank.failure(error))
    }

    pub const fn compared(&self) -> Option<u64> {
        self.records.compared()
    }
}

impl RecordSource for Stored<'_> {
    fn match_record(
        &mut self,
        record: Record,
        remaining: u64,
    ) -> Result<(std::collections::VecDeque<Record>, u64), Error> {
        self.matching
            .as_mut()
            .ok_or_else(|| {
                Error::new(
                    haematite_query::Refusal::QueryProtocol,
                    "match was not prepared",
                )
            })?
            .expand(record, remaining)
    }

    fn scan_reads(&self) -> Option<u64> {
        match &self.records {
            Records::Shaped(source) => Some(source.scanned()),
            Records::Scan(source) if self.matching.is_some() => Some(source.scanned()),
            Records::Scan(_) | Records::Search(_) => None,
        }
    }

    fn rerank(
        &mut self,
        index: usize,
        records: Vec<Record>,
        remaining: u64,
    ) -> Result<(std::collections::VecDeque<Record>, u64), Error> {
        self.rerank
            .as_mut()
            .ok_or_else(|| {
                Error::new(
                    haematite_query::Refusal::QueryProtocol,
                    "rerank was not prepared",
                )
            })?
            .rank(index, records, remaining)
    }

    fn similar(
        &mut self,
        index: usize,
        records: Vec<Record>,
        remaining: u64,
        memory: u64,
    ) -> Result<(std::collections::VecDeque<Record>, u64), Error> {
        self.rerank
            .as_mut()
            .ok_or_else(|| {
                Error::new(
                    haematite_query::Refusal::QueryProtocol,
                    "similar was not prepared",
                )
            })?
            .rank_with_memory(index, records, remaining, Some(memory))
    }

    fn count_origin(&self) -> Result<Origin, Error> {
        Ok(self.origin.clone())
    }
    fn next_record(&mut self) -> Result<Option<Record>, Error> {
        self.records.next_record()
    }
    fn passed(&self) -> Option<Origin> {
        self.records.passed()
    }
}

/// A scan resumes at its pending record; a search resumes in its scored hits.
pub(super) fn records<'a>(
    context: &Context<'a>,
    at: &'a Captured,
    executed: &Executed,
    resume: Option<&Resume>,
    scan: u64,
    service: Option<&Service>,
) -> Result<Stored<'a>, Failure> {
    let scope = &executed.scope;
    let roots = &executed.roots;
    let search = executed.search.as_ref();
    let start = resume.filter(|_| search.is_none()).map(|resume| Place {
        collection: resume.pending.collection.clone(),
        id: resume.pending.id.clone(),
        after: resume.passed,
    });
    let records = if roots.is_empty() {
        let source = AuthorizedSource::new(
            context.access,
            context.reader,
            at,
            scope,
            start.as_ref(),
            scan,
        )
        .map_err(|error| refused(&error))?;
        match (search, scope) {
            (None, _) => Records::Scan(source),
            (Some(search), Scope::Collection(collection)) => {
                let found = hits(
                    source,
                    &Read {
                        access: context.access,
                        reader: context.reader,
                        service,
                    },
                    at,
                    collection,
                    search,
                    &reach(scan, resume),
                )?;
                Records::Search(found)
            }
            (Some(_), Scope::All) => {
                return Err(Failure::new(
                    "query_protocol",
                    "a search was read without its collection",
                ));
            }
        }
    } else {
        if search.is_some() {
            return Err(Failure::new(
                "query_protocol",
                "a shaped source cannot carry a search",
            ));
        }
        Records::Shaped(Box::new(
            ShapedSource::new(context.access, context.reader, at, roots, start, scan)
                .map_err(|error| refused(&error))?,
        ))
    };
    let collection = match scope {
        Scope::Collection(name) => name.clone(),
        Scope::All => String::new(),
    };
    Ok(Stored {
        matching: None,
        rerank: None,
        records,
        origin: Origin {
            collection,
            id: String::new(),
            version: serde_json::value::to_raw_value(&at.version).map_err(Failure::storage)?,
            score: None,
        },
    })
}

fn reach(scan: u64, resume: Option<&Resume>) -> Reach {
    Reach {
        bound: scan,
        skip: resume
            .map_or(0, |resume| resume.records_read)
            .saturating_sub(1),
    }
}

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