Refactoring to support aggregate events

leandrodamascena · leandrodamascena · commit 84c3590b0979 · 2024-06-24T22:51:50.000+01:00
diff --git a/aws_lambda_powertools/event_handler/appsync.py b/aws_lambda_powertools/event_handler/appsync.py
@@ -3,7 +3,7 @@
 import warnings
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
-from aws_lambda_powertools.event_handler.graphql_appsync.exceptions import ResolverNotFoundError
+from aws_lambda_powertools.event_handler.graphql_appsync.exceptions import InvalidBatchResponse, ResolverNotFoundError
 from aws_lambda_powertools.event_handler.graphql_appsync.router import Router
 from aws_lambda_powertools.utilities.data_classes import AppSyncResolverEvent
 from aws_lambda_powertools.utilities.typing import LambdaContext
@@ -168,7 +168,12 @@ def _call_single_resolver(self, event: dict, data_model: Type[AppSyncResolverEve
             raise ValueError(f"No resolver found for '{self.current_event.type_name}.{self.current_event.field_name}'")
         return resolver["func"](**self.current_event.arguments)
 
-    def _call_sync_batch_resolver(self, resolver: Callable, raise_on_error: bool = False) -> List[Any]:
+    def _call_sync_batch_resolver(
+        self,
+        resolver: Callable,
+        raise_on_error: bool = False,
+        aggregate: bool = True,
+    ) -> List[Any]:
         """
         Calls a synchronous batch resolver function for each event in the current batch.
 
@@ -179,6 +184,10 @@ def _call_sync_batch_resolver(self, resolver: Callable, raise_on_error: bool = F
         raise_on_error: bool
             A flag indicating whether to raise an error when processing batches
             with failed items. Defaults to False, which means errors are handled without raising exceptions.
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Returns
         -------
@@ -188,6 +197,17 @@ def _call_sync_batch_resolver(self, resolver: Callable, raise_on_error: bool = F
 
         logger.debug(f"Graceful error handling flag {raise_on_error=}")
 
+        # Checks whether the entire batch should be processed at once
+        if aggregate:
+            # Process the entire batch
+            response = resolver(event=self.current_batch_event)
+
+            if not isinstance(response, List):
+                raise InvalidBatchResponse("The response must be a List when using batch resolvers")
+
+            return response
+
+        # Non aggregated events, so we call this event list x times
         # Stop on first exception we encounter
         if raise_on_error:
             return [
@@ -206,7 +226,12 @@ def _call_sync_batch_resolver(self, resolver: Callable, raise_on_error: bool = F
 
         return results
 
-    async def _call_async_batch_resolver(self, resolver: Callable, raise_on_error: bool = False) -> List[Any]:
+    async def _call_async_batch_resolver(
+        self,
+        resolver: Callable,
+        raise_on_error: bool = False,
+        aggregate: bool = True,
+    ) -> List[Any]:
         """
         Asynchronously call a batch resolver for each event in the current batch.
 
@@ -217,6 +242,10 @@ async def _call_async_batch_resolver(self, resolver: Callable, raise_on_error: b
         raise_on_error: bool
             A flag indicating whether to raise an error when processing batches
             with failed items. Defaults to False, which means errors are handled without raising exceptions.
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Returns
         -------
@@ -225,7 +254,17 @@ async def _call_async_batch_resolver(self, resolver: Callable, raise_on_error: b
         """
 
         logger.debug(f"Graceful error handling flag {raise_on_error=}")
-        response = []
+
+        response: List = []
+
+        # Checks whether the entire batch should be processed at once
+        if aggregate:
+            # Process the entire batch
+            response.extend(await asyncio.gather(resolver(event=self.current_batch_event)))
+            if not isinstance(response[0], List):
+                raise InvalidBatchResponse("The response must be a List when using batch resolvers")
+
+            return response[0]
 
         # Prime coroutines
         tasks = [resolver(event=e, **e.arguments) for e in self.current_batch_event]
@@ -286,14 +325,19 @@ def _call_batch_resolver(self, event: List[dict], data_model: Type[AppSyncResolv
 
         if resolver:
             logger.debug(f"Found sync resolver. {resolver=}, {field_name=}")
-            return self._call_sync_batch_resolver(resolver=resolver["func"], raise_on_error=resolver["raise_on_error"])
+            return self._call_sync_batch_resolver(
+                resolver=resolver["func"],
+                raise_on_error=resolver["raise_on_error"],
+                aggregate=resolver["aggregate"],
+            )
 
         if async_resolver:
             logger.debug(f"Found async resolver. {resolver=}, {field_name=}")
             return asyncio.run(
                 self._call_async_batch_resolver(
                     resolver=async_resolver["func"],
                     raise_on_error=async_resolver["raise_on_error"],
+                    aggregate=async_resolver["aggregate"],
                 ),
             )
 
@@ -371,6 +415,7 @@ def batch_resolver(
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         """Registers batch resolver function for GraphQL type and field name.
 
@@ -385,6 +430,10 @@ def batch_resolver(
             GraphQL field e.g., getTodo, createTodo, by default None
         raise_on_error : bool, optional
             Whether to fail entire batch upon error, or handle errors gracefully (None), by default False
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Returns
         -------
@@ -395,16 +444,19 @@ def batch_resolver(
             field_name=field_name,
             type_name=type_name,
             raise_on_error=raise_on_error,
+            aggregate=aggregate,
         )
 
     def async_batch_resolver(
         self,
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         return self._async_batch_resolver_registry.register(
             field_name=field_name,
             type_name=type_name,
             raise_on_error=raise_on_error,
+            aggregate=aggregate,
         )
diff --git a/aws_lambda_powertools/event_handler/graphql_appsync/_registry.py b/aws_lambda_powertools/event_handler/graphql_appsync/_registry.py
@@ -13,6 +13,7 @@ def register(
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         """Registers the resolver for field_name
 
@@ -25,6 +26,10 @@ def register(
         raise_on_error: bool
             A flag indicating whether to raise an error when processing batches
             with failed items. Defaults to False, which means errors are handled without raising exceptions.
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Return
         ----------
@@ -34,7 +39,11 @@ def register(
 
         def _register(func) -> Callable:
             logger.debug(f"Adding resolver `{func.__name__}` for field `{type_name}.{field_name}`")
-            self.resolvers[f"{type_name}.{field_name}"] = {"func": func, "raise_on_error": raise_on_error}
+            self.resolvers[f"{type_name}.{field_name}"] = {
+                "func": func,
+                "raise_on_error": raise_on_error,
+                "aggregate": aggregate,
+            }
             return func
 
         return _register
diff --git a/aws_lambda_powertools/event_handler/graphql_appsync/base.py b/aws_lambda_powertools/event_handler/graphql_appsync/base.py
@@ -49,6 +49,7 @@ def batch_resolver(
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         """
         Retrieve a batch resolver function for a specific type and field.
@@ -62,6 +63,10 @@ def batch_resolver(
         raise_on_error: bool
             A flag indicating whether to raise an error when processing batches
             with failed items. Defaults to False, which means errors are handled without raising exceptions.
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Examples
         --------
@@ -95,6 +100,7 @@ def async_batch_resolver(
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         """
         Retrieve a batch resolver function for a specific type and field and runs async.
@@ -108,6 +114,10 @@ def async_batch_resolver(
         raise_on_error: bool
             A flag indicating whether to raise an error when processing batches
             with failed items. Defaults to False, which means errors are handled without raising exceptions.
+        aggregate: bool
+            A flag indicating whether the batch items should be processed at once or individually.
+            If True (default), the batch resolver will process all items in the batch as a single event.
+            If False, the batch resolver will process each item in the batch individually.
 
         Examples
         --------
diff --git a/aws_lambda_powertools/event_handler/graphql_appsync/exceptions.py b/aws_lambda_powertools/event_handler/graphql_appsync/exceptions.py
@@ -2,3 +2,9 @@ class ResolverNotFoundError(Exception):
     """
     When a resolver is not found during a lookup.
     """
+
+
+class InvalidBatchResponse(Exception):
+    """
+    When a batch response something different from a List
+    """
diff --git a/aws_lambda_powertools/event_handler/graphql_appsync/router.py b/aws_lambda_powertools/event_handler/graphql_appsync/router.py
@@ -21,23 +21,27 @@ def batch_resolver(
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         return self._batch_resolver_registry.register(
             field_name=field_name,
             type_name=type_name,
             raise_on_error=raise_on_error,
+            aggregate=aggregate,
         )
 
     def async_batch_resolver(
         self,
         type_name: str = "*",
         field_name: Optional[str] = None,
         raise_on_error: bool = False,
+        aggregate: bool = True,
     ) -> Callable:
         return self._async_batch_resolver_registry.register(
             field_name=field_name,
             type_name=type_name,
             raise_on_error=raise_on_error,
+            aggregate=aggregate,
         )
 
     def append_context(self, **additional_context):
diff --git a/tests/e2e/event_handler_appsync/files/schema.graphql b/tests/e2e/event_handler_appsync/files/schema.graphql
@@ -17,4 +17,6 @@ type Post {
     downs: Int
     relatedPosts: [Post]
     relatedPostsAsync: [Post]
-}
+    relatedPostsAggregate: [Post]
+    relatedPostsAsyncAggregate: [Post]
+}
diff --git a/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py b/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py
@@ -76,6 +76,7 @@ class Post(BaseModel):
     downs: str
 
 
+# PROCESSING SINGLE RESOLVERS
 @app.resolver(type_name="Query", field_name="getPost")
 def get_post(post_id: str = "") -> dict:
     post = Post(**posts[post_id]).dict()
@@ -87,15 +88,27 @@ def all_posts() -> List[dict]:
     return list(posts.values())
 
 
-@app.batch_resolver(type_name="Post", field_name="relatedPosts")
+# PROCESSING BATCH WITHOUT AGGREGATION
+@app.batch_resolver(type_name="Post", field_name="relatedPosts", aggregate=False)
 def related_posts(event: AppSyncResolverEvent) -> Optional[list]:
     return posts_related[event.source["post_id"]] if event.source else None
 
 
-@app.async_batch_resolver(type_name="Post", field_name="relatedPostsAsync")
+@app.async_batch_resolver(type_name="Post", field_name="relatedPostsAsync", aggregate=False)
 async def related_posts_async(event: AppSyncResolverEvent) -> Optional[list]:
     return posts_related[event.source["post_id"]] if event.source else None
 
 
+# PROCESSING BATCH WITH AGGREGATION
+@app.batch_resolver(type_name="Post", field_name="relatedPostsAggregate")
+def related_posts_aggregate(event: List[AppSyncResolverEvent]) -> Optional[list]:
+    return [posts_related[record.source.get("post_id")] for record in event]
+
+
+@app.async_batch_resolver(type_name="Post", field_name="relatedPostsAsyncAggregate")
+async def related_posts_async_aggregate(event: List[AppSyncResolverEvent]) -> Optional[list]:
+    return [posts_related[record.source.get("post_id")] for record in event]
+
+
 def lambda_handler(event, context: LambdaContext) -> dict:
     return app.resolve(event, context)
diff --git a/tests/e2e/event_handler_appsync/infrastructure.py b/tests/e2e/event_handler_appsync/infrastructure.py
@@ -58,5 +58,19 @@ def _create_appsync_endpoint(self, function: Function):
             max_batch_size=10,
         )
 
+        lambda_datasource.create_resolver(
+            "QueryGetPostRelatedResolverAggregate",
+            type_name="Post",
+            field_name="relatedPostsAggregate",
+            max_batch_size=10,
+        )
+
+        lambda_datasource.create_resolver(
+            "QueryGetPostRelatedAsyncResolverAggregate",
+            type_name="Post",
+            field_name="relatedPostsAsyncAggregate",
+            max_batch_size=10,
+        )
+
         CfnOutput(self.stack, "GraphQLHTTPUrl", value=api.graphql_url)
         CfnOutput(self.stack, "GraphQLAPIKey", value=api.api_key)
diff --git a/tests/e2e/event_handler_appsync/test_appsync_resolvers.py b/tests/e2e/event_handler_appsync/test_appsync_resolvers.py
@@ -75,7 +75,7 @@ def test_appsync_get_post(appsync_endpoint, appsync_access_key):
 
 
 @pytest.mark.xdist_group(name="event_handler")
-def test_appsync_get_related_posts_batch(appsync_endpoint, appsync_access_key):
+def test_appsync_get_related_posts_batch_without_aggregate(appsync_endpoint, appsync_access_key):
     # GIVEN
     post_id = "2"
     related_posts_ids = ["3", "5"]
@@ -110,3 +110,41 @@ def test_appsync_get_related_posts_batch(appsync_endpoint, appsync_access_key):
         assert post["post_id"] in related_posts_ids
     for post in data["getPost"]["relatedPostsAsync"]:
         assert post["post_id"] in related_posts_ids
+
+
+@pytest.mark.xdist_group(name="event_handler")
+def test_appsync_get_related_posts_batch_with_aggregate(appsync_endpoint, appsync_access_key):
+    # GIVEN
+    post_id = "2"
+    related_posts_ids = ["3", "5"]
+
+    body = {
+        "query": f'query MyQuery {{ getPost(post_id: "{post_id}") \
+              {{ post_id relatedPostsAggregate {{ post_id }} relatedPostsAsyncAggregate {{ post_id }} }} }}',
+        "variables": None,
+        "operationName": "MyQuery",
+    }
+
+    # WHEN
+    response = data_fetcher.get_http_response(
+        Request(
+            method="POST",
+            url=appsync_endpoint,
+            json=body,
+            headers={"x-api-key": appsync_access_key, "Content-Type": "application/json"},
+        ),
+    )
+
+    # THEN expect a HTTP 200 response and content return Post id with dependent Posts id's
+    assert response.status_code == 200
+    assert response.content is not None
+
+    data = json.loads(response.content.decode("ascii"))["data"]
+
+    assert data["getPost"]["post_id"] == post_id
+    assert len(data["getPost"]["relatedPostsAggregate"]) == len(related_posts_ids)
+    assert len(data["getPost"]["relatedPostsAsyncAggregate"]) == len(related_posts_ids)
+    for post in data["getPost"]["relatedPostsAggregate"]:
+        assert post["post_id"] in related_posts_ids
+    for post in data["getPost"]["relatedPostsAsyncAggregate"]:
+        assert post["post_id"] in related_posts_ids
diff --git a/tests/functional/event_handler/required_dependencies/appsync/__init__.py b/tests/functional/event_handler/required_dependencies/appsync/__init__.py
diff --git a/tests/functional/event_handler/required_dependencies/appsync/test_appsync_batch_resolvers.py b/tests/functional/event_handler/required_dependencies/appsync/test_appsync_batch_resolvers.py
diff --git a/tests/functional/event_handler/required_dependencies/appsync/test_appsync_single_resolvers.py b/tests/functional/event_handler/required_dependencies/appsync/test_appsync_single_resolvers.py
