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

aws_durable_execution_sdk_python_testing/checkpoint/processors/wait.py

@@ -0,0 +1,95 @@
+"""Wait operation processor for handling WAIT operation updates."""
+
+from __future__ import annotations
+
+import logging
+import os
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+    WaitDetails,
+)
+
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import (
+    OperationProcessor,
+)
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
+
+
+class WaitProcessor(OperationProcessor):
+    """Processes WAIT operation updates with timer scheduling."""
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        notifier: ExecutionNotifier,
+        execution_arn: str,
+    ) -> Operation:
+        """Process WAIT operation update with scheduler integration for timers."""
+        match update.action:
+            case OperationAction.START:
+                wait_seconds = (
+                    update.wait_options.wait_seconds if update.wait_options else 0
+                )
+                time_scale = float(os.getenv("DURABLE_EXECUTION_TIME_SCALE", "1.0"))
+                logging.info("Using DURABLE_EXECUTION_TIME_SCALE: %f", time_scale)
+                scaled_wait_seconds = wait_seconds * time_scale
+
+                scheduled_end_timestamp = datetime.now(UTC) + timedelta(
+                    seconds=scaled_wait_seconds
+                )
+
+                # Create WaitDetails with scheduled timestamp
+                wait_details = WaitDetails(
+                    scheduled_end_timestamp=scheduled_end_timestamp
+                )
+
+                # Create new operation with wait details
+                wait_operation = Operation(
+                    operation_id=update.operation_id,
+                    operation_type=update.operation_type,
+                    status=OperationStatus.STARTED,
+                    parent_id=update.parent_id,
+                    name=update.name,
+                    start_timestamp=datetime.now(UTC),
+                    end_timestamp=None,
+                    sub_type=update.sub_type,
+                    execution_details=None,
+                    context_details=None,
+                    step_details=None,
+                    wait_details=wait_details,
+                    callback_details=None,
+                    chained_invoke_details=None,
+                )
+
+                # Schedule wait timer to complete after delay
+                notifier.notify_wait_timer_scheduled(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                    delay=scaled_wait_seconds,
+                )
+                return wait_operation
+            case OperationAction.CANCEL:
+                # TODO: need to cancel the WAIT in the executor
+                # TODO: increase sequence id
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.CANCELLED,
+                )
+            case _:
+                msg: str = "Invalid action for WAIT operation."
+
+                raise InvalidParameterValueException(msg)

aws_durable_execution_sdk_python_testing/checkpoint/transformer.py

@@ -0,0 +1,104 @@
+"""Operation transformer for converting OperationUpdates to Operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
+    OperationType,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.callback import (
+    CallbackProcessor,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.context import (
+    ContextProcessor,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.execution import (
+    ExecutionProcessor,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.step import (
+    StepProcessor,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.wait import (
+    WaitProcessor,
+)
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import (
+        OperationProcessor,
+    )
+
+from typing import ClassVar
+
+
+class OperationTransformer:
+    """Transforms OperationUpdates to Operations while maintaining order and triggering scheduler actions."""
+
+    _DEFAULT_PROCESSORS: ClassVar[dict[OperationType, OperationProcessor]] = {
+        OperationType.STEP: StepProcessor(),
+        OperationType.WAIT: WaitProcessor(),
+        OperationType.CONTEXT: ContextProcessor(),
+        OperationType.CALLBACK: CallbackProcessor(),
+        OperationType.EXECUTION: ExecutionProcessor(),
+    }
+
+    def __init__(
+        self,
+        processors: MutableMapping[OperationType, OperationProcessor] | None = None,
+    ):
+        self.processors = processors if processors else self._DEFAULT_PROCESSORS
+
+    def process_updates(
+        self,
+        updates: list[OperationUpdate],
+        current_operations: list[Operation],
+        notifier,
+        execution_arn: str,
+    ) -> tuple[list[Operation], list[OperationUpdate]]:
+        """Transform updates maintaining operation order and return (operations, updates)."""
+        op_map = {op.operation_id: op for op in current_operations}
+
+        # Start with copy of current operations list
+        result_operations = current_operations.copy()
+
+        for update in updates:
+            processor = self.processors.get(update.operation_type)
+            if processor:
+                current_op = op_map.get(update.operation_id)
+                updated_op = processor.process(
+                    update=update,
+                    current_op=current_op,
+                    notifier=notifier,
+                    execution_arn=execution_arn,
+                )
+
+                if updated_op is not None:
+                    if update.operation_id in op_map:
+                        # Update existing operation in-place
+                        for i, op in enumerate(result_operations):  # pragma: no branch
+                            # no branch coverage because result_operation empty not reachable here
+                            if op.operation_id == update.operation_id:
+                                result_operations[i] = updated_op
+                                break
+                    else:
+                        # Append new operation to end
+                        result_operations.append(updated_op)
+
+                    # Update map for future lookups
+                    op_map[update.operation_id] = updated_op
+            else:
+                msg: str = (
+                    f"Checkpoint for {update.operation_type} is not implemented yet."
+                )
+                raise InvalidParameterValueException(msg)
+
+        return result_operations, updates

aws_durable_execution_sdk_python_testing/checkpoint/validators/__init__.py

@@ -0,0 +1 @@
+"""Checkpoint validation module."""

aws_durable_execution_sdk_python_testing/checkpoint/validators/checkpoint.py

@@ -0,0 +1,242 @@
+"""Main checkpoint input validator."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    OperationAction,
+    OperationType,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.callback import (
+    CallbackOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.context import (
+    ContextOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.execution import (
+    ExecutionOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.invoke import (
+    ChainedInvokeOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.step import (
+    StepOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.wait import (
+    WaitOperationValidator,
+)
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.transitions import (
+    ValidActionsByOperationTypeValidator,
+)
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+
+MAX_ERROR_PAYLOAD_SIZE_BYTES = 32768
+
+
+class CheckpointValidator:
+    """Validates checkpoint input based on current state."""
+
+    @staticmethod
+    def validate_input(updates: list[OperationUpdate], execution: Execution) -> None:
+        """Perform validation on the given input based on the current state."""
+        if not updates:
+            return
+
+        CheckpointValidator._validate_conflicting_execution_update(updates)
+        CheckpointValidator._validate_parent_id_and_duplicate_id(updates, execution)
+
+        for update in updates:
+            CheckpointValidator._validate_operation_update(update, execution)
+
+    @staticmethod
+    def _validate_conflicting_execution_update(updates: list[OperationUpdate]) -> None:
+        """Validate that there are no conflicting execution updates."""
+        execution_updates = [
+            update
+            for update in updates
+            if update.operation_type == OperationType.EXECUTION
+        ]
+
+        if len(execution_updates) > 1:
+            msg_multiple_exec: str = "Cannot checkpoint multiple EXECUTION updates."
+
+            raise InvalidParameterValueException(msg_multiple_exec)
+
+        if execution_updates and updates[-1].operation_type != OperationType.EXECUTION:
+            msg_exec_last: str = "EXECUTION checkpoint must be the last update."
+
+            raise InvalidParameterValueException(msg_exec_last)
+
+    @staticmethod
+    def _validate_operation_update(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate a single operation update."""
+        CheckpointValidator._validate_inconsistent_operation_metadata(update, execution)
+        CheckpointValidator._validate_payload_sizes(update)
+        ValidActionsByOperationTypeValidator.validate(
+            update.operation_type, update.action
+        )
+        CheckpointValidator._validate_operation_status_transition(update, execution)
+
+    @staticmethod
+    def _validate_payload_sizes(update: OperationUpdate) -> None:
+        """Validate that operation payload sizes are not too large."""
+        if update.error is not None:
+            payload = json.dumps(update.error.to_dict())
+            if len(payload) > MAX_ERROR_PAYLOAD_SIZE_BYTES:
+                msg: str = f"Error object size must be less than {MAX_ERROR_PAYLOAD_SIZE_BYTES} bytes."
+                raise InvalidParameterValueException(msg)
+
+    @staticmethod
+    def _validate_operation_status_transition(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate that the operation status transition is valid."""
+        current_state = None
+        for operation in execution.operations:
+            if operation.operation_id == update.operation_id:
+                current_state = operation
+                break
+
+        match update.operation_type:
+            case OperationType.STEP:
+                StepOperationValidator.validate(current_state, update)
+            case OperationType.CONTEXT:
+                ContextOperationValidator.validate(current_state, update)
+            case OperationType.WAIT:
+                WaitOperationValidator.validate(current_state, update)
+            case OperationType.CALLBACK:
+                CallbackOperationValidator.validate(current_state, update)
+            case OperationType.CHAINED_INVOKE:
+                ChainedInvokeOperationValidator.validate(current_state, update)
+            case OperationType.EXECUTION:
+                ExecutionOperationValidator.validate(update)
+            case _:  # pragma: no cover
+                msg: str = "Invalid operation type."
+
+                raise InvalidParameterValueException(msg)
+
+    @staticmethod
+    def _validate_inconsistent_operation_metadata(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate that operation metadata is consistent with existing operation."""
+        current_state = None
+        for operation in execution.operations:
+            if operation.operation_id == update.operation_id:
+                current_state = operation
+                break
+
+        if current_state is not None:
+            if (
+                update.operation_type is not None
+                and update.operation_type != current_state.operation_type
+            ):
+                msg: str = "Inconsistent operation type."
+                raise InvalidParameterValueException(msg)
+
+            if (
+                update.sub_type is not None
+                and update.sub_type != current_state.sub_type
+            ):
+                msg_subtype: str = "Inconsistent operation subtype."
+                raise InvalidParameterValueException(msg_subtype)
+
+            if update.name is not None and update.name != current_state.name:
+                msg_name: str = "Inconsistent operation name."
+                raise InvalidParameterValueException(msg_name)
+
+            if (
+                update.parent_id is not None
+                and update.parent_id != current_state.parent_id
+            ):
+                msg_parent: str = "Inconsistent parent operation id."
+                raise InvalidParameterValueException(msg_parent)
+
+    @staticmethod
+    def _validate_parent_id_and_duplicate_id(
+        updates: list[OperationUpdate], execution: Execution
+    ) -> None:
+        """Validate parent IDs and check for duplicate operation IDs.
+
+        Validate that any provided parentId is valid, and also validate no duplicate operation is being
+        updated at the same time (unless it is a STEP/CONTEXT starting + performing one more non-START action).
+        """
+        operations_started: MutableMapping[str, OperationUpdate] = {}
+        last_updates_seen: MutableMapping[str, OperationUpdate] = {}
+
+        for update in updates:
+            if CheckpointValidator._is_invalid_duplicate_update(
+                update, last_updates_seen
+            ):
+                msg_duplicate: str = (
+                    "Cannot checkpoint multiple operations with the same ID."
+                )
+                raise InvalidParameterValueException(msg_duplicate)
+
+            if not CheckpointValidator._is_valid_parent_for_update(
+                execution, update, operations_started
+            ):
+                msg_parent: str = "Invalid parent operation id."
+                raise InvalidParameterValueException(msg_parent)
+
+            if update.action == OperationAction.START:
+                operations_started[update.operation_id] = update
+
+            last_updates_seen[update.operation_id] = update
+
+    @staticmethod
+    def _is_invalid_duplicate_update(
+        update: OperationUpdate, last_updates_seen: MutableMapping[str, OperationUpdate]
+    ) -> bool:
+        """Check if this is an invalid duplicate update."""
+        last_update = last_updates_seen.get(update.operation_id)
+        if last_update is None:
+            return False
+
+        if last_update.operation_type in (OperationType.STEP, OperationType.CONTEXT):
+            # Allow duplicate for STEP/CONTEXT if last was START and current is not START
+            allow_duplicate = (
+                last_update.action == OperationAction.START
+                and update.action != OperationAction.START
+            )
+            return not allow_duplicate
+
+        return True
+
+    @staticmethod
+    def _is_valid_parent_for_update(
+        execution: Execution,
+        update: OperationUpdate,
+        operations_started: MutableMapping[str, OperationUpdate],
+    ) -> bool:
+        """Check if the parent ID is valid for the update."""
+        parent_id = update.parent_id
+
+        if parent_id is None:
+            return True
+
+        # Check if parent is in operations started in this batch
+        if parent_id in operations_started:
+            parent_update = operations_started[parent_id]
+            return parent_update.operation_type == OperationType.CONTEXT
+
+        # Check if parent exists in current execution state
+        for operation in execution.operations:
+            if operation.operation_id == parent_id:
+                return operation.operation_type == OperationType.CONTEXT
+
+        return False

aws_durable_execution_sdk_python_testing/checkpoint/validators/operations/__init__.py

@@ -0,0 +1 @@
+"""Operation-specific validators."""

aws_durable_execution_sdk_python_testing/checkpoint/validators/operations/callback.py

@@ -0,0 +1,45 @@
+"""Callback operation validator."""
+
+from __future__ import annotations
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+VALID_ACTIONS_FOR_CALLBACK = frozenset(
+    [
+        OperationAction.START,
+    ]
+)
+
+
+class CallbackOperationValidator:
+    """Validates CALLBACK operation transitions."""
+
+    _ALLOWED_STATUS_TO_CANCEL = frozenset(
+        [
+            OperationStatus.STARTED,
+        ]
+    )
+
+    @staticmethod
+    def validate(current_state: Operation | None, update: OperationUpdate) -> None:
+        """Validate CALLBACK operation update."""
+        match update.action:
+            case OperationAction.START:
+                if current_state is not None:
+                    msg_callback_exists: str = (
+                        "Cannot start a CALLBACK that already exist."
+                    )
+                    raise InvalidParameterValueException(msg_callback_exists)
+            case _:
+                msg: str = "Invalid action for CALLBACK operation."
+                raise InvalidParameterValueException(msg)

aws_durable_execution_sdk_python_testing/checkpoint/validators/operations/context.py

@@ -0,0 +1,73 @@
+"""Context operation validator."""
+
+from __future__ import annotations
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+VALID_ACTIONS_FOR_CONTEXT = frozenset(
+    [
+        OperationAction.START,
+        OperationAction.FAIL,
+        OperationAction.SUCCEED,
+    ]
+)
+
+
+class ContextOperationValidator:
+    """Validates CONTEXT operation transitions."""
+
+    _ALLOWED_STATUS_TO_CLOSE = frozenset(
+        [
+            OperationStatus.STARTED,
+        ]
+    )
+
+    @staticmethod
+    def validate(current_state: Operation | None, update: OperationUpdate) -> None:
+        """Validate CONTEXT operation update."""
+        match update.action:
+            case OperationAction.START:
+                if current_state is not None:
+                    msg_context_exists: str = (
+                        "Cannot start a CONTEXT that already exist."
+                    )
+
+                    raise InvalidParameterValueException(msg_context_exists)
+            case OperationAction.FAIL | OperationAction.SUCCEED:
+                if (
+                    current_state is not None
+                    and current_state.status
+                    not in ContextOperationValidator._ALLOWED_STATUS_TO_CLOSE
+                ):
+                    msg_context_close: str = "Invalid current CONTEXT state to close."
+
+                    raise InvalidParameterValueException(msg_context_close)
+                if update.action == OperationAction.FAIL and update.payload is not None:
+                    msg_context_fail_payload: str = (
+                        "Cannot provide a Payload for FAIL action."
+                    )
+
+                    raise InvalidParameterValueException(msg_context_fail_payload)
+                if (
+                    update.action == OperationAction.SUCCEED
+                    and update.error is not None
+                ):
+                    msg_context_succeed_error: str = (
+                        "Cannot provide an Error for SUCCEED action."
+                    )
+
+                    raise InvalidParameterValueException(msg_context_succeed_error)
+            case _:
+                msg_context_invalid: str = "Invalid CONTEXT action."
+
+                raise InvalidParameterValueException(msg_context_invalid)

aws_durable_execution_sdk_python_testing/checkpoint/validators/operations/execution.py

@@ -0,0 +1,47 @@
+"""Execution operation validator."""
+
+from __future__ import annotations
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    OperationAction,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+VALID_ACTIONS_FOR_EXECUTION = frozenset(
+    [
+        OperationAction.SUCCEED,
+        OperationAction.FAIL,
+    ]
+)
+
+
+class ExecutionOperationValidator:
+    """Validates EXECUTION operation transitions."""
+
+    @staticmethod
+    def validate(update: OperationUpdate) -> None:
+        """Validate EXECUTION operation update."""
+        match update.action:
+            case OperationAction.SUCCEED:
+                if update.error is not None:
+                    msg_exec_succeed_error: str = (
+                        "Cannot provide an Error for SUCCEED action."
+                    )
+
+                    raise InvalidParameterValueException(msg_exec_succeed_error)
+            case OperationAction.FAIL:
+                if update.payload is not None:
+                    msg_exec_fail_payload: str = (
+                        "Cannot provide a Payload for FAIL action."
+                    )
+
+                    raise InvalidParameterValueException(msg_exec_fail_payload)
+            case _:
+                msg_exec_invalid: str = "Invalid EXECUTION action."
+
+                raise InvalidParameterValueException(msg_exec_invalid)

aws_durable_execution_sdk_python_testing/checkpoint/validators/operations/invoke.py

@@ -0,0 +1,56 @@
+"""Invoke operation validator."""
+
+from __future__ import annotations
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+VALID_ACTIONS_FOR_INVOKE = frozenset(
+    [
+        OperationAction.START,
+        OperationAction.CANCEL,
+    ]
+)
+
+
+class ChainedInvokeOperationValidator:
+    """Validates INVOKE operation transitions."""
+
+    _ALLOWED_STATUS_TO_CANCEL = frozenset(
+        [
+            OperationStatus.STARTED,
+        ]
+    )
+
+    @staticmethod
+    def validate(current_state: Operation | None, update: OperationUpdate) -> None:
+        """Validate INVOKE operation update."""
+        match update.action:
+            case OperationAction.START:
+                if current_state is not None:
+                    msg_invoke_exists: str = (
+                        "Cannot start an INVOKE that already exist."
+                    )
+
+                    raise InvalidParameterValueException(msg_invoke_exists)
+            case OperationAction.CANCEL:
+                if (
+                    current_state is None
+                    or current_state.status
+                    not in ChainedInvokeOperationValidator._ALLOWED_STATUS_TO_CANCEL
+                ):
+                    msg_invoke_cancel: str = "Cannot cancel an INVOKE that does not exist or has already completed."
+                    raise InvalidParameterValueException(msg_invoke_cancel)
+            case _:
+                msg_invoke_invalid: str = "Invalid INVOKE action."
+
+                raise InvalidParameterValueException(msg_invoke_invalid)