haematite-cli 0.21.8

//! Read each root through record access before filtering and projecting its body.

use std::collections::VecDeque;

use haematite::db::ReadOnlyDatabase;
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::{Origin, Record, RecordSource};
use serde_json::Value;

use super::parser::Root;
use super::source::{AuthorizedSource, Place, Scope};
use crate::serve::access::Access;
use crate::serve::storage::Captured;

struct Current<'a> {
    root: Root,
    source: AuthorizedSource<'a>,
}

pub struct ShapedSource<'a> {
    access: &'a Access,
    reader: &'a ReadOnlyDatabase,
    at: &'a Captured,
    roots: VecDeque<Root>,
    current: Option<Current<'a>>,
    start: Option<Place>,
    scan_limit: u64,
    completed: u64,
    last: Option<Origin>,
    stopped: Option<Origin>,
}

impl<'a> ShapedSource<'a> {
    pub fn new(
        access: &'a Access,
        reader: &'a ReadOnlyDatabase,
        at: &'a Captured,
        roots: &[Root],
        start: Option<Place>,
        scan_limit: u64,
    ) -> Result<Self, Error> {
        let mut roots: VecDeque<_> = roots.iter().cloned().collect();
        if let Some(place) = &start {
            let index = roots
                .iter()
                .position(|root| root.name == place.collection)
                .ok_or_else(|| protocol("the start place lies outside the shaped roots"))?;
            drop(roots.drain(..index));
        }
        Ok(Self {
            access,
            reader,
            at,
            roots,
            current: None,
            start,
            scan_limit,
            completed: 0,
            last: None,
            stopped: None,
        })
    }

    pub const fn scanned(&self) -> u64 {
        self.completed
            + match &self.current {
                Some(current) => current.source.scanned(),
                None => 0,
            }
    }

    fn open(&mut self) -> Result<bool, Error> {
        let Some(root) = self.roots.pop_front() else {
            return Ok(false);
        };
        let remaining = self
            .scan_limit
            .checked_sub(self.completed)
            .ok_or_else(|| protocol("shaped source exceeded its scan allowance"))?;
        let source = AuthorizedSource::new(
            self.access,
            self.reader,
            self.at,
            &Scope::Collection(root.name.clone()),
            self.start.take().as_ref(),
            remaining,
        )?;
        self.current = Some(Current { root, source });
        Ok(true)
    }

    fn finish(&mut self) -> Result<(), Error> {
        let current = self
            .current
            .take()
            .ok_or_else(|| protocol("shaped source lost its current root"))?;
        self.completed = self
            .completed
            .checked_add(current.source.scanned())
            .ok_or_else(|| protocol("shaped source scan count overflow"))?;
        if let Some(last) = current.source.last_place()? {
            self.last = Some(last);
        }
        Ok(())
    }
}

impl RecordSource for ShapedSource<'_> {
    fn next_record(&mut self) -> Result<Option<Record>, Error> {
        if self.stopped.is_some() {
            return Ok(None);
        }
        loop {
            if self.current.is_none() && !self.open()? {
                return Ok(None);
            }
            let current = self
                .current
                .as_mut()
                .ok_or_else(|| protocol("shaped source has no current root"))?;
            if self.completed == self.scan_limit {
                // At a root boundary, name the last charged place without taking a new candidate.
                if current.source.has_candidates()? {
                    self.stopped = Some(self.last.clone().ok_or_else(|| {
                        Error::new(
                            Refusal::QueryScanLimit,
                            "a page that may read no records moves nowhere",
                        )
                    })?);
                    return Ok(None);
                }
                self.finish()?;
                continue;
            }
            let Some(mut record) = current.source.next_record()? else {
                if let Some(stopped) = current.source.passed() {
                    self.stopped = Some(stopped);
                    return Ok(None);
                }
                self.finish()?;
                continue;
            };
            if !current.root.arguments.is_empty() {
                let value: Value = serde_json::from_str(record.value.get())
                    .map_err(|error| protocol(&error.to_string()))?;
                if !current.root.arguments.iter().all(|argument| {
                    value
                        .as_object()
                        .and_then(|object| object.get(&argument.field))
                        == Some(&argument.value)
                }) {
                    continue;
                }
            }
            record.value = current.root.selection.project(&record.value)?;
            return Ok(Some(record));
        }
    }

    fn scan_reads(&self) -> Option<u64> {
        Some(self.scanned())
    }

    fn passed(&self) -> Option<Origin> {
        self.stopped.clone()
    }
}

fn protocol(message: &str) -> Error {
    Error::new(Refusal::QueryProtocol, message)
}