feat(event_handler): adding status_code OpenAPI field

leandrodamascena · leandrodamascena · commit 3838a425f208 · 2026-04-08T15:08:39.000+08:00
diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py
@@ -387,6 +387,7 @@ def __init__(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Response]] | None = None,
     ):
         """
@@ -432,6 +433,9 @@ def __init__(
             Enable or disable validation for this specific route. If None, inherits from resolver setting.
         custom_response_validation_http_code: int | HTTPStatus | None, optional
             Whether to have custom http status code for this route if response validation fails
+        status_code: int
+            The default HTTP status code for successful responses. Used in both the OpenAPI schema
+            and the actual response when the handler returns a dict. Defaults to 200.
         middlewares: list[Callable[..., Response]] | None
             The list of route middlewares to be called in order.
         """
@@ -471,6 +475,7 @@ def __init__(
         self._body_field: ModelField | None = None
 
         self.custom_response_validation_http_code = custom_response_validation_http_code
+        self.status_code = status_code
 
         # Caches the name of any Request-typed parameter in the handler.
         # Avoids re-scanning the signature on every invocation.
@@ -652,6 +657,7 @@ def _get_openapi_path(
             response_description=self.response_description,
             body_field=self.body_field,
             custom_response_validation_http_code=self.custom_response_validation_http_code,
+            status_code=self.status_code,
             dependant=dependant,
             operation_ids=operation_ids,
             model_name_map=model_name_map,
@@ -808,6 +814,7 @@ def route(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         raise NotImplementedError()
@@ -871,6 +878,7 @@ def get(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Get route decorator with GET `method`
@@ -913,6 +921,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -934,6 +943,7 @@ def post(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Post route decorator with POST `method`
@@ -977,6 +987,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -998,6 +1009,7 @@ def put(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Put route decorator with PUT `method`
@@ -1041,6 +1053,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -1062,6 +1075,7 @@ def delete(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Delete route decorator with DELETE `method`
@@ -1104,6 +1118,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -1125,6 +1140,7 @@ def patch(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Patch route decorator with PATCH `method`
@@ -1170,6 +1186,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -1191,6 +1208,7 @@ def head(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Head route decorator with HEAD `method`
@@ -1235,6 +1253,7 @@ def lambda_handler(event, context):
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -2329,6 +2348,7 @@ def route(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Route decorator includes parameter `method`"""
@@ -2364,6 +2384,7 @@ def register_resolver(func: AnyCallableT) -> AnyCallableT:
                     deprecated,
                     enable_validation,
                     custom_response_validation_http_code,
+                    status_code,
                     middlewares,
                 )
 
@@ -2751,12 +2772,15 @@ def _to_response(self, result: dict | tuple | Response | BedrockResponse) -> Res
         - tuple[dict, int]: Same dict handling as above but with the option of including a status code
         - Response: returned as is, and allows for more flexibility
         """
-        status_code = HTTPStatus.OK
         if isinstance(result, (Response, BedrockResponse)):
             return result
         elif isinstance(result, tuple) and len(result) == 2:
             # Unpack result dict and status code from tuple
             result, status_code = result
+        else:
+            # Use the route's status_code if available, otherwise default to 200
+            route: Route | None = self.context.get("_route")
+            status_code = route.status_code if route else HTTPStatus.OK
 
         logger.debug("Simple response detected, serializing return before constructing final response")
         return Response(
@@ -2875,6 +2899,7 @@ def route(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         def register_route(func: AnyCallableT) -> AnyCallableT:
@@ -2903,6 +2928,7 @@ def register_route(func: AnyCallableT) -> AnyCallableT:
                 deprecated,
                 enable_validation,
                 custom_response_validation_http_code,
+                status_code,
             )
 
             # Collate Middleware for routes
@@ -2972,6 +2998,7 @@ def route(
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         # NOTE: see #1552 for more context.
@@ -2993,6 +3020,7 @@ def route(
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py
@@ -129,6 +129,7 @@ def get(  # type: ignore[override]
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
         security = None
@@ -150,6 +151,7 @@ def get(  # type: ignore[override]
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -172,6 +174,7 @@ def post(  # type: ignore[override]
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
         security = None
@@ -193,6 +196,7 @@ def post(  # type: ignore[override]
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -215,6 +219,7 @@ def put(  # type: ignore[override]
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
         security = None
@@ -236,6 +241,7 @@ def put(  # type: ignore[override]
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -258,6 +264,7 @@ def patch(  # type: ignore[override]
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable] | None = None,
     ):
         security = None
@@ -279,6 +286,7 @@ def patch(  # type: ignore[override]
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
@@ -301,6 +309,7 @@ def delete(  # type: ignore[override]
         deprecated: bool = False,
         enable_validation: bool | None = None,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
+        status_code: int = 200,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
         security = None
@@ -322,6 +331,7 @@ def delete(  # type: ignore[override]
             deprecated,
             enable_validation,
             custom_response_validation_http_code,
+            status_code,
             middlewares,
         )
 
diff --git a/aws_lambda_powertools/event_handler/openapi/schema_generator.py b/aws_lambda_powertools/event_handler/openapi/schema_generator.py
@@ -54,6 +54,7 @@ def generate_openapi_path(
     response_description: str | None,
     body_field: ModelField | None,
     custom_response_validation_http_code: HTTPStatus | None,
+    status_code: int = 200,
     dependant: Dependant,
     operation_ids: set[str],
     model_name_map: dict[TypeModelOrEnum, str],
@@ -108,6 +109,7 @@ def generate_openapi_path(
         responses=responses,
         response_description=response_description,
         custom_response_validation_http_code=custom_response_validation_http_code,
+        status_code=status_code,
         dependant=dependant,
         model_name_map=model_name_map,
         field_mapping=field_mapping,
@@ -220,6 +222,7 @@ def _build_responses(
     responses: dict[int, OpenAPIResponse] | None,
     response_description: str | None,
     custom_response_validation_http_code: HTTPStatus | None,
+    status_code: int = 200,
     dependant: Dependant,
     model_name_map: dict[TypeModelOrEnum, str],
     field_mapping: dict[tuple[ModelField, Literal["validation", "serialization"]], JsonSchemaValue],
@@ -237,9 +240,9 @@ def _build_responses(
     )
 
     if responses:
-        for status_code in list(responses):
-            operation_responses[status_code] = _build_custom_response(
-                response=copy.deepcopy(responses[status_code]),
+        for resp_code in list(responses):
+            operation_responses[resp_code] = _build_custom_response(
+                response=copy.deepcopy(responses[resp_code]),
                 dependant=dependant,
                 model_name_map=model_name_map,
                 field_mapping=field_mapping,
@@ -251,7 +254,7 @@ def _build_responses(
             field_mapping=field_mapping,
         )
 
-        operation_responses[200] = {
+        operation_responses[status_code] = {
             "description": response_description or DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
             "content": {DEFAULT_CONTENT_TYPE: response_schema},
         }
diff --git a/docs/core/event_handler/_openapi_customization_operations.md b/docs/core/event_handler/_openapi_customization_operations.md
@@ -13,4 +13,5 @@ Here's a breakdown of various customizable fields:
 | `tags`                 | `List[str]`                             | Tags are a way to categorize and group endpoints within the API documentation. They can help organize the operations by resources or other heuristic.                                                                                                                                                                      |
 | `operation_id`         | `str`                                   | A unique identifier for the operation, which can be used for referencing this operation in documentation or code. This ID must be unique across all operations described in the API.                                                                                                                                       |
 | `include_in_schema`    | `bool`                                  | A boolean value that determines whether or not this operation should be included in the OpenAPI schema. Setting it to `False` can hide the endpoint from generated documentation and schema exports, which might be useful for private or experimental endpoints.                                                          |
-| `deprecated`           | `bool`                                   | A boolean value that determines whether or not this operation should be marked as deprecated in the OpenAPI schema.                                                             |
+| `deprecated`           | `bool`                                  | A boolean value that determines whether or not this operation should be marked as deprecated in the OpenAPI schema.                                                                                                                                                                                                        |
+| `status_code`          | `int`                                   | The default HTTP status code for successful responses. Defaults to `200`. This value is used both in the OpenAPI schema and as the actual response status code when the handler returns a dictionary or plain value (not a `Response` object or tuple).                                                                    |
diff --git a/docs/core/event_handler/openapi.md b/docs/core/event_handler/openapi.md
@@ -73,7 +73,7 @@ To implement these customizations, include extra parameters when defining your r
 
 === "customizing_api_operations.py"
 
-    ```python hl_lines="11-20"
+    ```python hl_lines="11-20 29-36"
     --8<-- "examples/event_handler_rest/src/customizing_api_operations.py"
     ```
 
diff --git a/examples/event_handler_rest/src/customizing_api_operations.py b/examples/event_handler_rest/src/customizing_api_operations.py
@@ -26,5 +26,17 @@ def get_todo_title(todo_id: int) -> str:
     return todo.json()["title"]
 
 
+@app.post(
+    "/todos",
+    summary="Creates a new todo item",
+    description="Creates a new todo item and returns it",
+    response_description="The created todo object",
+    status_code=201,
+    tags=["Todos"],
+)
+def create_todo(title: str) -> dict:
+    return {"id": 1, "title": title}
+
+
 def lambda_handler(event: dict, context: LambdaContext) -> dict:
     return app.resolve(event, context)
diff --git a/tests/functional/event_handler/_pydantic/test_api_gateway.py b/tests/functional/event_handler/_pydantic/test_api_gateway.py
diff --git a/tests/functional/event_handler/_pydantic/test_openapi_responses.py b/tests/functional/event_handler/_pydantic/test_openapi_responses.py
