Distributed Transactions — The Saga Pattern.

# Saga Choreography Implementation

import asyncio
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Callable
from datetime import datetime, timezone
from enum import Enum
import json
import logging

logger = logging.getLogger(__name__)


# =============================================================================
# Domain Events
# =============================================================================

@dataclass
class DomainEvent:
    """Base class for domain events."""
    event_id: str
    correlation_id: str  # Saga ID
    # kw_only (Python 3.10+) lets subclasses add required fields after this defaulted one
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc), kw_only=True)
    
    def to_dict(self) -> dict:
        return {
            "event_id": self.event_id,
            "correlation_id": self.correlation_id,
            "timestamp": self.timestamp.isoformat(),
            "event_type": self.__class__.__name__,
            "data": self._event_data()
        }
    
    def _event_data(self) -> dict:
        raise NotImplementedError


@dataclass
class OrderCreated(DomainEvent):
    order_id: str
    customer_id: str
    items: List[dict]
    total_amount: float
    
    def _event_data(self) -> dict:
        return {
            "order_id": self.order_id,
            "customer_id": self.customer_id,
            "items": self.items,
            "total_amount": self.total_amount
        }


@dataclass
class InventoryReserved(DomainEvent):
    order_id: str
    reservation_id: str
    items: List[dict]
    
    def _event_data(self) -> dict:
        return {
            "order_id": self.order_id,
            "reservation_id": self.reservation_id,
            "items": self.items
        }


@dataclass
class InventoryReservationFailed(DomainEvent):
    order_id: str
    reason: str
    
    def _event_data(self) -> dict:
        return {"order_id": self.order_id, "reason": self.reason}


@dataclass
class PaymentCompleted(DomainEvent):
    order_id: str
    payment_id: str
    amount: float
    
    def _event_data(self) -> dict:
        return {
            "order_id": self.order_id,
            "payment_id": self.payment_id,
            "amount": self.amount
        }


@dataclass
class PaymentFailed(DomainEvent):
    order_id: str
    reason: str
    
    def _event_data(self) -> dict:
        return {"order_id": self.order_id, "reason": self.reason}


@dataclass  
class InventoryReleased(DomainEvent):
    order_id: str
    reservation_id: str
    
    def _event_data(self) -> dict:
        return {
            "order_id": self.order_id,
            "reservation_id": self.reservation_id
        }


@dataclass
class OrderCompleted(DomainEvent):
    order_id: str
    
    def _event_data(self) -> dict:
        return {"order_id": self.order_id}


@dataclass
class OrderFailed(DomainEvent):
    order_id: str
    reason: str
    
    def _event_data(self) -> dict:
        return {"order_id": self.order_id, "reason": self.reason}


# =============================================================================
# Event Publisher/Subscriber
# =============================================================================

class EventBus:
    """Simple in-memory event bus for demonstration."""
    
    def __init__(self):
        self._handlers: Dict[str, List[Callable]] = {}
    
    def subscribe(self, event_type: str, handler: Callable):
        if event_type not in self._handlers:
            self._handlers[event_type] = []
        self._handlers[event_type].append(handler)
    
    async def publish(self, event: DomainEvent):
        event_type = event.__class__.__name__
        logger.info(f"Publishing event: {event_type} for saga {event.correlation_id}")
        
        handlers = self._handlers.get(event_type, [])
        for handler in handlers:
            try:
                await handler(event)
            except Exception as e:
                logger.error(f"Handler failed for {event_type}: {e}")


# =============================================================================
# Services (Choreography Participants)
# =============================================================================

class OrderService:
    """Order service - initiates and tracks order saga."""
    
    def __init__(self, db, event_bus: EventBus):
        self.db = db
        self.events = event_bus
        
        # Subscribe to events
        self.events.subscribe("InventoryReserved", self._on_inventory_reserved)
        self.events.subscribe("InventoryReservationFailed", self._on_inventory_failed)
        self.events.subscribe("PaymentCompleted", self._on_payment_completed)
        self.events.subscribe("PaymentFailed", self._on_payment_failed)
    
    async def create_order(self, customer_id: str, items: List[dict]) -> str:
        """Create order and start saga via event."""
        order_id = generate_id()
        total = sum(item['price'] * item['quantity'] for item in items)
        
        # Save order in PENDING state
        await self.db.execute(
            """
            INSERT INTO orders (id, customer_id, items, total, status, created_at)
            VALUES ($1, $2, $3, $4, 'PENDING', NOW())
            """,
            order_id, customer_id, json.dumps(items), total
        )
        
        # Publish event to start saga
        event = OrderCreated(
            event_id=generate_id(),
            correlation_id=order_id,  # Order ID is the saga ID
            order_id=order_id,
            customer_id=customer_id,
            items=items,
            total_amount=total
        )
        await self.events.publish(event)
        
        return order_id
    
    async def _on_inventory_reserved(self, event: InventoryReserved):
        """Handle inventory reserved - update order state."""
        await self.db.execute(
            """
            UPDATE orders 
            SET status = 'INVENTORY_RESERVED', reservation_id = $2
            WHERE id = $1
            """,
            event.order_id, event.reservation_id
        )
    
    async def _on_inventory_failed(self, event: InventoryReservationFailed):
        """Handle inventory failure - fail order."""
        await self.db.execute(
            "UPDATE orders SET status = 'FAILED', failure_reason = $2 WHERE id = $1",
            event.order_id, event.reason
        )
        
        await self.events.publish(OrderFailed(
            event_id=generate_id(),
            correlation_id=event.correlation_id,
            order_id=event.order_id,
            reason=event.reason
        ))
    
    async def _on_payment_completed(self, event: PaymentCompleted):
        """Handle payment completed - complete order."""
        await self.db.execute(
            """
            UPDATE orders 
            SET status = 'COMPLETED', payment_id = $2
            WHERE id = $1
            """,
            event.order_id, event.payment_id
        )
        
        await self.events.publish(OrderCompleted(
            event_id=generate_id(),
            correlation_id=event.correlation_id,
            order_id=event.order_id
        ))
    
    async def _on_payment_failed(self, event: PaymentFailed):
        """Handle payment failure - will trigger compensation."""
        await self.db.execute(
            "UPDATE orders SET status = 'PAYMENT_FAILED', failure_reason = $2 WHERE id = $1",
            event.order_id, event.reason
        )


class InventoryService:
    """Inventory service - reserves and releases inventory."""
    
    def __init__(self, db, event_bus: EventBus):
        self.db = db
        self.events = event_bus
        
        # Subscribe to events
        self.events.subscribe("OrderCreated", self._on_order_created)
        self.events.subscribe("PaymentFailed", self._on_payment_failed)
    
    async def _on_order_created(self, event: OrderCreated):
        """React to order created - try to reserve inventory."""
        try:
            reservation_id = await self._reserve_inventory(
                event.order_id, 
                event.items
            )
            
            await self.events.publish(InventoryReserved(
                event_id=generate_id(),
                correlation_id=event.correlation_id,
                order_id=event.order_id,
                reservation_id=reservation_id,
                items=event.items
            ))
            
        except InsufficientInventoryError as e:
            await self.events.publish(InventoryReservationFailed(
                event_id=generate_id(),
                correlation_id=event.correlation_id,
                order_id=event.order_id,
                reason=str(e)
            ))
    
    async def _on_payment_failed(self, event: PaymentFailed):
        """
        React to payment failed - COMPENSATE by releasing inventory.
        
        This is the compensation transaction triggered by choreography.
        """
        # Find the reservation for this order
        reservation = await self.db.fetch_one(
            "SELECT id FROM inventory_reservations WHERE order_id = $1 AND status = 'active'",
            event.order_id
        )
        
        if reservation:
            await self._release_inventory(reservation['id'])
            
            await self.events.publish(InventoryReleased(
                event_id=generate_id(),
                correlation_id=event.correlation_id,
                order_id=event.order_id,
                reservation_id=reservation['id']
            ))
    
    async def _reserve_inventory(self, order_id: str, items: List[dict]) -> str:
        """Reserve inventory for order items."""
        reservation_id = generate_id()
        
        async with self.db.transaction(isolation="SERIALIZABLE"):
            for item in items:
                # Check and decrement inventory
                result = await self.db.execute(
                    """
                    UPDATE inventory 
                    SET quantity = quantity - $2
                    WHERE product_id = $1 AND quantity >= $2
                    """,
                    item['product_id'], item['quantity']
                )
                
                if result.rowcount == 0:
                    raise InsufficientInventoryError(
                        f"Insufficient inventory for {item['product_id']}"
                    )
            
            # Create reservation record
            await self.db.execute(
                """
                INSERT INTO inventory_reservations (id, order_id, items, status, created_at)
                VALUES ($1, $2, $3, 'active', NOW())
                """,
                reservation_id, order_id, json.dumps(items)
            )
        
        return reservation_id
    
    async def _release_inventory(self, reservation_id: str):
        """Release inventory reservation (compensation)."""
        async with self.db.transaction():
            reservation = await self.db.fetch_one(
                """
                SELECT items FROM inventory_reservations 
                WHERE id = $1 AND status = 'active'
                FOR UPDATE
                """,
                reservation_id
            )
            
            if not reservation:
                logger.warning(f"Reservation {reservation_id} not found or not active")
                return  # Idempotent - already released
            
            items = json.loads(reservation['items'])
            
            for item in items:
                await self.db.execute(
                    "UPDATE inventory SET quantity = quantity + $2 WHERE product_id = $1",
                    item['product_id'], item['quantity']
                )
            
            await self.db.execute(
                "UPDATE inventory_reservations SET status = 'released' WHERE id = $1",
                reservation_id
            )


class PaymentService:
    """Payment service - charges and refunds payments."""
    
    def __init__(self, db, event_bus: EventBus, payment_gateway):
        self.db = db
        self.events = event_bus
        self.gateway = payment_gateway
        
        # Subscribe to events
        self.events.subscribe("InventoryReserved", self._on_inventory_reserved)
    
    async def _on_inventory_reserved(self, event: InventoryReserved):
        """React to inventory reserved - charge payment."""
        # Get order details
        order = await self.db.fetch_one(
            "SELECT customer_id, total FROM orders WHERE id = $1",
            event.order_id
        )
        
        try:
            payment_id = await self._charge_customer(
                order['customer_id'],
                order['total'],
                event.order_id
            )
            
            await self.events.publish(PaymentCompleted(
                event_id=generate_id(),
                correlation_id=event.correlation_id,
                order_id=event.order_id,
                payment_id=payment_id,
                amount=order['total']
            ))
            
        except PaymentError as e:
            await self.events.publish(PaymentFailed(
                event_id=generate_id(),
                correlation_id=event.correlation_id,
                order_id=event.order_id,
                reason=str(e)
            ))
    
    async def _charge_customer(
        self,
        customer_id: str,
        amount: float,
        order_id: str
    ) -> str:
        """Charge customer via payment gateway."""
        payment_id = generate_id()
        
        # Get customer payment method
        customer = await self.db.fetch_one(
            "SELECT payment_method_id FROM customers WHERE id = $1",
            customer_id
        )
        
        # Charge via gateway
        result = await self.gateway.charge(
            payment_method_id=customer['payment_method_id'],
            amount=amount,
            idempotency_key=f"order-{order_id}"  # Idempotent!
        )
        
        if not result.success:
            raise PaymentError(result.error_message)
        
        # Record payment
        await self.db.execute(
            """
            INSERT INTO payments (id, order_id, customer_id, amount, status, gateway_ref)
            VALUES ($1, $2, $3, $4, 'completed', $5)
            """,
            payment_id, order_id, customer_id, amount, result.reference
        )
        
        return payment_id

# Saga Orchestration Implementation

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, Callable
from enum import Enum
from datetime import datetime, timedelta, timezone
import asyncio
import json
import logging

logger = logging.getLogger(__name__)


# =============================================================================
# Saga Definition
# =============================================================================

class SagaStepStatus(Enum):
    PENDING = "pending"
    EXECUTING = "executing"
    COMPLETED = "completed"
    FAILED = "failed"
    COMPENSATING = "compensating"
    COMPENSATED = "compensated"


class SagaStatus(Enum):
    STARTED = "started"
    EXECUTING = "executing"
    COMPLETED = "completed"
    COMPENSATING = "compensating"
    COMPENSATED = "compensated"
    FAILED = "failed"


@dataclass
class SagaStep:
    """Definition of a saga step."""
    name: str
    action: Callable  # The forward action
    compensation: Optional[Callable]  # The compensation action
    is_pivot: bool = False  # After pivot, no compensation - must complete
    retryable: bool = True
    max_retries: int = 3
    timeout: timedelta = field(default_factory=lambda: timedelta(seconds=30))


@dataclass
class SagaStepResult:
    """Result of executing a saga step."""
    success: bool
    data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


@dataclass
class SagaState:
    """Persistent state of a saga execution."""
    saga_id: str
    saga_type: str
    status: SagaStatus
    current_step: int
    input_data: Dict[str, Any]
    step_results: Dict[str, Any]  # Results from each step
    created_at: datetime
    updated_at: datetime
    error: Optional[str] = None


# =============================================================================
# Saga Definition DSL
# =============================================================================

class SagaDefinition:
    """
    Fluent interface for defining sagas.
    
    Usage:
        saga = (SagaDefinition("OrderSaga")
            .step("reserve_inventory", reserve_fn, release_fn)
            .step("charge_payment", charge_fn, refund_fn)
            .step("create_shipment", ship_fn, None, is_pivot=True)
            .step("send_notification", notify_fn, None)
            .build())
    """
    
    def __init__(self, name: str):
        self.name = name
        self.steps: List[SagaStep] = []
    
    def step(
        self,
        name: str,
        action: Callable,
        compensation: Optional[Callable] = None,
        is_pivot: bool = False,
        retryable: bool = True,
        max_retries: int = 3,
        timeout_seconds: int = 30
    ) -> 'SagaDefinition':
        """Add a step to the saga."""
        self.steps.append(SagaStep(
            name=name,
            action=action,
            compensation=compensation,
            is_pivot=is_pivot,
            retryable=retryable,
            max_retries=max_retries,
            timeout=timedelta(seconds=timeout_seconds)
        ))
        return self
    
    def build(self) -> 'Saga':
        """Build the saga."""
        return Saga(self.name, self.steps)


class Saga:
    """Executable saga definition."""
    
    def __init__(self, name: str, steps: List[SagaStep]):
        self.name = name
        self.steps = steps
        
        # Validate: steps after pivot should not have compensation
        pivot_found = False
        for step in steps:
            if step.is_pivot:
                pivot_found = True
            elif pivot_found and step.compensation:
                logger.warning(
                    f"Step {step.name} has compensation after pivot - will be ignored"
                )


# =============================================================================
# Saga Orchestrator
# =============================================================================

class SagaOrchestrator:
    """
    Central orchestrator for executing sagas.
    
    Responsibilities:
    - Execute saga steps in order
    - Persist saga state for recovery
    - Handle failures and compensation
    - Retry failed steps
    """
    
    def __init__(self, db, sagas: Dict[str, Saga]):
        self.db = db
        self.sagas = sagas  # Registered saga definitions
        self._tasks = set()  # strong references: the event loop only keeps weak ones

    def _spawn(self, coro):
        """Run a coroutine in the background without letting it be garbage-collected."""
        task = asyncio.create_task(coro)
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)
        return task
    
    async def start_saga(
        self,
        saga_type: str,
        input_data: Dict[str, Any]
    ) -> str:
        """Start a new saga execution."""
        saga = self.sagas.get(saga_type)
        if not saga:
            raise ValueError(f"Unknown saga type: {saga_type}")
        
        saga_id = generate_id()
        
        state = SagaState(
            saga_id=saga_id,
            saga_type=saga_type,
            status=SagaStatus.STARTED,
            current_step=0,
            input_data=input_data,
            step_results={},
            created_at=datetime.now(timezone.utc),
            updated_at=datetime.now(timezone.utc)
        )
        
        await self._save_state(state)
        
        # Execute asynchronously
        self._spawn(self._execute_saga(state))
        
        return saga_id
    
    async def _execute_saga(self, state: SagaState):
        """Execute the saga steps."""
        saga = self.sagas[state.saga_type]
        
        state.status = SagaStatus.EXECUTING
        await self._save_state(state)
        
        try:
            # Execute each step
            for i, step in enumerate(saga.steps):
                if step.name in state.step_results:
                    continue  # Already completed before a crash (recovery)
                state.current_step = i
                await self._save_state(state)
                
                result = await self._execute_step(step, state)
                
                if result.success:
                    state.step_results[step.name] = result.data
                    await self._save_step_result(state, i, step.name, result.data)
                    await self._save_state(state)
                else:
                    # Step failed
                    state.error = result.error
                    await self._save_state(state)
                    
                    if self._past_pivot(saga, i):
                        # Pivot already committed - must complete with forward recovery
                        # (a failed pivot never committed, so we still compensate)
                        await self._forward_recovery(saga, state, i)
                    else:
                        # Before pivot - compensate
                        await self._compensate(saga, state, i)
                    return
            
            # All steps completed successfully
            state.status = SagaStatus.COMPLETED
            await self._save_state(state)
            logger.info(f"Saga {state.saga_id} completed successfully")
            
        except Exception as e:
            logger.error(f"Saga {state.saga_id} failed with error: {e}")
            state.status = SagaStatus.FAILED
            state.error = str(e)
            await self._save_state(state)
    
    async def _execute_step(
        self,
        step: SagaStep,
        state: SagaState
    ) -> SagaStepResult:
        """Execute a single saga step with retry logic."""
        last_error = None
        
        for attempt in range(step.max_retries):
            try:
                # Build step context with all previous results
                context = {
                    "saga_id": state.saga_id,
                    "input": state.input_data,
                    "results": state.step_results
                }
                
                # Execute with timeout
                result = await asyncio.wait_for(
                    step.action(context),
                    timeout=step.timeout.total_seconds()
                )
                
                return SagaStepResult(success=True, data=result)
                
            except asyncio.TimeoutError:
                last_error = f"Step {step.name} timed out"
                logger.warning(f"{last_error}, attempt {attempt + 1}/{step.max_retries}")
                
            except Exception as e:
                last_error = str(e)
                logger.warning(
                    f"Step {step.name} failed: {e}, attempt {attempt + 1}/{step.max_retries}"
                )
            
            if attempt < step.max_retries - 1 and step.retryable:
                # Exponential backoff
                await asyncio.sleep(2 ** attempt)
        
        return SagaStepResult(success=False, error=last_error)
    
    async def _compensate(self, saga: Saga, state: SagaState, failed_step: int):
        """Execute compensation for completed steps in reverse order."""
        state.status = SagaStatus.COMPENSATING
        await self._save_state(state)
        
        logger.info(f"Starting compensation for saga {state.saga_id}")
        
        # Compensate in reverse order, excluding the failed step
        for i in range(failed_step - 1, -1, -1):
            step = saga.steps[i]
            
            if step.compensation is None:
                logger.info(f"No compensation for step {step.name}")
                continue
            
            if step.is_pivot:
                # Don't compensate pivot or steps before it
                logger.info(f"Reached pivot step {step.name}, stopping compensation")
                break
            
            await self._execute_compensation(step, state)
        
        state.status = SagaStatus.COMPENSATED
        await self._save_state(state)
        logger.info(f"Saga {state.saga_id} compensated")
    
    async def _execute_compensation(self, step: SagaStep, state: SagaState):
        """Execute compensation for a single step."""
        context = {
            "saga_id": state.saga_id,
            "input": state.input_data,
            "results": state.step_results,
            "step_result": state.step_results.get(step.name)
        }
        
        # Compensation should not fail - retry until success
        for attempt in range(10):  # More retries for compensation
            try:
                await asyncio.wait_for(
                    step.compensation(context),
                    timeout=step.timeout.total_seconds() * 2
                )
                logger.info(f"Compensated step {step.name}")
                await self._record_compensation(state, step.name, "completed")
                return
                
            except Exception as e:
                logger.error(
                    f"Compensation for {step.name} failed: {e}, "
                    f"attempt {attempt + 1}/10"
                )
                await asyncio.sleep(2 ** attempt)
        
        # Compensation failed - this is bad!
        logger.critical(
            f"COMPENSATION FAILED for step {step.name} in saga {state.saga_id}. "
            f"Manual intervention required!"
        )
        # Hand it to the recovery job (Chapter 12), alert operations team, etc.
        await self._record_compensation(state, step.name, "failed")
    
    async def _forward_recovery(
        self,
        saga: Saga,
        state: SagaState,
        failed_step: int
    ):
        """Retry steps after pivot until success."""
        logger.info(f"Forward recovery for saga {state.saga_id} from step {failed_step}")
        
        for i in range(failed_step, len(saga.steps)):
            step = saga.steps[i]
            
            # Retry indefinitely until success (with backoff)
            for attempt in range(100):  # Effectively unlimited
                result = await self._execute_step(step, state)
                
                if result.success:
                    state.step_results[step.name] = result.data
                    await self._save_step_result(state, i, step.name, result.data)
                    await self._save_state(state)
                    break
                
                wait_time = min(300, 2 ** attempt)  # Cap at 5 minutes
                logger.warning(
                    f"Forward recovery step {step.name} failed, "
                    f"retrying in {wait_time}s"
                )
                await asyncio.sleep(wait_time)
        
        state.status = SagaStatus.COMPLETED
        await self._save_state(state)
    
    def _past_pivot(self, saga: Saga, step_index: int) -> bool:
        """Check if we're past the pivot point."""
        for i in range(step_index):
            if saga.steps[i].is_pivot:
                return True
        return False
    
    async def _save_state(self, state: SagaState):
        """Persist saga state to database (schema in Chapter 12)."""
        state.updated_at = datetime.now(timezone.utc)
        
        await self.db.execute(
            """
            INSERT INTO saga_instances (saga_id, saga_type, status, current_step, 
                                       input_data, error_message, created_at, updated_at)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            ON CONFLICT (saga_id) DO UPDATE SET
                status = $3, current_step = $4, 
                error_message = $6, updated_at = $8
            """,
            state.saga_id, state.saga_type, state.status.value, state.current_step,
            json.dumps(state.input_data),
            state.error, state.created_at, state.updated_at
        )
    
    async def _save_step_result(self, state: SagaState, index: int, name: str, data):
        """Record a completed step; recovery rebuilds step results from these rows."""
        await self.db.execute(
            """
            INSERT INTO saga_step_executions
                (id, saga_id, step_name, step_index, status, output_data, completed_at)
            VALUES ($1, $2, $3, $4, 'completed', $5, NOW())
            """,
            generate_id(), state.saga_id, name, index, json.dumps(data)
        )
    
    async def _record_compensation(self, state: SagaState, step_name: str, status: str):
        """Record a compensation outcome; the recovery job retries 'failed' ones."""
        await self.db.execute(
            """
            INSERT INTO saga_compensations (id, saga_id, step_name, status, completed_at, next_retry_at)
            VALUES ($1, $2, $3, $4,
                    CASE WHEN $4 = 'completed' THEN NOW() END,
                    CASE WHEN $4 = 'failed' THEN NOW() END)
            ON CONFLICT (saga_id, step_name) DO UPDATE SET
                status = EXCLUDED.status, completed_at = EXCLUDED.completed_at,
                next_retry_at = EXCLUDED.next_retry_at
            """,
            generate_id(), state.saga_id, step_name, status
        )
    
    async def recover_sagas(self):
        """
        Recover incomplete sagas after restart.
        
        Called on startup to resume interrupted sagas.
        """
        incomplete = await self.db.fetch(
            """
            SELECT * FROM saga_instances 
            WHERE status IN ('started', 'executing', 'compensating')
            """
        )
        
        for row in incomplete:
            steps = await self.db.fetch(
                """
                SELECT step_name, output_data FROM saga_step_executions
                WHERE saga_id = $1 AND status = 'completed'
                """,
                row['saga_id']
            )
            state = SagaState(
                saga_id=row['saga_id'],
                saga_type=row['saga_type'],
                status=SagaStatus(row['status']),
                current_step=row['current_step'],
                input_data=json.loads(row['input_data']),
                step_results={s['step_name']: json.loads(s['output_data']) for s in steps},
                created_at=row['created_at'],
                updated_at=row['updated_at'],
                error=row['error_message']
            )
            
            logger.info(f"Recovering saga {state.saga_id} from step {state.current_step}")
            if state.status == SagaStatus.COMPENSATING:
                saga = self.sagas[state.saga_type]
                self._spawn(self._compensate(saga, state, state.current_step))
            else:
                self._spawn(self._execute_saga(state))


# =============================================================================
# Order Saga Definition
# =============================================================================

class OrderSagaSteps:
    """Step implementations for order saga."""
    
    def __init__(self, inventory_service, payment_service, shipping_service, notification_service):
        self.inventory = inventory_service
        self.payment = payment_service
        self.shipping = shipping_service
        self.notification = notification_service
    
    async def reserve_inventory(self, context: dict) -> dict:
        """Step 1: Reserve inventory."""
        order = context['input']
        
        reservation_id = await self.inventory.reserve(
            order_id=context['saga_id'],
            items=order['items']
        )
        
        return {"reservation_id": reservation_id}
    
    async def release_inventory(self, context: dict):
        """Compensation for reserve_inventory."""
        reservation_id = context['results'].get('reserve_inventory', {}).get('reservation_id')
        
        if reservation_id:
            await self.inventory.release(reservation_id)
    
    async def charge_payment(self, context: dict) -> dict:
        """Step 2: Charge payment."""
        order = context['input']
        
        payment_id = await self.payment.charge(
            customer_id=order['customer_id'],
            amount=order['total'],
            idempotency_key=f"saga-{context['saga_id']}"
        )
        
        return {"payment_id": payment_id}
    
    async def refund_payment(self, context: dict):
        """Compensation for charge_payment."""
        payment_id = context['results'].get('charge_payment', {}).get('payment_id')
        
        if payment_id:
            await self.payment.refund(
                payment_id=payment_id,
                idempotency_key=f"refund-{context['saga_id']}"
            )
    
    async def create_shipment(self, context: dict) -> dict:
        """Step 3: Create shipment (PIVOT - no compensation after this)."""
        order = context['input']
        
        shipment_id = await self.shipping.create_shipment(
            order_id=context['saga_id'],
            items=order['items'],
            address=order['shipping_address']
        )
        
        return {"shipment_id": shipment_id}
    
    async def send_notification(self, context: dict) -> dict:
        """Step 4: Send confirmation notification."""
        order = context['input']
        shipment_id = context['results']['create_shipment']['shipment_id']
        
        await self.notification.send_order_confirmation(
            customer_id=order['customer_id'],
            order_id=context['saga_id'],
            shipment_id=shipment_id
        )
        
        return {"notification_sent": True}


def create_order_saga(steps: OrderSagaSteps) -> Saga:
    """Create the order processing saga."""
    return (SagaDefinition("OrderSaga")
        .step(
            "reserve_inventory",
            steps.reserve_inventory,
            steps.release_inventory
        )
        .step(
            "charge_payment",
            steps.charge_payment,
            steps.refund_payment
        )
        .step(
            "create_shipment",
            steps.create_shipment,
            None,  # No compensation - this is the pivot
            is_pivot=True
        )
        .step(
            "send_notification",
            steps.send_notification,
            None  # Notification doesn't need compensation
        )
        .build())

-- Saga State Persistence Schema

-- Main saga state table
CREATE TABLE saga_instances (
    saga_id UUID PRIMARY KEY,
    saga_type VARCHAR(100) NOT NULL,
    status VARCHAR(50) NOT NULL,
    current_step INT NOT NULL DEFAULT 0,
    input_data JSONB NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMP,
    error_message TEXT,
    retry_count INT DEFAULT 0
);

-- Indexes for recovery queries (PostgreSQL: indexes are created separately)
CREATE INDEX idx_saga_status ON saga_instances (status);
CREATE INDEX idx_saga_type_status ON saga_instances (saga_type, status);
CREATE INDEX idx_saga_created ON saga_instances (created_at);

-- Step execution history
CREATE TABLE saga_step_executions (
    id UUID PRIMARY KEY,
    saga_id UUID NOT NULL REFERENCES saga_instances(saga_id),
    step_name VARCHAR(100) NOT NULL,
    step_index INT NOT NULL,
    status VARCHAR(50) NOT NULL,
    input_data JSONB,
    output_data JSONB,
    error_message TEXT,
    started_at TIMESTAMP NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMP,
    attempt_number INT NOT NULL DEFAULT 1,
    
    UNIQUE (saga_id, step_name, attempt_number)
);
CREATE INDEX idx_step_saga ON saga_step_executions (saga_id);

-- Compensation tracking
CREATE TABLE saga_compensations (
    id UUID PRIMARY KEY,
    saga_id UUID NOT NULL REFERENCES saga_instances(saga_id),
    step_name VARCHAR(100) NOT NULL,
    status VARCHAR(50) NOT NULL,  -- pending, executing, completed, failed
    started_at TIMESTAMP,
    completed_at TIMESTAMP,
    retry_count INT DEFAULT 0,
    last_error TEXT,
    next_retry_at TIMESTAMP,
    
    UNIQUE (saga_id, step_name)
);
CREATE INDEX idx_compensation_retry ON saga_compensations (status, next_retry_at);

-- Idempotency records
CREATE TABLE saga_idempotency_keys (
    idempotency_key VARCHAR(255) PRIMARY KEY,
    saga_id UUID NOT NULL,
    step_name VARCHAR(100) NOT NULL,
    response_data JSONB,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    expires_at TIMESTAMP NOT NULL
);
CREATE INDEX idx_idempotency_expires ON saga_idempotency_keys (expires_at);

# Background Job for Saga Recovery

import asyncio
from datetime import datetime, timedelta, timezone
from typing import List


class SagaRecoveryJob:
    """
    Background job that recovers stuck sagas.
    
    Handles:
    - Sagas interrupted by crashes
    - Stuck compensations
    - Timed-out steps
    """
    
    def __init__(self, db, orchestrator: SagaOrchestrator):
        self.db = db
        self.orchestrator = orchestrator
        self.running = False
    
    async def start(self):
        """Start the recovery job loop."""
        self.running = True
        
        while self.running:
            try:
                await self._recover_stuck_sagas()
                await self._retry_failed_compensations()
                await self._cleanup_expired_idempotency_keys()
            except Exception as e:
                logger.error(f"Recovery job error: {e}")
            
            await asyncio.sleep(30)  # Run every 30 seconds
    
    async def stop(self):
        """Stop the recovery job."""
        self.running = False
    
    async def _recover_stuck_sagas(self):
        """Find and recover sagas stuck in executing state."""
        # Find sagas stuck for more than 5 minutes
        cutoff = datetime.now(timezone.utc) - timedelta(minutes=5)
        
        stuck_sagas = await self.db.fetch(
            """
            SELECT * FROM saga_instances 
            WHERE status IN ('started', 'executing')
            AND updated_at < $1
            """,
            cutoff
        )
        
        for saga in stuck_sagas:
            logger.info(f"Recovering stuck saga: {saga['saga_id']}")
            
            # Reconstruct state and resume
            state = await self._reconstruct_state(saga)
            self.orchestrator._spawn(self.orchestrator._execute_saga(state))
    
    async def _retry_failed_compensations(self):
        """Retry compensations that failed but are due for retry."""
        now = datetime.now(timezone.utc)
        
        pending = await self.db.fetch(
            """
            SELECT c.*, s.saga_type
            FROM saga_compensations c
            JOIN saga_instances s ON c.saga_id = s.saga_id
            WHERE c.status = 'failed'
            AND c.next_retry_at <= $1
            """,
            now
        )
        
        for comp in pending:
            logger.info(
                f"Retrying compensation: {comp['step_name']} "
                f"for saga {comp['saga_id']}"
            )
            
            try:
                await self._execute_compensation(comp)
                
                await self.db.execute(
                    """
                    UPDATE saga_compensations 
                    SET status = 'completed', completed_at = NOW()
                    WHERE id = $1
                    """,
                    comp['id']
                )
                
            except Exception as e:
                retry_count = comp['retry_count'] + 1
                next_retry = now + timedelta(seconds=min(300, 2 ** retry_count))
                
                await self.db.execute(
                    """
                    UPDATE saga_compensations 
                    SET retry_count = $2, last_error = $3, next_retry_at = $4
                    WHERE id = $1
                    """,
                    comp['id'], retry_count, str(e), next_retry
                )
                
                if retry_count > 10:
                    logger.critical(
                        f"Compensation stuck after {retry_count} retries: "
                        f"{comp['saga_id']}/{comp['step_name']}"
                    )
                    # Alert operations team
    
    async def _cleanup_expired_idempotency_keys(self):
        """Clean up expired idempotency keys."""
        await self.db.execute(
            "DELETE FROM saga_idempotency_keys WHERE expires_at < NOW()"
        )
    
    async def _reconstruct_state(self, saga_row) -> SagaState:
        """Reconstruct saga state from database."""
        step_results = {}
        
        # Get completed step results
        steps = await self.db.fetch(
            """
            SELECT step_name, output_data 
            FROM saga_step_executions 
            WHERE saga_id = $1 AND status = 'completed'
            ORDER BY step_index
            """,
            saga_row['saga_id']
        )
        
        for step in steps:
            step_results[step['step_name']] = step['output_data']
        
        return SagaState(
            saga_id=str(saga_row['saga_id']),
            saga_type=saga_row['saga_type'],
            status=SagaStatus(saga_row['status']),
            current_step=saga_row['current_step'],
            input_data=saga_row['input_data'],
            step_results=step_results,
            created_at=saga_row['created_at'],
            updated_at=saga_row['updated_at'],
            error=saga_row['error_message']
        )
    
    async def _execute_compensation(self, comp: dict):
        """Execute a single compensation."""
        saga = self.orchestrator.sagas.get(comp['saga_type'])
        if not saga:
            raise ValueError(f"Unknown saga type: {comp['saga_type']}")
        
        # Find the step
        step = next(
            (s for s in saga.steps if s.name == comp['step_name']),
            None
        )
        
        if not step or not step.compensation:
            return
        
        # Get saga state for context
        saga_row = await self.db.fetch_one(
            "SELECT * FROM saga_instances WHERE saga_id = $1",
            comp['saga_id']
        )
        
        state = await self._reconstruct_state(saga_row)
        
        context = {
            "saga_id": state.saga_id,
            "input": state.input_data,
            "results": state.step_results,
            "step_result": state.step_results.get(comp['step_name'])
        }
        
        await step.compensation(context)

# Saga Monitoring and Metrics

from prometheus_client import Counter, Histogram, Gauge
from dataclasses import dataclass
from typing import Dict
import time


# Prometheus metrics
saga_started = Counter(
    'saga_started_total',
    'Total sagas started',
    ['saga_type']
)

saga_completed = Counter(
    'saga_completed_total',
    'Total sagas completed successfully',
    ['saga_type']
)

saga_failed = Counter(
    'saga_failed_total',
    'Total sagas that failed',
    ['saga_type', 'failed_step']
)

saga_compensated = Counter(
    'saga_compensated_total',
    'Total sagas that were compensated',
    ['saga_type']
)

saga_duration = Histogram(
    'saga_duration_seconds',
    'Time to complete saga',
    ['saga_type', 'outcome'],
    buckets=[0.1, 0.5, 1, 5, 10, 30, 60, 300, 600]
)

saga_step_duration = Histogram(
    'saga_step_duration_seconds',
    'Time to complete saga step',
    ['saga_type', 'step_name'],
    buckets=[0.01, 0.05, 0.1, 0.5, 1, 5, 10, 30]
)

saga_in_progress = Gauge(
    'saga_in_progress',
    'Number of sagas currently executing',
    ['saga_type']
)

compensation_retries = Counter(
    'saga_compensation_retries_total',
    'Total compensation retry attempts',
    ['saga_type', 'step_name']
)

saga_compensations_pending = Gauge(
    'saga_compensations_pending',
    'Compensations that failed and are waiting for a retry',
    ['saga_type']
)


class SagaMetricsCollector:
    """Collects and reports saga metrics."""
    
    def __init__(self):
        self.start_times: Dict[str, float] = {}
        self.step_start_times: Dict[str, float] = {}
    
    def on_saga_started(self, saga_id: str, saga_type: str):
        saga_started.labels(saga_type=saga_type).inc()
        saga_in_progress.labels(saga_type=saga_type).inc()
        self.start_times[saga_id] = time.time()
    
    def on_saga_completed(self, saga_id: str, saga_type: str):
        saga_completed.labels(saga_type=saga_type).inc()
        saga_in_progress.labels(saga_type=saga_type).dec()
        
        if saga_id in self.start_times:
            duration = time.time() - self.start_times[saga_id]
            saga_duration.labels(
                saga_type=saga_type,
                outcome='completed'
            ).observe(duration)
            del self.start_times[saga_id]
    
    def on_saga_failed(self, saga_id: str, saga_type: str, failed_step: str):
        saga_failed.labels(
            saga_type=saga_type,
            failed_step=failed_step
        ).inc()
        saga_in_progress.labels(saga_type=saga_type).dec()
        
        if saga_id in self.start_times:
            duration = time.time() - self.start_times[saga_id]
            saga_duration.labels(
                saga_type=saga_type,
                outcome='failed'
            ).observe(duration)
            del self.start_times[saga_id]
    
    def on_saga_compensated(self, saga_id: str, saga_type: str):
        saga_compensated.labels(saga_type=saga_type).inc()
    
    def on_step_started(self, saga_id: str, step_name: str):
        key = f"{saga_id}:{step_name}"
        self.step_start_times[key] = time.time()
    
    def on_step_completed(
        self,
        saga_id: str,
        saga_type: str,
        step_name: str
    ):
        key = f"{saga_id}:{step_name}"
        if key in self.step_start_times:
            duration = time.time() - self.step_start_times[key]
            saga_step_duration.labels(
                saga_type=saga_type,
                step_name=step_name
            ).observe(duration)
            del self.step_start_times[key]
    
    def on_compensation_retry(self, saga_type: str, step_name: str):
        compensation_retries.labels(
            saga_type=saga_type,
            step_name=step_name
        ).inc()

# Alerting Rules (Prometheus/Alertmanager format)

ALERTING_RULES = """
groups:
  - name: saga_alerts
    rules:
      # High saga failure rate
      - alert: SagaHighFailureRate
        expr: |
          sum by (saga_type) (rate(saga_failed_total[5m]))
            / sum by (saga_type) (rate(saga_started_total[5m])) > 0.1
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "High saga failure rate"
          description: "More than 10% of sagas are failing"
      
      # Saga taking too long
      - alert: SagaSlowExecution
        expr: |
          histogram_quantile(0.99, rate(saga_duration_seconds_bucket[5m])) > 60
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "Sagas taking too long"
          description: "p99 saga duration is above 60 seconds"
      
      # Stuck compensations
      - alert: SagaStuckCompensations
        expr: |
          saga_compensations_pending > 10
        for: 15m
        labels:
          severity: critical
        annotations:
          summary: "Stuck saga compensations"
          description: "{{ $value }} compensations have been pending for 15+ minutes"
      
      # Compensation retry storm
      - alert: SagaCompensationRetryStorm
        expr: |
          rate(saga_compensation_retries_total[5m]) > 1
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "High compensation retry rate"
          description: "Compensations are being retried frequently"
"""


# Dashboard query examples
DASHBOARD_QUERIES = {
    "saga_success_rate": """
        sum(rate(saga_completed_total[5m])) / 
        sum(rate(saga_started_total[5m]))
    """,
    
    "saga_p99_duration": """
        histogram_quantile(0.99, 
            sum(rate(saga_duration_seconds_bucket[5m])) by (le, saga_type)
        )
    """,
    
    "saga_failure_by_step": """
        sum(rate(saga_failed_total[1h])) by (saga_type, failed_step)
    """,
    
    "sagas_in_progress": """
        sum(saga_in_progress) by (saga_type)
    """,
    
    "compensation_pending": """
        count(saga_compensations_pending > 0) by (saga_type)
    """
}

# Testing Saga Steps

import pytest
from unittest.mock import AsyncMock, MagicMock


class TestOrderSagaSteps:
    """Unit tests for order saga steps."""
    
    @pytest.fixture
    def inventory_service(self):
        service = AsyncMock()
        service.reserve.return_value = "reservation-123"
        service.release.return_value = True
        return service
    
    @pytest.fixture
    def payment_service(self):
        service = AsyncMock()
        service.charge.return_value = "payment-456"
        service.refund.return_value = True
        return service
    
    @pytest.fixture
    def saga_steps(self, inventory_service, payment_service):
        return OrderSagaSteps(
            inventory_service=inventory_service,
            payment_service=payment_service,
            shipping_service=AsyncMock(),
            notification_service=AsyncMock()
        )
    
    @pytest.mark.asyncio
    async def test_reserve_inventory_success(self, saga_steps, inventory_service):
        """Test successful inventory reservation."""
        context = {
            'saga_id': 'saga-123',
            'input': {
                'items': [{'product_id': 'prod-1', 'quantity': 2}]
            },
            'results': {}
        }
        
        result = await saga_steps.reserve_inventory(context)
        
        assert result == {'reservation_id': 'reservation-123'}
        inventory_service.reserve.assert_called_once_with(
            order_id='saga-123',
            items=[{'product_id': 'prod-1', 'quantity': 2}]
        )
    
    @pytest.mark.asyncio
    async def test_release_inventory_compensation(self, saga_steps, inventory_service):
        """Test inventory release compensation."""
        context = {
            'saga_id': 'saga-123',
            'input': {},
            'results': {
                'reserve_inventory': {'reservation_id': 'reservation-123'}
            }
        }
        
        await saga_steps.release_inventory(context)
        
        inventory_service.release.assert_called_once_with('reservation-123')
    
    @pytest.mark.asyncio
    async def test_release_inventory_idempotent(self, saga_steps, inventory_service):
        """Test compensation is idempotent when no reservation exists."""
        context = {
            'saga_id': 'saga-123',
            'input': {},
            'results': {}  # No reservation result
        }
        
        # Should not raise, should be no-op
        await saga_steps.release_inventory(context)
        
        inventory_service.release.assert_not_called()
    
    @pytest.mark.asyncio
    async def test_charge_payment_uses_idempotency_key(
        self, 
        saga_steps, 
        payment_service
    ):
        """Test payment uses idempotency key."""
        context = {
            'saga_id': 'saga-123',
            'input': {
                'customer_id': 'cust-1',
                'total': 99.99
            },
            'results': {}
        }
        
        await saga_steps.charge_payment(context)
        
        payment_service.charge.assert_called_once()
        call_kwargs = payment_service.charge.call_args.kwargs
        assert call_kwargs['idempotency_key'] == 'saga-saga-123'

# Integration Testing for Sagas

import pytest
from testcontainers.postgres import PostgresContainer


class TestOrderSagaIntegration:
    """Integration tests for complete order saga."""
    
    @pytest.fixture(scope="class")
    def postgres(self):
        with PostgresContainer("postgres:14") as pg:
            yield pg
    
    @pytest.fixture
    async def db(self, postgres):
        conn = await asyncpg.connect(postgres.get_connection_url())
        await self._setup_schema(conn)
        yield conn
        await conn.close()
    
    @pytest.fixture
    def orchestrator(self, db):
        steps = OrderSagaSteps(
            inventory_service=RealInventoryService(db),
            payment_service=MockPaymentService(),
            shipping_service=MockShippingService(),
            notification_service=MockNotificationService()
        )
        
        saga = create_order_saga(steps)
        return SagaOrchestrator(db, {"OrderSaga": saga})
    
    @pytest.mark.asyncio
    async def test_successful_order_saga(self, orchestrator, db):
        """Test complete successful saga execution."""
        saga_id = await orchestrator.start_saga(
            "OrderSaga",
            {
                "customer_id": "cust-123",
                "items": [{"product_id": "prod-1", "quantity": 1}],
                "total": 99.99,
                "shipping_address": "123 Main St"
            }
        )
        
        # Wait for saga to complete
        await asyncio.sleep(2)
        
        # Check saga completed
        state = await db.fetch_one(
            "SELECT status FROM saga_instances WHERE saga_id = $1",
            saga_id
        )
        assert state['status'] == 'completed'
        
        # Check all steps completed
        steps = await db.fetch(
            """
            SELECT step_name, status FROM saga_step_executions 
            WHERE saga_id = $1 ORDER BY step_index
            """,
            saga_id
        )
        assert all(s['status'] == 'completed' for s in steps)
    
    @pytest.mark.asyncio
    async def test_saga_compensation_on_payment_failure(
        self, 
        orchestrator, 
        db
    ):
        """Test saga compensates when payment fails."""
        # Configure payment to fail
        orchestrator.sagas["OrderSaga"].steps[1].action = AsyncMock(
            side_effect=PaymentError("Card declined")
        )
        
        saga_id = await orchestrator.start_saga(
            "OrderSaga",
            {
                "customer_id": "cust-123",
                "items": [{"product_id": "prod-1", "quantity": 1}],
                "total": 99.99
            }
        )
        
        await asyncio.sleep(2)
        
        # Check saga was compensated
        state = await db.fetch_one(
            "SELECT status FROM saga_instances WHERE saga_id = $1",
            saga_id
        )
        assert state['status'] == 'compensated'
        
        # Check inventory was released
        compensation = await db.fetch_one(
            """
            SELECT status FROM saga_compensations 
            WHERE saga_id = $1 AND step_name = 'reserve_inventory'
            """,
            saga_id
        )
        assert compensation['status'] == 'completed'
    
    @pytest.mark.asyncio
    async def test_saga_recovery_after_crash(self, orchestrator, db):
        """Test saga recovers after simulated crash."""
        saga_id = await orchestrator.start_saga(
            "OrderSaga",
            {"customer_id": "cust-123", "items": [], "total": 50.00}
        )
        
        # Simulate crash by stopping orchestrator mid-execution
        await asyncio.sleep(0.1)
        
        # Mark saga as stuck
        await db.execute(
            """
            UPDATE saga_instances 
            SET updated_at = NOW() - INTERVAL '10 minutes'
            WHERE saga_id = $1
            """,
            saga_id
        )
        
        # Run recovery
        recovery_job = SagaRecoveryJob(db, orchestrator)
        await recovery_job._recover_stuck_sagas()
        
        await asyncio.sleep(2)
        
        # Verify saga eventually completes
        state = await db.fetch_one(
            "SELECT status FROM saga_instances WHERE saga_id = $1",
            saga_id
        )
        assert state['status'] in ('completed', 'compensated')