Repository: tailcallhq/tailcall (Rust)

Issue:
feat: add allowedHeaders setting in upstream when the api requires some headers for config generator.
## Present Behaviour:

In order  `album(p1: String!): Album @http(path: "/v1/albums/{{.args.p1}}")` this to work, we have to set the `allowedHeaders` setting with `Authorization` else it throws 401 error. 


```graphql
schema @server @upstream(baseURL: "https://api.spotify.com") {
  query: Query
}

type Query {
  album(p1: String!): Album @http(path: "/v1/albums/{{.args.p1}}")
}
```

## Expected

upstream should've the `allowedHeaders` setting, which should contain the headers which are required to resolve the http call i.e in following case it's `Authorization`.

```graphql
schema @server @upstream(baseURL: "https://api.spotify.com", allowedHeaders: ["Authorization"]) {
  query: Query
}

type Query {
  album(p1: String!): Album @http(path: "/v1/albums/{{.args.p1}}")
}
```

technical requirements:
- update the test cases if necessary.
- add integration test that cover's this case.


Discussion: **Action required:** Issue inactive for 30 days.
Status update or closure in 7 days.
Issue closed after 7 days of inactivity.



Fix (patch):
diff --git a/benches/from_json_bench.rs b/benches/from_json_bench.rs
index 4e43769d56..e953d9459d 100644
--- a/benches/from_json_bench.rs
+++ b/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);
diff --git a/src/cli/generator/config.rs b/src/cli/generator/config.rs
index a5f4859781..cc9296d5cc 100644
--- a/src/cli/generator/config.rs
+++ b/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
     }
diff --git a/src/cli/generator/generator.rs b/src/cli/generator/generator.rs
index f34ce02616..96333fbbdc 100644
--- a/src/cli/generator/generator.rs
+++ b/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 } => {
diff --git a/src/core/generator/from_json.rs b/src/core/generator/from_json.rs
index 538b7068e4..ce82cbe87d 100644
--- a/src/core/generator/from_json.rs
+++ b/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 =
diff --git a/src/core/generator/generator.rs b/src/core/generator/generator.rs
index 9420f87a6a..fee6674d91 100644
--- a/src/core/generator/generator.rs
+++ b/src/core/generator/generator.rs
@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use derive_setters::Setters;
 use prost_reflect::prost_types::FileDescriptorSet;
 use prost_reflect::DescriptorPool;
@@ -6,9 +8,10 @@ use url::Url;
 
 use super::from_proto::from_proto;
 use super::{FromJsonGenerator, NameGenerator, RequestSample};
-use crate::core::config::{self, Config, ConfigModule, GraphQLOperationType, Link, LinkType};
+use crate::core::config::{self, Config, ConfigModule, Link, LinkType};
 use crate::core::http::Method;
 use crate::core::merge_right::MergeRight;
+use crate::core::mustache::TemplateString;
 use crate::core::proto_reader::ProtoMetadata;
 use crate::core::transform::{Transform, TransformerOps};
 use crate::core::valid::Validator;
@@ -34,6 +37,7 @@ pub enum Input {
         res_body: Value,
         field_name: String,
         is_mutation: bool,
+        headers: Option<BTreeMap<String, TemplateString>>,
     },
     Proto(ProtoMetadata),
     Config {
@@ -103,29 +107,25 @@ impl Generator {
                 }
                 Input::Json {
                     url,
-                    res_body: response,
-                    field_name,
-                    is_mutation,
                     method,
                     req_body,
+                    res_body,
+                    field_name,
+                    is_mutation,
+                    headers,
                 } => {
-                    let operation_type = if *is_mutation {
-                        GraphQLOperationType::Mutation
-                    } else {
-                        GraphQLOperationType::Query
-                    };
-
-                    let request_sample = RequestSample::new(
+                    let req_sample = RequestSample::new(
                         url.to_owned(),
-                        method.to_owned(),
-                        req_body.to_owned(),
-                        response.to_owned(),
-                        field_name,
-                        operation_type.to_owned(),
-                    );
-                    config = config.merge_right(
-                        self.generate_from_json(&type_name_generator, &[request_sample])?,
-                    );
+                        res_body.to_owned(),
+                        field_name.to_owned(),
+                    )
+                    .with_method(method.to_owned())
+                    .with_headers(headers.to_owned())
+                    .with_is_mutation(is_mutation.to_owned())
+                    .with_req_body(req_body.to_owned());
+
+                    config = config
+                        .merge_right(self.generate_from_json(&type_name_generator, &[req_sample])?);
                 }
                 Input::Proto(proto_input) => {
                     config =
@@ -164,24 +164,39 @@ fn resolve_file_descriptor_set(
 
 #[cfg(test)]
 pub mod test {
+    use std::collections::BTreeMap;
+
     use prost_reflect::prost_types::FileDescriptorSet;
     use serde::{Deserialize, Deserializer};
     use serde_json::Value;
+    use url::Url;
 
     use super::Generator;
     use crate::core::config::transformer::Preset;
     use crate::core::generator::generator::Input;
-    use crate::core::generator::NameGenerator;
     use crate::core::http::Method;
+    use crate::core::mustache::TemplateString;
     use crate::core::proto_reader::ProtoMetadata;
 
     fn compile_protobuf(files: &[&str]) -> anyhow::Result<FileDescriptorSet> {
         Ok(protox::compile(files, [tailcall_fixtures::protobuf::SELF])?)
     }
 
+    #[derive(Deserialize)]
+    pub struct Request {
+        pub url: Url,
+        #[serde(default)]
+        pub method: Method,
+        #[serde(default)]
+        pub body: Option<Value>,
+        pub headers: Option<BTreeMap<String, TemplateString>>,
+    }
+
     pub struct JsonFixture {
-        pub url: String,
-        pub response: serde_json::Value,
+        pub request: Request,
+        pub response: Value,
+        pub is_mutation: bool,
+        pub field_name: String,
     }
 
     impl JsonFixture {
@@ -199,12 +214,11 @@ pub mod test {
         {
             let json_content: Value = Value::deserialize(deserializer)?;
 
-            let url = json_content
+            let req_value = json_content
                 .get("request")
-                .and_then(|req| req.get("url"))
-                .and_then(|url| url.as_str())
-                .ok_or_else(|| serde::de::Error::missing_field("request.url"))?
-                .to_string();
+                .ok_or_else(|| serde::de::Error::missing_field("request"))?;
+
+            let request = serde_json::from_value(req_value.to_owned()).unwrap();
 
             let response = json_content
                 .get("response")
@@ -212,7 +226,24 @@ pub mod test {
                 .cloned()
                 .ok_or_else(|| serde::de::Error::missing_field("response.body"))?;
 
-            Ok(JsonFixture { url, response })
+            // if is mutation isn't present, then mark it as false.
+            let is_mutation = json_content
+                .get("is_mutation")
+                .and_then(|is_mutation| is_mutation.as_bool().to_owned())
+                .unwrap_or_default();
+
+            let field_name = json_content
+                .get("fieldName")
+                .ok_or_else(|| serde::de::Error::missing_field("fieldName"))?
+                .as_str()
+                .unwrap_or_default();
+
+            Ok(JsonFixture {
+                request,
+                response,
+                is_mutation,
+                field_name: field_name.to_owned(),
+            })
         }
     }
 
@@ -247,18 +278,19 @@ pub mod test {
 
     #[tokio::test]
     async fn should_generate_config_from_json() -> anyhow::Result<()> {
-        let parsed_content = JsonFixture::read(
+        let JsonFixture { request, response, field_name, is_mutation } = JsonFixture::read(
             "src/core/generator/tests/fixtures/json/incompatible_properties.json",
         )
         .await?;
         let cfg_module = Generator::default()
             .inputs(vec![Input::Json {
-                url: parsed_content.url.parse()?,
-                method: Method::GET,
-                req_body: serde_json::Value::Null,
-                res_body: parsed_content.response,
-                field_name: "f1".to_string(),
-                is_mutation: false,
+                url: request.url,
+                method: request.method,
+                req_body: request.body.unwrap_or_default(),
+                res_body: response,
+                field_name,
+                is_mutation,
+                headers: request.headers,
             }])
             .transformers(vec![Box::new(Preset::default())])
             .generate(true)?;
@@ -283,17 +315,18 @@ pub mod test {
         };
 
         // Json Input
-        let parsed_content = JsonFixture::read(
+        let JsonFixture { request, response, field_name, is_mutation } = JsonFixture::read(
             "src/core/generator/tests/fixtures/json/incompatible_properties.json",
         )
         .await?;
         let json_input = Input::Json {
-            url: parsed_content.url.parse()?,
-            method: Method::GET,
-            req_body: serde_json::Value::Null,
-            res_body: parsed_content.response,
-            field_name: "f1".to_string(),
-            is_mutation: false,
+            url: request.url,
+            method: request.method,
+            req_body: request.body.unwrap_or_default(),
+            res_body: response,
+            field_name,
+            is_mutation,
+            headers: request.headers,
         };
 
         // Combine inputs
@@ -315,16 +348,17 @@ pub mod test {
             "src/core/generator/tests/fixtures/json/list_incompatible_object.json",
             "src/core/generator/tests/fixtures/json/list.json",
         ];
-        let field_name_generator = NameGenerator::new("f");
         for json_path in json_fixtures {
-            let parsed_content = JsonFixture::read(json_path).await?;
+            let JsonFixture { request, response, field_name, is_mutation } =
+                JsonFixture::read(json_path).await?;
             inputs.push(Input::Json {
-                url: parsed_content.url.parse()?,
-                method: Method::GET,
-                req_body: serde_json::Value::Null,
-                res_body: parsed_content.response,
-                field_name: field_name_generator.next(),
-                is_mutation: false,
+                url: request.url,
+                method: request.method,
+                req_body: request.body.unwrap_or_default(),
+                res_body: response,
+                field_name,
+                is_mutation,
+                headers: request.headers,
             });
         }
 
diff --git a/src/core/generator/json/operation_generator.rs b/src/core/generator/json/operation_generator.rs
index dad9302686..6097cd808a 100644
--- a/src/core/generator/json/operation_generator.rs
+++ b/src/core/generator/json/operation_generator.rs
@@ -69,23 +69,18 @@ mod test {
     use std::collections::BTreeMap;
 
     use super::OperationTypeGenerator;
-    use crate::core::config::{Config, Field, GraphQLOperationType, Type};
+    use crate::core::config::{Config, Field, Type};
     use crate::core::generator::{NameGenerator, RequestSample};
     use crate::core::http::Method;
     use crate::core::valid::Validator;
 
     #[test]
     fn test_query() {
-        let sample = RequestSample::new(
-            "https://jsonplaceholder.typicode.com/comments?postId=1"
-                .parse()
-                .unwrap(),
-            Method::GET,
-            serde_json::Value::Null,
-            serde_json::Value::Null,
-            "postComments",
-            GraphQLOperationType::Query,
-        );
+        let url = "https://jsonplaceholder.typicode.com/comments?postId=1"
+            .parse()
+            .unwrap();
+
+        let sample = RequestSample::new(url, Default::default(), "postComments".into());
         let config = Config::default();
         let config = OperationTypeGenerator
             .generate(&sample, "T44", &NameGenerator::new("Input"), config)
@@ -97,16 +92,11 @@ mod test {
 
     #[test]
     fn test_append_field_if_operation_type_exists() {
-        let sample = RequestSample::new(
-            "https://jsonplaceholder.typicode.com/comments?postId=1"
-                .parse()
-                .unwrap(),
-            Method::GET,
-            serde_json::Value::Null,
-            serde_json::Value::Null,
-            "postComments",
-            GraphQLOperationType::Query,
-        );
+        let url = "https://jsonplaceholder.typicode.com/comments?postId=1"
+            .parse()
+            .unwrap();
+
+        let sample = RequestSample::new(url, Default::default(), "postComments".into());
         let mut config = Config::default();
         let mut fields = BTreeMap::default();
         fields.insert(
@@ -136,16 +126,15 @@ mod test {
             }
         "#;
 
-        let sample = RequestSample::new(
-            "https://jsonplaceholder.typicode.com/posts"
-                .parse()
-                .unwrap(),
-            Method::POST,
-            serde_json::from_str(body).unwrap(),
-            serde_json::Value::Null,
-            "postComments",
-            GraphQLOperationType::Mutation,
-        );
+        let url = "https://jsonplaceholder.typicode.com/posts"
+            .parse()
+            .unwrap();
+
+        let sample = RequestSample::new(url, Default::default(), "postComments".into())
+            .with_method(Method::POST)
+            .with_req_body(serde_json::from_str(body).unwrap())
+            .with_is_mutation(true);
+
         let config = Config::default();
         let config = OperationTypeGenerator
             .generate(&sample, "T44", &NameGenerator::new("Input"), config)
diff --git a/src/core/generator/json/schema_generator.rs b/src/core/generator/json/schema_generator.rs
index c722590584..a2d5a778a3 100644
--- a/src/core/generator/json/schema_generator.rs
+++ b/src/core/generator/json/schema_generator.rs
@@ -1,3 +1,5 @@
+use std::collections::BTreeSet;
+
 use convert_case::{Case, Casing};
 
 use crate::core::config::{Config, GraphQLOperationType};
@@ -6,11 +8,15 @@ use crate::core::valid::Valid;
 
 pub struct SchemaGenerator<'a> {
     operation_type: &'a GraphQLOperationType,
+    header_keys: &'a Option<BTreeSet<String>>,
 }
 
 impl<'a> SchemaGenerator<'a> {
-    pub fn new(operation_type: &'a GraphQLOperationType) -> Self {
-        Self { operation_type }
+    pub fn new(
+        operation_type: &'a GraphQLOperationType,
+        header_keys: &'a Option<BTreeSet<String>>,
+    ) -> Self {
+        Self { operation_type, header_keys }
     }
 }
 
@@ -34,12 +40,18 @@ impl Transform for SchemaGenerator<'_> {
                 );
             }
         }
+
+        // Add allowed headers setting on upstream
+        config.upstream = config.upstream.allowed_headers(self.header_keys.to_owned());
+
         Valid::succeed(config)
     }
 }
 
 #[cfg(test)]
 mod test {
+    use std::collections::BTreeSet;
+
     use super::SchemaGenerator;
     use crate::core::config::GraphQLOperationType;
     use crate::core::transform::Transform;
@@ -47,7 +59,7 @@ mod test {
 
     #[test]
     fn test_schema_generator_with_mutation() {
-        let schema_gen = SchemaGenerator::new(&GraphQLOperationType::Mutation);
+        let schema_gen = SchemaGenerator::new(&GraphQLOperationType::Mutation, &None);
         let config = schema_gen
             .transform(Default::default())
             .to_result()
@@ -60,13 +72,28 @@ mod test {
 
     #[test]
     fn test_schema_generator_with_query() {
-        let schema_gen = SchemaGenerator::new(&GraphQLOperationType::Query);
+        let schema_gen = SchemaGenerator::new(&GraphQLOperationType::Query, &None);
+        let config = schema_gen
+            .transform(Default::default())
+            .to_result()
+            .unwrap();
+        assert!(config.schema.query.is_some());
+        assert_eq!(config.schema.query, Some("Query".to_owned()));
+
+        assert!(config.schema.mutation.is_none());
+    }
+
+    #[test]
+    fn test_schema_generator_with_headers() {
+        let expected_header_keys = Some(BTreeSet::from(["X-Custom-Header".to_owned()]));
+        let schema_gen = SchemaGenerator::new(&GraphQLOperationType::Query, &expected_header_keys);
         let config = schema_gen
             .transform(Default::default())
             .to_result()
             .unwrap();
         assert!(config.schema.query.is_some());
         assert_eq!(config.schema.query, Some("Query".to_owned()));
+        assert_eq!(config.upstream.allowed_headers, expected_header_keys);
 
         assert!(config.schema.mutation.is_none());
     }