Coverage for async-durable-execution/src/async_durable_execution/runner/local/execution.py: 100%
181 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1from __future__ import annotations
3from dataclasses import replace
4from datetime import datetime, timezone
5from enum import Enum
6from typing import Any
7from uuid import uuid4
9from ...execution import (
10 DurableExecutionInvocationOutput,
11 InvocationStatus,
12)
13from ...models import (
14 ErrorObject,
15 ExecutionDetails,
16 Operation,
17 OperationStatus,
18 OperationType,
19 OperationUpdate,
20)
21from ..exceptions import (
22 IllegalStateException,
23 InvalidParameterValueException,
24)
26# Import AWS exceptions
27from ..model import (
28 InvocationCompletedDetails,
29)
30from .model import (
31 StartDurableExecutionInput,
32 CheckpointToken,
33)
36class ExecutionStatus(Enum):
37 """Execution status for API responses."""
39 RUNNING = "RUNNING"
40 SUCCEEDED = "SUCCEEDED"
41 FAILED = "FAILED"
42 STOPPED = "STOPPED"
43 TIMED_OUT = "TIMED_OUT"
46class Execution:
47 """Execution state."""
49 def __init__(
50 self,
51 durable_execution_arn: str,
52 start_input: StartDurableExecutionInput,
53 operations: list[Operation],
54 ):
55 self.durable_execution_arn: str = durable_execution_arn
56 # operation is frozen, it won't mutate - no need to clone/deep-copy
57 self.start_input: StartDurableExecutionInput = start_input
58 self.operations: list[Operation] = operations
59 self.updates: list[OperationUpdate] = []
60 self.invocation_completions: list[InvocationCompletedDetails] = []
61 self.used_tokens: set[str] = set()
62 # TODO: this will need to persist/rehydrate depending on inmemory vs sqllite store
63 self._token_sequence: int = 0
64 self.is_complete: bool = False
65 self.result: DurableExecutionInvocationOutput | None = None
66 self.consecutive_failed_invocation_attempts: int = 0
67 self.close_status: ExecutionStatus | None = None
69 @property
70 def token_sequence(self) -> int:
71 """Get current token sequence value."""
72 return self._token_sequence
74 def current_status(self) -> ExecutionStatus:
75 """Get execution status."""
76 if not self.is_complete:
77 return ExecutionStatus.RUNNING
79 if not self.close_status:
80 msg: str = "close_status must be set when execution is complete"
81 raise IllegalStateException(msg)
83 return self.close_status
85 @staticmethod
86 def new(input: StartDurableExecutionInput) -> Execution: # noqa: A002
87 # make a nicer arn
88 # 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-_\.]+
89 # Example: arn:aws:lambda:us-east-1:123456789012:durable-execution:myDurableFunction:myDurableExecutionName:ce67da72-3701-4f83-9174-f4189d27b0a5
90 return Execution(
91 durable_execution_arn=str(uuid4())
92 + "/"
93 + (input.invocation_id or str(uuid4())),
94 start_input=input,
95 operations=[],
96 )
98 def to_json_dict(self) -> dict[str, Any]:
99 """Serialize execution to JSON-serializable dictionary"""
100 return {
101 "DurableExecutionArn": self.durable_execution_arn,
102 "StartInput": self.start_input.to_dict(),
103 "Operations": [op.to_json_dict() for op in self.operations],
104 "Updates": [update.to_dict() for update in self.updates],
105 "InvocationCompletions": [
106 completion.to_json_dict() for completion in self.invocation_completions
107 ],
108 "UsedTokens": list(self.used_tokens),
109 "TokenSequence": self._token_sequence,
110 "IsComplete": self.is_complete,
111 "Result": self.result.to_dict() if self.result else None,
112 "ConsecutiveFailedInvocationAttempts": self.consecutive_failed_invocation_attempts,
113 "CloseStatus": self.close_status.value if self.close_status else None,
114 }
116 @classmethod
117 def from_json_dict(cls, data: dict[str, Any]) -> Execution:
118 """Deserialize execution from dictionary."""
119 # Reconstruct start_input
120 start_input = StartDurableExecutionInput.from_dict(data["StartInput"])
122 # Reconstruct operations
123 operations = [
124 Operation.from_json_dict(op_data) for op_data in data["Operations"]
125 ]
127 # Create execution
128 execution = cls(
129 durable_execution_arn=data["DurableExecutionArn"],
130 start_input=start_input,
131 operations=operations,
132 )
134 # Set additional fields
135 execution.updates = [
136 OperationUpdate.from_dict(update_data) for update_data in data["Updates"]
137 ]
138 execution.invocation_completions = [
139 InvocationCompletedDetails.from_json_dict(item)
140 for item in data.get("InvocationCompletions", [])
141 ]
142 execution.used_tokens = set(data["UsedTokens"])
143 execution._token_sequence = data["TokenSequence"]
144 execution.is_complete = data["IsComplete"]
145 execution.result = (
146 DurableExecutionInvocationOutput.from_dict(data["Result"])
147 if data["Result"]
148 else None
149 )
150 execution.consecutive_failed_invocation_attempts = data[
151 "ConsecutiveFailedInvocationAttempts"
152 ]
153 close_status_str = data.get("CloseStatus")
154 execution.close_status = (
155 ExecutionStatus(close_status_str) if close_status_str else None
156 )
158 return execution
160 def start(self) -> None:
161 if self.start_input.invocation_id is None:
162 msg: str = "invocation_id is required"
163 raise InvalidParameterValueException(msg)
164 self.operations.append(
165 Operation(
166 operation_id=self.start_input.invocation_id,
167 parent_id=None,
168 name=self.start_input.execution_name,
169 start_timestamp=datetime.now(timezone.utc),
170 operation_type=OperationType.EXECUTION,
171 status=OperationStatus.STARTED,
172 execution_details=ExecutionDetails(
173 input_payload=self.start_input.get_normalized_input()
174 ),
175 )
176 )
178 def get_operation_execution_started(self) -> Operation:
179 if not self.operations:
180 msg: str = "execution not started."
182 raise IllegalStateException(msg)
184 return self.operations[0]
186 def get_new_checkpoint_token(self) -> str:
187 """Generate a new checkpoint token with incremented sequence"""
188 self._token_sequence += 1
189 new_token_sequence = self._token_sequence
190 token = CheckpointToken(
191 execution_arn=self.durable_execution_arn,
192 token_sequence=new_token_sequence,
193 )
194 token_str = token.to_str()
195 self.used_tokens.add(token_str)
196 return token_str
198 def get_navigable_operations(self) -> list[Operation]:
199 """Get list of operations, but exclude child operations where the parent has already completed."""
200 return self.operations
202 def get_assertable_operations(self) -> list[Operation]:
203 """Get list of operations, but exclude the EXECUTION operations"""
204 # TODO: this excludes EXECUTION at start, but can there be an EXECUTION at the end if there was a checkpoint with large payload?
205 return self.operations[1:]
207 def has_pending_operations(self) -> bool:
208 """True if execution has pending operations."""
210 for operation in self.operations:
211 if (
212 operation.operation_type == OperationType.STEP
213 and operation.status == OperationStatus.PENDING
214 ) or (
215 operation.operation_type
216 in [
217 OperationType.WAIT,
218 OperationType.CALLBACK,
219 OperationType.CHAINED_INVOKE,
220 ]
221 and operation.status == OperationStatus.STARTED
222 ):
223 return True
224 return False
226 def record_invocation_completion(
227 self, start_timestamp: datetime, end_timestamp: datetime, request_id: str
228 ) -> None:
229 """Record an invocation completion event."""
230 self.invocation_completions.append(
231 InvocationCompletedDetails(
232 start_timestamp=start_timestamp,
233 end_timestamp=end_timestamp,
234 request_id=request_id,
235 )
236 )
238 def complete_success(self, result: str | None) -> None:
239 """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""
240 self.result = DurableExecutionInvocationOutput(
241 status=InvocationStatus.SUCCEEDED, result=result
242 )
243 self.is_complete = True
244 self.close_status = ExecutionStatus.SUCCEEDED
245 self._end_execution(OperationStatus.SUCCEEDED)
247 def complete_fail(self, error: ErrorObject) -> None:
248 """Complete execution with failure (DecisionType.FAIL_WORKFLOW_EXECUTION)."""
249 self.result = DurableExecutionInvocationOutput(
250 status=InvocationStatus.FAILED, error=error
251 )
252 self.is_complete = True
253 self.close_status = ExecutionStatus.FAILED
254 self._end_execution(OperationStatus.FAILED)
256 def complete_timeout(self, error: ErrorObject) -> None:
257 """Complete execution with timeout."""
258 self.result = DurableExecutionInvocationOutput(
259 status=InvocationStatus.FAILED, error=error
260 )
261 self.is_complete = True
262 self.close_status = ExecutionStatus.TIMED_OUT
263 self._end_execution(OperationStatus.TIMED_OUT)
265 def complete_stopped(self, error: ErrorObject) -> None:
266 """Complete execution as terminated (TerminateWorkflowExecutionV2Request)."""
267 self.result = DurableExecutionInvocationOutput(
268 status=InvocationStatus.FAILED, error=error
269 )
270 self.is_complete = True
271 self.close_status = ExecutionStatus.STOPPED
272 self._end_execution(OperationStatus.STOPPED)
274 def find_operation(self, operation_id: str) -> tuple[int, Operation]:
275 """Find operation by ID, return index and operation."""
276 for i, operation in enumerate(self.operations):
277 if operation.operation_id == operation_id:
278 return i, operation
279 msg: str = f"Attempting to update state of an Operation [{operation_id}] that doesn't exist"
280 raise IllegalStateException(msg)
282 def find_callback_operation(self, callback_id: str) -> tuple[int, Operation]:
283 """Find callback operation by callback_id, return index and operation."""
284 for i, operation in enumerate(self.operations):
285 if (
286 operation.operation_type == OperationType.CALLBACK
287 and operation.callback_details
288 and operation.callback_details.callback_id == callback_id
289 ):
290 return i, operation
291 msg: str = f"Callback operation with callback_id [{callback_id}] not found"
292 raise IllegalStateException(msg)
294 def complete_wait(self, operation_id: str) -> Operation:
295 """Complete WAIT operation when timer fires."""
296 index, operation = self.find_operation(operation_id)
298 # Validate
299 if operation.status != OperationStatus.STARTED:
300 msg_wait_not_started: str = f"Attempting to transition a Wait Operation[{operation_id}] to SUCCEEDED when it's not STARTED"
301 raise IllegalStateException(msg_wait_not_started)
302 if operation.operation_type != OperationType.WAIT:
303 msg_not_wait: str = (
304 f"Expected WAIT operation, got {operation.operation_type}"
305 )
306 raise IllegalStateException(msg_not_wait)
308 self._token_sequence += 1
309 self.operations[index] = replace(
310 operation,
311 status=OperationStatus.SUCCEEDED,
312 end_timestamp=datetime.now(timezone.utc),
313 )
314 return self.operations[index]
316 def complete_retry(self, operation_id: str) -> Operation:
317 """Complete STEP retry when timer fires."""
318 index, operation = self.find_operation(operation_id)
320 # Validate
321 if operation.status != OperationStatus.PENDING:
322 msg_step_not_pending: str = f"Attempting to transition a Step Operation[{operation_id}] to READY when it's not PENDING"
323 raise IllegalStateException(msg_step_not_pending)
324 if operation.operation_type != OperationType.STEP:
325 msg_not_step: str = (
326 f"Expected STEP operation, got {operation.operation_type}"
327 )
328 raise IllegalStateException(msg_not_step)
330 self._token_sequence += 1
331 new_step_details = None
332 if operation.step_details:
333 new_step_details = replace(
334 operation.step_details, next_attempt_timestamp=None
335 )
337 updated_operation = replace(
338 operation, status=OperationStatus.READY, step_details=new_step_details
339 )
340 self.operations[index] = updated_operation
341 return updated_operation
343 def complete_callback_success(
344 self, callback_id: str, result: bytes | None = None
345 ) -> Operation:
346 """Complete CALLBACK operation with success."""
347 index, operation = self.find_callback_operation(callback_id)
348 if operation.status != OperationStatus.STARTED:
349 msg: str = f"Callback operation [{callback_id}] is not in STARTED state"
350 raise IllegalStateException(msg)
352 updated_callback_details = None
353 if operation.callback_details:
354 updated_callback_details = replace(
355 operation.callback_details,
356 result=result.decode() if result else None,
357 )
359 self.operations[index] = replace(
360 operation,
361 status=OperationStatus.SUCCEEDED,
362 end_timestamp=datetime.now(timezone.utc),
363 callback_details=updated_callback_details,
364 )
365 return self.operations[index]
367 def complete_callback_failure(
368 self, callback_id: str, error: ErrorObject
369 ) -> Operation:
370 """Complete CALLBACK operation with failure."""
371 index, operation = self.find_callback_operation(callback_id)
373 if operation.status != OperationStatus.STARTED:
374 msg: str = f"Callback operation [{callback_id}] is not in STARTED state"
375 raise IllegalStateException(msg)
377 updated_callback_details = None
378 if operation.callback_details:
379 updated_callback_details = replace(operation.callback_details, error=error)
381 self.operations[index] = replace(
382 operation,
383 status=OperationStatus.FAILED,
384 end_timestamp=datetime.now(timezone.utc),
385 callback_details=updated_callback_details,
386 )
387 return self.operations[index]
389 def complete_callback_timeout(
390 self, callback_id: str, error: ErrorObject
391 ) -> Operation:
392 """Complete CALLBACK operation with timeout."""
393 index, operation = self.find_callback_operation(callback_id)
395 if operation.status != OperationStatus.STARTED:
396 msg: str = f"Callback operation [{callback_id}] is not in STARTED state"
397 raise IllegalStateException(msg)
399 self._token_sequence += 1
400 updated_callback_details = None
401 if operation.callback_details:
402 updated_callback_details = replace(operation.callback_details, error=error)
404 self.operations[index] = replace(
405 operation,
406 status=OperationStatus.TIMED_OUT,
407 end_timestamp=datetime.now(timezone.utc),
408 callback_details=updated_callback_details,
409 )
410 return self.operations[index]
412 def _end_execution(self, status: OperationStatus) -> None:
413 """Set the end_timestamp on the main EXECUTION operation when execution completes."""
414 execution_op: Operation = self.get_operation_execution_started()
415 if execution_op.operation_type == OperationType.EXECUTION:
416 self.operations[0] = replace(
417 execution_op,
418 status=status,
419 end_timestamp=datetime.now(timezone.utc),
420 )