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

1"""Implement the Durable invoke operation.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7from typing import TYPE_CHECKING, TypeVar, cast 

8 

9# Import base classes for operation executor pattern 

10from .base import OperationExecutor 

11from .child import get_durable_context 

12from ..exceptions import ( 

13 CallableRuntimeError, 

14 ExecutionError, 

15 suspend_with_optional_resume_delay, 

16) 

17from ..models import ( 

18 ChainedInvokeOptions, 

19 Operation, 

20 OperationIdentifier, 

21 OperationStatus, 

22 OperationUpdate, 

23 OperationSubType, 

24) 

25from ..serdes import ( 

26 DEFAULT_JSON_SERDES, 

27) 

28from ..task import create_eager_task 

29 

30if TYPE_CHECKING: 

31 from .child import DurableContext 

32 from ..serdes import SerDes 

33 from ..state import ExecutionState 

34 

35P = TypeVar("P") # Payload type 

36R = TypeVar("R") # Result type 

37 

38logger = logging.getLogger(__name__) 

39 

40 

41class InvokeOperationExecutor(OperationExecutor[R]): 

42 """Executor for invoke operations.""" 

43 

44 def __init__( 

45 self, 

46 function_name: str, 

47 payload: P, 

48 state: ExecutionState, 

49 operation_identifier: OperationIdentifier, 

50 serdes_payload: SerDes[P] | None = None, 

51 serdes_result: SerDes[R] | None = None, 

52 tenant_id: str | None = None, 

53 ): 

54 """Initialize the invoke operation executor. 

55 

56 Args: 

57 function_name: Name of the function to invoke 

58 payload: The payload to pass to the invoked function 

59 state: The execution state 

60 operation_identifier: The operation identifier 

61 serdes_payload: Optional serializer for the invocation payload 

62 serdes_result: Optional deserializer for the invocation result 

63 tenant_id: Optional tenant identifier for the chained invocation 

64 """ 

65 super().__init__(state=state, operation_identifier=operation_identifier) 

66 self.function_name = function_name 

67 self.payload = payload 

68 self.serdes_payload = serdes_payload 

69 self.serdes_result = serdes_result 

70 self.tenant_id = tenant_id 

71 

72 async def start(self) -> R: 

73 """Start a new invoke operation.""" 

74 serialized_payload: str = await self.serialize_value( 

75 value=self.payload, 

76 serdes=self.serdes_payload or DEFAULT_JSON_SERDES, 

77 ) 

78 start_operation: OperationUpdate = OperationUpdate.create_invoke_start( 

79 identifier=self.operation_identifier, 

80 payload=serialized_payload, 

81 chained_invoke_options=ChainedInvokeOptions( 

82 function_name=self.function_name, 

83 tenant_id=self.tenant_id, 

84 ), 

85 ) 

86 await self.create_checkpoint(start_operation, is_sync=True) 

87 

88 logger.debug( 

89 "🚀 Invoke %s started, will suspend for completion", 

90 self.operation_name or self.function_name, 

91 ) 

92 

93 return await self.execute() 

94 

95 async def replay(self, operation: Operation) -> R: 

96 """Replay an existing invoke operation from its checkpoint.""" 

97 invoke_details = operation.chained_invoke_details 

98 if operation.status is OperationStatus.SUCCEEDED: 

99 result_data = invoke_details.result if invoke_details else None 

100 if result_data is None: 

101 return cast("R", None) 

102 

103 result: R = await self.deserialize_value( 

104 data=result_data, 

105 serdes=self.serdes_result or DEFAULT_JSON_SERDES, 

106 ) 

107 return result 

108 

109 # Terminal failures 

110 if ( 

111 operation.status is OperationStatus.FAILED 

112 or operation.status is OperationStatus.TIMED_OUT 

113 or operation.status is OperationStatus.STOPPED 

114 ): 

115 error = invoke_details.error if invoke_details else None 

116 if error is None: 

117 raise CallableRuntimeError( 

118 message="Unknown error. No ErrorObject exists on the Checkpoint Operation.", 

119 error_type=None, 

120 data=None, 

121 stack_trace=None, 

122 ) 

123 

124 raise CallableRuntimeError.from_error_object(error) 

125 

126 if operation.status is OperationStatus.STARTED: 

127 logger.debug( 

128 "⏳ Invoke %s still in progress, will suspend", 

129 self.operation_name or self.function_name, 

130 ) 

131 return await self.execute() 

132 

133 return await self.execute() 

134 

135 async def execute(self, operation: Operation | None = None) -> R: 

136 """Execute invoke operation by suspending to wait for async completion. 

137 

138 The invoke operation doesn't execute synchronously - it suspends and 

139 the backend executes the invoked function asynchronously. 

140 

141 Returns: 

142 Never returns - always suspends 

143 

144 Raises: 

145 Always suspends via suspend_with_optional_resume_delay 

146 ExecutionError: If suspend doesn't raise (should never happen) 

147 """ 

148 msg: str = f"Invoke {self.operation_identifier.operation_id} started, suspending for completion" 

149 suspend_with_optional_resume_delay(msg) 

150 # This line should never be reached since suspend_with_optional_resume_delay always raises 

151 error_msg: str = "suspend_with_optional_resume_delay should have raised an exception, but did not." 

152 raise ExecutionError(error_msg) from None 

153 

154 

155def invoke( 

156 function_name: str, 

157 payload: P, 

158 *, 

159 name: str | None = None, 

160 serdes_payload: SerDes[P] | None = None, 

161 serdes_result: SerDes[R] | None = None, 

162 tenant_id: str | None = None, 

163) -> asyncio.Task[R]: 

164 """Invoke another durable Lambda function and wait for its durable result. 

165 

166 Args: 

167 function_name: Qualified Lambda function name or ARN to invoke. 

168 payload: Payload to send to the invoked function. 

169 name: Optional durable operation name. 

170 serdes_payload: Optional serializer for the invocation payload. 

171 serdes_result: Optional deserializer for the invocation result. 

172 tenant_id: Optional tenant identifier for the chained invocation. 

173 """ 

174 context = get_durable_context("invoke") 

175 

176 with context._replay_aware(): 

177 operation_id = context.step_counter.create_step_id() 

178 operation_identifier = OperationIdentifier( 

179 operation_id=operation_id, 

180 sub_type=OperationSubType.CHAINED_INVOKE, 

181 parent_id=context.parent_id, 

182 name=name, 

183 ) 

184 

185 return create_eager_task( 

186 lambda: _invoke( 

187 function_name=function_name, 

188 payload=payload, 

189 context=context, 

190 operation_identifier=operation_identifier, 

191 serdes_payload=serdes_payload, 

192 serdes_result=serdes_result, 

193 tenant_id=tenant_id, 

194 ), 

195 ) 

196 

197 

198async def _invoke( 

199 function_name: str, 

200 payload: P, 

201 *, 

202 context: DurableContext, 

203 operation_identifier: OperationIdentifier, 

204 serdes_payload: SerDes[P] | None = None, 

205 serdes_result: SerDes[R] | None = None, 

206 tenant_id: str | None = None, 

207) -> R: 

208 executor: InvokeOperationExecutor[R] = InvokeOperationExecutor( 

209 function_name=function_name, 

210 payload=payload, 

211 state=context.execution_state, 

212 operation_identifier=operation_identifier, 

213 serdes_payload=serdes_payload, 

214 serdes_result=serdes_result, 

215 tenant_id=tenant_id, 

216 ) 

217 return await executor.process()