pyworkflow-engine 0.1.7__py3-none-any.whl

pyworkflow/scheduler/local.py

@@ -0,0 +1,248 @@
+"""
+Local scheduler for PyWorkflow.
+
+Provides a polling-based scheduler that runs in the same process as your application.
+This is the local runtime equivalent of Celery Beat.
+"""
+
+import asyncio
+from datetime import UTC, datetime
+
+from loguru import logger
+
+from pyworkflow.primitives.schedule import trigger_schedule
+from pyworkflow.storage.base import StorageBackend
+from pyworkflow.storage.schemas import OverlapPolicy, Schedule
+from pyworkflow.utils.schedule import calculate_next_run_time
+
+
+class LocalScheduler:
+    """
+    Local scheduler that polls storage for due schedules.
+
+    This is the local runtime equivalent of Celery Beat. It polls storage
+    for schedules that are due and triggers them using the configured runtime.
+
+    Usage:
+        from pyworkflow.scheduler import LocalScheduler
+        from pyworkflow.storage import InMemoryStorageBackend
+
+        storage = InMemoryStorageBackend()
+        scheduler = LocalScheduler(storage=storage, poll_interval=5.0)
+
+        # Run forever
+        await scheduler.run()
+
+        # Or run for a specific duration
+        await scheduler.run(duration=60.0)  # Run for 60 seconds
+
+    Example with CLI:
+        pyworkflow scheduler run --poll-interval 5
+
+    Args:
+        storage: Storage backend to use. If None, uses configured storage.
+        poll_interval: Seconds between storage polls (default: 5.0)
+    """
+
+    def __init__(
+        self,
+        storage: StorageBackend | None = None,
+        poll_interval: float = 5.0,
+    ):
+        """
+        Initialize the local scheduler.
+
+        Args:
+            storage: Storage backend to use. If None, uses configured storage.
+            poll_interval: Seconds between storage polls (default: 5.0)
+        """
+        self._storage = storage
+        self.poll_interval = poll_interval
+        self._running = False
+        self._start_time: datetime | None = None
+
+    @property
+    def storage(self) -> StorageBackend:
+        """Get the storage backend, resolving from config if needed."""
+        if self._storage is None:
+            from pyworkflow.config import get_config
+
+            config = get_config()
+            if config.storage is None:
+                raise ValueError(
+                    "Storage backend required. Configure storage or pass to constructor."
+                )
+            self._storage = config.storage
+        return self._storage
+
+    async def run(self, duration: float | None = None) -> None:
+        """
+        Run the scheduler loop.
+
+        Args:
+            duration: Optional duration in seconds to run. If None, runs forever.
+        """
+        self._running = True
+        self._start_time = datetime.now(UTC)
+
+        logger.info(
+            "Local scheduler started",
+            poll_interval=self.poll_interval,
+            duration=duration,
+        )
+
+        try:
+            while self._running:
+                # Check duration limit
+                if duration is not None:
+                    elapsed = (datetime.now(UTC) - self._start_time).total_seconds()
+                    if elapsed >= duration:
+                        logger.info("Scheduler duration limit reached")
+                        break
+
+                try:
+                    await self._tick()
+                except Exception as e:
+                    logger.error(f"Scheduler tick failed: {e}")
+
+                await asyncio.sleep(self.poll_interval)
+        finally:
+            self._running = False
+            logger.info("Local scheduler stopped")
+
+    async def _tick(self) -> None:
+        """Process due schedules in a single tick."""
+        now = datetime.now(UTC)
+        due_schedules = await self.storage.get_due_schedules(now)
+
+        if due_schedules:
+            logger.debug(f"Found {len(due_schedules)} due schedule(s)")
+
+        for schedule in due_schedules:
+            await self._process_schedule(schedule, now)
+
+    async def _process_schedule(self, schedule: Schedule, now: datetime) -> None:
+        """
+        Process a single due schedule.
+
+        Args:
+            schedule: The schedule to process
+            now: Current timestamp
+        """
+        # Check overlap policy
+        should_run, reason = await self._check_overlap_policy(schedule)
+
+        if not should_run:
+            logger.info(
+                f"Skipping schedule: {reason}",
+                schedule_id=schedule.schedule_id,
+            )
+            schedule.skipped_runs += 1
+            schedule.next_run_time = calculate_next_run_time(schedule.spec, last_run=now, now=now)
+            await self.storage.update_schedule(schedule)
+            return
+
+        # Trigger the schedule (uses runtime-agnostic start())
+        try:
+            logger.info(
+                "Triggering schedule",
+                schedule_id=schedule.schedule_id,
+                workflow_name=schedule.workflow_name,
+            )
+            run_id = await trigger_schedule(schedule.schedule_id, storage=self.storage)
+            logger.info(
+                "Schedule triggered successfully",
+                schedule_id=schedule.schedule_id,
+                run_id=run_id,
+            )
+        except Exception as e:
+            logger.error(
+                "Failed to trigger schedule",
+                schedule_id=schedule.schedule_id,
+                error=str(e),
+            )
+            # Update failed runs count
+            schedule.failed_runs += 1
+            schedule.next_run_time = calculate_next_run_time(schedule.spec, last_run=now, now=now)
+            await self.storage.update_schedule(schedule)
+
+    async def _check_overlap_policy(self, schedule: Schedule) -> tuple[bool, str]:
+        """
+        Check if schedule should run based on overlap policy.
+
+        Args:
+            schedule: The schedule to check
+
+        Returns:
+            Tuple of (should_run, reason if not running)
+        """
+        # No running workflows = always run
+        if not schedule.running_run_ids:
+            return True, ""
+
+        policy = schedule.overlap_policy
+
+        if policy == OverlapPolicy.SKIP:
+            return False, "previous run still active (SKIP policy)"
+
+        elif policy == OverlapPolicy.ALLOW_ALL:
+            return True, ""
+
+        elif policy == OverlapPolicy.BUFFER_ONE:
+            if schedule.buffered_count >= 1:
+                return False, "already buffered one run (BUFFER_ONE policy)"
+            schedule.buffered_count += 1
+            return True, ""
+
+        elif policy == OverlapPolicy.BUFFER_ALL:
+            schedule.buffered_count += 1
+            return True, ""
+
+        elif policy == OverlapPolicy.CANCEL_OTHER:
+            # Cancel running workflows before starting new one
+            from pyworkflow.engine.executor import cancel_workflow
+
+            for run_id in list(schedule.running_run_ids):
+                try:
+                    await cancel_workflow(run_id)
+                    logger.info(
+                        "Cancelled previous run",
+                        schedule_id=schedule.schedule_id,
+                        run_id=run_id,
+                    )
+                except Exception as e:
+                    logger.warning(
+                        "Failed to cancel previous run",
+                        schedule_id=schedule.schedule_id,
+                        run_id=run_id,
+                        error=str(e),
+                    )
+            return True, ""
+
+        # Default: allow
+        return True, ""
+
+    def stop(self) -> None:
+        """Stop the scheduler gracefully."""
+        logger.info("Stopping local scheduler...")
+        self._running = False
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the scheduler is currently running."""
+        return self._running
+
+    async def tick_once(self) -> int:
+        """
+        Process due schedules once (useful for testing).
+
+        Returns:
+            Number of schedules processed
+        """
+        now = datetime.now(UTC)
+        due_schedules = await self.storage.get_due_schedules(now)
+
+        for schedule in due_schedules:
+            await self._process_schedule(schedule, now)
+
+        return len(due_schedules)

pyworkflow/serialization/decoder.py

@@ -0,0 +1,146 @@
+"""
+Custom decoding for complex Python types.
+
+Reverses the encoding performed by encoder.py to reconstruct Python objects.
+"""
+
+import base64
+import json
+from datetime import date, datetime, timedelta
+from decimal import Decimal
+from typing import Any
+
+import cloudpickle
+
+
+def enhanced_json_decoder(dct: dict) -> Any:
+    """
+    Decode custom JSON types back to Python objects.
+
+    Args:
+        dct: Dictionary from JSON parsing
+
+    Returns:
+        Decoded Python object
+    """
+    if "__type__" not in dct:
+        return dct
+
+    type_name = dct["__type__"]
+
+    # Datetime types
+    if type_name == "datetime":
+        return datetime.fromisoformat(dct["value"])
+
+    if type_name == "date":
+        return date.fromisoformat(dct["value"])
+
+    if type_name == "timedelta":
+        return timedelta(seconds=dct["value"])
+
+    # Numeric types
+    if type_name == "decimal":
+        return Decimal(dct["value"])
+
+    # Enum types
+    if type_name == "enum":
+        # Dynamically import and reconstruct enum
+        module_name, class_name = dct["class"].rsplit(".", 1)
+        try:
+            module = __import__(module_name, fromlist=[class_name])
+            enum_class = getattr(module, class_name)
+            return enum_class(dct["value"])
+        except (ImportError, AttributeError, ValueError):
+            # If enum can't be reconstructed, return the dict
+            return dct
+
+    # Exception types
+    if type_name == "exception":
+        # Reconstruct exception
+        exc_class_name = dct["class"]
+        try:
+            # Try to get exception class from builtins
+            exc_class = getattr(__builtins__, exc_class_name, Exception)
+        except (AttributeError, TypeError):
+            exc_class = Exception
+
+        try:
+            return exc_class(*dct.get("args", []))
+        except Exception:
+            # If reconstruction fails, return generic exception
+            return Exception(dct.get("message", "Unknown error"))
+
+    # Binary data
+    if type_name == "bytes":
+        return base64.b64decode(dct["value"])
+
+    # Sets
+    if type_name == "set":
+        return set(dct["value"])
+
+    # Cloudpickle objects
+    if type_name == "cloudpickle":
+        try:
+            return cloudpickle.loads(base64.b64decode(dct["value"]))
+        except Exception:
+            # If unpickling fails, return the dict
+            return dct
+
+    # Unknown type - return as-is
+    return dct
+
+
+def deserialize(json_str: str) -> Any:
+    """
+    Deserialize JSON string to Python object.
+
+    Args:
+        json_str: JSON string
+
+    Returns:
+        Deserialized Python object
+
+    Examples:
+        >>> deserialize('{"name": "Alice", "age": 30}')
+        {'name': 'Alice', 'age': 30}
+
+        >>> deserialize('{"__type__": "datetime", "value": "2025-01-15T10:30:00"}')
+        datetime.datetime(2025, 1, 15, 10, 30)
+    """
+    return json.loads(json_str, object_hook=enhanced_json_decoder)
+
+
+def deserialize_args(json_str: str) -> tuple[Any, ...]:
+    """
+    Deserialize JSON string to tuple of arguments.
+
+    Args:
+        json_str: JSON string of list
+
+    Returns:
+        Tuple of arguments
+
+    Examples:
+        >>> deserialize_args('["hello", 42, true]')
+        ('hello', 42, True)
+    """
+    args_list = json.loads(json_str, object_hook=enhanced_json_decoder)
+    return tuple(args_list) if isinstance(args_list, list) else ()
+
+
+def deserialize_kwargs(json_str: str) -> dict:
+    """
+    Deserialize JSON string to dictionary of keyword arguments.
+
+    Args:
+        json_str: JSON string of dict
+
+    Returns:
+        Dictionary of keyword arguments
+
+    Examples:
+        >>> deserialize_kwargs('{"name": "Alice", "age": 30}')
+        {'name': 'Alice', 'age': 30}
+    """
+    kwargs_dict = json.loads(json_str, object_hook=enhanced_json_decoder)
+    return kwargs_dict if isinstance(kwargs_dict, dict) else {}

pyworkflow/serialization/encoder.py

@@ -0,0 +1,162 @@
+"""
+Custom encoding for complex Python types.
+
+Supports serialization of:
+- Primitives (int, str, bool, float, None)
+- Collections (list, dict, tuple, set)
+- Dates (datetime, date, timedelta)
+- Special types (Decimal, Enum, Exception, bytes)
+- Complex objects (via cloudpickle)
+"""
+
+import base64
+import json
+from datetime import date, datetime, timedelta
+from decimal import Decimal
+from enum import Enum
+from typing import Any
+
+import cloudpickle
+
+
+class EnhancedJSONEncoder(json.JSONEncoder):
+    """
+    JSON encoder with support for additional Python types.
+
+    Handles datetime, Decimal, Enum, Exception, bytes, and complex objects.
+    """
+
+    def default(self, obj: Any) -> Any:
+        """
+        Encode object to JSON-serializable form.
+
+        Args:
+            obj: Object to encode
+
+        Returns:
+            JSON-serializable representation
+        """
+        # Datetime types
+        if isinstance(obj, datetime):
+            return {
+                "__type__": "datetime",
+                "value": obj.isoformat(),
+            }
+
+        if isinstance(obj, date):
+            return {
+                "__type__": "date",
+                "value": obj.isoformat(),
+            }
+
+        if isinstance(obj, timedelta):
+            return {
+                "__type__": "timedelta",
+                "value": obj.total_seconds(),
+            }
+
+        # Numeric types
+        if isinstance(obj, Decimal):
+            return {
+                "__type__": "decimal",
+                "value": str(obj),
+            }
+
+        # Enum types
+        if isinstance(obj, Enum):
+            return {
+                "__type__": "enum",
+                "class": f"{obj.__class__.__module__}.{obj.__class__.__name__}",
+                "value": obj.value,
+            }
+
+        # Exception types
+        if isinstance(obj, Exception):
+            return {
+                "__type__": "exception",
+                "class": obj.__class__.__name__,
+                "message": str(obj),
+                "args": obj.args,
+            }
+
+        # Binary data
+        if isinstance(obj, bytes):
+            return {
+                "__type__": "bytes",
+                "value": base64.b64encode(obj).decode("ascii"),
+            }
+
+        # Sets (convert to list)
+        if isinstance(obj, set):
+            return {
+                "__type__": "set",
+                "value": list(obj),
+            }
+
+        # Complex objects - fall back to cloudpickle
+        try:
+            return {
+                "__type__": "cloudpickle",
+                "value": base64.b64encode(cloudpickle.dumps(obj)).decode("ascii"),
+            }
+        except Exception as e:
+            # If cloudpickle fails, raise serialization error
+            raise TypeError(
+                f"Object of type {type(obj).__name__} is not JSON serializable "
+                f"and could not be pickled: {e}"
+            ) from e
+
+
+def serialize(obj: Any) -> str:
+    """
+    Serialize Python object to JSON string.
+
+    Args:
+        obj: Object to serialize
+
+    Returns:
+        JSON string
+
+    Examples:
+        >>> serialize({"name": "Alice", "age": 30})
+        '{"name": "Alice", "age": 30}'
+
+        >>> from datetime import datetime
+        >>> serialize(datetime(2025, 1, 15, 10, 30))
+        '{"__type__": "datetime", "value": "2025-01-15T10:30:00"}'
+    """
+    return json.dumps(obj, cls=EnhancedJSONEncoder)
+
+
+def serialize_args(*args: Any) -> str:
+    """
+    Serialize positional arguments to JSON string.
+
+    Args:
+        *args: Positional arguments
+
+    Returns:
+        JSON string of list
+
+    Examples:
+        >>> serialize_args("hello", 42, True)
+        '["hello", 42, true]'
+    """
+    return json.dumps(list(args), cls=EnhancedJSONEncoder)
+
+
+def serialize_kwargs(**kwargs: Any) -> str:
+    """
+    Serialize keyword arguments to JSON string.
+
+    Args:
+        **kwargs: Keyword arguments
+
+    Returns:
+        JSON string of dict
+
+    Examples:
+        >>> serialize_kwargs(name="Alice", age=30)
+        '{"name": "Alice", "age": 30}'
+    """
+    return json.dumps(kwargs, cls=EnhancedJSONEncoder)

pyworkflow/storage/__init__.py

@@ -0,0 +1,54 @@
+"""
+Storage backends for PyWorkflow.
+
+Provides different storage implementations for workflow state persistence.
+"""
+
+from pyworkflow.storage.base import StorageBackend
+from pyworkflow.storage.config import config_to_storage, storage_to_config
+from pyworkflow.storage.file import FileStorageBackend
+from pyworkflow.storage.memory import InMemoryStorageBackend
+from pyworkflow.storage.schemas import (
+    Hook,
+    HookStatus,
+    RunStatus,
+    StepExecution,
+    StepStatus,
+    WorkflowRun,
+)
+
+# SQLite backend - optional import (requires sqlite3 in Python build)
+try:
+    from pyworkflow.storage.sqlite import SQLiteStorageBackend
+except ImportError:
+    SQLiteStorageBackend = None  # type: ignore
+
+# PostgreSQL backend - optional import (requires asyncpg)
+try:
+    from pyworkflow.storage.postgres import PostgresStorageBackend
+except ImportError:
+    PostgresStorageBackend = None  # type: ignore
+
+# DynamoDB backend - optional import (requires aiobotocore)
+try:
+    from pyworkflow.storage.dynamodb import DynamoDBStorageBackend
+except ImportError:
+    DynamoDBStorageBackend = None  # type: ignore
+
+__all__ = [
+    "StorageBackend",
+    "FileStorageBackend",
+    "InMemoryStorageBackend",
+    "SQLiteStorageBackend",
+    "PostgresStorageBackend",
+    "DynamoDBStorageBackend",
+    "WorkflowRun",
+    "StepExecution",
+    "Hook",
+    "RunStatus",
+    "StepStatus",
+    "HookStatus",
+    # Config utilities
+    "storage_to_config",
+    "config_to_storage",
+]