Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(event_handler): add exception handling mechanism for AppSyncResolver #5588

Merged
merged 4 commits into from
Nov 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 57 additions & 6 deletions aws_lambda_powertools/event_handler/appsync.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ def __init__(self):
"""
super().__init__()
self.context = {} # early init as customers might add context before event resolution
self._exception_handlers: dict[type, Callable] = {}

def __call__(
self,
Expand Down Expand Up @@ -142,12 +143,18 @@ def lambda_handler(event, context):
self.lambda_context = context
Router.lambda_context = context

if isinstance(event, list):
Router.current_batch_event = [data_model(e) for e in event]
response = self._call_batch_resolver(event=event, data_model=data_model)
else:
Router.current_event = data_model(event)
response = self._call_single_resolver(event=event, data_model=data_model)
try:
if isinstance(event, list):
Router.current_batch_event = [data_model(e) for e in event]
response = self._call_batch_resolver(event=event, data_model=data_model)
else:
Router.current_event = data_model(event)
response = self._call_single_resolver(event=event, data_model=data_model)
except Exception as exp:
response_builder = self._lookup_exception_handler(type(exp))
if response_builder:
return response_builder(exp)
raise

# We don't clear the context for coroutines because we don't have control over the event loop.
# If we clean the context immediately, it might not be available when the coroutine is actually executed.
Expand Down Expand Up @@ -470,3 +477,47 @@ def async_batch_resolver(
raise_on_error=raise_on_error,
aggregate=aggregate,
)

def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
"""
A decorator function that registers a handler for one or more exception types.

Parameters
----------
exc_class (type[Exception] | list[type[Exception]])
A single exception type or a list of exception types.

Returns
-------
Callable:
A decorator function that registers the exception handler.
"""

def register_exception_handler(func: Callable):
if isinstance(exc_class, list): # pragma: no cover
for exp in exc_class:
self._exception_handlers[exp] = func
else:
self._exception_handlers[exc_class] = func
return func

return register_exception_handler

def _lookup_exception_handler(self, exp_type: type) -> Callable | None:
"""
Looks up the registered exception handler for the given exception type or its base classes.

Parameters
----------
exp_type (type):
The exception type to look up the handler for.

Returns
-------
Callable | None:
The registered exception handler function if found, otherwise None.
"""
for cls in exp_type.__mro__:
if cls in self._exception_handlers:
return self._exception_handlers[cls]
return None
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@ def put_artifact(self, artifact_name: str, body: Any, content_type: str) -> None
key = artifact.location.s3_location.key

# boto3 doesn't support None to omit the parameter when using ServerSideEncryption and SSEKMSKeyId
# So we are using if/else instead.
# So we are using if/else instead.

if self.data.encryption_key:

Expand Down
13 changes: 13 additions & 0 deletions docs/core/event_handler/appsync.md
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,19 @@ You can use `append_context` when you want to share data between your App and Ro
--8<-- "examples/event_handler_graphql/src/split_operation_append_context_module.py"
```

### Exception handling

You can use **`exception_handler`** decorator with any Python exception. This allows you to handle a common exception outside your resolver, for example validation errors.

The `exception_handler` function also supports passing a list of exception types you wish to handle with one handler.

```python hl_lines="5-7 11" title="Exception handling"
--8<-- "examples/event_handler_graphql/src/exception_handling_graphql.py"
```

???+ warning
This is not supported when using async single resolvers.

### Batch processing

```mermaid
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from aws_lambda_powertools.event_handler import AppSyncResolver

app = AppSyncResolver()


@app.exception_handler(ValueError)
def handle_value_error(ex: ValueError):
return {"message": "error"}


@app.resolver(field_name="createSomething")
def create_something():
raise ValueError("Raising an exception")


def lambda_handler(event, context):
return app.resolve(event, context)
Original file line number Diff line number Diff line change
Expand Up @@ -981,3 +981,125 @@ async def get_user(event: List) -> List:
# THEN the resolver must be able to return a field in the batch_current_event
assert app.context == {}
assert ret[0] == "powertools"


def test_exception_handler_with_batch_resolver_and_raise_exception():

# GIVEN a AppSyncResolver instance
app = AppSyncResolver()

event = [
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": "1",
},
},
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": "2",
},
},
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": [3, 4],
},
},
]

# WHEN we configure exception handler for ValueError
@app.exception_handler(ValueError)
def handle_value_error(ex: ValueError):
return {"message": "error"}

# WHEN the sync batch resolver for the 'listLocations' field is defined with raise_on_error=True
@app.batch_resolver(field_name="listLocations", raise_on_error=True, aggregate=False)
def create_something(event: AppSyncResolverEvent) -> Optional[list]: # noqa AA03 VNE003
raise ValueError

# Call the implicit handler
result = app(event, {})

# THEN the return must be the Exception Handler error message
assert result["message"] == "error"


def test_exception_handler_with_batch_resolver_and_no_raise_exception():

# GIVEN a AppSyncResolver instance
app = AppSyncResolver()

event = [
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": "1",
},
},
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": "2",
},
},
{
"typeName": "Query",
"info": {
"fieldName": "listLocations",
"parentTypeName": "Post",
},
"fieldName": "listLocations",
"arguments": {},
"source": {
"id": [3, 4],
},
},
]

# WHEN we configure exception handler for ValueError
@app.exception_handler(ValueError)
def handle_value_error(ex: ValueError):
return {"message": "error"}

# WHEN the sync batch resolver for the 'listLocations' field is defined with raise_on_error=False
@app.batch_resolver(field_name="listLocations", raise_on_error=False, aggregate=False)
def create_something(event: AppSyncResolverEvent) -> Optional[list]: # noqa AA03 VNE003
raise ValueError

# Call the implicit handler
result = app(event, {})

# THEN the return must not trigger the Exception Handler, but instead return from the resolver
assert result == [None, None, None]
Original file line number Diff line number Diff line change
Expand Up @@ -329,3 +329,25 @@ async def get_async():
# THEN
assert asyncio.run(result) == "value"
assert app.context == {}


def test_exception_handler_with_single_resolver():
# GIVEN a AppSyncResolver instance
mock_event = load_event("appSyncDirectResolver.json")

app = AppSyncResolver()

# WHEN we configure exception handler for ValueError
@app.exception_handler(ValueError)
def handle_value_error(ex: ValueError):
return {"message": "error"}

@app.resolver(field_name="createSomething")
def create_something(id: str): # noqa AA03 VNE003
raise ValueError("Error")

# Call the implicit handler
result = app(mock_event, {})

# THEN the return must be the Exception Handler error message
assert result["message"] == "error"
Loading