Initial commit OpenAPI Extensions

Author: leandrodamascena

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,
     ):
         """
@@ -383,6 +384,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 +536,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 +945,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 +1005,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 +1044,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1053,6 +1062,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 +1102,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1109,6 +1120,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 +1160,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1165,6 +1178,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 +1217,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1220,6 +1235,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 +1277,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1278,6 +1295,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 +1336,7 @@ def lambda_handler(event, context):
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -1541,6 +1560,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.
@@ -1603,11 +1623,15 @@ def get_openapi_schema(
 
         info.update({field: value for field, value in optional_fields.items() if value})
 
+        if not openapi_extensions:
+            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 +1750,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
@@ -1778,6 +1803,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 +1831,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
@@ -1896,6 +1923,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 +1977,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 +2005,7 @@ def register_resolver(func: Callable):
                     operation_id,
                     include_in_schema,
                     security,
+                    openapi_extensions,
                     middlewares,
                 )
 
@@ -2489,6 +2519,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 +2528,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 +2544,7 @@ def register_route(func: Callable):
                 operation_id,
                 include_in_schema,
                 frozen_security,
+                fronzen_openapi_extensions,
             )
 
             # Collate Middleware for routes
@@ -2592,6 +2625,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 +2643,7 @@ def route(
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -102,6 +102,8 @@ def get(  # type: ignore[override]
         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 @@ def get(  # type: ignore[override]
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -137,6 +140,7 @@ def post(  # type: ignore[override]
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().post(
@@ -152,6 +156,7 @@ def post(  # type: ignore[override]
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -172,6 +177,7 @@ def put(  # type: ignore[override]
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().put(
@@ -187,6 +193,7 @@ def put(  # type: ignore[override]
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -207,6 +214,7 @@ def patch(  # type: ignore[override]
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().patch(
@@ -222,6 +230,7 @@ def patch(  # type: ignore[override]
             operation_id,
             include_in_schema,
             security,
+            openapi_extensions,
             middlewares,
         )
 
@@ -242,6 +251,7 @@ def delete(  # type: ignore[override]
         include_in_schema: bool = True,
         middlewares: Optional[List[Callable[..., Any]]] = None,
     ):
+        openapi_extensions = None
         security = None
 
         return super().delete(
@@ -257,6 +267,7 @@ def delete(  # type: ignore[override]
             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_serializer 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,