Coverage for async-durable-execution/src/async_durable_execution/context.py: 100%

1from __future__ import annotations 

2 

3from contextlib import contextmanager 

4from contextvars import ContextVar, Token 

5from typing import TYPE_CHECKING 

6 

7if TYPE_CHECKING: 

8 from .primitive.base import OperationContext 

9 from .serdes import SerDesContext 

10 

11 

12_current_context: ContextVar = ContextVar( 

13 "async_durable_execution.current_context", 

14 default=None, 

15) 

16 

17 

18def set_current_context(context) -> Token: 

19 """Bind the active durable context for the current async task.""" 

20 return _current_context.set(context) 

21 

22 

23def reset_current_context(token: Token) -> None: 

24 """Restore the previous durable context after a temporary override.""" 

25 _current_context.reset(token) 

26 

27 

28def get_current_context(): 

29 """Return the currently active durable execution context. 

30 

31 Raises: 

32 RuntimeError: If called outside a durable handler, step, callback submitter, 

33 or wait-for-condition checker. 

34 """ 

35 current_context = _current_context.get() 

36 if current_context is None: 

37 msg = ( 

38 "get_current_context() can only be used while a durable function, " 

39 "step function, wait_for_callback submitter, or " 

40 "wait_for_condition check, or SerDes operation is executing." 

41 ) 

42 raise RuntimeError(msg) 

43 return current_context 

44 

45 

46@contextmanager 

47def bind_current_context(context: OperationContext | SerDesContext): 

48 """Temporarily bind the supplied durable context while invoking user code.""" 

49 token = set_current_context(context) 

50 try: 

51 yield 

52 finally: 

53 reset_current_context(token)