Coverage for async_durable_execution/_operation/wait_for_callback.py: 100%

1"""wait_for_callback extension 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 .._core import ( 

11 Duration, 

12 OperationContext, 

13 OperationSubType, 

14 SerDes, 

15 bind_current_context, 

16 durable_callable, 

17 get_current_context, 

18) 

19from ..extension import get_extension_context 

20from .callback import Callback, create_callback as _create_callback 

21from .step import get_step_context, step as _step 

22 

23if TYPE_CHECKING: 

24 from collections.abc import Awaitable, Callable 

25 from .child import SummaryGenerator 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30def create_callback( 

31 *, 

32 name: str | None = None, 

33 timeout: Duration | None = None, 

34 heartbeat_timeout: Duration | None = None, 

35 serdes: SerDes | None = None, 

36) -> asyncio.Task[Callback]: 

37 """Create an SDK-owned callback through the stable operation SPI.""" 

38 return _create_callback( 

39 name=name, 

40 timeout=timeout, 

41 heartbeat_timeout=heartbeat_timeout, 

42 serdes=serdes, 

43 ) 

44 

45 

46def step( 

47 func: Callable[[], Awaitable[Any]], 

48 *, 

49 name: str | None = None, 

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

51 serdes: SerDes | None = None, 

52) -> asyncio.Task[Any]: 

53 """Run an SDK-owned submitter step through the stable operation SPI.""" 

54 return _step( 

55 func, 

56 name=name, 

57 retry_strategy=retry_strategy, 

58 serdes=serdes, 

59 ) 

60 

61 

62def _create_child_context_task( 

63 func: Callable[[], Awaitable[Any]], 

64 *, 

65 sub_type: OperationSubType, 

66 name: str | None = None, 

67 serdes: SerDes | None = None, 

68 summary_generator: SummaryGenerator | None = None, 

69 is_virtual: bool = False, 

70) -> asyncio.Task[Any]: 

71 """Run an SDK-owned callback scope through the stable operation SPI.""" 

72 return ( 

73 get_extension_context() 

74 ._reserve_sdk_operation(name) # noqa: SLF001 

75 ._run_in_child_context( # noqa: SLF001 

76 func, 

77 sub_type=sub_type, 

78 serdes=serdes, 

79 summary_generator=summary_generator, 

80 is_virtual=is_virtual, 

81 ) 

82 ) 

83 

84 

85@durable_callable 

86async def wait_for_callback_handler( 

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

88 name: str | None = None, 

89 timeout: Duration | None = None, 

90 heartbeat_timeout: Duration | None = None, 

91 serdes: SerDes | None = None, 

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

93) -> Any: 

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

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

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

97 callback: Callback = await create_callback( 

98 name=callback_name, 

99 timeout=timeout, 

100 heartbeat_timeout=heartbeat_timeout, 

101 serdes=serdes, 

102 ) 

103 

104 async def submitter_step() -> Any: 

105 step_context = get_step_context() 

106 callback_context = WaitForCallbackContext( 

107 callback_id=callback.callback_id, 

108 execution_state=step_context.execution_state, 

109 operation_identifier=step_context.operation_identifier, 

110 ) 

111 with bind_current_context(callback_context): 

112 return await submitter() 

113 

114 await step( 

115 func=submitter_step, 

116 name=submitter_step_name, 

117 retry_strategy=retry_strategy, 

118 serdes=serdes, 

119 ) 

120 

121 return await callback.result() 

122 

123 

124def wait_for_callback( 

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

126 *, 

127 name: str | None = None, 

128 timeout: Duration | None = None, 

129 heartbeat_timeout: Duration | None = None, 

130 serdes: SerDes | None = None, 

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

132) -> asyncio.Task[Any]: 

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

134 

135 Args: 

136 submitter: Async callable. Use get_wait_for_callback_context().callback_id 

137 inside the submitter to access the callback id. 

138 name: Optional durable operation name. 

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

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

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

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

143 """ 

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

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

146 

147 return _create_child_context_task( 

148 wait_for_callback_handler( 

149 submitter, 

150 name=context_name, 

151 timeout=timeout, 

152 heartbeat_timeout=heartbeat_timeout, 

153 serdes=serdes, 

154 retry_strategy=retry_strategy, 

155 ), 

156 sub_type=OperationSubType.WAIT_FOR_CALLBACK, 

157 name=context_name, 

158 serdes=serdes, 

159 ) 

160 

161 

162@dataclass(frozen=True) 

163class WaitForCallbackContext(OperationContext): 

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

165 

166 callback_id: str = "" 

167 

168 

169def get_wait_for_callback_context() -> WaitForCallbackContext: 

170 """Return the active `WaitForCallbackContext`.""" 

171 current_context = get_current_context() 

172 if not isinstance(current_context, WaitForCallbackContext): 

173 msg = ( 

174 "get_wait_for_callback_context() can only be used while a " 

175 "wait_for_callback submitter is executing." 

176 ) 

177 raise RuntimeError(msg) 

178 return current_context