pyworkflow-engine 0.1.7__py3-none-any.whl

pyworkflow/core/workflow.py

@@ -0,0 +1,294 @@
+"""
+@workflow decorator for defining durable workflows.
+
+Workflows are orchestration functions that coordinate steps. They are
+decorated with @workflow to enable:
+- Event sourcing and deterministic replay
+- Suspension and resumption (sleep, hooks)
+- Automatic state persistence
+- Fault tolerance
+"""
+
+import functools
+from collections.abc import Callable
+from typing import Any
+
+from loguru import logger
+
+from pyworkflow.context import LocalContext, reset_context, set_context
+from pyworkflow.core.exceptions import CancellationError, ContinueAsNewSignal, SuspensionSignal
+from pyworkflow.core.registry import register_workflow
+from pyworkflow.engine.events import (
+    create_workflow_cancelled_event,
+    create_workflow_completed_event,
+    create_workflow_failed_event,
+)
+from pyworkflow.serialization.encoder import serialize
+
+
+def workflow(
+    name: str | None = None,
+    durable: bool | None = None,
+    max_duration: str | None = None,
+    tags: list[str] | None = None,
+    recover_on_worker_loss: bool | None = None,
+    max_recovery_attempts: int | None = None,
+) -> Callable:
+    """
+    Decorator to mark async functions as workflows.
+
+    Workflows are orchestration functions that coordinate steps. They can be:
+    - Durable: Event-sourced, persistent, resumable (durable=True)
+    - Transient: Simple execution without persistence overhead (durable=False)
+
+    Args:
+        name: Optional workflow name (defaults to function name)
+        durable: Whether workflow is durable (None = use configured default)
+        max_duration: Optional max duration (e.g., "1h", "30m")
+        tags: Optional list of tags for categorization (max 3 tags)
+        recover_on_worker_loss: Whether to auto-recover on worker failure
+            (None = True for durable, False for transient)
+        max_recovery_attempts: Max recovery attempts on worker failure (default: 3)
+
+    Returns:
+        Decorated workflow function
+
+    Example (durable):
+        @workflow(name="process_order", durable=True)
+        async def process_order(order_id: str):
+            order = await validate_order(order_id)
+            payment = await charge_payment(order["total"])
+            await sleep("1h")  # Can suspend and resume
+            return payment
+
+    Example (transient):
+        @workflow(durable=False)
+        async def quick_task():
+            result = await my_step()
+            return result
+
+    Example (use configured default):
+        @workflow
+        async def simple_workflow():
+            result = await my_step()
+            return result
+
+    Example (fault tolerant):
+        @workflow(durable=True, recover_on_worker_loss=True, max_recovery_attempts=5)
+        async def critical_workflow():
+            # Will auto-recover if worker crashes
+            result = await important_step()
+            return result
+
+    Example (with tags):
+        @workflow(tags=["backend", "critical"])
+        async def tagged_workflow():
+            result = await my_step()
+            return result
+    """
+
+    # Validate tags
+    validated_tags = tags or []
+    if len(validated_tags) > 3:
+        raise ValueError(
+            f"Workflows can have at most 3 tags, got {len(validated_tags)}: {validated_tags}"
+        )
+
+    def decorator(func: Callable) -> Callable:
+        workflow_name = name or func.__name__
+
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # This wrapper is called during execution by the executor
+            # The actual workflow function runs with a context set up
+            return await func(*args, **kwargs)
+
+        # Register workflow
+        register_workflow(
+            name=workflow_name,
+            func=wrapper,
+            original_func=func,
+            max_duration=max_duration,
+            tags=validated_tags,
+        )
+
+        # Store metadata on wrapper
+        wrapper.__workflow__ = True  # type: ignore[attr-defined]
+        wrapper.__workflow_name__ = workflow_name  # type: ignore[attr-defined]
+        wrapper.__workflow_durable__ = durable  # type: ignore[attr-defined]  # None = use config default
+        wrapper.__workflow_max_duration__ = max_duration  # type: ignore[attr-defined]
+        wrapper.__workflow_tags__ = validated_tags  # type: ignore[attr-defined]
+        wrapper.__workflow_recover_on_worker_loss__ = (  # type: ignore[attr-defined]
+            recover_on_worker_loss  # None = use config default
+        )
+        wrapper.__workflow_max_recovery_attempts__ = (  # type: ignore[attr-defined]
+            max_recovery_attempts  # None = use config default
+        )
+
+        return wrapper
+
+    return decorator
+
+
+async def execute_workflow_with_context(
+    workflow_func: Callable,
+    run_id: str,
+    workflow_name: str,
+    storage: Any,  # StorageBackend or None for transient
+    args: tuple,
+    kwargs: dict,
+    event_log: list | None = None,
+    durable: bool = True,
+    cancellation_requested: bool = False,
+) -> Any:
+    """
+    Execute workflow function with proper context setup.
+
+    This is called by the executor to run a workflow with:
+    - Context initialization
+    - Event logging (durable mode only)
+    - Error handling
+    - Suspension handling (durable mode only)
+    - Cancellation handling
+
+    Args:
+        workflow_func: The workflow function to execute
+        run_id: Unique run identifier
+        workflow_name: Workflow name
+        storage: Storage backend instance (None for transient)
+        args: Positional arguments
+        kwargs: Keyword arguments
+        event_log: Optional existing event log for replay
+        durable: Whether this is a durable workflow
+        cancellation_requested: Whether cancellation was requested before execution
+
+    Returns:
+        Workflow result
+
+    Raises:
+        SuspensionSignal: When workflow needs to suspend (durable only)
+        CancellationError: When workflow is cancelled
+        Exception: On workflow failure
+    """
+    # Determine if we're actually durable (need both flag and storage)
+    is_durable = durable and storage is not None
+
+    # Create workflow context using new LocalContext
+    ctx = LocalContext(
+        run_id=run_id,
+        workflow_name=workflow_name,
+        storage=storage,
+        event_log=event_log or [],
+        durable=is_durable,
+    )
+
+    # Set cancellation state if requested before execution
+    if cancellation_requested:
+        ctx.request_cancellation(reason="Cancellation requested before execution")
+
+    # Set as current context using new API
+    token = set_context(ctx)
+
+    try:
+        # Note: Event replay is handled by LocalContext in its constructor
+        # when event_log is provided
+
+        logger.info(
+            f"Executing workflow: {workflow_name}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            durable=is_durable,
+            is_replay=bool(event_log),
+        )
+
+        # Execute workflow function
+        result = await workflow_func(*args, **kwargs)
+
+        # Record completion event (durable mode only)
+        if is_durable:
+            # Validate event limits before recording completion
+            await ctx.validate_event_limits()
+
+            completion_event = create_workflow_completed_event(
+                run_id, serialize(result), workflow_name
+            )
+            await storage.record_event(completion_event)
+
+        logger.info(
+            f"Workflow completed: {workflow_name}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            durable=is_durable,
+        )
+
+        return result
+
+    except SuspensionSignal as e:
+        # Workflow suspended (sleep/hook) - only happens in durable mode
+        logger.info(
+            f"Workflow suspended: {e.reason}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            reason=e.reason,
+        )
+        raise
+
+    except ContinueAsNewSignal as e:
+        # Workflow continuing as new execution
+        logger.info(
+            f"Workflow continuing as new: {workflow_name}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            new_args=e.workflow_args,
+            new_kwargs=e.workflow_kwargs,
+        )
+        # Re-raise for caller (executor) to handle continuation
+        raise
+
+    except CancellationError as e:
+        # Workflow was cancelled
+        logger.info(
+            f"Workflow cancelled: {workflow_name}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            reason=e.reason,
+        )
+
+        # Record cancellation event (durable mode only)
+        if is_durable:
+            cancelled_event = create_workflow_cancelled_event(
+                run_id=run_id,
+                reason=e.reason,
+                cleanup_completed=True,
+            )
+            await storage.record_event(cancelled_event)
+
+        raise
+
+    except Exception as e:
+        # Workflow failed
+        logger.error(
+            f"Workflow failed: {workflow_name}",
+            run_id=run_id,
+            workflow_name=workflow_name,
+            error=str(e),
+            exc_info=True,
+        )
+
+        # Record failure event (durable mode only)
+        if is_durable:
+            import traceback
+
+            failure_event = create_workflow_failed_event(
+                run_id=run_id,
+                error=str(e),
+                error_type=type(e).__name__,
+                traceback=traceback.format_exc(),
+            )
+            await storage.record_event(failure_event)
+
+        raise
+
+    finally:
+        # Clear context using new API
+        reset_context(token)

pyworkflow/discovery.py

@@ -0,0 +1,248 @@
+"""
+Workflow discovery utilities.
+
+This module provides functions to discover and register workflows by importing
+Python modules containing @workflow decorated functions.
+"""
+
+import importlib
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+
+
+class DiscoveryError(Exception):
+    """Raised when workflow discovery fails."""
+
+    pass
+
+
+def _find_project_root() -> Path | None:
+    """
+    Find the project root by looking for common project markers.
+
+    Searches upward from the current directory for:
+    - pyproject.toml
+    - setup.py
+    - .git directory
+
+    Returns:
+        Path to project root if found, None otherwise
+    """
+    current = Path.cwd()
+
+    for path in [current, *list(current.parents)]:
+        if (path / "pyproject.toml").exists():
+            return path
+        if (path / "setup.py").exists():
+            return path
+        if (path / ".git").exists():
+            return path
+
+    return None
+
+
+def _ensure_project_in_path() -> None:
+    """Add the project root and current directory to sys.path if not already present."""
+    # Add current directory first
+    cwd = str(Path.cwd())
+    if cwd not in sys.path:
+        sys.path.insert(0, cwd)
+        logger.debug(f"Added current directory to path: {cwd}")
+
+    # Add project root
+    project_root = _find_project_root()
+    if project_root:
+        root_str = str(project_root)
+        if root_str not in sys.path:
+            sys.path.insert(0, root_str)
+            logger.debug(f"Added project root to path: {root_str}")
+
+
+def _import_module(module_path: str) -> bool:
+    """
+    Import a single module to trigger workflow registration.
+
+    Args:
+        module_path: Python module path (e.g., "myapp.workflows")
+
+    Returns:
+        True if import succeeded, False otherwise
+    """
+    try:
+        importlib.import_module(module_path)
+        logger.info(f"Discovered workflows from module: {module_path}")
+        return True
+    except ModuleNotFoundError as e:
+        logger.error(f"Failed to import module '{module_path}': {e}")
+        return False
+    except Exception as e:
+        logger.error(f"Error importing module '{module_path}': {e}")
+        return False
+
+
+def _load_yaml_config() -> dict[str, Any] | None:
+    """
+    Load pyworkflow.config.yaml from current directory.
+
+    Returns:
+        Configuration dictionary if found, None otherwise
+    """
+    config_path = Path.cwd() / "pyworkflow.config.yaml"
+    if not config_path.exists():
+        return None
+
+    try:
+        import yaml
+
+        with open(config_path) as f:
+            config = yaml.safe_load(f)
+            logger.debug(f"Loaded config from: {config_path}")
+            return config
+    except ImportError:
+        logger.warning("PyYAML not installed, skipping YAML config")
+        return None
+    except Exception as e:
+        logger.error(f"Failed to load YAML config: {e}")
+        return None
+
+
+def discover_workflows(
+    module_path: str | None = None,
+    config: dict[str, Any] | None = None,
+    config_path: str | Path | None = None,
+) -> None:
+    """
+    Import Python modules to trigger workflow registration.
+
+    Workflows are registered when their module is imported, as the @workflow
+    decorator registers them in the global registry. This function handles
+    importing the appropriate module based on configuration priority.
+
+    Priority:
+    1. Explicit module_path argument
+    2. PYWORKFLOW_DISCOVER environment variable
+    3. Config dict (from YAML) with 'module' or 'modules' key
+    4. pyworkflow.config.yaml in current directory
+
+    Args:
+        module_path: Explicit module path to import
+        config: Configuration dict containing 'module' or 'modules' key
+        config_path: Path to the config file. If provided, the directory
+            containing the config file will be added to sys.path for imports.
+
+    Raises:
+        DiscoveryError: If specified modules cannot be imported
+
+    Examples:
+        # Explicit module
+        discover_workflows("myapp.workflows")
+
+        # From config dict with path context
+        discover_workflows(config={"module": "myapp.workflows"}, config_path="/app/pyworkflow.config.yaml")
+
+        # From environment variable
+        os.environ["PYWORKFLOW_DISCOVER"] = "myapp.workflows"
+        discover_workflows()
+
+        # From pyworkflow.config.yaml in cwd
+        discover_workflows()
+    """
+    # Ensure project root is in Python path for module imports
+    _ensure_project_in_path()
+
+    # If config_path provided, add its directory to sys.path
+    # This allows importing modules relative to the config file location
+    if config_path:
+        config_dir = Path(config_path).parent.resolve()
+        config_dir_str = str(config_dir)
+        if config_dir_str not in sys.path:
+            sys.path.insert(0, config_dir_str)
+            logger.debug(f"Added config directory to path: {config_dir_str}")
+
+    # Priority 1: Explicit module path
+    if module_path:
+        logger.debug(f"Discovering from explicit module: {module_path}")
+        if not _import_module(module_path):
+            raise DiscoveryError(
+                f"Cannot import module '{module_path}'. "
+                f"Make sure the module exists and is in your Python path."
+            )
+        return
+
+    # Priority 2: Environment variable
+    env_modules = os.getenv("PYWORKFLOW_DISCOVER", "")
+    if env_modules:
+        logger.debug(f"Discovering from PYWORKFLOW_DISCOVER: {env_modules}")
+        env_module_list = [m.strip() for m in env_modules.split(",") if m.strip()]
+        failed = []
+        for module in env_module_list:
+            if not _import_module(module):
+                failed.append(module)
+
+        if failed:
+            raise DiscoveryError(
+                f"Cannot import modules from PYWORKFLOW_DISCOVER: {', '.join(failed)}. "
+                f"Make sure the modules exist and are in your Python path."
+            )
+        return
+
+    # Priority 3: Config dict (passed from configure_from_yaml)
+    if config:
+        modules: list[str] = []
+
+        if "module" in config:
+            modules.append(config["module"])
+        if "modules" in config:
+            modules.extend(config["modules"])
+
+        if modules:
+            logger.debug(f"Discovering from config: {modules}")
+            failed = []
+            for module in modules:
+                if not _import_module(module):
+                    failed.append(module)
+
+            if failed:
+                project_root = _find_project_root()
+                raise DiscoveryError(
+                    f"Cannot import modules from config: {', '.join(failed)}. "
+                    f"Make sure the modules exist and are in your Python path.\n"
+                    f"  Current directory: {Path.cwd()}\n"
+                    f"  Project root: {project_root or 'not found'}"
+                )
+            return
+
+    # Priority 4: pyworkflow.config.yaml in cwd
+    yaml_config = _load_yaml_config()
+    if yaml_config:
+        # Support both 'module' (single) and 'modules' (list)
+        modules = []
+
+        if "module" in yaml_config:
+            modules.append(yaml_config["module"])
+        if "modules" in yaml_config:
+            modules.extend(yaml_config["modules"])
+
+        if modules:
+            logger.debug(f"Discovering from pyworkflow.config.yaml: {modules}")
+            failed = []
+            for module in modules:
+                if not _import_module(module):
+                    failed.append(module)
+
+            if failed:
+                project_root = _find_project_root()
+                raise DiscoveryError(
+                    f"Cannot import modules from pyworkflow.config.yaml: {', '.join(failed)}. "
+                    f"Make sure the modules exist and are in your Python path.\n"
+                    f"  Current directory: {Path.cwd()}\n"
+                    f"  Project root: {project_root or 'not found'}"
+                )
+            return
+
+    # No discovery source found
+    logger.debug("No workflow module specified for discovery")