from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.models.card_key import CardKey from app.schemas.agent_card_key import AgentCardKeyUpdate DEFAULT_PAGE_SIZE = 20 MAX_PAGE_SIZE = 200 class AgentCardKeyError(Exception): def __init__(self, message: str) -> None: self.message = message super().__init__(message) def _list_card_keys_filters( *, agent_id: int, status: str | None = None, username: str | None = None, ) -> list: filters = [CardKey.agent_id == agent_id] if status == "unused": filters.append(CardKey.activated_at.is_(None)) elif status == "used": filters.append(CardKey.activated_at.is_not(None)) if username and username.strip(): filters.append(CardKey.username.ilike(f"%{username.strip()}%")) return filters async def list_card_keys( db: AsyncSession, *, agent_id: int, page: int = 1, page_size: int = DEFAULT_PAGE_SIZE, status: str | None = None, username: str | None = None, ) -> tuple[list[CardKey], int]: page_size = min(max(page_size, 1), MAX_PAGE_SIZE) page = max(page, 1) offset = (page - 1) * page_size filters = _list_card_keys_filters( agent_id=agent_id, status=status, username=username, ) count_stmt = select(func.count()).select_from(CardKey) list_stmt = select(CardKey) for clause in filters: count_stmt = count_stmt.where(clause) list_stmt = list_stmt.where(clause) total = await db.scalar(count_stmt) or 0 result = await db.scalars( list_stmt.order_by(CardKey.id.desc()).offset(offset).limit(page_size) ) return list(result.all()), total async def _get_card_under_agent( db: AsyncSession, card_id: int, agent_id: int ) -> CardKey: card = await db.get(CardKey, card_id) if card is None or card.agent_id != agent_id: raise AgentCardKeyError("卡密不存在") return card async def update_card_key( db: AsyncSession, card_id: int, body: AgentCardKeyUpdate, *, agent_id: int, ) -> CardKey: card = await _get_card_under_agent(db, card_id, agent_id) if "agent_remark" in body.model_fields_set: card.agent_remark = body.agent_remark await db.flush() await db.refresh(card) return card