async_graphql/
look_ahead.rs

1use std::collections::HashMap;
2
3use crate::{
4    Context, Name, Positioned, SelectionField,
5    parser::types::{Field, FragmentDefinition, Selection, SelectionSet},
6};
7
8/// A selection performed by a query.
9pub struct Lookahead<'a> {
10    fragments: &'a HashMap<Name, Positioned<FragmentDefinition>>,
11    fields: Vec<&'a Field>,
12    context: &'a Context<'a>,
13}
14
15impl<'a> Lookahead<'a> {
16    pub(crate) fn new(
17        fragments: &'a HashMap<Name, Positioned<FragmentDefinition>>,
18        field: &'a Field,
19        context: &'a Context<'a>,
20    ) -> Self {
21        Self {
22            fragments,
23            fields: vec![field],
24            context,
25        }
26    }
27
28    /// Get the field of the selection set with the specified name. This will
29    /// ignore aliases.
30    ///
31    /// For example, calling `.field("a")` on `{ a { b } }` will return a
32    /// lookahead that represents `{ b }`.
33    #[must_use]
34    pub fn field(&self, name: &str) -> Self {
35        let mut fields = Vec::new();
36        for field in &self.fields {
37            filter(&mut fields, self.fragments, &field.selection_set.node, name)
38        }
39
40        Self {
41            fragments: self.fragments,
42            fields,
43            context: self.context,
44        }
45    }
46
47    /// Returns true if field exists otherwise return false.
48    #[inline]
49    pub fn exists(&self) -> bool {
50        !self.fields.is_empty()
51    }
52
53    /// Get the `SelectionField`s for each of the fields covered by this
54    /// `Lookahead`.
55    ///
56    /// There will be multiple fields in situations where the same field is
57    /// queried twice.
58    pub fn selection_fields(&self) -> Vec<SelectionField<'a>> {
59        self.fields
60            .iter()
61            .map(|field| SelectionField {
62                fragments: self.fragments,
63                field,
64                context: self.context,
65            })
66            .collect()
67    }
68}
69
70impl<'a> From<SelectionField<'a>> for Lookahead<'a> {
71    fn from(selection_field: SelectionField<'a>) -> Self {
72        Lookahead {
73            fragments: selection_field.fragments,
74            fields: vec![selection_field.field],
75            context: selection_field.context,
76        }
77    }
78}
79
80/// Convert a slice of `SelectionField`s to a `Lookahead`.
81/// Assumes all `SelectionField`s are from the same query and thus have the same
82/// fragments.
83///
84/// Fails if either no `SelectionField`s were provided.
85impl<'a> TryFrom<&[SelectionField<'a>]> for Lookahead<'a> {
86    type Error = ();
87
88    fn try_from(selection_fields: &[SelectionField<'a>]) -> Result<Self, Self::Error> {
89        if selection_fields.is_empty() {
90            Err(())
91        } else {
92            Ok(Lookahead {
93                fragments: selection_fields[0].fragments,
94                fields: selection_fields
95                    .iter()
96                    .map(|selection_field| selection_field.field)
97                    .collect(),
98                context: selection_fields[0].context,
99            })
100        }
101    }
102}
103
104fn filter<'a>(
105    fields: &mut Vec<&'a Field>,
106    fragments: &'a HashMap<Name, Positioned<FragmentDefinition>>,
107    selection_set: &'a SelectionSet,
108    name: &str,
109) {
110    for item in &selection_set.items {
111        match &item.node {
112            Selection::Field(field) => {
113                if field.node.name.node == name {
114                    fields.push(&field.node)
115                }
116            }
117            Selection::InlineFragment(fragment) => {
118                filter(fields, fragments, &fragment.node.selection_set.node, name)
119            }
120            Selection::FragmentSpread(spread) => {
121                if let Some(fragment) = fragments.get(&spread.node.fragment_name.node) {
122                    filter(fields, fragments, &fragment.node.selection_set.node, name)
123                }
124            }
125        }
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use crate::*;
132
133    #[tokio::test]
134    async fn test_look_ahead() {
135        #[derive(SimpleObject)]
136        #[graphql(internal)]
137        struct Detail {
138            c: i32,
139            d: i32,
140        }
141
142        #[derive(SimpleObject)]
143        #[graphql(internal)]
144        struct MyObj {
145            a: i32,
146            b: i32,
147            detail: Detail,
148        }
149
150        struct Query;
151
152        #[Object(internal)]
153        impl Query {
154            async fn obj(&self, ctx: &Context<'_>, n: i32) -> MyObj {
155                if ctx.look_ahead().field("a").exists() {
156                    // This is a query like `obj { a }`
157                    assert_eq!(n, 1);
158                } else if ctx.look_ahead().field("detail").field("c").exists()
159                    && ctx.look_ahead().field("detail").field("d").exists()
160                {
161                    // This is a query like `obj { detail { c } }`
162                    assert_eq!(n, 2);
163                } else if ctx.look_ahead().field("detail").field("c").exists() {
164                    // This is a query like `obj { detail { c } }`
165                    assert_eq!(n, 3);
166                } else {
167                    // This query doesn't have `a`
168                    assert_eq!(n, 4);
169                }
170                MyObj {
171                    a: 0,
172                    b: 0,
173                    detail: Detail { c: 0, d: 0 },
174                }
175            }
176        }
177
178        let schema = Schema::new(Query, EmptyMutation, EmptySubscription);
179
180        assert!(
181            schema
182                .execute(
183                    r#"{
184            obj(n: 1) {
185                a
186            }
187        }"#,
188                )
189                .await
190                .is_ok()
191        );
192
193        assert!(
194            schema
195                .execute(
196                    r#"{
197            obj(n: 1) {
198                k:a
199            }
200        }"#,
201                )
202                .await
203                .is_ok()
204        );
205
206        assert!(
207            schema
208                .execute(
209                    r#"{
210            obj(n: 3) {
211                detail {
212                    c
213                }
214            }
215        }"#,
216                )
217                .await
218                .is_ok()
219        );
220
221        assert!(
222            schema
223                .execute(
224                    r#"{
225            obj(n: 2) {
226                detail {
227                    d
228                }
229
230                detail {
231                    c
232                }
233            }
234        }"#,
235                )
236                .await
237                .is_ok()
238        );
239
240        assert!(
241            schema
242                .execute(
243                    r#"{
244            obj(n: 4) {
245                b
246            }
247        }"#,
248                )
249                .await
250                .is_ok()
251        );
252
253        assert!(
254            schema
255                .execute(
256                    r#"{
257            obj(n: 1) {
258                ... {
259                    a
260                }
261            }
262        }"#,
263                )
264                .await
265                .is_ok()
266        );
267
268        assert!(
269            schema
270                .execute(
271                    r#"{
272            obj(n: 3) {
273                ... {
274                    detail {
275                        c
276                    }
277                }
278            }
279        }"#,
280                )
281                .await
282                .is_ok()
283        );
284
285        assert!(
286            schema
287                .execute(
288                    r#"{
289            obj(n: 2) {
290                ... {
291                    detail {
292                        d
293                    }
294
295                    detail {
296                        c
297                    }
298                }
299            }
300        }"#,
301                )
302                .await
303                .is_ok()
304        );
305
306        assert!(
307            schema
308                .execute(
309                    r#"{
310            obj(n: 1) {
311                ... A
312            }
313        }
314        
315        fragment A on MyObj {
316            a
317        }"#,
318                )
319                .await
320                .is_ok()
321        );
322
323        assert!(
324            schema
325                .execute(
326                    r#"{
327            obj(n: 3) {
328                ... A
329            }
330        }
331        
332        fragment A on MyObj {
333            detail {
334                c
335            }
336        }"#,
337                )
338                .await
339                .is_ok()
340        );
341
342        assert!(
343            schema
344                .execute(
345                    r#"{
346            obj(n: 2) {
347                ... A
348                ... B
349            }
350        }
351        
352        fragment A on MyObj {
353            detail {
354                d
355            }
356        }
357        
358        fragment B on MyObj {
359            detail {
360                c
361            }
362        }"#,
363                )
364                .await
365                .is_ok()
366        );
367    }
368}