feat(event_handler): improved support for headers and cookies in v2 (#…

…1455)

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

.gitignore

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

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()

aws_lambda_powertools/event_handler/api_gateway.py

@@ -32,7 +32,7 @@
 _SAFE_URI = "-._~()'!*:@,;"  # https://www.ietf.org/rfc/rfc3986.txt
 # API GW/ALB decode non-safe URI chars; we must support them too
 _UNSAFE_URI = "%<> \[\]{}|^"  # noqa: W605
-_NAMED_GROUP_BOUNDARY_PATTERN = fr"(?P\1[{_SAFE_URI}{_UNSAFE_URI}\\w]+)"
+_NAMED_GROUP_BOUNDARY_PATTERN = rf"(?P\1[{_SAFE_URI}{_UNSAFE_URI}\\w]+)"
 
 
 class ProxyEventType(Enum):
@@ -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())

aws_lambda_powertools/event_handler/appsync.py

@@ -1,5 +1,4 @@
 import logging
-from abc import ABC
 from typing import Any, Callable, Optional, Type, TypeVar
 
 from aws_lambda_powertools.utilities.data_classes import AppSyncResolverEvent
@@ -10,7 +9,7 @@
 AppSyncResolverEventT = TypeVar("AppSyncResolverEventT", bound=AppSyncResolverEvent)
 
 
-class BaseRouter(ABC):
+class BaseRouter:
     current_event: AppSyncResolverEventT  # type: ignore[valid-type]
     lambda_context: LambdaContext
 

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

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()

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()

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

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"
     ```
 

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"]
+    )
+```