Skip to content

Commit

Permalink
feat(event_handler): add exception handling mechanism for AppSyncReso…
Browse files Browse the repository at this point in the history
…lver (#5588)

* Adding exception handler support

* Adding exception handler support - fix tests

* Adding exception handler support - fix tests + docs
  • Loading branch information
leandrodamascena authored Nov 19, 2024
1 parent 3ff5132 commit db8a338
Show file tree
Hide file tree
Showing 6 changed files with 232 additions and 7 deletions.
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
17 changes: 17 additions & 0 deletions examples/event_handler_graphql/src/exception_handling_graphql.py
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"

0 comments on commit db8a338

Please sign in to comment.