gq_core/query/
query_key.rs

1use std::{
2    borrow::Cow,
3    fmt::{self, Display, Formatter},
4    ops::Add,
5};
6
7use derive_getters::Getters;
8use derive_more::Constructor;
9use once_cell::sync::Lazy;
10use regex::Regex;
11use serde_json::Value;
12
13use super::{apply::InternalError, context::Context, query_arguments::QueryArguments};
14
15#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum RawKey {
17    Identifier(String),
18    // TODO: think of a better variant name
19    String(String),
20}
21
22// IMPORTANT: This regex must exactly match the identifier regex in the lexer. Also,
23// we have to add the '^' and '$' to make sure the whole string is matched.
24static IDENTIFIER_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"^[a-zA-Z_][\w-]*$").unwrap());
25
26impl From<&str> for RawKey {
27    fn from(value: &str) -> Self {
28        if IDENTIFIER_REGEX.is_match(value) {
29            RawKey::Identifier(value.to_string())
30        } else {
31            RawKey::String(value.to_string())
32        }
33    }
34}
35
36impl RawKey {
37    // TODO: `as_str` is not consistent with the `to_string` method of this struct.
38    // is this ok? `as_str` would return the inner string while `to_string` will return
39    // the escaped and quoted string.
40    pub fn as_str(&self) -> &str {
41        match self {
42            RawKey::Identifier(identifier) => identifier,
43            RawKey::String(escaped) => escaped,
44        }
45    }
46}
47
48impl Display for RawKey {
49    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
50        match self {
51            RawKey::Identifier(identifier) => identifier.fmt(f),
52            RawKey::String(unescaped_string) => {
53                let escaped_string = escape8259::escape(unescaped_string);
54                write!(f, "\"{escaped_string}\"")
55            }
56        }
57    }
58}
59
60#[derive(Debug, Clone, Constructor, Getters)]
61pub struct AtomicQueryKey {
62    // TODO: rename those attributes?
63    key: RawKey,
64    arguments: QueryArguments,
65}
66
67impl Display for AtomicQueryKey {
68    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
69        let key = self.key();
70        if self.arguments().0.is_empty() {
71            return key.fmt(f);
72        }
73
74        let arguments = self.arguments().to_string();
75        write!(f, "{key}({arguments})")
76    }
77}
78
79#[derive(Debug, Clone, Constructor, Getters, Default)]
80pub struct QueryKey {
81    pub keys: Vec<AtomicQueryKey>,
82}
83
84impl Display for QueryKey {
85    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
86        let keys = self
87            .keys()
88            .iter()
89            .map(ToString::to_string)
90            .collect::<Vec<_>>()
91            .join(".");
92        keys.fmt(f)
93    }
94}
95
96impl Add<QueryKey> for QueryKey {
97    type Output = QueryKey;
98
99    fn add(mut self, mut rhs: QueryKey) -> Self::Output {
100        self.keys.append(&mut rhs.keys);
101        self
102    }
103}
104impl<'a> QueryKey {
105    pub fn last_key(&self) -> &AtomicQueryKey {
106        self.keys().last().expect("query key cannot be empty")
107    }
108
109    // TODO: inspect should return the indexed context? in a lot of places we index and then create
110    // the indexed context
111    // TODO: maybe we should move the InternalError::KeyNotFound to this module? so we are not using something
112    // of the apply module here.
113    pub fn inspect<'b>(
114        &'a self,
115        value: &'b Value,
116        context: &Context<'a>,
117    ) -> Result<Cow<'b, Value>, InternalError<'a>> {
118        Self::do_inspect(
119            Cow::Borrowed(value),
120            self.keys(),
121            &QueryArguments::default(),
122            context,
123        )
124    }
125    pub fn inspect_owned(
126        &'a self,
127        value: Value,
128        context: &Context<'a>,
129    ) -> Result<Value, InternalError<'a>> {
130        self.inspect_owned_with_arguments(value, &QueryArguments::default(), context)
131    }
132
133    pub fn inspect_owned_with_arguments(
134        &'a self,
135        value: Value,
136        arguments: &QueryArguments,
137        context: &Context<'a>,
138    ) -> Result<Value, InternalError<'a>> {
139        Self::do_inspect(Cow::Owned(value), self.keys(), arguments, context).map(Cow::into_owned)
140    }
141
142    // TODO:
143    // Cow is used to handle both borrowed input data and owned input data, so there is no cloning
144    // when we are given an owned value (for example, inspecting in the root query).
145    // We still have the issue then the consumer wants to pass a reference as an input and only needs
146    // a reference to the inspected value, not an owned Value (argument filtering).
147    pub fn do_inspect<'b>(
148        value: Cow<'b, Value>,
149        keys: &'a [AtomicQueryKey],
150        parent_arguments: &QueryArguments,
151        context: &Context<'a>,
152    ) -> Result<Cow<'b, Value>, InternalError<'a>> {
153        let result = match value {
154            Cow::Owned(Value::Object(_)) | Cow::Borrowed(Value::Object(_)) => {
155                Self::do_inspect_object(value, keys, parent_arguments, context)?
156            }
157            Cow::Owned(Value::Array(_)) | Cow::Borrowed(Value::Array(_)) => {
158                Self::do_inspect_array(value, keys, parent_arguments, context)
159            }
160            value => Self::do_inspect_primitive(value, keys, parent_arguments, context)?,
161        };
162        Ok(result)
163    }
164
165    pub fn do_inspect_object<'b>(
166        value: Cow<'b, Value>,
167        keys: &'a [AtomicQueryKey],
168        parent_arguments: &QueryArguments,
169        context: &Context<'a>,
170    ) -> Result<Cow<'b, Value>, InternalError<'a>> {
171        if !parent_arguments.0.is_empty() {
172            // TODO: throw an error here or log a warning?
173            // in my opinion we should fail
174            return Err(InternalError::NonFiltrableValue(context.path().clone()));
175        }
176
177        let Some((atomic_query_key, rest)) = keys.split_first() else {
178            return Ok(value);
179        };
180
181        let raw_key = atomic_query_key.key();
182        let arguments = atomic_query_key.arguments();
183        let new_context = context.push_raw_key(raw_key);
184
185        let current = match value {
186            Cow::Owned(Value::Object(mut object)) => object
187                .get_mut(raw_key.as_str())
188                .map(Value::take)
189                .map(Cow::Owned),
190            Cow::Borrowed(Value::Object(object)) => object
191                // TODO: implement Borrow so we can do .get(raw_key)
192                .get(raw_key.as_str())
193                .map(Cow::Borrowed),
194            _ => unreachable!("In this match branch there are only Value::Object variants"),
195        }
196        .ok_or(InternalError::KeyNotFound(new_context.path().clone()))?;
197
198        Self::do_inspect(current, rest, arguments, &new_context)
199    }
200    pub fn do_inspect_array<'b>(
201        value: Cow<'b, Value>,
202        keys: &'a [AtomicQueryKey],
203        parent_arguments: &QueryArguments,
204        context: &Context<'a>,
205    ) -> Cow<'b, Value> {
206        let array_context = context.enter_array();
207
208        // TODO: think if there is a better way to do this, but I think this is the best we can do
209        let array_iter: Box<dyn Iterator<Item = Cow<Value>>> = match value {
210            Cow::Owned(Value::Array(array)) => Box::new(array.into_iter().map(Cow::Owned)),
211            Cow::Borrowed(Value::Array(array)) => Box::new(array.iter().map(Cow::Borrowed)),
212            _ => unreachable!("In this match branch there are only Value::Array variants"),
213        };
214
215        let result = array_iter
216            .enumerate()
217            .map(|(index, item)| (array_context.push_index(index), item))
218            .filter(|(item_context, item)| parent_arguments.satisfies(item, item_context))
219            .map(|(item_context, item)| {
220                let default_query_arguments = QueryArguments::default();
221                let arguments_to_propagate = match item {
222                    // Only propagate parent_arguments if the child is an array
223                    Cow::Owned(Value::Array(_)) | Cow::Borrowed(Value::Array(_)) => {
224                        parent_arguments
225                    }
226                    _ => &default_query_arguments,
227                };
228                Self::do_inspect(item, keys, arguments_to_propagate, &item_context)
229            })
230            .flat_map(|result| {
231                result
232                    .map_err(|error| {
233                        let array_error = InternalError::InsideArray(
234                            Box::new(error),
235                            array_context.path().clone(),
236                        );
237                        log::warn!("{array_error}");
238                    })
239                    .ok()
240            })
241            // We have to own the values if we want to return a Value::Array
242            .map(Cow::into_owned)
243            .collect();
244        Cow::Owned(Value::Array(result))
245    }
246
247    pub fn do_inspect_primitive<'b>(
248        value: Cow<'b, Value>,
249        keys: &'a [AtomicQueryKey],
250        parent_arguments: &QueryArguments,
251        context: &Context<'a>,
252    ) -> Result<Cow<'b, Value>, InternalError<'a>> {
253        if !parent_arguments.0.is_empty() {
254            return Err(InternalError::NonFiltrableValue(context.path().clone()));
255        }
256        if !keys.is_empty() {
257            return Err(InternalError::NonIndexableValue(context.path().clone()));
258        }
259        Ok(value)
260    }
261}