planar 0.5.0__py3-none-any.whl

planar/workflows/__init__.py

@@ -0,0 +1,42 @@
+from .decorators import (
+    WorkflowWrapper,
+    as_step,
+    step,
+    workflow,
+)
+from .execution import (
+    execute,
+)
+from .models import (
+    LockedResource,
+    StepStatus,
+    Workflow,
+    WorkflowStep,
+)
+from .notifications import (
+    Notification,
+    WorkflowNotification,
+    WorkflowNotificationCallback,
+    workflow_notification_context,
+)
+from .orchestrator import WorkflowOrchestrator, orchestrator_context
+from .step_core import suspend
+
+__all__ = [
+    "WorkflowWrapper",
+    "workflow",
+    "step",
+    "as_step",
+    "suspend",
+    "execute",
+    "orchestrator_context",
+    "workflow_notification_context",
+    "Workflow",
+    "WorkflowStep",
+    "LockedResource",
+    "StepStatus",
+    "WorkflowOrchestrator",
+    "WorkflowNotificationCallback",
+    "WorkflowNotification",
+    "Notification",
+]

planar/workflows/context.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass, field
+from uuid import UUID
+
+from planar.task_local import TaskLocal
+from planar.workflows.models import Workflow, WorkflowStep
+
+
+@dataclass(kw_only=True)
+class ExecutionContext:
+    workflow: Workflow
+    # This might seem redundant, but it is actually necessary to prevent
+    # implicit DB calls when accessing ctx.workflow.id which causes greenlet
+    # I/O errors with async SQLAlchemy. This can happen for example when a
+    # rollback is issued, which causes SQLAlchemy to expire the objects managed
+    # by the session.
+    workflow_id: UUID
+    current_step_id: int = 0
+    step_stack: list[WorkflowStep] = field(default_factory=list)
+    # The start_workflow helper (decorators.py) has
+    # the same parameters as the original function,
+    # so we can't pass this as an argument. Instead
+    # we use a context variable to signal that the
+    # started workflow should set this one as the
+    # parent.
+    bind_parent_workflow: bool = False
+
+
+data: TaskLocal[ExecutionContext] = TaskLocal()
+
+
+def in_context() -> bool:
+    return data.is_set()
+
+
+def get_context() -> ExecutionContext:
+    return data.get()
+
+
+def set_context(ctx: ExecutionContext):
+    return data.set(ctx)
+
+
+def delete_context():
+    return data.clear()

planar/workflows/contrib.py

@@ -0,0 +1,190 @@
+from datetime import datetime, timedelta
+from functools import wraps
+from typing import Any, Callable, Coroutine, Dict
+
+from planar.logging import get_logger
+from planar.session import get_session
+from planar.utils import P, T, U, utc_now
+from planar.workflows import step
+from planar.workflows.context import get_context
+from planar.workflows.events import check_event_exists, get_latest_event
+from planar.workflows.step_core import Suspend, suspend_workflow
+from planar.workflows.tracing import trace
+
+logger = get_logger(__name__)
+
+
+@step()
+async def get_deadline(max_wait_time: float) -> datetime:
+    return utc_now() + timedelta(seconds=max_wait_time)
+
+
+@step(display_name="Wait for event")
+async def wait_for_event(
+    event_key: str,
+    max_wait_time: float = -1,
+) -> Dict[str, Any]:
+    """
+    Creates a durable step that waits for a specific event to be emitted.
+
+    Args:
+        event_key: The event identifier to wait for
+
+    Returns:
+        The event payload as a dictionary
+    """
+    logger.debug("waiting for event", event_key=event_key, max_wait_time=max_wait_time)
+    await trace("enter", event_key=event_key)
+
+    # Get workflow context
+    ctx = get_context()
+    workflow_id = ctx.workflow.id
+
+    deadline = None
+    if max_wait_time >= 0:
+        deadline = await get_deadline(max_wait_time)
+        logger.debug(
+            "calculated deadline for event", event_key=event_key, deadline=deadline
+        )
+        await trace(
+            "deadline",
+            event_key=event_key,
+            max_wait_time=max_wait_time,
+            deadline=deadline,
+        )
+
+    async def transaction():
+        # Check if the event already exists
+        event_exists = await check_event_exists(event_key, workflow_id=workflow_id)
+        logger.debug(
+            "event exists check for workflow",
+            event_key=event_key,
+            exists=event_exists,
+        )
+        await trace("check-event-exists", event_key=event_key, exists=event_exists)
+
+        if event_exists:
+            # Event exists, get the event data and continue execution immediately
+            event = await get_latest_event(event_key, workflow_id=workflow_id)
+            logger.info(
+                "event already exists, proceeding with payload",
+                event_key=event_key,
+                payload=event.payload if event else None,
+            )
+            await trace("existing-event", event_key=event_key)
+            return event.payload if event and event.payload else {}
+
+        # If deadline has passed, raise an exception
+        now = utc_now()
+        if deadline is not None and now > deadline:
+            logger.warning(
+                "timeout waiting for event",
+                event_key=event_key,
+                deadline=deadline,
+                current_time=now,
+            )
+            await trace("deadline-timeout", event_key=event_key)
+            raise ValueError(f"Timeout waiting for event ${event_key}")
+
+        logger.info(
+            "event not found, suspending workflow",
+            event_key=event_key,
+        )
+        return suspend_workflow(
+            interval=timedelta(seconds=max_wait_time) if max_wait_time > 0 else None,
+            event_key=event_key,
+        )
+
+    session = get_session()
+    result = await session.run_transaction(transaction)
+    if isinstance(result, Suspend):
+        # Suspend until event is emitted
+        logger.debug(
+            "workflow suspended, waiting for event",
+            event_key=event_key,
+        )
+        await trace("suspend", event_key=event_key)
+        await (
+            result
+        )  # This will re-raise the Suspend object's exception or re-enter the generator
+        assert False, "Suspend should never return normally"  # Should not be reached
+    logger.info(
+        "event received or processed for workflow",
+        event_key=event_key,
+        result=result,
+    )
+    return result
+
+
+def wait(
+    poll_interval: float = 60.0,
+    max_wait_time: float = 3600.0,
+):
+    """
+    Creates a durable step that repeatedly checks a condition until it returns True.
+
+    This decorator wraps a function that returns a boolean. The function will be
+    called repeatedly until it returns True or until max_wait_time is reached.
+
+    Args:
+        poll_interval: How often to check the condition
+        max_wait_time: Maximum time to wait before failing
+
+    Returns:
+        A decorator that converts a boolean-returning function into a step
+        that waits for the condition to be true
+    """
+
+    def decorator(
+        func: Callable[P, Coroutine[T, U, bool]],
+    ) -> Callable[P, Coroutine[T, U, None]]:
+        @step()
+        @wraps(func)
+        async def wait_step(*args: P.args, **kwargs: P.kwargs) -> None:
+            logger.debug(
+                "wait step called",
+                func_name=func.__name__,
+                poll_interval=poll_interval,
+                max_wait_time=max_wait_time,
+            )
+            # Set up deadline for timeout
+            deadline = None
+            if max_wait_time >= 0:
+                deadline = await get_deadline(max_wait_time)
+                logger.debug(
+                    "calculated deadline for wait step",
+                    func_name=func.__name__,
+                    deadline=deadline,
+                )
+
+            # Check the condition
+            result = await func(*args, **kwargs)
+            logger.debug(
+                "condition check returned", func_name=func.__name__, result=result
+            )
+
+            # If condition is met, return immediately
+            if result:
+                logger.info("condition met, proceeding", func_name=func.__name__)
+                return
+
+            # If deadline has passed, raise an exception
+            if deadline is not None and utc_now() > deadline:
+                logger.warning(
+                    "timeout waiting for condition to be met",
+                    func_name=func.__name__,
+                    deadline=deadline,
+                )
+                raise ValueError("Timeout waiting for condition to be met")
+
+            # Otherwise, suspend the workflow to retry later
+            logger.info(
+                "condition not met, suspending",
+                func_name=func.__name__,
+                poll_interval_seconds=poll_interval,
+            )
+            await suspend_workflow(interval=timedelta(seconds=poll_interval))
+
+        return wait_step
+
+    return decorator

planar/workflows/decorators.py

@@ -0,0 +1,217 @@
+import inspect
+import weakref
+from datetime import timedelta
+from functools import wraps
+from typing import Callable, Coroutine, Type, cast
+from uuid import UUID
+from weakref import WeakKeyDictionary
+
+import planar.workflows.notifications as notifications
+from planar.logging import get_logger
+from planar.session import get_session
+from planar.utils import P, R, T, U
+from planar.workflows.context import (
+    get_context,
+    in_context,
+)
+from planar.workflows.execution import register_workflow_function
+from planar.workflows.models import StepType, Workflow, WorkflowStatus
+from planar.workflows.orchestrator import WorkflowOrchestrator
+from planar.workflows.serialization import serialize_args
+from planar.workflows.step_core import _step, suspend_workflow
+from planar.workflows.wrappers import StepWrapper, WorkflowWrapper
+
+logger = get_logger(__name__)
+
+
+def step(
+    *,
+    max_retries: int = 0,
+    step_type: StepType = StepType.COMPUTE,
+    return_type: Type | None = None,
+    display_name: str | None = None,
+):
+    """
+    Decorator to define a step in a workflow.
+
+    This decorator is used to define a step function.
+    It will register the function with the workflow engine and allow it to be executed.
+
+    """
+    step_decorator = _step(max_retries=max_retries, return_type=return_type)
+
+    def decorator(
+        func: Callable[P, Coroutine[T, U, R]],
+        step_type: StepType = step_type,
+        display_name: str | None = display_name,
+    ) -> StepWrapper[P, T, U, R]:
+        wrapper = step_decorator(func, step_type=step_type, display_name=display_name)
+
+        @workflow(name=func.__name__ + ".auto_workflow")
+        async def auto_workflow(*args: P.args, **kwargs: P.kwargs) -> R:
+            """
+            This is a special workflow that is used to run a step in a separate asyncio task
+            """
+            result = await wrapper(*args, **kwargs)
+            return result
+
+        @wraps(func)
+        def run_step(*args: P.args, **kwargs: P.kwargs) -> Coroutine[T, U, R]:
+            """
+            If not in workflow context, then simply call the function directly.
+
+            This allows users to use their workflow code within and outside of
+            workflows.
+            """
+            if not in_context():
+                return func(*args, **kwargs)
+            return wrapper(*args, **kwargs)
+
+        step_wrapper = StepWrapper(
+            original_fn=func,
+            wrapper=wrapper,
+            wrapped_fn=run_step,
+            auto_workflow=auto_workflow,
+        )
+        return step_wrapper
+
+    return decorator
+
+
+def workflow(*, name: str | None = None):
+    """
+    Decorator to define a workflow.
+
+    This decorator is used to define a workflow function.
+    It will register the function with the workflow engine and allow it to be executed.
+
+    """
+
+    def decorator(func: Callable[P, Coroutine[T, U, R]]) -> WorkflowWrapper[P, T, U, R]:
+        if not inspect.iscoroutinefunction(func):
+            raise TypeError("Workflow functions must be coroutines")
+
+        function_name = name or func.__name__
+
+        @wraps(func)
+        async def start_workflow(*args: P.args, **kwargs: P.kwargs) -> Workflow:
+            session = get_session()
+            serialized_args, serialized_kwargs = serialize_args(func, args, kwargs)
+            workflow = Workflow(
+                function_name=function_name,
+                args=serialized_args,
+                kwargs=serialized_kwargs,
+            )
+
+            if in_context():
+                ctx = get_context()
+                if ctx.bind_parent_workflow:
+                    logger.debug(
+                        "binding parent workflow", parent_workflow_id=ctx.workflow_id
+                    )
+                    workflow.parent_id = ctx.workflow.id
+
+            session.add(workflow)
+            await session.commit()
+            notifications.workflow_started(workflow)
+            if WorkflowOrchestrator.is_set():
+                orchestrator = WorkflowOrchestrator.get()
+                orchestrator.poll_soon()
+            return workflow
+
+        @_step()
+        async def start_workflow_step(*args: P.args, **kwargs: P.kwargs) -> UUID:
+            workflow = await start_workflow(*args, **kwargs)
+            return workflow.id
+
+        async def wait_for_completion(workflow_id: UUID):
+            orchestrator = WorkflowOrchestrator.get()
+            return cast(R, await orchestrator.wait_for_completion(workflow_id))
+
+        async def run_workflow_in_new_context(*args: P.args, **kwargs: P.kwargs) -> R:
+            async with WorkflowOrchestrator.ensure_started():
+                # Running outside of a workflow execution context.
+                # Start workflow normally and wait for completion.
+                workflow = await start_workflow(*args, **kwargs)
+                workflow_id = workflow.id
+                return await wait_for_completion(workflow_id)
+
+        async def run_child_workflow(*args: P.args, **kwargs: P.kwargs) -> R:
+            ctx = get_context()
+            # Since the parent will wait, we have to
+            # create the association when creating
+            # the child
+            ctx.bind_parent_workflow = True
+            workflow_id = await start_workflow_step(*args, **kwargs)
+            # Considering the current workflow will
+            # suspend (this lose the context) it is
+            # not necessary to reset the flag here.
+            # Leaving it only for documenting intent
+            # that this is a temporary setting.
+            ctx.bind_parent_workflow = False
+            session = get_session()
+            async with session.begin_read():
+                workflow = await session.get(Workflow, workflow_id)
+            assert workflow
+
+            if workflow.status == WorkflowStatus.PENDING:
+                # Suspend for 0 seconds. Since the poll query only selects
+                # workflows that have no children, suspending for 0 seconds
+                # mean this workflow will wakeup as soon as all children finish
+                await suspend_workflow(interval=timedelta(seconds=0))
+
+            return await wait_for_completion(workflow_id)
+
+        @wraps(func)
+        def run_workflow(*args: P.args, **kwargs: P.kwargs) -> Coroutine[T, U, R]:
+            if not in_context():
+                return run_workflow_in_new_context(*args, **kwargs)
+            return run_child_workflow(*args, **kwargs)
+
+        register_workflow_function(function_name, func)
+
+        wf_wrapper = WorkflowWrapper(
+            function_name=function_name,
+            original_fn=func,
+            start=start_workflow,
+            start_step=start_workflow_step,
+            wait_for_completion=wait_for_completion,
+            wrapped_fn=run_workflow,
+        )
+
+        return wf_wrapper
+
+    return decorator
+
+
+__AS_STEP_CACHE = WeakKeyDictionary()
+
+
+def __is_workflow_step(callable: Callable[P, Coroutine[T, U, R]]) -> bool:
+    return isinstance(callable, StepWrapper)
+
+
+def as_step(
+    func: Callable[P, Coroutine[T, U, R]],
+    step_type: StepType,
+    display_name: str | None = None,
+    return_type: Type[R] | None = None,
+) -> Callable[P, Coroutine[T, U, R]]:
+    """
+    This utility fn is for treating async functions as steps without modifying
+    the original function.
+
+    Only use this where it doesn't make sense to use the @step decorator (such as third
+    party functions)
+    """
+    if __is_workflow_step(func):
+        return func
+    # cache the result to avoid reapplying the step decorator for this callable in the future
+    weak_step_callable = __AS_STEP_CACHE.get(func, None)
+    step_callable = weak_step_callable() if weak_step_callable is not None else None
+    if step_callable is None:
+        step_callable = step(return_type=return_type)(
+            func, step_type=step_type, display_name=display_name
+        )
+        __AS_STEP_CACHE[func] = weakref.ref(step_callable)
+    return step_callable