feat(event_handler): improved support for headers and cookies in v2 (aws-powertools#1455)

Co-authored-by: Heitor Lessa <[email protected]>

Author: 2 people

Diff for: .gitignore

@@ -305,5 +305,8 @@ site/
 !404.html
 !docs/overrides/*.html
 
+# CDK
+.cdk
+
 !.github/workflows/lib
 examples/**/sam/.aws-sam

Diff for: aws_lambda_powertools/__init__.py

@@ -1,13 +1,15 @@
 # -*- coding: utf-8 -*-
 
-"""Top-level package for Lambda Python Powertools."""
-
+from pathlib import Path
 
+"""Top-level package for Lambda Python Powertools."""
 from .logging import Logger  # noqa: F401
 from .metrics import Metrics, single_metric  # noqa: F401
 from .package_logger import set_package_logger_handler
 from .tracing import Tracer  # noqa: F401
 
 __author__ = """Amazon Web Services"""
 
+PACKAGE_PATH = Path(__file__).parent
+
 set_package_logger_handler()

Diff for: aws_lambda_powertools/event_handler/api_gateway.py

@@ -124,10 +124,11 @@ def __init__(
 
     def to_dict(self) -> Dict[str, str]:
         """Builds the configured Access-Control http headers"""
-        headers = {
+        headers: Dict[str, str] = {
             "Access-Control-Allow-Origin": self.allow_origin,
             "Access-Control-Allow-Headers": ",".join(sorted(self.allow_headers)),
         }
+
         if self.expose_headers:
             headers["Access-Control-Expose-Headers"] = ",".join(self.expose_headers)
         if self.max_age is not None:
@@ -145,7 +146,8 @@ def __init__(
         status_code: int,
         content_type: Optional[str],
         body: Union[str, bytes, None],
-        headers: Optional[Dict] = None,
+        headers: Optional[Dict[str, Union[str, List[str]]]] = None,
+        cookies: Optional[List[str]] = None,
     ):
         """
 
@@ -158,13 +160,16 @@ 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
-            Optionally set specific http headers. Setting "Content-Type" hear would override the `content_type` value.
+        headers: dict[str, Union[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 = headers or {}
+        self.headers: Dict[str, Union[str, List[str]]] = headers if headers else {}
+        self.cookies = cookies or []
         if content_type:
             self.headers.setdefault("Content-Type", content_type)
 
@@ -196,11 +201,12 @@ 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"
+        cache_control = cache_control if self.response.status_code == 200 else "no-cache"
+        self.response.headers["Cache-Control"] = cache_control
 
     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")
@@ -226,11 +232,12 @@ def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dic
             logger.debug("Encoding bytes response with base64")
             self.response.base64_encoded = True
             self.response.body = base64.b64encode(self.response.body).decode()
+
         return {
             "statusCode": self.response.status_code,
-            "headers": self.response.headers,
             "body": self.response.body,
             "isBase64Encoded": self.response.base64_encoded,
+            **event.header_serializer().serialize(headers=self.response.headers, cookies=self.response.cookies),
         }
 
 
@@ -596,7 +603,7 @@ 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, Union[str, List[str]]] = {}
         if self._cors:
             logger.debug("CORS is enabled, updating headers.")
             headers.update(self._cors.to_dict())

Diff for: aws_lambda_powertools/shared/headers_serializer.py

@@ -0,0 +1,111 @@
+import warnings
+from collections import defaultdict
+from typing import Any, Dict, List, Union
+
+
+class BaseHeadersSerializer:
+    """
+    Helper class to correctly serialize headers and cookies for Amazon API Gateway,
+    ALB and Lambda Function URL response payload.
+    """
+
+    def serialize(self, headers: Dict[str, Union[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, List[str]]
+            A dictionary of headers to set in the response
+        cookies: List[str]
+            A list of cookies to set in the response
+        """
+        raise NotImplementedError()
+
+
+class HttpApiHeadersSerializer(BaseHeadersSerializer):
+    def serialize(self, headers: Dict[str, Union[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
+        runtime will automatically serialize them correctly on the output.
+
+        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
+        """
+
+        # 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():
+            if isinstance(values, str):
+                combined_headers[key] = values
+            else:
+                combined_headers[key] = ", ".join(values)
+
+        return {"headers": combined_headers, "cookies": cookies}
+
+
+class MultiValueHeadersSerializer(BaseHeadersSerializer):
+    def serialize(self, headers: Dict[str, Union[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.
+        The solution covers headers with just one key or multiple keys.
+
+        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]] = defaultdict(list)
+
+        for key, values in headers.items():
+            if isinstance(values, str):
+                payload[key].append(values)
+            else:
+                for value in values:
+                    payload[key].append(value)
+
+        if cookies:
+            payload.setdefault("Set-Cookie", [])
+            for cookie in cookies:
+                payload["Set-Cookie"].append(cookie)
+
+        return {"multiValueHeaders": payload}
+
+
+class SingleValueHeadersSerializer(BaseHeadersSerializer):
+    def serialize(self, headers: Dict[str, Union[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.
+
+        https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html#respond-to-load-balancer
+        """
+        payload: Dict[str, Dict[str, str]] = {}
+        payload.setdefault("headers", {})
+
+        if cookies:
+            if len(cookies) > 1:
+                warnings.warn(
+                    "Can't encode more than one cookie in the response. Sending the last cookie only. "
+                    "Did you enable multiValueHeaders on the ALB Target Group?"
+                )
+
+            # We can only send one cookie, send the last one
+            payload["headers"]["Set-Cookie"] = cookies[-1]
+
+        for key, values in headers.items():
+            if isinstance(values, str):
+                payload["headers"][key] = values
+            else:
+                if len(values) > 1:
+                    warnings.warn(
+                        f"Can't encode more than one header value for the same key ('{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

Diff for: aws_lambda_powertools/utilities/data_classes/alb_event.py

@@ -1,5 +1,10 @@
 from typing import Dict, List, Optional
 
+from aws_lambda_powertools.shared.headers_serializer import (
+    BaseHeadersSerializer,
+    MultiValueHeadersSerializer,
+    SingleValueHeadersSerializer,
+)
 from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent, DictWrapper
 
 
@@ -30,3 +35,11 @@ def multi_value_query_string_parameters(self) -> Optional[Dict[str, List[str]]]:
     @property
     def multi_value_headers(self) -> Optional[Dict[str, List[str]]]:
         return self.get("multiValueHeaders")
+
+    def header_serializer(self) -> BaseHeadersSerializer:
+        # When using the ALB integration, the `multiValueHeaders` feature can be disabled (default) or enabled.
+        # We can determine if the feature is enabled by looking if the event has a `multiValueHeaders` key.
+        if self.multi_value_headers:
+            return MultiValueHeadersSerializer()
+
+        return SingleValueHeadersSerializer()

Diff for: aws_lambda_powertools/utilities/data_classes/api_gateway_proxy_event.py

@@ -1,5 +1,10 @@
 from typing import Any, Dict, List, Optional
 
+from aws_lambda_powertools.shared.headers_serializer import (
+    BaseHeadersSerializer,
+    HttpApiHeadersSerializer,
+    MultiValueHeadersSerializer,
+)
 from aws_lambda_powertools.utilities.data_classes.common import (
     BaseProxyEvent,
     BaseRequestContext,
@@ -106,6 +111,9 @@ def path_parameters(self) -> Optional[Dict[str, str]]:
     def stage_variables(self) -> Optional[Dict[str, str]]:
         return self.get("stageVariables")
 
+    def header_serializer(self) -> BaseHeadersSerializer:
+        return MultiValueHeadersSerializer()
+
 
 class RequestContextV2AuthorizerIam(DictWrapper):
     @property
@@ -250,3 +258,6 @@ def path(self) -> str:
     def http_method(self) -> str:
         """The HTTP method used. Valid values include: DELETE, GET, HEAD, OPTIONS, PATCH, POST, and PUT."""
         return self.request_context.http.method
+
+    def header_serializer(self):
+        return HttpApiHeadersSerializer()

Diff for: aws_lambda_powertools/utilities/data_classes/common.py

@@ -2,6 +2,8 @@
 import json
 from typing import Any, Dict, Optional
 
+from aws_lambda_powertools.shared.headers_serializer import BaseHeadersSerializer
+
 
 class DictWrapper:
     """Provides a single read only access to a wrapper dict"""
@@ -127,6 +129,9 @@ def get_header_value(
         """
         return get_header_value(self.headers, name, default_value, case_sensitive)
 
+    def header_serializer(self) -> BaseHeadersSerializer:
+        raise NotImplementedError()
+
 
 class RequestContextClientCert(DictWrapper):
     @property

Diff for: docs/core/event_handler/api_gateway.md

@@ -312,11 +312,18 @@ For convenience, these are the default values when using `CORSConfig` to enable
 
 ### Fine grained responses
 
-You can use the `Response` class to have full control over the response, for example you might want to add additional headers or set a custom Content-type.
+You can use the `Response` class to have full control over the response. For example, you might want to add additional headers, cookies, or set a custom Content-type.
+
+???+ info
+    Powertools serializes headers and cookies according to the type of input event.
+    Some event sources require headers and cookies to be encoded as `multiValueHeaders`.
+
+???+ warning "Using multiple values for HTTP headers in ALB?"
+    Make sure you [enable the multi value headers feature](https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html#multi-value-headers) to serialize response headers correctly.
 
 === "fine_grained_responses.py"
 
-    ```python hl_lines="7 24-28"
+    ```python hl_lines="7 24-29"
     --8<-- "examples/event_handler_rest/src/fine_grained_responses.py"
     ```
 

Diff for: docs/upgrade.md

@@ -0,0 +1,57 @@
+---
+title: Upgrade guide
+description: Guide to update between major Powertools versions
+---
+
+<!-- markdownlint-disable MD043 -->
+
+## Migrate to v2 from v1
+
+The transition from Powertools for Python v1 to v2 is as painless as possible, as we aimed for minimal breaking changes.
+Changes at a glance:
+
+* The API for **event handler's `Response`** has minor changes to support multi value headers and cookies.
+
+???+ important
+    Powertools for Python v2 drops suport for Python 3.6, following the Python 3.6 End-Of-Life (EOL) reached on December 23, 2021.
+
+### Initial Steps
+
+Before you start, we suggest making a copy of your current working project or create a new branch with git.
+
+1. **Upgrade** Python to at least v3.7
+
+2. **Ensure** you have the latest `aws-lambda-powertools`
+
+    ```bash
+    pip install aws-lambda-powertools -U
+    ```
+
+3. **Review** the following sections to confirm whether they affect your code
+
+## Event Handler Response (headers and cookies)
+
+The `Response` class of the event handler utility changed slightly:
+
+1. The `headers` parameter now expects either a value or list of values per header (type `Union[str, Dict[str, List[str]]]`)
+2. We introduced a new `cookies` parameter (type `List[str]`)
+
+???+ note
+    Code that set headers as `Dict[str, str]` will still work unchanged.
+
+```python hl_lines="6 12 13"
+@app.get("/todos")
+def get_todos():
+    # Before
+    return Response(
+        # ...
+        headers={"Content-Type": "text/plain"}
+    )
+
+    # After
+    return Response(
+        # ...
+        headers={"Content-Type": ["text/plain"]},
+        cookies=["CookieName=CookieValue"]
+    )
+```

Diff for: examples/event_handler_rest/src/binary_responses_output.json

@@ -1,7 +1,7 @@
 {
     "body": "PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz4KPHN2ZyB3aWR0aD0iMjU2cHgiIGhlaWdodD0iMjU2cHgiIHZpZXdCb3g9IjAgMCAyNTYgMjU2IiB2ZXJzaW9uPSIxLjEiIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyIgeG1sbnM6eGxpbms9Imh0dHA6Ly93d3cudzMub3JnLzE5OTkveGxpbmsiIHByZXNlcnZlQXNwZWN0UmF0aW89InhNaWRZTWlkIj4KICAgIDx0aXRsZT5BV1MgTGFtYmRhPC90aXRsZT4KICAgIDxkZWZzPgogICAgICAgIDxsaW5lYXJHcmFkaWVudCB4MT0iMCUiIHkxPSIxMDAlIiB4Mj0iMTAwJSIgeTI9IjAlIiBpZD0ibGluZWFyR3JhZGllbnQtMSI+CiAgICAgICAgICAgIDxzdG9wIHN0b3AtY29sb3I9IiNDODUxMUIiIG9mZnNldD0iMCUiPjwvc3RvcD4KICAgICAgICAgICAgPHN0b3Agc3RvcC1jb2xvcj0iI0ZGOTkwMCIgb2Zmc2V0PSIxMDAlIj48L3N0b3A+CiAgICAgICAgPC9saW5lYXJHcmFkaWVudD4KICAgIDwvZGVmcz4KICAgIDxnPgogICAgICAgIDxyZWN0IGZpbGw9InVybCgjbGluZWFyR3JhZGllbnQtMSkiIHg9IjAiIHk9IjAiIHdpZHRoPSIyNTYiIGhlaWdodD0iMjU2Ij48L3JlY3Q+CiAgICAgICAgPHBhdGggZD0iTTg5LjYyNDExMjYsMjExLjIgTDQ5Ljg5MDMyNzcsMjExLjIgTDkzLjgzNTQ4MzIsMTE5LjM0NzIgTDExMy43NDcyOCwxNjAuMzM5MiBMODkuNjI0MTEyNiwyMTEuMiBaIE05Ni43MDI5MzU3LDExMC41Njk2IEM5Ni4xNjQwODU4LDEwOS40NjU2IDk1LjA0MTQ4MTMsMTA4Ljc2NDggOTMuODE2MjM4NCwxMDguNzY0OCBMOTMuODA2NjE2MywxMDguNzY0OCBDOTIuNTcxNzUxNCwxMDguNzY4IDkxLjQ0OTE0NjYsMTA5LjQ3NTIgOTAuOTE5OTE4NywxMTAuNTg1NiBMNDEuOTEzNDIwOCwyMTMuMDIwOCBDNDEuNDM4NzE5NywyMTQuMDEyOCA0MS41MDYwNzU4LDIxNS4xNzc2IDQyLjA5NjI0NTEsMjE2LjEwODggQzQyLjY3OTk5OTQsMjE3LjAzNjggNDMuNzA2MzgwNSwyMTcuNiA0NC44MDY1MzMxLDIxNy42IEw5MS42NTQ0MjMsMjE3LjYgQzkyLjg5NTcwMjcsMjE3LjYgOTQuMDIxNTE0OSwyMTYuODg2NCA5NC41NTM5NTAxLDIxNS43Njk2IEwxMjAuMjAzODU5LDE2MS42ODk2IEMxMjAuNjE3NjE5LDE2MC44MTI4IDEyMC42MTQ0MTIsMTU5Ljc5ODQgMTIwLjE4NzgyMiwxNTguOTI4IEw5Ni43MDI5MzU3LDExMC41Njk2IFogTTIwNy45ODUxMTcsMjExLjIgTDE2OC41MDc5MjgsMjExLjIgTDEwNS4xNzM3ODksNzguNjI0IEMxMDQuNjQ0NTYxLDc3LjUxMDQgMTAzLjUxNTU0MSw3Ni44IDEwMi4yNzc0NjksNzYuOCBMNzYuNDQ3OTQzLDc2LjggTDc2LjQ3NjgwOTksNDQuOCBMMTI3LjEwMzA2Niw0NC44IEwxOTAuMTQ1MzI4LDE3Ny4zNzI4IEMxOTAuNjc0NTU2LDE3OC40ODY0IDE5MS44MDM1NzUsMTc5LjIgMTkzLjA0MTY0NywxNzkuMiBMMjA3Ljk4NTExNywxNzkuMiBMMjA3Ljk4NTExNywyMTEuMiBaIE0yMTEuMTkyNTU4LDE3Mi44IEwxOTUuMDcxOTU4LDE3Mi44IEwxMzIuMDI5Njk2LDQwLjIyNzIgQzEzMS41MDA0NjgsMzkuMTEzNiAxMzAuMzcxNDQ5LDM4LjQgMTI5LjEzMDE2OSwzOC40IEw3My4yNzI1NzYsMzguNCBDNzEuNTA1Mjc1OCwzOC40IDcwLjA2ODM0MjEsMzkuODMwNCA3MC4wNjUxMzQ0LDQxLjU5NjggTDcwLjAyOTg1MjgsNzkuOTk2OCBDNzAuMDI5ODUyOCw4MC44NDggNzAuMzYzNDI2Niw4MS42NjA4IDcwLjk2OTYzMyw4Mi4yNjI0IEM3MS41Njk0MjQ2LDgyLjg2NCA3Mi4zODQxMTQ2LDgzLjIgNzMuMjM3Mjk0MSw4My4yIEwxMDAuMjUzNTczLDgzLjIgTDE2My41OTA5MiwyMTUuNzc2IEMxNjQuMTIzMzU1LDIxNi44ODk2IDE2NS4yNDU5NiwyMTcuNiAxNjYuNDg0MDMyLDIxNy42IEwyMTEuMTkyNTU4LDIxNy42IEMyMTIuOTY2Mjc0LDIxNy42IDIxNC40LDIxNi4xNjY0IDIxNC40LDIxNC40IEwyMTQuNCwxNzYgQzIxNC40LDE3NC4yMzM2IDIxMi45NjYyNzQsMTcyLjggMjExLjE5MjU1OCwxNzIuOCBMMjExLjE5MjU1OCwxNzIuOCBaIiBmaWxsPSIjRkZGRkZGIj48L3BhdGg+CiAgICA8L2c+Cjwvc3ZnPg==",
-    "headers": {
-        "Content-Type": "image/svg+xml"
+    "multiValueHeaders": {
+        "Content-Type": ["image/svg+xml"]
     },
     "isBase64Encoded": true,
     "statusCode": 200

Diff for: examples/event_handler_rest/src/compressing_responses_output.json

@@ -1,8 +1,8 @@
 {
     "statusCode": 200,
-    "headers": {
-        "Content-Type": "application/json",
-        "Content-Encoding": "gzip"
+    "multiValueHeaders": {
+        "Content-Type": ["application/json"],
+        "Content-Encoding": ["gzip"]
     },
     "body": "H4sIAAAAAAACE42STU4DMQyFrxJl3QXln96AMyAW7sSDLCVxiJ0Kqerd8TCCUOgii1EmP/783pOPXjmw+N3L0TfB+hz8brvxtC5KGtHvfMCIkzZx0HT5MPmNnziViIr2dIYoeNr8Q1x3xHsjcVadIbkZJoq2RXU8zzQROLseQ9505NzeCNQdMJNBE+UmY4zbzjAJhWtlZ57sB84BWtul+rteH2HPlVgWARwjqXkxpklK5gmEHAQqJBMtFsGVygcKmNVRjG0wxvuzGF2L0dpVUOKMC3bfJNjJgWMrCuZk7cUp02AiD72D6WKHHwUDKbiJs6AZ0VZXKOUx4uNvzdxT+E4mLcMA+6G8nzrLQkaxkNEVrFKW2VGbJCoCY7q2V3+tiv5kGThyxfTecDWbgGz/NfYXhL6ePgF9PnFdPgMAAA==",
     "isBase64Encoded": true