feat(event-handler): Add Response class

This will allow for fine grained control of the returning headers

Author: michaelbrewer

aws_lambda_powertools/event_handler/api_gateway.py

@@ -29,6 +29,42 @@ def __init__(
         self.cache_control = cache_control
 
 
+class Response:
+    def __init__(self, status_code: int, content_type: str, body: Union[str, bytes], headers: Dict = None):
+        self.status_code = status_code
+        self.body = body
+        self.base64_encoded = False
+        self.headers: Dict = headers if headers is not None else {}
+        if "Content-Type" not in self.headers:
+            self.headers["Content-Type"] = content_type
+
+    def add_cors(self, method: str):
+        self.headers["Access-Control-Allow-Origin"] = "*"
+        self.headers["Access-Control-Allow-Methods"] = method
+        self.headers["Access-Control-Allow-Credentials"] = "true"
+
+    def add_cache_control(self, cache_control: str):
+        self.headers["Cache-Control"] = cache_control if self.status_code == 200 else "no-cache"
+
+    def compress(self):
+        self.headers["Content-Encoding"] = "gzip"
+        if isinstance(self.body, str):
+            self.body = bytes(self.body, "utf-8")
+        gzip = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16)
+        self.body = gzip.compress(self.body) + gzip.flush()
+
+    def to_dict(self):
+        if isinstance(self.body, bytes):
+            self.base64_encoded = True
+            self.body = base64.b64encode(self.body).decode()
+        return {
+            "statusCode": self.status_code,
+            "headers": self.headers,
+            "body": self.body,
+            "isBase64Encoded": self.base64_encoded,
+        }
+
+
 class ApiGatewayResolver:
     current_event: BaseProxyEvent
     lambda_context: LambdaContext
@@ -65,35 +101,21 @@ def resolve(self, event, context) -> Dict[str, Any]:
         route, args = self._find_route(self.current_event.http_method, self.current_event.path)
         result = route.func(**args)
 
-        if isinstance(result, dict):
-            status_code = 200
-            content_type = "application/json"
-            body: Union[str, bytes] = json.dumps(result)
+        if isinstance(result, Response):
+            response = result
+        elif isinstance(result, dict):
+            response = Response(status_code=200, content_type="application/json", body=json.dumps(result))
         else:
-            status_code, content_type, body = result
-        headers = {"Content-Type": content_type}
+            response = Response(*result)
 
         if route.cors:
-            headers["Access-Control-Allow-Origin"] = "*"
-            headers["Access-Control-Allow-Methods"] = route.method
-            headers["Access-Control-Allow-Credentials"] = "true"
-
+            response.add_cors(route.method)
         if route.cache_control:
-            headers["Cache-Control"] = route.cache_control if status_code == 200 else "no-cache"
-
+            response.add_cache_control(route.cache_control)
         if route.compress and "gzip" in (self.current_event.get_header_value("accept-encoding") or ""):
-            headers["Content-Encoding"] = "gzip"
-            if isinstance(body, str):
-                body = bytes(body, "utf-8")
-            gzip = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16)
-            body = gzip.compress(body) + gzip.flush()
-
-        base64_encoded = False
-        if isinstance(body, bytes):
-            base64_encoded = True
-            body = base64.b64encode(body).decode()
-
-        return {"statusCode": status_code, "headers": headers, "body": body, "isBase64Encoded": base64_encoded}
+            response.compress()
+
+        return response.to_dict()
 
     @staticmethod
     def _build_rule_pattern(rule: str):

tests/functional/event_handler/test_api_gateway.py

@@ -2,10 +2,11 @@
 import json
 import zlib
 from pathlib import Path
+from typing import Dict, Tuple
 
 import pytest
 
-from aws_lambda_powertools.event_handler.api_gateway import ApiGatewayResolver, ProxyEventType
+from aws_lambda_powertools.event_handler.api_gateway import ApiGatewayResolver, ProxyEventType, Response
 from aws_lambda_powertools.utilities.data_classes import ALBEvent, APIGatewayProxyEvent, APIGatewayProxyEventV2
 
 
@@ -25,77 +26,96 @@ def read_media(file_name: str) -> bytes:
 
 
 def test_alb_event():
+    # GIVEN a Application Load Balancer proxy type event
     app = ApiGatewayResolver(proxy_type=ProxyEventType.alb_event)
 
     @app.get("/lambda")
-    def foo():
+    def foo() -> Tuple[int, str, str]:
         assert isinstance(app.current_event, ALBEvent)
         assert app.lambda_context == {}
         return 200, TEXT_HTML, "foo"
 
+    # WHEN
     result = app(load_event("albEvent.json"), {})
 
+    # THEN process event correctly
+    # AND set the current_event type as ALBEvent
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == TEXT_HTML
     assert result["body"] == "foo"
 
 
 def test_api_gateway_v1():
+    # GIVEN a Http API V1 proxy type event
     app = ApiGatewayResolver(proxy_type=ProxyEventType.http_api_v1)
 
     @app.get("/my/path")
-    def get_lambda():
+    def get_lambda() -> Tuple[int, str, str]:
         assert isinstance(app.current_event, APIGatewayProxyEvent)
         assert app.lambda_context == {}
         return 200, APPLICATION_JSON, json.dumps({"foo": "value"})
 
+    # WHEN
     result = app(LOAD_GW_EVENT, {})
 
+    # THEN process event correctly
+    # AND set the current_event type as APIGatewayProxyEvent
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
 
 
 def test_api_gateway():
+    # GIVEN a Rest API Gateway proxy type event
     app = ApiGatewayResolver(proxy_type=ProxyEventType.api_gateway)
 
     @app.get("/my/path")
-    def get_lambda():
+    def get_lambda() -> Tuple[int, str, str]:
         assert isinstance(app.current_event, APIGatewayProxyEvent)
         return 200, TEXT_HTML, "foo"
 
+    # WHEN
     result = app(LOAD_GW_EVENT, {})
 
+    # THEN process event correctly
+    # AND set the current_event type as APIGatewayProxyEvent
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == TEXT_HTML
     assert result["body"] == "foo"
 
 
 def test_api_gateway_v2():
+    # GIVEN a Http API V2 proxy type event
     app = ApiGatewayResolver(proxy_type=ProxyEventType.http_api_v2)
 
     @app.post("/my/path")
-    def my_path():
+    def my_path() -> Tuple[int, str, str]:
         assert isinstance(app.current_event, APIGatewayProxyEventV2)
         post_data = app.current_event.json_body
         return 200, "plain/text", post_data["username"]
 
+    # WHEN
     result = app(load_event("apiGatewayProxyV2Event.json"), {})
 
+    # THEN process event correctly
+    # AND set the current_event type as APIGatewayProxyEventV2
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == "plain/text"
     assert result["body"] == "tom"
 
 
 def test_include_rule_matching():
+    # GIVEN
     app = ApiGatewayResolver()
 
     @app.get("/<name>/<my_id>")
-    def get_lambda(my_id: str, name: str):
+    def get_lambda(my_id: str, name: str) -> Tuple[int, str, str]:
         assert name == "my"
         return 200, "plain/html", my_id
 
+    # WHEN
     result = app(LOAD_GW_EVENT, {})
 
+    # THEN
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == "plain/html"
     assert result["body"] == "path"
@@ -153,7 +173,7 @@ def test_cors():
     app = ApiGatewayResolver()
 
     @app.get("/my/path", cors=True)
-    def with_cors():
+    def with_cors() -> Tuple[int, str, str]:
         return 200, TEXT_HTML, "test"
 
     def handler(event, context):
@@ -176,7 +196,7 @@ def test_compress():
     app = ApiGatewayResolver()
 
     @app.get("/my/request", compress=True)
-    def with_compression():
+    def with_compression() -> Tuple[int, str, str]:
         return 200, APPLICATION_JSON, expected_value
 
     def handler(event, context):
@@ -197,7 +217,7 @@ def test_base64_encode():
     app = ApiGatewayResolver()
 
     @app.get("/my/path", compress=True)
-    def read_image():
+    def read_image() -> Tuple[int, str, bytes]:
         return 200, "image/png", read_media("idempotent_sequence_exception.png")
 
     mock_event = {"path": "/my/path", "httpMethod": "GET", "headers": {"Accept-Encoding": "deflate, gzip"}}
@@ -211,62 +231,100 @@ def read_image():
 
 
 def test_compress_no_accept_encoding():
+    # GIVEN a function with compress=True
+    # AND the request has no "Accept-Encoding" set to include gzip
     app = ApiGatewayResolver()
     expected_value = "Foo"
 
     @app.get("/my/path", compress=True)
-    def return_text():
+    def return_text() -> Tuple[int, str, str]:
         return 200, "text/plain", expected_value
 
+    # WHEN calling the event handler
     result = app({"path": "/my/path", "httpMethod": "GET", "headers": {}}, None)
 
+    # THEN don't perform any gzip compression
     assert result["isBase64Encoded"] is False
     assert result["body"] == expected_value
 
 
 def test_cache_control_200():
+    # GIVEN a function with cache_control set
     app = ApiGatewayResolver()
 
     @app.get("/success", cache_control="max-age=600")
-    def with_cache_control():
+    def with_cache_control() -> Tuple[int, str, str]:
         return 200, TEXT_HTML, "has 200 response"
 
     def handler(event, context):
         return app.resolve(event, context)
 
+    # WHEN calling the event handler
+    # AND the function returns a 200 status code
     result = handler({"path": "/success", "httpMethod": "GET"}, None)
 
+    # THEN return the set Cache-Control
     headers = result["headers"]
     assert headers["Content-Type"] == TEXT_HTML
     assert headers["Cache-Control"] == "max-age=600"
 
 
 def test_cache_control_non_200():
+    # GIVEN a function with cache_control set
     app = ApiGatewayResolver()
 
     @app.delete("/fails", cache_control="max-age=600")
-    def with_cache_control_has_500():
+    def with_cache_control_has_500() -> Tuple[int, str, str]:
         return 503, TEXT_HTML, "has 503 response"
 
     def handler(event, context):
         return app.resolve(event, context)
 
+    # WHEN calling the event handler
+    # AND the function returns a 503 status code
     result = handler({"path": "/fails", "httpMethod": "DELETE"}, None)
 
+    # THEN return a Cache-Control of "no-cache"
     headers = result["headers"]
     assert headers["Content-Type"] == TEXT_HTML
     assert headers["Cache-Control"] == "no-cache"
 
 
 def test_rest_api():
+    # GIVEN a function that returns a Dict
     app = ApiGatewayResolver(proxy_type=ProxyEventType.http_api_v1)
 
     @app.get("/my/path")
-    def rest_func():
+    def rest_func() -> Dict:
         return {"foo": "value"}
 
+    # WHEN calling the event handler
     result = app(LOAD_GW_EVENT, {})
 
+    # THEN automatically process this as a json rest api response
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
     assert result["body"] == json.dumps({"foo": "value"})
+
+
+def test_handling_response_type():
+    # GIVEN a function that returns Response
+    app = ApiGatewayResolver(proxy_type=ProxyEventType.http_api_v1)
+
+    @app.get("/my/path")
+    def rest_func() -> Response:
+        return Response(
+            status_code=404,
+            content_type="used-if-not-set-in-header",
+            body="Not found",
+            headers={"Content-Type": "header-content-type-wins", "custom": "value"},
+        )
+
+    # WHEN calling the event handler
+    result = app(LOAD_GW_EVENT, {})
+
+    # THEN the result can include some additional field control like overriding http headers
+    assert result["statusCode"] == 404
+    assert result["headers"]["Content-Type"] == "header-content-type-wins"
+    assert result["headers"]["custom"] == "value"
+    assert result["body"] == "Not found"