Polishing the PR with best practicies - using model_validator to be more specific

Author: leandrodamascena

aws_lambda_powertools/event_handler/api_gateway.py

@@ -1623,7 +1623,7 @@ def get_openapi_schema(
 
         info.update({field: value for field, value in optional_fields.items() if value})
 
-        if not openapi_extensions:
+        if not isinstance(openapi_extensions, Dict):
             openapi_extensions = {}
 
         output: Dict[str, Any] = {

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, model_serializer as parser_openapi_extension
+    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

aws_lambda_powertools/event_handler/openapi/models.py

@@ -13,6 +13,38 @@
 """
 
 
+class OpenapiExtensions(BaseModel):
+    openapi_extensions: Optional[Dict[str, Any]] = None
+
+    # This rule is valid for Pydantic v1 and v2
+    # If the 'openapi_extensions' field is present in the 'values' dictionary,
+    # update the 'values' dictionary with the contents of 'openapi_extensions',
+    # and then remove the 'openapi_extensions' field from the 'values' dictionary
+
+    if PYDANTIC_V2:
+        model_config = {"extra": "allow"}
+
+        @parser_openapi_extension(mode="before")
+        def serialize_openapi_extension(self):
+            if isinstance(self, dict) and self.get("openapi_extensions"):
+                self.update(self.get("openapi_extensions"))
+                self.pop("openapi_extensions", None)
+
+            return self
+
+    else:
+
+        @parser_openapi_extension(pre=False, allow_reuse=True)
+        def serialize_openapi_extension(cls, values):
+            if values.get("openapi_extensions"):
+                values.update(values["openapi_extensions"])
+                del values["openapi_extensions"]
+            return values
+
+        class Config:
+            extra = "allow"
+
+
 # https://swagger.io/specification/#contact-object
 class Contact(BaseModel):
     name: Optional[str] = None
@@ -77,33 +109,16 @@ class Config:
 
 
 # https://swagger.io/specification/#server-object
-class Server(BaseModel):
+class Server(OpenapiExtensions):
     url: Union[AnyUrl, str]
     description: Optional[str] = None
     variables: Optional[Dict[str, ServerVariable]] = None
-    openapi_extensions: Optional[Dict[str, Any]] = None
 
     if PYDANTIC_V2:
         model_config = {"extra": "allow"}
 
-        @parser_openapi_extension()
-        def serialize(self):
-            # If the 'openapi_extensions' field is not None, return it
-            if self.openapi_extensions:
-                return self.openapi_extensions
-
     else:
 
-        # If the 'openapi_extensions' field is present in the 'values' dictionary,
-        # update the 'values' dictionary with the contents of 'openapi_extensions',
-        # and then remove the 'openapi_extensions' field from the 'values' dictionary
-        @parser_openapi_extension(pre=False, allow_reuse=True)
-        def check_json(cls, values):
-            if values.get("openapi_extensions"):
-                values.update(values["openapi_extensions"])
-                del values["openapi_extensions"]
-            return values
-
         class Config:
             extra = "allow"
 
@@ -396,7 +411,7 @@ class Config:
 
 
 # https://swagger.io/specification/#operation-object
-class Operation(BaseModel):
+class Operation(OpenapiExtensions):
     tags: Optional[List[str]] = None
     summary: Optional[str] = None
     description: Optional[str] = None
@@ -410,23 +425,12 @@ class Operation(BaseModel):
     deprecated: Optional[bool] = None
     security: Optional[List[Dict[str, List[str]]]] = None
     servers: Optional[List[Server]] = None
-    openapi_extensions: Optional[Dict[str, Any]] = None
 
     if PYDANTIC_V2:
         model_config = {"extra": "allow"}
 
     else:
 
-        # If the 'openapi_extensions' field is present in the 'values' dictionary,
-        # update the 'values' dictionary with the contents of 'openapi_extensions',
-        # and then remove the 'openapi_extensions' field from the 'values' dictionary
-        @parser_openapi_extension(pre=False, allow_reuse=True)
-        def check_json(cls, values):
-            if values.get("openapi_extensions"):
-                values.update(values["openapi_extensions"])
-                del values["openapi_extensions"]
-            return values
-
         class Config:
             extra = "allow"
 
@@ -464,32 +468,15 @@ class SecuritySchemeType(Enum):
     openIdConnect = "openIdConnect"
 
 
-class SecurityBase(BaseModel):
+class SecurityBase(OpenapiExtensions):
     type_: SecuritySchemeType = Field(alias="type")
     description: Optional[str] = None
-    openapi_extensions: Optional[Dict[str, Any]] = None
 
     if PYDANTIC_V2:
         model_config = {"extra": "allow", "populate_by_name": True}
 
-        @parser_openapi_extension()
-        def serialize(self):
-            # If the 'openapi_extensions' field is not None, return it
-            if self.openapi_extensions:
-                return self.openapi_extensions
-
     else:
 
-        # If the 'openapi_extensions' field is present in the 'values' dictionary,
-        # update the 'values' dictionary with the contents of 'openapi_extensions',
-        # and then remove the 'openapi_extensions' field from the 'values' dictionary
-        @parser_openapi_extension(pre=False, allow_reuse=True)
-        def check_json(cls, values):
-            if values.get("openapi_extensions"):
-                values.update(values["openapi_extensions"])
-                del values["openapi_extensions"]
-            return values
-
         class Config:
             extra = "allow"
             allow_population_by_field_name = True
@@ -602,7 +589,7 @@ class Config:
 
 
 # https://swagger.io/specification/#openapi-object
-class OpenAPI(BaseModel):
+class OpenAPI(OpenapiExtensions):
     openapi: str
     info: Info
     jsonSchemaDialect: Optional[str] = None
@@ -614,23 +601,12 @@ class OpenAPI(BaseModel):
     security: Optional[List[Dict[str, List[str]]]] = None
     tags: Optional[List[Tag]] = None
     externalDocs: Optional[ExternalDocumentation] = None
-    openapi_extensions: Optional[Dict[str, Any]] = None
 
     if PYDANTIC_V2:
         model_config = {"extra": "allow"}
 
     else:
 
-        # If the 'openapi_extensions' field is present in the 'values' dictionary,
-        # update the 'values' dictionary with the contents of 'openapi_extensions',
-        # and then remove the 'openapi_extensions' field from the 'values' dictionary
-        @parser_openapi_extension(pre=False, allow_reuse=True)
-        def check_json(cls, values):
-            if values.get("openapi_extensions"):
-                values.update(values["openapi_extensions"])
-                del values["openapi_extensions"]
-            return values
-
         class Config:
             extra = "allow"
 

tests/functional/event_handler/_pydantic/test_openapi_extensions.py

@@ -65,6 +65,8 @@ def test_openapi_extension_server_level():
     # THEN the OpenAPI schema must contain the "x-amazon-apigateway-endpoint-configuration" at the server level
     assert "x-amazon-apigateway-endpoint-configuration" in schema["servers"][0]
     assert schema["servers"][0]["x-amazon-apigateway-endpoint-configuration"] == endpoint_config
+    assert schema["servers"][0]["url"] == server_config["url"]
+    assert schema["servers"][0]["description"] == server_config["description"]
 
 
 def test_openapi_extension_security_scheme_level_with_api_key():
@@ -102,6 +104,9 @@ def test_openapi_extension_security_scheme_level_with_api_key():
     assert "x-amazon-apigateway-authtype" in schema["components"]["securitySchemes"]["apiKey"]
     assert schema["components"]["securitySchemes"]["apiKey"]["x-amazon-apigateway-authtype"] == "custom"
     assert schema["components"]["securitySchemes"]["apiKey"]["x-amazon-apigateway-authorizer"] == authorizer_config
+    assert schema["components"]["securitySchemes"]["apiKey"]["name"] == api_key_config["name"]
+    assert schema["components"]["securitySchemes"]["apiKey"]["description"] == api_key_config["description"]
+    assert schema["components"]["securitySchemes"]["apiKey"]["in"] == "header"
 
 
 def test_openapi_extension_security_scheme_level_with_oauth2():
@@ -142,6 +147,19 @@ def test_openapi_extension_security_scheme_level_with_oauth2():
     # THEN the OpenAPI schema must contain the "x-amazon-apigateway-authorizer" extension at the security scheme level
     assert "x-amazon-apigateway-authorizer" in schema["components"]["securitySchemes"]["oauth2"]
     assert schema["components"]["securitySchemes"]["oauth2"]["x-amazon-apigateway-authorizer"] == authorizer_config
+    assert (
+        schema["components"]["securitySchemes"]["oauth2"]["x-amazon-apigateway-authorizer"]["identitySource"]
+        == "$request.header.Authorization"
+    )
+    assert schema["components"]["securitySchemes"]["oauth2"]["x-amazon-apigateway-authorizer"]["jwtConfiguration"][
+        "audience"
+    ] == ["test"]
+    assert (
+        schema["components"]["securitySchemes"]["oauth2"]["x-amazon-apigateway-authorizer"]["jwtConfiguration"][
+            "issuer"
+        ]
+        == "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_xxxxx/"
+    )
 
 
 def test_openapi_extension_operation_level(openapi_extension_integration_detail):
@@ -159,6 +177,7 @@ def lambda_handler():
     # THEN the OpenAPI schema must contain the "x-amazon-apigateway-integration" extension at the operation level
     assert "x-amazon-apigateway-integration" in schema["paths"]["/test"]["get"]
     assert schema["paths"]["/test"]["get"]["x-amazon-apigateway-integration"] == openapi_extension_integration_detail
+    assert schema["paths"]["/test"]["get"]["operationId"] == "lambda_handler_test_get"
 
 
 def test_openapi_extension_operation_level_multiple_paths(