feat(event-handler): apigwy cache_control option

Author: Michael Brewer

aws_lambda_powertools/event_handler/api_gateway.py

@@ -17,12 +17,15 @@ class ProxyEventType(Enum):
 
 
 class RouteEntry:
-    def __init__(self, method: str, rule: Any, func: Callable, cors: bool, compress: bool):
+    def __init__(
+        self, method: str, rule: Any, func: Callable, cors: bool, compress: bool, cache_control: Optional[str]
+    ):
         self.method = method.upper()
         self.rule = rule
         self.func = func
         self.cors = cors
         self.compress = compress
+        self.cache_control = cache_control
 
 
 class ApiGatewayResolver:
@@ -33,21 +36,21 @@ def __init__(self, proxy_type: Enum = ProxyEventType.http_api_v1):
         self._proxy_type = proxy_type
         self._routes: List[RouteEntry] = []
 
-    def get(self, rule: str, cors: bool = False, compress: bool = False):
-        return self.route(rule, "GET", cors, compress)
+    def get(self, rule: str, cors: bool = False, compress: bool = False, cache_control: str = None):
+        return self.route(rule, "GET", cors, compress, cache_control)
 
-    def post(self, rule: str, cors: bool = False, compress: bool = False):
-        return self.route(rule, "POST", cors, compress)
+    def post(self, rule: str, cors: bool = False, compress: bool = False, cache_control: str = None):
+        return self.route(rule, "POST", cors, compress, cache_control)
 
-    def put(self, rule: str, cors: bool = False, compress: bool = False):
-        return self.route(rule, "PUT", cors, compress)
+    def put(self, rule: str, cors: bool = False, compress: bool = False, cache_control: str = None):
+        return self.route(rule, "PUT", cors, compress, cache_control)
 
-    def delete(self, rule: str, cors: bool = False, compress: bool = False):
-        return self.route(rule, "DELETE", cors, compress)
+    def delete(self, rule: str, cors: bool = False, compress: bool = False, cache_control: str = None):
+        return self.route(rule, "DELETE", cors, compress, cache_control)
 
-    def route(self, rule: str, method: str, cors: bool = False, compress: bool = False):
+    def route(self, rule: str, method: str, cors: bool = False, compress: bool = False, cache_control: str = None):
         def register_resolver(func: Callable):
-            self._register(func, rule, method, cors, compress)
+            self._append(func, rule, method, cors, compress, cache_control)
             return func
 
         return register_resolver
@@ -58,31 +61,34 @@ def resolve(self, event: Dict, context: LambdaContext) -> Dict:
 
         route, args = self._find_route(self.current_event.http_method, self.current_event.path)
         result = route.func(**args)
+
+        status: int = result[0]
+        response: Dict[str, Any] = {"statusCode": status}
+
         headers = {"Content-Type": result[1]}
         if route.cors:
             headers["Access-Control-Allow-Origin"] = "*"
             headers["Access-Control-Allow-Methods"] = route.method
             headers["Access-Control-Allow-Credentials"] = "true"
+        if route.cache_control:
+            headers["Cache-Control"] = route.cache_control if status == 200 else "no-cache"
+        response["headers"] = headers
 
         body: Union[str, bytes] = result[2]
         if route.compress and "gzip" in (self.current_event.get_header_value("accept-encoding") or ""):
             gzip_compress = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16)
             if isinstance(body, str):
                 body = bytes(body, "utf-8")
             body = gzip_compress.compress(body) + gzip_compress.flush()
-
-        response = {"statusCode": result[0], "headers": headers}
-
         if isinstance(body, bytes):
             response["isBase64Encoded"] = True
             body = base64.b64encode(body).decode()
-
         response["body"] = body
 
         return response
 
-    def _register(self, func: Callable, rule: str, method: str, cors: bool, compress: bool):
-        self._routes.append(RouteEntry(method, self._build_rule_pattern(rule), func, cors, compress))
+    def _append(self, func: Callable, rule: str, method: str, cors: bool, compress: bool, cache_control: Optional[str]):
+        self._routes.append(RouteEntry(method, self._build_rule_pattern(rule), func, cors, compress, cache_control))
 
     @staticmethod
     def _build_rule_pattern(rule: str):

tests/functional/event_handler/test_api_gateway.py

@@ -154,3 +154,35 @@ def handler(event, context):
     assert isinstance(body, str)
     decompress = zlib.decompress(base64.b64decode(body), wbits=zlib.MAX_WBITS | 16).decode("UTF-8")
     assert decompress == expected_value
+
+
+def test_cache_control_200():
+    app = ApiGatewayResolver()
+
+    @app.get("/success", cache_control="max-age=600")
+    def with_cache_control():
+        return 200, "text/html", "has 200 response"
+
+    def handler(event, context):
+        return app.resolve(event, context)
+
+    result = handler({"path": "/success", "httpMethod": "GET"}, None)
+
+    headers = result["headers"]
+    assert headers["Cache-Control"] == "max-age=600"
+
+
+def test_cache_control_non_200():
+    app = ApiGatewayResolver()
+
+    @app.delete("/fails", cache_control="max-age=600")
+    def with_cache_control_has_500():
+        return 503, "text/html", "has 503 response"
+
+    def handler(event, context):
+        return app.resolve(event, context)
+
+    result = handler({"path": "/fails", "httpMethod": "DELETE"}, None)
+
+    headers = result["headers"]
+    assert headers["Cache-Control"] == "no-cache"