feat(event_handler): add OpenAPI extensions

aws_lambda_powertools/event_handler/api_gateway.py

@@ -323,6 +323,7 @@ def __init__(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Response]]] = None,
     ):
         """
@@ -360,6 +361,8 @@ def __init__(
             Whether or not to include this route in the OpenAPI schema
         security: List[Dict[str, List[str]]], optional
             The OpenAPI security for this route
+        openapi_extensions: Dict[str, Any], optional
+            Additional OpenAPI extensions as a dictionary.
         middlewares: Optional[List[Callable[..., Response]]]
             The list of route middlewares to be called in order.
         """
@@ -383,6 +386,7 @@ def __init__(
         self.tags = tags or []
         self.include_in_schema = include_in_schema
         self.security = security
+        self.openapi_extensions = openapi_extensions
         self.middlewares = middlewares or []
         self.operation_id = operation_id or self._generate_operation_id()
 
@@ -534,6 +538,10 @@ def _get_openapi_path(
         if self.security:
             operation["security"] = self.security
 
+        # Add OpenAPI extensions if present
+        if self.openapi_extensions:
+            operation.update(self.openapi_extensions)
+
         # Add the parameters to the OpenAPI operation
         if parameters:
             all_parameters = {(param["in"], param["name"]): param for param in parameters}
@@ -939,6 +947,7 @@ def route(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         raise NotImplementedError()
@@ -998,6 +1007,7 @@ def get(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         """Get route decorator with GET `method`
@@ -1036,6 +1046,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1053,6 +1064,7 @@ def post(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         """Post route decorator with POST `method`
@@ -1092,6 +1104,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1109,6 +1122,7 @@ def put(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         """Put route decorator with PUT `method`
@@ -1148,6 +1162,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1165,6 +1180,7 @@ def delete(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         """Delete route decorator with DELETE `method`
@@ -1203,6 +1219,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1220,6 +1237,7 @@ def patch(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable]] = None,
     ):
         """Patch route decorator with PATCH `method`
@@ -1261,6 +1279,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1278,6 +1297,7 @@ def head(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable]] = None,
     ):
         """Head route decorator with HEAD `method`
@@ -1318,6 +1338,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1541,6 +1562,7 @@ def get_openapi_schema(
         license_info: Optional["License"] = None,
         security_schemes: Optional[Dict[str, "SecurityScheme"]] = None,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
     ) -> "OpenAPI":
         """
         Returns the OpenAPI schema as a pydantic model.
@@ -1571,6 +1593,8 @@ def get_openapi_schema(
             A declaration of the security schemes available to be used in the specification.
         security: List[Dict[str, List[str]]], optional
             A declaration of which security mechanisms are applied globally across the API.
+        openapi_extensions: Dict[str, Any], optional
+            Additional OpenAPI extensions as a dictionary.
 
         Returns
         -------
@@ -1603,11 +1627,15 @@ def get_openapi_schema(
 
         info.update({field: value for field, value in optional_fields.items() if value})
 
+        if not isinstance(openapi_extensions, Dict):
+            openapi_extensions = {}
+
         output: Dict[str, Any] = {
             "openapi": openapi_version,
             "info": info,
             "servers": self._get_openapi_servers(servers),
             "security": self._get_openapi_security(security, security_schemes),
+            **openapi_extensions,
         }
 
         components: Dict[str, Dict[str, Any]] = {}
@@ -1726,6 +1754,7 @@ def get_openapi_json_schema(
         license_info: Optional["License"] = None,
         security_schemes: Optional[Dict[str, "SecurityScheme"]] = None,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
     ) -> str:
         """
         Returns the OpenAPI schema as a JSON serializable dict
@@ -1756,6 +1785,8 @@ def get_openapi_json_schema(
             A declaration of the security schemes available to be used in the specification.
         security: List[Dict[str, List[str]]], optional
             A declaration of which security mechanisms are applied globally across the API.
+        openapi_extensions: Dict[str, Any], optional
+            Additional OpenAPI extensions as a dictionary.
 
         Returns
         -------
@@ -1778,6 +1809,7 @@ def get_openapi_json_schema(
                 license_info=license_info,
                 security_schemes=security_schemes,
                 security=security,
+                openapi_extensions=openapi_extensions,
             ),
             by_alias=True,
             exclude_none=True,
@@ -1805,6 +1837,7 @@ def enable_swagger(
         security: Optional[List[Dict[str, List[str]]]] = None,
         oauth2_config: Optional["OAuth2Config"] = None,
         persist_authorization: bool = False,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
     ):
         """
         Returns the OpenAPI schema as a JSON serializable dict
@@ -1847,6 +1880,8 @@ def enable_swagger(
             The OAuth2 configuration for the Swagger UI.
         persist_authorization: bool, optional
             Whether to persist authorization data on browser close/refresh.
+        openapi_extensions: Dict[str, Any], optional
+            Additional OpenAPI extensions as a dictionary.
         """
         from aws_lambda_powertools.event_handler.openapi.compat import model_json
         from aws_lambda_powertools.event_handler.openapi.models import Server
@@ -1896,6 +1931,7 @@ def swagger_handler():
                 license_info=license_info,
                 security_schemes=security_schemes,
                 security=security,
+                openapi_extensions=openapi_extensions,
             )
 
             # The .replace('</', '<\\/') part is necessary to prevent a potential issue where the JSON string contains
@@ -1949,6 +1985,7 @@ def route(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         """Route decorator includes parameter `method`"""
@@ -1976,6 +2013,7 @@ def register_resolver(func: Callable):
                     operation_id,
                     include_in_schema,
                     security,
+                    openapi_extensions,
                     middlewares,
                 )
 
@@ -2489,6 +2527,7 @@ def route(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         def register_route(func: Callable):
@@ -2497,6 +2536,7 @@ def register_route(func: Callable):
             frozen_responses = _FrozenDict(responses) if responses else None
             frozen_tags = frozenset(tags) if tags else None
             frozen_security = _FrozenListDict(security) if security else None
+            fronzen_openapi_extensions = _FrozenDict(openapi_extensions) if openapi_extensions else None
 
             route_key = (
                 rule,
@@ -2512,6 +2552,7 @@ def register_route(func: Callable):
                 operation_id,
                 include_in_schema,
                 frozen_security,
+                fronzen_openapi_extensions,
             )
 
             # Collate Middleware for routes
@@ -2592,6 +2633,7 @@ def route(
         operation_id: Optional[str] = None,
         include_in_schema: bool = True,
         security: Optional[List[Dict[str, List[str]]]] = None,
+        openapi_extensions: Optional[Dict[str, Any]] = None,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
         # NOTE: see #1552 for more context.
@@ -2609,6 +2651,7 @@ def route(
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -102,6 +102,8 @@
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+
+        openapi_extensions = None
         security = None
 
         return super(BedrockAgentResolver, self).get(
@@ -117,6 +119,7 @@
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -137,6 +140,7 @@
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().post(
@@ -152,6 +156,7 @@
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -172,6 +177,7 @@
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().put(
@@ -187,6 +193,7 @@
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -207,6 +214,7 @@
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().patch(
@@ -222,6 +230,7 @@
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -242,6 +251,7 @@
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().delete(
@@ -257,6 +267,7 @@
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 

aws_lambda_powertools/event_handler/openapi/compat.py

@@ -41,7 +41,7 @@
 RequestErrorModel: Type[BaseModel] = create_model("Request")
 
 if PYDANTIC_V2:  # pragma: no cover # false positive; dropping in v3
-    from pydantic import TypeAdapter, ValidationError
+    from pydantic import TypeAdapter, ValidationError, model_validator as parser_openapi_extension
     from pydantic._internal._typing_extra import eval_type_lenient
     from pydantic.fields import FieldInfo
     from pydantic._internal._utils import lenient_issubclass
@@ -217,7 +217,7 @@ def model_json(model: BaseModel, **kwargs: Any) -> Any:
         return model.model_dump_json(**kwargs)
 
 else:
-    from pydantic import BaseModel, ValidationError
+    from pydantic import BaseModel, ValidationError, root_validator as parser_openapi_extension
     from pydantic.fields import (
         ModelField,
         Required,