tailcallhq · tusharmath · Mar 13, 2024 · Mar 12, 2024 · Mar 12, 2024 · Mar 12, 2024
diff --git a/generated/.tailcallrc.graphql b/generated/.tailcallrc.graphql
@@ -286,12 +286,6 @@ directive @server(
   """
   queryValidation: Boolean
   """
-  The `responseHeaders` are key-value pairs included in every server response. Useful 
-  for setting headers like `Access-Control-Allow-Origin` for cross-origin requests 
-  or additional headers for downstream services.
-  """
-  responseHeaders: [KeyValue]
-  """
   `responseValidation` Tailcall automatically validates responses from upstream services 
   using inferred schema. @default `false`.
   """
@@ -577,6 +571,12 @@ input Headers {
   value is the least of the values received from upstream services. @default `false`.
   """
   cacheControl: Boolean
+  """
+  `headers` are key-value pairs included in every server response. Useful for setting 
+  headers like `Access-Control-Allow-Origin` for cross-origin requests or additional 
+  headers for downstream services.
+  """
+  custom: [KeyValue]
 }
 """
 The @http operator indicates that a field or node is backed by a REST API.For instance, 

diff --git a/generated/.tailcallrc.schema.json b/generated/.tailcallrc.schema.json
@@ -1215,6 +1215,13 @@
             "boolean",
             "null"
           ]
+        },
+        "custom": {
+          "description": "`headers` are key-value pairs included in every server response. Useful for setting headers like `Access-Control-Allow-Origin` for cross-origin requests or additional headers for downstream services.",
+          "type": "array",
+          "items": {
+            "$ref": "#/definitions/KeyValue"
+          }
         }
       }
     },
@@ -1588,13 +1595,6 @@
             "null"
           ]
         },
-        "responseHeaders": {
-          "description": "The `responseHeaders` are key-value pairs included in every server response. Useful for setting headers like `Access-Control-Allow-Origin` for cross-origin requests or additional headers for downstream services.",
-          "type": "array",
-          "items": {
-            "$ref": "#/definitions/KeyValue"
-          }
-        },
         "responseValidation": {
           "description": "`responseValidation` Tailcall automatically validates responses from upstream services using inferred schema. @default `false`.",
           "type": [

diff --git a/src/blueprint/server.rs b/src/blueprint/server.rs
@@ -160,7 +160,7 @@ fn validate_hostname(hostname: String) -> Valid<IpAddr, String> {
     }
 }
 
-fn handle_response_headers(resp_headers: BTreeMap<String, String>) -> Valid<HeaderMap, String> {
+fn handle_response_headers(resp_headers: Vec<(String, String)>) -> Valid<HeaderMap, String> {
     Valid::from_iter(resp_headers.iter(), |(k, v)| {
         let name = Valid::from(
             HeaderName::from_bytes(k.as_bytes())
@@ -173,7 +173,8 @@ fn handle_response_headers(resp_headers: BTreeMap<String, String>) -> Valid<Head
         name.zip(value)
     })
     .map(|headers| headers.into_iter().collect::<HeaderMap>())
-    .trace("responseHeaders")
+    .trace("custom")
+    .trace("headers")
     .trace("@server")
     .trace("schema")
 }

diff --git a/src/config/headers.rs b/src/config/headers.rs
@@ -1,5 +1,6 @@
 use serde::{Deserialize, Serialize};
 
+use crate::config::KeyValue;
 use crate::is_default;
 
 #[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq, Eq, schemars::JsonSchema)]
@@ -10,10 +11,33 @@ pub struct Headers {
     /// activated. The `max-age` value is the least of the values received from
     /// upstream services. @default `false`.
     pub cache_control: Option<bool>,
+
+    #[serde(default, skip_serializing_if = "is_default")]
+    /// `headers` are key-value pairs included in every server
+    /// response. Useful for setting headers like `Access-Control-Allow-Origin`
+    /// for cross-origin requests or additional headers for downstream services.
+    pub custom: Vec<KeyValue>,
 }
 
 impl Headers {
     pub fn enable_cache_control(&self) -> bool {
         self.cache_control.unwrap_or(false)
     }
 }
+
+pub fn merge_headers(current: Option<Headers>, other: Option<Headers>) -> Option<Headers> {
+    let mut headers = current.clone();
+
+    if let Some(other_headers) = other {
+        if let Some(mut self_headers) = current.clone() {
+            self_headers.cache_control = other_headers.cache_control.or(self_headers.cache_control);
+            self_headers.custom.extend(other_headers.custom);
+
+            headers = Some(self_headers);
+        } else {
+            headers = Some(other_headers);
+        }
+    }
+
+    headers
+}
diff --git a/src/config/key_values.rs b/src/config/key_values.rs
@@ -26,6 +26,19 @@ pub struct KeyValue {
     pub value: String,
 }
 
+pub fn merge_key_value_vecs(current: &[KeyValue], other: &[KeyValue]) -> Vec<KeyValue> {
+    let mut acc: BTreeMap<&String, &String> =
+        current.iter().map(|kv| (&kv.key, &kv.value)).collect();
+
+    for kv in other {
+        acc.insert(&kv.key, &kv.value);
+    }
+
+    acc.iter()
+        .map(|(k, v)| KeyValue { key: k.to_string(), value: v.to_string() })
+        .collect()
+}
+
 impl Serialize for KeyValues {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -97,4 +110,54 @@ mod tests {
         // Using the deref trait
         assert_eq!(kv["a"], "b");
     }
+
+    #[test]
+    fn test_merge_with_both_empty() {
+        let current = vec![];
+        let other = vec![];
+        let result = merge_key_value_vecs(&current, &other);
+        assert!(result.is_empty());
+    }
+
+    #[test]
+    fn test_merge_with_current_empty() {
+        let current = vec![];
+        let other = vec![KeyValue { key: "key1".to_string(), value: "value1".to_string() }];
+        let result = merge_key_value_vecs(&current, &other);
+        assert_eq!(result.len(), 1);
+        assert_eq!(result[0].key, "key1");
+        assert_eq!(result[0].value, "value1");
+    }
+
+    #[test]
+    fn test_merge_with_other_empty() {
+        let current = vec![KeyValue { key: "key1".to_string(), value: "value1".to_string() }];
+        let other = vec![];
+        let result = merge_key_value_vecs(&current, &other);
+        assert_eq!(result.len(), 1);
+        assert_eq!(result[0].key, "key1");
+        assert_eq!(result[0].value, "value1");
+    }
+
+    #[test]
+    fn test_merge_with_unique_keys() {
+        let current = vec![KeyValue { key: "key1".to_string(), value: "value1".to_string() }];
+        let other = vec![KeyValue { key: "key2".to_string(), value: "value2".to_string() }];
+        let result = merge_key_value_vecs(&current, &other);
+        assert_eq!(result.len(), 2);
+        assert_eq!(result[0].key, "key1");
+        assert_eq!(result[0].value, "value1");
+        assert_eq!(result[1].key, "key2");
+        assert_eq!(result[1].value, "value2");
+    }
+
+    #[test]
+    fn test_merge_with_overlapping_keys() {
+        let current = vec![KeyValue { key: "key1".to_string(), value: "value1".to_string() }];
+        let other = vec![KeyValue { key: "key1".to_string(), value: "value2".to_string() }];
+        let result = merge_key_value_vecs(&current, &other);
+        assert_eq!(result.len(), 1);
+        assert_eq!(result[0].key, "key1");
+        assert_eq!(result[0].value, "value2");
+    }
 }
diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -1,6 +1,7 @@
 pub use config::*;
 pub use config_module::*;
 pub use expr::*;
+pub use headers::*;
 pub use key_values::*;
 pub use link::*;
 pub use reader_context::*;

diff --git a/src/config/server.rs b/src/config/server.rs
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
 
+use super::{merge_headers, merge_key_value_vecs};
 use crate::config::headers::Headers;
 use crate::config::KeyValue;
 use crate::is_default;
@@ -64,12 +65,6 @@ pub struct Server {
     /// @default `false`.
     pub query_validation: Option<bool>,
 
-    #[serde(skip_serializing_if = "is_default", default)]
-    /// The `responseHeaders` are key-value pairs included in every server
-    /// response. Useful for setting headers like `Access-Control-Allow-Origin`
-    /// for cross-origin requests or additional headers for downstream services.
-    pub response_headers: Vec<KeyValue>,
-
     #[serde(default, skip_serializing_if = "is_default")]
     /// `responseValidation` Tailcall automatically validates responses from
     /// upstream services using inferred schema. @default `false`.
@@ -165,12 +160,16 @@ impl Server {
             .collect()
     }
 
-    pub fn get_response_headers(&self) -> BTreeMap<String, String> {
-        self.response_headers
-            .clone()
-            .iter()
-            .map(|kv| (kv.key.clone(), kv.value.clone()))
-            .collect()
+    pub fn get_response_headers(&self) -> Vec<(String, String)> {
+        self.headers
+            .as_ref()
+            .map(|h| h.custom.clone())
+            .map_or(Vec::new(), |headers| {
+                headers
+                    .iter()
+                    .map(|kv| (kv.key.clone(), kv.value.clone()))
+                    .collect()
+            })
     }
 
     pub fn get_version(self) -> HttpVersion {
@@ -183,7 +182,7 @@ impl Server {
 
     pub fn merge_right(mut self, other: Self) -> Self {
         self.apollo_tracing = other.apollo_tracing.or(self.apollo_tracing);
-        self.headers = other.headers.or(self.headers);
+        self.headers = merge_headers(self.headers, other.headers);
         self.graphiql = other.graphiql.or(self.graphiql);
         self.introspection = other.introspection.or(self.introspection);
         self.query_validation = other.query_validation.or(self.query_validation);
@@ -196,7 +195,7 @@ impl Server {
         self.workers = other.workers.or(self.workers);
         self.port = other.port.or(self.port);
         self.hostname = other.hostname.or(self.hostname);
-        self.vars = other.vars.iter().fold(self.vars, |mut acc, kv| {
+        self.vars = other.vars.iter().fold(self.vars.to_vec(), |mut acc, kv| {
             let position = acc.iter().position(|x| x.key == kv.key);
             if let Some(pos) = position {
                 acc[pos] = kv.clone();
@@ -205,19 +204,7 @@ impl Server {
             };
             acc
         });
-        self.response_headers =
-            other
-                .response_headers
-                .iter()
-                .fold(self.response_headers, |mut acc, kv| {
-                    let position = acc.iter().position(|x| x.key == kv.key);
-                    if let Some(pos) = position {
-                        acc[pos] = kv.clone();
-                    } else {
-                        acc.push(kv.clone());
-                    };
-                    acc
-                });
+        self.vars = merge_key_value_vecs(&self.vars, &other.vars);
         self.version = other.version.or(self.version);
         self.pipeline_flush = other.pipeline_flush.or(self.pipeline_flush);
         self.script = other.script.or(self.script);

diff --git a/tests/execution/custom-headers.md b/tests/execution/custom-headers.md
@@ -3,16 +3,18 @@
 ```json @server
 {
   "server": {
-    "responseHeaders": [
-      {
-        "key": "x-id",
-        "value": "1"
-      },
-      {
-        "key": "x-name",
-        "value": "John Doe"
-      }
-    ]
+    "headers": {
+      "custom": [
+        {
+          "key": "x-id",
+          "value": "1"
+        },
+        {
+          "key": "x-name",
+          "value": "John Doe"
+        }
+      ]
+    }
   },
   "upstream": {},
   "schema": {

diff --git a/tests/execution/test-response-header-merge.md b/tests/execution/test-response-header-merge.md
@@ -0,0 +1,31 @@
+# test-response-header-value
+
+```graphql @server
+schema @server(headers: {custom: [{key: "a", value: "a"}]}) {
+  query: Query
+}
+
+type User {
+  name: String
+  age: Int
+}
+
+type Query {
+  user: User @const(data: {name: "John"})
+}
+```
+
+```graphql @server
+schema @server(headers: {custom: [{key: "a", value: "b"}]}) {
+  query: Query
+}
+
+type User {
+  name: String
+  age: Int
+}
+
+type Query {
+  user: User @const(data: {name: "John"})
+}
+```
diff --git a/tests/execution/test-response-header-value.md b/tests/execution/test-response-header-value.md
@@ -3,7 +3,7 @@
 ###### sdl error
 
 ```graphql @server
-schema @server(responseHeaders: [{key: "a", value: "a \n b"}]) {
+schema @server(headers: {custom: [{key: "a", value: "a \n b"}]}) {
   query: Query
 }
 

diff --git a/tests/execution/test-response-headers-multi.md b/tests/execution/test-response-headers-multi.md
@@ -3,7 +3,7 @@
 ###### sdl error
 
 ```graphql @server
-schema @server(responseHeaders: [{key: "a b", value: "a \n b"}, {key: "a c", value: "a \n b"}]) {
+schema @server(headers: {custom: [{key: "a b", value: "a \n b"}, {key: "a c", value: "a \n b"}]}) {
   query: Query
 }
 

diff --git a/tests/execution/test-response-headers-name.md b/tests/execution/test-response-headers-name.md
@@ -3,7 +3,7 @@
 ###### sdl error
 
 ```graphql @server
-schema @server(responseHeaders: [{key: "🤣", value: "a"}]) {
+schema @server(headers: {custom: [{key: "🤣", value: "a"}]}) {
   query: Query
 }
 

diff --git a/tests/snapshots/execution_spec__custom-headers.md_merged.snap b/tests/snapshots/execution_spec__custom-headers.md_merged.snap
@@ -2,7 +2,7 @@
 source: tests/execution_spec.rs
 expression: merged
 ---
-schema @server(responseHeaders: [{key: "x-id", value: "1"}, {key: "x-name", value: "John Doe"}]) @upstream {
+schema @server(headers: {custom: [{key: "x-id", value: "1"}, {key: "x-name", value: "John Doe"}]}) @upstream {
   query: Query
 }
 

diff --git a/tests/snapshots/execution_spec__test-response-header-merge.md_merged.snap b/tests/snapshots/execution_spec__test-response-header-merge.md_merged.snap
@@ -0,0 +1,16 @@
+---
+source: tests/execution_spec.rs
+expression: merged
+---
+schema @server(headers: {custom: [{key: "a", value: "a"}, {key: "a", value: "b"}]}) @upstream {
+  query: Query
+}
+
+type Query {
+  user: User @const(data: {name: "John"})
+}
+
+type User {
+  age: Int
+  name: String
+}