Skip to main content

haematite_cli/serve/query/
parser.rs

1//! Read a query's source and ordered stages without opening any records.
2
3mod group_by;
4mod r#match;
5#[path = "run/percentile_parse.rs"]
6mod percentile;
7mod rerank;
8mod rg;
9mod shape;
10mod source;
11mod spans;
12mod stages;
13
14#[cfg(test)]
15mod rerank_tests;
16pub(in crate::serve) mod window;
17
18use std::fmt;
19use std::num::NonZeroU32;
20
21pub use r#match::{Direction, Node, Relation};
22pub use source::{Argument, Root, Source};
23pub use stages::{Question, Stage};
24
25use crate::serve::protocol::Metric;
26
27/// A byte range in the original query, excluding surrounding whitespace.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct Span {
30    pub start: usize,
31    pub end: usize,
32}
33
34/// Keep a parsed operation beside its location for query checking.
35#[derive(Debug, PartialEq, Eq)]
36pub struct Spanned<T> {
37    pub value: T,
38    pub span: Span,
39}
40
41#[derive(Debug, PartialEq, Eq)]
42pub struct Plan {
43    pub source: Source,
44    pub stages: Vec<Spanned<Stage>>,
45}
46
47/// The stored-vector search understood by the existing query runner.
48#[derive(Debug, PartialEq, Eq)]
49pub struct Search {
50    pub field: String,
51    pub record: String,
52    pub measure: Metric,
53    pub limit: NonZeroU32,
54}
55
56impl Plan {
57    #[must_use]
58    pub fn plain(text: &str) -> Self {
59        Self {
60            source: Source::All,
61            stages: vec![Spanned {
62                value: Stage::Jq(text.to_owned()),
63                span: Span {
64                    start: 0,
65                    end: text.len(),
66                },
67            }],
68        }
69    }
70}
71
72/// A named refusal and the byte range to mark in the query text.
73#[derive(Debug, PartialEq, Eq)]
74pub enum Refused {
75    NotYet(&'static str, Span),
76    Unserved(&'static str, Span),
77    Parse(String, Span),
78    WholeRecords(String, Span),
79    Percentile(String, Span),
80    Flag(String, Span),
81    Window(&'static str, String, Span),
82}
83
84impl std::error::Error for Refused {}
85
86impl Refused {
87    #[must_use]
88    pub const fn code(&self) -> &'static str {
89        match self {
90            Self::NotYet(..) | Self::Unserved(..) | Self::WholeRecords(..) => {
91                "query_stage_not_available"
92            }
93            Self::Parse(..) => "query_parse",
94            Self::Percentile(..) => "query_percentile_range",
95            Self::Flag(..) => "query_flag",
96            Self::Window(code, ..) => code,
97        }
98    }
99
100    #[must_use]
101    pub const fn span(&self) -> Span {
102        match self {
103            Self::NotYet(_, span)
104            | Self::Unserved(_, span)
105            | Self::Parse(_, span)
106            | Self::WholeRecords(_, span)
107            | Self::Percentile(_, span)
108            | Self::Flag(_, span)
109            | Self::Window(_, _, span) => *span,
110        }
111    }
112}
113
114impl fmt::Display for Refused {
115    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
116        match self {
117            Self::NotYet(stage, _) => write!(formatter, "{stage} is not available yet")?,
118            Self::Unserved(words, _) => formatter.write_str(words)?,
119            Self::Parse(message, _)
120            | Self::WholeRecords(message, _)
121            | Self::Percentile(message, _)
122            | Self::Flag(message, _)
123            | Self::Window(_, message, _) => formatter.write_str(message)?,
124        }
125        let span = self.span();
126        write!(formatter, " at bytes {}..{}", span.start, span.end)
127    }
128}
129
130/// Read a source and stages in exactly the order written.
131///
132/// # Errors
133/// Returns a named refusal with its original byte span for invalid or unavailable syntax.
134pub fn pipeline(text: &str) -> Result<Plan, Refused> {
135    pipeline_with(text, |_, _| {})
136}
137
138/// Report each completed operation while applying the same pipeline rules.
139pub(super) fn pipeline_with(
140    text: &str,
141    mut completed: impl FnMut(&str, Span),
142) -> Result<Plan, Refused> {
143    let mut input = Input { text, rest: text };
144    input.trim();
145    let (mut plan, complete) = start_pipeline(&mut input, &mut completed)?;
146    if complete {
147        return Ok(plan);
148    }
149    let mut saw_jq = false;
150    let mut saw_count = false;
151    loop {
152        input.trim();
153        let start = input.position();
154        stages::position(&input, &plan, saw_count, saw_jq)?;
155        let value = stages::read(&mut input)?;
156        let span = Span {
157            start,
158            end: input.position(),
159        };
160        r#match::position(&value, &plan.source, plan.stages.len(), span)?;
161        if matches!(value, Stage::Jq(_)) {
162            if saw_jq {
163                return Err(Refused::NotYet("a second jq stage", span));
164            }
165            saw_jq = true;
166        }
167        let name = match &value {
168            Stage::Similar { .. } => "similar",
169            Stage::Rerank { .. } => "rerank",
170            Stage::Limit(_) => "limit",
171            Stage::Sort { .. } => "sort",
172            Stage::Count => "count",
173            Stage::Percentile { .. } => "percentile",
174            Stage::GroupBy { .. } => "group_by",
175            Stage::Jq(_) => "jq",
176            Stage::Shape(_) => "{",
177            Stage::Rg { .. } => "rg",
178            Stage::Match { .. } => "match",
179            Stage::Window(_) => "window",
180        };
181        completed(name, span);
182        saw_count = matches!(value, Stage::Count);
183        plan.stages.push(Spanned { value, span });
184        input.trim();
185        if input.rest.is_empty() {
186            return Ok(plan);
187        }
188        if saw_count && !input.rest.starts_with('|') {
189            return Err(Refused::Parse(
190                "nothing follows count".to_owned(),
191                spans::stage(input.rest, input.position()),
192            ));
193        }
194        input.require("|", "a stage is followed only by `|`")?;
195    }
196}
197
198fn start_pipeline(
199    input: &mut Input<'_>,
200    completed: &mut impl FnMut(&str, Span),
201) -> Result<(Plan, bool), Refused> {
202    let source_start = input.position();
203    let source_word = input.word().to_owned();
204    let source_word = if input.rest.starts_with('{') {
205        "{".to_owned()
206    } else {
207        source_word
208    };
209    let mut initial_shape = None;
210    let source = if input.rest.starts_with('\'') {
211        Source::All
212    } else if input.rest.starts_with('{') {
213        let (source, span, selection) = shape::source(input.rest, input.position())?;
214        input.advance(span.end - span.start);
215        initial_shape = selection.map(|selection| Spanned {
216            value: Stage::Shape(selection),
217            span,
218        });
219        source
220    } else {
221        source::read(input)?
222    };
223    let mut plan = Plan {
224        source,
225        stages: Vec::new(),
226    };
227    if let Some(stage) = initial_shape {
228        plan.stages.push(stage);
229    }
230    if plan.source != Source::All {
231        let consumed = input
232            .text
233            .split_at(source_start)
234            .1
235            .split_at(input.position() - source_start)
236            .0;
237        let end = source_start + consumed.trim_end().len();
238        completed(
239            &source_word,
240            Span {
241                start: source_start,
242                end,
243            },
244        );
245        input.trim();
246        if input.rest.is_empty() {
247            return Ok((plan, true));
248        }
249        input.require("|", "a collection is followed only by `|`")?;
250    }
251    Ok((plan, false))
252}
253
254struct Input<'a> {
255    text: &'a str,
256    rest: &'a str,
257}
258
259impl Input<'_> {
260    const fn position(&self) -> usize {
261        self.text.len() - self.rest.len()
262    }
263
264    fn trim(&mut self) {
265        self.rest = self.rest.trim_start();
266    }
267
268    const fn advance(&mut self, bytes: usize) {
269        self.rest = self.rest.split_at(bytes).1;
270    }
271
272    fn take(&mut self, token: &str) -> bool {
273        if let Some(rest) = self.rest.strip_prefix(token) {
274            self.rest = rest;
275            true
276        } else {
277            false
278        }
279    }
280
281    fn require(&mut self, token: &str, message: &str) -> Result<(), Refused> {
282        self.trim();
283        if self.take(token) {
284            Ok(())
285        } else {
286            Err(self.parse(message))
287        }
288    }
289
290    fn word(&self) -> &str {
291        let length = self
292            .rest
293            .find(|character| !bare(character))
294            .unwrap_or(self.rest.len());
295        self.rest.split_at(length).0
296    }
297
298    fn token_span(&self) -> Span {
299        let length = self.word().len();
300        Span {
301            start: self.position(),
302            end: self.position()
303                + if length == 0 {
304                    self.rest.chars().next().map_or(0, char::len_utf8)
305                } else {
306                    length
307                },
308        }
309    }
310
311    fn parse(&self, message: impl Into<String>) -> Refused {
312        Refused::Parse(message.into(), self.token_span())
313    }
314
315    fn not_yet(&self, name: &'static str) -> Refused {
316        Refused::NotYet(name, self.token_span())
317    }
318
319    fn reserved(&self) -> Result<(), Refused> {
320        let name = if self.rest.starts_with('{') {
321            "{"
322        } else {
323            self.word()
324        };
325        super::words::reserved(name).map_or(Ok(()), |word| Err(self.not_yet(word)))
326    }
327
328    fn string(&mut self, what: &str) -> Result<String, Refused> {
329        self.trim();
330        if !self.rest.starts_with('"') {
331            return Err(self.parse(format!("{what} is written in double quotes")));
332        }
333        let mut escaped = false;
334        for (index, character) in self.rest.char_indices().skip(1) {
335            if escaped {
336                escaped = false;
337            } else if character == '\\' {
338                escaped = true;
339            } else if character == '"' {
340                let (literal, rest) = self.rest.split_at(index + 1);
341                let value = serde_json::from_str(literal).map_err(|error| {
342                    Refused::Parse(
343                        format!("{what} is not a string: {error}"),
344                        Span {
345                            start: self.position(),
346                            end: self.position() + literal.len(),
347                        },
348                    )
349                })?;
350                self.rest = rest;
351                return Ok(value);
352            }
353        }
354        Err(self.parse(format!("{what}'s quotes are not closed")))
355    }
356}
357
358const fn first(character: char) -> bool {
359    character.is_ascii_alphanumeric() || character == '_'
360}
361
362const fn bare(character: char) -> bool {
363    first(character) || matches!(character, '.' | '-')
364}
365
366#[cfg(test)]
367#[path = "parser_tests.rs"]
368mod tests;
369
370#[cfg(test)]
371#[path = "parser_pipeline_tests.rs"]
372mod pipeline_tests;
373
374#[cfg(test)]
375#[path = "parser/shape_tests.rs"]
376mod shape_tests;