Bug: Using `discriminator` field with TypeAdapter instances

[xkortex]

Use case

I happened upon using the discriminator to allow my lambdas to handle multiple different event sources, e.g. aws.events and aws.s3. Using Pydantic's TypeAdapter and Field classes, we can create an adapter which looks at a specific field (in this case source) to determine how to parse the data, without having to do the expensive try-validate-fail-try_next_model pattern. This allows us to decorate a lambda in such a way that it can automatically cast different one of several event types to stricter subtypes

EventBridgeSource = Annotated[
    Union[S3EventNotificationEventBridgeModel, ScheduledNotificationEventBridgeModel],
    Field(discriminator="source"),
]

EventBridgeModelAdapter = TypeAdapter(Union[EventBridgeSource, EventBridgeModel])

Solution/User Experience

This fails with the stock event_parser so I had to bodge event_parser to get it to work. I think the issue is trying to cast the input TypeAdapter into a TypeAdapter again, throwing PydanticSchemaGenerationError. It should be straightforward to allow _retrieve_or_set_model_from_cache to check if the input is already a TypeAdapter and return that from the cache rather than trying to wrap it.

Setup

from typing import Any, Callable, Literal, Optional, TypeVar, Union

from typing_extensions import Annotated

from pydantic import BaseModel, Field, ValidationError, TypeAdapter

# from aws_lambda_powertools.utilities.parser import event_parser
from aws_lambda_powertools.utilities.parser.envelopes.base import Envelope

from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.parser.models.event_bridge import EventBridgeModel
from aws_lambda_powertools.utilities.parser.models.s3 import (
    S3EventNotificationEventBridgeDetailModel,
)


T = TypeVar("T")


class S3EventNotificationEventBridgeModel(EventBridgeModel):
    detail: S3EventNotificationEventBridgeDetailModel
    source: Literal["aws.s3"]


class ScheduledNotificationEventBridgeModel(EventBridgeModel):
    source: Literal["aws.events"]


EventBridgeSource = Annotated[
    Union[S3EventNotificationEventBridgeModel, ScheduledNotificationEventBridgeModel],
    Field(discriminator="source"),
]

EventBridgeModelAdapter = TypeAdapter(Union[EventBridgeSource, EventBridgeModel])


class LambdaContext:
    function_name: str = "test-function"
    memory_limit_in_mb: int = 128
    invoked_function_arn: str = f"arn:aws:lambda:us-east-1:123456789012:test-function"
    aws_request_id: str = 'e7a150fa-ed3d-4fbb-8158-e470b60621da'

@lambda_handler_decorator
def event_parser(
    handler: Callable[..., 'EventParserReturnType'],
    event: dict[str, Any],
    context: LambdaContext = None,
    model: Optional[type[T]] = None,
    envelope: Optional[type[Envelope]] = None,
    **kwargs: Any,
) -> 'EventParserReturnType':
    parsed_event = EventBridgeModelAdapter.validate_python(event)
    return handler(parsed_event, context, **kwargs)


@event_parser(model=EventBridgeModelAdapter)
def handler(event: EventBridgeModelAdapter, context = None):
    print(f'Got event, parsed into {type(event)=}')

Testing

object_event = {
  "version": "0",
  "id": "55ec5937-9835-fb01-3122-a054f478882b",
  "detail-type": "Object Created",
  "source": "aws.s3",
  "account": "539101081866",
  "time": "2024-10-16T22:58:07Z",
  "region": "us-east-1",
  "resources": [
    "arn:aws:s3:::bpa-sandbox-mcdermott"
  ],
  "detail": {
    "version": "0",
    "bucket": {
      "name": "bpa-sandbox-mcdermott"
    },
    "object": {
      "key": "testbed/hello3.txt",
      "size": 12,
    },
    "requester": "539101081866",
  }
}

scheduled_event = {
  "version": "0",
  "id": "ac8bd595-adcc-cd98-cc06-5e8fad5d9565",
  "detail-type": "Scheduled Event",
  "source": "aws.events",
  "account": "084286224071",
  "time": "2024-10-23T09:00:00Z",
  "region": "us-east-1",
  "resources": [
    "arn:aws:events:us-east-1:084286224071:rule/rl-live-StatesIntegration-RuleDailyDownloadWorkflo-1KZH76DGASA21"
  ],
  "detail": {}
}

fallback_event = {
  "version": "0",
  "id": "ac8bd595-adcc-cd98-cc06-5e8fad5d9565",
  "detail-type": "Scheduled Event",
  "source": "aws.other",
  "account": "084286224071",
  "time": "2024-10-23T09:00:00Z",
  "region": "us-east-1",
  "resources": [
    "arn:aws:events:us-east-1:084286224071:rule/rl-live-StatesIntegration-RuleDailyDownloadWorkflo-1KZH76DGASA21"
  ],
  "detail": {}
}

handler(object_event, LambdaContext())
handler(scheduled_event, LambdaContext())
handler(fallback_event, LambdaContext())

results

Got event, parsed into type(event)=<class '__main__.S3EventNotificationEventBridgeModel'>
Got event, parsed into type(event)=<class '__main__.ScheduledNotificationEventBridgeModel'>
Got event, parsed into type(event)=<class 'aws_lambda_powertools.utilities.parser.models.event_bridge.EventBridgeModel'>

Alternative solutions

Not sure how else to implement this.

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

[xkortex]

Your _retrieve_or_set_model_from_cache chokes on being passed aTypeAdapter, but the fix is simple.

```python In [13]: EventBridgeModelAdapter = TypeAdapter(Union[EventBridgeSource, EventBridgeModel])

In [14]: _retrieve_or_set_model_from_cache(EventBridgeModelAdapter)

AttributeError Traceback (most recent call last)
File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:270, in TypeAdapter._init_core_attrs(self, rebuild_mocks)
269 try:
--> 270 self._core_schema = _getattr_no_parents(self._type, 'pydantic_core_schema')
271 self._validator = _getattr_no_parents(self._type, 'pydantic_validator')

File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:112, in _getattr_no_parents(obj, attribute)
111 else:
--> 112 raise AttributeError(attribute)

AttributeError: pydantic_core_schema

During handling of the above exception, another exception occurred:

PydanticSchemaGenerationError Traceback (most recent call last)
Cell In[14], line 1
----> 1 _retrieve_or_set_model_from_cache(EventBridgeModelAdapter)

File /usr/local/lib/python3.11/site-packages/aws_lambda_powertools/utilities/parser/functions.py:43, in _retrieve_or_set_model_from_cache(model)
40 if id_model in CACHE_TYPE_ADAPTER:
41 return CACHE_TYPE_ADAPTER[id_model]
---> 43 CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)
44 return CACHE_TYPE_ADAPTER[id_model]

File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:257, in TypeAdapter.init(self, type, config, _parent_depth, module)
252 if not self._defer_build():
253 # Immediately initialize the core schema, validator and serializer
254 with self._with_frame_depth(1): # +1 frame depth for this init
255 # Model itself may be using deferred building. For backward compatibility we don't rebuild model mocks
256 # here as part of init even though TypeAdapter itself is not using deferred building.
--> 257 self._init_core_attrs(rebuild_mocks=False)

File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:135, in _frame_depth..wrapper..wrapped(self, *args, **kwargs)
132 @wraps(func)
133 def wrapped(self: TypeAdapterT, *args: P.args, **kwargs: P.kwargs) -> R:
134 with self._with_frame_depth(depth + 1): # depth + 1 for the wrapper function
--> 135 return func(self, *args, **kwargs)

File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:277, in TypeAdapter._init_core_attrs(self, rebuild_mocks)
274 config_wrapper = _config.ConfigWrapper(self._config)
275 core_config = config_wrapper.core_config(None)
--> 277 self._core_schema = _get_schema(self._type, config_wrapper, parent_depth=self._parent_depth)
278 self._validator = create_schema_validator(
279 schema=self._core_schema,
280 schema_type=self._type,
(...)
285 plugin_settings=config_wrapper.plugin_settings,
286 )
287 self._serializer = SchemaSerializer(self._core_schema, core_config)

File /usr/local/lib/python3.11/site-packages/pydantic/type_adapter.py:95, in get_schema(type, config_wrapper, parent_depth)
91 global_ns.update(local_ns or {})
92 gen = (config_wrapper.schema_generator or generate_schema.GenerateSchema)(
93 config_wrapper, types_namespace=global_ns, typevars_map={}
94 )
---> 95 schema = gen.generate_schema(type)
96 schema = gen.clean_schema(schema)
97 return schema

File /usr/local/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py:655, in GenerateSchema.generate_schema(self, obj, from_dunder_get_core_schema)
652 schema = from_property
654 if schema is None:
--> 655 schema = self._generate_schema_inner(obj)
657 metadata_js_function = _extract_get_pydantic_json_schema(obj, schema)
658 if metadata_js_function is not None:

File /usr/local/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py:929, in GenerateSchema._generate_schema_inner(self, obj)
926 if isinstance(obj, PydanticRecursiveRef):
927 return core_schema.definition_reference_schema(schema_ref=obj.type_ref)
--> 929 return self.match_type(obj)

File /usr/local/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py:1038, in GenerateSchema.match_type(self, obj)
1036 if self._arbitrary_types:
1037 return self._arbitrary_type_schema(obj)
-> 1038 return self._unknown_type_schema(obj)

File /usr/local/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py:558, in GenerateSchema._unknown_type_schema(self, obj)
557 def _unknown_type_schema(self, obj: Any) -> CoreSchema:
--> 558 raise PydanticSchemaGenerationError(
559 f'Unable to generate pydantic-core schema for {obj!r}. '
560 'Set arbitrary_types_allowed=True in the model_config to ignore this error'
561 ' or implement __get_pydantic_core_schema__ on your type to fully support it.'
562 '\n\nIf you got this error by calling handler() within'
563 ' __get_pydantic_core_schema__ then you likely need to call'
564 ' handler.generate_schema(<some type>) since we do not call'
565 ' __get_pydantic_core_schema__ on <some type> otherwise to avoid infinite recursion.'
566 )

PydanticSchemaGenerationError: Unable to generate pydantic-core schema for <pydantic.type_adapter.TypeAdapter object at 0x7ffffc34ff10>. Set arbitrary_types_allowed=True in the model_config to ignore this error or implement __get_pydantic_core_schema__ on your type to fully support it.

If you got this error by calling handler() within __get_pydantic_core_schema__ then you likely need to call handler.generate_schema(<some type>) since we do not call __get_pydantic_core_schema__ on <some type> otherwise to avoid infinite recursion.

For further information visit https://errors.pydantic.dev/2.9/u/schema-for-unknown-type

</details>

### the fix 
<details>
```python
from aws_lambda_powertools.shared.cache_dict import LRUDict

from aws_lambda_powertools.utilities.parser.types import T

CACHE_TYPE_ADAPTER = LRUDict(max_items=1024)


def _retrieve_or_set_model_from_cache(model: type[T]) -> TypeAdapter:

    id_model = id(model)

    if id_model in CACHE_TYPE_ADAPTER:
        return CACHE_TYPE_ADAPTER[id_model]

    if isinstance(model, TypeAdapter):
        CACHE_TYPE_ADAPTER[id_model] = model
    else:
        CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)
    return CACHE_TYPE_ADAPTER[id_model]

[leandrodamascena]

Hi @xkortex! Thanks a lot for reporting this bug, we really need to fix this bug when passing a TypeAdapter to parser instead of generic models. But just to give more context here, we already support the discriminator field when using Union tags, that's why in Powertools v3 we switched from direct validation (model_validate) to TypeAdapter, you can see the PR here. It also open the door for customer uses Dataclass and other types of primitives.

In theory, you don't need to use TypeAdapter in your code because you are using classes that inherit from BaseModel, so there is no benefit in converting them to TypeAdapter before sending to the parser. I mean, even on the performance side, it's pretty much the same execution time as you would get without doing this (I did a performance test here), this code should work:

class S3EventNotificationEventBridgeModel(EventBridgeModel):
    detail: S3EventNotificationEventBridgeDetailModel
    source: Literal["aws.s3"]


class ScheduledNotificationEventBridgeModel(EventBridgeModel):
    source: Literal["aws.events"]


EventBridgeModelAdapter = Annotated[
    Union[S3EventNotificationEventBridgeModel, ScheduledNotificationEventBridgeModel, EventBridgeModel],
    Field(discriminator="source"),
]

@event_parser(model=EventBridgeModelAdapter)
def handler(event: EventBridgeModelAdapter, context = None):
    print(f'Got event, parsed into {type(event)=}')

Another thing is that considering you are using the code you sent here, you should not use this event: EventBridgeModelAdapter because you cannot use TypeAdapter instances as a type annotation.

That said, would you be available to submit a PR to fix this? I'd love to have your contribution. The solution you sent works perfectly.