small changes

Author: leandrodamascena

aws_lambda_powertools/event_handler/openapi/compat.py

@@ -4,9 +4,6 @@
 from collections import deque
 from collections.abc import Mapping, Sequence
 from copy import copy
-
-# MAINTENANCE: remove when deprecating Pydantic v1. Mypy doesn't handle two different code paths that import different
-# versions of a module, so we need to ignore errors here.
 from dataclasses import dataclass, is_dataclass
 from typing import TYPE_CHECKING, Any, Deque, FrozenSet, List, Set, Tuple, Union
 

aws_lambda_powertools/event_handler/openapi/params.py

@@ -1042,7 +1042,7 @@ def _has_discriminator(field_info: FieldInfo) -> bool:
     return hasattr(field_info, "discriminator") and field_info.discriminator is not None
 
 
-def _handle_discriminator_with_body(
+def _handle_discriminator_with_param(
     annotations: list[FieldInfo],
     annotation: Any,
 ) -> tuple[FieldInfo | None, Any, bool]:
@@ -1111,10 +1111,10 @@ def get_field_info_annotated_type(annotation, value, is_path_param: bool) -> tup
 
     # Determine which annotation to use
     powertools_annotation: FieldInfo | None = None
-    has_discriminator_with_body = False
+    has_discriminator_with_param = False
 
     if len(powertools_annotations) == 2:
-        powertools_annotation, type_annotation, has_discriminator_with_body = _handle_discriminator_with_body(
+        powertools_annotation, type_annotation, has_discriminator_with_param = _handle_discriminator_with_param(
             powertools_annotations,
             annotation,
         )
@@ -1126,7 +1126,7 @@ def get_field_info_annotated_type(annotation, value, is_path_param: bool) -> tup
     # Process the annotation if it exists
     field_info: FieldInfo | None = None
     if isinstance(powertools_annotation, FieldInfo):
-        field_info = _create_field_info(powertools_annotation, type_annotation, has_discriminator_with_body)
+        field_info = _create_field_info(powertools_annotation, type_annotation, has_discriminator_with_param)
         _set_field_default(field_info, value, is_path_param)
 
         # Preserve full annotated type for discriminated unions

docs/core/event_handler/api_gateway.md

@@ -428,6 +428,17 @@ We use the `Annotated` and OpenAPI `Body` type to instruct Event Handler that ou
     --8<-- "examples/event_handler_rest/src/validating_payload_subset_output.json"
     ```
 
+##### Discriminated unions
+
+You can use Pydantic's `Field(discriminator="...")` with union types to create discriminated unions (also known as tagged unions). This allows the Event Handler to automatically determine which model to use based on a discriminator field in the request body.
+
+```python hl_lines="3 4 8 31 36" title="discriminated_unions.py"
+--8<-- "examples/event_handler_rest/src/discriminated_unions.py"
+```
+
+1. `Field(discriminator="action")` tells Pydantic to use the `action` field to determine which model to instantiate
+2. `Body()` annotation tells the Event Handler to parse the request body using the discriminated union
+
 #### Validating responses
 
 You can use `response_validation_error_http_code` to set a custom HTTP code for failed response validation. When this field is set, we will raise a `ResponseValidationError` instead of a `RequestValidationError`.
@@ -568,23 +579,6 @@ You can use the `Form` type to tell the Event Handler that a parameter expects f
     --8<-- "examples/event_handler_rest/src/working_with_form_data.py"
     ```
 
-#### Discriminated unions
-
-!!! info "You must set `enable_validation=True` to use discriminated unions via type annotation."
-
-You can use Pydantic's `Field(discriminator="...")` with union types to create discriminated unions (also known as tagged unions). This allows the Event Handler to automatically determine which model to use based on a discriminator field in the request body.
-
-In the following example, we define two action types (`FooAction` and `BarAction`) that share a common discriminator field `action`. The Event Handler will automatically parse the request body and instantiate the correct model based on the `action` field value:
-
-```python hl_lines="3 4 8 31 36" title="discriminated_unions.py"
---8<-- "examples/event_handler_rest/src/discriminated_unions.py"
-```
-
-1. `Field(discriminator="action")` tells Pydantic to use the `action` field to determine which model to instantiate
-2. `Body()` annotation tells the Event Handler to parse the request body using the discriminated union
-
-When you send a request with `{"action": "foo", "foo_data": "example"}`, the Event Handler will automatically create a `FooAction` instance. Similarly, `{"action": "bar", "bar_data": 42}` will create a `BarAction` instance.
-
 #### Supported types for response serialization
 
 With data validation enabled, we natively support serializing the following data types to JSON:

tests/e2e/event_handler/handlers/data_validation_with_fields.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from aws_lambda_powertools.event_handler import APIGatewayRestResolver
+from aws_lambda_powertools.event_handler.openapi.params import Body
+
+app = APIGatewayRestResolver(enable_validation=True)
+app.enable_swagger()
+
+
+class FooAction(BaseModel):
+    action: Literal["foo"]
+    foo_data: str
+
+
+class BarAction(BaseModel):
+    action: Literal["bar"]
+    bar_data: int
+
+
+Action = Annotated[FooAction | BarAction, Field(discriminator="action")]
+
+
+@app.post("/data_validation_with_fields")
+def create_action(action: Annotated[Action, Body(discriminator="action")]):
+    return {"message": "Powertools e2e API"}
+
+
+def lambda_handler(event, context):
+    print(event)
+    return app.resolve(event, context)

tests/e2e/event_handler/infrastructure.py

@@ -24,6 +24,7 @@ def create_resources(self):
                 functions["OpenapiHandler"],
                 functions["OpenapiHandlerWithPep563"],
                 functions["DataValidationAndMiddleware"],
+                functions["DataValidationWithFields"],
             ],
         )
         self._create_api_gateway_http(function=functions["ApiGatewayHttpHandler"])
@@ -105,6 +106,9 @@ def _create_api_gateway_rest(self, function: list[Function]):
         openapi_schema = apigw.root.add_resource("data_validation_middleware")
         openapi_schema.add_method("GET", apigwv1.LambdaIntegration(function[3], proxy=True))
 
+        openapi_schema = apigw.root.add_resource("data_validation_with_fields")
+        openapi_schema.add_method("POST", apigwv1.LambdaIntegration(function[4], proxy=True))
+
         CfnOutput(self.stack, "APIGatewayRestUrl", value=apigw.url)
 
     def _create_lambda_function_url(self, function: Function):

tests/e2e/event_handler/test_openapi.py

@@ -59,3 +59,20 @@ def test_get_openapi_validation_and_middleware(apigw_rest_endpoint):
     )
 
     assert response.status_code == 202
+
+
+def test_openapi_with_fields_discriminator(apigw_rest_endpoint):
+    # GIVEN
+    url = f"{apigw_rest_endpoint}data_validation_with_fields"
+
+    # WHEN
+    response = data_fetcher.get_http_response(
+        Request(
+            method="POST",
+            url=url,
+            json={"action": "foo", "foo_data": "foo data working"},
+        ),
+    )
+
+    assert "Powertools e2e API" in response.text
+    assert response.status_code == 200

tests/functional/event_handler/_pydantic/test_openapi_validation_middleware.py

@@ -2030,35 +2030,3 @@ def create_action(action: Annotated[action_type, Body()]):
 
     result = app(gw_event, {})
     assert result["statusCode"] == 422
-
-
-def test_field_other_features_still_work(gw_event):
-    """Test that other Field features still work after discriminator fix"""
-    app = APIGatewayRestResolver(enable_validation=True)
-
-    class UserInput(BaseModel):
-        name: Annotated[str, Field(min_length=2, max_length=50, description="User name")]
-        age: Annotated[int, Field(ge=18, le=120, description="User age")]
-        email: Annotated[str, Field(pattern=r".+@.+\..+", description="User email")]
-
-    @app.post("/users")
-    def create_user(user: UserInput):
-        return {"created": user.model_dump()}
-
-    gw_event["path"] = "/users"
-    gw_event["httpMethod"] = "POST"
-    gw_event["headers"]["content-type"] = "application/json"
-    gw_event["body"] = '{"name": "John", "age": 25, "email": "[email protected]"}'
-
-    result = app(gw_event, {})
-    assert result["statusCode"] == 200
-
-    response_body = json.loads(result["body"])
-    assert response_body["created"]["name"] == "John"
-    assert response_body["created"]["age"] == 25
-    assert response_body["created"]["email"] == "[email protected]"
-
-    gw_event["body"] = '{"name": "John", "age": 16, "email": "[email protected]"}'
-
-    result = app(gw_event, {})
-    assert result["statusCode"] == 422