feat(event_handler): add support for additional response models

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -37,6 +37,9 @@
 from aws_lambda_powertools.event_handler.openapi.types import (
     COMPONENT_REF_PREFIX,
     METHODS_WITH_BODY,
+    OpenAPIResponse,
+    OpenAPIResponseContentModel,
+    OpenAPIResponseContentSchema,
     validation_error_definition,
     validation_error_response_definition,
 )
@@ -273,7 +276,7 @@ def __init__(
         cache_control: Optional[str],
         summary: Optional[str],
         description: Optional[str],
-        responses: Optional[Dict[int, Dict[str, Any]]],
+        responses: Optional[Dict[int, OpenAPIResponse]],
         response_description: Optional[str],
         tags: Optional[List[str]],
         operation_id: Optional[str],
@@ -303,7 +306,7 @@ def __init__(
             The OpenAPI summary for this route
         description: Optional[str]
             The OpenAPI description for this route
-        responses: Optional[Dict[int, Dict[str, Any]]]
+        responses: Optional[Dict[int, OpenAPIResponse]]
             The OpenAPI responses for this route
         response_description: Optional[str]
             The OpenAPI response description for this route
@@ -501,11 +504,54 @@ def _get_openapi_path(
 
         # Add the response to the OpenAPI operation
         if self.responses:
-            # If the user supplied responses, we use them and don't set a default 200 response
+            for status_code in list(self.responses):
+                response = self.responses[status_code]
+
+                # Case 1: there is not 'content' key
+                if "content" not in response:
+                    response["content"] = {
+                        "application/json": self._openapi_operation_return(
+                            param=dependant.return_param,
+                            model_name_map=model_name_map,
+                            field_mapping=field_mapping,
+                        ),
+                    }
+
+                # Case 2: there is a 'content' key
+                else:
+                    # Need to iterate to transform any 'model' into a 'schema'
+                    for content_type, payload in response["content"].items():
+                        new_payload: OpenAPIResponseContentSchema
+
+                        # Case 2.1: the 'content' has a model
+                        if "model" in payload:
+                            from aws_lambda_powertools.event_handler.openapi.params import analyze_param
+
+                            return_field = analyze_param(
+                                param_name="return",
+                                annotation=cast(OpenAPIResponseContentModel, payload)["model"],
+                                value=None,
+                                is_path_param=False,
+                                is_response_param=True,
+                            )
+
+                            new_payload = self._openapi_operation_return(
+                                param=return_field,
+                                model_name_map=model_name_map,
+                                field_mapping=field_mapping,
+                            )
+
+                        # Case 2.2: the 'content' has a schema
+                        else:
+                            # Do nothing! We already have what we need!
+                            new_payload = payload
+
+                        response["content"][content_type] = new_payload
+
             operation["responses"] = self.responses
         else:
             # Set the default 200 response
-            responses = operation.setdefault("responses", self.responses or {})
+            responses = operation.setdefault("responses", {})
             success_response = responses.setdefault(200, {})
             success_response["description"] = self.response_description or _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION
             success_response["content"] = {"application/json": {"schema": {}}}
@@ -682,7 +728,7 @@ def _openapi_operation_return(
             Tuple["ModelField", Literal["validation", "serialization"]],
             "JsonSchemaValue",
         ],
-    ) -> Dict[str, Any]:
+    ) -> OpenAPIResponseContentSchema:
         """
         Returns the OpenAPI operation return.
         """
@@ -832,7 +878,7 @@ def route(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -890,7 +936,7 @@ def get(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -943,7 +989,7 @@ def post(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -997,7 +1043,7 @@ def put(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -1051,7 +1097,7 @@ def delete(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -1104,7 +1150,7 @@ def patch(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -1657,7 +1703,7 @@ def route(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -2127,7 +2173,7 @@ def route(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: Optional[str] = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,
@@ -2216,7 +2262,7 @@ def route(
         cache_control: Optional[str] = None,
         summary: Optional[str] = None,
         description: Optional[str] = None,
-        responses: Optional[Dict[int, Dict[str, Any]]] = None,
+        responses: Optional[Dict[int, OpenAPIResponse]] = None,
         response_description: str = _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
         tags: Optional[List[str]] = None,
         operation_id: Optional[str] = None,

aws_lambda_powertools/event_handler/openapi/types.py

@@ -2,6 +2,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type, Union
 
+from aws_lambda_powertools.shared.types import NotRequired, TypedDict
+
 if TYPE_CHECKING:
     from pydantic import BaseModel  # noqa: F401
 
@@ -43,3 +45,16 @@
         },
     },
 }
+
+
+class OpenAPIResponseContentSchema(TypedDict, total=False):
+    schema: Dict
+
+
+class OpenAPIResponseContentModel(TypedDict):
+    model: Any
+
+
+class OpenAPIResponse(TypedDict):
+    description: str
+    content: NotRequired[Dict[str, Union[OpenAPIResponseContentSchema, OpenAPIResponseContentModel]]]

tests/functional/event_handler/test_openapi_responses.py

@@ -1,4 +1,9 @@
-from aws_lambda_powertools.event_handler import APIGatewayRestResolver
+from random import random
+from typing import Union
+
+from pydantic import BaseModel
+
+from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response
 
 
 def test_openapi_default_response():
@@ -47,3 +52,69 @@ def handler():
 
     assert 200 not in responses.keys()
     assert 422 not in responses.keys()
+
+
+def test_openapi_union_response():
+    app = APIGatewayRestResolver(enable_validation=True)
+
+    class User(BaseModel):
+        pass
+
+    class Order(BaseModel):
+        pass
+
+    @app.get(
+        "/",
+        responses={
+            200: {"description": "200 Response", "content": {"application/json": {"model": User}}},
+            202: {"description": "202 Response", "content": {"application/json": {"model": Order}}},
+        },
+    )
+    def handler() -> Response[Union[User, Order]]:
+        if random() > 0.5:
+            return Response(status_code=200, body=User())
+        else:
+            return Response(status_code=202, body=Order())
+
+    schema = app.get_openapi_schema()
+    responses = schema.paths["/"].get.responses
+    assert 200 in responses.keys()
+    assert responses[200].description == "200 Response"
+    assert responses[200].content["application/json"].schema_.ref == "#/components/schemas/User"
+
+    assert 202 in responses.keys()
+    assert responses[202].description == "202 Response"
+    assert responses[202].content["application/json"].schema_.ref == "#/components/schemas/Order"
+
+
+def test_openapi_union_partial_response():
+    app = APIGatewayRestResolver(enable_validation=True)
+
+    class User(BaseModel):
+        pass
+
+    class Order(BaseModel):
+        pass
+
+    @app.get(
+        "/",
+        responses={
+            200: {"description": "200 Response"},
+            202: {"description": "202 Response", "content": {"application/json": {"model": Order}}},
+        },
+    )
+    def handler() -> Response[Union[User, Order]]:
+        if random() > 0.5:
+            return Response(status_code=200, body=User())
+        else:
+            return Response(status_code=202, body=Order())
+
+    schema = app.get_openapi_schema()
+    responses = schema.paths["/"].get.responses
+    assert 200 in responses.keys()
+    assert responses[200].description == "200 Response"
+    assert responses[200].content["application/json"].schema_.anyOf is not None
+
+    assert 202 in responses.keys()
+    assert responses[202].description == "202 Response"
+    assert responses[202].content["application/json"].schema_.ref == "#/components/schemas/Order"