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

aws_durable_execution_sdk_python_testing/scheduler.py

@@ -0,0 +1,246 @@
+"""A Scheduler that can run awaitables or standard sync callables on a schedule once or repeatedly."""
+
+from __future__ import annotations
+
+import asyncio
+import itertools
+import logging
+import threading
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from concurrent.futures import Future
+
+logger = logging.getLogger(__name__)
+
+
+class Event:
+    """An event created by Scheduler that will block on wait until it's set."""
+
+    def __init__(self, scheduler: Scheduler, asyncio_event: asyncio.Event) -> None:
+        self._scheduler: Scheduler = scheduler
+        self._asyncio_event: asyncio.Event = asyncio_event
+        self._exception: Exception | None = None
+
+    def set(self):
+        """Set the event with this to unblock wait."""
+        self._scheduler.set_event(self._asyncio_event)
+
+    def set_exception(self, exception: Exception):
+        """Set exception and unblock waiters."""
+        self._exception = exception
+        self._scheduler.set_event(self._asyncio_event)
+
+    def wait(self, timeout: float | None = None, *, clear_on_set: bool = True) -> bool:
+        """Wait until the event is set.
+
+        Args:
+            timeout (int | float | None): Wait for event to set until this timeout.
+            clear_on_set (bool): Remove the event from the Scheduler on completion.
+                                 Use this if you won't re-use the event.
+
+        Returns:
+            True when set. False if the event timed out without being set.
+
+        Raises:
+            Exception: If an exception was stored via set_exception().
+        """
+        result = self._scheduler.wait_for_event(self._asyncio_event, timeout)
+        if clear_on_set:
+            self._scheduler.remove_event(self._asyncio_event)
+        if result and self._exception:
+            raise self._exception
+        return result
+
+    def remove(self):
+        """Remove the event from the Scheduler. Do this to avoid build-up of many events in the scheduler."""
+        self._scheduler.remove_event(self._asyncio_event)
+
+
+class Scheduler:
+    """A Scheduler to run callables later, repeatedly or raise events."""
+
+    def __init__(self) -> None:
+        self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
+        self._ready_event: threading.Event = threading.Event()
+        self._thread: threading.Thread = threading.Thread(
+            target=self._start_loop, daemon=True
+        )
+        self._running: bool = False
+        self._events: set[asyncio.Event] = set()
+
+    # region context manager
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+
+    # endregion context manager
+
+    # region event loop
+    def start(self):
+        """Start the scheduler. Not thread-safe."""
+        if self._running:
+            return
+
+        self._running = True
+
+        self._thread.start()
+        # Wait for inside of loop to notify it's ready (meaning _start_loop has completed)
+        self._ready_event.wait()
+
+    def stop(self):
+        """Stop the scheduler, releasing resources. Not thread-safe."""
+        if not self._running:
+            return
+
+        self._running = False
+        self._loop.call_soon_threadsafe(self._cleanup_and_stop)
+        self._thread.join()
+
+    def is_started(self) -> bool:
+        """Return True if the scheduler is started."""
+        return self._running
+
+    def event_count(self) -> int:
+        """Return the number of events in the scheduler."""
+        return len(self._events)
+
+    def task_count(self) -> int:
+        """Return the number of tasks in the scheduler."""
+        if not self._running:
+            return 0
+        return len(asyncio.all_tasks(self._loop))
+
+    def _cleanup_and_stop(self):
+        """Cancel all tasks and clear all events. Stop the event-loop."""
+        # Cancel all tasks
+        for task in asyncio.all_tasks(self._loop):
+            task.cancel()
+
+        # Clear events (don't set them)
+        self._events.clear()
+
+        self._loop.stop()
+
+    def _start_loop(self):
+        """Initialize the event-loop. The ready event notifies that the loop is started."""
+        asyncio.set_event_loop(self._loop)
+        # signal that loop is ready from within the loop
+        self._loop.call_soon(self._ready_event.set)
+        # block indefinitely - call_soon with the read_event will run soon as the loop starts
+        self._loop.run_forever()
+
+    # endregion event loop
+    # region Tasks
+    def call_later(
+        self,
+        func: Callable[[], Any],
+        delay: float = 0,
+        count: int | None = 1,
+        completion_event: Event | None = None,
+    ) -> Future[Any]:
+        """Call func after the delay.
+
+        If func is async it runs inside a thread-safe coroutine. If func is sync it runs in its own
+        threadpool, so it won't block the event loop.
+
+        Args:
+            func (Callable[[], Any]): The function to call later. This can be an async or a standard
+                                      sync function.
+            delay (float | int): Delay in seconds before calling func.
+            count (int | None): Number of times to call func. Default is 1 (call once).
+                               Use None for infinite repeats.
+            completion_event (Event | None): Event to notify on exception.
+
+        Returns: Future that completes when the scheduled work is done.
+        """
+        # infinite counter if count = None, else it maxes out at count
+        loop_iter: itertools.count[int] | range = (
+            itertools.count() if count is None else range(count)
+        )
+
+        async def delayed_func() -> Any:
+            try:
+                for _ in loop_iter:
+                    await asyncio.sleep(delay)
+
+                    try:
+                        if asyncio.iscoroutinefunction(func):
+                            result = await func()
+                        else:
+                            result = await asyncio.to_thread(func)
+                        return result  # noqa: TRY300
+                    except Exception as err:
+                        if completion_event:
+                            completion_event.set_exception(err)
+                        else:
+                            msg: str = "error in scheduled task"
+                            logger.exception(msg)
+                        raise
+            except asyncio.CancelledError:  # noqa: TRY302
+                # might want to handle more things here
+                raise
+
+        future: Future[Any] = asyncio.run_coroutine_threadsafe(
+            delayed_func(), self._loop
+        )
+        return future
+
+    # endregion Tasks
+
+    # region Events
+
+    def create_event(self) -> Event:
+        """Create an event controlled by the Scheduler to signal between threads and coroutines."""
+        # create event inside the Scheduler event-loop
+        future: Future[asyncio.Event] = asyncio.run_coroutine_threadsafe(
+            self._create_event(), self._loop
+        )
+
+        # Add timeout to prevent surprising "hangs" if for whatever reason event fails to create.
+        # result with block. Do NOT call anything in _create_event that calls back into scheduler
+        # methods because it could create a circular depdendency which will deadlock.
+        event = future.result(timeout=5.0)
+        return Event(self, event)
+
+    def wait_for_event(
+        self, event: asyncio.Event, timeout: float | None = None
+    ) -> bool:
+        """Run event's wait inside the Scheduler event-loop."""
+        if event not in self._events:
+            return False
+
+        future: Future[bool] = asyncio.run_coroutine_threadsafe(
+            asyncio.wait_for(event.wait(), timeout), self._loop
+        )
+
+        try:
+            return future.result()
+        except TimeoutError:
+            return False
+
+    def set_event(self, event: asyncio.Event):
+        """Set event inside the Scheduler event-loop."""
+        if event in self._events:
+            self._loop.call_soon_threadsafe(event.set)
+
+    def remove_event(self, event: asyncio.Event):
+        """Remove event from Scheduler in the Scheduler event-loop."""
+
+        def _remove():
+            self._events.discard(event)
+
+        self._loop.call_soon_threadsafe(_remove)
+
+    async def _create_event(self) -> asyncio.Event:
+        """Create event and add it to the scheduler events list."""
+        event = asyncio.Event()
+        self._events.add(event)
+        return event
+
+    # endregion Events

aws_durable_execution_sdk_python_testing/stores/__init__.py

@@ -0,0 +1 @@
+

aws_durable_execution_sdk_python_testing/stores/base.py

@@ -0,0 +1,147 @@
+"""Base classes and protocols for execution stores."""
+
+from __future__ import annotations
+
+from datetime import UTC
+from enum import Enum
+from typing import TYPE_CHECKING, Protocol
+
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.lambda_service import Operation
+
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+
+
+class StoreType(Enum):
+    """Supported execution store types."""
+
+    MEMORY = "memory"
+    FILESYSTEM = "filesystem"
+    SQLITE = "sqlite"
+
+
+class ExecutionStore(Protocol):
+    """Protocol for execution storage implementations."""
+
+    # ignore cover because coverage doesn't understand elipses
+    def save(self, execution: Execution) -> None: ...  # pragma: no cover
+    def load(self, execution_arn: str) -> Execution: ...  # pragma: no cover
+    def update(self, execution: Execution) -> None: ...  # pragma: no cover
+    def query(
+        self,
+        function_name: str | None = None,
+        execution_name: str | None = None,
+        status_filter: str | None = None,
+        started_after: str | None = None,
+        started_before: str | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+        reverse_order: bool = False,  # noqa: FBT001, FBT002
+    ) -> tuple[list[Execution], str | None]: ...  # pragma: no cover
+    def list_all(
+        self,
+    ) -> list[Execution]: ...  # pragma: no cover  # Keep for backward compatibility
+
+
+class BaseExecutionStore(ExecutionStore):
+    """Base implementation for execution stores with shared query logic."""
+
+    @staticmethod
+    def process_query(
+        executions: list[Execution],
+        function_name: str | None = None,
+        execution_name: str | None = None,
+        status_filter: str | None = None,
+        started_after: str | None = None,
+        started_before: str | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+        reverse_order: bool = False,  # noqa: FBT001, FBT002
+    ) -> tuple[list[Execution], str | None]:
+        """Apply filtering, sorting, and pagination to executions."""
+        # Apply filters
+        filtered: list[Execution] = []
+        for execution in executions:
+            if function_name and execution.start_input.function_name != function_name:
+                continue
+            if (
+                execution_name
+                and execution.start_input.execution_name != execution_name
+            ):
+                continue
+
+            # Status filtering
+            if status_filter and execution.current_status().value != status_filter:
+                continue
+
+            # Time filtering
+            if started_after or started_before:
+                try:
+                    operation: Operation = execution.get_operation_execution_started()
+                    if operation.start_timestamp:
+                        timestamp: float = (
+                            operation.start_timestamp.timestamp()
+                            if hasattr(operation.start_timestamp, "timestamp")
+                            else operation.start_timestamp.replace(
+                                tzinfo=UTC
+                            ).timestamp()
+                        )
+                        if started_after and timestamp < float(started_after):
+                            continue
+                        if started_before and timestamp > float(started_before):
+                            continue
+                except (ValueError, AttributeError):
+                    continue
+
+            filtered.append(execution)
+
+        # Sort by start timestamp
+        def get_sort_key(exe: Execution):
+            try:
+                op: Operation = exe.get_operation_execution_started()
+                if op.start_timestamp:
+                    return (
+                        op.start_timestamp.timestamp()
+                        if hasattr(op.start_timestamp, "timestamp")
+                        else op.start_timestamp.replace(tzinfo=UTC).timestamp()
+                    )
+            except Exception:  # noqa: BLE001, S110
+                pass
+            return 0
+
+        filtered.sort(key=get_sort_key, reverse=reverse_order)
+
+        # Apply pagination
+        if limit is not None and limit > 0:
+            end_idx: int = offset + limit
+            paginated: list[Execution] = filtered[offset:end_idx]
+            has_more: bool = end_idx < len(filtered)
+            next_marker: str | None = str(end_idx) if has_more else None
+            return paginated, next_marker
+        return filtered[offset:], None
+
+    def query(
+        self,
+        function_name: str | None = None,
+        execution_name: str | None = None,
+        status_filter: str | None = None,
+        started_after: str | None = None,
+        started_before: str | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+        reverse_order: bool = False,  # noqa: FBT001, FBT002
+    ) -> tuple[list[Execution], str | None]:
+        """Apply filtering, sorting, and pagination to executions."""
+        executions: list[Execution] = self.list_all()
+        return self.process_query(
+            executions,
+            function_name=function_name,
+            execution_name=execution_name,
+            status_filter=status_filter,
+            started_after=started_after,
+            started_before=started_before,
+            limit=limit,
+            offset=offset,
+            reverse_order=reverse_order,
+        )

aws_durable_execution_sdk_python_testing/stores/filesystem.py

@@ -0,0 +1,104 @@
+"""File system-based execution store implementation."""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import UTC, datetime
+from pathlib import Path
+
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    ResourceNotFoundException,
+)
+from aws_durable_execution_sdk_python_testing.execution import Execution
+from aws_durable_execution_sdk_python_testing.stores.base import (
+    BaseExecutionStore,
+)
+
+
+class DateTimeEncoder(json.JSONEncoder):
+    """Custom JSON encoder that handles datetime objects."""
+
+    def default(self, obj):
+        if isinstance(obj, datetime):
+            return obj.timestamp()
+        return super().default(obj)
+
+
+def datetime_object_hook(obj):
+    """JSON object hook to convert unix timestamps back to datetime objects."""
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if isinstance(value, int | float) and key.endswith(
+                ("_timestamp", "_time", "Timestamp", "Time")
+            ):
+                try:  # noqa: SIM105
+                    obj[key] = datetime.fromtimestamp(value, tz=UTC)
+                except (ValueError, OSError):
+                    # Leave as number if not a valid timestamp
+                    pass
+    return obj
+
+
+class FileSystemExecutionStore(BaseExecutionStore):
+    """File system-based execution store for persistence."""
+
+    def __init__(self, storage_dir: Path) -> None:
+        self._storage_dir = storage_dir
+
+    @classmethod
+    def create(cls, storage_dir: str | Path | None = None) -> FileSystemExecutionStore:
+        """Create a FileSystemExecutionStore with directory creation.
+
+        Args:
+            storage_dir: Directory path for storage. Defaults to '.durable_executions'
+
+        Returns:
+            FileSystemExecutionStore instance with created directory
+        """
+        path = Path(storage_dir) if storage_dir else Path(".durable_executions")
+        path.mkdir(exist_ok=True)
+        return cls(storage_dir=path)
+
+    def _get_file_path(self, execution_arn: str) -> Path:
+        """Get file path for execution ARN."""
+        # Use ARN as filename with .json extension, replacing unsafe characters
+        safe_filename = execution_arn.replace(":", "_").replace("/", "_")
+        return self._storage_dir / f"{safe_filename}.json"
+
+    def save(self, execution: Execution) -> None:
+        """Save execution to file system."""
+        file_path = self._get_file_path(execution.durable_execution_arn)
+        data = execution.to_dict()
+
+        with open(file_path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, cls=DateTimeEncoder)
+
+    def load(self, execution_arn: str) -> Execution:
+        """Load execution from file system."""
+        file_path = self._get_file_path(execution_arn)
+        if not file_path.exists():
+            msg = f"Execution {execution_arn} not found"
+            raise ResourceNotFoundException(msg)
+
+        with open(file_path, encoding="utf-8") as f:
+            data = json.load(f, object_hook=datetime_object_hook)
+
+        return Execution.from_dict(data)
+
+    def update(self, execution: Execution) -> None:
+        """Update execution in file system (same as save)."""
+        self.save(execution)
+
+    def list_all(self) -> list[Execution]:
+        """List all executions from file system."""
+        executions = []
+        for file_path in self._storage_dir.glob("*.json"):
+            try:
+                with open(file_path, encoding="utf-8") as f:
+                    data = json.load(f, object_hook=datetime_object_hook)
+                executions.append(Execution.from_dict(data))
+            except (json.JSONDecodeError, KeyError, OSError) as e:
+                logging.warning("Skipping corrupted file %s: %s", file_path, e)
+                continue
+        return executions

aws_durable_execution_sdk_python_testing/stores/memory.py

@@ -0,0 +1,38 @@
+"""In-memory execution store implementation."""
+
+from __future__ import annotations
+
+from threading import Lock
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python_testing.stores.base import (
+    BaseExecutionStore,
+)
+
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+
+
+class InMemoryExecutionStore(BaseExecutionStore):
+    """Dict-based storage for testing."""
+
+    def __init__(self) -> None:
+        self._store: dict[str, Execution] = {}
+        self._lock: Lock = Lock()
+
+    def save(self, execution: Execution) -> None:
+        with self._lock:
+            self._store[execution.durable_execution_arn] = execution
+
+    def load(self, execution_arn: str) -> Execution:
+        with self._lock:
+            return self._store[execution_arn]
+
+    def update(self, execution: Execution) -> None:
+        with self._lock:
+            self._store[execution.durable_execution_arn] = execution
+
+    def list_all(self) -> list[Execution]:
+        with self._lock:
+            return list(self._store.values())