Skip to content

Commit fadeef9

Browse files
leandrodamascenaleandrodamascena
committed
Adding exception handler support
1 parent 5426a7a commit fadeef9

File tree

3 files changed

+207
-7
lines changed

3 files changed

+207
-7
lines changed

‎aws_lambda_powertools/event_handler/appsync.py

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ def __init__(self):
5353
"""
5454
super().__init__()
5555
self.context = {} # early init as customers might add context before event resolution
56+
self._exception_handlers: dict[type, Callable] = {}
5657

5758
def __call__(
5859
self,
@@ -142,12 +143,18 @@ def lambda_handler(event, context):
142143
self.lambda_context = context
143144
Router.lambda_context = context
144145

145-
if isinstance(event, list):
146-
Router.current_batch_event = [data_model(e) for e in event]
147-
response = self._call_batch_resolver(event=event, data_model=data_model)
148-
else:
149-
Router.current_event = data_model(event)
150-
response = self._call_single_resolver(event=event, data_model=data_model)
146+
try:
147+
if isinstance(event, list):
148+
Router.current_batch_event = [data_model(e) for e in event]
149+
response = self._call_batch_resolver(event=event, data_model=data_model)
150+
else:
151+
Router.current_event = data_model(event)
152+
response = self._call_single_resolver(event=event, data_model=data_model)
153+
except Exception as exp:
154+
response_builder = self._lookup_exception_handler(type(exp))
155+
if response_builder:
156+
return response_builder(exp)
157+
raise
151158

152159
# We don't clear the context for coroutines because we don't have control over the event loop.
153160
# If we clean the context immediately, it might not be available when the coroutine is actually executed.
@@ -470,3 +477,47 @@ def async_batch_resolver(
470477
raise_on_error=raise_on_error,
471478
aggregate=aggregate,
472479
)
480+
481+
def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
482+
"""
483+
A decorator function that registers a handler for one or more exception types.
484+
485+
Parameters
486+
----------
487+
exc_class (type[Exception] | list[type[Exception]])
488+
A single exception type or a list of exception types.
489+
490+
Returns
491+
-------
492+
Callable:
493+
A decorator function that registers the exception handler.
494+
"""
495+
496+
def register_exception_handler(func: Callable):
497+
if isinstance(exc_class, list): # pragma: no cover
498+
for exp in exc_class:
499+
self._exception_handlers[exp] = func
500+
else:
501+
self._exception_handlers[exc_class] = func
502+
return func
503+
504+
return register_exception_handler
505+
506+
def _lookup_exception_handler(self, exp_type: type) -> Callable | None:
507+
"""
508+
Looks up the registered exception handler for the given exception type or its base classes.
509+
510+
Parameters
511+
----------
512+
exp_type (type):
513+
The exception type to look up the handler for.
514+
515+
Returns
516+
-------
517+
Callable | None:
518+
The registered exception handler function if found, otherwise None.
519+
"""
520+
for cls in exp_type.__mro__:
521+
if cls in self._exception_handlers:
522+
return self._exception_handlers[cls]
523+
return None

‎aws_lambda_powertools/utilities/data_classes/code_pipeline_job_event.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ def put_artifact(self, artifact_name: str, body: Any, content_type: str) -> None
311311
key = artifact.location.s3_location.key
312312

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

316316
if self.data.encryption_key:
317317

‎tests/functional/event_handler/required_dependencies/appsync/test_appsync.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
from typing import Optional
2+
3+
from aws_lambda_powertools.event_handler import AppSyncResolver
4+
from aws_lambda_powertools.utilities.data_classes import AppSyncResolverEvent
5+
from tests.functional.utils import load_event
6+
7+
8+
def test_exception_handler_with_single_resolver():
9+
# GIVEN a AppSyncResolver instance
10+
mock_event = load_event("appSyncDirectResolver.json")
11+
12+
app = AppSyncResolver()
13+
14+
# WHEN we configure exception handler for ValueError
15+
@app.exception_handler(ValueError)
16+
def handle_value_error(ex: ValueError):
17+
return {"message": "error"}
18+
19+
@app.resolver(field_name="createSomething")
20+
def create_something(id: str): # noqa AA03 VNE003
21+
raise ValueError("Error")
22+
23+
# Call the implicit handler
24+
result = app(mock_event, {})
25+
26+
# THEN the return must be the Exception Handler error message
27+
assert result["message"] == "error"
28+
29+
30+
def test_exception_handler_with_batch_resolver_and_raise_exception():
31+
32+
# GIVEN a AppSyncResolver instance
33+
app = AppSyncResolver()
34+
35+
event = [
36+
{
37+
"typeName": "Query",
38+
"info": {
39+
"fieldName": "listLocations",
40+
"parentTypeName": "Post",
41+
},
42+
"fieldName": "listLocations",
43+
"arguments": {},
44+
"source": {
45+
"id": "1",
46+
},
47+
},
48+
{
49+
"typeName": "Query",
50+
"info": {
51+
"fieldName": "listLocations",
52+
"parentTypeName": "Post",
53+
},
54+
"fieldName": "listLocations",
55+
"arguments": {},
56+
"source": {
57+
"id": "2",
58+
},
59+
},
60+
{
61+
"typeName": "Query",
62+
"info": {
63+
"fieldName": "listLocations",
64+
"parentTypeName": "Post",
65+
},
66+
"fieldName": "listLocations",
67+
"arguments": {},
68+
"source": {
69+
"id": [3, 4],
70+
},
71+
},
72+
]
73+
74+
# WHEN we configure exception handler for ValueError
75+
@app.exception_handler(ValueError)
76+
def handle_value_error(ex: ValueError):
77+
return {"message": "error"}
78+
79+
# WHEN the sync batch resolver for the 'listLocations' field is defined with raise_on_error=True
80+
@app.batch_resolver(field_name="listLocations", raise_on_error=True, aggregate=False)
81+
def create_something(event: AppSyncResolverEvent) -> Optional[list]: # noqa AA03 VNE003
82+
raise ValueError
83+
84+
# Call the implicit handler
85+
result = app(event, {})
86+
87+
# THEN the return must be the Exception Handler error message
88+
assert result["message"] == "error"
89+
90+
91+
def test_exception_handler_with_batch_resolver_and_no_raise_exception():
92+
93+
# GIVEN a AppSyncResolver instance
94+
app = AppSyncResolver()
95+
96+
event = [
97+
{
98+
"typeName": "Query",
99+
"info": {
100+
"fieldName": "listLocations",
101+
"parentTypeName": "Post",
102+
},
103+
"fieldName": "listLocations",
104+
"arguments": {},
105+
"source": {
106+
"id": "1",
107+
},
108+
},
109+
{
110+
"typeName": "Query",
111+
"info": {
112+
"fieldName": "listLocations",
113+
"parentTypeName": "Post",
114+
},
115+
"fieldName": "listLocations",
116+
"arguments": {},
117+
"source": {
118+
"id": "2",
119+
},
120+
},
121+
{
122+
"typeName": "Query",
123+
"info": {
124+
"fieldName": "listLocations",
125+
"parentTypeName": "Post",
126+
},
127+
"fieldName": "listLocations",
128+
"arguments": {},
129+
"source": {
130+
"id": [3, 4],
131+
},
132+
},
133+
]
134+
135+
# WHEN we configure exception handler for ValueError
136+
@app.exception_handler(ValueError)
137+
def handle_value_error(ex: ValueError):
138+
return {"message": "error"}
139+
140+
# WHEN the sync batch resolver for the 'listLocations' field is defined with raise_on_error=False
141+
@app.batch_resolver(field_name="listLocations", raise_on_error=False, aggregate=False)
142+
def create_something(event: AppSyncResolverEvent) -> Optional[list]: # noqa AA03 VNE003
143+
raise ValueError
144+
145+
# Call the implicit handler
146+
result = app(event, {})
147+
148+
# THEN the return must not trigger the Exception Handler, but instead return from the resolver
149+
assert result == [None, None, None]

0 commit comments

Comments
 (0)