feat(event_handler): add route-level custom response validation in OpenAPI utility (#6341)

* feat(api-gateway-resolver): Add option for custom response validation error status code.

* feat(docs): Added doc for custom response validation error responses.

* feat(unit-test): Add tests for custom response validation error.

* fix: Formatting.

* fix(unit-test): fix failed CI.

* feat(unit-test): add tests for incorrect types and invalid configs

* refactor: rename response_validation_error_http_status to response_validation_error_http_code

* refactor(tests): move unit tests into openapi_validation functional test file

* feat: add route-specific custom response validation and tests

* fix: except Route implementation

* fix: put custom_response_validation_http_code before middleware

* feat: route's custom response validation must take precedence over app's.

* feat: added more tests.

* refactor: improved error messagee and tests' descriptions.

* feat: updated docs.

* move veritifcation method of route custom http code to BaseRouter.

* fix: add validate function for route http code to APIGatewayResolver not Router

* feat: add custom_response_validation_http_code to the routes of Bedrock

* fix: make mypy happy

* fix: address comments

* fix(openapi): add response for response validation error and definition for it

* minor changes

* minor changes

---------

Co-authored-by: Amin Farjadi <[email protected]>
Co-authored-by: Leandro Damascena <[email protected]>

Author: 3 people

aws_lambda_powertools/event_handler/api_gateway.py

@@ -35,6 +35,7 @@
     OpenAPIResponse,
     OpenAPIResponseContentModel,
     OpenAPIResponseContentSchema,
+    response_validation_error_response_definition,
     validation_error_definition,
     validation_error_response_definition,
 )
@@ -319,6 +320,7 @@ def __init__(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: HTTPStatus | None = None,
         middlewares: list[Callable[..., Response]] | None = None,
     ):
         """
@@ -360,11 +362,13 @@ def __init__(
             Additional OpenAPI extensions as a dictionary.
         deprecated: bool
             Whether or not to mark this route as deprecated in the OpenAPI schema
+        custom_response_validation_http_code: int | HTTPStatus | None, optional
+            Whether to have custom http status code for this route if response validation fails
         middlewares: list[Callable[..., Response]] | None
             The list of route middlewares to be called in order.
         """
         self.method = method.upper()
-        self.path = "/" if path.strip() == "" else path
+        self.path = path if path.strip() else "/"
 
         # OpenAPI spec only understands paths with { }. So we'll have to convert Powertools' < >.
         # https://swagger.io/specification/#path-templating
@@ -397,6 +401,8 @@ def __init__(
         # _body_field is used to cache the dependant model for the body field
         self._body_field: ModelField | None = None
 
+        self.custom_response_validation_http_code = custom_response_validation_http_code
+
     def __call__(
         self,
         router_middlewares: list[Callable],
@@ -565,6 +571,16 @@ def _get_openapi_path(
             },
         }
 
+        # Add custom response validation response, if exists
+        if self.custom_response_validation_http_code:
+            http_code = self.custom_response_validation_http_code.value
+            operation_responses[http_code] = {
+                "description": "Response Validation Error",
+                "content": {"application/json": {"schema": {"$ref": f"{COMPONENT_REF_PREFIX}ResponseValidationError"}}},
+            }
+            # Add model definition
+            definitions["ResponseValidationError"] = response_validation_error_response_definition
+
         # Add the response to the OpenAPI operation
         if self.responses:
             for status_code in list(self.responses):
@@ -942,6 +958,7 @@ def route(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         raise NotImplementedError()
@@ -1003,6 +1020,7 @@ def get(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Get route decorator with GET `method`
@@ -1043,6 +1061,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1062,6 +1081,7 @@ def post(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Post route decorator with POST `method`
@@ -1103,6 +1123,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1122,6 +1143,7 @@ def put(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Put route decorator with PUT `method`
@@ -1163,6 +1185,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1182,6 +1205,7 @@ def delete(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Delete route decorator with DELETE `method`
@@ -1222,6 +1246,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1241,6 +1266,7 @@ def patch(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Patch route decorator with PATCH `method`
@@ -1284,6 +1310,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1303,6 +1330,7 @@ def head(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Head route decorator with HEAD `method`
@@ -1345,6 +1373,7 @@ def lambda_handler(event, context):
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )
 
@@ -1571,6 +1600,7 @@ def _validate_response_validation_error_http_code(
         response_validation_error_http_code: HTTPStatus | int | None,
         enable_validation: bool,
     ) -> HTTPStatus:
+
         if response_validation_error_http_code and not enable_validation:
             msg = "'response_validation_error_http_code' cannot be set when enable_validation is False."
             raise ValueError(msg)
@@ -1588,6 +1618,33 @@ def _validate_response_validation_error_http_code(
 
         return response_validation_error_http_code or HTTPStatus.UNPROCESSABLE_ENTITY
 
+    def _add_resolver_response_validation_error_response_to_route(
+        self,
+        route_openapi_path: tuple[dict[str, Any], dict[str, Any]],
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Adds resolver response validation error response to route's operations."""
+        path, path_definitions = route_openapi_path
+        if self._has_response_validation_error and "ResponseValidationError" not in path_definitions:
+            response_validation_error_response = {
+                "description": "Response Validation Error",
+                "content": {
+                    "application/json": {
+                        "schema": {"$ref": f"{COMPONENT_REF_PREFIX}ResponseValidationError"},
+                    },
+                },
+            }
+            http_code = self._response_validation_error_http_code.value
+            for operation in path.values():
+                operation["responses"][http_code] = response_validation_error_response
+        return path, path_definitions
+
+    def _generate_schemas(self, definitions: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
+        schemas = {k: definitions[k] for k in sorted(definitions)}
+        # add response validation error definition
+        if self._response_validation_error_http_code:
+            schemas.setdefault("ResponseValidationError", response_validation_error_response_definition)
+        return schemas
+
     def get_openapi_schema(
         self,
         *,
@@ -1739,14 +1796,14 @@ def get_openapi_schema(
                 field_mapping=field_mapping,
             )
             if result:
-                path, path_definitions = result
+                path, path_definitions = self._add_resolver_response_validation_error_response_to_route(result)
                 if path:
                     paths.setdefault(route.openapi_path, {}).update(path)
                 if path_definitions:
                     definitions.update(path_definitions)
 
         if definitions:
-            components["schemas"] = {k: definitions[k] for k in sorted(definitions)}
+            components["schemas"] = self._generate_schemas(definitions)
         if security_schemes:
             components["securitySchemes"] = security_schemes
         if components:
@@ -2108,6 +2165,29 @@ def swagger_handler():
                 body=body,
             )
 
+    def _validate_route_response_validation_error_http_code(
+        self,
+        custom_response_validation_http_code: int | HTTPStatus | None,
+    ) -> HTTPStatus | None:
+        if custom_response_validation_http_code and not self._enable_validation:
+            msg = (
+                "'custom_response_validation_http_code' cannot be set for route when enable_validation is False "
+                "on resolver."
+            )
+            raise ValueError(msg)
+
+        if (
+            not isinstance(custom_response_validation_http_code, HTTPStatus)
+            and custom_response_validation_http_code is not None
+        ):
+            try:
+                custom_response_validation_http_code = HTTPStatus(custom_response_validation_http_code)
+            except ValueError:
+                msg = f"'{custom_response_validation_http_code}' must be an integer representing an HTTP status code or an enum of type HTTPStatus."  # noqa: E501
+                raise ValueError(msg) from None
+
+        return custom_response_validation_http_code
+
     def route(
         self,
         rule: str,
@@ -2125,10 +2205,15 @@ def route(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Route decorator includes parameter `method`"""
 
+        custom_response_validation_http_code = self._validate_route_response_validation_error_http_code(
+            custom_response_validation_http_code,
+        )
+
         def register_resolver(func: AnyCallableT) -> AnyCallableT:
             methods = (method,) if isinstance(method, str) else method
             logger.debug(f"Adding route using rule {rule} and methods: {','.join(m.upper() for m in methods)}")
@@ -2154,6 +2239,7 @@ def register_resolver(func: AnyCallableT) -> AnyCallableT:
                     security,
                     openapi_extensions,
                     deprecated,
+                    custom_response_validation_http_code,
                     middlewares,
                 )
 
@@ -2523,15 +2609,17 @@ def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuild
             )
 
         # OpenAPIValidationMiddleware will only raise ResponseValidationError when
-        # 'self._response_validation_error_http_code' is not None
+        # 'self._response_validation_error_http_code' is not None or
+        # when route has custom_response_validation_http_code
         if isinstance(exp, ResponseValidationError):
-            http_code = self._response_validation_error_http_code
+            # route validation must take precedence over app validation
+            http_code = route.custom_response_validation_http_code or self._response_validation_error_http_code
             errors = [{"loc": e["loc"], "type": e["type"]} for e in exp.errors()]
             return self._response_builder_class(
                 response=Response(
                     status_code=http_code.value,
                     content_type=content_types.APPLICATION_JSON,
-                    body={"statusCode": self._response_validation_error_http_code, "detail": errors},
+                    body={"statusCode": http_code, "detail": errors},
                 ),
                 serializer=self._serializer,
                 route=route,
@@ -2682,6 +2770,7 @@ def route(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         def register_route(func: AnyCallableT) -> AnyCallableT:
@@ -2708,6 +2797,7 @@ def register_route(func: AnyCallableT) -> AnyCallableT:
                 frozen_security,
                 frozen_openapi_extensions,
                 deprecated,
+                custom_response_validation_http_code,
             )
 
             # Collate Middleware for routes
@@ -2794,6 +2884,7 @@ def route(
         security: list[dict[str, list[str]]] | None = None,
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         # NOTE: see #1552 for more context.
@@ -2813,6 +2904,7 @@ def route(
             security,
             openapi_extensions,
             deprecated,
+            custom_response_validation_http_code,
             middlewares,
         )