async_graphql/extensions/
analyzer.rs

1use std::sync::Arc;
2
3use futures_util::lock::Mutex;
4
5use crate::{
6    Response, ServerError, ValidationResult,
7    extensions::{Extension, ExtensionContext, ExtensionFactory, NextRequest, NextValidation},
8    value,
9};
10
11/// Analyzer extension
12///
13/// This extension will output the `analyzer` field containing `complexity` and
14/// `depth` in the response extension of each query.
15pub struct Analyzer;
16
17impl ExtensionFactory for Analyzer {
18    fn create(&self) -> Arc<dyn Extension> {
19        Arc::new(AnalyzerExtension::default())
20    }
21}
22
23#[derive(Default)]
24struct AnalyzerExtension {
25    validation_result: Mutex<Option<ValidationResult>>,
26}
27
28#[async_trait::async_trait]
29impl Extension for AnalyzerExtension {
30    async fn request(&self, ctx: &ExtensionContext<'_>, next: NextRequest<'_>) -> Response {
31        let mut resp = next.run(ctx).await;
32        let validation_result = self.validation_result.lock().await.take();
33        if let Some(validation_result) = validation_result {
34            resp = resp.extension(
35                "analyzer",
36                value! ({
37                    "complexity": validation_result.complexity,
38                    "depth": validation_result.depth,
39                }),
40            );
41        }
42        resp
43    }
44
45    async fn validation(
46        &self,
47        ctx: &ExtensionContext<'_>,
48        next: NextValidation<'_>,
49    ) -> Result<ValidationResult, Vec<ServerError>> {
50        let res = next.run(ctx).await?;
51        *self.validation_result.lock().await = Some(res);
52        Ok(res)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use crate::*;
59
60    struct Query;
61
62    #[derive(Copy, Clone)]
63    struct MyObj;
64
65    #[Object(internal)]
66    impl MyObj {
67        async fn value(&self) -> i32 {
68            1
69        }
70
71        async fn obj(&self) -> MyObj {
72            MyObj
73        }
74    }
75
76    #[Object(internal)]
77    impl Query {
78        async fn value(&self) -> i32 {
79            1
80        }
81
82        async fn obj(&self) -> MyObj {
83            MyObj
84        }
85
86        #[graphql(complexity = "count * child_complexity")]
87        async fn objs(&self, count: usize) -> Vec<MyObj> {
88            vec![MyObj; count]
89        }
90    }
91
92    #[tokio::test]
93    async fn analyzer() {
94        let schema = Schema::build(Query, EmptyMutation, EmptySubscription)
95            .extension(extensions::Analyzer)
96            .finish();
97
98        let res = schema
99            .execute(
100                r#"{
101            value obj {
102                value obj {
103                    value
104                }
105            }
106            objs(count: 10) { value }
107        }"#,
108            )
109            .await
110            .into_result()
111            .unwrap()
112            .extensions
113            .remove("analyzer");
114        assert_eq!(
115            res,
116            Some(value!({
117                "complexity": 5 + 10,
118                "depth": 3,
119            }))
120        );
121    }
122}