# database.py
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker

# 主数据库配置（写操作）
MASTER_DATABASE_URL = "postgresql+asyncpg://user:password@master-host/dbname"
master_engine = create_async_engine(MASTER_DATABASE_URL, pool_size=10)

# 从数据库配置（读操作） 
REPLICA_DATABASE_URL = "postgresql+asyncpg://user:password@replica-host/dbname"
replica_engine = create_async_engine(REPLICA_DATABASE_URL, pool_size=20)

# 创建会话工厂
MasterSession = sessionmaker(master_engine, class_=AsyncSession, expire_on_commit=False)
ReplicaSession = sessionmaker(replica_engine, class_=AsyncSession, expire_on_commit=False)

# dependencies.py
from fastapi import Request, Depends
from database import MasterSession, ReplicaSession


async def get_db(request: Request):
    """智能路由数据库连接"""
    # 写操作路由到主库
    if request.method in ['POST', 'PUT', 'DELETE']:
        db = MasterSession()
    else:  # 读操作使用从库
        db = ReplicaSession()

    try:
        yield db
    finally:
        await db.close()


# 在路由中使用
@app.post("/orders")
async def create_order(
        order: OrderSchema,
        db: AsyncSession = Depends(get_db)
):
# 业务逻辑

# services/transaction_coordinator.py
from typing import List
from fastapi import HTTPException


class SagaCoordinator:
    def __init__(self):
        self.compensation_actions = []

    async def execute_transaction(self, steps: List[callable]):
        """执行Saga事务"""
        try:
            for step in steps:
                await step()
        except Exception as e:
            await self.compensate()
            raise HTTPException(500, "Transaction failed")

    async def compensate(self):
        """补偿操作执行"""
        for action in reversed(self.compensation_actions):
            try:
                await action()
            except Exception as compen_e:
                # 记录补偿失败日志
                logger.error(f"Compensation failed: {compen_e}")


# 使用示例
async def create_order_transaction():
    coordinator = SagaCoordinator()

    async def deduct_inventory():
        # 预留库存
        coordinator.compensation_actions.append(restore_inventory)

    async def create_order_record():
        # 创建订单记录
        coordinator.compensation_actions.append(delete_order_record)

    await coordinator.execute_transaction([
        deduct_inventory,
        create_order_record
    ])

# models.py
from pydantic import BaseModel


class OrderCreate(BaseModel):
    user_id: int
    product_id: str
    quantity: int


# services/order_service.py
from sqlalchemy import text
from motor.motor_asyncio import AsyncIOMotorClient


class OrderService:
    def __init__(self):
        # 初始化各数据库连接
        self.pg_pool = MasterSession
        self.mongo_client = AsyncIOMotorClient(MONGO_URI)
        self.mysql_pool = create_async_engine(MYSQL_URI)

    async def create_order(self, order_data: OrderCreate):
        """创建订单事务"""
        async with self.pg_pool() as pg_session,
                self.mysql_pool.begin() as mysql_conn:
            # 步骤1：扣减MySQL库存
            mysql_update = text("""
                UPDATE inventory 
                SET stock = stock - :quantity 
                WHERE product_id = :product_id
                AND stock >= :quantity
            """)
            await mysql_conn.execute(
                mysql_update,
                product_id=order_data.product_id,
                quantity=order_data.quantity
            )

            # 步骤2：创建PostgreSQL订单
            pg_insert = text("""
                INSERT INTO orders (user_id, product_id, quantity)
                VALUES (:user_id, :product_id, :quantity)
            """)
            await pg_session.execute(pg_insert, order_data.dict())

            # 步骤3：更新MongoDB用户行为
            mongo_db = self.mongo_client.user_behavior
            await mongo_db.events.insert_one({
                "user_id": order_data.user_id,
                "event_type": "order_created",
                "timestamp": datetime.now()
            })

            # 提交PostgreSQL事务
            await pg_session.commit()

# 正确使用方式
async def get_db():
    async with Session() as session:
        yield session

async def safe_transaction():
    async with session.begin():
        try:
            # 业务操作
            await session.commit()
        except:
            await session.rollback()
            raise