Feature request: event_parser envelopes should handle unions of BaseModels

[dcheno]

Use case

Pydantic supports parsing unions of BaseModel types. This can be pretty handy when a lambda might accept a few different event types.

Example: a callback handler lambda might process an error response or a success response type. These might have different fields.

The current implementation of BaseEnvelope will error if you pass a Union or Annotated Union to the model field. This is because neither of those types have a parse_raw or parse_obj method which BaseModel expects.

Solution/User Experience

I believe this can mostly be taken care of by switching from using model.parse_raw and model.parse_obj here in BaseEnvelope to use pydantic.parse_obj_as and pydantic.parse_raw_as (available from pydantic 1.7 onwards so should match this projects dependencies). Docs here.

Happy to work on a PR, just wanted to check in here before I did too much work.

Alternative solutions

1. I've gotten around this in the past by just parsing the parent type in `event_parser`. For example, parsing `EventBridgeModel`, then parsing the `detail` myself using `parse_obj_as`.
2. This can also be solved on a case-by-case basis by rearchitecting the event structure so that the Union takes place within the model being parsed, but changing the event fields isn't always an option in existing codebases.

Acknowledgment

• This feature request meets Powertools for AWS Lambda (Python) Tenets
• Should this be considered in other Powertools for AWS Lambda languages? i.e. Java, TypeScript, and .NET

[boring-cyborg]

Thanks for opening your first issue here! We'll come back to you as soon as we can.
In the meantime, check out the #python channel on our Powertools for AWS Lambda Discord: Invite link

[leandrodamascena]

Looking at this now.

[leandrodamascena]

Hi @dcheno! Do you have a snippet of the code you are using to reproduce this feature request? I understand your point about our current codebase, but I'm trying to think of situations where this might happen.

Thank you

[dcheno]

hey @leandrodamascena , thanks for taking a look.

Sure, here's the kind of thing I run into occassionaly:

from aws_lambda_powertools.utilities.parser import event_parser
from pydantic import BaseModel, Field
from typing import Annotated, Any, Literal, Union


class SuccessfulCallback(BaseModel):
    status: Literal["succeeded"]
    name: str
    breed: Literal["Husky", "Labrador"]


class FailedCallback(BaseModel):
    status: Literal["failed"]
    error: str


DogCallback = Annotated[
    Union[SuccessfulCallback, FailedCallback], Field(discriminator="status")
]


@event_parser(model=DogCallback)
def lambda_handler(event: DogCallback, _: Any) -> str:
    if isinstance(event, FailedCallback): # alternatively `if event.status == "failed"` (if your type checker is on top of things)
        return f"Uh oh. Had a problem: {event.error}"

    return f"Successfully retrieved {event.breed} named {event.name}"


def main() -> None:
    print(
        lambda_handler(
            {"status": "succeeded", "name": "Clifford", "breed": "Labrador"}, None
        )
    )


if __name__ == "__main__":
    main()

Which returns error:

Traceback (most recent call last):
  .../python3.10/typing.py", line 375, in __getattr__
    raise AttributeError(item)
AttributeError: parse_obj

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  ...site-packages/aws_lambda_powertools/utilities/parser/parser.py", line 82, in event_parser
    parsed_event = parse(event=event, model=model, envelope=envelope) if envelope else parse(event=event, model=model)
 ...aws_lambda_powertools.utilities.parser.exceptions.InvalidModelTypeError: Input model must implement BaseModel, model=typing.Annotated[typing.Union[__main__.SuccessfulCallback, __main__.FailedCallback], FieldInfo(default=PydanticUndefined, discriminator='status', extra={})]

There are definitely ways around this, but I've had lambdas like this a few times so thought it would be worth checking in.

[leandrodamascena]

Hello @dcheno! I'm back to this issue and sorry for a long time without an update.

I was looking at our base code and trying to come up with some scenarios to see how we can update it but I think, we can't do that right now. Powertools now supports Pydantic v1 and v2 and keep Pydantic functions that work in both versions without code change. I see the parse_raw_as function has been removed in Pydantic v2.

We already have a Roadmap to drop Pydanticv1 in Powertools v3 and in this new version, we will refactor the parser utility with Pydanticv2 and can add it. We are estimating (tentative) Powertools v3 for later this year/early next year.

I'll update our roadmap to add this feature to support unions of BaseModels in an Annotated type.

Thanks.