Add QueryBuilder::extension function. #87

src/query.rs

@@ -1,6 +1,6 @@
 use crate::context::{Data, DeferList, ResolveId};
 use crate::error::ParseRequestError;
-use crate::extensions::Extension;
+use crate::extensions::{BoxExtension, Extension};
 use crate::mutation_resolver::do_mutation_resolve;
 use crate::registry::CacheControl;
 use crate::{
@@ -141,6 +141,7 @@ pub struct QueryBuilder {
     pub(crate) operation_name: Option<String>,
     pub(crate) variables: Variables,
     pub(crate) ctx_data: Option<Data>,
+    extensions: Vec<Box<dyn Fn() -> BoxExtension + Send + Sync>>,
 }
 
 impl QueryBuilder {
@@ -151,6 +152,7 @@ impl QueryBuilder {
             operation_name: None,
             variables: Default::default(),
             ctx_data: None,
+            extensions: Default::default(),
         }
     }
 
@@ -167,6 +169,16 @@ impl QueryBuilder {
         QueryBuilder { variables, ..self }
     }
 
+    /// Add an extension
+    pub fn extension<F: Fn() -> E + Send + Sync + 'static, E: Extension>(
+        mut self,
+        extension_factory: F,
+    ) -> Self {
+        self.extensions
+            .push(Box::new(move || Box::new(extension_factory())));
+        self
+    }
+
     /// Add a context data that can be accessed in the `Context`, you access it with `Context::data`.
     ///
     /// **This data is only valid for this query**
@@ -259,7 +271,8 @@ impl QueryBuilder {
         Mutation: ObjectType + Send + Sync + 'static,
         Subscription: SubscriptionType + Send + Sync + 'static,
     {
-        let (mut document, cache_control, extensions) = schema.prepare_query(&self.query_source)?;
+        let (mut document, cache_control, extensions) =
+            schema.prepare_query(&self.query_source, &self.extensions)?;
 
         // execute
         let inc_resolve_id = AtomicUsize::default();

src/schema.rs

@@ -296,15 +296,6 @@ where
         Self::build(query, mutation, subscription).finish()
     }
 
-    pub(crate) fn create_extensions(&self) -> Extensions {
-        Extensions(
-            self.extensions
-                .iter()
-                .map(|factory| factory())
-                .collect_vec(),
-        )
-    }
-
     /// Execute query without create the `QueryBuilder`.
     pub async fn execute(&self, query_source: &str) -> Result<QueryResponse> {
         QueryBuilder::new(query_source).execute(self).await
@@ -320,9 +311,17 @@ where
     pub(crate) fn prepare_query(
         &self,
         source: &str,
+        query_extensions: &[Box<dyn Fn() -> BoxExtension + Send + Sync>],
     ) -> Result<(Document, CacheControl, Extensions)> {
         // create extension instances
-        let extensions = self.create_extensions();
+        let extensions = Extensions(
+            self.0
+                .extensions
+                .iter()
+                .chain(query_extensions)
+                .map(|factory| factory())
+                .collect_vec(),
+        );
 
         extensions.parse_start(source);
         let document = extensions.log_error(parse_query(source).map_err(Into::<Error>::into))?;
@@ -366,7 +365,7 @@ where
         variables: Variables,
         ctx_data: Option<Arc<Data>>,
     ) -> Result<impl Stream<Item = Result<serde_json::Value>> + Send> {
-        let (mut document, _, extensions) = self.prepare_query(source)?;
+        let (mut document, _, extensions) = self.prepare_query(source, &Vec::new())?;
 
         if !document.retain_operation(operation_name) {
             return extensions.log_error(if let Some(name) = operation_name {