Coverage for async_durable_execution/_primitive/child.py: 90%
124 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Implementation for run_in_child_context."""
3from __future__ import annotations
5import asyncio
6import logging
7from collections.abc import Callable
8from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast
10from .base import OperationExecutor
11from .._core import (
12 CallableRuntimeError,
13 ContextOptions,
14 DurableContext,
15 ErrorObject,
16 ExecutionError,
17 ExecutionState,
18 InvocationError,
19 Operation,
20 OperationIdentifier,
21 OperationStatus,
22 OperationSubTypeValue,
23 OperationType,
24 OperationUpdate,
25 SerDes,
26 _encode_sdk_control_error_data,
27 _restore_sdk_control_error,
28 bind_current_context,
29 create_eager_task,
30 deserialize,
31 get_durable_context,
32 serialize,
33)
35if TYPE_CHECKING:
36 from collections.abc import Awaitable
38logger = logging.getLogger(__name__)
40T = TypeVar("T")
41C_contra = TypeVar("C_contra", contravariant=True)
43SummaryGenerator: TypeAlias = Callable[[C_contra], str]
44"""Create a compact JSON summary for an oversized child context result."""
46# Checkpoint size limit in bytes (256KB)
47CHECKPOINT_SIZE_LIMIT = 256 * 1024
50class ChildOperationExecutor(OperationExecutor[T]):
51 """Executor for child context operations."""
53 SERDES_OPERATION_TYPE = OperationType.CONTEXT
55 def __init__(
56 self,
57 func: Callable[[], Awaitable[T]],
58 state: ExecutionState,
59 operation_identifier: OperationIdentifier,
60 *,
61 serdes: SerDes | None = None,
62 summary_generator: SummaryGenerator | None = None,
63 is_virtual: bool = False,
64 ) -> None:
65 """Initialize the child operation executor.
67 Args:
68 func: The child context function to execute
69 state: The execution state
70 operation_identifier: The operation identifier
71 serdes: Optional serializer for the child context result.
72 summary_generator: Optional summary generator for large child results.
73 is_virtual: Whether this child context should skip lifecycle checkpoints.
74 """
75 super().__init__(state=state, operation_identifier=operation_identifier)
76 self.func = func
77 self.serdes = serdes
78 self.summary_generator = summary_generator
79 self.is_virtual = is_virtual
81 async def start(self) -> T:
82 """Start a new child context operation."""
83 if not self.is_virtual:
84 start_operation: OperationUpdate = OperationUpdate.create_context_start(
85 identifier=self.operation_identifier,
86 sub_type=self.operation_identifier.sub_type,
87 )
88 await self.create_checkpoint(start_operation, is_sync=False)
90 return await self.execute(None)
92 async def replay(self, operation: Operation) -> T:
93 """Replay an existing child context operation from its checkpoint."""
94 if (
95 operation.status is OperationStatus.SUCCEEDED
96 and not self._is_replay_children(operation)
97 ):
98 logger.debug(
99 "Child context already completed, skipping execution for id: %s, name: %s",
100 self.operation_identifier.operation_id,
101 self.operation_name,
102 )
103 result_payload = self._get_result(operation)
104 if result_payload is None:
105 return cast("T", None)
107 result: T = await deserialize(
108 serdes=self.serdes,
109 data=result_payload,
110 operation_id=self.operation_id,
111 durable_execution_arn=self.durable_execution_arn,
112 recursive_level=self.state.recursive_level,
113 operation_name=operation.name,
114 parent_id=operation.parent_id,
115 operation_type=operation.operation_type,
116 operation_sub_type=operation.sub_type,
117 )
118 return result
120 if operation.status is OperationStatus.SUCCEEDED and self._is_replay_children(
121 operation
122 ):
123 return await self.execute(operation)
125 if operation.status is OperationStatus.FAILED:
126 self._raise_callable_error(operation)
128 return await self.execute(operation)
130 async def execute(self, operation: Operation | None) -> T:
131 """Execute child context function with error handling and large payload support.
133 Args:
134 operation: The checkpointed operation state, if any
136 Returns:
137 The result of executing the child context function
139 Raises:
140 SuspendExecution: Re-raised without checkpointing
141 InvocationError: Re-raised without checkpointing when retryable
142 ExecutionError: Re-raised after checkpointing FAIL
143 CallableRuntimeError: Raised for other exceptions after checkpointing FAIL
144 """
145 logger.debug(
146 "▶️ Executing child context for id: %s, name: %s",
147 self.operation_identifier.operation_id,
148 self.operation_identifier.name,
149 )
150 try:
151 replaying_children = self._is_replay_children(operation)
152 raw_result: T = await self.func()
154 if self.is_virtual:
155 logger.debug(
156 "Virtual context: Exiting child context without creating another checkpoint. id: %s, name: %s",
157 self.operation_identifier.operation_id,
158 self.operation_identifier.name,
159 )
160 return raw_result
162 # If in replay_children mode, return without checkpointing
163 if replaying_children:
164 logger.debug(
165 "ReplayChildren mode: Executed child context again on replay due to large payload. Exiting child context without creating another checkpoint. id: %s, name: %s",
166 self.operation_identifier.operation_id,
167 self.operation_identifier.name,
168 )
169 return raw_result
171 # Serialize result
172 serialized_result: str = await serialize(
173 serdes=self.serdes,
174 value=raw_result,
175 operation_id=self.operation_id,
176 durable_execution_arn=self.durable_execution_arn,
177 recursive_level=self.state.recursive_level,
178 operation_name=self.operation_identifier.name,
179 parent_id=self.operation_identifier.parent_id,
180 operation_type=self.SERDES_OPERATION_TYPE,
181 operation_sub_type=self.operation_identifier.sub_type,
182 )
184 # Check payload size and use ReplayChildren mode if needed
185 # Summary Generator Logic:
186 # When the serialized result exceeds 256KB, we use ReplayChildren mode to avoid
187 # checkpointing large payloads. Instead, we checkpoint a compact summary and mark
188 # the operation for replay. This matches the TypeScript implementation behavior.
189 #
190 # See TypeScript reference:
191 # - aws-durable-execution-sdk-js/src/handlers/run-in-child-context-handler/run-in-child-context-handler.ts (lines ~200-220)
192 #
193 # The summary generator creates a JSON summary with metadata (type, counts, status)
194 # instead of the full BatchResult. During replay, the child context is re-executed
195 # to reconstruct the full result rather than deserializing from the checkpoint.
196 replay_children: bool = False
197 if len(serialized_result) > CHECKPOINT_SIZE_LIMIT:
198 logger.debug(
199 "Large payload detected, using ReplayChildren mode: id: %s, name: %s, payload_size: %d, limit: %d",
200 self.operation_identifier.operation_id,
201 self.operation_identifier.name,
202 len(serialized_result),
203 CHECKPOINT_SIZE_LIMIT,
204 )
205 replay_children = True
206 # Use summary generator if provided, otherwise use empty string (matches TypeScript)
207 serialized_result = (
208 self.summary_generator(raw_result) if self.summary_generator else ""
209 )
211 # Checkpoint SUCCEED
212 success_operation: OperationUpdate = OperationUpdate.create_context_succeed(
213 identifier=self.operation_identifier,
214 payload=serialized_result,
215 sub_type=self.operation_identifier.sub_type,
216 context_options=ContextOptions(replay_children=replay_children),
217 )
218 # Checkpoint child context SUCCEED with blocking (is_sync=True, default).
219 # Must ensure the child context result is persisted before returning to the parent.
220 # This guarantees the result is durable and child operations won't be re-executed on replay
221 # (unless replay_children=True for large payloads).
222 await self.create_checkpoint(success_operation)
224 logger.debug(
225 "✅ Successfully completed child context for id: %s, name: %s",
226 self.operation_identifier.operation_id,
227 self.operation_identifier.name,
228 )
229 except Exception as e:
230 if isinstance(e, InvocationError) and e.is_retryable():
231 raise
233 error_object = ErrorObject.from_exception(e)
234 sdk_error_data = _encode_sdk_control_error_data(e)
235 if sdk_error_data is not None:
236 error_object = ErrorObject(
237 message=error_object.message,
238 type=error_object.type,
239 data=sdk_error_data,
240 stack_trace=error_object.stack_trace,
241 )
243 # Virtual deliberately does not write checkpoints, but exception still propagates below
244 if not self.is_virtual:
245 fail_operation: OperationUpdate = OperationUpdate.create_context_fail(
246 identifier=self.operation_identifier,
247 error=error_object,
248 sub_type=self.operation_identifier.sub_type,
249 )
250 # Checkpoint child context FAIL with blocking (is_sync=True, default).
251 # Must ensure the failure state is persisted before raising the exception.
252 # This guarantees the error is durable and child operations won't be re-executed on replay.
253 await self.create_checkpoint(fail_operation)
255 # Preserve SDK control errors for the top-level execution handler
256 # after checkpointing their failure.
257 if isinstance(e, InvocationError | ExecutionError):
258 raise
259 if sdk_error_data is not None: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 control_error = _restore_sdk_control_error(
261 error_object.message or str(e),
262 error_object.type,
263 sdk_error_data,
264 )
265 if control_error is not None:
266 raise control_error from e
267 raise CallableRuntimeError.from_error_object(error_object) from e
269 if replay_children:
270 return raw_result
272 # SUCCEED is already durable. Deserialization failures must propagate
273 # without attempting a terminal FAIL transition for the same context.
274 return await deserialize(
275 serdes=self.serdes,
276 data=serialized_result,
277 operation_id=self.operation_id,
278 durable_execution_arn=self.durable_execution_arn,
279 recursive_level=self.state.recursive_level,
280 operation_name=self.operation_identifier.name,
281 parent_id=self.operation_identifier.parent_id,
282 operation_type=self.SERDES_OPERATION_TYPE,
283 operation_sub_type=self.operation_identifier.sub_type,
284 )
286 @staticmethod
287 def _is_replay_children(operation: Operation | None) -> bool:
288 if operation is None or operation.context_details is None:
289 return False
290 return operation.context_details.replay_children
292 @staticmethod
293 def _get_result(operation: Operation) -> str | None:
294 if operation.context_details is None:
295 return None
296 return operation.context_details.result
298 @staticmethod
299 def _raise_callable_error(operation: Operation) -> None:
300 error = operation.context_details.error if operation.context_details else None
301 if error is None:
302 msg = "Unknown error. No ErrorObject exists on the Checkpoint Operation."
303 raise CallableRuntimeError(
304 message=msg,
305 error_type=None,
306 data=None,
307 stack_trace=None,
308 )
310 control_error = _restore_sdk_control_error(
311 error.message or "Child context failed",
312 error.type,
313 error.data,
314 )
315 if control_error is not None:
316 raise control_error
318 raise CallableRuntimeError.from_error_object(error)
321def run_in_child_context(
322 func: Callable[[], Awaitable[T]],
323 *,
324 name: str | None = None,
325 serdes: SerDes | None = None,
326 summary_generator: SummaryGenerator | None = None,
327 is_virtual: bool = False,
328) -> asyncio.Task[T]:
329 """Execute a durable sub-workflow inside its own child context.
331 Args:
332 func: The child context function to execute.
333 name: Optional durable operation name.
334 serdes: Optional serializer for the child context result.
335 summary_generator: Optional summary generator for large child results.
336 is_virtual: Whether the child context should skip lifecycle checkpoints.
337 """
339 from .._operation.child import run_in_child_context as operation_child_context
341 return operation_child_context(
342 func,
343 name=name,
344 serdes=serdes,
345 summary_generator=summary_generator,
346 is_virtual=is_virtual,
347 )
350def _create_child_context_task(
351 func: Callable[[], Awaitable[T]],
352 *,
353 sub_type: OperationSubTypeValue,
354 name: str | None = None,
355 serdes: SerDes | None = None,
356 summary_generator: SummaryGenerator | None = None,
357 is_virtual: bool = False,
358) -> asyncio.Task[T]:
359 context = get_durable_context()
361 operation_id = context.step_counter.create_step_id()
362 child_context = context.create_child_context(
363 operation_id=operation_id,
364 is_virtual=is_virtual,
365 )
366 operation_identifier = OperationIdentifier(
367 operation_id=operation_id,
368 sub_type=sub_type,
369 parent_id=context.parent_id,
370 name=name,
371 )
373 async def run_child_context() -> T:
374 with context._replay_aware():
375 return await _run_child_context(
376 func,
377 context=context,
378 child_context=child_context,
379 operation_identifier=operation_identifier,
380 serdes=serdes,
381 summary_generator=summary_generator,
382 is_virtual=is_virtual,
383 )
385 return create_eager_task(run_child_context)
388async def _run_in_child_context(
389 func: Callable[[], Awaitable[T]],
390 *,
391 sub_type: OperationSubTypeValue,
392 name: str | None = None,
393 serdes: SerDes | None = None,
394 summary_generator: SummaryGenerator | None = None,
395 is_virtual: bool = False,
396) -> T:
397 """Execute a durable sub-workflow with an explicit operation subtype."""
398 context = get_durable_context()
399 with context._replay_aware():
400 operation_id = context.step_counter.create_step_id()
402 child_context = context.create_child_context(
403 operation_id=operation_id,
404 is_virtual=is_virtual,
405 )
406 operation_identifier = OperationIdentifier(
407 operation_id=operation_id,
408 sub_type=sub_type,
409 parent_id=context.parent_id,
410 name=name,
411 )
413 return await _run_child_context(
414 func,
415 context=context,
416 child_context=child_context,
417 operation_identifier=operation_identifier,
418 serdes=serdes,
419 summary_generator=summary_generator,
420 is_virtual=is_virtual,
421 )
424async def _run_child_context(
425 func: Callable[[], Awaitable[T]],
426 *,
427 context: DurableContext,
428 child_context: DurableContext,
429 operation_identifier: OperationIdentifier,
430 serdes: SerDes | None = None,
431 summary_generator: SummaryGenerator | None = None,
432 is_virtual: bool = False,
433) -> T:
434 async def callable_with_child_context() -> T:
435 with bind_current_context(child_context):
436 return await func()
438 executor: ChildOperationExecutor[T] = ChildOperationExecutor(
439 callable_with_child_context,
440 context.execution_state,
441 operation_identifier,
442 serdes=serdes,
443 summary_generator=summary_generator,
444 is_virtual=is_virtual,
445 )
446 return await executor.process()