Polyglot Streaming Platform with SAGA, CQRS, and Multi-Cloud Portability: A Complete Integration Tutorial¶

"""SAGA pattern implementation for distributed transactions."""
import uuid
import asyncio
import logging
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
import json

import grpc
from prometheus_client import Counter, Histogram, Gauge

from proto.order_pb2 import (
    SagaState, SagaStep, SagaStatus, SagaStepStatus,
    CreateOrderRequest, ProcessPaymentRequest
)
from proto.order_pb2_grpc import OrderServiceStub
from proto.payment_pb2_grpc import PaymentServiceStub
from proto.inventory_pb2_grpc import InventoryServiceStub

saga_metrics = {
    "saga_started": Counter("saga_started_total", "SAGAs started", ["saga_type"]),
    "saga_completed": Counter("saga_completed_total", "SAGAs completed", ["saga_type", "status"]),
    "saga_duration": Histogram("saga_duration_seconds", "SAGA duration", ["saga_type"]),
    "saga_compensations": Counter("saga_compensations_total", "SAGA compensations", ["saga_type"]),
    "active_sagas": Gauge("active_sagas", "Active SAGAs", ["saga_type"]),
}


class SagaStepResult(Enum):
    """SAGA step execution result."""
    SUCCESS = "success"
    FAILURE = "failure"
    RETRY = "retry"


@dataclass
class SagaStepDefinition:
    """SAGA step definition."""
    step_id: str
    service_name: str
    action: str
    compensation_action: str
    timeout_seconds: int = 30
    retry_count: int = 3
    retry_delay_seconds: float = 1.0
    payload: Dict[str, Any] = field(default_factory=dict)


@dataclass
class SagaExecutionResult:
    """SAGA execution result."""
    saga_id: str
    status: SagaStatus
    completed_steps: List[str]
    failed_step: Optional[str] = None
    compensation_applied: bool = False
    duration_seconds: float = 0.0
    error: Optional[str] = None


class SagaOrchestrator:
    """Orchestrates SAGA transactions across polyglot services."""

    def __init__(
        self,
        order_service_stub: OrderServiceStub,
        payment_service_stub: PaymentServiceStub,
        inventory_service_stub: InventoryServiceStub,
        state_store: Optional[Any] = None  # For persistence
    ):
        self.order_stub = order_service_stub
        self.payment_stub = payment_service_stub
        self.inventory_stub = inventory_service_stub
        self.state_store = state_store
        self.active_sagas: Dict[str, SagaState] = {}

    async def execute_saga(
        self,
        saga_type: str,
        steps: List[SagaStepDefinition],
        initial_payload: Dict[str, Any]
    ) -> SagaExecutionResult:
        """Execute a SAGA transaction."""
        saga_id = str(uuid.uuid4())
        start_time = datetime.utcnow()

        saga_metrics["saga_started"].labels(saga_type=saga_type).inc()
        saga_metrics["active_sagas"].labels(saga_type=saga_type).inc()

        # Initialize SAGA state
        saga_state = SagaState(
            saga_id=saga_id,
            order_id=initial_payload.get("order_id", ""),
            status=SagaStatus.SAGA_STATUS_IN_PROGRESS,
            started_at=int(start_time.timestamp() * 1000)
        )

        self.active_sagas[saga_id] = saga_state

        completed_steps = []
        failed_step = None

        try:
            # Execute steps sequentially
            for step_def in steps:
                step = SagaStep(
                    step_id=step_def.step_id,
                    service_name=step_def.service_name,
                    action=step_def.action,
                    status=SagaStepStatus.SAGA_STEP_STATUS_PENDING,
                    payload=json.dumps(step_def.payload).encode()
                )

                saga_state.steps.append(step)

                # Execute step with retries
                result = await self._execute_step_with_retry(
                    step_def,
                    initial_payload,
                    saga_id
                )

                if result == SagaStepResult.SUCCESS:
                    step.status = SagaStepStatus.SAGA_STEP_STATUS_COMPLETED
                    completed_steps.append(step_def.step_id)
                elif result == SagaStepResult.RETRY:
                    # Retry logic handled in _execute_step_with_retry
                    step.status = SagaStepStatus.SAGA_STEP_STATUS_FAILED
                    failed_step = step_def.step_id
                    break
                else:  # FAILURE
                    step.status = SagaStepStatus.SAGA_STEP_STATUS_FAILED
                    failed_step = step_def.step_id
                    break

            # Check if all steps completed
            if failed_step is None:
                saga_state.status = SagaStatus.SAGA_STATUS_COMPLETED
                saga_state.completed_at = int(datetime.utcnow().timestamp() * 1000)

                saga_metrics["saga_completed"].labels(
                    saga_type=saga_type,
                    status="success"
                ).inc()
            else:
                # Compensate for completed steps
                await self._compensate_saga(saga_state, saga_type)
                saga_state.status = SagaStatus.SAGA_STATUS_COMPENSATED

                saga_metrics["saga_compensations"].labels(saga_type=saga_type).inc()
                saga_metrics["saga_completed"].labels(
                    saga_type=saga_type,
                    status="compensated"
                ).inc()

        except Exception as e:
            logging.error(f"SAGA {saga_id} failed: {e}", exc_info=True)
            saga_state.status = SagaStatus.SAGA_STATUS_FAILED
            await self._compensate_saga(saga_state, saga_type)

            saga_metrics["saga_completed"].labels(
                saga_type=saga_type,
                status="failed"
            ).inc()

        finally:
            duration = (datetime.utcnow() - start_time).total_seconds()
            saga_metrics["saga_duration"].labels(saga_type=saga_type).observe(duration)
            saga_metrics["active_sagas"].labels(saga_type=saga_type).dec()

            # Persist state
            if self.state_store:
                await self._persist_saga_state(saga_state)

        return SagaExecutionResult(
            saga_id=saga_id,
            status=saga_state.status,
            completed_steps=completed_steps,
            failed_step=failed_step,
            compensation_applied=saga_state.status == SagaStatus.SAGA_STATUS_COMPENSATED,
            duration_seconds=duration
        )

    async def _execute_step_with_retry(
        self,
        step_def: SagaStepDefinition,
        payload: Dict[str, Any],
        saga_id: str
    ) -> SagaStepResult:
        """Execute a SAGA step with retry logic."""
        for attempt in range(step_def.retry_count):
            try:
                # Route to appropriate service based on service_name
                if step_def.service_name == "order-service":
                    result = await self._execute_order_step(step_def, payload, saga_id)
                elif step_def.service_name == "payment-service":
                    result = await self._execute_payment_step(step_def, payload, saga_id)
                elif step_def.service_name == "inventory-service":
                    result = await self._execute_inventory_step(step_def, payload, saga_id)
                else:
                    raise ValueError(f"Unknown service: {step_def.service_name}")

                if result:
                    return SagaStepResult.SUCCESS
                else:
                    if attempt < step_def.retry_count - 1:
                        await asyncio.sleep(step_def.retry_delay_seconds * (attempt + 1))
                        continue
                    return SagaStepResult.FAILURE

            except asyncio.TimeoutError:
                if attempt < step_def.retry_count - 1:
                    await asyncio.sleep(step_def.retry_delay_seconds * (attempt + 1))
                    continue
                return SagaStepResult.RETRY

            except Exception as e:
                logging.error(f"Step {step_def.step_id} failed: {e}", exc_info=True)
                if attempt < step_def.retry_count - 1:
                    await asyncio.sleep(step_def.retry_delay_seconds * (attempt + 1))
                    continue
                return SagaStepResult.FAILURE

        return SagaStepResult.FAILURE

    async def _execute_order_step(
        self,
        step_def: SagaStepDefinition,
        payload: Dict[str, Any],
        saga_id: str
    ) -> bool:
        """Execute order service step."""
        if step_def.action == "create_order":
            request = CreateOrderRequest(
                customer_id=payload["customer_id"],
                shipping_address=payload["shipping_address"],
                payment_method=payload["payment_method"]
            )
            # Add items from payload
            for item in payload.get("items", []):
                order_item = request.items.add()
                order_item.product_id = item["product_id"]
                order_item.quantity = item["quantity"]
                # Set price from item

            response = await asyncio.to_thread(
                self.order_stub.CreateOrder,
                request
            )
            return response.order_id != ""

        return False

    async def _execute_payment_step(
        self,
        step_def: SagaStepDefinition,
        payload: Dict[str, Any],
        saga_id: str
    ) -> bool:
        """Execute payment service step."""
        if step_def.action == "process_payment":
            from proto.payment_pb2 import ProcessPaymentRequest, Money

            request = ProcessPaymentRequest(
                order_id=payload["order_id"],
                amount=Money(
                    currency=payload["amount"]["currency"],
                    amount_cents=payload["amount"]["amount_cents"]
                ),
                payment_method=payload["payment_method"],
                saga_id=saga_id
            )

            response = await asyncio.to_thread(
                self.payment_stub.ProcessPayment,
                request
            )
            return response.status == 1  # COMPLETED

        elif step_def.action == "refund_payment":
            from proto.payment_pb2 import RefundPaymentRequest

            request = RefundPaymentRequest(
                payment_id=payload["payment_id"],
                amount=payload.get("amount"),
                reason=payload.get("reason", "SAGA compensation"),
                saga_id=saga_id
            )

            response = await asyncio.to_thread(
                self.payment_stub.RefundPayment,
                request
            )
            return response.success

        return False

    async def _execute_inventory_step(
        self,
        step_def: SagaStepDefinition,
        payload: Dict[str, Any],
        saga_id: str
    ) -> bool:
        """Execute inventory service step."""
        # Similar pattern for inventory service
        # In production, implement actual gRPC calls
        return True

    async def _compensate_saga(
        self,
        saga_state: SagaState,
        saga_type: str
    ):
        """Compensate a failed SAGA by executing compensation actions in reverse order."""
        saga_state.status = SagaStatus.SAGA_STATUS_COMPENSATING

        # Execute compensations in reverse order
        for step in reversed(saga_state.steps):
            if step.status == SagaStepStatus.SAGA_STEP_STATUS_COMPLETED:
                # Execute compensation action
                try:
                    await self._execute_compensation(step, saga_state.saga_id)
                    step.status = SagaStepStatus.SAGA_STEP_STATUS_COMPENSATED
                except Exception as e:
                    logging.error(
                        f"Compensation failed for step {step.step_id}: {e}",
                        exc_info=True
                    )
                    # Log but continue with other compensations

    async def _execute_compensation(
        self,
        step: SagaStep,
        saga_id: str
    ):
        """Execute compensation for a step."""
        # Route compensation to appropriate service
        if step.service_name == "payment-service":
            payload = json.loads(step.payload.decode())
            await self._execute_payment_step(
                SagaStepDefinition(
                    step_id=f"{step.step_id}_compensation",
                    service_name=step.service_name,
                    action="refund_payment",
                    compensation_action="",
                    payload=payload
                ),
                payload,
                saga_id
            )
        # Add other service compensations

    async def _persist_saga_state(self, saga_state: SagaState):
        """Persist SAGA state for recovery."""
        # In production, persist to database or event store
        pass


class OrderSagaOrchestrator:
    """Specialized orchestrator for order processing SAGAs."""

    def __init__(self, base_orchestrator: SagaOrchestrator):
        self.orchestrator = base_orchestrator

    async def create_order_saga(
        self,
        customer_id: str,
        items: List[Dict],
        shipping_address: Dict,
        payment_method: str
    ) -> SagaExecutionResult:
        """Execute order creation SAGA."""
        steps = [
            SagaStepDefinition(
                step_id="create_order",
                service_name="order-service",
                action="create_order",
                compensation_action="cancel_order",
                payload={
                    "customer_id": customer_id,
                    "items": items,
                    "shipping_address": shipping_address
                }
            ),
            SagaStepDefinition(
                step_id="reserve_inventory",
                service_name="inventory-service",
                action="reserve_items",
                compensation_action="release_reservation",
                payload={"items": items}
            ),
            SagaStepDefinition(
                step_id="process_payment",
                service_name="payment-service",
                action="process_payment",
                compensation_action="refund_payment",
                payload={
                    "order_id": "",  # Will be set by first step
                    "amount": self._calculate_total(items),
                    "payment_method": payment_method
                }
            ),
        ]

        initial_payload = {
            "customer_id": customer_id,
            "items": items,
            "shipping_address": shipping_address,
            "payment_method": payment_method
        }

        return await self.orchestrator.execute_saga(
            saga_type="create_order",
            steps=steps,
            initial_payload=initial_payload
        )

    def _calculate_total(self, items: List[Dict]) -> Dict:
        """Calculate order total."""
        total_cents = sum(
            item["quantity"] * item["price_cents"]
            for item in items
        )
        return {
            "currency": "USD",
            "amount_cents": total_cents
        }

"""CQRS pattern implementation for read/write separation."""
import asyncio
import logging
from typing import Dict, List, Optional, Any
from datetime import datetime
from dataclasses import dataclass

import asyncpg
import duckdb
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from prometheus_client import Counter, Histogram, Gauge

from proto.order_pb2 import Order, OrderStatus

cqrs_metrics = {
    "commands_executed": Counter("cqrs_commands_total", "Commands executed", ["command_type"]),
    "queries_executed": Counter("cqrs_queries_total", "Queries executed", ["query_type"]),
    "event_projections": Counter("cqrs_projections_total", "Event projections", ["projection_type"]),
    "read_model_lag": Gauge("cqrs_read_model_lag_seconds", "Read model lag", ["model"]),
    "write_latency": Histogram("cqrs_write_latency_seconds", "Write latency", ["command_type"]),
    "read_latency": Histogram("cqrs_read_latency_seconds", "Read latency", ["query_type"]),
}


@dataclass
class Command:
    """Command for write side."""
    command_id: str
    command_type: str
    payload: Dict[str, Any]
    timestamp: datetime


@dataclass
class Query:
    """Query for read side."""
    query_id: str
    query_type: str
    filters: Dict[str, Any]
    timestamp: datetime


class WriteSide:
    """CQRS write side - handles commands and generates events."""

    def __init__(self, database_url: str, event_store):
        self.engine = create_async_engine(database_url)
        self.Session = sessionmaker(
            self.engine, class_=AsyncSession, expire_on_commit=False
        )
        self.event_store = event_store

    async def execute_command(self, command: Command) -> Dict[str, Any]:
        """Execute a command on the write side."""
        start_time = datetime.utcnow()

        try:
            async with self.Session() as session:
                if command.command_type == "create_order":
                    result = await self._create_order_command(session, command)
                elif command.command_type == "update_order_status":
                    result = await self._update_order_status_command(session, command)
                elif command.command_type == "cancel_order":
                    result = await self._cancel_order_command(session, command)
                else:
                    raise ValueError(f"Unknown command type: {command.command_type}")

                await session.commit()

                # Generate and store event
                event = self._generate_event(command, result)
                await self.event_store.append_event(event)

                duration = (datetime.utcnow() - start_time).total_seconds()
                cqrs_metrics["write_latency"].labels(
                    command_type=command.command_type
                ).observe(duration)
                cqrs_metrics["commands_executed"].labels(
                    command_type=command.command_type
                ).inc()

                return result

        except Exception as e:
            logging.error(f"Command execution failed: {e}", exc_info=True)
            raise

    async def _create_order_command(
        self,
        session: AsyncSession,
        command: Command
    ) -> Dict[str, Any]:
        """Handle create order command."""
        # Insert into write database
        result = await session.execute(
            """
            INSERT INTO orders (order_id, customer_id, status, total_amount, created_at)
            VALUES (:order_id, :customer_id, :status, :total_amount, :created_at)
            RETURNING order_id
            """,
            {
                "order_id": command.payload["order_id"],
                "customer_id": command.payload["customer_id"],
                "status": "pending",
                "total_amount": command.payload["total_amount"],
                "created_at": datetime.utcnow()
            }
        )

        order_id = result.scalar()

        # Insert order items
        for item in command.payload.get("items", []):
            await session.execute(
                """
                INSERT INTO order_items (order_id, product_id, quantity, price)
                VALUES (:order_id, :product_id, :quantity, :price)
                """,
                {
                    "order_id": order_id,
                    "product_id": item["product_id"],
                    "quantity": item["quantity"],
                    "price": item["price"]
                }
            )

        return {"order_id": order_id, "status": "pending"}

    async def _update_order_status_command(
        self,
        session: AsyncSession,
        command: Command
    ) -> Dict[str, Any]:
        """Handle update order status command."""
        await session.execute(
            """
            UPDATE orders
            SET status = :status, updated_at = :updated_at
            WHERE order_id = :order_id
            """,
            {
                "order_id": command.payload["order_id"],
                "status": command.payload["status"],
                "updated_at": datetime.utcnow()
            }
        )

        return {"order_id": command.payload["order_id"], "status": command.payload["status"]}

    async def _cancel_order_command(
        self,
        session: AsyncSession,
        command: Command
    ) -> Dict[str, Any]:
        """Handle cancel order command."""
        await session.execute(
            """
            UPDATE orders
            SET status = 'cancelled', cancelled_at = :cancelled_at, cancellation_reason = :reason
            WHERE order_id = :order_id
            """,
            {
                "order_id": command.payload["order_id"],
                "cancelled_at": datetime.utcnow(),
                "reason": command.payload.get("reason", "")
            }
        )

        return {"order_id": command.payload["order_id"], "status": "cancelled"}

    def _generate_event(self, command: Command, result: Dict[str, Any]) -> Dict[str, Any]:
        """Generate event from command result."""
        return {
            "event_id": str(uuid.uuid4()),
            "event_type": f"{command.command_type}_completed",
            "aggregate_id": result.get("order_id"),
            "aggregate_type": "order",
            "payload": result,
            "timestamp": datetime.utcnow().isoformat(),
            "command_id": command.command_id
        }


class ReadSide:
    """CQRS read side - optimized for queries."""

    def __init__(self, parquet_path: str, duckdb_path: str = ":memory:"):
        self.parquet_path = parquet_path
        self.conn = duckdb.connect(duckdb_path)
        self._initialize_read_model()

    def _initialize_read_model(self):
        """Initialize read model from Parquet files."""
        # Create views from Parquet files
        self.conn.execute("""
            CREATE VIEW orders_read AS
            SELECT * FROM read_parquet('orders.parquet')
        """)

        self.conn.execute("""
            CREATE VIEW order_items_read AS
            SELECT * FROM read_parquet('order_items.parquet')
        """)

    async def execute_query(self, query: Query) -> List[Dict[str, Any]]:
        """Execute a query on the read side."""
        start_time = datetime.utcnow()

        try:
            if query.query_type == "get_order":
                result = await self._get_order_query(query)
            elif query.query_type == "list_orders":
                result = await self._list_orders_query(query)
            elif query.query_type == "get_order_history":
                result = await self._get_order_history_query(query)
            else:
                raise ValueError(f"Unknown query type: {query.query_type}")

            duration = (datetime.utcnow() - start_time).total_seconds()
            cqrs_metrics["read_latency"].labels(
                query_type=query.query_type
            ).observe(duration)
            cqrs_metrics["queries_executed"].labels(
                query_type=query.query_type
            ).inc()

            return result

        except Exception as e:
            logging.error(f"Query execution failed: {e}", exc_info=True)
            raise

    async def _get_order_query(self, query: Query) -> List[Dict[str, Any]]:
        """Get order by ID."""
        order_id = query.filters.get("order_id")

        result = self.conn.execute("""
            SELECT 
                o.order_id,
                o.customer_id,
                o.status,
                o.total_amount,
                o.created_at,
                o.updated_at,
                json_agg(
                    json_object(
                        'product_id', oi.product_id,
                        'quantity', oi.quantity,
                        'price', oi.price
                    )
                ) as items
            FROM orders_read o
            LEFT JOIN order_items_read oi ON o.order_id = oi.order_id
            WHERE o.order_id = ?
            GROUP BY o.order_id, o.customer_id, o.status, o.total_amount, o.created_at, o.updated_at
        """, [order_id]).fetchall()

        return [dict(row) for row in result]

    async def _list_orders_query(self, query: Query) -> List[Dict[str, Any]]:
        """List orders with filters."""
        filters = query.filters
        where_clauses = []
        params = []

        if filters.get("customer_id"):
            where_clauses.append("customer_id = ?")
            params.append(filters["customer_id"])

        if filters.get("status"):
            where_clauses.append("status = ?")
            params.append(filters["status"])

        where_sql = " AND ".join(where_clauses) if where_clauses else "1=1"

        result = self.conn.execute(f"""
            SELECT * FROM orders_read
            WHERE {where_sql}
            ORDER BY created_at DESC
            LIMIT ?
        """, params + [filters.get("limit", 100)]).fetchall()

        return [dict(row) for row in result]

    async def _get_order_history_query(self, query: Query) -> List[Dict[str, Any]]:
        """Get order history with events."""
        order_id = query.filters.get("order_id")

        # Query from event store or read model
        result = self.conn.execute("""
            SELECT 
                event_type,
                payload,
                timestamp
            FROM order_events_read
            WHERE order_id = ?
            ORDER BY timestamp ASC
        """, [order_id]).fetchall()

        return [dict(row) for row in result]

    async def project_event(self, event: Dict[str, Any]):
        """Project event into read model."""
        event_type = event["event_type"]

        if event_type == "create_order_completed":
            await self._project_order_created(event)
        elif event_type == "update_order_status_completed":
            await self._project_order_status_updated(event)
        elif event_type == "cancel_order_completed":
            await self._project_order_cancelled(event)

        cqrs_metrics["event_projections"].labels(
            projection_type=event_type
        ).inc()

    async def _project_order_created(self, event: Dict[str, Any]):
        """Project order created event."""
        # In production, update Parquet files or materialized views
        # For demo, update DuckDB in-memory view
        pass

    async def _project_order_status_updated(self, event: Dict[str, Any]):
        """Project order status updated event."""
        pass

    async def _project_order_cancelled(self, event: Dict[str, Any]):
        """Project order cancelled event."""
        pass


class EventStore:
    """Event store for CQRS event sourcing."""

    def __init__(self, database_url: str):
        self.engine = create_async_engine(database_url)
        self.Session = sessionmaker(
            self.engine, class_=AsyncSession, expire_on_commit=False
        )

    async def append_event(self, event: Dict[str, Any]):
        """Append event to event store."""
        async with self.Session() as session:
            await session.execute(
                """
                INSERT INTO events (
                    event_id, event_type, aggregate_id, aggregate_type,
                    payload, timestamp, command_id
                ) VALUES (
                    :event_id, :event_type, :aggregate_id, :aggregate_type,
                    :payload, :timestamp, :command_id
                )
                """,
                {
                    "event_id": event["event_id"],
                    "event_type": event["event_type"],
                    "aggregate_id": event["aggregate_id"],
                    "aggregate_type": event["aggregate_type"],
                    "payload": json.dumps(event["payload"]),
                    "timestamp": datetime.fromisoformat(event["timestamp"]),
                    "command_id": event.get("command_id")
                }
            )
            await session.commit()

    async def get_events(
        self,
        aggregate_id: str,
        aggregate_type: str,
        from_version: int = 0
    ) -> List[Dict[str, Any]]:
        """Get events for an aggregate."""
        async with self.Session() as session:
            result = await session.execute(
                """
                SELECT event_id, event_type, payload, timestamp, version
                FROM events
                WHERE aggregate_id = :aggregate_id
                  AND aggregate_type = :aggregate_type
                  AND version > :from_version
                ORDER BY version ASC
                """,
                {
                    "aggregate_id": aggregate_id,
                    "aggregate_type": aggregate_type,
                    "from_version": from_version
                }
            )

            events = []
            for row in result:
                events.append({
                    "event_id": row.event_id,
                    "event_type": row.event_type,
                    "payload": json.loads(row.payload),
                    "timestamp": row.timestamp.isoformat(),
                    "version": row.version
                })

            return events

"""Outbox pattern for reliable event publishing."""
import asyncio
import json
import logging
from typing import List, Dict, Any, Optional
from datetime import datetime
from dataclasses import dataclass

import asyncpg
from confluent_kafka import Producer
from prometheus_client import Counter, Gauge, Histogram

outbox_metrics = {
    "outbox_events_created": Counter("outbox_events_created_total", "Outbox events created"),
    "outbox_events_published": Counter("outbox_events_published_total", "Outbox events published", ["status"]),
    "outbox_events_failed": Counter("outbox_events_failed_total", "Outbox events failed"),
    "outbox_lag": Gauge("outbox_lag_events", "Outbox lag (unpublished events)"),
    "outbox_publish_duration": Histogram("outbox_publish_duration_seconds", "Outbox publish duration"),
}


@dataclass
class OutboxEvent:
    """Outbox event record."""
    event_id: str
    aggregate_id: str
    aggregate_type: str
    event_type: str
    payload: Dict[str, Any]
    created_at: datetime
    published_at: Optional[datetime] = None
    retry_count: int = 0
    status: str = "pending"  # pending, published, failed


class OutboxManager:
    """Manages outbox table and event publishing."""

    def __init__(
        self,
        database_url: str,
        kafka_producer: Producer,
        kafka_topic: str,
        poll_interval_seconds: float = 1.0
    ):
        self.database_url = database_url
        self.producer = kafka_producer
        self.topic = kafka_topic
        self.poll_interval = poll_interval_seconds
        self.running = False

    async def create_outbox_event(
        self,
        aggregate_id: str,
        aggregate_type: str,
        event_type: str,
        payload: Dict[str, Any],
        transaction: Optional[Any] = None
    ) -> str:
        """Create outbox event within a transaction."""
        event_id = str(uuid.uuid4())

        if transaction:
            # Use existing transaction
            await transaction.execute(
                """
                INSERT INTO outbox (
                    event_id, aggregate_id, aggregate_type, event_type,
                    payload, created_at, status
                ) VALUES (
                    :event_id, :aggregate_id, :aggregate_type, :event_type,
                    :payload, :created_at, 'pending'
                )
                """,
                {
                    "event_id": event_id,
                    "aggregate_id": aggregate_id,
                    "aggregate_type": aggregate_type,
                    "event_type": event_type,
                    "payload": json.dumps(payload),
                    "created_at": datetime.utcnow()
                }
            )
        else:
            # Create new connection
            conn = await asyncpg.connect(self.database_url)
            try:
                await conn.execute(
                    """
                    INSERT INTO outbox (
                        event_id, aggregate_id, aggregate_type, event_type,
                        payload, created_at, status
                    ) VALUES ($1, $2, $3, $4, $5, $6, 'pending')
                    """,
                    event_id,
                    aggregate_id,
                    aggregate_type,
                    event_type,
                    json.dumps(payload),
                    datetime.utcnow()
                )
            finally:
                await conn.close()

        outbox_metrics["outbox_events_created"].inc()
        return event_id

    async def start_poller(self):
        """Start outbox event poller."""
        self.running = True

        while self.running:
            try:
                await self._poll_and_publish()
                await asyncio.sleep(self.poll_interval)
            except Exception as e:
                logging.error(f"Outbox poller error: {e}", exc_info=True)
                await asyncio.sleep(self.poll_interval * 2)  # Back off on error

    async def stop_poller(self):
        """Stop outbox event poller."""
        self.running = False

    async def _poll_and_publish(self):
        """Poll outbox table and publish events."""
        conn = await asyncpg.connect(self.database_url)

        try:
            # Fetch unpublished events
            rows = await conn.fetch("""
                SELECT 
                    event_id, aggregate_id, aggregate_type, event_type,
                    payload, created_at, retry_count
                FROM outbox
                WHERE status = 'pending'
                  AND retry_count < 5
                ORDER BY created_at ASC
                LIMIT 100
                FOR UPDATE SKIP LOCKED
            """)

            if not rows:
                # Update lag metric
                lag_count = await conn.fetchval("""
                    SELECT COUNT(*) FROM outbox WHERE status = 'pending'
                """)
                outbox_metrics["outbox_lag"].set(lag_count)
                return

            # Publish each event
            for row in rows:
                try:
                    await self._publish_event(row)

                    # Mark as published
                    await conn.execute("""
                        UPDATE outbox
                        SET status = 'published',
                            published_at = $1
                        WHERE event_id = $2
                    """, datetime.utcnow(), row['event_id'])

                    outbox_metrics["outbox_events_published"].labels(status="success").inc()

                except Exception as e:
                    logging.error(f"Failed to publish event {row['event_id']}: {e}")

                    # Increment retry count
                    await conn.execute("""
                        UPDATE outbox
                        SET retry_count = retry_count + 1,
                            status = CASE 
                                WHEN retry_count + 1 >= 5 THEN 'failed'
                                ELSE 'pending'
                            END
                        WHERE event_id = $1
                    """, row['event_id'])

                    outbox_metrics["outbox_events_published"].labels(status="failed").inc()
                    outbox_metrics["outbox_events_failed"].inc()

        finally:
            await conn.close()

    async def _publish_event(self, row: Dict[str, Any]):
        """Publish event to Kafka."""
        start_time = datetime.utcnow()

        event_data = {
            "event_id": row['event_id'],
            "aggregate_id": row['aggregate_id'],
            "aggregate_type": row['aggregate_type'],
            "event_type": row['event_type'],
            "payload": json.loads(row['payload']),
            "timestamp": row['created_at'].isoformat()
        }

        # Publish to Kafka
        future = self.producer.produce(
            topic=self.topic,
            key=row['aggregate_id'],
            value=json.dumps(event_data).encode('utf-8'),
            callback=self._delivery_callback
        )

        # Wait for delivery
        self.producer.poll(1)

        duration = (datetime.utcnow() - start_time).total_seconds()
        outbox_metrics["outbox_publish_duration"].observe(duration)

    def _delivery_callback(self, err, msg):
        """Kafka delivery callback."""
        if err:
            logging.error(f"Message delivery failed: {err}")
        else:
            logging.debug(f"Message delivered to {msg.topic()} [{msg.partition()}]")


class TransactionalOutbox:
    """Transactional outbox pattern implementation."""

    def __init__(self, database_url: str, outbox_manager: OutboxManager):
        self.database_url = database_url
        self.outbox_manager = outbox_manager

    async def execute_with_outbox(
        self,
        operation: Callable,
        aggregate_id: str,
        aggregate_type: str,
        event_type: str,
        event_payload: Dict[str, Any]
    ) -> Any:
        """Execute operation with outbox event in same transaction."""
        conn = await asyncpg.connect(self.database_url)
        tr = conn.transaction()

        try:
            await tr.start()

            # Execute business operation
            result = await operation(conn, tr)

            # Create outbox event in same transaction
            await self.outbox_manager.create_outbox_event(
                aggregate_id=aggregate_id,
                aggregate_type=aggregate_type,
                event_type=event_type,
                payload=event_payload,
                transaction=tr
            )

            await tr.commit()
            return result

        except Exception as e:
            await tr.rollback()
            raise
        finally:
            await conn.close()