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

1"""Composite wait_for_callback operation built from callback, step, and child context.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7from dataclasses import dataclass 

8from typing import TYPE_CHECKING, Any 

9 

10from ..config import Duration 

11from ..context import bind_current_context, get_current_context 

12from ..execution import durable_callable 

13from ..primitive.base import OperationContext 

14from ..primitive.callback import Callback, create_callback 

15from ..primitive.child import run_in_child_context 

16from ..primitive.step import step 

17 

18if TYPE_CHECKING: 

19 from collections.abc import Awaitable, Callable 

20 

21 from ..serdes import SerDes 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26@durable_callable 

27async def wait_for_callback_handler( 

28 submitter: Callable[[], Awaitable[Any]], 

29 name: str | None = None, 

30 timeout: Duration | None = None, 

31 heartbeat_timeout: Duration | None = None, 

32 serdes: SerDes | None = None, 

33 retry_strategy: Callable[[Exception, int], Duration | None] | None = None, 

34) -> Any: 

35 """Create a callback, run a submitter, and wait for callback completion.""" 

36 callback_name = f"{name}-callback" if name is not None else "callback" 

37 submitter_step_name = f"{name}-submitter" if name is not None else "submitter" 

38 callback: Callback = await create_callback( 

39 name=callback_name, 

40 timeout=timeout, 

41 heartbeat_timeout=heartbeat_timeout, 

42 serdes=serdes, 

43 ) 

44 

45 async def submitter_step(): 

46 step_context = get_current_context() 

47 callback_context = WaitForCallbackContext( 

48 callback_id=callback.callback_id, 

49 execution_state=step_context.execution_state, 

50 operation_identifier=step_context.operation_identifier, 

51 ) 

52 with bind_current_context(callback_context): 

53 return await submitter() 

54 

55 await step( 

56 func=submitter_step, 

57 name=submitter_step_name, 

58 retry_strategy=retry_strategy, 

59 serdes=serdes, 

60 ) 

61 

62 return await callback.result() 

63 

64 

65def wait_for_callback( 

66 submitter: Callable[[], Awaitable[Any]], 

67 *, 

68 name: str | None = None, 

69 timeout: Duration | None = None, 

70 heartbeat_timeout: Duration | None = None, 

71 serdes: SerDes | None = None, 

72 retry_strategy: Callable[[Exception, int], Duration | None] | None = None, 

73) -> asyncio.Task[Any]: 

74 """Create a callback, run a submitter, then suspend until the callback resolves. 

75 

76 Args: 

77 submitter: Async callable. Use get_current_context().callback_id inside the 

78 submitter to access the callback id. 

79 name: Optional durable operation name. 

80 timeout: Optional maximum time to wait for callback completion. 

81 heartbeat_timeout: Optional maximum time to wait between callback heartbeats. 

82 serdes: Optional serializer for callback results and submitter results. 

83 retry_strategy: Optional strategy that returns a retry delay or None to stop. 

84 """ 

85 context_name = name if name is not None else getattr(submitter, "__name__", None) 

86 logger.debug("wait_for_callback name: %s", context_name) 

87 

88 return run_in_child_context( 

89 wait_for_callback_handler( 

90 submitter, 

91 name=context_name, 

92 timeout=timeout, 

93 heartbeat_timeout=heartbeat_timeout, 

94 serdes=serdes, 

95 retry_strategy=retry_strategy, 

96 ), 

97 name=context_name, 

98 ) 

99 

100 

101@dataclass(frozen=True) 

102class WaitForCallbackContext(OperationContext): 

103 """Context available during wait_for_callback submitter execution.""" 

104 

105 callback_id: str = ""