Skip to content

Commit e0566b0

Browse files
fix(event_handler): isolate local ASGI request state (#8452)
* fix(event_handler): isolate local ASGI request state * fix(event_handler): share request state with included routers * test(event_handler): cover router state delegation --------- Co-authored-by: Leandro <lcdama@amazon.pt>
1 parent 51c89a3 commit e0566b0

3 files changed

Lines changed: 286 additions & 20 deletions

File tree

‎aws_lambda_powertools/event_handler/api_gateway.py‎

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3035,21 +3035,18 @@ def include_router(self, router: Router, prefix: str | None = None) -> None:
30353035
An optional prefix to be added to the originally defined rule
30363036
"""
30373037

3038-
# Add reference to parent ApiGatewayResolver to support use cases where people subclass it to add custom logic
3039-
router.api_resolver = self
3040-
30413038
logger.debug("Merging App context with Router context")
30423039
self.context.update(**router.context)
30433040

3041+
# Delegate request state to the resolver after preserving the router context.
3042+
router.api_resolver = self
3043+
30443044
logger.debug("Appending Router middlewares into App middlewares.")
30453045
self._router_middlewares = self._router_middlewares + router._router_middlewares
30463046

30473047
logger.debug("Appending Router exception_handler into App exception_handler.")
30483048
self.exception_handler_manager.update_exception_handlers(router._exception_handlers)
30493049

3050-
# use pointer to allow context clearance after event is processed e.g., resolve(evt, ctx)
3051-
router.context = self.context
3052-
30533050
# Iterate through the routes defined in the router to configure and apply middlewares for each route
30543051
for route, func in router._routes.items():
30553052
new_route = route
@@ -3112,9 +3109,48 @@ def __init__(self):
31123109
self._routes: dict[tuple, Callable] = {}
31133110
self._routes_with_middleware: dict[tuple, list[Callable]] = {}
31143111
self.api_resolver: BaseRouter | None = None
3115-
self.context = {} # early init as customers might add context before event resolution
3112+
self._context: dict = {} # early init as customers might add context before event resolution
31163113
self._exception_handlers: dict[type, Callable] = {}
31173114

3115+
@property
3116+
def current_event(self) -> BaseProxyEvent:
3117+
if self.api_resolver is not None:
3118+
return self.api_resolver.current_event
3119+
return BaseRouter.current_event
3120+
3121+
@current_event.setter
3122+
def current_event(self, value: BaseProxyEvent) -> None:
3123+
if self.api_resolver is not None:
3124+
self.api_resolver.current_event = value
3125+
else:
3126+
BaseRouter.current_event = value
3127+
3128+
@property
3129+
def lambda_context(self) -> LambdaContext:
3130+
if self.api_resolver is not None:
3131+
return self.api_resolver.lambda_context
3132+
return BaseRouter.lambda_context
3133+
3134+
@lambda_context.setter
3135+
def lambda_context(self, value: LambdaContext) -> None:
3136+
if self.api_resolver is not None:
3137+
self.api_resolver.lambda_context = value
3138+
else:
3139+
BaseRouter.lambda_context = value
3140+
3141+
@property
3142+
def context(self) -> dict:
3143+
if self.api_resolver is not None:
3144+
return self.api_resolver.context
3145+
return self._context
3146+
3147+
@context.setter
3148+
def context(self, value: dict) -> None:
3149+
if self.api_resolver is not None:
3150+
self.api_resolver.context = value
3151+
else:
3152+
self._context = value
3153+
31183154
def route(
31193155
self,
31203156
rule: str,

‎aws_lambda_powertools/event_handler/http_resolver.py‎

Lines changed: 68 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
from __future__ import annotations
22

33
import base64
4+
from contextvars import ContextVar
5+
from dataclasses import dataclass, field
46
from typing import TYPE_CHECKING, Any, Callable
57
from urllib.parse import parse_qs
68

@@ -13,6 +15,8 @@
1315
from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
1416

1517
if TYPE_CHECKING:
18+
from collections.abc import Mapping, MutableMapping
19+
1620
from aws_lambda_powertools.shared.cookies import Cookie
1721

1822

@@ -95,7 +99,7 @@ def _from_dict(cls, data: dict[str, Any]) -> HttpProxyEvent:
9599
return instance
96100

97101
@classmethod
98-
def from_asgi(cls, scope: dict[str, Any], body: bytes | None = None) -> HttpProxyEvent:
102+
def from_asgi(cls, scope: Mapping[str, Any], body: bytes | None = None) -> HttpProxyEvent:
99103
"""
100104
Create an HttpProxyEvent from an ASGI scope dict.
101105
@@ -159,6 +163,14 @@ def get_remaining_time_in_millis(self) -> int: # pragma: no cover
159163
return 300000 # 5 minutes
160164

161165

166+
@dataclass
167+
class _RequestState:
168+
event: BaseProxyEvent | None = None
169+
lambda_context: Any = None
170+
context: dict = field(default_factory=dict)
171+
processed_stack_frames: list[str] = field(default_factory=list)
172+
173+
162174
class HttpResolverLocal(ApiGatewayResolver):
163175
"""
164176
ASGI-compatible HTTP resolver.
@@ -204,6 +216,8 @@ def __init__(
204216
strip_prefixes: list[str | Any] | None = None,
205217
enable_validation: bool = False,
206218
):
219+
self._startup_state = _RequestState()
220+
self._request_state: ContextVar[_RequestState | None] = ContextVar("local_http_request", default=None)
207221
super().__init__(
208222
proxy_type=ProxyEventType.APIGatewayProxyEvent, # Use REST API format internally
209223
cors=cors,
@@ -212,7 +226,46 @@ def __init__(
212226
strip_prefixes=strip_prefixes,
213227
enable_validation=enable_validation,
214228
)
215-
self._is_async_mode = False
229+
230+
@property
231+
def _state(self) -> _RequestState:
232+
return self._request_state.get() or self._startup_state
233+
234+
# Powertools declares these as mutable attributes. Properties preserve that
235+
# interface while directing each task to its own state. asyncio.to_thread
236+
# propagates the ContextVar, so middleware sees the same request dictionary.
237+
@property
238+
def current_event(self) -> BaseProxyEvent:
239+
# Preserve the inherited synchronous resolve() path outside ASGI calls.
240+
return self._state.event or BaseRouter.current_event
241+
242+
@current_event.setter
243+
def current_event(self, value: BaseProxyEvent) -> None:
244+
self._state.event = value
245+
246+
@property
247+
def lambda_context(self) -> Any:
248+
return self._state.lambda_context or BaseRouter.lambda_context
249+
250+
@lambda_context.setter
251+
def lambda_context(self, value: Any) -> None:
252+
self._state.lambda_context = value
253+
254+
@property
255+
def context(self) -> dict:
256+
return self._state.context
257+
258+
@context.setter
259+
def context(self, value: dict) -> None:
260+
self._state.context = value
261+
262+
@property
263+
def processed_stack_frames(self) -> list[str]:
264+
return self._state.processed_stack_frames
265+
266+
@processed_stack_frames.setter
267+
def processed_stack_frames(self, value: list[str]) -> None:
268+
self._state.processed_stack_frames = value
216269

217270
def _to_proxy_event(self, event: dict) -> BaseProxyEvent:
218271
"""Convert event dict to HttpProxyEvent."""
@@ -234,7 +287,7 @@ async def _resolve_async(self) -> dict: # type: ignore[override]
234287
response_builder = await super()._resolve_async()
235288
return response_builder.build(self.current_event, self._cors)
236289

237-
async def asgi_handler(self, scope: dict, receive: Callable, send: Callable) -> None:
290+
async def asgi_handler(self, scope: MutableMapping[str, Any], receive: Callable, send: Callable) -> None:
238291
"""
239292
ASGI interface - allows running with uvicorn/hypercorn/etc.
240293
@@ -274,25 +327,27 @@ async def asgi_handler(self, scope: dict, receive: Callable, send: Callable) ->
274327
# Create mock Lambda context
275328
context: Any = MockLambdaContext()
276329

277-
# Set up resolver state (similar to resolve())
278-
BaseRouter.current_event = self._to_proxy_event(event._data)
279-
BaseRouter.lambda_context = context
280-
281-
self._is_async_mode = True
282-
330+
# Never write BaseRouter's class attributes: another ASGI request may
331+
# enter while validation or the handler is awaiting I/O.
332+
state = _RequestState(
333+
event=self._to_proxy_event(event._data),
334+
lambda_context=context,
335+
context=self._startup_state.context.copy(),
336+
)
337+
token = self._request_state.set(state)
283338
try:
284-
# Use async resolve
285339
response = await self._resolve_async()
286340
finally:
287-
self._is_async_mode = False
288-
self.clear_context()
341+
# Reset only this task's binding. Middleware threads may still be
342+
# unwinding after cancellation and retain their request's state.
343+
self._request_state.reset(token)
289344

290345
# Send HTTP response
291346
await self._send_response(send, response)
292347

293348
async def __call__( # type: ignore[override]
294349
self,
295-
scope: dict,
350+
scope: MutableMapping[str, Any],
296351
receive: Callable,
297352
send: Callable,
298353
) -> None:

0 commit comments

Comments
 (0)