aws-durable-execution-sdk-python 1.0.0__py3-none-any.whl

aws_durable_execution_sdk_python/logger.py

@@ -0,0 +1,131 @@
+"""Custom logging."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.types import LoggerInterface
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, MutableMapping
+
+    from aws_durable_execution_sdk_python.context import ExecutionState
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+
+
+@dataclass(frozen=True)
+class LogInfo:
+    execution_state: ExecutionState
+    parent_id: str | None = None
+    operation_id: str | None = None
+    name: str | None = None
+    attempt: int | None = None
+
+    @classmethod
+    def from_operation_identifier(
+        cls,
+        execution_state: ExecutionState,
+        op_id: OperationIdentifier,
+        attempt: int | None = None,
+    ) -> LogInfo:
+        """Create new log info from an execution arn, OperationIdentifier and attempt."""
+        return cls(
+            execution_state=execution_state,
+            parent_id=op_id.parent_id,
+            operation_id=op_id.operation_id,
+            name=op_id.name,
+            attempt=attempt,
+        )
+
+    def with_parent_id(self, parent_id: str) -> LogInfo:
+        """Clone the log info with a new parent id."""
+        return LogInfo(
+            execution_state=self.execution_state,
+            parent_id=parent_id,
+            operation_id=self.operation_id,
+            name=self.name,
+            attempt=self.attempt,
+        )
+
+
+class Logger(LoggerInterface):
+    def __init__(
+        self,
+        logger: LoggerInterface,
+        default_extra: Mapping[str, object],
+        execution_state: ExecutionState,
+    ) -> None:
+        self._logger = logger
+        self._default_extra = default_extra
+        self._execution_state = execution_state
+
+    @classmethod
+    def from_log_info(cls, logger: LoggerInterface, info: LogInfo) -> Logger:
+        """Create a new logger with the given LogInfo."""
+        extra: MutableMapping[str, object] = {
+            "executionArn": info.execution_state.durable_execution_arn
+        }
+        if info.parent_id:
+            extra["parentId"] = info.parent_id
+        if info.name:
+            # Use 'operation_name' instead of 'name' as key because the stdlib LogRecord internally reserved 'name' parameter
+            extra["operationName"] = info.name
+        if info.attempt is not None:
+            extra["attempt"] = info.attempt + 1
+        if info.operation_id:
+            extra["operationId"] = info.operation_id
+        return cls(
+            logger=logger, default_extra=extra, execution_state=info.execution_state
+        )
+
+    def with_log_info(self, info: LogInfo) -> Logger:
+        """Clone the existing logger with new LogInfo."""
+        return Logger.from_log_info(
+            logger=self._logger,
+            info=info,
+        )
+
+    def get_logger(self) -> LoggerInterface:
+        """Get the underlying logger."""
+        return self._logger
+
+    def debug(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._log(self._logger.debug, msg, *args, extra=extra)
+
+    def info(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._log(self._logger.info, msg, *args, extra=extra)
+
+    def warning(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._log(self._logger.warning, msg, *args, extra=extra)
+
+    def error(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._log(self._logger.error, msg, *args, extra=extra)
+
+    def exception(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._log(self._logger.exception, msg, *args, extra=extra)
+
+    def _log(
+        self,
+        log_func: Callable,
+        msg: object,
+        *args: object,
+        extra: Mapping[str, object] | None = None,
+    ):
+        if not self._should_log():
+            return
+        merged_extra = {**self._default_extra, **(extra or {})}
+        log_func(msg, *args, extra=merged_extra)
+
+    def _should_log(self) -> bool:
+        return not self._execution_state.is_replaying()

aws_durable_execution_sdk_python/operation/__init__.py

@@ -0,0 +1 @@
+"""Operation modules."""

aws_durable_execution_sdk_python/operation/callback.py

@@ -0,0 +1,123 @@
+"""Implementation for the Durable create_callback and wait_for_callback operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from aws_durable_execution_sdk_python.config import StepConfig
+from aws_durable_execution_sdk_python.exceptions import CallbackError
+from aws_durable_execution_sdk_python.lambda_service import (
+    CallbackOptions,
+    OperationUpdate,
+)
+from aws_durable_execution_sdk_python.types import WaitForCallbackContext
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python.config import (
+        CallbackConfig,
+        WaitForCallbackConfig,
+    )
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+    from aws_durable_execution_sdk_python.state import (
+        CheckpointedResult,
+        ExecutionState,
+    )
+    from aws_durable_execution_sdk_python.types import (
+        Callback,
+        DurableContext,
+        StepContext,
+    )
+
+
+def create_callback_handler(
+    state: ExecutionState,
+    operation_identifier: OperationIdentifier,
+    config: CallbackConfig | None = None,
+) -> str:
+    """Create the callback checkpoint and return the callback id."""
+    callback_options: CallbackOptions = (
+        CallbackOptions(
+            timeout_seconds=config.timeout_seconds,
+            heartbeat_timeout_seconds=config.heartbeat_timeout_seconds,
+        )
+        if config
+        else CallbackOptions()
+    )
+
+    checkpointed_result: CheckpointedResult = state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
+    if checkpointed_result.is_failed():
+        # have to throw the exact same error on replay as the checkpointed failure
+        checkpointed_result.raise_callable_error()
+
+    if (
+        checkpointed_result.is_started()
+        or checkpointed_result.is_succeeded()
+        or checkpointed_result.is_timed_out()
+    ):
+        # callback id should already exist
+        if (
+            not checkpointed_result.operation
+            or not checkpointed_result.operation.callback_details
+        ):
+            msg = f"Missing callback details for operation: {operation_identifier.operation_id}"
+            raise CallbackError(msg)
+
+        return checkpointed_result.operation.callback_details.callback_id
+
+    create_callback_operation = OperationUpdate.create_callback(
+        identifier=operation_identifier,
+        callback_options=callback_options,
+    )
+    # Checkpoint callback START with blocking (is_sync=True, default).
+    # Must wait for the API to generate and return the callback ID before proceeding.
+    # The callback ID is needed immediately by the caller to pass to external systems.
+    state.create_checkpoint(operation_update=create_callback_operation)
+
+    result: CheckpointedResult = state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
+
+    if not result.operation or not result.operation.callback_details:
+        msg = f"Missing callback details for operation: {operation_identifier.operation_id}"
+        raise CallbackError(msg)
+
+    return result.operation.callback_details.callback_id
+
+
+def wait_for_callback_handler(
+    context: DurableContext,
+    submitter: Callable[[str, WaitForCallbackContext], None],
+    name: str | None = None,
+    config: WaitForCallbackConfig | None = None,
+) -> Any:
+    """Wait for a callback to be invoked by an external system.
+
+    This is a helper function that is used to create a callback and wait for it to be invoked by an external system.
+    """
+    name_with_space: str = f"{name} " if name else ""
+    callback: Callback = context.create_callback(
+        name=f"{name_with_space}create callback id", config=config
+    )
+
+    def submitter_step(step_context: StepContext):
+        return submitter(
+            callback.callback_id, WaitForCallbackContext(logger=step_context.logger)
+        )
+
+    step_config = (
+        StepConfig(
+            retry_strategy=config.retry_strategy,
+            serdes=config.serdes,
+        )
+        if config
+        else None
+    )
+    context.step(
+        func=submitter_step, name=f"{name_with_space}submitter", config=step_config
+    )
+
+    return callback.result()

aws_durable_execution_sdk_python/operation/child.py

@@ -0,0 +1,162 @@
+"""Implementation for run_in_child_context."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, TypeVar
+
+from aws_durable_execution_sdk_python.config import ChildConfig
+from aws_durable_execution_sdk_python.exceptions import (
+    InvocationError,
+    SuspendExecution,
+)
+from aws_durable_execution_sdk_python.lambda_service import (
+    ContextOptions,
+    ErrorObject,
+    OperationSubType,
+    OperationUpdate,
+)
+from aws_durable_execution_sdk_python.serdes import deserialize, serialize
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+    from aws_durable_execution_sdk_python.state import ExecutionState
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+# Checkpoint size limit in bytes (256KB)
+CHECKPOINT_SIZE_LIMIT = 256 * 1024
+
+
+def child_handler(
+    func: Callable[[], T],
+    state: ExecutionState,
+    operation_identifier: OperationIdentifier,
+    config: ChildConfig | None,
+) -> T:
+    logger.debug(
+        "▶️ Executing child context for id: %s, name: %s",
+        operation_identifier.operation_id,
+        operation_identifier.name,
+    )
+
+    if not config:
+        config = ChildConfig()
+
+    checkpointed_result = state.get_checkpoint_result(operation_identifier.operation_id)
+    if (
+        checkpointed_result.is_succeeded()
+        and not checkpointed_result.is_replay_children()
+    ):
+        logger.debug(
+            "Child context already completed, skipping execution for id: %s, name: %s",
+            operation_identifier.operation_id,
+            operation_identifier.name,
+        )
+        if checkpointed_result.result is None:
+            return None  # type: ignore
+        return deserialize(
+            serdes=config.serdes,
+            data=checkpointed_result.result,
+            operation_id=operation_identifier.operation_id,
+            durable_execution_arn=state.durable_execution_arn,
+        )
+    if checkpointed_result.is_failed():
+        checkpointed_result.raise_callable_error()
+    sub_type = config.sub_type or OperationSubType.RUN_IN_CHILD_CONTEXT
+
+    if not checkpointed_result.is_existent():
+        start_operation = OperationUpdate.create_context_start(
+            identifier=operation_identifier,
+            sub_type=sub_type,
+        )
+        # Checkpoint child context START with non-blocking (is_sync=False).
+        # This is a fire-and-forget operation for performance - we don't need to wait for
+        # persistence before executing the child context. The START checkpoint is purely
+        # for observability and tracking the operation hierarchy.
+        state.create_checkpoint(operation_update=start_operation, is_sync=False)
+
+    try:
+        raw_result: T = func()
+        if checkpointed_result.is_replay_children():
+            logger.debug(
+                "ReplayChildren mode: Executed child context again on replay due to large payload. Exiting child context without creating another checkpoint. id: %s, name: %s",
+                operation_identifier.operation_id,
+                operation_identifier.name,
+            )
+            return raw_result
+        serialized_result: str = serialize(
+            serdes=config.serdes,
+            value=raw_result,
+            operation_id=operation_identifier.operation_id,
+            durable_execution_arn=state.durable_execution_arn,
+        )
+        # Summary Generator Logic:
+        # When the serialized result exceeds 256KB, we use ReplayChildren mode to avoid
+        # checkpointing large payloads. Instead, we checkpoint a compact summary and mark
+        # the operation for replay. This matches the TypeScript implementation behavior.
+        #
+        # See TypeScript reference:
+        # - aws-durable-execution-sdk-js/src/handlers/run-in-child-context-handler/run-in-child-context-handler.ts (lines ~200-220)
+        #
+        # The summary generator creates a JSON summary with metadata (type, counts, status)
+        # instead of the full BatchResult. During replay, the child context is re-executed
+        # to reconstruct the full result rather than deserializing from the checkpoint.
+        replay_children: bool = False
+        if len(serialized_result) > CHECKPOINT_SIZE_LIMIT:
+            logger.debug(
+                "Large payload detected, using ReplayChildren mode: id: %s, name: %s, payload_size: %d, limit: %d",
+                operation_identifier.operation_id,
+                operation_identifier.name,
+                len(serialized_result),
+                CHECKPOINT_SIZE_LIMIT,
+            )
+            replay_children = True
+            # Use summary generator if provided, otherwise use empty string (matches TypeScript)
+            serialized_result = (
+                config.summary_generator(raw_result) if config.summary_generator else ""
+            )
+
+        success_operation = OperationUpdate.create_context_succeed(
+            identifier=operation_identifier,
+            payload=serialized_result,
+            sub_type=sub_type,
+            context_options=ContextOptions(replay_children=replay_children),
+        )
+        # Checkpoint child context SUCCEED with blocking (is_sync=True, default).
+        # Must ensure the child context result is persisted before returning to the parent.
+        # This guarantees the result is durable and child operations won't be re-executed on replay
+        # (unless replay_children=True for large payloads).
+        state.create_checkpoint(operation_update=success_operation)
+
+        logger.debug(
+            "✅ Successfully completed child context for id: %s, name: %s",
+            operation_identifier.operation_id,
+            operation_identifier.name,
+        )
+        return raw_result  # noqa: TRY300
+    except SuspendExecution:
+        # Don't checkpoint SuspendExecution - let it bubble up
+        raise
+    except Exception as e:
+        error_object = ErrorObject.from_exception(e)
+        fail_operation = OperationUpdate.create_context_fail(
+            identifier=operation_identifier, error=error_object, sub_type=sub_type
+        )
+        # Checkpoint child context FAIL with blocking (is_sync=True, default).
+        # Must ensure the failure state is persisted before raising the exception.
+        # This guarantees the error is durable and child operations won't be re-executed on replay.
+        state.create_checkpoint(operation_update=fail_operation)
+
+        # InvocationError and its derivatives can be retried
+        # When we encounter an invocation error (in all of its forms), we bubble that
+        # error upwards (with the checkpoint in place) such that we reach the
+        # execution handler at the very top, which will then induce a retry from the
+        # dataplane.
+        if isinstance(e, InvocationError):
+            raise
+        raise error_object.to_callable_runtime_error() from e

aws_durable_execution_sdk_python/operation/invoke.py

@@ -0,0 +1,119 @@
+"""Implement the Durable invoke operation."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, TypeVar
+
+from aws_durable_execution_sdk_python.config import InvokeConfig
+from aws_durable_execution_sdk_python.exceptions import ExecutionError
+from aws_durable_execution_sdk_python.lambda_service import (
+    ChainedInvokeOptions,
+    OperationUpdate,
+)
+from aws_durable_execution_sdk_python.serdes import (
+    DEFAULT_JSON_SERDES,
+    deserialize,
+    serialize,
+)
+from aws_durable_execution_sdk_python.suspend import suspend_with_optional_resume_delay
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+    from aws_durable_execution_sdk_python.state import ExecutionState
+
+P = TypeVar("P")  # Payload type
+R = TypeVar("R")  # Result type
+
+logger = logging.getLogger(__name__)
+
+
+def invoke_handler(
+    function_name: str,
+    payload: P,
+    state: ExecutionState,
+    operation_identifier: OperationIdentifier,
+    config: InvokeConfig[P, R] | None,
+) -> R:
+    """Invoke another Durable Function."""
+    logger.debug(
+        "🔗 Invoke %s (%s)",
+        operation_identifier.name or function_name,
+        operation_identifier.operation_id,
+    )
+
+    if not config:
+        config = InvokeConfig[P, R]()
+    tenant_id = config.tenant_id
+
+    # Check if we have existing step data
+    checkpointed_result = state.get_checkpoint_result(operation_identifier.operation_id)
+
+    if checkpointed_result.is_succeeded():
+        # Return persisted result - no need to check for errors in successful operations
+        if (
+            checkpointed_result.operation
+            and checkpointed_result.operation.chained_invoke_details
+            and checkpointed_result.operation.chained_invoke_details.result
+        ):
+            return deserialize(
+                serdes=config.serdes_result or DEFAULT_JSON_SERDES,
+                data=checkpointed_result.operation.chained_invoke_details.result,
+                operation_id=operation_identifier.operation_id,
+                durable_execution_arn=state.durable_execution_arn,
+            )
+        return None  # type: ignore
+
+    if (
+        checkpointed_result.is_failed()
+        or checkpointed_result.is_timed_out()
+        or checkpointed_result.is_stopped()
+    ):
+        # Operation failed, throw the exact same error on replay as the checkpointed failure
+        checkpointed_result.raise_callable_error()
+
+    if checkpointed_result.is_started():
+        # Operation is still running, suspend until completion
+        logger.debug(
+            "⏳ Invoke %s still in progress, suspending",
+            operation_identifier.name or function_name,
+        )
+        msg = f"Invoke {operation_identifier.operation_id} still in progress"
+        suspend_with_optional_resume_delay(msg, config.timeout_seconds)
+
+    serialized_payload: str = serialize(
+        serdes=config.serdes_payload or DEFAULT_JSON_SERDES,
+        value=payload,
+        operation_id=operation_identifier.operation_id,
+        durable_execution_arn=state.durable_execution_arn,
+    )
+
+    # the backend will do the invoke once it gets this checkpoint
+    start_operation: OperationUpdate = OperationUpdate.create_invoke_start(
+        identifier=operation_identifier,
+        payload=serialized_payload,
+        chained_invoke_options=ChainedInvokeOptions(
+            function_name=function_name,
+            tenant_id=tenant_id,
+        ),
+    )
+
+    # Checkpoint invoke START with blocking (is_sync=True, default).
+    # Must ensure the chained invocation is recorded before suspending execution.
+    # This guarantees the invoke operation is durable and will be tracked by the backend.
+    state.create_checkpoint(operation_update=start_operation)
+
+    logger.debug(
+        "🚀 Invoke %s started, suspending for async execution",
+        operation_identifier.name or function_name,
+    )
+
+    # Suspend so invoke executes asynchronously without consuming cpu here
+    msg = (
+        f"Invoke {operation_identifier.operation_id} started, suspending for completion"
+    )
+    suspend_with_optional_resume_delay(msg, config.timeout_seconds)
+    # This line should never be reached since suspend_with_optional_resume_delay always raises
+    # if it is ever reached, we will crash in a non-retryable manner via ExecutionError
+    msg = "suspend_with_optional_resume_delay should have raised an exception, but did not."
+    raise ExecutionError(msg) from None

aws_durable_execution_sdk_python/operation/map.py

@@ -0,0 +1,137 @@
+"""Implementation for Durable Map operation."""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor
+from aws_durable_execution_sdk_python.concurrency.models import (
+    BatchResult,
+    Executable,
+)
+from aws_durable_execution_sdk_python.config import MapConfig
+from aws_durable_execution_sdk_python.lambda_service import OperationSubType
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.context import DurableContext
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+    from aws_durable_execution_sdk_python.serdes import SerDes
+    from aws_durable_execution_sdk_python.state import (
+        CheckpointedResult,
+        ExecutionState,
+    )
+    from aws_durable_execution_sdk_python.types import SummaryGenerator
+
+logger = logging.getLogger(__name__)
+
+# Input item type
+T = TypeVar("T")
+# Result type
+R = TypeVar("R")
+
+
+class MapExecutor(Generic[T, R], ConcurrentExecutor[Callable, R]):  # noqa: PYI059
+    def __init__(
+        self,
+        executables: list[Executable[Callable]],
+        items: Sequence[T],
+        max_concurrency: int | None,
+        completion_config,
+        top_level_sub_type: OperationSubType,
+        iteration_sub_type: OperationSubType,
+        name_prefix: str,
+        serdes: SerDes | None,
+        summary_generator: SummaryGenerator | None = None,
+        item_serdes: SerDes | None = None,
+    ):
+        super().__init__(
+            executables=executables,
+            max_concurrency=max_concurrency,
+            completion_config=completion_config,
+            sub_type_top=top_level_sub_type,
+            sub_type_iteration=iteration_sub_type,
+            name_prefix=name_prefix,
+            serdes=serdes,
+            summary_generator=summary_generator,
+            item_serdes=item_serdes,
+        )
+        self.items = items
+
+    @classmethod
+    def from_items(
+        cls,
+        items: Sequence[T],
+        func: Callable,
+        config: MapConfig,
+    ) -> MapExecutor[T, R]:
+        """Create MapExecutor from items and a callable."""
+        executables: list[Executable[Callable]] = [
+            Executable(index=i, func=func) for i in range(len(items))
+        ]
+
+        return cls(
+            executables=executables,
+            items=items,
+            max_concurrency=config.max_concurrency,
+            completion_config=config.completion_config,
+            top_level_sub_type=OperationSubType.MAP,
+            iteration_sub_type=OperationSubType.MAP_ITERATION,
+            name_prefix="map-item-",
+            serdes=config.serdes,
+            summary_generator=config.summary_generator,
+            item_serdes=config.item_serdes,
+        )
+
+    def execute_item(self, child_context, executable: Executable[Callable]) -> R:
+        logger.debug("🗺️ Processing map item: %s", executable.index)
+        item = self.items[executable.index]
+        result: R = executable.func(child_context, item, executable.index, self.items)
+        logger.debug("✅ Processed map item: %s", executable.index)
+        return result
+
+
+def map_handler(
+    items: Sequence[T],
+    func: Callable,
+    config: MapConfig | None,
+    execution_state: ExecutionState,
+    map_context: DurableContext,
+    operation_identifier: OperationIdentifier,
+) -> BatchResult[R]:
+    """Execute a callable for each item in parallel."""
+    # Summary Generator Construction (matches TypeScript implementation):
+    # Construct the summary generator at the handler level, just like TypeScript does in map-handler.ts.
+    # This matches the pattern where handlers are responsible for configuring operation-specific behavior.
+    #
+    # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/map-handler/map-handler.ts (~line 79)
+
+    executor: MapExecutor[T, R] = MapExecutor.from_items(
+        items=items,
+        func=func,
+        config=config or MapConfig(summary_generator=MapSummaryGenerator()),
+    )
+
+    checkpoint: CheckpointedResult = execution_state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
+    if checkpoint.is_succeeded():
+        # if we've reached this point, then not only is the step succeeded, but it is also `replay_children`.
+        return executor.replay(execution_state, map_context)
+    # we are making it explicit that we are now executing within the map_context
+    return executor.execute(execution_state, executor_context=map_context)
+
+
+class MapSummaryGenerator:
+    def __call__(self, result: BatchResult) -> str:
+        fields = {
+            "totalCount": result.total_count,
+            "successCount": result.success_count,
+            "failureCount": result.failure_count,
+            "completionReason": result.completion_reason.value,
+            "status": result.status.value,
+            "type": "MapResult",
+        }
+        return json.dumps(fields)