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

* Adding exception handler support

* Adding exception handler support - fix tests

* Adding exception handler support - fix tests + docs

Author: leandrodamascena

aws_lambda_powertools/event_handler/appsync.py

@@ -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,
@@ -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.
@@ -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

aws_lambda_powertools/utilities/data_classes/code_pipeline_job_event.py

@@ -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:
 

docs/core/event_handler/appsync.md

@@ -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

examples/event_handler_graphql/src/exception_handling_graphql.py

@@ -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)

tests/functional/event_handler/required_dependencies/appsync/test_appsync_batch_resolvers.py

@@ -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]

tests/functional/event_handler/required_dependencies/appsync/test_appsync_single_resolvers.py

@@ -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"