feat: add allowed headers in configuration (#2706)

Co-authored-by: Tushar Mathur <[email protected]>

benches/from_json_bench.rs

@@ -27,6 +27,7 @@ pub fn benchmark_from_json_method(c: &mut Criterion) {
         res_body: reqs[0].clone(),
         field_name: "f1".to_string(),
         is_mutation: false,
+        headers: None,
     }];
 
     let config_generator = Generator::default().inputs(cfg_gen_reqs);

src/cli/generator/config.rs

@@ -196,6 +196,10 @@ impl Location<UnResolved> {
 }
 
 impl<A> Headers<A> {
+    pub fn into_btree_map(self) -> Option<BTreeMap<String, TemplateString>> {
+        self.0
+    }
+
     pub fn as_btree_map(&self) -> &Option<BTreeMap<String, TemplateString>> {
         &self.0
     }

src/cli/generator/generator.rs

@@ -135,6 +135,7 @@ impl Generator {
                         res_body: serde_json::from_str(&response.content)?,
                         field_name,
                         is_mutation,
+                        headers: headers.into_btree_map(),
                     });
                 }
                 Source::Proto { src } => {

src/core/generator/from_json.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 
 use convert_case::{Case, Casing};
 use serde_json::Value;
@@ -10,6 +10,7 @@ use crate::core::config::transformer::RenameTypes;
 use crate::core::config::{Config, GraphQLOperationType};
 use crate::core::http::Method;
 use crate::core::merge_right::MergeRight;
+use crate::core::mustache::TemplateString;
 use crate::core::transform::{Transform, TransformerOps};
 use crate::core::valid::{Valid, Validator};
 
@@ -20,27 +21,46 @@ pub struct RequestSample {
     pub res_body: Value,
     pub field_name: String,
     pub operation_type: GraphQLOperationType,
+    pub headers: Option<BTreeMap<String, TemplateString>>,
 }
 
 impl RequestSample {
-    #[allow(clippy::too_many_arguments)]
-    pub fn new<T: Into<String>>(
-        url: Url,
-        method: Method,
-        body: serde_json::Value,
-        resp: Value,
-        field_name: T,
-        operation_type: GraphQLOperationType,
-    ) -> Self {
+    pub fn new(url: Url, response_body: Value, field_name: String) -> Self {
         Self {
             url,
-            method,
-            req_body: body,
-            res_body: resp,
-            field_name: field_name.into(),
-            operation_type,
+            field_name,
+            res_body: response_body,
+            method: Default::default(),
+            req_body: Default::default(),
+            headers: Default::default(),
+            operation_type: Default::default(),
         }
     }
+
+    pub fn with_method(mut self, method: Method) -> Self {
+        self.method = method;
+        self
+    }
+
+    pub fn with_req_body(mut self, req_body: Value) -> Self {
+        self.req_body = req_body;
+        self
+    }
+
+    pub fn with_headers(mut self, headers: Option<BTreeMap<String, TemplateString>>) -> Self {
+        self.headers = headers;
+        self
+    }
+
+    pub fn with_is_mutation(mut self, is_mutation: bool) -> Self {
+        let operation_type = if is_mutation {
+            GraphQLOperationType::Mutation
+        } else {
+            GraphQLOperationType::Query
+        };
+        self.operation_type = operation_type;
+        self
+    }
 }
 
 pub struct FromJsonGenerator<'a> {
@@ -89,12 +109,23 @@ impl Transform for FromJsonGenerator<'_> {
                 ),
             };
 
+            // collect the required header keys
+            let header_keys = sample.headers.as_ref().map(|headers_inner| {
+                headers_inner
+                    .iter()
+                    .map(|(k, _)| k.to_owned())
+                    .collect::<BTreeSet<_>>()
+            });
+
             let mut rename_types = HashMap::new();
             rename_types.insert(existing_name, suggested_name);
 
             // these transformations are required in order to generate a base config.
             GraphQLTypesGenerator::new(sample, type_name_gen)
-                .pipe(json::SchemaGenerator::new(&sample.operation_type))
+                .pipe(json::SchemaGenerator::new(
+                    &sample.operation_type,
+                    &header_keys,
+                ))
                 .pipe(json::FieldBaseUrlGenerator::new(
                     &sample.url,
                     &sample.operation_type,
@@ -113,10 +144,8 @@ impl Transform for FromJsonGenerator<'_> {
 #[cfg(test)]
 mod tests {
     use crate::core::config::transformer::Preset;
-    use crate::core::config::GraphQLOperationType;
     use crate::core::generator::generator::test::JsonFixture;
     use crate::core::generator::{FromJsonGenerator, NameGenerator, RequestSample};
-    use crate::core::http::Method;
     use crate::core::transform::TransformerOps;
     use crate::core::valid::Validator;
 
@@ -130,17 +159,16 @@ mod tests {
             "src/core/generator/tests/fixtures/json/nested_same_properties.json",
             "src/core/generator/tests/fixtures/json/incompatible_root_object.json",
         ];
-        let field_name_generator = NameGenerator::new("f");
         for fixture in fixtures {
-            let JsonFixture { url, response } = JsonFixture::read(fixture).await?;
-            request_samples.push(RequestSample::new(
-                url.parse()?,
-                Method::GET,
-                serde_json::Value::Null,
-                response,
-                field_name_generator.next(),
-                GraphQLOperationType::Query,
-            ));
+            let JsonFixture { request, response, is_mutation, field_name } =
+                JsonFixture::read(fixture).await?;
+            let req_sample = RequestSample::new(request.url, response, field_name)
+                .with_method(request.method)
+                .with_headers(request.headers)
+                .with_is_mutation(is_mutation)
+                .with_req_body(request.body.unwrap_or_default());
+
+            request_samples.push(req_sample);
         }
 
         let config =