from fastapi import Request
from strawberry.fastapi import BaseContext
from app.db import SessionLocal
from app.auth.deps import get_current_user
from app.models import User
from typing import AsyncGenerator

class CustomContext(BaseContext):
    def __init__(self, request: Request):
        super().__init__(request)
        self.db = SessionLocal()
        self.request = request

    async def get_current_user(self) -> User | None:
        # 复用现有的 auth 逻辑
        try:
            return await get_current_user(self.request)
        except Exception:
            return None

async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]:
    context = CustomContext(request)
    try:
        yield context
    finally:
        context.db.close()  # 确保请求结束时关闭数据库连接