Coverage for async_durable_execution/_runner/local/execution.py: 99%

1from __future__ import annotations 

2 

3from dataclasses import replace 

4from datetime import datetime, timezone 

5from uuid import uuid4 

6 

7from ..._core import ( 

8 DurableExecutionInvocationOutput, 

9 ErrorObject, 

10 ExecutionDetails, 

11 InvocationStatus, 

12 Operation, 

13 OperationStatus, 

14 OperationType, 

15 OperationUpdate, 

16) 

17from ..exceptions import ( 

18 IllegalStateException, 

19 InvalidParameterValueException, 

20) 

21 

22# Import AWS exceptions 

23from ..model import ( 

24 InvocationCompletedDetails, 

25) 

26from .model import ( 

27 StartDurableExecutionInput, 

28 CheckpointToken, 

29) 

30 

31 

32class Execution: 

33 """Execution state.""" 

34 

35 def __init__( 

36 self, 

37 durable_execution_arn: str, 

38 start_input: StartDurableExecutionInput, 

39 operations: list[Operation], 

40 ) -> None: 

41 self.durable_execution_arn: str = durable_execution_arn 

42 # operation is frozen, it won't mutate - no need to clone/deep-copy 

43 self.start_input: StartDurableExecutionInput = start_input 

44 self.operations: list[Operation] = operations 

45 self.updates: list[OperationUpdate] = [] 

46 self.invocation_completions: list[InvocationCompletedDetails] = [] 

47 self.used_tokens: set[str] = set() 

48 self._token_sequence: int = 0 

49 self.is_complete: bool = False 

50 self.result: DurableExecutionInvocationOutput | None = None 

51 self.consecutive_failed_invocation_attempts: int = 0 

52 

53 @property 

54 def token_sequence(self) -> int: 

55 """Get current token sequence value.""" 

56 return self._token_sequence 

57 

58 @staticmethod 

59 def new(input: StartDurableExecutionInput) -> Execution: # noqa: A002 

60 # make a nicer arn 

61 # Pattern: arn:(aws[a-zA-Z-]*)?:lambda:[a-z]{2}(-gov)?-[a-z]+-\d{1}:\d{12}:durable-execution:[a-zA-Z0-9-_\.]+:[a-zA-Z0-9-_\.]+:[a-zA-Z0-9-_\.]+ 

62 # Example: arn:aws:lambda:us-east-1:123456789012:durable-execution:myDurableFunction:myDurableExecutionName:ce67da72-3701-4f83-9174-f4189d27b0a5 

63 return Execution( 

64 durable_execution_arn=str(uuid4()) 

65 + "/" 

66 + (input.invocation_id or str(uuid4())), 

67 start_input=input, 

68 operations=[], 

69 ) 

70 

71 def start(self) -> None: 

72 if self.start_input.invocation_id is None: 

73 msg: str = "invocation_id is required" 

74 raise InvalidParameterValueException(msg) 

75 self.operations.append( 

76 Operation( 

77 operation_id=self.start_input.invocation_id, 

78 parent_id=None, 

79 name=self.start_input.execution_name, 

80 start_timestamp=datetime.now(timezone.utc), 

81 operation_type=OperationType.EXECUTION, 

82 status=OperationStatus.STARTED, 

83 execution_details=ExecutionDetails( 

84 input_payload=self.start_input.get_normalized_input() 

85 ), 

86 ) 

87 ) 

88 

89 def get_operation_execution_started(self) -> Operation: 

90 if not self.operations: 

91 msg: str = "execution not started." 

92 

93 raise IllegalStateException(msg) 

94 

95 return self.operations[0] 

96 

97 def get_new_checkpoint_token(self) -> str: 

98 """Generate a new checkpoint token with incremented sequence""" 

99 self._token_sequence += 1 

100 new_token_sequence = self._token_sequence 

101 token = CheckpointToken( 

102 execution_arn=self.durable_execution_arn, 

103 token_sequence=new_token_sequence, 

104 ) 

105 token_str = token.to_str() 

106 self.used_tokens.add(token_str) 

107 return token_str 

108 

109 def get_navigable_operations(self) -> list[Operation]: 

110 """Get list of operations, but exclude child operations where the parent has already completed.""" 

111 return self.operations 

112 

113 def get_assertable_operations(self) -> list[Operation]: 

114 """Get list of operations, but exclude the EXECUTION operations""" 

115 # TODO: this excludes EXECUTION at start, but can there be an EXECUTION at the end if there was a checkpoint with large payload? 

116 return self.operations[1:] 

117 

118 def has_pending_operations(self) -> bool: 

119 """True if execution has pending operations.""" 

120 

121 for operation in self.operations: 

122 if ( 

123 operation.operation_type == OperationType.STEP 

124 and operation.status == OperationStatus.PENDING 

125 ) or ( 

126 operation.operation_type 

127 in [ 

128 OperationType.WAIT, 

129 OperationType.CALLBACK, 

130 OperationType.CHAINED_INVOKE, 

131 ] 

132 and operation.status == OperationStatus.STARTED 

133 ): 

134 return True 

135 return False 

136 

137 def record_invocation_completion( 

138 self, start_timestamp: datetime, end_timestamp: datetime, request_id: str 

139 ) -> None: 

140 """Record an invocation completion event.""" 

141 self.invocation_completions.append( 

142 InvocationCompletedDetails( 

143 start_timestamp=start_timestamp, 

144 end_timestamp=end_timestamp, 

145 request_id=request_id, 

146 ) 

147 ) 

148 

149 def complete_success(self, result: str | None) -> None: 

150 """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION).""" 

151 self.result = DurableExecutionInvocationOutput( 

152 status=InvocationStatus.SUCCEEDED, result=result 

153 ) 

154 self.is_complete = True 

155 self._end_execution(OperationStatus.SUCCEEDED) 

156 

157 def complete_fail(self, error: ErrorObject) -> None: 

158 """Complete execution with failure (DecisionType.FAIL_WORKFLOW_EXECUTION).""" 

159 self.result = DurableExecutionInvocationOutput( 

160 status=InvocationStatus.FAILED, error=error 

161 ) 

162 self.is_complete = True 

163 self._end_execution(OperationStatus.FAILED) 

164 

165 def complete_timeout(self, error: ErrorObject) -> None: 

166 """Complete execution with timeout.""" 

167 self.result = DurableExecutionInvocationOutput( 

168 status=InvocationStatus.FAILED, error=error 

169 ) 

170 self.is_complete = True 

171 self._end_execution(OperationStatus.TIMED_OUT) 

172 

173 def find_operation(self, operation_id: str) -> tuple[int, Operation]: 

174 """Find operation by ID, return index and operation.""" 

175 for i, operation in enumerate(self.operations): 

176 if operation.operation_id == operation_id: 

177 return i, operation 

178 msg: str = f"Attempting to update state of an Operation [{operation_id}] that doesn't exist" 

179 raise IllegalStateException(msg) 

180 

181 def find_callback_operation(self, callback_id: str) -> tuple[int, Operation]: 

182 """Find callback operation by callback_id, return index and operation.""" 

183 for i, operation in enumerate(self.operations): 

184 if ( 

185 operation.operation_type == OperationType.CALLBACK 

186 and operation.callback_details 

187 and operation.callback_details.callback_id == callback_id 

188 ): 

189 return i, operation 

190 msg: str = f"Callback operation with callback_id [{callback_id}] not found" 

191 raise IllegalStateException(msg) 

192 

193 def complete_wait(self, operation_id: str) -> Operation: 

194 """Complete WAIT operation when timer fires.""" 

195 index, operation = self.find_operation(operation_id) 

196 

197 # Validate 

198 if operation.status != OperationStatus.STARTED: 

199 msg_wait_not_started: str = f"Attempting to transition a Wait Operation[{operation_id}] to SUCCEEDED when it's not STARTED" 

200 raise IllegalStateException(msg_wait_not_started) 

201 if operation.operation_type != OperationType.WAIT: 

202 msg_not_wait: str = ( 

203 f"Expected WAIT operation, got {operation.operation_type}" 

204 ) 

205 raise IllegalStateException(msg_not_wait) 

206 

207 self._token_sequence += 1 

208 self.operations[index] = replace( 

209 operation, 

210 status=OperationStatus.SUCCEEDED, 

211 end_timestamp=datetime.now(timezone.utc), 

212 ) 

213 return self.operations[index] 

214 

215 def complete_retry(self, operation_id: str) -> Operation: 

216 """Complete STEP retry when timer fires.""" 

217 index, operation = self.find_operation(operation_id) 

218 

219 # Validate 

220 if operation.status != OperationStatus.PENDING: 

221 msg_step_not_pending: str = f"Attempting to transition a Step Operation[{operation_id}] to READY when it's not PENDING" 

222 raise IllegalStateException(msg_step_not_pending) 

223 if operation.operation_type != OperationType.STEP: 

224 msg_not_step: str = ( 

225 f"Expected STEP operation, got {operation.operation_type}" 

226 ) 

227 raise IllegalStateException(msg_not_step) 

228 

229 self._token_sequence += 1 

230 new_step_details = None 

231 if operation.step_details: 

232 new_step_details = replace( 

233 operation.step_details, next_attempt_timestamp=None 

234 ) 

235 

236 updated_operation = replace( 

237 operation, status=OperationStatus.READY, step_details=new_step_details 

238 ) 

239 self.operations[index] = updated_operation 

240 return updated_operation 

241 

242 def complete_callback_success( 

243 self, callback_id: str, result: bytes | None = None 

244 ) -> Operation: 

245 """Complete CALLBACK operation with success.""" 

246 index, operation = self.find_callback_operation(callback_id) 

247 if operation.status != OperationStatus.STARTED: 

248 msg: str = f"Callback operation [{callback_id}] is not in STARTED state" 

249 raise IllegalStateException(msg) 

250 

251 updated_callback_details = None 

252 if operation.callback_details: 

253 updated_callback_details = replace( 

254 operation.callback_details, 

255 result=result.decode() if result else None, 

256 ) 

257 

258 self.operations[index] = replace( 

259 operation, 

260 status=OperationStatus.SUCCEEDED, 

261 end_timestamp=datetime.now(timezone.utc), 

262 callback_details=updated_callback_details, 

263 ) 

264 return self.operations[index] 

265 

266 def complete_callback_failure( 

267 self, callback_id: str, error: ErrorObject 

268 ) -> Operation: 

269 """Complete CALLBACK operation with failure.""" 

270 index, operation = self.find_callback_operation(callback_id) 

271 

272 if operation.status != OperationStatus.STARTED: 

273 msg: str = f"Callback operation [{callback_id}] is not in STARTED state" 

274 raise IllegalStateException(msg) 

275 

276 updated_callback_details = None 

277 if operation.callback_details: 

278 updated_callback_details = replace(operation.callback_details, error=error) 

279 

280 self.operations[index] = replace( 

281 operation, 

282 status=OperationStatus.FAILED, 

283 end_timestamp=datetime.now(timezone.utc), 

284 callback_details=updated_callback_details, 

285 ) 

286 return self.operations[index] 

287 

288 def complete_callback_timeout( 

289 self, callback_id: str, error: ErrorObject 

290 ) -> Operation: 

291 """Complete CALLBACK operation with timeout.""" 

292 index, operation = self.find_callback_operation(callback_id) 

293 

294 if operation.status != OperationStatus.STARTED: 

295 msg: str = f"Callback operation [{callback_id}] is not in STARTED state" 

296 raise IllegalStateException(msg) 

297 

298 self._token_sequence += 1 

299 updated_callback_details = None 

300 if operation.callback_details: 300 ↛ 303line 300 didn't jump to line 303 because the condition on line 300 was always true

301 updated_callback_details = replace(operation.callback_details, error=error) 

302 

303 self.operations[index] = replace( 

304 operation, 

305 status=OperationStatus.TIMED_OUT, 

306 end_timestamp=datetime.now(timezone.utc), 

307 callback_details=updated_callback_details, 

308 ) 

309 return self.operations[index] 

310 

311 def _end_execution(self, status: OperationStatus) -> None: 

312 """Set the end_timestamp on the main EXECUTION operation when execution completes.""" 

313 execution_op: Operation = self.get_operation_execution_started() 

314 if execution_op.operation_type == OperationType.EXECUTION: 

315 self.operations[0] = replace( 

316 execution_op, 

317 status=status, 

318 end_timestamp=datetime.now(timezone.utc), 

319 )