stabilize 0.9.2__py3-none-any.whl

stabilize/handlers/start_workflow.py

@@ -0,0 +1,136 @@
+"""
+StartWorkflowHandler - handles pipeline execution startup.
+
+This handler is triggered when a new pipeline execution is started.
+It finds initial stages (those with no dependencies) and queues them
+for execution.
+"""
+
+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 (
+    CancelWorkflow,
+    StartStage,
+    StartWorkflow,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.workflow import Workflow
+
+logger = logging.getLogger(__name__)
+
+
+class StartWorkflowHandler(StabilizeHandler[StartWorkflow]):
+    """
+    Handler for StartWorkflow messages.
+
+    When a pipeline execution starts:
+    1. Check if execution should be queued (concurrent limits)
+    2. Find initial stages (no dependencies)
+    3. Push StartStage for each initial stage
+    4. Mark execution as RUNNING
+    """
+
+    @property
+    def message_type(self) -> type[StartWorkflow]:
+        return StartWorkflow
+
+    def handle(self, message: StartWorkflow) -> None:
+        """Handle the StartWorkflow message."""
+
+        def on_execution(execution: Workflow) -> None:
+            # Check if already started or canceled
+            if execution.status != WorkflowStatus.NOT_STARTED:
+                logger.warning(
+                    "Execution %s already has status %s, ignoring StartWorkflow",
+                    execution.id,
+                    execution.status,
+                )
+                return
+
+            if execution.is_canceled:
+                logger.info("Execution %s was canceled before start", execution.id)
+                self._terminate(execution)
+                return
+
+            # Check if start time has expired
+            if self._is_after_start_time_expiry(execution):
+                logger.warning("Execution %s start time expired, canceling", execution.id)
+                self.queue.push(
+                    CancelWorkflow(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        user="system",
+                        reason="Could not begin execution before start time expiry",
+                    )
+                )
+                return
+
+            # TODO: Check if should queue (concurrent execution limits)
+            # if execution.should_queue():
+            #     self.pending_execution_service.enqueue(execution.pipeline_config_id, message)
+            #     return
+
+            self._start(execution, message)
+
+        self.with_execution(message, on_execution)
+
+    def _start(
+        self,
+        execution: Workflow,
+        message: StartWorkflow,
+    ) -> None:
+        """Start the execution."""
+        initial_stages = execution.initial_stages()
+
+        if not initial_stages:
+            logger.warning("No initial stages found for execution %s", execution.id)
+            execution.status = WorkflowStatus.TERMINAL
+            self.repository.update_status(execution)
+            # Publish ExecutionComplete event
+            return
+
+        # Mark as running
+        execution.status = WorkflowStatus.RUNNING
+        execution.start_time = self.current_time_millis()
+        self.repository.update_status(execution)
+
+        # Queue all initial stages
+        for stage in initial_stages:
+            logger.debug(
+                "Queuing initial stage %s (%s) for execution %s",
+                stage.name,
+                stage.id,
+                execution.id,
+            )
+            self.queue.push(
+                StartStage(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=stage.id,
+                )
+            )
+
+        logger.info(
+            "Started execution %s with %d initial stage(s)",
+            execution.id,
+            len(initial_stages),
+        )
+
+    def _terminate(self, execution: Workflow) -> None:
+        """Terminate a canceled execution."""
+        # Publish ExecutionComplete event
+        if execution.pipeline_config_id:
+            # Queue start waiting executions
+            pass
+
+    def _is_after_start_time_expiry(self, execution: Workflow) -> bool:
+        """Check if current time is past start time expiry."""
+        if execution.start_time_expiry is None:
+            return False
+        return self.current_time_millis() > execution.start_time_expiry

stabilize/launcher.py

@@ -0,0 +1,307 @@
+"""
+WorkflowLauncher - creates and starts pipeline executions.
+
+This module provides the high-level interface for launching pipelines.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from stabilize.models.stage import StageExecution
+from stabilize.models.workflow import (
+    Trigger,
+    Workflow,
+)
+from stabilize.stages.builder import StageDefinitionBuilderFactory
+
+if TYPE_CHECKING:
+    from stabilize.orchestrator import Orchestrator
+    from stabilize.persistence.store import WorkflowStore
+    from stabilize.tasks.registry import TaskRegistry
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowLauncher:
+    """
+    Launcher for pipeline executions.
+
+    Creates executions from configuration and starts them via the runner.
+    """
+
+    def __init__(
+        self,
+        repository: WorkflowStore,
+        runner: Orchestrator,
+        stage_builder_factory: StageDefinitionBuilderFactory | None = None,
+        task_registry: TaskRegistry | None = None,
+    ) -> None:
+        """
+        Initialize the launcher.
+
+        Args:
+            repository: The execution repository
+            runner: The execution runner
+            stage_builder_factory: Factory for stage builders
+            task_registry: Registry for task implementations
+        """
+        self.repository = repository
+        self.runner = runner
+        self.stage_builder_factory = stage_builder_factory or StageDefinitionBuilderFactory()
+        self.task_registry = task_registry
+
+    def start(
+        self,
+        pipeline_config: dict[str, Any],
+        trigger: dict[str, Any] | None = None,
+    ) -> Workflow:
+        """
+        Start a pipeline execution from configuration.
+
+        Args:
+            pipeline_config: Pipeline configuration dictionary
+            trigger: Optional trigger information
+
+        Returns:
+            The created execution
+        """
+        # Parse the execution
+        execution = self.parse_execution(pipeline_config, trigger)
+
+        # Store it
+        self.repository.store(execution)
+
+        # Start it
+        self.runner.start(execution)
+
+        logger.info(f"Launched execution {execution.id} for pipeline {execution.name}")
+
+        return execution
+
+    def parse_execution(
+        self,
+        config: dict[str, Any],
+        trigger_config: dict[str, Any] | None = None,
+    ) -> Workflow:
+        """
+        Parse a pipeline configuration into an execution.
+
+        Args:
+            config: Pipeline configuration
+            trigger_config: Optional trigger configuration
+
+        Returns:
+            A Workflow ready to run
+        """
+        # Parse trigger
+        trigger = Trigger()
+        if trigger_config:
+            trigger = Trigger(
+                type=trigger_config.get("type", "manual"),
+                user=trigger_config.get("user", "anonymous"),
+                parameters=trigger_config.get("parameters", {}),
+                artifacts=trigger_config.get("artifacts", []),
+                payload=trigger_config.get("payload", {}),
+            )
+
+        # Parse stages
+        stages = self._parse_stages(config.get("stages", []))
+
+        # Create execution
+        execution = Workflow.create(
+            application=config.get("application", "unknown"),
+            name=config.get("name", "Unnamed Pipeline"),
+            stages=stages,
+            trigger=trigger,
+            pipeline_config_id=config.get("id"),
+        )
+
+        # Set additional properties
+        execution.is_limit_concurrent = config.get("limitConcurrent", False)
+        execution.max_concurrent_executions = config.get("maxConcurrentExecutions", 0)
+        execution.keep_waiting_pipelines = config.get("keepWaitingPipelines", False)
+
+        return execution
+
+    def _parse_stages(
+        self,
+        stage_configs: list[dict[str, Any]],
+    ) -> list[StageExecution]:
+        """
+        Parse stage configurations into StageExecution objects.
+
+        Args:
+            stage_configs: List of stage configuration dictionaries
+
+        Returns:
+            List of StageExecution objects
+        """
+        stages = []
+
+        for config in stage_configs:
+            stage = self._parse_stage(config)
+            stages.append(stage)
+
+        return stages
+
+    def _parse_stage(self, config: dict[str, Any]) -> StageExecution:
+        """
+        Parse a single stage configuration.
+
+        Args:
+            config: Stage configuration dictionary
+
+        Returns:
+            A StageExecution object
+        """
+        # Get requisite stages
+        requisite_ids: set[str] = set()
+        if "requisiteStageRefIds" in config:
+            requisite_ids = set(config["requisiteStageRefIds"])
+
+        # Create context from config (excluding metadata fields)
+        context = {
+            k: v
+            for k, v in config.items()
+            if k
+            not in {
+                "id",
+                "refId",
+                "type",
+                "name",
+                "requisiteStageRefIds",
+                "parentStageId",
+                "syntheticStageOwner",
+            }
+        }
+
+        stage = StageExecution.create(
+            type=config.get("type", "unknown"),
+            name=config.get("name", config.get("type", "Unknown")),
+            ref_id=config.get("refId", config.get("id", "")),
+            context=context,
+            requisite_stage_ref_ids=requisite_ids,
+        )
+
+        # Build tasks using stage definition builder
+        builder = self.stage_builder_factory.get(stage.type)
+        stage.tasks = builder.build_tasks(stage)
+
+        # Mark first/last tasks
+        if stage.tasks:
+            stage.tasks[0].stage_start = True
+            stage.tasks[-1].stage_end = True
+
+        return stage
+
+    def create_orchestration(
+        self,
+        application: str,
+        name: str,
+        stages: list[dict[str, Any]],
+    ) -> Workflow:
+        """
+        Create an ad-hoc orchestration (single execution not from a pipeline).
+
+        Args:
+            application: Application name
+            name: Orchestration name
+            stages: Stage configurations
+
+        Returns:
+            The created execution
+        """
+        parsed_stages = self._parse_stages(stages)
+
+        execution = Workflow.create_orchestration(
+            application=application,
+            name=name,
+            stages=parsed_stages,
+        )
+
+        self.repository.store(execution)
+        self.runner.start(execution)
+
+        return execution
+
+
+def create_simple_pipeline(
+    name: str,
+    application: str,
+    stages: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """
+    Create a simple pipeline configuration.
+
+    Helper function for building pipeline configs programmatically.
+
+    Args:
+        name: Pipeline name
+        application: Application name
+        stages: List of stage configurations
+
+    Returns:
+        A pipeline configuration dictionary
+
+    Example:
+        config = create_simple_pipeline(
+            name="Deploy to Prod",
+            application="myapp",
+            stages=[
+                {"refId": "1", "type": "wait", "name": "Wait", "waitTime": 30},
+                {"refId": "2", "type": "deploy", "name": "Deploy",
+                 "requisiteStageRefIds": ["1"]},
+            ],
+        )
+    """
+    return {
+        "name": name,
+        "application": application,
+        "stages": stages,
+    }
+
+
+def create_stage_config(
+    ref_id: str,
+    stage_type: str,
+    name: str,
+    requisites: list[str] | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """
+    Create a stage configuration.
+
+    Helper function for building stage configs.
+
+    Args:
+        ref_id: Stage reference ID
+        stage_type: Stage type
+        name: Stage name
+        requisites: Prerequisite stage ref IDs
+        **kwargs: Additional stage context
+
+    Returns:
+        A stage configuration dictionary
+
+    Example:
+        stage = create_stage_config(
+            ref_id="1",
+            stage_type="wait",
+            name="Wait for approval",
+            waitTime=3600,
+        )
+    """
+    config: dict[str, Any] = {
+        "refId": ref_id,
+        "type": stage_type,
+        "name": name,
+    }
+
+    if requisites:
+        config["requisiteStageRefIds"] = requisites
+
+    config.update(kwargs)
+
+    return config

stabilize/migrations/01KDQ4N9QPJ6Q4MCV3V9GHWPV4_initial_schema.sql

@@ -0,0 +1,97 @@
+-- migration: initial_schema
+-- id: 01KDQ4N9QPJ6Q4MCV3V9GHWPV4
+
+-- migrate: up
+
+CREATE TABLE pipeline_executions (
+    id VARCHAR(26) PRIMARY KEY,
+    type VARCHAR(50) NOT NULL,
+    application VARCHAR(255) NOT NULL,
+    name VARCHAR(255),
+    status VARCHAR(50) NOT NULL,
+    start_time BIGINT,
+    end_time BIGINT,
+    start_time_expiry BIGINT,
+    trigger JSONB,
+    context JSONB,
+    is_canceled BOOLEAN DEFAULT FALSE,
+    canceled_by VARCHAR(255),
+    cancellation_reason TEXT,
+    paused JSONB,
+    pipeline_config_id VARCHAR(255),
+    is_limit_concurrent BOOLEAN DEFAULT FALSE,
+    max_concurrent_executions INT DEFAULT 0,
+    keep_waiting_pipelines BOOLEAN DEFAULT FALSE,
+    origin VARCHAR(50) DEFAULT 'unknown',
+    created_at TIMESTAMP DEFAULT NOW()
+);
+
+CREATE TABLE stage_executions (
+    id VARCHAR(26) PRIMARY KEY,
+    execution_id VARCHAR(26) NOT NULL REFERENCES pipeline_executions(id) ON DELETE CASCADE,
+    ref_id VARCHAR(50) NOT NULL,
+    type VARCHAR(100) NOT NULL,
+    name VARCHAR(255),
+    status VARCHAR(50) NOT NULL,
+    context JSONB,
+    outputs JSONB,
+    requisite_stage_ref_ids TEXT[],
+    parent_stage_id VARCHAR(26),
+    synthetic_stage_owner VARCHAR(20),
+    start_time BIGINT,
+    end_time BIGINT,
+    start_time_expiry BIGINT,
+    scheduled_time BIGINT,
+    UNIQUE(execution_id, ref_id)
+);
+
+CREATE TABLE task_executions (
+    id VARCHAR(26) PRIMARY KEY,
+    stage_id VARCHAR(26) NOT NULL REFERENCES stage_executions(id) ON DELETE CASCADE,
+    name VARCHAR(255) NOT NULL,
+    implementing_class VARCHAR(255) NOT NULL,
+    status VARCHAR(50) NOT NULL,
+    start_time BIGINT,
+    end_time BIGINT,
+    stage_start BOOLEAN DEFAULT FALSE,
+    stage_end BOOLEAN DEFAULT FALSE,
+    loop_start BOOLEAN DEFAULT FALSE,
+    loop_end BOOLEAN DEFAULT FALSE,
+    task_exception_details JSONB
+);
+
+CREATE TABLE queue_messages (
+    id SERIAL PRIMARY KEY,
+    message_id VARCHAR(36) NOT NULL UNIQUE,
+    message_type VARCHAR(100) NOT NULL,
+    payload JSONB NOT NULL,
+    deliver_at TIMESTAMP NOT NULL DEFAULT NOW(),
+    attempts INT DEFAULT 0,
+    max_attempts INT DEFAULT 10,
+    locked_until TIMESTAMP,
+    created_at TIMESTAMP DEFAULT NOW()
+);
+
+CREATE INDEX idx_executions_application ON pipeline_executions(application);
+CREATE INDEX idx_executions_config_id ON pipeline_executions(pipeline_config_id);
+CREATE INDEX idx_executions_status ON pipeline_executions(status);
+CREATE INDEX idx_stage_execution ON stage_executions(execution_id);
+CREATE INDEX idx_stage_parent ON stage_executions(parent_stage_id);
+CREATE INDEX idx_task_stage ON task_executions(stage_id);
+CREATE INDEX idx_queue_deliver ON queue_messages(deliver_at) WHERE attempts < 10;
+CREATE INDEX idx_queue_locked ON queue_messages(locked_until);
+
+-- migrate: down
+
+DROP INDEX IF EXISTS idx_queue_locked;
+DROP INDEX IF EXISTS idx_queue_deliver;
+DROP INDEX IF EXISTS idx_task_stage;
+DROP INDEX IF EXISTS idx_stage_parent;
+DROP INDEX IF EXISTS idx_stage_execution;
+DROP INDEX IF EXISTS idx_executions_status;
+DROP INDEX IF EXISTS idx_executions_config_id;
+DROP INDEX IF EXISTS idx_executions_application;
+DROP TABLE IF EXISTS queue_messages;
+DROP TABLE IF EXISTS task_executions;
+DROP TABLE IF EXISTS stage_executions;
+DROP TABLE IF EXISTS pipeline_executions;

stabilize/migrations/01KDRK3TXW4R2GERC1WBCQYJGG_rag_embeddings.sql

@@ -0,0 +1,25 @@
+-- migration: rag_embeddings
+-- id: 01KDRK3TXW4R2GERC1WBCQYJGG
+
+-- migrate: up
+
+CREATE TABLE rag_embeddings (
+    id SERIAL PRIMARY KEY,
+    doc_id VARCHAR(255) NOT NULL,
+    content TEXT NOT NULL,
+    embedding JSONB NOT NULL,
+    embedding_model VARCHAR(100) NOT NULL,
+    chunk_index INT DEFAULT 0,
+    created_at TIMESTAMP DEFAULT NOW(),
+    updated_at TIMESTAMP DEFAULT NOW(),
+    UNIQUE(doc_id, chunk_index, embedding_model)
+);
+
+CREATE INDEX idx_rag_embeddings_doc ON rag_embeddings(doc_id);
+CREATE INDEX idx_rag_embeddings_model ON rag_embeddings(embedding_model);
+
+-- migrate: down
+
+DROP INDEX IF EXISTS idx_rag_embeddings_model;
+DROP INDEX IF EXISTS idx_rag_embeddings_doc;
+DROP TABLE IF EXISTS rag_embeddings;

stabilize/migrations/__init__.py

@@ -0,0 +1 @@
+"""Stabilize database migrations for PostgreSQL."""

stabilize/models/__init__.py

@@ -0,0 +1,15 @@
+"""Core data models for pipeline execution."""
+
+from stabilize.models.stage import StageExecution, SyntheticStageOwner
+from stabilize.models.status import WorkflowStatus
+from stabilize.models.task import TaskExecution
+from stabilize.models.workflow import Workflow, WorkflowType
+
+__all__ = [
+    "WorkflowStatus",
+    "TaskExecution",
+    "StageExecution",
+    "SyntheticStageOwner",
+    "Workflow",
+    "WorkflowType",
+]