python-saga-orchestrator 0.1.0__py3-none-any.whl

saga_orchestrator/core/orchestrator.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Any, Generic, TypeVar
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from ..domain.mixins import SagaStateMixin
+from ..domain.models import SagaDefinition, SagaSnapshot
+from .engine import SagaEngine
+from .repository import SagaRepository
+
+ModelT = TypeVar("ModelT", bound=SagaStateMixin)
+
+
+class SagaOrchestrator(Generic[ModelT]):
+    """Provide the public runtime API for saga execution."""
+
+    def __init__(
+        self,
+        *,
+        model_class: type[ModelT],
+        session_maker: async_sessionmaker[AsyncSession],
+        execution_lease: timedelta = timedelta(minutes=5),
+    ) -> None:
+        """Initialize the orchestrator facade."""
+        self._engine = SagaEngine(
+            model_class=model_class,
+            session_maker=session_maker,
+            execution_lease=execution_lease,
+        )
+
+    @property
+    def engine(self) -> SagaEngine[ModelT]:
+        """Return the engine used by the orchestrator."""
+        return self._engine
+
+    @property
+    def repository(self) -> SagaRepository[ModelT]:
+        """Return the repository used by the engine."""
+        return self._engine.repository
+
+    def register(self, name: str, saga_definition: SagaDefinition) -> None:
+        """Register a saga definition under a runtime name."""
+        self._engine.register(name, saga_definition)
+
+    async def start(
+        self,
+        *,
+        saga_name: str,
+        initial_data: BaseModel | dict[str, Any] | Any,
+        aggregation_id: str,
+        trace_id: str | None = None,
+    ) -> UUID:
+        """Create a new saga instance and start executing it."""
+        return await self._engine.start(
+            saga_name=saga_name,
+            initial_data=initial_data,
+            aggregation_id=aggregation_id,
+            trace_id=trace_id,
+        )
+
+    async def notify(
+        self, *, saga_id: UUID, token: UUID, event: Any | None = None
+    ) -> bool:
+        """Resume a suspended saga when the provided execution token matches."""
+        return await self._engine.notify(saga_id=saga_id, token=token, event=event)
+
+    async def run_due(self, *, limit: int = 100) -> int:
+        """Resume due running, suspended, and compensating sagas."""
+        return await self._engine.run_due(limit=limit)
+
+    async def get_snapshot(self, saga_id: UUID) -> SagaSnapshot:
+        """Return the snapshot view of one saga."""
+        return await self._engine.get_snapshot(saga_id)
+
+    async def resume(self, saga_id: UUID) -> None:
+        """Resume forward execution of one saga."""
+        await self._engine.resume(saga_id)

saga_orchestrator/core/repository.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Generic, TypeVar
+from uuid import UUID
+
+from sqlalchemy import Select, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ..domain.exceptions import ActiveSagaAlreadyExistsError, SagaNotFoundError
+from ..domain.mixins import SagaStateMixin
+from ..domain.models.enums import SagaStatus
+
+ModelT = TypeVar("ModelT", bound=SagaStateMixin)
+
+
+class SagaRepository(Generic[ModelT]):
+    """Provide persistence operations for saga state rows."""
+
+    ACTIVE_STATUSES: tuple[SagaStatus, ...] = (
+        SagaStatus.RUNNING,
+        SagaStatus.SUSPENDED,
+        SagaStatus.COMPENSATING,
+    )
+
+    def __init__(self, model_class: type[ModelT]) -> None:
+        """Initialize the repository for one saga state model."""
+        self.model_class = model_class
+
+    async def get(self, session: AsyncSession, saga_id: UUID) -> ModelT:
+        """Return one saga row by id."""
+        stmt: Select[tuple[ModelT]] = select(self.model_class).where(
+            self.model_class.id == saga_id
+        )
+        result = await session.execute(stmt)
+        saga = result.scalar_one_or_none()
+        if saga is None:
+            raise SagaNotFoundError(f"Saga '{saga_id}' not found")
+        return saga
+
+    async def get_for_update(self, session: AsyncSession, saga_id: UUID) -> ModelT:
+        """Return one saga row by id and lock it for update."""
+        stmt: Select[tuple[ModelT]] = (
+            select(self.model_class)
+            .where(self.model_class.id == saga_id)
+            .with_for_update(nowait=False)
+        )
+        result = await session.execute(stmt)
+        saga = result.scalar_one_or_none()
+        if saga is None:
+            raise SagaNotFoundError(f"Saga '{saga_id}' not found")
+        return saga
+
+    async def get_active_by_aggregation_id_for_update(
+        self,
+        session: AsyncSession,
+        aggregation_id: str,
+    ) -> ModelT | None:
+        """Return an active saga for the aggregation id and lock it for update."""
+        stmt: Select[tuple[ModelT]] = (
+            select(self.model_class)
+            .where(
+                self.model_class.aggregation_id == aggregation_id,
+                self.model_class.status.in_(self.ACTIVE_STATUSES),
+            )
+            .with_for_update(nowait=False)
+        )
+        result = await session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def ensure_no_active_aggregation_conflict(
+        self,
+        session: AsyncSession,
+        aggregation_id: str,
+    ) -> None:
+        """Raise when an active saga exists for the aggregation id."""
+        existing = await self.get_active_by_aggregation_id_for_update(
+            session,
+            aggregation_id,
+        )
+        if existing is not None:
+            raise ActiveSagaAlreadyExistsError(
+                "Active saga already exists for aggregation_id "
+                f"'{aggregation_id}' (saga_id={existing.id})"
+            )
+
+    async def create(self, session: AsyncSession, saga: ModelT) -> ModelT:
+        """Add a new saga row to the session and flush it."""
+        session.add(saga)
+        await session.flush()
+        return saga
+
+    async def due_suspended(
+        self,
+        session: AsyncSession,
+        now: datetime,
+        limit: int,
+    ) -> list[ModelT]:
+        """Return suspended sagas whose deadlines are due."""
+        return await self._due_by_status(
+            session=session,
+            status=SagaStatus.SUSPENDED,
+            now=now,
+            limit=limit,
+        )
+
+    async def due_running(
+        self,
+        session: AsyncSession,
+        now: datetime,
+        limit: int,
+    ) -> list[ModelT]:
+        """Return running sagas whose deadlines are due."""
+        return await self._due_by_status(
+            session=session,
+            status=SagaStatus.RUNNING,
+            now=now,
+            limit=limit,
+        )
+
+    async def due_compensating(
+        self,
+        session: AsyncSession,
+        now: datetime,
+        limit: int,
+    ) -> list[ModelT]:
+        """Return compensating sagas whose deadlines are due."""
+        return await self._due_by_status(
+            session=session,
+            status=SagaStatus.COMPENSATING,
+            now=now,
+            limit=limit,
+        )
+
+    async def _due_by_status(
+        self,
+        *,
+        session: AsyncSession,
+        status: SagaStatus,
+        now: datetime,
+        limit: int,
+    ) -> list[ModelT]:
+        """Return due saga rows for one status ordered by deadline."""
+        stmt = (
+            select(self.model_class)
+            .where(
+                self.model_class.status == status,
+                self.model_class.deadline_at.is_not(None),
+                self.model_class.deadline_at <= now,
+            )
+            .order_by(self.model_class.deadline_at.asc())
+            .limit(limit)
+        )
+        if self._supports_skip_locked(session):
+            stmt = stmt.with_for_update(skip_locked=True)
+        else:
+            stmt = stmt.with_for_update(nowait=False)
+
+        result = await session.execute(stmt)
+        return list(result.scalars().all())
+
+    @staticmethod
+    def _supports_skip_locked(session: AsyncSession) -> bool:
+        """Return whether the current dialect supports ``SKIP LOCKED``."""
+        bind = session.get_bind()
+        return bind is not None and bind.dialect.name == "postgresql"

saga_orchestrator/domain/__init__.py

@@ -0,0 +1 @@
+"""Domain module."""

saga_orchestrator/domain/exceptions/__init__.py

@@ -0,0 +1,17 @@
+"""Domain exceptions module."""
+
+from .saga import (
+    ActiveSagaAlreadyExistsError,
+    SagaDefinitionError,
+    SagaNotFoundError,
+    SagaStateError,
+    TypeValidationError,
+)
+
+__all__ = [
+    "ActiveSagaAlreadyExistsError",
+    "SagaDefinitionError",
+    "TypeValidationError",
+    "SagaNotFoundError",
+    "SagaStateError",
+]

saga_orchestrator/domain/exceptions/saga.py

@@ -0,0 +1,22 @@
+class SagaError(Exception):
+    """Base exception for orchestrator errors."""
+
+
+class SagaDefinitionError(SagaError):
+    """Invalid saga definition or registration."""
+
+
+class TypeValidationError(SagaDefinitionError):
+    """Type mismatch in step definitions."""
+
+
+class SagaNotFoundError(SagaError):
+    """Saga instance is not found in persistence."""
+
+
+class ActiveSagaAlreadyExistsError(SagaError):
+    """An active saga already exists for the provided aggregation key."""
+
+
+class SagaStateError(SagaError):
+    """Invalid saga state transition or operation."""

saga_orchestrator/domain/mixins/__init__.py

@@ -0,0 +1,7 @@
+"""Domain mixins module."""
+
+from .saga_state import SagaStateMixin
+
+__all__ = [
+    "SagaStateMixin",
+]

saga_orchestrator/domain/mixins/saga_state.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import uuid
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import JSON, DateTime, Enum, Integer, String, Text, func
+from sqlalchemy.dialects.postgresql import JSONB, UUID
+from sqlalchemy.ext.mutable import MutableDict, MutableList
+from sqlalchemy.orm import Mapped, declarative_mixin, mapped_column
+
+from ..models.enums import SagaStatus
+
+
+def _json_type() -> JSON:
+    return JSON().with_variant(JSONB, "postgresql")
+
+
+@declarative_mixin
+class SagaStateMixin:
+    id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
+    )
+    aggregation_id: Mapped[str] = mapped_column(String(255), index=True)
+    trace_id: Mapped[str] = mapped_column(String(255), index=True)
+    status: Mapped[SagaStatus] = mapped_column(
+        Enum(SagaStatus),
+        default=SagaStatus.RUNNING,
+        index=True,
+    )
+    current_step_index: Mapped[int] = mapped_column(Integer, default=0)
+    step_execution_token: Mapped[uuid.UUID | None] = mapped_column(
+        UUID(as_uuid=True),
+        nullable=True,
+    )
+    context: Mapped[dict[str, Any]] = mapped_column(
+        MutableDict.as_mutable(_json_type()), default=dict
+    )
+    step_history: Mapped[list[dict[str, Any]]] = mapped_column(
+        MutableList.as_mutable(_json_type()), default=list
+    )
+    deadline_at: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True),
+        nullable=True,
+        index=True,
+    )
+    retry_counter: Mapped[int] = mapped_column(Integer, default=0)
+    last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), nullable=False, server_default=func.now()
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True),
+        nullable=False,
+        server_default=func.now(),
+        onupdate=func.now(),
+    )

saga_orchestrator/domain/models/__init__.py

@@ -0,0 +1,21 @@
+"""Domain models module."""
+
+from .builder import SagaDefinition
+from .retry import ExponentialRetry, FixedRetry, NoRetry, RetryPolicy
+from .saga_snapshot import SagaAdminSnapshot, SagaSnapshot
+from .step import BaseStep, InputContext, StepDefinition, StepInputMap, StepRef
+
+__all__ = [
+    "SagaDefinition",
+    "RetryPolicy",
+    "NoRetry",
+    "FixedRetry",
+    "ExponentialRetry",
+    "SagaAdminSnapshot",
+    "SagaSnapshot",
+    "StepRef",
+    "InputContext",
+    "StepInputMap",
+    "StepDefinition",
+    "BaseStep",
+]

saga_orchestrator/domain/models/builder.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .step import StepDefinition
+
+
+@dataclass(frozen=True)
+class SagaDefinition:
+    steps: tuple[StepDefinition[Any, Any], ...]
+    compensate_on_failure: bool = True

saga_orchestrator/domain/models/enums/__init__.py

@@ -0,0 +1,7 @@
+"""Domain enum models module."""
+
+from .saga_status import SagaStatus
+
+__all__ = [
+    "SagaStatus",
+]

saga_orchestrator/domain/models/enums/saga_status.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class SagaStatus(str, Enum):
+    RUNNING = "RUNNING"
+    SUSPENDED = "SUSPENDED"
+    FAILED = "FAILED"
+    COMPENSATING = "COMPENSATING"
+    COMPLETED = "COMPLETED"
+
+    @property
+    def is_terminal(self) -> bool:
+        return self in {self.FAILED, self.COMPLETED}

saga_orchestrator/domain/models/retry.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+
+
+@dataclass(frozen=True)
+class RetryPolicy:
+    max_attempts: int
+
+    def next_delay(self, attempt_number: int) -> timedelta | None:
+        if attempt_number > self.max_attempts:
+            return None
+        return self._delay_for_attempt(attempt_number)
+
+    def _delay_for_attempt(self, attempt_number: int) -> timedelta:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class NoRetry(RetryPolicy):
+    def __init__(self) -> None:
+        super().__init__(max_attempts=0)
+
+    def _delay_for_attempt(self, attempt_number: int) -> timedelta:
+        return timedelta(seconds=0)
+
+
+@dataclass(frozen=True)
+class FixedRetry(RetryPolicy):
+    delay: timedelta
+
+    def _delay_for_attempt(self, attempt_number: int) -> timedelta:
+        return self.delay
+
+
+@dataclass(frozen=True)
+class ExponentialRetry(RetryPolicy):
+    base_delay: timedelta
+    multiplier: float = 2.0
+    max_delay: timedelta | None = None
+
+    def _delay_for_attempt(self, attempt_number: int) -> timedelta:
+        seconds = self.base_delay.total_seconds() * (
+            self.multiplier ** max(attempt_number - 1, 0)
+        )
+        delay = timedelta(seconds=seconds)
+        if self.max_delay is None:
+            return delay
+        return min(delay, self.max_delay)

saga_orchestrator/domain/models/saga_snapshot.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+from .enums import SagaStatus
+
+
+@dataclass(frozen=True)
+class SagaSnapshot:
+    id: UUID
+    aggregation_id: str
+    status: SagaStatus
+    current_step_index: int
+    retry_counter: int
+    deadline_at: datetime | None
+    trace_id: str
+    step_execution_token: UUID | None
+    last_error: str | None
+
+
+@dataclass(frozen=True)
+class SagaAdminSnapshot:
+    id: UUID
+    aggregation_id: str
+    trace_id: str
+    status: SagaStatus
+    current_step_index: int
+    step_execution_token: UUID | None
+    retry_counter: int
+    deadline_at: datetime | None
+    last_error: str | None
+    context: dict[str, Any]
+    step_history: list[dict[str, Any]]

saga_orchestrator/domain/models/step.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, Generic, TypeAlias, TypeVar, get_type_hints
+
+from pydantic import BaseModel
+
+from ..exceptions import TypeValidationError
+from .retry import RetryPolicy
+
+InputModelT = TypeVar("InputModelT", bound=BaseModel)
+OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
+DepModelT = TypeVar("DepModelT", bound=BaseModel)
+
+
+@dataclass(frozen=True)
+class StepRef(Generic[OutputModelT]):
+    step_id: str
+    output_model: type[OutputModelT]
+
+
+@dataclass
+class InputContext:
+    initial_data: Any
+    context: dict[str, Any]
+    step_outputs: dict[str, Any]
+    latest_event: Any | None = None
+    events: list[Any] | None = None
+
+
+RootInputMap: TypeAlias = Callable[[InputContext], InputModelT | dict[str, Any]]
+StepInputMap: TypeAlias = RootInputMap | Callable[[Any], InputModelT | dict[str, Any]]
+
+
+@dataclass
+class StepDefinition(Generic[InputModelT, OutputModelT]):
+    step_id: str
+    step: BaseStep[InputModelT, OutputModelT]
+    input_map: StepInputMap[InputModelT]
+    timeout: timedelta | None
+    retry_policy: RetryPolicy
+    depends_on: StepRef[Any] | None = None
+
+    @property
+    def input_model(self) -> type[InputModelT]:
+        return self.step.input_model
+
+    @property
+    def output_model(self) -> type[OutputModelT]:
+        return self.step.output_model
+
+
+class BaseStep(Generic[InputModelT, OutputModelT]):
+    input_model: type[InputModelT]
+    output_model: type[OutputModelT]
+
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+        if cls is BaseStep:
+            return
+
+        hints = get_type_hints(cls.execute)
+        if "inp" not in hints or "return" not in hints:
+            raise TypeValidationError(
+                f"Step '{cls.__name__}' must type annotate execute(inp) and return type"
+            )
+        input_model = hints["inp"]
+        output_model = hints["return"]
+        if not (inspect.isclass(input_model) and issubclass(input_model, BaseModel)):
+            raise TypeValidationError(
+                f"Step '{cls.__name__}' input must inherit from pydantic BaseModel"
+            )
+        if not (inspect.isclass(output_model) and issubclass(output_model, BaseModel)):
+            raise TypeValidationError(
+                f"Step '{cls.__name__}' output must inherit from pydantic BaseModel"
+            )
+        cls.input_model = input_model
+        cls.output_model = output_model
+
+    async def execute(self, inp: InputModelT) -> OutputModelT:
+        raise NotImplementedError
+
+    async def compensate(self, inp: InputModelT, out: OutputModelT) -> None:
+        raise NotImplementedError