feat(event_handler): add support for multiple headers with same key

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -7,6 +7,7 @@
 import warnings
 import zlib
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from enum import Enum
 from functools import partial
 from http import HTTPStatus
@@ -122,18 +123,18 @@ def __init__(
         self.max_age = max_age
         self.allow_credentials = allow_credentials
 
-    def to_dict(self) -> Dict[str, str]:
+    def to_dict(self) -> Dict[str, List[str]]:
         """Builds the configured Access-Control http headers"""
-        headers = {
-            "Access-Control-Allow-Origin": self.allow_origin,
-            "Access-Control-Allow-Headers": ",".join(sorted(self.allow_headers)),
-        }
+        headers: Dict[str, List[str]] = defaultdict(list)
+        headers["Access-Control-Allow-Origin"].append(self.allow_origin)
+        headers["Access-Control-Allow-Headers"].append(",".join(sorted(self.allow_headers)))
+
         if self.expose_headers:
-            headers["Access-Control-Expose-Headers"] = ",".join(self.expose_headers)
+            headers["Access-Control-Expose-Headers"].append(",".join(self.expose_headers))
         if self.max_age is not None:
-            headers["Access-Control-Max-Age"] = str(self.max_age)
+            headers["Access-Control-Max-Age"].append(str(self.max_age))
         if self.allow_credentials is True:
-            headers["Access-Control-Allow-Credentials"] = "true"
+            headers["Access-Control-Allow-Credentials"].append("true")
         return headers
 
 
@@ -145,7 +146,7 @@ def __init__(
         status_code: int,
         content_type: Optional[str],
         body: Union[str, bytes, None],
-        headers: Optional[Dict[str, str]] = None,
+        headers: Optional[Dict[str, List[str]]] = None,
         cookies: Optional[List[str]] = None,
     ):
         """
@@ -159,18 +160,18 @@ def __init__(
             provided http headers
         body: Union[str, bytes, None]
             Optionally set the response body. Note: bytes body will be automatically base64 encoded
-        headers: dict[str, str]
+        headers: dict[str, List[str]]
             Optionally set specific http headers. Setting "Content-Type" here would override the `content_type` value.
         cookies: list[str]
             Optionally set cookies.
         """
         self.status_code = status_code
         self.body = body
         self.base64_encoded = False
-        self.headers: Dict[str, str] = headers or {}
+        self.headers: Dict[str, List[str]] = defaultdict(list, **headers) if headers else defaultdict(list)
         self.cookies = cookies or []
         if content_type:
-            self.headers.setdefault("Content-Type", content_type)
+            self.headers.setdefault("Content-Type", [content_type])
 
 
 class Route:
@@ -200,11 +201,11 @@ def _add_cors(self, cors: CORSConfig):
 
     def _add_cache_control(self, cache_control: str):
         """Set the specified cache control headers for 200 http responses. For non-200 `no-cache` is used."""
-        self.response.headers["Cache-Control"] = cache_control if self.response.status_code == 200 else "no-cache"
+        self.response.headers["Cache-Control"].append(cache_control if self.response.status_code == 200 else "no-cache")
 
     def _compress(self):
         """Compress the response body, but only if `Accept-Encoding` headers includes gzip."""
-        self.response.headers["Content-Encoding"] = "gzip"
+        self.response.headers["Content-Encoding"].append("gzip")
         if isinstance(self.response.body, str):
             logger.debug("Converting string response to bytes before compressing it")
             self.response.body = bytes(self.response.body, "utf-8")
@@ -602,14 +603,14 @@ def _path_starts_with(path: str, prefix: str):
 
     def _not_found(self, method: str) -> ResponseBuilder:
         """Called when no matching route was found and includes support for the cors preflight response"""
-        headers = {}
+        headers: Dict[str, List[str]] = defaultdict(list)
         if self._cors:
             logger.debug("CORS is enabled, updating headers.")
             headers.update(self._cors.to_dict())
 
             if method == "OPTIONS":
                 logger.debug("Pre-flight request detected. Returning CORS with null response")
-                headers["Access-Control-Allow-Methods"] = ",".join(sorted(self._cors_methods))
+                headers["Access-Control-Allow-Methods"].append(",".join(sorted(self._cors_methods)))
                 return ResponseBuilder(Response(status_code=204, content_type=None, headers=headers, body=""))
 
         handler = self._lookup_exception_handler(NotFoundError)

aws_lambda_powertools/shared/headers_serializer.py

@@ -1,5 +1,6 @@
 import warnings
 from abc import ABC
+from collections import defaultdict
 from typing import Any, Dict, List
 
 
@@ -8,14 +9,14 @@ class BaseHeadersSerializer(ABC):
     Helper class to correctly serialize headers and cookies on the response payload.
     """
 
-    def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, List[str]], cookies: List[str]) -> Dict[str, Any]:
         """
         Serializes headers and cookies according to the request type.
         Returns a dict that can be merged with the response payload.
 
         Parameters
         ----------
-        headers: Dict[str, str]
+        headers: Dict[str, List[str]]
             A dictionary of headers to set in the response
         cookies: List[str]
             A list of cookies to set in the response
@@ -24,7 +25,7 @@ def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, An
 
 
 class HttpApiSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, List[str]], cookies: List[str]) -> Dict[str, Any]:
         """
         When using HTTP APIs or LambdaFunctionURLs, everything is taken care automatically for us.
         We can directly assign a list of cookies and a dict of headers to the response payload, and the
@@ -33,11 +34,18 @@ def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, An
         https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.proxy-format
         https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.response
         """
-        return {"headers": headers, "cookies": cookies}
+
+        # Format 2.0 doesn't have multiValueHeaders or multiValueQueryStringParameters fields.
+        # Duplicate headers are combined with commas and included in the headers field.
+        combined_headers: Dict[str, str] = {}
+        for key, values in headers.items():
+            combined_headers[key] = ",".join(values)
+
+        return {"headers": combined_headers, "cookies": cookies}
 
 
 class MultiValueHeadersSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, List[str]], cookies: List[str]) -> Dict[str, Any]:
         """
         When using REST APIs, headers can be encoded using the `multiValueHeaders` key on the response.
         This is also the case when using an ALB integration with the `multiValueHeaders` option enabled.
@@ -46,10 +54,11 @@ def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, An
         https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-output-format
         https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html#multi-value-headers-response
         """
-        payload: Dict[str, List[str]] = {}
+        payload: Dict[str, List[str]] = defaultdict(list)
 
-        for key, value in headers.items():
-            payload[key] = [value]
+        for key, values in headers.items():
+            for value in values:
+                payload[key].append(value)
 
         if cookies:
             payload.setdefault("Set-Cookie", [])
@@ -60,7 +69,7 @@ def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, An
 
 
 class SingleValueHeadersSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, List[str]], cookies: List[str]) -> Dict[str, Any]:
         """
         The ALB integration has `multiValueHeaders` disabled by default.
         If we try to set multiple headers with the same key, or more than one cookie, print a warning.
@@ -80,7 +89,14 @@ def serialize(self, headers: Dict[str, str], cookies: List[str]) -> Dict[str, An
             # We can only send one cookie, send the last one
             payload["headers"]["Set-Cookie"] = cookies[-1]
 
-        for key, value in headers.items():
-            payload["headers"][key] = value
+        for key, values in headers.items():
+            if len(values) > 1:
+                warnings.warn(
+                    "Can't encode more than one header value for the same key in the response. "
+                    "Did you enable multiValueHeaders on the ALB Target Group?"
+                )
+
+            # We can only set one header per key, send the last one
+            payload["headers"][key] = values[-1]
 
         return payload

examples/event_handler_rest/src/fine_grained_responses.py

@@ -19,7 +19,7 @@ def get_todos():
     todos: requests.Response = requests.get("https://jsonplaceholder.typicode.com/todos")
     todos.raise_for_status()
 
-    custom_headers = {"X-Transaction-Id": f"{uuid4()}"}
+    custom_headers = {"X-Transaction-Id": [f"{uuid4()}"]}
 
     return Response(
         status_code=HTTPStatus.OK.value,  # 200

tests/functional/event_handler/test_api_gateway.py

@@ -441,7 +441,7 @@ def rest_func() -> Response:
             status_code=404,
             content_type="used-if-not-set-in-header",
             body="Not found",
-            headers={"Content-Type": "header-content-type-wins", "custom": "value"},
+            headers={"Content-Type": ["header-content-type-wins"], "custom": ["value"]},
         )
 
     # WHEN calling the event handler
@@ -573,7 +573,7 @@ def custom_preflight():
             status_code=200,
             content_type=content_types.TEXT_HTML,
             body="Foo",
-            headers={"Access-Control-Allow-Methods": "CUSTOM"},
+            headers={"Access-Control-Allow-Methods": ["CUSTOM"]},
         )
 
     @app.route(method="CUSTOM", rule="/some-call", cors=True)

tests/functional/event_handler/test_headers_serializer.py

@@ -13,13 +13,13 @@ def test_headers_serializer_http_api():
     payload = serializer.serialize(cookies=[], headers={})
     assert payload == {"cookies": [], "headers": {}}
 
-    payload = serializer.serialize(cookies=[], headers={"Content-Type": "text/html"})
+    payload = serializer.serialize(cookies=[], headers={"Content-Type": ["text/html"]})
     assert payload == {"cookies": [], "headers": {"Content-Type": "text/html"}}
 
     payload = serializer.serialize(cookies=["UUID=12345"], headers={})
     assert payload == {"cookies": ["UUID=12345"], "headers": {}}
 
-    payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": "bar,zbr"})
+    payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": ["bar,zbr"]})
     assert payload == {"cookies": ["UUID=12345", "SSID=0xdeadbeef"], "headers": {"Foo": "bar,zbr"}}
 
 
@@ -29,13 +29,13 @@ def test_headers_serializer_multi_value_headers():
     payload = serializer.serialize(cookies=[], headers={})
     assert payload == {"multiValueHeaders": {}}
 
-    payload = serializer.serialize(cookies=[], headers={"Content-Type": "text/html"})
+    payload = serializer.serialize(cookies=[], headers={"Content-Type": ["text/html"]})
     assert payload == {"multiValueHeaders": {"Content-Type": ["text/html"]}}
 
     payload = serializer.serialize(cookies=["UUID=12345"], headers={})
     assert payload == {"multiValueHeaders": {"Set-Cookie": ["UUID=12345"]}}
 
-    payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": "bar,zbr"})
+    payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": ["bar,zbr"]})
     assert payload == {"multiValueHeaders": {"Set-Cookie": ["UUID=12345", "SSID=0xdeadbeef"], "Foo": ["bar,zbr"]}}
 
 
@@ -45,7 +45,7 @@ def test_headers_serializer_single_value_headers():
     payload = serializer.serialize(cookies=[], headers={})
     assert payload == {"headers": {}}
 
-    payload = serializer.serialize(cookies=[], headers={"Content-Type": "text/html"})
+    payload = serializer.serialize(cookies=[], headers={"Content-Type": ["text/html"]})
     assert payload == {"headers": {"Content-Type": "text/html"}}
 
     payload = serializer.serialize(cookies=["UUID=12345"], headers={})
@@ -54,7 +54,7 @@ def test_headers_serializer_single_value_headers():
     with warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("default")
 
-        payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": "bar,zbr"})
+        payload = serializer.serialize(cookies=["UUID=12345", "SSID=0xdeadbeef"], headers={"Foo": ["bar,zbr"]})
         assert payload == {"headers": {"Set-Cookie": "SSID=0xdeadbeef", "Foo": "bar,zbr"}}
 
         assert len(w) == 1