Skip to main content

haematite_cli/serve/query/parser/
source.rs

1use haematite_query::stage::shape::Selection;
2use serde_json::Value;
3
4use super::{Input, Refused, Span, first};
5
6#[derive(Debug, PartialEq, Eq)]
7pub enum Source {
8    All,
9    Shaped(Vec<Root>),
10    Collection {
11        name: String,
12        arguments: Vec<Argument>,
13    },
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct Root {
18    pub name: String,
19    pub arguments: Vec<Argument>,
20    pub selection: Selection,
21    pub span: Span,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct Argument {
26    pub field: String,
27    pub value: Value,
28}
29
30pub(super) fn read(input: &mut Input<'_>) -> Result<Source, Refused> {
31    if input.word() == "match" {
32        return Ok(Source::All);
33    }
34    input.reserved()?;
35    match input.word() {
36        "percentile" => {
37            return Err(input.parse("`percentile` is a stage and must follow a collection"));
38        }
39        "similar" | "sort" | "count" | "limit" | "window" | "rerank" => {
40            return Err(input.parse(format!(
41                "`{}` is a stage and must follow a collection",
42                input.word()
43            )));
44        }
45        "rg" | "grep" => {
46            return Err(input.parse("`rg` is a stage and must follow a collection"));
47        }
48        "group_by" => {
49            return Err(input.parse("`group_by` is a stage and must follow a collection"));
50        }
51        _ => {}
52    }
53    collection(input)
54}
55
56/// Root braces identify a collection, even when its name is a stage word.
57pub(super) fn collection(input: &mut Input<'_>) -> Result<Source, Refused> {
58    let name = input.word().to_owned();
59    if !name.starts_with(first) {
60        return Err(input.parse("a query that is not jq starts with a collection"));
61    }
62    input.advance(name.len());
63    input.trim();
64    if !input.take("(") {
65        return Ok(Source::Collection {
66            name,
67            arguments: Vec::new(),
68        });
69    }
70    if name == "collection" {
71        input.trim();
72        let start = input.position();
73        let name = input.string("a collection name")?;
74        if name.is_empty() || name.contains('\0') {
75            return Err(Refused::Parse(
76                "a collection name has 1 or more characters and no NUL".to_owned(),
77                Span {
78                    start,
79                    end: input.position(),
80                },
81            ));
82        }
83        input.require(")", "`collection(` takes one quoted name and a `)`")?;
84        return Ok(Source::Collection {
85            name,
86            arguments: Vec::new(),
87        });
88    }
89    let mut arguments = Vec::new();
90    loop {
91        input.trim();
92        let length = input
93            .rest
94            .find(|character| !first(character))
95            .unwrap_or(input.rest.len());
96        let field = input.rest.split_at(length).0.to_owned();
97        if field.is_empty() {
98            return Err(input.parse("a collection argument names a top-level field before `:`"));
99        }
100        input.advance(length);
101        input.require(
102            ":",
103            "a collection argument separates its field and JSON value with `:`",
104        )?;
105        input.trim();
106        let value = value(input)?;
107        arguments.push(Argument { field, value });
108        input.trim();
109        if input.take(")") {
110            return Ok(Source::Collection { name, arguments });
111        }
112        if input.word() != "and" {
113            return Err(input.parse("collection arguments are joined by `and` and closed by `)`"));
114        }
115        input.advance(3);
116    }
117}
118
119/// Find argument boundaries outside JSON strings, arrays and objects.
120fn value(input: &mut Input<'_>) -> Result<Value, Refused> {
121    let mut quoted = false;
122    let mut escaped = false;
123    let mut depth = 0_usize;
124    let mut end = input.rest.len();
125    for (index, character) in input.rest.char_indices() {
126        if quoted {
127            if escaped {
128                escaped = false;
129            } else if character == '\\' {
130                escaped = true;
131            } else if character == '"' {
132                quoted = false;
133            }
134            continue;
135        }
136        let (before, after) = input.rest.split_at(index);
137        match character {
138            '"' => quoted = true,
139            '[' | '{' => depth += 1,
140            ']' | '}' => depth = depth.saturating_sub(1),
141            ')' if depth == 0 => {
142                end = index;
143                break;
144            }
145            'a' if depth == 0
146                && before.ends_with(char::is_whitespace)
147                && after
148                    .strip_prefix("and")
149                    .is_some_and(|rest| rest.starts_with(char::is_whitespace)) =>
150            {
151                end = index;
152                break;
153            }
154            _ => {}
155        }
156    }
157    let literal = input.rest.split_at(end).0.trim_end();
158    let span = Span {
159        start: input.position(),
160        end: input.position() + literal.len(),
161    };
162    let value = serde_json::from_str(literal).map_err(|error| {
163        Refused::Parse(
164            format!("a collection argument needs a JSON value: {error}"),
165            span,
166        )
167    })?;
168    input.advance(end);
169    Ok(value)
170}
171
172#[cfg(test)]
173#[path = "source_tests.rs"]
174mod tests;