stabilize 0.9.2__py3-none-any.whl

stabilize/models/workflow.py

@@ -0,0 +1,317 @@
+"""
+Workflow model.
+
+A pipeline execution represents a running instance of a pipeline, containing
+all stages and their runtime state. The execution tracks:
+- Overall status
+- All stages (including synthetic stages)
+- Trigger information
+- Timing data
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+from stabilize.models.stage import StageExecution
+from stabilize.models.status import WorkflowStatus
+
+
+def _generate_execution_id() -> str:
+    """Generate a unique execution ID using ULID."""
+    import ulid
+
+    return str(ulid.new())
+
+
+class WorkflowType(Enum):
+    """
+    Type of execution.
+
+    PIPELINE: A full pipeline execution
+    ORCHESTRATION: An ad-hoc orchestration (single stage)
+    """
+
+    PIPELINE = "PIPELINE"
+    ORCHESTRATION = "ORCHESTRATION"
+
+
+@dataclass
+class Trigger:
+    """
+    Trigger information for a pipeline execution.
+
+    Contains details about what triggered the pipeline (manual, webhook, cron, etc.)
+    and any parameters passed to the execution.
+    """
+
+    type: str = "manual"
+    user: str = "anonymous"
+    parameters: dict[str, Any] = field(default_factory=dict)
+    artifacts: list[dict[str, Any]] = field(default_factory=list)
+    payload: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert trigger to dictionary for storage."""
+        return {
+            "type": self.type,
+            "user": self.user,
+            "parameters": self.parameters,
+            "artifacts": self.artifacts,
+            "payload": self.payload,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Trigger:
+        """Create trigger from dictionary."""
+        return cls(
+            type=data.get("type", "manual"),
+            user=data.get("user", "anonymous"),
+            parameters=data.get("parameters", {}),
+            artifacts=data.get("artifacts", []),
+            payload=data.get("payload", {}),
+        )
+
+
+@dataclass
+class PausedDetails:
+    """
+    Details about a paused execution.
+    """
+
+    paused_by: str = ""
+    pause_time: int | None = None
+    resume_time: int | None = None
+    paused_ms: int = 0
+
+    @property
+    def is_paused(self) -> bool:
+        """Check if currently paused."""
+        return self.pause_time is not None and self.resume_time is None
+
+
+@dataclass
+class Workflow:
+    """
+    Represents a pipeline execution.
+
+    This is the top-level container for all execution state. It holds all stages
+    and tracks the overall execution status.
+
+    Attributes:
+        id: Unique identifier (ULID)
+        type: PIPELINE or ORCHESTRATION
+        application: Application name this pipeline belongs to
+        name: Pipeline name
+        status: Current execution status
+        stages: All stages in this execution (including synthetic)
+        trigger: Trigger information
+        start_time: Epoch milliseconds when execution started
+        end_time: Epoch milliseconds when execution completed
+        start_time_expiry: If not started by this time, cancel
+        is_canceled: Whether execution has been canceled
+        canceled_by: User who canceled the execution
+        cancellation_reason: Reason for cancellation
+        paused: Pause details if execution is paused
+        pipeline_config_id: ID of the pipeline configuration
+        is_limit_concurrent: Whether to limit concurrent executions
+        max_concurrent_executions: Max concurrent executions allowed
+        keep_waiting_pipelines: Keep queued pipelines on cancel
+        origin: Origin of the execution (e.g., "api", "deck")
+    """
+
+    id: str = field(default_factory=_generate_execution_id)
+    type: WorkflowType = WorkflowType.PIPELINE
+    application: str = ""
+    name: str = ""
+    status: WorkflowStatus = WorkflowStatus.NOT_STARTED
+    stages: list[StageExecution] = field(default_factory=list)
+    trigger: Trigger = field(default_factory=Trigger)
+    start_time: int | None = None
+    end_time: int | None = None
+    start_time_expiry: int | None = None
+    is_canceled: bool = False
+    canceled_by: str | None = None
+    cancellation_reason: str | None = None
+    paused: PausedDetails | None = None
+    pipeline_config_id: str | None = None
+    is_limit_concurrent: bool = False
+    max_concurrent_executions: int = 0
+    keep_waiting_pipelines: bool = False
+    origin: str = "unknown"
+
+    def __post_init__(self) -> None:
+        """Set execution reference on all stages after construction."""
+        for stage in self.stages:
+            stage._execution = self
+
+    def add_stage(self, stage: StageExecution) -> None:
+        """Add a stage to this execution."""
+        stage._execution = self
+        self.stages.append(stage)
+
+    def remove_stage(self, stage_id: str) -> None:
+        """Remove a stage from this execution."""
+        self.stages = [s for s in self.stages if s.id != stage_id]
+
+    def stage_by_id(self, stage_id: str) -> StageExecution:
+        """
+        Get a stage by its ID.
+
+        Raises:
+            ValueError: If stage not found
+        """
+        for stage in self.stages:
+            if stage.id == stage_id:
+                return stage
+        raise ValueError(f"Stage {stage_id} not found")
+
+    def stage_by_ref_id(self, ref_id: str) -> StageExecution | None:
+        """Get a stage by its reference ID."""
+        for stage in self.stages:
+            if stage.ref_id == ref_id:
+                return stage
+        return None
+
+    # ========== Stage Queries ==========
+
+    def initial_stages(self) -> list[StageExecution]:
+        """
+        Get all initial stages (no dependencies, not synthetic).
+
+        These are the stages that can start immediately when execution begins.
+        """
+        return [stage for stage in self.stages if stage.is_initial() and not stage.is_synthetic()]
+
+    def top_level_stages(self) -> list[StageExecution]:
+        """Get all top-level stages (not synthetic)."""
+        return [stage for stage in self.stages if not stage.is_synthetic()]
+
+    # ========== Context Aggregation ==========
+
+    def get_context(self) -> dict[str, Any]:
+        """
+        Get aggregated context from all stages.
+
+        Returns merged outputs from all stages in topological order.
+        Collections are concatenated, latest value wins for non-collections.
+        """
+        from stabilize.dag.topological import topological_sort
+
+        result: dict[str, Any] = {}
+
+        for stage in topological_sort(self.stages):
+            for key, value in stage.outputs.items():
+                if key in result and isinstance(result[key], list) and isinstance(value, list):
+                    # Concatenate lists, avoiding duplicates
+                    existing = result[key]
+                    for item in value:
+                        if item not in existing:
+                            existing.append(item)
+                else:
+                    result[key] = value
+
+        return result
+
+    # ========== Status Methods ==========
+
+    def update_status(self, status: WorkflowStatus) -> None:
+        """Update the execution status."""
+        self.status = status
+
+    def cancel(self, user: str, reason: str) -> None:
+        """Mark this execution as canceled."""
+        self.is_canceled = True
+        self.canceled_by = user
+        self.cancellation_reason = reason
+
+    def pause(self, user: str) -> None:
+        """Pause this execution."""
+        import time
+
+        self.paused = PausedDetails(
+            paused_by=user,
+            pause_time=int(time.time() * 1000),
+        )
+        self.status = WorkflowStatus.PAUSED
+
+    def resume(self) -> None:
+        """Resume this execution."""
+        import time
+
+        if self.paused and self.paused.pause_time:
+            self.paused.resume_time = int(time.time() * 1000)
+            self.paused.paused_ms = self.paused.resume_time - self.paused.pause_time
+        self.status = WorkflowStatus.RUNNING
+
+    def paused_duration_relative_to(self, instant_ms: int) -> int:
+        """
+        Get paused duration relative to a given instant.
+
+        Returns 0 if not paused or pause was before the instant.
+        """
+        if self.paused and self.paused.pause_time:
+            if self.paused.pause_time > instant_ms:
+                return self.paused.paused_ms
+        return 0
+
+    # ========== Factory Methods ==========
+
+    @classmethod
+    def create(
+        cls,
+        application: str,
+        name: str,
+        stages: list[StageExecution],
+        trigger: Trigger | None = None,
+        pipeline_config_id: str | None = None,
+    ) -> Workflow:
+        """
+        Factory method to create a new pipeline execution.
+
+        Args:
+            application: Application name
+            name: Pipeline name
+            stages: List of stages
+            trigger: Optional trigger info
+            pipeline_config_id: Optional config ID
+
+        Returns:
+            A new Workflow instance
+        """
+        execution = cls(
+            application=application,
+            name=name,
+            stages=stages,
+            trigger=trigger or Trigger(),
+            pipeline_config_id=pipeline_config_id,
+        )
+        return execution
+
+    @classmethod
+    def create_orchestration(
+        cls,
+        application: str,
+        name: str,
+        stages: list[StageExecution],
+    ) -> Workflow:
+        """
+        Factory method to create an orchestration (ad-hoc execution).
+
+        Args:
+            application: Application name
+            name: Orchestration name
+            stages: List of stages
+
+        Returns:
+            A new Workflow with type ORCHESTRATION
+        """
+        execution = cls(
+            type=WorkflowType.ORCHESTRATION,
+            application=application,
+            name=name,
+            stages=stages,
+        )
+        return execution

stabilize/orchestrator.py

@@ -0,0 +1,113 @@
+"""
+Orchestrator - starts and manages pipeline executions.
+
+This module provides the main entry point for running pipelines.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from stabilize.queue.messages import (
+    CancelWorkflow,
+    RestartStage,
+    ResumeStage,
+    StartWorkflow,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.workflow import Workflow
+    from stabilize.queue.queue import Queue
+
+
+class Orchestrator:
+    """
+    Runner for pipeline executions.
+
+    Provides methods to start, cancel, restart, and resume executions
+    by pushing appropriate messages to the queue.
+    """
+
+    def __init__(self, queue: Queue) -> None:
+        """
+        Initialize the runner.
+
+        Args:
+            queue: The message queue
+        """
+        self.queue = queue
+
+    def start(self, execution: Workflow) -> None:
+        """
+        Start a pipeline execution.
+
+        Args:
+            execution: The execution to start
+        """
+        self.queue.push(
+            StartWorkflow(
+                execution_type=execution.type.value,
+                execution_id=execution.id,
+            )
+        )
+
+    def cancel(
+        self,
+        execution: Workflow,
+        user: str,
+        reason: str,
+    ) -> None:
+        """
+        Cancel a running execution.
+
+        Args:
+            execution: The execution to cancel
+            user: Who is canceling
+            reason: Why it's being canceled
+        """
+        self.queue.push(
+            CancelWorkflow(
+                execution_type=execution.type.value,
+                execution_id=execution.id,
+                user=user,
+                reason=reason,
+            )
+        )
+
+    def restart(
+        self,
+        execution: Workflow,
+        stage_id: str,
+    ) -> None:
+        """
+        Restart a stage in an execution.
+
+        Args:
+            execution: The execution
+            stage_id: The stage to restart
+        """
+        self.queue.push(
+            RestartStage(
+                execution_type=execution.type.value,
+                execution_id=execution.id,
+                stage_id=stage_id,
+            )
+        )
+
+    def unpause(self, execution: Workflow) -> None:
+        """
+        Resume a paused execution.
+
+        Args:
+            execution: The execution to resume
+        """
+        # Resume all paused stages
+        for stage in execution.stages:
+            if stage.status.name == "PAUSED":
+                self.queue.push(
+                    ResumeStage(
+                        execution_type=execution.type.value,
+                        execution_id=execution.id,
+                        stage_id=stage.id,
+                    )
+                )

stabilize/persistence/__init__.py

@@ -0,0 +1,28 @@
+"""Persistence layer for pipeline execution."""
+
+from stabilize.persistence.factory import (
+    create_queue,
+    create_repository,
+    detect_backend,
+)
+from stabilize.persistence.memory import InMemoryWorkflowStore
+from stabilize.persistence.postgres import PostgresWorkflowStore
+from stabilize.persistence.sqlite import SqliteWorkflowStore
+from stabilize.persistence.store import (
+    WorkflowNotFoundError,
+    WorkflowStore,
+)
+
+__all__ = [
+    # Abstract interface
+    "WorkflowStore",
+    "WorkflowNotFoundError",
+    # Implementations
+    "PostgresWorkflowStore",
+    "SqliteWorkflowStore",
+    "InMemoryWorkflowStore",
+    # Factory functions
+    "create_repository",
+    "create_queue",
+    "detect_backend",
+]

stabilize/persistence/connection.py

@@ -0,0 +1,185 @@
+"""
+Singleton connection manager for database connections.
+
+Provides centralized connection pooling for PostgreSQL and thread-local
+connections for SQLite, ensuring efficient resource usage across all
+repository and queue instances.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import threading
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from psycopg_pool import ConnectionPool
+
+
+class SingletonMeta(type):
+    """
+    Thread-safe metaclass for singleton pattern.
+
+    Ensures only one instance of a class exists, even when accessed
+    from multiple threads simultaneously.
+    """
+
+    _instances: dict[type, Any] = {}
+    _lock: threading.Lock = threading.Lock()
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        if cls not in cls._instances:
+            with cls._lock:
+                # Double-check locking pattern
+                if cls not in cls._instances:
+                    instance = super().__call__(*args, **kwargs)
+                    cls._instances[cls] = instance
+        return cls._instances[cls]
+
+    @classmethod
+    def reset(mcs, cls: type) -> None:
+        """Reset singleton instance (for testing)."""
+        with mcs._lock:
+            if cls in mcs._instances:
+                instance = mcs._instances.pop(cls)
+                if hasattr(instance, "close_all"):
+                    instance.close_all()
+
+
+class ConnectionManager(metaclass=SingletonMeta):
+    """
+    Singleton connection manager for all database connections.
+
+    Manages:
+    - PostgreSQL connection pools (one pool per connection string)
+    - SQLite thread-local connections (one connection per thread per db path)
+
+    Usage:
+        manager = ConnectionManager()
+        pool = manager.get_postgres_pool("postgresql://...")
+        conn = manager.get_sqlite_connection("sqlite:///./db.sqlite")
+    """
+
+    def __init__(self) -> None:
+        self._postgres_pools: dict[str, ConnectionPool] = {}
+        self._postgres_lock = threading.Lock()
+
+        self._sqlite_local = threading.local()
+        self._sqlite_lock = threading.Lock()
+        self._sqlite_paths: set[str] = set()
+
+    def get_postgres_pool(
+        self,
+        connection_string: str,
+        min_size: int = 5,
+        max_size: int = 15,
+    ) -> ConnectionPool:
+        """
+        Get or create a PostgreSQL connection pool.
+
+        Args:
+            connection_string: PostgreSQL connection string
+            min_size: Minimum pool size
+            max_size: Maximum pool size
+
+        Returns:
+            Shared ConnectionPool instance for this connection string
+        """
+        if connection_string not in self._postgres_pools:
+            with self._postgres_lock:
+                if connection_string not in self._postgres_pools:
+                    from psycopg.rows import dict_row
+                    from psycopg_pool import ConnectionPool
+
+                    pool = ConnectionPool(
+                        connection_string,
+                        min_size=min_size,
+                        max_size=max_size,
+                        open=True,
+                        kwargs={"row_factory": dict_row},
+                    )
+                    self._postgres_pools[connection_string] = pool
+        return self._postgres_pools[connection_string]
+
+    def get_sqlite_connection(self, connection_string: str) -> sqlite3.Connection:
+        """
+        Get or create a thread-local SQLite connection.
+
+        Args:
+            connection_string: SQLite connection string
+
+        Returns:
+            Thread-local Connection instance for this database
+        """
+        db_path = self._parse_sqlite_path(connection_string)
+
+        # Track paths for cleanup
+        with self._sqlite_lock:
+            self._sqlite_paths.add(db_path)
+
+        # Get thread-local storage
+        if not hasattr(self._sqlite_local, "connections"):
+            self._sqlite_local.connections = {}
+
+        connections: dict[str, sqlite3.Connection] = self._sqlite_local.connections
+
+        if db_path not in connections or connections[db_path] is None:
+            conn = sqlite3.connect(
+                db_path,
+                timeout=30,
+                check_same_thread=False,
+            )
+            conn.row_factory = sqlite3.Row
+            conn.execute("PRAGMA foreign_keys = ON")
+            if db_path != ":memory:":
+                conn.execute("PRAGMA journal_mode = WAL")
+                conn.execute("PRAGMA busy_timeout = 30000")
+            connections[db_path] = conn
+
+        return connections[db_path]
+
+    def _parse_sqlite_path(self, connection_string: str) -> str:
+        """Parse SQLite connection string to extract database path."""
+        if connection_string.startswith("sqlite:///"):
+            return connection_string[10:]
+        elif connection_string.startswith("sqlite://"):
+            return connection_string[9:]
+        return connection_string
+
+    def close_postgres_pool(self, connection_string: str) -> None:
+        """Close a specific PostgreSQL pool."""
+        with self._postgres_lock:
+            if connection_string in self._postgres_pools:
+                pool = self._postgres_pools.pop(connection_string)
+                pool.close()
+
+    def close_sqlite_connection(self, connection_string: str) -> None:
+        """Close SQLite connection for current thread."""
+        db_path = self._parse_sqlite_path(connection_string)
+        if hasattr(self._sqlite_local, "connections"):
+            connections: dict[str, sqlite3.Connection | None] = self._sqlite_local.connections
+            conn = connections.get(db_path)
+            if conn is not None:
+                conn.close()
+                connections[db_path] = None
+
+    def close_all(self) -> None:
+        """Close all connections (for shutdown/testing)."""
+        # Close all PostgreSQL pools
+        with self._postgres_lock:
+            for pool in self._postgres_pools.values():
+                pool.close()
+            self._postgres_pools.clear()
+
+        # Close SQLite connections for current thread
+        if hasattr(self._sqlite_local, "connections"):
+            connections: dict[str, sqlite3.Connection] = self._sqlite_local.connections
+            for conn in connections.values():
+                if conn is not None:
+                    conn.close()
+            connections.clear()
+
+
+def get_connection_manager() -> ConnectionManager:
+    """Get the singleton ConnectionManager instance."""
+    return ConnectionManager()