From ff9f05599af64e6c9eeeda050c868c7729706a34 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Sun, 6 Oct 2024 00:17:35 +0100 Subject: [PATCH 1/3] Making contextual data accessible for async functions --- .../event_handler/appsync.py | 7 +++- docs/core/event_handler/appsync.md | 4 +- .../handlers/appsync_resolver_handler.py | 3 +- .../appsync/test_appsync_batch_resolvers.py | 38 ++++++++++++++++++ .../appsync/test_appsync_single_resolvers.py | 40 +++++++++++++++++++ 5 files changed, 87 insertions(+), 5 deletions(-) diff --git a/aws_lambda_powertools/event_handler/appsync.py b/aws_lambda_powertools/event_handler/appsync.py index c60256ca706..606dbb19ef0 100644 --- a/aws_lambda_powertools/event_handler/appsync.py +++ b/aws_lambda_powertools/event_handler/appsync.py @@ -149,7 +149,12 @@ def lambda_handler(event, context): Router.current_event = data_model(event) response = self._call_single_resolver(event=event, data_model=data_model) - self.clear_context() + # 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. + # For single async operations, the context should be cleaned up manually after the coroutine completes. + # See: https://github.com/aws-powertools/powertools-lambda-python/issues/5290 + if not asyncio.iscoroutine(response): + self.clear_context() return response diff --git a/docs/core/event_handler/appsync.md b/docs/core/event_handler/appsync.md index cb5f26da724..a2f29e5dba5 100644 --- a/docs/core/event_handler/appsync.md +++ b/docs/core/event_handler/appsync.md @@ -270,8 +270,8 @@ Let's assume you have `split_operation.py` as your Lambda function entrypoint an You can use `append_context` when you want to share data between your App and Router instances. Any data you share will be available via the `context` dictionary available in your App or Router context. -???+ info - For safety, we always clear any data available in the `context` dictionary after each invocation. +???+ warning + For safety, we clear the context after each invocation, except for async single resolvers. For these, use `app.context.clear()` before returning the function. ???+ tip This can also be useful for middlewares injecting contextual information before a request is processed. diff --git a/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py b/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py index 594290f478d..b49cee90741 100644 --- a/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py +++ b/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py @@ -79,8 +79,7 @@ class Post(BaseModel): # PROCESSING SINGLE RESOLVERS @app.resolver(type_name="Query", field_name="getPost") def get_post(post_id: str = "") -> dict: - post = Post(**posts[post_id]).dict() - return post + return Post(**posts[post_id]).model_dump_json() @app.resolver(type_name="Query", field_name="allPosts") 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 index a6452ee683d..c594be54a5b 100644 --- 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 @@ -943,3 +943,41 @@ def get_user(event: List) -> List: # THEN the resolver must be able to return a field in the batch_current_event assert ret[0] == mock_event[0]["identity"]["sub"] + + +def test_context_is_accessible_in_sync_batch_resolver(): + mock_event = load_event("appSyncBatchEvent.json") + + # GIVEN An instance of AppSyncResolver and a resolver function registered with the app + app = AppSyncResolver() + + @app.batch_resolver(field_name="createSomething") + def get_user(event: List) -> List: + return [app.context.get("project_name")] + + # WHEN we resolve the event + app.append_context(project_name="powertools") + ret = app.resolve(mock_event, {}) + + # THEN the resolver must be able to return a field in the batch_current_event + assert app.context == {} + assert ret[0] == "powertools" + + +def test_context_is_accessible_in_async_batch_resolver(): + mock_event = load_event("appSyncBatchEvent.json") + + # GIVEN An instance of AppSyncResolver and a resolver function registered with the app + app = AppSyncResolver() + + @app.async_batch_resolver(field_name="createSomething") + async def get_user(event: List) -> List: + return [app.context.get("project_name")] + + # WHEN we resolve the event + app.append_context(project_name="powertools") + ret = app.resolve(mock_event, {}) + + # THEN the resolver must be able to return a field in the batch_current_event + assert app.context == {} + assert ret[0] == "powertools" 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 index 966e3a7a650..df44793f33b 100644 --- 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 @@ -289,3 +289,43 @@ def get_user(id: str) -> dict: # noqa AA03 VNE003 # THEN the resolver must be able to return a field in the current_event assert ret == mock_event["identity"]["sub"] + + +def test_route_context_is_not_cleared_after_resolve_async(): + # GIVEN + app = AppSyncResolver() + event = {"typeName": "Query", "fieldName": "listLocations", "arguments": {"name": "value"}} + + @app.resolver(field_name="listLocations") + async def get_locations(name: str): + return f"get_locations#{name}" + + # WHEN event resolution kicks in + app.append_context(is_admin=True) + app.resolve(event, {}) + + # THEN context should be empty + assert app.context == {"is_admin": True} + + +def test_route_context_is_manually_cleared_after_resolve_async(): + # GIVEN + # GIVEN + app = AppSyncResolver() + + mock_event = {"typeName": "Customer", "fieldName": "field", "arguments": {}} + + @app.resolver(field_name="field") + async def get_async(): + app.context.clear() + await asyncio.sleep(0.0001) + return "value" + + # WHEN + mock_context = LambdaContext() + app.append_context(is_admin=True) + result = app.resolve(mock_event, mock_context) + + # THEN + assert asyncio.run(result) == "value" + assert app.context == {} From ef06a204921dd63e58430469f32649714679e2d7 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 24 Oct 2024 11:37:44 +0100 Subject: [PATCH 2/3] V4 comment --- aws_lambda_powertools/event_handler/appsync.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aws_lambda_powertools/event_handler/appsync.py b/aws_lambda_powertools/event_handler/appsync.py index 606dbb19ef0..6f1cb72d067 100644 --- a/aws_lambda_powertools/event_handler/appsync.py +++ b/aws_lambda_powertools/event_handler/appsync.py @@ -153,6 +153,7 @@ def lambda_handler(event, context): # If we clean the context immediately, it might not be available when the coroutine is actually executed. # For single async operations, the context should be cleaned up manually after the coroutine completes. # See: https://github.com/aws-powertools/powertools-lambda-python/issues/5290 + # REVIEW: Review this support in Powertools V4 if not asyncio.iscoroutine(response): self.clear_context() From 16965263a93c904a72c075e1f2421fc48438fc5f Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 24 Oct 2024 14:03:36 +0100 Subject: [PATCH 3/3] Reverting test --- .../event_handler_appsync/handlers/appsync_resolver_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py b/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py index b49cee90741..594290f478d 100644 --- a/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py +++ b/tests/e2e/event_handler_appsync/handlers/appsync_resolver_handler.py @@ -79,7 +79,8 @@ class Post(BaseModel): # PROCESSING SINGLE RESOLVERS @app.resolver(type_name="Query", field_name="getPost") def get_post(post_id: str = "") -> dict: - return Post(**posts[post_id]).model_dump_json() + post = Post(**posts[post_id]).dict() + return post @app.resolver(type_name="Query", field_name="allPosts")