stabilize 0.9.2__py3-none-any.whl

stabilize/handlers/base.py

@@ -0,0 +1,226 @@
+"""
+Base message handler classes.
+
+This module provides the base classes for all message handlers in the
+pipeline execution engine.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from datetime import timedelta
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from stabilize.models.stage import StageExecution
+from stabilize.models.task import TaskExecution
+from stabilize.models.workflow import Workflow
+from stabilize.queue.messages import (
+    CompleteWorkflow,
+    ContinueParentStage,
+    InvalidStageId,
+    InvalidTaskId,
+    InvalidWorkflowId,
+    Message,
+    StageLevel,
+    StartStage,
+    TaskLevel,
+    WorkflowLevel,
+)
+
+if TYPE_CHECKING:
+    from stabilize.persistence.store import WorkflowStore
+    from stabilize.queue.queue import Queue
+
+logger = logging.getLogger(__name__)
+
+M = TypeVar("M", bound=Message)
+
+
+class MessageHandler(ABC, Generic[M]):
+    """
+    Base class for message handlers.
+
+    Each handler processes a specific type of message.
+    """
+
+    @property
+    @abstractmethod
+    def message_type(self) -> type[M]:
+        """Return the type of message this handler processes."""
+        pass
+
+    @abstractmethod
+    def handle(self, message: M) -> None:
+        """Handle a message."""
+        pass
+
+
+class StabilizeHandler(MessageHandler[M], ABC):
+    """
+    Base handler with common utilities.
+
+    Provides helper methods for retrieving executions, stages, and tasks,
+    as well as the startNext() implementation.
+    """
+
+    def __init__(
+        self,
+        queue: Queue,
+        repository: WorkflowStore,
+        retry_delay: timedelta = timedelta(seconds=15),
+    ) -> None:
+        self.queue = queue
+        self.repository = repository
+        self.retry_delay = retry_delay
+
+    # ========== Execution Retrieval ==========
+
+    def with_execution(
+        self,
+        message: WorkflowLevel,
+        block: Callable[[Workflow], None],
+    ) -> None:
+        """
+        Execute a block with the execution for a message.
+
+        Args:
+            message: Message containing execution ID
+            block: Function to call with the execution
+        """
+        try:
+            execution = self.repository.retrieve(message.execution_id)
+            block(execution)
+        except Exception as e:
+            logger.error("Failed to retrieve execution %s: %s", message.execution_id, e)
+            self.queue.push(
+                InvalidWorkflowId(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                )
+            )
+
+    def with_stage(
+        self,
+        message: StageLevel,
+        block: Callable[[StageExecution], None],
+    ) -> None:
+        """
+        Execute a block with the stage for a message.
+
+        Args:
+            message: Message containing stage ID
+            block: Function to call with the stage
+        """
+
+        def on_execution(execution: Workflow) -> None:
+            try:
+                stage = execution.stage_by_id(message.stage_id)
+                block(stage)
+            except ValueError:
+                logger.error("Stage not found: %s", message.stage_id)
+                self.queue.push(
+                    InvalidStageId(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                    )
+                )
+
+        self.with_execution(message, on_execution)
+
+    def with_task(
+        self,
+        message: TaskLevel,
+        block: Callable[[StageExecution, TaskExecution], None],
+    ) -> None:
+        """
+        Execute a block with the stage and task for a message.
+
+        Args:
+            message: Message containing task ID
+            block: Function to call with (stage, task)
+        """
+
+        def on_stage(stage: StageExecution) -> None:
+            task = self._find_task(stage, message.task_id)
+            if task is None:
+                logger.error("Task not found: %s", message.task_id)
+                self.queue.push(
+                    InvalidTaskId(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                        task_id=message.task_id,
+                    )
+                )
+            else:
+                block(stage, task)
+
+        self.with_stage(message, on_stage)
+
+    def _find_task(
+        self,
+        stage: StageExecution,
+        task_id: str,
+    ) -> TaskExecution | None:
+        """Find a task by ID in a stage."""
+        for task in stage.tasks:
+            if task.id == task_id:
+                return task
+        return None
+
+    # ========== Stage Navigation ==========
+
+    def start_next(self, stage: StageExecution) -> None:
+        """
+        Start the next stage(s) after a stage completes.
+
+        This is the critical method for DAG traversal:
+        1. Find downstream stages (those that depend on this stage)
+        2. Push StartStage for each downstream stage
+        3. If this is a synthetic stage, notify parent
+        4. If no downstream and not synthetic, complete execution
+        """
+        execution = stage.execution
+        downstream_stages = stage.downstream_stages()
+        phase = stage.synthetic_stage_owner
+
+        if downstream_stages:
+            # Start all downstream stages
+            for downstream in downstream_stages:
+                self.queue.push(
+                    StartStage(
+                        execution_type=execution.type.value,
+                        execution_id=execution.id,
+                        stage_id=downstream.id,
+                    )
+                )
+        elif phase is not None:
+            # Synthetic stage - notify parent
+            parent = stage.parent()
+            self.queue.ensure(
+                ContinueParentStage(
+                    execution_type=execution.type.value,
+                    execution_id=execution.id,
+                    stage_id=parent.id,
+                    phase=phase,
+                ),
+                timedelta(seconds=0),
+            )
+        else:
+            # Top-level stage with no downstream - complete execution
+            self.queue.push(
+                CompleteWorkflow(
+                    execution_type=execution.type.value,
+                    execution_id=execution.id,
+                )
+            )
+
+    # ========== Utility Methods ==========
+
+    def current_time_millis(self) -> int:
+        """Get current time in milliseconds."""
+        return int(time.time() * 1000)

stabilize/handlers/complete_stage.py

@@ -0,0 +1,209 @@
+"""
+CompleteStageHandler - handles stage completion.
+
+This is a critical handler that:
+1. Determines stage status from tasks and synthetic stages
+2. Plans and starts after stages
+3. Plans and starts on-failure stages
+4. Triggers downstream stages via startNext()
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.models.status import WorkflowStatus
+from stabilize.queue.messages import (
+    CancelStage,
+    CompleteStage,
+    CompleteWorkflow,
+    StartStage,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+logger = logging.getLogger(__name__)
+
+
+class CompleteStageHandler(StabilizeHandler[CompleteStage]):
+    """
+    Handler for CompleteStage messages.
+
+    Execution flow:
+    1. Check if stage already complete
+    2. Determine status from synthetic stages and tasks
+    3. If success: Plan and start after stages
+    4. If failure: Plan and start on-failure stages
+    5. Update stage status and end time
+    6. If status allows continuation: startNext()
+    7. Otherwise: CancelStage + CompleteWorkflow
+    """
+
+    @property
+    def message_type(self) -> type[CompleteStage]:
+        return CompleteStage
+
+    def handle(self, message: CompleteStage) -> None:
+        """Handle the CompleteStage message."""
+
+        def on_stage(stage: StageExecution) -> None:
+            # Check if already complete
+            if stage.status not in {WorkflowStatus.RUNNING, WorkflowStatus.NOT_STARTED}:
+                logger.debug(
+                    "Stage %s already has status %s, ignoring CompleteStage",
+                    stage.name,
+                    stage.status,
+                )
+                return
+
+            try:
+                # Determine status from tasks and synthetic stages
+                status = stage.determine_status()
+
+                # Handle after stages
+                if status.is_complete and not status.is_halt:
+                    after_stages = stage.first_after_stages()
+                    if not after_stages:
+                        self._plan_after_stages(stage)
+                        after_stages = stage.first_after_stages()
+
+                    not_started = [s for s in after_stages if s.status == WorkflowStatus.NOT_STARTED]
+                    if not_started:
+                        for s in not_started:
+                            self.queue.push(
+                                StartStage(
+                                    execution_type=message.execution_type,
+                                    execution_id=message.execution_id,
+                                    stage_id=s.id,
+                                )
+                            )
+                        return
+
+                    # If status is NOT_STARTED with no after stages, it's weird
+                    if status == WorkflowStatus.NOT_STARTED:
+                        logger.warning("Stage %s had no tasks or synthetic stages", stage.name)
+                        status = WorkflowStatus.SKIPPED
+
+                # Handle failure - plan on-failure stages
+                elif status.is_failure:
+                    has_on_failure = self._plan_on_failure_stages(stage)
+                    if has_on_failure:
+                        after_stages = stage.first_after_stages()
+                        for s in after_stages:
+                            self.queue.push(
+                                StartStage(
+                                    execution_type=message.execution_type,
+                                    execution_id=message.execution_id,
+                                    stage_id=s.id,
+                                )
+                            )
+                        return
+
+                # Update stage status
+                stage.status = status
+                stage.end_time = self.current_time_millis()
+                self.repository.store_stage(stage)
+
+                logger.info("Stage %s completed with status %s", stage.name, status)
+
+                # Handle FAILED_CONTINUE propagation to parent
+                if (
+                    status == WorkflowStatus.FAILED_CONTINUE
+                    and stage.synthetic_stage_owner is not None
+                    and not stage.allow_sibling_stages_to_continue_on_failure
+                    and stage.parent_stage_id is not None
+                ):
+                    # Propagate failure to parent
+                    self.queue.push(
+                        CompleteStage(
+                            execution_type=message.execution_type,
+                            execution_id=message.execution_id,
+                            stage_id=stage.parent_stage_id,
+                        )
+                    )
+                elif status in {
+                    WorkflowStatus.SUCCEEDED,
+                    WorkflowStatus.FAILED_CONTINUE,
+                    WorkflowStatus.SKIPPED,
+                }:
+                    # Continue to downstream stages
+                    self.start_next(stage)
+                else:
+                    # Failure - cancel and complete execution
+                    self.queue.push(
+                        CancelStage(
+                            execution_type=message.execution_type,
+                            execution_id=message.execution_id,
+                            stage_id=message.stage_id,
+                        )
+                    )
+                    if stage.synthetic_stage_owner is None or stage.parent_stage_id is None:
+                        self.queue.push(
+                            CompleteWorkflow(
+                                execution_type=message.execution_type,
+                                execution_id=message.execution_id,
+                            )
+                        )
+                    else:
+                        # Propagate to parent
+                        self.queue.push(
+                            CompleteStage(
+                                execution_type=message.execution_type,
+                                execution_id=message.execution_id,
+                                stage_id=stage.parent_stage_id,
+                            )
+                        )
+
+            except Exception as e:
+                logger.error(
+                    "Error completing stage %s: %s",
+                    stage.name,
+                    e,
+                    exc_info=True,
+                )
+                stage.context["exception"] = {
+                    "details": {"error": str(e)},
+                }
+                stage.status = WorkflowStatus.TERMINAL
+                stage.end_time = self.current_time_millis()
+                self.repository.store_stage(stage)
+
+                self.queue.push(
+                    CancelStage(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                    )
+                )
+                self.queue.push(
+                    CompleteWorkflow(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                    )
+                )
+
+        self.with_stage(message, on_stage)
+
+    def _plan_after_stages(self, stage: StageExecution) -> None:
+        """
+        Plan after stages using the stage definition builder.
+
+        TODO: Implement StageDefinitionBuilder integration
+        """
+        # For now, do nothing - after stages should be pre-defined
+        pass
+
+    def _plan_on_failure_stages(self, stage: StageExecution) -> bool:
+        """
+        Plan on-failure stages using the stage definition builder.
+
+        TODO: Implement StageDefinitionBuilder integration
+
+        Returns:
+            True if on-failure stages were added
+        """
+        # For now, return False - no on-failure stages
+        return False

stabilize/handlers/complete_task.py

@@ -0,0 +1,75 @@
+"""
+CompleteTaskHandler - handles task completion.
+
+This handler updates task status and triggers either the next task
+or stage completion.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.queue.messages import (
+    CompleteStage,
+    CompleteTask,
+    StartTask,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+    from stabilize.models.task import TaskExecution
+
+logger = logging.getLogger(__name__)
+
+
+class CompleteTaskHandler(StabilizeHandler[CompleteTask]):
+    """
+    Handler for CompleteTask messages.
+
+    Execution flow:
+    1. Update task status and end time
+    2. If there's a next task: Push StartTask
+    3. Otherwise: Push CompleteStage
+    """
+
+    @property
+    def message_type(self) -> type[CompleteTask]:
+        return CompleteTask
+
+    def handle(self, message: CompleteTask) -> None:
+        """Handle the CompleteTask message."""
+
+        def on_task(stage: StageExecution, task: TaskExecution) -> None:
+            # Update task status
+            task.status = message.status
+            task.end_time = self.current_time_millis()
+
+            # Save the stage
+            self.repository.store_stage(stage)
+
+            logger.debug("Task %s completed with status %s", task.name, message.status)
+
+            # Check for next task
+            next_task = stage.next_task(task)
+            if next_task is not None:
+                self.queue.push(
+                    StartTask(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                        task_id=next_task.id,
+                    )
+                )
+            else:
+                # No more tasks - complete stage
+                self.queue.push(
+                    CompleteStage(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                    )
+                )
+
+        self.with_task(message, on_task)

stabilize/handlers/complete_workflow.py

@@ -0,0 +1,150 @@
+"""
+CompleteWorkflowHandler - handles execution completion.
+
+This handler determines the final execution status based on all
+top-level stages and marks the execution as complete.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.models.status import CONTINUABLE_STATUSES, WorkflowStatus
+from stabilize.queue.messages import (
+    CancelStage,
+    CompleteWorkflow,
+    StartWaitingWorkflows,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+    from stabilize.models.workflow import Workflow
+
+logger = logging.getLogger(__name__)
+
+
+class CompleteWorkflowHandler(StabilizeHandler[CompleteWorkflow]):
+    """
+    Handler for CompleteWorkflow messages.
+
+    Execution flow:
+    1. Check if execution already complete
+    2. Determine final status from top-level stages
+    3. Update execution status
+    4. Cancel any running stages if failed
+    5. Start waiting executions if queue is enabled
+    """
+
+    @property
+    def message_type(self) -> type[CompleteWorkflow]:
+        return CompleteWorkflow
+
+    def handle(self, message: CompleteWorkflow) -> None:
+        """Handle the CompleteWorkflow message."""
+
+        def on_execution(execution: Workflow) -> None:
+            # Check if already complete
+            if execution.status.is_complete:
+                logger.debug(
+                    "Execution %s already complete with status %s",
+                    execution.id,
+                    execution.status,
+                )
+                return
+
+            # Determine final status
+            status = self._determine_final_status(execution, message)
+            if status is None:
+                # Not ready to complete - stages still running
+                return
+
+            # Update execution
+            execution.status = status
+            execution.end_time = self.current_time_millis()
+            self.repository.update_status(execution)
+
+            logger.info("Execution %s completed with status %s", execution.id, status)
+
+            # Cancel any running stages if not successful
+            if status != WorkflowStatus.SUCCEEDED:
+                running_stages = [s for s in execution.top_level_stages() if s.status == WorkflowStatus.RUNNING]
+                for stage in running_stages:
+                    self.queue.push(
+                        CancelStage(
+                            execution_type=message.execution_type,
+                            execution_id=message.execution_id,
+                            stage_id=stage.id,
+                        )
+                    )
+
+            # Start waiting executions if configured
+            if execution.status != WorkflowStatus.RUNNING and execution.pipeline_config_id:
+                self.queue.push(
+                    StartWaitingWorkflows(
+                        pipeline_config_id=execution.pipeline_config_id,
+                        purge_queue=not execution.keep_waiting_pipelines,
+                    )
+                )
+
+        self.with_execution(message, on_execution)
+
+    def _determine_final_status(
+        self,
+        execution: Workflow,
+        message: CompleteWorkflow,
+    ) -> WorkflowStatus | None:
+        """
+        Determine the final execution status.
+
+        Returns None if execution is not ready to complete.
+        """
+        stages = execution.top_level_stages()
+        statuses = [s.status for s in stages]
+
+        # All succeeded/skipped/failed_continue -> SUCCEEDED
+        if all(s in CONTINUABLE_STATUSES for s in statuses):
+            return WorkflowStatus.SUCCEEDED
+
+        # Any TERMINAL -> TERMINAL
+        if WorkflowStatus.TERMINAL in statuses:
+            return WorkflowStatus.TERMINAL
+
+        # Any CANCELED -> CANCELED
+        if WorkflowStatus.CANCELED in statuses:
+            return WorkflowStatus.CANCELED
+
+        # Any STOPPED and no other branches incomplete
+        if WorkflowStatus.STOPPED in statuses:
+            if not self._other_branches_incomplete(stages):
+                # Check for override
+                if self._should_override_success(execution):
+                    return WorkflowStatus.TERMINAL
+                return WorkflowStatus.SUCCEEDED
+
+        # Still running - re-queue
+        logger.debug(
+            "Re-queuing CompleteWorkflow for %s - stages not complete. Statuses: %s",
+            execution.id,
+            statuses,
+        )
+        self.queue.push(message, self.retry_delay)
+        return None
+
+    def _other_branches_incomplete(self, stages: list[StageExecution]) -> bool:
+        """Check if any other branches are incomplete."""
+        for stage in stages:
+            if stage.status == WorkflowStatus.RUNNING:
+                return True
+            if stage.status == WorkflowStatus.NOT_STARTED and stage.all_upstream_stages_complete():
+                return True
+        return False
+
+    def _should_override_success(self, execution: Workflow) -> bool:
+        """Check if success should be overridden to failure."""
+        for stage in execution.stages:
+            if stage.status == WorkflowStatus.STOPPED:
+                if stage.context.get("completeOtherBranchesThenFail"):
+                    return True
+        return False