Coverage for async-durable-execution/src/async_durable_execution/runner/local/processor.py: 98%
133 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Checkpoint validation and operation transformation helpers."""
3from __future__ import annotations
5import json
6from collections.abc import Mapping
7from typing import TYPE_CHECKING
9from ...models import (
10 Operation,
11 OperationAction,
12 OperationType,
13 OperationUpdate,
14)
15from ..exceptions import (
16 InvalidParameterValueException,
17)
18from .processors import (
19 create_default_processors,
20)
21from .processors.base import (
22 OperationProcessor,
23)
24from .processors.invoke import ChainedInvokeProcessor
26if TYPE_CHECKING:
27 from collections.abc import MutableMapping
29 from .execution import Execution
31MAX_ERROR_PAYLOAD_SIZE_BYTES = 32768
34class OperationTransformer:
35 """Transforms OperationUpdates to Operations while maintaining order and triggering scheduler actions."""
37 def __init__(
38 self,
39 processors: MutableMapping[OperationType, OperationProcessor] | None = None,
40 ):
41 self.processors = processors or create_default_processors()
43 def mock_invoke_result(self, function_name: str, result: object) -> None:
44 """Register a local mock result for a chained invoke function."""
45 processor = self.processors.get(OperationType.CHAINED_INVOKE)
46 if isinstance(processor, ChainedInvokeProcessor):
47 processor.mock_result(function_name=function_name, result=result)
49 def process_updates(
50 self,
51 updates: list[OperationUpdate],
52 current_operations: list[Operation],
53 runner,
54 execution_arn: str,
55 ) -> tuple[list[Operation], list[OperationUpdate]]:
56 """Transform updates maintaining operation order and return (operations, updates)."""
57 op_map = {op.operation_id: op for op in current_operations}
59 # Start with copy of current operations list
60 result_operations = current_operations.copy()
62 for update in updates:
63 processor = self.processors.get(update.operation_type)
64 if processor:
65 current_op = op_map.get(update.operation_id)
66 updated_op = processor.process(
67 update=update,
68 current_op=current_op,
69 notifier=runner,
70 execution_arn=execution_arn,
71 )
73 if updated_op is not None:
74 if update.operation_id in op_map:
75 # Update existing operation in-place
76 for i, op in enumerate(result_operations): # pragma: no branch
77 # no branch coverage because result_operation empty not reachable here
78 if op.operation_id == update.operation_id:
79 result_operations[i] = updated_op
80 break
81 else:
82 # Append new operation to end
83 result_operations.append(updated_op)
85 # Update map for future lookups
86 op_map[update.operation_id] = updated_op
87 else:
88 msg: str = (
89 f"Checkpoint for {update.operation_type} is not implemented yet."
90 )
91 raise InvalidParameterValueException(msg)
93 return result_operations, updates
96class CheckpointValidator:
97 """Validates checkpoint input based on current state."""
99 @staticmethod
100 def validate_input(
101 updates: list[OperationUpdate],
102 execution: Execution,
103 processors: Mapping[OperationType, OperationProcessor] | None = None,
104 ) -> None:
105 """Perform validation on the given input based on the current state."""
106 if not updates:
107 return
109 CheckpointValidator._validate_conflicting_execution_update(updates)
110 CheckpointValidator._validate_parent_id_and_duplicate_id(updates, execution)
112 for update in updates:
113 CheckpointValidator._validate_operation_update(
114 update, execution, processors
115 )
117 @staticmethod
118 def _validate_conflicting_execution_update(updates: list[OperationUpdate]) -> None:
119 """Validate that there are no conflicting execution updates."""
120 execution_updates = [
121 update
122 for update in updates
123 if update.operation_type == OperationType.EXECUTION
124 ]
126 if len(execution_updates) > 1:
127 msg_multiple_exec: str = "Cannot checkpoint multiple EXECUTION updates."
129 raise InvalidParameterValueException(msg_multiple_exec)
131 if execution_updates and updates[-1].operation_type != OperationType.EXECUTION:
132 msg_exec_last: str = "EXECUTION checkpoint must be the last update."
134 raise InvalidParameterValueException(msg_exec_last)
136 @staticmethod
137 def _validate_operation_update(
138 update: OperationUpdate,
139 execution: Execution,
140 processors: Mapping[OperationType, OperationProcessor] | None,
141 ) -> None:
142 """Validate a single operation update."""
143 CheckpointValidator._validate_inconsistent_operation_metadata(update, execution)
144 CheckpointValidator._validate_payload_sizes(update)
145 CheckpointValidator._validate_operation_status_transition(
146 update, execution, processors
147 )
149 @staticmethod
150 def _validate_payload_sizes(update: OperationUpdate) -> None:
151 """Validate that operation payload sizes are not too large."""
152 if update.error is not None:
153 payload = json.dumps(update.error.to_dict())
154 if len(payload) > MAX_ERROR_PAYLOAD_SIZE_BYTES:
155 msg: str = f"Error object size must be less than {MAX_ERROR_PAYLOAD_SIZE_BYTES} bytes."
156 raise InvalidParameterValueException(msg)
158 @staticmethod
159 def _validate_operation_status_transition(
160 update: OperationUpdate,
161 execution: Execution,
162 processors: Mapping[OperationType, OperationProcessor] | None,
163 ) -> None:
164 """Validate that the operation status transition is valid."""
165 current_state = None
166 for operation in execution.operations:
167 if operation.operation_id == update.operation_id:
168 current_state = operation
169 break
171 processor = (processors or create_default_processors()).get(
172 update.operation_type
173 )
174 if processor is None:
175 msg: str = "Invalid operation type."
176 raise InvalidParameterValueException(msg)
178 processor.validate(current_state, update)
180 @staticmethod
181 def _validate_inconsistent_operation_metadata(
182 update: OperationUpdate, execution: Execution
183 ) -> None:
184 """Validate that operation metadata is consistent with existing operation."""
185 current_state = None
186 for operation in execution.operations:
187 if operation.operation_id == update.operation_id:
188 current_state = operation
189 break
191 if current_state is not None:
192 if (
193 update.operation_type is not None
194 and update.operation_type != current_state.operation_type
195 ):
196 msg: str = "Inconsistent operation type."
197 raise InvalidParameterValueException(msg)
199 if (
200 update.sub_type is not None
201 and update.sub_type != current_state.sub_type
202 ):
203 msg_subtype: str = "Inconsistent operation subtype."
204 raise InvalidParameterValueException(msg_subtype)
206 if update.name is not None and update.name != current_state.name:
207 msg_name: str = "Inconsistent operation name."
208 raise InvalidParameterValueException(msg_name)
210 if (
211 update.parent_id is not None
212 and update.parent_id != current_state.parent_id
213 ):
214 msg_parent: str = "Inconsistent parent operation id."
215 raise InvalidParameterValueException(msg_parent)
217 @staticmethod
218 def _validate_parent_id_and_duplicate_id(
219 updates: list[OperationUpdate], execution: Execution
220 ) -> None:
221 """Validate parent IDs and check for duplicate operation IDs.
223 Validate that any provided parentId is valid, and also validate no duplicate operation is being
224 updated at the same time (unless it is a STEP/CONTEXT starting + performing one more non-START action).
225 """
226 operations_started: MutableMapping[str, OperationUpdate] = {}
227 last_updates_seen: MutableMapping[str, OperationUpdate] = {}
229 for update in updates:
230 if CheckpointValidator._is_invalid_duplicate_update(
231 update, last_updates_seen
232 ):
233 msg_duplicate: str = (
234 "Cannot checkpoint multiple operations with the same ID."
235 )
236 raise InvalidParameterValueException(msg_duplicate)
238 if not CheckpointValidator._is_valid_parent_for_update(
239 execution, update, operations_started
240 ):
241 msg_parent: str = "Invalid parent operation id."
242 raise InvalidParameterValueException(msg_parent)
244 if update.action == OperationAction.START:
245 operations_started[update.operation_id] = update
247 last_updates_seen[update.operation_id] = update
249 @staticmethod
250 def _is_invalid_duplicate_update(
251 update: OperationUpdate, last_updates_seen: MutableMapping[str, OperationUpdate]
252 ) -> bool:
253 """Check if this is an invalid duplicate update."""
254 last_update = last_updates_seen.get(update.operation_id)
255 if last_update is None:
256 return False
258 if last_update.operation_type in (OperationType.STEP, OperationType.CONTEXT):
259 # Allow duplicate for STEP/CONTEXT if last was START and current is not START
260 allow_duplicate = (
261 last_update.action == OperationAction.START
262 and update.action != OperationAction.START
263 )
264 return not allow_duplicate
266 return True
268 @staticmethod
269 def _is_valid_parent_for_update(
270 execution: Execution,
271 update: OperationUpdate,
272 operations_started: MutableMapping[str, OperationUpdate],
273 ) -> bool:
274 """Check if the parent ID is valid for the update."""
275 parent_id = update.parent_id
277 if parent_id is None:
278 return True
280 # Check if parent is in operations started in this batch
281 if parent_id in operations_started:
282 parent_update = operations_started[parent_id]
283 return parent_update.operation_type == OperationType.CONTEXT
285 # Check if parent exists in current execution state
286 for operation in execution.operations:
287 if operation.operation_id == parent_id:
288 return operation.operation_type == OperationType.CONTEXT
290 return False