feat(api-gateway): add debug mode

aws_lambda_powertools/event_handler/api_gateway.py

@@ -1,11 +1,15 @@
 import base64
 import json
 import logging
+import os
 import re
+import traceback
 import zlib
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Set, Union
 
+from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared.functions import resolve_truthy_env_var_choice
 from aws_lambda_powertools.shared.json_encoder import Encoder
 from aws_lambda_powertools.utilities.data_classes import ALBEvent, APIGatewayProxyEvent, APIGatewayProxyEventV2
 from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
@@ -237,20 +241,28 @@ def lambda_handler(event, context):
     current_event: BaseProxyEvent
     lambda_context: LambdaContext
 
-    def __init__(self, proxy_type: Enum = ProxyEventType.APIGatewayProxyEvent, cors: CORSConfig = None):
+    def __init__(
+        self,
+        proxy_type: Enum = ProxyEventType.APIGatewayProxyEvent,
+        cors: CORSConfig = None,
+        debug: Optional[bool] = None,
+    ):
         """
         Parameters
         ----------
         proxy_type: ProxyEventType
             Proxy request type, defaults to API Gateway V1
         cors: CORSConfig
             Optionally configure and enabled CORS. Not each route will need to have to cors=True
+        debug: Optional[bool]
+            Enables debug mode, by default False. Can be enabled by "POWERTOOLS_API_DEBUG" environment variable
         """
         self._proxy_type = proxy_type
         self._routes: List[Route] = []
         self._cors = cors
         self._cors_enabled: bool = cors is not None
         self._cors_methods: Set[str] = {"OPTIONS"}
+        self._debug = resolve_truthy_env_var_choice(choice=debug, env=os.getenv(constants.API_DEBUG_ENV, "false"))
 
     def get(self, rule: str, cors: bool = None, compress: bool = False, cache_control: str = None):
         """Get route decorator with GET `method`
@@ -469,16 +481,29 @@ def _not_found(self, method: str) -> ResponseBuilder:
                 status_code=404,
                 content_type="application/json",
                 headers=headers,
-                body=json.dumps({"message": "Not found"}),
+                body=self._json_dump({"message": "Not found"}),
             )
         )
 
     def _call_route(self, route: Route, args: Dict[str, str]) -> ResponseBuilder:
         """Actually call the matching route with any provided keyword arguments."""
-        return ResponseBuilder(self._to_response(route.func(**args)), route)
-
-    @staticmethod
-    def _to_response(result: Union[Dict, Response]) -> Response:
+        try:
+            return ResponseBuilder(self._to_response(route.func(**args)), route)
+        except Exception:
+            if self._debug:
+                # If the user has turned on debug mode,
+                # we'll let the original exception propagate so
+                # they get more information about what went wrong.
+                return ResponseBuilder(
+                    Response(
+                        status_code=500,
+                        content_type="text/plain",
+                        body="".join(traceback.format_exc()),
+                    )
+                )
+            raise
+
+    def _to_response(self, result: Union[Dict, Response]) -> Response:
         """Convert the route's result to a Response
 
          2 main result types are supported:
@@ -494,5 +519,12 @@ def _to_response(result: Union[Dict, Response]) -> Response:
         return Response(
             status_code=200,
             content_type="application/json",
-            body=json.dumps(result, separators=(",", ":"), cls=Encoder),
+            body=self._json_dump(result),
         )
+
+    def _json_dump(self, obj: Any) -> str:
+        """Does a concise json serialization"""
+        if self._debug:
+            return json.dumps(obj, indent=4, cls=Encoder)
+        else:
+            return json.dumps(obj, separators=(",", ":"), cls=Encoder)

aws_lambda_powertools/shared/constants.py

@@ -18,3 +18,5 @@
 
 XRAY_SDK_MODULE = "aws_xray_sdk"
 XRAY_SDK_CORE_MODULE = "aws_xray_sdk.core"
+
+API_DEBUG_ENV: str = "POWERTOOLS_API_DEBUG"

tests/functional/event_handler/test_api_gateway.py

@@ -5,13 +5,16 @@
 from pathlib import Path
 from typing import Dict
 
+import pytest
+
 from aws_lambda_powertools.event_handler.api_gateway import (
     ApiGatewayResolver,
     CORSConfig,
     ProxyEventType,
     Response,
     ResponseBuilder,
 )
+from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.json_encoder import Encoder
 from aws_lambda_powertools.utilities.data_classes import ALBEvent, APIGatewayProxyEvent, APIGatewayProxyEventV2
 from tests.functional.utils import load_event
@@ -490,3 +493,70 @@ def custom_method():
     assert headers["Content-Type"] == TEXT_HTML
     assert "Access-Control-Allow-Origin" in result["headers"]
     assert headers["Access-Control-Allow-Methods"] == "CUSTOM"
+
+
+def test_unhandled_exceptions_debug_on():
+    # GIVEN debug is enabled
+    # AND an unhandled exception is raised
+    app = ApiGatewayResolver(debug=True)
+    assert app._debug
+
+    @app.get("/raises-error")
+    def raises_error():
+        raise RuntimeError("Foo")
+
+    # WHEN calling the handler
+    result = app({"path": "/raises-error", "httpMethod": "GET"}, None)
+
+    # THEN return a 500
+    # AND Content-Type is set to text/plain
+    # AND include the exception traceback in the response
+    assert result["statusCode"] == 500
+    assert "Traceback (most recent call last)" in result["body"]
+    headers = result["headers"]
+    assert headers["Content-Type"] == "text/plain"
+
+
+def test_unhandled_exceptions_debug_off():
+    # GIVEN debug is disabled
+    # AND an unhandled exception is raised
+    app = ApiGatewayResolver(debug=False)
+    assert not app._debug
+
+    @app.get("/raises-error")
+    def raises_error():
+        raise RuntimeError("Foo")
+
+    # WHEN calling the handler
+    # THEN raise the original exception
+    with pytest.raises(RuntimeError) as e:
+        app({"path": "/raises-error", "httpMethod": "GET"}, None)
+
+    # AND include the original error
+    assert e.value.args == ("Foo",)
+
+
+def test_debug_mode_environment_variable(monkeypatch):
+    # GIVEN a debug mode environment variable is set
+    monkeypatch.setenv(constants.API_DEBUG_ENV, "true")
+    app = ApiGatewayResolver()
+
+    # WHEN calling app._debug
+    # THEN the debug mode is enabled
+    assert app._debug
+
+
+def test_debug_json_formatting():
+    # GIVEN debug is True
+    app = ApiGatewayResolver(debug=True)
+    response = {"message": "Foo"}
+
+    @app.get("/foo")
+    def foo():
+        return response
+
+    # WHEN calling the handler
+    result = app({"path": "/foo", "httpMethod": "GET"}, None)
+
+    # THEN return a pretty print json in the body
+    assert result["body"] == json.dumps(response, indent=4)