planar 0.5.0__py3-none-any.whl

planar/workflows/events.py

@@ -0,0 +1,185 @@
+"""
+Event system for workflow engine.
+
+This module provides functions for emitting and checking events that workflows
+might be waiting for.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, Optional, cast
+from uuid import UUID, uuid4
+
+from sqlmodel import col, select, update
+
+from planar.logging import get_logger
+from planar.session import get_session
+from planar.workflows.models import Workflow, WorkflowEvent
+from planar.workflows.orchestrator import WorkflowOrchestrator
+from planar.workflows.tracing import trace
+
+logger = get_logger(__name__)
+
+
+async def emit_event(
+    event_key: str,
+    payload: Optional[Dict[str, Any]] = None,
+    workflow_id: Optional[UUID] = None,
+) -> tuple[WorkflowEvent, int]:
+    """
+    Emit a new event that workflows might be waiting for.
+
+    Args:
+        event_key: The event identifier
+        payload: Optional data to include with the event
+        workflow_id: Optional workflow ID if the event is targeted at a specific workflow
+
+    Returns:
+        The created event record
+    """
+    logger.debug(
+        "emitting event",
+        event_key=event_key,
+        workflow_id=str(workflow_id),
+        payload_keys=list(payload.keys()) if payload else None,
+    )
+    await trace("enter", event_key=event_key)
+    session = get_session()
+
+    select_condition = col(Workflow.waiting_for_event) == event_key
+    if workflow_id:
+        select_condition &= col(Workflow.id) == workflow_id
+    update_query = (
+        update(Workflow)
+        .where(select_condition)
+        .values(waiting_for_event=None, wakeup_at=None)
+        .returning(col(Workflow.id))
+    )
+
+    async def transaction():
+        # Update affected events
+        workflow_ids = (await session.exec(cast(Any, update_query))).all()
+        logger.info(
+            "event woke up workflows", event_key=event_key, count=len(workflow_ids)
+        )
+        await trace(
+            "wake-affected-workflows", event_key=event_key, count=len(workflow_ids)
+        )
+        # Create the event record
+        event = WorkflowEvent(
+            id=uuid4(),
+            event_key=event_key,
+            workflow_id=workflow_id,
+            payload=payload or {},
+        )
+        session.add(event)
+        logger.debug("event record created", event_key=event_key, event_id=event.id)
+        await trace("add-event-record", event_key=event_key)
+
+        return event, workflow_ids
+
+    event, workflow_ids = await session.run_transaction(transaction)
+    await trace("commit", event_key=event_key)
+    logger.info("event committed to database", event_key=event_key, event_id=event.id)
+
+    if workflow_ids and WorkflowOrchestrator.is_set():
+        logger.debug("requesting orchestrator poll due to event", event_key=event_key)
+        WorkflowOrchestrator.get().poll_soon()
+
+    await trace("return", event_key=event_key)
+    return event, len(workflow_ids)
+
+
+async def check_event_exists(
+    event_key: str, since: Optional[datetime] = None, workflow_id: Optional[UUID] = None
+) -> bool:
+    """
+    Check if an event with the given key exists, optionally after a specific time.
+
+    Args:
+        event_key: The event identifier
+        since: Only consider events after this time
+        workflow_id: Optional workflow ID to check for workflow-specific events
+
+    Returns:
+        True if a matching event exists, False otherwise
+    """
+    logger.debug(
+        "checking if event exists",
+        event_key=event_key,
+        since=since,
+    )
+    session = get_session()
+
+    # Start building the query
+    query = select(WorkflowEvent).where(WorkflowEvent.event_key == event_key)
+
+    # If a timestamp is provided, only check for events after that time
+    if since:
+        query = query.where(WorkflowEvent.timestamp > since)
+
+    # If a workflow ID is provided, check for events specific to that workflow
+    # or global events (no workflow ID)
+    if workflow_id:
+        query = query.where(
+            (col(WorkflowEvent.workflow_id) == workflow_id)
+            | (col(WorkflowEvent.workflow_id).is_(None))
+        )
+
+    # Execute the query and check if any result exists
+    event = (await session.exec(query)).first()
+    exists = event is not None
+    logger.debug("event exists check result", event_key=event_key, exists=exists)
+    return exists
+
+
+async def get_latest_event(
+    event_key: str, since: Optional[datetime] = None, workflow_id: Optional[UUID] = None
+) -> Optional[WorkflowEvent]:
+    """
+    Get the most recent event with the given key.
+
+    Args:
+        event_key: The event identifier
+        since: Only consider events after this time
+        workflow_id: Optional workflow ID to check for workflow-specific events
+
+    Returns:
+        The most recent matching event, or None if no match found
+    """
+    logger.debug(
+        "getting latest event",
+        event_key=event_key,
+        since=since,
+    )
+    session = get_session()
+
+    # Start building the query
+    query = select(WorkflowEvent).where(WorkflowEvent.event_key == event_key)
+
+    # If a timestamp is provided, only check for events after that time
+    if since:
+        query = query.where(WorkflowEvent.timestamp > since)
+
+    # If a workflow ID is provided, check for events specific to that workflow
+    # or global events (no workflow ID)
+    if workflow_id:
+        query = query.where(
+            (col(WorkflowEvent.workflow_id) == workflow_id)
+            | (col(WorkflowEvent.workflow_id).is_(None))
+        )
+
+    # Order by timestamp descending and get the first (most recent) result
+    query = query.order_by(col(WorkflowEvent.timestamp).desc())
+
+    # Execute the query and return the first result (or None)
+    event = (await session.exec(query)).first()
+    if event:
+        logger.debug(
+            "latest event found",
+            event_key=event_key,
+            event_id=event.id,
+            timestamp=event.timestamp,
+        )
+    else:
+        logger.debug("no event found with given criteria", event_key=event_key)
+    return event

planar/workflows/exceptions.py

@@ -0,0 +1,34 @@
+import builtins
+from typing import Any
+
+
+class StepError(Exception):
+    def __init__(self, type: str, message: str):
+        self.type = type
+        self.original_message = message
+        super().__init__(f"{type}: {message}")
+
+
+class WorkflowError(StepError):
+    pass
+
+
+class LockResourceFailed(Exception):
+    pass
+
+
+class NonDeterministicStepCallError(Exception):
+    """Raised when a step call is not deterministic compared to previous executions."""
+
+    pass
+
+
+def try_restore_exception(exception: dict[str, Any]) -> Exception:
+    exc_type = exception["type"]
+    exc_message = exception["message"]
+    # Try to get the exception class from the builtins module
+    exc_class = getattr(builtins, exc_type, None)
+    if isinstance(exc_class, type) and issubclass(exc_class, Exception):
+        return exc_class(exc_message)
+    # Fallback to a custom exception if not found
+    return StepError(exc_type, exc_message)

planar/workflows/execution.py

@@ -0,0 +1,198 @@
+import traceback
+from datetime import timedelta
+from typing import Any, Callable, Coroutine, Generic
+from weakref import WeakValueDictionary
+
+import planar.workflows.notifications as notifications
+from planar.logging import get_logger
+from planar.session import get_session
+from planar.utils import R, T, U
+from planar.workflows.context import ExecutionContext, delete_context, set_context
+from planar.workflows.exceptions import try_restore_exception
+from planar.workflows.lock import lock_workflow
+from planar.workflows.models import Workflow, WorkflowStatus
+from planar.workflows.serialization import (
+    deserialize_args,
+    deserialize_result,
+    serialize_result,
+)
+from planar.workflows.step_core import Suspend
+from planar.workflows.tracing import trace
+
+_DEFAULT_LOCK_DURATION = timedelta(minutes=10)
+_WORKFLOW_FUNCTION_REGISTRY: WeakValueDictionary[str, Callable[..., Any]] = (
+    WeakValueDictionary()
+)
+
+logger = get_logger(__name__)
+
+
+class YieldWrapper:
+    def __init__(self, value):
+        self.value = value
+
+    def __await__(self):
+        return (yield self.value)
+
+
+def register_workflow_function(name: str, func: Callable[..., Any]):
+    _WORKFLOW_FUNCTION_REGISTRY[name] = func
+
+
+class StepperResult(Generic[R]):
+    def __init__(self, *, value: R | None, suspend: Suspend | None):
+        self.value = value
+        self.suspend = suspend
+
+
+async def stepper(coro: Coroutine[T, U, R]) -> StepperResult[R]:
+    logger.debug("stepper started")
+    coroutine_iterator = coro.__await__()
+    try:
+        yielded = next(coroutine_iterator)  # Start the coroutine
+        while True:
+            if isinstance(yielded, Suspend):
+                logger.debug("stepper encountered suspend")
+                return StepperResult(value=None, suspend=yielded)
+            else:
+                try:
+                    result = await YieldWrapper(yielded)
+                except BaseException as e:
+                    # if an exception is raised by the event loop
+                    # (most likely a cancellation), propagate it to the coroutine
+                    logger.debug(
+                        "stepper propagating exception to coroutine",
+                        exception_type=type(e).__name__,
+                    )
+                    yielded = coroutine_iterator.throw(e)
+                    continue
+                # send the result back to the coroutine
+                yielded = coroutine_iterator.send(result)
+    except StopIteration as e:
+        logger.debug("stepper finished with stopiteration")
+        return StepperResult(value=e.value, suspend=None)
+
+
+def workflow_result(workflow: Workflow):
+    if workflow.status == WorkflowStatus.SUCCEEDED:
+        original_fn = _WORKFLOW_FUNCTION_REGISTRY[workflow.function_name]
+        return deserialize_result(original_fn, workflow.result)
+    elif workflow.status == WorkflowStatus.FAILED:
+        assert workflow.error
+        raise try_restore_exception(workflow.error)
+    assert False, "May only be called on finished workflows"
+
+
+async def execute(workflow: Workflow):
+    logger.debug(
+        "executing workflow",
+        workflow_id=workflow.id,
+        function_name=workflow.function_name,
+    )
+    session = get_session()
+    original_fn = _WORKFLOW_FUNCTION_REGISTRY[workflow.function_name]
+    serialized_args = workflow.args or []
+    serialized_kwargs = workflow.kwargs or {}
+    args, kwargs = deserialize_args(original_fn, serialized_args, serialized_kwargs)
+
+    # Cache the workflow id here to avoid "Was IO attempted in an unexpected
+    # place?" SQLAlchemy errors when acessing expired attributes in an
+    # AsyncSession.
+    #
+    # Even though we unconditionally set expire_on_commit=True on
+    # PlanarSession, this is still necessary because SQLAlchemy will expire all
+    # attributes of the workflow object on a session rollback. More details:
+    # https://github.com/sqlalchemy/sqlalchemy/discussions/8282#discussioncomment-3213994
+    workflow_id = workflow.id
+    set_context(
+        ExecutionContext(
+            workflow=workflow,
+            workflow_id=workflow_id,
+        )
+    )
+    logger.debug("execution context set for workflow", workflow_id=workflow_id)
+
+    try:
+        stepper_result = await stepper(original_fn(*args, **kwargs))
+        logger.debug(
+            "stepper result for workflow",
+            workflow_id=workflow_id,
+            has_suspend=stepper_result.suspend is not None,
+            has_value=stepper_result.value is not None,
+        )
+        if stepper_result.suspend:
+            if stepper_result.suspend.exception:
+                logger.error(
+                    "workflow suspended due to an exception from stepper",
+                    workflow_id=workflow_id,
+                    exception=str(stepper_result.suspend.exception),
+                )
+                raise stepper_result.suspend.exception
+            workflow.status = WorkflowStatus.PENDING
+            logger.info(
+                "workflow suspended",
+                workflow_id=workflow_id,
+                wakeup_at=workflow.wakeup_at,
+                event_key=workflow.waiting_for_event,
+            )
+            return stepper_result.suspend
+        workflow.status = WorkflowStatus.SUCCEEDED
+        workflow.result = serialize_result(original_fn, stepper_result.value)
+        logger.info(
+            "workflow succeeded", workflow_id=workflow_id, result=workflow.result
+        )
+        return stepper_result.value
+    except Exception as e:
+        logger.exception("exception during workflow execution", workflow_id=workflow_id)
+        workflow.status = WorkflowStatus.FAILED
+        workflow.error = {
+            "type": type(e).__name__,
+            "message": str(e),
+            "traceback": str(traceback.format_exc()),
+        }
+        raise e
+    finally:
+        delete_context()
+        logger.debug("execution context deleted for workflow", workflow_id=workflow_id)
+        await session.commit()
+        # notify after committing to the db
+        if workflow.status == WorkflowStatus.SUCCEEDED:
+            notifications.workflow_succeeded(workflow)
+        elif workflow.status == WorkflowStatus.FAILED:
+            notifications.workflow_failed(workflow)
+        else:
+            notifications.workflow_suspended(workflow)
+
+
+async def lock_and_execute(
+    workflow: Workflow,
+    lock_duration: timedelta = _DEFAULT_LOCK_DURATION,
+):
+    logger.debug("attempting to lock and execute workflow", workflow_id=workflow.id)
+    session = get_session()
+
+    async with lock_workflow(workflow, lock_duration):
+        logger.debug("lock acquired for workflow", workflow_id=workflow.id)
+        async with session.begin_read():
+            await session.refresh(workflow)
+
+        if workflow.status != WorkflowStatus.PENDING:
+            logger.info(
+                "workflow is not pending, returning existing result",
+                workflow_id=workflow.id,
+                status=workflow.status,
+            )
+            return workflow_result(workflow)
+
+        notifications.workflow_resumed(workflow)
+        logger.info("workflow resumed", workflow_id=workflow.id)
+
+        # Execute until the next suspend or completion
+        result = await execute(workflow)
+        await trace("return", workflow_id=workflow.id)
+        logger.debug(
+            "execution finished for workflow",
+            workflow_id=workflow.id,
+            result_type=type(result),
+        )
+        return result

planar/workflows/lock.py

@@ -0,0 +1,229 @@
+import asyncio
+import traceback
+from contextlib import asynccontextmanager
+from datetime import timedelta
+
+from sqlalchemy.exc import IntegrityError, OperationalError
+from sqlalchemy.orm.exc import StaleDataError, UnmappedInstanceError
+
+from planar.db import new_session
+from planar.logging import get_logger
+from planar.session import engine_var, get_session
+from planar.utils import utc_now
+from planar.workflows.exceptions import LockResourceFailed
+from planar.workflows.models import LockedResource, Workflow, workflow_exec_lock_key
+from planar.workflows.tracing import trace
+
+_DEFAULT_RETRY_INTERVAL_SECONDS = 5
+_DEFAULT_RETRY_COUNT = 30  # with default retry interval, wait 2.5 minutes
+
+logger = get_logger(__name__)
+
+
+async def lock_heartbeat(
+    lock: LockedResource,
+    lock_duration: timedelta,
+    lock_owner: asyncio.Task,
+):
+    lock_key = lock.lock_key
+    engine = engine_var.get()
+    heartbeat_interval = lock_duration / 3
+    await trace(
+        "enter",
+        lock_key=lock_key,
+        lock_duration=lock_duration.total_seconds(),
+        lock_heartbeat_interval=heartbeat_interval.total_seconds(),
+    )
+    logger.debug(
+        "lock heartbeat started",
+        lock_key=lock_key,
+        duration_seconds=lock_duration.total_seconds(),
+        interval_seconds=heartbeat_interval.total_seconds(),
+    )
+    async with new_session(engine) as session:
+        try:
+            async with session.begin():
+                session.add(lock)
+            while True:
+                try:
+                    sleep_seconds = heartbeat_interval.total_seconds()
+                    await trace("sleep", lock_key=lock_key, sleep_seconds=sleep_seconds)
+                    await asyncio.sleep(sleep_seconds)
+                    # Renew the lock
+                    async with session.begin():
+                        lock.lock_until = utc_now() + lock_duration
+                        await trace(
+                            "renew-lock",
+                            lock_key=lock_key,
+                            lock_until=lock.lock_until,
+                        )
+                        logger.debug(
+                            "lock renewed",
+                            lock_key=lock_key,
+                            lock_until=lock.lock_until,
+                        )
+                    await trace("commit", lock_key=lock_key)
+                except StaleDataError:
+                    logger.exception(
+                        "stale data error in lock heartbeat, cancelling owner task",
+                        lock_key=lock_key,
+                    )
+                    await trace("stale-data-error", lock_key=lock_key)
+                    # This would happen if the process paused for too long and some
+                    # other worker acquired the lock. Some possible causes:
+                    #
+                    # - "stop the world" GC that took too long
+                    # - Some network call took too long. For example, imagine that
+                    #   `await session.commit()` took minutes to return after the
+                    #   changes were actually committed
+                    #
+                    # No matter what the cause was (and however unlikely it is to
+                    # happen in practice), it is no longer safe to continue
+                    # processing this resource. Kill the main task.
+                    lock_owner.cancel()
+                    break
+                except asyncio.CancelledError:
+                    logger.debug("lock heartbeat cancelled by owner", lock_key=lock_key)
+                    # Cancelled by the lock owner
+                    break
+                except Exception:
+                    logger.exception(
+                        "exception in lock heartbeat, cancelling owner task",
+                        lock_key=lock_key,
+                    )
+
+                    await trace(
+                        "exception",
+                        lock_key=lock_key,
+                        traceback=traceback.format_exc(),
+                    )
+                    # similarly to the `StaleDataError, kill the main task
+                    lock_owner.cancel()
+                    break
+        finally:
+            # ensure the lock object is detached from the session
+            try:
+                session.expunge(lock)
+            except UnmappedInstanceError:
+                # it is possible that the lock was not added to the session yet
+                pass
+    await trace("exit", lock_key=lock_key)
+    logger.debug("lock heartbeat stopped", lock_key=lock_key)
+
+
+@asynccontextmanager
+async def lock_resource(
+    lock_key: str,
+    lock_duration: timedelta,
+    retry_count: int = _DEFAULT_RETRY_COUNT,
+    retry_interval_seconds: int = _DEFAULT_RETRY_INTERVAL_SECONDS,
+):
+    assert retry_count >= 0
+    await trace("enter", lock_key=lock_key)
+    logger.debug(
+        "attempting to lock resource",
+        lock_key=lock_key,
+        duration_seconds=lock_duration.total_seconds(),
+        retries=retry_count,
+    )
+    session = get_session()
+
+    lock = None
+    for remaining in range(retry_count, -1, -1):
+        try:
+            async with session.begin():
+                lock = LockedResource(
+                    lock_key=lock_key,
+                    lock_until=utc_now() + lock_duration,
+                )
+                session.add(lock)
+                await trace("add-locked-resource", lock_key=lock_key)
+            await trace("commit")
+            logger.info(
+                "resource locked", lock_key=lock_key, lock_until=lock.lock_until
+            )
+            # This LockedResource instance will be passed to the heartbeat task
+            # which will use a different session to manage it.
+            session.expunge(lock)
+            break
+        except (OperationalError, IntegrityError) as e:
+            logger.exception(
+                "failed to acquire lock for resource on attempt",
+                lock_key=lock_key,
+                attempt=retry_count - remaining + 1,
+            )
+            await trace("add-locked-resource-error", lock_key=lock_key, error=str(e))
+        finally:
+            # ensure the session is ready for re-use after an exception
+            await session.rollback()
+        lock = None
+        await trace("retry", lock_key=lock_key, remaining_retry_count=remaining)
+        if remaining > 0:
+            logger.debug(
+                "retrying lock for resource",
+                lock_key=lock_key,
+                retry_interval_seconds=retry_interval_seconds,
+                retries_left=remaining,
+            )
+            await asyncio.sleep(retry_interval_seconds)
+
+    if lock is None:
+        logger.warning(
+            "failed to lock resource after all attempts",
+            lock_key=lock_key,
+            attempts=retry_count + 1,
+        )
+        await trace("no-remaining-retries", lock_key=lock_key)
+        raise LockResourceFailed(f'Failed to lock resource "{lock_key}"')
+
+    # Start the heartbeat to renew the lock periodically
+    await trace("start-heartbeat", lock_key=lock_key)
+    current_task = asyncio.current_task()
+    assert current_task
+    heartbeat_task = asyncio.create_task(
+        lock_heartbeat(lock, lock_duration, current_task)
+    )
+
+    try:
+        await trace("yield", lock_key=lock_key)
+        yield
+    finally:
+        # Stop the heartbeat
+        await trace("cancel-heartbeat", lock_key=lock_key)
+        heartbeat_task.cancel()
+        try:
+            await heartbeat_task
+        except asyncio.CancelledError:
+            pass
+        if session.in_transaction():
+            # Session should not be in a transaction here. This is probably the
+            # result of an exception which was not handled by calling rollback.
+            # We'll do it here because we need to release the lock, which has
+            # to be done in another transaction, but leave a warning in the
+            # logs
+            logger.warning(
+                "session is still in transaction, rolling back", lock_key=lock_key
+            )
+            await session.rollback()
+        async with session.begin():
+            session.add(lock)
+            await session.delete(lock)
+            await trace("release-lock", lock_key=lock_key)
+        await trace("exit", lock_key=lock_key)
+
+
+@asynccontextmanager
+async def lock_workflow(
+    workflow: Workflow,
+    lock_duration: timedelta,
+    retry_count: int = _DEFAULT_RETRY_COUNT,
+    retry_interval_seconds: int = _DEFAULT_RETRY_INTERVAL_SECONDS,
+):
+    lock_key = workflow_exec_lock_key(workflow.id)
+    async with lock_resource(
+        lock_key,
+        lock_duration,
+        retry_count=retry_count,
+        retry_interval_seconds=retry_interval_seconds,
+    ):
+        yield

planar/workflows/misc.py

@@ -0,0 +1,5 @@
+from typing import Any, Callable
+
+
+def func_full_name(func: Callable[..., Any]) -> str:
+    return f"{func.__module__}.{func.__qualname__}"